Implement snapshotting for the acoustic wave equation #2474

malfarhan7 · 2024-10-26T09:26:05Z

Implement snapshotting to save snapshots of the forward wavefield used to compute the gradient to reduce memory usage.

mloubout

Thank you for the contribution!

I have left some comments as it needs some changes to be mergeable. Some kind of test also needs to be added so that it is maintainable.

mloubout · 2024-10-28T16:49:10Z

examples/seismic/acoustic/operators.py

+    # Build operator equations
+    equations = eqn + src_term + rec_term
+
+    if factor:


This needs to be wrapped into a utility function as it's duplicated below

I've created a function to construct usnaps.

mloubout · 2024-10-28T16:49:34Z

examples/seismic/acoustic/operators.py

+        nsnaps = (geometry.nt + factor - 1) // factor
+        time_subsampled = ConditionalDimension(
+            't_sub', parent=model.grid.time_dim, factor=factor)
+        usnaps = TimeFunction(name='usnaps', grid=model.grid,


You still have u with full time saved line 135 you can't have both

mloubout · 2024-10-28T16:50:16Z

examples/seismic/acoustic/operators.py

-                    name='Forward', **kwargs)
+    op = Operator(equations, subs=model.spacing_map, name='Forward', **kwargs)
+    if usnaps is not None:
+        return op, usnaps


No the operator build cannot return objects like that. This is an abstract operator with placeholders that might not be correct for runtime.

Fixed. The operator build only returns op now.

mloubout · 2024-10-28T16:50:38Z

examples/seismic/acoustic/operators.py

+
+    if factor is not None:
+        # Condition to apply gradient update only at snapshot times
+        condition = Eq(time % factor, 0)


No you don't need that usnap already contains the conditon

mloubout · 2024-10-28T16:51:04Z

examples/seismic/acoustic/operators.py

+        u = TimeFunction(name='u', grid=model.grid,
+                         save=geometry.nt if save else None,
+                         time_order=2, space_order=space_order)
+        if kernel == 'OT2':


Unnecessary duplicate, u contains the information you should not need separate cases for gradient_update

Fixed. No cases are used.

mloubout · 2024-10-28T16:51:24Z

examples/seismic/acoustic/wavesolver.py

@@ -1,6 +1,6 @@
 from devito import Function, TimeFunction, DevitoCheckpoint, CheckpointOperator, Revolver
 from devito.tools import memoized_meth
-from examples.seismic.acoustic.operators import (
+from devitofwi.devito.acoustic.operators import (


You are right. I did not catch it.

mloubout · 2024-10-28T16:51:43Z

examples/seismic/acoustic/wavesolver.py

@@ -108,12 +111,24 @@ def forward(self, src=None, rec=None, u=None, model=None, save=None, **kwargs):
        model = model or self.model
        # Pick vp from model unless explicitly provided
        kwargs.update(model.physical_params(**kwargs))
+        # Get the operator
+        op_fwd = self.op_fwd(save=save, factor=factor)
+        # Prepare parameters for operator apply


Don't know what this is for.

mloubout · 2024-10-28T16:52:12Z

examples/seismic/acoustic/wavesolver.py

-                                          dt=kwargs.pop('dt', self.dt), **kwargs)
-
-        return rec, u, summary
+        if factor:


No, usnap needs to be create here like u then passed as argument

Fixed. usnaps is created now.

mloubout · 2024-10-28T16:52:22Z

examples/seismic/acoustic/wavesolver.py

+            op_args['usnaps'] = usnaps
+            summary = op.apply(**op_args)
+
+        else:


shouldn't need if else only kwargs

mloubout · 2024-10-28T16:52:50Z

examples/seismic/acoustic/wavesolver.py

@@ -209,8 +236,17 @@ def jacobian_adjoint(self, rec, u, src=None, v=None, grad=None, model=None,
            wrp.apply_forward()
            summary = wrp.apply_reverse()
        else:
-            summary = self.op_grad().apply(rec=rec, grad=grad, v=v, u=u, dt=dt,
-                                           **kwargs)
+            if factor is not None:


Again, not needed, input u should contain all metada needed

…apshot_acoustic

review-notebook-app · 2024-11-05T13:55:08Z

Check out this pull request on

See visual diffs & provide feedback on Jupyter Notebooks.

Powered by ReviewNB

malfarhan7 · 2024-11-29T10:08:51Z

Hi Mathias, thank you for your feedback. I have reviewed and cleaned the code to the best of my understanding. I have included a notebook to compare computing the FWI gradient with and without snapshotting and two scripts to calculate the memory usage of both methods. After updating the code, the memory usage for calculating the gradient with snapshotting is more than twice that of the older code version. This reduced memory usage (I guess) because I was passing 'usnaps' with the operator (which is not good practice). I am wondering, is it possible to improve the code more to reduce the memory usage?

FabioLuporini · 2024-12-09T15:37:05Z

examples/seismic/acoustic/operators.py

                     time_order=2, space_order=space_order)
    rec = geometry.rec

    s = model.grid.stepping_dim.spacing
    eqn = iso_stencil(v, model, kernel, forward=False)
-
+    


revert change, pep8 violation

FabioLuporini · 2024-12-09T15:37:20Z

examples/seismic/acoustic/operators.py

    receivers = rec.inject(field=v.backward, expr=rec * s**2 / m)
-
+    


same as above

FabioLuporini · 2024-12-09T15:38:22Z

examples/seismic/acoustic/operators.py

-    u = TimeFunction(name='u', grid=model.grid, save=geometry.nt if save
-                     else None, time_order=2, space_order=space_order)
-    v = TimeFunction(name='v', grid=model.grid, save=None,
+    if factor: # Apply the imaging condition at the snapshots of the full wavefield


Leave a blank line between the grad = and this if factor

Move the comment inside the body of the if

FabioLuporini · 2024-12-09T15:38:47Z

examples/seismic/acoustic/operators.py

-    v = TimeFunction(name='v', grid=model.grid, save=None,
+    if factor: # Apply the imaging condition at the snapshots of the full wavefield
+        u = create_snapshot_time_function(model, 'u', geometry, space_order, factor)
+    else:# Apply the imaging condition at every time step of the full wavefield


Move the comment inside the body of the else

FabioLuporini · 2024-12-09T15:39:13Z

examples/seismic/acoustic/wavesolver.py

@@ -90,30 +91,38 @@ def forward(self, src=None, rec=None, u=None, model=None, save=None, **kwargs):
            The time-constant velocity.
        save : bool, optional
            Whether or not to save the entire (unrolled) wavefield.
+        factor : int, optional
+        Downsampling factor to save snapshots of the wavefield.


FabioLuporini · 2024-12-09T15:45:38Z

examples/seismic/tutorials/Gradient_full_marmousi.py

+import matplotlib.pyplot as plt
+from scipy.ndimage import gaussian_filter
+import scipy
+from memory_profiler import memory_usage


imports from stdlib at the very top

then blank line

then imports from third parties (eg scipy)

then blank line

then examples imports
then devito imports

FabioLuporini · 2024-12-09T15:48:30Z

examples/seismic/acoustic/operators.py

+    nsnaps = (geometry.nt + factor - 1) // factor
+    time_subsampled = ConditionalDimension('t_sub', 
+                                            parent=model.grid.time_dim, factor=factor)
+    u_ = TimeFunction(name=name, grid=model.grid, 


"usnaps" for homogeneity

FabioLuporini · 2024-12-09T15:50:54Z

examples/seismic/acoustic/operators.py

    """
    m = model.m

    # Create symbols for forward wavefield, source and receivers
    u = TimeFunction(name='u', grid=model.grid,
-                     save=geometry.nt if save else None,
+                     save=geometry.nt if save and factor is None else None,


I'm not a big fan of this composite conditional involving both save and factor, which is also repeated across other modules

Moved the conditional statement outside of the u definition but I do not know if there is a better way to avoid the composite conditional statement.

FabioLuporini · 2024-12-09T15:51:18Z

examples/seismic/acoustic/operators.py

    # Substitute spacing terms to reduce flops
    return Operator(eqn + receivers + [gradient_update], subs=model.spacing_map,
-                    name='Gradient', **kwargs)
+                  name='Gradient', **kwargs)


FabioLuporini · 2024-12-09T15:52:36Z

examples/seismic/acoustic/wavesolver.py

-                                          dt=kwargs.pop('dt', self.dt), **kwargs)
-
-        return rec, u, summary
+        if factor: # Return snapshots of the forward wavefield


since factor is passed down to op_fwd, I don't think we need the extra if factor : .... else: ... here, somehow it should be avoided and/or it's avoidable

Is there a better way to do this, as the code did not run correctly without the condition? I made the return statement conditional so as not to break people's code, so I kept the number of returned objects at three.

malfarhan7 · 2024-12-12T11:02:50Z

Hi Fabio, thank you for your feedback. I have reviewed and cleaned the code to the best of my understanding.

Implement snapshotting for the acoustic wave equation

44694f5

mloubout requested changes Oct 28, 2024

View reviewed changes

mloubout added the examples examples label Oct 28, 2024

malfarhan7 added 6 commits November 5, 2024 10:08

Clean up snapshotting code

02be5b3

Add script to compute memory usage for full gradient.

d909155

Add script to compute memory usage for snapshotted gradient.

c0791d9

Add notebook to compare full gradient with snapshotted gradient.

71a99cc

Add Marm.bin to data folder

46bbf56

Merge branch 'master' of https://github.com/malfarhan7/devito into sn…

0353a78

…apshot_acoustic

FabioLuporini requested changes Dec 9, 2024

View reviewed changes

malfarhan7 added 3 commits December 12, 2024 13:34

Fix PEP8 violations

810ddd3

Fix PEP8 violations

6c1bc75

Add one script for gradient computation

7f12068

		receivers = rec.inject(field=v.backward, expr=rec * s**2 / m)

Implement snapshotting for the acoustic wave equation #2474

Are you sure you want to change the base?

Implement snapshotting for the acoustic wave equation #2474

Conversation

malfarhan7 commented Oct 26, 2024

mloubout left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

malfarhan7 Nov 5, 2024 • edited Loading

Choose a reason for hiding this comment

review-notebook-app bot commented Nov 5, 2024

malfarhan7 commented Nov 29, 2024

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

malfarhan7 commented Dec 12, 2024

malfarhan7 Nov 5, 2024 •

edited

Loading