diff --git a/tests/test_im_calc.py b/tests/test_im_calc.py new file mode 100644 index 00000000..2042abfb --- /dev/null +++ b/tests/test_im_calc.py @@ -0,0 +1,179 @@ +"""Pure-function tests for the intensity measure calculation. + +Nothing here runs OpenQuake or an IM kernel: what is tested is the metadata +plumbing that carries the SW4 supergrid (absorbing layer) penetration from the +waveform file into the intensity measure file, which is arithmetic-free and +where the silent failures live. +""" + +import numpy as np +import pytest +import xarray as xr + +from workflow.scripts import im_calc + + +def _waveform_dataset( + n_stations: int = 3, + supergrid_depth: list[float] | None = None, + attrs: dict[str, float] | None = None, +) -> xr.Dataset: + """A minimal stand-in for an opened broadband/LF waveform file.""" + stations = [f"ST{index:02d}" for index in range(n_stations)] + coords: dict[str, object] = {"station": stations} + if supergrid_depth is not None: + coords["supergrid_depth"] = ( + "station", + np.array(supergrid_depth, dtype=np.float32), + ) + coords["supergrid_depth_gp"] = ( + "station", + np.array(supergrid_depth, dtype=np.float32) / 400.0, + ) + return xr.Dataset( + {"waveform": (("station",), np.ones(n_stations, dtype=np.float32))}, + coords=coords, + attrs=attrs or {}, + ) + + +def test_a_solver_with_no_absorbing_layer_still_gets_the_coordinate() -> None: + """Every IM file carries the coordinate, whatever the solver produced it. + + EMOD3D and the high-frequency simulation have no supergrid at all, so the + value has to be NaN -- "not applicable / not reported" -- and never `0.0`, + which would assert the station had been checked and found in the interior. + """ + supergrid = im_calc.supergrid_coordinates(_waveform_dataset(n_stations=4)) + + assert set(supergrid) == set(im_calc.SUPERGRID_COORDINATES) + for values in supergrid.values(): + assert values.dims == ("station",) + assert values.dtype == np.float32 + assert np.isnan(values.values).all() + + +def test_a_reported_penetration_is_passed_through_unchanged() -> None: + """The three states survive verbatim: clean, flagged, and unknown.""" + dataset = _waveform_dataset(supergrid_depth=[0.0, 5750.0, np.nan]) + + supergrid = im_calc.supergrid_coordinates(dataset) + + depth = supergrid["supergrid_depth"] + assert depth.values[0] == 0.0 + assert depth.values[1] == pytest.approx(5750.0) + assert np.isnan(depth.values[2]) + # `> 0` is the documented threshold; NaN must not satisfy it. + np.testing.assert_array_equal(depth.values > 0, [False, True, False]) + + +def test_the_coordinate_is_read_eagerly() -> None: + """`im-calc` opens the waveform file chunked, so the coordinate arrives as + a dask array whose "auto" station chunking differs from the waveform's. + Loading it here keeps that mismatch out of the attached coordinates, the + way `vs30` is loaded eagerly for the same reason. + """ + dataset = _waveform_dataset(supergrid_depth=[0.0, 1.0, 2.0]).chunk({"station": 1}) + + supergrid = im_calc.supergrid_coordinates(dataset) + + assert supergrid["supergrid_depth"].chunks is None + + +def test_nothing_is_claimed_about_a_run_that_reported_nothing() -> None: + """An all-NaN flag must not put `absorbing_layer` in the root attributes. + + Writing it would claim the run had an absorbing layer that somebody + measured, on the strength of a coordinate that says only "unknown". + """ + dataset = _waveform_dataset(attrs={"SGWIDTH": 12000.0}) + supergrid = im_calc.supergrid_coordinates(dataset) + + assert im_calc.supergrid_attributes(dataset, supergrid) == {} + + +def test_the_sponge_width_comes_from_the_waveform_file() -> None: + """Not from the realisation configuration. + + The configuration is editable after a run; the waveform file is what SW4 + actually wrote. Reading the config here would let the IM file's + self-description drift away from the run it describes. + """ + dataset = _waveform_dataset( + supergrid_depth=[0.0, 5750.0, 0.0], + attrs={"SGWIDTH": 12000.0, "SGWIDTHGP": 30.0}, + ) + supergrid = im_calc.supergrid_coordinates(dataset) + + attributes = im_calc.supergrid_attributes(dataset, supergrid) + + assert attributes["absorbing_layer"] == "sw4_supergrid" + assert attributes["absorbing_layer_width_m"] == pytest.approx(12000.0) + assert attributes["absorbing_layer_width_gp"] == pytest.approx(30.0) + + +def test_a_flag_without_a_width_still_names_the_layer() -> None: + """A station file written with penetrations but no file-level width (or a + broadband file whose width did not survive) must not lose the layer name. + """ + dataset = _waveform_dataset(supergrid_depth=[0.0, 5750.0, 0.0]) + supergrid = im_calc.supergrid_coordinates(dataset) + + attributes = im_calc.supergrid_attributes(dataset, supergrid) + + assert attributes == {"absorbing_layer": "sw4_supergrid"} + + +def test_the_flag_is_attached_as_a_coordinate_on_every_leaf() -> None: + """A data variable here would die at `bb_sim`'s `combined` dataset on the + next run through the pipeline, and cannot be selected alongside an IM. + The root deliberately has no station dimension, so it stays untouched. + """ + dataset = _waveform_dataset(supergrid_depth=[0.0, 5750.0, np.nan]) + supergrid = im_calc.supergrid_coordinates(dataset) + dtree = xr.DataTree.from_dict( + { + "PGA": xr.Dataset( + {"rotd50": (("station",), np.ones(3))}, + coords={"station": dataset.station}, + ), + "pSA": xr.Dataset( + {"rotd50": (("station",), np.ones(3))}, + coords={"station": dataset.station}, + ), + } + ) + + parameterised = im_calc.add_station_parameters(dtree, supergrid) + + for group in ("PGA", "pSA"): + leaf = parameterised[group].dataset + for name in im_calc.SUPERGRID_COORDINATES: + assert name in leaf.coords + assert name not in leaf.data_vars + assert "supergrid_depth" not in parameterised.dataset.coords + + +def test_the_recommended_threshold_travels_with_the_data() -> None: + """The coordinate's `description` is the only documentation a downstream + user ever sees, so it has to carry the threshold and say what the trace + is, or every tool invents its own cut. + """ + for name in im_calc.SUPERGRID_COORDINATES: + assert name in im_calc.COORDINATE_METADATA + description = im_calc.COORDINATE_METADATA[name]["description"] + assert "> 0" in description + assert "NaN" in description + + dataset = xr.Dataset( + {"rotd50": (("station",), np.ones(2))}, + coords={ + "station": ["ST00", "ST01"], + "supergrid_depth": ("station", np.array([0.0, 1.0], dtype=np.float32)), + }, + ) + annotated = im_calc.add_units(xr.DataTree.from_dict({"PGA": dataset})) + + attrs = annotated["PGA"].dataset["supergrid_depth"].attrs + assert attrs["units"] == "m" + assert "> 0" in attrs["description"] diff --git a/tests/test_lf_to_xarray.py b/tests/test_lf_to_xarray.py new file mode 100644 index 00000000..66a8c46d --- /dev/null +++ b/tests/test_lf_to_xarray.py @@ -0,0 +1,212 @@ +"""Tests for reading SW4 station recordings, and in particular for the +supergrid (absorbing layer) penetration SW4 reports per station. + +The station file fixture below is the first synthetic SW4 recording in the +suite; it is deliberately written against the layout documented in Section +12.9 of the SW4 User Guide (a root `DELTA`, one group per station holding +`NPTS`, `STLA,STLO,STDP` and the three geographic components) so that it is +reusable for any other SW4-read test. +""" + +from pathlib import Path + +import h5py +import numpy as np +import pytest +import xarray as xr + +from workflow.scripts import lf_to_xarray + + +def write_sw4_station_file( + path: Path, + stations: dict[str, dict[str, float] | None], + npts: int = 8, + dt: float = 0.05, + widths: dict[str, float] | None = None, +) -> Path: + """Write a synthetic SW4 HDF5 station recording. + + Parameters + ---------- + path : Path + Where to write the file. + stations : dict + Map from station name to either `None` (an old-style station, with no + supergrid datasets at all) or a dict which may hold `SGDEPTH` and + `SGDEPTHGP`. A dict holding only one of the two produces a + deliberately corrupt station. + npts : int + Number of samples per component. + dt : float + Sample spacing, written to the root `DELTA`. + widths : dict, optional + File-level scalars (`SGWIDTH`, `SGWIDTHGP`) written beside `DELTA`. + + Returns + ------- + Path + The path written, for convenience. + """ + with h5py.File(path, "w") as handle: + handle.create_dataset("DELTA", data=np.array([dt])) + for name, value in (widths or {}).items(): + handle.create_dataset(name, data=np.array([value])) + for index, (station, supergrid) in enumerate(stations.items()): + group = handle.create_group(station) + group.create_dataset("NPTS", data=np.array([npts])) + group.create_dataset( + "STLA,STLO,STDP", + data=np.array([-43.5 + index, 172.6 + index, 0.0]), + ) + for component in ("EW", "NS", "UP"): + group.create_dataset( + component, data=np.arange(npts, dtype=np.float32) + index + ) + for key, value in (supergrid or {}).items(): + group.create_dataset(key, data=np.array([value])) + return path + + +def test_supergrid_penetration_arrives_as_float32_coordinates(tmp_path: Path) -> None: + """The flag must be a *coordinate*, and it must be floating point. + + Both halves are load bearing. A station-dimension coordinate rides + through `bb-sim` and `im-calc` untouched, whereas a data variable is + silently dropped by `bb_sim._process_bb_chunk`, so a data variable here + would mean the flag never reaches the intensity measures. And the + downstream consumer opens IM files with `mask_and_scale=False`, so an + integer with a `_FillValue` would read back raw and become a plausible + penetration depth; only a real float NaN survives that. + """ + ffp = write_sw4_station_file( + tmp_path / "stations.h5", + { + "AAAA": {"SGDEPTH": 0.0, "SGDEPTHGP": 0.0}, + "BBBB": {"SGDEPTH": 5750.0, "SGDEPTHGP": 14.375}, + }, + ) + + dset = lf_to_xarray.read_station_metadata(ffp) + + for name in ("supergrid_depth", "supergrid_depth_gp"): + assert name in dset.coords + assert name not in dset.data_vars + assert dset.coords[name].dims == ("station",) + assert dset.coords[name].dtype == np.float32 + + ordered = dset.sortby("station") + np.testing.assert_array_equal(ordered["supergrid_depth"].values, [0.0, 5750.0]) + np.testing.assert_allclose( + ordered["supergrid_depth_gp"].values, [0.0, 14.375], rtol=1e-6 + ) + + +def test_an_old_station_file_converts_with_an_all_nan_flag(tmp_path: Path) -> None: + """A file written before SW4 reported the supergrid must not raise. + + This is the common case for every recording made so far, so it has to be + a no-op rather than an error. The value must be NaN and never `0.0`: `0.0` + is the positive claim "this station was checked and is in the interior", + which nobody checked here. + """ + ffp = write_sw4_station_file( + tmp_path / "old.h5", {"AAAA": None, "BBBB": None, "CCCC": None} + ) + + dset = lf_to_xarray.read_station_metadata(ffp) + + assert dset.sizes["station"] == 3 + for name in ("supergrid_depth", "supergrid_depth_gp"): + assert name in dset.coords + assert dset.coords[name].dtype == np.float32 + assert np.isnan(dset.coords[name].values).all() + assert "SGWIDTH" not in dset.attrs + assert "SGWIDTHGP" not in dset.attrs + + +def test_stations_missing_the_flag_are_nan_not_zero(tmp_path: Path) -> None: + """Mixed groups: only the stations SW4 reported on get a number.""" + ffp = write_sw4_station_file( + tmp_path / "mixed.h5", + { + "AAAA": {"SGDEPTH": 0.0, "SGDEPTHGP": 0.0}, + "BBBB": None, + "CCCC": {"SGDEPTH": 1200.0, "SGDEPTHGP": 3.0}, + }, + ) + + depth = lf_to_xarray.read_station_metadata(ffp).sortby("station")["supergrid_depth"] + + assert depth.values[0] == 0.0 + assert np.isnan(depth.values[1]) + assert depth.values[2] == 1200.0 + + +def test_one_dataset_without_the_other_is_a_corrupt_file(tmp_path: Path) -> None: + """`SGDEPTHGP` missing while `SGDEPTH` is present is corruption, not age. + + The back-compatibility guard is deliberately on `SGDEPTH` alone, so this + raises rather than quietly reporting a metre depth with no grid-point + depth beside it. + """ + ffp = write_sw4_station_file(tmp_path / "corrupt.h5", {"AAAA": {"SGDEPTH": 900.0}}) + + with pytest.raises(KeyError): + lf_to_xarray.read_station_metadata(ffp) + + +def test_the_sponge_width_is_lifted_into_the_dataset_attributes( + tmp_path: Path, +) -> None: + """`SGWIDTH`/`SGWIDTHGP` make the file self-describing. + + They are what turns the penetration into a severity fraction downstream, + and taking them from the file rather than from the realisation + configuration is the point: the configuration can be edited after the run. + """ + ffp = write_sw4_station_file( + tmp_path / "width.h5", + {"AAAA": {"SGDEPTH": 0.0, "SGDEPTHGP": 0.0}}, + widths={"SGWIDTH": 12000.0, "SGWIDTHGP": 30.0}, + ) + + dset = lf_to_xarray.read_station_metadata(ffp) + + assert dset.attrs["SGWIDTH"] == pytest.approx(12000.0) + assert dset.attrs["SGWIDTHGP"] == pytest.approx(30.0) + # The pre-existing attributes must survive alongside them. + assert dset.attrs["nt"] == 8 + assert dset.attrs["dt"] == pytest.approx(0.05) + + +def test_the_flag_survives_a_netcdf_round_trip(tmp_path: Path) -> None: + """As a coordinate, and as NaN -- checked the way a consumer reads it. + + `mask_and_scale=False` is what `eqvis`'s `open_ims` passes, so this is the + exact read path the flag has to survive: no fill-value decoding, NaN read + straight off disk. + """ + ffp = write_sw4_station_file( + tmp_path / "roundtrip.h5", + { + "AAAA": {"SGDEPTH": 0.0, "SGDEPTHGP": 0.0}, + "BBBB": None, + "CCCC": {"SGDEPTH": 5750.0, "SGDEPTHGP": 14.0}, + }, + widths={"SGWIDTH": 12000.0, "SGWIDTHGP": 30.0}, + ) + dset = lf_to_xarray.convert_sw4_station_recording(ffp) + output = tmp_path / "lf.nc" + # The same engine `lf-to-xarray` itself writes with. + dset.to_netcdf(output, engine="h5netcdf") + + with xr.open_dataset(output, mask_and_scale=False) as reopened: + assert "supergrid_depth" in reopened.coords + assert "supergrid_depth" not in reopened.data_vars + assert reopened["supergrid_depth"].dtype == np.float32 + depth = reopened.sortby("station")["supergrid_depth"].values + assert depth[0] == 0.0 + assert np.isnan(depth[1]) + assert depth[2] == 5750.0 + assert reopened.attrs["SGWIDTH"] == pytest.approx(12000.0) diff --git a/workflow/scripts/im_calc.py b/workflow/scripts/im_calc.py index 54bb3904..f73e93dd 100644 --- a/workflow/scripts/im_calc.py +++ b/workflow/scripts/im_calc.py @@ -134,6 +134,98 @@ } +SUPERGRID_COORDINATES = ("supergrid_depth", "supergrid_depth_gp") +"""Per-station absorbing-layer penetration, as written by `lf-to-xarray`. + +Station-dimension *coordinates*, never data variables: `bb-sim` carries +coordinates through untouched but drops data variables +(`bb_sim._process_bb_chunk`). +""" + +SUPERGRID_WIDTH_ATTRIBUTES = { + "SGWIDTH": "absorbing_layer_width_m", + "SGWIDTHGP": "absorbing_layer_width_gp", +} +"""Map from the LF file's own supergrid-width attributes to the IM root attrs. + +The width is taken from the waveform file rather than from the realisation +configuration on purpose: the configuration can be edited after the run, so +reading it would make the IM file's self-description disagree with the run +that produced it. +""" + + +def supergrid_coordinates(dataset: xr.Dataset) -> dict[str, xr.DataArray]: + """Extract the absorbing-layer penetration coordinates from a waveform file. + + The coordinates are present on SW4 low-frequency output (and, riding along + as coordinates, on the broadband file derived from it). They are absent for + every other solver, and absent on SW4 output written before SW4 reported + the supergrid. Both cases give an all-NaN coordinate: the penetration is + *unknown*, which is not the same claim as `0.0` ("checked, in the + interior"). Materialising it unconditionally means every IM file carries + the coordinate, whatever the solver. + + Parameters + ---------- + dataset : xr.Dataset + The waveform dataset, with a `station` dimension. + + Returns + ------- + dict[str, xr.DataArray] + A map from coordinate name to per-station values, loaded eagerly. + """ + + def absent() -> xr.DataArray: # numpydoc ignore=GL08 + return xr.DataArray( + np.full(dataset.sizes["station"], np.nan, dtype=np.float32), + dims="station", + coords={"station": dataset.station}, + ) + + return { + name: ( + dataset.coords[name].astype(np.float32).compute() + if name in dataset.coords + else absent() + ) + for name in SUPERGRID_COORDINATES + } + + +def supergrid_attributes( + dataset: xr.Dataset, supergrid: dict[str, xr.DataArray] +) -> dict[str, str | float]: + """Describe the absorbing layer at the root of the IM file. + + Nothing is said unless at least one station has a *reported* penetration: + an all-NaN coordinate means no solver reported one, and claiming an + absorbing layer then would be a claim about a run nobody measured. + + Parameters + ---------- + dataset : xr.Dataset + The waveform dataset, whose attributes carry the sponge width. + supergrid : dict[str, xr.DataArray] + The output of `supergrid_coordinates`. + + Returns + ------- + dict[str, str | float] + Root attributes naming the absorbing layer and its width, or an empty + dict if the run does not report one. + """ + if not bool(np.isfinite(supergrid["supergrid_depth"]).any()): + return {} + + attributes: dict[str, str | float] = {"absorbing_layer": "sw4_supergrid"} + for source_name, attribute_name in SUPERGRID_WIDTH_ATTRIBUTES.items(): + if source_name in dataset.attrs: + attributes[attribute_name] = float(dataset.attrs[source_name]) + return attributes + + def add_station_parameters( dtree: xr.DataTree, station_parameters: dict[str, xr.DataArray] ) -> xr.DataTree: @@ -544,6 +636,11 @@ def calculate_intensity_measures( station=broadband.station.str.match(r"^(\w{4})$").values ) + # Read once the station set is final. `im-calc` runs on a raw SW4 LF file + # as well as on `realisation.bb`, and the coordinates ride through `bb-sim` + # untouched, so both paths carry the flag; every other solver gets NaN. + supergrid = supergrid_coordinates(broadband) + intensity_measures = override_ims or intensity_measure_parameters.ims psa_periods = np.array(intensity_measure_parameters.valid_periods, dtype=np.float64) @@ -617,7 +714,7 @@ def calculate_intensity_measures( "ztor": source_parameters.avg_ztor, "zbot": source_parameters.avg_zbot, "hypo_depth": source_parameters.hypo_depth, - } + } | supergrid_attributes(broadband, supergrid) dtree = xr.DataTree.from_dict(im_results, nested=True) @@ -628,8 +725,10 @@ def calculate_intensity_measures( distances.as_dict() # Belt and braces: these already ride along as coordinates on every # leaf, so re-attaching them is idempotent -- but it makes the - # guarantee independent of xarray's coordinate propagation. - | {"latitude": broadband["latitude"], "longitude": broadband["longitude"]}, + # guarantee independent of xarray's coordinate propagation, and + # supplies the all-NaN fallback for solvers that report nothing. + | {"latitude": broadband["latitude"], "longitude": broadband["longitude"]} + | supergrid, ) dtree = add_units(dtree) diff --git a/workflow/scripts/lf_to_xarray.py b/workflow/scripts/lf_to_xarray.py index 2a69bb3b..a1bdf4b7 100644 --- a/workflow/scripts/lf_to_xarray.py +++ b/workflow/scripts/lf_to_xarray.py @@ -25,30 +25,305 @@ See the output of `lf-to-xarray --help`. """ +from enum import StrEnum, auto from pathlib import Path +from typing import Annotated +import dask.array as da +import h5py +import numpy as np import typer +import xarray as xr from qcore import cli, timeseries from workflow import log_utils app = typer.Typer() +CMS = 100.0 +# Unit to convert m/s to cm/s + +TARGET_CHUNK_BYTES = 128 * 2**20 +# Target size of a dask chunk (all components for a batch of stations). + + +def _read_station_batch( + stations: xr.DataArray, + sw4_ffp: Path, + time: xr.DataArray, + component: xr.DataArray, +) -> xr.DataArray: + """Read waveforms for a batch of stations from an SW4 recording file. + + Parameters + ---------- + stations : xr.DataArray + Names of the station groups to read. + sw4_ffp : Path + Path to the SW4 HDF5 station recording file. + time : xr.DataArray + Time coordinates of the recording. + component : xr.DataArray + Component coordinates of the recording. + + Returns + ------- + xr.DataArray + Velocity waveforms in m/s with shape (3, len(stations), npts). + Components are ordered to match the EMOD3D LF convention: + x = east-west, y = north-south, z = down. + + Notes + ----- + The datasets carry SW4's displacement-mode names (EW/NS/UP), but for + SRF rupture sources the time function SW4 receives is the slip *rate*, + so the nominal displacement output is physically velocity (see the + note under the rupture command in the SW4 User's Guide). + + Raises + ------ + RuntimeError + If a station group lacks the EW/NS/UP datasets. + """ + waveforms = np.empty((len(component), len(stations), len(time)), dtype=np.float32) + with h5py.File(sw4_ffp, "r") as handle: + for i, station_name in enumerate(stations): + group = handle[station_name.item()] + if "NS" not in group: + raise RuntimeError( + f"Station {station_name.item()} has no EW/NS/UP datasets." + " The SW4 rechdf5 command must output geographic (NSEW)" + " displacement-mode components (grid X/Y output is not" + " supported: it would need de-rotation by the grid azimuth)." + ) + waveforms[0, i] = group["EW"][:] + waveforms[1, i] = group["NS"][:] + waveforms[2, i] = group["UP"][:] + + return xr.DataArray( + waveforms, + dims=["component", "station", "time"], + coords=dict(time=time, component=component, station=stations.values), + ) + + +def read_station_metadata(sw4_ffp: Path) -> xr.Dataset: + """Initialise an xarray dataset using metadata read from the station recording file. + + Parameters + ---------- + sw4_ffp: Path + Path to SW4 recording file. + + Returns + ------- + xr.Dataset + Xarray dataset with initialised coordinate arrays and attributes. + + The supergrid penetration SW4 reports per station is returned as the + station-dimension *coordinates* `supergrid_depth` (metres) and + `supergrid_depth_gp` (grid points), not as data variables. That is + load-bearing: station-dimension coordinates ride through `bb-sim` and + `im-calc` untouched, whereas data variables are dropped by + `bb_sim._process_bb_chunk`. See the note at `bb_sim.py`'s `combined` + dataset. + + Raises + ------ + RuntimeError + If the HDF5 file is not in the format expected for an SW4 recording file + (see Section 12.9 of the SW4 User Guide). + """ + global_npts = None + stations = [] + latitudes = [] + longitudes = [] + supergrid_depths = [] + supergrid_depths_gp = [] + + with h5py.File(sw4_ffp, "r") as handle: + dt = np.float32(handle["DELTA"][:].squeeze()) + # SW4 writes the supergrid (absorbing layer) width once per file, + # beside DELTA. Guarded the same way as SGDEPTH below: station files + # written before SW4 reported the supergrid have neither. + attrs = dict(dt=dt) + for width_name in ("SGWIDTH", "SGWIDTHGP"): + if width_name in handle: + attrs[width_name] = float(handle[width_name][:].squeeze()) + for station_name, group in handle.items(): + if "NPTS" not in group: + continue + npts = int(group["NPTS"][:].squeeze()) + if global_npts is not None and npts != global_npts: + raise RuntimeError( + f"SW4 output is corrupted: {npts=} but {global_npts=}" + ) + global_npts = npts + stations.append(station_name) + + latitude, longitude, _ = group["STLA,STLO,STDP"][:] + latitudes.append(latitude) + longitudes.append(longitude) + + if "SGDEPTH" in group: + supergrid_depths.append(float(group["SGDEPTH"][:].squeeze())) + # Deliberately read under the SGDEPTH guard rather than its + # own: one present without the other is a corrupt file, not an + # old one, and a KeyError is then the right outcome. + supergrid_depths_gp.append(float(group["SGDEPTHGP"][:].squeeze())) + else: + # An old station file, or a solver with no absorbing layer: + # unknown, *not* clean. Never 0.0 here -- that would assert + # this station was checked and found in the interior. + supergrid_depths.append(np.nan) + supergrid_depths_gp.append(np.nan) + + if global_npts is None: + raise RuntimeError( + "No valid station recordings found in file. Are you sure this is an SW4 station file? Use `h5ls` to check the file structure." + ) + + time = np.arange(global_npts) * dt + return xr.Dataset( + dict( + lat=("station", latitudes), + lon=("station", longitudes), + ), + coords=dict( + station=stations, + component=["x", "y", "z"], + time=time, + # float32 with NaN for "unknown", never an integer with a + # _FillValue: downstream readers open these files with + # `mask_and_scale=False`, so a sentinel would read back raw and + # become a plausible penetration depth. + supergrid_depth=( + "station", + np.array(supergrid_depths, dtype=np.float32), + ), + supergrid_depth_gp=( + "station", + np.array(supergrid_depths_gp, dtype=np.float32), + ), + ), + attrs=attrs | dict(nt=global_npts), + ) + + +def _template_waveform(dset: xr.Dataset, batch_size: int) -> xr.DataArray: + ncomponent = len(dset.coords["component"]) + nstation = len(dset.coords["station"]) + ntime = len(dset.coords["time"]) + return xr.DataArray( + # Chunks must match what _read_station_batch returns (all components + # and timesteps for one batch of stations). If dask is left to pick + # chunks itself it splits the station and time axes, and map_blocks + # then advertises output keys it never produces. + da.empty( + (ncomponent, nstation, ntime), + dtype=np.float32, + chunks=(ncomponent, batch_size, ntime), + ), + dims=["component", "station", "time"], + # Dimension coordinates only. `map_blocks` cross-checks the user + # function's output against the template and raises if the template + # advertises a coordinate the function does not return, and + # `_read_station_batch` returns only the three dimension coordinates. + # The station-dimension coordinates the dataset also carries (such as + # `supergrid_depth`) are re-attached when the result is assigned back + # onto `dset`. + coords={dim: dset.coords[dim] for dim in ("component", "station", "time")}, + ) + + +def convert_sw4_station_recording(sw4_ffp: Path) -> xr.Dataset: + """Convert SW4 station recording to an xarray dataset. + + Parameters + ---------- + sw4_ffp : Path + Path to the SW4 HDF5 station recording file. + + Returns + ------- + xr.Dataset + An xarray dataset lazily constructed from the HDF5 file. Waveform data + is read in batches of stations when the dataset is computed or written. + """ + dset = read_station_metadata(sw4_ffp) + batch_size = max( + 1, + TARGET_CHUNK_BYTES + // ( + len(dset.coords["component"]) + * len(dset.coords["time"]) + * np.float32().itemsize + ), + ) + # The dimension must be named "station" so map_blocks can line the input + # batches up with the station axis of the template. It is deliberately left + # without a station coordinate, as an index coordinate cannot be chunked. + chunked_stations = xr.DataArray(dset["station"].values, dims=["station"]).chunk( + {"station": batch_size} + ) + waveform = xr.map_blocks( + _read_station_batch, + chunked_stations, + kwargs=dict(time=dset["time"], component=dset["component"], sw4_ffp=sw4_ffp), + template=_template_waveform(dset, batch_size), + ) + + waveform = (waveform * CMS).differentiate("time") + dset["waveform"] = waveform + dset.attrs["units"] = "cm/s^2" + # SW4 station recordings begin at simulation time zero. + dset.attrs["start_sec"] = 0.0 + + return dset + + +class Format(StrEnum): + """Input low frequency file format.""" + + SW4 = auto() + """SW4 HDF5 station recording.""" + EMOD3D = auto() + """EMOD3D LFSeis directory.""" + @cli.from_docstring(app) @log_utils.log_call() -def convert_lf_to_xarray_dataset(lfseis_directory: Path, output_ffp: Path) -> None: +def convert_lf_to_xarray_dataset( + low_frequency_path: Annotated[Path, typer.Argument(exists=True)], + output_ffp: Annotated[Path, typer.Argument(writable=True, dir_okay=False)], + format: Format = Format.EMOD3D, +) -> None: """Merge low-frequency outputs into an xarray dataset. Parameters ---------- - lfseis_directory : Path + low_frequency_path : Path Directory containing station seismogram outputs. output_ffp : Path Path to write the xarray dataset + format : Format, optional + Format for the low-frequency inputs (EMOD3D or SW4). If format is SW4, + the low frequency path should be an HDF5 file in the SW4 station format + (Section 12.9 of the SW4 User Guide). If format is instead EMOD3D, the + low frequency path should be a directory containing LFSeis files. + Defaults to EMOD3D. """ - lf_dataset = timeseries.read_lfseis_directory(lfseis_directory) - lf_dataset.to_netcdf(output_ffp, engine="h5netcdf") + match format: + case Format.EMOD3D if low_frequency_path.is_dir(): + lf_dataset = timeseries.read_lfseis_directory(low_frequency_path) + lf_dataset.to_netcdf(output_ffp, engine="h5netcdf") + case Format.EMOD3D: + raise ValueError("EMOD3D format requires directory containing LFSeis files") + case Format.SW4 if low_frequency_path.is_file(): + lf_dataset = convert_sw4_station_recording(low_frequency_path) + lf_dataset.to_netcdf(output_ffp, engine="h5netcdf") + case Format.SW4: + raise ValueError("SW4 format requires station recording file.") if __name__ == "__main__":