Enable sweep runs over paganin filter to use hdf5 file as backing

httomo/data/param_sweep_store.py

@@ -1,8 +1,13 @@
 from __future__ import annotations
+from os import PathLike
+from pathlib import Path
 from typing import Literal, Optional, Tuple
 
+from mpi4py import MPI
 import numpy as np
+import h5py
 
+from httomo.runner.dataset_store_backing import DataSetStoreBacking
 from httomo.runner.auxiliary_data import AuxiliaryData
 from httomo.sweep_runner.param_sweep_block import ParamSweepBlock
 from httomo.sweep_runner.param_sweep_store_interfaces import ParamSweepSource
@@ -23,6 +28,17 @@ def __init__(self, source: ParamSweepWriter) -> None:
         ), "Reader should be created from writer with data not `None`"
         self._data = source._data
         self._aux_data = source.aux_data
+        self._h5filename: Optional[Path] = None
+        self.h5file: Optional[h5py.File] = None
+
+        if source.is_file_based:
+            self._h5filename = source.filename
+            self.h5file = h5py.File(source.filename, "r")
+            self._data = self.h5file["/data"]
+        else:
+            self._data = source._data
+
+        source.finalize()
 
     @property
     def no_of_sweeps(self) -> int:
@@ -56,18 +72,52 @@ def read_sweep_results(self) -> ParamSweepBlock:
             slicing_dim=self.extract_dim,
         )
 
+    def finalize(self):
+        self._data = None
+        if self.h5file is not None:
+            self.h5file.close()
+            self.h5file = None
+        if self._h5filename is not None:
+            self._h5filename.unlink()
+            self._h5filename = None
+
+    @property
+    def is_file_based(self) -> bool:
+        return self._h5filename is not None
+
+    @property
+    def filename(self) -> Optional[Path]:
+        return self._h5filename
+
 
 class ParamSweepWriter:
     """Write parameter sweep results, concatenating them along the `detector_y` dim"""
 
-    def __init__(self, no_of_sweeps: int) -> None:
+    def __init__(
+        self,
+        no_of_sweeps: int,
+        comm: MPI.Comm,
+        temppath: PathLike,
+        backing: DataSetStoreBacking,
+    ) -> None:
         self._concat_dim: Literal[1] = 1
         self._no_of_sweeps = no_of_sweeps
         self._no_of_sweeps_written: int = 0
         self._single_shape: Optional[Tuple[int, int, int]] = None
         self._total_shape: Optional[Tuple[int, int, int]] = None
-        self._data: Optional[np.ndarray] = None
+        self._data: Optional[np.ndarray | h5py.Dataset] = None
         self._slices_per_sweep: Optional[int] = None
+        self._comm = comm
+        self._backing = backing
+        self._temppath = temppath
+        self._h5filename: Optional[Path] = None
+        self.h5file: Optional[h5py.File] = None
+
+        if backing == DataSetStoreBacking.File:
+            self._h5filename = (
+                Path(self._temppath) / f"sweep-output-rank-{self._comm.rank}.h5"
+            )
+            self.h5file = h5py.File(self._h5filename, "w")
 
     @property
     def no_of_sweeps(self) -> int:
@@ -128,8 +178,14 @@ def write_sweep_result(self, block: ParamSweepBlock):
                 self.single_shape[2],
             )
             self._slices_per_sweep = self.single_shape[self.concat_dim]
-            self._data = np.empty(shape=self.total_shape, dtype=block.data.dtype)
             self._aux_data = block.aux_data
+            if self._backing == DataSetStoreBacking.RAM:
+                self._data = np.empty(shape=self.total_shape, dtype=block.data.dtype)
+            elif self._backing == DataSetStoreBacking.File:
+                assert self.h5file is not None
+                self._data = self.h5file.create_dataset(
+                    "/data", shape=self.total_shape, dtype=block.data.dtype
+                )
 
         slices = [slice(None, None, 1)] * 3
         sweep_res_start = self.no_of_sweeps_written * self.slices_per_sweep
@@ -138,5 +194,22 @@ def write_sweep_result(self, block: ParamSweepBlock):
             sweep_res_start + self.slices_per_sweep,
             1,
         )
+        assert (
+            self._data is not None
+        ), "Numpy array or hdf5 file should be created prior to writing sweep result"
         self._data[slices[0], slices[1], slices[2]] = block.data
         self.increment_no_of_sweeps_written()
+
+    def finalize(self):
+        self._data = None
+        if self.h5file is not None:
+            self.h5file.close()
+            self.h5file = None
+
+    @property
+    def is_file_based(self) -> bool:
+        return self._backing is DataSetStoreBacking.File
+
+    @property
+    def filename(self) -> Optional[Path]:
+        return self._h5filename

httomo/sweep_runner/param_sweep_runner.py

@@ -7,10 +7,12 @@
 import numpy as np
 
 import httomo
-from httomo.data.param_sweep_store import ParamSweepReader, ParamSweepWriter
+from httomo import globals
+from httomo.data.param_sweep_store import ParamSweepWriter
 from httomo.method_wrappers.images import ImagesWrapper
 from httomo.method_wrappers.save_intermediate import SaveIntermediateFilesWrapper
 from httomo.runner.block_split import BlockSplitter
+from httomo.runner.dataset_store_backing import DataSetStoreBacking
 from httomo.runner.method_wrapper import MethodWrapper
 from httomo.runner.pipeline import Pipeline
 from httomo.sweep_runner.param_sweep_block import ParamSweepBlock
@@ -21,6 +23,10 @@
 from httomo.preview import PreviewConfig, PreviewDimConfig
 from httomo.runner.dataset_store_interfaces import DataSetSource
 
+PAGANIN_SWEEP_VALUE_THRESHOLD = 50
+"""The number of sweep values for the paganin filter above which the sweep results should be
+written to an hdf5 file as opposed to kept in RAM."""
+
 
 class ParamSweepRunner:
     def __init__(
@@ -42,6 +48,7 @@ def __init__(
             self._start_sweep_idx : self._stop_sweep_idx
         ]
         self._set_image_saver_params()
+        self._comm = comm
 
     @property
     def block(self) -> ParamSweepBlock:
@@ -212,8 +219,24 @@ def execute_after_sweep(self):
 
     def execute_sweep(self):
         """Execute all param variations of the same method in the sweep"""
-        writer = ParamSweepWriter(len(self._sweep_values))
         method = self._stages.sweep.method
+        if (
+            method.method_name == "paganin_filter"
+            and len(self._stages.sweep.values) > PAGANIN_SWEEP_VALUE_THRESHOLD
+        ):
+            writer = ParamSweepWriter(
+                len(self._sweep_values),
+                self._comm,
+                globals.run_out_dir,
+                DataSetStoreBacking.File,
+            )
+        else:
+            writer = ParamSweepWriter(
+                len(self._sweep_values),
+                self._comm,
+                globals.run_out_dir,
+                DataSetStoreBacking.RAM,
+            )
 
         log_once(f"Running {method.method_name} ({method.package_name})")
         log_once("    Beginning parameter sweep")
@@ -254,8 +277,10 @@ def execute_sweep(self):
 
         log_once("    Finished parameter sweep")
 
-        reader = ParamSweepReader(writer)
+        reader = writer.make_reader()
         self._block = reader.read_sweep_results()
+        if reader.is_file_based:
+            reader.finalize()
 
     def execute(self):
         """Load input data and execute all stages (before sweep, sweep, after sweep)"""

tests/data/test_param_sweep_store.py

@@ -1,45 +1,55 @@
+from os import PathLike
+
 import pytest
 import numpy as np
+from mpi4py.MPI import COMM_WORLD
 from pytest_mock import MockerFixture
 
 from httomo.data.param_sweep_store import ParamSweepWriter
 from httomo.runner.auxiliary_data import AuxiliaryData
+from httomo.runner.dataset_store_backing import DataSetStoreBacking
 from httomo.sweep_runner.param_sweep_block import ParamSweepBlock
 
 
-def make_param_sweep_writer() -> ParamSweepWriter:
+def make_param_sweep_writer(tmp_path: PathLike) -> ParamSweepWriter:
     NO_OF_SWEEPS = 5
-    return ParamSweepWriter(NO_OF_SWEEPS)
+    return ParamSweepWriter(NO_OF_SWEEPS, COMM_WORLD, tmp_path, DataSetStoreBacking.RAM)
 
 
-def test_param_sweep_writer_get_no_of_sweeps():
-    writer = make_param_sweep_writer()
+def test_param_sweep_writer_get_no_of_sweeps(tmp_path: PathLike):
+    writer = make_param_sweep_writer(tmp_path)
     assert writer.no_of_sweeps == 5
 
 
-def test_param_sweep_writer_get_concat_dim():
-    writer = make_param_sweep_writer()
+def test_param_sweep_writer_get_concat_dim(tmp_path: PathLike):
+    writer = make_param_sweep_writer(tmp_path)
     assert writer.concat_dim == 1
 
 
-def test_param_sweep_writer_get_single_shape_before_write_raises_error():
-    writer = make_param_sweep_writer()
+def test_param_sweep_writer_get_single_shape_before_write_raises_error(
+    tmp_path: PathLike,
+):
+    writer = make_param_sweep_writer(tmp_path)
     with pytest.raises(ValueError) as e:
         writer.single_shape
     assert (
         "Shape of single sweep result isn't known until the first write has occurred"
     ) in str(e)
 
 
-def test_param_sweep_writer_get_slices_per_sweep_before_write_raises_error():
-    writer = make_param_sweep_writer()
+def test_param_sweep_writer_get_slices_per_sweep_before_write_raises_error(
+    tmp_path: PathLike,
+):
+    writer = make_param_sweep_writer(tmp_path)
     with pytest.raises(ValueError) as e:
         writer.slices_per_sweep
     assert ("Slices per sweep isn't known until the first write has occurred") in str(e)
 
 
-def test_param_sweep_writer_get_total_shape_before_write_raises_error():
-    writer = make_param_sweep_writer()
+def test_param_sweep_writer_get_total_shape_before_write_raises_error(
+    tmp_path: PathLike,
+):
+    writer = make_param_sweep_writer(tmp_path)
     with pytest.raises(ValueError) as e:
         writer.total_shape
     assert (
@@ -48,19 +58,19 @@ def test_param_sweep_writer_get_total_shape_before_write_raises_error():
     ) in str(e)
 
 
-def test_param_sweep_write_make_reader_errors_if_data_none():
-    writer = make_param_sweep_writer()
+def test_param_sweep_write_make_reader_errors_if_data_none(tmp_path: PathLike):
+    writer = make_param_sweep_writer(tmp_path)
     with pytest.raises(ValueError) as e:
         writer.make_reader()
 
     assert "no data has been written yet" in str(e)
 
 
 def test_param_sweep_writer_write_sweep_result_transfers_gpu_arr_to_cpu(
-    mocker: MockerFixture,
+    mocker: MockerFixture, tmp_path: PathLike
 ):
     SWEEP_RES_SHAPE = (180, 3, 160)
-    writer = make_param_sweep_writer()
+    writer = make_param_sweep_writer(tmp_path)
     block = ParamSweepBlock(
         data=np.ones(SWEEP_RES_SHAPE, dtype=np.float32),
         aux_data=AuxiliaryData(angles=np.ones(SWEEP_RES_SHAPE[0], dtype=np.float32)),
@@ -77,7 +87,7 @@ def test_param_sweep_writer_write_sweep_result_transfers_gpu_arr_to_cpu(
     to_cpu_spy.assert_called_once()
 
 
-def test_param_sweep_writer_reader_write_res_and_read():
+def test_param_sweep_writer_reader_write_res_and_read(tmp_path: PathLike):
     CONCAT_DIM = 1
     SWEEP_RES_SHAPE = (180, 3, 160)
     NO_OF_SWEEPS = 2
@@ -86,7 +96,9 @@ def test_param_sweep_writer_reader_write_res_and_read():
         SWEEP_RES_SHAPE[1] * NO_OF_SWEEPS,
         SWEEP_RES_SHAPE[2],
     )
-    writer = ParamSweepWriter(NO_OF_SWEEPS)
+    writer = ParamSweepWriter(
+        NO_OF_SWEEPS, COMM_WORLD, tmp_path, DataSetStoreBacking.RAM
+    )
 
     # Define an array that will contain data representing the fake result of two parameter
     # sweep executions.
@@ -131,3 +143,58 @@ def test_param_sweep_writer_reader_write_res_and_read():
     np.testing.assert_array_equal(
         block_with_middle_slices.data[:, 1, :], SWEEP_RESULT_TWO_MID_SLICE
     )
+
+
+def test_writer_closes_file_on_finalize(mocker: MockerFixture, tmp_path: PathLike):
+    CONCAT_DIM = 1
+    BLOCK_SHAPE = (180, 3, 160)
+    block_data = np.arange(np.prod(BLOCK_SHAPE), dtype=np.float32).reshape(BLOCK_SHAPE)
+    aux_data = AuxiliaryData(np.ones(BLOCK_SHAPE[0], dtype=np.float32))
+    writer = ParamSweepWriter(
+        no_of_sweeps=1,
+        comm=COMM_WORLD,
+        temppath=tmp_path,
+        backing=DataSetStoreBacking.File,
+    )
+    writer.write_sweep_result(
+        ParamSweepBlock(
+            data=block_data,
+            aux_data=aux_data,
+            slicing_dim=CONCAT_DIM,
+        )
+    )
+    mock_fileclose = mocker.patch.object(writer.h5file, "close")
+    writer.finalize()
+    mock_fileclose.assert_called_once()
+
+
+def test_making_reader_closes_file_and_deletes(tmp_path: PathLike):
+    CONCAT_DIM = 1
+    BLOCK_SHAPE = (180, 3, 160)
+    block_data = np.arange(np.prod(BLOCK_SHAPE), dtype=np.float32).reshape(BLOCK_SHAPE)
+    aux_data = AuxiliaryData(np.ones(BLOCK_SHAPE[0], dtype=np.float32))
+    writer = ParamSweepWriter(
+        no_of_sweeps=1,
+        comm=COMM_WORLD,
+        temppath=tmp_path,
+        backing=DataSetStoreBacking.File,
+    )
+    writer.write_sweep_result(
+        ParamSweepBlock(
+            data=block_data,
+            aux_data=aux_data,
+            slicing_dim=CONCAT_DIM,
+        )
+    )
+    reader = writer.make_reader()
+
+    assert writer.h5file is None
+    assert writer.filename is not None
+    assert writer.filename.exists()
+    assert reader.filename == writer.filename
+    assert reader.h5file is not None
+    assert reader.h5file.get("data", None) is not None
+
+    reader.finalize()
+
+    assert not writer.filename.exists()