diff --git a/devito/operations/interpolators.py b/devito/operations/interpolators.py index 52501c54c4..00b7dd6c34 100644 --- a/devito/operations/interpolators.py +++ b/devito/operations/interpolators.py @@ -486,8 +486,7 @@ def interpolation_coeffs(self): sf = SubFunction(name="wsinc%s" % r.name, dtype=self.sfunction.dtype, shape=shape, dimensions=dimensions, space_order=0, alias=self.sfunction.alias, - distributor=self.sfunction._distributor, - parent=self.sfunction) + parent=None) coeffs[r] = sf return coeffs diff --git a/devito/operator/operator.py b/devito/operator/operator.py index 3b3290f621..c9e84ac1c0 100644 --- a/devito/operator/operator.py +++ b/devito/operator/operator.py @@ -27,7 +27,7 @@ from devito.tools import (DAG, OrderedSet, Signer, ReducerMap, as_tuple, flatten, filter_sorted, frozendict, is_integer, split, timed_pass, timed_region, contains_val) -from devito.types import Grid, Evaluable, SubFunction +from devito.types import Grid, Evaluable __all__ = ['Operator'] @@ -660,12 +660,7 @@ def _postprocess_errors(self, retval): def _postprocess_arguments(self, args, **kwargs): """Process runtime arguments upon returning from ``.apply()``.""" for p in self.parameters: - try: - subfuncs = (args[getattr(p, s).name] for s in p._sub_functions) - p._arg_apply(args[p.name], *subfuncs, alias=kwargs.get(p.name)) - except AttributeError: - if not (isinstance(p, SubFunction) and p.parent in self.parameters): - p._arg_apply(args[p.name], alias=kwargs.get(p.name)) + p._arg_apply(args[p.name], alias=kwargs.get(p.name)) @cached_property def _known_arguments(self): diff --git a/devito/types/sparse.py b/devito/types/sparse.py index 094ec7dd98..186ec03935 100644 --- a/devito/types/sparse.py +++ b/devito/types/sparse.py @@ -1,9 +1,4 @@ from collections import OrderedDict -try: - from collections.abc import Iterable -except ImportError: - # Before python 3.10 - from collections import Iterable from itertools import product import sympy @@ -34,6 +29,14 @@ _default_radius = {'linear': 1, 'sinc': 4} +class SparseSubFunction(SubFunction): + + def _arg_apply(self, dataobj, **kwargs): + if self.parent is not None: + return self.parent._dist_subfunc_gather(dataobj, self) + return super()._arg_apply(dataobj, **kwargs) + + class AbstractSparseFunction(DiscreteFunction): """ @@ -49,7 +52,8 @@ class AbstractSparseFunction(DiscreteFunction): _sub_functions = () """SubFunctions encapsulated within this AbstractSparseFunction.""" - __rkwargs__ = DiscreteFunction.__rkwargs__ + ('npoint_global', 'space_order') + __rkwargs__ = (DiscreteFunction.__rkwargs__ + + ('dimensions', 'npoint_global', 'space_order')) def __init_finalize__(self, *args, **kwargs): super().__init_finalize__(*args, **kwargs) @@ -61,10 +65,22 @@ def __init_finalize__(self, *args, **kwargs): @classmethod def __indices_setup__(cls, *args, **kwargs): - dimensions = as_tuple(kwargs.get('dimensions')) - if not dimensions: - dimensions = (Dimension(name='p_%s' % kwargs["name"]),) + # Need this not to break MatrixSparseFunction + try: + _sub_funcs = tuple(cls._sub_functions) + except TypeError: + _sub_funcs = () + # If a subfunction provided use the sparse dimension + for f in _sub_funcs: + try: + sparse_dim = kwargs[f].indices[0] + break + except (KeyError, AttributeError): + continue + else: + sparse_dim = Dimension(name='p_%s' % kwargs["name"]) + dimensions = as_tuple(kwargs.get('dimensions', sparse_dim)) if args: return tuple(dimensions), tuple(args) else: @@ -120,6 +136,22 @@ def __subfunc_setup__(self, key, suffix, dtype=None): if key is not None and not isinstance(key, SubFunction): key = np.array(key) + # Check if already a SubFunction + if isinstance(key, SubFunction): + d = self.indices[self._sparse_position] + if d in key.indices: + # Can use as is, dimension already matches + if self.alias: + return key._rebuild(alias=self.alias, name=name) + else: + return key + else: + # Need to rebuild so the dimensions match the parent + # SparseFunction, for example we end up here via `.subs(d, new_d)` + indices = (d, *key.indices[1:]) + return key._rebuild(*indices, name=name, alias=self.alias) + + # Given an array or nothing, create dimension and SubFunction if key is not None: dimensions = (self._sparse_dim, Dimension(name='d')) if key.ndim > 2: @@ -132,22 +164,13 @@ def __subfunc_setup__(self, key, suffix, dtype=None): dimensions = (self._sparse_dim, Dimension(name='d')) shape = (self.npoint, self.grid.dim) - # Check if already a SubFunction - if isinstance(key, SubFunction): - # Need to rebuild so the dimensions match the parent SparseFunction - indices = (self.indices[self._sparse_position], *key.indices[1:]) - return key._rebuild(*indices, name=name, shape=shape, - alias=self.alias, halo=None) - elif key is not None and not isinstance(key, Iterable): - raise ValueError("`%s` must be either SubFunction " - "or iterable (e.g., list, np.ndarray)" % key) - if key is None: # Fallback to default behaviour dtype = dtype or self.dtype else: - if (shape != key.shape and key.shape != (shape[1],)) and \ - self._distributor.nprocs == 1: + if shape != key.shape and \ + key.shape != (shape[1],) and \ + self._distributor.nprocs == 1: raise ValueError("Incompatible shape for %s, `%s`; expected `%s`" % (suffix, key.shape[:2], shape)) @@ -157,7 +180,7 @@ def __subfunc_setup__(self, key, suffix, dtype=None): else: dtype = dtype or self.dtype - sf = SubFunction( + sf = SparseSubFunction( name=name, dtype=dtype, dimensions=dimensions, shape=shape, space_order=0, initializer=key, alias=self.alias, distributor=self._distributor, parent=self @@ -584,12 +607,6 @@ def _dist_scatter(self, data=None): mapper.update(self._dist_subfunc_scatter(getattr(self, i))) return mapper - def _dist_gather(self, data, *subfunc): - self._dist_data_gather(data) - for (sg, s) in zip(subfunc, self._sub_functions): - if getattr(self, s) is not None: - self._dist_subfunc_gather(sg, getattr(self, s)) - def _eval_at(self, func): return self @@ -637,11 +654,11 @@ def _arg_values(self, **kwargs): return values - def _arg_apply(self, dataobj, *subfuncs, alias=None): + def _arg_apply(self, dataobj, alias=None): key = alias if alias is not None else self if isinstance(key, AbstractSparseFunction): # Gather into `self.data` - key._dist_gather(dataobj, *subfuncs) + key._dist_data_gather(dataobj) elif self._distributor.nprocs > 1: raise NotImplementedError("Don't know how to gather data from an " "object of type `%s`" % type(key)) @@ -698,7 +715,7 @@ def __indices_setup__(cls, *args, **kwargs): dimensions = as_tuple(kwargs.get('dimensions')) if not dimensions: dimensions = (kwargs['grid'].time_dim, - Dimension(name='p_%s' % kwargs["name"]),) + *super().__indices_setup__(*args, **kwargs)[0]) if args: return tuple(dimensions), tuple(args) diff --git a/examples/seismic/elastic/operators.py b/examples/seismic/elastic/operators.py index dd9d793dfe..1576ed133a 100644 --- a/examples/seismic/elastic/operators.py +++ b/examples/seismic/elastic/operators.py @@ -1,6 +1,5 @@ from devito import Eq, Operator, VectorTimeFunction, TensorTimeFunction from devito import div, grad, diag, solve -from examples.seismic import PointSource, Receiver def src_rec(v, tau, model, geometry): @@ -9,12 +8,9 @@ def src_rec(v, tau, model, geometry): """ s = model.grid.time_dim.spacing # Source symbol with input wavelet - src = PointSource(name='src', grid=model.grid, time_range=geometry.time_axis, - npoint=geometry.nsrc) - rec1 = Receiver(name='rec1', grid=model.grid, time_range=geometry.time_axis, - npoint=geometry.nrec) - rec2 = Receiver(name='rec2', grid=model.grid, time_range=geometry.time_axis, - npoint=geometry.nrec) + src = geometry.src + rec1 = geometry.new_rec(name="rec1") + rec2 = geometry.new_rec(name="rec2") # The source injection term src_expr = src.inject(tau.forward.diagonal(), expr=src * s) diff --git a/examples/seismic/inversion/fwi.py b/examples/seismic/inversion/fwi.py index c1893f379d..831b981763 100644 --- a/examples/seismic/inversion/fwi.py +++ b/examples/seismic/inversion/fwi.py @@ -2,7 +2,7 @@ from devito import configuration, Function, norm, mmax, mmin -from examples.seismic import demo_model, AcquisitionGeometry, Receiver +from examples.seismic import demo_model, AcquisitionGeometry from examples.seismic.acoustic import AcousticWaveSolver from inversion_utils import compute_residual, update_with_box @@ -55,33 +55,30 @@ source_locations[:, 0] = 20. source_locations[:, 1] = np.linspace(0., 1000, num=nshots) +# Create placeholders for the data residual and data +residual = geometry.new_rec(name='residual') +d_obs = geometry.new_rec(name='d_obs', coordinates=residual.coordinates) +d_syn = geometry.new_rec(name='d_syn', coordinates=residual.coordinates) + +src = solver.geometry.src + def fwi_gradient(vp_in): # Create symbols to hold the gradient grad = Function(name="grad", grid=model.grid) objective = 0. for i in range(nshots): - # Create placeholders for the data residual and data - residual = Receiver(name='residual', grid=model.grid, - time_range=geometry.time_axis, - coordinates=geometry.rec_positions) - d_obs = Receiver(name='d_obs', grid=model.grid, - time_range=geometry.time_axis, - coordinates=geometry.rec_positions) - d_syn = Receiver(name='d_syn', grid=model.grid, - time_range=geometry.time_axis, - coordinates=geometry.rec_positions) # Update source location - solver.geometry.src_positions[0, :] = source_locations[i, :] + src.coordinates.data[0, :] = source_locations[i, :] # Generate synthetic data from true model - solver.forward(vp=model.vp, rec=d_obs) + solver.forward(vp=model.vp, rec=d_obs, src=src) # Compute smooth data and full forward wavefield u0 - _, u0, _ = solver.forward(vp=vp_in, save=True, rec=d_syn) + _, u0, _ = solver.forward(vp=vp_in, save=True, rec=d_syn, src=src) # Compute gradient from data residual and update objective function - residual = compute_residual(residual, d_obs, d_syn) + compute_residual(residual, d_obs, d_syn) objective += .5*norm(residual)**2 solver.jacobian_adjoint(rec=residual, u=u0, vp=vp_in, grad=grad) diff --git a/examples/seismic/inversion/inversion_utils.py b/examples/seismic/inversion/inversion_utils.py index 9f5709b315..2d51275f92 100644 --- a/examples/seismic/inversion/inversion_utils.py +++ b/examples/seismic/inversion/inversion_utils.py @@ -13,12 +13,9 @@ def compute_residual(res, dobs, dsyn): assert np.allclose(dobs.coordinates.data[:], dsyn.coordinates.data) assert np.allclose(res.coordinates.data[:], dsyn.coordinates.data) # Create a difference operator - diff_eq = Eq(res, dsyn.subs({dsyn.dimensions[-1]: res.dimensions[-1]}) - - dobs.subs({dobs.dimensions[-1]: res.dimensions[-1]})) + diff_eq = Eq(res, dsyn - dobs) Operator(diff_eq)() - return res - def update_with_box(vp, alpha, dm, vmin=2.0, vmax=3.5): """ diff --git a/examples/seismic/test_seismic_utils.py b/examples/seismic/test_seismic_utils.py index ee0aca9da8..f02ae48cb9 100644 --- a/examples/seismic/test_seismic_utils.py +++ b/examples/seismic/test_seismic_utils.py @@ -94,3 +94,12 @@ def test_geom(shape): assert geometry.new_rec(name="bonjour").name == "bonjour" assert geometry.new_src(name="bonjour").name == "bonjour" + + src1 = geometry.src + src2 = geometry.src + assert src1.coordinates is not src2.coordinates + assert src1._sparse_dim is src2._sparse_dim + + src3 = geometry.new_src(name="src3", coordinates=src1.coordinates) + assert src1.coordinates is src3.coordinates + assert src1._sparse_dim is src3._sparse_dim diff --git a/examples/seismic/tti/operators.py b/examples/seismic/tti/operators.py index 9fac0c2047..13871b9067 100644 --- a/examples/seismic/tti/operators.py +++ b/examples/seismic/tti/operators.py @@ -1,6 +1,5 @@ from devito import (Eq, Operator, Function, TimeFunction, NODE, Inc, solve, cos, sin, sqrt, div, grad) -from examples.seismic import PointSource, Receiver from examples.seismic.acoustic.operators import freesurface @@ -444,10 +443,8 @@ def ForwardOperator(model, geometry, space_order=4, v = TimeFunction(name='v', grid=model.grid, staggered=stagg_v, save=geometry.nt if save else None, time_order=time_order, space_order=space_order) - src = PointSource(name='src', grid=model.grid, time_range=geometry.time_axis, - npoint=geometry.nsrc) - rec = Receiver(name='rec', grid=model.grid, time_range=geometry.time_axis, - npoint=geometry.nrec) + src = geometry.src + rec = geometry.rec # FD kernels of the PDE FD_kernel = kernels[(kernel, len(model.shape))] @@ -493,10 +490,8 @@ def AdjointOperator(model, geometry, space_order=4, time_order=time_order, space_order=space_order) r = TimeFunction(name='r', grid=model.grid, staggered=stagg_r, time_order=time_order, space_order=space_order) - srca = PointSource(name='srca', grid=model.grid, time_range=geometry.time_axis, - npoint=geometry.nsrc) - rec = Receiver(name='rec', grid=model.grid, time_range=geometry.time_axis, - npoint=geometry.nrec) + srca = geometry.new_src(name='srca', src_type=None) + rec = geometry.rec # FD kernels of the PDE FD_kernel = kernels[(kernel, len(model.shape))] @@ -535,11 +530,8 @@ def JacobianOperator(model, geometry, space_order=4, time_order = 2 # Create source and receiver symbols - src = Receiver(name='src', grid=model.grid, time_range=geometry.time_axis, - npoint=geometry.nsrc) - - rec = Receiver(name='rec', grid=model.grid, time_range=geometry.time_axis, - npoint=geometry.nrec) + src = geometry.src + rec = geometry.rec # Create wavefields and a dm field u0 = TimeFunction(name='u0', grid=model.grid, save=None, time_order=time_order, @@ -607,8 +599,7 @@ def JacobianAdjOperator(model, geometry, space_order=4, dm = Function(name="dm", grid=model.grid) - rec = Receiver(name='rec', grid=model.grid, time_range=geometry.time_axis, - npoint=geometry.nrec) + rec = geometry.rec # FD kernels of the PDE FD_kernel = kernels[('centered', len(model.shape))] diff --git a/examples/seismic/utils.py b/examples/seismic/utils.py index 835961fa02..6491d9ca5a 100644 --- a/examples/seismic/utils.py +++ b/examples/seismic/utils.py @@ -93,6 +93,10 @@ def __init__(self, model, rec_positions, src_positions, t0, tn, **kwargs): self._interpolation = kwargs.get('interpolation', 'linear') self._r = kwargs.get('r', _default_radius[self.interpolation]) + # Initialize to empty, created at new src/rec + self._src_coordinates = None + self._rec_coordinates = None + def resample(self, dt): self._dt = dt return self @@ -159,22 +163,24 @@ def interpolation(self): def rec(self): return self.new_rec() - def new_rec(self, name='rec'): - return Receiver(name=name, grid=self.grid, - time_range=self.time_axis, npoint=self.nrec, - coordinates=self.rec_positions, - interpolation=self.interpolation, r=self._r) + def new_rec(self, name='rec', coordinates=None): + coords = coordinates or self.rec_positions + rec = Receiver(name=name, grid=self.grid, + time_range=self.time_axis, npoint=self.nrec, + interpolation=self.interpolation, r=self._r, + coordinates=coords) + + return rec @property def adj_src(self): if self.src_type is None: - warning("No source type defined, returning uninitiallized (zero) shot record") return self.new_rec() + coords = self.rec_positions adj_src = sources[self.src_type](name='rec', grid=self.grid, f0=self.f0, time_range=self.time_axis, npoint=self.nrec, - coordinates=self.rec_positions, - t0=self._t0w, a=self._a, - interpolation=self.interpolation, r=self._r) + interpolation=self.interpolation, r=self._r, + coordinates=coords, t0=self._t0w, a=self._a) # Revert time axis to have a proper shot record and not compute on zeros for i in range(self.nrec): adj_src.data[:, i] = adj_src.wavelet[::-1] @@ -184,19 +190,22 @@ def adj_src(self): def src(self): return self.new_src() - def new_src(self, name='src', src_type='self'): + def new_src(self, name='src', src_type='self', coordinates=None): + coords = coordinates or self.src_positions if self.src_type is None or src_type is None: warning("No source type defined, returning uninitiallized (zero) source") - return PointSource(name=name, grid=self.grid, - time_range=self.time_axis, npoint=self.nsrc, - coordinates=self.src_positions, - interpolation=self.interpolation, r=self._r) + src = PointSource(name=name, grid=self.grid, + time_range=self.time_axis, npoint=self.nsrc, + coordinates=coords, + interpolation=self.interpolation, r=self._r) else: - return sources[self.src_type](name=name, grid=self.grid, f0=self.f0, - time_range=self.time_axis, npoint=self.nsrc, - coordinates=self.src_positions, - t0=self._t0w, a=self._a, - interpolation=self.interpolation, r=self._r) + src = sources[self.src_type](name=name, grid=self.grid, f0=self.f0, + time_range=self.time_axis, npoint=self.nsrc, + coordinates=coords, + t0=self._t0w, a=self._a, + interpolation=self.interpolation, r=self._r) + + return src sources = {'Wavelet': WaveletSource, 'Ricker': RickerSource, 'Gabor': GaborSource} diff --git a/examples/seismic/viscoacoustic/operators.py b/examples/seismic/viscoacoustic/operators.py index 02bd7245d1..1253b61686 100755 --- a/examples/seismic/viscoacoustic/operators.py +++ b/examples/seismic/viscoacoustic/operators.py @@ -3,7 +3,6 @@ from devito import (Eq, Operator, VectorTimeFunction, TimeFunction, Function, NODE, div, grad, solve) -from examples.seismic import PointSource, Receiver def src_rec(p, model, geometry, **kwargs): @@ -14,10 +13,8 @@ def src_rec(p, model, geometry, **kwargs): dt = model.grid.time_dim.spacing m = model.m # Source symbol with input wavelet - src = PointSource(name="src", grid=model.grid, time_range=geometry.time_axis, - npoint=geometry.nsrc) - rec = Receiver(name='rec', grid=model.grid, time_range=geometry.time_axis, - npoint=geometry.nrec) + src = geometry.src + rec = geometry.rec forward = kwargs.get('forward', True) time_order = p.time_order diff --git a/examples/seismic/viscoacoustic/wavesolver.py b/examples/seismic/viscoacoustic/wavesolver.py index 05125807b8..e491f8dc4b 100755 --- a/examples/seismic/viscoacoustic/wavesolver.py +++ b/examples/seismic/viscoacoustic/wavesolver.py @@ -1,7 +1,6 @@ from devito import (VectorTimeFunction, TimeFunction, Function, NODE, DevitoCheckpoint, CheckpointOperator, Revolver) from devito.tools import memoized_meth -from examples.seismic import PointSource from examples.seismic.viscoacoustic.operators import ( ForwardOperator, AdjointOperator, GradientOperator, BornOperator ) @@ -181,9 +180,7 @@ def adjoint(self, rec, srca=None, va=None, pa=None, r=None, model=None, **kwargs Adjoint source, wavefield and performance summary. """ # Create a new adjoint source and receiver symbol - srca = srca or PointSource(name='srca', grid=self.model.grid, - time_range=self.geometry.time_axis, - coordinates=self.geometry.src_positions) + srca = srca or self.geometry.new_src(name='srca', src_type=None) if self.time_order == 1: va = va or VectorTimeFunction(name="va", grid=self.model.grid, diff --git a/tests/test_adjoint.py b/tests/test_adjoint.py index 759b86a3ae..619e292d43 100644 --- a/tests/test_adjoint.py +++ b/tests/test_adjoint.py @@ -3,7 +3,7 @@ from devito import Operator, norm, Function, Grid, SparseFunction, inner from devito.logger import info -from examples.seismic import demo_model, Receiver +from examples.seismic import demo_model from examples.seismic.acoustic import acoustic_setup from examples.seismic.tti import tti_setup from examples.seismic.viscoacoustic import viscoacoustic_setup @@ -105,9 +105,7 @@ def test_adjoint_F(self, mkey, shape, kernel, space_order, time_order, setup_fun **(presets[mkey]), dtype=np.float64) # Create adjoint receiver symbol - srca = Receiver(name='srca', grid=solver.model.grid, - time_range=solver.geometry.time_axis, - coordinates=solver.geometry.src_positions) + srca = solver.geometry.new_src(name="srca", src_type=None) # Run forward and adjoint operators rec = solver.forward(save=False)[0] diff --git a/tests/test_mpi.py b/tests/test_mpi.py index 4602810292..ddcc0de27f 100644 --- a/tests/test_mpi.py +++ b/tests/test_mpi.py @@ -526,10 +526,14 @@ def test_scatter_gather(self, mode): loc_data = loc_data*2 # Gather - sf._dist_gather(loc_data, loc_coords) + sf._dist_data_gather(loc_data) assert len(sf.data) == 1 assert np.all(sf.data == data[sf.local_indices]*2) + sf._dist_subfunc_gather(loc_coords, sf.coordinates) + assert sf.coordinates.data.shape == (1, 2) + assert np.all(sf.coordinates.data == coords[sf.local_indices[0], :]) + @pytest.mark.parallel(mode=4) @switchconfig(condition=isinstance(configuration['compiler'], (OneapiCompiler)), safe_math=True) diff --git a/tests/test_sparse.py b/tests/test_sparse.py index 136d1ef9e3..7c94de69ce 100644 --- a/tests/test_sparse.py +++ b/tests/test_sparse.py @@ -423,8 +423,7 @@ def test_rebuild(self, sptype): # Check new subfunction for subf in sp2._sub_functions: if getattr(sp2, subf) is not None: - assert getattr(sp2, subf).name.startswith("sr_") - assert np.all(getattr(sp2, subf).data == 0) + assert getattr(sp2, subf) == getattr(sp, subf) # Rebuild with different name as an alias sp2 = sp._rebuild(name="sr2", alias=True)