From a05f9549a1f75ebb076badbf2e36894aa845f70c Mon Sep 17 00:00:00 2001 From: mloubout Date: Mon, 11 May 2026 16:55:36 -0400 Subject: [PATCH 1/2] api: add option to apply interp indices to function coeffs at inject --- devito/operations/interpolators.py | 18 ++++++++++------- devito/types/sparse.py | 6 ++++-- tests/test_interpolation.py | 32 ++++++++++++++++++++++++++++++ 3 files changed, 47 insertions(+), 9 deletions(-) diff --git a/devito/operations/interpolators.py b/devito/operations/interpolators.py index 0ed8756fa6..f5777e4d02 100644 --- a/devito/operations/interpolators.py +++ b/devito/operations/interpolators.py @@ -164,19 +164,21 @@ class Injection(UnevaluatedSparseOperation): __rargs__ = ('field', 'expr', 'implicit_dims') + UnevaluatedSparseOperation.__rargs__ - def __new__(cls, field, expr, implicit_dims, interpolator): + def __new__(cls, field, expr, implicit_dims, interpolator, interp_expr=False): obj = super().__new__(cls, interpolator) # TODO: unused now, but will be necessary to compute the adjoint obj.field = field obj.expr = expr obj.implicit_dims = implicit_dims + obj.interp_expr = interp_expr return obj def operation(self, **kwargs): return self.interpolator._inject(expr=self.expr, field=self.field, - implicit_dims=self.implicit_dims) + implicit_dims=self.implicit_dims, + interp_expr=self.interp_expr) def __repr__(self): return f"Injection({repr(self.expr)} into {repr(self.field)})" @@ -366,7 +368,7 @@ def interpolate(self, expr, increment=False, self_subs=None, implicit_dims=None) @check_radius @check_coords - def inject(self, field, expr, implicit_dims=None): + def inject(self, field, expr, implicit_dims=None, interp_expr=False): """ Generate equations injecting an arbitrary expression into a field. @@ -381,7 +383,7 @@ def inject(self, field, expr, implicit_dims=None): injection expression, but that should be honored when constructing the operator. """ - return Injection(field, expr, implicit_dims, self) + return Injection(field, expr, implicit_dims, self, interp_expr=interp_expr) def _interpolate(self, expr, increment=False, self_subs=None, implicit_dims=None): """ @@ -433,7 +435,7 @@ def _interpolate(self, expr, increment=False, self_subs=None, implicit_dims=None return temps + summands + last - def _inject(self, field, expr, implicit_dims=None): + def _inject(self, field, expr, implicit_dims=None, interp_expr=False): """ Generate equations injecting an arbitrary expression into a field. @@ -479,8 +481,10 @@ def _inject(self, field, expr, implicit_dims=None): self._rdim(subdomain=subdomain)) # List of indirection indices for all adjacent grid points - idx_subs, temps = self._interp_idx(fields, implicit_dims=implicit_dims, - pos_only=variables, subdomain=subdomain) + finterp = fields + as_tuple(variables) if interp_expr else fields + pos_only = () if interp_expr else variables + idx_subs, temps = self._interp_idx(finterp, implicit_dims=implicit_dims, + pos_only=pos_only, subdomain=subdomain) # Substitute coordinate base symbols into the interpolation coefficients eqns = [Inc(_field.xreplace(idx_subs), diff --git a/devito/types/sparse.py b/devito/types/sparse.py index aeb3c6c544..3ca5a79223 100644 --- a/devito/types/sparse.py +++ b/devito/types/sparse.py @@ -1089,7 +1089,8 @@ def interpolate(self, expr, u_t=None, p_t=None, increment=False, implicit_dims=N return super().interpolate(expr, increment=increment, self_subs=subs, implicit_dims=implicit_dims) - def inject(self, field, expr, u_t=None, p_t=None, implicit_dims=None): + def inject(self, field, expr, u_t=None, p_t=None, implicit_dims=None, + interp_expr=False): """ Generate equations injecting an arbitrary expression into a field. @@ -1114,7 +1115,8 @@ def inject(self, field, expr, u_t=None, p_t=None, implicit_dims=None): if p_t is not None: expr = expr.subs({self.time_dim: p_t}) - return super().inject(field, expr, implicit_dims=implicit_dims) + return super().inject(field, expr, implicit_dims=implicit_dims, + interp_expr=interp_expr) @property def forward(self): diff --git a/tests/test_interpolation.py b/tests/test_interpolation.py index 8ee0f6c8dd..37f28f949e 100644 --- a/tests/test_interpolation.py +++ b/tests/test_interpolation.py @@ -11,6 +11,7 @@ SparseTimeFunction, SubDomain, TimeFunction, switchconfig ) from devito.operations.interpolators import LinearInterpolator, SincInterpolator +from devito.symbolics import retrieve_functions from examples.seismic import ( AcquisitionGeometry, Receiver, RickerSource, TimeAxis, demo_model ) @@ -569,6 +570,37 @@ def test_inject_from_field(shape, coords, result, npoints=19): assert np.allclose(a.data[indices], result, rtol=1.e-5) +def test_inject_interp_expr(): + """ + Test that the Function coefficient gets interpolated too. + """ + coords = [(.05, .95), (.45, .45)] + a = unit_box(shape=(11, 11)) + a.data[:] = 0. + p = points(a.grid, ranges=coords, npoints=19) + m = Function(name='m', grid=a.grid) + m.data_with_halo[:] = 1. + + expr = p.inject(a, m, interp_expr=True) + op = Operator(expr) + + op(a=a) + + indices = [slice(4, 6, 1) for _ in coords] + indices[0] = slice(1, -1, 1) + assert np.allclose(a.data[indices], 1, rtol=1.e-5) + + # Extract interp expr to check indices + e_expr = expr.evaluate + funcs = retrieve_functions(e_expr[-1]) + assert m in {f.function for f in funcs} + # All funcs should have the same indices wit radius + # includint the coefficient m + indices = {f.indices for f in funcs} + assert len(indices) == 1 + assert str(indices.pop()) == '(rp_pointsx + posx, rp_pointsy + posy)' + + @pytest.mark.parametrize('shape', [(50, 50, 50)]) def test_position(shape): t0 = 0.0 # Start time From a0b7fc0257739378f09bdbfe0e4eadea54595f1f Mon Sep 17 00:00:00 2001 From: mloubout Date: Tue, 12 May 2026 09:23:23 -0400 Subject: [PATCH 2/2] always interp whole expr for inject as doesn't make sense --- devito/operations/interpolators.py | 29 ++++++++------------------- devito/types/sparse.py | 6 ++---- tests/test_interpolation.py | 32 ------------------------------ 3 files changed, 10 insertions(+), 57 deletions(-) diff --git a/devito/operations/interpolators.py b/devito/operations/interpolators.py index f5777e4d02..defc1b5656 100644 --- a/devito/operations/interpolators.py +++ b/devito/operations/interpolators.py @@ -164,21 +164,19 @@ class Injection(UnevaluatedSparseOperation): __rargs__ = ('field', 'expr', 'implicit_dims') + UnevaluatedSparseOperation.__rargs__ - def __new__(cls, field, expr, implicit_dims, interpolator, interp_expr=False): + def __new__(cls, field, expr, implicit_dims, interpolator): obj = super().__new__(cls, interpolator) # TODO: unused now, but will be necessary to compute the adjoint obj.field = field obj.expr = expr obj.implicit_dims = implicit_dims - obj.interp_expr = interp_expr return obj def operation(self, **kwargs): return self.interpolator._inject(expr=self.expr, field=self.field, - implicit_dims=self.implicit_dims, - interp_expr=self.interp_expr) + implicit_dims=self.implicit_dims) def __repr__(self): return f"Injection({repr(self.expr)} into {repr(self.field)})" @@ -309,7 +307,7 @@ def _positions(self, implicit_dims): return [Eq(v, INT(floor(k)), implicit_dims=implicit_dims) for k, v in self.sfunction._position_map.items()] - def _interp_idx(self, variables, implicit_dims=None, pos_only=(), subdomain=None): + def _interp_idx(self, variables, implicit_dims=None, subdomain=None): """ Generate interpolation indices for the DiscreteFunctions in ``variables``. """ @@ -333,16 +331,6 @@ def _interp_idx(self, variables, implicit_dims=None, pos_only=(), subdomain=None idx_subs = {v: v.subs(subs) for v in variables} - # Position only replacement, not radius dependent. - # E.g src.inject(vp(x)*src) needs to use vp[posx] at all points - # not vp[posx + rx] - idx_subs.update({ - v: v.subs({ - k: p - for (k, p) in zip(mapper, pos, strict=True) - }) for v in pos_only - }) - return idx_subs, temps @check_radius @@ -368,7 +356,7 @@ def interpolate(self, expr, increment=False, self_subs=None, implicit_dims=None) @check_radius @check_coords - def inject(self, field, expr, implicit_dims=None, interp_expr=False): + def inject(self, field, expr, implicit_dims=None): """ Generate equations injecting an arbitrary expression into a field. @@ -383,7 +371,7 @@ def inject(self, field, expr, implicit_dims=None, interp_expr=False): injection expression, but that should be honored when constructing the operator. """ - return Injection(field, expr, implicit_dims, self, interp_expr=interp_expr) + return Injection(field, expr, implicit_dims, self) def _interpolate(self, expr, increment=False, self_subs=None, implicit_dims=None): """ @@ -435,7 +423,7 @@ def _interpolate(self, expr, increment=False, self_subs=None, implicit_dims=None return temps + summands + last - def _inject(self, field, expr, implicit_dims=None, interp_expr=False): + def _inject(self, field, expr, implicit_dims=None): """ Generate equations injecting an arbitrary expression into a field. @@ -481,10 +469,9 @@ def _inject(self, field, expr, implicit_dims=None, interp_expr=False): self._rdim(subdomain=subdomain)) # List of indirection indices for all adjacent grid points - finterp = fields + as_tuple(variables) if interp_expr else fields - pos_only = () if interp_expr else variables + finterp = fields + as_tuple(variables) idx_subs, temps = self._interp_idx(finterp, implicit_dims=implicit_dims, - pos_only=pos_only, subdomain=subdomain) + subdomain=subdomain) # Substitute coordinate base symbols into the interpolation coefficients eqns = [Inc(_field.xreplace(idx_subs), diff --git a/devito/types/sparse.py b/devito/types/sparse.py index 3ca5a79223..aeb3c6c544 100644 --- a/devito/types/sparse.py +++ b/devito/types/sparse.py @@ -1089,8 +1089,7 @@ def interpolate(self, expr, u_t=None, p_t=None, increment=False, implicit_dims=N return super().interpolate(expr, increment=increment, self_subs=subs, implicit_dims=implicit_dims) - def inject(self, field, expr, u_t=None, p_t=None, implicit_dims=None, - interp_expr=False): + def inject(self, field, expr, u_t=None, p_t=None, implicit_dims=None): """ Generate equations injecting an arbitrary expression into a field. @@ -1115,8 +1114,7 @@ def inject(self, field, expr, u_t=None, p_t=None, implicit_dims=None, if p_t is not None: expr = expr.subs({self.time_dim: p_t}) - return super().inject(field, expr, implicit_dims=implicit_dims, - interp_expr=interp_expr) + return super().inject(field, expr, implicit_dims=implicit_dims) @property def forward(self): diff --git a/tests/test_interpolation.py b/tests/test_interpolation.py index 37f28f949e..8ee0f6c8dd 100644 --- a/tests/test_interpolation.py +++ b/tests/test_interpolation.py @@ -11,7 +11,6 @@ SparseTimeFunction, SubDomain, TimeFunction, switchconfig ) from devito.operations.interpolators import LinearInterpolator, SincInterpolator -from devito.symbolics import retrieve_functions from examples.seismic import ( AcquisitionGeometry, Receiver, RickerSource, TimeAxis, demo_model ) @@ -570,37 +569,6 @@ def test_inject_from_field(shape, coords, result, npoints=19): assert np.allclose(a.data[indices], result, rtol=1.e-5) -def test_inject_interp_expr(): - """ - Test that the Function coefficient gets interpolated too. - """ - coords = [(.05, .95), (.45, .45)] - a = unit_box(shape=(11, 11)) - a.data[:] = 0. - p = points(a.grid, ranges=coords, npoints=19) - m = Function(name='m', grid=a.grid) - m.data_with_halo[:] = 1. - - expr = p.inject(a, m, interp_expr=True) - op = Operator(expr) - - op(a=a) - - indices = [slice(4, 6, 1) for _ in coords] - indices[0] = slice(1, -1, 1) - assert np.allclose(a.data[indices], 1, rtol=1.e-5) - - # Extract interp expr to check indices - e_expr = expr.evaluate - funcs = retrieve_functions(e_expr[-1]) - assert m in {f.function for f in funcs} - # All funcs should have the same indices wit radius - # includint the coefficient m - indices = {f.indices for f in funcs} - assert len(indices) == 1 - assert str(indices.pop()) == '(rp_pointsx + posx, rp_pointsy + posy)' - - @pytest.mark.parametrize('shape', [(50, 50, 50)]) def test_position(shape): t0 = 0.0 # Start time