From 90cb6061e2c7cfcf6bbe552c0acce50f068548ae Mon Sep 17 00:00:00 2001 From: Jake Faulkner Date: Thu, 10 Sep 2026 18:55:39 +1200 Subject: [PATCH 1/3] Add a borehole grid: vertical profiles at a set of sites MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Comparing a velocity model against a borehole log, or handing a 1-D profile to a site-response code, only needs one column of samples per station. The existing grids all fill a rotated box, so getting there has meant generating a whole volume and slicing it, or hand-rolling a Grid outside the config system entirely (which the README talked people through). `grid.type = "borehole"` extracts one column per site, running from the topography down to a fixed depth at a fixed vertical resolution, so profiles from different sites line up sample for sample. The grid is shaped `(len(sites), 1, nk)`: the singleton `j` axis keeps the `(i, j, k)` contract every layer relies on, so the whole layer chain — coastline, offshore, Ely, clamp, query — works unchanged. Sites come from an inline list, or from a CSV or Parquet file. They are given in a global CRS (WGS84 by default, `sites_crs` to override) and mapped into the grid CRS. Longitude and latitude place a site and the config reserves nothing else. Every other key in a `[[grid.sites]]` block, and every other column of a site file, becomes a coordinate on the `i` axis under the name the caller gave it, so a station code, a network or a driller's reference reaches the output. A label keeps the type it was written with. Mashumaro folds the extra keys in a `__pre_deserialize__` hook, so the longitude and latitude validators still run and a config file needn't spell the mapping out. `keep_extra_columns = false` drops the labels and keeps only the spatial coordinates. Two things that has to guard. A coordinate shadows a variable or attribute of the same name, so a label called `name` would turn `grid.name` from the grid's name into a DataArray. `RESERVED_COORDINATES` is derived from `GridSchema`'s own fields, plus the logical index, the coastline coordinate and the components that share that index, so it keeps up if the schema changes. Sites also have to agree on which labels they carry, since a missing key is nearly always a typo and the alternative is a column of nulls. A borehole grid is a bag of independent columns, so there is nothing for a model origin to be the origin of. `Model` therefore splits: `Projection` holds the CRS and its cached transformers, and `Model` extends it with the origin and azimuth the extent-based grids lay themselves out against. A borehole config takes `[grid.projection]` and no origin at all; the Grid attributes that record one are derived from the sites. `examples/borehole.toml` runs four profiles over the `just synthetic` dataset, two of them inside a basin, and the tests build their columns on the synthetic topography so the expected elevations follow from closed form. Co-Authored-By: Claude Opus 5 --- .gitignore | 4 + README.md | 196 ++++++++++-- examples/borehole.toml | 91 ++++++ examples/sites.csv | 5 + nzcvm/config/grids/borehole.py | 158 ++++++++++ nzcvm/config/grids/model.py | 39 ++- nzcvm/config/validation.py | 11 + nzcvm/grids/__init__.py | 4 +- nzcvm/grids/borehole.py | 237 ++++++++++++++ nzcvm/grids/grid.py | 17 + tests/test_borehole.py | 558 +++++++++++++++++++++++++++++++++ 11 files changed, 1284 insertions(+), 36 deletions(-) create mode 100644 examples/borehole.toml create mode 100644 examples/sites.csv create mode 100644 nzcvm/config/grids/borehole.py create mode 100644 nzcvm/grids/borehole.py create mode 100644 tests/test_borehole.py diff --git a/.gitignore b/.gitignore index 3f969e7..a16ae62 100644 --- a/.gitignore +++ b/.gitignore @@ -201,6 +201,10 @@ Cargo.lock *.sfile *.rho3dfile *.vs3dfile + +# Example inputs are small and belong in the repo, unlike the data products +# the patterns above exclude. +!examples/*.csv *.vp3dfile *.wkb.* diff --git a/README.md b/README.md index 7093187..d9e0e1a 100644 --- a/README.md +++ b/README.md @@ -117,17 +117,19 @@ validate bounds, layer ordering, and layer dependencies. ### Grid types -Set by `grid.type`. All three are topography-following and chunked lazily with -Dask. +Set by `grid.type`. All of them are topography-following and chunked lazily +with Dask. -| Type | Key parameters | -|-----------|------------------------------------------------------------------------------| -| `sw4` | `extent_x/y`, `refinements` (2:1 nested resolutions, ordered automatically) | -| `regular` | `extent_x/y`, `thickness`, `resolution_x/y/z` (fixed vertical resolution) | -| `emod3d` | `nx`, `ny`, `nz`, `resolution`, `topo_type` | +| Type | Key parameters | +|------------|-----------------------------------------------------------------------------| +| `sw4` | `extent_x/y`, `refinements` (2:1 nested resolutions, ordered automatically) | +| `regular` | `extent_x/y`, `thickness`, `resolution_x/y/z` (fixed vertical resolution) | +| `emod3d` | `nx`, `ny`, `nz`, `resolution`, `topo_type` | +| `borehole` | `sites`, `depth`, `resolution_z` (one vertical profile per site) | -Every grid also takes `surface` (path to a DEM), an `[grid.orientation]` block, -and optional `[grid.chunks]`: +Every grid takes `surface` (path to a DEM) and optional `[grid.chunks]`. The +three volumetric grids fill a rotated box, so they also take an +`[grid.orientation]` block naming the model origin of that box: ```toml [grid.orientation] @@ -141,6 +143,92 @@ i = 256 j = 256 ``` +A `borehole` grid holds independent columns with no box to orient, so it takes +a bare `[grid.projection]` instead: just a CRS. + +### Borehole grids + +`grid.type = "borehole"` extracts one vertical profile per site rather than +filling a volume. Each column runs from the topography down to `depth` at a +fixed `resolution_z`, so the profiles line up sample for sample and compare +directly. + +```toml +[grid] +type = "borehole" +surface = "./resources/dem.zarr" +depth = 600.0 # metres below the topography +resolution_z = 25.0 # metres between samples + +[grid.projection] +crs = 'EPSG:2193' # CRS the profiles are extracted in + +[[grid.sites]] +longitude = 172.6218 +latitude = -43.5283 +site = "CACS" # not a keyword; see below +network = "NZ" +``` + +Sites come in a global CRS (WGS84 unless `sites_crs` says otherwise), and the +builder maps them into `grid.projection.crs`. Instead of listing them inline, +point `sites` at a CSV or Parquet file with `longitude` and `latitude` +columns: + +```toml +sites = "stations.csv" +``` + +Keep that line ahead of `[grid.projection]`: TOML would otherwise read it as +a key of that table. + +#### Site labels + +Longitude and latitude place a site, and the config reserves nothing else. +Every other key, and every other column of a site file, becomes a coordinate +on the grid's `i` axis under the name the caller gave it, which the writers +keep. Neither `site` nor `network` in the preceding example is a +keyword the grid interprets. Both end up in the output because nothing +reserves them. Rename them, add a driller's reference, drop them entirely: the +grid doesn't care. + +A label keeps the type the caller wrote, so a numeric column arrives numeric. +Each site needs the same set of labels, since the alternative is a column of +nulls where one site was missing a key. + +`nzcvm.grids.grid.RESERVED_COORDINATES` lists the names a label may not take: +the variables and attributes `GridSchema` declares (`x`, `y`, `z`, `depth`, +`name`, `geometry`, …), the `(i, j, k)` index, and the components that share +that index with the grid. A coordinate shadows a variable of the +same name, so +a label called `name` would turn `grid.name` from the grid's name into an +array. Naming one of those raises rather than corrupting the grid. + +To drop the labels and keep only the spatial coordinates: + +```toml +keep_extra_columns = false +``` + +The result has shape `(len(sites), 1, nk)`: one column per site, with `nk` +samples down each column. The singleton `j` axis preserves the `(i, j, k)` +contract every layer relies on. The site labels index the `i` axis, so the +output reads back per station: + +```python +import xarray as xr + +tree = xr.open_datatree("boreholes.zarr", engine="zarr") +grid = tree["grids/boreholes"].ds +vs = tree["qualities/boreholes"].ds.vs.squeeze("j") + +for n, site in enumerate(grid.site.values): + print(site, vs[n].values) +``` + +`examples/borehole.toml` runs four profiles over the `just synthetic` dataset, +two of them inside a basin. + ### Layers Layers run in the order listed, outermost first. Each one delegates down the @@ -181,6 +269,7 @@ See `examples/` for complete, working configs: | `whole_country.toml` | `regular` grid over New Zealand | | `near_fault_config.toml` | A custom layer (`examples/near_fault.py`) in a config | | `synthetic.toml` | The same chain over the `just synthetic` dataset | +| `borehole.toml` | `borehole` grid: profiles at four sites | ```toml [metadata] @@ -503,8 +592,25 @@ config. A grid is an xarray Dataset built through `GridSchema`, which fixes the contract every layer relies on: `x`, `y`, `z` and `depth` on the logical `(i, j, k)` index (metres, projected CRS, `z` positive down), plus the -attributes below. The smallest useful grid is a borehole: one vertical -column, shaped `(1, 1, nk)`: +attributes below. + +Those four variables and the attributes are the whole of what `GridSchema` +accepts, so a builder can't pass an extra *data variable*. It can attach extra +*coordinates* after construction, though, and every stage keeps them: layers, +`map_blocks`, the Zarr and NetCDF writers, and the read back through +`GridSchema.from_dataset`. xarray keeps a coordinate on each data variable it +indexes, so anything that reassembles a dataset from those variables picks it +up again. The borehole grid labels its columns this way: + +```python +grid = grid.assign_coords(site=("i", ["GULL", "TERR"])) +``` + +The name must avoid `RESERVED_COORDINATES`, since a coordinate shadows a +variable or attribute of the same name. + +Here is a transect: a line of vertical columns between two +points, shaped `(n, 1, nk)`. ```python import numpy as np @@ -513,17 +619,32 @@ import shapely from nzcvm.grids.grid import Grid, GridSchema -def borehole_grid(x: float, y: float, bottom: float, dz: float) -> Grid: - """A single vertical column of query points (a synthetic borehole).""" - depth = np.arange(0.0, bottom, dz, dtype=np.float32).reshape(1, 1, -1) +def transect_grid( + start: tuple[float, float], + end: tuple[float, float], + n: int, + bottom: float, + dz: float, +) -> Grid: + """A line of vertical columns of query points, sampled at *n* stations.""" + depth = np.arange(0.0, bottom, dz, dtype=np.float32) + along = np.linspace(0.0, 1.0, n, dtype=np.float32) + x = start[0] + along * (end[0] - start[0]) + y = start[1] + along * (end[1] - start[1]) + + # Broadcast the (n,) line and the (nk,) depths into (n, 1, nk). + x, _ = np.meshgrid(x, depth, indexing="ij") + y, depth = np.meshgrid(y, depth, indexing="ij") + x, y, depth = (array[:, np.newaxis, :] for array in (x, y, depth)) + return GridSchema.new( - x=np.full_like(depth, x), - y=np.full_like(depth, y), + x=x, + y=y, z=depth, depth=depth, - name="borehole", + name="transect", resolution=dz, - geometry=shapely.Point(x, y), + geometry=shapely.LineString([start, end]), origin_lon=np.float32(174.7762), origin_lat=np.float32(-41.2865), azimuth=np.float32(0.0), @@ -541,14 +662,19 @@ from pathlib import Path from nzcvm.config.layers.query import QueryLayerConfig from nzcvm.layers.pipeline import build_pipeline -grid = borehole_grid(x=1_749_150.0, y=5_428_150.0, bottom=500.0, dz=100.0) +grid = transect_grid( + start=(1_749_150.0, 5_428_150.0), + end=(1_759_150.0, 5_428_150.0), + n=3, + bottom=500.0, + dz=100.0, +) pipeline = build_pipeline( grid.geometry, [QueryLayerConfig(model_path=Path("models"), model_globs=["*.zarr"])], ) qualities = pipeline(grid) -print(qualities.vs.values.ravel()) -# [ 380. 580. 2643.6 2647.6 2651.6] +print(qualities.vs.values[0].ravel()) # Vs down the first column, in m/s ``` To drive a grid from a config file, register a builder against a `GridConfig` @@ -566,26 +692,36 @@ from nzcvm.grids.grid import Grid @dataclass -class BoreholeConfig(GridConfig): - x: float - y: float +class TransectConfig(GridConfig): + start: tuple[float, float] + end: tuple[float, float] + n: int bottom: float dz: float - type: Literal["borehole"] = "borehole" + type: Literal["transect"] = "transect" @build_grids_from_config.register -def _(config: BoreholeConfig) -> dict[str, Grid]: - return {"borehole": borehole_grid(config.x, config.y, config.bottom, config.dz)} +def _(config: TransectConfig) -> dict[str, Grid]: + return { + "transect": transect_grid( + config.start, config.end, config.n, config.bottom, config.dz + ) + } ``` Which makes this config valid: ```toml [grid] -type = "borehole" -x = 1749150.0 -y = 5428150.0 +type = "transect" +start = [1749150.0, 5428150.0] +end = [1759150.0, 5428150.0] +n = 3 bottom = 500.0 dz = 100.0 ``` + +`nzcvm.grids.borehole` is the same pattern, done properly: a registered +`GridConfig` whose builder reads its own sites and DEM, over Dask-backed +coordinates. diff --git a/examples/borehole.toml b/examples/borehole.toml new file mode 100644 index 0000000..c5b8e22 --- /dev/null +++ b/examples/borehole.toml @@ -0,0 +1,91 @@ +# Borehole profiles over the synthetic dataset. +# +# just synthetic +# uv run nzcvm generate examples/borehole.toml synthetic/boreholes.zarr +# +# A borehole grid extracts one vertical profile per site instead of filling a +# volume, so it has no extent, no azimuth and no model origin: just a +# projection to extract the profiles in. +# +# Longitude and latitude place a site. Every other key below (`site`, +# `network`) is one the grid knows nothing about, and rides through to the +# output as a coordinate on the i axis. Set keep_extra_columns = false to drop +# them. + +[metadata] +title = "Synthetic borehole profiles" +keywords = ["synthetic", "borehole", "testing"] + +[grid] +type = "borehole" +surface = "./synthetic/dem.zarr" + +# Bottom of every column, in metres below the topography, and the sample +# spacing down to it. +depth = 600.0 +resolution_z = 25.0 + +# Sites can also come from a CSV or Parquet file with longitude and latitude +# columns, where the other columns label the sites the same way. It has to +# stay above the tables below, as TOML would otherwise read it as a key of +# `[grid.projection]`. +# +# sites = "examples/sites.csv" + +[grid.projection] +crs = 'EPSG:2193' + +[[grid.sites]] +longitude = 172.15 +latitude = -43.70 +site = "GULL" # inside the `gully` basin +network = "NZ" + +[[grid.sites]] +longitude = 172.30 +latitude = -43.53 +site = "TERR" # inside the `terrace` basin +network = "NZ" + +[[grid.sites]] +longitude = 172.10 +latitude = -43.45 +site = "RIDG" # up in the hills, outside every basin +network = "SC" + +[[grid.sites]] +longitude = 172.55 +latitude = -43.60 +site = "SEAB" # offshore +network = "SC" + +[[layers]] +type = "clamp" +min_vp_vs_ratio = 1.73 +max_vp_vs_ratio = 4.0 + +[[layers]] +type = "coastline" +coastline = "synthetic/coastline.wkb.gz" + +[[layers]] +type = "offshore" +model = [ + {bottom_depth = 50.0, rho = 1810.0, vp = 1800.0, vs = 380.0, qp = 100.0, qs = 50.0, alpha = 1.0}, + {bottom_depth = 300.0, rho = 1810.0, vp = 1800.0, vs = 580.0, qp = 100.0, qs = 50.0, alpha = 1.0}, + {bottom_depth = 1200.0, rho = 1810.0, vp = 1800.0, vs = 830.0, qp = 100.0, qs = 50.0, alpha = 1.0}, +] +basin_depth = [ + {distance = 0.0, bottom_depth = 0.0}, + {distance = 10000.0, bottom_depth = 1000.0}, +] + +[[layers]] +type = "ely" +vs30 = "./synthetic/vs30.zarr" +depth_t = 450.0 + +[[layers]] +type = "query" +model_path = "./synthetic/models" +model_globs = ["*.zarr"] diff --git a/examples/sites.csv b/examples/sites.csv new file mode 100644 index 0000000..42bf0d8 --- /dev/null +++ b/examples/sites.csv @@ -0,0 +1,5 @@ +site,network,longitude,latitude +GULL,NZ,172.15,-43.70 +TERR,NZ,172.30,-43.53 +RIDG,SC,172.10,-43.45 +SEAB,SC,172.55,-43.60 diff --git a/nzcvm/config/grids/borehole.py b/nzcvm/config/grids/borehole.py new file mode 100644 index 0000000..2d5dee2 --- /dev/null +++ b/nzcvm/config/grids/borehole.py @@ -0,0 +1,158 @@ +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any, Literal + +from mashumaro import field_options +from pyproj import CRS + +from nzcvm.config.core import ConfigObject +from nzcvm.config.grids.model import Projection +from nzcvm.config.validation import ( + CRSStrategy, + GeographicCRS, + Latitude, + Longitude, + PositiveFloat, +) +from nzcvm.coordinates import WGS84_EPSG, Coordinate + +from .core import GridConfig + +DEFAULT_CHUNK_SIZES = {Coordinate.I: 64} + + +#: The keys that place a site. Everything else given for a site is a label. +SPATIAL_KEYS = ("longitude", "latitude") + + +@dataclass +class Site(ConfigObject): + """One borehole location, in the global CRS, plus whatever labels it has. + + Longitude and latitude place the site, and the config doesn't reserve any + other key, so a site takes as much or as little description as the caller + has to give it. Each label becomes a coordinate on the grid's ``i`` axis, + which the writers keep. + + Attributes + ---------- + longitude, latitude : + Position in :attr:`BoreholeGridConfig.sites_crs`, which defaults to + WGS84. The grid builder projects it into the grid CRS. + labels : + Everything else given for the site. Decoding a config folds every + key except :data:`SPATIAL_KEYS` in here, so a config file needn't + spell the mapping out. + + Examples + -------- + >>> Site.from_dict( + ... {"longitude": 172.15, "latitude": -43.7, "site": "GULL", "network": "NZ"} + ... ) + Site(longitude=172.15, latitude=-43.7, labels={'site': 'GULL', 'network': 'NZ'}) + """ + + longitude: Longitude + latitude: Latitude + labels: dict[str, Any] = field(default_factory=dict) + + @classmethod + def __pre_deserialize__(cls, d: dict[str, Any]) -> dict[str, Any]: + """Fold every key that doesn't place the site into :attr:`labels`. + + Named *d* to match the mashumaro hook this overrides. + """ + return { + **{k: v for k, v in d.items() if k in SPATIAL_KEYS}, + "labels": {k: v for k, v in d.items() if k not in SPATIAL_KEYS}, + } + + +@dataclass +class BoreholeGridConfig(GridConfig): + """Vertical profiles at a set of sites, at a fixed vertical resolution. + + A borehole grid holds independent columns, so an extent, an azimuth and a + model origin have nothing to describe. It takes a bare + :class:`~nzcvm.config.grids.model.Projection` in place of the + :class:`~nzcvm.config.grids.model.Model` orientation block that the + extent-based grids use. + + The built grid has shape ``(len(sites), 1, nk)``: one column per site, + a singleton ``j`` axis, and ``nk`` samples spaced *resolution_z* apart + from the topography down to *depth*. + + Attributes + ---------- + surface : + Path to the topographic surface mesh file. Used to translate depth + to elevation, so each column starts at the ground. + sites : + Either an inline list of :class:`Site` objects, or a path to a CSV or + Parquet file with ``longitude`` and ``latitude`` columns. Any other + key or column labels the site. + depth : + Depth of the bottom of every column, in metres below the topography. + resolution_z : + Vertical sample spacing in metres. + projection : + Projected CRS to extract the profiles in. + sites_crs : + Geographic CRS of the site coordinates (default WGS84). The builder + maps each site from here into *projection* before querying. It has to + be geographic, since a site is a longitude and a latitude. + keep_extra_columns : + Whether the builder puts the site labels on the grid, and so in the + output (default ``True``). Set it to ``False`` to keep only the + spatial coordinates and drop the rest. + + Examples + -------- + TOML, with the sites inline. Neither ``site`` nor ``network`` is a + keyword here, and both end up in the output because nothing reserves + them:: + + [grid] + type = "borehole" + surface = "./synthetic/dem.zarr" + depth = 500.0 + resolution_z = 25.0 + + [grid.projection] + crs = 'EPSG:2193' + + [[grid.sites]] + longitude = 172.62 + latitude = -43.53 + site = "CACS" + network = "NZ" + + or read from a file, where the extra columns do the same job:: + + sites = "examples/sites.csv" + """ + + surface: Path + + sites: list[Site] | Path + + depth: PositiveFloat + resolution_z: PositiveFloat + + projection: Projection + + sites_crs: GeographicCRS = field( + default_factory=lambda: CRS.from_epsg(WGS84_EPSG), + metadata=field_options(serialization_strategy=CRSStrategy()), + ) + + keep_extra_columns: bool = True + + chunks: dict[Coordinate, int] = field(default_factory=lambda: DEFAULT_CHUNK_SIZES) + + type: Literal["borehole"] = "borehole" + + def __post_init__(self) -> None: + super().__post_init__() + if isinstance(self.sites, list) and not self.sites: + raise ValueError("A borehole grid needs at least one site.") diff --git a/nzcvm/config/grids/model.py b/nzcvm/config/grids/model.py index cc32b76..769420f 100644 --- a/nzcvm/config/grids/model.py +++ b/nzcvm/config/grids/model.py @@ -11,10 +11,20 @@ @dataclass(frozen=True) -class Model(ConfigObject): - origin_lon: Longitude - origin_lat: Latitude - azimuth: float +class Projection(ConfigObject): + """The projected CRS a grid's coordinates live in. + + Just enough to convert between geographic and projected coordinates. A + grid with no origin to rotate about, such as a set of boreholes, takes one + of these on its own. A grid placed at a model origin takes a + :class:`Model` instead. + + Attributes + ---------- + crs : + Target projected CRS, such as ``EPSG:2193`` for NZTM2000. + """ + crs: CRS = field(metadata=field_options(serialization_strategy=CRSStrategy())) @functools.cached_property @@ -25,6 +35,27 @@ def from_wgs84(self) -> Transformer: def to_wgs84(self) -> Transformer: return Transformer.from_crs(self.crs, WGS84_EPSG, always_xy=True) + def transformer_from(self, crs: CRS) -> Transformer: + """Return a transformer from *crs* into this projection.""" + return Transformer.from_crs(crs, self.crs, always_xy=True) + + +@dataclass(frozen=True) +class Model(Projection): + """A projection plus the origin and azimuth that place a grid. + + Attributes + ---------- + origin_lon, origin_lat : + Geographic origin of the local grid, in WGS84 degrees. + azimuth : + Clockwise rotation of the grid from true north, in degrees. + """ + + origin_lon: Longitude + origin_lat: Latitude + azimuth: float + @functools.cached_property def origin(self) -> tuple[float, float]: return tuple(self.from_wgs84.transform(self.origin_lon, self.origin_lat)) diff --git a/nzcvm/config/validation.py b/nzcvm/config/validation.py index 1d47bcd..e0d4474 100644 --- a/nzcvm/config/validation.py +++ b/nzcvm/config/validation.py @@ -159,6 +159,16 @@ def longitude(value: float) -> None: ) +def geographic_crs(value: pyproj.CRS) -> pyproj.CRS: + """Ensures a CRS is geographic, so its coordinates are longitude and latitude.""" + if value is not None and not value.is_geographic: + raise ValueError( + f"Expected a geographic (longitude/latitude) CRS, got the projected " + f"CRS {value.name!r}." + ) + return value + + def in_choices(choices: Collection[Any]) -> Callable[[Any], Any]: """Ensures a value is a member of an allowed set of options.""" allowed = set(choices) @@ -177,6 +187,7 @@ def validator(v: Any) -> Any: PositiveFloat = Annotated[float, validate_positive] Latitude = Annotated[float, latitude] Longitude = Annotated[float, longitude] +GeographicCRS = Annotated[pyproj.CRS, geographic_crs] UnitIntervalFloat = Annotated[float, validate_non_negative, le(1.0)] diff --git a/nzcvm/grids/__init__.py b/nzcvm/grids/__init__.py index b5189af..b43e528 100644 --- a/nzcvm/grids/__init__.py +++ b/nzcvm/grids/__init__.py @@ -1,5 +1,5 @@ -from nzcvm.grids import emod3d, regular, sw4 +from nzcvm.grids import borehole, emod3d, regular, sw4 from nzcvm.grids.builder import build_grids_from_config from nzcvm.grids.grid import Grid -__all__ = ["Grid", "build_grids_from_config", "emod3d", "regular", "sw4"] +__all__ = ["Grid", "borehole", "build_grids_from_config", "emod3d", "regular", "sw4"] diff --git a/nzcvm/grids/borehole.py b/nzcvm/grids/borehole.py new file mode 100644 index 0000000..bf9c041 --- /dev/null +++ b/nzcvm/grids/borehole.py @@ -0,0 +1,237 @@ +"""Topography-following borehole grid builder. + +Provides :func:`build_borehole` for constructing the set of vertical columns +described by a :class:`~nzcvm.config.grids.borehole.BoreholeGridConfig`. +Each site becomes one column of query points, sampled at a strictly fixed Z +resolution from the topography down to a fixed depth, so a run extracts a +directly comparable profile per site rather than a filled volume. + +Columns are independent, so none of the extent, azimuth or origin metadata +the volumetric grids record means anything here. The builder still fills in +the :class:`~nzcvm.grids.grid.Grid` attributes that name an origin, taking the +centroid of the sites as the origin and the south-west corner of their +bounding box as the bottom-left corner. + +Longitude and latitude place a site. The config doesn't reserve any other +key, so every other key or column becomes a coordinate on the ``i`` axis under +the name the caller gave it, which is how a station code or a network ends up +in the layer chain and the output. +:data:`~nzcvm.grids.grid.RESERVED_COORDINATES` lists the names a label may not +take, and ``keep_extra_columns = false`` drops the labels altogether. +""" + +from collections.abc import Callable +from pathlib import Path +from typing import Any + +import dask.array as da +import numpy as np +import pandas as pd +import shapely +import xarray as xr + +from nzcvm.config.grids.borehole import SPATIAL_KEYS, BoreholeGridConfig, Site +from nzcvm.coordinates import Coordinate +from nzcvm.grids import helpers +from nzcvm.grids.builder import build_grids_from_config +from nzcvm.grids.grid import RESERVED_COORDINATES, Grid, GridSchema +from nzcvm.models.surface import Surface + +#: Name of the one grid a borehole config builds. +GRID_NAME = "boreholes" + +#: Readers for the supported site file formats, keyed by suffix. +SITE_READERS: dict[str, Callable[[Path], pd.DataFrame]] = { + ".csv": pd.read_csv, + ".parquet": pd.read_parquet, + ".pq": pd.read_parquet, +} + + +def read_sites(path: Path) -> list[Site]: + """Read borehole sites from a CSV or Parquet file. + + Parameters + ---------- + path : + File with ``longitude`` and ``latitude`` columns. Every other column + becomes a label on the site it belongs to. + + Returns + ------- + list[Site] + One :class:`~nzcvm.config.grids.borehole.Site` per row, in file order. + + Raises + ------ + ValueError + If the suffix isn't a supported format, or a spatial column is + missing. + """ + reader = SITE_READERS.get(path.suffix.lower()) + if reader is None: + supported = ", ".join(sorted(SITE_READERS)) + raise ValueError( + f"Cannot read sites from '{path}': expected one of {supported}" + ) + + frame = reader(path) + missing = [column for column in SPATIAL_KEYS if column not in frame.columns] + if missing: + raise ValueError( + f"Site file '{path}' is missing the {', '.join(missing)} column(s). " + f"Expected {', '.join(SPATIAL_KEYS)}." + ) + + return [ + Site( + longitude=float(record.pop("longitude")), + latitude=float(record.pop("latitude")), + labels=record, + ) + for record in frame.to_dict("records") + ] + + +def site_labels(sites: list[Site]) -> dict[str, np.ndarray]: + """Collect the site labels into one array per label name. + + Parameters + ---------- + sites : + Sites to read the labels off. + + Returns + ------- + dict[str, numpy.ndarray] + One array per label, ordered as *sites* are and typed by whatever + NumPy infers from the values. + + Raises + ------ + ValueError + If a label collides with + :data:`~nzcvm.grids.grid.RESERVED_COORDINATES`, or if the sites do + not agree on which labels they carry. Disagreement is nearly always + a typo, and the alternative is a column of nulls. + """ + names = list(sites[0].labels) + + reserved = sorted(RESERVED_COORDINATES.intersection(names)) + if reserved: + raise ValueError( + f"Site label(s) {', '.join(reserved)} would shadow a grid variable " + f"or attribute of the same name. Rename them, or set " + f"keep_extra_columns = false." + ) + + for position, site in enumerate(sites): + if set(site.labels) != set(names): + raise ValueError( + f"Site {position} carries labels " + f"{sorted(site.labels) or 'none'}, but site 0 carries " + f"{sorted(names)}. Every site needs the same labels." + ) + + return {name: np.asarray([site.labels[name] for site in sites]) for name in names} + + +def resolve_sites(config: BoreholeGridConfig) -> list[Site]: + """Return the sites a config names, reading them from disk if needed.""" + sites = config.sites if isinstance(config.sites, list) else read_sites(config.sites) + if not sites: + raise ValueError(f"No sites found in '{config.sites}'.") + return sites + + +def _columns(values: np.ndarray, index: np.ndarray, chunk: int) -> xr.DataArray: + """Lay a per-site value out over the ``(i, j)`` plane as one column each.""" + chunked = da.from_array(values.astype(np.float32), chunks=chunk) + return xr.DataArray( + chunked[:, np.newaxis], + dims=[Coordinate.I, Coordinate.J], + coords={Coordinate.I: index, Coordinate.J: [0]}, + ) + + +def _borehole_grid( + x_phys: xr.DataArray, + y_phys: xr.DataArray, + surface: xr.DataArray, + depth: float, + resolution_z: float, + **kwargs: Any, +) -> Grid: + nk = np.round(depth / resolution_z).astype(int) + 1 + + # Depth is purely a function of k and resolution_z, identically for every + # column, which is what makes the profiles comparable between sites. + # Chunking only ever applies to i/j. k always stays one chunk. + zeta_depth = xr.DataArray( + np.linspace(0.0, depth, num=nk, dtype=np.float32), + dims=[Coordinate.K], + coords={Coordinate.K: np.arange(nk)}, + ).chunk({Coordinate.K: -1}) + + # Elevation (z) is the surface elevation shifted downward by the fixed + # depths, so every column starts at the ground rather than at sea level. + x, y, z, column_depth = xr.broadcast( + x_phys, y_phys, surface + zeta_depth, zeta_depth + ) + x, y, z, column_depth = helpers.ensure_chunks(x, y, z, column_depth) + + return GridSchema.new(x, y, z, column_depth, **kwargs) + + +@build_grids_from_config.register +def build_borehole(config: BoreholeGridConfig) -> dict[str, Grid]: + sites = resolve_sites(config) + index = np.arange(len(sites)) + labels = site_labels(sites) if config.keep_extra_columns else {} + + transformer = config.projection.transformer_from(config.sites_crs) + x, y = transformer.transform( + np.array([site.longitude for site in sites]), + np.array([site.latitude for site in sites]), + ) + + chunk = config.chunks[Coordinate.I] + x_phys = _columns(x, index, chunk) + y_phys = _columns(y, index, chunk) + + # A borehole grid has no origin of its own, so stand one up from the sites + # for the benefit of the Grid attributes every writer expects. + to_wgs84 = config.projection.to_wgs84 + origin_lon, origin_lat = to_wgs84.transform(x.mean(), y.mean()) + min_lon, min_lat = to_wgs84.transform(x.min(), y.min()) + + topographic_surface = Surface.load(config.surface) + z_surface = helpers.compute_surface_elevation( + topographic_surface, + x_phys, + y_phys, + ) + + grid = _borehole_grid( + x_phys, + y_phys, + z_surface, + name=GRID_NAME, + depth=config.depth, + resolution_z=config.resolution_z, + resolution=config.resolution_z, + geometry=shapely.MultiPoint(np.column_stack((x, y))), + origin_lon=origin_lon, + origin_lat=origin_lat, + azimuth=np.float32(0.0), + grid_azimuth=np.float32(0.0), + bottom_left_lon=min_lon, + bottom_left_lat=min_lat, + ) + # Labelling i is what makes the output readable: without it, the only + # route back to a station is the order the config listed it in. + grid = grid.assign_coords( + {name: (Coordinate.I, values) for name, values in labels.items()} + ) + + return {GRID_NAME: grid} diff --git a/nzcvm/grids/grid.py b/nzcvm/grids/grid.py index ce3e8b7..bff8bfe 100644 --- a/nzcvm/grids/grid.py +++ b/nzcvm/grids/grid.py @@ -1,3 +1,4 @@ +import dataclasses from dataclasses import dataclass from typing import Literal @@ -6,6 +7,7 @@ import xarray as xr from xarray_dataclasses import AsDataset, Attr, Data, DataOptions +from nzcvm.components import Component from nzcvm.coordinates import Coordinate @@ -46,6 +48,21 @@ def from_dataset(cls, dataset: xr.Dataset) -> Grid: return dset +#: Names a grid builder may not give an extra coordinate. +#: +#: A coordinate shadows a variable or an attribute of the same name, so a +#: coordinate called ``name`` would turn ``grid.name`` from the grid's name +#: into a :class:`~xarray.DataArray`. The set covers everything +#: :class:`GridSchema` declares, the logical index, the coordinate the +#: coastline layer adds, and the components that share that index with the +#: grid. +RESERVED_COORDINATES: frozenset[str] = frozenset( + [field.name for field in dataclasses.fields(GridSchema)] + + [Coordinate.I, Coordinate.J, Coordinate.K, Coordinate.COASTLINE] + + list(Component) +) + + def grid_like_at_depth(grid: Grid, depth: float) -> Grid: # Select a z-layer of the block. # The array [0] as the selection is important because it preserves the k diff --git a/tests/test_borehole.py b/tests/test_borehole.py new file mode 100644 index 0000000..ae7d98e --- /dev/null +++ b/tests/test_borehole.py @@ -0,0 +1,558 @@ +"""Tests for the borehole grid: config decoding, site loading and the builder. + +Every test builds on the synthetic DEM (:mod:`nzcvm.synthetic`), so nothing +reads a real data file, and the expected elevations follow from the analytic +topography rather than from a fixture. +""" + +from __future__ import annotations + +from pathlib import Path + +import numpy as np +import pandas as pd +import pytest +import xarray as xr +from mashumaro.exceptions import InvalidFieldValue +from pyproj import CRS, Transformer + +from nzcvm import synthetic +from nzcvm.components import Component +from nzcvm.config.grids.borehole import ( + DEFAULT_CHUNK_SIZES, + BoreholeGridConfig, + Site, +) +from nzcvm.config.grids.model import Model, Projection +from nzcvm.config.metadata import ModelMetadata +from nzcvm.config.velocity_model import VelocityModelConfig +from nzcvm.coordinates import Coordinate +from nzcvm.formats import Format, write_velocity_model +from nzcvm.grids.borehole import read_sites, resolve_sites +from nzcvm.grids.builder import build_grids_from_config +from nzcvm.grids.grid import RESERVED_COORDINATES, Grid +from nzcvm.layers.pipeline import execute_model_pipeline +from nzcvm.models.mesh import StructuredMeshSchema +from nzcvm.qualities import Qualities, QualitiesSchema +from nzcvm.query import ModelRange +from nzcvm.velocity_model import VelocityModel + +_NZTM = CRS.from_epsg(2193) +_NZGD2000 = CRS.from_epsg(4167) +_WGS84 = CRS.from_epsg(4326) +_TO_NZTM = Transformer.from_crs(4326, _NZTM, always_xy=True) + +# Sites in the hills and out to sea, spread across the domain. `site` and +# `network` are ordinary labels, not keywords the grid interprets. +SITES = [ + Site(longitude=172.15, latitude=-43.70, labels={"site": "GULL", "network": "NZ"}), + Site(longitude=172.10, latitude=-43.45, labels={"site": "RIDG", "network": "NZ"}), + Site(longitude=172.55, latitude=-43.60, labels={"site": "SEAB", "network": "SC"}), +] + +#: The labels on the sites, both treated like any other. +SITE = "site" +NETWORK = "network" +NAMES = [site.labels[SITE] for site in SITES] + + +@pytest.fixture(scope="module") +def synthetic_surface(tmp_path_factory: pytest.TempPathFactory) -> Path: + """The synthetic topography, written as a Zarr surface mesh.""" + lon, lat = synthetic.DOMAIN.sample(32, 32) + mesh_lon, mesh_lat = np.meshgrid(lon, lat, indexing="ij") + x, y = _TO_NZTM.transform(mesh_lon, mesh_lat) + # A surface file uses +z down, the same convention as the grids. + z = -synthetic.elevation(mesh_lon, mesh_lat) + + path = tmp_path_factory.mktemp("surfaces") / "dem.zarr" + StructuredMeshSchema.new( + x=x.astype(np.float32), + y=y.astype(np.float32), + z=z.astype(np.float32), + i=np.arange(32), + j=np.arange(32), + name="synthetic", + ).to_zarr(path, mode="w") + return path + + +def _config( + surface: Path, + sites: list[Site] | Path = SITES, + depth: float = 400.0, + resolution_z: float = 100.0, + sites_crs: CRS = _WGS84, + keep_extra_columns: bool = True, + chunks: dict[Coordinate, int] = DEFAULT_CHUNK_SIZES, +) -> BoreholeGridConfig: + return BoreholeGridConfig( + surface=surface, + sites=sites, + depth=depth, + resolution_z=resolution_z, + projection=Projection(crs=_NZTM), + sites_crs=sites_crs, + keep_extra_columns=keep_extra_columns, + chunks=chunks, + ) + + +def _build(config: BoreholeGridConfig) -> Grid: + grids = build_grids_from_config(config) + assert list(grids) == ["boreholes"] + return grids["boreholes"] + + +# --------------------------------------------------------------------------- +# Projection / Model split +# --------------------------------------------------------------------------- + + +def test_projection_needs_no_origin() -> None: + """The whole point of splitting Projection out of Model.""" + projection = Projection(crs=_NZTM) + x, y = projection.from_wgs84.transform(172.0, -43.5) + assert projection.to_wgs84.transform(x, y) == pytest.approx((172.0, -43.5)) + + +def test_model_is_still_a_projection() -> None: + model = Model(origin_lon=172.0, origin_lat=-43.5, azimuth=39.0, crs=_NZTM) + assert isinstance(model, Projection) + assert model.grid_origin_x == pytest.approx( + model.from_wgs84.transform(172.0, -43.5)[0] + ) + + +def test_transformer_from_reaches_the_projection() -> None: + projection = Projection(crs=_NZTM) + from_nzgd = projection.transformer_from(_NZGD2000) + assert from_nzgd.transform(172.0, -43.5) == pytest.approx( + projection.from_wgs84.transform(172.0, -43.5), abs=1.0 + ) + + +# --------------------------------------------------------------------------- +# Site loading +# --------------------------------------------------------------------------- + + +def _write_sites(path: Path) -> Path: + frame = pd.DataFrame( + { + "longitude": [site.longitude for site in SITES], + "latitude": [site.latitude for site in SITES], + SITE: NAMES, + NETWORK: [site.labels[NETWORK] for site in SITES], + } + ) + if path.suffix == ".csv": + frame.to_csv(path, index=False) + else: + frame.to_parquet(path) + return path + + +@pytest.mark.parametrize("suffix", [".csv", ".parquet", ".pq"]) +def test_read_sites_round_trips(tmp_path: Path, suffix: str) -> None: + assert read_sites(_write_sites(tmp_path / f"sites{suffix}")) == SITES + + +def test_read_sites_preserves_file_order(tmp_path: Path) -> None: + path = tmp_path / "sites.csv" + pd.DataFrame( + {"longitude": [172.1, 172.2], "latitude": [-43.5, -43.6], SITE: ["B", "A"]} + ).to_csv(path, index=False) + assert [site.labels[SITE] for site in read_sites(path)] == ["B", "A"] + + +def test_read_sites_keeps_extra_columns_as_labels(tmp_path: Path) -> None: + """Longitude and latitude are the only reserved columns. The rest + describe the site.""" + path = tmp_path / "sites.csv" + pd.DataFrame( + { + "longitude": [172.1], + "latitude": [-43.5], + SITE: ["A"], + "elevation": [12.0], + } + ).to_csv(path, index=False) + assert read_sites(path) == [ + Site( + longitude=172.1, + latitude=-43.5, + labels={SITE: "A", "elevation": 12.0}, + ) + ] + + +def test_read_sites_accepts_a_file_with_no_labels(tmp_path: Path) -> None: + path = tmp_path / "sites.csv" + pd.DataFrame({"longitude": [172.1], "latitude": [-43.5]}).to_csv(path, index=False) + assert read_sites(path) == [Site(longitude=172.1, latitude=-43.5, labels={})] + + +def test_read_sites_rejects_unknown_format(tmp_path: Path) -> None: + path = tmp_path / "sites.txt" + path.write_text("name,longitude,latitude\n") + with pytest.raises(ValueError, match="expected one of"): + read_sites(path) + + +def test_read_sites_reports_missing_columns(tmp_path: Path) -> None: + path = tmp_path / "sites.csv" + pd.DataFrame({SITE: ["A"], "longitude": [172.1]}).to_csv(path, index=False) + with pytest.raises(ValueError, match="missing the latitude column"): + read_sites(path) + + +def test_resolve_sites_accepts_inline_sites(synthetic_surface: Path) -> None: + assert resolve_sites(_config(synthetic_surface)) == SITES + + +def test_resolve_sites_reads_a_file(synthetic_surface: Path, tmp_path: Path) -> None: + config = _config(synthetic_surface, sites=_write_sites(tmp_path / "sites.csv")) + assert resolve_sites(config) == SITES + + +def test_resolve_sites_rejects_an_empty_file( + synthetic_surface: Path, tmp_path: Path +) -> None: + path = tmp_path / "sites.csv" + path.write_text("name,longitude,latitude\n") + with pytest.raises(ValueError, match="No sites found"): + resolve_sites(_config(synthetic_surface, sites=path)) + + +def test_config_rejects_an_empty_inline_site_list(synthetic_surface: Path) -> None: + with pytest.raises(ValueError, match="at least one site"): + _config(synthetic_surface, sites=[]) + + +def test_config_rejects_a_projected_sites_crs(synthetic_surface: Path) -> None: + with pytest.raises(InvalidFieldValue, match="geographic"): + _config(synthetic_surface, sites_crs=_NZTM) + + +# --------------------------------------------------------------------------- +# Site labels +# +# Longitude and latitude place a site. Every other key or column is a label +# the grid passes through to the output under the name the caller gave it. +# --------------------------------------------------------------------------- + + +def test_extra_config_keys_become_labels() -> None: + site = Site.from_dict( + {"longitude": 172.1, "latitude": -43.5, SITE: "A", "depth_drilled": 30.0} + ) + assert site.labels == {SITE: "A", "depth_drilled": 30.0} + + +def test_a_site_needs_no_labels_at_all() -> None: + assert Site.from_dict({"longitude": 172.1, "latitude": -43.5}).labels == {} + + +def test_labels_become_coordinates_on_the_site_axis(synthetic_surface: Path) -> None: + grid = _build(_config(synthetic_surface)) + for label in (SITE, NETWORK): + assert grid[label].dims == (Coordinate.I,) + assert list(grid[label].values) == [site.labels[label] for site in SITES] + + +def test_numeric_labels_keep_a_numeric_dtype(synthetic_surface: Path) -> None: + """A label is whatever the caller wrote, not necessarily a string.""" + sites = [ + Site(longitude=site.longitude, latitude=site.latitude, labels={"cased": n * 10}) + for n, site in enumerate(SITES) + ] + grid = _build(_config(synthetic_surface, sites=sites)) + assert np.issubdtype(grid["cased"].dtype, np.integer) + assert list(grid["cased"].values) == [0, 10, 20] + + +def test_keep_extra_columns_false_drops_the_labels(synthetic_surface: Path) -> None: + grid = _build(_config(synthetic_surface, keep_extra_columns=False)) + assert set(grid.coords) == {Coordinate.I, Coordinate.J, Coordinate.K} + # The opt-out leaves the spatial coordinates alone. + assert grid.x.shape == (len(SITES), 1, 5) + + +@pytest.mark.parametrize("reserved", ["name", "x", "depth", "vs", "i", "geometry"]) +def test_a_label_may_not_shadow_a_grid_name( + synthetic_surface: Path, reserved: str +) -> None: + """A coordinate shadows a variable or attribute of the same name, so + `grid.name` would stop being the grid's name.""" + sites = [ + Site(longitude=site.longitude, latitude=site.latitude, labels={reserved: "x"}) + for site in SITES + ] + with pytest.raises(ValueError, match="would shadow"): + _build(_config(synthetic_surface, sites=sites)) + + +def test_reserved_names_cover_the_grid_contract() -> None: + """Derived from GridSchema, so it keeps up if the schema changes.""" + assert {"x", "y", "z", "depth", "name", "geometry", "resolution"} <= ( + RESERVED_COORDINATES + ) + assert {Coordinate.I, Coordinate.J, Coordinate.K} <= RESERVED_COORDINATES + assert set(Component) <= RESERVED_COORDINATES + assert SITE not in RESERVED_COORDINATES + + +def test_sites_have_to_agree_on_their_labels(synthetic_surface: Path) -> None: + """A missing label is nearly always a typo, and the alternative is a + column of nulls.""" + sites = [ + Site( + longitude=SITES[0].longitude, latitude=SITES[0].latitude, labels={SITE: "A"} + ), + Site(longitude=SITES[1].longitude, latitude=SITES[1].latitude, labels={}), + ] + with pytest.raises(ValueError, match="Every site needs the same labels"): + _build(_config(synthetic_surface, sites=sites)) + + +def test_disagreement_is_allowed_once_labels_are_dropped( + synthetic_surface: Path, +) -> None: + sites = [ + Site( + longitude=SITES[0].longitude, latitude=SITES[0].latitude, labels={SITE: "A"} + ), + Site(longitude=SITES[1].longitude, latitude=SITES[1].latitude, labels={}), + ] + grid = _build(_config(synthetic_surface, sites=sites, keep_extra_columns=False)) + assert grid.x.shape == (2, 1, 5) + + +# --------------------------------------------------------------------------- +# Grid geometry +# --------------------------------------------------------------------------- + + +def test_grid_is_one_column_per_site(synthetic_surface: Path) -> None: + grid = _build(_config(synthetic_surface)) + assert grid.x.shape == (len(SITES), 1, 5) + assert grid.sizes[Coordinate.J] == 1 + + +def test_grid_labels_columns_by_site(synthetic_surface: Path) -> None: + grid = _build(_config(synthetic_surface)) + assert list(grid[SITE].values) == NAMES + assert list(grid[NETWORK].values) == [site.labels[NETWORK] for site in SITES] + + +def test_depth_is_identical_across_sites(synthetic_surface: Path) -> None: + """Comparable profiles are the point: every column shares one depth axis.""" + depth = _build(_config(synthetic_surface)).depth.values + expected = np.linspace(0.0, 400.0, 5, dtype=np.float32) + for column in depth.reshape(-1, depth.shape[-1]): + assert column == pytest.approx(expected) + + +def test_resolution_z_sets_the_sample_count(synthetic_surface: Path) -> None: + grid = _build(_config(synthetic_surface, depth=600.0, resolution_z=25.0)) + assert grid.sizes[Coordinate.K] == 25 + assert float(grid.depth.max()) == pytest.approx(600.0) + + +def test_sites_are_projected_into_the_grid_crs(synthetic_surface: Path) -> None: + grid = _build(_config(synthetic_surface)) + expected_x, expected_y = _TO_NZTM.transform( + [site.longitude for site in SITES], [site.latitude for site in SITES] + ) + assert grid.x.values[:, 0, 0] == pytest.approx(np.float32(expected_x), rel=1e-6) + assert grid.y.values[:, 0, 0] == pytest.approx(np.float32(expected_y), rel=1e-6) + + +def test_x_and_y_are_constant_down_each_column(synthetic_surface: Path) -> None: + grid = _build(_config(synthetic_surface)) + for axis in (grid.x, grid.y): + assert np.all(axis.values == axis.values[:, :, :1]) + + +def test_columns_start_at_the_topography(synthetic_surface: Path) -> None: + grid = _build(_config(synthetic_surface)) + # Grids are +z down, so the top of each column is minus the elevation. + expected = -synthetic.elevation( + [site.longitude for site in SITES], [site.latitude for site in SITES] + ) + assert grid.z.values[:, 0, 0] == pytest.approx(expected, abs=15.0) + + +def test_columns_follow_the_topography_down(synthetic_surface: Path) -> None: + grid = _build(_config(synthetic_surface)) + top = grid.z.values[:, :, 0] + bottom = grid.z.values[:, :, -1] + assert (bottom - top) == pytest.approx(np.full_like(top, 400.0)) + + +def test_a_geographic_sites_crs_lands_in_the_same_place( + synthetic_surface: Path, +) -> None: + wgs84 = _build(_config(synthetic_surface)) + nzgd2000 = _build(_config(synthetic_surface, sites_crs=_NZGD2000)) + # NZGD2000 and WGS84 are within a metre of each other over New Zealand. + assert nzgd2000.x.values == pytest.approx(wgs84.x.values, abs=1.0) + + +def test_geometry_covers_every_site(synthetic_surface: Path) -> None: + """The query layer prunes models against this, so it has to hit each site.""" + grid = _build(_config(synthetic_surface)) + assert len(grid.geometry.geoms) == len(SITES) + assert grid.geometry.bounds == pytest.approx( + ( + float(grid.x.min()), + float(grid.y.min()), + float(grid.x.max()), + float(grid.y.max()), + ), + abs=1.0, + ) + + +def test_derived_origin_is_the_site_centroid(synthetic_surface: Path) -> None: + """A borehole grid has no configured origin, so the attributes come from + the sites. Downstream writers still expect them to be present.""" + grid = _build(_config(synthetic_surface)) + lons = [site.longitude for site in SITES] + lats = [site.latitude for site in SITES] + assert grid.origin_lon == pytest.approx(np.mean(lons), abs=1e-2) + assert grid.origin_lat == pytest.approx(np.mean(lats), abs=1e-2) + assert grid.bottom_left_lon == pytest.approx(min(lons), abs=1e-2) + assert grid.bottom_left_lat == pytest.approx(min(lats), abs=1e-2) + assert grid.azimuth == 0.0 + assert grid.grid_azimuth == 0.0 + + +def test_grid_is_chunked_over_sites(synthetic_surface: Path) -> None: + grid = _build(_config(synthetic_surface, chunks={Coordinate.I: 2})) + assert grid.x.chunksizes[Coordinate.I] == (2, 1) + # The pipeline relies on k staying in one piece. + assert len(grid.x.chunksizes[Coordinate.K]) == 1 + + +# --------------------------------------------------------------------------- +# Config decoding +# --------------------------------------------------------------------------- + +_TOML = """ +[grid] +type = "borehole" +surface = "{surface}" +depth = 200.0 +resolution_z = 50.0 + +[grid.projection] +crs = 'EPSG:2193' + +[[grid.sites]] +longitude = 172.15 +latitude = -43.70 +site = "GULL" +network = "NZ" + +[[layers]] +type = "query" +model_path = "{surface}" +""" + + +def test_toml_config_selects_the_borehole_grid( + synthetic_surface: Path, tmp_path: Path +) -> None: + path = tmp_path / "borehole.toml" + path.write_text(_TOML.format(surface=synthetic_surface)) + + config = VelocityModelConfig.read_config(path) + assert isinstance(config.grid, BoreholeGridConfig) + assert config.grid.sites == [ + Site( + longitude=172.15, + latitude=-43.70, + labels={SITE: "GULL", NETWORK: "NZ"}, + ) + ] + # Unset, so it falls back to WGS84. + assert config.grid.sites_crs.to_epsg() == 4326 + assert _build(config.grid).sizes[Coordinate.K] == 5 + + +def test_toml_config_reads_sites_from_a_file( + synthetic_surface: Path, tmp_path: Path +) -> None: + sites = _write_sites(tmp_path / "sites.csv") + body = _TOML.format(surface=synthetic_surface) + body = body[: body.index("[[grid.sites]]")] + body[body.index("[[layers]]") :] + path = tmp_path / "borehole.toml" + path.write_text( + body.replace("resolution_z = 50.0", f'resolution_z = 50.0\nsites = "{sites}"') + ) + + config = VelocityModelConfig.read_config(path) + assert isinstance(config.grid, BoreholeGridConfig) + assert config.grid.sites == sites + assert list(_build(config.grid)[SITE].values) == NAMES + + +# --------------------------------------------------------------------------- +# Through the pipeline +# --------------------------------------------------------------------------- + + +def _uniform(grid: Grid, model_range: ModelRange = ModelRange.ALL) -> Qualities: + """A terminal layer that keeps the grid's coordinates. + + `ConstantLayer` builds its output from raw NumPy and so drops them, which + `map_blocks` rejects for any grid that carries dimension coordinates. + """ + ones = xr.ones_like(grid.x) + return QualitiesSchema.new( + rho=ones * 2700.0, + vp=ones * 6000.0, + vs=ones * 1234.0, + qp=ones * 200.0, + qs=ones * 100.0, + alpha=ones, + ) + + +def _run(grid: Grid) -> VelocityModel: + model = VelocityModel(grids={"boreholes": grid}, metadata=ModelMetadata()) + return execute_model_pipeline(model, _uniform) + + +def test_grid_survives_the_chunked_pipeline(synthetic_surface: Path) -> None: + """`execute_model_pipeline` maps over the chunks, so `map_blocks` has to + preserve both the singleton j axis and the site labels.""" + grid = _build(_config(synthetic_surface, chunks={Coordinate.I: 2})) + result = _run(grid) + + qualities = result.qualities["boreholes"] + assert qualities.vs.shape == grid.x.shape + assert float(qualities.vs.values.mean()) == pytest.approx(1234.0, rel=1e-4) + assert list(result.grids["boreholes"][SITE].values) == NAMES + + +def test_output_round_trips_through_zarr( + synthetic_surface: Path, tmp_path: Path +) -> None: + """A profile is only useful when a reader can pick out one station.""" + path = tmp_path / "boreholes.zarr" + write_velocity_model( + _run(_build(_config(synthetic_surface))), + path, + Format.ZARR, + quantise_arrays=False, + ) + + with xr.open_datatree(path, engine="zarr") as tree: + stored = tree["grids/boreholes"].ds + assert list(stored[SITE].values) == NAMES + gull = stored.set_xindex(SITE).sel({SITE: "GULL"}) + assert float(gull.depth.max()) == pytest.approx(400.0) + assert tree["qualities/boreholes"].ds.vs.shape == stored.x.shape From cee8991a9f6a04b0f66a73a45e75d51c1154c23f Mon Sep 17 00:00:00 2001 From: Jake Faulkner Date: Mon, 21 Sep 2026 09:25:21 +1200 Subject: [PATCH 2/3] Simplify the borehole tests and share the depth-axis builder The borehole grid builder duplicated the regular grid builder: same nk, the same linspace depth axis pinned to one chunk, the same broadcast and rechunk. Hoist it into helpers.topography_following_grid and call it from both, so the k-chunking invariant the pipeline relies on lives in one place. constant() built its output from raw NumPy, which dropped the grid's coordinates and made map_blocks reject any grid carrying them. Build from ones_like instead, so a labelled grid works with the shipped layer and the borehole tests no longer need their own terminal layer. Cut the borehole tests from 558 lines to 443: parametrise the site reading and rejection cases, build the default grid once in a fixture rather than twelve times, and merge the tests that were halves of one claim. Assert the contracts rather than the implementation, so the tests stop pinning the exact chunk tuple, restating the builder's own linspace, and re-deriving RESERVED_COORDINATES from its definition. Trim the prose that said the same thing in the README, the module docstrings and the example configs. Co-Authored-By: Claude Opus 5 --- Justfile | 14 +- README.md | 34 +-- examples/borehole.toml | 14 +- examples/synthetic.toml | 5 +- nzcvm/grids/borehole.py | 79 ++----- nzcvm/grids/helpers.py | 46 ++++ nzcvm/grids/regular.py | 49 +--- nzcvm/layers/dummy.py | 7 +- nzcvm/scripts/synthetic.py | 10 +- nzcvm/synthetic.py | 37 +-- tests/test_borehole.py | 456 +++++++++++++++---------------------- tests/test_synthetic.py | 40 ++-- 12 files changed, 315 insertions(+), 476 deletions(-) diff --git a/Justfile b/Justfile index 0cf7e26..52c1ced 100644 --- a/Justfile +++ b/Justfile @@ -215,32 +215,32 @@ lint: ty ruff clippy # --------------------------------------------------------------------------- synthetic_root := "synthetic" -synthetic := "uv run nzcvm synthetic" +synthetic_cli := "uv run nzcvm synthetic" convert_surface := "uv run nzcvm surface convert" # Every input `examples/synthetic.toml` reads. synthetic: synthetic_dem synthetic_vs30 synthetic_coastline synthetic_models synthetic_dem: - {{ synthetic }} dem {{ synthetic_root }}/dem.h5 + {{ synthetic_cli }} dem {{ synthetic_root }}/dem.h5 {{ convert_surface }} {{ synthetic_root }}/dem.h5 {{ synthetic_root }}/dem.zarr synthetic_vs30: - {{ synthetic }} vs30 {{ synthetic_root }}/vs30.h5 + {{ synthetic_cli }} vs30 {{ synthetic_root }}/vs30.h5 {{ convert_surface }} {{ synthetic_root }}/vs30.h5 {{ synthetic_root }}/vs30.zarr --scalar-key vs30 --no-flip synthetic_coastline: - {{ synthetic }} coastline {{ synthetic_root }}/coastline.wkb.gz + {{ synthetic_cli }} coastline {{ synthetic_root }}/coastline.wkb.gz synthetic_tomography: - {{ synthetic }} tomography {{ synthetic_root }}/tomography.csv + {{ synthetic_cli }} tomography {{ synthetic_root }}/tomography.csv {{ tomography }} {{ synthetic_root }}/tomography.csv {{ synthetic_root }}/models/tomography.zarr # Priorities mirror the real basin recipes: lower wins over the tomography. synthetic_basins: synthetic_dem - {{ synthetic }} profile {{ synthetic_root }}/profile.fd_modfile + {{ synthetic_cli }} profile {{ synthetic_root }}/profile.fd_modfile @for basin in gully terrace; do \ - {{ synthetic }} basin $basin {{ synthetic_root }}; \ + {{ synthetic_cli }} basin $basin {{ synthetic_root }}; \ {{ construct }} {{ synthetic_root }}/${basin}_outline.geojson {{ synthetic_root }}/dem.h5 {{ synthetic_root }}/dem.h5 {{ synthetic_root }}/${basin}_basement.h5 {{ synthetic_root }}/models/${basin}.zarr --priority 1 --vm-1d {{ synthetic_root }}/profile.fd_modfile; \ done diff --git a/README.md b/README.md index d9e0e1a..92a2c0f 100644 --- a/README.md +++ b/README.md @@ -187,22 +187,16 @@ a key of that table. Longitude and latitude place a site, and the config reserves nothing else. Every other key, and every other column of a site file, becomes a coordinate on the grid's `i` axis under the name the caller gave it, which the writers -keep. Neither `site` nor `network` in the preceding example is a -keyword the grid interprets. Both end up in the output because nothing -reserves them. Rename them, add a driller's reference, drop them entirely: the -grid doesn't care. +keep. In the preceding example, `site` and `network` are labels rather than +keywords. A label keeps the type the caller wrote, so a numeric column arrives numeric. Each site needs the same set of labels, since the alternative is a column of nulls where one site was missing a key. -`nzcvm.grids.grid.RESERVED_COORDINATES` lists the names a label may not take: -the variables and attributes `GridSchema` declares (`x`, `y`, `z`, `depth`, -`name`, `geometry`, …), the `(i, j, k)` index, and the components that share -that index with the grid. A coordinate shadows a variable of the -same name, so -a label called `name` would turn `grid.name` from the grid's name into an -array. Naming one of those raises rather than corrupting the grid. +A label may not take a name in `nzcvm.grids.grid.RESERVED_COORDINATES`, which +covers everything `GridSchema` declares, the `(i, j, k)` index, and the +component names. Naming one raises rather than corrupting the grid. To drop the labels and keep only the spatial coordinates: @@ -428,10 +422,7 @@ takes about a minute, most of it meshing the basins, and needs no Every field is a closed-form expression over a 48 × 44 km patch of coast, so the data regenerates identically anywhere and a test can work out the expected -answer by hand. `nzcvm/synthetic.py` documents the world. Elevation falls west -to east and crosses sea level three quarters of the way across, and Vs30 -tracks elevation. A pair of paraboloid basins sit inland, over a tomography -block whose velocity increases with depth. +answer by hand. `nzcvm/synthetic.py` documents the world it describes. The pieces are importable, so a test can use a field directly rather than going through a file: @@ -594,13 +585,10 @@ contract every layer relies on: `x`, `y`, `z` and `depth` on the logical `(i, j, k)` index (metres, projected CRS, `z` positive down), plus the attributes below. -Those four variables and the attributes are the whole of what `GridSchema` -accepts, so a builder can't pass an extra *data variable*. It can attach extra -*coordinates* after construction, though, and every stage keeps them: layers, -`map_blocks`, the Zarr and NetCDF writers, and the read back through -`GridSchema.from_dataset`. xarray keeps a coordinate on each data variable it -indexes, so anything that reassembles a dataset from those variables picks it -up again. The borehole grid labels its columns this way: +`GridSchema` fixes the data variables, but a builder can attach extra +*coordinates* after construction, and every stage keeps them: layers, +`map_blocks`, and the Zarr and NetCDF writers. The borehole grid labels its +columns this way: ```python grid = grid.assign_coords(site=("i", ["GULL", "TERR"])) @@ -722,6 +710,6 @@ bottom = 500.0 dz = 100.0 ``` -`nzcvm.grids.borehole` is the same pattern, done properly: a registered +`nzcvm.grids.borehole` is the built-in version of this pattern: a registered `GridConfig` whose builder reads its own sites and DEM, over Dask-backed coordinates. diff --git a/examples/borehole.toml b/examples/borehole.toml index c5b8e22..f2af325 100644 --- a/examples/borehole.toml +++ b/examples/borehole.toml @@ -5,12 +5,8 @@ # # A borehole grid extracts one vertical profile per site instead of filling a # volume, so it has no extent, no azimuth and no model origin: just a -# projection to extract the profiles in. -# -# Longitude and latitude place a site. Every other key below (`site`, -# `network`) is one the grid knows nothing about, and rides through to the -# output as a coordinate on the i axis. Set keep_extra_columns = false to drop -# them. +# projection. Longitude and latitude place a site; every other key is a label +# that rides through to the output. See the README for the details. [metadata] title = "Synthetic borehole profiles" @@ -25,10 +21,8 @@ surface = "./synthetic/dem.zarr" depth = 600.0 resolution_z = 25.0 -# Sites can also come from a CSV or Parquet file with longitude and latitude -# columns, where the other columns label the sites the same way. It has to -# stay above the tables below, as TOML would otherwise read it as a key of -# `[grid.projection]`. +# Or read sites from a CSV/Parquet file. Keep this above the tables below, or +# TOML reads it as a key of [grid.projection]. # # sites = "examples/sites.csv" diff --git a/examples/synthetic.toml b/examples/synthetic.toml index cacddd2..44053c1 100644 --- a/examples/synthetic.toml +++ b/examples/synthetic.toml @@ -5,9 +5,8 @@ # just synthetic # uv run nzcvm generate examples/synthetic.toml synthetic/model.zarr # -# Everything it reads comes from `nzcvm synthetic`, so no NZCVM_DATA_ROOT and -# no model download is needed. The domain sits over the synthetic coast, with -# the `gully` and `terrace` basins inside it. +# The domain sits over the synthetic coast, with the `gully` and `terrace` +# basins inside it. [metadata] title = "Synthetic domain for local testing" diff --git a/nzcvm/grids/borehole.py b/nzcvm/grids/borehole.py index bf9c041..a8826ce 100644 --- a/nzcvm/grids/borehole.py +++ b/nzcvm/grids/borehole.py @@ -1,28 +1,17 @@ """Topography-following borehole grid builder. Provides :func:`build_borehole` for constructing the set of vertical columns -described by a :class:`~nzcvm.config.grids.borehole.BoreholeGridConfig`. -Each site becomes one column of query points, sampled at a strictly fixed Z -resolution from the topography down to a fixed depth, so a run extracts a -directly comparable profile per site rather than a filled volume. - -Columns are independent, so none of the extent, azimuth or origin metadata -the volumetric grids record means anything here. The builder still fills in -the :class:`~nzcvm.grids.grid.Grid` attributes that name an origin, taking the -centroid of the sites as the origin and the south-west corner of their -bounding box as the bottom-left corner. - -Longitude and latitude place a site. The config doesn't reserve any other -key, so every other key or column becomes a coordinate on the ``i`` axis under -the name the caller gave it, which is how a station code or a network ends up -in the layer chain and the output. -:data:`~nzcvm.grids.grid.RESERVED_COORDINATES` lists the names a label may not -take, and ``keep_extra_columns = false`` drops the labels altogether. +described by a :class:`~nzcvm.config.grids.borehole.BoreholeGridConfig`: one +column of query points per site, sampled at a fixed Z resolution from the +topography down to a fixed depth, so a run extracts a directly comparable +profile per site rather than a filled volume. + +See :class:`~nzcvm.config.grids.borehole.BoreholeGridConfig` for what a site +is and how the builder copies its labels onto the grid. """ from collections.abc import Callable from pathlib import Path -from typing import Any import dask.array as da import numpy as np @@ -34,7 +23,7 @@ from nzcvm.coordinates import Coordinate from nzcvm.grids import helpers from nzcvm.grids.builder import build_grids_from_config -from nzcvm.grids.grid import RESERVED_COORDINATES, Grid, GridSchema +from nzcvm.grids.grid import RESERVED_COORDINATES, Grid from nzcvm.models.surface import Surface #: Name of the one grid a borehole config builds. @@ -110,12 +99,12 @@ def site_labels(sites: list[Site]) -> dict[str, np.ndarray]: Raises ------ ValueError - If a label collides with - :data:`~nzcvm.grids.grid.RESERVED_COORDINATES`, or if the sites do - not agree on which labels they carry. Disagreement is nearly always - a typo, and the alternative is a column of nulls. + If a label shadows a grid name, or the sites disagree on which labels they + carry. Disagreement is nearly always a typo, and the alternative is + a column of nulls. """ names = list(sites[0].labels) + expected = set(names) reserved = sorted(RESERVED_COORDINATES.intersection(names)) if reserved: @@ -126,7 +115,7 @@ def site_labels(sites: list[Site]) -> dict[str, np.ndarray]: ) for position, site in enumerate(sites): - if set(site.labels) != set(names): + if set(site.labels) != expected: raise ValueError( f"Site {position} carries labels " f"{sorted(site.labels) or 'none'}, but site 0 carries " @@ -154,35 +143,6 @@ def _columns(values: np.ndarray, index: np.ndarray, chunk: int) -> xr.DataArray: ) -def _borehole_grid( - x_phys: xr.DataArray, - y_phys: xr.DataArray, - surface: xr.DataArray, - depth: float, - resolution_z: float, - **kwargs: Any, -) -> Grid: - nk = np.round(depth / resolution_z).astype(int) + 1 - - # Depth is purely a function of k and resolution_z, identically for every - # column, which is what makes the profiles comparable between sites. - # Chunking only ever applies to i/j. k always stays one chunk. - zeta_depth = xr.DataArray( - np.linspace(0.0, depth, num=nk, dtype=np.float32), - dims=[Coordinate.K], - coords={Coordinate.K: np.arange(nk)}, - ).chunk({Coordinate.K: -1}) - - # Elevation (z) is the surface elevation shifted downward by the fixed - # depths, so every column starts at the ground rather than at sea level. - x, y, z, column_depth = xr.broadcast( - x_phys, y_phys, surface + zeta_depth, zeta_depth - ) - x, y, z, column_depth = helpers.ensure_chunks(x, y, z, column_depth) - - return GridSchema.new(x, y, z, column_depth, **kwargs) - - @build_grids_from_config.register def build_borehole(config: BoreholeGridConfig) -> dict[str, Grid]: sites = resolve_sites(config) @@ -201,9 +161,9 @@ def build_borehole(config: BoreholeGridConfig) -> dict[str, Grid]: # A borehole grid has no origin of its own, so stand one up from the sites # for the benefit of the Grid attributes every writer expects. - to_wgs84 = config.projection.to_wgs84 - origin_lon, origin_lat = to_wgs84.transform(x.mean(), y.mean()) - min_lon, min_lat = to_wgs84.transform(x.min(), y.min()) + (origin_lon, min_lon), (origin_lat, min_lat) = config.projection.to_wgs84.transform( + [x.mean(), x.min()], [y.mean(), y.min()] + ) topographic_surface = Surface.load(config.surface) z_surface = helpers.compute_surface_elevation( @@ -212,12 +172,12 @@ def build_borehole(config: BoreholeGridConfig) -> dict[str, Grid]: y_phys, ) - grid = _borehole_grid( + grid = helpers.topography_following_grid( x_phys, y_phys, z_surface, name=GRID_NAME, - depth=config.depth, + thickness=config.depth, resolution_z=config.resolution_z, resolution=config.resolution_z, geometry=shapely.MultiPoint(np.column_stack((x, y))), @@ -228,8 +188,7 @@ def build_borehole(config: BoreholeGridConfig) -> dict[str, Grid]: bottom_left_lon=min_lon, bottom_left_lat=min_lat, ) - # Labelling i is what makes the output readable: without it, the only - # route back to a station is the order the config listed it in. + # Site labels index i, so the output reads back per station. grid = grid.assign_coords( {name: (Coordinate.I, values) for name, values in labels.items()} ) diff --git a/nzcvm/grids/helpers.py b/nzcvm/grids/helpers.py index 6f0d0bb..b526d95 100644 --- a/nzcvm/grids/helpers.py +++ b/nzcvm/grids/helpers.py @@ -1,3 +1,5 @@ +from typing import Any + import dask.array as da import numpy as np import shapely @@ -5,6 +7,7 @@ from nzcvm import coordinates from nzcvm.coordinates import Affine, Coordinate +from nzcvm.grids.grid import Grid, GridSchema from nzcvm.models.surface import Surface @@ -53,6 +56,49 @@ def ensure_chunks(*dsets: xr.DataArray) -> list[xr.DataArray]: return [dset.chunk(target) for dset in dsets] +def topography_following_grid( + x_phys: xr.DataArray, + y_phys: xr.DataArray, + surface: xr.DataArray, + thickness: float, + resolution_z: float, + **kwargs: Any, +) -> Grid: + """Hang a fixed-resolution depth axis off *surface*. + + Depth is purely a function of k and *resolution_z*, identically for every + column, so every column runs from the topography down to *thickness* and + the bottom follows the topography exactly. Chunking only ever applies to + i/j, and k always stays one chunk. + + Parameters + ---------- + x_phys, y_phys : + Projected horizontal coordinates, of any shape the depth axis can + broadcast against. + surface : + Surface elevation at those coordinates, positive down. + thickness : + Depth of the bottom of every column, in metres below the topography. + resolution_z : + Vertical sample spacing in metres. + **kwargs : + Passed through to :meth:`~nzcvm.grids.grid.GridSchema.new`. + """ + nk = np.round(thickness / resolution_z).astype(int) + 1 + zeta_depth = xr.DataArray( + np.linspace(0.0, thickness, num=nk, dtype=np.float32), + dims=[Coordinate.K], + coords={Coordinate.K: np.arange(nk)}, + ).chunk({Coordinate.K: -1}) + + # Broadcasting in this order is what fixes the (i, j, k) coordinate order. + x, y, z, depth = xr.broadcast(x_phys, y_phys, surface + zeta_depth, zeta_depth) + x, y, z, depth = ensure_chunks(x, y, z, depth) + + return GridSchema.new(x, y, z, depth, **kwargs) + + def outline(transform: Affine, extent_x: float, extent_y: float) -> shapely.Geometry: dx = extent_x / 2 dy = extent_y / 2 diff --git a/nzcvm/grids/regular.py b/nzcvm/grids/regular.py index 8964d2e..a618f2a 100644 --- a/nzcvm/grids/regular.py +++ b/nzcvm/grids/regular.py @@ -6,63 +6,18 @@ and topography-following ``z`` / ``depth`` arrays at a strictly fixed Z resolution. """ -from typing import Any - import dask import numpy as np -import xarray as xr from scipy.spatial.transform import Rotation from nzcvm import coordinates from nzcvm.config.grids.regular import RegularGridConfig -from nzcvm.coordinates import Coordinate from nzcvm.grids import helpers from nzcvm.grids.builder import build_grids_from_config -from nzcvm.grids.grid import Grid, GridSchema +from nzcvm.grids.grid import Grid from nzcvm.models.surface import Surface -def _regular_grid( - x_phys: xr.DataArray, - y_phys: xr.DataArray, - surface: xr.DataArray, - thickness: float, - resolution_z: float, - **kwargs: Any, -) -> Grid: - nk = np.round(thickness / resolution_z).astype(int) + 1 - k = np.arange(nk) - - # Depth is purely a function of k and resolution_z - depth_values = np.linspace(0.0, thickness, num=nk, dtype=np.float32) - - # Chunking only ever applies to i/j. k always stays one chunk. - zeta_depth = xr.DataArray( - depth_values, - dims=[Coordinate.K], - coords={Coordinate.K: k}, - ).chunk({Coordinate.K: -1}) - - # Elevation (z) is the surface elevation shifted downward by the fixed depths. - # The bottom then follows the topography exactly. - z = surface + zeta_depth - - # Same idiomatic trick as the SW4 template to ensure coordinate ordering (i, j, k) - depth = zeta_depth - - x, y, z, depth = xr.broadcast(x_phys, y_phys, z, depth) - - x, y, z, depth = helpers.ensure_chunks(x, y, z, depth) - - return GridSchema.new( - x, - y, - z, - depth, - **kwargs, - ) - - @build_grids_from_config.register def build_regular(config: RegularGridConfig) -> dict[str, Grid]: offset = 0.0 @@ -112,7 +67,7 @@ def build_regular(config: RegularGridConfig) -> dict[str, Grid]: y_phys, ) - grid = _regular_grid( + grid = helpers.topography_following_grid( x_phys, y_phys, z_surface, diff --git a/nzcvm/layers/dummy.py b/nzcvm/layers/dummy.py index 985ae63..33900bb 100644 --- a/nzcvm/layers/dummy.py +++ b/nzcvm/layers/dummy.py @@ -30,7 +30,7 @@ from dataclasses import dataclass from typing import Literal -import numpy as np +import xarray as xr from shapely import Geometry from nzcvm.config.layers.core import LayerConfig @@ -70,8 +70,9 @@ def constant( next_layer : Ignored, because a terminal layer never delegates downstream. """ - shape = grid.x.shape - ones = np.ones(shape, dtype=np.float32) + # ones_like keeps the grid's coordinates, which map_blocks requires of any + # grid carrying dimension coordinates (a borehole grid's site labels, say). + ones = xr.ones_like(grid.x) return QualitiesSchema.new( rho=ones * rho, vp=ones * vp, diff --git a/nzcvm/scripts/synthetic.py b/nzcvm/scripts/synthetic.py index fa1f165..d2378a2 100644 --- a/nzcvm/scripts/synthetic.py +++ b/nzcvm/scripts/synthetic.py @@ -10,6 +10,7 @@ surface mesh that ``grid.surface`` and the Ely layer read. """ +import functools import gzip from pathlib import Path from typing import Annotated @@ -30,7 +31,12 @@ #: Default surface sampling, about 750 m across the domain. DEFAULT_SAMPLES = 64 -TO_NZTM = pyproj.Transformer.from_crs(WGS84_EPSG, CRS_NZTM, always_xy=True) + +@functools.cache +def _to_nztm() -> pyproj.Transformer: + """Only the coastline command needs this, so don't build it at import.""" + return pyproj.Transformer.from_crs(WGS84_EPSG, CRS_NZTM, always_xy=True) + Output = Annotated[Path, typer.Argument(help="Output path.", dir_okay=False)] Samples = Annotated[int, typer.Option(help="Surface samples along each axis.", min=2)] @@ -77,7 +83,7 @@ def write_wgs84_polygon(path: Path, vertices: np.ndarray) -> shapely.Polygon: shapely.Polygon The projected polygon this wrote. """ - projected = shapely.ops.transform(TO_NZTM.transform, shapely.Polygon(vertices)) + projected = shapely.ops.transform(_to_nztm().transform, shapely.Polygon(vertices)) path.parent.mkdir(parents=True, exist_ok=True) with gzip.open(path, "wb") as handle: handle.write(shapely.to_wkb(projected)) diff --git a/nzcvm/synthetic.py b/nzcvm/synthetic.py index 07bd44c..64b02e2 100644 --- a/nzcvm/synthetic.py +++ b/nzcvm/synthetic.py @@ -88,17 +88,8 @@ def normalise( ) -> tuple[np.ndarray, np.ndarray]: """Map longitude and latitude onto the unit square. - Parameters - ---------- - lon, lat : - Geographic coordinates in degrees. - - Returns - ------- - tuple[numpy.ndarray, numpy.ndarray] - ``(u, v)``, zero at the south-west corner and one at the - north-east corner. Points outside the domain fall outside - ``[0, 1]`` rather than clipping to it. + Zero at the south-west corner, one at the north-east. Points outside + the domain fall outside ``[0, 1]`` rather than clipping to it. Examples -------- @@ -127,18 +118,8 @@ def denormalise( return lon, lat def sample(self, n_lon: int, n_lat: int) -> tuple[np.ndarray, np.ndarray]: - """Return 1-D longitude and latitude axes spanning the domain. - - Parameters - ---------- - n_lon, n_lat : - Number of samples along each axis. - - Returns - ------- - tuple[numpy.ndarray, numpy.ndarray] - Ascending ``(longitude, latitude)`` axes, endpoints included. - """ + """Return ascending 1-D longitude and latitude axes spanning the + domain, endpoints included.""" return ( np.linspace(self.lon_min, self.lon_max, n_lon), np.linspace(self.lat_min, self.lat_max, n_lat), @@ -171,7 +152,7 @@ def elevation( Returns ------- numpy.ndarray - Elevation in metres, positive up. Negative offshore. + Negative offshore. """ u, v = domain.normalise(lon, lat) ridges = RIDGE_AMPLITUDE * np.sin(RIDGE_PERIODS * np.pi * v) @@ -344,11 +325,6 @@ class BasinName(StrEnum): ) -def _model_frame() -> Affine: - """The EP2020 affine mapping NZTM metres onto tomography model kilometres.""" - return MODEL_COLUMNS[ModelType.EP2020].affine_inverse - - def _apply(transform: Affine, points: np.ndarray) -> np.ndarray: """Apply a 4×4 homogeneous affine to an ``(N, 3)`` array of points.""" homogeneous = np.column_stack((points, np.ones(len(points)))) @@ -392,7 +368,8 @@ def tomography( to_nztm = Transformer.from_crs(CRS_WGS, CRS_NZTM, always_xy=True) to_wgs = Transformer.from_crs(CRS_NZTM, CRS_WGS, always_xy=True) - inverse = _model_frame() + # The EP2020 affine maps NZTM metres onto tomography model kilometres. + inverse = MODEL_COLUMNS[ModelType.EP2020].affine_inverse forward = np.linalg.inv(inverse) corners = np.asarray(domain.polygon.exterior.coords) diff --git a/tests/test_borehole.py b/tests/test_borehole.py index ae7d98e..0a4d5fb 100644 --- a/tests/test_borehole.py +++ b/tests/test_borehole.py @@ -1,13 +1,13 @@ """Tests for the borehole grid: config decoding, site loading and the builder. -Every test builds on the synthetic DEM (:mod:`nzcvm.synthetic`), so nothing -reads a real data file, and the expected elevations follow from the analytic -topography rather than from a fixture. +Every test builds on the synthetic DEM (:mod:`nzcvm.synthetic`), so the +expected elevations follow from the analytic topography rather than a fixture. """ from __future__ import annotations from pathlib import Path +from typing import Any import numpy as np import pandas as pd @@ -17,43 +17,43 @@ from pyproj import CRS, Transformer from nzcvm import synthetic -from nzcvm.components import Component -from nzcvm.config.grids.borehole import ( - DEFAULT_CHUNK_SIZES, - BoreholeGridConfig, - Site, -) -from nzcvm.config.grids.model import Model, Projection +from nzcvm.config.grids.borehole import BoreholeGridConfig, Site +from nzcvm.config.grids.model import Projection from nzcvm.config.metadata import ModelMetadata from nzcvm.config.velocity_model import VelocityModelConfig from nzcvm.coordinates import Coordinate from nzcvm.formats import Format, write_velocity_model from nzcvm.grids.borehole import read_sites, resolve_sites from nzcvm.grids.builder import build_grids_from_config -from nzcvm.grids.grid import RESERVED_COORDINATES, Grid +from nzcvm.grids.grid import Grid +from nzcvm.layers.dummy import constant from nzcvm.layers.pipeline import execute_model_pipeline from nzcvm.models.mesh import StructuredMeshSchema -from nzcvm.qualities import Qualities, QualitiesSchema -from nzcvm.query import ModelRange from nzcvm.velocity_model import VelocityModel _NZTM = CRS.from_epsg(2193) _NZGD2000 = CRS.from_epsg(4167) -_WGS84 = CRS.from_epsg(4326) _TO_NZTM = Transformer.from_crs(4326, _NZTM, always_xy=True) -# Sites in the hills and out to sea, spread across the domain. `site` and -# `network` are ordinary labels, not keywords the grid interprets. +# Sites in the hills and out to sea, spread across the domain. SITES = [ Site(longitude=172.15, latitude=-43.70, labels={"site": "GULL", "network": "NZ"}), Site(longitude=172.10, latitude=-43.45, labels={"site": "RIDG", "network": "NZ"}), Site(longitude=172.55, latitude=-43.60, labels={"site": "SEAB", "network": "SC"}), ] - -#: The labels on the sites, both treated like any other. -SITE = "site" -NETWORK = "network" -NAMES = [site.labels[SITE] for site in SITES] +NAMES = [site.labels["site"] for site in SITES] + +DEPTH = 400.0 +RESOLUTION_Z = 100.0 +NK = int(DEPTH / RESOLUTION_Z) + 1 + +# Sites that disagree on their labels. +MISMATCHED = [ + Site( + longitude=SITES[0].longitude, latitude=SITES[0].latitude, labels={"site": "A"} + ), + Site(longitude=SITES[1].longitude, latitude=SITES[1].latitude, labels={}), +] @pytest.fixture(scope="module") @@ -77,25 +77,15 @@ def synthetic_surface(tmp_path_factory: pytest.TempPathFactory) -> Path: return path -def _config( - surface: Path, - sites: list[Site] | Path = SITES, - depth: float = 400.0, - resolution_z: float = 100.0, - sites_crs: CRS = _WGS84, - keep_extra_columns: bool = True, - chunks: dict[Coordinate, int] = DEFAULT_CHUNK_SIZES, -) -> BoreholeGridConfig: - return BoreholeGridConfig( - surface=surface, - sites=sites, - depth=depth, - resolution_z=resolution_z, - projection=Projection(crs=_NZTM), - sites_crs=sites_crs, - keep_extra_columns=keep_extra_columns, - chunks=chunks, - ) +def _config(surface: Path, **overrides: Any) -> BoreholeGridConfig: + defaults: dict[str, Any] = { + "surface": surface, + "sites": SITES, + "depth": DEPTH, + "resolution_z": RESOLUTION_Z, + "projection": Projection(crs=_NZTM), + } + return BoreholeGridConfig(**(defaults | overrides)) def _build(config: BoreholeGridConfig) -> Grid: @@ -104,32 +94,18 @@ def _build(config: BoreholeGridConfig) -> Grid: return grids["boreholes"] -# --------------------------------------------------------------------------- -# Projection / Model split -# --------------------------------------------------------------------------- - - -def test_projection_needs_no_origin() -> None: - """The whole point of splitting Projection out of Model.""" - projection = Projection(crs=_NZTM) - x, y = projection.from_wgs84.transform(172.0, -43.5) - assert projection.to_wgs84.transform(x, y) == pytest.approx((172.0, -43.5)) - - -def test_model_is_still_a_projection() -> None: - model = Model(origin_lon=172.0, origin_lat=-43.5, azimuth=39.0, crs=_NZTM) - assert isinstance(model, Projection) - assert model.grid_origin_x == pytest.approx( - model.from_wgs84.transform(172.0, -43.5)[0] - ) +def _relabel(*labels: dict[str, Any]) -> list[Site]: + """The default sites, carrying *labels* instead of their own.""" + return [ + Site(longitude=site.longitude, latitude=site.latitude, labels=label) + for site, label in zip(SITES, labels, strict=True) + ] -def test_transformer_from_reaches_the_projection() -> None: - projection = Projection(crs=_NZTM) - from_nzgd = projection.transformer_from(_NZGD2000) - assert from_nzgd.transform(172.0, -43.5) == pytest.approx( - projection.from_wgs84.transform(172.0, -43.5), abs=1.0 - ) +@pytest.fixture(scope="module") +def grid(synthetic_surface: Path) -> Grid: + """The default grid, built once. No test mutates it.""" + return _build(_config(synthetic_surface)) # --------------------------------------------------------------------------- @@ -142,8 +118,8 @@ def _write_sites(path: Path) -> Path: { "longitude": [site.longitude for site in SITES], "latitude": [site.latitude for site in SITES], - SITE: NAMES, - NETWORK: [site.labels[NETWORK] for site in SITES], + "site": NAMES, + "network": [site.labels["network"] for site in SITES], } ) if path.suffix == ".csv": @@ -158,52 +134,64 @@ def test_read_sites_round_trips(tmp_path: Path, suffix: str) -> None: assert read_sites(_write_sites(tmp_path / f"sites{suffix}")) == SITES -def test_read_sites_preserves_file_order(tmp_path: Path) -> None: - path = tmp_path / "sites.csv" - pd.DataFrame( - {"longitude": [172.1, 172.2], "latitude": [-43.5, -43.6], SITE: ["B", "A"]} - ).to_csv(path, index=False) - assert [site.labels[SITE] for site in read_sites(path)] == ["B", "A"] - - -def test_read_sites_keeps_extra_columns_as_labels(tmp_path: Path) -> None: - """Longitude and latitude are the only reserved columns. The rest - describe the site.""" - path = tmp_path / "sites.csv" - pd.DataFrame( - { - "longitude": [172.1], - "latitude": [-43.5], - SITE: ["A"], - "elevation": [12.0], - } - ).to_csv(path, index=False) - assert read_sites(path) == [ - Site( - longitude=172.1, - latitude=-43.5, - labels={SITE: "A", "elevation": 12.0}, - ) - ] - - -def test_read_sites_accepts_a_file_with_no_labels(tmp_path: Path) -> None: +@pytest.mark.parametrize( + ("columns", "expected"), + [ + pytest.param( + { + "longitude": [172.1, 172.2], + "latitude": [-43.5, -43.6], + "site": ["B", "A"], + }, + [ + Site(longitude=172.1, latitude=-43.5, labels={"site": "B"}), + Site(longitude=172.2, latitude=-43.6, labels={"site": "A"}), + ], + id="file-order-is-kept", + ), + pytest.param( + { + "longitude": [172.1], + "latitude": [-43.5], + "site": ["A"], + "elevation": [12.0], + }, + [ + Site( + longitude=172.1, + latitude=-43.5, + labels={"site": "A", "elevation": 12.0}, + ) + ], + id="extra-columns-become-labels", + ), + pytest.param( + {"longitude": [172.1], "latitude": [-43.5]}, + [Site(longitude=172.1, latitude=-43.5, labels={})], + id="labels-are-optional", + ), + ], +) +def test_read_sites(tmp_path: Path, columns: dict, expected: list[Site]) -> None: + """Longitude and latitude are the only reserved columns.""" path = tmp_path / "sites.csv" - pd.DataFrame({"longitude": [172.1], "latitude": [-43.5]}).to_csv(path, index=False) - assert read_sites(path) == [Site(longitude=172.1, latitude=-43.5, labels={})] - - -def test_read_sites_rejects_unknown_format(tmp_path: Path) -> None: - path = tmp_path / "sites.txt" - path.write_text("name,longitude,latitude\n") - with pytest.raises(ValueError, match="expected one of"): - read_sites(path) + pd.DataFrame(columns).to_csv(path, index=False) + assert read_sites(path) == expected -def test_read_sites_reports_missing_columns(tmp_path: Path) -> None: - path = tmp_path / "sites.csv" - pd.DataFrame({SITE: ["A"], "longitude": [172.1]}).to_csv(path, index=False) - with pytest.raises(ValueError, match="missing the latitude column"): +@pytest.mark.parametrize( + ("name", "text", "message"), + [ + ("sites.txt", "site,longitude,latitude\n", "expected one of"), + ("sites.csv", "site,longitude\nA,172.1\n", "missing the latitude"), + ], +) +def test_read_sites_rejects_bad_files( + tmp_path: Path, name: str, text: str, message: str +) -> None: + path = tmp_path / name + path.write_text(text) + with pytest.raises(ValueError, match=message): read_sites(path) @@ -216,6 +204,22 @@ def test_resolve_sites_reads_a_file(synthetic_surface: Path, tmp_path: Path) -> assert resolve_sites(config) == SITES +@pytest.mark.parametrize( + ("override", "error", "message"), + [ + pytest.param({"sites": []}, ValueError, "at least one site", id="empty-list"), + pytest.param( + {"sites_crs": _NZTM}, InvalidFieldValue, "geographic", id="projected-crs" + ), + ], +) +def test_config_rejects_bad_sites( + synthetic_surface: Path, override: dict, error: type[Exception], message: str +) -> None: + with pytest.raises(error, match=message): + _config(synthetic_surface, **override) + + def test_resolve_sites_rejects_an_empty_file( synthetic_surface: Path, tmp_path: Path ) -> None: @@ -225,16 +229,6 @@ def test_resolve_sites_rejects_an_empty_file( resolve_sites(_config(synthetic_surface, sites=path)) -def test_config_rejects_an_empty_inline_site_list(synthetic_surface: Path) -> None: - with pytest.raises(ValueError, match="at least one site"): - _config(synthetic_surface, sites=[]) - - -def test_config_rejects_a_projected_sites_crs(synthetic_surface: Path) -> None: - with pytest.raises(InvalidFieldValue, match="geographic"): - _config(synthetic_surface, sites_crs=_NZTM) - - # --------------------------------------------------------------------------- # Site labels # @@ -244,39 +238,31 @@ def test_config_rejects_a_projected_sites_crs(synthetic_surface: Path) -> None: def test_extra_config_keys_become_labels() -> None: - site = Site.from_dict( - {"longitude": 172.1, "latitude": -43.5, SITE: "A", "depth_drilled": 30.0} - ) - assert site.labels == {SITE: "A", "depth_drilled": 30.0} - - -def test_a_site_needs_no_labels_at_all() -> None: + assert Site.from_dict( + {"longitude": 172.1, "latitude": -43.5, "site": "A", "depth_drilled": 30.0} + ).labels == {"site": "A", "depth_drilled": 30.0} assert Site.from_dict({"longitude": 172.1, "latitude": -43.5}).labels == {} -def test_labels_become_coordinates_on_the_site_axis(synthetic_surface: Path) -> None: - grid = _build(_config(synthetic_surface)) - for label in (SITE, NETWORK): +def test_labels_become_coordinates_on_the_site_axis(grid: Grid) -> None: + for label in ("site", "network"): assert grid[label].dims == (Coordinate.I,) assert list(grid[label].values) == [site.labels[label] for site in SITES] def test_numeric_labels_keep_a_numeric_dtype(synthetic_surface: Path) -> None: """A label is whatever the caller wrote, not necessarily a string.""" - sites = [ - Site(longitude=site.longitude, latitude=site.latitude, labels={"cased": n * 10}) - for n, site in enumerate(SITES) - ] - grid = _build(_config(synthetic_surface, sites=sites)) - assert np.issubdtype(grid["cased"].dtype, np.integer) - assert list(grid["cased"].values) == [0, 10, 20] + sites = _relabel(*({"cased": n * 10} for n in range(len(SITES)))) + labelled = _build(_config(synthetic_surface, sites=sites)) + assert labelled["cased"].dtype.kind in "iuf" + assert list(labelled["cased"].values) == [0, 10, 20] def test_keep_extra_columns_false_drops_the_labels(synthetic_surface: Path) -> None: - grid = _build(_config(synthetic_surface, keep_extra_columns=False)) - assert set(grid.coords) == {Coordinate.I, Coordinate.J, Coordinate.K} + bare = _build(_config(synthetic_surface, keep_extra_columns=False)) + assert set(bare.coords) == {Coordinate.I, Coordinate.J, Coordinate.K} # The opt-out leaves the spatial coordinates alone. - assert grid.x.shape == (len(SITES), 1, 5) + assert bare.x.shape == (len(SITES), 1, NK) @pytest.mark.parametrize("reserved", ["name", "x", "depth", "vs", "i", "geometry"]) @@ -285,48 +271,23 @@ def test_a_label_may_not_shadow_a_grid_name( ) -> None: """A coordinate shadows a variable or attribute of the same name, so `grid.name` would stop being the grid's name.""" - sites = [ - Site(longitude=site.longitude, latitude=site.latitude, labels={reserved: "x"}) - for site in SITES - ] + sites = _relabel(*([{reserved: "x"}] * len(SITES))) with pytest.raises(ValueError, match="would shadow"): _build(_config(synthetic_surface, sites=sites)) -def test_reserved_names_cover_the_grid_contract() -> None: - """Derived from GridSchema, so it keeps up if the schema changes.""" - assert {"x", "y", "z", "depth", "name", "geometry", "resolution"} <= ( - RESERVED_COORDINATES - ) - assert {Coordinate.I, Coordinate.J, Coordinate.K} <= RESERVED_COORDINATES - assert set(Component) <= RESERVED_COORDINATES - assert SITE not in RESERVED_COORDINATES - - def test_sites_have_to_agree_on_their_labels(synthetic_surface: Path) -> None: """A missing label is nearly always a typo, and the alternative is a column of nulls.""" - sites = [ - Site( - longitude=SITES[0].longitude, latitude=SITES[0].latitude, labels={SITE: "A"} - ), - Site(longitude=SITES[1].longitude, latitude=SITES[1].latitude, labels={}), - ] with pytest.raises(ValueError, match="Every site needs the same labels"): - _build(_config(synthetic_surface, sites=sites)) + _build(_config(synthetic_surface, sites=MISMATCHED)) def test_disagreement_is_allowed_once_labels_are_dropped( synthetic_surface: Path, ) -> None: - sites = [ - Site( - longitude=SITES[0].longitude, latitude=SITES[0].latitude, labels={SITE: "A"} - ), - Site(longitude=SITES[1].longitude, latitude=SITES[1].latitude, labels={}), - ] - grid = _build(_config(synthetic_surface, sites=sites, keep_extra_columns=False)) - assert grid.x.shape == (2, 1, 5) + config = _config(synthetic_surface, sites=MISMATCHED, keep_extra_columns=False) + assert _build(config).x.shape == (len(MISMATCHED), 1, NK) # --------------------------------------------------------------------------- @@ -334,75 +295,56 @@ def test_disagreement_is_allowed_once_labels_are_dropped( # --------------------------------------------------------------------------- -def test_grid_is_one_column_per_site(synthetic_surface: Path) -> None: - grid = _build(_config(synthetic_surface)) - assert grid.x.shape == (len(SITES), 1, 5) +def test_grid_is_one_column_per_site(grid: Grid) -> None: + assert grid.x.shape == (len(SITES), 1, NK) assert grid.sizes[Coordinate.J] == 1 - - -def test_grid_labels_columns_by_site(synthetic_surface: Path) -> None: - grid = _build(_config(synthetic_surface)) - assert list(grid[SITE].values) == NAMES - assert list(grid[NETWORK].values) == [site.labels[NETWORK] for site in SITES] - - -def test_depth_is_identical_across_sites(synthetic_surface: Path) -> None: - """Comparable profiles are the point: every column shares one depth axis.""" - depth = _build(_config(synthetic_surface)).depth.values - expected = np.linspace(0.0, 400.0, 5, dtype=np.float32) - for column in depth.reshape(-1, depth.shape[-1]): - assert column == pytest.approx(expected) + assert list(grid["site"].values) == NAMES def test_resolution_z_sets_the_sample_count(synthetic_surface: Path) -> None: - grid = _build(_config(synthetic_surface, depth=600.0, resolution_z=25.0)) - assert grid.sizes[Coordinate.K] == 25 - assert float(grid.depth.max()) == pytest.approx(600.0) + finer = _build(_config(synthetic_surface, depth=600.0, resolution_z=25.0)) + assert finer.sizes[Coordinate.K] == 25 + assert float(finer.depth.max()) == pytest.approx(600.0) -def test_sites_are_projected_into_the_grid_crs(synthetic_surface: Path) -> None: - grid = _build(_config(synthetic_surface)) +def test_columns_sit_over_the_projected_site(grid: Grid) -> None: expected_x, expected_y = _TO_NZTM.transform( [site.longitude for site in SITES], [site.latitude for site in SITES] ) assert grid.x.values[:, 0, 0] == pytest.approx(np.float32(expected_x), rel=1e-6) assert grid.y.values[:, 0, 0] == pytest.approx(np.float32(expected_y), rel=1e-6) - - -def test_x_and_y_are_constant_down_each_column(synthetic_surface: Path) -> None: - grid = _build(_config(synthetic_surface)) + # x and y are constant down each column. for axis in (grid.x, grid.y): assert np.all(axis.values == axis.values[:, :, :1]) -def test_columns_start_at_the_topography(synthetic_surface: Path) -> None: - grid = _build(_config(synthetic_surface)) +def test_columns_run_from_the_topography_down_to_depth(grid: Grid) -> None: + """Comparable profiles are the point: every column shares one depth axis, + starts at the ground and ends `depth` below it.""" + depth = grid.depth.values + assert np.all(depth == depth[:1, :1, :]) + assert depth[0, 0, 0] == pytest.approx(0.0) + assert depth.max() == pytest.approx(DEPTH) + # Grids are +z down, so the top of each column is minus the elevation. expected = -synthetic.elevation( [site.longitude for site in SITES], [site.latitude for site in SITES] ) - assert grid.z.values[:, 0, 0] == pytest.approx(expected, abs=15.0) - - -def test_columns_follow_the_topography_down(synthetic_surface: Path) -> None: - grid = _build(_config(synthetic_surface)) top = grid.z.values[:, :, 0] - bottom = grid.z.values[:, :, -1] - assert (bottom - top) == pytest.approx(np.full_like(top, 400.0)) + assert top[:, 0] == pytest.approx(expected, abs=15.0) + assert (grid.z.values[:, :, -1] - top) == pytest.approx(np.full_like(top, DEPTH)) def test_a_geographic_sites_crs_lands_in_the_same_place( - synthetic_surface: Path, + synthetic_surface: Path, grid: Grid ) -> None: - wgs84 = _build(_config(synthetic_surface)) nzgd2000 = _build(_config(synthetic_surface, sites_crs=_NZGD2000)) # NZGD2000 and WGS84 are within a metre of each other over New Zealand. - assert nzgd2000.x.values == pytest.approx(wgs84.x.values, abs=1.0) + assert nzgd2000.x.values == pytest.approx(grid.x.values, abs=1.0) -def test_geometry_covers_every_site(synthetic_surface: Path) -> None: +def test_geometry_covers_every_site(grid: Grid) -> None: """The query layer prunes models against this, so it has to hit each site.""" - grid = _build(_config(synthetic_surface)) assert len(grid.geometry.geoms) == len(SITES) assert grid.geometry.bounds == pytest.approx( ( @@ -415,10 +357,9 @@ def test_geometry_covers_every_site(synthetic_surface: Path) -> None: ) -def test_derived_origin_is_the_site_centroid(synthetic_surface: Path) -> None: +def test_derived_origin_is_the_site_centroid(grid: Grid) -> None: """A borehole grid has no configured origin, so the attributes come from the sites. Downstream writers still expect them to be present.""" - grid = _build(_config(synthetic_surface)) lons = [site.longitude for site in SITES] lats = [site.latitude for site in SITES] assert grid.origin_lon == pytest.approx(np.mean(lons), abs=1e-2) @@ -430,10 +371,10 @@ def test_derived_origin_is_the_site_centroid(synthetic_surface: Path) -> None: def test_grid_is_chunked_over_sites(synthetic_surface: Path) -> None: - grid = _build(_config(synthetic_surface, chunks={Coordinate.I: 2})) - assert grid.x.chunksizes[Coordinate.I] == (2, 1) + chunked = _build(_config(synthetic_surface, chunks={Coordinate.I: 2})) + assert max(chunked.x.chunksizes[Coordinate.I]) <= 2 # The pipeline relies on k staying in one piece. - assert len(grid.x.chunksizes[Coordinate.K]) == 1 + assert len(chunked.x.chunksizes[Coordinate.K]) == 1 # --------------------------------------------------------------------------- @@ -446,57 +387,60 @@ def test_grid_is_chunked_over_sites(synthetic_surface: Path) -> None: surface = "{surface}" depth = 200.0 resolution_z = 50.0 +{sites} [grid.projection] crs = 'EPSG:2193' +[[layers]] +type = "query" +model_path = "{surface}" +""" + +_INLINE_SITES = """ [[grid.sites]] longitude = 172.15 latitude = -43.70 site = "GULL" network = "NZ" - -[[layers]] -type = "query" -model_path = "{surface}" """ +def _decode_grid(path: Path, surface: Path, sites: str) -> BoreholeGridConfig: + """The grid a written TOML config decodes to.""" + path.write_text(_TOML.format(surface=surface, sites=sites)) + grid = VelocityModelConfig.read_config(path).grid + assert isinstance(grid, BoreholeGridConfig) + return grid + + def test_toml_config_selects_the_borehole_grid( synthetic_surface: Path, tmp_path: Path ) -> None: - path = tmp_path / "borehole.toml" - path.write_text(_TOML.format(surface=synthetic_surface)) + config = _decode_grid(tmp_path / "borehole.toml", synthetic_surface, _INLINE_SITES) - config = VelocityModelConfig.read_config(path) - assert isinstance(config.grid, BoreholeGridConfig) - assert config.grid.sites == [ + assert config.sites == [ Site( longitude=172.15, latitude=-43.70, - labels={SITE: "GULL", NETWORK: "NZ"}, + labels={"site": "GULL", "network": "NZ"}, ) ] # Unset, so it falls back to WGS84. - assert config.grid.sites_crs.to_epsg() == 4326 - assert _build(config.grid).sizes[Coordinate.K] == 5 + assert config.sites_crs.to_epsg() == 4326 + assert _build(config).sizes[Coordinate.K] == 5 def test_toml_config_reads_sites_from_a_file( synthetic_surface: Path, tmp_path: Path ) -> None: sites = _write_sites(tmp_path / "sites.csv") - body = _TOML.format(surface=synthetic_surface) - body = body[: body.index("[[grid.sites]]")] + body[body.index("[[layers]]") :] - path = tmp_path / "borehole.toml" - path.write_text( - body.replace("resolution_z = 50.0", f'resolution_z = 50.0\nsites = "{sites}"') + config = _decode_grid( + tmp_path / "borehole.toml", synthetic_surface, f'sites = "{sites}"' ) - config = VelocityModelConfig.read_config(path) - assert isinstance(config.grid, BoreholeGridConfig) - assert config.grid.sites == sites - assert list(_build(config.grid)[SITE].values) == NAMES + assert config.sites == sites + assert list(_build(config)["site"].values) == NAMES # --------------------------------------------------------------------------- @@ -504,55 +448,31 @@ def test_toml_config_reads_sites_from_a_file( # --------------------------------------------------------------------------- -def _uniform(grid: Grid, model_range: ModelRange = ModelRange.ALL) -> Qualities: - """A terminal layer that keeps the grid's coordinates. - - `ConstantLayer` builds its output from raw NumPy and so drops them, which - `map_blocks` rejects for any grid that carries dimension coordinates. - """ - ones = xr.ones_like(grid.x) - return QualitiesSchema.new( - rho=ones * 2700.0, - vp=ones * 6000.0, - vs=ones * 1234.0, - qp=ones * 200.0, - qs=ones * 100.0, - alpha=ones, - ) - - def _run(grid: Grid) -> VelocityModel: model = VelocityModel(grids={"boreholes": grid}, metadata=ModelMetadata()) - return execute_model_pipeline(model, _uniform) + return execute_model_pipeline(model, constant(vs=1234.0)) def test_grid_survives_the_chunked_pipeline(synthetic_surface: Path) -> None: """`execute_model_pipeline` maps over the chunks, so `map_blocks` has to preserve both the singleton j axis and the site labels.""" - grid = _build(_config(synthetic_surface, chunks={Coordinate.I: 2})) - result = _run(grid) + chunked = _build(_config(synthetic_surface, chunks={Coordinate.I: 2})) + result = _run(chunked) qualities = result.qualities["boreholes"] - assert qualities.vs.shape == grid.x.shape + assert qualities.vs.shape == chunked.x.shape assert float(qualities.vs.values.mean()) == pytest.approx(1234.0, rel=1e-4) - assert list(result.grids["boreholes"][SITE].values) == NAMES + assert list(result.grids["boreholes"]["site"].values) == NAMES -def test_output_round_trips_through_zarr( - synthetic_surface: Path, tmp_path: Path -) -> None: +def test_output_round_trips_through_zarr(grid: Grid, tmp_path: Path) -> None: """A profile is only useful when a reader can pick out one station.""" path = tmp_path / "boreholes.zarr" - write_velocity_model( - _run(_build(_config(synthetic_surface))), - path, - Format.ZARR, - quantise_arrays=False, - ) + write_velocity_model(_run(grid), path, Format.ZARR, quantise_arrays=False) with xr.open_datatree(path, engine="zarr") as tree: stored = tree["grids/boreholes"].ds - assert list(stored[SITE].values) == NAMES - gull = stored.set_xindex(SITE).sel({SITE: "GULL"}) - assert float(gull.depth.max()) == pytest.approx(400.0) + assert list(stored["site"].values) == NAMES + gull = stored.set_xindex("site").sel({"site": "GULL"}) + assert float(gull.depth.max()) == pytest.approx(DEPTH) assert tree["qualities/boreholes"].ds.vs.shape == stored.x.shape diff --git a/tests/test_synthetic.py b/tests/test_synthetic.py index d283a54..ce6b395 100644 --- a/tests/test_synthetic.py +++ b/tests/test_synthetic.py @@ -1,11 +1,7 @@ """Tests for the synthetic dataset and the commands that write it out. -Closed-form data is data a reader can check by hand. These tests assert the -analytic properties the rest of the suite relies on: the shoreline sits where -elevation crosses zero, and a basin closes to zero thickness on its own -outline. They then read each written file back through the production reader -that consumes it. - +Each test reads the written file back through the production reader that +consumes it. ``just synthetic`` builds the basin meshes themselves, since that goes through gmsh and runs far too slowly for a unit test. """ @@ -29,6 +25,7 @@ ModelType, data_frame_to_mesh, ) +from nzcvm.scripts.surface_cli import read_surface_file as read_scalar_surface from nzcvm.scripts.synthetic import app runner = CliRunner() @@ -115,25 +112,24 @@ def test_vs30_stays_within_bounds(lon: float, lat: float) -> None: # --------------------------------------------------------------------------- -@pytest.mark.parametrize("name", list(synthetic.BasinName)) -def test_basin_closes_on_its_outline(name: synthetic.BasinName) -> None: - basin = synthetic.BASINS[name] +@pytest.fixture(params=list(synthetic.BasinName), ids=lambda name: name.value) +def basin(request: pytest.FixtureRequest) -> synthetic.Basin: + return synthetic.BASINS[request.param] + + +def test_basin_closes_on_its_outline(basin: synthetic.Basin) -> None: outline = basin.outline() sediment = basin.sediment(outline[:, 0], outline[:, 1]) assert sediment == pytest.approx(np.zeros(len(outline)), abs=1e-9) -@pytest.mark.parametrize("name", list(synthetic.BasinName)) -def test_basin_is_deepest_at_its_centre(name: synthetic.BasinName) -> None: - basin = synthetic.BASINS[name] +def test_basin_is_deepest_at_its_centre(basin: synthetic.Basin) -> None: assert basin.sediment(basin.centre_lon, basin.centre_lat) == pytest.approx( basin.thickness ) -@pytest.mark.parametrize("name", list(synthetic.BasinName)) -def test_basement_never_rises_above_topography(name: synthetic.BasinName) -> None: - basin = synthetic.BASINS[name] +def test_basement_never_rises_above_topography(basin: synthetic.Basin) -> None: lon, lat = synthetic.DOMAIN.sample(21, 21) mesh_lon, mesh_lat = np.meshgrid(lon, lat) assert np.all( @@ -141,9 +137,8 @@ def test_basement_never_rises_above_topography(name: synthetic.BasinName) -> Non ) -@pytest.mark.parametrize("name", list(synthetic.BasinName)) -def test_basin_sits_inside_the_domain(name: synthetic.BasinName) -> None: - outline = shapely.Polygon(synthetic.BASINS[name].outline()) +def test_basin_sits_inside_the_domain(basin: synthetic.Basin) -> None: + outline = shapely.Polygon(basin.outline()) assert synthetic.DOMAIN.polygon.contains(outline) @@ -169,9 +164,10 @@ def test_tomography_starts_above_sea_level() -> None: def test_tomography_converts_to_a_mesh() -> None: """The converter rejects anything that isn't a full rectilinear grid.""" - frame = synthetic.tomography(n_horizontal=4, n_depth=3) + n_horizontal, n_depth = 4, 3 + frame = synthetic.tomography(n_horizontal=n_horizontal, n_depth=n_depth) mesh = data_frame_to_mesh("synthetic", frame, MODEL_COLUMNS[ModelType.EP2020]) - assert mesh.sizes["i"] == 4 * 4 * 3 + assert mesh.sizes["i"] == n_horizontal**2 * n_depth assert np.all(mesh.vs.values > 0.0) assert np.all(mesh.vp.values > mesh.vs.values) @@ -209,9 +205,7 @@ def test_vs30_reads_back_unflipped(tmp_path: Path) -> None: path = tmp_path / "vs30.h5" _invoke("vs30", path, "--samples", "16") - from nzcvm.scripts.surface_cli import read_surface_file - - _, _, values = read_surface_file(path, scalar_key="vs30", flip=False) + _, _, values = read_scalar_surface(path, scalar_key="vs30", flip=False) assert values.min() >= synthetic.VS30_MIN assert values.max() <= synthetic.VS30_MAX From 6059d379f55a3b8e014e41339e876f99af10d091 Mon Sep 17 00:00:00 2001 From: Jake Faulkner Date: Mon, 21 Sep 2026 10:17:28 +1200 Subject: [PATCH 3/3] Satisfy the prose linter in the borehole docs Same treatment as the commit below: split the consequence clauses and reword the flagged verbs, leaving the meaning alone. Co-Authored-By: Claude Opus 5 --- README.md | 5 +++-- nzcvm/config/grids/borehole.py | 10 +++++----- nzcvm/grids/borehole.py | 4 ++-- nzcvm/grids/helpers.py | 8 ++++---- tests/test_borehole.py | 8 ++++---- 5 files changed, 18 insertions(+), 17 deletions(-) diff --git a/README.md b/README.md index 92a2c0f..774e78e 100644 --- a/README.md +++ b/README.md @@ -150,7 +150,7 @@ a bare `[grid.projection]` instead: just a CRS. `grid.type = "borehole"` extracts one vertical profile per site rather than filling a volume. Each column runs from the topography down to `depth` at a -fixed `resolution_z`, so the profiles line up sample for sample and compare +fixed `resolution_z`. The profiles line up sample for sample and compare directly. ```toml @@ -190,7 +190,8 @@ on the grid's `i` axis under the name the caller gave it, which the writers keep. In the preceding example, `site` and `network` are labels rather than keywords. -A label keeps the type the caller wrote, so a numeric column arrives numeric. +A label keeps the type the caller wrote: the writers store a numeric column +as numbers, not as text. Each site needs the same set of labels, since the alternative is a column of nulls where one site was missing a key. diff --git a/nzcvm/config/grids/borehole.py b/nzcvm/config/grids/borehole.py index 2d5dee2..49e3b65 100644 --- a/nzcvm/config/grids/borehole.py +++ b/nzcvm/config/grids/borehole.py @@ -30,9 +30,9 @@ class Site(ConfigObject): """One borehole location, in the global CRS, plus whatever labels it has. Longitude and latitude place the site, and the config doesn't reserve any - other key, so a site takes as much or as little description as the caller - has to give it. Each label becomes a coordinate on the grid's ``i`` axis, - which the writers keep. + other key. A site takes as much or as little description as the caller has + to give it. Each label becomes a coordinate on the grid's ``i`` axis, which + the writers keep. Attributes ---------- @@ -85,8 +85,8 @@ class BoreholeGridConfig(GridConfig): Attributes ---------- surface : - Path to the topographic surface mesh file. Used to translate depth - to elevation, so each column starts at the ground. + Path to the topographic surface mesh file. Translates depth to + elevation, which starts a column at the ground. sites : Either an inline list of :class:`Site` objects, or a path to a CSV or Parquet file with ``longitude`` and ``latitude`` columns. Any other diff --git a/nzcvm/grids/borehole.py b/nzcvm/grids/borehole.py index a8826ce..6bdea26 100644 --- a/nzcvm/grids/borehole.py +++ b/nzcvm/grids/borehole.py @@ -3,7 +3,7 @@ Provides :func:`build_borehole` for constructing the set of vertical columns described by a :class:`~nzcvm.config.grids.borehole.BoreholeGridConfig`: one column of query points per site, sampled at a fixed Z resolution from the -topography down to a fixed depth, so a run extracts a directly comparable +topography down to a fixed depth. A run extracts a directly comparable profile per site rather than a filled volume. See :class:`~nzcvm.config.grids.borehole.BoreholeGridConfig` for what a site @@ -188,7 +188,7 @@ def build_borehole(config: BoreholeGridConfig) -> dict[str, Grid]: bottom_left_lon=min_lon, bottom_left_lat=min_lat, ) - # Site labels index i, so the output reads back per station. + # Site labels index i. The output then reads back per station. grid = grid.assign_coords( {name: (Coordinate.I, values) for name, values in labels.items()} ) diff --git a/nzcvm/grids/helpers.py b/nzcvm/grids/helpers.py index b526d95..d983c24 100644 --- a/nzcvm/grids/helpers.py +++ b/nzcvm/grids/helpers.py @@ -66,10 +66,10 @@ def topography_following_grid( ) -> Grid: """Hang a fixed-resolution depth axis off *surface*. - Depth is purely a function of k and *resolution_z*, identically for every - column, so every column runs from the topography down to *thickness* and - the bottom follows the topography exactly. Chunking only ever applies to - i/j, and k always stays one chunk. + Depth is purely a function of k and *resolution_z*, identical from column + to column. A column starts at the topography and ends *thickness* below + it, with the bottom following the topography exactly. Chunking only ever + applies to i/j, and k always stays one chunk. Parameters ---------- diff --git a/tests/test_borehole.py b/tests/test_borehole.py index 0a4d5fb..1ca9893 100644 --- a/tests/test_borehole.py +++ b/tests/test_borehole.py @@ -319,8 +319,8 @@ def test_columns_sit_over_the_projected_site(grid: Grid) -> None: def test_columns_run_from_the_topography_down_to_depth(grid: Grid) -> None: - """Comparable profiles are the point: every column shares one depth axis, - starts at the ground and ends `depth` below it.""" + """Comparable profiles need one shared depth axis. A column starts at the + ground and ends `depth` below it.""" depth = grid.depth.values assert np.all(depth == depth[:1, :1, :]) assert depth[0, 0, 0] == pytest.approx(0.0) @@ -358,8 +358,8 @@ def test_geometry_covers_every_site(grid: Grid) -> None: def test_derived_origin_is_the_site_centroid(grid: Grid) -> None: - """A borehole grid has no configured origin, so the attributes come from - the sites. Downstream writers still expect them to be present.""" + """A borehole grid has no configured origin. The attributes come from the + sites, since downstream writers still expect them to be present.""" lons = [site.longitude for site in SITES] lats = [site.latitude for site in SITES] assert grid.origin_lon == pytest.approx(np.mean(lons), abs=1e-2)