From eab1faf4f46714fe58b6ffc88fbbfbab3a83d83e Mon Sep 17 00:00:00 2001 From: Jake Faulkner Date: Mon, 31 Aug 2026 20:31:02 +1200 Subject: [PATCH] Generate NZCVM velocity model configurations from a realisation `create-nzvm-input` writes the config `nzcvm generate` consumes. One realisation can produce several distinct velocity models, along two independent axes: --format sw4 | emod3d the grid the model is sampled onto: a mesh-refined SW4 grid written as an sfile, or a uniform EMOD3D binary grid --layers full | tomography the layer stack queried to fill it: the whole thing (basins, coastline, offshore, Ely GTL taper) or just the background tomography and the numerical clamps, which is the reference no-bells-and-whistles model Splitting those two axes is the point: comparing a full model against a tomography-only one on the *same* grid is how a basin's contribution is isolated, and it should not require hand-editing a config. For SW4 the model is padded by more than the grid is -- one supergrid width, as `create-sw4-input` uses, plus a few grid points of slack -- so SW4 never queries outside the sfile. `test_model_padding_contains_the_padded_grid` asserts that ordering directly, because if it ever inverts the symptom is a solver reading off the end of the model rather than an error. The `nzcvm` realisation section carries the layer stack as nzcvm's own `LayerConfig` objects and the DEM surface as a `Path`, which is what the `path_serialiser` added at the bottom of this stack exists to write. The container image installs nzcvm explicitly. It is a bare requirement resolved through `[tool.uv.sources]`, which pip does not read, so `pip install workflow` alone cannot find it. The other git-sourced dependencies are PEP 508 direct references in `dependencies` and pip resolves those itself. Blocked on an nzcvm release: this makes nzcvm a hard import of `workflow.schemas`, so every workflow install pulls it in, not just the SW4 path. Co-Authored-By: Claude Opus 5 --- container/Dockerfile | 7 +- pyproject.toml | 1 + tests/test_realisation.py | 1 + tests/test_sw4_template.py | 21 +- .../default_parameters/root/defaults.yaml | 296 ++++++++++++++ workflow/realisations.py | 15 + workflow/schemas.py | 10 + workflow/scripts/nzvm_input_template.py | 363 ++++++++++++++++++ 8 files changed, 712 insertions(+), 2 deletions(-) create mode 100755 workflow/scripts/nzvm_input_template.py diff --git a/container/Dockerfile b/container/Dockerfile index 6d727517..17efa2f3 100644 --- a/container/Dockerfile +++ b/container/Dockerfile @@ -77,7 +77,12 @@ RUN curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- -y # openquake.engine 3.26+ requires gdal~=3.13.1, which is sdist-only and refuses to # build against the libgdal 3.12.2 that Ubuntu ships. Cap it until libgdal catches up. ARG WORKFLOW_BRANCH=pegasus +# nzcvm is a bare requirement resolved through `[tool.uv.sources]`, which pip does +# not read, so it has to be named explicitly here. The other git-sourced +# dependencies (im-calculation, hf-simulation) are PEP 508 direct references in +# `dependencies` and pip resolves those itself. +ARG NZCVM_BRANCH=master RUN python3 -m venv $VIRTUAL_ENV && \ pip install --no-cache-dir --upgrade pip setuptools wheel && \ - pip install --no-cache-dir "git+https://github.com/ucgmsim/workflow@${WORKFLOW_BRANCH}" "openquake.engine<3.26" --no-binary=h5py && \ + pip install --no-cache-dir "git+https://github.com/lispandfound/nzcvm@${NZCVM_BRANCH}" "git+https://github.com/ucgmsim/workflow@${WORKFLOW_BRANCH}" "openquake.engine<3.26" --no-binary=h5py && \ chmod -R a+rX /sw diff --git a/pyproject.toml b/pyproject.toml index 013bf13c..f56defd9 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -72,6 +72,7 @@ copy-domain-parameters = "workflow.scripts.copy_velocity_model_parameters:app" srf-to-hdf5 = "workflow.scripts.srf_to_hdf5:app" create-e3d-par = "workflow.scripts.create_e3d_par:app" create-sw4-input = "workflow.scripts.sw4_template:app" +create-nzvm-input = "workflow.scripts.nzvm_input_template:app" generate-stoch = "workflow.scripts.generate_stoch:app" merge-ts = "workflow.scripts.merge_ts:app" hf-sim = "workflow.scripts.hf_sim:app" diff --git a/tests/test_realisation.py b/tests/test_realisation.py index 18da2867..635bcdb7 100644 --- a/tests/test_realisation.py +++ b/tests/test_realisation.py @@ -1086,6 +1086,7 @@ def test_sources(tmp_path: Path) -> None: realisations.Resolution, realisations.RuptureVelocity, realisations.Refinements, + realisations.NZCVMSettings, ], ) @pytest.mark.parametrize("defaults_version", list(defaults.DefaultsVersion)) diff --git a/tests/test_sw4_template.py b/tests/test_sw4_template.py index 14a97c51..a82088a5 100644 --- a/tests/test_sw4_template.py +++ b/tests/test_sw4_template.py @@ -16,7 +16,7 @@ from velocity_modelling.bounding_box import BoundingBox from workflow import defaults, sw4 from workflow.realisations import Refinement, Refinements, SW4Parameters -from workflow.scripts import sw4_template +from workflow.scripts import nzvm_input_template, sw4_template SPONGE_KM = 12.0 """The v26_7_1Hz sponge width, in kilometres.""" @@ -114,6 +114,25 @@ def test_bottom_refinement_holds_the_sponge(depth_km: float) -> None: assert thickness / coarsest > supergrid_width / coarsest +def test_model_padding_contains_the_padded_grid(domain: BoundingBox) -> None: + """`create-nzvm-input` pads by more than `create-sw4-input` does. + + If this ordering ever inverts, SW4 queries outside the sfile. + """ + model_padding_km = ( + SPONGE_KM + nzvm_input_template.SW4_MODEL_SLACK_GRIDPOINTS * 400.0 / 1000.0 + ) + grid = domain.pad(pad_x=(SPONGE_KM, SPONGE_KM), pad_y=(SPONGE_KM, SPONGE_KM)) + model = domain.pad( + pad_x=(model_padding_km, model_padding_km), + pad_y=(model_padding_km, model_padding_km), + ) + + assert model.polygon.contains_properly(grid.polygon) + assert model.extent_x > grid.extent_x + assert model.extent_y > grid.extent_y + + def test_adjust_for_topography_leaves_resolutions_alone() -> None: """The coarsest resolution can be read before or after adjustment. diff --git a/workflow/default_parameters/root/defaults.yaml b/workflow/default_parameters/root/defaults.yaml index 0bc55587..90b256f6 100644 --- a/workflow/default_parameters/root/defaults.yaml +++ b/workflow/default_parameters/root/defaults.yaml @@ -225,6 +225,302 @@ srf: risetimefac: 1.0 risetimedep: 0.0 inittime: 0.0 +nzcvm: + surface: /nzcvm/resources/dem.zarr + chunks: + i: 256 + j: 256 + layers: + - clamps: + vs: + max: null + min: 500.0 + qs: + min: 0.05 + max: 0.05 + min_ref: "vs" + max_ref: "vs" + qp: + min: 2.0 + max: 2.0 + min_ref: "qs" + max_ref: "qs" + + max_vp_vs_ratio: 4.0 + min_vp_vs_ratio: 1.73 + provides: [] + requires: [] + type: clamp + - type: backus + samples: 5 + - coastline: /nzcvm/resources/coastline.wkb.gz + provides: + - coastline + requires: [] + type: coastline + - basin_depth: + - bottom_depth: 0 + distance: 0 + - bottom_depth: 2000 + distance: 20000 + - bottom_depth: 3000 + distance: 50000 + model: + - alpha: 1.0 + bottom_depth: 50 + qp: 38.0 + qs: 19.0 + rho: 1810.0 + vp: 1800.0 + vs: 380.0 + - alpha: 1.0 + bottom_depth: 150 + qp: 48.0 + qs: 24.0 + rho: 1810.0 + vp: 1800.0 + vs: 480.0 + - alpha: 1.0 + bottom_depth: 300 + qp: 58.0 + qs: 29.0 + rho: 1810.0 + vp: 1800.0 + vs: 580.0 + - alpha: 1.0 + bottom_depth: 500 + qp: 68.0 + qs: 34.0 + rho: 1810.0 + vp: 1800.0 + vs: 680.0 + - alpha: 1.0 + bottom_depth: 800 + qp: 75.0 + qs: 37.5 + rho: 1810.0 + vp: 1800.0 + vs: 750.0 + - alpha: 1.0 + bottom_depth: 1200 + qp: 83.0 + qs: 41.5 + rho: 1810.0 + vp: 1800.0 + vs: 830.0 + - alpha: 1.0 + bottom_depth: 1800 + qp: 90.0 + qs: 45.0 + rho: 1860.0 + vp: 1900.0 + vs: 900.0 + - alpha: 1.0 + bottom_depth: 2600 + qp: 100.0 + qs: 50.0 + rho: 1920.0 + vp: 2030.0 + vs: 1000.0 + - alpha: 1.0 + bottom_depth: 3600 + qp: 105.0 + qs: 52.5 + rho: 1970.0 + vp: 2140.0 + vs: 1050.0 + - alpha: 1.0 + bottom_depth: 4800 + qp: 110.0 + qs: 55.0 + rho: 1990.0 + vp: 2200.0 + vs: 1100.0 + - alpha: 1.0 + bottom_depth: 6200 + qp: 115.0 + qs: 57.5 + rho: 2060.0 + vp: 2400.0 + vs: 1150.0 + - alpha: 1.0 + bottom_depth: 7800 + qp: 120.0 + qs: 60.0 + rho: 2150.0 + vp: 2700.0 + vs: 1200.0 + - alpha: 1.0 + bottom_depth: 9600 + qp: 143.0 + qs: 71.5 + rho: 2220.0 + vp: 3000.0 + vs: 1430.0 + - alpha: 1.0 + bottom_depth: 11600 + qp: 164.0 + qs: 82.0 + rho: 2280.0 + vp: 3270.0 + vs: 1640.0 + - alpha: 1.0 + bottom_depth: 13800 + qp: 186.0 + qs: 93.0 + rho: 2320.0 + vp: 3530.0 + vs: 1860.0 + - alpha: 1.0 + bottom_depth: 16200 + qp: 207.0 + qs: 103.5 + rho: 2360.0 + vp: 3800.0 + vs: 2070.0 + - alpha: 1.0 + bottom_depth: 18800 + qp: 228.0 + qs: 114.0 + rho: 2400.0 + vp: 4070.0 + vs: 2280.0 + - alpha: 1.0 + bottom_depth: 21600 + qp: 249.0 + qs: 124.5 + rho: 2440.0 + vp: 4330.0 + vs: 2490.0 + - alpha: 1.0 + bottom_depth: 24600 + qp: 270.0 + qs: 135.0 + rho: 2480.0 + vp: 4600.0 + vs: 2700.0 + - alpha: 1.0 + bottom_depth: 27800 + qp: 277.0 + qs: 138.5 + rho: 2490.0 + vp: 4710.0 + vs: 2770.0 + - alpha: 1.0 + bottom_depth: 31200 + qp: 284.0 + qs: 142.0 + rho: 2510.0 + vp: 4820.0 + vs: 2840.0 + - alpha: 1.0 + bottom_depth: 34800 + qp: 291.0 + qs: 145.5 + rho: 2520.0 + vp: 4930.0 + vs: 2910.0 + - alpha: 1.0 + bottom_depth: 38600 + qp: 298.0 + qs: 149.0 + rho: 2540.0 + vp: 5040.0 + vs: 2980.0 + - alpha: 1.0 + bottom_depth: 42600 + qp: 305.0 + qs: 152.5 + rho: 2560.0 + vp: 5150.0 + vs: 3050.0 + - alpha: 1.0 + bottom_depth: 46800 + qp: 312.0 + qs: 156.0 + rho: 2580.0 + vp: 5260.0 + vs: 3120.0 + - alpha: 1.0 + bottom_depth: 51200 + qp: 319.0 + qs: 159.5 + rho: 2600.0 + vp: 5370.0 + vs: 3190.0 + - alpha: 1.0 + bottom_depth: 55800 + qp: 326.0 + qs: 163.0 + rho: 2610.0 + vp: 5480.0 + vs: 3260.0 + - alpha: 1.0 + bottom_depth: 60600 + qp: 333.0 + qs: 166.5 + rho: 2630.0 + vp: 5590.0 + vs: 3330.0 + - alpha: 1.0 + bottom_depth: 65600 + qp: 340.0 + qs: 170.0 + rho: 2660.0 + vp: 5700.0 + vs: 3400.0 + - alpha: 1.0 + bottom_depth: 73600 + qp: 360.0 + qs: 180.0 + rho: 2720.0 + vp: 6000.0 + vs: 3600.0 + - alpha: 1.0 + bottom_depth: 85600 + qp: 360.0 + qs: 180.0 + rho: 2720.0 + vp: 6000.0 + vs: 3600.0 + - alpha: 1.0 + bottom_depth: 112600 + qp: 370.0 + qs: 185.0 + rho: 2830.0 + vp: 6500.0 + vs: 3700.0 + - alpha: 1.0 + bottom_depth: 151600 + qp: 430.0 + qs: 215.0 + rho: 3120.0 + vp: 7500.0 + vs: 4300.0 + - alpha: 1.0 + bottom_depth: 10151599 + qp: 460.0 + qs: 230.0 + rho: 3330.0 + vp: 8100.0 + vs: 4600.0 + provides: [] + requires: + - coastline + type: offshore + - depth_t: 450.0 + provides: [] + requires: + - coastline + type: ely + vs30: /nzcvm/resources/vs30.zarr + - model_globs: + - "*.zarr" + model_path: /nzcvm/models + provides: [] + requires: [] + type: query + velocity_model: min_vs: 0.5 ds_multiplier: 1.2 diff --git a/workflow/realisations.py b/workflow/realisations.py index ccc1e9cf..53a104b4 100644 --- a/workflow/realisations.py +++ b/workflow/realisations.py @@ -23,6 +23,8 @@ import numpy as np import numpy.typing as npt import pandas as pd +from nzcvm.config.layers import LayerConfig +from nzcvm.coordinates import Coordinate from schema import Schema from IM import im_calculation @@ -998,6 +1000,19 @@ def refinements_for_depth(self, depth: float) -> list[Refinement]: return refinements +@dataclasses.dataclass +class NZCVMSettings(RealisationConfiguration): + _config_key: ClassVar[str] = "nzcvm" + _schema: ClassVar[Schema] = schemas.NZCVM_SCHEMA + + layers: list[LayerConfig] + """nzcvm layer config""" + chunks: dict[Coordinate, int] + """nzcvm chunk configuration""" + surface: Path + """nzcvm DEM surface""" + + @dataclasses.dataclass class VelocityModelParameters(RealisationConfiguration): """Parameters defining the velocity model.""" diff --git a/workflow/schemas.py b/workflow/schemas.py index fb98dbbd..ef7cb128 100644 --- a/workflow/schemas.py +++ b/workflow/schemas.py @@ -8,9 +8,12 @@ import dataclasses from enum import IntEnum, StrEnum +from pathlib import Path import numpy as np import pandas as pd +from nzcvm.config.layers import LayerConfig +from nzcvm.coordinates import Coordinate from schema import And, Literal, Optional, Or, Schema, Use from IM import im_calculation @@ -960,6 +963,13 @@ def _corners_to_array(corners_spec: list[dict[str, float]]) -> np.ndarray: } ) +NZCVM_SCHEMA = Schema( + { + Literal("layers"): [Use(LayerConfig.from_dict)], + Literal("chunks"): Or({}, {Use(Coordinate): int}), + Literal("surface"): Use(Path), + } +) SEED_SCHEMA = Schema( { Literal("genslip_seed", description="The random seed passed to genslip."): int, diff --git a/workflow/scripts/nzvm_input_template.py b/workflow/scripts/nzvm_input_template.py new file mode 100755 index 00000000..4fe17089 --- /dev/null +++ b/workflow/scripts/nzvm_input_template.py @@ -0,0 +1,363 @@ +#!/usr/bin/env python3 +"""NZCVM Velocity Model Configuration. + +Description +----------- +Generate an NZCVM velocity model configuration from a realisation. The +configuration describes both the *grid* the model is sampled onto and the +*layers* that are queried to fill it, so one realisation can produce several +distinct velocity models: + +- ``--format sw4`` writes an :class:`~nzcvm.config.grids.sw4.SW4GridConfig`, + a mesh-refined grid written out as an sfile. +- ``--format emod3d`` writes an :class:`~nzcvm.config.grids.emod3d.EMOD3DGrid`, + a uniform grid written out as EMOD3D binary files. +- ``--layers full`` keeps the realisation's whole layer stack (basins, + coastline, offshore, Ely GTL taper, ...). +- ``--layers tomography`` keeps only the background tomography query and the + numerical clamps, which is the reference "no bells and whistles" model. + +Inputs +------ +1. A realisation with domain parameters, resolution, and an `nzcvm` section. + +Outputs +------- +1. An NZCVM velocity model configuration (JSON), ready for `nzcvm generate`. + +Environment +----------- +Can be run in the cybershake container. Can also be run from your own computer +using the `create-nzvm-input` command after +`pip install workflow@git+https://github.com/ucgmsim/workflow`. + +Usage +----- +`create-nzvm-input [OPTIONS] REALISATION_FFP OUTPUT_PATH` + +For More Help +------------- +See the output of `create-nzvm-input --help`. +""" + +import dataclasses +import functools +import json +from enum import StrEnum, auto +from pathlib import Path +from typing import Annotated + +import pyproj +import typer +from nzcvm.config import VelocityModelConfig +from nzcvm.config.grids.emod3d import EMOD3DGrid, TopographyType +from nzcvm.config.grids.model import Model +from nzcvm.config.grids.sw4 import MeshRefinement, SW4GridConfig +from nzcvm.config.layers.core import LayerConfig +from nzcvm.config.layers.query import QueryLayerConfig +from nzcvm.coordinates import Coordinate + +from qcore import cli +from workflow import sw4 +from workflow.realisations import ( + DomainParameters, + NZCVMSettings, + RealisationMetadata, + Refinements, + Resolution, + SW4Parameters, + VelocityModelParameters, +) + +app = typer.Typer() + +NZTM_EPSG = 2193 +"""The CRS every NZCVM grid is built in.""" + +SW4_DEPTH_OFFSET_KM = 10.0 +"""Extra depth (km) modelled below the domain so SW4 refinement adjustment has room.""" + +SW4_MODEL_SLACK_GRIDPOINTS = 4 +"""Gridpoints of velocity model kept beyond the padded SW4 grid, per face. + +`create-sw4-input` pads the SW4 grid by one supergrid sponge width on every +lateral face and on the bottom, so the requested domain is the grid's *interior* +rather than being eaten into by the absorbing layer. The velocity model has to +cover that padded grid or SW4 queries outside the sfile, so this module pads the +model by the same sponge width plus this slack. The slack absorbs the rounding +that turns an extent into a gridpoint count in NZCVM and in SW4, which need not +agree exactly. +""" + +EMOD3D_FREE_SURFACE_PADDING = 1 +"""Extra gridpoints in z for EMOD3D's free surface shift. + +EMOD3D shifts the model down one gridpoint for the free surface and so never +reads the deepest layer (see `genmodel.c`), which means the model must carry +one row more than the domain has. `create-e3d-par` assumes the same padding +when it checks the velocity model's size, so the two must agree. +""" + +TOMOGRAPHY_LAYER_TYPES = frozenset({"clamp", "query"}) +"""Layer types kept by `--layers tomography`. + +The tomography-only model is the background 3D tomography and nothing else: +no basins, no offshore model, no Ely near-surface taper. The clamps stay +because they are numerical guards on the simulation (minimum Vs, Vp/Vs +ratio bounds), not a feature of the geology. +""" + +DEFAULT_TOMOGRAPHY_GLOB = "ep2020.zarr" +"""Model glob matching only the tomography mesh in the NZCVM model directory.""" + + +class GridFormat(StrEnum): + """The simulator whose grid the velocity model is sampled onto.""" + + SW4 = auto() + EMOD3D = auto() + + +class LayerSelection(StrEnum): + """How much of the realisation's NZCVM layer stack to keep.""" + + FULL = auto() + """Every layer the realisation configures.""" + TOMOGRAPHY = auto() + """Background tomography and numerical clamps only.""" + + +def select_layers( + layers: list[LayerConfig], + selection: LayerSelection, + tomography_glob: str, +) -> list[LayerConfig]: + """Filter a realisation's layer stack down to `selection`. + + Parameters + ---------- + layers : list[LayerConfig] + The realisation's configured layers, in pipeline order. + selection : LayerSelection + Which layers to keep. + tomography_glob : str + Glob matching the tomography mesh, used to narrow the query layer when + `selection` is `LayerSelection.TOMOGRAPHY`. + + Returns + ------- + list[LayerConfig] + The kept layers, in their original order. + + Raises + ------ + ValueError + If the selection leaves no query layer, since a velocity model with + nothing to query is always empty. + """ + if selection == LayerSelection.FULL: + return layers + + selected = [] + for layer in layers: + if getattr(layer, "type", None) not in TOMOGRAPHY_LAYER_TYPES: + continue + if isinstance(layer, QueryLayerConfig): + layer = dataclasses.replace(layer, model_globs=[tomography_glob]) + selected.append(layer) + + if not any(isinstance(layer, QueryLayerConfig) for layer in selected): + raise ValueError( + "Layer selection left no query layer: the realisation's nzcvm " + "section must configure one for a velocity model to be generated." + ) + return selected + + +def sw4_grid( + domain_parameters: DomainParameters, + refinements: Refinements, + sw4_params: SW4Parameters, + nzcvm_settings: NZCVMSettings, +) -> SW4GridConfig: + """Build the SW4 mesh-refined grid configuration. + + The model is deliberately larger than the domain. `create-sw4-input` pads the + SW4 grid by one supergrid sponge width on every face, so the model has to be + padded by at least as much or SW4 queries outside the sfile. + + Parameters + ---------- + domain_parameters : DomainParameters + The domain to model. + refinements : Refinements + The theoretical mesh refinements, resolved against the domain depth. + sw4_params : SW4Parameters + The SW4 parameters, which fix the supergrid sponge width the model has + to cover. + nzcvm_settings : NZCVMSettings + Supplies the topographic surface and chunking. + + Returns + ------- + SW4GridConfig + The grid configuration. + """ + domain = domain_parameters.domain + domain_refinements = refinements.refinements_for_depth( + domain_parameters.depth + SW4_DEPTH_OFFSET_KM + ) + + # NOTE: The sponge width must be measured on the refinements *SW4* will use, + # i.e. resolved against the bare domain depth. Resolving against + # `depth + SW4_DEPTH_OFFSET_KM` can land on a coarser bottom layer than SW4 + # actually gets (a 400 m bottom layer gives a 12 km sponge where SW4's 200 m + # one gives 6 km), which would overstate the padding needed here and, worse, + # disagree with `create-sw4-input`. + coarsest_resolution = sw4.coarsest_resolution(refinements, domain_parameters.depth) + supergrid_width = sw4.supergrid_width(sw4_params, coarsest_resolution) + model_padding = supergrid_width + SW4_MODEL_SLACK_GRIDPOINTS * coarsest_resolution + + domain_refinements[-1].bottom += model_padding + + padding_km = model_padding / 1000.0 + # `BoundingBox.pad` takes kilometres and pads along the box's own rotated + # axes, keeping the centroid and the azimuth. Padding symmetrically keeps the + # model concentric with the padded SW4 grid, which is what makes + # `create-sw4-input`'s extent-versus-footprint check a containment check. + domain = domain.pad(pad_x=(padding_km, padding_km), pad_y=(padding_km, padding_km)) + + origin_lat, origin_lon = domain.origin + return SW4GridConfig( + extent_x=domain.extent_x * 1000.0, + extent_y=domain.extent_y * 1000.0, + orientation=Model( + origin_lon=origin_lon, + origin_lat=origin_lat, + crs=pyproj.CRS(NZTM_EPSG), + azimuth=domain.great_circle_bearing, + ), + surface=nzcvm_settings.surface, + chunks=nzcvm_settings.chunks, + refinements={ + f"layer_{refinement.resolution}m": MeshRefinement( + resolution=refinement.resolution, bottom=refinement.bottom + ) + for refinement in domain_refinements + }, + ) + + +def emod3d_grid( + domain_parameters: DomainParameters, + resolution: Resolution, + velocity_model_parameters: VelocityModelParameters, + nzcvm_settings: NZCVMSettings, +) -> EMOD3DGrid: + """Build the EMOD3D uniform grid configuration. + + Parameters + ---------- + domain_parameters : DomainParameters + The domain to model. + resolution : Resolution + The uniform grid spacing, in kilometres. + velocity_model_parameters : VelocityModelParameters + Supplies the topography type. + nzcvm_settings : NZCVMSettings + Supplies the topographic surface and chunking. + + Returns + ------- + EMOD3DGrid + The grid configuration, sized to match what `create-e3d-par` expects. + """ + domain = domain_parameters.domain + origin_lat, origin_lon = domain.origin + return EMOD3DGrid( + surface=nzcvm_settings.surface, + nx=domain_parameters.nx(resolution.resolution), + ny=domain_parameters.ny(resolution.resolution), + nz=domain_parameters.nz(resolution.resolution) + EMOD3D_FREE_SURFACE_PADDING, + # NOTE: NZCVM works in metres, the realisation's resolution in kilometres. + resolution=resolution.resolution * 1000.0, + orientation=Model( + origin_lon=origin_lon, + origin_lat=origin_lat, + crs=pyproj.CRS(NZTM_EPSG), + azimuth=domain.great_circle_bearing, + ), + topo_type=TopographyType(velocity_model_parameters.topo_type.lower()), + # The EMOD3D grid is only ever chunked horizontally: depth is a single + # chunk by construction, so a k chunk size would be silently ignored. + chunks={ + coordinate: size + for coordinate, size in nzcvm_settings.chunks.items() + if coordinate in (Coordinate.I, Coordinate.J) + }, + ) + + +@cli.from_docstring(app) +def generate_template( + realisation_ffp: Path, + output_path: Path, + format: Annotated[GridFormat, typer.Option()] = GridFormat.SW4, + layers: Annotated[LayerSelection, typer.Option()] = LayerSelection.FULL, + tomography_glob: Annotated[str, typer.Option()] = DEFAULT_TOMOGRAPHY_GLOB, +) -> None: + """Generate an NZCVM velocity model configuration from a realisation file. + + Parameters + ---------- + realisation_ffp : Path + Path to the realisation file containing domain parameters. + output_path : Path + Path where the generated configuration will be written. + format : GridFormat + The simulator whose grid to sample the velocity model onto. + layers : LayerSelection + How much of the realisation's layer stack to keep. `tomography` drops + the basins, offshore model and near-surface taper. + tomography_glob : str + Glob matching the tomography mesh, used when `layers` is `tomography`. + """ + metadata = RealisationMetadata.read_from_realisation(realisation_ffp) + domain_parameters = DomainParameters.read_from_realisation(realisation_ffp) + nzcvm_settings = NZCVMSettings.read_from_realisation(realisation_ffp) + + if not nzcvm_settings.layers: + raise ValueError("NZCVM requires at least one defined layer.") + + match format: + case GridFormat.SW4: + grid = sw4_grid( + domain_parameters, + Refinements.read_from_realisation_or_defaults( + realisation_ffp, metadata.defaults_version + ), + SW4Parameters.read_from_realisation_or_defaults( + realisation_ffp, metadata.defaults_version + ), + nzcvm_settings, + ) + case GridFormat.EMOD3D: + grid = emod3d_grid( + domain_parameters, + Resolution.read_from_realisation_or_defaults( + realisation_ffp, metadata.defaults_version + ), + VelocityModelParameters.read_from_realisation_or_defaults( + realisation_ffp, metadata.defaults_version + ), + nzcvm_settings, + ) + + config = VelocityModelConfig( + grid=grid, + layers=select_layers(nzcvm_settings.layers, layers, tomography_glob), + ) + + output_path.write_text( + config.to_json(encoder=functools.partial(json.dumps, indent=4)) # ty: ignore + )