Add stoch writing routine

source_modelling/stoch.py

@@ -12,8 +12,9 @@
 .. [0] https://wiki.canterbury.ac.nz/display/QuakeCore/File+Formats+Used+In+Ground+Motion+Simulation#FileFormatsUsedInGroundMotionSimulation-Stochformat
 """
 
+import dataclasses
 from pathlib import Path
-from typing import NamedTuple, TextIO, TypeAlias, cast
+from typing import IO, NamedTuple, Self, TextIO, TypeAlias, cast
 
 import numpy as np
 from numpy.typing import NDArray
@@ -149,37 +150,18 @@ def _read_stoch_plane(handle: TextIO) -> StochPlane:
     return StochPlane(header, slip, rise, trup)
 
 
+@dataclasses.dataclass
 class StochFile:
     """
     Class for reading and accessing stochastic slip model files.
 
     This class handles the parsing of stochastic slip model files and provides
     access to the contained planes and their properties.
 
-    Parameters
-    ----------
-    filename : Path
-        Path to the stochastic slip model file.
-
-    Attributes
-    ----------
-    data : list[StochPlane]
-        Structured raw data read from the stoch file.
-
-    Raises
-    ------
-    ValueError
-        If the number of planes specified in the file is not a positive integer.
-
-    Notes
-    -----
-    Stochastic slip model files contain information about fault planes and their
-    properties such as slip, rise time, and rupture time.
-
     Examples
     --------
     >>> # Assuming 'stoch_model.stoch' exists with valid stochastic slip model data
-    >>> stoch_file = StochFile('stoch_model.stoch')
+    >>> stoch_file = StochFile.from_file('stoch_model.stoch')
     >>> planes = stoch_file.planes
     >>> slips = stoch_file.slip
     >>> rise_times = stoch_file.rise
@@ -190,15 +172,23 @@ class StochFile:
     ...     print(f"Slip array shape for the first plane: {slips[0].shape}")
     """
 
-    def __init__(self, filename: Path):
+    data: list[StochPlane]
+
+    @classmethod
+    def from_file(cls, filename: Path) -> Self:
         """
-        Initialize a StochFile instance by reading data from the specified file.
+        Initialise a StochFile instance by reading data from the specified file.
 
         Parameters
         ----------
         filename : Path
             Path to the stochastic slip model file.
 
+        Returns
+        -------
+        StochFile
+            The stoch file parsed from the filename.
+
         Raises
         ------
         ValueError
@@ -209,10 +199,36 @@ def __init__(self, filename: Path):
             n_planes = parse_utils.read_int(handle, "n_planes")
             if n_planes <= 0:
                 raise parse_utils.ParseError(
-                    f"Expected non-negative integer number of planes, received: {n_planes}."
+                    f"Expected positive integer number of planes, received: {n_planes}."
                 )
             planes = [_read_stoch_plane(handle) for _ in range(n_planes)]
-            self.data: list[StochPlane] = planes
+            return cls(planes)
+
+    def dump(self, handle: IO[str]) -> None:
+        """Write a stoch file to a file-like object.
+
+        Parameters
+        ----------
+        handle : File-like object
+          The object to write the stoch file output to.
+        """
+        if not self.data:
+            raise ValueError("Cannot dump empty stoch file")
+        handle.write(f"{len(self.data)}\n")
+        for plane in self.data:
+            header = plane.header
+            handle.write(
+                f"{header.longitude:10.4f} {header.latitude:10.4f} {header.nx:5d} {header.ny:5d} {header.dx:8.2f} {header.dy:8.2f}\n"
+            )
+            handle.write(
+                f"{header.strike:4.0f} {header.dip:4.0f} {header.average_rake:4.0f} {header.dtop:8.2f} {header.shypo:8.2f} {header.dhypo:8.2f}\n"
+            )
+            np.savetxt(
+                handle,
+                np.vstack((plane.slip, plane.rise, plane.trup)),
+                fmt="%13.5e",
+                delimiter="",
+            )
 
     @property
     def planes(self) -> list[Plane]:
@@ -226,7 +242,7 @@ def planes(self) -> list[Plane]:
 
         Examples
         --------
-        >>> stoch_file = StochFile('stoch_model.stoch')
+        >>> stoch_file = StochFile.from_file('stoch_model.stoch')
         >>> planes = stoch_file.planes
         >>> print(f"Number of planes: {len(planes)}")
         """

tests/test_stoch.py

@@ -9,6 +9,7 @@
 from source_modelling.stoch import (
     StochFile,
     StochHeader,
+    StochPlane,
     _read_stoch_header,
     _read_stoch_plane,
 )
@@ -124,7 +125,7 @@ def test_read_stoch_plane_from_file(sample_stoch_file_plane: Path):
 def test_stoch_file_initialization(sample_stoch_file: Path):
     """Test initializing a StochFile from a real file."""
     # Initialize the StochFile with the sample file
-    stoch_file = StochFile(sample_stoch_file)
+    stoch_file = StochFile.from_file(sample_stoch_file)
 
     # Verify the file was read correctly
     assert len(stoch_file.data) == 1
@@ -149,7 +150,7 @@ def test_stoch_file_initialization(sample_stoch_file: Path):
 def test_stoch_file_properties(sample_stoch_file: Path):
     """Test the properties of a StochFile."""
     # Initialize the StochFile with the sample file
-    stoch_file = StochFile(sample_stoch_file)
+    stoch_file = StochFile.from_file(sample_stoch_file)
 
     # Test the slip property
     slip_arrays = stoch_file.slip
@@ -182,7 +183,7 @@ def test_stoch_file_properties(sample_stoch_file: Path):
 def test_multiple_planes(sample_stoch_file_multi_plane: Path):
     """Test reading a file with multiple planes."""
     # Initialize the StochFile with the multi-plane sample file
-    stoch_file = StochFile(sample_stoch_file_multi_plane)
+    stoch_file = StochFile.from_file(sample_stoch_file_multi_plane)
 
     # Verify the file was read correctly
     assert len(stoch_file.data) == 2
@@ -212,7 +213,7 @@ def test_multiple_planes(sample_stoch_file_multi_plane: Path):
 
 def test_real_world_stoch():
     """Test reading a real-world stochastic file."""
-    stoch_file = StochFile(STOCH_PATH)
+    stoch_file = StochFile.from_file(STOCH_PATH)
 
     # Check the plane headers
     headers = [plane.header for plane in stoch_file.data]
@@ -277,6 +278,135 @@ def test_stoch_file_invalid_planes(bad_header_file: Path):
 
     # Test that ValueError is raised
     with pytest.raises(
-        parse_utils.ParseError, match="Expected non-negative integer number of planes"
+        parse_utils.ParseError, match="Expected positive integer number of planes"
     ):
-        StochFile(file_path)
+        StochFile.from_file(file_path)
+
+
+def dump_and_reload(stoch_file: StochFile, tmp_path: Path) -> StochFile:
+    """Dump a StochFile and read the result back in."""
+    dumped = tmp_path / "dumped.stoch"
+    with open(dumped, "w") as handle:
+        stoch_file.dump(handle)
+    return StochFile.from_file(dumped)
+
+
+def assert_stoch_files_equal(left: StochFile, right: StochFile) -> None:
+    """Assert two StochFile objects hold equivalent data."""
+    assert len(left.data) == len(right.data)
+    for left_plane, right_plane in zip(left.data, right.data):
+        assert left_plane.header == right_plane.header
+        assert left_plane.slip == pytest.approx(right_plane.slip, rel=1e-5)
+        assert left_plane.rise == pytest.approx(right_plane.rise, rel=1e-5)
+        assert left_plane.trup == pytest.approx(right_plane.trup, rel=1e-5)
+
+
+@pytest.fixture
+def simple_stoch() -> StochFile:
+    """A single-plane StochFile built in memory."""
+    header = StochHeader(174.5, -41.3, 3, 2, 1.0, 1.0, 45, 60, 90, 0.5, 2.5, 1.5)
+    slip = np.array([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]], dtype=np.float32)
+    rise = np.array([[0.1, 0.2, 0.3], [0.4, 0.5, 0.6]], dtype=np.float32)
+    trup = np.array([[0.01, 0.02, 0.03], [0.04, 0.05, 0.06]], dtype=np.float32)
+    return StochFile([StochPlane(header, slip, rise, trup)])
+
+
+def test_dump_exact_format(simple_stoch: StochFile):
+    """Test that dump writes the expected fixed-width columns."""
+    handle = io.StringIO()
+    simple_stoch.dump(handle)
+    lines = handle.getvalue().split("\n")
+
+    # One plane count line, two header lines, and 3 * ny data lines.
+    assert lines[-1] == ""  # trailing newline
+    assert len(lines) == 1 + 1 + 2 + 3 * 2
+
+    assert lines[0] == "1"
+    assert lines[1] == "  174.5000   -41.3000     3     2     1.00     1.00"
+    assert lines[2] == "  45   60   90     0.50     2.50     1.50"
+    assert lines[3] == "  1.00000e+00  2.00000e+00  3.00000e+00"
+    assert lines[4] == "  4.00000e+00  5.00000e+00  6.00000e+00"
+    assert lines[5] == "  1.00000e-01  2.00000e-01  3.00000e-01"
+    assert lines[6] == "  4.00000e-01  5.00000e-01  6.00000e-01"
+    assert lines[7] == "  1.00000e-02  2.00000e-02  3.00000e-02"
+    assert lines[8] == "  4.00000e-02  5.00000e-02  6.00000e-02"
+
+
+def test_dump_round_trip_in_memory(simple_stoch: StochFile, tmp_path: Path):
+    """Test that dumping and re-reading an in-memory StochFile is lossless."""
+    assert_stoch_files_equal(simple_stoch, dump_and_reload(simple_stoch, tmp_path))
+
+
+def test_dump_round_trip(sample_stoch_file: Path, tmp_path: Path):
+    """Test that a parsed stoch file survives a dump/read round trip."""
+    stoch_file = StochFile.from_file(sample_stoch_file)
+    assert_stoch_files_equal(stoch_file, dump_and_reload(stoch_file, tmp_path))
+
+
+def test_dump_round_trip_multi_plane(
+    sample_stoch_file_multi_plane: Path, tmp_path: Path
+):
+    """Test that a multi-plane stoch file survives a dump/read round trip."""
+    stoch_file = StochFile.from_file(sample_stoch_file_multi_plane)
+    reloaded = dump_and_reload(stoch_file, tmp_path)
+    assert len(reloaded.data) == 2
+    assert_stoch_files_equal(stoch_file, reloaded)
+
+
+def test_dump_round_trip_real_world(tmp_path: Path):
+    """Test that the real-world stoch file survives a dump/read round trip."""
+    stoch_file = StochFile.from_file(STOCH_PATH)
+    reloaded = dump_and_reload(stoch_file, tmp_path)
+    assert_stoch_files_equal(stoch_file, reloaded)
+    # Planes have differing dimensions, so check the shapes are preserved too.
+    assert [(plane.header.ny, plane.header.nx) for plane in reloaded.data] == [
+        slip.shape for slip in reloaded.slip
+    ]
+
+
+def test_dump_is_idempotent(tmp_path: Path):
+    """Test that dumping an already-dumped file produces identical output."""
+    stoch_file = StochFile.from_file(STOCH_PATH)
+    first = io.StringIO()
+    stoch_file.dump(first)
+    second = io.StringIO()
+    dump_and_reload(stoch_file, tmp_path).dump(second)
+    assert first.getvalue() == second.getvalue()
+
+
+def test_dump_negative_values_remain_separated(tmp_path: Path):
+    """Test that negative array values do not run their columns together."""
+    header = StochHeader(174.5, -41.3, 2, 1, 1.0, 1.0, 45, 60, 90, 0.5, 2.5, 1.5)
+    negative = np.array([[-1.5, -2.5]], dtype=np.float32)
+    stoch_file = StochFile([StochPlane(header, negative, negative, negative)])
+
+    handle = io.StringIO()
+    stoch_file.dump(handle)
+    assert " -1.50000e+00 -2.50000e+00" in handle.getvalue()
+
+    assert_stoch_files_equal(stoch_file, dump_and_reload(stoch_file, tmp_path))
+
+
+def test_dump_sentinel_hypocentre_values(tmp_path: Path):
+    """Test that the -999 hypocentre sentinels survive a round trip."""
+    header = StochHeader(
+        172.1255, -43.5728, 2, 1, 2.0, 2.0, 35, 70, 39, 0.97, -999.0, -999.0
+    )
+    data = np.array([[1.0, 2.0]], dtype=np.float32)
+    stoch_file = StochFile([StochPlane(header, data, data, data)])
+
+    reloaded = dump_and_reload(stoch_file, tmp_path)
+    assert reloaded.data[0].header.shypo == -999.0
+    assert reloaded.data[0].header.dhypo == -999.0
+    assert_stoch_files_equal(stoch_file, reloaded)
+
+
+def test_dump_plane_count_matches_data(sample_stoch_file_multi_plane: Path):
+    """Test that the plane count written matches the number of planes dumped."""
+    stoch_file = StochFile.from_file(sample_stoch_file_multi_plane)
+    handle = io.StringIO()
+    stoch_file.dump(handle)
+    lines = handle.getvalue().splitlines()
+    assert int(lines[0]) == len(stoch_file.data)
+    expected_lines = 1 + sum(2 + 3 * plane.header.ny for plane in stoch_file.data)
+    assert len(lines) == expected_lines