From c189335f589a55fd59047951b1fd7ddfae20a074 Mon Sep 17 00:00:00 2001 From: Jake Faulkner Date: Thu, 10 Sep 2026 22:19:00 +1200 Subject: [PATCH] Treat any extra site key or column as a label, not just a name `site` was a privileged coordinate: `Site` declared a required `name`, the builder wrote it to a fixed `Coordinate.SITE`, and a station's network code or a driller's reference had nowhere to go. Nothing about a borehole grid earns that. The grid needs a position; everything else is the caller's business. Longitude and latitude now place a site and the config reserves nothing else. Every other key in a `[[grid.sites]]` block, and every other column of a CSV or Parquet site file, becomes a coordinate on the `i` axis under the name the caller gave it, and so a column in table output. A label keeps the type it was written with, so a numeric column arrives numeric. `Coordinate.SITE` is gone; `examples/borehole.toml` calls its labels `site` and `network`, and the code privileges neither. Two things the change 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 and break `{grid.name: grid}`. `RESERVED_COORDINATES` in `nzcvm.grids.grid` is derived from `GridSchema`'s own fields, plus the logical index, the coastline coordinate and the component names a writer merges alongside the grid, so it keeps up if the schema changes. A label naming one of those raises. Sites also have to agree on which labels they carry. A missing key is nearly always a typo, and the alternative is a column of nulls. `keep_extra_columns = false` drops the labels and keeps only the spatial coordinates. 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 `labels` mapping out. Co-Authored-By: Claude Opus 5 --- README.md | 46 +++++++-- examples/borehole.toml | 24 +++-- examples/sites.csv | 10 +- nzcvm/config/grids/borehole.py | 59 ++++++++--- nzcvm/coordinates.py | 1 - nzcvm/grids/borehole.py | 87 +++++++++++++---- nzcvm/grids/grid.py | 17 ++++ tests/test_borehole.py | 168 +++++++++++++++++++++++++++----- tests/test_layer_coordinates.py | 19 ++-- tests/test_table.py | 15 +-- 10 files changed, 360 insertions(+), 86 deletions(-) diff --git a/README.md b/README.md index b48c978..0ffbb21 100644 --- a/README.md +++ b/README.md @@ -164,14 +164,15 @@ resolution_z = 25.0 # metres between samples crs = 'EPSG:2193' # CRS the profiles are extracted in [[grid.sites]] -name = "CACS" 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 `name`, `longitude` and `latitude` +point `sites` at a CSV or Parquet file with `longitude` and `latitude` columns: ```toml @@ -181,10 +182,37 @@ 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, and a column in +table output. 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 a writer +puts alongside 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. A `site` coordinate labels the `i` axis, so -the output reads back per station: +contract every layer relies on. The site labels index the `i` axis, so the +output reads back per station: ```python import xarray as xr @@ -205,9 +233,9 @@ uv run nzcvm generate examples/borehole.toml boreholes.csv ``` ``` -grid,site,i,j,k,x,y,z,depth,rho,vp,vs,qp,qs,alpha -boreholes,GULL,0,0,0,1531509.5,5161095.5,-641.124146,0,1810,1800.00012,500,100,50,1 -boreholes,GULL,0,0,1,1531509.5,5161095.5,-616.124146,25,1810,1800,500,100,50,1 +grid,site,network,i,j,k,x,y,z,depth,rho,vp,vs,qp,qs,alpha +boreholes,GULL,NZ,0,0,0,1531509.5,5161095.5,-641.124146,0,1810,1800.00012,500,100,50,1 +boreholes,GULL,NZ,0,0,1,1531509.5,5161095.5,-616.124146,25,1810,1800,500,100,50,1 ``` which `pandas` groups straight back into profiles: @@ -607,7 +635,9 @@ up again. The borehole grid labels its columns this way: grid = grid.assign_coords(site=("i", ["GULL", "TERR"])) ``` -The `csv` and `parquet` writers turn any such coordinate into a label column. +The name must avoid `RESERVED_COORDINATES`, since a coordinate shadows a +variable or attribute of the same name. The `csv` and `parquet` writers turn +any coordinate outside the contract into a label column. Here is a transect: a line of vertical columns between two points, shaped `(n, 1, nk)`. diff --git a/examples/borehole.toml b/examples/borehole.toml index 4572131..c5b8e22 100644 --- a/examples/borehole.toml +++ b/examples/borehole.toml @@ -6,6 +6,11 @@ # 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" @@ -20,9 +25,10 @@ surface = "./synthetic/dem.zarr" depth = 600.0 resolution_z = 25.0 -# Sites can also come from a CSV or Parquet file with name, longitude and -# latitude columns. It has to stay above the tables below, as TOML would -# otherwise read it as a key of `[grid.projection]`. +# 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" @@ -30,24 +36,28 @@ resolution_z = 25.0 crs = 'EPSG:2193' [[grid.sites]] -name = "GULL" # inside the `gully` basin longitude = 172.15 latitude = -43.70 +site = "GULL" # inside the `gully` basin +network = "NZ" [[grid.sites]] -name = "TERR" # inside the `terrace` basin longitude = 172.30 latitude = -43.53 +site = "TERR" # inside the `terrace` basin +network = "NZ" [[grid.sites]] -name = "RIDG" # up in the hills, outside every basin longitude = 172.10 latitude = -43.45 +site = "RIDG" # up in the hills, outside every basin +network = "SC" [[grid.sites]] -name = "SEAB" # offshore longitude = 172.55 latitude = -43.60 +site = "SEAB" # offshore +network = "SC" [[layers]] type = "clamp" diff --git a/examples/sites.csv b/examples/sites.csv index 2b34f4b..42bf0d8 100644 --- a/examples/sites.csv +++ b/examples/sites.csv @@ -1,5 +1,5 @@ -name,longitude,latitude -GULL,172.15,-43.70 -TERR,172.30,-43.53 -RIDG,172.10,-43.45 -SEAB,172.55,-43.60 +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 index 097fad3..57d7a6d 100644 --- a/nzcvm/config/grids/borehole.py +++ b/nzcvm/config/grids/borehole.py @@ -1,6 +1,6 @@ from dataclasses import dataclass, field from pathlib import Path -from typing import Literal +from typing import Any, Literal from mashumaro import field_options from pyproj import CRS @@ -12,7 +12,6 @@ GeographicCRS, Latitude, Longitude, - NonEmptyStr, PositiveFloat, ) from nzcvm.coordinates import WGS84_EPSG, Coordinate @@ -22,23 +21,51 @@ 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. + """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 + and a column in table output. Attributes ---------- - name : - Label for the site. The builder keeps it on the ``site`` coordinate - of the grid, so the output reads back per station. 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'}) """ - name: NonEmptyStr 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 @@ -62,7 +89,8 @@ class BoreholeGridConfig(GridConfig): 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 ``name``, ``longitude`` and ``latitude`` columns. + 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 : @@ -73,10 +101,16 @@ class BoreholeGridConfig(GridConfig): 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:: + 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" @@ -88,11 +122,12 @@ class BoreholeGridConfig(GridConfig): crs = 'EPSG:2193' [[grid.sites]] - name = "CACS" longitude = 172.62 latitude = -43.53 + site = "CACS" + network = "NZ" - or read from a file:: + or read from a file, where the extra columns do the same job:: sites = "examples/sites.csv" """ @@ -111,6 +146,8 @@ class BoreholeGridConfig(GridConfig): 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" diff --git a/nzcvm/coordinates.py b/nzcvm/coordinates.py index cac7694..04cf614 100644 --- a/nzcvm/coordinates.py +++ b/nzcvm/coordinates.py @@ -49,7 +49,6 @@ class Coordinate(StrEnum): Z = auto() DEPTH = auto() COASTLINE = auto() - SITE = auto() I = auto() J = auto() K = auto() diff --git a/nzcvm/grids/borehole.py b/nzcvm/grids/borehole.py index e221c3f..bf9c041 100644 --- a/nzcvm/grids/borehole.py +++ b/nzcvm/grids/borehole.py @@ -11,6 +11,13 @@ 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 @@ -23,15 +30,15 @@ import shapely import xarray as xr -from nzcvm.config.grids.borehole import BoreholeGridConfig, Site +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 Grid, GridSchema +from nzcvm.grids.grid import RESERVED_COORDINATES, Grid, GridSchema from nzcvm.models.surface import Surface -#: Columns a site file has to provide. -SITE_COLUMNS = ("name", "longitude", "latitude") +#: 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]] = { @@ -47,8 +54,8 @@ def read_sites(path: Path) -> list[Site]: Parameters ---------- path : - File with ``name``, ``longitude`` and ``latitude`` columns. The - reader skips any other column. + File with ``longitude`` and ``latitude`` columns. Every other column + becomes a label on the site it belongs to. Returns ------- @@ -58,7 +65,7 @@ def read_sites(path: Path) -> list[Site]: Raises ------ ValueError - If the suffix isn't a supported format, or a required column is + If the suffix isn't a supported format, or a spatial column is missing. """ reader = SITE_READERS.get(path.suffix.lower()) @@ -69,21 +76,66 @@ def read_sites(path: Path) -> list[Site]: ) frame = reader(path) - missing = [column for column in SITE_COLUMNS if column not in frame.columns] + 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(SITE_COLUMNS)}." + f"Expected {', '.join(SPATIAL_KEYS)}." ) return [ - Site(name=str(name), longitude=float(longitude), latitude=float(latitude)) - for name, longitude, latitude in frame[list(SITE_COLUMNS)].itertuples( - index=False + 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) @@ -135,6 +187,7 @@ def _borehole_grid( 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( @@ -163,7 +216,7 @@ def build_borehole(config: BoreholeGridConfig) -> dict[str, Grid]: x_phys, y_phys, z_surface, - name="boreholes", + name=GRID_NAME, depth=config.depth, resolution_z=config.resolution_z, resolution=config.resolution_z, @@ -175,10 +228,10 @@ def build_borehole(config: BoreholeGridConfig) -> dict[str, Grid]: bottom_left_lon=min_lon, bottom_left_lat=min_lat, ) - # Labelling i by site is what makes the output readable: without it, the - # only route back to a station is the order the config listed it in. + # 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( - {Coordinate.SITE: (Coordinate.I, [site.name for site in sites])} + {name: (Coordinate.I, values) for name, values in labels.items()} ) - return {grid.name: grid} + return {GRID_NAME: grid} diff --git a/nzcvm/grids/grid.py b/nzcvm/grids/grid.py index ce3e8b7..256b342 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 a writer merges alongside 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 index 6008b8b..ae7d98e 100644 --- a/tests/test_borehole.py +++ b/tests/test_borehole.py @@ -17,6 +17,7 @@ from pyproj import CRS, Transformer from nzcvm import synthetic +from nzcvm.components import Component from nzcvm.config.grids.borehole import ( DEFAULT_CHUNK_SIZES, BoreholeGridConfig, @@ -29,7 +30,7 @@ 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 Grid +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 @@ -41,13 +42,19 @@ _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. +# 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(name="GULL", longitude=172.15, latitude=-43.70), - Site(name="RIDG", longitude=172.10, latitude=-43.45), - Site(name="SEAB", longitude=172.55, latitude=-43.60), + 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: @@ -76,6 +83,7 @@ def _config( 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( @@ -85,6 +93,7 @@ def _config( resolution_z=resolution_z, projection=Projection(crs=_NZTM), sites_crs=sites_crs, + keep_extra_columns=keep_extra_columns, chunks=chunks, ) @@ -131,9 +140,10 @@ def test_transformer_from_reaches_the_projection() -> None: def _write_sites(path: Path) -> Path: frame = pd.DataFrame( { - "name": [site.name for site in SITES], "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": @@ -151,22 +161,36 @@ def test_read_sites_round_trips(tmp_path: Path, suffix: str) -> None: def test_read_sites_preserves_file_order(tmp_path: Path) -> None: path = tmp_path / "sites.csv" pd.DataFrame( - {"name": ["B", "A"], "longitude": [172.1, 172.2], "latitude": [-43.5, -43.6]} + {"longitude": [172.1, 172.2], "latitude": [-43.5, -43.6], SITE: ["B", "A"]} ).to_csv(path, index=False) - assert [site.name for site in read_sites(path)] == ["B", "A"] + assert [site.labels[SITE] for site in read_sites(path)] == ["B", "A"] -def test_read_sites_ignores_extra_columns(tmp_path: Path) -> None: +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( { - "name": ["A"], "longitude": [172.1], "latitude": [-43.5], + SITE: ["A"], "elevation": [12.0], } ).to_csv(path, index=False) - assert read_sites(path) == [Site(name="A", longitude=172.1, latitude=-43.5)] + 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: @@ -178,7 +202,7 @@ def test_read_sites_rejects_unknown_format(tmp_path: Path) -> None: def test_read_sites_reports_missing_columns(tmp_path: Path) -> None: path = tmp_path / "sites.csv" - pd.DataFrame({"name": ["A"], "longitude": [172.1]}).to_csv(path, index=False) + pd.DataFrame({SITE: ["A"], "longitude": [172.1]}).to_csv(path, index=False) with pytest.raises(ValueError, match="missing the latitude column"): read_sites(path) @@ -211,6 +235,100 @@ def test_config_rejects_a_projected_sites_crs(synthetic_surface: Path) -> None: _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 # --------------------------------------------------------------------------- @@ -224,7 +342,8 @@ def test_grid_is_one_column_per_site(synthetic_surface: Path) -> None: def test_grid_labels_columns_by_site(synthetic_surface: Path) -> None: grid = _build(_config(synthetic_surface)) - assert list(grid[Coordinate.SITE].values) == [site.name for site in SITES] + 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: @@ -332,9 +451,10 @@ def test_grid_is_chunked_over_sites(synthetic_surface: Path) -> None: crs = 'EPSG:2193' [[grid.sites]] -name = "GULL" longitude = 172.15 latitude = -43.70 +site = "GULL" +network = "NZ" [[layers]] type = "query" @@ -350,7 +470,13 @@ def test_toml_config_selects_the_borehole_grid( config = VelocityModelConfig.read_config(path) assert isinstance(config.grid, BoreholeGridConfig) - assert config.grid.sites == [Site(name="GULL", longitude=172.15, latitude=-43.70)] + 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 @@ -370,9 +496,7 @@ def test_toml_config_reads_sites_from_a_file( config = VelocityModelConfig.read_config(path) assert isinstance(config.grid, BoreholeGridConfig) assert config.grid.sites == sites - assert list(_build(config.grid)[Coordinate.SITE].values) == [ - site.name for site in SITES - ] + assert list(_build(config.grid)[SITE].values) == NAMES # --------------------------------------------------------------------------- @@ -411,9 +535,7 @@ def test_grid_survives_the_chunked_pipeline(synthetic_surface: Path) -> None: 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"][Coordinate.SITE].values) == [ - site.name for site in SITES - ] + assert list(result.grids["boreholes"][SITE].values) == NAMES def test_output_round_trips_through_zarr( @@ -430,7 +552,7 @@ def test_output_round_trips_through_zarr( with xr.open_datatree(path, engine="zarr") as tree: stored = tree["grids/boreholes"].ds - assert list(stored[Coordinate.SITE].values) == [site.name for site in SITES] - gull = stored.set_xindex(Coordinate.SITE).sel({Coordinate.SITE: "GULL"}) + 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 diff --git a/tests/test_layer_coordinates.py b/tests/test_layer_coordinates.py index 75efe88..f2b4868 100644 --- a/tests/test_layer_coordinates.py +++ b/tests/test_layer_coordinates.py @@ -47,7 +47,6 @@ ) from nzcvm.config.layers.query import QueryLayerConfig from nzcvm.config.metadata import ModelMetadata -from nzcvm.coordinates import Coordinate from nzcvm.grids.builder import build_grids_from_config from nzcvm.grids.grid import Grid from nzcvm.layers.core import Layer, layer_from_config @@ -67,10 +66,14 @@ # One site inland and one offshore, so the coastline-dependent layers see both # sides of the shoreline and can't skip their work. SITES = [ - Site(name="GULL", longitude=172.15, latitude=-43.70), - Site(name="SEAB", longitude=172.55, latitude=-43.60), + Site(longitude=172.15, latitude=-43.70, labels={"site": "GULL"}), + Site(longitude=172.55, latitude=-43.60, labels={"site": "SEAB"}), ] +#: The label on the sites, treated like any other. +SITE = "site" +NAMES = [site.labels[SITE] for site in SITES] + # --------------------------------------------------------------------------- # The synthetic data each layer needs, in the formats the layers read @@ -303,8 +306,8 @@ def test_layer_preserves_the_site_label( layer = layer_from_config(config)(config, grid.geometry, _Terminal()) qualities = layer(grid) - assert Coordinate.SITE in qualities.coords, layer_type - assert list(qualities[Coordinate.SITE].values) == [site.name for site in SITES] + assert SITE in qualities.coords, layer_type + assert list(qualities[SITE].values) == NAMES # A dropped label would show up as a reindex to NaN rather than an error. assert not np.isnan(qualities.vs.values).any(), layer_type @@ -346,7 +349,7 @@ def test_full_chain_preserves_the_site_label( ) qualities = pipeline(concrete_grid.copy()) - assert list(qualities[Coordinate.SITE].values) == [site.name for site in SITES] + assert list(qualities[SITE].values) == NAMES assert qualities.vs.shape == concrete_grid.x.shape @@ -374,7 +377,7 @@ def test_full_chain_survives_map_blocks( ) qualities = model.qualities["boreholes"] - assert list(qualities[Coordinate.SITE].values) == [site.name for site in SITES] + assert list(qualities[SITE].values) == NAMES assert not np.isnan(qualities.vs.values).any() @@ -392,4 +395,4 @@ def test_a_numpy_terminal_drops_the_label(concrete_grid: Grid) -> None: from nzcvm.layers.dummy import ConstantLayer qualities = ConstantLayer(vs=1234.0)(concrete_grid) - assert Coordinate.SITE not in qualities.coords + assert SITE not in qualities.coords diff --git a/tests/test_table.py b/tests/test_table.py index 196762a..1326d49 100644 --- a/tests/test_table.py +++ b/tests/test_table.py @@ -29,6 +29,9 @@ from nzcvm.qualities import QualitiesSchema from nzcvm.velocity_model import VelocityModel +#: The label the fixture grid puts on `i`, treated like any other. +SITE = "site" + SHAPE = (2, 1, 3) @@ -53,7 +56,7 @@ def _grid(name: str = "boreholes", sites: list[str] | None = None) -> Grid: bottom_left_lat=np.float32(-43.5), ) if sites is not None: - grid = grid.assign_coords({Coordinate.SITE: (Coordinate.I, sites)}) + grid = grid.assign_coords(site=(Coordinate.I, sites)) return grid @@ -131,7 +134,7 @@ def test_write_velocity_model_dispatches_to_parquet(tmp_path: Path) -> None: def test_extra_coordinates_become_label_columns(labelled: pd.DataFrame) -> None: assert labelled.columns.tolist() == [ GRID_COLUMN, - Coordinate.SITE, + SITE, Coordinate.I, Coordinate.J, Coordinate.K, @@ -148,18 +151,18 @@ def test_a_grid_without_extra_coordinates_has_no_label_column( ) -> None: path = tmp_path / "plain.csv" to_csv(_model(_grid()), path) - assert Coordinate.SITE not in pd.read_csv(path).columns + assert SITE not in pd.read_csv(path).columns def test_labels_follow_their_own_axis(labelled: pd.DataFrame) -> None: """`site` lives on i, so every row of a column shares one label.""" - by_site = labelled.groupby(Coordinate.SITE.value)[Coordinate.I.value].unique() + by_site = labelled.groupby(SITE)[Coordinate.I.value].unique() assert by_site["GULL"].tolist() == [0] assert by_site["TERR"].tolist() == [1] def test_every_label_gets_a_full_profile(labelled: pd.DataFrame) -> None: - counts = labelled[Coordinate.SITE].value_counts() + counts = labelled[SITE].value_counts() assert counts.to_dict() == {"GULL": SHAPE[2], "TERR": SHAPE[2]} @@ -221,7 +224,7 @@ def test_parquet_keeps_the_float32_dtype(tmp_path: Path) -> None: to_parquet(_model(_grid(sites=["GULL", "TERR"])), path) table = pd.read_parquet(path) assert table.vs.dtype == np.float32 - assert table[Coordinate.SITE].tolist() == [ + assert table[SITE].tolist() == [ "GULL", "GULL", "GULL",