From 765465e2f8bdee9bee183c6774e0e60dabbccf82 Mon Sep 17 00:00:00 2001 From: barkure <43804451+barkure@users.noreply.github.com> Date: Fri, 17 Jul 2026 16:46:17 +0800 Subject: [PATCH 1/3] Fix elastic callback wavefields and harden location/dt checks Elastic backward callbacks dropped the last memory wavefield via grad_wavefields[:-1]. The Python forward path also mis-zipped shear and memory names against empty matrix placeholders. Skip empty locations in extent checks, reject zero dt, and reject storage compression on the acoustic Python backend. --- src/deepwave/acoustic.py | 4 + src/deepwave/common.py | 4 + src/deepwave/elastic.py | 52 ++++++++--- tests/test_acoustic.py | 56 ++++++++++++ tests/test_callbacks_elastic.py | 152 ++++++++++++++++++++++++++++++++ tests/test_common.py | 32 +++++++ 6 files changed, 290 insertions(+), 10 deletions(-) diff --git a/src/deepwave/acoustic.py b/src/deepwave/acoustic.py index dbb1540..48ac84f 100644 --- a/src/deepwave/acoustic.py +++ b/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) diff --git a/src/deepwave/common.py b/src/deepwave/common.py index 1f8e302..1f14b9c 100644 --- a/src/deepwave/common.py +++ b/src/deepwave/common.py @@ -1928,6 +1928,8 @@ def check_locations_within_extents( for dim in range(location.shape[-1]): dim_location = location[..., dim] dim_location = dim_location[dim_location != IGNORE_LOCATION] + if dim_location.numel() == 0: + continue if ( dim_location.min() < extents[dim][0] or extents[dim][1] <= dim_location.max() @@ -2375,6 +2377,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: diff --git a/src/deepwave/elastic.py b/src/deepwave/elastic.py index 3de44c5..d77d0c3 100644 --- a/src/deepwave/elastic.py +++ b/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, diff --git a/tests/test_acoustic.py b/tests/test_acoustic.py index 5bd1bc6..56dfb3d 100644 --- a/tests/test_acoustic.py +++ b/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, + ) diff --git a/tests/test_callbacks_elastic.py b/tests/test_callbacks_elastic.py index be8a3f1..5f51c4f 100644 --- a/tests/test_callbacks_elastic.py +++ b/tests/test_callbacks_elastic.py @@ -547,3 +547,155 @@ 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", [True, False]) +@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 + + 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 + } + ) + + recorders = {} + for python_backend in [False, True]: + 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[True].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:] + ) diff --git a/tests/test_common.py b/tests/test_common.py index ccd7959..cc7dbfb 100644 --- a/tests/test_common.py +++ b/tests/test_common.py @@ -5,8 +5,10 @@ import pytest import torch +from deepwave import IGNORE_LOCATION from deepwave.common import ( cfl_condition_n, + check_locations_within_extents, check_points_per_wavelength, check_source_amplitudes_locations_match, cosine_taper_end, @@ -1201,6 +1203,36 @@ def test_cfl_condition_n_negative_dt() -> None: assert inner_dt == pytest.approx(-0.001) +def test_cfl_condition_n_zero_dt() -> None: + """Test cfl_condition_n with zero dt (should raise error).""" + with pytest.raises( + ValueError, + match=re.escape("dt must be non-zero."), + ): + cfl_condition_n([5.0, 5.0], 0.0, 1500.0) + + +# Tests for check_locations_within_extents +def test_check_locations_within_extents_empty_locations() -> None: + """Test check_locations_within_extents with no valid locations.""" + extents = [(2, 8), (2, 8)] + empty_location = torch.zeros(1, 0, 2, dtype=torch.long) + ignored_location = torch.full((1, 2, 2), IGNORE_LOCATION) + # Should not raise + check_locations_within_extents(extents, [empty_location, ignored_location]) + + +def test_check_locations_within_extents_outside() -> None: + """Test check_locations_within_extents with a location outside extents.""" + extents = [(2, 8), (2, 8)] + location = torch.tensor([[[1, 5]]]) + with pytest.raises( + RuntimeError, + match=re.escape("Locations are not within survey extents."), + ): + check_locations_within_extents(extents, [location]) + + def test_cfl_condition_n_different_grid_spacing() -> None: """Test cfl_condition_n with different grid spacings in dimensions.""" grid_spacing = [5.0, 10.0] From dd195270002d218780de753303bb89c83c123518 Mon Sep 17 00:00:00 2001 From: barkure <43804451+barkure@users.noreply.github.com> Date: Fri, 17 Jul 2026 16:46:17 +0800 Subject: [PATCH 2/3] Make test_elasticfunc use eager Python backend python_backend=True maps to torch.compile, whose codegen can depend on process-global dynamo state left by earlier tests. Use eager as the independent reference so the C/CUDA comparison is isolation-stable. --- tests/test_elastic.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/tests/test_elastic.py b/tests/test_elastic.py index 1e5889d..112323a 100644 --- a/tests/test_elastic.py +++ b/tests/test_elastic.py @@ -1199,8 +1199,14 @@ def run_elasticfunc(nt: int = 3) -> None: pml_profiles[i] = torch.randn_like(profile) * mask def wrap(python): + # Use "eager" (not True/"compile") as the independent Python reference. + # python_backend=True maps to torch.compile, whose codegen can depend + # on process-global dynamo cache state (e.g. automatic dynamic shapes + # left by earlier tests), producing ~1e-6 relative FP differences that + # break default allclose against the C/CUDA backend. + backend = "eager" if python else False inputs = ( - python, + backend, pml_profiles, [ dy, From 8ab90930d368963429741e78a4e02abb656087a5 Mon Sep 17 00:00:00 2001 From: barkure <43804451+barkure@users.noreply.github.com> Date: Fri, 17 Jul 2026 17:13:23 +0800 Subject: [PATCH 3/3] Address review feedback on extents checks and callback tests Skip location tensors only when empty or fully IGNORE_LOCATION, then validate remaining points per dimension. Use python_backend="eager" in elastic callback tests so comparisons are not order-dependent on torch.compile state. --- src/deepwave/common.py | 40 ++++++++++++++++++--------------- tests/test_callbacks_elastic.py | 10 ++++++--- tests/test_common.py | 23 +++++++++++++++++++ 3 files changed, 52 insertions(+), 21 deletions(-) diff --git a/src/deepwave/common.py b/src/deepwave/common.py index 1f14b9c..7e1a112 100644 --- a/src/deepwave/common.py +++ b/src/deepwave/common.py @@ -1924,24 +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.numel() == 0: - continue - 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( diff --git a/tests/test_callbacks_elastic.py b/tests/test_callbacks_elastic.py index 5f51c4f..9dd66c4 100644 --- a/tests/test_callbacks_elastic.py +++ b/tests/test_callbacks_elastic.py @@ -602,7 +602,7 @@ def _elastic_survey(ndim, nt, device=None): return source_amplitudes_y, source_locations_y -@pytest.mark.parametrize("python_backend", [True, False]) +@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.""" @@ -625,6 +625,8 @@ def __call__(self, state: deepwave.common.CallbackState) -> None: 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, @@ -667,8 +669,10 @@ def __call__(self, state: deepwave.common.CallbackState) -> None: } ) + # 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, True]: + for python_backend in [False, "eager"]: recorder = Recorder() deepwave.elastic( lamb, @@ -685,7 +689,7 @@ def __call__(self, state: deepwave.common.CallbackState) -> None: recorders[python_backend] = recorder compiled_records = recorders[False].records - python_records = recorders[True].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): diff --git a/tests/test_common.py b/tests/test_common.py index cc7dbfb..bf25ab9 100644 --- a/tests/test_common.py +++ b/tests/test_common.py @@ -1233,6 +1233,29 @@ def test_check_locations_within_extents_outside() -> None: check_locations_within_extents(extents, [location]) +def test_check_locations_within_extents_partial_ignore() -> None: + """Partially ignored points are still checked on valid coordinates. + + A point with only some coordinates set to IGNORE_LOCATION is not a + fully ignored location, so its remaining coordinates must still lie + within the survey extents. + """ + extents = [(2, 8), (2, 8)] + # One fully ignored point (skipped) and one valid in-bounds point. + mixed_ok = torch.tensor( + [[[IGNORE_LOCATION, IGNORE_LOCATION], [3, 4]]], + dtype=torch.long, + ) + check_locations_within_extents(extents, [mixed_ok]) + # One coordinate ignored, the other out of bounds → still an error. + partial_bad = torch.tensor([[[IGNORE_LOCATION, 0]]], dtype=torch.long) + with pytest.raises( + RuntimeError, + match=re.escape("Locations are not within survey extents."), + ): + check_locations_within_extents(extents, [partial_bad]) + + def test_cfl_condition_n_different_grid_spacing() -> None: """Test cfl_condition_n with different grid spacings in dimensions.""" grid_spacing = [5.0, 10.0]