api: fix interpolate with complex dtype

Author: mloubout

devito/finite_differences/finite_difference.py

@@ -157,8 +157,6 @@ def first_derivative(expr, dim, fd_order, **kwargs):
 
 def make_derivative(expr, dim, fd_order, deriv_order, side, matvec, x0, coefficients,
                     expand, weights=None):
-    if deriv_order == 0 and not expr.is_Add:
-        print(expr, dim, fd_order)
     # Always expand time derivatives to avoid issue with buffering and streaming.
     # Time derivative are almost always short stencils and won't benefit from
     # unexpansion in the rare case the derivative is not evaluated for time stepping.

devito/ir/equations/equation.py

@@ -234,6 +234,19 @@ def __new__(cls, *args, **kwargs):
                 shift = 0
             expr = uxreplace(expr, {d: IntDiv(index, d.symbolic_factor) + shift})
 
+        # Merge conditionals when possible. E.g if we have an implicit_dim
+        # and there is a dimension with the same parent, we ca merged
+        # its condition
+        for d in input_expr.implicit_dims:
+            if d not in conditionals:
+                continue
+            for cd in dict(conditionals):
+                if cd.parent == d.parent and cd != d:
+                    cond = conditionals.pop(d)
+                    mode = cd.relation and d.relation
+                    conditionals[cd] = mode(cond, conditionals[cd])
+                    break
+
         conditionals = frozendict(conditionals)
 
         # Lower all Differentiable operations into SymPy operations

devito/passes/clusters/cse.py

@@ -36,6 +36,18 @@ def retrieve_ctemps(exprs, mode='all'):
     return search(exprs, lambda expr: isinstance(expr, CTemp), mode, 'dfs')
 
 
+def cse_dtype(exprdtype, cdtype):
+    """
+    Return the dtype of a CSE temporary given the dtype of the expression to be
+    captured and the cluster's dtype.
+    """
+    if np.issubdtype(cdtype, np.complexfloating):
+        return np.promote_types(exprdtype, cdtype(0).real.__class__).type
+    else:
+        # Real cluster, can safely promote to the largest precision
+        return np.promote_types(exprdtype, cdtype).type
+
+
 @cluster_pass
 def cse(cluster, sregistry=None, options=None, **kwargs):
     """
@@ -86,7 +98,7 @@ def cse(cluster, sregistry=None, options=None, **kwargs):
     if cluster.is_fence:
         return cluster
 
-    make_dtype = lambda e: np.promote_types(e.dtype, dtype).type
+    make_dtype = lambda e: cse_dtype(e.dtype, dtype)
     make = lambda e: CTemp(name=sregistry.make_name(), dtype=make_dtype(e))
 
     exprs = _cse(cluster, make, min_cost=min_cost, mode=mode)

tests/test_interpolation.py

@@ -855,6 +855,21 @@ def test_point_symbol_types(dtype, expected):
     assert point_symbol.dtype is expected
 
 
+@pytest.mark.parametrize('dtype', [np.complex64, np.complex128])
+def test_interp_complex(dtype):
+    grid = Grid((11, 11, 11))
+
+    sc = SparseFunction(name="sc", grid=grid, npoint=1, dtype=dtype)
+    sc.coordinates.data[:] = [.5, .5, .5]
+
+    fc = Function(name="fc", grid=grid, npoint=2, dtype=dtype)
+    fc.data[:] = np.random.randn(*grid.shape) + 1j * np.random.randn(*grid.shape)
+    opC = Operator([sc.interpolate(expr=fc)], name="OpC")
+    opC()
+
+    assert np.isclose(sc.data[0], fc.data[5, 5, 5])
+
+
 class SD0(SubDomain):
     name = 'sd0'