ConditionalDimension tutorial

This tutorial explains how to create equations that only get executed under given conditions.

from devito import Dimension, Function, Grid
import numpy as np

# We define a 10x10 grid, dimensions are x, y
shape = (10, 10)
grid = Grid(shape = shape)
x, y = grid.dimensions

# Define function 𝑓. We will initialize f's data with ones on its diagonal.

f = Function(name='f', grid=grid)
f.data[:] = np.eye(10)
f.data

Data([[1., 0., 0., 0., 0., 0., 0., 0., 0., 0.],
      [0., 1., 0., 0., 0., 0., 0., 0., 0., 0.],
      [0., 0., 1., 0., 0., 0., 0., 0., 0., 0.],
      [0., 0., 0., 1., 0., 0., 0., 0., 0., 0.],
      [0., 0., 0., 0., 1., 0., 0., 0., 0., 0.],
      [0., 0., 0., 0., 0., 1., 0., 0., 0., 0.],
      [0., 0., 0., 0., 0., 0., 1., 0., 0., 0.],
      [0., 0., 0., 0., 0., 0., 0., 1., 0., 0.],
      [0., 0., 0., 0., 0., 0., 0., 0., 1., 0.],
      [0., 0., 0., 0., 0., 0., 0., 0., 0., 1.]], dtype=float32)

We begin by constructing an Operator that increases the values of f, a Function defined on a two-dimensional Grid, by one.

#NBVAL_IGNORE_OUTPUT
from devito import Eq, Operator
op0 = Operator(Eq(f, f + 1))
op0.apply()
f.data

Operator `Kernel` ran in 0.01 s

Data([[2., 1., 1., 1., 1., 1., 1., 1., 1., 1.],
      [1., 2., 1., 1., 1., 1., 1., 1., 1., 1.],
      [1., 1., 2., 1., 1., 1., 1., 1., 1., 1.],
      [1., 1., 1., 2., 1., 1., 1., 1., 1., 1.],
      [1., 1., 1., 1., 2., 1., 1., 1., 1., 1.],
      [1., 1., 1., 1., 1., 2., 1., 1., 1., 1.],
      [1., 1., 1., 1., 1., 1., 2., 1., 1., 1.],
      [1., 1., 1., 1., 1., 1., 1., 2., 1., 1.],
      [1., 1., 1., 1., 1., 1., 1., 1., 2., 1.],
      [1., 1., 1., 1., 1., 1., 1., 1., 1., 2.]], dtype=float32)

Every value has been updated by one. You should be able to see twos on the diagonal and ones everywhere else.

#print(op0.ccode) # Print the generated code

Devito API for relationals

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:

• Le (less-than or equal)
• Lt (less-than)
• Ge (greater-than or equal)
• Gt (greater-than)
• Ne (not equal)

Relationals are used to define a condition and are passed as an argument to ConditionalDimension at construction time.

Devito API for ConditionalDimension

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.

A ConditionalDimension is used in two typical cases:

• Use case A: To constrain the execution of loop iterations to make sure certain conditions are honoured

• Use case B: To subsample a Function

from devito import ConditionalDimension
print(ConditionalDimension.__doc__)

    Symbol defining 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.

    Parameters
    ----------
    name : str
        Name of the dimension.
    parent : Dimension
        The parent Dimension.
    factor : int, optional, default=None
        The number of iterations between two executions of the if-branch. If None
        (default), ``condition`` must be provided.
    condition : expr-like, optional, default=None
        An arbitrary SymPy expression, typically involving the ``parent``
        Dimension. When it evaluates to True, the if-branch is executed. If None
        (default), ``factor`` must be provided.
    indirect : bool, optional, default=False
        If True, use `self`, rather than the parent Dimension, to
        index into arrays. A typical use case is when arrays are accessed
        indirectly via the ``condition`` expression.

    Examples
    --------
    Among the other things, ConditionalDimensions are indicated to implement
    Function subsampling. In the following example, an Operator evaluates the
    Function ``g`` and saves its content into ``f`` every ``factor=4`` iterations.

    >>> from devito import Dimension, ConditionalDimension, Function, Eq, Operator
    >>> size, factor = 16, 4
    >>> i = Dimension(name='i')
    >>> ci = ConditionalDimension(name='ci', parent=i, factor=factor)
    >>> g = Function(name='g', shape=(size,), dimensions=(i,))
    >>> f = Function(name='f', shape=(int(size/factor),), dimensions=(ci,))
    >>> op = Operator([Eq(g, 1), Eq(f, g)])

    The Operator generates the following for-loop (pseudocode)

    .. code-block:: C

        for (int i = i_m; i <= i_M; i += 1) {
          g[i] = 1;
          if (i%4 == 0) {
            f[i / 4] = g[i];
          }
        }

    Another typical use case is when one needs to constrain the execution of
    loop iterations so that certain conditions are honoured. The following
    artificial example uses ConditionalDimension to guard against out-of-bounds
    accesses in indirectly accessed arrays.

    >>> from sympy import And
    >>> ci = ConditionalDimension(name='ci', parent=i,
    ...                           condition=And(g[i] > 0, g[i] < 4, evaluate=False))
    >>> f = Function(name='f', shape=(int(size/factor),), dimensions=(ci,))
    >>> op = Operator(Eq(f[g[i]], f[g[i]] + 1))

    The Operator generates the following for-loop (pseudocode)

    .. code-block:: C

        for (int i = i_m; i <= i_M; i += 1) {
          if (g[i] > 0 && g[i] < 4) {
            f[g[i]] = f[g[i]] + 1;
          }
        }

    

Use case A - Honour a condition

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. 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.

#NBVAL_IGNORE_OUTPUT
from devito import Gt

f.data[:] = np.eye(10)

# Define the condition f(x, y) > 0
condition = Gt(f, 0)

# Construct a ConditionalDimension
ci = ConditionalDimension(name='ci', parent=y, condition=condition)

op1 = Operator(Eq(f, f + 1))
op1.apply()
f.data
# print(op1.ccode) # Uncomment to view code

Operator `Kernel` ran in 0.01 s

Data([[2., 1., 1., 1., 1., 1., 1., 1., 1., 1.],
      [1., 2., 1., 1., 1., 1., 1., 1., 1., 1.],
      [1., 1., 2., 1., 1., 1., 1., 1., 1., 1.],
      [1., 1., 1., 2., 1., 1., 1., 1., 1., 1.],
      [1., 1., 1., 1., 2., 1., 1., 1., 1., 1.],
      [1., 1., 1., 1., 1., 2., 1., 1., 1., 1.],
      [1., 1., 1., 1., 1., 1., 2., 1., 1., 1.],
      [1., 1., 1., 1., 1., 1., 1., 2., 1., 1.],
      [1., 1., 1., 1., 1., 1., 1., 1., 2., 1.],
      [1., 1., 1., 1., 1., 1., 1., 1., 1., 2.]], dtype=float32)

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.

#NBVAL_IGNORE_OUTPUT

f.data[:] = np.eye(10)

op2 = Operator(Eq(f, f + 1, implicit_dims=ci))
print(op2.body.body[-1])
op2.apply()

assert (np.count_nonzero(f.data - np.diag(np.diagonal(f.data)))==0)
assert (np.count_nonzero(f.data) == 10)

f.data

Operator `Kernel` ran in 0.01 s

START(section0)
for (int x = x_m; x <= x_M; x += 1)
{
  #pragma omp simd aligned(f:32)
  for (int y = y_m; y <= y_M; y += 1)
  {
    if (f[x + 1][y + 1] > 0)
    {
      f[x + 1][y + 1] = f[x + 1][y + 1] + 1;
    }
  }
}
STOP(section0,timers)

Data([[2., 0., 0., 0., 0., 0., 0., 0., 0., 0.],
      [0., 2., 0., 0., 0., 0., 0., 0., 0., 0.],
      [0., 0., 2., 0., 0., 0., 0., 0., 0., 0.],
      [0., 0., 0., 2., 0., 0., 0., 0., 0., 0.],
      [0., 0., 0., 0., 2., 0., 0., 0., 0., 0.],
      [0., 0., 0., 0., 0., 2., 0., 0., 0., 0.],
      [0., 0., 0., 0., 0., 0., 2., 0., 0., 0.],
      [0., 0., 0., 0., 0., 0., 0., 2., 0., 0.],
      [0., 0., 0., 0., 0., 0., 0., 0., 2., 0.],
      [0., 0., 0., 0., 0., 0., 0., 0., 0., 2.]], dtype=float32)

The generated code is as expected and only the elements that were greater than zero were incremented by one.

Let’s now create a new Function g, initialized with ones along its diagonal. 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.

#NBVAL_IGNORE_OUTPUT

from sympy import And
from devito import Ne, Lt

f.data[:] = np.eye(10)


g = Function(name='g', grid=grid)
g.data[:] = np.eye(10)

ci = ConditionalDimension(name='ci', parent=y, condition=And(Ne(g, 0), Lt(y, 5)))

op3 = Operator(Eq(f, f + g, implicit_dims=ci))
op3.apply()

print(op3.body.body[-1])

assert (np.count_nonzero(f.data - np.diag(np.diagonal(f.data)))==0)
assert (np.count_nonzero(f.data) == 10)
assert np.all(f.data[np.nonzero(f.data[:5,:5])] == 2)
assert np.all(f.data[5:,5:] == np.eye(5))

f.data

Operator `Kernel` ran in 0.01 s

START(section0)
for (int x = x_m; x <= x_M; x += 1)
{
  #pragma omp simd aligned(f,g:32)
  for (int y = y_m; y <= y_M; y += 1)
  {
    if (y < 5 && g[x + 1][y + 1] != 0)
    {
      f[x + 1][y + 1] = f[x + 1][y + 1] + g[x + 1][y + 1];
    }
  }
}
STOP(section0,timers)

Data([[2., 0., 0., 0., 0., 0., 0., 0., 0., 0.],
      [0., 2., 0., 0., 0., 0., 0., 0., 0., 0.],
      [0., 0., 2., 0., 0., 0., 0., 0., 0., 0.],
      [0., 0., 0., 2., 0., 0., 0., 0., 0., 0.],
      [0., 0., 0., 0., 2., 0., 0., 0., 0., 0.],
      [0., 0., 0., 0., 0., 1., 0., 0., 0., 0.],
      [0., 0., 0., 0., 0., 0., 1., 0., 0., 0.],
      [0., 0., 0., 0., 0., 0., 0., 1., 0., 0.],
      [0., 0., 0., 0., 0., 0., 0., 0., 1., 0.],
      [0., 0., 0., 0., 0., 0., 0., 0., 0., 1.]], dtype=float32)

You can see that f has been updated only for the first five columns with the f+g expression.

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.

#NBVAL_IGNORE_OUTPUT

h = Function(name='h', shape=grid.shape, dimensions=(x, ci))

op4 = Operator(Eq(h, h + g))
op4.apply()

print(op4.body.body[-1])

assert (np.count_nonzero(h.data) == 5)

h.data

Operator `Kernel` ran in 0.01 s

START(section0)
for (int x = x_m; x <= x_M; x += 1)
{
  for (int y = y_m; y <= y_M; y += 1)
  {
    if (y < 5 && g[x + 1][y + 1] != 0)
    {
      h[x + 1][y + 1] = g[x + 1][y + 1] + h[x + 1][y + 1];
    }
  }
}
STOP(section0,timers)

Data([[1., 0., 0., 0., 0., 0., 0., 0., 0., 0.],
      [0., 1., 0., 0., 0., 0., 0., 0., 0., 0.],
      [0., 0., 1., 0., 0., 0., 0., 0., 0., 0.],
      [0., 0., 0., 1., 0., 0., 0., 0., 0., 0.],
      [0., 0., 0., 0., 1., 0., 0., 0., 0., 0.],
      [0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],
      [0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],
      [0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],
      [0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],
      [0., 0., 0., 0., 0., 0., 0., 0., 0., 0.]], dtype=float32)

Use case B - Function subsampling

ConditionalDimensions 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.

#NBVAL_IGNORE_OUTPUT
size, factor = 16, 4
i = Dimension(name='i')
ci = ConditionalDimension(name='ci', parent=i, factor=factor)

g = Function(name='g', shape=(size,), dimensions=(i,))
# Intialize g 
g.data[:,]= list(range(size))
f = Function(name='f', shape=(int(size/factor),), dimensions=(ci,))

op5 = Operator([Eq(f, g)])
print(op5.body.body[-1])

op5.apply()

assert np.all(f.data[:] == g.data[0:-1:4])

print("\n Data in g \n", g.data)
print("\n Subsampled data in f \n", f.data)

Operator `Kernel` ran in 0.01 s

START(section0)
for (int i = i_m; i <= i_M; i += 1)
{
  if ((i)%(cif) == 0)
  {
    f[i / cif] = g[i];
  }
}
STOP(section0,timers)

 Data in g 
 [ 0.  1.  2.  3.  4.  5.  6.  7.  8.  9. 10. 11. 12. 13. 14. 15.]

 Subsampled data in f 
 [ 0.  4.  8. 12.]

Finally, we can notice that f holds as expected the subsampled values of g.

ConditionalDimension can be used as an alternative to save snaps to disk as illustrated in 08_snapshotting. The example subsamples the TimeDimension, capturing snaps of the our wavefield at equally spaced snaps.

Example A - Count nonzero elements

Here we present an example of using ConditionalDimension by counting the nonzero elements of an array.

from devito.types import Scalar
from devito import Inc

grid = Grid(shape=shape)
# Define function 𝑓. We will initialize f's data with ones on its diagonal.
f = Function(name='f', shape=shape, dimensions=grid.dimensions, space_order=0)
f.data[:] = np.eye(shape[0])
f.data[:]

Data([[1., 0., 0., 0., 0., 0., 0., 0., 0., 0.],
      [0., 1., 0., 0., 0., 0., 0., 0., 0., 0.],
      [0., 0., 1., 0., 0., 0., 0., 0., 0., 0.],
      [0., 0., 0., 1., 0., 0., 0., 0., 0., 0.],
      [0., 0., 0., 0., 1., 0., 0., 0., 0., 0.],
      [0., 0., 0., 0., 0., 1., 0., 0., 0., 0.],
      [0., 0., 0., 0., 0., 0., 1., 0., 0., 0.],
      [0., 0., 0., 0., 0., 0., 0., 1., 0., 0.],
      [0., 0., 0., 0., 0., 0., 0., 0., 1., 0.],
      [0., 0., 0., 0., 0., 0., 0., 0., 0., 1.]], dtype=float32)

#NBVAL_IGNORE_OUTPUT
# ConditionalDimension with Ne(f, 0) condition
ci = ConditionalDimension(name='ci', parent=f.dimensions[-1],
                          condition=Ne(f, 0))

# g is our counter. g needs to be a Function to write in
g = Function(name='g', shape=(1,), dimensions=(ci,), dtype=np.int32)

# Extra Scalar used to avoid reduction in gcc-5
eq0 = Inc(g[0], 1, implicit_dims=(f.dimensions + (ci,)))

op0 = Operator([eq0])
op0.apply()

assert (np.count_nonzero(f.data) == g.data[0])

print(op0.body.body[-1])
g.data[0]

Operator `Kernel` ran in 0.01 s

START(section0)
for (int x = x_m; x <= x_M; x += 1)
{
  for (int y = y_m; y <= y_M; y += 1)
  {
    if (f[x][y] != 0)
    {
      g[1] += 1;
    }
  }
}
STOP(section0,timers)

10

Example B - Find nonzero positions

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.

#NBVAL_IGNORE_OUTPUT
count = g.data[0] # Number of nonzeros

# Dimension used only to nest different size of Functions under the same dim
id_dim = Dimension(name='id_dim')

# Conditional for nonzero element
ci = ConditionalDimension(name='ci', parent=f.dimensions[-1],
                             condition=Ne(f, 0))
g = Function(name='g', shape=(count, len(f.dimensions)),
                dimensions=(ci, id_dim), dtype=np.int32, space_order=0)

eqs = []

# Extra Scalar used to avoid reduction in gcc-5
k = Scalar(name='k', dtype=np.int32)
eqi = Eq(k, -1)
eqs.append(eqi)
eqii = Inc(k, 1, implicit_dims=(f.dimensions + (ci,)))
eqs.append(eqii)

for n, i in enumerate(f.dimensions):
    eqs.append(Eq(g[k, n], f.dimensions[n], implicit_dims=(f.dimensions + (ci,))))

# TODO: Must be language='C' for now due to issue #2061
op0 = Operator(eqs, language='C')
op0.apply()
print(op0.body.body[-1])

g.data[:]

ids_np = np.nonzero(f.data[:])
for i in range(len(shape)-1):
    assert all(g.data[:, i] == ids_np[i])

Operator `Kernel` ran in 0.01 s

int k = -1;

START(section0)
for (int x = x_m; x <= x_M; x += 1)
{
  for (int y = y_m; y <= y_M; y += 1)
  {
    if (f[x][y] != 0)
    {
      k += 1;
      g[k][0] = x;
      g[k][1] = y;
    }
  }
}
STOP(section0,timers)

References

• ConditionalDimension documentation.