From fcf9906f9f797d322fb440e9a6db51351535045b Mon Sep 17 00:00:00 2001 From: Jake Faulkner Date: Mon, 31 Aug 2026 13:55:33 +1200 Subject: [PATCH] Port srf2stoch to Python, conserving moment Replaces the `srf2stoch` binary with an in-process conversion, and drops it from the container build. The conversion is a box average expressed as a sparse fractional-overlap matrix: row j of the kernel holds the overlap weights between coarse bin j and the fine cells it spans. The special case where the two grids span exactly the same length is what srf2stoch.c implements; the sparse form generalises it without materialising the empty overlaps, and handles the padded case where they do not. That padded case is why this was worth porting. The HF code demands a single dx/dy across every SRF segment, and `nx = ceil(len / dx)` means the coarse grid is generally *longer* than the plane it covers. The overhang is split evenly between the two ends, so the grids stay concentric -- the stoch format records only a centre point and an `nx * dx` extent, so off-centre padding would sit the slip distribution in the wrong place on the plane HF reconstructs. Edge bins then carry weights summing to the covered fraction rather than to 1, which is what conserves total slip*area rather than cell value. Quantities that are not spread over a cell are handled separately. Rupture time is a time, so it is divided by the covered fraction; every cell is partially covered because nx and ny round up, so this never divides by zero. Rise time is slip-averaged, and there the coverage factor cancels, so it must *not* be divided out again. Rake is averaged as a circular mean weighted by slip -- the arithmetic mean of -179 and 179 degrees is 0, which is the opposite of the right answer. A plane with no slip anywhere has no slip-weighted mean, so it falls back to an unweighted average. This changes every stoch file, and therefore every high-frequency result. The previous binary's outputs were not moment-conserving for segments whose length was not an exact multiple of the stoch dx. Co-Authored-By: Claude Opus 5 --- container/runner.def | 2 +- tests/test_generate_stoch.py | 408 +++++++++++++++++++++++++++++ workflow/scripts/generate_stoch.py | 272 ++++++++++++++++--- 3 files changed, 640 insertions(+), 42 deletions(-) create mode 100644 tests/test_generate_stoch.py diff --git a/container/runner.def b/container/runner.def index 52ef4706..cf6484b6 100644 --- a/container/runner.def +++ b/container/runner.def @@ -8,7 +8,7 @@ From: earthquakesuc/workflow-bootstrap:latest cd /EMOD3D && \ mkdir build && cd build && \ cmake -DHF_TIME_WINDOW=off -DUSE_FFTW=on -DCMAKE_C_FLAGS="-march=x86-64 -mtune=generic -std=gnu89" .. && \ - make -j $(nproc) genslip_v5.6.2 srf2stoch generic_slip2srf \ + make -j $(nproc) genslip_v5.6.2 generic_slip2srf \ hb_high_binmod_v6.0.3 cd / diff --git a/tests/test_generate_stoch.py b/tests/test_generate_stoch.py new file mode 100644 index 00000000..1087f7af --- /dev/null +++ b/tests/test_generate_stoch.py @@ -0,0 +1,408 @@ +import os +from pathlib import Path + +import numpy as np +import pandas as pd +import pytest +import scipy.sparse as sp +from typer.testing import CliRunner + +from source_modelling import sources, srf +from source_modelling.srf import SrfFile +from source_modelling.stoch import StochFile +from workflow import defaults, realisations +from workflow.scripts.generate_stoch import ( + _box_average_matrix, + app, + circular_mean, + convert_srf_to_stoch, +) + +# (nstk, ndip, len, wid) for each plane of the synthetic SRF. The first +# plane divides evenly into the 2km stoch grid used below (dstk = ddip = +# 0.5), the second does not (dstk = ddip = 0.3). +PLANE_SHAPES = [(13, 7, 6.5, 3.5), (9, 5, 2.7, 1.5)] + +DT = 0.1 + +# A real (multi-segment) SRF, if a source_modelling checkout is available. +REAL_SRF_FFP = ( + Path( + os.environ.get( + "SOURCE_MODELLING_PATH", Path.home() / "src" / "source_modelling" + ) + ) + / "tests" + / "srfs" + / "3366146.srf" +) + + +def make_srf(seed: int = 1) -> SrfFile: + """Build a small synthetic (version 1.0) SRF with random slip.""" + rng = np.random.default_rng(seed) + header = pd.DataFrame( + [ + { + "elon": 172.0 + i, + "elat": -43.5 - i, + "nstk": nstk, + "ndip": ndip, + "len": length, + "wid": width, + "stk": 45.0 + 90 * i, + "dip": 60.0, + "dtop": 1.0, + "shyp": 0.5, + "dhyp": 1.0, + } + for i, (nstk, ndip, length, width) in enumerate(PLANE_SHAPES) + ] + ) + n_points = int((header["nstk"] * header["ndip"]).sum()) + # The rise time of each point is a whole number of timesteps so that + # the SRF round-trips through disk exactly (on reading, rise = nt * dt). + nt = rng.integers(1, 6, n_points) + points = pd.DataFrame( + { + "lon": rng.uniform(171, 173, n_points), + "lat": rng.uniform(-44, -43, n_points), + "dep": rng.uniform(1, 10, n_points), + "stk": np.repeat( + header["stk"].to_numpy(), (header["nstk"] * header["ndip"]).to_numpy() + ), + "dip": 60.0, + "area": 1e6, + "tinit": rng.uniform(0, 10, n_points), + "dt": DT, + "rake": rng.uniform(0, 360, n_points), + "slip": rng.uniform(0, 100, n_points), + "rise": nt * DT, + } + ) + # Slip velocity time series: nt samples per point, integrating to the + # total slip of that point. + indptr = np.concatenate([[0], np.cumsum(nt)]) + indices = np.concatenate([np.arange(n) for n in nt]) + data = np.repeat(points["slip"].to_numpy() / (nt * DT), nt) + slipt1 = sp.csr_array( + (data, indices, indptr), shape=(n_points, int(nt.max())), dtype=np.float32 + ) + return SrfFile("1.0", header, points, slipt1) + + +@pytest.fixture +def synthetic_srf() -> SrfFile: + return make_srf() + + +def covered_fraction(n_coarse: int, coarse_dx: float, extent: float) -> np.ndarray: + """Fraction of each coarse cell that lies over a plane of length `extent`. + + The coarse grid is centred on the plane, so the overhang is split + evenly between the first and last cells. + """ + overhang = (n_coarse * coarse_dx - extent) / 2 + edges = np.arange(n_coarse + 1) * coarse_dx - overhang + return (np.minimum(edges[1:], extent) - np.maximum(edges[:-1], 0)) / coarse_dx + + +def fine_moment(srf_file: SrfFile, i: int) -> float: + """Sum of slip * patch area (in km^2) for plane `i` of an SRF.""" + plane = srf_file.header.iloc[i] + patch_area = (plane["len"] / plane["nstk"]) * (plane["wid"] / plane["ndip"]) + return float(srf_file.segments[i]["slip"].sum() * patch_area) + + +# --- _box_average_matrix ----------------------------------------------------- + + +def test_box_average_matrix_matches_docstring_example() -> None: + """Five fine cells of width 3 pooled into three coarse cells of width 5.""" + matrix = _box_average_matrix(5, 3, 3.0, 5.0).toarray() + assert matrix == pytest.approx( + np.array( + [ + [3 / 5, 2 / 5, 0, 0, 0], + [0, 1 / 5, 3 / 5, 1 / 5, 0], + [0, 0, 0, 2 / 5, 3 / 5], + ] + ) + ) + + +def test_box_average_matrix_is_identity_when_grids_agree() -> None: + matrix = _box_average_matrix(7, 7, 0.5, 0.5).toarray() + assert matrix == pytest.approx(np.eye(7)) + + +@pytest.mark.parametrize( + ("n_fine", "fine_dx", "coarse_dx"), + [(100, 0.1, 2.0), (13, 0.5, 2.0), (9, 0.3, 2.0), (37, 0.2, 1.7), (5, 3.0, 5.0)], +) +def test_box_average_matrix_rows_are_weighted_averages( + n_fine: int, fine_dx: float, coarse_dx: float +) -> None: + """Every coarse bin is an average of the fine cells it covers. + + The weights of a bin sum to one, except for the first and last bins, + which hang off either end of the fine grid by half the overhang each + and sum to the covered fraction of the bin. + """ + n_coarse = int(np.ceil(n_fine * fine_dx / coarse_dx)) + assert n_coarse >= 2, "the covered fraction below assumes two distinct end bins" + matrix = _box_average_matrix(n_fine, n_coarse, fine_dx, coarse_dx).toarray() + assert matrix.shape == (n_coarse, n_fine) + assert (matrix >= 0).all() + + row_sums = matrix.sum(axis=1) + overhang = (n_coarse * coarse_dx - n_fine * fine_dx) / 2 + covered = (coarse_dx - overhang) / coarse_dx + assert row_sums[1:-1] == pytest.approx(np.ones(n_coarse - 2)) + assert row_sums[0] == pytest.approx(covered) + assert row_sums[-1] == pytest.approx(covered) + + +@pytest.mark.parametrize( + ("n_fine", "fine_dx", "coarse_dx"), + [(100, 0.1, 2.0), (13, 0.5, 2.0), (9, 0.3, 2.0), (37, 0.2, 1.7), (5, 3.0, 5.0)], +) +def test_box_average_matrix_is_centred( + n_fine: int, fine_dx: float, coarse_dx: float +) -> None: + """The coarse grid is centred on the fine grid, not aligned to its start. + + The stoch format records a centre point and an ``nx * dx`` extent, so a + coarse grid longer than the plane has to overhang both ends equally. + Otherwise the slip would sit off-centre on the plane the HF code + reconstructs from the header. + """ + n_coarse = int(np.ceil(n_fine * fine_dx / coarse_dx)) + matrix = _box_average_matrix(n_fine, n_coarse, fine_dx, coarse_dx).toarray() + # Reversing both the bins and the cells they cover is the same grid. + assert matrix == pytest.approx(matrix[::-1, ::-1]) + + +@pytest.mark.parametrize( + ("n_fine", "fine_dx", "coarse_dx"), + [(100, 0.1, 2.0), (13, 0.5, 2.0), (9, 0.3, 2.0), (37, 0.2, 1.7), (5, 3.0, 5.0)], +) +def test_box_average_matrix_conserves_mass( + n_fine: int, fine_dx: float, coarse_dx: float +) -> None: + """Averaging then re-integrating over the coarse cells preserves the integral.""" + n_coarse = int(np.ceil(n_fine * fine_dx / coarse_dx)) + matrix = _box_average_matrix(n_fine, n_coarse, fine_dx, coarse_dx) + values = np.random.default_rng(2).uniform(0, 10, n_fine) + coarse = matrix @ values + assert (coarse.sum() * coarse_dx) == pytest.approx(values.sum() * fine_dx) + + +# --- Moment preservation ----------------------------------------------------- + + +@pytest.mark.parametrize(("dx", "dy"), [(2.0, 2.0), (1.0, 1.0), (0.7, 1.3), (0.5, 0.5)]) +def test_convert_srf_to_stoch_preserves_moment( + synthetic_srf: SrfFile, dx: float, dy: float +) -> None: + """Total moment (slip x area) of each plane survives the down-sampling. + + The stoch cells are physically larger than the SRF patches, so the + box average must be weighted by the overlap between the two grids for + the sum of slip x area to be unchanged. + """ + stoch_file = convert_srf_to_stoch(synthetic_srf, dx, dy) + assert len(stoch_file.data) == len(PLANE_SHAPES) + for i, plane in enumerate(stoch_file.data): + coarse_moment = float(plane.slip.sum()) * dx * dy + assert coarse_moment == pytest.approx(fine_moment(synthetic_srf, i), rel=1e-5) + + +@pytest.mark.slow +@pytest.mark.skipif( + not REAL_SRF_FFP.exists(), reason=f"{REAL_SRF_FFP} is not available" +) +def test_convert_srf_to_stoch_preserves_moment_real_srf() -> None: + """Moment is preserved for a real multi-segment rupture.""" + srf_file = srf.read_srf(REAL_SRF_FFP) + stoch_file = convert_srf_to_stoch(srf_file, 2.0, 2.0) + for i, plane in enumerate(stoch_file.data): + coarse_moment = float(plane.slip.sum()) * plane.header.dx * plane.header.dy + assert coarse_moment == pytest.approx(fine_moment(srf_file, i), rel=1e-5) + + +def test_convert_srf_to_stoch_preserves_uniform_slip(synthetic_srf: SrfFile) -> None: + """A uniform slip distribution down-samples to the same uniform slip.""" + dx = dy = 2.0 + synthetic_srf.points["slip"] = 42.0 + stoch_file = convert_srf_to_stoch(synthetic_srf, dx, dy) + for i, plane in enumerate(stoch_file.data): + header = synthetic_srf.header.iloc[i] + # Cells the plane only partially covers are scaled down by the + # covered fraction of the cell, which is what keeps the moment + # (rather than the slip value) constant. + covered_x = covered_fraction(plane.header.nx, dx, header["len"]) + covered_y = covered_fraction(plane.header.ny, dy, header["wid"]) + assert plane.slip == pytest.approx( + 42.0 * np.outer(covered_y, covered_x), rel=1e-5 + ) + # The partial cells are the two ends, not just the far end. + assert plane.slip == pytest.approx(plane.slip[::-1, ::-1], rel=1e-5) + + +def test_convert_srf_to_stoch_grid_covers_the_plane(synthetic_srf: SrfFile) -> None: + """The stoch grid is the smallest dx by dy grid covering the SRF plane.""" + dx, dy = 2.0, 2.0 + stoch_file = convert_srf_to_stoch(synthetic_srf, dx, dy) + for i, plane in enumerate(stoch_file.data): + header = synthetic_srf.header.iloc[i] + assert plane.header.nx == int(np.ceil(header["len"] / dx)) + assert plane.header.ny == int(np.ceil(header["wid"] / dy)) + assert plane.slip.shape == (plane.header.ny, plane.header.nx) + assert plane.rise.shape == plane.slip.shape + assert plane.trup.shape == plane.slip.shape + + +def test_convert_srf_to_stoch_rise_is_slip_weighted(synthetic_srf: SrfFile) -> None: + """Rise time is averaged in proportion to slip, not by area.""" + # One plane, one stoch cell, two patches: all of the slip is on the + # patch with a rise time of 3s, so the cell rise time must be 3s. + srf_file = synthetic_srf + srf_file.header = srf_file.header.iloc[:1].copy() + srf_file.header.loc[0, ["nstk", "ndip", "len", "wid"]] = [2, 1, 1.0, 0.5] + srf_file.points = srf_file.points.iloc[:2].copy() + srf_file.points["slip"] = [0.0, 10.0] + srf_file.points["rise"] = [7.0, 3.0] + + (plane,) = convert_srf_to_stoch(srf_file, 2.0, 2.0).data + assert plane.slip.shape == (1, 1) + assert plane.rise.item() == pytest.approx(3.0) + + +@pytest.mark.parametrize(("dx", "dy"), [(2.0, 2.0), (1.0, 1.0), (0.7, 1.3), (0.5, 0.5)]) +def test_convert_srf_to_stoch_trup_is_not_scaled_by_coverage( + synthetic_srf: SrfFile, dx: float, dy: float +) -> None: + """Rupture time is a time, so partially covered cells must not dilute it. + + Slip is deliberately scaled down in the cells at the edge of the plane + to conserve the moment. Applying the same scaling to the rupture time + would make the rupture arrive early at the edges of every plane. + """ + synthetic_srf.points["tinit"] = 5.0 + stoch_file = convert_srf_to_stoch(synthetic_srf, dx, dy) + for plane in stoch_file.data: + assert plane.trup == pytest.approx(np.full(plane.trup.shape, 5.0), rel=1e-5) + + +def test_convert_srf_to_stoch_trup_matches_a_uniform_average( + synthetic_srf: SrfFile, +) -> None: + """A cell covering the whole plane gets the mean rupture time of the plane.""" + srf_file = synthetic_srf + srf_file.header = srf_file.header.iloc[:1].copy() + srf_file.header.loc[0, ["nstk", "ndip", "len", "wid"]] = [2, 1, 1.0, 0.5] + srf_file.points = srf_file.points.iloc[:2].copy() + srf_file.points["tinit"] = [4.0, 6.0] + + (plane,) = convert_srf_to_stoch(srf_file, 2.0, 2.0).data + assert plane.trup.item() == pytest.approx(5.0) + + +def test_convert_srf_to_stoch_zero_slip_rise(synthetic_srf: SrfFile) -> None: + """Cells with no slip get a nominal (non-zero) rise time.""" + synthetic_srf.points["slip"] = 0.0 + stoch_file = convert_srf_to_stoch(synthetic_srf, 2.0, 2.0) + for plane in stoch_file.data: + assert (plane.slip == 0).all() + assert plane.rise == pytest.approx(np.full(plane.rise.shape, 1e-5)) + + +# --- circular_mean ----------------------------------------------------------- + + +def test_circular_mean_wraps_around_zero() -> None: + mean = circular_mean(np.array([350.0, 10.0]), np.array([1.0, 1.0])) + # 0 and 360 are the same bearing. + assert min(mean, 360 - mean) == pytest.approx(0.0, abs=1e-9) + + +def test_circular_mean_is_weighted() -> None: + # Three quarters of the weight sits at 0 degrees, one quarter at 90. + expected = np.degrees(np.arctan2(0.25, 0.75)) + assert circular_mean(np.array([0.0, 90.0]), np.array([3.0, 1.0])) == ( + pytest.approx(expected) + ) + + +def test_average_rake_is_in_degrees(synthetic_srf: SrfFile) -> None: + """The stoch header rake is a bearing in degrees, not radians.""" + synthetic_srf.points["rake"] = 185.0 + stoch_file = convert_srf_to_stoch(synthetic_srf, 2.0, 2.0) + for plane in stoch_file.data: + assert plane.header.average_rake == pytest.approx(185.0, abs=1e-3) + + +# --- Integration ------------------------------------------------------------- + + +@pytest.fixture +def realisation_ffp(tmp_path: Path, synthetic_srf: SrfFile) -> Path: + """A realisation whose sources match the planes of the synthetic SRF.""" + realisation_ffp = tmp_path / "realisation.json" + realisations.RealisationMetadata( + name="generate stoch test", + version="1", + defaults_version=defaults.DefaultsVersion.v24_2_2_1, + ).write_to_realisation(realisation_ffp) + realisations.SourceConfig( + source_geometries={ + f"plane_{i}": sources.Plane.from_centroid_strike_dip( + np.array([plane["elat"], plane["elon"]]), + plane["dip"], + plane["len"], + plane["wid"], + dtop=plane["dtop"], + strike=plane["stk"], + ) + for i, plane in synthetic_srf.header.iterrows() + } + ).write_to_realisation(realisation_ffp) + return realisation_ffp + + +def test_generate_stoch_smoke( + tmp_path: Path, realisation_ffp: Path, synthetic_srf: SrfFile +) -> None: + """An SRF file on disk converts into a readable stoch file.""" + srf_ffp = tmp_path / "realisation.srf" + stoch_ffp = tmp_path / "realisation.stoch" + srf.write_srf(srf_ffp, synthetic_srf) + + result = CliRunner().invoke( + app, [str(realisation_ffp), str(srf_ffp), str(stoch_ffp)] + ) + assert result.exit_code == 0, result.output + assert stoch_ffp.exists() + + stoch_file = StochFile.from_file(stoch_ffp) + assert len(stoch_file.data) == len(PLANE_SHAPES) + + srf_file = srf.read_srf(srf_ffp) + for i, plane in enumerate(stoch_file.data): + header = srf_file.header.iloc[i] + # Every plane uses the configured stoch dx/dy, as the HF code + # requires. Planes smaller than a cell round up to a single cell + # rather than down-sampling to an empty grid. + assert plane.header.dx == pytest.approx(2.0) + assert plane.header.dy == pytest.approx(2.0) + assert plane.slip.shape == (plane.header.ny, plane.header.nx) + assert plane.header.dtop == pytest.approx(header["dtop"]) + assert plane.header.dip == pytest.approx(header["dip"]) + assert plane.header.strike == pytest.approx(header["stk"] % 360) + assert (plane.slip >= 0).all() + assert (plane.rise > 0).all() + # The written file preserves the moment to the precision of the + # %e formatting used by the stoch format. + coarse_moment = float(plane.slip.sum()) * plane.header.dx * plane.header.dy + assert coarse_moment == pytest.approx(fine_moment(srf_file, i), rel=1e-4) diff --git a/workflow/scripts/generate_stoch.py b/workflow/scripts/generate_stoch.py index a9ebfaaf..3f041b78 100644 --- a/workflow/scripts/generate_stoch.py +++ b/workflow/scripts/generate_stoch.py @@ -18,7 +18,7 @@ Environment ----------- -Can be run in the cybershake container. Can also be run from your own computer using the `generate-stoch` command which is installed after running `pip install workflow@git+https://github.com/ucgmsim/workflow`. If you are executing on your own computer you also need to specify the `srf2stoch` path (`--srf2stoch-path`). +Can be run in the cybershake container. Can also be run from your own computer using the `generate-stoch` command which is installed after running `pip install workflow@git+https://github.com/ucgmsim/workflow`. For More Help ------------- @@ -28,30 +28,245 @@ from pathlib import Path from typing import Annotated +import numpy as np +import scipy.sparse as sp import typer from qcore import cli -from source_modelling import sources, srf +from source_modelling import srf +from source_modelling.srf import SrfFile +from source_modelling.stoch import StochFile, StochHeader, StochPlane from workflow import log_utils, realisations -from workflow.realisations import RealisationMetadata, SourceConfig, StochConfig +from workflow.realisations import RealisationMetadata, StochConfig app = typer.Typer() +def _box_average_matrix( + n_fine: int, n_coarse: int, fine_dx: float, coarse_dx: float +) -> sp.csr_matrix: + """Build an area-pooling kernel for averaging high-resolution data into lower-resolution data. + + Assuming we have `n_fine` fine gridpoints, and `n_coarse` coarse gridpoints, + Row j of the returned matrix gives the fractional-overlap weights (summing + to 1) between coarse bin j and the fine cells it spans. This is equivalent + to the ``adaptive_avg_pool2d`` kernel in pytorch with padding. + + The two grids are centred on each other. If the coarse grid is longer than + the fine grid, the bins at either end hang off the fine grid and their + weights sum to the covered fraction rather than to 1. This is what makes + the kernel conserve the total (slip * area) rather than the cell value. See + ``convert_srf_to_stoch`` for how we handle trise where this is not what we want. + + Parameters + ---------- + n_fine : int + The number of elements in the original, high-resolution grid dimension. + n_coarse : int + The number of elements in the target, downsampled grid dimension. + fine_dx : float + The physical resolution of the fine cells. + coarse_dx : float + The physical resolution of the coarse cells. + + Returns + ------- + scipy.sparse.csr_matrix + A sparse matrix of shape (n_coarse, n_fine) containing the area-weighted + fractional overlap coefficients. + + Notes + ----- + The weights correspond exactly to the fractional overlap of coarse bins over + fine bins. This is mathematically equivalent to upsampling both grids to + their Least Common Multiple (LCM) base units, padding the geometry with + zeros in these units and computing a standard block average. The special + case where the fine grid and coarse grid span the same length is implemented + by srf2stoch.c. The main advantage of our approach is that a sparse matrix + does not have to materialise all the empty cell overlaps in memory. A + secondary advantage is that we handle the padded case, which lets us set a + uniform dx/dy for all SRF segments as the HF code demands without changing + total moment. Because the stoch format records only a centre point and an + ``nx * dx`` extent, the padding has to be centred too, or the slip + distribution would sit off-centre on the plane the HF code reconstructs. + + For example, downsampling 5 fine cells to 3 coarse cells implies an LCM of + 15 base units. The 5 fine cells (A-E) take up 3 units each, while the 3 + coarse cells (C0-C2) take up 5 units each. + + The visual alignment of this overlap is as follows: + + :: + + THE LCM GRID (15 Base Units) + |---|---|---|---|---|---|---|---|---|---|---|---|---|---|---| + + FINE INPUT GRID (5 cells, each = 3 base units) + |-----------|-----------|-----------|-----------|-----------| + | A | B | C | D | E | + |-----------|-----------|-----------|-----------|-----------| + + COARSE OUTPUT GRID (3 cells, each = 5 base units) + |-------------------|-------------------|-------------------| + | C0 | C1 | C2 | + |-------------------|-------------------|-------------------| + + Each row in the returned sparse matrix corresponds to the fractional + makeup of a single coarse bin: + + * **Row 0 (Coarse Bin 0):** Spans 5 base units. Covers all 3 units of A + (3/5) and 2 units of B (2/5). + * **Row 1 (Coarse Bin 1):** Spans 5 base units. Covers the remaining 1 + unit of B (1/5), all 3 units of C (3/5), and 1 unit of D (1/5). + * **Row 2 (Coarse Bin 2):** Spans 5 base units. Covers the remaining 2 + units of D (2/5) and all 3 units of E (3/5). + + """ + bin_width = coarse_dx / fine_dx + # The coarse grid may be longer than the fine grid it covers. Split the + # excess evenly between the two ends so that both grids stay centred on + # the same point (see the Notes above). + overhang = (n_coarse * bin_width - n_fine) / 2 + edges = np.arange(n_coarse + 1) * bin_width - overhang + rows, cols, weights = [], [], [] + for j in range(n_coarse): + lo, hi = edges[j], edges[j + 1] + idx = np.arange(max(int(np.floor(lo)), 0), min(int(np.ceil(hi)), n_fine)) + weights.append((np.minimum(idx + 1, hi) - np.maximum(idx, lo)) / bin_width) + rows.append(np.full(len(idx), j)) + cols.append(idx) + return sp.csr_matrix( + (np.concatenate(weights), (np.concatenate(rows), np.concatenate(cols))), + shape=(n_coarse, n_fine), + ) + + +def circular_mean(angles: np.ndarray, weights: np.ndarray) -> float: + """Take the circular mean of `angles` with respect to `weights`. + + Parameters + ---------- + angles : array of floats + The angles to average, in degrees. + weights : array of floats + The weights to apply to the average. + + Returns + ------- + float + The weighted circular mean of angles. + """ + + rad = np.radians(np.ravel(angles)) + x = np.cos(rad) + y = np.sin(rad) + weights = np.ravel(weights) + if not weights.any(): + # A plane with no slip anywhere has no slip-weighted mean, so + # fall back to an unweighted average of the angles. + weights = np.ones_like(weights) + avg_vector = np.average(np.c_[x, y], weights=weights, axis=0) + return np.degrees(np.arctan2(avg_vector[1], avg_vector[0])).item() % 360.0 + + +def convert_srf_to_stoch(srf_file: SrfFile, dx: float, dy: float) -> StochFile: + """Convert an SRF file into a Stoch file by box-averaging slip, tinit and tinit * rise. + + Parameters + ---------- + srf_file : SrfFile + The SRF file to convert. + dx : float + The desired strike-resolution for the output stoch file. + dy : float + The desired dip-resolution for the output stoch file. + + Returns + ------- + StochFile + An output stoch file downsampled from ``srf_file``. + """ + planes = [] + for i, segment in enumerate(srf_file.segments): + header = srf_file.header.iloc[i].astype(np.float32) + nstk, ndip = int(header["nstk"]), int(header["ndip"]) + + slip = segment["slip"].to_numpy(dtype=np.float32).reshape(ndip, nstk) + rake = segment["rake"].to_numpy(dtype=np.float32).reshape( + ndip, nstk + ) % np.float32(360.0) + rise = segment["rise"].to_numpy(dtype=np.float32).reshape(ndip, nstk) + tinit = segment["tinit"].to_numpy(dtype=np.float32).reshape(ndip, nstk) + + dstk = float(header["len"]) / nstk + ddip = float(header["wid"]) / ndip + + nx = int(np.ceil(float(header["len"]) / dx)) + ny = int(np.ceil(float(header["wid"]) / dy)) + wx = _box_average_matrix(nstk, nx, dstk, dx).astype(np.float32) + wy = _box_average_matrix(ndip, ny, ddip, dy).astype(np.float32) + + def box_average( + values: np.ndarray, + wy: sp.csr_matrix = wy, + wx: sp.csr_matrix = wx, + ) -> np.ndarray: + return wy @ values @ wx.T + + slip_grid = box_average(slip) + # Cells at the edge of the plane are only partially covered by the SRF, + # so their weights sum to less than one. That is what conserves the + # moment for slip, but rupture time is a time rather than a quantity + # spread over the cell, so it has to be divided by the covered + # fraction. Every cell is partially covered because nx and ny are + # rounded up, so this never divides by zero. + # + # NOTE: This does materialise an array of order (ny, nx) but it is the + # *coarse* ny, nx. Unless we have ruptures larger than Hikurangi this is + # unlikely to ever be an issue. + coverage = np.outer( + np.asarray(wy.sum(axis=1)).ravel(), np.asarray(wx.sum(axis=1)).ravel() + ) + trup_grid = box_average(tinit) / coverage + + # The rise grid is a slip-averaged rise in each cell. Note that because + # we are dividing rise by slip the coverage factor conveniently cancels + # out and we do not have to (should not) similarly divide for the rise + # sum. + rise_sum = box_average(rise * slip) + rise_grid = np.where( + slip_grid > 0, + rise_sum / np.where(slip_grid > 0, slip_grid, 1), + np.float32(1e-5), + ) + + stoch_header = StochHeader( + longitude=header["elon"], + latitude=header["elat"], + nx=nx, + ny=ny, + dx=dx, + dy=dy, + strike=header["stk"] % np.float32(360.0), + dip=header["dip"], + average_rake=circular_mean(rake, slip), + dtop=header["dtop"], + shypo=header["shyp"], + dhypo=header["dhyp"], + ) + planes.append(StochPlane(stoch_header, slip_grid, rise_grid, trup_grid)) + return StochFile(planes) + + @cli.from_docstring(app) @log_utils.log_call() def generate_stoch( realisation_ffp: Annotated[Path, typer.Argument(exists=True, dir_okay=False)], srf_ffp: Annotated[Path, typer.Argument(exists=True, dir_okay=False)], stoch_ffp: Annotated[Path, typer.Argument(dir_okay=False)], - srf2stoch_path: Annotated[Path, typer.Option(exists=True)] = Path( - "/EMOD3D/tools/srf2stoch" - ), ) -> None: """Generate a stoch file from an SRF file. - This function uses the `srf2stoch` binary to generate a stoch file from the provided SRF file. - Parameters ---------- realisation_ffp : Path @@ -60,43 +275,18 @@ def generate_stoch( Path to the SRF file which is used as input for the stoch file generation. stoch_ffp : Path Path to the output file where the generated stoch file will be saved. - srf2stoch_path : Path, optional - Path to the `srf2stoch` binary used for the conversion. """ metadata = RealisationMetadata.read_from_realisation(realisation_ffp) stoch_config = StochConfig.read_from_realisation_or_defaults( realisation_ffp, metadata.defaults_version ) - source_config = SourceConfig.read_from_realisation(realisation_ffp) - - if all( - isinstance(fault, sources.Point) - for fault in source_config.source_geometries.values() - ): - srf_file = srf.read_srf(srf_ffp) - source = srf_file.header.iloc[0] - srf_nstk = int(source["nstk"]) - srf_len = float(source["len"]) - dx = srf_len / srf_nstk - srf_ndip = int(source["ndip"]) - srf_wid = float(source["wid"]) - dy = srf_wid / srf_ndip - else: - geometries = list(source_config.source_geometries.values()) - min_length = min(fault.length for fault in geometries) - min_width = min(fault.width for fault in geometries) - # If the stoch dx is greater than the length (resp. dy and width), we might get an empty stoch file - dx = min(stoch_config.stoch_dx, min_length / 2) - dy = min(stoch_config.stoch_dy, min_width / 2) - - log_utils.log_check_call( - [ - str(srf2stoch_path), - f"dx={dx}", - f"dy={dy}", - f"infile={srf_ffp}", - f"outfile={stoch_ffp}", - ] - ) + srf_file = srf.read_srf(srf_ffp) + dx = stoch_config.stoch_dx + dy = stoch_config.stoch_dy + + stoch_file = convert_srf_to_stoch(srf_file, dx, dy) + with open(stoch_ffp, "w") as f: + stoch_file.dump(f) + realisations.append_log_entry(realisation_ffp)