Fix: Stop mutation of source and sensor by kspaceFirstOrder

kwave/kWaveSimulation.py

@@ -1,5 +1,6 @@
 import logging
 import warnings
+from copy import deepcopy
 from dataclasses import dataclass
 
 import numpy as np
@@ -38,8 +39,10 @@ def __init__(
         self.precision = None
         self.kgrid = kgrid
         self.medium = medium
-        self.source = source
-        self.sensor = sensor
+        self.original_source = source
+        self.original_sensor = sensor
+        self.source = deepcopy(source)
+        self.sensor = deepcopy(sensor)
         self.options = simulation_options
 
         # =========================================================================
@@ -1517,7 +1520,7 @@ def _is_binary_sensor_mask(self, kgrid_dim: int) -> bool:
             return True
 
         # If the sensor mask is larger by PML size in each dimension, it's still a binary mask
-        if pml_size := self.options.pml_size is None:
+        if (pml_size := self.options.pml_size) is None:
             return False
 
         if len(pml_size) == 1:

tests/test_kspaceFirstOrder3D_state.py

@@ -0,0 +1,126 @@
+from copy import deepcopy
+from pathlib import Path
+from tempfile import TemporaryDirectory
+
+import numpy as np
+import pytest
+
+from kwave.data import Vector
+from kwave.kgrid import kWaveGrid
+from kwave.kmedium import kWaveMedium
+from kwave.ksensor import kSensor
+from kwave.ksource import kSource
+from kwave.kspaceFirstOrder3D import kspaceFirstOrder3D
+from kwave.options.simulation_execution_options import SimulationExecutionOptions
+from kwave.options.simulation_options import SimulationOptions
+from kwave.utils.filters import smooth
+from kwave.utils.mapgen import make_ball
+
+
+def make_simulation_parameters(directory: Path):
+    # create the computational grid
+    PML_size = 10  # size of the PML in grid points
+    N = Vector([32, 64, 64]) - 2 * PML_size  # number of grid points
+    d = Vector([0.2e-3, 0.2e-3, 0.2e-3])  # grid point spacing [m]
+    kgrid = kWaveGrid(N, d)
+
+    # define the properties of the propagation medium
+    medium = kWaveMedium(sound_speed=1500)  # [m/s]
+
+    # create initial pressure distribution using makeBall
+    ball_magnitude = 10  # [Pa]
+    ball_radius = 3  # [grid points]
+    p0_array = ball_magnitude * make_ball(N, N / 2, ball_radius)
+    p0_array = smooth(p0_array, restore_max=True)
+
+    source = kSource()
+    source.p0 = p0_array
+
+    # define a binary planar sensor
+    sensor = kSensor()
+    sensor_mask_array = np.zeros(N)
+    sensor_mask_array[0, :, :] = 1 # Corrected to be a plane for 3D
+    sensor.mask = sensor_mask_array
+
+    input_filename = directory / "kwave_input.h5"
+    output_filename = directory / "kwave_output.h5"
+    checkpoint_filename = directory / "kwave_checkpoint.h5"
+
+    simulation_options = SimulationOptions(
+        save_to_disk=True, # Must be true for kspaceFirstOrder3D
+        pml_size=PML_size,
+        pml_inside=False,
+        smooth_p0=False, # p0 is already smoothed
+        data_cast="single",
+        input_filename=input_filename,
+        output_filename=output_filename,
+    )
+
+    checkpoint_timesteps = 300
+
+    execution_options = SimulationExecutionOptions(
+        is_gpu_simulation=False,  # Assuming CPU for basic test
+        checkpoint_file=checkpoint_filename,
+        checkpoint_timesteps=checkpoint_timesteps,
+        verbose_level=0,  # Keep test output clean
+    )
+    return kgrid, medium, source, sensor, simulation_options, execution_options
+
+
+def test_kspaceFirstOrder3D_input_state_preservation():
+    with TemporaryDirectory() as tmpdir_str:
+        tmpdir = Path(tmpdir_str)
+        kgrid, medium, source, sensor, simulation_options, execution_options = make_simulation_parameters(tmpdir)
+
+        # Store original states of critical attributes for comparison
+        original_source_p0 = deepcopy(source.p0)
+        original_sensor_mask = deepcopy(sensor.mask)
+
+
+        # If source.p or source.u were time-varying, store their initial states too.
+        # For this test, p0 is the main source attribute.
+
+        # First run
+        try:
+            _ = kspaceFirstOrder3D(
+                kgrid=kgrid,
+                medium=medium,
+                source=source,
+                sensor=sensor,
+                simulation_options=simulation_options,
+                execution_options=execution_options,
+            )
+        except Exception as e:
+            pytest.fail(f"First call to kspaceFirstOrder3D failed: {e}")
+
+        # Check if original source and sensor attributes are unchanged
+        assert np.array_equal(source.p0, original_source_p0), "source.p0 was modified after first run"
+        assert np.array_equal(sensor.mask, original_sensor_mask), "sensor.mask was modified after first run"
+
+        # Second run (should not fail if state is preserved)
+        # For the second run, we need new input/output filenames or to ensure the C++ code can overwrite.
+        # Easiest is to use new filenames for the test.
+        simulation_options_run2 = deepcopy(simulation_options)
+        simulation_options_run2.input_filename = tmpdir / "kwave_input_run2.h5"
+        simulation_options_run2.output_filename = tmpdir / "kwave_output_run2.h5"
+
+        execution_options_run2 = deepcopy(execution_options)
+        if execution_options_run2.checkpoint_file:  # Only change if it exists
+            execution_options_run2.checkpoint_file = tmpdir / "kwave_checkpoint_run2.h5"
+
+        try:
+            _ = kspaceFirstOrder3D(
+                kgrid=kgrid,
+                medium=medium,
+                source=source,
+                sensor=sensor,
+                simulation_options=simulation_options_run2,
+                execution_options=execution_options_run2,
+            )
+        except Exception as e:
+            pytest.fail(f"Second call to kspaceFirstOrder3D with original objects failed: {e}")
+
+        # Final check that attributes are still the same as the initial state
+        assert np.array_equal(source.p0, original_source_p0), "source.p0 was modified after second run"
+        assert np.array_equal(sensor.mask, original_sensor_mask), "sensor.mask was modified after second run"
+