api: fix printer for complex dtype

Author: mloubout

Diff for: devito/finite_differences/differentiable.py

@@ -14,8 +14,7 @@
 from devito.logger import warning
 from devito.tools import (as_tuple, filter_ordered, flatten, frozendict,
                           infer_dtype, is_integer, split)
-from devito.types import (Array, DimensionTuple, Evaluable, Indexed,
-                          StencilDimension)
+from devito.types import Array, DimensionTuple, Evaluable, StencilDimension
 
 __all__ = ['Differentiable', 'DiffDerivative', 'IndexDerivative', 'EvalDerivative',
            'Weights']

Diff for: devito/passes/iet/misc.py

@@ -3,7 +3,6 @@
 import cgen
 import numpy as np
 import sympy
-import numpy as np
 
 from devito import configuration
 from devito.finite_differences import Max, Min

Diff for: devito/symbolics/inspection.py

@@ -3,7 +3,6 @@
 import numpy as np
 from sympy import (Function, Indexed, Integer, Mul, Number,
                    Pow, S, Symbol, Tuple)
-from sympy.core.operations import AssocOp
 from sympy.core.numbers import ImaginaryUnit
 
 from devito.finite_differences import Derivative

Diff for: devito/symbolics/printer.py

@@ -103,7 +103,8 @@ def _print_math_func(self, expr, nest=False, known=None):
         dtype = sympy_dtype(expr)
         if dtype is np.float32:
             cname += 'f'
-
+        if np.issubdtype(self.dtype, np.complexfloating):
+            cname = 'c%s' % cname
         args = ', '.join((self._print(arg) for arg in expr.args))
 
         return '%s(%s)' % (cname, args)
@@ -250,6 +251,8 @@ def _print_TrigonometricFunction(self, expr):
         func_name = str(expr.func)
         if self.dtype == np.float32:
             func_name += 'f'
+        if np.issubdtype(self.dtype, np.complexfloating):
+            func_name = 'c%s' % func_name
         return '%s(%s)' % (func_name, self._print(*expr.args))
 
     def _print_DefFunction(self, expr):

Diff for: devito/types/basic.py

@@ -1440,7 +1440,7 @@ def _C_ctype(self):
         try:
             if np.issubdtype(self.dtype, np.complexfloating):
                 rtype = self.dtype(0).real.__class__
-                ctname = '%s _Complex' % dtype_to_cstr(rtype)
+                ctname = '%s complex' % dtype_to_cstr(rtype)
                 ctype = dtype_to_ctype(rtype)
                 r = type(ctname, (ctype,), {})
                 return POINTER(r)

Diff for: tests/test_operator.py

@@ -9,7 +9,7 @@
                     SparseFunction, SparseTimeFunction, Dimension, error, SpaceDimension,
                     NODE, CELL, dimensions, configuration, TensorFunction,
                     TensorTimeFunction, VectorFunction, VectorTimeFunction,
-                    div, grad, switchconfig)
+                    div, grad, switchconfig, exp)
 from devito import  Inc, Le, Lt, Ge, Gt  # noqa
 from devito.exceptions import InvalidOperator
 from devito.finite_differences.differentiable import diff2sympy
@@ -640,6 +640,24 @@ def test_tensor(self, func1):
         op2 = Operator([Eq(f, f.dx) for f in f1.values()])
         assert str(op1.ccode) == str(op2.ccode)
 
+    def test_complex(self):
+        grid = Grid((5, 5))
+        x, y = grid.dimensions
+        # Float32 complex is called complex64 in numpy
+        u = Function(name="u", grid=grid, dtype=np.complex64)
+
+        eq = Eq(u, x + 1j*y + exp(1j + x.spacing))
+        # Currently wrong alias type
+        op = Operator(eq, opt='noop')
+        op()
+
+        # Check against numpy
+        dx = grid.spacing_map[x.spacing]
+        xx, yy = np.meshgrid(np.linspace(0, 4, 5), np.linspace(0, 4, 5))
+        npres = xx + 1j*yy + np.exp(1j + dx)
+
+        assert np.allclose(u.data, npres.T, rtol=1e-7, atol=0)
+
 
 class TestAllocation(object):
 
@@ -724,10 +742,10 @@ def verify_parameters(self, parameters, expected):
         """
         boilerplate = ['timers']
         parameters = [p.name for p in parameters]
-        for exp in expected:
-            if exp not in parameters + boilerplate:
-                error("Missing parameter: %s" % exp)
-            assert exp in parameters + boilerplate
+        for expi in expected:
+            if expi not in parameters + boilerplate:
+                error("Missing parameter: %s" % expi)
+            assert expi in parameters + boilerplate
         extra = [p for p in parameters if p not in expected and p not in boilerplate]
         if len(extra) > 0:
             error("Redundant parameters: %s" % str(extra))