From 765f5df22f05716e4e1dc224445f2999694d3f3b Mon Sep 17 00:00:00 2001 From: Carlos Cueto Date: Mon, 3 Jun 2024 14:28:17 +0100 Subject: [PATCH 1/9] Introduced forced attenuation mode --- stride/physics/iso_acoustic/devito.py | 25 +++++++++++++++++++------ 1 file changed, 19 insertions(+), 6 deletions(-) diff --git a/stride/physics/iso_acoustic/devito.py b/stride/physics/iso_acoustic/devito.py index 0136f1d..751d36a 100644 --- a/stride/physics/iso_acoustic/devito.py +++ b/stride/physics/iso_acoustic/devito.py @@ -92,7 +92,7 @@ class IsoAcousticDevito(ProblemTypeBase): Type of source/receiver interpolation (``linear`` for bi-/tri-linear or ``hicks`` for sinc interpolation), defaults to ``linear``. attenuation_power : int, optional - Power of the attenuation law if attenuation is given (``0`` or ``2``), + Power of the attenuation law if attenuation is given (``0``, ``2``, or None), defaults to ``0``. drp : bool, optional Whether or not to use dispersion-relation preserving coefficients (only @@ -272,7 +272,7 @@ async def before_forward(self, wavelets, vp, rho=None, alpha=None, **kwargs): Type of source/receiver interpolation (``linear`` for bi-/tri-linear or ``hicks`` for sinc interpolation), defaults to ``linear``. attenuation_power : int, optional - Power of the attenuation law if attenuation is given (``0`` or ``2``), + Power of the attenuation law if attenuation is given (``0``, ``2``, or None), defaults to ``0``. drp : bool, optional Whether or not to use dispersion-relation preserving coefficients (only @@ -456,9 +456,13 @@ async def before_forward(self, wavelets, vp, rho=None, alpha=None, **kwargs): self.dev_grid.vars.buoy.data_with_halo[:] = 1/rho_with_halo if alpha is not None: - self.logger.perf('(ShotID %d) Using attenuation with power %d' % (problem.shot_id, self.attenuation_power)) + att_pwr = str(self.attenuation_power) if self.attenuation_power is not None else 'None' + self.logger.perf('(ShotID %d) Using attenuation with power %s' % (problem.shot_id, att_pwr)) - db_to_neper = 100 * (1e-6 / (2*np.pi))**self.attenuation_power / (20 * np.log10(np.exp(1))) + if self.attenuation_power is not None: + db_to_neper = 100 * (1e-6 / (2*np.pi))**self.attenuation_power / (20 * np.log10(np.exp(1))) + else: + db_to_neper = 1 alpha_with_halo = self.dev_grid.with_halo(alpha.extended_data)*db_to_neper self.dev_grid.vars.alpha.data_with_halo[:] = alpha_with_halo @@ -812,7 +816,10 @@ async def before_adjoint(self, adjoint_source, wavelets, vp, rho=None, alpha=Non self.dev_grid.vars.buoy.data_with_halo[:] = 1/rho_with_halo if alpha is not None: - db_to_neper = 100 * (1e-6 / (2*np.pi))**self.attenuation_power / (20 * np.log10(np.exp(1))) + if self.attenuation_power is not None: + db_to_neper = 100 * (1e-6 / (2*np.pi))**self.attenuation_power / (20 * np.log10(np.exp(1))) + else: + db_to_neper = 1 alpha_with_halo = self.dev_grid.with_halo(alpha.extended_data)*db_to_neper self.dev_grid.vars.alpha.data_with_halo[:] = alpha_with_halo @@ -1385,7 +1392,7 @@ def _stencil(self, field, wavelets, vp, rho=None, alpha=None, direction='forward subs = None # Get the attenuation term - if alpha_fun is not None: + if alpha_fun is not None and self.attenuation_power is not None: if self.attenuation_power == 0: u = field elif self.attenuation_power == 2: @@ -1458,6 +1465,12 @@ def _stencil(self, field, wavelets, vp, rho=None, alpha=None, direction='forward coefficients=subs) for dom in boundary] stencils += stencil_boundary + # Forced attenuation update + if alpha_fun is not None and self.attenuation_power is None: + stencils += [devito.Eq(u_next, alpha_fun*u_next, + subdomain=None, + coefficients=subs)] + return sub_befores + eq_before + stencils + eq_after + sub_afters def _medium_functions(self, vp, rho=None, alpha=None, **kwargs): From 5f2cb606c272b27e56e16d75738942cbc0a5c819 Mon Sep 17 00:00:00 2001 From: Carlos Cueto Date: Tue, 11 Jun 2024 16:36:15 +0100 Subject: [PATCH 2/9] Improvements to prec scale and dumping --- stride/optimisation/optimisers/optimiser.py | 12 +++++++++-- stride/physics/iso_acoustic/devito.py | 22 ++++++++++++++++++--- stride/problem/data.py | 16 +++++++++++++-- 3 files changed, 43 insertions(+), 7 deletions(-) diff --git a/stride/optimisation/optimisers/optimiser.py b/stride/optimisation/optimisers/optimiser.py index 9d31340..145a4a4 100644 --- a/stride/optimisation/optimisers/optimiser.py +++ b/stride/optimisation/optimisers/optimiser.py @@ -77,13 +77,14 @@ async def pre_process(self, grad=None, processed_grad=None, **kwargs): logger = mosaic.logger() logger.perf('Updating variable %s,' % self.variable.name) + problem = kwargs.pop('problem', None) + iteration = kwargs.pop('iteration', None) + if processed_grad is None: if grad is None: if hasattr(self.variable, 'is_proxy') and self.variable.is_proxy: await self.variable.pull(attr='grad') - problem = kwargs.pop('problem', None) - iteration = kwargs.pop('iteration', None) dump_grad = kwargs.pop('dump_grad', self.dump_grad) dump_prec = kwargs.pop('dump_prec', self.dump_prec) if dump_grad and problem is not None: @@ -112,6 +113,13 @@ async def pre_process(self, grad=None, processed_grad=None, **kwargs): processed_grad = await self._process_grad(grad, variable=self.variable, **kwargs) + dump_processed_grad = kwargs.pop('dump_processed_grad', self.dump_grad) + if dump_processed_grad and problem is not None: + processed_grad.dump(path=problem.output_folder, + project_name=problem.name, + parameter='processed_%s' % self.variable.grad.name, + version=iteration.abs_id + 1) + test_step_size = kwargs.pop('test_step_size', self.test_step_size) processed_grad.data[:] *= test_step_size diff --git a/stride/physics/iso_acoustic/devito.py b/stride/physics/iso_acoustic/devito.py index 751d36a..1a1bac0 100644 --- a/stride/physics/iso_acoustic/devito.py +++ b/stride/physics/iso_acoustic/devito.py @@ -719,7 +719,7 @@ async def before_adjoint(self, adjoint_source, wavelets, vp, rho=None, alpha=Non if not fw3d_mode: rec_term = rec.inject(field=p_a.backward, expr=-rec.subs({t: t-1}) * self.time.step**2 * vp2) else: - rec_term = rec.inject(field=p_a.backward, expr=-rec * self.time.step ** 2 * vp2) + rec_term = rec.inject(field=p_a.backward, expr=-rec * self.time.step**2 * vp2) if wavelets.needs_grad: src_term = src.interpolate(expr=p_a) @@ -911,7 +911,7 @@ async def after_adjoint(self, adjoint_source, wavelets, vp, rho=None, alpha=None Tuple with the gradients of the variables that need them """ - problem = kwargs.pop('problem') + problem = kwargs.get('problem') shot = problem.shot dump_adjoint_wavefield = kwargs.pop('dump_adjoint_wavefield', False) @@ -1033,7 +1033,7 @@ async def get_grad_vp(self, vp, **kwargs): variable_prec = np.asarray(variable_prec.data[self.space.inner], dtype=np.float32) variable_grad *= -2 / vp.data**3 - variable_prec *= +4 / vp.data**6 * self.time.step**2 + # variable_prec *= +4 / vp.data**6 * self.time.step**2 deallocate = kwargs.pop('deallocate', False) if deallocate: @@ -1049,6 +1049,22 @@ async def get_grad_vp(self, vp, **kwargs): extended_shape=variable_prec.shape, inner=None) + problem = kwargs.pop('problem', None) + iteration = kwargs.pop('iteration', None) + dump_local_grad = kwargs.pop('dump_local_grad', False) + dump_local_prec = kwargs.pop('dump_local_prec', False) + if dump_local_grad and problem is not None: + grad.dump(path=problem.output_folder, + project_name=problem.name, + parameter='vp_local_grad-Shot%05d' % problem.shot_id, + version=iteration.abs_id + 1) + + if dump_local_prec and problem is not None: + grad.prec.dump(path=problem.output_folder, + project_name=problem.name, + parameter='vp_local_src_prec-Shot%05d' % problem.shot_id, + version=iteration.abs_id + 1) + return grad async def prepare_grad_rho(self, rho, **kwargs): diff --git a/stride/problem/data.py b/stride/problem/data.py index 037089a..ff018ac 100644 --- a/stride/problem/data.py +++ b/stride/problem/data.py @@ -365,7 +365,7 @@ def process_grad(self, global_prec=True, **kwargs): self.grad.apply_prec(**kwargs) return self.grad - def apply_prec(self, prec_scale=0.25, prec_op=None, prec=None, **kwargs): + def apply_prec(self, prec_scale=4.0, prec_smooth=1.5, prec_op=None, prec=None, **kwargs): """ Apply a pre-conditioner to the current field. @@ -375,6 +375,8 @@ def apply_prec(self, prec_scale=0.25, prec_op=None, prec=None, **kwargs): Condition scaling for the preconditioner. prec_op : callable, optional Additional operation to apply to the preconditioner. + prec_smooth : float, optional + Smoothing to apply to the preconditioner. prec : StructuredData, optional Pre-conditioner to apply. Defaults to self.prec. @@ -382,9 +384,18 @@ def apply_prec(self, prec_scale=0.25, prec_op=None, prec=None, **kwargs): ------- """ + # tmpnrm = sum(abs(a)) + # sumall(j) = tmpnrm/nn + # eps(j) = sumall(j) * precstabfactor + # val(:) = 1.0/eps(:) + # u(:,ix,iy,iz) = u(:,ix,iy,iz)/(a(:,ix,iy,iz)*val(:)+1.0) + prec = self.prec if prec is None else prec if prec is not None: + if prec_smooth is not None: + prec.data[:] = scipy.ndimage.gaussian_filter(prec.data, prec_smooth) + prec_factor = np.sum(np.abs(prec.data)) if prec_factor > 1e-31: @@ -395,7 +406,8 @@ def apply_prec(self, prec_scale=0.25, prec_op=None, prec=None, **kwargs): prec.data[:] = prec_op(prec.data) non_zero = np.abs(prec.data) > 0. - self.data[non_zero] /= prec.data[non_zero] + prec.data[non_zero] = 1/prec.data[non_zero] + self.data[non_zero] *= prec.data[non_zero] return self From 7b8e0763008eae5174dcae6b028f8edd72f6f00d Mon Sep 17 00:00:00 2001 From: Carlos Cueto Date: Fri, 14 Jun 2024 11:40:55 +0100 Subject: [PATCH 3/9] Added rampoff to gradient mask --- .../pipelines/steps/mask_field.py | 34 +++++++++++++++++-- .../pipelines/steps/norm_field.py | 2 +- .../pipelines/steps/shift_traces.py | 3 +- stride/physics/common/devito.py | 8 ++++- stride/problem/data.py | 8 +---- 5 files changed, 43 insertions(+), 12 deletions(-) diff --git a/stride/optimisation/pipelines/steps/mask_field.py b/stride/optimisation/pipelines/steps/mask_field.py index f09f416..79cd04c 100644 --- a/stride/optimisation/pipelines/steps/mask_field.py +++ b/stride/optimisation/pipelines/steps/mask_field.py @@ -5,6 +5,31 @@ from ....core import Operator +def _rampoff_mask(shape, ramp_size): + mask = np.ones(shape, dtype=np.float32) + + for dim_i in range(len(shape)): + for index in range(ramp_size): + pos = np.abs((ramp_size - index - 1) / float(ramp_size - 1)) + val = 1 - np.cos(np.pi / 2 * (1 - pos)) + + # : slices + all_ind = [slice(index, s - index + 1) for s in shape] + + # Left slice + all_ind[dim_i] = index + mask[tuple(all_ind)] = val + + # : slices + all_ind = [slice(index, s - index + 1) for s in shape] + + # right slice + all_ind[dim_i] = -index + mask[tuple(all_ind)] = val + + return mask + + class MaskField(Operator): """ Mask a StructuredData object to remove values outside inner domain. @@ -17,13 +42,18 @@ class MaskField(Operator): def __init__(self, **kwargs): super().__init__(**kwargs) + self.mask_rampoff = kwargs.pop('mask_rampoff', 10) self._mask = kwargs.pop('mask', None) def forward(self, field, **kwargs): - mask = kwargs.pop('mask', self._mask) + mask = kwargs.pop('mask', None) + mask_rampoff = kwargs.pop('mask_rampoff', self.mask_rampoff) + mask = self._mask if mask is None else mask if mask is None: - mask = np.zeros(field.extended_shape) + mask = np.zeros(field.extended_shape, dtype=np.float32) mask[field.inner] = 1 + mask *= _rampoff_mask(mask.shape, mask_rampoff) + self._mask = mask out_field = field.alike(name=name_from_op_name(self, field)) out_field.extended_data[:] = field.extended_data diff --git a/stride/optimisation/pipelines/steps/norm_field.py b/stride/optimisation/pipelines/steps/norm_field.py index a0b869d..f963fd7 100644 --- a/stride/optimisation/pipelines/steps/norm_field.py +++ b/stride/optimisation/pipelines/steps/norm_field.py @@ -28,7 +28,7 @@ def forward(self, field, **kwargs): out_field = field.alike(name=name_from_op_name(self, field)) out_field.extended_data[:] = field.extended_data - out_field.extended_data[:] /= self.norm_value + out_field.extended_data[:] *= 1. / self.norm_value return out_field diff --git a/stride/optimisation/pipelines/steps/shift_traces.py b/stride/optimisation/pipelines/steps/shift_traces.py index 58a8eb9..763175c 100644 --- a/stride/optimisation/pipelines/steps/shift_traces.py +++ b/stride/optimisation/pipelines/steps/shift_traces.py @@ -62,7 +62,8 @@ def _apply(self, traces, **kwargs): out_data = traces.extended_data.copy() out_traces = traces.alike(name=name_from_op_name(self, traces)) - out_data[:, :-shift] = out_data[:, shift:] + if shift > 0: + out_data[:, :-shift] = out_data[:, shift:] out_traces.extended_data[:] = out_data diff --git a/stride/physics/common/devito.py b/stride/physics/common/devito.py index e01fd69..a939098 100644 --- a/stride/physics/common/devito.py +++ b/stride/physics/common/devito.py @@ -588,7 +588,7 @@ def sparse_time_function(self, name, num=1, space_order=None, time_order=None, Spatial coordinates of the sparse points (num points, dimensions), only needed when interpolation is not linear. interpolation_type : str, optional - Type of interpolation to perform (``linear`` or ``hicks``), defaults + Type of interpolation to perform (``linear``, ``sinc``, or ``hicks``), defaults to ``linear``, computationally more efficient but less accurate. kwargs Additional arguments for the Devito constructor. @@ -620,6 +620,12 @@ def sparse_time_function(self, name, num=1, space_order=None, time_order=None, if interpolation_type == 'linear': fun = devito.SparseTimeFunction(**sparse_kwargs) + elif interpolation_type == 'sinc': + r = sparse_kwargs.pop('r', 7) + fun = devito.SparseTimeFunction(interpolation='sinc', r=r, + coordinates=coordinates, + **sparse_kwargs) + elif interpolation_type == 'hicks': r = sparse_kwargs.pop('r', 7) diff --git a/stride/problem/data.py b/stride/problem/data.py index ff018ac..0b5d9ea 100644 --- a/stride/problem/data.py +++ b/stride/problem/data.py @@ -365,7 +365,7 @@ def process_grad(self, global_prec=True, **kwargs): self.grad.apply_prec(**kwargs) return self.grad - def apply_prec(self, prec_scale=4.0, prec_smooth=1.5, prec_op=None, prec=None, **kwargs): + def apply_prec(self, prec_scale=4.0, prec_smooth=1.0, prec_op=None, prec=None, **kwargs): """ Apply a pre-conditioner to the current field. @@ -384,12 +384,6 @@ def apply_prec(self, prec_scale=4.0, prec_smooth=1.5, prec_op=None, prec=None, * ------- """ - # tmpnrm = sum(abs(a)) - # sumall(j) = tmpnrm/nn - # eps(j) = sumall(j) * precstabfactor - # val(:) = 1.0/eps(:) - # u(:,ix,iy,iz) = u(:,ix,iy,iz)/(a(:,ix,iy,iz)*val(:)+1.0) - prec = self.prec if prec is None else prec if prec is not None: From 6af89ad21f8f558a94601ce6ecd3e5c6124c480a Mon Sep 17 00:00:00 2001 From: Carlos Cueto Date: Fri, 14 Jun 2024 14:22:44 +0100 Subject: [PATCH 4/9] Reduce core object state changes --- mosaic/cli/mprof.py | 4 ++-- mosaic/core/base.py | 1 - mosaic/core/tessera.py | 8 +++----- stride/physics/iso_acoustic/devito.py | 1 - 4 files changed, 5 insertions(+), 9 deletions(-) diff --git a/mosaic/cli/mprof.py b/mosaic/cli/mprof.py index 681936f..5563e61 100644 --- a/mosaic/cli/mprof.py +++ b/mosaic/cli/mprof.py @@ -25,11 +25,11 @@ def __init__(self, label_id, label): # EVENTS # # * Tessera -# Remote: init --> (listening --> running) --> collected +# Remote: init --> listening --> collected # Proxy: pending --> init --> listening --> collected # # * Task -# Remote: init --> pending --> ready --> running --> done --> collected +# Remote: init --> pending --> done/failed --> collected # Proxy: pending --> init --> queued --> (done --> result) --> collected diff --git a/mosaic/core/base.py b/mosaic/core/base.py index c3ebf5a..261684f 100644 --- a/mosaic/core/base.py +++ b/mosaic/core/base.py @@ -439,7 +439,6 @@ def __deepcopy__(self, memo): async def deregister(self): try: self.logger.debug('Garbage collected object %s' % self) - self.state_changed('collected') except AttributeError: pass diff --git a/mosaic/core/tessera.py b/mosaic/core/tessera.py index 1ab1056..41ac1dc 100644 --- a/mosaic/core/tessera.py +++ b/mosaic/core/tessera.py @@ -262,6 +262,8 @@ async def listen_async(self): if self._state != 'init': return + self.state_changed('listening') + while True: sender_id, task = await self._task_queue.get() # Make sure that the loop does not keep implicit references to the task until the @@ -288,12 +290,10 @@ async def run_async(self): ------- """ - if self._state != 'init': + if self._state != 'listening': return while True: - self.state_changed('listening') - sender_id, task, future = await self._run_queue.get() # Make sure that the loop does not keep implicit references to the task until the # next task arrives in the queue @@ -319,8 +319,6 @@ async def run_async(self): self.obj.__class__.__name__)) await asyncio.sleep(0) - self.state_changed('running') - task.state_changed('running') await self.logger.send() await self.call_safe(sender_id, method, task) diff --git a/stride/physics/iso_acoustic/devito.py b/stride/physics/iso_acoustic/devito.py index 1a1bac0..6d78d67 100644 --- a/stride/physics/iso_acoustic/devito.py +++ b/stride/physics/iso_acoustic/devito.py @@ -1033,7 +1033,6 @@ async def get_grad_vp(self, vp, **kwargs): variable_prec = np.asarray(variable_prec.data[self.space.inner], dtype=np.float32) variable_grad *= -2 / vp.data**3 - # variable_prec *= +4 / vp.data**6 * self.time.step**2 deallocate = kwargs.pop('deallocate', False) if deallocate: From 472f412c6c15501280ac576ae2decc91e2812ed2 Mon Sep 17 00:00:00 2001 From: Carlos Cueto Date: Mon, 17 Jun 2024 14:34:35 +0100 Subject: [PATCH 5/9] Improve comms debug print --- mosaic/comms/comms.py | 17 +++++++++-------- 1 file changed, 9 insertions(+), 8 deletions(-) diff --git a/mosaic/comms/comms.py b/mosaic/comms/comms.py index 7b26350..dc3ac16 100644 --- a/mosaic/comms/comms.py +++ b/mosaic/comms/comms.py @@ -771,19 +771,20 @@ def _process_send(self, method, cmd=None, reply=False, **kwargs): 'cmd': cmd, } + msg = serialise(msg) + msg_size = sizeof(msg) + if not method.startswith('log') and not method.startswith('update_monitored_node'): if method == 'cmd': method = '%s:%s.%s' % (method, cmd['type'], cmd['method']) - self.logger.debug('Sending cmd %s %s to %s (%s) from %s' % (method, cmd['method'], - self.uid, cmd['uid'], - self._runtime.uid)) + self.logger.debug('Sending cmd %s %s to %s (%s) from %s ' + '(size %.2f MB)' % (method, cmd['method'], self.uid, cmd['uid'], + self._runtime.uid, msg_size/1024**2)) else: - self.logger.debug('Sending msg %s to %s from %s' % (method, self.uid, - self._runtime.uid)) - - msg = serialise(msg) - msg_size = sizeof(msg) + self.logger.debug('Sending msg %s to %s from %s ' + '(size %.2f MB)' % (method, self.uid, self._runtime.uid, + msg_size/1024**2)) compression = [] compressed_msg = [] From 78a4432b565a13feec3bc762dfb225d124d03cfe Mon Sep 17 00:00:00 2001 From: Carlos Cueto Date: Thu, 20 Jun 2024 10:08:00 +0100 Subject: [PATCH 6/9] Deactivate source shifting in fw3dmode --- stride/physics/iso_acoustic/devito.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/stride/physics/iso_acoustic/devito.py b/stride/physics/iso_acoustic/devito.py index 6d78d67..554e631 100644 --- a/stride/physics/iso_acoustic/devito.py +++ b/stride/physics/iso_acoustic/devito.py @@ -470,8 +470,8 @@ async def before_forward(self, wavelets, vp, rho=None, alpha=None, **kwargs): # Set geometry and wavelet wavelets = wavelets.data - if fw3d_mode: - wavelets[:, 1:] = wavelets[:, :-1] + # if fw3d_mode: + # wavelets[:, 1:] = wavelets[:, :-1] if diff_source: wavelets = np.gradient(wavelets, self.time.step, axis=-1) @@ -827,8 +827,8 @@ async def before_adjoint(self, adjoint_source, wavelets, vp, rho=None, alpha=Non # Set geometry and adjoint source adjoint_source = adjoint_source.data - if fw3d_mode: - adjoint_source[:, 1:] = adjoint_source[:, :-1] + # if fw3d_mode: + # adjoint_source[:, 1:] = adjoint_source[:, :-1] window = scipy.signal.get_window(('tukey', 0.001), time_bounds[1]-time_bounds[0], False) window = np.pad(window, ((time_bounds[0], self.time.num-time_bounds[1]),), mode='constant', constant_values=0.) From d2586c2e4f7755e32282d5b0e0b93a7b0d65588f Mon Sep 17 00:00:00 2001 From: Carlos Cueto Date: Thu, 18 Jul 2024 13:48:03 +0100 Subject: [PATCH 7/9] Work consistently in slowness --- docker/README.md | 6 +- mosaic/cli/mrun.py | 2 +- mosaic/comms/serialisation.py | 2 +- stride/__init__.py | 4 +- stride/core.py | 4 + stride/optimisation/loss/l2_distance.py | 4 +- .../optimisers/gradient_descent.py | 6 +- stride/optimisation/optimisers/optimiser.py | 23 ++- .../optimisation/step_length/line_search.py | 6 +- stride/physics/common/devito.py | 14 +- stride/physics/iso_acoustic/devito.py | 155 +++++++++++------- stride/physics/iso_elastic/devito.py | 18 +- stride/physics/marmottant/devito.py | 6 +- stride/physics/problem_type.py | 42 +++-- stride/problem/data.py | 29 +++- 15 files changed, 199 insertions(+), 122 deletions(-) diff --git a/docker/README.md b/docker/README.md index 27f7c71..32bc49b 100644 --- a/docker/README.md +++ b/docker/README.md @@ -81,19 +81,19 @@ docker build --network=host --file docker/Dockerfile.stride --tag stride . And to build the GPU image with `openacc` offloading and the `nvc` compiler, simply run: ```bash -docker build --build-arg base=devitocodes/base:nvidia-nvc --network=host --file docker/Dockerfile.stride --tag stride . +docker build --build-arg base=devitocodes/bases:nvidia-nvc --network=host --file docker/Dockerfile.stride --tag stride . ``` or if you wish to use the `llvm-15` (clang) compiler with `openmp` offlaoding: ```bash -docker build --build-arg base=devitocodes/base:nvidia-clang --network=host --file docker/Dockerfile.stride --tag stride . +docker build --build-arg base=devitocodes/bases:nvidia-clang --network=host --file docker/Dockerfile.stride --tag stride . ``` and finally for AMD architectures: ```bash -docker build --build-arg base=devitocodes/base:amd --network=host --file docker/Dockerfile.stride --tag stride . +docker build --build-arg base=devitocodes/bases:amd --network=host --file docker/Dockerfile.stride --tag stride . ``` diff --git a/mosaic/cli/mrun.py b/mosaic/cli/mrun.py index 38baead..917127b 100644 --- a/mosaic/cli/mrun.py +++ b/mosaic/cli/mrun.py @@ -146,7 +146,7 @@ def start_runtime(*args, **extra_kwargs): _runtime.init_file(runtime_config) def run_head(): - process = cmd_subprocess.run(cmd, + process = cmd_subprocess.run(' '.join(cmd), shell=True, stdout=_stdout, stderr=_stderr) diff --git a/mosaic/comms/serialisation.py b/mosaic/comms/serialisation.py index 4d3f6f9..348f983 100644 --- a/mosaic/comms/serialisation.py +++ b/mosaic/comms/serialisation.py @@ -35,7 +35,7 @@ def serialise(data): """ try: return pickle5_dumps(data) - except pickle.PicklingError: + except (pickle.PicklingError, AttributeError): return cloudpickle.dumps(data), [] diff --git a/stride/__init__.py b/stride/__init__.py index f8e4178..fab679b 100644 --- a/stride/__init__.py +++ b/stride/__init__.py @@ -108,7 +108,7 @@ async def forward(problem, pde, *args, **kwargs): published_args = await asyncio.gather(*published_args) platform = kwargs.get('platform', 'cpu') - using_gpu = platform in ['nvidia-acc', 'gpu'] + using_gpu = 'nvidia' in platform or 'gpu' in platform if using_gpu: devices = kwargs.pop('devices', None) num_gpus = gpu_count() if devices is None else len(devices) @@ -251,7 +251,7 @@ async def adjoint(problem, pde, loss, optimisation_loop, optimiser, *args, **kwa keep_residual = isinstance(step_size, LineSearch) platform = kwargs.get('platform', 'cpu') - using_gpu = platform in ['nvidia-acc', 'gpu'] + using_gpu = 'nvidia' in platform or 'gpu' in platform if using_gpu: devices = kwargs.pop('devices', None) num_gpus = gpu_count() if devices is None else len(devices) diff --git a/stride/core.py b/stride/core.py index 2f667b0..173ca64 100644 --- a/stride/core.py +++ b/stride/core.py @@ -288,6 +288,7 @@ def __init__(self, *args, **kwargs): self.grad = None self.prec = None + self.transform = kwargs.pop('transform', None) self.graph = Graph() self.prev_op = None @@ -384,6 +385,7 @@ def detach(self, *args, **kwargs): """ kwargs['name'] = kwargs.pop('name', self._init_name) kwargs['needs_grad'] = kwargs.pop('needs_grad', self.needs_grad) + kwargs['transform'] = kwargs.pop('transform', self.transform) if hasattr(self, 'has_tessera') and self.has_tessera: cpy = self.__class__.parameter(*args, **kwargs) @@ -411,6 +413,7 @@ def as_parameter(self, *args, **kwargs): """ kwargs['name'] = kwargs.pop('name', self._init_name) kwargs['needs_grad'] = kwargs.pop('needs_grad', self.needs_grad) + kwargs['transform'] = kwargs.pop('transform', self.transform) cpy = self.__class__.parameter(*args, **kwargs) @@ -437,6 +440,7 @@ def copy(self, *args, **kwargs): """ kwargs['name'] = kwargs.pop('name', self._init_name) kwargs['needs_grad'] = kwargs.pop('needs_grad', self.needs_grad) + kwargs['transform'] = kwargs.pop('transform', self.transform) propagate_tessera = kwargs.pop('propagate_tessera', True) diff --git a/stride/optimisation/loss/l2_distance.py b/stride/optimisation/loss/l2_distance.py index 01ab098..fd5b69f 100644 --- a/stride/optimisation/loss/l2_distance.py +++ b/stride/optimisation/loss/l2_distance.py @@ -24,7 +24,7 @@ def __init__(self, **kwargs): self.residual = None - async def forward(self, modelled, observed, **kwargs): + def forward(self, modelled, observed, **kwargs): problem = kwargs.pop('problem', None) shot_id = problem.shot_id if problem is not None \ else kwargs.pop('shot_id', 0) @@ -38,7 +38,7 @@ async def forward(self, modelled, observed, **kwargs): return fun - async def adjoint(self, d_fun, modelled, observed, **kwargs): + def adjoint(self, d_fun, modelled, observed, **kwargs): grad_modelled = None if modelled.needs_grad: grad_modelled = +np.asarray(d_fun) * self.residual.copy(name='modelledresidual') diff --git a/stride/optimisation/optimisers/gradient_descent.py b/stride/optimisation/optimisers/gradient_descent.py index d91cc78..21a7dd8 100644 --- a/stride/optimisation/optimisers/gradient_descent.py +++ b/stride/optimisation/optimisers/gradient_descent.py @@ -30,6 +30,6 @@ async def pre_process(self, grad=None, processed_grad=None, **kwargs): **kwargs) return processed_grad - async def update_variable(self, step_size, direction): - self.variable.data[:] -= step_size * direction.data - return self.variable + def update_variable(self, step_size, variable, direction): + variable.data[:] -= step_size * direction.data + return variable diff --git a/stride/optimisation/optimisers/optimiser.py b/stride/optimisation/optimisers/optimiser.py index 145a4a4..87e0553 100644 --- a/stride/optimisation/optimisers/optimiser.py +++ b/stride/optimisation/optimisers/optimiser.py @@ -96,6 +96,7 @@ async def pre_process(self, grad=None, processed_grad=None, **kwargs): if dump_prec and self.variable.grad.prec is not None and problem is not None: self.variable.grad.prec.dump(path=problem.output_folder, project_name=problem.name, + parameter='raw_%s' % self.variable.grad.prec.name, version=iteration.abs_id+1) grad = self.variable.process_grad(**kwargs) @@ -105,6 +106,11 @@ async def pre_process(self, grad=None, processed_grad=None, **kwargs): project_name=problem.name, version=iteration.abs_id+1) + if dump_prec and self.variable.grad.prec is not None and problem is not None: + self.variable.grad.prec.dump(path=problem.output_folder, + project_name=problem.name, + version=iteration.abs_id+1) + min_dir = np.min(grad.data) max_dir = np.max(grad.data) @@ -171,7 +177,7 @@ async def step(self, step_size=None, grad=None, processed_grad=None, **kwargs): step_size = self.step_size if step_size is None else step_size step_loop = kwargs.pop('step_loop', None) if isinstance(step_size, LineSearch): - await step_size.init_search( + step_size.init_search( variable=self.variable, direction=direction, **kwargs @@ -185,7 +191,7 @@ async def step(self, step_size=None, grad=None, processed_grad=None, **kwargs): next_step = 1. done_search = True else: - next_step, done_search = await step_size.next_step( + next_step, done_search = step_size.next_step( variable=self.variable, direction=direction, **kwargs @@ -216,7 +222,14 @@ async def step(self, step_size=None, grad=None, processed_grad=None, **kwargs): self.variable.data[:] = variable_before.data.copy() # update variable - await self.update_variable(next_step, direction) + if self.variable.transform is not None: + variable = self.variable.transform(self.variable) + else: + variable = self.variable + upd_variable = self.update_variable(next_step, variable, direction) + if self.variable.transform is not None: + upd_variable = self.variable.transform(upd_variable) + self.variable.data[:] = upd_variable.data.copy() # post-process variable after update await self.post_process(**kwargs) @@ -262,13 +275,15 @@ async def post_process(self, **kwargs): self.variable.release_grad() @abstractmethod - async def update_variable(self, step_size, direction): + def update_variable(self, step_size, variable, direction): """ Parameters ---------- step_size : float Step size to use for updating the variable. + variable : Data + Variable to update. direction : Data Direction in which to update the variable. diff --git a/stride/optimisation/step_length/line_search.py b/stride/optimisation/step_length/line_search.py index 65a47d9..84831ee 100644 --- a/stride/optimisation/step_length/line_search.py +++ b/stride/optimisation/step_length/line_search.py @@ -8,12 +8,12 @@ class LineSearch: @abstractmethod - async def init_search(self, variable, direction, **kwargs): + def init_search(self, variable, direction, **kwargs): pass @abstractmethod - async def next_step(self, variable, direction, **kwargs): + def next_step(self, variable, direction, **kwargs): pass - async def forward_test(self, variable, direction, **kwargs): + def forward_test(self, variable, direction, **kwargs): pass diff --git a/stride/physics/common/devito.py b/stride/physics/common/devito.py index a939098..a5e688d 100644 --- a/stride/physics/common/devito.py +++ b/stride/physics/common/devito.py @@ -906,14 +906,14 @@ def set_operator(self, op, **kwargs): default_config = { 'name': self.name, 'subs': subs, - 'opt': 'advanced-fsg', + 'opt': 'advanced', } elif platform == 'cpu-icc': default_config = { 'name': self.name, 'subs': subs, - 'opt': 'advanced-fsg', + 'opt': 'advanced', 'compiler': 'icc', } @@ -921,7 +921,7 @@ def set_operator(self, op, **kwargs): default_config = { 'name': self.name, 'subs': subs, - 'opt': 'advanced-fsg', + 'opt': 'advanced', 'compiler': 'nvc', } @@ -929,7 +929,7 @@ def set_operator(self, op, **kwargs): default_config = { 'name': self.name, 'subs': subs, - 'opt': 'advanced-fsg', + 'opt': 'advanced', 'autotuning': 'off', 'compiler': 'nvc', 'language': 'openacc', @@ -940,7 +940,7 @@ def set_operator(self, op, **kwargs): default_config = { 'name': self.name, 'subs': subs, - 'opt': 'advanced-fsg', + 'opt': 'advanced', 'compiler': 'cuda', 'language': 'cuda', 'platform': 'nvidiaX', @@ -950,7 +950,7 @@ def set_operator(self, op, **kwargs): default_config = { 'name': self.name, 'subs': subs, - 'opt': 'advanced-fsg', + 'opt': 'advanced', 'compiler': 'hip', 'language': 'hip', 'platform': 'amdgpuX', @@ -964,7 +964,7 @@ def set_operator(self, op, **kwargs): context = {'log-level': 'DEBUG' if log_level in ['perf', 'debug'] else 'INFO'} compiler_config = {} for key, value in default_config.items(): - if key in devito.configuration: + if key in devito.configuration and key != 'opt': context[key] = value else: compiler_config[key] = value diff --git a/stride/physics/iso_acoustic/devito.py b/stride/physics/iso_acoustic/devito.py index 554e631..10b480b 100644 --- a/stride/physics/iso_acoustic/devito.py +++ b/stride/physics/iso_acoustic/devito.py @@ -220,7 +220,7 @@ def subdomains(self): # forward - async def before_forward(self, wavelets, vp, rho=None, alpha=None, **kwargs): + def before_forward(self, wavelets, vp, rho=None, alpha=None, **kwargs): """ Prepare the problem type to run the state or forward problem. @@ -318,6 +318,7 @@ async def before_forward(self, wavelets, vp, rho=None, alpha=None, **kwargs): save_wavefield |= alpha.needs_grad platform = kwargs.get('platform', 'cpu') + stream_wavefield = kwargs.pop('stream_wavefield', True) is_nvidia = platform is not None and 'nvidia' in platform is_nvc = platform is not None and (is_nvidia or 'nvc' in platform) @@ -365,7 +366,10 @@ async def before_forward(self, wavelets, vp, rho=None, alpha=None, **kwargs): # Define the saving of the wavefield if save_wavefield is True: space_order = None if self._needs_grad(rho, alpha) else 0 - layers = devito.HostDevice if is_nvidia else devito.NoLayers + if stream_wavefield: + layers = devito.HostDevice if is_nvidia else devito.NoLayers + else: + layers = devito.Device if is_nvidia else devito.NoLayers p_saved = self.dev_grid.undersampled_time_function('p_saved', time_bounds=time_bounds, factor=self.undersampling_factor, @@ -470,9 +474,6 @@ async def before_forward(self, wavelets, vp, rho=None, alpha=None, **kwargs): # Set geometry and wavelet wavelets = wavelets.data - # if fw3d_mode: - # wavelets[:, 1:] = wavelets[:, :-1] - if diff_source: wavelets = np.gradient(wavelets, self.time.step, axis=-1) @@ -486,7 +487,7 @@ async def before_forward(self, wavelets, vp, rho=None, alpha=None, **kwargs): self.dev_grid.vars.src.coordinates.data[:] = shot.source_coordinates self.dev_grid.vars.rec.coordinates.data[:] = shot.receiver_coordinates - async def run_forward(self, wavelets, vp, rho=None, alpha=None, **kwargs): + def run_forward(self, wavelets, vp, rho=None, alpha=None, **kwargs): """ Run the state or forward problem. @@ -554,7 +555,7 @@ async def run_forward(self, wavelets, vp, rho=None, alpha=None, **kwargs): **functions, **devito_args) - async def after_forward(self, wavelets, vp, rho=None, alpha=None, **kwargs): + def after_forward(self, wavelets, vp, rho=None, alpha=None, **kwargs): """ Clean up after the state run and retrieve the time traces. @@ -660,7 +661,7 @@ def _rm_tmpdir(): # adjoint - async def before_adjoint(self, adjoint_source, wavelets, vp, rho=None, alpha=None, **kwargs): + def before_adjoint(self, adjoint_source, wavelets, vp, rho=None, alpha=None, **kwargs): """ Prepare the problem type to run the adjoint problem. @@ -727,7 +728,7 @@ async def before_adjoint(self, adjoint_source, wavelets, vp, rho=None, alpha=Non src_term = [] # Define gradient - gradient_update = await self.prepare_grad(wavelets, vp, rho, alpha, **kwargs) + gradient_update = self.prepare_grad(wavelets, vp, rho, alpha, **kwargs) # Maybe save wavefield dump_adjoint_wavefield = kwargs.pop('dump_adjoint_wavefield', False) @@ -780,7 +781,7 @@ async def before_adjoint(self, adjoint_source, wavelets, vp, rho=None, alpha=Non self.dev_grid.vars.src.data_with_halo.fill(0.) self.dev_grid.vars.p_a.data_with_halo.fill(0.) self.boundary.clear() - await self.init_grad(wavelets, vp, rho, alpha, **kwargs) + self.init_grad(wavelets, vp, rho, alpha, **kwargs) # Set wavefield if necessary cache_forward = kwargs.pop('cache_forward', False) @@ -827,9 +828,6 @@ async def before_adjoint(self, adjoint_source, wavelets, vp, rho=None, alpha=Non # Set geometry and adjoint source adjoint_source = adjoint_source.data - # if fw3d_mode: - # adjoint_source[:, 1:] = adjoint_source[:, :-1] - window = scipy.signal.get_window(('tukey', 0.001), time_bounds[1]-time_bounds[0], False) window = np.pad(window, ((time_bounds[0], self.time.num-time_bounds[1]),), mode='constant', constant_values=0.) window = window.reshape((self.time.num, 1)) @@ -840,7 +838,7 @@ async def before_adjoint(self, adjoint_source, wavelets, vp, rho=None, alpha=Non self.dev_grid.vars.src.coordinates.data[:] = shot.source_coordinates self.dev_grid.vars.rec.coordinates.data[:] = shot.receiver_coordinates - async def run_adjoint(self, adjoint_source, wavelets, vp, rho=None, alpha=None, **kwargs): + def run_adjoint(self, adjoint_source, wavelets, vp, rho=None, alpha=None, **kwargs): """ Run the adjoint problem. @@ -886,7 +884,7 @@ async def run_adjoint(self, adjoint_source, wavelets, vp, rho=None, alpha=None, **functions, **devito_args) - async def after_adjoint(self, adjoint_source, wavelets, vp, rho=None, alpha=None, **kwargs): + def after_adjoint(self, adjoint_source, wavelets, vp, rho=None, alpha=None, **kwargs): """ Clean up after the adjoint run and retrieve the time gradients (if needed). @@ -922,7 +920,7 @@ async def after_adjoint(self, adjoint_source, wavelets, vp, rho=None, alpha=None if dump_adjoint_wavefield and dump_wavefield_id == shot.id: self.logger.perf('(ShotID %d) Dumping adjoint wavefield' % problem.shot_id) - iteration = kwargs.pop('iteration', None) + iteration = kwargs.get('iteration', None) version = iteration.abs_id+1 if iteration is not None else 0 p_dump_data = np.asarray(self.dev_grid.vars.p_a_dump.data, dtype=np.float32) p_dump = StructuredData(name='adjoint_wavefield-Shot%05d' % shot.id, @@ -945,11 +943,11 @@ async def after_adjoint(self, adjoint_source, wavelets, vp, rho=None, alpha=None self.dev_grid.deallocate('buoy') self.dev_grid.deallocate('alpha', collect=True) - return await self.get_grad(wavelets, vp, rho, alpha, **kwargs) + return self.get_grad(wavelets, vp, rho, alpha, **kwargs) # gradients - async def prepare_grad_vp(self, vp, **kwargs): + def prepare_grad_vp(self, vp, **kwargs): """ Prepare the problem type to calculate the gradients wrt Vp. @@ -990,7 +988,7 @@ async def prepare_grad_vp(self, vp, **kwargs): return p_dt_update + (grad_update, prec_update) - async def init_grad_vp(self, vp, **kwargs): + def init_grad_vp(self, vp, **kwargs): """ Initialise buffers in the problem type to calculate the gradients wrt Vp. @@ -1009,7 +1007,7 @@ async def init_grad_vp(self, vp, **kwargs): prec = self.dev_grid.function('prec_vp') prec.data_with_halo.fill(0.) - async def get_grad_vp(self, vp, **kwargs): + def get_grad_vp(self, vp, **kwargs): """ Retrieve the gradients calculated wrt to the input. @@ -1032,7 +1030,20 @@ async def get_grad_vp(self, vp, **kwargs): variable_prec = self.dev_grid.vars.prec_vp variable_prec = np.asarray(variable_prec.data[self.space.inner], dtype=np.float32) - variable_grad *= -2 / vp.data**3 + is_slowness = False + if vp.transform is not None: + # try to figure out if the user wants slowness by testing the provided transform + try: + is_slowness = 1/vp.transform(10) == 10 + except Exception: + pass + + if is_slowness: + variable_grad *= +1 / vp.data + variable_prec *= (1 / vp.data)**2 + else: + variable_grad *= -2 / vp.data**3 + variable_prec *= (2 / vp.data)**3**2 deallocate = kwargs.pop('deallocate', False) if deallocate: @@ -1066,7 +1077,7 @@ async def get_grad_vp(self, vp, **kwargs): return grad - async def prepare_grad_rho(self, rho, **kwargs): + def prepare_grad_rho(self, rho, **kwargs): """ Prepare the problem type to calculate the gradients wrt rho. @@ -1103,7 +1114,7 @@ async def prepare_grad_rho(self, rho, **kwargs): return grad_term_update + (grad_update, prec_update) - async def init_grad_rho(self, rho, **kwargs): + def init_grad_rho(self, rho, **kwargs): """ Initialise buffers in the problem type to calculate the gradients wrt rho. @@ -1122,7 +1133,7 @@ async def init_grad_rho(self, rho, **kwargs): prec = self.dev_grid.function('prec_rho') prec.data_with_halo.fill(0.) - async def get_grad_rho(self, rho, **kwargs): + def get_grad_rho(self, rho, **kwargs): """ Retrieve the gradients calculated wrt to rho. @@ -1365,8 +1376,14 @@ def _check_conditions(self, wavelets, vp, rho=None, alpha=None, def _stencil(self, field, wavelets, vp, rho=None, alpha=None, direction='forward', save_wavefield=False, **kwargs): + stencils = [] + # Prepare medium functions vp_fun, vp2_fun, inv_vp2_fun, rho_fun, buoy_fun, alpha_fun = self._medium_functions(vp, rho, alpha, **kwargs) + if rho is not None: + rho_constant = np.isclose(np.min(rho.extended_data), np.max(rho.extended_data)) + else: + rho_constant = False # Forward or backward forward = direction == 'forward' @@ -1381,30 +1398,51 @@ def _stencil(self, field, wavelets, vp, rho=None, alpha=None, direction='forward **kwargs) # Get the spatial FD - laplacian = self.dev_grid.function('laplacian', - coefficients='symbolic' if self.drp else 'standard') - laplacian_update = self._laplacian(field, laplacian, vp_fun, vp2_fun, inv_vp2_fun, - rho=rho_fun, buoy=buoy_fun, alpha=alpha_fun, - **kwargs) - if self.kernel == 'OT2': - laplacian_term = self._diff_op(laplacian_update, + # get the subs + if self.drp: + extra_functions = () + subs = self._symbolic_coefficients(field, *extra_functions) + else: + subs = None + + laplacian_term = self._diff_op(field, vp_fun, vp2_fun, inv_vp2_fun, rho=rho_fun, buoy=buoy_fun, alpha=alpha_fun, + rho_constant=rho_constant, **kwargs) else: - laplacian_term = self._diff_op(laplacian, - vp_fun, vp2_fun, inv_vp2_fun, - rho=rho_fun, buoy=buoy_fun, alpha=alpha_fun, - **kwargs) + laplacian_2 = self.dev_grid.function('laplacian', + coefficients='symbolic' if self.drp else 'standard') - # Get the subs - if self.drp: - extra_functions = () - subs = self._symbolic_coefficients(field, laplacian, - *extra_functions) - else: - subs = None + # get the subs + if self.drp: + extra_functions = () + subs = self._symbolic_coefficients(field, laplacian_2, *extra_functions) + else: + subs = None + + # first laplacian application - L2 + laplacian_term_2 = self._diff_op(field, + vp_fun, vp2_fun, inv_vp2_fun, + rho=rho_fun, buoy=buoy_fun, alpha=alpha_fun, + rho_constant=rho_constant, + **kwargs) + + stencil_laplacian = devito.Eq(laplacian_2, laplacian_term_2, + subdomain=abox, + coefficients=subs) + stencils.append(stencil_laplacian) + + # second laplacian application - L4 + laplacian_term_4 = self._diff_op(laplacian_2, + vp_fun, vp2_fun, inv_vp2_fun, + rho=rho_fun, buoy=buoy_fun, alpha=alpha_fun, + rho_constant=rho_constant, + **kwargs) + + # final term + laplacian_term = self._laplacian(laplacian_2, laplacian_term_4) # Get the attenuation term if alpha_fun is not None and self.attenuation_power is not None: @@ -1421,7 +1459,7 @@ def _stencil(self, field, wavelets, vp, rho=None, alpha=None, direction='forward attenuation_term = 0 # Set up the boundary - boundary_field = laplacian if self.kernel != 'OT2' and 'PML' in self.boundary_type else field + boundary_field = laplacian_2 if self.kernel != 'OT2' and 'PML' in self.boundary_type else field boundary_term, eq_before, eq_after = self.boundary.apply(boundary_field, vp.extended_data, velocity_fun=vp_fun, direction=direction, subs=subs, @@ -1454,14 +1492,6 @@ def _stencil(self, field, wavelets, vp, rho=None, alpha=None, direction='forward - vp2_fun*sub_exprs, u_next) # Time-stepping stencil - stencils = [] - - if self.kernel != 'OT2': - stencil_laplacian = devito.Eq(laplacian, laplacian_update, - subdomain=abox, - coefficients=subs) - stencils.append(stencil_laplacian) - if 'hybrid' in self.boundary_type: domain = abox else: @@ -1510,7 +1540,7 @@ def _medium_functions(self, vp, rho=None, alpha=None, **kwargs): return vp_fun, vp2_fun, inv_vp2_fun, rho_fun, buoy_fun, alpha_fun - def _laplacian(self, field, laplacian, vp, vp2, inv_vp2, **kwargs): + def _laplacian(self, laplacian_2, laplacian_4, **kwargs): if self.kernel not in ['OT2', 'OT4']: raise ValueError("Unrecognized kernel") @@ -1518,22 +1548,33 @@ def _laplacian(self, field, laplacian, vp, vp2, inv_vp2, **kwargs): bi_harmonic = 0 else: - bi_harmonic = self.time.step**2/12 * self._diff_op(field, - vp, vp2, inv_vp2, - **kwargs) + bi_harmonic = self.time.step**2/12 * laplacian_4 - laplacian_update = field + bi_harmonic + laplacian_update = laplacian_2 + bi_harmonic return laplacian_update def _diff_op(self, field, vp, vp2, inv_vp2, **kwargs): rho = kwargs.pop('rho', None) buoy = kwargs.pop('buoy', None) + rho_constant = kwargs.pop('rho_constant', False) if buoy is None: return vp2 * field.laplace else: - return vp2 * rho * devito.div(buoy * devito.grad(field, shift=+0.5), shift=-0.5) + if rho_constant: + return vp2 * self._div_op(self._grad_op(field, shift=+1), shift=-1) + else: + return vp2 * rho * self._div_op(self._mul_buoy(buoy, self._grad_op(field, shift=+1)), shift=-1) + + def _grad_op(self, f, shift=+1): + return devito.grad(f, shift=shift * 0.5) + + def _div_op(self, fs, shift=-1): + return devito.div(fs, shift=shift * 0.5) + + def _mul_buoy(self, buoy, fs): + return buoy * fs def _subdomains(self, *args, **kwargs): problem = kwargs.get('problem') diff --git a/stride/physics/iso_elastic/devito.py b/stride/physics/iso_elastic/devito.py index bf8c8f0..157a5af 100644 --- a/stride/physics/iso_elastic/devito.py +++ b/stride/physics/iso_elastic/devito.py @@ -81,7 +81,7 @@ def clear_operators(self): # forward - async def before_forward(self, wavelets, vp, vs, rho, **kwargs): + def before_forward(self, wavelets, vp, vs, rho, **kwargs): """ Prepare the problem type to run the state or forward problem. @@ -204,7 +204,7 @@ async def before_forward(self, wavelets, vp, vs, rho, **kwargs): self.dev_grid.vars.src.coordinates.data[:] = shot.source_coordinates self.dev_grid.vars.rec_tau.coordinates.data[:] = shot.receiver_coordinates - async def run_forward(self, wavelets, vp, vs, rho, **kwargs): + def run_forward(self, wavelets, vp, vs, rho, **kwargs): """ Run the state or forward problem. @@ -234,7 +234,7 @@ async def run_forward(self, wavelets, vp, vs, rho, **kwargs): **functions, **kwargs.pop('devito_args', {})) - async def after_forward(self, wavelets, vp, vs, rho, **kwargs): + def after_forward(self, wavelets, vp, vs, rho, **kwargs): """ Clean up after the state run and retrieve the time traces. @@ -277,19 +277,19 @@ async def after_forward(self, wavelets, vp, vs, rho, **kwargs): # adjoint - async def before_adjoint(self, adjoint_source, wavelets, vp, rho=None, alpha=None, **kwargs): + def before_adjoint(self, adjoint_source, wavelets, vp, rho=None, alpha=None, **kwargs): """ Not implemented """ pass - async def run_adjoint(self, adjoint_source, wavelets, vp, rho=None, alpha=None, **kwargs): + def run_adjoint(self, adjoint_source, wavelets, vp, rho=None, alpha=None, **kwargs): """ Not implemented """ pass - async def after_adjoint(self, **kwargs): + def after_adjoint(self, **kwargs): """ Not implemented """ @@ -297,19 +297,19 @@ async def after_adjoint(self, **kwargs): # gradients - async def prepare_grad_vp(self, **kwargs): + def prepare_grad_vp(self, **kwargs): """ Not implemented """ pass - async def init_grad_vp(self, **kwargs): + def init_grad_vp(self, **kwargs): """ Not implemented """ pass - async def get_grad_vp(self, **kwargs): + def get_grad_vp(self, **kwargs): """ Not implemented """ diff --git a/stride/physics/marmottant/devito.py b/stride/physics/marmottant/devito.py index 4eb1955..bbd61b7 100644 --- a/stride/physics/marmottant/devito.py +++ b/stride/physics/marmottant/devito.py @@ -101,7 +101,7 @@ def clear_operators(self): # forward - async def before_forward(self, r_0, x_0=None, + def before_forward(self, r_0, x_0=None, vp=1540., rho=997, sigma=0.073, mu=0.002, p_0=101325, p=0., kappa=1.07, kappa_s=5E-9, chi=0.4, r_buckle=None, r_break=None, **kwargs): @@ -159,7 +159,7 @@ async def before_forward(self, r_0, x_0=None, self.state_operator.set_operator(init_terms + stencil, **kwargs) self.state_operator.compile() - async def run_forward(self, *args, **kwargs): + def run_forward(self, *args, **kwargs): """ Run the state or forward problem. @@ -178,7 +178,7 @@ async def run_forward(self, *args, **kwargs): **functions, **kwargs.pop('devito_args', {})) - async def after_forward(self, *args, **kwargs): + def after_forward(self, *args, **kwargs): """ Clean up after the state run and retrieve the time traces. diff --git a/stride/physics/problem_type.py b/stride/physics/problem_type.py index 26ffed4..92cd7bf 100644 --- a/stride/physics/problem_type.py +++ b/stride/physics/problem_type.py @@ -1,4 +1,3 @@ - from abc import ABC import mosaic @@ -6,7 +5,6 @@ from ..core import Operator from ..problem.base import Gridded - __all__ = ['ProblemTypeBase'] @@ -56,7 +54,7 @@ def __init__(self, **kwargs): Gridded.__init__(self, **kwargs) Operator.__init__(self, **kwargs) - async def forward(self, *args, **kwargs): + def forward(self, *args, **kwargs): """ Run the state or forward problem. @@ -73,18 +71,18 @@ async def forward(self, *args, **kwargs): pre_str = '(ShotID %d) ' % problem.shot_id self.logger.perf('%sPreparing to run state for shot' % pre_str) - await self.before_forward(*args, **kwargs) + self.before_forward(*args, **kwargs) self.logger.perf('%sRunning state equation for shot' % pre_str) - await self.run_forward(*args, **kwargs) + self.run_forward(*args, **kwargs) self.logger.perf('%sCompleting state equation run for shot' % pre_str) - output = await self.after_forward(*args, **kwargs) + output = self.after_forward(*args, **kwargs) self.logger.perf('%sCompleted state equation run for shot' % pre_str) return output - async def adjoint(self, *args, **kwargs): + def adjoint(self, *args, **kwargs): """ Run the adjoint problem. @@ -101,18 +99,18 @@ async def adjoint(self, *args, **kwargs): pre_str = '(ShotID %d) ' % problem.shot_id self.logger.perf('%sPreparing to run adjoint for shot' % pre_str) - await self.before_adjoint(*args, **kwargs) + self.before_adjoint(*args, **kwargs) self.logger.perf('%sRunning adjoint equation for shot' % pre_str) - await self.run_adjoint(*args, **kwargs) + self.run_adjoint(*args, **kwargs) self.logger.perf('%sCompleting adjoint equation run for shot' % pre_str) - output = await self.after_adjoint(*args, **kwargs) + output = self.after_adjoint(*args, **kwargs) self.logger.perf('%sCompleted adjoint equation run for shot' % pre_str) return output - async def before_forward(self, *args, **kwargs): + def before_forward(self, *args, **kwargs): """ Prepare the problem type to run the state or forward problem. @@ -126,7 +124,7 @@ async def before_forward(self, *args, **kwargs): raise NotImplementedError('before_forward has not been implemented ' 'for objects of type %s' % self.__class__.__name__) - async def run_forward(self, *args, **kwargs): + def run_forward(self, *args, **kwargs): """ Run the state or forward problem. @@ -140,7 +138,7 @@ async def run_forward(self, *args, **kwargs): raise NotImplementedError('run_forward has not been implemented ' 'for objects of type %s' % self.__class__.__name__) - async def after_forward(self, *args, **kwargs): + def after_forward(self, *args, **kwargs): """ Clean up after the state run and retrieve the outputs. @@ -156,7 +154,7 @@ async def after_forward(self, *args, **kwargs): raise NotImplementedError('after_forward has not been implemented ' 'for objects of type %s' % self.__class__.__name__) - async def before_adjoint(self, *args, **kwargs): + def before_adjoint(self, *args, **kwargs): """ Prepare the problem type to run the adjoint problem. @@ -170,7 +168,7 @@ async def before_adjoint(self, *args, **kwargs): raise NotImplementedError('before_adjoint has not been implemented ' 'for objects of type %s' % self.__class__.__name__) - async def run_adjoint(self, *args, **kwargs): + def run_adjoint(self, *args, **kwargs): """ Run the adjoint problem. @@ -184,7 +182,7 @@ async def run_adjoint(self, *args, **kwargs): raise NotImplementedError('run_adjoint has not been implemented ' 'for objects of type %s' % self.__class__.__name__) - async def after_adjoint(self, *args, **kwargs): + def after_adjoint(self, *args, **kwargs): """ Clean up after the adjoint run and retrieve the gradients (if needed). @@ -200,7 +198,7 @@ async def after_adjoint(self, *args, **kwargs): raise NotImplementedError('after_adjoint has not been implemented ' 'for objects of type %s' % self.__class__.__name__) - async def prepare_grad(self, *wrt, **kwargs): + def prepare_grad(self, *wrt, **kwargs): """ Prepare the problem type to calculate the gradients wrt the inputs. @@ -229,7 +227,7 @@ async def prepare_grad(self, *wrt, **kwargs): if method is None: raise ValueError('Variable %s not implemented' % variable.name) - update = await method(variable, wrt=wrt, **kwargs) + update = method(variable, wrt=wrt, **kwargs) if not isinstance(update, tuple): update = (update,) @@ -238,7 +236,7 @@ async def prepare_grad(self, *wrt, **kwargs): return gradient_update - async def init_grad(self, *wrt, **kwargs): + def init_grad(self, *wrt, **kwargs): """ Initialise buffers in the problem type to calculate the gradients wrt the inputs. @@ -263,9 +261,9 @@ async def init_grad(self, *wrt, **kwargs): if method is None: raise ValueError('Variable %s not implemented' % variable.name) - await method(variable, wrt=wrt, **kwargs) + method(variable, wrt=wrt, **kwargs) - async def get_grad(self, *wrt, **kwargs): + def get_grad(self, *wrt, **kwargs): """ Retrieve the gradients calculated wrt to the inputs. @@ -294,6 +292,6 @@ async def get_grad(self, *wrt, **kwargs): if method is None: raise ValueError('Variable %s not implemented' % variable.name) - grads.append(await method(variable, wrt=wrt, **kwargs)) + grads.append(method(variable, wrt=wrt, **kwargs)) return tuple(grads) diff --git a/stride/problem/data.py b/stride/problem/data.py index 0b5d9ea..3e024ad 100644 --- a/stride/problem/data.py +++ b/stride/problem/data.py @@ -120,6 +120,10 @@ class StructuredData(Data): """ def __init__(self, **kwargs): + # hacky, but does the trick for now + name = kwargs.get('name', None) + if name is not None and 'vp' in name: + kwargs['transform'] = kwargs.pop('transform', lambda x: 1 / x) super().__init__(**kwargs) shape = kwargs.pop('shape', None) @@ -365,7 +369,7 @@ def process_grad(self, global_prec=True, **kwargs): self.grad.apply_prec(**kwargs) return self.grad - def apply_prec(self, prec_scale=4.0, prec_smooth=1.0, prec_op=None, prec=None, **kwargs): + def apply_prec(self, prec_scale=4.0, prec_smooth=None, prec_op=None, prec=None, **kwargs): """ Apply a pre-conditioner to the current field. @@ -400,8 +404,7 @@ def apply_prec(self, prec_scale=4.0, prec_smooth=1.0, prec_op=None, prec=None, * prec.data[:] = prec_op(prec.data) non_zero = np.abs(prec.data) > 0. - prec.data[non_zero] = 1/prec.data[non_zero] - self.data[non_zero] *= prec.data[non_zero] + self.data[non_zero] *= 1/prec.data[non_zero] return self @@ -593,6 +596,15 @@ def __truediv__(self, other): return res + def __rtruediv__(self, other): + res, other_data = self._prepare_op(other) + res.extended_data[:] = res.extended_data.__rtruediv__(other_data) + + self._op_grad(res, other, '__rtruediv__') + self._op_prec(res, other, '__rtruediv__') + + return res + def __floordiv__(self, other): res, other_data = self._prepare_op(other) res.extended_data[:] = res.extended_data.__floordiv__(other_data) @@ -602,6 +614,15 @@ def __floordiv__(self, other): return res + def __rfloordiv__(self, other): + res, other_data = self._prepare_op(other) + res.extended_data[:] = res.extended_data.__rfloordiv__(other_data) + + self._op_grad(res, other, '__rfloordiv__') + self._op_prec(res, other, '__rfloordiv__') + + return res + def __iadd__(self, other): res = self other_data = self._prepare_other(other) @@ -661,8 +682,6 @@ def __ifloordiv__(self, other): __radd__ = __add__ __rsub__ = __sub__ __rmul__ = __mul__ - __rtruediv__ = __truediv__ - __rfloordiv__ = __floordiv__ def __get_desc__(self, **kwargs): if self._data is None: From 7f3c94def69ca5e36ce228bd798023c80f90387c Mon Sep 17 00:00:00 2001 From: Carlos Cueto Date: Fri, 19 Jul 2024 09:22:41 +0100 Subject: [PATCH 8/9] Remove deprecation warning --- stride/physics/common/import_devito.py | 2 +- stride/physics/iso_acoustic/devito.py | 6 +----- 2 files changed, 2 insertions(+), 6 deletions(-) diff --git a/stride/physics/common/import_devito.py b/stride/physics/common/import_devito.py index f29d95f..043740c 100644 --- a/stride/physics/common/import_devito.py +++ b/stride/physics/common/import_devito.py @@ -7,7 +7,7 @@ from devitopro import * # noqa: F401 from devitopro.types.enriched import (DiskHostDevice, DiskHost, DiskDevice, # noqa: F401 HostDevice, Host, Device, NoLayers) - from devitopro.types.dynamic import CompressedTimeFunction + from devitopro.types.compressed import CompressedTimeFunction pro_available = True except ImportError: diff --git a/stride/physics/iso_acoustic/devito.py b/stride/physics/iso_acoustic/devito.py index 10b480b..6a35753 100644 --- a/stride/physics/iso_acoustic/devito.py +++ b/stride/physics/iso_acoustic/devito.py @@ -3,7 +3,6 @@ import glob import shutil import tempfile -import warnings import numpy as np import scipy.signal @@ -21,9 +20,6 @@ __all__ = ['IsoAcousticDevito'] -warnings.filterwarnings('ignore', category=np.VisibleDeprecationWarning) - - @mosaic.tessera class IsoAcousticDevito(ProblemTypeBase): """ @@ -1034,7 +1030,7 @@ def get_grad_vp(self, vp, **kwargs): if vp.transform is not None: # try to figure out if the user wants slowness by testing the provided transform try: - is_slowness = 1/vp.transform(10) == 10 + is_slowness = np.isclose(1/vp.transform(10), 10) except Exception: pass From a0818bcf8c4d614bd1e5aada3416139e3e85dfef Mon Sep 17 00:00:00 2001 From: Carlos Cueto Date: Fri, 19 Jul 2024 11:06:08 +0100 Subject: [PATCH 9/9] Rebalance update range --- .../optimisation/pipelines/steps/mask_field.py | 2 ++ .../optimisation/pipelines/steps/norm_field.py | 17 ++++++++++++++++- 2 files changed, 18 insertions(+), 1 deletion(-) diff --git a/stride/optimisation/pipelines/steps/mask_field.py b/stride/optimisation/pipelines/steps/mask_field.py index 79cd04c..695ac96 100644 --- a/stride/optimisation/pipelines/steps/mask_field.py +++ b/stride/optimisation/pipelines/steps/mask_field.py @@ -9,6 +9,8 @@ def _rampoff_mask(shape, ramp_size): mask = np.ones(shape, dtype=np.float32) for dim_i in range(len(shape)): + if 2*ramp_size > shape[dim_i]: + continue for index in range(ramp_size): pos = np.abs((ramp_size - index - 1) / float(ramp_size - 1)) val = 1 - np.cos(np.pi / 2 * (1 - pos)) diff --git a/stride/optimisation/pipelines/steps/norm_field.py b/stride/optimisation/pipelines/steps/norm_field.py index f963fd7..84e5577 100644 --- a/stride/optimisation/pipelines/steps/norm_field.py +++ b/stride/optimisation/pipelines/steps/norm_field.py @@ -18,17 +18,32 @@ def __init__(self, **kwargs): super().__init__(**kwargs) self.global_norm = kwargs.pop('global_norm', False) + self.norm_guess_change = kwargs.pop('norm_guess_change', 0.5) self.norm_value = None def forward(self, field, **kwargs): + variable = kwargs.pop('variable', None) global_norm = kwargs.pop('global_norm', self.global_norm) + norm_guess_change = kwargs.pop('norm_guess_change', self.norm_guess_change) if self.norm_value is None or not global_norm: self.norm_value = np.max(np.abs(field.extended_data)) + 1e-31 + # work out guess change based on field value + if variable is not None: + min_val = np.min(variable.extended_data) + max_val = np.max(variable.extended_data) + + mid_val = (max_val + min_val) / 2.0 + if variable.transform is not None: + mid_val = variable.transform(mid_val) + var_corr = mid_val * norm_guess_change / 100 + else: + var_corr = 1. + out_field = field.alike(name=name_from_op_name(self, field)) out_field.extended_data[:] = field.extended_data - out_field.extended_data[:] *= 1. / self.norm_value + out_field.extended_data[:] *= var_corr / self.norm_value return out_field