dsl: Add Re and Im operators for taking real and imaginary parts of an expression

Author: EdCaunt

devito/finite_differences/differentiable.py

@@ -22,7 +22,7 @@
 from devito.types.basic import AbstractFunction
 
 __all__ = ['Differentiable', 'DiffDerivative', 'IndexDerivative', 'EvalDerivative',
-           'Weights']
+           'Weights', 'Re', 'Im']
 
 
 class Differentiable(sympy.Expr, Evaluable):
@@ -645,6 +645,39 @@ def __str__(self):
     __repr__ = __str__
 
 
+class ComplexPart(Differentiable, sympy.core.function.Application):
+    """Abstract class for `Re` or `Im` of an expression"""
+
+    def __new__(cls, *args, **kwargs):
+        if len(args) != 1:
+            raise ValueError(f"{cls.__name__} is constructed with exactly one arg;"
+                             f" {len(args)} were supplied.")
+
+        # Diffify any Add, Mul, etc which might be in the expression
+        new_args = (diffify(args[0]),)
+
+        if not np.issubdtype(new_args[0].dtype, np.complexfloating):
+            raise ValueError(f"{cls.__name__} requires a complex dtype,"
+                             f" not {new_args[0].dtype}.")
+
+        return super().__new__(cls, *new_args, **kwargs)
+
+    def __str__(self):
+        return f"{self.__class__.__name__}({self.args[0]})"
+
+    __repr__ = __str__
+
+
+class Re(ComplexPart):
+    """Get the real part of an expression"""
+    pass
+
+
+class Im(ComplexPart):
+    """Get the imaginary part of an expression"""
+    pass
+
+
 class IndexSum(sympy.Expr, Evaluable):
 
     """

devito/ir/cgen/printer.py

@@ -217,6 +217,14 @@ def _print_SafeInv(self, expr):
         val = self._print(expr.val)
         return f'SAFEINV({val}, {base})'
 
+    def _print_Re(self, expr):
+        """Print an Re as an access into the second entry of a float array."""
+        return f"Re({self._print(expr.args[0])})"
+
+    def _print_Im(self, expr):
+        """Print an Im as an access into the second entry of a float array."""
+        return f"Im({self._print(expr.args[0])})"
+
     def _print_Mod(self, expr):
         """Print a Mod as a C-like %-based operation."""
         args = [f'({self._print(a)})' for a in expr.args]

devito/ir/equations/equation.py

@@ -8,7 +8,7 @@
 from devito.ir.support import (GuardFactor, Interval, IntervalGroup, IterationSpace,
                                Stencil, detect_io, detect_accesses)
 from devito.symbolics import IntDiv, limits_mapper, uxreplace
-from devito.tools import Pickable, Tag, frozendict
+from devito.tools import Pickable, Tag, frozendict, infer_dtype
 from devito.types import Eq, Inc, ReduceMax, ReduceMin, relational_min
 
 __all__ = ['LoweredEq', 'ClusterizedEq', 'DummyEq', 'OpInc', 'OpMin', 'OpMax',
@@ -48,7 +48,12 @@ def directions(self):
 
     @property
     def dtype(self):
-        return self.lhs.dtype
+        try:
+            rhs_dtype = self.rhs.dtype
+        except AttributeError:
+            rhs_dtype = None
+
+        return infer_dtype({self.lhs.dtype, rhs_dtype} - {None})
 
     @property
     def state(self):

devito/passes/iet/misc.py

@@ -5,7 +5,7 @@
 import sympy
 
 from devito.finite_differences import Max, Min
-from devito.finite_differences.differentiable import SafeInv
+from devito.finite_differences.differentiable import SafeInv, Re, Im
 from devito.logger import warning
 from devito.ir import (Any, Forward, DummyExpr, Iteration, List, Prodder,
                        FindApplications, FindNodes, FindSymbols, Transformer,
@@ -245,6 +245,22 @@ def _(expr, langbb):
              f'(((a) < {eps}F || ({b}) < {eps}F) ? (0.0F) : ((1.0F) / (a)))'),), {}
 
 
+@_lower_macro_math.register(Re)
+def _(expr, langbb):
+    return (('Re(x)',
+             '(_Generic((x), '
+             'float _Complex : (float *) &(x), '
+             'double _Complex : (double *) &(x))[0])'),), {}
+
+
+@_lower_macro_math.register(Im)
+def _(expr, langbb):
+    return (('Im(x)',
+             '(_Generic((x), '
+             'float _Complex : (float *) &(x), '
+             'double _Complex : (double *) &(x))[1])'),), {}
+
+
 @iet_pass
 def minimize_symbols(iet):
     """

tests/test_symbolics.py

@@ -7,7 +7,7 @@
 from sympy import Expr, Symbol
 from devito import (Constant, Dimension, Grid, Function, solve, TimeFunction, Eq,  # noqa
                     Operator, SubDimension, norm, Le, Ge, Gt, Lt, Abs, sin, cos,
-                    Min, Max)
+                    Min, Max, Re, Im)
 from devito.finite_differences.differentiable import SafeInv, Weights
 from devito.ir import Expression, FindNodes, ccode
 from devito.symbolics import (retrieve_functions, retrieve_indexed, evalrel,  # noqa
@@ -818,3 +818,108 @@ def test_assumptions(self, op, expr, assumptions, expected):
         assumptions = eval(assumptions)
         expected = eval(expected)
         assert evalrel(op, eqn, assumptions) == expected
+
+
+class TestComplexParts:
+    def setup_basic(self, dtype):
+        grid = Grid(shape=(5,), extent=(4.,))
+        f = Function(name='f', grid=grid, dtype=dtype)
+        f.data_with_halo[:] = np.arange(7) + 1j*np.arange(7, 14)[::-1]
+
+        f_real = Function(name='f_real', grid=grid)
+        f_imag = Function(name='f_imag', grid=grid)
+        return f, f_real, f_imag
+
+    def test_printing(self):
+        f, f_real, f_imag = self.setup_basic(np.complex64)
+
+        eq_re = Eq(f_real, Re(f))
+        eq_im = Eq(f_imag, Im(f))
+
+        op = Operator([eq_re, eq_im])
+
+        assert ("#define Im(x) (_Generic((x), float _Complex : (float *) &(x), "
+                "double _Complex : (double *) &(x))[1])") in str(op.ccode)
+        assert ("#define Re(x) (_Generic((x), float _Complex : (float *) &(x), "
+                "double _Complex : (double *) &(x))[0])") in str(op.ccode)
+
+        assert "f_real[x + 1] = Re(f[x + 1])" in str(op.ccode)
+        assert "f_imag[x + 1] = Im(f[x + 1])" in str(op.ccode)
+
+    @pytest.mark.parametrize('dtype', [np.complex64, np.complex128])
+    def test_trivial(self, dtype):
+        f, f_real, f_imag = self.setup_basic(dtype)
+
+        eq_re = Eq(f_real, Re(f+1))
+        eq_im = Eq(f_imag, Im(f+1))
+
+        Operator([eq_re, eq_im])()
+
+        rcheck = np.array([2., 3., 4., 5., 6.])
+        icheck = np.array([12., 11., 10., 9., 8.])
+        assert np.all(np.isclose(f_real.data, rcheck))
+        assert np.all(np.isclose(f_imag.data, icheck))
+
+    @pytest.mark.parametrize('dtype', [np.complex64, np.complex128])
+    def test_trivial_imag(self, dtype):
+        f, f_real, f_imag = self.setup_basic(dtype)
+
+        eq_re = Eq(f_real, Re(f+1j))
+        eq_im = Eq(f_imag, Im(f+1j))
+
+        Operator([eq_re, eq_im])()
+
+        rcheck = np.array([1., 2., 3., 4., 5.])
+        icheck = np.array([13., 12., 11., 10., 9.])
+        assert np.all(np.isclose(f_real.data, rcheck))
+        assert np.all(np.isclose(f_imag.data, icheck))
+
+    def test_deriv(self):
+        f, f_real, f_imag = self.setup_basic(np.complex64)
+
+        eq_re = Eq(f_real, Re(f.dx))
+        eq_im = Eq(f_imag, Im(f.dx))
+
+        Operator([eq_re, eq_im])()
+
+        assert np.all(np.isclose(f_real.data, 1.))
+        assert np.all(np.isclose(f_imag.data, -1.))
+
+    def test_outer_deriv(self):
+        f, f_real, f_imag = self.setup_basic(np.complex64)
+
+        eq_re = Eq(f_real, Re(f).dx)
+        eq_im = Eq(f_imag, Im(f).dx)
+
+        Operator([eq_re, eq_im])()
+
+        assert np.all(np.isclose(f_real.data, 1.))
+        assert np.all(np.isclose(f_imag.data, -1.))
+
+    def test_mul(self):
+        grid = Grid(shape=(5,))
+
+        f = Function(name='f', grid=grid, dtype=np.complex64)
+        g = Function(name='g', grid=grid)
+        h = Function(name='h', grid=grid, dtype=np.complex64)
+        f.data[:] = 1 + 1j
+        g.data[:] = 2
+        h.data[:] = 2j
+
+        fg_re = Function(name='fg_re', grid=grid)
+        fg_im = Function(name='fg_im', grid=grid)
+        fh_re = Function(name='fh_re', grid=grid)
+        fh_im = Function(name='fh_im', grid=grid)
+
+        eq_fg_re = Eq(fg_re, Re(f*g))
+        eq_fg_im = Eq(fg_im, Im(f*g))
+        eq_fh_re = Eq(fh_re, Re(f*h))
+        eq_fh_im = Eq(fh_im, Im(f*h))
+
+        Operator([eq_fg_re, eq_fg_im, eq_fh_re, eq_fh_im])()
+
+        assert np.all(np.isclose(fg_re.data, 2.))
+        assert np.all(np.isclose(fg_im.data, 2.))
+
+        assert np.all(np.isclose(fh_re.data, -2.))
+        assert np.all(np.isclose(fh_im.data, 2.))