02.Expressions.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Expressions containing NDArray objects\n",
    "\n",
    "Python-Blosc2 implements a powerful way to operate with NDArray (and other flavors) objects.  In this section, we will see how to do computations with NDArray arrays in a simple way.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-12-02T06:19:47.839123Z",
     "start_time": "2024-12-02T06:19:47.812673Z"
    }
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "import blosc2"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## A simple example\n",
    "First, let's create a couple of NDArrays."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-12-02T06:19:47.971242Z",
     "start_time": "2024-12-02T06:19:47.852322Z"
    }
   },
   "outputs": [],
   "source": [
    "shape = (500, 1000)\n",
    "a = blosc2.linspace(0, 1, np.prod(shape), dtype=np.float32, shape=shape, urlpath=\"a.b2nd\", mode=\"w\")\n",
    "b = blosc2.linspace(1, 2, np.prod(shape), dtype=np.float64, shape=shape, urlpath=\"b.b2nd\", mode=\"w\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, let's create an expression that involves `a` and `b`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-12-02T06:19:47.989421Z",
     "start_time": "2024-12-02T06:19:47.980856Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "type       : LazyExpr\n",
      "expression : ((((o0 ** 2) + (o1 ** 2)) + (2 * sin((o0 * o1)))) + 1)\n",
      "operands   : {'o0': 'a.b2nd', 'o1': 'b.b2nd'}\n",
      "shape      : (500, 1000)\n",
      "dtype      : float64\n",
      "\n"
     ]
    }
   ],
   "source": [
    "c = a**2 + b**2 + 2 * blosc2.sin(a * b) + 1\n",
    "print(c.info)  # at this stage, the expression has not been computed yet"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We see that the outcome of the expression is a `LazyExpr` object.  This object is a placeholder for the actual computation that will be done when we compute it.  This is a powerful feature because it allows us to build complex expressions without actually computing them until we really need the result.\n",
    "\n",
    "Also, note that you can throw [many math functions](https://www.blosc.org/python-blosc2/reference/array_operations.html) at your expressions. These are mainly the ones supported by [numexpr](https://github.com/pydata/numexpr), plus different reduction operations.\n",
    "\n",
    "Now, let's compute it. `LazyExpr` objects follow the [LazyArray interface](https://www.blosc.org/python-blosc2/reference/lazyarray.html), and this provides several ways for performing the computation, depending on the object we want as the desired output.\n",
    "\n",
    "First, let's use the `compute` method.  The result will be another NDArray array:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-12-02T06:19:48.028533Z",
     "start_time": "2024-12-02T06:19:48.008790Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Type: <class 'blosc2.ndarray.NDArray'>\n",
      "Compression ratio: 1.89x\n"
     ]
    }
   ],
   "source": [
    "d = c.compute()  # compute the expression\n",
    "print(f\"Type: {type(d)}\")\n",
    "print(f\"Compression ratio: {d.schunk.cratio:.2f}x\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Or, we can store the result in a file:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-12-02T06:19:48.497741Z",
     "start_time": "2024-12-02T06:19:48.036008Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "-rw-r--r--  1 faltet  staff   2.0M Dec  3 13:11 result.b2nd\n"
     ]
    }
   ],
   "source": [
    "d = c.compute(urlpath=\"result.b2nd\", mode=\"w\")\n",
    "!ls -lh result.b2nd"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note that all the output is stored in the file as computation proceeds; this is an efficient way to store large results on disk.  Incidentally, both operands and results are stored on disk here, so you can operate with very large arrays in a very small memory footprint."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, let's compute the expression and store the result in a NumPy array.  For this, we will use the `__getitem__` method:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-12-02T06:19:48.515438Z",
     "start_time": "2024-12-02T06:19:48.509272Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Type: <class 'numpy.ndarray'>\n"
     ]
    }
   ],
   "source": [
    "npd = d[:]\n",
    "print(f\"Type: {type(npd)}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As you can see, the result is a NumPy array now.\n",
    "\n",
    "Depending on your needs, you can choose to get the result as a NDArray array or as a NumPy array.  The former is more storage efficient, but the latter is more flexible when interacting with other libraries that do not support NDArray arrays."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You can also compute just *part* of the expression by passing an item argument to the lazy array:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([2.        , 2.000008  , 2.000016  , 2.000024  , 2.000032  ,\n",
       "       2.00004   , 2.000048  , 2.000056  , 2.000064  , 2.000072  ,\n",
       "       2.00008   , 2.000088  , 2.000096  , 2.000104  , 2.000112  ,\n",
       "       2.00012   , 2.000128  , 2.000136  , 2.00014401, 2.00015201])"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "d[0, :20]  # just computes row 0 and cols 0 to 20"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Saving expressions to disk"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You can save literal expressions to disk.  For this, use the `save` method of ``LazyArray`` objects.  For example, let's save the expression `c` to disk:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-12-02T06:19:48.533906Z",
     "start_time": "2024-12-02T06:19:48.523628Z"
    }
   },
   "outputs": [],
   "source": [
    "c = a**2 + b**2 + 2 * blosc2.sin(a * b) + 1\n",
    "c.save(urlpath=\"expr.b2nd\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And you can load it back with the `open` function:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-12-02T06:19:48.551156Z",
     "start_time": "2024-12-02T06:19:48.543375Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "type       : LazyExpr\n",
      "expression : o0 ** 2 + o1 ** 2 + 2 * sin(o0 * o1) + 1\n",
      "operands   : {'o0': 'a.b2nd', 'o1': 'b.b2nd'}\n",
      "shape      : (500, 1000)\n",
      "dtype      : float64\n",
      "\n"
     ]
    }
   ],
   "source": [
    "c2 = blosc2.open(\"expr.b2nd\")\n",
    "print(c2.info)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, you can compute it as before:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-12-02T06:19:48.597949Z",
     "start_time": "2024-12-02T06:19:48.573821Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Compression ratio: 1.89x\n"
     ]
    }
   ],
   "source": [
    "d2 = c2.compute()\n",
    "print(f\"Compression ratio: {d2.schunk.cratio:.2f}x\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Reductions\n",
    "\n",
    "We can also perform reductions as part of expressions.  Let's see an example:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-12-02T06:19:48.617767Z",
     "start_time": "2024-12-02T06:19:48.604334Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "np.float64(999999.9999999473)"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "c = (a + b).sum()\n",
    "c"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As we can see, the result is a scalar. That means that reductions in expressions always perform the computation immediately."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can also specify the axis for the reduction:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-12-02T06:19:48.637975Z",
     "start_time": "2024-12-02T06:19:48.625011Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Shape of c: (500,)\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "array([1001.998004  , 1005.998012  , 1009.99802   , 1013.99802799])"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "c = (a + b).sum(axis=1)\n",
    "print(f\"Shape of c: {c.shape}\")\n",
    "# Show the first 4 elements of the result\n",
    "c[:4]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Reductions can also be part of more complex expressions:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-12-02T06:19:48.666470Z",
     "start_time": "2024-12-02T06:19:48.646093Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Shape of c: (500, 1000)\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "array([1000.0010009 , 1000.00300336, 1000.00500598, 1000.00700854])"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "c = (a + b).sum(axis=0) + 2 * a + 1\n",
    "print(f\"Shape of c: {c.shape}\")\n",
    "# Show the first 4 elements of the result\n",
    "c[0, 0:4]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In particular, note that the result of the reduction above has a different shape than `a`, but the expression is still computed correctly.  This is because the shape of the reduction is *compatible* with the shape of the operands."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Querying NDArray arrays\n",
    "\n",
    "A powerful feature of Blosc2 compute engine is its ability to do queries on NDArray arrays with structured types.  Let's see an example."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-12-02T06:19:49.167542Z",
     "start_time": "2024-12-02T06:19:48.674012Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "First 3 rows:\n",
      " [( 1, -2., 0.34558419) ( 0, -1., 0.82161814) (-1,  0., 0.33043708)]\n"
     ]
    }
   ],
   "source": [
    "N = 1000_000\n",
    "rng = np.random.default_rng(seed=1)\n",
    "it = ((-x + 1, x - 2, rng.normal()) for x in range(N))\n",
    "sa = blosc2.fromiter(\n",
    "    it, dtype=[(\"A\", \"i4\"), (\"B\", \"f4\"), (\"C\", \"f8\")], shape=(N,),\n",
    "    urlpath=\"sa-1M_tutorial2.b2nd\", mode=\"w\"\n",
    ")\n",
    "print(\"First 3 rows:\\n\", sa[:3])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, we can select rows depending on the value of different fields:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-12-02T06:19:49.200013Z",
     "start_time": "2024-12-02T06:19:49.173304Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([(1, -2., 0.34558419), (0, -1., 0.82161814)],\n",
       "      dtype=[('A', '<i4'), ('B', '<f4'), ('C', '<f8')])"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "A = sa[\"A\"]\n",
    "B = sa[\"B\"]\n",
    "C = sa[\"C\"]\n",
    "expr = sa[A > B]\n",
    "expr[:]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can do the same on a more compact way by using an expression in string form inside the brackets:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-12-02T06:19:49.240700Z",
     "start_time": "2024-12-02T06:19:49.217779Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([(1, -2., 0.34558419), (0, -1., 0.82161814)],\n",
       "      dtype=[('A', '<i4'), ('B', '<f4'), ('C', '<f8')])"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "sa[\"A > B\"][:]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The expression can also be a complex one:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-12-02T06:19:49.279062Z",
     "start_time": "2024-12-02T06:19:49.247940Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([(0, -1., 0.82161814)],\n",
       "      dtype=[('A', '<i4'), ('B', '<f4'), ('C', '<f8')])"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "sa[\"(A > B) & (sin(C) > .5)\"][:]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can also query and extract a single field:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-12-02T06:19:49.386143Z",
     "start_time": "2024-12-02T06:19:49.366402Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([0.34558419, 0.82161814])"
      ]
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "sa[\"C\"][\"A > B\"][:]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And perform reductions on queries on a single field:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-12-02T06:22:23.743178Z",
     "start_time": "2024-12-02T06:22:23.689217Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "np.float64(0.7955493778255271)"
      ]
     },
     "execution_count": 21,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "sa[\"C\"][\"(A < B) & (C > 0)\"].mean()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Combining all this weaponry allows to query your data on a simple and efficient way. As the computation is lazy, all the operations are grouped and executed together for maximum performance. The only exception is that, when a reduction is found, it is computed eagerly; but it can still be part of more general expressions, as well as being able to be saved and loaded from disk."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Summary\n",
    "\n",
    "In this section, we have seen how to perform computations with NDArray arrays, and more in particular, how to create expressions, compute them, and save them to disk. Also, we have looked at performing reductions, broadcasting, selections and combinations of both. Lazy expressions allow you to build and compute complex computations from operands that can be in-memory, on-disk or remote (`C2Array`) in a simple and effective way."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}