api: fix interpolate with complex dtype

Author: mloubout

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'