Skip to content

Commit

Permalink
api: dix sympy assumptions for complex valued objects
Browse files Browse the repository at this point in the history
  • Loading branch information
mloubout committed Mar 10, 2025
1 parent 2729050 commit f49d03a
Show file tree
Hide file tree
Showing 6 changed files with 198 additions and 15 deletions.
2 changes: 1 addition & 1 deletion devito/mpi/routines.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down
6 changes: 5 additions & 1 deletion devito/passes/iet/misc.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)))'),), {}
Expand Down
17 changes: 10 additions & 7 deletions devito/types/basic.py
Original file line number Diff line number Diff line change
Expand Up @@ -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')
Expand Down Expand Up @@ -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 ()
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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):
"""
Expand Down
158 changes: 152 additions & 6 deletions examples/seismic/tutorials/17_fourier_mode.ipynb
Original file line number Diff line number Diff line change
Expand Up @@ -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",
Expand All @@ -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"
}
Expand All @@ -222,7 +368,7 @@
},
{
"cell_type": "code",
"execution_count": 10,
"execution_count": 11,
"metadata": {},
"outputs": [
{
Expand Down Expand Up @@ -250,7 +396,7 @@
},
{
"cell_type": "code",
"execution_count": 11,
"execution_count": 12,
"metadata": {},
"outputs": [],
"source": [
Expand Down
17 changes: 17 additions & 0 deletions tests/test_builtins.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
13 changes: 13 additions & 0 deletions tests/test_symbolics.py
Original file line number Diff line number Diff line change
Expand Up @@ -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')

Expand Down

0 comments on commit f49d03a

Please sign in to comment.