#592: Add checkpoint flags and some test logic. (#599)

precicely · waltsims · web-flow · commit df48759c2b61 · 2025-05-24T08:50:39.000-07:00
* #592: added checkpoint flags and some test logic. * #592: added minimal example of how to use checkpointing. * #592: added windows guard in tests. * #592: improved test coverage. * #592: moved checkpoint validation to instantiation. Added setters/getters. * #592: added additional "valid" test for checkpointing. * #592: tweaked tests. * #592: addressed gap in testing coverage. --------- Co-authored-by: Walter Simson <walter.a.simson@gmail.com>
diff --git a/examples/checkpointing/README.md b/examples/checkpointing/README.md
@@ -0,0 +1,5 @@
+# Checkpointing
+
+This example demonstrates how to checkpoint simulations.
+
+It uses the same simulation parameters as found in "3D FFT Reconstruction For A Planar Sensor".
diff --git a/examples/checkpointing/checkpoint.py b/examples/checkpointing/checkpoint.py
@@ -0,0 +1,107 @@
+from copy import copy
+from pathlib import Path
+from tempfile import TemporaryDirectory
+
+import numpy as np
+
+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
+
+# This script demonstrates how to use the checkpointing feature of k-Wave.
+# It runs the same simulation twice, with the second run starting from a
+# checkpoint file created during the first run.
+# This implementation checkpoints based on the number of completed timesteps.
+# Alternatively, you can checkpoint based on the simulation time. Below is a
+# snippet demonstrating how:
+#
+# checkpoint_interval = 1  # 1 second
+# execution_options = SimulationExecutionOptions(
+#   checkpoint_file=checkpoint_filename,
+#   checkpoint_interval=checkpoint_interval
+# )
+# Note: checkpoint timesteps and checkpoint interval must be integers.
+
+
+def make_simulation_parameters(directory: Path, checkpoint_timesteps: int):
+    """
+    See the 3D FFT Reconstruction For A Planar Sensor example for context.
+    """
+    scale = 1
+
+    # create the computational grid
+    PML_size = 10  # size of the PML in grid points
+    N = Vector([32, 64, 64]) * scale - 2 * PML_size  # number of grid points
+    d = Vector([0.2e-3, 0.2e-3, 0.2e-3]) / scale  # 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 * scale  # [grid points]
+    p0 = ball_magnitude * make_ball(N, N / 2, ball_radius)
+    p0 = smooth(p0, restore_max=True)
+
+    source = kSource()
+    source.p0 = p0
+
+    # define a binary planar sensor
+    sensor = kSensor()
+    sensor_mask = np.zeros(N)
+    sensor_mask[0] = 1
+    sensor.mask = sensor_mask
+
+    input_filename = directory / "kwave_input.h5"
+    output_filename = directory / "kwave_output.h5"
+    checkpoint_filename = directory / "kwave_checkpoint.h5"
+
+    # set the input arguments
+    simulation_options = SimulationOptions(
+        save_to_disk=True,
+        pml_size=PML_size,
+        pml_inside=False,
+        smooth_p0=False,
+        data_cast="single",
+        input_filename=input_filename,
+        output_filename=output_filename,
+    )
+
+    execution_options = SimulationExecutionOptions(checkpoint_file=checkpoint_filename, checkpoint_timesteps=checkpoint_timesteps)
+    return kgrid, medium, source, sensor, simulation_options, execution_options
+
+
+def main():
+    # This simulation has 212 timesteps, therefore, the checkpoint is taken at
+    # the halfway point.
+    checkpoint_timesteps = 106
+
+    with TemporaryDirectory() as tmpdir:
+        tmpdir = Path(tmpdir)
+        # create the simulation parameters for the 1st run
+        kgrid, medium, source, sensor, simulation_options, execution_options = make_simulation_parameters(
+            directory=tmpdir, checkpoint_timesteps=checkpoint_timesteps
+        )
+        kspaceFirstOrder3D(kgrid, source, sensor, medium, simulation_options, execution_options)
+        print("Temporary directory contains 3 files (including checkpoint):")
+        print("\t-", "\n\t- ".join([f.name for f in tmpdir.glob("*.h5")]))
+
+        # create the simulation parameters for the 2nd run
+        kgrid, medium, source, sensor, simulation_options, execution_options = make_simulation_parameters(
+            directory=tmpdir, checkpoint_timesteps=checkpoint_timesteps
+        )
+        kspaceFirstOrder3D(kgrid, source, sensor, medium, simulation_options, execution_options)
+        print("Temporary directory contains 2 files (checkpoint has been deleted):")
+        print("\t-", "\n\t- ".join([f.name for f in tmpdir.glob("*.h5")]))
+
+
+if __name__ == "__main__":
+    main()
diff --git a/kwave/options/simulation_execution_options.py b/kwave/options/simulation_execution_options.py
@@ -1,11 +1,14 @@
 import os
 import warnings
+from logging import getLogger
 from pathlib import Path
 from typing import Optional, Union
 
 from kwave import BINARY_DIR, PLATFORM
 from kwave.ksensor import kSensor
 
+logger = getLogger(__name__)
+
 
 class SimulationExecutionOptions:
     """
@@ -27,6 +30,9 @@ def __init__(
         verbose_level: int = 0,
         auto_chunking: Optional[bool] = True,
         show_sim_log: bool = True,
+        checkpoint_interval: Optional[int] = None,  # [seconds]
+        checkpoint_timesteps: Optional[int] = None,  # [timestep integer]
+        checkpoint_file: Optional[Path | str] = None,  # [path to hdf5 file]
     ):
         self.is_gpu_simulation = is_gpu_simulation
         self._binary_path = binary_path
@@ -41,6 +47,13 @@ def __init__(
         self.verbose_level = verbose_level
         self.auto_chunking = auto_chunking
         self.show_sim_log = show_sim_log
+        self.checkpoint_interval = checkpoint_interval
+        self.checkpoint_timesteps = checkpoint_timesteps
+        self.checkpoint_file = checkpoint_file
+
+        if self.checkpoint_file is not None:
+            if self.checkpoint_interval is None and self.checkpoint_timesteps is None:
+                raise ValueError("One of checkpoint_interval or checkpoint_timesteps must be set when checkpoint_file is set.")
 
     @property
     def num_threads(self) -> Union[int, str]:
@@ -164,6 +177,47 @@ def device_num(self, value: Optional[int]):
             raise ValueError("Device number must be non-negative")
         self._device_num = value
 
+    @property
+    def checkpoint_interval(self) -> Optional[int]:
+        return self._checkpoint_interval
+
+    @checkpoint_interval.setter
+    def checkpoint_interval(self, value: Optional[int]):
+        if value is not None:
+            if not isinstance(value, int) or value < 0:
+                raise ValueError("Checkpoint interval must be a positive integer")
+        self._checkpoint_interval = value
+
+    @property
+    def checkpoint_timesteps(self) -> Optional[int]:
+        return self._checkpoint_timesteps
+
+    @checkpoint_timesteps.setter
+    def checkpoint_timesteps(self, value: Optional[int]):
+        if value is not None:
+            if not isinstance(value, int) or value < 0:
+                raise ValueError("Checkpoint timesteps must be a positive integer")
+        self._checkpoint_timesteps = value
+
+    @property
+    def checkpoint_file(self) -> Optional[Path]:
+        if self._checkpoint_file is None:
+            return None
+        return self._checkpoint_file
+
+    @checkpoint_file.setter
+    def checkpoint_file(self, value: Optional[Path | str]):
+        if value is not None:
+            if not isinstance(value, (str, Path)):
+                raise ValueError("Checkpoint file must be a string or Path object.")
+            if isinstance(value, str):
+                value = Path(value)
+            if not value.parent.is_dir():
+                raise FileNotFoundError(f"Checkpoint folder {value.parent} does not exist.")
+            if value.suffix != ".h5":
+                raise ValueError(f"Checkpoint file {value} must have .h5 extension.")
+        self._checkpoint_file = value
+
     def as_list(self, sensor: kSensor) -> list[str]:
         options_list = []
 
@@ -179,6 +233,17 @@ def as_list(self, sensor: kSensor) -> list[str]:
             options_list.append("--verbose")
             options_list.append(str(self.verbose_level))
 
+        if (self.checkpoint_interval is not None or self.checkpoint_timesteps is not None) and self.checkpoint_file is not None:
+            if self.checkpoint_timesteps is not None:
+                options_list.append("--checkpoint_timesteps")
+                options_list.append(str(self.checkpoint_timesteps))
+            if self.checkpoint_interval is not None:
+                options_list.append("--checkpoint_interval")
+                options_list.append(str(self.checkpoint_interval))
+
+            options_list.append("--checkpoint_file")
+            options_list.append(str(self.checkpoint_file))
+
         record_options_map = {
             "p": "p_raw",
             "p_max": "p_max",
diff --git a/tests/test_simulation_execution_options.py b/tests/test_simulation_execution_options.py
@@ -1,5 +1,7 @@
 import os
 import unittest
+from pathlib import Path
+from tempfile import TemporaryDirectory
 from unittest.mock import Mock, patch
 
 from kwave import PLATFORM
@@ -31,6 +33,8 @@ def test_default_initialization(self):
         self.assertEqual(options.verbose_level, 0)
         self.assertTrue(options.auto_chunking)
         self.assertTrue(options.show_sim_log)
+        self.assertIsNone(options.checkpoint_interval)
+        self.assertIsNone(options.checkpoint_file)
 
     def test_num_threads_setter_valid(self):
         """Test setting a valid number of threads."""
@@ -214,6 +218,113 @@ def test_list_compared_to_string(self):
         options_string = options.get_options_string(self.mock_sensor)
         self.assertEqual(" ".join(options_list), options_string)
 
+    def test_as_list_with_valid_checkpoint_interval_options(self):
+        """Test the checkpoint interval options."""
+        options = self.default_options
+        options.num_threads = 1
+        options.checkpoint_interval = 10
+        options.checkpoint_file = "checkpoint.h5"
+
+        options_list = options.as_list(self.mock_sensor)
+        expected_elements = [
+            "--checkpoint_interval",
+            "10",
+            "--checkpoint_file",
+            "checkpoint.h5",
+            "--p_raw",
+            "--u_max",
+            "-s",
+            "10",
+        ]
+        if not PLATFORM == "windows":
+            expected_elements.insert(0, "-t")
+            expected_elements.insert(1, f"1")
+        self.assertListEqual(expected_elements, options_list)
+
+    def test_as_list_with_valid_checkpoint_timesteps_options(self):
+        """Test the checkpoint interval options."""
+        options = self.default_options
+        options.num_threads = 1
+        options.checkpoint_timesteps = 10
+        options.checkpoint_file = "checkpoint.h5"
+
+        options_list = options.as_list(self.mock_sensor)
+        expected_elements = [
+            "--checkpoint_timesteps",
+            "10",
+            "--checkpoint_file",
+            "checkpoint.h5",
+            "--p_raw",
+            "--u_max",
+            "-s",
+            "10",
+        ]
+        if not PLATFORM == "windows":
+            expected_elements.insert(0, "-t")
+            expected_elements.insert(1, f"1")
+        self.assertListEqual(expected_elements, options_list)
+
+    def test_initialization_with_invalid_checkpoint_options(self):
+        """Test initialization with invalid checkpoint options."""
+
+        # Test with valid checkpoint_file but unset checkpoint_timesteps and checkpoint_interval
+        with TemporaryDirectory() as temp_dir:
+            checkpoint_file = Path(temp_dir) / "checkpoint.h5"
+            with self.assertRaises(ValueError):
+                SimulationExecutionOptions(checkpoint_file=checkpoint_file)
+
+        # Test with invalid checkpoint_file type
+        with self.assertRaises(ValueError):
+            SimulationExecutionOptions(checkpoint_file=12345, checkpoint_timesteps=10)
+
+        # Test with invalid checkpoint_timesteps
+        with self.assertRaises(ValueError):
+            SimulationExecutionOptions(checkpoint_timesteps=-1, checkpoint_file="checkpoint.h5")
+
+        with self.assertRaises(ValueError):
+            SimulationExecutionOptions(checkpoint_timesteps="not_an_integer", checkpoint_file="checkpoint.h5")
+
+        # Test with invalid checkpoint_interval
+        with self.assertRaises(ValueError):
+            SimulationExecutionOptions(checkpoint_interval=-1, checkpoint_file="checkpoint.h5")
+
+        with self.assertRaises(ValueError):
+            SimulationExecutionOptions(checkpoint_interval="not_an_integer", checkpoint_file="checkpoint.h5")
+
+        # Test with invalid checkpoint_file
+        with self.assertRaises(ValueError):
+            SimulationExecutionOptions(
+                checkpoint_interval=10,
+                checkpoint_file="checkpoint.txt",  # Wrong extension
+            )
+
+        with self.assertRaises(FileNotFoundError):
+            SimulationExecutionOptions(
+                checkpoint_interval=10,
+                checkpoint_file="nonexistent_dir/checkpoint.h5",  # Non-existent directory
+            )
+
+    def test_initialization_with_valid_checkpoint_options(self):
+        """Test initialization with valid checkpoint options."""
+        with TemporaryDirectory() as temp_dir:
+            checkpoint_file = Path(temp_dir) / "checkpoint.h5"
+
+            # Test with valid checkpoint_timesteps
+            options = SimulationExecutionOptions(checkpoint_timesteps=10, checkpoint_file=checkpoint_file)
+            self.assertEqual(options.checkpoint_timesteps, 10)
+            self.assertEqual(options.checkpoint_file, checkpoint_file)
+
+            # Test with valid checkpoint_interval
+            options = SimulationExecutionOptions(checkpoint_interval=20, checkpoint_file=checkpoint_file)
+            self.assertEqual(options.checkpoint_interval, 20)
+            self.assertEqual(options.checkpoint_file, checkpoint_file)
+
+            # Test with both checkpoint options - should be valid
+            options = SimulationExecutionOptions(checkpoint_timesteps=10, checkpoint_interval=20, checkpoint_file=checkpoint_file)
+            self.assertEqual(options.checkpoint_timesteps, 10)
+            self.assertEqual(options.checkpoint_interval, 20)
+            self.assertEqual(options.checkpoint_file, checkpoint_file)
+
 
 if __name__ == "__main__":
     unittest.main()
