Fix elastic callback wavefield mapping and related input-validation edge cases

src/deepwave/acoustic.py

@@ -1362,6 +1362,10 @@ def acoustic_python(
         raise RuntimeError(
             "Specifying storage mode is not supported in Python backend."
         )
+    if storage_compression:
+        raise RuntimeError(
+            "Storage compression is not supported in the Python backend."
+        )
 
     is_batched = any(deepwave.common.is_inside_vmap(x) for x in args)
     ndim = len(grid_spacing)

src/deepwave/common.py

@@ -1924,22 +1924,28 @@ def check_locations_within_extents(
 
     """
     for location in locations:
-        if location is not None:
-            for dim in range(location.shape[-1]):
-                dim_location = location[..., dim]
-                dim_location = dim_location[dim_location != IGNORE_LOCATION]
-                if (
-                    dim_location.min() < extents[dim][0]
-                    or extents[dim][1] <= dim_location.max()
-                ):
-                    raise RuntimeError(
-                        "Locations are not within "
-                        "survey extents. This probably occurred "
-                        "because you provided an "
-                        "initial wavefield that does not cover all "
-                        "source and receiver locations, given the "
-                        "specified origin.",
-                    )
+        if location is None or location.numel() == 0:
+            continue
+        # Drop points whose every coordinate is IGNORE_LOCATION. Skip the
+        # whole tensor when nothing remains (empty or fully ignored).
+        flat = location.reshape(-1, location.shape[-1])
+        valid = flat[~(flat == IGNORE_LOCATION).all(dim=-1)]
+        if valid.shape[0] == 0:
+            continue
+        for dim in range(valid.shape[-1]):
+            dim_location = valid[:, dim]
+            if (
+                dim_location.min() < extents[dim][0]
+                or extents[dim][1] <= dim_location.max()
+            ):
+                raise RuntimeError(
+                    "Locations are not within "
+                    "survey extents. This probably occurred "
+                    "because you provided an "
+                    "initial wavefield that does not cover all "
+                    "source and receiver locations, given the "
+                    "specified origin.",
+                )
 
 
 def check_extents_match_wavefields_shape(
@@ -2375,6 +2381,8 @@ def cfl_condition_n(
         dt = float(dt)
     except (TypeError, ValueError) as e:
         raise TypeError("dt must be a float.") from e
+    if dt == 0:
+        raise ValueError("dt must be non-zero.")
     try:
         max_abs_vel = float(max_abs_vel)
     except (TypeError, ValueError) as e:

src/deepwave/elastic.py

@@ -1919,7 +1919,7 @@ def backward(
 
                 if backward_callback is not None:
                     callback_wavefields = dict(
-                        zip(wavefield_names, grad_wavefields[:-1])
+                        zip(wavefield_names, grad_wavefields)
                     )
                     backward_callback(
                         deepwave.common.CallbackState(
@@ -2504,20 +2504,52 @@ def elastic_python(
         if forward_callback is not None and step % callback_frequency == 0:
             callback_wavefields = dict(zip(v_names, velocities))
             callback_wavefields.update(dict(zip(normal_s_names, normal_stresses)))
-            for i in range(len(shear_s_names)):
-                callback_wavefields.update(
-                    dict(zip(shear_s_names[i], shear_stresses[i]))
+            # The shear names are flattened in the same order as the
+            # upper-triangular (i < j) entries of the symmetric shear
+            # stress matrix, whose diagonal contains empty placeholders.
+            callback_wavefields.update(
+                dict(
+                    zip(
+                        [name for names in shear_s_names for name in names],
+                        [
+                            shear_stresses[i][j]
+                            for i in range(ndim)
+                            for j in range(i + 1, ndim)
+                        ],
+                    )
                 )
+            )
             callback_wavefields.update(dict(zip(v_mem_normal_names, v_mem_vars_normal)))
-            for i in range(len(v_mem_shear_names)):
-                callback_wavefields.update(
-                    dict(zip(v_mem_shear_names[i], v_mem_vars_shear[i]))
+            # The memory variable names are flattened row by row, matching
+            # the non-placeholder entries of the shear matrices (the
+            # diagonal entries are empty placeholders).
+            callback_wavefields.update(
+                dict(
+                    zip(
+                        [name for names in v_mem_shear_names for name in names],
+                        [
+                            var
+                            for row in v_mem_vars_shear
+                            for var in row
+                            if var.numel() > 0
+                        ],
+                    )
                 )
+            )
             callback_wavefields.update(dict(zip(s_mem_normal_names, s_mem_vars_normal)))
-            for i in range(len(s_mem_shear_names)):
-                callback_wavefields.update(
-                    dict(zip(s_mem_shear_names[i], s_mem_vars_shear[i]))
+            callback_wavefields.update(
+                dict(
+                    zip(
+                        [name for names in s_mem_shear_names for name in names],
+                        [
+                            var
+                            for row in s_mem_vars_shear
+                            for var in row
+                            if var.numel() > 0
+                        ],
+                    )
                 )
+            )
             forward_callback(
                 deepwave.common.CallbackState(
                     dt,

tests/test_acoustic.py

@@ -1357,3 +1357,59 @@ def call(model: torch.Tensor) -> Tuple[torch.Tensor, ...]:
         )
 
     run_checkpoint_equivalence(v, call, use_reentrant=use_reentrant)
+
+
+def test_initial_wavefield_empty_receiver_locations() -> None:
+    """Check initial wavefields with no valid receiver locations.
+
+    Receiver locations that are empty or all IGNORE_LOCATION should not
+    cause an error when initial wavefields are provided.
+    """
+    torch.manual_seed(0)
+    v = 1500 * torch.ones(10, 10)
+    rho = 1000 * torch.ones(10, 10)
+    pressure_0 = torch.zeros(1, 10, 10)
+    pressure_0[0, 5, 5] = 1
+    ignored_receiver_locations = torch.full((1, 1, 2), IGNORE_LOCATION)
+    empty_receiver_locations = torch.zeros(1, 0, 2, dtype=torch.long)
+
+    for receiver_locations_p in [
+        ignored_receiver_locations,
+        empty_receiver_locations,
+    ]:
+        acousticprop(
+            v,
+            rho,
+            5.0,
+            0.001,
+            receiver_locations_p=receiver_locations_p,
+            pressure_0=pressure_0,
+            nt=5,
+            pml_width=0,
+        )
+
+
+def test_python_backend_storage_compression() -> None:
+    """Check that storage_compression with the Python backend raises."""
+    v = 1500 * torch.ones(10, 10)
+    rho = 1000 * torch.ones(10, 10)
+    source_amplitudes = torch.zeros(1, 1, 10)
+    source_amplitudes[0, 0, 5] = 1
+    source_locations = torch.zeros(1, 1, 2, dtype=torch.long)
+    source_locations[0, 0, 0] = 5
+    source_locations[0, 0, 1] = 5
+
+    with pytest.raises(
+        RuntimeError,
+        match="Storage compression is not supported in the Python backend.",
+    ):
+        acoustic(
+            v,
+            rho,
+            5.0,
+            0.001,
+            source_amplitudes_p=source_amplitudes,
+            source_locations_p=source_locations,
+            python_backend=True,
+            storage_compression=True,
+        )

tests/test_callbacks_elastic.py

@@ -547,3 +547,159 @@ def __call__(self, state: deepwave.common.CallbackState) -> None:
     )
     out[-2].sum().backward()
     assert backward_counter.count == (nt + 2) // 3
+
+
+EXPECTED_WAVEFIELD_NAMES_2D = {
+    "vy_0",
+    "sigmayy_0",
+    "sigmaxy_0",
+    "m_vyy_0",
+    "m_vyx_0",
+    "m_vxy_0",
+    "m_sigmayyy_0",
+    "m_sigmaxyy_0",
+    "m_sigmaxyx_0",
+    "vx_0",
+    "sigmaxx_0",
+    "m_vxx_0",
+    "m_sigmaxxx_0",
+}
+
+EXPECTED_WAVEFIELD_NAMES_3D = {
+    "vz_0",
+    "sigmazz_0",
+    "sigmayz_0",
+    "sigmaxz_0",
+    "m_vzz_0",
+    "m_vzy_0",
+    "m_vzx_0",
+    "m_vyz_0",
+    "m_vxz_0",
+    "m_sigmazzz_0",
+    "m_sigmayzy_0",
+    "m_sigmaxzx_0",
+    "m_sigmayzz_0",
+    "m_sigmaxzz_0",
+    *EXPECTED_WAVEFIELD_NAMES_2D,
+}
+
+
+def _elastic_models(ndim, device=None):
+    """Create constant elastic models for the given number of dimensions."""
+    shape = (10,) * ndim
+    lamb = torch.ones(*shape, device=device) * 2200
+    mu = torch.ones(*shape, device=device) * 1000
+    buoyancy = torch.ones(*shape, device=device) * 1 / 2200
+    return lamb, mu, buoyancy
+
+
+def _elastic_survey(ndim, nt, device=None):
+    """Create a central y-directed point source for the given ndim."""
+    source_amplitudes_y = torch.zeros(1, 1, nt, device=device)
+    source_amplitudes_y[0, 0, 5] = 1
+    source_locations_y = torch.zeros(1, 1, ndim, dtype=torch.long, device=device)
+    source_locations_y[0, 0, :] = 5
+    return source_amplitudes_y, source_locations_y
+
+
+@pytest.mark.parametrize("python_backend", [False, "eager"])
+@pytest.mark.parametrize("ndim", [2, 3])
+def test_elastic_callback_wavefield_names(python_backend, ndim) -> None:
+    """Check that callbacks expose the full set of non-empty wavefields."""
+    lamb, mu, buoyancy = _elastic_models(ndim)
+    lamb.requires_grad_()
+    mu.requires_grad_()
+    buoyancy.requires_grad_()
+    dx = 5.0
+    dt = 0.004
+    nt = 10
+    source_amplitudes_y, source_locations_y = _elastic_survey(ndim, nt)
+    expected = EXPECTED_WAVEFIELD_NAMES_3D if ndim == 3 else EXPECTED_WAVEFIELD_NAMES_2D
+
+    class Checker:
+        """Checks the names and contents of the callback wavefields."""
+
+        def __call__(self, state: deepwave.common.CallbackState) -> None:
+            """Check that all expected wavefields are present and non-empty."""
+            assert set(state._wavefields.keys()) == expected  # noqa: SLF001
+            for name in expected:
+                assert state.get_wavefield(name).numel() > 0
+
+    # Use "eager" rather than True/"compile" so name checks are not coupled
+    # to process-global torch.compile state.
+    out = deepwave.elastic(
+        lamb,
+        mu,
+        buoyancy,
+        dx,
+        dt,
+        source_amplitudes_y=source_amplitudes_y,
+        source_locations_y=source_locations_y,
+        forward_callback=Checker(),
+        backward_callback=None if python_backend else Checker(),
+        callback_frequency=1,
+        python_backend=python_backend,
+    )
+    if not python_backend:
+        out[-1].sum().backward()
+
+
+@pytest.mark.parametrize("ndim", [2, 3])
+def test_elastic_callback_wavefield_backend_consistency(ndim) -> None:
+    """Check that both backends expose the same wavefields in callbacks."""
+    lamb, mu, buoyancy = _elastic_models(ndim)
+    dx = 5.0
+    dt = 0.004
+    nt = 10
+    source_amplitudes_y, source_locations_y = _elastic_survey(ndim, nt)
+    expected = EXPECTED_WAVEFIELD_NAMES_3D if ndim == 3 else EXPECTED_WAVEFIELD_NAMES_2D
+
+    class Recorder:
+        """Records the inner view of every wavefield at each callback."""
+
+        def __init__(self) -> None:
+            self.records = []
+
+        def __call__(self, state: deepwave.common.CallbackState) -> None:
+            """Record a clone of every wavefield."""
+            self.records.append(
+                {
+                    name: state.get_wavefield(name).clone()
+                    for name in state._wavefields  # noqa: SLF001
+                }
+            )
+
+    # Compare the C/CUDA backend against the eager Python reference, not
+    # python_backend=True (torch.compile), which can be order-dependent.
+    recorders = {}
+    for python_backend in [False, "eager"]:
+        recorder = Recorder()
+        deepwave.elastic(
+            lamb,
+            mu,
+            buoyancy,
+            dx,
+            dt,
+            source_amplitudes_y=source_amplitudes_y,
+            source_locations_y=source_locations_y,
+            forward_callback=recorder,
+            callback_frequency=3,
+            python_backend=python_backend,
+        )
+        recorders[python_backend] = recorder
+
+    compiled_records = recorders[False].records
+    python_records = recorders["eager"].records
+    assert len(compiled_records) == len(python_records)
+    assert len(compiled_records) > 1
+    for compiled_step, python_step in zip(compiled_records, python_records):
+        assert set(compiled_step.keys()) == expected
+        assert set(python_step.keys()) == expected
+        for name in expected:
+            assert torch.allclose(compiled_step[name], python_step[name]), name
+    # The simulation should have propagated non-zero shear wavefields, so
+    # the comparison above is not trivially comparing zeros.
+    shear_name = "sigmaxz_0" if ndim == 3 else "sigmaxy_0"
+    assert any(
+        step[shear_name].abs().max() > 0 for step in python_records[1:]
+    )