diff --git a/source_modelling/stoch.py b/source_modelling/stoch.py index cf72b2b6..32368bd8 100644 --- a/source_modelling/stoch.py +++ b/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,6 +150,7 @@ 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. @@ -156,30 +158,10 @@ class StochFile: 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)}") """ diff --git a/tests/test_stoch.py b/tests/test_stoch.py index c345cee3..2ccb6191 100644 --- a/tests/test_stoch.py +++ b/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