diff --git a/pyproject.toml b/pyproject.toml index 15d482eb..d5bd1055 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -11,7 +11,7 @@ requires-python = ">=3.12" dynamic = ["version"] dependencies = [ # UCGMSim Dependencies - "im-calculation>=2025.12.5", + "im-calculation @ git+https://github.com/ucgmsim/IM_calculation@no_parallel", "velocity-modelling>=2026.8.1", "nshmdb>=2026.09.1", "oq_wrapper>=2025.12.3", @@ -37,6 +37,7 @@ dependencies = [ "h5py>=3.15.1", "parse>=1.21.0", "rich>=14.3.2", + "dask>=2026.6.0", ] [project.optional-dependencies] @@ -171,3 +172,6 @@ exclude = [ # don't report on objects that match any of these regex '\.undocumented_method$', '\.__repr__$', ] + +[tool.uv.sources] +im-calculation = { git = "https://github.com/ucgmsim/im_calculation", branch = "no_parallel" } diff --git a/uv.lock b/uv.lock index f165710f..2a43ba87 100644 --- a/uv.lock +++ b/uv.lock @@ -481,6 +481,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/73/86/43fa9f15c5b9fb6e82620428827cd3c284aa933431405d1bcf5231ae3d3e/cligj-0.7.2-py3-none-any.whl", hash = "sha256:c1ca117dbce1fe20a5809dc96f01e1c2840f6dcc939b3ddbb1111bf330ba82df", size = 7069, upload-time = "2021-05-28T21:23:26.877Z" }, ] +[[package]] +name = "cloudpickle" +version = "3.1.2" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/27/fb/576f067976d320f5f0114a8d9fa1215425441bb35627b1993e5afd8111e5/cloudpickle-3.1.2.tar.gz", hash = "sha256:7fda9eb655c9c230dab534f1983763de5835249750e85fbcef43aaa30a9a2414", size = 22330, upload-time = "2025-11-03T09:25:26.604Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/88/39/799be3f2f0f38cc727ee3b4f1445fe6d5e4133064ec2e4115069418a5bb6/cloudpickle-3.1.2-py3-none-any.whl", hash = "sha256:9acb47f6afd73f60dc1df93bb801b472f05ff42fa6c84167d25cb206be1fbf4a", size = 22228, upload-time = "2025-11-03T09:25:25.534Z" }, +] + [[package]] name = "colorama" version = "0.4.6" @@ -664,6 +673,24 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/e7/05/c19819d5e3d95294a6f5947fb9b9629efb316b96de511b418c53d245aae6/cycler-0.12.1-py3-none-any.whl", hash = "sha256:85cef7cff222d8644161529808465972e51340599459b8ac3ccbac5a854e0d30", size = 8321, upload-time = "2023-10-07T05:32:16.783Z" }, ] +[[package]] +name = "dask" +version = "2026.8.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "click" }, + { name = "cloudpickle" }, + { name = "fsspec" }, + { name = "packaging" }, + { name = "partd" }, + { name = "pyyaml" }, + { name = "toolz" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/33/a7/6b3c7ac32b642fbbe0821111654e0bd8cfbe88f68560bcf23cc78ab35c71/dask-2026.8.0.tar.gz", hash = "sha256:8a94c37b5de6d869343340dc26c3c3acca7ec48a3abdabe00ea3abb1125884d5", size = 11561752, upload-time = "2026-08-24T19:21:25.906Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/f8/3a/4fc99e788bcfa1b3b3f21abf57da45898d807d007e7f6fd1c7300904eb70/dask-2026.8.0-py3-none-any.whl", hash = "sha256:ccc0c83a189b0398602435189771d28dad7b5773b6089bb8dce14ae732dd782c", size = 1492182, upload-time = "2026-08-24T19:21:23.997Z" }, +] + [[package]] name = "decorator" version = "5.3.1" @@ -1166,8 +1193,8 @@ wheels = [ [[package]] name = "im-calculation" -version = "2026.7.2" -source = { registry = "https://pypi.org/simple" } +version = "2026.7.3.dev4+g5a0c1a5e1" +source = { git = "https://github.com/ucgmsim/im_calculation?branch=no_parallel#5a0c1a5e1b6ae9728a8e68806b220aa65c9ff111" } dependencies = [ { name = "numpy" }, { name = "pandas", extra = ["hdf5"] }, @@ -1175,25 +1202,9 @@ dependencies = [ { name = "pyfftw" }, { name = "qcore-utils" }, { name = "scipy" }, - { name = "tqdm" }, { name = "typer" }, { name = "xarray", extra = ["io"] }, ] -sdist = { url = "https://files.pythonhosted.org/packages/d3/3a/0cbe62d24188c3f19b70e001bb5f67d2652b11c8e418120f2d34dce6a19e/im_calculation-2026.7.2.tar.gz", hash = "sha256:5ff212f205cb497e72a23715d5d08a9e4bc8460b5a3c3fa1dec85a7ed9cd85ee", size = 36853, upload-time = "2026-07-24T01:37:37.363Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/44/41/19df181c1b329737abe77fc2ed4b7387aa2d6254aa5d1ae77b48224a9f56/im_calculation-2026.7.2-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:b3dbe685ac7c899cb1ccbb85c254820fe3eff8d74bb040ac8b0e63c6bb8720e8", size = 335354, upload-time = "2026-07-24T01:37:21.588Z" }, - { url = "https://files.pythonhosted.org/packages/ef/a3/b8da5118812dc422fb57ba90fe937a50b795086d3ced08b0b34089b5c5e4/im_calculation-2026.7.2-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:207f80a357c369602b4bc225905f518489d526ee48f61d408703b730bdf41b06", size = 324311, upload-time = "2026-07-24T01:37:22.957Z" }, - { url = "https://files.pythonhosted.org/packages/0c/49/e9d8f63d39ba119a78fa55f6b66331c07d1a8c652a969b04d15fbe913fe0/im_calculation-2026.7.2-cp312-cp312-manylinux_2_28_x86_64.whl", hash = "sha256:b18552e381ecc2b66e83e65184f8bc96eb3846ee27eb5ad2c495432f54355571", size = 379780, upload-time = "2026-07-24T01:37:24.397Z" }, - { url = "https://files.pythonhosted.org/packages/02/e0/7497a21bcab4ec718900b30fbf06b4fcc0a72106d3531d8a0b85c359518b/im_calculation-2026.7.2-cp312-cp312-win_amd64.whl", hash = "sha256:058792616f43eecf3d476e8fce2b7877df08f10ffb960b0668b98ad9aab618a5", size = 214835, upload-time = "2026-07-24T01:37:25.888Z" }, - { url = "https://files.pythonhosted.org/packages/b9/1a/1665f6b1537044ea37a487efbb974d97fc0473678a510647430283b1db2f/im_calculation-2026.7.2-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:c1cd23663eb0c7c31a6a519a02b31c49cdc6b93b19dc5c0ec681573bbe86bc62", size = 335722, upload-time = "2026-07-24T01:37:27.138Z" }, - { url = "https://files.pythonhosted.org/packages/af/75/5aaae7c0cbeac1aacd0d12fd63fe16f2adebabd6a81b25c3270e89d68028/im_calculation-2026.7.2-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:975ac472625f8e89243433469d80a16e9965e2b25d43d425bb769f1af1336c03", size = 324461, upload-time = "2026-07-24T01:37:28.504Z" }, - { url = "https://files.pythonhosted.org/packages/a7/c8/8dcbd6d3c9437c7142f4b3ac3d30c88788a88a6b9afe2e94fe7c97cc7038/im_calculation-2026.7.2-cp313-cp313-manylinux_2_28_x86_64.whl", hash = "sha256:40e482eb7f27c2997811638be9226a600840174cbeb45154f92e15dce0b8c9e9", size = 379874, upload-time = "2026-07-24T01:37:29.701Z" }, - { url = "https://files.pythonhosted.org/packages/6e/69/7359ae8205fa123081e50e5f6e52961413e9bbe978282738f55700819e06/im_calculation-2026.7.2-cp313-cp313-win_amd64.whl", hash = "sha256:8e33932b8c161c6ecb07337a2a344b7eee3687bcf93f29120e764678416255f7", size = 214804, upload-time = "2026-07-24T01:37:31.027Z" }, - { url = "https://files.pythonhosted.org/packages/c7/9a/9cad373586173e0cae969948e79abd7856c1dd2c7a36b07f9c844ce955bf/im_calculation-2026.7.2-cp314-cp314-macosx_10_15_x86_64.whl", hash = "sha256:ada009ccdad40eff4d7c5372e6d505aafa5edae457bb3c8e011286fea9f4cc14", size = 336098, upload-time = "2026-07-24T01:37:32.152Z" }, - { url = "https://files.pythonhosted.org/packages/4a/0e/053d6158e0824d10b03b17a06d5894274e94fae897de807211aa7704e068/im_calculation-2026.7.2-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:67b7f592565fa387cafd5bf96467baf9c681530f70a94806e11c45f7ba52788c", size = 325159, upload-time = "2026-07-24T01:37:33.302Z" }, - { url = "https://files.pythonhosted.org/packages/aa/e4/15c044238dfb35ad449b9928b04fb5e769afe8cc15b31146149dca7d8f2f/im_calculation-2026.7.2-cp314-cp314-manylinux_2_28_x86_64.whl", hash = "sha256:8be67bbc8a238fcceccf9bb1994f92eacd1b77c31e81c43f3e4ae0364fd1d868", size = 382310, upload-time = "2026-07-24T01:37:34.574Z" }, - { url = "https://files.pythonhosted.org/packages/35/f2/65a50321e5f1f1ee4dce5d099315d7deb7cd31fd9ccce3002f36cf35b62f/im_calculation-2026.7.2-cp314-cp314-win_amd64.whl", hash = "sha256:84821d7715992d9f332d309ac764967bcc22f48cc878649807fc0e80c68ad06c", size = 223493, upload-time = "2026-07-24T01:37:36.243Z" }, -] [[package]] name = "imagesize" @@ -1345,6 +1356,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/d8/c6/32d68bfbf1d0c36888530ef6fd72864861af23dc546302b41033471a8c3a/llvmlite-0.49.0-cp314-cp314t-win_amd64.whl", hash = "sha256:be637e465010bc9c50f070468f7f1cf5385e92fee364d192dd5e6cea790ecba9", size = 42986602, upload-time = "2026-08-11T16:25:57.69Z" }, ] +[[package]] +name = "locket" +version = "1.0.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/2f/83/97b29fe05cb6ae28d2dbd30b81e2e402a3eed5f460c26e9eaa5895ceacf5/locket-1.0.0.tar.gz", hash = "sha256:5c0d4c052a8bbbf750e056a8e65ccd309086f4f0f18a2eac306a8dfa4112a632", size = 4350, upload-time = "2022-04-20T22:04:44.312Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/db/bc/83e112abc66cd466c6b83f99118035867cecd41802f8d044638aa78a106e/locket-1.0.0-py2.py3-none-any.whl", hash = "sha256:b6c819a722f7b6bd955b80781788e4a66a55628b858d347536b7e81325a3a5e3", size = 4398, upload-time = "2022-04-20T22:04:42.23Z" }, +] + [[package]] name = "lxml" version = "6.1.2" @@ -2104,6 +2124,19 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/6f/c5/7c16e99869e1f422629092cfd23e3b58e461988c3f9c36fd3624bb4142e6/parse-1.22.1-py2.py3-none-any.whl", hash = "sha256:20f0925a46f06602485ac90d751764d0697fd8455aaa97489ba8953a4b66de32", size = 20925, upload-time = "2026-05-26T03:44:51.156Z" }, ] +[[package]] +name = "partd" +version = "1.4.2" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "locket" }, + { name = "toolz" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/b2/3a/3f06f34820a31257ddcabdfafc2672c5816be79c7e353b02c1f318daa7d4/partd-1.4.2.tar.gz", hash = "sha256:d022c33afbdc8405c226621b015e8067888173d85f7f5ecebb3cafed9a20f02c", size = 21029, upload-time = "2024-05-06T19:51:41.945Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/71/e7/40fb618334dcdf7c5a316c0e7343c5cd82d3d866edc100d98e29bc945ecd/partd-1.4.2-py3-none-any.whl", hash = "sha256:978e4ac767ec4ba5b86c6eaa52e5a2a3bc748a2ca839e8cc798f1cc6ce6efb0f", size = 18905, upload-time = "2024-05-06T19:51:39.271Z" }, +] + [[package]] name = "pillow" version = "12.3.0" @@ -3287,6 +3320,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/7b/61/cceae43728b7de99d9b847560c262873a1f6c98202171fd5ed62640b494b/tomli-2.4.1-py3-none-any.whl", hash = "sha256:0d85819802132122da43cb86656f8d1f8c6587d54ae7dcaf30e90533028b49fe", size = 14583, upload-time = "2026-03-25T20:22:03.012Z" }, ] +[[package]] +name = "toolz" +version = "1.1.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/11/d6/114b492226588d6ff54579d95847662fc69196bdeec318eb45393b24c192/toolz-1.1.0.tar.gz", hash = "sha256:27a5c770d068c110d9ed9323f24f1543e83b2f300a687b7891c1a6d56b697b5b", size = 52613, upload-time = "2025-10-17T04:03:21.661Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/fb/12/5911ae3eeec47800503a238d971e51722ccea5feb8569b735184d5fcdbc0/toolz-1.1.0-py3-none-any.whl", hash = "sha256:15ccc861ac51c53696de0a5d6d4607f99c210739caf987b5d2054f3efed429d8", size = 58093, upload-time = "2025-10-17T04:03:20.435Z" }, +] + [[package]] name = "tqdm" version = "4.70.0" @@ -3470,6 +3512,7 @@ wheels = [ name = "workflow" source = { editable = "." } dependencies = [ + { name = "dask" }, { name = "geopandas" }, { name = "h5py" }, { name = "im-calculation" }, @@ -3517,11 +3560,12 @@ types = [ [package.metadata] requires-dist = [ { name = "coverage", extras = ["toml"], marker = "extra == 'test'" }, + { name = "dask", specifier = ">=2026.6.0" }, { name = "deptry", marker = "extra == 'dev'" }, { name = "geopandas" }, { name = "h5py", specifier = ">=3.15.1" }, { name = "hypothesis", extras = ["numpy"], marker = "extra == 'test'", specifier = ">=6.0.0" }, - { name = "im-calculation", specifier = ">=2025.12.5" }, + { name = "im-calculation", git = "https://github.com/ucgmsim/im_calculation?branch=no_parallel" }, { name = "nshmdb", specifier = ">=2026.9.1" }, { name = "numpy" }, { name = "numpydoc", marker = "extra == 'dev'" }, diff --git a/workflow/realisations.py b/workflow/realisations.py index b4d8b432..364b64f2 100644 --- a/workflow/realisations.py +++ b/workflow/realisations.py @@ -672,6 +672,23 @@ class Rakes(RealisationConfiguration): rakes: dict[str, float] """A map from faults to their rake angles.""" + def as_vectors(self) -> dict[str, npt.NDArray[np.float64]]: + """Represent each rake angle as a unit vector. + + Rakes are angles, so they cannot be averaged directly (the mean of + -179 and 179 degrees is 0, not 180). Averaging the unit vectors and + recovering the angle with `arctan2` avoids this. + + Returns + ------- + dict + A map from faults to the unit vector of their rake angle. + """ + return { + k: np.array([np.cos(np.radians(r)), np.sin(np.radians(r))]) + for k, r in self.rakes.items() + } + def __getitem__(self, key: str) -> float: """Get the rake for a fault name. @@ -713,6 +730,38 @@ def __getitem__(self, key: str) -> BoldM: """ return self.magnitudes[key] + @property + def moments(self) -> dict[str, float]: # numpydoc ignore=RT01 + """dict: a map from faults to their moment.""" + return { + k: moment.magnitude_to_moment(mag, bold_m=True) + for k, mag in self.magnitudes.items() + } + + def moment_averaged(self, values: dict[str, Any]) -> Any: + """Average per-fault quantities, weighted by fault moment. + + Parameters + ---------- + values : dict + A map from faults to the quantity to average. Every fault in + this realisation must be present. Values may be scalars or + arrays, provided they all share the same shape. + + Returns + ------- + Any + The moment-weighted average of `values`, with the same shape as + the individual values. + """ + keys = list(self.magnitudes) + moments = self.moments + return np.average( + [values[key] for key in keys], + weights=[moments[key] for key in keys], + axis=0, + ) + @property def total_moment(self) -> float: # numpydoc ignore=RT01 """float: total moment of realisation""" diff --git a/workflow/scripts/im_calc.py b/workflow/scripts/im_calc.py index 6d16ccfc..54bb3904 100644 --- a/workflow/scripts/im_calc.py +++ b/workflow/scripts/im_calc.py @@ -27,38 +27,192 @@ See the output of `im-calc --help`. """ +import dataclasses import functools from pathlib import Path from typing import Annotated import numpy as np -import pandas as pd import shapely -import tqdm import typer import xarray as xr -from IM import im_reader, ims -from IM.im_calculation import IM +from IM import ims +from IM.ims import IM from qcore import cli, coordinates from source_modelling import sources from source_modelling.sources import IsSource -from workflow import realisations, utils +from workflow import realisations from workflow.realisations import ( DomainParameters, IntensityMeasureCalculationParameters, Magnitudes, + Rakes, RealisationMetadata, - Resolution, RupturePropagationConfig, SourceConfig, ) -PSA_STEP = 10000 - app = typer.Typer() +COORDINATE_METADATA = { + "station": {"description": "Station identifiers"}, + "vs30": { + "description": "Average shear-wave velocity to 30m depth", + "units": "m/s", + }, + "z1pt0": { + "description": "Depth to the 1.0 km/s shear-wave velocity horizon", + "units": "km", + }, + "z2pt5": { + "description": "Depth to the 2.5 km/s shear-wave velocity horizon", + "units": "km", + }, + "epi": {"description": "Epicentral distance", "units": "km"}, + "hyp": {"description": "Hypocentral distance", "units": "km"}, + "rrup": {"description": "Rupture distance", "units": "km"}, + "rjb": {"description": "Joyner-Boore distance", "units": "km"}, + "rx": {"description": "Generalised strike-parallel distance", "units": "km"}, + "ry": {"description": "Generalised strike-normal distance", "units": "km"}, + "supergrid_depth": { + "description": ( + "Penetration of the station into the solver's absorbing layer. " + "0 = interior (valid ground motion); > 0 = inside the damped, " + "coordinate-stretched sponge, where the trace is NOT a " + "ground-motion prediction and should be excluded (recommended " + "threshold: supergrid_depth > 0); NaN = not reported by this " + "solver, i.e. unknown rather than clean." + ), + "units": "m", + }, + "supergrid_depth_gp": { + "description": ( + "Penetration of the station into the solver's absorbing layer, " + "in grid points. Same three states as supergrid_depth (0 = " + "interior; > 0 = inside the sponge, the trace is NOT a " + "ground-motion prediction; NaN = not reported by this solver) " + "and the same recommended threshold of > 0." + ), + "units": "gridpoints", + }, + "latitude": {"description": "Station latitude", "units": "degrees"}, + "longitude": {"description": "Station longitude", "units": "degrees"}, + "frequency": {"description": "Frequency of motion", "units": "Hz"}, + "period": {"description": "Period of motion", "units": "s"}, +} + +IM_METADATA = { + IM.PGA: "Peak ground acceleration", + IM.PGV: "Peak ground velocity", + IM.PGD: "Peak ground displacement", + IM.CAV: "Cumulative absolute velocity", + IM.CAV5: "Cumulative absolute velocity (above 5 cm/s)", + IM.AI: "Arias intensity", + IM.Ds575: "Significant duration (5-75%)", + IM.Ds595: "Significant duration (5-95%)", + IM.pSA: "Pseudo-spectral acceleration", + IM.FAS: "Fourier amplitude spectrum", +} + + +# The 'g0' unit is used for acceleration and is equivalent to 9.81 m/s^2. The +# reason for this is that 'g' is reserved for 'grams'. This is a decision +# made by the `pint` library, which is used to handle the units. +IM_UNITS = { + IM.PGA: "g0", + IM.PGV: "cm/s", + IM.PGD: "cm", + IM.CAV: "m/s", + IM.CAV5: "m/s", + IM.AI: "m/s", + IM.Ds575: "s", + IM.Ds595: "s", + IM.FAS: "g0 * s", + IM.pSA: "g0", +} + + +def add_station_parameters( + dtree: xr.DataTree, station_parameters: dict[str, xr.DataArray] +) -> xr.DataTree: + """Attach per-station parameters as coordinates on every leaf of the tree. + + Parameters + ---------- + dtree : xr.DataTree + The tree of intensity measure datasets. + station_parameters : dict + A map from parameter name (distance and site measures) to the + per-station values. + + Returns + ------- + xr.DataTree + The tree, with the parameters attached to every dataset containing + data. + """ + + def parameterise(dataset: xr.Dataset) -> xr.Dataset: # numpydoc ignore=GL08 + if not dataset.data_vars: + return dataset + + dataset = dataset.copy(deep=False) + dataset.coords.update(station_parameters) + return dataset + + dtree = dtree.map_over_datasets(parameterise) + + return dtree + + +def add_units(dtree: xr.DataTree) -> xr.DataTree: + """Annotate coordinates and intensity measures with units and descriptions. + + Empirical datasets are left alone, because they are annotated as they + are calculated (their values are in log-space, so they do not share the + units of the simulated intensity measures). + + Parameters + ---------- + dtree : xr.DataTree + The tree of intensity measure datasets. + + Returns + ------- + xr.DataTree + The tree, with unit and description metadata attached. + """ + + def unitify(dataset: xr.Dataset) -> xr.Dataset: # numpydoc ignore=GL08 + if not dataset.data_vars: + return dataset + + dataset = dataset.copy(deep=False) + + for name, description in COORDINATE_METADATA.items(): + if name not in dataset.coords: + continue + dataset.coords[name].attrs.update(description) + + if "name" not in dataset.attrs: + return dataset + + name = dataset.attrs["name"] + + for data_var in dataset.data_vars.values(): + data_var.attrs["units"] = IM_UNITS[name] + + description = IM_METADATA[name] + dataset.attrs["description"] = description + return dataset + + dtree = dtree.map_over_datasets(unitify) + + return dtree + + def _source_polygon(source_geometries: dict[str, IsSource]) -> shapely.Geometry: """Extract source polygon in longitude, latitude format. @@ -111,6 +265,221 @@ def _trace_polygon(source_geometries: dict[str, IsSource]) -> shapely.Geometry: return shapely.normalize(shapely.union_all(geometries)) +@dataclasses.dataclass +class Distances: + """Source-to-site distance measures, in kilometres.""" + + rrup: xr.DataArray + """Shortest distance to the rupture plane.""" + rjb: xr.DataArray + """Shortest distance to the surface projection of the rupture.""" + hyp: xr.DataArray + """Distance to the hypocentre.""" + epi: xr.DataArray + """Distance to the epicentre.""" + rx: xr.DataArray | None = None + """Strike-parallel distance. Only defined for planar sources.""" + ry: xr.DataArray | None = None + """Strike-normal distance. Only defined for planar sources.""" + + def as_dict(self) -> dict[str, xr.DataArray]: + """Map distance measure name to distances, omitting undefined measures. + + Returns + ------- + dict + A map from distance measure name to per-station distances. + """ + return { + field.name: value + for field in dataclasses.fields(self) + if (value := getattr(self, field.name)) is not None + } + + +def calculate_distances( + source_geometries: SourceConfig, + hypocentre: np.ndarray, + broadband: xr.Dataset, +) -> Distances: + """Calculate source-to-site distances for every station in the broadband. + + Parameters + ---------- + source_geometries : SourceConfig + The source geometries of the realisation. + hypocentre : np.ndarray + The hypocentre, in latitude, longitude, depth format. + broadband : xr.Dataset + The broadband waveform dataset, supplying the station locations. + + Returns + ------- + Distances + The distance measures for each station, in kilometres. `rx` and `ry` + are only calculated if every source in the realisation is planar. + """ + latitude = broadband.latitude.values + longitude = broadband.longitude.values + station_locations = np.stack((latitude, longitude), axis=-1) + + rrup = xr.DataArray( + np.array( + [ + min( + source.rrup_distance(np.append(station, 0)) + for source in source_geometries.source_geometries.values() + ) + for station in station_locations + ] + ) + / 1000, + dims=["station"], + coords=dict(station=broadband.station), + ) + rjb = xr.DataArray( + np.array( + [ + min( + source.rjb_distance(np.append(station, 0)) + for source in source_geometries.source_geometries.values() + ) + for station in station_locations + ] + ) + / 1000, + dims=["station"], + coords=dict(station=broadband.station), + ) + + hyp = xr.DataArray( + coordinates.distance_between_wgs_depth_coordinates( + np.c_[station_locations, np.zeros_like(latitude)], + hypocentre, + ) + / 1000, + dims=["station"], + coords=dict(station=broadband.station), + ) + epi = xr.DataArray( + coordinates.distance_between_wgs_depth_coordinates( + station_locations, + hypocentre[:2], + ) + / 1000, + dims=["station"], + coords=dict(station=broadband.station), + ) + + distances = Distances(rrup=rrup, rjb=rjb, hyp=hyp, epi=epi) + all_faults_have_rx_ry = all( + isinstance(source, sources.Plane | sources.Fault) + for source in source_geometries.source_geometries.values() + ) + if all_faults_have_rx_ry: + rx, ry = sources.multi_fault_rx_ry_distance( + list(source_geometries.source_geometries.values()), # ty: ignore[invalid-argument-type] + station_locations, + ) + rx /= 1000.0 + ry /= 1000.0 + distances.rx = xr.DataArray( + rx, dims="station", coords=dict(station=broadband.station) + ) + distances.ry = xr.DataArray( + ry, dims="station", coords=dict(station=broadband.station) + ) + return distances + + +@dataclasses.dataclass +class SourceParameters: + """Rupture parameters describing the realisation as a single source.""" + + mag: float + """The total moment magnitude of the rupture.""" + avg_rake: float + """The moment-averaged rake angle (degrees).""" + avg_dip: float + """The moment-averaged dip angle (degrees).""" + avg_ztor: float + """The moment-averaged depth to the top of the rupture (km).""" + avg_zbot: float + """The moment-averaged depth to the bottom of the rupture (km).""" + hypo_depth: float + """The depth of the hypocentre (km).""" + + +def calculate_source_parameters( + source_config: SourceConfig, + magnitudes: Magnitudes, + rakes: Rakes, + hypocentre: np.ndarray, +) -> SourceParameters: + """Reduce a multi-fault realisation to a single set of rupture parameters. + + Ground motion models describe a rupture with a single magnitude, rake, + dip and depth. Multi-fault realisations are collapsed into these by + averaging each fault's contribution, weighted by its moment. + + Parameters + ---------- + source_config : SourceConfig + The source geometries of the realisation. + magnitudes : Magnitudes + The per-fault magnitudes, used for the moment weighting. + rakes : Rakes + The per-fault rake angles. + hypocentre : np.ndarray + The hypocentre, in latitude, longitude, depth (metres) format. + + Returns + ------- + SourceParameters + The rupture parameters of the realisation as a whole. + """ + mag = magnitudes.total_magnitude + + avg_rake_vector = magnitudes.moment_averaged(rakes.as_vectors()) + avg_rake = np.degrees(np.arctan2(avg_rake_vector[1], avg_rake_vector[0])) + + avg_dip_vector = magnitudes.moment_averaged( + { + k: np.array([np.cos(np.radians(f.dip)), np.sin(np.radians(f.dip))]) + for k, f in source_config.source_geometries.items() + } + ) + avg_dip = np.degrees(np.arctan2(avg_dip_vector[1], avg_dip_vector[0])) + + if all( + hasattr(f, "top_m") and hasattr(f, "bottom_m") + for f in source_config.source_geometries.values() + ): + avg_ztor = magnitudes.moment_averaged( + {k: f.top_m / 1000.0 for k, f in source_config.source_geometries.items()} # ty: ignore[unresolved-attribute] + ) + avg_zbot = magnitudes.moment_averaged( + {k: f.bottom_m / 1000.0 for k, f in source_config.source_geometries.items()} # ty: ignore[unresolved-attribute] + ) + else: + avg_ztor = magnitudes.moment_averaged( + { + k: f.centroid[-1] / 1000.0 + for k, f in source_config.source_geometries.items() + } + ) + avg_zbot = avg_ztor + + return SourceParameters( + mag=mag, + avg_rake=float(avg_rake), + avg_dip=float(avg_dip), + avg_ztor=avg_ztor, + avg_zbot=avg_zbot, + hypo_depth=float(hypocentre[2]) / 1000.0, + ) + + @cli.from_docstring(app) def calculate_intensity_measures( realisation_ffp: Annotated[ @@ -121,12 +490,10 @@ def calculate_intensity_measures( ], output_path: Annotated[Path, typer.Argument(dir_okay=False, writable=True)], simulated_stations: Annotated[bool, typer.Option()] = True, - psa_step: Annotated[int, typer.Option()] = PSA_STEP, ko_directory: Annotated[ Path | None, typer.Option(exists=True, file_okay=False) ] = None, override_ims: Annotated[list[IM] | None, typer.Option("-i", "--im")] = None, - cores: Annotated[int | None, typer.Option(min=1)] = None, ) -> None: """Calculate intensity measures for simulation data. @@ -140,23 +507,12 @@ def calculate_intensity_measures( Output directory for IM calc summary statistics. simulated_stations : bool, default True If passed, calculate for simulated stations. - psa_step : int - Maximum number of stations to read from disk at once for pSA calculation ko_directory : Path Directory containing the KO matrix files for FAS calculation. Not required for other IMs. override_ims : list of str Intensity measures to calculate. If not set, reads from the realisation file. - cores : int or None - Set the number of cores for parallel processing of IMs. If set - to `None`, will default to the available cores from - `utils.get_available_cores`. """ - cores = cores or utils.get_available_cores() - metadata = RealisationMetadata.read_from_realisation(realisation_ffp) - resolution = Resolution.read_from_realisation_or_defaults( - realisation_ffp, metadata.defaults_version - ) intensity_measure_parameters = ( IntensityMeasureCalculationParameters.read_from_realisation_or_defaults( realisation_ffp, metadata.defaults_version @@ -167,39 +523,45 @@ def calculate_intensity_measures( rup_prop_config = RupturePropagationConfig.read_from_realisation(realisation_ffp) magnitudes = Magnitudes.read_from_realisation(realisation_ffp) - broadband = xr.open_dataset(broadband_simulation_ffp) + # Chunk over stations only: `component` and `time` must each be a single + # chunk because the IM kernels take them as core dimensions, and the + # file's own on-disk chunking may split either. Dask sizes the station + # chunks from its `array.chunk-size` config. + broadband = xr.open_dataset(broadband_simulation_ffp).chunk( + {"component": -1, "time": -1, "station": "auto"} + ) + # SW4 low-frequency files name these `lat`/`lon`; EMOD3D's name them + # `latitude`/`longitude`. Normalise once, here. + if "latitude" not in broadband and "lat" in broadband: + broadband = broadband.rename({"lat": "latitude", "lon": "longitude"}) + + dt = broadband.attrs["dt"] + if broadband.attrs["units"] == "cm/s^2": + broadband["waveform"] = broadband["waveform"] / 981.0 if not simulated_stations: - broadband = broadband.where( - broadband.station.str.match(r"^(\w{4})$"), drop=True + broadband = broadband.isel( + station=broadband.station.str.match(r"^(\w{4})$").values ) intensity_measures = override_ims or intensity_measure_parameters.ims - nyquist_frequency = 1 / (2 * resolution.dt) + psa_periods = np.array(intensity_measure_parameters.valid_periods, dtype=np.float64) + + nyquist_frequency = 1 / (2 * dt) im_function_map = { - IM.PGA: functools.partial(ims.peak_ground_acceleration, cores=cores), - IM.PGV: functools.partial( - ims.peak_ground_velocity, dt=resolution.dt, cores=cores - ), - IM.PGD: functools.partial( - ims.peak_ground_displacement, dt=resolution.dt, cores=cores - ), - IM.CAV: functools.partial( - ims.cumulative_absolute_velocity, dt=resolution.dt, cores=cores - ), - IM.AI: functools.partial(ims.arias_intensity, dt=resolution.dt, cores=cores), - IM.Ds575: functools.partial(ims.ds575, dt=resolution.dt, cores=cores), - IM.Ds595: functools.partial(ims.ds595, dt=resolution.dt, cores=cores), + IM.PGA: (ims.peak_ground_acceleration), + IM.PGV: functools.partial(ims.peak_ground_velocity, dt=dt), + IM.PGD: functools.partial(ims.peak_ground_displacement, dt=dt), + IM.CAV: functools.partial(ims.cumulative_absolute_velocity, dt=dt), + IM.AI: functools.partial(ims.arias_intensity, dt=dt), + IM.Ds575: functools.partial(ims.ds575, dt=dt), + IM.Ds595: functools.partial(ims.ds595, dt=dt), IM.pSA: functools.partial( ims.pseudo_spectral_acceleration, - periods=np.array( - intensity_measure_parameters.valid_periods, dtype=np.float64 - ), - dt=np.float64(resolution.dt), - step=psa_step, - cores=cores, + periods=psa_periods, + dt=dt, ), } @@ -212,119 +574,64 @@ def calculate_intensity_measures( ) im_function_map[IM.FAS] = functools.partial( ims.fourier_amplitude_spectra, - dt=resolution.dt, + dt=dt, freqs=intensity_measure_parameters.fas_frequencies[ intensity_measure_parameters.fas_frequencies <= nyquist_frequency ], ko_directory=ko_directory, - cores=cores, ) - latitude = broadband.latitude.values - longitude = broadband.longitude.values - station_locations = np.stack((latitude, longitude), axis=-1) - - rrup = ( - np.array( - [ - min( - source.rrup_distance(np.append(station, 0)) - for source in source_geometries.source_geometries.values() - ) - for station in station_locations - ] - ) - / 1000 - ) - rjb = ( - np.array( - [ - min( - source.rjb_distance(np.append(station, 0)) - for source in source_geometries.source_geometries.values() - ) - for station in station_locations - ] - ) - / 1000 - ) hypocentre = source_geometries.source_geometries[ rup_prop_config.initial_fault ].fault_coordinates_to_wgs_depth_coordinates(rup_prop_config.hypocentre) - - hyp = ( - coordinates.distance_between_wgs_depth_coordinates( - np.c_[station_locations, np.zeros_like(latitude)], - hypocentre, - ) - / 1000 - ) - epi = ( - coordinates.distance_between_wgs_depth_coordinates( - station_locations, - hypocentre[:2], - ) - / 1000 - ) - stations = broadband.station.values - dataset = xr.Dataset( - coords={ - "station": ("station", stations), - "component": ( - "component", - ["000", "090", "ver", "geom", "rotd0", "rotd50", "rotd100", "eas"], - ), - "rrup": ("station", rrup), - "rjb": ("station", rjb), - "hyp": ("station", hyp), - "epi": ("station", epi), - }, - attrs={ - "hypo_lat": hypocentre[0], - "hypo_lon": hypocentre[1], - "source": shapely.to_wkt( - _source_polygon(source_geometries.source_geometries) - ), - "trace": shapely.to_wkt( - _trace_polygon(source_geometries.source_geometries) - ), - "domain": shapely.to_wkt( - shapely.transform( - domain_parameters.domain.polygon, - lambda c: coordinates.nztm_to_wgs_depth(c)[:, ::-1], - ) - ), - "magnitude": magnitudes.total_magnitude, - "event": metadata.name, - }, + distances = calculate_distances(source_geometries, hypocentre, broadband) + rakes = Rakes.read_from_realisation(realisation_ffp) + source_parameters = calculate_source_parameters( + source_geometries, magnitudes, rakes, hypocentre ) - all_faults_have_rx_ry = all( - isinstance(source, sources.Plane | sources.Fault) - for source in source_geometries.source_geometries.values() - ) - if all_faults_have_rx_ry: - rx, ry = sources.multi_fault_rx_ry_distance( - list(source_geometries.source_geometries.values()), # ty: ignore[invalid-argument-type] - station_locations, - ) - dataset["rx"] = xr.DataArray(rx, dims="station", coords={"station": stations}) - dataset["ry"] = xr.DataArray(ry, dims="station", coords={"station": stations}) + # Each IM function is dask-native: it accepts the lazy `waveform` DataArray + # and returns a lazy Dataset with the same `station` chunking, one data + # variable per component. Nothing is computed until `dtree.to_netcdf` + # below, which streams the result chunk by chunk. + im_results: dict[str, xr.Dataset] = { + im_name: im_function_map[im_name](broadband.waveform) + for im_name in intensity_measures + } - waveform = broadband.waveform.values.astype(np.float64) + attributes = { + "hypo_lat": hypocentre[0], + "hypo_lon": hypocentre[1], + "source": shapely.to_wkt(_source_polygon(source_geometries.source_geometries)), + "trace": shapely.to_wkt(_trace_polygon(source_geometries.source_geometries)), + "domain": shapely.to_wkt( + shapely.transform( + domain_parameters.domain.polygon, + lambda c: coordinates.nztm_to_wgs_depth(c)[:, ::-1], + ) + ), + "magnitude": magnitudes.total_magnitude, + "event": metadata.name, + "rake": source_parameters.avg_rake, + "dip": source_parameters.avg_dip, + "ztor": source_parameters.avg_ztor, + "zbot": source_parameters.avg_zbot, + "hypo_depth": source_parameters.hypo_depth, + } - for im_name in (pbar := tqdm.tqdm(intensity_measures)): - pbar.set_description(im_name) - im_fn = im_function_map[im_name] + dtree = xr.DataTree.from_dict(im_results, nested=True) - result = im_fn(waveform) + dtree.attrs = attributes - if isinstance(result, pd.DataFrame): - result["station"] = broadband.station.values - result = result.set_index("station").to_xarray().to_array(dim="component") - elif isinstance(result, xr.DataArray): - result = result.assign_coords(station=broadband.station) - dataset[im_name] = result - im_reader.write_intensity_measures(dataset, output_path) + dtree = add_station_parameters( + dtree, + 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"]}, + ) + dtree = add_units(dtree) + dtree.to_netcdf(output_path) realisations.append_log_entry(realisation_ffp)