From 17d8dd1d57a1ecb928a109f7632a269fb9368754 Mon Sep 17 00:00:00 2001 From: Jake Faulkner Date: Mon, 31 Aug 2026 14:00:12 +1200 Subject: [PATCH] Rewrite bb-sim on dask and switch to the site-calculation models Two changes to the same function body, which is why they are one commit. **Larger than memory.** `bb-sim` read both waveform files entirely into memory and looped over the three components in Python. Now LF and HF are opened lazily, the common stations are selected on the *backend* arrays before chunking -- selecting after chunking makes station reordering an all-to-all dask shuffle, which materialises the whole array -- and the recombination runs per chunk under `map_blocks`. Chunking is over stations only, so each chunk holds complete traces for resampling, alignment and filtering. `resample_signal` and `align_datasets` replace the old pad-and-align: the two legs no longer have to share a timestep, so an SW4 LF run at one dt can be combined with HF at another. `relabel_hf_components` maps HF's 090/000/ver onto LF's x/y/z once, up front. **Site amplification.** Replaces `qcore.siteamp_models.cb_amp_multi` with the `site-calculation` models, selected by `bb.site_amp_version`: CB2014 as before, or BA2018. The amplification is now constrained to an explicit [fmin, fmax] band with logarithmic tapers at both ends, which is what the two new `fhightop` and `fmax` parameters are for -- the old code had a lowpass taper only. `site_amp_version` becomes an enum rather than a free string. It was previously declared, defaulted to "2014", and read by nothing at all. The shared broadband parameters move from each defaults version into root; the versions now carry only `flo`, which is the one value that genuinely differs between them. Station-dimension coordinates -- `supergrid_depth` among them -- ride through `map_blocks` untouched, which is why the supergrid penetration is carried as a coordinate rather than a data variable. This changes broadband results: different site amplification model, and a highpass taper where there was none. Co-Authored-By: Claude Opus 5 --- pyproject.toml | 2 + tests/test_realisation.py | 13 +- uv.lock | 27 ++ .../default_parameters/root/defaults.yaml | 6 + .../v24_2_2_1/defaults.yaml | 3 - .../v24_2_2_2/defaults.yaml | 3 - .../v24_2_2_4/defaults.yaml | 3 - .../v26_7_1Hz/defaults.yaml | 3 - workflow/realisations.py | 11 +- workflow/schemas.py | 15 +- workflow/scripts/bb_sim.py | 444 ++++++++++++++---- 11 files changed, 408 insertions(+), 122 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 256dfabf..177255a3 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -16,6 +16,7 @@ dependencies = [ "nshmdb>=2026.09.1", "oq_wrapper>=2025.12.3", "qcore-utils>=2025.12.2", + "site-calculation>=2026.7.1", "source_modelling>=2026.08.1", # Data Formats "geopandas", @@ -141,6 +142,7 @@ known-first-party = [ "IM", "mera", "velocity_modelling", + "site_calculation", ] [tool.ruff.lint.per-file-ignores] diff --git a/tests/test_realisation.py b/tests/test_realisation.py index 635bcdb7..e736fe95 100644 --- a/tests/test_realisation.py +++ b/tests/test_realisation.py @@ -604,7 +604,12 @@ def test_emod3d(tmp_path: Path) -> None: def test_broadband_parameters(tmp_path: Path) -> None: test_realisation = tmp_path / "realisation.json" broadband_parameters = realisations.BroadbandParameters( - flo=0.5, fmidbot=0.5, fmin=0.25, site_amp_version="2014" + flo=0.5, + fmidbot=0.5, + fmin=0.2, + fhightop=100.0, + fmax=101.0, + site_amp_version=schemas.SiteAmpModel.BA2018, ) broadband_parameters.write_to_realisation(test_realisation) with open(test_realisation, "r") as realisation_handle: @@ -612,8 +617,10 @@ def test_broadband_parameters(tmp_path: Path) -> None: "bb": { "flo": 0.5, "fmidbot": 0.5, - "fmin": 0.25, - "site_amp_version": "2014", + "fmin": 0.2, + "fhightop": 100.0, + "fmax": 101.0, + "site_amp_version": "ba2018", } } assert ( diff --git a/uv.lock b/uv.lock index c0762b94..a306a31b 100644 --- a/uv.lock +++ b/uv.lock @@ -3018,6 +3018,31 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/e0/f9/0595336914c5619e5f28a1fb793285925a8cd4b432c9da0a987836c7f822/shellingham-1.5.4-py2.py3-none-any.whl", hash = "sha256:7ecfff8f2fd72616f7481040475a65b2bf8af90a56c89140852d1120324e8686", size = 9755, upload-time = "2023-10-24T04:13:38.866Z" }, ] +[[package]] +name = "site-calculation" +version = "2026.7.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "numpy" }, + { name = "pyfftw" }, + { name = "scipy" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/ea/0e/d8c2513b130fd5e23cbc9455f10ac85541dd897b88563c2002c330509176/site_calculation-2026.7.1.tar.gz", hash = "sha256:9e32970aa5b5f4b2650d94441c75bf0dbe83f82cfe57d38ea258c2ddfa8857e7", size = 127578, upload-time = "2026-07-31T04:46:27.647Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/1d/64/f7cabee255e847e01a1fca39933aa5a9ace6ba0ac552c2fd40b73154a07b/site_calculation-2026.7.1-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:9a66ea2e32ed2b73ef91dddb091117775e0f7894526382caa5a7b5883c1ed325", size = 254857, upload-time = "2026-07-31T04:46:12.745Z" }, + { url = "https://files.pythonhosted.org/packages/16/9c/774d399059501ad14293e3a373fe5608b87696f4cad9606a7a32a5a5422f/site_calculation-2026.7.1-cp313-cp313-macosx_11_0_x86_64.whl", hash = "sha256:7dbc8d4bb442af684593279e05f749f152d307458955b56ab689d85e5006d5bf", size = 259469, upload-time = "2026-07-31T04:46:14.067Z" }, + { url = "https://files.pythonhosted.org/packages/4c/9d/b1fcdfc26fcc202d7800cbd13f63977ddb774091dd10985f92afe02db2bc/site_calculation-2026.7.1-cp313-cp313-manylinux_2_28_x86_64.whl", hash = "sha256:1d13f6f7a178162a6d81770d9ce1c8a4e3a0caaa90da9f717200fa5db2a644d9", size = 294668, upload-time = "2026-07-31T04:46:15.313Z" }, + { url = "https://files.pythonhosted.org/packages/ee/25/1c154149e6b6b0967fc28e919b19b455836f05a3372d979f8405dfe5a5b6/site_calculation-2026.7.1-cp313-cp313-win_amd64.whl", hash = "sha256:c06018d59f8dd4fdc6c89c30b125e6e671e2fa50a54fada663f0a17c8f11b2de", size = 149183, upload-time = "2026-07-31T04:46:16.575Z" }, + { url = "https://files.pythonhosted.org/packages/be/7e/5335a1b9ce66c060ce6f089c7271ca21d488ea03265df2f620c220039399/site_calculation-2026.7.1-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:26db03728f865c738e65f2aae027e545f6e8f25f47497d9fcd7e370bcbd24905", size = 254473, upload-time = "2026-07-31T04:46:18.134Z" }, + { url = "https://files.pythonhosted.org/packages/0c/69/6487c0bc361c82214e032d916ea125b28fdc0c0c8672a0215444500d646c/site_calculation-2026.7.1-cp314-cp314-macosx_11_0_x86_64.whl", hash = "sha256:23b2443f599ba737a6158a6742b27c051bb362355915115fdfb7a07ab3415439", size = 259136, upload-time = "2026-07-31T04:46:19.666Z" }, + { url = "https://files.pythonhosted.org/packages/da/52/49553e8d6194dba19439703f25e628e2bb9388aaec14f06e48d3b5c62ece/site_calculation-2026.7.1-cp314-cp314-manylinux_2_28_x86_64.whl", hash = "sha256:67855ec4c7c3288e753f349b6f6d915e86a05c42ddbca9bdead24c780639a79e", size = 294532, upload-time = "2026-07-31T04:46:20.854Z" }, + { url = "https://files.pythonhosted.org/packages/55/b3/6d9670475b007be219b249300b9f2c6167bd4f65ecd34afc42e5351d76f0/site_calculation-2026.7.1-cp314-cp314-win_amd64.whl", hash = "sha256:2c428c6d491728f9ab502819a1ae1e4cb95a84293e6da938dc8dfd9dfc1313bb", size = 153606, upload-time = "2026-07-31T04:46:21.963Z" }, + { url = "https://files.pythonhosted.org/packages/0c/36/5f24de7b7cf370e552108f180e4403f85253257db3df278dde86c1a31cc0/site_calculation-2026.7.1-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:5cd24e0c69905436b86073edcde12f2b4981d0e1b091366b435bf5b82788ed69", size = 254443, upload-time = "2026-07-31T04:46:23.089Z" }, + { url = "https://files.pythonhosted.org/packages/67/c0/8981b96ffbef6da17d9f85f946839dc344b914a9db9a343a6a76f0cfeaee/site_calculation-2026.7.1-cp314-cp314t-macosx_11_0_x86_64.whl", hash = "sha256:01804b221c4ed7332052b27693c41f147d0beaf6b48e7def4dc8904af5050340", size = 258922, upload-time = "2026-07-31T04:46:24.307Z" }, + { url = "https://files.pythonhosted.org/packages/31/4c/03bbd62ccdea3780dd154ee078b15a2d3fac8cda3c1f98bbf9d0976c5511/site_calculation-2026.7.1-cp314-cp314t-manylinux_2_28_x86_64.whl", hash = "sha256:1c86e27e6d985a827fafef1bbe841060c411df4af00d62ba8ec077ac0815e2ef", size = 294490, upload-time = "2026-07-31T04:46:25.426Z" }, + { url = "https://files.pythonhosted.org/packages/e4/16/1c6cd0b224abf7e88361679657243265eb3c89d863daa9d2d7c37e09cf37/site_calculation-2026.7.1-cp314-cp314t-win_amd64.whl", hash = "sha256:81e9bffdc595d85d25044a1d2119e32266d0886c387373b5da1a5c6535f8a897", size = 153507, upload-time = "2026-07-31T04:46:26.504Z" }, +] + [[package]] name = "six" version = "1.17.0" @@ -3512,6 +3537,7 @@ dependencies = [ { name = "schema" }, { name = "scipy" }, { name = "shapely" }, + { name = "site-calculation" }, { name = "source-modelling" }, { name = "structlog" }, { name = "tqdm" }, @@ -3569,6 +3595,7 @@ requires-dist = [ { name = "scipy" }, { name = "scipy-stubs", marker = "extra == 'types'" }, { name = "shapely" }, + { name = "site-calculation", specifier = ">=2026.7.1" }, { name = "source-modelling", specifier = ">=2026.8.1" }, { name = "structlog" }, { name = "tqdm" }, diff --git a/workflow/default_parameters/root/defaults.yaml b/workflow/default_parameters/root/defaults.yaml index 90b256f6..893bd35e 100644 --- a/workflow/default_parameters/root/defaults.yaml +++ b/workflow/default_parameters/root/defaults.yaml @@ -1526,3 +1526,9 @@ hf_velocity_model_1d: rho: 3.33 Qp: 394.80 Qs: 197.40 +bb: + fmin: 0.2 + fmidbot: 0.5 + fhightop: 100.0 + fmax: 101.0 + site_amp_version: "ba2018" diff --git a/workflow/default_parameters/v24_2_2_1/defaults.yaml b/workflow/default_parameters/v24_2_2_1/defaults.yaml index 8c990b3c..514a3707 100644 --- a/workflow/default_parameters/v24_2_2_1/defaults.yaml +++ b/workflow/default_parameters/v24_2_2_1/defaults.yaml @@ -3,6 +3,3 @@ resolution: resolution: 0.1 bb: flo: 1.0 - fmidbot: 0.5 - fmin: 0.25 - site_amp_version: "2014" diff --git a/workflow/default_parameters/v24_2_2_2/defaults.yaml b/workflow/default_parameters/v24_2_2_2/defaults.yaml index ec662d84..47072a4d 100644 --- a/workflow/default_parameters/v24_2_2_2/defaults.yaml +++ b/workflow/default_parameters/v24_2_2_2/defaults.yaml @@ -3,6 +3,3 @@ resolution: resolution: 0.2 bb: flo: 0.5 - fmidbot: 0.5 - fmin: 0.25 - site_amp_version: "2014" diff --git a/workflow/default_parameters/v24_2_2_4/defaults.yaml b/workflow/default_parameters/v24_2_2_4/defaults.yaml index beb9c260..ddecf8bb 100644 --- a/workflow/default_parameters/v24_2_2_4/defaults.yaml +++ b/workflow/default_parameters/v24_2_2_4/defaults.yaml @@ -3,6 +3,3 @@ resolution: resolution: 0.4 bb: flo: 0.25 - fmidbot: 0.5 - fmin: 0.25 - site_amp_version: "2014" diff --git a/workflow/default_parameters/v26_7_1Hz/defaults.yaml b/workflow/default_parameters/v26_7_1Hz/defaults.yaml index 635c53d6..25db0c72 100644 --- a/workflow/default_parameters/v26_7_1Hz/defaults.yaml +++ b/workflow/default_parameters/v26_7_1Hz/defaults.yaml @@ -1,9 +1,6 @@ --- bb: flo: 1.0 - fmidbot: 0.5 - fmin: 0.2 - site_amp_version: "ba2018" velocity_model: # Overrides the root default of 2.0 km, which is an EMOD3D value and has no # absorbing layer to clear. Two criteria bear on this number, and only the diff --git a/workflow/realisations.py b/workflow/realisations.py index 53a104b4..3617570f 100644 --- a/workflow/realisations.py +++ b/workflow/realisations.py @@ -1366,10 +1366,15 @@ class BroadbandParameters(RealisationConfiguration): flo: float """low/high frequency cutoff.""" fmidbot: float - """fmidbot for site amplification""" + """fmidbot for the site amplification lowpass taper.""" fmin: float - """fmin for site amplification.""" - site_amp_version: str + """fmin for the site amplification lowpass taper.""" + fhightop: float + """fhightop for the site amplification highpass taper.""" + fmax: float + """fmax for the site amplification highpass taper.""" + site_amp_version: schemas.SiteAmpModel + """The site amplification model to apply.""" @dataclasses.dataclass diff --git a/workflow/schemas.py b/workflow/schemas.py index ef7cb128..a7b5ae18 100644 --- a/workflow/schemas.py +++ b/workflow/schemas.py @@ -35,6 +35,13 @@ class Stype(StrEnum): seki = "seki" +class SiteAmpModel(StrEnum): + """Site amplification models for broadband simulation.""" + + CB2014 = "cb2014" + BA2018 = "ba2018" + + class KModel(IntEnum): """Correlation length models for genslip.""" @@ -1168,7 +1175,13 @@ def _corners_to_array(corners_spec: list[dict[str, float]]) -> np.ndarray: Literal("fmin", description="fmin for site amplification"): And( NUMBER, _is_non_negative ), - "site_amp_version": str, + Literal("fhightop", description="fhightop for site amplification"): And( + NUMBER, _is_non_negative + ), + Literal("fmax", description="fmax for site amplification"): And( + NUMBER, _is_non_negative + ), + "site_amp_version": Use(SiteAmpModel), } ) diff --git a/workflow/scripts/bb_sim.py b/workflow/scripts/bb_sim.py index 9b03ae3e..037c3955 100644 --- a/workflow/scripts/bb_sim.py +++ b/workflow/scripts/bb_sim.py @@ -39,78 +39,285 @@ from typing import Annotated import numpy as np -import numpy.typing as npt import pandas as pd +import pyfftw +import scipy as sp import typer import xarray as xr -from qcore import cli, siteamp_models, timeseries +from qcore import cli, timeseries +from site_calculation import amplification from workflow import log_utils, realisations -from workflow.realisations import ( - BroadbandParameters, - RealisationMetadata, - Resolution, -) +from workflow.realisations import BroadbandParameters, RealisationMetadata +from workflow.schemas import SiteAmpModel + +# Site amplification model -> (amplification function, model frequencies). +# Both models share the same (vs30, vs30_sim, pga) calling convention. +SITE_AMP_MODELS = { + SiteAmpModel.CB2014: ( + amplification.campbell_bozorgnia_2014, + amplification.CAMPBELL_BOZORGNIA_2014_FREQUENCIES, + ), + SiteAmpModel.BA2018: ( + amplification.bayless_abrahamson_2018, + amplification.BAYLESS_ABRAHAMSON_2018_FREQUENCIES, + ), +} app = typer.Typer() G = 1 / 981.0 +TARGET_CHUNK_BYTES = 256 * 2**20 -def align_waveforms( - lf_waveform: npt.NDArray[np.floating], - hf_waveform: npt.NDArray[np.floating], - lf_start: float, - hf_start: float, - dt: float, -) -> tuple[ - npt.NDArray[np.floating], npt.NDArray[np.floating], npt.NDArray[np.floating] -]: - """Align LF and HF waveforms to a common time axis. +def align_datasets( + lf: xr.Dataset, hf: xr.Dataset, dt: float +) -> tuple[xr.DataArray, xr.DataArray]: + """Lazily align LF and HF waveforms onto a common time axis. + + Both waveforms are zero-padded to span the same time domain, + running from the earliest start to the latest end of the two + simulations. Parameters ---------- - lf_waveform : array of floats - The low-frequency waveform to align. - hf_waveform : array of floats - The high-frequency waveform to align. - lf_start : float - The start of the LF simulation. - hf_start : float - The start of the HF simulation. + lf : xr.Dataset + The low-frequency dataset, with a 'start_sec' attribute. + hf : xr.Dataset + The high-frequency dataset, with a 'start_sec' attribute. dt : float - The timestep for the simulation. + The shared timestep of both datasets. Returns ------- - array of floats - The aligned low-frequency results. - array of floats - The aligned high-frequency results. - array of floats - The new time values. + xr.DataArray + The aligned low-frequency waveform. + xr.DataArray + The aligned high-frequency waveform. """ + lf_start = lf.attrs["start_sec"] + hf_start = hf.attrs["start_sec"] + start = min(lf_start, hf_start) + lf_offset = round((lf_start - start) / dt) + hf_offset = round((hf_start - start) / dt) + common_nt = max(lf_offset + lf.sizes["time"], hf_offset + hf.sizes["time"]) + common_time = start + np.arange(common_nt) * dt + + def pad_waveform(waveform: xr.DataArray, offset: int) -> xr.DataArray: + padded = waveform.pad( + time=(offset, common_nt - offset - waveform.sizes["time"]), + constant_values=0.0, + ) + return padded.assign_coords(time=common_time) + + return ( + pad_waveform(lf["waveform"], lf_offset), + pad_waveform(hf["waveform"], hf_offset), + ) - lf_nt = lf_waveform.shape[1] - hf_nt = hf_waveform.shape[1] - lf_time = lf_start + np.arange(lf_nt) * dt - hf_time = hf_start + np.arange(hf_nt) * dt +def resample_signal(dset: xr.Dataset, dt: float) -> xr.Dataset: + """Resample waveform dataset to a new time step. - start = min(lf_time[0], hf_time[0]) - end = max(lf_time[-1], hf_time[-1]) - common_time = np.arange(start, end + dt / 2, dt) + Parameters + ---------- + dset : xr.Dataset + Input dataset with dimensions (component, station, time) and + attributes 'dt'. + dt : float + Desired time step in seconds. - lf_aligned = np.zeros((lf_waveform.shape[0], common_time.size), dtype=np.float32) - hf_aligned = np.zeros((hf_waveform.shape[0], common_time.size), dtype=np.float32) + Returns + ------- + xr.Dataset + Resampled dataset with updated time coordinates and dt attribute. + """ + duration = dset["waveform"].sizes["time"] * dset.attrs["dt"] + nt = round(duration / dt) + + # NOTE: I am not providing a default start second because we consider it an + # error not to provide one (no implicit magic behaviour). + new_time = np.arange(nt) * dt + dset.attrs["start_sec"] + + resampled_waveform = xr.apply_ufunc( + sp.signal.resample, + dset["waveform"], + # This tells xarray that resample expects an array with all of the time component intact. + # So it will be passed arrays of shape (n_component, n_stations, n_time) = (i, j, nt) + input_core_dims=[["time"]], + # This tells xarray that the time dimension is going to be returned in + # its entirety by scipy resample. + output_core_dims=[["time"]], + # This tells xarray that the time coordinates from the dset dataset are no + # longer any good. They will be dropped from the output array. + exclude_dims=set(["time"]), + # Array passed to resample will have time in the inner-most axis and the + # default axis for resample is 0. + kwargs=dict(num=nt, axis=-1), + dask="parallelized", + # The size of the resampled time dimension cannot be inferred by + # dask, so it must be given explicitly. + dask_gufunc_kwargs=dict(output_sizes={"time": nt}), + ).chunk({"time": -1, "component": -1, "station": dset.chunksizes["station"]}) + + resampled_waveform = resampled_waveform.assign_coords(time=new_time) + # Must drop both waveform variable and time dimension to avoid xarray + # automatically reindexing the waveform according to the new axes. + new_dset = dset.drop_vars(["waveform", "time"]).assign(waveform=resampled_waveform) + new_dset.attrs["dt"] = dt + return new_dset + + +# Reference Vs30 (m/s) of the high-frequency simulation, i.e. the Vs30 +# the waveforms are amplified *from* towards each station's target Vs30. +VS30_SIM = 500.0 + + +LF_COMPONENTS = ["x", "y", "z"] +"""Component labels every low-frequency output uses. + +Written by `lf-to-xarray` for SW4 and by `qcore.timeseries` for EMOD3D: +x = east-west, y = north-south, z = vertical. +""" + +HF_COMPONENT_TO_LF = {"090": "x", "000": "y", "ver": "z"} +"""Component relabelling from the high-frequency simulation's convention. + +`hf_simulation` names components by azimuth -- 090 for east, 000 for north, +and `ver` for the vertical -- where the low-frequency outputs name the same +three x/y/z. Same components, same order, different labels. +""" - lf_indices = ((lf_time - start) / dt).round().astype(int) - hf_indices = ((hf_time - start) / dt).round().astype(int) - lf_aligned[:, lf_indices] = lf_waveform - hf_aligned[:, hf_indices] = hf_waveform +def relabel_hf_components(hf: xr.Dataset) -> xr.Dataset: + """Put an HF dataset's components on the low-frequency naming. + + Parameters + ---------- + hf : xr.Dataset + The high-frequency dataset, as written by `hf-sim`. + + Returns + ------- + xr.Dataset + The same dataset with its `component` coordinate relabelled. Labels + already in the low-frequency convention are left alone, so this is + idempotent. + + Raises + ------ + ValueError + If the components still do not match the low-frequency set. Without + this check the mismatch is silent: xarray aligns the two datasets on + `component`, finds no labels in common, and fills every high-frequency + sample with NaN. The failure then surfaces a long way downstream, as a + NaN somewhere inside the site amplification model. + """ + components = [ + HF_COMPONENT_TO_LF.get(str(component), str(component)) + for component in hf.component.values + ] + if set(components) != set(LF_COMPONENTS): + raise ValueError( + f"High-frequency components {list(hf.component.values)} do not " + f"correspond to the low-frequency components {LF_COMPONENTS}. " + f"Add the mapping to HF_COMPONENT_TO_LF." + ) + return hf.assign_coords(component=components) + + +def _process_bb_chunk( + dset: xr.Dataset, + dt: float, + flo: float, + fmin: float, + fmidbot: float, + fhightop: float, + fmax: float, + site_amp_model: SiteAmpModel, +) -> xr.Dataset: + """Compute broadband waveforms for a chunk of stations. + + Applies the selected site amplification model to the high-frequency + waveforms, then merges them with the low-frequency waveforms using a + matched pair of high-pass and low-pass Butterworth filters. + + Parameters + ---------- + dset : xr.Dataset + Dataset with variables ``lf_waveform`` and ``hf_waveform`` + (dims component, station, time) on a common time axis, and + ``vs30`` (dims station). + dt : float + Broadband timestep. + flo : float + The frequency (Hz) at which the low-frequency and + high-frequency waveforms are merged. + fmin : float + Frequency (Hz) below which the site amplification is tapered + out (lowpass end of the amplification band). + fmidbot : float + Frequency (Hz) above which the site amplification is applied in + full at the lowpass end. + fhightop : float + Frequency (Hz) below which the site amplification is applied in + full at the highpass end. + fmax : float + Frequency (Hz) above which the site amplification is tapered out + (highpass end of the amplification band). + site_amp_model : SiteAmpModel + The site amplification model to apply. + + Returns + ------- + xr.Dataset + Dataset with a single ``waveform`` variable containing the + broadband waveforms in units of g. + """ + lf_waveform = dset["lf_waveform"].values + hf_waveform = dset["hf_waveform"].values + nt = lf_waveform.shape[-1] + + amp_model_fn, amp_model_freqs = SITE_AMP_MODELS[site_amp_model] + + # Zero-pad to a length pyfftw can transform efficiently, and + # pre-compute the FFT output frequencies the amplification is + # sampled at. + n_fft = pyfftw.next_fast_len(nt) + fft_freqs = np.fft.rfftfreq(n_fft, dt) + + # The amplification models require float64 inputs. + vs30 = dset["vs30"].values.astype(np.float64) + vs30_sim = np.full_like(vs30, VS30_SIM) + + bb_waveform = np.empty(lf_waveform.shape, dtype=np.float32) + # Site amplification depends on each component's PGA, so amplify + # component-by-component (vectorised over stations). + for i in range(bb_waveform.shape[0]): + pga = np.abs(hf_waveform[i]).max(axis=-1).astype(np.float64) * G + + amp = amp_model_fn(vs30, vs30_sim, pga) + amp = amplification.interpolate_frequencies(amp_model_freqs, fft_freqs, amp) + # Constrain the amplification to the [fmin, fmax] band, tapering + # logarithmically at either end. + amplification.amp_lowpass(fft_freqs, amp, fmin, fmidbot) + amplification.amp_highpass(fft_freqs, amp, fhightop, fmax) + + # Taper the tail of the HF waveform (5%) to limit spectral + # leakage before amplification. + hf_component = hf_waveform[i].copy() + amplification.taper(hf_component, 0.05) + hf_amped = amplification.amplify_waveform(hf_component, amp, n_fft) + + hf_filtered = timeseries.bwfilter(hf_amped, dt, flo, timeseries.Band.HIGHPASS) + lf_filtered = timeseries.bwfilter( + lf_waveform[i], dt, flo, timeseries.Band.LOWPASS + ) + bb_waveform[i] = (hf_filtered + lf_filtered) * G - return lf_aligned, hf_aligned, common_time + return dset.drop_vars(["lf_waveform", "hf_waveform", "vs30"]).assign( + waveform=(("component", "station", "time"), bb_waveform) + ) @cli.from_docstring(app) @@ -145,15 +352,35 @@ def combine_hf_and_lf( broadband_config = BroadbandParameters.read_from_realisation_or_defaults( realisation_ffp, metadata.defaults_version ) - resolution = Resolution.read_from_realisation_or_defaults( - realisation_ffp, metadata.defaults_version - ) - bb_dt = resolution.dt - - # load data stores + # Open lazily (no dask) and select the common stations *before* chunking. + # The LF and HF files store stations in different orders, so selecting after + # chunking is an all-to-all dask shuffle in which every output chunk depends + # on every input chunk. This will result in materialising the whole array + # in-memory. Selecting on the lazy backend arrays instead lets each dask + # chunk read just its own stations from disk. lf = xr.open_dataset(low_frequency_waveform_file) - hf = xr.open_dataset(high_frequency_waveform_file) - common_stations = list( + lf = lf.drop_duplicates("station", keep="first") + hf = relabel_hf_components(xr.open_dataset(high_frequency_waveform_file)) + hf = hf.drop_duplicates("station", keep="first") + + common_stations = sorted( + set(map(str, hf.station.values)) & set(map(str, lf.station.values)) + ) + # Chunk over stations only, so every chunk holds complete time + # series for resampling, alignment and filtering. + nt = max(len(lf["time"]), len(hf["time"])) + n_stations = round(TARGET_CHUNK_BYTES / (3 * nt * np.float64().itemsize)) + chunking = {"component": -1, "station": n_stations, "time": -1} + lf = lf.sel(station=common_stations).chunk(chunking) + hf = hf.sel(station=common_stations).chunk(chunking) + + bb_dt = min(lf.attrs["dt"], hf.attrs["dt"]) + + if not np.isclose(lf.attrs["dt"], bb_dt): + lf = resample_signal(lf, bb_dt) + if not np.isclose(hf.attrs["dt"], bb_dt): + hf = resample_signal(hf, bb_dt) + common_stations = sorted( set(map(str, hf.station.values)) & set(map(str, lf.station.values)) ) hf = hf.sel(station=common_stations) @@ -166,62 +393,73 @@ def combine_hf_and_lf( ).set_index("station") vs30_df["vsite"] = vs30_df["vsite"].astype(np.float32) vs30_df = vs30_df.loc[common_stations] - vs30_df["vref"] = 500.0 - vs30_df["vpga"] = 500.0 - - bb_waveforms = [] - new_time_coords = None - - for i, (lf_component, hf_component) in enumerate(zip(lf.component, hf.component)): - hf_waveform_raw = hf.sel(component=hf_component).waveform.values - lf_waveform_raw = lf.sel(component=lf_component).waveform.values - - temp_lf_padded, temp_hf_padded, new_time_coords = align_waveforms( - lf_waveform_raw, - hf_waveform_raw, - lf.attrs["start_sec"], - hf.attrs["start_sec"], - bb_dt, - ) - bb_nt = temp_lf_padded.shape[1] - vs30_df["pga"] = np.abs(temp_hf_padded).max(axis=1) * G - - assert isinstance(vs30_df, pd.DataFrame) - hf_amp_val = siteamp_models.cb_amp_multi(vs30_df) - hf_amp_fas_vals = siteamp_models.cb2014_to_fas_amplification_factors( - hf_amp_val, - bb_dt, - bb_nt, - ) - hf_waveform_amped = timeseries.ampdeamp( - temp_hf_padded, hf_amp_fas_vals, amplify=True - ) - hf_filtered = timeseries.bwfilter( - hf_waveform_amped, bb_dt, broadband_config.flo, timeseries.Band.HIGHPASS - ) - lf_filtered = timeseries.bwfilter( - temp_lf_padded, bb_dt, broadband_config.flo, timeseries.Band.LOWPASS - ) - bb_waveforms.append((hf_filtered + lf_filtered) * G) - - bb_waveform = np.stack(bb_waveforms, dtype=np.float32) - - xr.Dataset( - {"waveform": (["component", "station", "time"], bb_waveform)}, + lf_aligned, hf_aligned = align_datasets(lf, hf, bb_dt) + + # Station-dimension *coordinates* on `lf_aligned` (`supergrid_depth`, and + # EMOD3D's `x`/`y`) ride from here to the intensity measure file with no + # help: they survive `map_blocks` below, `_process_bb_chunk` drops data + # variables only, and `IM.ims` keeps the input's non-dimension + # coordinates. Anything that must reach the IM file and is *not* a + # coordinate has to be hand-carried, the way `vs30` is. + combined = xr.Dataset( + { + "lf_waveform": lf_aligned, + # reset_coords drops the HF lat/lon coordinates, which would + # otherwise conflict with the LF-derived latitude/longitude. + "hf_waveform": hf_aligned.reset_coords(drop=True), + "vs30": vs30_df["vsite"].to_xarray(), + }, coords={ "component": ("component", ["x", "y", "z"]), "station": ("station", common_stations), - "time": ("time", new_time_coords), - "x": ("station", lf.x.values), - "y": ("station", lf.y.values), + "time": lf_aligned.time, "latitude": ("station", lf.lat.values), "longitude": ("station", lf.lon.values), }, - attrs={ - "units": "g", - }, - ).to_netcdf( + attrs={"units": "g"}, + ).chunk(chunking) + + combined = combined.unify_chunks() + template = ( + combined["lf_waveform"].astype(np.float32).rename("waveform").to_dataset() + ) + template.attrs = combined.attrs + + bb = xr.map_blocks( + _process_bb_chunk, + combined, + kwargs=dict( + dt=bb_dt, + flo=broadband_config.flo, + fmin=broadband_config.fmin, + fmidbot=broadband_config.fmidbot, + fhightop=broadband_config.fhightop, + fmax=broadband_config.fmax, + site_amp_model=broadband_config.site_amp_version, + ), + template=template, + ) + bb["vs30"] = combined["vs30"] + attributes = dict( + dt=bb_dt, + flo=broadband_config.flo, + fmin=broadband_config.fmin, + fmidbot=broadband_config.fmidbot, + fhightop=broadband_config.fhightop, + fmax=broadband_config.fmax, + site_amp_model=str(broadband_config.site_amp_version), + ) + # Attributes, unlike station coordinates, are *not* carried through + # map_blocks: `template` above only has `combined`'s. The LF file's + # supergrid width describes the run that produced the waveforms, and + # `im-calc` writes it into the IM file's root attributes, so pass it on. + attributes |= { + name: lf.attrs[name] for name in ("SGWIDTH", "SGWIDTHGP") if name in lf.attrs + } + bb.attrs.update(attributes) + + bb.to_netcdf( output_ffp, engine="h5netcdf", encoding={