diff --git a/pyproject.toml b/pyproject.toml index 67d7ee9b..120904db 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -16,6 +16,7 @@ dependencies = [ "nshmdb>=2026.09.1", "oq_wrapper>=2026.05.2", "qcore-utils>=2025.12.2", + "hf-simulation @ git+https://github.com/ucgmsim/high-frequency", "site-calculation>=2026.7.1", "source_modelling>=2026.08.1", # Data Formats diff --git a/tests/test_hf.py b/tests/test_hf.py index c006ca3c..ded26447 100644 --- a/tests/test_hf.py +++ b/tests/test_hf.py @@ -1,60 +1,37 @@ -from pathlib import Path from types import SimpleNamespace import numpy as np -import pytest +from hf_simulation import PathDurationModel, Ray from hypothesis import given from hypothesis import strategies as st from workflow.realisations import ( HFConfig, - Resolution, RuptureVelocity, ) from workflow.scripts import hf_sim -def test_build_hf_input_serialisation() -> None: - stoch_ffp = Path("/path/to/stoch") - velocity_model = Path("/path/to/vmodel") +def test_build_config_mirrors_the_realisation() -> None: + """The realisation's `hf` section reaches `hf_simulation.HfConfig` unchanged. + The two structures mirror each other group for group, so `build_config` is a splat plus + the two values the realisation deliberately does not carry: the record duration, which + the domain computes, and the rupture-velocity multipliers, which live in their own + section because SRF generation reads them too. This pins both halves of that. + """ hf_config = HFConfig( - sdrop=50.0, - rayset=[1, 2], - no_siteamp=False, - nbu=1, - ift=0, - flo=0.1, - fhi=10.0, - fmax=20.0, - kappa=0.045, - qfexp=0.6, - czero=2.5, - calpha=0.0, - mom=None, - rupv=1.2, - vs_moho=3.5, - nl_skip=0, - vp_sig=0.1, - vsh_sig=0.1, - rho_sig=0.1, - qs_sig=0.1, - ic_flag=True, - velocity_name="test_vel", - fa_sig1=0.2, - fa_sig2=0.2, - rv_sig1=0.1, - path_dur=11, - t_sec=0.0, - site_specific=False, - dpath_pert=0, - stress_parameter_adjustment_fault_area=None, - stress_parameter_adjustment_target_magnitude=None, - stress_parameter_adjustment_tect_type=0, + source={ + "stress_drop_bars": 50.0, + "corner_frequency_constant": 2.5, + "corner_frequency_alpha": 0.1, + "rupture_velocity": {"sigma": 0.1}, + }, + path={"rayset": [1, 2], "q_frequency_exponent": 0.6, "path_duration_model": 11}, + site={"kappa_s": 0.045, "fmax_hz": 20.0}, + record={"dt": 0.005}, ) - - res = Resolution(resolution=0.1) - rv = RuptureVelocity( + rupture_velocity = RuptureVelocity( rvfrac=0.8, rvfrac_shal=0.7, rvfrac_deep=0.9, @@ -64,32 +41,26 @@ def test_build_hf_input_serialisation() -> None: deep_transition_range=1, rvfrac_slip_sig=None, ) - # Rather than create DomainParameters with a bounding box, we simplify with a mock object + # A bounding box is not needed to read one field off it. domain = SimpleNamespace(duration=100.0) - result = hf_sim.build_hf_input( - stoch_ffp, - velocity_model, - res, - hf_config, - rv, - domain, # ty: ignore[invalid-argument-type] - ) - - lines = result.split("\n") - - assert lines[1] == "50.0" # sdrop - assert lines[2] == "{station_input_file}" # placeholder for station file - assert lines[3] == "{output_file}" # placeholder for output file - assert lines[4] == "2 1 2" # rayset count + rays - assert lines[5] == "1" # int(not no_siteamp) -> int(not False) -> 1 - assert lines[7] == "{seed}" # seed placeholder - assert lines[9] == "100.0 0.005 20.0 0.045 0.6" # Domain and resolution parameters - assert lines[10] == "0.8 0.7 0.9 2.5 0.0" # rupture velocity + czero,alpha - assert lines[11] == "-1 1.2" # mom (None -> -1) and rupv - assert lines[12] == str(stoch_ffp) # Stoch file path - assert lines[15] == "0 0.1 0.1 0.1 0.1 1" # Sigs and ic_flag (True -> 1) - assert lines[20] == "-1 -1 -1" # Optional stress parameters + config = hf_sim.build_config(hf_config, rupture_velocity, domain) # ty: ignore[invalid-argument-type] + + # Splatted through unchanged. + assert config.source.stress_drop_bars == 50.0 + assert config.source.corner_frequency_constant == 2.5 + assert config.site.fmax_hz == 20.0 + assert config.record.dt == 0.005 + # Ints become the enums the simulation takes. + assert config.path.rayset == (Ray.DIRECT, Ray.MOHO_REFLECTION) + assert config.path.path_duration_model is PathDurationModel.BOORE_THOMPSON_2014 + # Injected, because the `hf` section does not carry them. + assert config.record.duration_s == 100.0 + assert config.source.rupture_velocity.fraction == 0.8 + assert config.source.rupture_velocity.shallow == 0.7 + assert config.source.rupture_velocity.deep == 0.9 + # ... but the sigma does come from the `hf` section. + assert config.source.rupture_velocity.sigma == 0.1 STATION_STRATEGY = st.text( @@ -99,7 +70,7 @@ def test_build_hf_input_serialisation() -> None: def test_station_seeds() -> None: seed = hf_sim.station_seeds(0, ["station"]) - assert seed.dtype == np.int32 + assert seed.dtype == np.uint64 assert seed.shape == (1,) # Seeds should be referentially transparent: i.e. depend only on the seed and station name seed_1 = hf_sim.station_seeds(0, ["station"]) @@ -107,7 +78,8 @@ def test_station_seeds() -> None: @given( - seed=st.integers(min_value=-(1 << 31), max_value=(1 << 31) - 1), + # Non-negative: SeedSequence rejects negative entropy, which `station_seeds` says. + seed=st.integers(min_value=0, max_value=(1 << 31) - 1), stations=st.lists(STATION_STRATEGY, min_size=1, unique=True), ) def test_station_seeds_on_name_only(seed: int, stations: list[str]) -> None: @@ -123,48 +95,3 @@ def test_station_seeds_on_name_only(seed: int, stations: list[str]) -> None: # one. for station, expected_seed in zip(stations, station_seeds): assert hf_sim.station_seeds(seed, [station]).item() == expected_seed - - -def test_create_hf_dataset_structure() -> None: - # 1. Setup Mock Data - n_stations = 2 - n_components = 3 # Fixed by function logic - n_time = 100 - - names = ["station_a", "station_b"] - waveform = np.random.rand(n_components, n_stations, n_time).astype(np.float32) - lat = np.array([-43.5, -43.6]) - lon = np.array([172.6, 172.7]) - dist = np.array([10.5, 20.1]) - seeds = np.array([123, 456]) - vrefs = np.array([300.0, 350.0]) - dt = 0.02 - start_sec = 0.0 - - ds = hf_sim.create_hf_dataset( - waveform=waveform, - latitude=lat, - longitude=lon, - names=names, - epicentre_distance=dist, - seed=seeds, - vref=vrefs, - dt=dt, - start_sec=start_sec, - ) - - assert ds.sizes == {"component": 3, "station": 2, "time": 100} - - np.testing.assert_array_equal(ds.station.values, names) - np.testing.assert_array_equal(ds.component.values, ["x", "y", "z"]) - assert ds.time.values[1] == pytest.approx(0.02) - assert ds.lat.dims == ("station",) - assert ds.lon.dims == ("station",) - - assert "waveform" in ds.data_vars - assert ds.waveform.dims == ("component", "station", "time") - np.testing.assert_allclose(ds.waveform.values, waveform) - - assert ds.attrs["dt"] == dt - assert ds.attrs["nt"] == n_time - assert ds.attrs["units"] == "cm/s^2" diff --git a/uv.lock b/uv.lock index 747a1a91..c9ca99e6 100644 --- a/uv.lock +++ b/uv.lock @@ -1040,6 +1040,14 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/4e/14/32cf2aae083c74b95678875e11d25d0cb9e51a33d27f4218eb9aeacfcd70/hdf5plugin-7.0.0-py3-none-win_amd64.whl", hash = "sha256:2e052af8d7848e8bac92646584617503a08bb9b466cfa810a49ecd93e89b7ffa", size = 3523827, upload-time = "2026-06-25T20:59:37.758Z" }, ] +[[package]] +name = "hf-simulation" +version = "0.1.dev178+g15175ed01" +source = { git = "https://github.com/ucgmsim/high-frequency#15175ed012fabc37c164d9ee49c65f2ecf0e4186" } +dependencies = [ + { name = "numpy" }, +] + [[package]] name = "hpack" version = "4.2.0" @@ -3522,6 +3530,7 @@ dependencies = [ { name = "dask" }, { name = "geopandas" }, { name = "h5py" }, + { name = "hf-simulation" }, { name = "im-calculation" }, { name = "netcdf4" }, { name = "nshmdb" }, @@ -3574,6 +3583,7 @@ requires-dist = [ { name = "deptry", marker = "extra == 'dev'" }, { name = "geopandas" }, { name = "h5py", specifier = ">=3.15.1" }, + { name = "hf-simulation", git = "https://github.com/ucgmsim/high-frequency" }, { name = "hypothesis", extras = ["numpy"], marker = "extra == 'test'", specifier = ">=6.0.0" }, { name = "im-calculation", git = "https://github.com/ucgmsim/im_calculation?branch=no_parallel" }, { name = "netcdf4" }, diff --git a/workflow/default_parameters/root/defaults.yaml b/workflow/default_parameters/root/defaults.yaml index 8f0940c9..2e96ce60 100644 --- a/workflow/default_parameters/root/defaults.yaml +++ b/workflow/default_parameters/root/defaults.yaml @@ -60,39 +60,31 @@ emod3d: yseis: 0 zseis: 0 pertbfile: none.pertb +# Mirrors `hf_simulation.HfConfig` group for group and field for field, so this block can +# be deserialised straight into it rather than translated. Two of that class's inputs are +# absent on purpose: `record.duration_s` is computed from the domain, and the three +# `source.rupture_velocity` multipliers live in `rupture_velocity:` below because SRF +# generation reads the same physical values. hf-sim injects both. hf: - nbu: 4 - ift: 0 - flo: 0.02 - fhi: 19.9 - nl_skip: -99 - vp_sig: 0.0 - vsh_sig: 0.0 - rho_sig: 0.0 - qs_sig: 0.0 - ic_flag: true - velocity_name: "-1" - t_sec: 0.0 - sdrop: 50.0 - rayset: [1] - no_siteamp: false - fmax: 10.0 - kappa: 0.045 - qfexp: 0.6 - czero: 2.1 - calpha: -99.0 - mom: null - rupv: null - site_specific: false - vs_moho: 999.9 - fa_sig1: 0.0 - fa_sig2: 0.0 - rv_sig1: 0.1 - path_dur: 11 - dpath_pert: 0.0 - stress_parameter_adjustment_tect_type: 0 - stress_parameter_adjustment_target_magnitude: null - stress_parameter_adjustment_fault_area: null + source: + stress_drop_bars: 50.0 + corner_frequency_constant: 2.1 + # 0.1, not -99.0. The Fortran read any value below -1.0 as "use the built-in default" + # (`if(Calpha.lt.-1.0) Calpha = Calpha_default`), so -99.0 here MEANT 0.1. That decode + # lived in the deck reader, which is gone -- hf-simulation takes the value literally + # now, and -99.0 makes alpha_T negative for anything but a vertical strike-slip fault. + corner_frequency_alpha: 0.1 + rupture_velocity: + sigma: 0.1 + path: + rayset: [1] + q_frequency_exponent: 0.6 + path_duration_model: 11 + site: + kappa_s: 0.045 + fmax_hz: 10.0 + record: + dt: 0.005 stoch: stoch_dx: 2.0 stoch_dy: 2.0 @@ -1530,6 +1522,8 @@ hf_velocity_model_1d: rho: 3.33 Qp: 394.80 Qs: 197.40 + # Truncate the model where Vs reaches this. 999.9 means no layer does. + vs_moho: 999.9 bb: fmin: 0.2 fmidbot: 0.5 diff --git a/workflow/default_parameters/v24_2_2_1/defaults.yaml b/workflow/default_parameters/v24_2_2_1/defaults.yaml index 514a3707..62fb8bdf 100644 --- a/workflow/default_parameters/v24_2_2_1/defaults.yaml +++ b/workflow/default_parameters/v24_2_2_1/defaults.yaml @@ -3,3 +3,6 @@ resolution: resolution: 0.1 bb: flo: 1.0 +hf: + record: + dt: 0.005 diff --git a/workflow/default_parameters/v24_2_2_2/defaults.yaml b/workflow/default_parameters/v24_2_2_2/defaults.yaml index 47072a4d..3961323c 100644 --- a/workflow/default_parameters/v24_2_2_2/defaults.yaml +++ b/workflow/default_parameters/v24_2_2_2/defaults.yaml @@ -3,3 +3,6 @@ resolution: resolution: 0.2 bb: flo: 0.5 +hf: + record: + dt: 0.005 diff --git a/workflow/default_parameters/v24_2_2_4/defaults.yaml b/workflow/default_parameters/v24_2_2_4/defaults.yaml index ddecf8bb..52c36321 100644 --- a/workflow/default_parameters/v24_2_2_4/defaults.yaml +++ b/workflow/default_parameters/v24_2_2_4/defaults.yaml @@ -3,3 +3,6 @@ resolution: resolution: 0.4 bb: flo: 0.25 +hf: + record: + dt: 0.005 diff --git a/workflow/default_parameters/v26_7_1Hz/defaults.yaml b/workflow/default_parameters/v26_7_1Hz/defaults.yaml index 25db0c72..a7b2ad9c 100644 --- a/workflow/default_parameters/v26_7_1Hz/defaults.yaml +++ b/workflow/default_parameters/v26_7_1Hz/defaults.yaml @@ -105,3 +105,6 @@ sw4: - name: "imagehdf5" parameters: { mode: "vmax", z: 0.0, file: "surf_vmax", precision: "float" } +hf: + record: + dt: 0.005 diff --git a/workflow/realisations.py b/workflow/realisations.py index 11eaf148..827a6115 100644 --- a/workflow/realisations.py +++ b/workflow/realisations.py @@ -18,7 +18,7 @@ from collections.abc import Sequence from importlib import metadata from pathlib import Path -from typing import Any, ClassVar, Literal, Self +from typing import Any, ClassVar, Self import numpy as np import numpy.typing as npt @@ -1091,10 +1091,15 @@ class HFVelocityModel1D(VelocityModel1D): """1D Velocity Model for SRF and HF. Differs from the VelocityModel1D class in the default case with a minimum - Vs of 500 m/s.""" + Vs of 500 m/s, and in carrying `vs_moho` -- which is a property of the model, not of + the simulation, and truncates it. + """ _config_key: ClassVar[str] = "hf_velocity_model_1d" - _schema: ClassVar[Schema] = schemas.VELOCITY_MODEL_1D_SCHEMA + _schema: ClassVar[Schema] = schemas.HF_VELOCITY_MODEL_1D_SCHEMA + + vs_moho: float = 999.9 + """Shear velocity at which to truncate the model at the Moho, km/s.""" @dataclasses.dataclass @@ -1142,80 +1147,35 @@ class RuptureVelocity(RealisationConfiguration): @dataclasses.dataclass class HFConfig(RealisationConfiguration): - """High frequency simulation configuration.""" + """High frequency simulation configuration. + + **Mirrors `hf_simulation.HfConfig` group for group and field for field**, so this + section can be deserialised straight into it rather than translated. Keep it that way: + a translation layer between two descriptions of the same physics is the thing that + drifts. + + Two of that class's inputs are deliberately absent. `record.duration_s` is computed + from the domain, and the three `source.rupture_velocity` multipliers live in + :class:`RuptureVelocity` because SRF generation reads the same physical values. Both + are injected by `hf-sim`. + """ _config_key: ClassVar[str] = "hf" _schema: ClassVar[Schema] = schemas.HF_CONFIG_SCHEMA - nbu: int - """Unknown!""" - ift: int - """Unknown!""" - flo: float - """Unknown!""" - fhi: float - """Unknown!""" - nl_skip: int - """Skip empty lines in input?""" - vp_sig: float - """Unknown!""" - vsh_sig: float - """Unknown!""" - qs_sig: float - """Unknown!""" - rho_sig: float - """Unknown!""" - ic_flag: bool - """Unknown!""" - velocity_name: str - """Unknown""" - t_sec: float - """High frequency output start time.""" - sdrop: float - """Stress drop average (bars)""" - rayset: list[Literal[1, 2]] - """ray types 1: direct, 2: moho""" - no_siteamp: bool - """Disable BJ97 site amplification factors""" - fmax: float - """Max simulation frequency""" - kappa: float - """Unknown!""" - qfexp: float - """Q frequency exponent""" - czero: float - """C0 coefficient""" - calpha: float - """Ca coefficient""" - mom: float | None - """Seismic moment for HF simulation (or None, to infer value)""" - rupv: float | None - """Rupture velocity (or binary default)""" - site_specific: bool - """Enable site-specific calculation""" - vs_moho: float - """vs of moho layer""" - fa_sig1: float - """Fourier amplitude uncertainty (1)""" - fa_sig2: float - """Fourier amplitude uncertainty (2)""" - rv_sig1: float - """Rupture velocity uncertainty""" - path_dur: Literal[0, 1, 2, 11, 12] - """path duration model. - - 0: GP2010 - - 1: WUS modification trail/error - - 2: ENA modification trial/error - - 11: WUS formulation of BT2014 - - 12: ENA formulation of BT2015. Models 11 and 12 over predict for multiple rays.""" - dpath_pert: float - """Log of path duration multiplier""" - stress_parameter_adjustment_tect_type: Literal[0, 1, 2] - """Adjustment option 0 = off, 1 = active tectonic, 2 = stable continent""" - stress_parameter_adjustment_target_magnitude: float | None - """Target magnitude (or inferred if None)""" - stress_parameter_adjustment_fault_area: float | None - """Target magnitude (or inferred if None)""" + source: dict[str, Any] + """Radiation strength and rupture speed. See `hf_simulation.SourceParameters`.""" + path: dict[str, Any] + """Rays and attenuation. See `hf_simulation.PathParameters`.""" + site: dict[str, Any] + """The near-surface. See `hf_simulation.SiteParameters`.""" + record: dict[str, Any] + """Duration and sample interval. See `hf_simulation.RecordParameters`.""" + + @property + def dt(self) -> float: + """float: Sample interval, seconds.""" + return self.record["dt"] @dataclasses.dataclass @@ -1407,7 +1367,7 @@ def render(self) -> str: parts.append(f"{key}={value}") return " ".join(parts) - def merged(self, **overrides: str | int | float | bool | None) -> "SW4Command": + def merged(self, **overrides: str | float | bool | None) -> "SW4Command": """Return a copy of this command with `overrides` merged into its parameters. Parameters diff --git a/workflow/schemas.py b/workflow/schemas.py index 543c59cb..e6fa7bb3 100644 --- a/workflow/schemas.py +++ b/workflow/schemas.py @@ -1013,6 +1013,16 @@ def _corners_to_array(corners_spec: list[dict[str, float]]) -> np.ndarray: } ) +HF_VELOCITY_MODEL_1D_SCHEMA = Schema( + { + **VELOCITY_MODEL_1D_SCHEMA.schema, + Literal( + "vs_moho", + description="Shear velocity at which to truncate the model at the Moho (km/s)", + ): And(NUMBER, _is_positive), + } +) + REALISATION_METADATA_SCHEMA = Schema( { Literal("name", description="The name of the realisation"): str, @@ -1033,62 +1043,43 @@ def _corners_to_array(corners_spec: list[dict[str, float]]) -> np.ndarray: HF_CONFIG_SCHEMA = Schema( { - Literal("nbu", description="Unknown!"): int, - Literal("ift", description="Unknown!"): int, - Literal("flo", description="Unknown!"): NUMBER, - Literal("fhi", description="Unknown!"): NUMBER, - Literal("nl_skip", description="Skip empty lines in input?"): int, - Literal("vp_sig", description="Unknown!"): NUMBER, - Literal("vsh_sig", description="Unknown!"): NUMBER, - Literal("rho_sig", description="Unknown!"): NUMBER, - Literal("qs_sig", description="Unknown!"): NUMBER, - Literal("ic_flag", description="Unknown!"): bool, - Literal("velocity_name", description="Unknown!"): str, - Literal("t_sec", description="High frequency output start time."): And( - NUMBER, _is_non_negative - ), - Literal("sdrop", description="Stress drop average (bars)"): NUMBER, - Literal("rayset", description="ray types 1: direct, 2: moho"): [Or(1, 2)], - Literal( - "no_siteamp", description="Disable BJ97 site amplification factors" - ): bool, - Literal("fmax", description="Max simulation frequency"): And( - NUMBER, _is_positive - ), - Literal("kappa", description="Unknown!"): NUMBER, - Literal("qfexp", description="Q frequency exponent"): NUMBER, - Literal("czero", description="C0 coefficient"): NUMBER, - Literal("calpha", description="Ca coefficient"): NUMBER, - Literal("mom", description="Seimic moment (or null, to infer value)"): Or( - NUMBER, None - ), - Literal("rupv", description="Rupture velocity (or binary default)"): Or( - NUMBER, None - ), - Literal("site_specific", description="Enable site-specific calculation"): bool, - Literal("vs_moho", description="vs of moho layer"): NUMBER, - Literal("fa_sig1", "Fourier amplitude uncertainty (1)"): NUMBER, - Literal("fa_sig2", description="Fourier amplitude uncertainty (2)"): NUMBER, - Literal("rv_sig1", description="Rupture velocity uncertainty"): And( - NUMBER, _is_non_negative - ), - Literal( - "path_dur", - description="path duration model. 0: GP2010, 1: WUS modification trail/errol, 2: ENA modificiation trial/error" - ", 11: WUS formutian of BT2014, 12: ENA formulation of BT2015. Models 11 and 12 overpredict for multiple rays.", - ): Or(0, 1, 2, 11, 12), - Literal("dpath_pert", description="Log of path duration multiplier"): NUMBER, Literal( - "stress_parameter_adjustment_tect_type", - description="Adjustment option 0 = off, 1 = active tectonic, 2 = stable continent", - ): Or(0, 1, 2), - Literal( - "stress_parameter_adjustment_target_magnitude", - description="Target magnitude (or inferred if null)", - ): Or(NUMBER, None), - Literal( - "stress_parameter_adjustment_fault_area", "Fault area (or inferred if null)" - ): Or(NUMBER, None), + "source", + description="The earthquake source: radiation strength and rupture speed", + ): { + Literal("stress_drop_bars", description="Brune stress parameter"): And( + NUMBER, _is_positive + ), + Literal( + "corner_frequency_constant", description="c0 of Graves & Pitarka eq. 13" + ): NUMBER, + Literal( + "corner_frequency_alpha", + description="c_alpha of the alpha_T adjustment", + ): NUMBER, + Literal("rupture_velocity", description="Depth-dependent rupture taper"): { + Literal( + "sigma", description="Log-normal scatter on the rupture factor" + ): And(NUMBER, _is_non_negative), + }, + }, + Literal("path", description="Which rays, and how the medium attenuates"): { + Literal("rayset", description="ray types 1: direct, 2: moho"): [Or(1, 2)], + Literal("q_frequency_exponent", description="x in Q(f) = Q0 f^x"): NUMBER, + Literal( + "path_duration_model", + description="0: GP2010, 1: WUS, 2: ENA, 11: BT2014, 12: BT2015", + ): Or(0, 1, 2, 11, 12), + }, + Literal("site", description="The near-surface"): { + Literal("kappa_s", description="Near-surface attenuation (s)"): NUMBER, + Literal("fmax_hz", description="High-frequency cutoff"): And( + NUMBER, _is_positive + ), + }, + Literal("record", description="The shape of the record to produce"): { + Literal("dt", description="Sample interval (s)"): And(NUMBER, _is_positive), + }, } ) diff --git a/workflow/scripts/hf_sim.py b/workflow/scripts/hf_sim.py index 2ffd495f..db88f036 100644 --- a/workflow/scripts/hf_sim.py +++ b/workflow/scripts/hf_sim.py @@ -17,10 +17,11 @@ Environment ----------- -Can be run in the cybershake container. Can also be run from your own computer using the `hf-sim` command which is installed after running `pip install workflow@git+https://github.com/ucgmsim/workflow`. If you do run this on your own computer, you need a version of `hb_high_binmod` installed. +Can be run in the cybershake container. Can also be run from your own computer using the +`hf-sim` command after `pip install workflow@git+https://github.com/ucgmsim/workflow`. -> [!NOTE] -> The high-frequency code is very brittle. It is recommended you have both versions 6.0.3 and 5.4.5 built to run with. Sometimes it is necessary to switch between versions if one does not work. +Unlike previous versions this needs no `hb_high_binmod` binary and no writable work +directory: the simulation is the `hf-simulation` package, called in-process. Usage ----- @@ -31,326 +32,216 @@ See the output of `hf-sim --help`. """ -import concurrent.futures -import subprocess -import tempfile -from collections.abc import Iterable -from concurrent.futures.thread import ThreadPoolExecutor from pathlib import Path from typing import Annotated +import dask +import dask.array as da import numpy as np -import numpy.typing as npt import pandas as pd import typer import xarray as xr +from hf_simulation import ( + COMPONENTS, + FaultSegment, + HfConfig, + PathDurationModel, + PathParameters, + Ray, + RecordParameters, + Simulator, + SiteParameters, + SlipModel, + SourceParameters, + VelocityModel1D, + station_seeds, +) +from hf_simulation import ( + RuptureVelocity as RuptureVelocityTaper, +) +from tqdm.dask import TqdmCallback from qcore import cli +from source_modelling.stoch import StochFile from workflow import log_utils, realisations, utils from workflow.realisations import ( DomainParameters, - HFConfig, HFVelocityModel1D, RealisationMetadata, - Resolution, RuptureVelocity, Seeds, ) +from workflow.realisations import ( + HFConfig as HFConfigDefaults, +) app = typer.Typer() +TARGET_CHUNK_BYTES = 128 * 2**20 +"""Target size of a dask chunk (all components for a batch of stations).""" -def rupture_velocity_hf_transition_bands( - rupture_velocity: RuptureVelocity, -) -> tuple[float, float, float, float]: - """Produce transition bands for rupture velocity parameters. - Converts median-centred description into bounds description. - - Parameters - ---------- - rupture_velocity : RuptureVelocity - Rupture velocity configuration - - - Returns - ------- - tuple[float, float, float, float] - The shallow min/max, deep min/max transition depths. - """ - deep = rupture_velocity.deep_depth - deep_range = rupture_velocity.deep_transition_range - shallow = rupture_velocity.shallow_depth - shallow_range = rupture_velocity.shallow_transition_range - deep_min = deep - deep_range - deep_max = deep + deep_range - shallow_min = shallow - shallow_range - shallow_max = shallow + shallow_range - return shallow_min, shallow_max, deep_min, deep_max - - -def build_hf_input( - stoch_ffp: Path, - velocity_model: Path, - resolution: Resolution, - hf_config: HFConfig, +def build_config( + hf_config: HFConfigDefaults, rupture_velocity: RuptureVelocity, domain_parameters: DomainParameters, -) -> str: - """Build a high-frequency input template string. +) -> HfConfig: + """Translate the realisation's configuration into the simulation's. Parameters ---------- - stoch_ffp : Path - The path to the stoch file. - velocity_model : Path - The path to the velocity model. - resolution : Resolution - HF simulation resolution. - hf_config : HFConfig - The high-frequency config. + hf_config : HFConfigDefaults + The realisation's high-frequency configuration. rupture_velocity : RuptureVelocity - The rupture velocity settings. + The realisation's rupture velocity settings. domain_parameters : DomainParameters - The simulation domain parameters. + Supplies the record duration. Returns ------- - str - A template HF input, this template has two format placeholders - `station_input_file` and `output_file` which can be - substituted to yield a high-frequency input in for each - station. + HfConfig + The simulation configuration. """ - # Underscore-prefixed because the line consuming them is commented out - # below, pending the EMOD3D PR noted there. - _shallow_min, _shallow_max, _deep_min, _deep_max = ( - rupture_velocity_hf_transition_bands(rupture_velocity) - ) - hf_sim_input = [ - "", - hf_config.sdrop, - "{station_input_file}", - "{output_file}", - f"{len(hf_config.rayset)} {' '.join(str(ray) for ray in hf_config.rayset)}", - int(not hf_config.no_siteamp), - f"{hf_config.nbu} {hf_config.ift} {hf_config.flo} {hf_config.fhi}", - "{seed}", - 1, # one station in the input - f"{domain_parameters.duration} {resolution.dt} {hf_config.fmax} {hf_config.kappa} {hf_config.qfexp}", - f"{rupture_velocity.rvfrac} {rupture_velocity.rvfrac_shal} {rupture_velocity.rvfrac_deep} {hf_config.czero} {hf_config.calpha}", - # TODO: This requires PR from EMOD3D to merge before we can do this! - # f"{_shallow_min} {_shallow_max} {_deep_min} {_deep_max}", - f"{hf_config.mom or -1} {hf_config.rupv or -1}", - stoch_ffp, - velocity_model, - hf_config.vs_moho, - f"{hf_config.nl_skip} {hf_config.vp_sig} {hf_config.vsh_sig} {hf_config.rho_sig} {hf_config.qs_sig} {int(hf_config.ic_flag)}", - hf_config.velocity_name, - f"{hf_config.fa_sig1} {hf_config.fa_sig2} {hf_config.rv_sig1}", - hf_config.path_dur, - 0, # maybe don't need this? - # If running v5.4.5 it stops reading input here and so - # these parameters are unused. It is harmless to add them - # regardless of version - ( - f"{hf_config.stress_parameter_adjustment_fault_area or -1} " - f"{hf_config.stress_parameter_adjustment_target_magnitude or -1} " - f"{hf_config.stress_parameter_adjustment_tect_type or -1}" + # The realisation's `hf` section mirrors `HfConfig` group for group, so each group is + # splatted straight in. The two things it does NOT carry are injected here: the record + # duration, which the domain computes, and the rupture-velocity multipliers, which live + # in their own section because SRF generation reads the same values. + return HfConfig( + source=SourceParameters( + **hf_config.source + | { + "rupture_velocity": RuptureVelocityTaper( + fraction=rupture_velocity.rvfrac, + shallow=rupture_velocity.rvfrac_shal, + deep=rupture_velocity.rvfrac_deep, + **hf_config.source["rupture_velocity"], + ) + } + ), + path=PathParameters( + **hf_config.path + | { + "rayset": tuple(Ray(ray) for ray in hf_config.path["rayset"]), + "path_duration_model": PathDurationModel( + hf_config.path["path_duration_model"] + ), + } ), - 0, # seek bytes to 0 (no binary offset for this output) - "", - ] - return "\n".join(str(line) for line in hf_sim_input) - - -def hf_simulate_station( - hf_sim_path: Path, - hf_stdin_template: str, - station_latitude: float, - station_longitude: float, - station_name: str, - seed: int, -) -> tuple[str, float, np.ndarray]: - """Simulate a seismic station using the HF (High-Frequency) simulation tool. + site=SiteParameters(**hf_config.site), + record=RecordParameters( + duration_s=domain_parameters.duration, **hf_config.record + ), + ) + + +def build_slip_model(stoch_ffp: Path) -> SlipModel: + """Read a stoch file into a simulation slip model. Parameters ---------- - hf_sim_path : Path - The path to the HF simulation binary. - hf_stdin_template : str - The stdin input template for the HF simulation binary. - station_latitude : float - The station latitude. - station_longitude : float - The station longitude. - station_name : str - The station name. - seed : int - The seed for this HF simulation. + stoch_ffp : Path + Path to the stoch file. Returns ------- - str - The completed station name. - float - The epicentre distance obtained from the simulation output. - array of floats - The simulation waveform. - - Raises - ------ - ValueError - If the output does not contain exactly one epicentre distance value. - CalledProcessError - If the HF binary throws an error. A note to the exception is - added with the stderr. + SlipModel + The slip model, one segment per stoch plane. """ - with ( - tempfile.NamedTemporaryFile(mode="w") as input_file, - tempfile.NamedTemporaryFile() as output_file, - ): - input_file.write(f"{station_longitude} {station_latitude} {station_name}\n") - input_file.flush() - - hf_sim_input_str = hf_stdin_template.format( - station_input_file=input_file.name, output_file=output_file.name, seed=seed - ) - - logger = log_utils.get_logger(__name__) - logger.info("running hf", station=station_name, input=hf_sim_input_str) - - try: - output = subprocess.run( - str(hf_sim_path), - input=hf_sim_input_str, - check=True, - text=True, - stderr=subprocess.PIPE, - ) - except subprocess.CalledProcessError as e: - logger.error( - "hf failed", station=station_name, stdout=e.stdout, stderr=e.stderr + stoch = StochFile.from_file(stoch_ffp) + return SlipModel( + [ + FaultSegment( + longitude_deg=plane.header.longitude, + latitude_deg=plane.header.latitude, + strike_deg=plane.header.strike, + dip_deg=plane.header.dip, + rake_deg=plane.header.average_rake, + top_depth_km=plane.header.dtop, + subfault_length_km=plane.header.dx, + subfault_width_km=plane.header.dy, + hypocentre_along_strike_km=plane.header.shypo, + hypocentre_down_dip_km=plane.header.dhypo, + # (down-dip, along-strike), which is how the stoch format stores them. + slip=plane.slip.astype(np.float32), + rise_time_s=plane.rise.astype(np.float32), + rupture_time_s=plane.trup.astype(np.float32), ) - e.add_note(e.stderr) - raise - - epicentre_distance = float(output.stderr.strip()) - - logger.info( - "hf succeeded", - station=station_name, - epicentre_distance=epicentre_distance, - stderr=output.stderr, - ) - - station_waveform = np.fromfile(output_file, dtype=np.float32).reshape((-1, 3)) - - return station_name, epicentre_distance, station_waveform + for plane in stoch.data + ] + ) -def station_seeds(seed: int, stations: Iterable[str]) -> npt.NDArray[np.int32]: - """Create a list of per-station seeds in an order-invariant fashion with a root seed. +def build_velocity_model(velocity_model: HFVelocityModel1D) -> VelocityModel1D: + """Convert the realisation's 1D velocity model into the simulation's. Parameters ---------- - seed : int - The root seed. - stations : Iterable[str] - The stations to seed. The order and number of stations should - not matter. The station seeds are based on their name only. + velocity_model : HFVelocityModel1D + The realisation's layered model, including the Moho truncation velocity. Returns ------- - npt.NDArray[np.int32] - A list of station seeds. + VelocityModel1D + The velocity model, already truncated at the Moho. """ - station_hashes = np.array( - [utils.stable_hash(name) for name in stations], dtype=np.int32 + model = velocity_model.model + return VelocityModel1D( + thickness_km=model["thickness"].to_numpy(np.float32), + vp_km_s=model["Vp"].to_numpy(np.float64), + vsh_km_s=model["Vs"].to_numpy(np.float64), + density_g_cm3=model["rho"].to_numpy(np.float64), + quality_factor_p=model["Qp"].to_numpy(np.float32), + quality_factor_s=model["Qs"].to_numpy(np.float32), + vs_moho_km_s=velocity_model.vs_moho, ) - # Rather than add (which could overflow and cause annoying numpy - # warnings), we just xor the hf seed with the station hashes. - # Since this is invertible, we ensure that the same hf seed gives - # the same station seeds. - return np.int32(seed) ^ station_hashes - - -def create_hf_dataset( - # array-like used here to reduce the number of times we have to - # change the types if the downstream function inputs change. - waveform: npt.ArrayLike, - latitude: npt.ArrayLike, - longitude: npt.ArrayLike, - names: npt.ArrayLike, - epicentre_distance: npt.ArrayLike, - seed: npt.ArrayLike, - vref: npt.ArrayLike, - dt: float, - start_sec: float, -) -> xr.Dataset: - """ - Create a structured xarray Dataset for HF simulation data. + + +def simulate_chunk( + station_chunk: xr.Dataset, + time: np.ndarray, + simulator: Simulator, +) -> xr.DataArray: + """Simulate one dask block's worth of stations in a single call. Parameters ---------- - waveform : ArrayLike - The waveform data. Expected shape is (3, n_stations, nt), - representing the three components (x, y, z). - latitude : ArrayLike - Latitude coordinates for each station. Shape (n_stations,). - longitude : ArrayLike - Longitude coordinates for each station. Shape (n_stations,). - names : ArrayLike - Names/IDs for each station. Shape (n_stations,). Used as the - primary index for the 'station' dimension. - epicentre_distance : ArrayLike - Distance from the station to the epicentre. Shape (n_stations,). - seed : ArrayLike - Random seed values associated with each station. Shape (n_stations,). - vref : ArrayLike - Reference velocity (Vs30 or similar) for each station. Shape (n_stations,). - dt : float - Time step increment in seconds. - start_sec : float - The start time of the simulation in seconds. + station_chunk : xr.Dataset + Stations in this block, with `latitude`, `longitude` and `seed`. + time : np.ndarray + The shared time axis. + simulator : Simulator + The prepared simulation, shared across every block. Returns ------- - xr.Dataset - A dataset containing the waveforms and associated station metadata, - indexed by station, component, and time. + xr.DataArray + Waveforms over (component, station, time). Notes ----- - The dataset follows specific dimensional mapping: - * **waveform**: mapped to (component, station, time). - * **coordinates**: 'lat' and 'lon' are non-index coordinates tied to - the 'station' dimension. - * **attributes**: global metadata includes 'units' (fixed to cm/s^2), - 'nt', and 'dt'. + The simulator is built once and shared rather than rebuilt per block, so the + station-independent work -- the air layer, the slip-model normalisation, the moment + scaling -- is done once for the whole run. It does not mutate, which is what makes + sharing it across dask's threads safe. **A process-based scheduler will not work**: + neither `Simulator` nor `SlipModel` is picklable, and that was already true of the + models this used to take. """ - waveform = np.asarray(waveform) - nt = waveform.shape[-1] - time = np.arange(nt) * dt - return xr.Dataset( - { - "waveform": (["component", "station", "time"], waveform), - "epicentre_distance": (["station"], epicentre_distance), - "seed": (["station"], seed), - "vref": (["station"], vref), - }, + # The block's own station order, not a sorted one: map_blocks requires the output to + # line up with the template block. Station order does not affect any waveform -- the + # simulation guarantees that and tests it -- but the LABELS still have to match. + station_names = station_chunk["station"].values + waveform = simulator.run_stations( + latitude_deg=station_chunk["latitude"].values.astype(np.float32), + longitude_deg=station_chunk["longitude"].values.astype(np.float32), + station_seed=station_chunk["seed"].values.astype(np.uint64), + ) + return xr.DataArray( + waveform, + dims=["component", "station", "time"], coords={ - "station": ("station", names), - "component": ("component", ["x", "y", "z"]), - "time": ("time", time), - "lat": (["station"], latitude), - "lon": (["station"], longitude), - }, - attrs={ - "start_sec": start_sec, - "nt": nt, - "dt": dt, - "units": "cm/s^2", + "component": list(COMPONENTS), + "station": station_names, + "time": time, }, ) @@ -359,28 +250,11 @@ def create_hf_dataset( @log_utils.log_call() def run_hf( realisation_ffp: Annotated[Path, typer.Argument()], - stoch_ffp: Annotated[ - Path, - typer.Argument(exists=True), - ], + stoch_ffp: Annotated[Path, typer.Argument(exists=True)], station_file: Annotated[Path, typer.Argument(exists=True)], out_file: Annotated[Path, typer.Argument()], - hf_sim_path: Annotated[Path, typer.Option()] = Path( - "/EMOD3D/tools/hb_high_binmod_v6.0.3" - ), - work_directory: Annotated[ - Path, - typer.Option(exists=True, writable=True, file_okay=False), - ] = Path("/out"), ) -> None: - """Run the HF (High-Frequency) simulation and generate the HF output file. - - This function performs the following steps: - 1. Reads configuration and domain parameters from the realisation file. - 2. Filters stations based on their location relative to the domain. - 3. Uses multiprocessing to simulate each station and calculate epicentre distances. - 4. Reads the velocity model and calculates the `vs` value. - 5. Writes the HF output file, including header and station-specific data. + """Run the HF simulation and write the HF output file. Parameters ---------- @@ -392,30 +266,17 @@ def run_hf( Path to the file containing station locations and names. out_file : Path Filepath where the HF output will be saved. - hf_sim_path : Path, optional - Path to the HF simulation binary. - work_directory : Path, optional - Directory for intermediate files. Must be writable. - - Returns - ------- - None - The function does not return any value. It writes the HF output directly to `out_file`. """ + metadata = RealisationMetadata.read_from_realisation(realisation_ffp) seeds = Seeds.read_from_realisation_or_random(realisation_ffp) - domain_parameters = DomainParameters.read_from_realisation(realisation_ffp) - metadata = RealisationMetadata.read_from_realisation(realisation_ffp) - velocity_model = HFVelocityModel1D.read_from_realisation_or_defaults( - realisation_ffp, metadata.defaults_version - ) - hf_config = HFConfig.read_from_realisation_or_defaults( + hf_config = HFConfigDefaults.read_from_realisation_or_defaults( realisation_ffp, metadata.defaults_version ) rupture_velocity = RuptureVelocity.read_from_realisation_or_defaults( realisation_ffp, metadata.defaults_version ) - resolution = Resolution.read_from_realisation_or_defaults( + velocity_model_1d = HFVelocityModel1D.read_from_realisation_or_defaults( realisation_ffp, metadata.defaults_version ) @@ -423,61 +284,79 @@ def run_hf( station_file, delimiter=r"\s+", header=None, - names=["longitude", "latitude", "name"], - ).set_index("name") + names=["longitude", "latitude", "station"], + ).set_index("station") + + # Name-derived and order-invariant, so adding a station leaves every other station's + # waveform untouched and re-running a subset reproduces it exactly. stations["seed"] = station_seeds(seeds.hf_seed, stations.index) - velocity_model_path = work_directory / "velocity_model" - velocity_model.write_velocity_model(velocity_model_path) - nt = int( - np.float32(domain_parameters.duration) / np.float32(resolution.dt) - ) # Match Fortran's single-precision for consistent nt calculation - waveform = np.empty((3, len(stations), nt), dtype=np.float32) - - hf_input_template = build_hf_input( - stoch_ffp, - velocity_model_path, - resolution, - hf_config, - rupture_velocity, - domain_parameters, + # That invariance makes station order free to choose, and it is worth + # choosing. Runtime scales with subfault-to-station distance, and station + # files are spatially sorted, so the far stations all land in the last + # chunks and the run ends on a straggler. Sorting by the seed is a + # deterministic pseudorandom permutation. + stations = stations.sort_values("seed") + + simulator = Simulator( + build_slip_model(stoch_ffp), + build_velocity_model(velocity_model_1d), + build_config(hf_config, rupture_velocity, domain_parameters), ) - stations["epicentre_distance"] = np.nan - - with ThreadPoolExecutor(max_workers=utils.get_available_cores()) as executor: - station_index = {station: i for i, station in enumerate(stations.index)} - futures = [ - executor.submit( - hf_simulate_station, - hf_sim_path, - hf_input_template, - station["latitude"], - station["longitude"], - str(name), - int(station["seed"]), - ) - for name, station in stations.iterrows() - ] - for future in concurrent.futures.as_completed(futures): - station, epicentre, station_waveform = future.result() - stations.loc[station, "epicentre_distance"] = epicentre - i = station_index[station] - - for component in range(3): - waveform[component, i] = station_waveform[:, component] - - vs = velocity_model.model["Vs"].iloc[0] * 1000 - stations["vs"] = vs - - ds = create_hf_dataset( - waveform=waveform, - latitude=stations["latitude"], - longitude=stations["longitude"], - names=stations.index, - epicentre_distance=stations["epicentre_distance"], - seed=stations["seed"], - vref=stations["vs"], - dt=resolution.dt, - start_sec=hf_config.t_sec, + + # float32 throughout: this mirrors how the simulation truncates duration/dt to a + # sample count, so the dask template matches what comes back. + nt = int(np.float32(domain_parameters.duration) / np.float32(hf_config.dt)) + # The record starts at the origin time. This was a configurable `t_sec` that every + # realisation set to zero. + time = np.arange(nt) * hf_config.dt + + # Also bound by parallelism: chunk size set from memory alone gives 3 tasks for a + # 900-station run, so most of the allocation idles. Peak memory is + # num_workers * chunk_bytes, which this only ever lowers. + num_workers = utils.get_available_cores() + memory_chunk = TARGET_CHUNK_BYTES // (len(COMPONENTS) * nt * np.float32().itemsize) + chunk_size = max(1, min(memory_chunk, -(-len(stations) // (4 * num_workers)))) + logger = log_utils.get_logger(__name__) + logger.info( + "concurrency settings", + num_workers=num_workers, + memory_bound_stations=memory_chunk, + chunk_size=chunk_size, ) - ds.to_netcdf(out_file, engine="h5netcdf") + with ( + dask.config.set(scheduler="threads", num_workers=num_workers), + TqdmCallback(desc="Station chunks"), + ): + template = xr.DataArray( + da.empty( + (len(COMPONENTS), len(stations), nt), + dtype=np.float32, + chunks=(len(COMPONENTS), chunk_size, nt), + ), + dims=["component", "station", "time"], + coords={ + "component": list(COMPONENTS), + "station": stations.index, + "time": time, + }, + ) + + station_inputs = stations.to_xarray().chunk({"station": chunk_size}) + waveform = station_inputs.map_blocks( + simulate_chunk, + template=template, + kwargs={"time": time, "simulator": simulator}, + ).rename("waveform") + + station_inputs["vs"] = xr.full_like( + station_inputs["latitude"], velocity_model_1d.model["Vs"].iloc[0] * 1000 + ) + dataset = xr.merge([waveform, station_inputs]) + dataset.attrs = { + "start_sec": 0.0, + "dt": hf_config.dt, + "nt": nt, + "units": "cm/s^2", + } + dataset.to_netcdf(out_file, engine="h5netcdf") realisations.append_log_entry(realisation_ffp)