diff --git a/README.md b/README.md index 774e78e..08940b2 100644 --- a/README.md +++ b/README.md @@ -221,6 +221,29 @@ 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. @@ -252,6 +275,13 @@ Inferred from the output path, or forced with `--format`. | `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 | + +`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 @@ -596,7 +626,8 @@ 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. +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/nzcvm/formats/__init__.py b/nzcvm/formats/__init__.py index ce9d4c4..78c627a 100644 --- a/nzcvm/formats/__init__.py +++ b/nzcvm/formats/__init__.py @@ -10,7 +10,7 @@ from nzcvm.velocity_model import VelocityModel -from . import datatree, emod3d, sfile +from . import datatree, emod3d, sfile, table class Format(StrEnum): @@ -30,6 +30,8 @@ class Format(StrEnum): SFILE = auto() NETCDF = auto() ZARR = auto() + CSV = auto() + PARQUET = auto() def from_path(path: Path) -> Format: @@ -59,7 +61,14 @@ def from_path(path: Path) -> Format: >>> from_path(Path("model.h5")) """ - format_map = {".sfile": Format.SFILE, ".h5": Format.NETCDF, ".zarr": Format.ZARR} + format_map = { + ".sfile": Format.SFILE, + ".h5": Format.NETCDF, + ".zarr": Format.ZARR, + ".csv": Format.CSV, + ".parquet": Format.PARQUET, + ".pq": Format.PARQUET, + } ext = path.suffix if ext in format_map: @@ -106,3 +115,7 @@ def write_velocity_model( datatree.to_netcdf(velocity_model, path, quantise_arrays) case Format.ZARR: datatree.to_zarr(velocity_model, path) + case Format.CSV: + table.to_csv(velocity_model, path) + case Format.PARQUET: + table.to_parquet(velocity_model, path) diff --git a/nzcvm/formats/table.py b/nzcvm/formats/table.py new file mode 100644 index 0000000..e5000ab --- /dev/null +++ b/nzcvm/formats/table.py @@ -0,0 +1,126 @@ +"""Flat table velocity-model writers, for CSV and Parquet. + +Both write one row per grid point. Each row gives the logical index, the +position, and then the components. A grid may hold coordinates beyond the +``(i, j, k)`` index, and each of those becomes a leading label column. That +puts a borehole grid's ``site`` labels in the table, so a reader can tell one +profile from another without counting rows. + +:func:`flatten` builds the table and the two writers encode it, so the column +set is the same either way. Parquet keeps the float32 arrays typed and exact +and compresses a large table well, while CSV renders every number as text, +which is what :data:`FLOAT_FORMAT` is for. Per-column metadata makes a small +Parquet file the larger of the two, so the choice is about the reader rather +than about size. + +The whole table goes through :mod:`pandas`, which keeps every point in memory +at once. That suits the outputs a person reads: boreholes, transects, a few +profiles. Volumetric grids belong in Zarr or NetCDF. +""" + +from pathlib import Path + +import pandas as pd + +from nzcvm.components import Component +from nzcvm.coordinates import Coordinate +from nzcvm.grids.grid import Grid +from nzcvm.qualities import Qualities +from nzcvm.velocity_model import VelocityModel + +#: CSV precision. Every grid and quality array is float32, and nine digits +#: round-trip a float32 exactly. Left to pandas, a seven-digit easting comes +#: out as ``1.5315095e+06`` instead. Parquet keeps the float32 typed and +#: needs none of this. +FLOAT_FORMAT = "%.9g" + +#: Column for the name of the grid a row came from. An SW4 domain writes one +#: grid per refinement level, so the name is what separates them in one table. +GRID_COLUMN = "grid" + +#: Columns every grid has, in the order they appear after the label columns. +FIXED_COLUMNS: tuple[str, ...] = ( + Coordinate.I, + Coordinate.J, + Coordinate.K, + Coordinate.X, + Coordinate.Y, + Coordinate.Z, + Coordinate.DEPTH, + *Component, +) + + +def _table(name: str, grid: Grid, qualities: Qualities) -> pd.DataFrame: + """Flatten one grid and its qualities into a row-per-point table. + + Parameters + ---------- + name : + Name of the grid, written into the :data:`GRID_COLUMN` column. + grid : + Grid holding the ``x``, ``y``, ``z`` and ``depth`` positions. + qualities : + Components sampled on *grid*. + + Returns + ------- + pandas.DataFrame + One row per grid point, columns ordered + :data:`GRID_COLUMN`, labels, then :data:`FIXED_COLUMNS`. + """ + table = grid.assign(qualities).to_dataframe().reset_index() + # Any coordinate past the logical index labels the rows: `site` on a + # borehole grid, and anything a custom grid builder adds. + labels = [column for column in table.columns if column not in FIXED_COLUMNS] + table.insert(0, GRID_COLUMN, name) + return table[[GRID_COLUMN, *labels, *FIXED_COLUMNS]] + + +def flatten(velocity_model: VelocityModel) -> pd.DataFrame: + """Flatten every grid in *velocity_model* into one row-per-point table. + + Parameters + ---------- + velocity_model : + Model whose grids the query pipeline has already populated. + + Returns + ------- + pandas.DataFrame + The grids concatenated in order, one row per point. + """ + tables = [ + _table(name, grid, qualities) + for name, (grid, qualities) in velocity_model.pairwise.items() + ] + return pd.concat(tables, ignore_index=True) + + +def to_csv(velocity_model: VelocityModel, path: Path) -> None: + """Write *velocity_model* to *path* as one CSV table. + + Parameters + ---------- + velocity_model : + Model whose grids the query pipeline has already populated. + path : + Destination file. + """ + flatten(velocity_model).to_csv(path, index=False, float_format=FLOAT_FORMAT) + + +def to_parquet(velocity_model: VelocityModel, path: Path) -> None: + """Write *velocity_model* to *path* as one Parquet table. + + The same columns as :func:`to_csv`, with the float32 arrays kept typed + rather than rendered as text. + + Parameters + ---------- + velocity_model : + Model whose grids the query pipeline has already populated. + path : + Destination file. + """ + flatten(velocity_model).to_parquet(path, index=False) diff --git a/tests/test_layer_coordinates.py b/tests/test_layer_coordinates.py new file mode 100644 index 0000000..dfe110a --- /dev/null +++ b/tests/test_layer_coordinates.py @@ -0,0 +1,395 @@ +"""A layer has to hand back the coordinates of the grid it received. + +A layer takes a :class:`~nzcvm.grids.grid.Grid` and returns +:class:`~nzcvm.qualities.Qualities`. Nothing in the :class:`Layer` contract +says the result keeps the grid's coordinates, and a layer that assembles its +output from raw NumPy drops them: the array has no coordinates to keep. Two +things break at once. ``map_blocks`` rejects a chunk whose coordinates differ +from the template, and the ``site`` labels a borehole grid puts on ``i`` never +reach a writer. + +So each test here passes a labelled borehole grid to one registered layer, +and to the whole chain, then checks the label is still on the result. +:data:`COVERED` has to match +:attr:`~nzcvm.layers.core.Layer.registry`, so a new layer fails here until +someone covers it. +""" + +from __future__ import annotations + +import gzip +import importlib +import pkgutil +from dataclasses import dataclass, fields +from pathlib import Path +from typing import Literal + +import numpy as np +import pytest +import shapely +import shapely.ops +import xarray as xr +from pyproj import CRS, Transformer + +from nzcvm import synthetic +from nzcvm.components import Component +from nzcvm.config.grids.borehole import BoreholeGridConfig, Site +from nzcvm.config.grids.model import Projection +from nzcvm.config.layers.backus import BackusAveragedLayerConfig +from nzcvm.config.layers.clamp import Bound, ClampLayerConfig +from nzcvm.config.layers.coastline import CoastlineConfig +from nzcvm.config.layers.core import LayerConfig +from nzcvm.config.layers.ely import ElyLayerConfig +from nzcvm.config.layers.offshore import ( + DepthModel, + OffshoreBasinConfig, + VelocityModel1D, +) +from nzcvm.config.layers.query import QueryLayerConfig +from nzcvm.config.metadata import ModelMetadata +from nzcvm.grids.builder import build_grids_from_config +from nzcvm.grids.grid import Grid +from nzcvm.layers.core import Layer, layer_from_config +from nzcvm.layers.pipeline import build_pipeline, execute_model_pipeline +from nzcvm.models.mesh import StructuredMeshSchema +from nzcvm.qualities import Qualities, QualitiesSchema +from nzcvm.query import ModelRange +from nzcvm.scripts.convert_tomography import ( + MODEL_COLUMNS, + ModelType, + data_frame_to_mesh, +) + +_NZTM = CRS.from_epsg(2193) +_TO_NZTM = Transformer.from_crs(4326, _NZTM, always_xy=True) + +# 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(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 +# --------------------------------------------------------------------------- + + +def _surface(path: Path, values: np.ndarray, samples: int) -> Path: + """Write a Zarr surface mesh holding *values* over the synthetic domain.""" + lon, lat = synthetic.DOMAIN.sample(samples, samples) + mesh_lon, mesh_lat = np.meshgrid(lon, lat, indexing="ij") + x, y = _TO_NZTM.transform(mesh_lon, mesh_lat) + StructuredMeshSchema.new( + x=x.astype(np.float32), + y=y.astype(np.float32), + z=values.astype(np.float32), + i=np.arange(samples), + j=np.arange(samples), + name=path.stem, + ).to_zarr(path, mode="w") + return path + + +@pytest.fixture(scope="module") +def resources(tmp_path_factory: pytest.TempPathFactory) -> dict[str, Path]: + """A synthetic DEM, Vs30 map, coastline and model directory.""" + root = tmp_path_factory.mktemp("resources") + samples = 24 + lon, lat = synthetic.DOMAIN.sample(samples, samples) + mesh_lon, mesh_lat = np.meshgrid(lon, lat, indexing="ij") + + # A surface file and a grid agree on +z down, so elevation flips. + dem = _surface(root / "dem.zarr", -synthetic.elevation(mesh_lon, mesh_lat), samples) + vs30 = _surface(root / "vs30.zarr", synthetic.vs30(mesh_lon, mesh_lat), samples) + + coastline = root / "coastline.wkb.gz" + projected = shapely.ops.transform( + _TO_NZTM.transform, shapely.Polygon(synthetic.land()) + ) + with gzip.open(coastline, "wb") as handle: + handle.write(shapely.to_wkb(projected)) + + models = root / "models" + models.mkdir() + mesh = data_frame_to_mesh( + "tomography", + synthetic.tomography(n_horizontal=6, n_depth=5), + MODEL_COLUMNS[ModelType.EP2020], + ) + mesh.to_zarr(models / "tomography.zarr", mode="w") + + return {"dem": dem, "vs30": vs30, "coastline": coastline, "models": models} + + +@pytest.fixture(scope="module") +def borehole_grid(resources: dict[str, Path]) -> Grid: + """A two-site borehole grid, labelled on ``i`` by site.""" + return build_grids_from_config( + BoreholeGridConfig( + surface=resources["dem"], + sites=SITES, + depth=400.0, + resolution_z=100.0, + projection=Projection(crs=_NZTM), + ) + )["boreholes"] + + +@pytest.fixture() +def concrete_grid(borehole_grid: Grid) -> Grid: + """The same grid, computed. + + `execute_model_pipeline` hoists the chunked dispatch into one + `map_blocks` per grid, so a layer is only ever handed a concrete chunk. + Calling a layer on a Dask-backed grid raises inside `apply_ufunc`. + """ + return borehole_grid.compute() + + +# --------------------------------------------------------------------------- +# One config per registered layer +# --------------------------------------------------------------------------- + + +def _configs(resources: dict[str, Path]) -> dict[str, LayerConfig]: + return { + "backus": BackusAveragedLayerConfig(samples=3), + "clamp": ClampLayerConfig(clamps={Component.VS: Bound(min=4000.0)}), + "coastline": CoastlineConfig(coastline=resources["coastline"]), + "ely": ElyLayerConfig(vs30=resources["vs30"], depth_t=450.0), + "offshore": OffshoreBasinConfig( + basin_depth=[ + DepthModel(distance=0.0, bottom_depth=0.0), + DepthModel(distance=10_000.0, bottom_depth=1000.0), + ], + model=[ + VelocityModel1D( + bottom_depth=50.0, + rho=1810.0, + vp=1800.0, + vs=380.0, + qp=100.0, + qs=50.0, + alpha=1.0, + ), + VelocityModel1D( + bottom_depth=300.0, + rho=1810.0, + vp=1800.0, + vs=580.0, + qp=100.0, + qs=50.0, + alpha=1.0, + ), + VelocityModel1D( + bottom_depth=1200.0, + rho=1810.0, + vp=1800.0, + vs=830.0, + qp=100.0, + qs=50.0, + alpha=1.0, + ), + ], + ), + "query": QueryLayerConfig( + model_path=resources["models"], model_globs=["*.zarr"] + ), + } + + +#: The layers these tests pass a grid to. +COVERED = frozenset({"backus", "clamp", "coastline", "ely", "offshore", "query"}) + + +def _shipped_layer_types() -> set[str]: + """The ``type`` discriminator of every layer config in the package. + + Read off :mod:`nzcvm.config.layers` rather than + :attr:`~nzcvm.layers.core.Layer.registry`, which picks up the dummy and + sentinel layers the rest of the suite defines. The discriminator is a + dataclass field default, which avoids building a config instance: several + of them take a required path. + """ + package = importlib.import_module("nzcvm.config.layers") + types = set() + for _loader, module_name, _is_pkg in pkgutil.walk_packages( + package.__path__, package.__name__ + "." + ): + module = importlib.import_module(module_name) + for value in vars(module).values(): + if ( + isinstance(value, type) + and issubclass(value, LayerConfig) + and value is not LayerConfig + and value.__module__ == module_name + ): + types.update( + field.default + for field in fields(value) + if field.name == "type" and isinstance(field.default, str) + ) + return types + + +def test_every_shipped_layer_is_covered() -> None: + assert _shipped_layer_types() == COVERED + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +@dataclass +class _TerminalConfig(LayerConfig): + type: Literal["_terminal"] = "_terminal" + + +class _Terminal(Layer[_TerminalConfig]): + """A terminal built with xarray, so it keeps the grid's coordinates. + + `nzcvm.layers.dummy.constant` would be the obvious choice, but it builds + its output from `np.ones` and so hands back no coordinates at all, which + is the property under test here. + """ + + def __init__(self) -> None: + super().__init__(_TerminalConfig(), None, None) # ty: ignore[invalid-argument-type] + + def __call__( + self, grid: Grid, model_range: ModelRange = ModelRange.ALL + ) -> Qualities: + ones = xr.ones_like(grid.x) + return QualitiesSchema.new( + rho=ones * 2700.0, + vp=ones * 6000.0, + vs=ones * 3500.0, + qp=ones * 200.0, + qs=ones * 100.0, + alpha=ones, + ) + + +def _with_coastline(grid: Grid, resources: dict[str, Path]) -> Grid: + """Add the ``coastline`` coordinate that `ely` and `offshore` require. + + The coastline layer writes the coordinate onto the grid it receives, so + calling it for that side effect is the honest way to get one. + """ + config = CoastlineConfig(coastline=resources["coastline"]) + populated = grid.copy() + layer_from_config(config)(config, grid.geometry, _Terminal())(populated) + return populated + + +# --------------------------------------------------------------------------- +# Each layer on its own +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize("layer_type", sorted(COVERED)) +def test_layer_preserves_the_site_label( + layer_type: str, concrete_grid: Grid, resources: dict[str, Path] +) -> None: + config = _configs(resources)[layer_type] + grid = concrete_grid.copy() + if "coastline" in config.requires: + grid = _with_coastline(grid, resources) + + layer = layer_from_config(config)(config, grid.geometry, _Terminal()) + qualities = layer(grid) + + 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 + + +@pytest.mark.parametrize("layer_type", sorted(COVERED)) +def test_layer_returns_the_grid_shape( + layer_type: str, concrete_grid: Grid, resources: dict[str, Path] +) -> None: + """A singleton j axis is easy to squeeze away by accident.""" + config = _configs(resources)[layer_type] + grid = concrete_grid.copy() + if "coastline" in config.requires: + grid = _with_coastline(grid, resources) + + layer = layer_from_config(config)(config, grid.geometry, _Terminal()) + assert layer(grid).vs.shape == concrete_grid.x.shape, layer_type + + +# --------------------------------------------------------------------------- +# The whole chain, chunked +# --------------------------------------------------------------------------- + + +def test_full_chain_preserves_the_site_label( + concrete_grid: Grid, resources: dict[str, Path] +) -> None: + """Ordered as a real config: outermost first, `query` last.""" + configs = _configs(resources) + pipeline = build_pipeline( + concrete_grid.geometry, + [ + configs["clamp"], + configs["coastline"], + configs["offshore"], + configs["ely"], + configs["backus"], + configs["query"], + ], + ) + qualities = pipeline(concrete_grid.copy()) + + assert list(qualities[SITE].values) == NAMES + assert qualities.vs.shape == concrete_grid.x.shape + + +def test_full_chain_survives_map_blocks( + borehole_grid: Grid, resources: dict[str, Path] +) -> None: + """`execute_model_pipeline` compares each chunk against a template built + from the grid, so a layer that drops the label fails the whole run.""" + from nzcvm.velocity_model import VelocityModel + + configs = _configs(resources) + pipeline = build_pipeline( + borehole_grid.geometry, + [ + configs["clamp"], + configs["coastline"], + configs["offshore"], + configs["ely"], + configs["query"], + ], + ) + model = execute_model_pipeline( + VelocityModel(grids={"boreholes": borehole_grid}, metadata=ModelMetadata()), + pipeline, + ) + + qualities = model.qualities["boreholes"] + assert list(qualities[SITE].values) == NAMES + assert not np.isnan(qualities.vs.values).any() + + +# --------------------------------------------------------------------------- +# The constant terminal layer +# --------------------------------------------------------------------------- + + +def test_the_constant_terminal_keeps_the_label(concrete_grid: Grid) -> None: + """`constant` builds its output with `ones_like`, so it keeps the grid's + coordinates and `map_blocks` accepts the result.""" + from nzcvm.layers.dummy import ConstantLayer + + qualities = ConstantLayer(vs=1234.0)(concrete_grid) + assert list(qualities[SITE].values) == list(concrete_grid[SITE].values) diff --git a/tests/test_table.py b/tests/test_table.py new file mode 100644 index 0000000..1326d49 --- /dev/null +++ b/tests/test_table.py @@ -0,0 +1,251 @@ +"""Tests for the flat table writers, CSV and Parquet. + +The interesting property is the header. A grid may hold coordinates past the +``(i, j, k)`` index, and a writer has to turn each one into a label column, +which is what lets a reader tell one borehole profile from another. The rest +checks that the table lists each grid point once, and that a round trip +recovers the float32 values exactly. + +Both encodings come off one :func:`~nzcvm.formats.table.flatten` call, so the +header tests run against CSV alone and the round trip runs against both. +""" + +from __future__ import annotations + +from pathlib import Path + +import numpy as np +import pandas as pd +import pytest +import shapely +import xarray as xr + +from nzcvm.components import Component +from nzcvm.config.metadata import ModelMetadata +from nzcvm.coordinates import Coordinate +from nzcvm.formats import Format, from_path, write_velocity_model +from nzcvm.formats.table import GRID_COLUMN, to_csv, to_parquet +from nzcvm.grids.grid import Grid, GridSchema +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) + + +def _grid(name: str = "boreholes", sites: list[str] | None = None) -> Grid: + """A small concrete grid, optionally labelled with a ``site`` coordinate.""" + ni, nj, nk = SHAPE + i, j, k = np.meshgrid(np.arange(ni), np.arange(nj), np.arange(nk), indexing="ij") + grid = GridSchema.new( + # Values distinct per point, so a misordered table shows up. + x=(1_500_000.0 + i).astype(np.float32), + y=(5_100_000.0 + j).astype(np.float32), + z=(100.0 * i + k).astype(np.float32), + depth=(25.0 * k).astype(np.float32), + name=name, + resolution=25.0, + geometry=shapely.box(171.9, -43.6, 172.1, -43.4), + origin_lon=np.float32(172.0), + origin_lat=np.float32(-43.5), + azimuth=np.float32(0.0), + grid_azimuth=np.float32(0.0), + bottom_left_lon=np.float32(172.0), + bottom_left_lat=np.float32(-43.5), + ) + if sites is not None: + grid = grid.assign_coords(site=(Coordinate.I, sites)) + return grid + + +def _model(*grids: Grid) -> VelocityModel: + """Pair each grid with qualities that vary point by point.""" + qualities = {} + for n, grid in enumerate(grids): + ramp = xr.zeros_like(grid.x) + grid.depth + 1000.0 * n + qualities[grid.name] = QualitiesSchema.new( + rho=ramp + 1.0, + vp=ramp + 2.0, + vs=ramp + 3.0, + qp=ramp + 4.0, + qs=ramp + 5.0, + alpha=xr.ones_like(grid.x), + ) + return VelocityModel( + grids={grid.name: grid for grid in grids}, + metadata=ModelMetadata(), + qualities=qualities, + ) + + +@pytest.fixture() +def labelled(tmp_path: Path) -> pd.DataFrame: + """The table written for a grid that holds site labels.""" + path = tmp_path / "boreholes.csv" + to_csv(_model(_grid(sites=["GULL", "TERR"])), path) + return pd.read_csv(path) + + +# --------------------------------------------------------------------------- +# Format selection +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize( + "suffix, expected", + [ + (".csv", Format.CSV), + (".parquet", Format.PARQUET), + (".pq", Format.PARQUET), + ], +) +def test_format_is_inferred_from_the_extension(suffix: str, expected: Format) -> None: + assert from_path(Path("boreholes").with_suffix(suffix)) is expected + + +@pytest.mark.parametrize("format", [Format.CSV, Format.PARQUET]) +def test_quantisation_is_rejected(format: Format, tmp_path: Path) -> None: + """ZFP applies to the array stores, so it can't mean anything here.""" + with pytest.raises(ValueError, match="quantisation"): + write_velocity_model( + _model(_grid()), tmp_path / "out.table", format, quantise_arrays=True + ) + + +def test_write_velocity_model_dispatches_to_csv(tmp_path: Path) -> None: + path = tmp_path / "out.csv" + write_velocity_model(_model(_grid()), path, Format.INFERRED, quantise_arrays=False) + assert path.read_text().startswith(f"{GRID_COLUMN},") + + +def test_write_velocity_model_dispatches_to_parquet(tmp_path: Path) -> None: + path = tmp_path / "out.parquet" + write_velocity_model(_model(_grid()), path, Format.INFERRED, quantise_arrays=False) + assert pd.read_parquet(path).columns[0] == GRID_COLUMN + + +# --------------------------------------------------------------------------- +# The header +# --------------------------------------------------------------------------- + + +def test_extra_coordinates_become_label_columns(labelled: pd.DataFrame) -> None: + assert labelled.columns.tolist() == [ + GRID_COLUMN, + SITE, + Coordinate.I, + Coordinate.J, + Coordinate.K, + Coordinate.X, + Coordinate.Y, + Coordinate.Z, + Coordinate.DEPTH, + *Component, + ] + + +def test_a_grid_without_extra_coordinates_has_no_label_column( + tmp_path: Path, +) -> None: + path = tmp_path / "plain.csv" + to_csv(_model(_grid()), path) + 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(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[SITE].value_counts() + assert counts.to_dict() == {"GULL": SHAPE[2], "TERR": SHAPE[2]} + + +# --------------------------------------------------------------------------- +# The table +# --------------------------------------------------------------------------- + + +def test_one_row_per_grid_point(labelled: pd.DataFrame) -> None: + assert len(labelled) == np.prod(SHAPE) + index = [Coordinate.I.value, Coordinate.J.value, Coordinate.K.value] + assert not labelled.duplicated(subset=index).any() + + +@pytest.mark.parametrize( + "writer, reader, suffix", + [ + (to_csv, pd.read_csv, ".csv"), + (to_parquet, pd.read_parquet, ".parquet"), + ], +) +def test_values_round_trip_losslessly( + writer, reader, suffix: str, tmp_path: Path +) -> None: + """CSV recovers the float32 from its written precision. Parquet stores it.""" + grid = _grid(sites=["GULL", "TERR"]) + model = _model(grid) + path = (tmp_path / "boreholes").with_suffix(suffix) + writer(model, path) + table = reader(path) + + qualities = model.qualities[grid.name] + for name, expected in [*grid.data_vars.items(), *qualities.data_vars.items()]: + assert np.array_equal( + table[name].to_numpy().astype(np.float32), expected.values.ravel() + ), name + + +def test_grid_column_separates_multiple_grids(tmp_path: Path) -> None: + """An SW4 domain writes one grid per refinement into the same table.""" + path = tmp_path / "refinements.csv" + to_csv(_model(_grid("coarse"), _grid("fine")), path) + + table = pd.read_csv(path) + assert table[GRID_COLUMN].value_counts().to_dict() == { + "coarse": np.prod(SHAPE), + "fine": np.prod(SHAPE), + } + # The second grid's qualities are offset by 1000, so the split is real. + assert table.groupby(GRID_COLUMN).vs.min().to_dict() == { + "coarse": 3.0, + "fine": 1003.0, + } + + +def test_parquet_keeps_the_float32_dtype(tmp_path: Path) -> None: + """Text rendering is a CSV problem, so Parquet shouldn't inherit it.""" + path = tmp_path / "boreholes.parquet" + to_parquet(_model(_grid(sites=["GULL", "TERR"])), path) + table = pd.read_parquet(path) + assert table.vs.dtype == np.float32 + assert table[SITE].tolist() == [ + "GULL", + "GULL", + "GULL", + "TERR", + "TERR", + "TERR", + ] + + +def test_parquet_and_csv_agree_on_columns(tmp_path: Path) -> None: + model = _model(_grid(sites=["GULL", "TERR"])) + to_csv(model, tmp_path / "out.csv") + to_parquet(model, tmp_path / "out.parquet") + assert ( + pd.read_csv(tmp_path / "out.csv").columns.tolist() + == pd.read_parquet(tmp_path / "out.parquet").columns.tolist() + ) + + +def test_eastings_avoid_scientific_notation(tmp_path: Path) -> None: + """A seven-digit float32 easting reads as 1.5315095e+06 unless asked not to.""" + path = tmp_path / "boreholes.csv" + to_csv(_model(_grid(sites=["GULL", "TERR"])), path) + assert "e+06" not in path.read_text()