tests: Update Re and Im tests to be more comprehensive

Author: EdCaunt

devito/ir/equations/equation.py

@@ -49,11 +49,9 @@ def directions(self):
     @property
     def dtype(self):
         try:
-            rhs_dtype = self.rhs.dtype
+            return infer_dtype({self.lhs.dtype, self.rhs.dtype} - {None})
         except AttributeError:
-            rhs_dtype = None
-
-        return infer_dtype({self.lhs.dtype, rhs_dtype} - {None})
+            return self.lhs.dtype
 
     @property
     def state(self):

devito/passes/iet/languages/C.py

@@ -69,11 +69,9 @@ def _print_ListInitializer(self, expr):
             return li
 
     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

@@ -106,7 +106,6 @@ 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])})'

tests/test_symbolics.py

@@ -915,7 +915,6 @@ def test_print_div():
 
 
 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)
@@ -925,87 +924,90 @@ def setup_basic(self, dtype):
         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:
+    def run_operator(self, eqs, language):
+        with switchconfig(language=language):
             Operator(eqs)()
 
-    @pytest.mark.parametrize('cxx', [False, True])
-    def test_printing(self, cxx):
+    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))'
+
+    @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))
 
-        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)
+        with switchconfig(language=language):
+            op = Operator([eq_re, eq_im])
+
+        if language in ('CXX', 'CXXopenmp'):
+            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('language', ['C', 'CXX', 'CXXopenmp'])
     @pytest.mark.parametrize('dtype', [np.complex64, np.complex128])
-    def test_trivial(self, cxx, dtype):
+    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.))
 
-        self.run_operator([eq_re, eq_im], cxx)
+        self.run_operator([eq_re, eq_im], language)
 
         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('language', ['C', 'CXX', 'CXXopenmp'])
     @pytest.mark.parametrize('dtype', [np.complex64, np.complex128])
-    def test_trivial_imag(self, cxx, 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))
 
-        self.run_operator([eq_re, eq_im], cxx)
+        self.run_operator([eq_re, eq_im], language)
 
         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):
+    @pytest.mark.parametrize('language', ['C', 'CXX', 'CXXopenmp'])
+    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))
 
-        self.run_operator([eq_re, eq_im], cxx)
+        self.run_operator([eq_re, eq_im], language)
 
         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):
+    @pytest.mark.parametrize('language', ['C', 'CXX', 'CXXopenmp'])
+    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)
 
-        self.run_operator([eq_re, eq_im], cxx)
+        self.run_operator([eq_re, eq_im], language)
 
         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):
+    @pytest.mark.parametrize('language', ['C', 'CXX', 'CXXopenmp'])
+    def test_mul(self, language):
         grid = Grid(shape=(5,))
 
         f = Function(name='f', grid=grid, dtype=np.complex64)
@@ -1025,7 +1027,7 @@ def test_mul(self, cxx):
         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)
+        self.run_operator([eq_fg_re, eq_fg_im, eq_fh_re, eq_fh_im], language)
 
         assert np.all(np.isclose(fg_re.data, 2.))
         assert np.all(np.isclose(fg_im.data, 2.))