#!/usr/bin/env python
# coding: utf-8

# # `ConditionalDimension` tutorial

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

# In[1]:


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


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

# In[2]:


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


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

# In[3]:


#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](https://docs.sympy.org/latest/modules/core.html#module-sympy.core.relational):
# 
# - 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`

# In[4]:


from devito import ConditionalDimension
print(ConditionalDimension.__doc__)


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

# In[5]:


#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


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

# In[6]:


#NBVAL_IGNORE_OUTPUT

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

op2 = Operator(Eq(f, f + 1, implicit_dims=ci))
print(op2.ccode)
op2.apply()

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

f.data


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

# In[7]:


#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.ccode)

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


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

# In[8]:


#NBVAL_IGNORE_OUTPUT

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

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

print(op4.ccode)

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

h.data


# # Use case B (Function subsampling)
# 
# `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.

# In[9]:


#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.ccode)

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)


# 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](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. 

# # References
# 
# - [ConditionalDimension](https://github.com/devitocodes/devito/blob/e918e26757a4f0ea3c3c416b6e0f48f4e5023c37/devito/types/dimension.py#L635 "class ConditionalDimension(DerivedDimension):") documentation.