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)