{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# `ConditionalDimension` tutorial"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This tutorial explains how to create equations that only get executed under given conditions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Data([[1., 0., 0., 0., 0., 0., 0., 0., 0., 0.],\n",
       "      [0., 1., 0., 0., 0., 0., 0., 0., 0., 0.],\n",
       "      [0., 0., 1., 0., 0., 0., 0., 0., 0., 0.],\n",
       "      [0., 0., 0., 1., 0., 0., 0., 0., 0., 0.],\n",
       "      [0., 0., 0., 0., 1., 0., 0., 0., 0., 0.],\n",
       "      [0., 0., 0., 0., 0., 1., 0., 0., 0., 0.],\n",
       "      [0., 0., 0., 0., 0., 0., 1., 0., 0., 0.],\n",
       "      [0., 0., 0., 0., 0., 0., 0., 1., 0., 0.],\n",
       "      [0., 0., 0., 0., 0., 0., 0., 0., 1., 0.],\n",
       "      [0., 0., 0., 0., 0., 0., 0., 0., 0., 1.]], dtype=float32)"
      ]
     },
     "execution_count": 1,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from devito import Dimension, Function, Grid\n",
    "import numpy as np\n",
    "\n",
    "# We define a 10x10 grid, dimensions are x, y\n",
    "shape = (10, 10)\n",
    "grid = Grid(shape = shape)\n",
    "x, y = grid.dimensions\n",
    "\n",
    "# Define function 𝑓. We will initialize f's data with ones on its diagonal.\n",
    "\n",
    "f = Function(name='f', grid=grid)\n",
    "f.data[:] = np.eye(10)\n",
    "f.data"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We begin by constructing an `Operator` that increases the values of `f`, a `Function` defined on a two-dimensional `Grid`, by one."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Operator `Kernel` ran in 0.01 s\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "Data([[2., 1., 1., 1., 1., 1., 1., 1., 1., 1.],\n",
       "      [1., 2., 1., 1., 1., 1., 1., 1., 1., 1.],\n",
       "      [1., 1., 2., 1., 1., 1., 1., 1., 1., 1.],\n",
       "      [1., 1., 1., 2., 1., 1., 1., 1., 1., 1.],\n",
       "      [1., 1., 1., 1., 2., 1., 1., 1., 1., 1.],\n",
       "      [1., 1., 1., 1., 1., 2., 1., 1., 1., 1.],\n",
       "      [1., 1., 1., 1., 1., 1., 2., 1., 1., 1.],\n",
       "      [1., 1., 1., 1., 1., 1., 1., 2., 1., 1.],\n",
       "      [1., 1., 1., 1., 1., 1., 1., 1., 2., 1.],\n",
       "      [1., 1., 1., 1., 1., 1., 1., 1., 1., 2.]], dtype=float32)"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "#NBVAL_IGNORE_OUTPUT\n",
    "from devito import Eq, Operator\n",
    "op0 = Operator(Eq(f, f + 1))\n",
    "op0.apply()\n",
    "f.data"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Every value has been updated by one. You should be able to see twos on the diagonal and ones everywhere else."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "#print(op0.ccode) # Print the generated code"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Devito API for relationals\n",
    "\n",
    "In order to construct a `ConditionalDimension`, one needs to get familiar with relationals. The Devito API for relationals currently supports the following classes of relation, which inherit from their [SymPy counterparts](https://docs.sympy.org/latest/modules/core.html#module-sympy.core.relational):\n",
    "\n",
    "- Le (less-than or equal)\n",
    "- Lt (less-than)\n",
    "- Ge (greater-than or equal)\n",
    "- Gt (greater-than)\n",
    "- Ne (not equal)\n",
    "\n",
    "Relationals are used to define a condition and are passed as an argument to `ConditionalDimension` at construction time.\n",
    "\n",
    "\n",
    "# Devito API for ConditionalDimension\n",
    "\n",
    "A `ConditionalDimension` defines a non-convex iteration sub-space derived from a ``parent`` Dimension, implemented by the compiler generating conditional \"if-then\" code within the parent Dimension's iteration space.\n",
    "\n",
    "A `ConditionalDimension` is used in two typical cases: \n",
    "\n",
    "- Use case A:\n",
    "To constrain the execution of loop iterations to make sure certain conditions are honoured\n",
    "\n",
    "- Use case B:\n",
    "To subsample a `Function`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "    Symbol defining a non-convex iteration sub-space derived from a ``parent``\n",
      "    Dimension, implemented by the compiler generating conditional \"if-then\" code\n",
      "    within the parent Dimension's iteration space.\n",
      "\n",
      "    Parameters\n",
      "    ----------\n",
      "    name : str\n",
      "        Name of the dimension.\n",
      "    parent : Dimension\n",
      "        The parent Dimension.\n",
      "    factor : int, optional, default=None\n",
      "        The number of iterations between two executions of the if-branch. If None\n",
      "        (default), ``condition`` must be provided.\n",
      "    condition : expr-like, optional, default=None\n",
      "        An arbitrary SymPy expression, typically involving the ``parent``\n",
      "        Dimension. When it evaluates to True, the if-branch is executed. If None\n",
      "        (default), ``factor`` must be provided.\n",
      "    indirect : bool, optional, default=False\n",
      "        If True, use `self`, rather than the parent Dimension, to\n",
      "        index into arrays. A typical use case is when arrays are accessed\n",
      "        indirectly via the ``condition`` expression.\n",
      "\n",
      "    Examples\n",
      "    --------\n",
      "    Among the other things, ConditionalDimensions are indicated to implement\n",
      "    Function subsampling. In the following example, an Operator evaluates the\n",
      "    Function ``g`` and saves its content into ``f`` every ``factor=4`` iterations.\n",
      "\n",
      "    >>> from devito import Dimension, ConditionalDimension, Function, Eq, Operator\n",
      "    >>> size, factor = 16, 4\n",
      "    >>> i = Dimension(name='i')\n",
      "    >>> ci = ConditionalDimension(name='ci', parent=i, factor=factor)\n",
      "    >>> g = Function(name='g', shape=(size,), dimensions=(i,))\n",
      "    >>> f = Function(name='f', shape=(int(size/factor),), dimensions=(ci,))\n",
      "    >>> op = Operator([Eq(g, 1), Eq(f, g)])\n",
      "\n",
      "    The Operator generates the following for-loop (pseudocode)\n",
      "\n",
      "    .. code-block:: C\n",
      "\n",
      "        for (int i = i_m; i <= i_M; i += 1) {\n",
      "          g[i] = 1;\n",
      "          if (i%4 == 0) {\n",
      "            f[i / 4] = g[i];\n",
      "          }\n",
      "        }\n",
      "\n",
      "    Another typical use case is when one needs to constrain the execution of\n",
      "    loop iterations so that certain conditions are honoured. The following\n",
      "    artificial example uses ConditionalDimension to guard against out-of-bounds\n",
      "    accesses in indirectly accessed arrays.\n",
      "\n",
      "    >>> from sympy import And\n",
      "    >>> ci = ConditionalDimension(name='ci', parent=i,\n",
      "    ...                           condition=And(g[i] > 0, g[i] < 4, evaluate=False))\n",
      "    >>> f = Function(name='f', shape=(int(size/factor),), dimensions=(ci,))\n",
      "    >>> op = Operator(Eq(f[g[i]], f[g[i]] + 1))\n",
      "\n",
      "    The Operator generates the following for-loop (pseudocode)\n",
      "\n",
      "    .. code-block:: C\n",
      "\n",
      "        for (int i = i_m; i <= i_M; i += 1) {\n",
      "          if (g[i] > 0 && g[i] < 4) {\n",
      "            f[g[i]] = f[g[i]] + 1;\n",
      "          }\n",
      "        }\n",
      "\n",
      "    \n"
     ]
    }
   ],
   "source": [
    "from devito import ConditionalDimension\n",
    "print(ConditionalDimension.__doc__)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Use case A - Honour a condition\n",
    "\n",
    "Now it's time to show a more descriptive example. We define a conditional that increments by one all values of a Function that are larger than zero. We name our `ConditionalDimension` `ci`.\n",
    "In this example we want to update again by one all the elements in `f` that are larger than zero. Before updating the elements we reinitialize our data to the eye function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Operator `Kernel` ran in 0.01 s\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "Data([[2., 1., 1., 1., 1., 1., 1., 1., 1., 1.],\n",
       "      [1., 2., 1., 1., 1., 1., 1., 1., 1., 1.],\n",
       "      [1., 1., 2., 1., 1., 1., 1., 1., 1., 1.],\n",
       "      [1., 1., 1., 2., 1., 1., 1., 1., 1., 1.],\n",
       "      [1., 1., 1., 1., 2., 1., 1., 1., 1., 1.],\n",
       "      [1., 1., 1., 1., 1., 2., 1., 1., 1., 1.],\n",
       "      [1., 1., 1., 1., 1., 1., 2., 1., 1., 1.],\n",
       "      [1., 1., 1., 1., 1., 1., 1., 2., 1., 1.],\n",
       "      [1., 1., 1., 1., 1., 1., 1., 1., 2., 1.],\n",
       "      [1., 1., 1., 1., 1., 1., 1., 1., 1., 2.]], dtype=float32)"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "#NBVAL_IGNORE_OUTPUT\n",
    "from devito import Gt\n",
    "\n",
    "f.data[:] = np.eye(10)\n",
    "\n",
    "# Define the condition f(x, y) > 0\n",
    "condition = Gt(f, 0)\n",
    "\n",
    "# Construct a ConditionalDimension\n",
    "ci = ConditionalDimension(name='ci', parent=y, condition=condition)\n",
    "\n",
    "op1 = Operator(Eq(f, f + 1))\n",
    "op1.apply()\n",
    "f.data\n",
    "# print(op1.ccode) # Uncomment to view code\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We've constructed `ci`, but it's still unused, so `op1` is still identical to `op0`. We need to pass `ci` as an \"implicit `Dimension`\" to the equation that we want to be conditionally executed as shown in the next cell."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Operator `Kernel` ran in 0.01 s\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "START(section0)\n",
      "for (int x = x_m; x <= x_M; x += 1)\n",
      "{\n",
      "  #pragma omp simd aligned(f:32)\n",
      "  for (int y = y_m; y <= y_M; y += 1)\n",
      "  {\n",
      "    if (f[x + 1][y + 1] > 0)\n",
      "    {\n",
      "      f[x + 1][y + 1] = f[x + 1][y + 1] + 1;\n",
      "    }\n",
      "  }\n",
      "}\n",
      "STOP(section0,timers)\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "Data([[2., 0., 0., 0., 0., 0., 0., 0., 0., 0.],\n",
       "      [0., 2., 0., 0., 0., 0., 0., 0., 0., 0.],\n",
       "      [0., 0., 2., 0., 0., 0., 0., 0., 0., 0.],\n",
       "      [0., 0., 0., 2., 0., 0., 0., 0., 0., 0.],\n",
       "      [0., 0., 0., 0., 2., 0., 0., 0., 0., 0.],\n",
       "      [0., 0., 0., 0., 0., 2., 0., 0., 0., 0.],\n",
       "      [0., 0., 0., 0., 0., 0., 2., 0., 0., 0.],\n",
       "      [0., 0., 0., 0., 0., 0., 0., 2., 0., 0.],\n",
       "      [0., 0., 0., 0., 0., 0., 0., 0., 2., 0.],\n",
       "      [0., 0., 0., 0., 0., 0., 0., 0., 0., 2.]], dtype=float32)"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "#NBVAL_IGNORE_OUTPUT\n",
    "\n",
    "f.data[:] = np.eye(10)\n",
    "\n",
    "op2 = Operator(Eq(f, f + 1, implicit_dims=ci))\n",
    "print(op2.body.body[-1])\n",
    "op2.apply()\n",
    "\n",
    "assert (np.count_nonzero(f.data - np.diag(np.diagonal(f.data)))==0)\n",
    "assert (np.count_nonzero(f.data) == 10)\n",
    "\n",
    "f.data"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The generated code is as expected and only the elements that were greater than zero were incremented by one.\n",
    "\n",
    "Let's now create a new `Function` `g`, initialized with ones along its diagonal.\n",
    "We want to add `f(x,y)` to `g(x,y)` only if (i) `g != 0` and (ii) `y < 5` (i.e., over the first five columns). To join the two conditions we can use `sympy.And`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Operator `Kernel` ran in 0.01 s\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "START(section0)\n",
      "for (int x = x_m; x <= x_M; x += 1)\n",
      "{\n",
      "  #pragma omp simd aligned(f,g:32)\n",
      "  for (int y = y_m; y <= y_M; y += 1)\n",
      "  {\n",
      "    if (y < 5 && g[x + 1][y + 1] != 0)\n",
      "    {\n",
      "      f[x + 1][y + 1] = f[x + 1][y + 1] + g[x + 1][y + 1];\n",
      "    }\n",
      "  }\n",
      "}\n",
      "STOP(section0,timers)\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "Data([[2., 0., 0., 0., 0., 0., 0., 0., 0., 0.],\n",
       "      [0., 2., 0., 0., 0., 0., 0., 0., 0., 0.],\n",
       "      [0., 0., 2., 0., 0., 0., 0., 0., 0., 0.],\n",
       "      [0., 0., 0., 2., 0., 0., 0., 0., 0., 0.],\n",
       "      [0., 0., 0., 0., 2., 0., 0., 0., 0., 0.],\n",
       "      [0., 0., 0., 0., 0., 1., 0., 0., 0., 0.],\n",
       "      [0., 0., 0., 0., 0., 0., 1., 0., 0., 0.],\n",
       "      [0., 0., 0., 0., 0., 0., 0., 1., 0., 0.],\n",
       "      [0., 0., 0., 0., 0., 0., 0., 0., 1., 0.],\n",
       "      [0., 0., 0., 0., 0., 0., 0., 0., 0., 1.]], dtype=float32)"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "#NBVAL_IGNORE_OUTPUT\n",
    "\n",
    "from sympy import And\n",
    "from devito import Ne, Lt\n",
    "\n",
    "f.data[:] = np.eye(10)\n",
    "\n",
    "\n",
    "g = Function(name='g', grid=grid)\n",
    "g.data[:] = np.eye(10)\n",
    "\n",
    "ci = ConditionalDimension(name='ci', parent=y, condition=And(Ne(g, 0), Lt(y, 5)))\n",
    "\n",
    "op3 = Operator(Eq(f, f + g, implicit_dims=ci))\n",
    "op3.apply()\n",
    "\n",
    "print(op3.body.body[-1])\n",
    "\n",
    "assert (np.count_nonzero(f.data - np.diag(np.diagonal(f.data)))==0)\n",
    "assert (np.count_nonzero(f.data) == 10)\n",
    "assert np.all(f.data[np.nonzero(f.data[:5,:5])] == 2)\n",
    "assert np.all(f.data[5:,5:] == np.eye(5))\n",
    "\n",
    "f.data"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You can see that `f` has been updated only for the first five columns with the `f+g` expression."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "A `ConditionalDimension` can be also used at `Function` construction time. Let's use `ci` from the previous cell to explicitly construct a `Function` `h`. You will notice that in this case there is no need to pass `implicit_dims` to the `Eq` as the `ConditionalDimension` `ci` is now a dimension in `h`. `h` still has the size of the full domain, not the size of the domain that satisfies the condition."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Operator `Kernel` ran in 0.01 s\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "START(section0)\n",
      "for (int x = x_m; x <= x_M; x += 1)\n",
      "{\n",
      "  for (int y = y_m; y <= y_M; y += 1)\n",
      "  {\n",
      "    if (y < 5 && g[x + 1][y + 1] != 0)\n",
      "    {\n",
      "      h[x + 1][y + 1] = g[x + 1][y + 1] + h[x + 1][y + 1];\n",
      "    }\n",
      "  }\n",
      "}\n",
      "STOP(section0,timers)\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "Data([[1., 0., 0., 0., 0., 0., 0., 0., 0., 0.],\n",
       "      [0., 1., 0., 0., 0., 0., 0., 0., 0., 0.],\n",
       "      [0., 0., 1., 0., 0., 0., 0., 0., 0., 0.],\n",
       "      [0., 0., 0., 1., 0., 0., 0., 0., 0., 0.],\n",
       "      [0., 0., 0., 0., 1., 0., 0., 0., 0., 0.],\n",
       "      [0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],\n",
       "      [0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],\n",
       "      [0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],\n",
       "      [0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],\n",
       "      [0., 0., 0., 0., 0., 0., 0., 0., 0., 0.]], dtype=float32)"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "#NBVAL_IGNORE_OUTPUT\n",
    "\n",
    "h = Function(name='h', shape=grid.shape, dimensions=(x, ci))\n",
    "\n",
    "op4 = Operator(Eq(h, h + g))\n",
    "op4.apply()\n",
    "\n",
    "print(op4.body.body[-1])\n",
    "\n",
    "assert (np.count_nonzero(h.data) == 5)\n",
    "\n",
    "h.data"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Use case B - Function subsampling\n",
    "\n",
    "`ConditionalDimension`s are additionally indicated to implement `Function` subsampling. In the following example, an `Operator` subsamples `Function` `g` and saves its content into `f` every `factor=4` iterations."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Operator `Kernel` ran in 0.01 s\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "START(section0)\n",
      "for (int i = i_m; i <= i_M; i += 1)\n",
      "{\n",
      "  if ((i)%(cif) == 0)\n",
      "  {\n",
      "    f[i / cif] = g[i];\n",
      "  }\n",
      "}\n",
      "STOP(section0,timers)\n",
      "\n",
      " Data in g \n",
      " [ 0.  1.  2.  3.  4.  5.  6.  7.  8.  9. 10. 11. 12. 13. 14. 15.]\n",
      "\n",
      " Subsampled data in f \n",
      " [ 0.  4.  8. 12.]\n"
     ]
    }
   ],
   "source": [
    "#NBVAL_IGNORE_OUTPUT\n",
    "size, factor = 16, 4\n",
    "i = Dimension(name='i')\n",
    "ci = ConditionalDimension(name='ci', parent=i, factor=factor)\n",
    "\n",
    "g = Function(name='g', shape=(size,), dimensions=(i,))\n",
    "# Intialize g \n",
    "g.data[:,]= list(range(size))\n",
    "f = Function(name='f', shape=(int(size/factor),), dimensions=(ci,))\n",
    "\n",
    "op5 = Operator([Eq(f, g)])\n",
    "print(op5.body.body[-1])\n",
    "\n",
    "op5.apply()\n",
    "\n",
    "assert np.all(f.data[:] == g.data[0:-1:4])\n",
    "\n",
    "print(\"\\n Data in g \\n\", g.data)\n",
    "print(\"\\n Subsampled data in f \\n\", f.data)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally, we can notice that `f` holds as expected the subsampled values of `g`."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`ConditionalDimension` can be used as an alternative to save snaps to disk as illustrated in [08_snapshotting](https://github.com/devitocodes/devito/blob/master/examples/seismic/tutorials/08_snapshotting.ipynb \"08_snapshotting\"). The example subsamples the `TimeDimension`, capturing snaps of the our wavefield at equally spaced snaps. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Example A - Count nonzero elements\n",
    "\n",
    "Here we present an example of using `ConditionalDimension` by counting the nonzero elements of an array."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Data([[1., 0., 0., 0., 0., 0., 0., 0., 0., 0.],\n",
       "      [0., 1., 0., 0., 0., 0., 0., 0., 0., 0.],\n",
       "      [0., 0., 1., 0., 0., 0., 0., 0., 0., 0.],\n",
       "      [0., 0., 0., 1., 0., 0., 0., 0., 0., 0.],\n",
       "      [0., 0., 0., 0., 1., 0., 0., 0., 0., 0.],\n",
       "      [0., 0., 0., 0., 0., 1., 0., 0., 0., 0.],\n",
       "      [0., 0., 0., 0., 0., 0., 1., 0., 0., 0.],\n",
       "      [0., 0., 0., 0., 0., 0., 0., 1., 0., 0.],\n",
       "      [0., 0., 0., 0., 0., 0., 0., 0., 1., 0.],\n",
       "      [0., 0., 0., 0., 0., 0., 0., 0., 0., 1.]], dtype=float32)"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from devito.types import Scalar\n",
    "from devito import Inc\n",
    "\n",
    "grid = Grid(shape=shape)\n",
    "# Define function 𝑓. We will initialize f's data with ones on its diagonal.\n",
    "f = Function(name='f', shape=shape, dimensions=grid.dimensions, space_order=0)\n",
    "f.data[:] = np.eye(shape[0])\n",
    "f.data[:]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Operator `Kernel` ran in 0.01 s\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "START(section0)\n",
      "for (int x = x_m; x <= x_M; x += 1)\n",
      "{\n",
      "  for (int y = y_m; y <= y_M; y += 1)\n",
      "  {\n",
      "    if (f[x][y] != 0)\n",
      "    {\n",
      "      g[1] += 1;\n",
      "    }\n",
      "  }\n",
      "}\n",
      "STOP(section0,timers)\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "10"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "#NBVAL_IGNORE_OUTPUT\n",
    "# ConditionalDimension with Ne(f, 0) condition\n",
    "ci = ConditionalDimension(name='ci', parent=f.dimensions[-1],\n",
    "                          condition=Ne(f, 0))\n",
    "\n",
    "# g is our counter. g needs to be a Function to write in\n",
    "g = Function(name='g', shape=(1,), dimensions=(ci,), dtype=np.int32)\n",
    "\n",
    "# Extra Scalar used to avoid reduction in gcc-5\n",
    "eq0 = Inc(g[0], 1, implicit_dims=(f.dimensions + (ci,)))\n",
    "\n",
    "op0 = Operator([eq0])\n",
    "op0.apply()\n",
    "\n",
    "assert (np.count_nonzero(f.data) == g.data[0])\n",
    "\n",
    "print(op0.body.body[-1])\n",
    "g.data[0]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Example B - Find nonzero positions \n",
    "\n",
    "Here we present another example of using `ConditionalDimension`. An `Operator` saves the nonzero positions of an array. We know that we have `g.data[0]` nonzero elements from Example A."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Operator `Kernel` ran in 0.01 s\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "int k = -1;\n",
      "\n",
      "START(section0)\n",
      "for (int x = x_m; x <= x_M; x += 1)\n",
      "{\n",
      "  for (int y = y_m; y <= y_M; y += 1)\n",
      "  {\n",
      "    if (f[x][y] != 0)\n",
      "    {\n",
      "      k += 1;\n",
      "      g[k][0] = x;\n",
      "      g[k][1] = y;\n",
      "    }\n",
      "  }\n",
      "}\n",
      "STOP(section0,timers)\n"
     ]
    }
   ],
   "source": [
    "#NBVAL_IGNORE_OUTPUT\n",
    "count = g.data[0] # Number of nonzeros\n",
    "\n",
    "# Dimension used only to nest different size of Functions under the same dim\n",
    "id_dim = Dimension(name='id_dim')\n",
    "\n",
    "# Conditional for nonzero element\n",
    "ci = ConditionalDimension(name='ci', parent=f.dimensions[-1],\n",
    "                             condition=Ne(f, 0))\n",
    "g = Function(name='g', shape=(count, len(f.dimensions)),\n",
    "                dimensions=(ci, id_dim), dtype=np.int32, space_order=0)\n",
    "\n",
    "eqs = []\n",
    "\n",
    "# Extra Scalar used to avoid reduction in gcc-5\n",
    "k = Scalar(name='k', dtype=np.int32)\n",
    "eqi = Eq(k, -1)\n",
    "eqs.append(eqi)\n",
    "eqii = Inc(k, 1, implicit_dims=(f.dimensions + (ci,)))\n",
    "eqs.append(eqii)\n",
    "\n",
    "for n, i in enumerate(f.dimensions):\n",
    "    eqs.append(Eq(g[k, n], f.dimensions[n], implicit_dims=(f.dimensions + (ci,))))\n",
    "\n",
    "# TODO: Must be openmp=False for now due to issue #2061\n",
    "op0 = Operator(eqs, openmp=False)\n",
    "op0.apply()\n",
    "print(op0.body.body[-1])\n",
    "\n",
    "g.data[:]\n",
    "\n",
    "ids_np = np.nonzero(f.data[:])\n",
    "for i in range(len(shape)-1):\n",
    "    assert all(g.data[:, i] == ids_np[i])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# References\n",
    "\n",
    "- [ConditionalDimension](https://github.com/devitocodes/devito/blob/e918e26757a4f0ea3c3c416b6e0f48f4e5023c37/devito/types/dimension.py#L635 \"class ConditionalDimension(DerivedDimension):\") documentation."
   ]
  }
 ],
 "metadata": {
  "celltoolbar": "Tags",
  "kernelspec": {
   "display_name": "Python 3",
   "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.11.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}