Realisation config plumbing, and SW4 station coordinates

pyproject.toml

@@ -34,6 +34,7 @@ dependencies = [
   "schema",        # For loading realisations
   "structlog",     # Logging.
   "psutil",        # To get the CPU affinity for jobs
+  "h5py>=3.15.1",
   "parse>=1.21.0",
   "rich>=14.3.2",
 ]
@@ -66,6 +67,7 @@ generate-station-coordinates = "workflow.scripts.generate_station_coordinates:ap
 generate-model-coordinates = "workflow.scripts.generate_model_coordinates:app"
 generate-rupture-propagation = "workflow.scripts.generate_rupture_propagation:app"
 copy-domain-parameters = "workflow.scripts.copy_velocity_model_parameters:app"
+srf-to-hdf5 = "workflow.scripts.srf_to_hdf5:app"
 create-e3d-par = "workflow.scripts.create_e3d_par:app"
 generate-stoch = "workflow.scripts.generate_stoch:app"
 merge-ts = "workflow.scripts.merge_ts:app"

tests/test_hf.py

@@ -48,8 +48,6 @@ def test_build_hf_input_serialisation() -> None:
         t_sec=0.0,
         site_specific=False,
         dpath_pert=0,
-        stoch_dx=2.0,
-        stoch_dy=2.0,
         stress_parameter_adjustment_fault_area=None,
         stress_parameter_adjustment_target_magnitude=None,
         stress_parameter_adjustment_tect_type=0,

tests/test_realisation.py

@@ -115,6 +115,7 @@ def test_srf_config_example(tmp_path: Path) -> None:
     )
     srf_config = realisations.SRFConfig(
         resolution=0.1,
+        dt=0.005,
         point_source_params=schemas.PointSourceParams(
             stype=schemas.Stype.cos,
             risetime=0.5,
@@ -249,6 +250,7 @@ def test_srf_config_example(tmp_path: Path) -> None:
             },
             "srf": {
                 "resolution": 0.1,
+                "dt": 0.005,
                 "point_source_params": {
                     "stype": "cos",
                     "risetime": 0.5,

tests/test_realisation_to_srf.py

@@ -8,6 +8,7 @@
 def test_build_genslip_command_static_args() -> None:
     srf_config = SRFConfig(
         resolution=0.1,
+        dt=0.005,
         point_source_params=schemas.PointSourceParams(
             stype=schemas.Stype.cos,
             risetime=0.5,

workflow/default_parameters/root/defaults.yaml

@@ -93,6 +93,7 @@ hf:
   stress_parameter_adjustment_tect_type: 0
   stress_parameter_adjustment_target_magnitude: null
   stress_parameter_adjustment_fault_area: null
+stoch:
   stoch_dx: 2.0
   stoch_dy: 2.0
 rupture_velocity:
@@ -163,6 +164,7 @@ srf:
   read_erf: false
   read_gsf: true
   resolution: 0.1
+  dt: 0.005
   risetime_coef: 1.6
   risetimedep: 6.5
   risetimedep_range: 1.5

workflow/realisations.py

@@ -77,6 +77,25 @@ class RealisationParseError(Exception):
     """Realisation JSON parse error."""
 
 
+def path_serialiser(obj: Any) -> Any:
+    """Serialise `Path` values, which `json` does not handle natively.
+
+    Parameters
+    ----------
+    obj : Any
+        The object `json` could not serialise.
+
+    Returns
+    -------
+    Any
+        The string form of `obj` if it is a `Path`, and `obj` unchanged
+        otherwise.
+    """
+    if isinstance(obj, Path):
+        return str(obj)
+    return obj
+
+
 @dataclasses.dataclass
 class RealisationConfiguration(ABC):
     """Abstract base class for RealisationConfiguration."""
@@ -220,7 +239,12 @@ def write_to_realisation(
                 realisation_configuration = json.load(realisation_file_handle)
         realisation_configuration.update({self._config_key: self.to_dict()})
         with open(realisation_ffp, "w", encoding="utf-8") as realisation_file_handle:
-            json.dump(realisation_configuration, realisation_file_handle, indent=4)
+            json.dump(
+                realisation_configuration,
+                realisation_file_handle,
+                indent=4,
+                default=path_serialiser,
+            )
 
 
 @dataclasses.dataclass
@@ -373,6 +397,9 @@ class SRFConfig(RealisationConfiguration):
     resolution: float
     """The resolution of the SRF discretisation (different, in general, from the simulation resolution)."""
 
+    dt: float
+    """SRF temporal resolution (timestep)."""
+
     point_source_params: schemas.PointSourceParams | None
     """Parameters for point source approximation, if applicable."""
 
@@ -1049,11 +1076,23 @@ class HFConfig(RealisationConfiguration):
     """Target magnitude (or inferred if None)"""
     stress_parameter_adjustment_fault_area: float | None
     """Target magnitude (or inferred if None)"""
-    # these are used in stoch generation, rather than HF invocation
+
+
+@dataclasses.dataclass
+class StochConfig(RealisationConfiguration):
+    """Stoch file generation.
+
+    Not part of :class:`HFConfig`: these size the stoch grid, which is an input to the
+    high-frequency simulation rather than one of its parameters.
+    """
+
+    _config_key: ClassVar[str] = "stoch"
+    _schema: ClassVar[Schema] = schemas.STOCH_CONFIG_SCHEMA
+
     stoch_dx: float
     """stoch file resolution in x."""
     stoch_dy: float
-    """stoch file resolution in x."""
+    """stoch file resolution in y."""
 
 
 @dataclasses.dataclass

workflow/schemas.py

@@ -428,6 +428,9 @@ def _corners_to_array(corners_spec: list[dict[str, float]]) -> np.ndarray:
         Literal(
             "resolution", description="The resolution of the SRF discretisation."
         ): And(NUMBER, _is_positive),
+        Literal("dt", description="SRF temporal resolution (timestep)."): And(
+            NUMBER, _is_positive
+        ),
         Optional(
             Literal(
                 "point_source_params",
@@ -1069,6 +1072,11 @@ def _corners_to_array(corners_spec: list[dict[str, float]]) -> np.ndarray:
         Literal(
             "stress_parameter_adjustment_fault_area", "Fault area (or inferred if null)"
         ): Or(NUMBER, None),
+    }
+)
+
+STOCH_CONFIG_SCHEMA = Schema(
+    {
         Literal("stoch_dx", description="Stoch file dx"): And(NUMBER, _is_positive),
         Literal("stoch_dy", description="Stoch file dy"): And(NUMBER, _is_positive),
     }

workflow/scripts/generate_station_coordinates.py

@@ -28,10 +28,14 @@
 See the output of `generate-station-coordinates --help`.
 """
 
+from enum import StrEnum, auto
 from pathlib import Path
 from typing import Annotated
 
+import h5py
+import numpy as np
 import pandas as pd
+import shapely
 import typer
 
 from qcore import cli, coordinates
@@ -41,29 +45,66 @@
 app = typer.Typer()
 
 
-@cli.from_docstring(app)
-@log_utils.log_call()
-def generate_fd_files(
-    realisation_ffp: Annotated[Path, typer.Argument(readable=True, dir_okay=False)],
-    stat_file: Annotated[Path, typer.Argument(readable=True, dir_okay=False)],
-    output_path: Annotated[Path, typer.Argument(file_okay=False, writable=True)],
+class Format(StrEnum):
+    """The solver the station files are written for."""
+
+    EMOD3D = auto()
+    """EMOD3D `.statcords` grid-point files."""
+    SW4 = auto()
+    """SW4 HDF5 station file."""
+
+
+def write_ascii_station_locations(
+    stations: pd.DataFrame,
+    ll_out: Path,
+    lon_col: str = "lon",
+    lat_col: str = "lat",
+    name_col: str = "name",
 ) -> None:
-    """Generate station coordinate files.
+    """Write stations to a whitespace-separated `lon lat name` file.
+
+    Both solvers take the same `.ll` station list, so it is written once
+    here from whichever coordinate columns the caller has.
 
     Parameters
     ----------
-    realisation_ffp : Path
-        Path to realisation json file.
-    stat_file : Path
-        The location of the station files.
+    stations : pd.DataFrame
+        The stations to write.
+    ll_out : Path
+        Path of the file to write.
+    lon_col, lat_col, name_col : str, optional
+        Columns of `stations` holding the longitude, latitude and name.
+    """
+    with open(ll_out, "w", encoding="utf-8") as llf:
+        stations.apply(
+            lambda station: llf.write(
+                f"{station[lon_col]:11.5f} {station[lat_col]:11.5f} {station[name_col]}\n"
+            ),
+            axis=1,
+        )
+
+
+def write_emod3d_station_format(
+    domain_parameters: DomainParameters,
+    resolution_parameters: Resolution,
+    stations: pd.DataFrame,
+    output_path: Path,
+) -> None:
+    """Write station coordinates in EMOD3D format to two output files.
+
+    Parameters
+    ----------
+    domain_parameters : DomainParameters
+        Object containing domain definition and methods to compute grid dimensions.
+    resolution_parameters : Resolution
+        Object containing the grid resolution.
+    stations : pd.DataFrame
+        DataFrame with columns `lat`, `lon`, and `name`. Latitude and longitude
+        are in degrees.
     output_path : Path
-        Output path for station files.
+        Directory path where the output files will be written.
     """
-    output_path.mkdir(exist_ok=True)
-    domain_parameters = DomainParameters.read_from_realisation(realisation_ffp)
-    resolution_parameters = Resolution.read_from_realisation(realisation_ffp)
     domain = domain_parameters.domain
-
     nx = domain_parameters.nx(resolution_parameters.resolution)
     ny = domain_parameters.ny(resolution_parameters.resolution)
     mlat, mlon = domain.origin
@@ -74,14 +115,6 @@ def generate_fd_files(
     gp_out = output_path / "stations.statcords"
     ll_out = output_path / "stations.ll"
 
-    # retrieve in station names, latitudes and longitudes
-    stations = pd.read_csv(
-        stat_file,
-        delimiter=r"\s+",
-        comment="#",
-        names=["lon", "lat", "name"],
-    )
-
     x, y = proj(
         lat=stations["lat"].to_numpy(float), lon=stations["lon"].to_numpy(float)
     ).T
@@ -127,13 +160,93 @@ def generate_fd_files(
             axis=1,
         )
 
-    # create ll file
-    with open(ll_out, "w", encoding="utf-8") as llf:
-        stations.apply(
-            lambda station: llf.write(
-                f"{station['grid_lon']:11.5f} {station['grid_lat']:11.5f} {station['name']}\n"
-            ),
-            axis=1,
-        )
+    write_ascii_station_locations(
+        stations, ll_out, lon_col="grid_lon", lat_col="grid_lat"
+    )
+
+
+def write_sw4_station_format(
+    domain_parameters: DomainParameters, stations: pd.DataFrame, output_path: Path
+) -> None:
+    """Write station coordinates in SW4 format.
+
+    Stations outside the domain are dropped rather than clamped to its
+    edge: unlike EMOD3D's grid-point files there is no nearest gridpoint
+    to snap to, and a station SW4 cannot place is not a recording.
+
+    Parameters
+    ----------
+    domain_parameters : DomainParameters
+        Domain definition, used to test which stations fall inside it.
+    stations : pd.DataFrame
+        DataFrame with columns `lat`, `lon`, and `name`. Latitude and
+        longitude are in degrees.
+    output_path : Path
+        Directory to write `stations.h5` and `stations.ll` into.
+    """
+    lat_lon = stations[["lat", "lon"]].to_numpy()
+    nzvm_coordinates = coordinates.wgs_depth_to_nztm(lat_lon)
+    poly = domain_parameters.domain.polygon
+    mask = shapely.contains_xy(poly, nzvm_coordinates[:, 0], nzvm_coordinates[:, 1])
+    stations = stations.loc[mask]
+
+    if len(stations) == 0:
+        raise ValueError("No stations in domain.")
+
+    with h5py.File(output_path / "stations.h5", "w") as f:
+        for station_name, position in stations.set_index("name").iterrows():
+            station_dset = f.create_group(station_name)
+            location = station_dset.create_dataset(
+                "STLA,STLO,STDP", (3,), dtype=np.float64
+            )
+            location[0] = position["lat"]
+            location[1] = position["lon"]
+            location[2] = 0.0
+
+    write_ascii_station_locations(stations, output_path / "stations.ll")
+
+
+@cli.from_docstring(app)
+@log_utils.log_call()
+def generate_fd_files(
+    realisation_ffp: Annotated[Path, typer.Argument(readable=True, dir_okay=False)],
+    stat_file: Annotated[Path, typer.Argument(readable=True, dir_okay=False)],
+    output_path: Annotated[Path, typer.Argument(file_okay=False, writable=True)],
+    format: Format = Format.EMOD3D,
+) -> None:
+    """Generate station coordinate files.
+
+    Parameters
+    ----------
+    realisation_ffp : Path
+        Path to realisation json file.
+    stat_file : Path
+        The location of the station files.
+    output_path : Path
+        Output path for station files.
+    format : Format, optional
+        The solver to write station files for. EMOD3D writes
+        `stations.statcords` and `stations.ll`; SW4 writes `stations.h5`
+        and `stations.ll`. Defaults to EMOD3D.
+    """
+    output_path.mkdir(exist_ok=True)
+    domain_parameters = DomainParameters.read_from_realisation(realisation_ffp)
+    resolution_parameters = Resolution.read_from_realisation(realisation_ffp)
+
+    # retrieve in station names, latitudes and longitudes
+    stations = pd.read_csv(
+        stat_file,
+        delimiter=r"\s+",
+        comment="#",
+        names=["lon", "lat", "name"],
+    )
+
+    match format:
+        case Format.EMOD3D:
+            write_emod3d_station_format(
+                domain_parameters, resolution_parameters, stations, output_path
+            )
+        case Format.SW4:
+            write_sw4_station_format(domain_parameters, stations, output_path)
 
     realisations.append_log_entry(realisation_ffp)