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.__name__}.")
+
+        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/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/languages/C.py

@@ -55,3 +55,13 @@ class CPrinter(BasePrinter, C99CodePrinter):
 
     def _print_ImaginaryUnit(self, expr):
         return '_Complex_I'
+
+    def _print_Re(self, expr):
+        """Print an Re as an access into the second entry of a float array."""
+        return (f'{self.func_prefix(expr)}real{self.func_literal(expr).lower()}'
+                f'({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'{self.func_prefix(expr)}imag{self.func_literal(expr).lower()}'
+                f'({self._print(expr.args[0])})')

devito/passes/iet/languages/CXX.py

@@ -103,6 +103,13 @@ class CXXPrinter(BasePrinter, CXX11CodePrinter):
 
     def _print_ImaginaryUnit(self, expr):
         return f'1i{self.prec_literal(expr).lower()}'
+        # return '1i'
+
+    def _print_Re(self, expr):
+        return f'{self._ns}real({self._print(expr.args[0])})'
+
+    def _print_Im(self, expr):
+        return f'{self._ns}imag({self._print(expr.args[0])})'
 
     def _print_Cast(self, expr):
         # The CXX recommended way to cast is to use static_cast

tests/test_symbolics.py

@@ -7,7 +7,7 @@
 from sympy import Expr, Number, 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, switchconfig)
 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
@@ -874,3 +874,123 @@ def test_assumptions(self, op, expr, assumptions, expected):
         assumptions = eval(assumptions)
         expected = eval(expected)
         assert evalrel(op, eqn, assumptions) == expected
+
+
+class TestComplexParts:
+    # TODO: Add a cxx switchconfig
+    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 run_operator(self, eqs, cxx):
+        if cxx:
+            with switchconfig(language='CXX'):
+                Operator(eqs)()
+        else:
+            Operator(eqs)()
+
+    @pytest.mark.parametrize('cxx', [False, True])
+    def test_printing(self, cxx):
+        f, f_real, f_imag = self.setup_basic(np.complex64)
+
+        eq_re = Eq(f_real, Re(f))
+        eq_im = Eq(f_imag, Im(f))
+
+        if cxx:
+            with switchconfig(language='CXX'):
+                op = Operator([eq_re, eq_im])
+                assert "f_real[x + 1] = std::real(f[x + 1])" in str(op.ccode)
+                assert "f_imag[x + 1] = std::imag(f[x + 1])" in str(op.ccode)
+
+        else:
+            op = Operator([eq_re, eq_im])
+            assert "f_real[x + 1] = crealf(f[x + 1])" in str(op.ccode)
+            assert "f_imag[x + 1] = cimagf(f[x + 1])" in str(op.ccode)
+
+    @pytest.mark.parametrize('cxx', [False, True])
+    @pytest.mark.parametrize('dtype', [np.complex64, np.complex128])
+    def test_trivial(self, cxx, 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.))
+
+        self.run_operator([eq_re, eq_im], cxx)
+
+        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('cxx', [False, True])
+    @pytest.mark.parametrize('dtype', [np.complex64, np.complex128])
+    def test_trivial_imag(self, cxx, 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))
+
+        self.run_operator([eq_re, eq_im], cxx)
+
+        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))
+
+    @pytest.mark.parametrize('cxx', [False, True])
+    def test_deriv(self, cxx):
+        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))
+
+        self.run_operator([eq_re, eq_im], cxx)
+
+        assert np.all(np.isclose(f_real.data, 1.))
+        assert np.all(np.isclose(f_imag.data, -1.))
+
+    @pytest.mark.parametrize('cxx', [False, True])
+    def test_outer_deriv(self, cxx):
+        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)
+
+        self.run_operator([eq_re, eq_im], cxx)
+
+        assert np.all(np.isclose(f_real.data, 1.))
+        assert np.all(np.isclose(f_imag.data, -1.))
+
+    @pytest.mark.parametrize('cxx', [False, True])
+    def test_mul(self, cxx):
+        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))
+
+        self.run_operator([eq_fg_re, eq_fg_im, eq_fh_re, eq_fh_im], cxx)
+
+        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.))