Minor documentation fixes (correcting internal links)

Author: ar4

docs/example_callback_animation.rst

@@ -9,7 +9,7 @@ In this example, we will create an animation that shows three things side-by-sid
 2. The backward-propagating wavefield (the adjoint wavefield).
 3. The formation of the model gradient during backpropagation.
 
-We set-up the propagation as normal, using a simple single layer model. Next, we define two callback functions to store the forward and backward wavefields and the current value of the gradient with respect to the velocity model as backpropagation progresses. We can use regular Python functions for these, but, to show you how it is done, I will use a callable class for the forward callback. The `ForwardCallback` instance will be called during the forward pass, and the `backward_callback` function will be called during the backward pass. Each receives a `CallbackState <class-deepwave.common.CallbackState>` object when it is called that provides access to the simulation data. We use these to save the relevant tensors to our snapshot storage. The wavefield names are the same as those in the propagator signature, so the current wavefield is `wavefield_0`::
+We set-up the propagation as normal, using a simple single layer model. Next, we define two callback functions to store the forward and backward wavefields and the current value of the gradient with respect to the velocity model as backpropagation progresses. We can use regular Python functions for these, but, to show you how it is done, I will use a callable class for the forward callback. The `ForwardCallback` instance will be called during the forward pass, and the `backward_callback` function will be called during the backward pass. Each receives a :py:class:`CallbackState <deepwave.common.CallbackState>` object when it is called that provides access to the simulation data. We use these to save the relevant tensors to our snapshot storage. The wavefield names are the same as those in the propagator signature, so the current wavefield is `wavefield_0`::
 
         # Storage for snapshots from callbacks
         callback_frequency = 1

docs/example_target_wavefield.rst

@@ -58,7 +58,7 @@ The optimisation loop is then not very dissimilar to earlier examples. We'll set
     for i in range(50):
         optimiser.step(closure)
 
-Lastly, we will save the wave propagation time steps so that we can make them into an animated GIF. The full example code shows two ways of achieving this. The simplest is to use Deepwave's `callback feature <example_callback_animation>`. The other uses the approach discussed in :doc:`the checkpointing example <example_checkpointing>`.
+Lastly, we will save the wave propagation time steps so that we can make them into an animated GIF. The full example code shows two ways of achieving this. The simplest is to use Deepwave's :doc:`callback feature <example_callback_animation>`. The other uses the approach discussed in :doc:`the checkpointing example <example_checkpointing>`.
 
 Using `FFmpeg <https://ffmpeg.org>`_ to join these individual time steps into an animated GIF (although you could use other tools, such as Matplotlib)::
 

src/deepwave/scalar.py

@@ -1529,7 +1529,6 @@ def backward(
         )
 
         if wfc.numel() > 0 and nt > 0:
-            start_t = 0
             forward(
                 v.data_ptr(),
                 ggv.data_ptr(),
@@ -1581,7 +1580,7 @@ def backward(
                 False,
                 v_batched,
                 ggv_batched,
-                start_t,
+                0,
                 pml_y0,
                 pml_y1,
                 pml_x0,
@@ -1743,7 +1742,6 @@ def backward(
         gwfp = -gwfp
 
         if wfc.numel() > 0 and nt > 0 and v.requires_grad:
-            start_t = 0
             backward(
                 v.data_ptr(),
                 ggv.data_ptr(),
@@ -1804,7 +1802,7 @@ def backward(
                 False,
                 v_batched,
                 ggv_batched,
-                start_t,
+                nt,
                 pml_y0,
                 pml_y1,
                 pml_x0,
@@ -1848,6 +1846,9 @@ def backward(
                 None,
                 None,
                 None,
+                None,
+                None,
+                None,
             )
         return (
             ggwfp[s],
@@ -1883,6 +1884,9 @@ def backward(
             None,
             None,
             None,
+            None,
+            None,
+            None,
         )
 
 

tests/test_scalar.py

@@ -395,7 +395,7 @@ def test_unused_source_receiver(
     assert torch.allclose(modelf.grad, modeli.grad)
 
 
-def run_scalarfunc(nt: int = 3, gradgrad: bool = False) -> None:
+def run_scalarfunc(nt: int = 3) -> None:
     """Runs scalar_func for testing purposes."""
     from deepwave.scalar import scalar_func
 
@@ -500,39 +500,38 @@ def run_scalarfunc(nt: int = 3, gradgrad: bool = False) -> None:
             1,
         ),
     )
-    if gradgrad:
-        torch.autograd.gradgradcheck(
-            scalar_func,
-            (
-                c,
-                source_amplitudes,
-                wfc,
-                wfp,
-                psiy,
-                psix,
-                zetay,
-                zetax,
-                ay,
-                ax,
-                by,
-                bx,
-                dbydy,
-                dbxdx,
-                sources_i,
-                receivers_i,
-                dy,
-                dx,
-                dt,
-                nt,
-                step_ratio,
-                accuracy,
-                pml_width,
-                n_batch,
-                None,
-                None,
-                1,
-            ),
-        )
+    torch.autograd.gradgradcheck(
+        scalar_func,
+        (
+            c,
+            source_amplitudes,
+            wfc,
+            wfp,
+            psiy,
+            psix,
+            zetay,
+            zetax,
+            ay,
+            ax,
+            by,
+            bx,
+            dbydy,
+            dbxdx,
+            sources_i,
+            receivers_i,
+            dy,
+            dx,
+            dt,
+            nt,
+            step_ratio,
+            accuracy,
+            pml_width,
+            n_batch,
+            None,
+            None,
+            1,
+        ),
+    )
 
 
 def test_scalarfunc() -> None:
@@ -541,11 +540,6 @@ def test_scalarfunc() -> None:
     run_scalarfunc(nt=5)
 
 
-def test_scalarfunc_gradgrad() -> None:
-    """Test scalar_func with gradgrad enabled."""
-    run_scalarfunc(nt=4, gradgrad=True)
-
-
 def test_gradcheck_2d() -> None:
     """Test gradcheck in a 2D model."""
     run_gradcheck_2d(propagator=scalarprop)