dsl: Update Re and Im to Real and Imag, add Conj

Author: EdCaunt

devito/finite_differences/differentiable.py

@@ -22,7 +22,7 @@
 from devito.types.basic import AbstractFunction
 
 __all__ = ['Differentiable', 'DiffDerivative', 'IndexDerivative', 'EvalDerivative',
-           'Weights', 'Re', 'Im']
+           'Weights', 'Real', 'Imag', 'Conj']
 
 
 class Differentiable(sympy.Expr, Evaluable):
@@ -646,7 +646,7 @@ def __str__(self):
 
 
 class ComplexPart(Differentiable, sympy.core.function.Application):
-    """Abstract class for `Re` or `Im` of an expression"""
+    """Abstract class for `Real`, `Imag`, or `Conj` of an expression"""
 
     def __new__(cls, *args, **kwargs):
         if len(args) != 1:
@@ -668,16 +668,21 @@ def __str__(self):
     __repr__ = __str__
 
 
-class Re(ComplexPart):
+class Real(ComplexPart):
     """Get the real part of an expression"""
     pass
 
 
-class Im(ComplexPart):
+class Imag(ComplexPart):
     """Get the imaginary part of an expression"""
     pass
 
 
+class Conj(ComplexPart):
+    """Get the complex conjugate of an expression"""
+    pass
+
+
 class IndexSum(sympy.Expr, Evaluable):
 
     """

devito/passes/iet/languages/C.py

@@ -56,10 +56,14 @@ class CPrinter(BasePrinter, C99CodePrinter):
     def _print_ImaginaryUnit(self, expr):
         return '_Complex_I'
 
-    def _print_Re(self, expr):
+    def _print_Real(self, expr):
         return (f'{self.func_prefix(expr)}real{self.func_literal(expr).lower()}'
                 f'({self._print(expr.args[0])})')
 
-    def _print_Im(self, expr):
+    def _print_Imag(self, expr):
         return (f'{self.func_prefix(expr)}imag{self.func_literal(expr).lower()}'
                 f'({self._print(expr.args[0])})')
+
+    def _print_Conj(self, expr):
+        return (f'conj{self.func_literal(expr).lower()}'
+                f'({self._print(expr.args[0])})')

devito/passes/iet/languages/CXX.py

@@ -104,12 +104,15 @@ class CXXPrinter(BasePrinter, CXX11CodePrinter):
     def _print_ImaginaryUnit(self, expr):
         return f'1i{self.prec_literal(expr).lower()}'
 
-    def _print_Re(self, expr):
+    def _print_Real(self, expr):
         return f'{self._ns}real({self._print(expr.args[0])})'
 
-    def _print_Im(self, expr):
+    def _print_Imag(self, expr):
         return f'{self._ns}imag({self._print(expr.args[0])})'
 
+    def _print_Conj(self, expr):
+        return f'{self._ns}conj({self._print(expr.args[0])})'
+
     def _print_Cast(self, expr):
         # The CXX recommended way to cast is to use static_cast
         tstr = self._print(expr._C_ctype)

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, Re, Im, switchconfig)
+                    Min, Max, Real, Imag, Conj, 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
@@ -837,15 +837,15 @@ def run_operator(self, eqs, language):
     def test_devito_print(self):
         f, _, _ = self.setup_basic(np.complex64)
 
-        assert str(Re(f)) == 'Re(f(x))'
-        assert str(Im(f)) == 'Im(f(x))'
+        assert str(Real(f)) == 'Real(f(x))'
+        assert str(Imag(f)) == 'Imag(f(x))'
 
     @pytest.mark.parametrize('language', ['C', 'CXX', 'CXXopenmp'])
     def test_printing(self, language):
         f, f_real, f_imag = self.setup_basic(np.complex64)
 
-        eq_re = Eq(f_real, Re(f))
-        eq_im = Eq(f_imag, Im(f))
+        eq_re = Eq(f_real, Real(f))
+        eq_im = Eq(f_imag, Imag(f))
 
         with switchconfig(language=language):
             op = Operator([eq_re, eq_im])
@@ -863,8 +863,8 @@ def test_printing(self, language):
     def test_trivial(self, language, 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.))
+        eq_re = Eq(f_real, Real(f+1.))
+        eq_im = Eq(f_imag, Imag(f+1.))
 
         self.run_operator([eq_re, eq_im], language)
 
@@ -878,8 +878,8 @@ def test_trivial(self, language, dtype):
     def test_trivial_imag(self, language, 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))
+        eq_re = Eq(f_real, Real(f+1j))
+        eq_im = Eq(f_imag, Imag(f+1j))
 
         self.run_operator([eq_re, eq_im], language)
 
@@ -892,8 +892,8 @@ def test_trivial_imag(self, language, dtype):
     def test_deriv(self, language):
         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))
+        eq_re = Eq(f_real, Real(f.dx))
+        eq_im = Eq(f_imag, Imag(f.dx))
 
         self.run_operator([eq_re, eq_im], language)
 
@@ -904,8 +904,8 @@ def test_deriv(self, language):
     def test_outer_deriv(self, language):
         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)
+        eq_re = Eq(f_real, Real(f).dx)
+        eq_im = Eq(f_imag, Imag(f).dx)
 
         self.run_operator([eq_re, eq_im], language)
 
@@ -928,10 +928,10 @@ def test_mul(self, language):
         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))
+        eq_fg_re = Eq(fg_re, Real(f*g))
+        eq_fg_im = Eq(fg_im, Imag(f*g))
+        eq_fh_re = Eq(fh_re, Real(f*h))
+        eq_fh_im = Eq(fh_im, Imag(f*h))
 
         self.run_operator([eq_fg_re, eq_fg_im, eq_fh_re, eq_fh_im], language)
 
@@ -940,3 +940,15 @@ def test_mul(self, language):
 
         assert np.all(np.isclose(fh_re.data, -2.))
         assert np.all(np.isclose(fh_im.data, 2.))
+
+    @pytest.mark.parametrize('language', ['C', 'CXX', 'CXXopenmp'])
+    def test_conj(self, language):
+        grid = Grid(shape=(5,))
+        f = Function(name='f', grid=grid, dtype=np.complex64)
+        g = Function(name='g', grid=grid, dtype=np.complex64)
+
+        f.data[:] = np.arange(5) + 1j*np.arange(5)[::-1]
+
+        self.run_operator([Eq(g, Conj(f))], language)
+
+        assert np.all(np.isclose(g.data, np.conj(f.data)))