api: dix sympy assumptions for complex valued objects

Author: mloubout

devito/mpi/routines.py

@@ -306,7 +306,7 @@ def _make_bundles(self, hs):
             except ValueError:
                 for i in candidates:
                     name = "bag_%s" % i.name
-                    bag = Bag(name=name, components=i)
+                    bag = Bag(name=name, components=(i,))
                     halo_scheme = halo_scheme.add(bag, hse)
 
         hs = hs._rebuild(halo_scheme=halo_scheme)

devito/passes/iet/misc.py

@@ -235,7 +235,11 @@ def _(expr, lang):
 
 @_lower_macro_math.register(SafeInv)
 def _(expr, lang):
-    eps = np.finfo(expr.base.dtype).resolution**2
+    try:
+        eps = np.finfo(expr.base.dtype).resolution**2
+    except ValueError:
+        print(f"Warning: dtype not recognized in SafeInv for {expr.base}")
+        eps = np.finfo(np.float32).resolution**2
     b = Cast('b', dtype=np.float32)
     return (('SAFEINV(a, b)',
              f'(((a) < {eps}F || ({b}) < {eps}F) ? (0.0F) : ((1.0F) / (a)))'),), {}

devito/types/basic.py

@@ -349,8 +349,6 @@ class AbstractSymbol(sympy.Symbol, Basic, Pickable, Evaluable):
     is_Symbol = True
 
     # SymPy default assumptions
-    is_real = True
-    is_imaginary = False
     is_commutative = True
 
     __rkwargs__ = ('name', 'dtype', 'is_const')
@@ -411,6 +409,12 @@ def _hashable_content(self):
     def dtype(self):
         return self._dtype
 
+    def _eval_is_real(self):
+        return not self.is_imaginary
+
+    def _eval_is_imaginary(self):
+        return np.iscomplexobj(self.dtype(0))
+
     @property
     def indices(self):
         return ()
@@ -859,7 +863,6 @@ class AbstractFunction(sympy.Function, Basic, Pickable, Evaluable):
     is_AbstractFunction = True
 
     # SymPy default assumptions
-    is_imaginary = False
     is_commutative = True
 
     # Devito default assumptions
@@ -888,6 +891,8 @@ def __new__(cls, *args, **kwargs):
         # Extract the `indices`, as perhaps they're explicitly provided
         dimensions, indices = cls.__indices_setup__(*args, **kwargs)
 
+        # Sympy assumptions
+
         # If it's an alias or simply has a different name, ignore `function`.
         # These cases imply the construction of a new AbstractFunction off
         # an existing one! This relieves the pressure on the caller by not
@@ -955,6 +960,8 @@ def _sympystr(self, printer, **kwargs):
         return str(self)
 
     _latex = _sympystr
+    _eval_is_real = AbstractSymbol._eval_is_real
+    _eval_is_imaginary = AbstractSymbol._eval_is_imaginary
 
     def _pretty(self, printer, **kwargs):
         return printer._print_Function(self, func_name=self.name)
@@ -1315,10 +1322,6 @@ def is_const(self):
     def is_transient(self):
         return self._is_transient
 
-    @property
-    def is_real(self):
-        return not np.iscomplex(self.dtype(0))
-
     @property
     def is_persistent(self):
         """

examples/seismic/tutorials/17_fourier_mode.ipynb

@@ -191,6 +191,152 @@
    "cell_type": "code",
    "execution_count": 9,
    "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "#define _POSIX_C_SOURCE 200809L\n",
+      "#define START(S) struct timeval start_ ## S , end_ ## S ; gettimeofday(&start_ ## S , NULL);\n",
+      "#define STOP(S,T) gettimeofday(&end_ ## S, NULL); T->S += (double)(end_ ## S .tv_sec-start_ ## S.tv_sec)+(double)(end_ ## S .tv_usec-start_ ## S .tv_usec)/1000000;\n",
+      "#define MAX(a,b) (((a) > (b)) ? (a) : (b))\n",
+      "\n",
+      "#include \"stdlib.h\"\n",
+      "#include \"math.h\"\n",
+      "#include \"sys/time.h\"\n",
+      "#include \"omp.h\"\n",
+      "#include \"complex.h\"\n",
+      "\n",
+      "struct dataobj\n",
+      "{\n",
+      "  void *restrict data;\n",
+      "  unsigned long * size;\n",
+      "  unsigned long * npsize;\n",
+      "  unsigned long * dsize;\n",
+      "  int * hsize;\n",
+      "  int * hofs;\n",
+      "  int * oofs;\n",
+      "  void * dmap;\n",
+      "} ;\n",
+      "\n",
+      "struct profiler\n",
+      "{\n",
+      "  double section0;\n",
+      "  double section1;\n",
+      "  double section2;\n",
+      "} ;\n",
+      "\n",
+      "\n",
+      "int Kernel(struct dataobj *restrict damp_vec, struct dataobj *restrict freq_modes_vec, struct dataobj *restrict rec_vec, struct dataobj *restrict rec_coords_vec, struct dataobj *restrict src_vec, struct dataobj *restrict src_coords_vec, struct dataobj *restrict u_vec, struct dataobj *restrict vp_vec, const int x_M, const int x_m, const int y_M, const int y_m, const float dt, const float o_x, const float o_y, const int p_rec_M, const int p_rec_m, const int p_src_M, const int p_src_m, const int time_M, const int time_m, const int nthreads, const int nthreads_nonaffine, struct profiler * timers)\n",
+      "{\n",
+      "  float (*restrict damp)[damp_vec->size[1]] __attribute__ ((aligned (64))) = (float (*)[damp_vec->size[1]]) damp_vec->data;\n",
+      "  float _Complex (*restrict freq_modes)[freq_modes_vec->size[1]] __attribute__ ((aligned (64))) = (float _Complex (*)[freq_modes_vec->size[1]]) freq_modes_vec->data;\n",
+      "  float (*restrict rec)[rec_vec->size[1]] __attribute__ ((aligned (64))) = (float (*)[rec_vec->size[1]]) rec_vec->data;\n",
+      "  float (*restrict rec_coords)[rec_coords_vec->size[1]] __attribute__ ((aligned (64))) = (float (*)[rec_coords_vec->size[1]]) rec_coords_vec->data;\n",
+      "  float (*restrict src)[src_vec->size[1]] __attribute__ ((aligned (64))) = (float (*)[src_vec->size[1]]) src_vec->data;\n",
+      "  float (*restrict src_coords)[src_coords_vec->size[1]] __attribute__ ((aligned (64))) = (float (*)[src_coords_vec->size[1]]) src_coords_vec->data;\n",
+      "  float (*restrict u)[u_vec->size[1]][u_vec->size[2]] __attribute__ ((aligned (64))) = (float (*)[u_vec->size[1]][u_vec->size[2]]) u_vec->data;\n",
+      "  float (*restrict vp)[vp_vec->size[1]] __attribute__ ((aligned (64))) = (float (*)[vp_vec->size[1]]) vp_vec->data;\n",
+      "\n",
+      "  float _Complex r2 = 1.0F/(dt*dt);\n",
+      "  float _Complex r3 = 1.0F/dt;\n",
+      "\n",
+      "  for (int time = time_m, t0 = (time)%(3), t1 = (time + 2)%(3), t2 = (time + 1)%(3); time <= time_M; time += 1, t0 = (time)%(3), t1 = (time + 2)%(3), t2 = (time + 1)%(3))\n",
+      "  {\n",
+      "    float _Complex r1 = cexpf(6.28318530717959e-2F*time*_Complex_I*dt);\n",
+      "    START(section0)\n",
+      "    #pragma omp parallel num_threads(nthreads)\n",
+      "    {\n",
+      "      #pragma omp for schedule(dynamic,1)\n",
+      "      for (int x = x_m; x <= x_M; x += 1)\n",
+      "      {\n",
+      "        #pragma omp simd aligned(damp,freq_modes,u,vp:16)\n",
+      "        for (int y = y_m; y <= y_M; y += 1)\n",
+      "        {\n",
+      "          float _Complex r4 = 1.0F/(vp[x + 2][y + 2]*vp[x + 2][y + 2]);\n",
+      "          u[t2][x + 2][y + 2] = (-r4*(-2.0F*r2*u[t0][x + 2][y + 2] + r2*u[t1][x + 2][y + 2]) + r3*damp[x + 2][y + 2]*u[t0][x + 2][y + 2] + 1.0e-2F*(u[t0][x + 1][y + 2] + u[t0][x + 2][y + 1] + u[t0][x + 2][y + 3] + u[t0][x + 3][y + 2]) - 3.99999991e-2F*u[t0][x + 2][y + 2])/(r4*r2 + r3*damp[x + 2][y + 2]);\n",
+      "          freq_modes[x][y] += r1*u[t0][x + 2][y + 2];\n",
+      "        }\n",
+      "      }\n",
+      "    }\n",
+      "    STOP(section0,timers)\n",
+      "\n",
+      "    START(section1)\n",
+      "    #pragma omp parallel num_threads(nthreads_nonaffine)\n",
+      "    {\n",
+      "      int chunk_size = (int)(MAX(1, (int)((1.0/3.0)*(p_src_M - p_src_m + 1)/nthreads_nonaffine)));\n",
+      "      #pragma omp for schedule(dynamic,chunk_size)\n",
+      "      for (int p_src = p_src_m; p_src <= p_src_M; p_src += 1)\n",
+      "      {\n",
+      "        for (int rsrcx = 0; rsrcx <= 1; rsrcx += 1)\n",
+      "        {\n",
+      "          for (int rsrcy = 0; rsrcy <= 1; rsrcy += 1)\n",
+      "          {\n",
+      "            int posx = (int)(floorf(1.0e-1*(-o_x + src_coords[p_src][0])));\n",
+      "            int posy = (int)(floorf(1.0e-1*(-o_y + src_coords[p_src][1])));\n",
+      "            float px = 1.0e-1F*(-o_x + src_coords[p_src][0]) - floorf(1.0e-1F*(-o_x + src_coords[p_src][0]));\n",
+      "            float py = 1.0e-1F*(-o_y + src_coords[p_src][1]) - floorf(1.0e-1F*(-o_y + src_coords[p_src][1]));\n",
+      "            if (rsrcx + posx >= x_m - 1 && rsrcy + posy >= y_m - 1 && rsrcx + posx <= x_M + 1 && rsrcy + posy <= y_M + 1)\n",
+      "            {\n",
+      "              float r0 = 3.06250F*(vp[posx + 2][posy + 2]*vp[posx + 2][posy + 2])*(rsrcx*px + (1 - rsrcx)*(1 - px))*(rsrcy*py + (1 - rsrcy)*(1 - py))*src[time][p_src];\n",
+      "              #pragma omp atomic update\n",
+      "              u[t2][rsrcx + posx + 2][rsrcy + posy + 2] += r0;\n",
+      "            }\n",
+      "          }\n",
+      "        }\n",
+      "      }\n",
+      "    }\n",
+      "    STOP(section1,timers)\n",
+      "\n",
+      "    START(section2)\n",
+      "    #pragma omp parallel num_threads(nthreads_nonaffine)\n",
+      "    {\n",
+      "      int chunk_size = (int)(MAX(1, (int)((1.0/3.0)*(p_rec_M - p_rec_m + 1)/nthreads_nonaffine)));\n",
+      "      #pragma omp for schedule(dynamic,chunk_size)\n",
+      "      for (int p_rec = p_rec_m; p_rec <= p_rec_M; p_rec += 1)\n",
+      "      {\n",
+      "        float r7 = 1.0e-1F*(-o_x + rec_coords[p_rec][0]);\n",
+      "        float r5 = floorf(r7);\n",
+      "        int posx = (int)r5;\n",
+      "        float r8 = 1.0e-1F*(-o_y + rec_coords[p_rec][1]);\n",
+      "        float r6 = floorf(r8);\n",
+      "        int posy = (int)r6;\n",
+      "        float px = -r5 + r7;\n",
+      "        float py = -r6 + r8;\n",
+      "        float sum = 0.0F;\n",
+      "\n",
+      "        for (int rrecx = 0; rrecx <= 1; rrecx += 1)\n",
+      "        {\n",
+      "          for (int rrecy = 0; rrecy <= 1; rrecy += 1)\n",
+      "          {\n",
+      "            if (rrecx + posx >= x_m - 1 && rrecy + posy >= y_m - 1 && rrecx + posx <= x_M + 1 && rrecy + posy <= y_M + 1)\n",
+      "            {\n",
+      "              sum += (rrecx*px + (1 - rrecx)*(1 - px))*(rrecy*py + (1 - rrecy)*(1 - py))*u[t2][rrecx + posx + 2][rrecy + posy + 2];\n",
+      "            }\n",
+      "          }\n",
+      "        }\n",
+      "\n",
+      "        rec[time][p_rec] = sum;\n",
+      "      }\n",
+      "    }\n",
+      "    STOP(section2,timers)\n",
+      "  }\n",
+      "\n",
+      "  return 0;\n",
+      "}\n",
+      "\n"
+     ]
+    }
+   ],
+   "source": [
+    "#NBVAL_IGNORE_OUTPUT\n",
+    "print(op)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 10,
+   "metadata": {},
    "outputs": [
     {
      "name": "stderr",
@@ -203,14 +349,14 @@
      "data": {
       "text/plain": [
        "PerformanceSummary([(PerfKey(name='section0', rank=None),\n",
-       "                     PerfEntry(time=0.01551399999999999, gflopss=0.0, gpointss=0.0, oi=0.0, ops=0, itershapes=[])),\n",
+       "                     PerfEntry(time=0.016613000000000006, gflopss=0.0, gpointss=0.0, oi=0.0, ops=0, itershapes=[])),\n",
        "                    (PerfKey(name='section1', rank=None),\n",
-       "                     PerfEntry(time=0.0023819999999999913, gflopss=0.0, gpointss=0.0, oi=0.0, ops=0, itershapes=[])),\n",
+       "                     PerfEntry(time=0.002539999999999992, gflopss=0.0, gpointss=0.0, oi=0.0, ops=0, itershapes=[])),\n",
        "                    (PerfKey(name='section2', rank=None),\n",
-       "                     PerfEntry(time=0.002333999999999994, gflopss=0.0, gpointss=0.0, oi=0.0, ops=0, itershapes=[]))])"
+       "                     PerfEntry(time=0.0025529999999999923, gflopss=0.0, gpointss=0.0, oi=0.0, ops=0, itershapes=[]))])"
       ]
      },
-     "execution_count": 9,
+     "execution_count": 10,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -222,7 +368,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 10,
+   "execution_count": 11,
    "metadata": {},
    "outputs": [
     {
@@ -250,7 +396,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 11,
+   "execution_count": 12,
    "metadata": {},
    "outputs": [],
    "source": [

tests/test_builtins.py

@@ -374,6 +374,23 @@ def test_inner_sparse(self):
         term2 = np.inner(rec0.data.reshape(-1), rec1.data.reshape(-1))
         assert np.isclose(term1/term2 - 1, 0.0, rtol=0.0, atol=1e-5)
 
+    @pytest.mark.parametrize('dtype', [np.float32, np.complex64])
+    def test_norm_dense(self, dtype):
+        """
+        Test that norm produces the correct result against NumPy
+        """
+        grid = Grid((101, 101), extent=(1000., 1000.))
+
+        f = Function(name='f', grid=grid, dtype=dtype)
+
+        f.data[:] = 1 + np.random.randn(*f.shape).astype(grid.dtype)
+        if np.iscomplexobj(f.data):
+            f.data[:] += 1j*np.random.randn(*f.shape).astype(grid.dtype)
+        term1 = np.linalg.norm(f.data)
+        term2 = norm(f)
+        assert np.isreal(term2)
+        assert np.isclose(term1/term2 - 1, 0.0, rtol=0.0, atol=1e-5)
+
     def test_norm_sparse(self):
         """
         Test that norm produces the correct result against NumPy

tests/test_symbolics.py

@@ -112,6 +112,19 @@ def test_modified_sympy_assumptions():
     assert s2 == s1
 
 
+def test_real():
+    for dtype in [np.float32, np.complex64]:
+        c = Constant(name='c', dtype=dtype)
+        assert c.is_real is not np.iscomplexobj(dtype(0))
+        assert c.is_imaginary is np.iscomplexobj(dtype(0))
+        f = Function(name='f', dtype=dtype, grid=Grid((11,)))
+        assert f.is_real is not np.iscomplexobj(dtype(0))
+        assert f.is_imaginary is np.iscomplexobj(dtype(0))
+        s = dSymbol(name='s', dtype=dtype)
+        assert s.is_real is not np.iscomplexobj(dtype(0))
+        assert s.is_imaginary is np.iscomplexobj(dtype(0))
+
+
 def test_constant():
     c = Constant(name='c')