Metadata-Version: 2.4
Name: nzcvm
Version: 2026.10.5
Summary: New Zealand Community Velocity Model
Author: ucgmsim
Classifier: Programming Language :: Rust
Classifier: Programming Language :: Python :: Implementation :: CPython
Requires-Python: >=3.13
Description-Content-Type: text/markdown
Requires-Dist: dask[array,diagnostics,distributed]>=2026.3.0
Requires-Dist: h5py>=3.16.0
Requires-Dist: hdf5plugin>=6.0.0
Requires-Dist: mashumaro[toml,yaml]>=3.20
Requires-Dist: numcodecs>=0.16.5
Requires-Dist: numpy>=2.4.4
Requires-Dist: pandas[parquet]>=2.0
Requires-Dist: psutil>=7.2.2
Requires-Dist: pyproj>=3.7.2
Requires-Dist: rich>=14.3.3
Requires-Dist: scipy>=1.17.1
Requires-Dist: shapely>=2.1.2
Requires-Dist: tqdm>=4.67.3
Requires-Dist: typer>=0.24.2
Requires-Dist: xarray-dataclasses>=1.10.0
Requires-Dist: xarray[io]>=2026.2.0
Requires-Dist: zarr>=3.1.6
Requires-Dist: zfpy>=1.0.1
Provides-Extra: vis
Requires-Dist: pyvista>=0.47.2; extra == "vis"
Provides-Extra: mesh
Requires-Dist: gmsh>=4.15.2; extra == "mesh"
Provides-Extra: raster
Requires-Dist: rioxarray>=0.22.0; extra == "raster"
Provides-Extra: bench
Requires-Dist: geopandas>=1.1.3; extra == "bench"
Requires-Dist: joblib>=1.5.3; extra == "bench"
Requires-Dist: platformdirs>=4.9.6; extra == "bench"

# NZCVM

New Zealand Community Velocity Model: tools for building and querying
tetrahedral velocity models.

A velocity model is a collection of tetrahedral meshes. Each mesh holds cell
data for seismic velocity (Vp, Vs), density (rho) and the quality factors
(Qp, Qs), plus a priority that controls blending where meshes overlap.
`nzcvm generate` samples those meshes onto a structured 3D grid. It then pushes
the result through a chain of layers and writes it in a format a simulation
code can read.

![Fence diagram of shear-wave velocity through four west-east sections of a generated Canterbury velocity model, standing on shaded relief](docs/images/canterbury_fence.png)

*One `nzcvm generate` run over a 260 × 200 km Canterbury domain. EP2020
tomography blended with fourteen basin and volcanic models, a 1-D offshore
profile, the Ely Vs30 taper, a Vp/Vs clamp, and Backus averaging over each
depth cell. The low-velocity wedge thickening eastward is the Canterbury
Plains sediment sequence. The high-Vs lens interrupting it on the second
section is the Banks Peninsula volcanics.*

---

## Build

The core query engine is a Rust extension built with
[setuptools-rust](https://github.com/PyO3/setuptools-rust). Requires Python
3.13 or newer.

`uv` is the preferred build tool for this repo:

```sh
uv sync   # creates a venv and builds the Rust extension
```

Or build and install the wheel by hand:

```sh
pip install build
python -m build --wheel
pip install dist/*.whl --force-reinstall
```

The extension compiles with the release profile by default. Set
`SETUPTOOLS_RUST_CARGO_PROFILE=dev` for debug builds.

### Non-pip dependencies

| Dependency              | Purpose                                           |
|-------------------------|---------------------------------------------------|
| Rust toolchain (stable) | Compiling the extension (if building from source) |
| HDF5 ≥ 1.12             | Needed at runtime by h5py                         |

`pyproject.toml` declares the Python dependencies. The base install covers
`nzcvm generate`. Commands with heavy, single-purpose dependencies need an
optional extra, and say which one when it's missing:

| Extra    | Installs                         | Needed for           |
|----------|----------------------------------|----------------------|
| `vis`    | pyvista                          | `nzcvm view`         |
| `mesh`   | gmsh                             | `nzcvm basin`        |
| `raster` | rioxarray (rasterio, GDAL)       | `nzcvm convert-tiff` |
| `bench`  | geopandas, joblib, platformdirs  | `nzcvm tree-stats`   |

```sh
pip install 'nzcvm[vis]'     # also installs pyvista
uv sync --extra mesh         # in a checkout
uv sync --all-extras         # everything
```

---

## Quick start

```sh
uv run nzcvm generate examples/2014p240655.toml output.zarr
```

The config selects the grid, the layer chain, and the models to query. The
output extension selects the writer. The preceding example needs `resources/`
(DEM, Vs30, coastline) and `models/` to be present.

With no data to hand, build the synthetic dataset instead, then run against
it:

```sh
just synthetic
uv run nzcvm generate examples/synthetic.toml synthetic/model.zarr
```

---

## Command line

| Command              | Purpose                                                       |
|----------------------|---------------------------------------------------------------|
| `nzcvm generate`     | Generate a velocity model from a config file                  |
| `nzcvm view`         | Interactive 3D PyVista viewer for model output               |
| `nzcvm basin`        | Construct a tetrahedral mesh for a basin model                |
| `nzcvm tomography`   | Convert a CSV-like tomography model to a tetrahedral mesh     |
| `nzcvm surface`      | Convert an HDF5 topography surface to a VTK unstructured grid |
| `nzcvm convert-tiff` | Convert a GeoTIFF raster to a surface                         |
| `nzcvm tree-stats`   | Benchmark BVH query performance                               |
| `nzcvm synthetic`    | Write synthetic DEM, Vs30, basin, and tomography inputs       |
| `nzcvm plugins`      | List available layers, grids, and formats, and their packages |

Useful `generate` options:

| Option           | Effect                                                          |
|------------------|-----------------------------------------------------------------|
| `--n-threads`    | Query threads (defaults to the process CPU affinity)            |
| `--format`       | Force an output format instead of inferring it from the path    |
| `--config-format`| Force `toml` / `yaml` / `json` instead of inferring             |
| `--quantise`     | ZFP-compress the arrays in NetCDF/Zarr output                   |
| `--distributed`  | Run on a local Dask distributed cluster                         |
| `--progress`     | tqdm progress bar (not compatible with `--distributed`)         |
| `--monitor`      | Log CPU/memory usage while running                              |
| `--log-level`    | `DEBUG` / `INFO` / `WARNING` (default `WARNING`)                |

---

## Configuration

A config is a TOML, YAML or JSON file with three sections: `metadata`, `grid`,
and an ordered list of `layers`. Config objects are plain dataclasses
deserialised by mashumaro, so you can also build them in pure Python. They
validate bounds, layer ordering, and layer dependencies.

### Grid types

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`, `decay`                                     |
| `borehole` | `sites`, `depth`, `resolution_z` (one vertical profile per site)            |

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]
crs = 'EPSG:2193'     # target projected CRS
azimuth = 39.0        # clockwise from north, degrees
origin_lon = 176.00145
origin_lat = -39.65225

[grid.chunks]
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.

### Terrain decay

The `sw4` and `emod3d` grids share one terrain-following coordinate. A level
at nominal depth `A` (its depth in a flat box) lies at

```
z = A + B(A, S) · S        (S = surface, positive down)
```

The decay weight `B` sets how quickly the levels stop following the
topography. `B = 1` follows it exactly and `B = 0` is flat. Choose the decay
with a `[grid.decay]` table:

| `type`     | Weight `B`                                     | Parameters                                    |
|------------|------------------------------------------------|-----------------------------------------------|
| `squashed` | `1`: the surface shifted down                  | none                                          |
| `linear`   | `1 - A/L`, flat below `L`                      | `length` (default: bottom of the first block) |
| `tapered`  | `1 - A/(ratio · elevation)` on land, `1` at sea | `ratio` (default 1)                          |
| `sleve`    | `sinh((L - A)/s) / sinh(L/s)`, flat below `L`  | `scale`, `length` (default: bottom of the first block) |

An `emod3d` grid requires a decay. EMOD3D's `squashed` and `squashed_tapered`
topography types are `squashed` and `tapered` (with the default `ratio = 1`).

The builder rejects a decay that folds a column over itself. This can happen
with `sleve` under deep water: keep `scale · tanh(length / scale)` deeper than
the deepest seafloor.

#### Solver coordinates

Every grid samples the earth at its true `z`. The grid's `solver` attribute
says where the solver places each sample, and the sfile writer writes its
interfaces in those coordinates. It follows from the decay:

- **No `decay`** (`sw4` default): **physical**, at the true `z`, for SW4 with
  a `topography` command. The builder stretches the first refinement linearly
  between the surface and its flat bottom, so each column is evenly spaced, as
  SW4's sfile reader assumes.
- **Any `decay`** (and every `emod3d` grid): **nominal**, at the nominal depth
  `A` in a flat box. EMOD3D works this way, and SW4 does too without a
  `topography` command, so this is SW4's EMOD3D-compatible mode. The decay is
  a deliberate distortion of the earth into the box. Each block's nominal
  spacing is its refinement resolution.

```toml
[grid]
type = "sw4"

[grid.decay]   # nominal: run SW4 without topography
type = "tapered"
```

The `sw4`, `emod3d`, `regular` and `borehole` grids store their nominal depths
as a `nominal_depth` coordinate on `k`, so NetCDF, Zarr, CSV, and Parquet
outputs include both sets of coordinates.

### 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`. 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. In the preceding example, `site` and `network` are labels rather than
keywords.

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.

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:

```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)
```

Writing to a `*.csv` or `*.parquet` path instead puts those labels in the
header, one row per sample:

```sh
uv run nzcvm generate examples/borehole.toml boreholes.csv
```

```
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:

```python
import pandas as pd

table = pd.read_csv("boreholes.csv")  # or read_parquet("boreholes.parquet")
for site, profile in table.groupby("site", sort=False):
    print(site, profile.vs.to_numpy())
```

`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
chain, so the **last** entry must be `query`. Some layers declare
dependencies: `ely` and `offshore` both need the `coastline` coordinate, so a
`coastline` layer has to appear before them. The config raises a validation
error if a dependency is missing.

| Type        | Description                                                             |
|-------------|--------------------------------------------------------------------------|
| `query`     | Queries the tetrahedral model tree (required, and always last)          |
| `ely`       | Ely et al. (2010) near-surface Vs taper from a Vs30 map                 |
| `offshore`  | 1-D offshore/coastal velocity profile (requires `coastline`)            |
| `coastline` | Computes signed distance to the coastline, provides `coastline`         |
| `clamp`     | Clamps components and the Vp/Vs ratio to physical bounds                |
| `backus`    | Alpha-weighted super-sampling over each depth cell                     |
| custom      | Any layer registered with `@functional_layer` or as a `Layer` subclass  |

### Output formats

Inferred from the output path, or forced with `--format`.

| Format   | Path      | Description                                                     |
|----------|-----------|-----------------------------------------------------------------|
| `zarr`   | `*.zarr`  | Chunked array store, keeps all metadata, good for debugging     |
| `netcdf` | `*.h5`    | NetCDF4/HDF5 via xarray                                         |
| `sfile`  | `*.sfile` | sfile HDF5 format for driving [SW4](github.com/geodynamics/sw4) |
| `emod3d` | directory | `rho3dfile.d`, `vp3dfile.p`, `vs3dfile.s` binaries suitable for driving [EMOD3D](https://doi.org/10.1785/BSSA0860041091)              |
| `csv`     | `*.csv`   | Flat table, one row per point, labelled by grid, and by site   |
| `parquet` | `*.parquet`, `*.pq` | The same table, with the float32 columns kept typed |

Plugin packages can add more formats (see [Plugins](#plugins)).

`csv` and `parquet` share one flattening step, so the columns are the same
either way. Both hold the whole table in memory, which suits the outputs a
person reads: boreholes, transects, a few profiles. Volumetric grids belong in
Zarr or NetCDF.

### Example configuration

See `examples/` for complete, working configs:

| File                          | Shows                                              |
|-------------------------------|-----------------------------------------------------|
| `2014p240655.toml`            | SW4 grid, full layer chain                          |
| `2014p240655_emod3d.toml`     | The same domain as an EMOD3D grid                   |
| `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]
title = "Domain for geonet event 2014p240655"

[grid]
type = "sw4"
surface = "./resources/dem.zarr"
extent_x = 210000.0
extent_y = 330000.0

[grid.orientation]
crs = 'EPSG:2193'
azimuth = 39.0
origin_lon = 176.00145
origin_lat = -39.65225

[grid.refinements.top_layer]
resolution = 200.0
bottom = 5000.0     # bottom elevation, +z down

[[layers]]
type = "clamp"
min_vp_vs_ratio = 1.73
max_vp_vs_ratio = 4.0

[[layers]]
type = "coastline"
coastline = "resources/coastline.wkb.gz"

[[layers]]
type = "ely"
vs30 = "./resources/vs30.zarr"
depth_t = 450.0

[[layers]]
type = "query"
model_path = "./models"
model_globs = ["*.zarr"]
```

---

## Reading the output

Zarr and NetCDF output is a standard xarray `DataTree` with two top-level
groups, each with one node per grid:

```
/
├── grids/<name>       x, y, z, depth   (i, j, k)
└── qualities/<name>   vp, vs, rho, qp, qs, alpha   (i, j, k)
```

`x`, `y`, `z` and `depth` are *data variables* on the logical `(i, j, k)`
index, not dimension coordinates. The grid is curvilinear so each coordinate is a 3D array.

```python
import matplotlib.pyplot as plt
import xarray as xr

dt = xr.open_datatree("output.zarr", engine="zarr")
grid = dt["grids/grid"].ds
qual = dt["qualities/grid"].ds

# Horizontal Vs slice 200 m down (k = 2 at 100 m vertical spacing)
plt.pcolormesh(grid.x[:, :, 2], grid.y[:, :, 2], qual.vs[:, :, 2])

# West–east vertical cross-section at j = 100
plt.pcolormesh(grid.x[:, 100, :], -grid.z[:, 100, :], qual.vs[:, 100, :])
```

This block generates a depth slice and a cross-section through an 80 × 80 × 20
km Wellington domain at 400 m horizontal resolution. Both sit high in the
section because the Wellington basins are shallow: the low-velocity cap runs
100 to 400 m thick over most of the domain, so a slice at 500 m cuts almost
entirely below it.

![Vs through a generated Wellington-region velocity model: a map slice 200 m below the topography and a west-east cross-section](docs/images/wellington_vs.png)

For interactive 3D visualisation with PyVista (requires `nzcvm[vis]`):

```sh
nzcvm view model output.zarr --scalar vs --coastline resources/coastline.wkb.gz
```

`nzcvm view model` also diffs two models (`--compare-to`, `--diff-mode`) and
can render off-screen to an image file (`--off-screen --screenshot out.png`).

---

## Python API

Query a model tree directly, without building a grid:

```python
from pathlib import Path
from nzcvm.models.model import ModelTree

tree = ModelTree.load_models(
    [Path("models/ep2020.zarr"), Path("models/Wellington.zarr")]
)
quality = tree.query(x=1_749_150.0, y=5_428_150.0, z=500.0)
print(quality.vp, quality.vs)  # None if the point is outside every mesh
```

`load_models` takes an iterable of mesh paths, anything `xarray` can open
(the meshes distributed in `models/` are Zarr). Coordinates are in the model's
projected CRS with `z` positive downwards.

`query_many` is the vectorised form and returns a `Qualities` dataset.
`explain` shows which models contributed to the blend:

```python
>>> tree.explain(1_749_150.0, 5_428_150.0, 100.0)
(ρ=1810.00, Vp=1800.00, Vs=580.00, Qp=58.00, Qs=29.00, ɑ=1.00)
├── Model 0 (priority = 38)
│   └── Quality: (ρ=1810.00, Vp=1800.00, Vs=580.00, ...)
└── Model 1 (priority = 255)
    └── Quality: (ρ=2539.03, Vp=4564.84, Vs=2630.92, ...)
```

Lower priority numbers win. Overlapping models are alpha-composited until the
cumulative alpha reaches 1.0. `ModelRange` restricts a query by priority band:
`BASINS` is 0-127, `TOMOGRAPHY` is 128-255, `ALL` is both.

---

## Testing

```sh
uv run pytest tests/
uv run pytest --doctest-modules nzcvm/   # doctests
uv run ruff check nzcvm/ tests/          # linting
uv run ty check nzcvm/                   # type checking
```

Tests marked `real_data` need a model directory supplied via `MODEL_PATH`.

### Synthetic data

`just synthetic` writes a self-contained data root under `synthetic/`, so a
run of `nzcvm generate` covers the whole path without the real NZCVM data. It
takes about a minute, most of it meshing the basins, and needs no
`NZCVM_DATA_ROOT`.

| Recipe                 | Writes                                                     |
|------------------------|------------------------------------------------------------|
| `synthetic_dem`        | `dem.h5`, `dem.zarr` (topography for `grid.surface`)       |
| `synthetic_vs30`       | `vs30.h5`, `vs30.zarr` (Vs30 map for the `ely` layer)      |
| `synthetic_coastline`  | `coastline.wkb.gz` (land polygon in NZTM)                  |
| `synthetic_tomography` | `models/tomography.zarr` (background tomography mesh)      |
| `synthetic_basins`     | `models/gully.zarr`, `models/terrace.zarr`                 |
| `clean_synthetic`      | Removes the lot                                            |

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 it describes.

The pieces are importable, so a test can use a field directly rather than
going through a file:

```python
from nzcvm import synthetic

synthetic.elevation(172.0, -43.6)  # 550.0 m above sea level, inland
synthetic.vs30(172.6, -43.6)  # 200.0 m/s, offshore
```

---

## Code architecture

Four subpackages, each with a narrow responsibility.

### `nzcvm.models`

Geospatial Rust wrappers and mesh I/O. `MeshModel` wraps one tetrahedral
mesh. `ModelTree` combines many meshes into a priority-ordered BVH tree and
handles alpha-composited queries. `Surface` interpolates values from a 2-D
triangular surface mesh (used for the DEM and the Vs30 map). `mesh` provides
the tetrahedral and structured mesh dataclasses and their I/O.

### `nzcvm.layers`

A `Layer` accepts a `Grid` (an xarray Dataset of 3D coordinates) and a
`ModelRange`, and returns `Qualities`. Layers chain via constructor injection
(`next_layer`), and register themselves against a config class through an
`__init_subclass__` hook on `Layer`.

Built-in layers: `QueryLayer`, `ElyLayer`, `OffshoreBasinLayer`,
`CoastlineLayer`, `ClampLayer`, `BackusAveragedLayer`.

`build_pipeline` assembles the chain from a config list, and
`execute_model_pipeline` applies it to every grid with one `map_blocks` per
grid. That hoists the chunked dispatch out of the layers: each layer always
receives a fully concrete chunk and can use plain NumPy.

### `nzcvm.config`

Grid and layer configuration. Each layer comes with a companion `LayerConfig`
dataclass and each grid with a companion `GridConfig`, dispatched on the `type`
discriminator. `VelocityModelConfig` is the top-level object and validates
layer ordering and dependencies.

### `nzcvm.grids`

Builds 3D curvilinear meshes (`sw4`, `regular` or `emod3d`) as xarray
`DataTree` nodes, chunked lazily with Dask and assembled from a `GridConfig`
by the `build_grids_from_config` single-dispatch function.

`Qualities` (in `nzcvm.qualities`) is an `xr.Dataset` subclass holding the
typed velocity, density and quality-factor arrays returned by every layer.

---

## Extending with custom grids, layers, and formats

Grids, layers, and output formats are all extension points.

### Functional layers (simple case)

The `@functional_layer` decorator turns a plain function into a registered
layer, generating a matching `LayerConfig` from its keyword parameters.

```python
from nzcvm.grids.grid import Grid
from nzcvm.layers.core import Layer
from nzcvm.layers.functional import functional_layer
from nzcvm.query import ModelRange


@functional_layer
def scale_vs(
    grid: Grid,
    model_range: ModelRange = ModelRange.ALL,
    *,
    next_layer: Layer,
    factor: float = 1.0,
):
    """Multiply Vs throughout the model by *factor*."""
    qualities = next_layer(grid, model_range)
    qualities["vs"] = qualities["vs"] * factor
    return qualities
```

Add it to a config:

```toml
[[layers]]
type = "scale_vs"
factor = 0.9
```

See `examples/near_fault.py` for a fuller example that uses a spatial distance
mask to perturb Vs near a fault zone.

### Class-based layers

For layers that need state, caching, or a non-trivial config, subclass `Layer`
and pass a matching `LayerConfig` via the `config_cls` keyword. A layer's
`__init__` receives `(config, geometry, next_layer)`. The `geometry` is the
domain footprint, useful for pruning resources at construction time.

```python
from dataclasses import dataclass

import numpy as np
from shapely import Geometry

from nzcvm.config.layers.core import LayerConfig
from nzcvm.grids.grid import Grid
from nzcvm.layers.core import Layer
from nzcvm.qualities import Qualities
from nzcvm.query import ModelRange


@dataclass
class DepthFloorConfig(LayerConfig):
    """Raise Vs to a floor that grows linearly with depth."""

    surface_floor: float = 500.0
    gradient: float = 0.05  # m/s of floor per metre of depth
    type: str = "depth_floor"


class DepthFloorLayer(Layer[DepthFloorConfig], config_cls=DepthFloorConfig):
    def __init__(
        self, config: DepthFloorConfig, geometry: Geometry, next_layer: Layer
    ) -> None:
        super().__init__(config, geometry, next_layer)

    def __call__(
        self, grid: Grid, model_range: ModelRange = ModelRange.ALL
    ) -> Qualities:
        qualities = self.next_layer(grid, model_range)
        floor = self.config.surface_floor + self.config.gradient * grid.depth
        qualities["vs"] = np.maximum(qualities["vs"], floor)
        return qualities
```

Use it in TOML:

```toml
[[layers]]
type = "depth_floor"
surface_floor = 500.0
gradient = 0.05
```

Importing the module registers the layer, so import it before decoding the
config.

### Custom grid types

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.

`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"]))
```

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)`.

```python
import numpy as np
import shapely

from nzcvm.grids.grid import Grid, GridSchema


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=x,
        y=y,
        z=depth,
        depth=depth,
        name="transect",
        resolution=dz,
        geometry=shapely.LineString([start, end]),
        origin_lon=np.float32(174.7762),
        origin_lat=np.float32(-41.2865),
        azimuth=np.float32(0.0),
        grid_azimuth=np.float32(0.0),
        bottom_left_lon=np.float32(174.7762),
        bottom_left_lat=np.float32(-41.2865),
    )
```

Pass the result straight to a pipeline:

```python
from pathlib import Path

from nzcvm.config.layers.query import QueryLayerConfig
from nzcvm.layers.pipeline import build_pipeline

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[0].ravel())  # Vs down the first column, in m/s
```

To drive a grid from a config file, register a builder against a `GridConfig`
subclass. `build_grids_from_config` is a `functools.singledispatch` function
returning a `dict[str, Grid]`. One entry per grid, since an SW4 domain
produces one mesh per refinement level.

```python
from dataclasses import dataclass
from typing import Literal

from nzcvm.config.grids import GridConfig
from nzcvm.grids.builder import build_grids_from_config
from nzcvm.grids.grid import Grid


@dataclass
class TransectConfig(GridConfig):
    start: tuple[float, float]
    end: tuple[float, float]
    n: int
    bottom: float
    dz: float
    type: Literal["transect"] = "transect"


@build_grids_from_config.register
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 = "transect"
start = [1749150.0, 5428150.0]
end = [1759150.0, 5428150.0]
n = 3
bottom = 500.0
dz = 100.0
```

`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.

### Custom output formats

Register a writer with `@register_format`. It receives the populated velocity
model and the output path. Any `extensions` let `nzcvm generate` infer the
format from the output path. Without them, select the format with `--format`.

```python
from pathlib import Path

from nzcvm.formats import register_format
from nzcvm.velocity_model import VelocityModel


@register_format("vs_summary", extensions=(".vs.txt",))
def to_vs_summary(velocity_model: VelocityModel, path: Path) -> None:
    with path.open("w") as f:
        for name, (_grid, qualities) in velocity_model.pairwise.items():
            f.write(f"{name}: mean Vs {float(qualities['vs'].mean()):.1f} m/s\n")
```

Pass `supports_quantisation=True` if the writer accepts a `quantise_arrays`
keyword for `--quantise`.

### Plugins

To make custom layers, grids, or formats available to `nzcvm generate` without
importing them yourself, declare entry points in the package that provides
them. NZCVM imports each group as it loads the matching subsystem:

| Entry-point group     | Module contents                                       |
|-----------------------|-------------------------------------------------------|
| `nzcvm.layer_configs` | `LayerConfig` subclasses                              |
| `nzcvm.layer_impls`   | `Layer` subclasses or `@functional_layer` functions   |
| `nzcvm.grid_configs`  | `GridConfig` subclasses                               |
| `nzcvm.grid_impls`    | `@build_grids_from_config.register` builders          |
| `nzcvm.formats`       | `@register_format` writers                            |

```toml
# pyproject.toml of the plugin package
[project.entry-points."nzcvm.layer_configs"]
depth_floor = "nzcvm_depth_floor.config"

[project.entry-points."nzcvm.layer_impls"]
depth_floor = "nzcvm_depth_floor.layer"
```

Configs and implementations load separately, so reading a config doesn't
import a plugin's runtime dependencies. A `@functional_layer` module defines
both at once, so list it only under `nzcvm.layer_impls`. Its configs then
decode once something imports `nzcvm.layers`, which `nzcvm generate` always
does. A plugin that fails to import logs a warning rather than breaking NZCVM.
`nzcvm plugins` lists every available layer, grid, and format, which package
provides it, and any plugin that failed to load.
