Miroslav Simko <ms@iolabs.ch> 2026-09-02T07:50:14+02:00
Commit #92 ยท 117 snippets
src/iolabs/common/color_intensity_data.py | 12 +- src/iolabs/common/segment_points_io.py | 110 +++---- tests/test_color_intensity_data.py | 34 +- tests/test_segment_points_io.py | 529 +++++++++++++++--------------- 4 files changed, 346 insertions(+), 339 deletions(-)
| 710 | Returns: | 710 | Returns: |
| 711 | ``(rows, streamable)``. ``rows`` is ``-1`` when the header is | 711 | ``(rows, streamable)``. ``rows`` is ``-1`` when the header is |
| 712 | unreadable (missing member, corrupt archive, unsupported NPY version). | 712 | unreadable (missing member, corrupt archive, unsupported NPY version). |
| 713 | """ | 713 | """ |
| 714 | npz_path = Path(path) | 714 | npz_path = pathlib.Path(path) |
| 715 | try: | 715 | try: |
| 716 | with zipfile.ZipFile(npz_path) as archive: | 716 | with zipfile.ZipFile(npz_path) as archive: |
| 717 | info = archive.getinfo("points.npy") | 717 | info = archive.getinfo("points.npy") |
| 718 | stored = info.compress_type == zipfile.ZIP_STORED | 718 | stored = info.compress_type == zipfile.ZIP_STORED |
| 733 | # caught here so a partial write degrades to "unreadable" instead. | 733 | # caught here so a partial write degrades to "unreadable" instead. |
| 734 | return -1, False | 734 | return -1, False |
| 735 | 735 | ||
| 736 | 736 | ||
| 737 | def _iter_points_chunks_streamed(path: Path, chunk_points: int) -> Iterator[np.ndarray]: | 737 | def _iter_points_chunks_streamed(path: pathlib.Path, chunk_points: int) -> abc.Iterator[np.ndarray]: |
| 738 | """Stream the ``points.npy`` member of a ZIP_STORED npz in row chunks.""" | 738 | """Stream the ``points.npy`` member of a ZIP_STORED npz in row chunks.""" |
| 739 | with zipfile.ZipFile(path) as archive, archive.open("points.npy") as handle: | 739 | with zipfile.ZipFile(path) as archive, archive.open("points.npy") as handle: |
| 740 | version = np.lib.format.read_magic(handle) | 740 | version = np.lib.format.read_magic(handle) |
| 741 | if version == (1, 0): | 741 | if version == (1, 0): |
| 7 | field is left to that subclass's ``__post_init__`` to derive, as the | 7 | field is left to that subclass's ``__post_init__`` to derive, as the |
| 8 | constructor could not accept it. | 8 | constructor could not accept it. |
| 9 | """ | 9 | """ |
| 10 | 10 | ||
| 11 | from collections.abc import Callable | 11 | import dataclasses |
| 12 | from dataclasses import dataclass, field, fields | 12 | from collections import abc |
| 13 | from typing import TypeVar | 13 | from typing import TypeVar |
| 14 | 14 | ||
| 15 | import numpy as np | 15 | import numpy as np |
| 16 | 16 |
| 28 | #: ``type(self)``, so a subclass stays its own type through them. | 28 | #: ``type(self)``, so a subclass stays its own type through them. |
| 29 | ColorIntensityDataT = TypeVar("ColorIntensityDataT", bound="ColorIntensityData") | 29 | ColorIntensityDataT = TypeVar("ColorIntensityDataT", bound="ColorIntensityData") |
| 30 | 30 | ||
| 31 | 31 | ||
| 32 | @dataclass | 32 | @dataclasses.dataclass |
| 33 | class ColorIntensityData: | 33 | class ColorIntensityData: |
| 34 | """Stores per-point RGB, intensity, scan angle, and return-count arrays. | 34 | """Stores per-point RGB, intensity, scan angle, and return-count arrays. |
| 35 | 35 | ||
| 36 | Attributes: | 36 | Attributes: |
| 53 | green: np.ndarray | 53 | green: np.ndarray |
| 54 | blue: np.ndarray | 54 | blue: np.ndarray |
| 55 | intensity: np.ndarray | 55 | intensity: np.ndarray |
| 56 | scan_angle_rank: np.ndarray | 56 | scan_angle_rank: np.ndarray |
| 57 | number_of_returns: np.ndarray | None = field(default=None, kw_only=True) | 57 | number_of_returns: np.ndarray | None = dataclasses.field(default=None, kw_only=True) |
| 58 | 58 | ||
| 59 | def __post_init__(self) -> None: | 59 | def __post_init__(self) -> None: |
| 60 | """Zero-fill an omitted ``number_of_returns``, or coerce a supplied one to uint8. | 60 | """Zero-fill an omitted ``number_of_returns``, or coerce a supplied one to uint8. |
| 61 | 61 |
| 79 | source=type(self).__name__, | 79 | source=type(self).__name__, |
| 80 | ) | 80 | ) |
| 81 | 81 | ||
| 82 | def _map_fields( | 82 | def _map_fields( |
| 83 | self: ColorIntensityDataT, transform: Callable[[str], np.ndarray] | 83 | self: ColorIntensityDataT, transform: abc.Callable[[str], np.ndarray] |
| 84 | ) -> ColorIntensityDataT: | 84 | ) -> ColorIntensityDataT: |
| 85 | """Rebuild this instance's type by applying *transform* to every field name. | 85 | """Rebuild this instance's type by applying *transform* to every field name. |
| 86 | 86 | ||
| 87 | Only ``init=True`` fields are passed to the constructor: a subclass may | 87 | Only ``init=True`` fields are passed to the constructor: a subclass may |
| 91 | """ | 91 | """ |
| 92 | return type(self)( | 92 | return type(self)( |
| 93 | **{ | 93 | **{ |
| 94 | data_field.name: transform(data_field.name) | 94 | data_field.name: transform(data_field.name) |
| 95 | for data_field in fields(self) | 95 | for data_field in dataclasses.fields(self) |
| 96 | if data_field.init | 96 | if data_field.init |
| 97 | } | 97 | } |
| 98 | ) | 98 | ) |
| 99 | 99 |
| 34 | and consumers working in the local frame ignore it. This module never | 34 | and consumers working in the local frame ignore it. This module never |
| 35 | applies, negates, or bakes in a sign. | 35 | applies, negates, or bakes in a sign. |
| 36 | """ | 36 | """ |
| 37 | 37 | ||
| 38 | import dataclasses | ||
| 39 | import fnmatch | ||
| 38 | import functools | 40 | import functools |
| 39 | import json | 41 | import json |
| 40 | import logging | 42 | import logging |
| 43 | import pathlib | ||
| 44 | import types | ||
| 41 | import zipfile | 45 | import zipfile |
| 42 | from collections.abc import Callable, Iterator, Mapping, Sequence | 46 | from collections import abc |
| 43 | from dataclasses import dataclass | ||
| 44 | from fnmatch import fnmatch | ||
| 45 | from pathlib import Path | ||
| 46 | from types import MappingProxyType | ||
| 47 | 47 | ||
| 48 | import numpy as np | 48 | import numpy as np |
| 49 | 49 | ||
| 50 | from . import _dtype_coercion | 50 | from . import _dtype_coercion |
| 51 | 51 | ||
| 52 | logger = logging.getLogger(__name__) | 52 | logger = logging.getLogger(__name__) |
| 53 | 53 | ||
| 54 | #: Builds a member for a record that predates its key, given the point count. | 54 | #: Builds a member for a record that predates its key, given the point count. |
| 55 | FillFactory = Callable[[int], np.ndarray] | 55 | FillFactory = abc.Callable[[int], np.ndarray] |
| 56 | 56 | ||
| 57 | #: The point-record key whose row count defines ``N`` for every other member. | 57 | #: The point-record key whose row count defines ``N`` for every other member. |
| 58 | POINTS_KEY = "points" | 58 | POINTS_KEY = "points" |
| 59 | 59 |
| 63 | #: Storage dtype of :data:`NUMBER_OF_RETURNS_KEY` (LAS carries 3 bits, values 1-7). | 63 | #: Storage dtype of :data:`NUMBER_OF_RETURNS_KEY` (LAS carries 3 bits, values 1-7). |
| 64 | NUMBER_OF_RETURNS_DTYPE = np.uint8 | 64 | NUMBER_OF_RETURNS_DTYPE = np.uint8 |
| 65 | 65 | ||
| 66 | 66 | ||
| 67 | @dataclass(frozen=True) | 67 | @dataclasses.dataclass(frozen=True) |
| 68 | class PointFieldSpec: | 68 | class PointFieldSpec: |
| 69 | """How one member of the point record is validated, stored and synthesised. | 69 | """How one member of the point record is validated, stored and synthesised. |
| 70 | 70 | ||
| 71 | Attributes: | 71 | Attributes: |
| 108 | 108 | ||
| 109 | #: The point-record contract: ordered key -> spec. Adding a field to the NPZ | 109 | #: The point-record contract: ordered key -> spec. Adding a field to the NPZ |
| 110 | #: schema means adding an entry here (plus teaching producers to emit it); the | 110 | #: schema means adding an entry here (plus teaching producers to emit it); the |
| 111 | #: load/save/merge machinery below is entirely registry-driven. | 111 | #: load/save/merge machinery below is entirely registry-driven. |
| 112 | POINT_RECORD_SCHEMA: Mapping[str, PointFieldSpec] = MappingProxyType( | 112 | POINT_RECORD_SCHEMA: abc.Mapping[str, PointFieldSpec] = types.MappingProxyType( |
| 113 | { | 113 | { |
| 114 | POINTS_KEY: PointFieldSpec(required=True, columns=3, noun="coordinate"), | 114 | POINTS_KEY: PointFieldSpec(required=True, columns=3, noun="coordinate"), |
| 115 | "red": PointFieldSpec(required=True), | 115 | "red": PointFieldSpec(required=True), |
| 116 | "green": PointFieldSpec(required=True), | 116 | "green": PointFieldSpec(required=True), |
| 127 | ) | 127 | ) |
| 128 | 128 | ||
| 129 | 129 | ||
| 130 | def _schema_keys( | 130 | def _schema_keys( |
| 131 | schema: Mapping[str, PointFieldSpec], *, required: bool | None = None | 131 | schema: abc.Mapping[str, PointFieldSpec], *, required: bool | None = None |
| 132 | ) -> tuple[str, ...]: | 132 | ) -> tuple[str, ...]: |
| 133 | """List the schema's keys in registry order, optionally by required-ness.""" | 133 | """List the schema's keys in registry order, optionally by required-ness.""" |
| 134 | return tuple( | 134 | return tuple( |
| 135 | key for key, spec in schema.items() if required is None or spec.required is required | 135 | key for key, spec in schema.items() if required is None or spec.required is required |
| 190 | return "(N,)" if spec.columns is None else f"(N, {spec.columns})" | 190 | return "(N,)" if spec.columns is None else f"(N, {spec.columns})" |
| 191 | 191 | ||
| 192 | 192 | ||
| 193 | def _validate_record_shapes( | 193 | def _validate_record_shapes( |
| 194 | record: Mapping[str, np.ndarray], | 194 | record: abc.Mapping[str, np.ndarray], |
| 195 | schema: Mapping[str, PointFieldSpec], | 195 | schema: abc.Mapping[str, PointFieldSpec], |
| 196 | *, | 196 | *, |
| 197 | source: str, | 197 | source: str, |
| 198 | ) -> int: | 198 | ) -> int: |
| 199 | """Check every present member against its spec's shape and return ``N``. | 199 | """Check every present member against its spec's shape and return ``N``. |
| 227 | ) | 227 | ) |
| 228 | return n_points | 228 | return n_points |
| 229 | 229 | ||
| 230 | 230 | ||
| 231 | def load_points_npz(path: str | Path) -> PointRecord: | 231 | def load_points_npz(path: str | pathlib.Path) -> PointRecord: |
| 232 | """Load one ``*_run3_points.npz``-style file and validate the schema. | 232 | """Load one ``*_run3_points.npz``-style file and validate the schema. |
| 233 | 233 | ||
| 234 | Every required key of :data:`POINT_RECORD_SCHEMA` must be stored; ``points`` | 234 | Every required key of :data:`POINT_RECORD_SCHEMA` must be stored; ``points`` |
| 235 | must be shape ``(N, 3)`` and every ancillary array (``red``, ``green``, | 235 | must be shape ``(N, 3)`` and every ancillary array (``red``, ``green``, |
| 254 | ValueError: A required key is missing, an array has the wrong shape, or | 254 | ValueError: A required key is missing, an array has the wrong shape, or |
| 255 | a member with a declared storage dtype is stored with a | 255 | a member with a declared storage dtype is stored with a |
| 256 | non-integer/bool dtype or values outside that dtype's range. | 256 | non-integer/bool dtype or values outside that dtype's range. |
| 257 | """ | 257 | """ |
| 258 | npz_path = Path(path) | 258 | npz_path = pathlib.Path(path) |
| 259 | schema = POINT_RECORD_SCHEMA | 259 | schema = POINT_RECORD_SCHEMA |
| 260 | required_keys = _schema_keys(schema, required=True) | 260 | required_keys = _schema_keys(schema, required=True) |
| 261 | with np.load(npz_path) as data: | 261 | with np.load(npz_path) as data: |
| 262 | missing = [key for key in required_keys if key not in data.files] | 262 | missing = [key for key in required_keys if key not in data.files] |
| 285 | record[key] = spec.fill(n_points) | 285 | record[key] = spec.fill(n_points) |
| 286 | return {key: record[key] for key in schema} | 286 | return {key: record[key] for key in schema} |
| 287 | 287 | ||
| 288 | 288 | ||
| 289 | def save_points_npz(path: str | Path, record: Mapping[str, np.ndarray]) -> Path: | 289 | def save_points_npz(path: str | pathlib.Path, record: abc.Mapping[str, np.ndarray]) -> pathlib.Path: |
| 290 | """Write a point record as a compressed NPZ (inverse of :func:`load_points_npz`). | 290 | """Write a point record as a compressed NPZ (inverse of :func:`load_points_npz`). |
| 291 | 291 | ||
| 292 | Every key in :data:`POINT_RECORD_KEYS` must be present, ``number_of_returns`` | 292 | Every key in :data:`POINT_RECORD_KEYS` must be present, ``number_of_returns`` |
| 293 | included: it is optional on load only, to read datasets that predate it. | 293 | included: it is optional on load only, to read datasets that predate it. |
| 303 | Raises: | 303 | Raises: |
| 304 | ValueError: A key is missing, an array has the wrong shape, or a member | 304 | ValueError: A key is missing, an array has the wrong shape, or a member |
| 305 | with a declared storage dtype cannot be cast to it losslessly. | 305 | with a declared storage dtype cannot be cast to it losslessly. |
| 306 | """ | 306 | """ |
| 307 | npz_path = Path(path) | 307 | npz_path = pathlib.Path(path) |
| 308 | schema = POINT_RECORD_SCHEMA | 308 | schema = POINT_RECORD_SCHEMA |
| 309 | keys = _schema_keys(schema) | 309 | keys = _schema_keys(schema) |
| 310 | missing = [key for key in keys if key not in record] | 310 | missing = [key for key in keys if key not in record] |
| 311 | if missing: | 311 | if missing: |
| 323 | np.savez_compressed(npz_path, **payload) | 323 | np.savez_compressed(npz_path, **payload) |
| 324 | return npz_path | 324 | return npz_path |
| 325 | 325 | ||
| 326 | 326 | ||
| 327 | def mask_record(record: Mapping[str, np.ndarray], mask: np.ndarray) -> PointRecord: | 327 | def mask_record(record: abc.Mapping[str, np.ndarray], mask: np.ndarray) -> PointRecord: |
| 328 | """Select the same points from every member of a point record. | 328 | """Select the same points from every member of a point record. |
| 329 | 329 | ||
| 330 | Schema-agnostic: whatever keys the record carries are all indexed with | 330 | Schema-agnostic: whatever keys the record carries are all indexed with |
| 331 | *mask* along axis 0, so ``points`` keeps its ``(n, 3)`` rows while the | 331 | *mask* along axis 0, so ``points`` keeps its ``(n, 3)`` rows while the |
| 367 | ) | 367 | ) |
| 368 | return {key: array[mask] for key, array in arrays.items()} | 368 | return {key: array[mask] for key, array in arrays.items()} |
| 369 | 369 | ||
| 370 | 370 | ||
| 371 | def concat_records(records: Sequence[Mapping[str, np.ndarray]]) -> PointRecord: | 371 | def concat_records(records: abc.Sequence[abc.Mapping[str, np.ndarray]]) -> PointRecord: |
| 372 | """Concatenate point records member-by-member along axis 0. | 372 | """Concatenate point records member-by-member along axis 0. |
| 373 | 373 | ||
| 374 | Schema-agnostic: every key of the first record is concatenated across all | 374 | Schema-agnostic: every key of the first record is concatenated across all |
| 375 | records, so ``points`` grows by rows and the ancillary members by elements. | 375 | records, so ``points`` grows by rows and the ancillary members by elements. |
| 408 | } | 408 | } |
| 409 | 409 | ||
| 410 | 410 | ||
| 411 | def load_segment_points( | 411 | def load_segment_points( |
| 412 | files: Sequence[str | Path], | 412 | files: abc.Sequence[str | pathlib.Path], |
| 413 | ) -> tuple[PointRecord, np.ndarray, list[str]]: | 413 | ) -> tuple[PointRecord, np.ndarray, list[str]]: |
| 414 | """Load and concatenate multiple point NPZs for one segment. | 414 | """Load and concatenate multiple point NPZs for one segment. |
| 415 | 415 | ||
| 416 | Returns ``(merged_record, point_file_ids, file_stems)`` where | 416 | Returns ``(merged_record, point_file_ids, file_stems)`` where |
| 417 | ``point_file_ids[i]`` is the index into *files* for merged point ``i``, | 417 | ``point_file_ids[i]`` is the index into *files* for merged point ``i``, |
| 418 | and ``file_stems`` holds ``Path(f).stem`` for each input in order. | 418 | and ``file_stems`` holds ``Path(f).stem`` for each input in order. |
| 419 | """ | 419 | """ |
| 420 | npz_files = [Path(path) for path in files] | 420 | npz_files = [pathlib.Path(path) for path in files] |
| 421 | if not npz_files: | 421 | if not npz_files: |
| 422 | raise FileNotFoundError("No *_points.npz files provided for segment load") | 422 | raise FileNotFoundError("No *_points.npz files provided for segment load") |
| 423 | 423 | ||
| 424 | records: list[PointRecord] = [] | 424 | records: list[PointRecord] = [] |
| 441 | return merged, point_file_ids, file_stems | 441 | return merged, point_file_ids, file_stems |
| 442 | 442 | ||
| 443 | 443 | ||
| 444 | def geoshift_from_mapping( | 444 | def geoshift_from_mapping( |
| 445 | mapping: Mapping[str, object], | 445 | mapping: abc.Mapping[str, object], |
| 446 | *, | 446 | *, |
| 447 | source: str = "geoshift mapping", | 447 | source: str = "geoshift mapping", |
| 448 | ) -> np.ndarray: | 448 | ) -> np.ndarray: |
| 449 | """Convert an already-parsed geoshift mapping to a shape ``(3,)`` float64 array. | 449 | """Convert an already-parsed geoshift mapping to a shape ``(3,)`` float64 array. |
| 464 | Raises: | 464 | Raises: |
| 465 | ValueError: If *mapping* is not an object, a key is missing, or a | 465 | ValueError: If *mapping* is not an object, a key is missing, or a |
| 466 | value is not numeric. | 466 | value is not numeric. |
| 467 | """ | 467 | """ |
| 468 | if not isinstance(mapping, Mapping): | 468 | if not isinstance(mapping, abc.Mapping): |
| 469 | raise ValueError( | 469 | raise ValueError( |
| 470 | f"{source}: geoshift JSON must be an object with keys x/y/z, " | 470 | f"{source}: geoshift JSON must be an object with keys x/y/z, " |
| 471 | f"got {type(mapping).__name__}" | 471 | f"got {type(mapping).__name__}" |
| 472 | ) | 472 | ) |
| 473 | values: Mapping[str, object] = mapping | 473 | values: abc.Mapping[str, object] = mapping |
| 474 | nested = mapping.get("geoshift") | 474 | nested = mapping.get("geoshift") |
| 475 | if isinstance(nested, Mapping): | 475 | if isinstance(nested, abc.Mapping): |
| 476 | values = nested | 476 | values = nested |
| 477 | 477 | ||
| 478 | missing = [key for key in ("x", "y", "z") if key not in values] | 478 | missing = [key for key in ("x", "y", "z") if key not in values] |
| 479 | if missing: | 479 | if missing: |
| 489 | f"{[values['x'], values['y'], values['z']]!r}" | 489 | f"{[values['x'], values['y'], values['z']]!r}" |
| 490 | ) from exc | 490 | ) from exc |
| 491 | 491 | ||
| 492 | 492 | ||
| 493 | def load_geoshift(path: str | Path) -> np.ndarray: | 493 | def load_geoshift(path: str | pathlib.Path) -> np.ndarray: |
| 494 | """Load a Step-3 ``run3_geoshift.json`` as a shape ``(3,)`` float64 array. | 494 | """Load a Step-3 ``run3_geoshift.json`` as a shape ``(3,)`` float64 array. |
| 495 | 495 | ||
| 496 | Expected JSON layout (written by segmentation-trajectory SegmentMapper):: | 496 | Expected JSON layout (written by segmentation-trajectory SegmentMapper):: |
| 497 | 497 | ||
| 498 | {"x": float, "y": float, "z": float} | 498 | {"x": float, "y": float, "z": float} |
| 499 | 499 | ||
| 500 | The value is returned as recorded; no sign is applied (see module docstring). | 500 | The value is returned as recorded; no sign is applied (see module docstring). |
| 501 | """ | 501 | """ |
| 502 | geoshift_path = Path(path) | 502 | geoshift_path = pathlib.Path(path) |
| 503 | with geoshift_path.open(encoding="utf-8") as handle: | 503 | with geoshift_path.open(encoding="utf-8") as handle: |
| 504 | data = json.load(handle) | 504 | data = json.load(handle) |
| 505 | return geoshift_from_mapping(data, source=str(geoshift_path)) | 505 | return geoshift_from_mapping(data, source=str(geoshift_path)) |
| 506 | 506 | ||
| 507 | 507 | ||
| 508 | def geoshift_candidate_paths(directory: str | Path) -> list[Path]: | 508 | def geoshift_candidate_paths(directory: str | pathlib.Path) -> list[pathlib.Path]: |
| 509 | """List the ``run3_geoshift.json`` paths the fleet's three conventions use. | 509 | """List the ``run3_geoshift.json`` paths the fleet's three conventions use. |
| 510 | 510 | ||
| 511 | In lookup order: | 511 | In lookup order: |
| 512 | 512 |
| 523 | 523 | ||
| 524 | Returns: | 524 | Returns: |
| 525 | The candidate paths in lookup order (existence is not checked). | 525 | The candidate paths in lookup order (existence is not checked). |
| 526 | """ | 526 | """ |
| 527 | base = Path(directory) | 527 | base = pathlib.Path(directory) |
| 528 | return [ | 528 | return [ |
| 529 | base / GEOSHIFT_NAME, | 529 | base / GEOSHIFT_NAME, |
| 530 | base / LANE_POINTS_DIR_NAME / GEOSHIFT_NAME, | 530 | base / LANE_POINTS_DIR_NAME / GEOSHIFT_NAME, |
| 531 | base.parent / GEOSHIFT_NAME, | 531 | base.parent / GEOSHIFT_NAME, |
| 532 | ] | 532 | ] |
| 533 | 533 | ||
| 534 | 534 | ||
| 535 | def find_geoshift_or_none(directory: str | Path) -> np.ndarray | None: | 535 | def find_geoshift_or_none(directory: str | pathlib.Path) -> np.ndarray | None: |
| 536 | """Find and load a geoshift near *directory*, or return ``None``. | 536 | """Find and load a geoshift near *directory*, or return ``None``. |
| 537 | 537 | ||
| 538 | Searches :func:`geoshift_candidate_paths` in order and loads the first | 538 | Searches :func:`geoshift_candidate_paths` in order and loads the first |
| 539 | existing file. Datasets processed without a geoshift have no such file and | 539 | existing file. Datasets processed without a geoshift have no such file and |
| 561 | logger.debug("Using geoshift %s for %s", existing[0], directory) | 561 | logger.debug("Using geoshift %s for %s", existing[0], directory) |
| 562 | return load_geoshift(existing[0]) | 562 | return load_geoshift(existing[0]) |
| 563 | 563 | ||
| 564 | 564 | ||
| 565 | def find_geoshift(directory: str | Path) -> np.ndarray: | 565 | def find_geoshift(directory: str | pathlib.Path) -> np.ndarray: |
| 566 | """Find and load a geoshift near *directory*, raising when absent. | 566 | """Find and load a geoshift near *directory*, raising when absent. |
| 567 | 567 | ||
| 568 | Args: | 568 | Args: |
| 569 | directory: Dataset root, ``lane_points`` directory, or segment directory. | 569 | directory: Dataset root, ``lane_points`` directory, or segment directory. |
| 603 | ) from exc | 603 | ) from exc |
| 604 | 604 | ||
| 605 | 605 | ||
| 606 | def _normalize_file_patterns(file_patterns: object, *, source: str) -> list[str]: | 606 | def _normalize_file_patterns(file_patterns: object, *, source: str) -> list[str]: |
| 607 | raw_patterns: Sequence[object] | 607 | raw_patterns: abc.Sequence[object] |
| 608 | if isinstance(file_patterns, str): | 608 | if isinstance(file_patterns, str): |
| 609 | raw_patterns = [file_patterns] | 609 | raw_patterns = [file_patterns] |
| 610 | elif isinstance(file_patterns, Sequence): | 610 | elif isinstance(file_patterns, abc.Sequence): |
| 611 | raw_patterns = file_patterns | 611 | raw_patterns = file_patterns |
| 612 | else: | 612 | else: |
| 613 | raise ValueError( | 613 | raise ValueError( |
| 614 | f"{source} must be a string or a sequence of strings, got " | 614 | f"{source} must be a string or a sequence of strings, got " |
| 623 | return normalized | 623 | return normalized |
| 624 | 624 | ||
| 625 | 625 | ||
| 626 | def normalize_segment_file_blacklist( | 626 | def normalize_segment_file_blacklist( |
| 627 | mapping: Mapping[object, object] | None, | 627 | mapping: abc.Mapping[object, object] | None, |
| 628 | ) -> dict[int, list[str]]: | 628 | ) -> dict[int, list[str]]: |
| 629 | """Normalize a per-segment fnmatch blacklist mapping. | 629 | """Normalize a per-segment fnmatch blacklist mapping. |
| 630 | 630 | ||
| 631 | Keys may be ints or ``segment_<idx>`` strings. Values are a string or a | 631 | Keys may be ints or ``segment_<idx>`` strings. Values are a string or a |
| 633 | the same segment are merged and sorted. | 633 | the same segment are merged and sorted. |
| 634 | """ | 634 | """ |
| 635 | if mapping is None: | 635 | if mapping is None: |
| 636 | return {} | 636 | return {} |
| 637 | if not isinstance(mapping, Mapping): | 637 | if not isinstance(mapping, abc.Mapping): |
| 638 | raise ValueError("npz_blacklist_by_segment must be a mapping of segment to files") | 638 | raise ValueError("npz_blacklist_by_segment must be a mapping of segment to files") |
| 639 | 639 | ||
| 640 | normalized: dict[int, set[str]] = {} | 640 | normalized: dict[int, set[str]] = {} |
| 641 | for segment_key, file_patterns in mapping.items(): | 641 | for segment_key, file_patterns in mapping.items(): |
| 653 | } | 653 | } |
| 654 | 654 | ||
| 655 | 655 | ||
| 656 | def filter_segment_files( | 656 | def filter_segment_files( |
| 657 | files: Sequence[str | Path], | 657 | files: abc.Sequence[str | pathlib.Path], |
| 658 | segment_index: int, | 658 | segment_index: int, |
| 659 | blacklist: Mapping[int, Sequence[str]], | 659 | blacklist: abc.Mapping[int, abc.Sequence[str]], |
| 660 | ) -> tuple[list[Path], list[Path]]: | 660 | ) -> tuple[list[pathlib.Path], list[pathlib.Path]]: |
| 661 | """Split input files into kept and blacklisted buckets via fnmatch. | 661 | """Split input files into kept and blacklisted buckets via fnmatch. |
| 662 | 662 | ||
| 663 | Patterns from ``blacklist[segment_index]`` are matched against both the | 663 | Patterns from ``blacklist[segment_index]`` are matched against both the |
| 664 | file name and the full POSIX path. Returns ``(kept_files, excluded_files)``. | 664 | file name and the full POSIX path. Returns ``(kept_files, excluded_files)``. |
| 665 | """ | 665 | """ |
| 666 | kept_files: list[Path] = [] | 666 | kept_files: list[pathlib.Path] = [] |
| 667 | excluded_files: list[Path] = [] | 667 | excluded_files: list[pathlib.Path] = [] |
| 668 | file_patterns = list(blacklist.get(segment_index, [])) | 668 | file_patterns = list(blacklist.get(segment_index, [])) |
| 669 | 669 | ||
| 670 | for raw_path in files: | 670 | for raw_path in files: |
| 671 | npz_path = Path(raw_path) | 671 | npz_path = pathlib.Path(raw_path) |
| 672 | path_text = npz_path.as_posix() | 672 | path_text = npz_path.as_posix() |
| 673 | file_name = npz_path.name | 673 | file_name = npz_path.name |
| 674 | is_blacklisted = any( | 674 | is_blacklisted = any( |
| 675 | fnmatch(file_name, pattern) or fnmatch(path_text, pattern) | 675 | fnmatch.fnmatch(file_name, pattern) or fnmatch.fnmatch(path_text, pattern) |
| 676 | for pattern in file_patterns | 676 | for pattern in file_patterns |
| 677 | ) | 677 | ) |
| 678 | if is_blacklisted: | 678 | if is_blacklisted: |
| 679 | excluded_files.append(npz_path) | 679 | excluded_files.append(npz_path) |
| 695 | ) | 695 | ) |
| 696 | return kept_files, excluded_files | 696 | return kept_files, excluded_files |
| 697 | 697 | ||
| 698 | 698 | ||
| 699 | def read_points_header(path: str | Path) -> tuple[int, bool]: | 699 | def read_points_header(path: str | pathlib.Path) -> tuple[int, bool]: |
| 700 | """Return ``(row_count, chunk_streamable)`` for a record's ``points.npy`` member. | 700 | """Return ``(row_count, chunk_streamable)`` for a record's ``points.npy`` member. |
| 701 | 701 | ||
| 702 | Reads only the NPY header inside the NPZ ZIP container, so the point data | 702 | Reads only the NPY header inside the NPZ ZIP container, so the point data |
| 703 | is never materialised. Streaming needs an uncompressed (``ZIP_STORED``) | 703 | is never materialised. Streaming needs an uncompressed (``ZIP_STORED``) |
| 753 | buffer = handle.read(count * row_bytes) | 753 | buffer = handle.read(count * row_bytes) |
| 754 | yield np.frombuffer(buffer, dtype=dtype).reshape(count, cols).copy() | 754 | yield np.frombuffer(buffer, dtype=dtype).reshape(count, cols).copy() |
| 755 | 755 | ||
| 756 | 756 | ||
| 757 | def iter_points_chunks(path: str | Path, chunk_points: int) -> Iterator[np.ndarray]: | 757 | def iter_points_chunks(path: str | pathlib.Path, chunk_points: int) -> abc.Iterator[np.ndarray]: |
| 758 | """Yield a record's ``points`` rows, chunked when the member is streamable. | 758 | """Yield a record's ``points`` rows, chunked when the member is streamable. |
| 759 | 759 | ||
| 760 | Records at or below *chunk_points* rows -- and any record whose | 760 | Records at or below *chunk_points* rows -- and any record whose |
| 761 | ``points.npy`` member is compressed, Fortran-ordered or otherwise | 761 | ``points.npy`` member is compressed, Fortran-ordered or otherwise |
| 774 | Yields: | 774 | Yields: |
| 775 | ``(rows_i, 3)`` point chunks in file order; the final chunk holds the | 775 | ``(rows_i, 3)`` point chunks in file order; the final chunk holds the |
| 776 | remainder and may be shorter. | 776 | remainder and may be shorter. |
| 777 | """ | 777 | """ |
| 778 | npz_path = Path(path) | 778 | npz_path = pathlib.Path(path) |
| 779 | rows, streamable = read_points_header(npz_path) | 779 | rows, streamable = read_points_header(npz_path) |
| 780 | if streamable and chunk_points > 0 and rows > chunk_points: | 780 | if streamable and chunk_points > 0 and rows > chunk_points: |
| 781 | logger.info( | 781 | logger.info( |
| 782 | "%s: streaming %d points in chunks of %d", | 782 | "%s: streaming %d points in chunks of %d", |
| 789 | with np.load(npz_path) as data: | 789 | with np.load(npz_path) as data: |
| 790 | yield np.asarray(data["points"]) | 790 | yield np.asarray(data["points"]) |
| 791 | 791 | ||
| 792 | 792 | ||
| 793 | @dataclass(frozen=True) | 793 | @dataclasses.dataclass(frozen=True) |
| 794 | class RecordSpan: | 794 | class RecordSpan: |
| 795 | """One input record's identity and row range within a concatenated cloud. | 795 | """One input record's identity and row range within a concatenated cloud. |
| 796 | 796 | ||
| 797 | Attributes: | 797 | Attributes: |
| 809 | """One past the last global row index contributed by this file.""" | 809 | """One past the last global row index contributed by this file.""" |
| 810 | return self.offset + self.count | 810 | return self.offset + self.count |
| 811 | 811 | ||
| 812 | 812 | ||
| 813 | def discover_run3_files(segment_dir: str | Path) -> list[Path]: | 813 | def discover_run3_files(segment_dir: str | pathlib.Path) -> list[pathlib.Path]: |
| 814 | """List a segment directory's ``*_run3_points.npz`` records in load order. | 814 | """List a segment directory's ``*_run3_points.npz`` records in load order. |
| 815 | 815 | ||
| 816 | Partial writes carry a trailing marker suffix (``*.npz.part`` and friends); | 816 | Partial writes carry a trailing marker suffix (``*.npz.part`` and friends); |
| 817 | those are skipped rather than handed to a loader that would crash on them. | 817 | those are skipped rather than handed to a loader that would crash on them. |
| 822 | Returns: | 822 | Returns: |
| 823 | Sorted, complete record paths. Empty when the directory is missing or | 823 | Sorted, complete record paths. Empty when the directory is missing or |
| 824 | holds no records. | 824 | holds no records. |
| 825 | """ | 825 | """ |
| 826 | seg_dir = Path(segment_dir) | 826 | seg_dir = pathlib.Path(segment_dir) |
| 827 | if not seg_dir.is_dir(): | 827 | if not seg_dir.is_dir(): |
| 828 | logger.warning("Segment directory missing: %s", seg_dir) | 828 | logger.warning("Segment directory missing: %s", seg_dir) |
| 829 | return [] | 829 | return [] |
| 830 | 830 | ||
| 831 | found: list[Path] = [] | 831 | found: list[pathlib.Path] = [] |
| 832 | for path in sorted(seg_dir.glob(f"{RUN3_POINTS_GLOB}*")): | 832 | for path in sorted(seg_dir.glob(f"{RUN3_POINTS_GLOB}*")): |
| 833 | if not path.name.endswith(RUN3_POINTS_SUFFIX): | 833 | if not path.name.endswith(RUN3_POINTS_SUFFIX): |
| 834 | logger.info("Skipping incomplete/partial run3 record: %s", path.name) | 834 | logger.info("Skipping incomplete/partial run3 record: %s", path.name) |
| 835 | continue | 835 | continue |
| 837 | return found | 837 | return found |
| 838 | 838 | ||
| 839 | 839 | ||
| 840 | def concat_points_npz( | 840 | def concat_points_npz( |
| 841 | files: Sequence[str | Path], | 841 | files: abc.Sequence[str | pathlib.Path], |
| 842 | *, | 842 | *, |
| 843 | target_dtypes: Mapping[str, np.dtype] | None = None, | 843 | target_dtypes: abc.Mapping[str, np.dtype] | None = None, |
| 844 | ) -> tuple[PointRecord, list[RecordSpan]]: | 844 | ) -> tuple[PointRecord, list[RecordSpan]]: |
| 845 | """Concatenate point records after validating a consistent schema. | 845 | """Concatenate point records after validating a consistent schema. |
| 846 | 846 | ||
| 847 | Each file must satisfy the :func:`load_points_npz` contract. By default | 847 | Each file must satisfy the :func:`load_points_npz` contract. By default |
| 872 | FileNotFoundError: If *files* is empty. | 872 | FileNotFoundError: If *files* is empty. |
| 873 | ValueError: If a file violates the point-record contract or its dtypes | 873 | ValueError: If a file violates the point-record contract or its dtypes |
| 874 | disagree with the first file's (for keys without a target dtype). | 874 | disagree with the first file's (for keys without a target dtype). |
| 875 | """ | 875 | """ |
| 876 | npz_files = [Path(path) for path in files] | 876 | npz_files = [pathlib.Path(path) for path in files] |
| 877 | if not npz_files: | 877 | if not npz_files: |
| 878 | raise FileNotFoundError("No *_points.npz files provided for concatenation") | 878 | raise FileNotFoundError("No *_points.npz files provided for concatenation") |
| 879 | 879 | ||
| 880 | casts: dict[str, np.dtype] = { | 880 | casts: dict[str, np.dtype] = { |
| 911 | return concat_records(records), spans | 911 | return concat_records(records), spans |
| 912 | 912 | ||
| 913 | 913 | ||
| 914 | def load_run3_segment( | 914 | def load_run3_segment( |
| 915 | segment_dir: str | Path, | 915 | segment_dir: str | pathlib.Path, |
| 916 | *, | 916 | *, |
| 917 | target_dtypes: Mapping[str, np.dtype] | None = None, | 917 | target_dtypes: abc.Mapping[str, np.dtype] | None = None, |
| 918 | ) -> tuple[PointRecord, list[RecordSpan]]: | 918 | ) -> tuple[PointRecord, list[RecordSpan]]: |
| 919 | """Discover and concatenate one segment's run3 records. | 919 | """Discover and concatenate one segment's run3 records. |
| 920 | 920 | ||
| 921 | Combines :func:`discover_run3_files` (sorted glob, ``.part`` skipped) with | 921 | Combines :func:`discover_run3_files` (sorted glob, ``.part`` skipped) with |
| 931 | Raises: | 931 | Raises: |
| 932 | FileNotFoundError: If the directory holds no complete run3 record. | 932 | FileNotFoundError: If the directory holds no complete run3 record. |
| 933 | ValueError: If the records' schema or dtypes disagree. | 933 | ValueError: If the records' schema or dtypes disagree. |
| 934 | """ | 934 | """ |
| 935 | seg_dir = Path(segment_dir) | 935 | seg_dir = pathlib.Path(segment_dir) |
| 936 | files = discover_run3_files(seg_dir) | 936 | files = discover_run3_files(seg_dir) |
| 937 | if not files: | 937 | if not files: |
| 938 | raise FileNotFoundError(f"No {RUN3_POINTS_GLOB} in {seg_dir}") | 938 | raise FileNotFoundError(f"No {RUN3_POINTS_GLOB} in {seg_dir}") |
| 939 | return concat_points_npz(files, target_dtypes=target_dtypes) | 939 | return concat_points_npz(files, target_dtypes=target_dtypes) |
| 1 | """Tests for ColorIntensityData selection and concatenation operations.""" | 1 | """Tests for ColorIntensityData selection and concatenation operations.""" |
| 2 | from dataclasses import dataclass, field | 2 | import dataclasses |
| 3 | 3 | ||
| 4 | import numpy as np | 4 | import numpy as np |
| 5 | import pytest | 5 | import pytest |
| 6 | 6 | ||
| 7 | from iolabs.common.color_intensity_data import ColorIntensityData | 7 | from iolabs.common import color_intensity_data |
| 8 | 8 | ||
| 9 | 9 | ||
| 10 | def _make_sample(n: int = 5, offset: int = 0) -> ColorIntensityData: | 10 | def _make_sample(n: int = 5, offset: int = 0) -> color_intensity_data.ColorIntensityData: |
| 11 | """Create a sample ColorIntensityData with n points and optional array offset.""" | 11 | """Create a sample ColorIntensityData with n points and optional array offset.""" |
| 12 | return ColorIntensityData( | 12 | return color_intensity_data.ColorIntensityData( |
| 13 | red=np.arange(offset, offset + n, dtype=np.uint8), | 13 | red=np.arange(offset, offset + n, dtype=np.uint8), |
| 14 | green=np.arange(offset + 10, offset + 10 + n, dtype=np.uint8), | 14 | green=np.arange(offset + 10, offset + 10 + n, dtype=np.uint8), |
| 15 | blue=np.arange(offset + 20, offset + 20 + n, dtype=np.uint8), | 15 | blue=np.arange(offset + 20, offset + 20 + n, dtype=np.uint8), |
| 16 | intensity=np.arange(offset + 100, offset + 100 + n, dtype=np.float64), | 16 | intensity=np.arange(offset + 100, offset + 100 + n, dtype=np.float64), |
| 85 | """Tests for the AI3D-382 number_of_returns field.""" | 85 | """Tests for the AI3D-382 number_of_returns field.""" |
| 86 | 86 | ||
| 87 | def test_defaults_to_zeros_when_omitted(self): | 87 | def test_defaults_to_zeros_when_omitted(self): |
| 88 | """Callers predating AI3D-382 get 0 (unknown) per point, never 1.""" | 88 | """Callers predating AI3D-382 get 0 (unknown) per point, never 1.""" |
| 89 | data = ColorIntensityData( | 89 | data = color_intensity_data.ColorIntensityData( |
| 90 | red=np.zeros(4, dtype=np.uint8), | 90 | red=np.zeros(4, dtype=np.uint8), |
| 91 | green=np.zeros(4, dtype=np.uint8), | 91 | green=np.zeros(4, dtype=np.uint8), |
| 92 | blue=np.zeros(4, dtype=np.uint8), | 92 | blue=np.zeros(4, dtype=np.uint8), |
| 93 | intensity=np.zeros(4, dtype=np.float64), | 93 | intensity=np.zeros(4, dtype=np.float64), |
| 99 | np.testing.assert_array_equal(data.number_of_returns, np.zeros(4, dtype=np.uint8)) | 99 | np.testing.assert_array_equal(data.number_of_returns, np.zeros(4, dtype=np.uint8)) |
| 100 | 100 | ||
| 101 | def test_zero_fill_survives_mask_and_append(self): | 101 | def test_zero_fill_survives_mask_and_append(self): |
| 102 | """A defaulted field stays aligned through the mask/concat paths.""" | 102 | """A defaulted field stays aligned through the mask/concat paths.""" |
| 103 | legacy = ColorIntensityData( | 103 | legacy = color_intensity_data.ColorIntensityData( |
| 104 | red=np.zeros(3, dtype=np.uint8), | 104 | red=np.zeros(3, dtype=np.uint8), |
| 105 | green=np.zeros(3, dtype=np.uint8), | 105 | green=np.zeros(3, dtype=np.uint8), |
| 106 | blue=np.zeros(3, dtype=np.uint8), | 106 | blue=np.zeros(3, dtype=np.uint8), |
| 107 | intensity=np.zeros(3, dtype=np.float64), | 107 | intensity=np.zeros(3, dtype=np.float64), |
| 123 | 123 | ||
| 124 | def test_extra_field_flows_through_mask_and_append(self): | 124 | def test_extra_field_flows_through_mask_and_append(self): |
| 125 | """A subclass field is masked and concatenated by the generic transforms.""" | 125 | """A subclass field is masked and concatenated by the generic transforms.""" |
| 126 | 126 | ||
| 127 | @dataclass | 127 | @dataclasses.dataclass |
| 128 | class WithClassification(ColorIntensityData): | 128 | class WithClassification(color_intensity_data.ColorIntensityData): |
| 129 | classification: np.ndarray | None = None | 129 | classification: np.ndarray | None = None |
| 130 | 130 | ||
| 131 | def _make(n: int, offset: int) -> WithClassification: | 131 | def _make(n: int, offset: int) -> WithClassification: |
| 132 | base = _make_sample(n, offset) | 132 | base = _make_sample(n, offset) |
| 160 | 160 | ||
| 161 | def test_non_init_subclass_field_is_not_passed_to_the_constructor(self): | 161 | def test_non_init_subclass_field_is_not_passed_to_the_constructor(self): |
| 162 | """A derived ``init=False`` field must not break the generic transforms.""" | 162 | """A derived ``init=False`` field must not break the generic transforms.""" |
| 163 | 163 | ||
| 164 | @dataclass | 164 | @dataclasses.dataclass |
| 165 | class WithPointCount(ColorIntensityData): | 165 | class WithPointCount(color_intensity_data.ColorIntensityData): |
| 166 | point_count: int = field(init=False, default=0) | 166 | point_count: int = dataclasses.field(init=False, default=0) |
| 167 | 167 | ||
| 168 | def __post_init__(self) -> None: | 168 | def __post_init__(self) -> None: |
| 169 | super().__post_init__() | 169 | super().__post_init__() |
| 170 | self.point_count = len(self.red) | 170 | self.point_count = len(self.red) |
| 195 | 195 | ||
| 196 | def test_subclass_may_declare_a_required_field(self): | 196 | def test_subclass_may_declare_a_required_field(self): |
| 197 | """A trailing defaulted field would make a required subclass field a TypeError.""" | 197 | """A trailing defaulted field would make a required subclass field a TypeError.""" |
| 198 | 198 | ||
| 199 | @dataclass | 199 | @dataclasses.dataclass |
| 200 | class WithRequiredClassification(ColorIntensityData): | 200 | class WithRequiredClassification(color_intensity_data.ColorIntensityData): |
| 201 | classification: np.ndarray | 201 | classification: np.ndarray |
| 202 | 202 | ||
| 203 | data = WithRequiredClassification( | 203 | data = WithRequiredClassification( |
| 204 | red=np.zeros(3, dtype=np.uint8), | 204 | red=np.zeros(3, dtype=np.uint8), |
| 217 | assert len(merged.classification) == 5 | 217 | assert len(merged.classification) == 5 |
| 218 | 218 | ||
| 219 | def test_pre_ai3d_382_fields_still_take_positional_args(self): | 219 | def test_pre_ai3d_382_fields_still_take_positional_args(self): |
| 220 | """The five original fields keep their positional order for old call sites.""" | 220 | """The five original fields keep their positional order for old call sites.""" |
| 221 | data = ColorIntensityData( | 221 | data = color_intensity_data.ColorIntensityData( |
| 222 | np.zeros(2, dtype=np.uint8), | 222 | np.zeros(2, dtype=np.uint8), |
| 223 | np.zeros(2, dtype=np.uint8), | 223 | np.zeros(2, dtype=np.uint8), |
| 224 | np.zeros(2, dtype=np.uint8), | 224 | np.zeros(2, dtype=np.uint8), |
| 225 | np.zeros(2, dtype=np.float64), | 225 | np.zeros(2, dtype=np.float64), |
| 230 | 230 | ||
| 231 | def test_number_of_returns_is_not_positional(self): | 231 | def test_number_of_returns_is_not_positional(self): |
| 232 | """Passing it as a sixth positional arg is a TypeError, not a silent mismatch.""" | 232 | """Passing it as a sixth positional arg is a TypeError, not a silent mismatch.""" |
| 233 | with pytest.raises(TypeError): | 233 | with pytest.raises(TypeError): |
| 234 | ColorIntensityData( | 234 | color_intensity_data.ColorIntensityData( |
| 235 | np.zeros(2, dtype=np.uint8), | 235 | np.zeros(2, dtype=np.uint8), |
| 236 | np.zeros(2, dtype=np.uint8), | 236 | np.zeros(2, dtype=np.uint8), |
| 237 | np.zeros(2, dtype=np.uint8), | 237 | np.zeros(2, dtype=np.uint8), |
| 238 | np.zeros(2, dtype=np.float64), | 238 | np.zeros(2, dtype=np.float64), |
| 244 | class TestNumberOfReturnsDtypeContract: | 244 | class TestNumberOfReturnsDtypeContract: |
| 245 | """The constructor enforces the same uint8 contract as segment_points_io.""" | 245 | """The constructor enforces the same uint8 contract as segment_points_io.""" |
| 246 | 246 | ||
| 247 | @staticmethod | 247 | @staticmethod |
| 248 | def _make(number_of_returns: np.ndarray) -> ColorIntensityData: | 248 | def _make(number_of_returns: np.ndarray) -> color_intensity_data.ColorIntensityData: |
| 249 | n = len(number_of_returns) | 249 | n = len(number_of_returns) |
| 250 | return ColorIntensityData( | 250 | return color_intensity_data.ColorIntensityData( |
| 251 | red=np.zeros(n, dtype=np.uint8), | 251 | red=np.zeros(n, dtype=np.uint8), |
| 252 | green=np.zeros(n, dtype=np.uint8), | 252 | green=np.zeros(n, dtype=np.uint8), |
| 253 | blue=np.zeros(n, dtype=np.uint8), | 253 | blue=np.zeros(n, dtype=np.uint8), |
| 254 | intensity=np.zeros(n, dtype=np.float64), | 254 | intensity=np.zeros(n, dtype=np.float64), |
| 2 | 2 | ||
| 3 | import functools | 3 | import functools |
| 4 | import json | 4 | import json |
| 5 | import logging | 5 | import logging |
| 6 | from pathlib import Path | 6 | import pathlib |
| 7 | from types import MappingProxyType | 7 | import types |
| 8 | 8 | ||
| 9 | import numpy as np | 9 | import numpy as np |
| 10 | import pytest | 10 | import pytest |
| 11 | 11 | ||
| 12 | from iolabs.common import segment_points_io | 12 | from iolabs.common import segment_points_io |
| 13 | from iolabs.common.segment_points_io import ( | 13 | |
| 14 | GEOSHIFT_NAME, | 14 | # Import-time snapshot of the shipped registry object. Monkeypatching the |
| 15 | NUMBER_OF_RETURNS_KEY, | 15 | # module attribute must not mutate this binding; the leak test below checks |
| 16 | OPTIONAL_POINT_RECORD_KEYS, | 16 | # both this snapshot and the live ``segment_points_io.POINT_RECORD_SCHEMA``. |
| 17 | POINT_RECORD_KEYS, | 17 | _POINT_RECORD_SCHEMA_AT_IMPORT = segment_points_io.POINT_RECORD_SCHEMA |
| 18 | POINT_RECORD_SCHEMA, | 18 | _POINT_RECORD_KEYS_AT_IMPORT = segment_points_io.POINT_RECORD_KEYS |
| 19 | REQUIRED_POINT_RECORD_KEYS, | ||
| 20 | PointFieldSpec, | ||
| 21 | RecordSpan, | ||
| 22 | concat_points_npz, | ||
| 23 | concat_records, | ||
| 24 | discover_run3_files, | ||
| 25 | filter_segment_files, | ||
| 26 | find_geoshift, | ||
| 27 | find_geoshift_or_none, | ||
| 28 | geoshift_candidate_paths, | ||
| 29 | geoshift_from_mapping, | ||
| 30 | iter_points_chunks, | ||
| 31 | load_geoshift, | ||
| 32 | load_points_npz, | ||
| 33 | load_run3_segment, | ||
| 34 | load_segment_points, | ||
| 35 | mask_record, | ||
| 36 | normalize_segment_file_blacklist, | ||
| 37 | parse_segment_key, | ||
| 38 | read_points_header, | ||
| 39 | save_points_npz, | ||
| 40 | ) | ||
| 41 | 19 | ||
| 42 | 20 | ||
| 43 | def _make_record(n: int, *, seed: int = 0) -> dict[str, np.ndarray]: | 21 | def _make_record(n: int, *, seed: int = 0) -> dict[str, np.ndarray]: |
| 44 | rng = np.random.default_rng(seed) | 22 | rng = np.random.default_rng(seed) |
| 55 | 33 | ||
| 56 | def _make_legacy_record(n: int, *, seed: int = 0) -> dict[str, np.ndarray]: | 34 | def _make_legacy_record(n: int, *, seed: int = 0) -> dict[str, np.ndarray]: |
| 57 | """Build a pre-AI3D-382 record: every key except ``number_of_returns``.""" | 35 | """Build a pre-AI3D-382 record: every key except ``number_of_returns``.""" |
| 58 | record = _make_record(n, seed=seed) | 36 | record = _make_record(n, seed=seed) |
| 59 | del record[NUMBER_OF_RETURNS_KEY] | 37 | del record[segment_points_io.NUMBER_OF_RETURNS_KEY] |
| 60 | return record | 38 | return record |
| 61 | 39 | ||
| 62 | 40 | ||
| 63 | def test_save_load_points_npz_round_trip(tmp_path: Path) -> None: | 41 | def test_save_load_points_npz_round_trip(tmp_path: pathlib.Path) -> None: |
| 64 | record = _make_record(7, seed=1) | 42 | record = _make_record(7, seed=1) |
| 65 | path = tmp_path / "scan_a_run3_points.npz" | 43 | path = tmp_path / "scan_a_run3_points.npz" |
| 66 | save_points_npz(path, record) | 44 | segment_points_io.save_points_npz(path, record) |
| 67 | 45 | ||
| 68 | loaded = load_points_npz(path) | 46 | loaded = segment_points_io.load_points_npz(path) |
| 69 | assert list(loaded.keys()) == list(POINT_RECORD_KEYS) | 47 | assert list(loaded.keys()) == list(segment_points_io.POINT_RECORD_KEYS) |
| 70 | for key in POINT_RECORD_KEYS: | 48 | for key in segment_points_io.POINT_RECORD_KEYS: |
| 71 | np.testing.assert_array_equal(loaded[key], record[key]) | 49 | np.testing.assert_array_equal(loaded[key], record[key]) |
| 72 | 50 | ||
| 73 | 51 | ||
| 74 | def test_load_points_npz_rejects_missing_keys(tmp_path: Path) -> None: | 52 | def test_load_points_npz_rejects_missing_keys(tmp_path: pathlib.Path) -> None: |
| 75 | path = tmp_path / "broken.npz" | 53 | path = tmp_path / "broken.npz" |
| 76 | np.savez_compressed(path, points=np.zeros((2, 3), dtype=np.float32), red=np.zeros(2)) | 54 | np.savez_compressed(path, points=np.zeros((2, 3), dtype=np.float32), red=np.zeros(2)) |
| 77 | with pytest.raises(ValueError, match="missing required key"): | 55 | with pytest.raises(ValueError, match="missing required key"): |
| 78 | load_points_npz(path) | 56 | segment_points_io.load_points_npz(path) |
| 79 | 57 | ||
| 80 | 58 | ||
| 81 | def test_load_points_npz_rejects_bad_points_shape(tmp_path: Path) -> None: | 59 | def test_load_points_npz_rejects_bad_points_shape(tmp_path: pathlib.Path) -> None: |
| 82 | record = _make_record(3, seed=2) | 60 | record = _make_record(3, seed=2) |
| 83 | record["points"] = np.zeros((3, 2), dtype=np.float32) | 61 | record["points"] = np.zeros((3, 2), dtype=np.float32) |
| 84 | path = tmp_path / "bad_shape.npz" | 62 | path = tmp_path / "bad_shape.npz" |
| 85 | np.savez_compressed(path, **record) | 63 | np.savez_compressed(path, **record) |
| 86 | with pytest.raises(ValueError, match=r"shape \(N, 3\)"): | 64 | with pytest.raises(ValueError, match=r"shape \(N, 3\)"): |
| 87 | load_points_npz(path) | 65 | segment_points_io.load_points_npz(path) |
| 88 | 66 | ||
| 89 | 67 | ||
| 90 | def test_load_points_npz_rejects_row_count_mismatch(tmp_path: Path) -> None: | 68 | def test_load_points_npz_rejects_row_count_mismatch(tmp_path: pathlib.Path) -> None: |
| 91 | record = _make_record(4, seed=3) | 69 | record = _make_record(4, seed=3) |
| 92 | record["intensity"] = record["intensity"][:2] | 70 | record["intensity"] = record["intensity"][:2] |
| 93 | path = tmp_path / "mismatch.npz" | 71 | path = tmp_path / "mismatch.npz" |
| 94 | np.savez_compressed(path, **record) | 72 | np.savez_compressed(path, **record) |
| 95 | with pytest.raises(ValueError, match="intensity"): | 73 | with pytest.raises(ValueError, match="intensity"): |
| 96 | load_points_npz(path) | 74 | segment_points_io.load_points_npz(path) |
| 97 | 75 | ||
| 98 | 76 | ||
| 99 | def test_load_points_npz_rejects_scalar_ancillary(tmp_path: Path) -> None: | 77 | def test_load_points_npz_rejects_scalar_ancillary(tmp_path: pathlib.Path) -> None: |
| 100 | record = _make_record(3, seed=4) | 78 | record = _make_record(3, seed=4) |
| 101 | record["red"] = np.array(42, dtype=np.uint16) | 79 | record["red"] = np.array(42, dtype=np.uint16) |
| 102 | path = tmp_path / "scalar_ancillary.npz" | 80 | path = tmp_path / "scalar_ancillary.npz" |
| 103 | np.savez_compressed(path, **record) | 81 | np.savez_compressed(path, **record) |
| 104 | with pytest.raises(ValueError, match=r"'red' must have shape \(N,\), got \(\)"): | 82 | with pytest.raises(ValueError, match=r"'red' must have shape \(N,\), got \(\)"): |
| 105 | load_points_npz(path) | 83 | segment_points_io.load_points_npz(path) |
| 106 | 84 | ||
| 107 | 85 | ||
| 108 | def test_load_points_npz_rejects_column_vector_ancillary(tmp_path: Path) -> None: | 86 | def test_load_points_npz_rejects_column_vector_ancillary(tmp_path: pathlib.Path) -> None: |
| 109 | record = _make_record(3, seed=5) | 87 | record = _make_record(3, seed=5) |
| 110 | record["intensity"] = record["intensity"].reshape(3, 1) | 88 | record["intensity"] = record["intensity"].reshape(3, 1) |
| 111 | path = tmp_path / "column_ancillary.npz" | 89 | path = tmp_path / "column_ancillary.npz" |
| 112 | np.savez_compressed(path, **record) | 90 | np.savez_compressed(path, **record) |
| 113 | with pytest.raises( | 91 | with pytest.raises( |
| 114 | ValueError, match=r"'intensity' must have shape \(N,\), got \(3, 1\)" | 92 | ValueError, match=r"'intensity' must have shape \(N,\), got \(3, 1\)" |
| 115 | ): | 93 | ): |
| 116 | load_points_npz(path) | 94 | segment_points_io.load_points_npz(path) |
| 117 | 95 | ||
| 118 | 96 | ||
| 119 | def test_save_points_npz_rejects_scalar_ancillary(tmp_path: Path) -> None: | 97 | def test_save_points_npz_rejects_scalar_ancillary(tmp_path: pathlib.Path) -> None: |
| 120 | record = _make_record(2, seed=6) | 98 | record = _make_record(2, seed=6) |
| 121 | record["green"] = np.array(7, dtype=np.uint16) | 99 | record["green"] = np.array(7, dtype=np.uint16) |
| 122 | with pytest.raises(ValueError, match=r"'green' must have shape \(N,\), got \(\)"): | 100 | with pytest.raises(ValueError, match=r"'green' must have shape \(N,\), got \(\)"): |
| 123 | save_points_npz(tmp_path / "out.npz", record) | 101 | segment_points_io.save_points_npz(tmp_path / "out.npz", record) |
| 124 | 102 | ||
| 125 | 103 | ||
| 126 | def test_save_points_npz_rejects_column_vector_ancillary(tmp_path: Path) -> None: | 104 | def test_save_points_npz_rejects_column_vector_ancillary(tmp_path: pathlib.Path) -> None: |
| 127 | record = _make_record(2, seed=7) | 105 | record = _make_record(2, seed=7) |
| 128 | record["scan_angle"] = record["scan_angle"].reshape(2, 1) | 106 | record["scan_angle"] = record["scan_angle"].reshape(2, 1) |
| 129 | with pytest.raises( | 107 | with pytest.raises( |
| 130 | ValueError, match=r"'scan_angle' must have shape \(N,\), got \(2, 1\)" | 108 | ValueError, match=r"'scan_angle' must have shape \(N,\), got \(2, 1\)" |
| 131 | ): | 109 | ): |
| 132 | save_points_npz(tmp_path / "out.npz", record) | 110 | segment_points_io.save_points_npz(tmp_path / "out.npz", record) |
| 133 | 111 | ||
| 134 | 112 | ||
| 135 | def test_save_points_npz_rejects_incomplete_record(tmp_path: Path) -> None: | 113 | def test_save_points_npz_rejects_incomplete_record(tmp_path: pathlib.Path) -> None: |
| 136 | with pytest.raises(ValueError, match="missing required key"): | 114 | with pytest.raises(ValueError, match="missing required key"): |
| 137 | save_points_npz(tmp_path / "out.npz", {"points": np.zeros((1, 3))}) | 115 | segment_points_io.save_points_npz(tmp_path / "out.npz", {"points": np.zeros((1, 3))}) |
| 138 | 116 | ||
| 139 | 117 | ||
| 140 | def test_number_of_returns_is_part_of_the_written_contract() -> None: | 118 | def test_number_of_returns_is_part_of_the_written_contract() -> None: |
| 141 | assert NUMBER_OF_RETURNS_KEY in POINT_RECORD_KEYS | 119 | assert segment_points_io.NUMBER_OF_RETURNS_KEY in segment_points_io.POINT_RECORD_KEYS |
| 142 | assert NUMBER_OF_RETURNS_KEY not in REQUIRED_POINT_RECORD_KEYS | 120 | assert ( |
| 143 | assert set(REQUIRED_POINT_RECORD_KEYS) < set(POINT_RECORD_KEYS) | 121 | segment_points_io.NUMBER_OF_RETURNS_KEY |
| 122 | not in segment_points_io.REQUIRED_POINT_RECORD_KEYS | ||
| 123 | ) | ||
| 124 | assert set(segment_points_io.REQUIRED_POINT_RECORD_KEYS) < set( | ||
| 125 | segment_points_io.POINT_RECORD_KEYS | ||
| 126 | ) | ||
| 144 | 127 | ||
| 145 | 128 | ||
| 146 | def test_save_points_npz_always_writes_number_of_returns(tmp_path: Path) -> None: | 129 | def test_save_points_npz_always_writes_number_of_returns(tmp_path: pathlib.Path) -> None: |
| 147 | record = _make_record(6, seed=60) | 130 | record = _make_record(6, seed=60) |
| 148 | path = save_points_npz(tmp_path / "with_returns_run3_points.npz", record) | 131 | path = segment_points_io.save_points_npz(tmp_path / "with_returns_run3_points.npz", record) |
| 149 | 132 | ||
| 150 | with np.load(path) as data: | 133 | with np.load(path) as data: |
| 151 | assert NUMBER_OF_RETURNS_KEY in data.files | 134 | assert segment_points_io.NUMBER_OF_RETURNS_KEY in data.files |
| 152 | stored = np.asarray(data[NUMBER_OF_RETURNS_KEY]) | 135 | stored = np.asarray(data[segment_points_io.NUMBER_OF_RETURNS_KEY]) |
| 153 | assert stored.dtype == np.uint8 | 136 | assert stored.dtype == np.uint8 |
| 154 | np.testing.assert_array_equal(stored, record[NUMBER_OF_RETURNS_KEY]) | 137 | np.testing.assert_array_equal(stored, record[segment_points_io.NUMBER_OF_RETURNS_KEY]) |
| 155 | 138 | ||
| 156 | 139 | ||
| 157 | def test_save_points_npz_casts_number_of_returns_to_uint8(tmp_path: Path) -> None: | 140 | def test_save_points_npz_casts_number_of_returns_to_uint8(tmp_path: pathlib.Path) -> None: |
| 158 | record = _make_record(4, seed=61) | 141 | record = _make_record(4, seed=61) |
| 159 | record[NUMBER_OF_RETURNS_KEY] = record[NUMBER_OF_RETURNS_KEY].astype(np.int64) | 142 | record[segment_points_io.NUMBER_OF_RETURNS_KEY] = record[ |
| 160 | path = save_points_npz(tmp_path / "cast_run3_points.npz", record) | 143 | segment_points_io.NUMBER_OF_RETURNS_KEY |
| 144 | ].astype(np.int64) | ||
| 145 | path = segment_points_io.save_points_npz(tmp_path / "cast_run3_points.npz", record) | ||
| 161 | 146 | ||
| 162 | loaded = load_points_npz(path) | 147 | loaded = segment_points_io.load_points_npz(path) |
| 163 | assert loaded[NUMBER_OF_RETURNS_KEY].dtype == np.uint8 | 148 | assert loaded[segment_points_io.NUMBER_OF_RETURNS_KEY].dtype == np.uint8 |
| 164 | 149 | ||
| 165 | 150 | ||
| 166 | def test_save_points_npz_rejects_out_of_range_number_of_returns(tmp_path: Path) -> None: | 151 | def test_save_points_npz_rejects_out_of_range_number_of_returns(tmp_path: pathlib.Path) -> None: |
| 167 | record = _make_record(3, seed=62) | 152 | record = _make_record(3, seed=62) |
| 168 | record[NUMBER_OF_RETURNS_KEY] = np.array([1, -1, 3], dtype=np.int16) | 153 | record[segment_points_io.NUMBER_OF_RETURNS_KEY] = np.array([1, -1, 3], dtype=np.int16) |
| 169 | with pytest.raises(ValueError, match="number_of_returns"): | 154 | with pytest.raises(ValueError, match="number_of_returns"): |
| 170 | save_points_npz(tmp_path / "out.npz", record) | 155 | segment_points_io.save_points_npz(tmp_path / "out.npz", record) |
| 171 | 156 | ||
| 172 | 157 | ||
| 173 | def test_save_points_npz_rejects_missing_number_of_returns(tmp_path: Path) -> None: | 158 | def test_save_points_npz_rejects_missing_number_of_returns(tmp_path: pathlib.Path) -> None: |
| 174 | with pytest.raises(ValueError, match="number_of_returns"): | 159 | with pytest.raises(ValueError, match="number_of_returns"): |
| 175 | save_points_npz(tmp_path / "out.npz", _make_legacy_record(3, seed=63)) | 160 | segment_points_io.save_points_npz(tmp_path / "out.npz", _make_legacy_record(3, seed=63)) |
| 176 | 161 | ||
| 177 | 162 | ||
| 178 | def test_load_points_npz_fills_zeros_for_legacy_records(tmp_path: Path) -> None: | 163 | def test_load_points_npz_fills_zeros_for_legacy_records(tmp_path: pathlib.Path) -> None: |
| 179 | """Datasets written before AI3D-382 lack the key; 0 means unknown, never 1.""" | 164 | """Datasets written before AI3D-382 lack the key; 0 means unknown, never 1.""" |
| 180 | legacy = _make_legacy_record(5, seed=64) | 165 | legacy = _make_legacy_record(5, seed=64) |
| 181 | path = tmp_path / "legacy_run3_points.npz" | 166 | path = tmp_path / "legacy_run3_points.npz" |
| 182 | np.savez_compressed(path, **legacy) | 167 | np.savez_compressed(path, **legacy) |
| 183 | 168 | ||
| 184 | loaded = load_points_npz(path) | 169 | loaded = segment_points_io.load_points_npz(path) |
| 185 | 170 | ||
| 186 | assert list(loaded.keys()) == list(POINT_RECORD_KEYS) | 171 | assert list(loaded.keys()) == list(segment_points_io.POINT_RECORD_KEYS) |
| 187 | returns = loaded[NUMBER_OF_RETURNS_KEY] | 172 | returns = loaded[segment_points_io.NUMBER_OF_RETURNS_KEY] |
| 188 | assert returns.shape == (5,) | 173 | assert returns.shape == (5,) |
| 189 | assert returns.dtype == np.uint8 | 174 | assert returns.dtype == np.uint8 |
| 190 | np.testing.assert_array_equal(returns, np.zeros(5, dtype=np.uint8)) | 175 | np.testing.assert_array_equal(returns, np.zeros(5, dtype=np.uint8)) |
| 191 | 176 | ||
| 192 | 177 | ||
| 193 | def test_load_points_npz_rejects_bad_number_of_returns_shape(tmp_path: Path) -> None: | 178 | def test_load_points_npz_rejects_bad_number_of_returns_shape(tmp_path: pathlib.Path) -> None: |
| 194 | record = _make_record(4, seed=65) | 179 | record = _make_record(4, seed=65) |
| 195 | record[NUMBER_OF_RETURNS_KEY] = record[NUMBER_OF_RETURNS_KEY][:2] | 180 | record[segment_points_io.NUMBER_OF_RETURNS_KEY] = record[ |
| 181 | segment_points_io.NUMBER_OF_RETURNS_KEY | ||
| 182 | ][:2] | ||
| 196 | path = tmp_path / "bad_returns.npz" | 183 | path = tmp_path / "bad_returns.npz" |
| 197 | np.savez_compressed(path, **record) | 184 | np.savez_compressed(path, **record) |
| 198 | with pytest.raises( | 185 | with pytest.raises( |
| 199 | ValueError, match=r"'number_of_returns' must have shape \(N,\), got \(2,\)" | 186 | ValueError, match=r"'number_of_returns' must have shape \(N,\), got \(2,\)" |
| 200 | ): | 187 | ): |
| 201 | load_points_npz(path) | 188 | segment_points_io.load_points_npz(path) |
| 202 | 189 | ||
| 203 | 190 | ||
| 204 | def test_save_points_npz_rejects_bool_number_of_returns(tmp_path: Path) -> None: | 191 | def test_save_points_npz_rejects_bool_number_of_returns(tmp_path: pathlib.Path) -> None: |
| 205 | """A bool mask is not a return count; casting it would fabricate 0/1 counts.""" | 192 | """A bool mask is not a return count; casting it would fabricate 0/1 counts.""" |
| 206 | record = _make_record(3, seed=68) | 193 | record = _make_record(3, seed=68) |
| 207 | record[NUMBER_OF_RETURNS_KEY] = np.array([True, False, True]) | 194 | record[segment_points_io.NUMBER_OF_RETURNS_KEY] = np.array([True, False, True]) |
| 208 | with pytest.raises(ValueError, match="number_of_returns.*bool"): | 195 | with pytest.raises(ValueError, match="number_of_returns.*bool"): |
| 209 | save_points_npz(tmp_path / "out.npz", record) | 196 | segment_points_io.save_points_npz(tmp_path / "out.npz", record) |
| 210 | 197 | ||
| 211 | 198 | ||
| 212 | def test_load_points_npz_rejects_bool_number_of_returns(tmp_path: Path) -> None: | 199 | def test_load_points_npz_rejects_bool_number_of_returns(tmp_path: pathlib.Path) -> None: |
| 213 | record = _make_record(3, seed=69) | 200 | record = _make_record(3, seed=69) |
| 214 | record[NUMBER_OF_RETURNS_KEY] = np.array([True, False, True]) | 201 | record[segment_points_io.NUMBER_OF_RETURNS_KEY] = np.array([True, False, True]) |
| 215 | path = tmp_path / "bool_returns_run3_points.npz" | 202 | path = tmp_path / "bool_returns_run3_points.npz" |
| 216 | np.savez_compressed(path, **record) | 203 | np.savez_compressed(path, **record) |
| 217 | with pytest.raises(ValueError, match="number_of_returns.*bool"): | 204 | with pytest.raises(ValueError, match="number_of_returns.*bool"): |
| 218 | load_points_npz(path) | 205 | segment_points_io.load_points_npz(path) |
| 219 | 206 | ||
| 220 | 207 | ||
| 221 | def test_load_points_npz_casts_stored_number_of_returns_to_uint8(tmp_path: Path) -> None: | 208 | def test_load_points_npz_casts_stored_number_of_returns_to_uint8(tmp_path: pathlib.Path) -> None: |
| 222 | """A hand-rolled producer's wider dtype must not leak into merges.""" | 209 | """A hand-rolled producer's wider dtype must not leak into merges.""" |
| 223 | record = _make_record(5, seed=70) | 210 | record = _make_record(5, seed=70) |
| 224 | stored = record[NUMBER_OF_RETURNS_KEY].astype(np.int32) | 211 | stored = record[segment_points_io.NUMBER_OF_RETURNS_KEY].astype(np.int32) |
| 225 | record[NUMBER_OF_RETURNS_KEY] = stored | 212 | record[segment_points_io.NUMBER_OF_RETURNS_KEY] = stored |
| 226 | path = tmp_path / "int32_returns_run3_points.npz" | 213 | path = tmp_path / "int32_returns_run3_points.npz" |
| 227 | np.savez_compressed(path, **record) | 214 | np.savez_compressed(path, **record) |
| 228 | 215 | ||
| 229 | loaded = load_points_npz(path) | 216 | loaded = segment_points_io.load_points_npz(path) |
| 230 | 217 | ||
| 231 | assert loaded[NUMBER_OF_RETURNS_KEY].dtype == np.uint8 | 218 | assert loaded[segment_points_io.NUMBER_OF_RETURNS_KEY].dtype == np.uint8 |
| 232 | np.testing.assert_array_equal(loaded[NUMBER_OF_RETURNS_KEY], stored) | 219 | np.testing.assert_array_equal(loaded[segment_points_io.NUMBER_OF_RETURNS_KEY], stored) |
| 233 | 220 | ||
| 234 | 221 | ||
| 235 | def test_load_points_npz_rejects_out_of_range_stored_number_of_returns( | 222 | def test_load_points_npz_rejects_out_of_range_stored_number_of_returns( |
| 236 | tmp_path: Path, | 223 | tmp_path: pathlib.Path, |
| 237 | ) -> None: | 224 | ) -> None: |
| 238 | record = _make_record(3, seed=71) | 225 | record = _make_record(3, seed=71) |
| 239 | record[NUMBER_OF_RETURNS_KEY] = np.array([1, 300, 3], dtype=np.int32) | 226 | record[segment_points_io.NUMBER_OF_RETURNS_KEY] = np.array([1, 300, 3], dtype=np.int32) |
| 240 | path = tmp_path / "out_of_range_returns_run3_points.npz" | 227 | path = tmp_path / "out_of_range_returns_run3_points.npz" |
| 241 | np.savez_compressed(path, **record) | 228 | np.savez_compressed(path, **record) |
| 242 | with pytest.raises(ValueError, match="number_of_returns.*fit in uint8"): | 229 | with pytest.raises(ValueError, match="number_of_returns.*fit in uint8"): |
| 243 | load_points_npz(path) | 230 | segment_points_io.load_points_npz(path) |
| 244 | 231 | ||
| 245 | 232 | ||
| 246 | def test_load_points_npz_rejects_float_stored_number_of_returns(tmp_path: Path) -> None: | 233 | def test_load_points_npz_rejects_float_stored_number_of_returns(tmp_path: pathlib.Path) -> None: |
| 247 | record = _make_record(3, seed=72) | 234 | record = _make_record(3, seed=72) |
| 248 | record[NUMBER_OF_RETURNS_KEY] = np.array([1.0, 2.0, 3.0], dtype=np.float32) | 235 | record[segment_points_io.NUMBER_OF_RETURNS_KEY] = np.array([1.0, 2.0, 3.0], dtype=np.float32) |
| 249 | path = tmp_path / "float_returns_run3_points.npz" | 236 | path = tmp_path / "float_returns_run3_points.npz" |
| 250 | np.savez_compressed(path, **record) | 237 | np.savez_compressed(path, **record) |
| 251 | with pytest.raises(ValueError, match="number_of_returns.*integer array"): | 238 | with pytest.raises(ValueError, match="number_of_returns.*integer array"): |
| 252 | load_points_npz(path) | 239 | segment_points_io.load_points_npz(path) |
| 253 | 240 | ||
| 254 | 241 | ||
| 255 | def test_concat_points_npz_merges_mixed_stored_return_dtypes(tmp_path: Path) -> None: | 242 | def test_concat_points_npz_merges_mixed_stored_return_dtypes(tmp_path: pathlib.Path) -> None: |
| 256 | """int32-stored and uint8-stored records concatenate after the load coercion.""" | 243 | """int32-stored and uint8-stored records concatenate after the load coercion.""" |
| 257 | wide = _make_record(3, seed=73) | 244 | wide = _make_record(3, seed=73) |
| 258 | wide[NUMBER_OF_RETURNS_KEY] = wide[NUMBER_OF_RETURNS_KEY].astype(np.int32) | 245 | wide[segment_points_io.NUMBER_OF_RETURNS_KEY] = wide[ |
| 246 | segment_points_io.NUMBER_OF_RETURNS_KEY | ||
| 247 | ].astype(np.int32) | ||
| 259 | wide_path = tmp_path / "alpha_run3_points.npz" | 248 | wide_path = tmp_path / "alpha_run3_points.npz" |
| 260 | np.savez_compressed(wide_path, **wide) | 249 | np.savez_compressed(wide_path, **wide) |
| 261 | narrow_path = save_points_npz(tmp_path / "beta_run3_points.npz", _make_record(2, seed=74)) | 250 | narrow_path = segment_points_io.save_points_npz( |
| 251 | tmp_path / "beta_run3_points.npz", _make_record(2, seed=74) | ||
| 252 | ) | ||
| 262 | 253 | ||
| 263 | merged, _ = concat_points_npz([wide_path, narrow_path]) | 254 | merged, _ = segment_points_io.concat_points_npz([wide_path, narrow_path]) |
| 264 | 255 | ||
| 265 | assert merged[NUMBER_OF_RETURNS_KEY].dtype == np.uint8 | 256 | assert merged[segment_points_io.NUMBER_OF_RETURNS_KEY].dtype == np.uint8 |
| 266 | assert merged[NUMBER_OF_RETURNS_KEY].shape == (5,) | 257 | assert merged[segment_points_io.NUMBER_OF_RETURNS_KEY].shape == (5,) |
| 267 | 258 | ||
| 268 | 259 | ||
| 269 | def test_load_segment_points_merges_legacy_and_new_records(tmp_path: Path) -> None: | 260 | def test_load_segment_points_merges_legacy_and_new_records(tmp_path: pathlib.Path) -> None: |
| 270 | legacy = _make_legacy_record(3, seed=66) | 261 | legacy = _make_legacy_record(3, seed=66) |
| 271 | modern = _make_record(4, seed=67) | 262 | modern = _make_record(4, seed=67) |
| 272 | legacy_path = tmp_path / "alpha_run3_points.npz" | 263 | legacy_path = tmp_path / "alpha_run3_points.npz" |
| 273 | np.savez_compressed(legacy_path, **legacy) | 264 | np.savez_compressed(legacy_path, **legacy) |
| 274 | modern_path = tmp_path / "beta_run3_points.npz" | 265 | modern_path = tmp_path / "beta_run3_points.npz" |
| 275 | save_points_npz(modern_path, modern) | 266 | segment_points_io.save_points_npz(modern_path, modern) |
| 276 | 267 | ||
| 277 | merged, _, _ = load_segment_points([legacy_path, modern_path]) | 268 | merged, _, _ = segment_points_io.load_segment_points([legacy_path, modern_path]) |
| 278 | 269 | ||
| 279 | np.testing.assert_array_equal( | 270 | np.testing.assert_array_equal( |
| 280 | merged[NUMBER_OF_RETURNS_KEY], | 271 | merged[segment_points_io.NUMBER_OF_RETURNS_KEY], |
| 281 | np.concatenate( | 272 | np.concatenate( |
| 282 | [np.zeros(3, dtype=np.uint8), modern[NUMBER_OF_RETURNS_KEY]], axis=0 | 273 | [np.zeros(3, dtype=np.uint8), modern[segment_points_io.NUMBER_OF_RETURNS_KEY]], axis=0 |
| 283 | ), | 274 | ), |
| 284 | ) | 275 | ) |
| 285 | 276 | ||
| 286 | 277 | ||
| 287 | def test_concat_points_npz_carries_number_of_returns(tmp_path: Path) -> None: | 278 | def test_concat_points_npz_carries_number_of_returns(tmp_path: pathlib.Path) -> None: |
| 288 | first = _make_record(3, seed=68) | 279 | first = _make_record(3, seed=68) |
| 289 | second = _make_legacy_record(2, seed=69) | 280 | second = _make_legacy_record(2, seed=69) |
| 290 | save_points_npz(tmp_path / "a_run3_points.npz", first) | 281 | segment_points_io.save_points_npz(tmp_path / "a_run3_points.npz", first) |
| 291 | np.savez_compressed(tmp_path / "b_run3_points.npz", **second) | 282 | np.savez_compressed(tmp_path / "b_run3_points.npz", **second) |
| 292 | 283 | ||
| 293 | merged, spans = concat_points_npz( | 284 | merged, spans = segment_points_io.concat_points_npz( |
| 294 | [tmp_path / "a_run3_points.npz", tmp_path / "b_run3_points.npz"] | 285 | [tmp_path / "a_run3_points.npz", tmp_path / "b_run3_points.npz"] |
| 295 | ) | 286 | ) |
| 296 | 287 | ||
| 297 | assert [span.count for span in spans] == [3, 2] | 288 | assert [span.count for span in spans] == [3, 2] |
| 298 | assert merged[NUMBER_OF_RETURNS_KEY].dtype == np.uint8 | 289 | assert merged[segment_points_io.NUMBER_OF_RETURNS_KEY].dtype == np.uint8 |
| 299 | np.testing.assert_array_equal( | 290 | np.testing.assert_array_equal( |
| 300 | merged[NUMBER_OF_RETURNS_KEY], | 291 | merged[segment_points_io.NUMBER_OF_RETURNS_KEY], |
| 301 | np.concatenate( | 292 | np.concatenate( |
| 302 | [first[NUMBER_OF_RETURNS_KEY], np.zeros(2, dtype=np.uint8)], axis=0 | 293 | [first[segment_points_io.NUMBER_OF_RETURNS_KEY], np.zeros(2, dtype=np.uint8)], axis=0 |
| 303 | ), | 294 | ), |
| 304 | ) | 295 | ) |
| 305 | 296 | ||
| 306 | 297 | ||
| 307 | def test_load_run3_segment_carries_number_of_returns(tmp_path: Path) -> None: | 298 | def test_load_run3_segment_carries_number_of_returns(tmp_path: pathlib.Path) -> None: |
| 308 | seg_dir = tmp_path / "segment_011" | 299 | seg_dir = tmp_path / "segment_011" |
| 309 | seg_dir.mkdir() | 300 | seg_dir.mkdir() |
| 310 | first = _make_record(2, seed=70) | 301 | first = _make_record(2, seed=70) |
| 311 | second = _make_record(3, seed=71) | 302 | second = _make_record(3, seed=71) |
| 312 | save_points_npz(seg_dir / "a_run3_points.npz", first) | 303 | segment_points_io.save_points_npz(seg_dir / "a_run3_points.npz", first) |
| 313 | save_points_npz(seg_dir / "b_run3_points.npz", second) | 304 | segment_points_io.save_points_npz(seg_dir / "b_run3_points.npz", second) |
| 314 | 305 | ||
| 315 | merged, _ = load_run3_segment(seg_dir) | 306 | merged, _ = segment_points_io.load_run3_segment(seg_dir) |
| 316 | 307 | ||
| 317 | np.testing.assert_array_equal( | 308 | np.testing.assert_array_equal( |
| 318 | merged[NUMBER_OF_RETURNS_KEY], | 309 | merged[segment_points_io.NUMBER_OF_RETURNS_KEY], |
| 319 | np.concatenate( | 310 | np.concatenate( |
| 320 | [first[NUMBER_OF_RETURNS_KEY], second[NUMBER_OF_RETURNS_KEY]], axis=0 | 311 | [ |
| 312 | first[segment_points_io.NUMBER_OF_RETURNS_KEY], | ||
| 313 | second[segment_points_io.NUMBER_OF_RETURNS_KEY], | ||
| 314 | ], | ||
| 315 | axis=0, | ||
| 321 | ), | 316 | ), |
| 322 | ) | 317 | ) |
| 323 | 318 | ||
| 324 | 319 | ||
| 325 | def test_load_segment_points_merges_and_tracks_file_ids(tmp_path: Path) -> None: | 320 | def test_load_segment_points_merges_and_tracks_file_ids(tmp_path: pathlib.Path) -> None: |
| 326 | records = [_make_record(3, seed=10), _make_record(5, seed=11)] | 321 | records = [_make_record(3, seed=10), _make_record(5, seed=11)] |
| 327 | paths = [ | 322 | paths = [ |
| 328 | tmp_path / "alpha_run3_points.npz", | 323 | tmp_path / "alpha_run3_points.npz", |
| 329 | tmp_path / "beta_run3_points.npz", | 324 | tmp_path / "beta_run3_points.npz", |
| 330 | ] | 325 | ] |
| 331 | for path, record in zip(paths, records, strict=True): | 326 | for path, record in zip(paths, records, strict=True): |
| 332 | save_points_npz(path, record) | 327 | segment_points_io.save_points_npz(path, record) |
| 333 | 328 | ||
| 334 | merged, point_file_ids, file_stems = load_segment_points(paths) | 329 | merged, point_file_ids, file_stems = segment_points_io.load_segment_points(paths) |
| 335 | 330 | ||
| 336 | assert file_stems == ["alpha_run3_points", "beta_run3_points"] | 331 | assert file_stems == ["alpha_run3_points", "beta_run3_points"] |
| 337 | assert merged["points"].shape == (8, 3) | 332 | assert merged["points"].shape == (8, 3) |
| 338 | assert point_file_ids.dtype == np.int32 | 333 | assert point_file_ids.dtype == np.int32 |
| 351 | 346 | ||
| 352 | 347 | ||
| 353 | def test_load_segment_points_empty_raises() -> None: | 348 | def test_load_segment_points_empty_raises() -> None: |
| 354 | with pytest.raises(FileNotFoundError, match="No \\*_points.npz"): | 349 | with pytest.raises(FileNotFoundError, match="No \\*_points.npz"): |
| 355 | load_segment_points([]) | 350 | segment_points_io.load_segment_points([]) |
| 356 | 351 | ||
| 357 | 352 | ||
| 358 | def test_load_geoshift_parses_xyz_json(tmp_path: Path) -> None: | 353 | def test_load_geoshift_parses_xyz_json(tmp_path: pathlib.Path) -> None: |
| 359 | path = tmp_path / "run3_geoshift.json" | 354 | path = tmp_path / "run3_geoshift.json" |
| 360 | path.write_text(json.dumps({"x": 725883.5, "y": 5422097.8, "z": 390.8}), encoding="utf-8") | 355 | path.write_text(json.dumps({"x": 725883.5, "y": 5422097.8, "z": 390.8}), encoding="utf-8") |
| 361 | geoshift = load_geoshift(path) | 356 | geoshift = segment_points_io.load_geoshift(path) |
| 362 | assert geoshift.shape == (3,) | 357 | assert geoshift.shape == (3,) |
| 363 | assert geoshift.dtype == np.float64 | 358 | assert geoshift.dtype == np.float64 |
| 364 | np.testing.assert_allclose(geoshift, [725883.5, 5422097.8, 390.8]) | 359 | np.testing.assert_allclose(geoshift, [725883.5, 5422097.8, 390.8]) |
| 365 | 360 | ||
| 366 | 361 | ||
| 367 | def test_load_geoshift_rejects_missing_keys(tmp_path: Path) -> None: | 362 | def test_load_geoshift_rejects_missing_keys(tmp_path: pathlib.Path) -> None: |
| 368 | path = tmp_path / "run3_geoshift.json" | 363 | path = tmp_path / "run3_geoshift.json" |
| 369 | path.write_text(json.dumps({"x": 1.0, "y": 2.0}), encoding="utf-8") | 364 | path.write_text(json.dumps({"x": 1.0, "y": 2.0}), encoding="utf-8") |
| 370 | with pytest.raises(ValueError, match="missing geoshift key"): | 365 | with pytest.raises(ValueError, match="missing geoshift key"): |
| 371 | load_geoshift(path) | 366 | segment_points_io.load_geoshift(path) |
| 372 | 367 | ||
| 373 | 368 | ||
| 374 | def test_parse_segment_key_accepts_int_and_prefixed_forms() -> None: | 369 | def test_parse_segment_key_accepts_int_and_prefixed_forms() -> None: |
| 375 | assert parse_segment_key(32) == 32 | 370 | assert segment_points_io.parse_segment_key(32) == 32 |
| 376 | assert parse_segment_key("32") == 32 | 371 | assert segment_points_io.parse_segment_key("32") == 32 |
| 377 | assert parse_segment_key("segment_33") == 33 | 372 | assert segment_points_io.parse_segment_key("segment_33") == 33 |
| 378 | assert parse_segment_key("segment_066") == 66 | 373 | assert segment_points_io.parse_segment_key("segment_066") == 66 |
| 379 | 374 | ||
| 380 | 375 | ||
| 381 | def test_parse_segment_key_rejects_invalid() -> None: | 376 | def test_parse_segment_key_rejects_invalid() -> None: |
| 382 | with pytest.raises(ValueError, match="Segment key must not be empty"): | 377 | with pytest.raises(ValueError, match="Segment key must not be empty"): |
| 383 | parse_segment_key("segment_") | 378 | segment_points_io.parse_segment_key("segment_") |
| 384 | with pytest.raises(ValueError, match="integer or segment_<idx>"): | 379 | with pytest.raises(ValueError, match="integer or segment_<idx>"): |
| 385 | parse_segment_key("lane_a") | 380 | segment_points_io.parse_segment_key("lane_a") |
| 386 | with pytest.raises(ValueError, match="integer or segment_<idx>"): | 381 | with pytest.raises(ValueError, match="integer or segment_<idx>"): |
| 387 | parse_segment_key(True) | 382 | segment_points_io.parse_segment_key(True) |
| 388 | 383 | ||
| 389 | 384 | ||
| 390 | def test_normalize_segment_file_blacklist_accepts_multiple_key_formats() -> None: | 385 | def test_normalize_segment_file_blacklist_accepts_multiple_key_formats() -> None: |
| 391 | blacklist = normalize_segment_file_blacklist( | 386 | blacklist = segment_points_io.normalize_segment_file_blacklist( |
| 392 | { | 387 | { |
| 393 | "32": ["scan_a_run3_points.npz", "nested/scan_b_run3_points.npz"], | 388 | "32": ["scan_a_run3_points.npz", "nested/scan_b_run3_points.npz"], |
| 394 | "segment_33": "scan_c_run3_points.npz", | 389 | "segment_33": "scan_c_run3_points.npz", |
| 395 | 33: ["scan_d_run3_points.npz"], | 390 | 33: ["scan_d_run3_points.npz"], |
| 402 | 397 | ||
| 403 | 398 | ||
| 404 | def test_normalize_segment_file_blacklist_rejects_invalid_values() -> None: | 399 | def test_normalize_segment_file_blacklist_rejects_invalid_values() -> None: |
| 405 | with pytest.raises(ValueError, match="mapping of segment to files"): | 400 | with pytest.raises(ValueError, match="mapping of segment to files"): |
| 406 | normalize_segment_file_blacklist(["32:scan_a_run3_points.npz"]) # type: ignore[arg-type] | 401 | segment_points_io.normalize_segment_file_blacklist(["32:scan_a_run3_points.npz"]) # type: ignore[arg-type] |
| 407 | with pytest.raises(ValueError, match="at least one non-empty"): | 402 | with pytest.raises(ValueError, match="at least one non-empty"): |
| 408 | normalize_segment_file_blacklist({"32": [" ", ""]}) | 403 | segment_points_io.normalize_segment_file_blacklist({"32": [" ", ""]}) |
| 409 | assert normalize_segment_file_blacklist(None) == {} | 404 | assert segment_points_io.normalize_segment_file_blacklist(None) == {} |
| 410 | 405 | ||
| 411 | 406 | ||
| 412 | def test_filter_segment_files_respects_segment_specific_exact_and_glob_rules( | 407 | def test_filter_segment_files_respects_segment_specific_exact_and_glob_rules( |
| 413 | caplog: pytest.LogCaptureFixture, | 408 | caplog: pytest.LogCaptureFixture, |
| 414 | ) -> None: | 409 | ) -> None: |
| 415 | segment_32_files = [ | 410 | segment_32_files = [ |
| 416 | Path("/tmp/lane_points/segment_32/scan_a_run3_points.npz"), | 411 | pathlib.Path("/tmp/lane_points/segment_32/scan_a_run3_points.npz"), |
| 417 | Path("/tmp/lane_points/segment_32/nested/scan_b_run3_points.npz"), | 412 | pathlib.Path("/tmp/lane_points/segment_32/nested/scan_b_run3_points.npz"), |
| 418 | Path("/tmp/lane_points/segment_32/scan_c_run3_points.npz"), | 413 | pathlib.Path("/tmp/lane_points/segment_32/scan_c_run3_points.npz"), |
| 419 | ] | 414 | ] |
| 420 | blacklist = normalize_segment_file_blacklist( | 415 | blacklist = segment_points_io.normalize_segment_file_blacklist( |
| 421 | { | 416 | { |
| 422 | "32": [ | 417 | "32": [ |
| 423 | "scan_a_run3_points.npz", | 418 | "scan_a_run3_points.npz", |
| 424 | "*/nested/scan_b_run3_points.npz", | 419 | "*/nested/scan_b_run3_points.npz", |
| 427 | } | 422 | } |
| 428 | ) | 423 | ) |
| 429 | 424 | ||
| 430 | with caplog.at_level(logging.INFO, logger="iolabs.common.segment_points_io"): | 425 | with caplog.at_level(logging.INFO, logger="iolabs.common.segment_points_io"): |
| 431 | kept_32, blacklisted_32 = filter_segment_files( | 426 | kept_32, blacklisted_32 = segment_points_io.filter_segment_files( |
| 432 | segment_32_files, | 427 | segment_32_files, |
| 433 | segment_index=32, | 428 | segment_index=32, |
| 434 | blacklist=blacklist, | 429 | blacklist=blacklist, |
| 435 | ) | 430 | ) |
| 436 | kept_33, blacklisted_33 = filter_segment_files( | 431 | kept_33, blacklisted_33 = segment_points_io.filter_segment_files( |
| 437 | segment_32_files, | 432 | segment_32_files, |
| 438 | segment_index=33, | 433 | segment_index=33, |
| 439 | blacklist=blacklist, | 434 | blacklist=blacklist, |
| 440 | ) | 435 | ) |
| 451 | assert [path.name for path in blacklisted_33] == ["scan_c_run3_points.npz"] | 446 | assert [path.name for path in blacklisted_33] == ["scan_c_run3_points.npz"] |
| 452 | assert any("excluding 2 input NPZ files" in record.message for record in caplog.records) | 447 | assert any("excluding 2 input NPZ files" in record.message for record in caplog.records) |
| 453 | 448 | ||
| 454 | 449 | ||
| 455 | def _write_geoshift(path: Path, xyz: tuple[float, float, float]) -> Path: | 450 | def _write_geoshift(path: pathlib.Path, xyz: tuple[float, float, float]) -> pathlib.Path: |
| 456 | path.parent.mkdir(parents=True, exist_ok=True) | 451 | path.parent.mkdir(parents=True, exist_ok=True) |
| 457 | path.write_text( | 452 | path.write_text( |
| 458 | json.dumps({"x": xyz[0], "y": xyz[1], "z": xyz[2]}), encoding="utf-8" | 453 | json.dumps({"x": xyz[0], "y": xyz[1], "z": xyz[2]}), encoding="utf-8" |
| 459 | ) | 454 | ) |
| 460 | return path | 455 | return path |
| 461 | 456 | ||
| 462 | 457 | ||
| 463 | def _save_stored_npz(path: Path, record: dict[str, np.ndarray]) -> Path: | 458 | def _save_stored_npz(path: pathlib.Path, record: dict[str, np.ndarray]) -> pathlib.Path: |
| 464 | """Write an uncompressed (ZIP_STORED) npz, the chunk-streamable layout.""" | 459 | """Write an uncompressed (ZIP_STORED) npz, the chunk-streamable layout.""" |
| 465 | np.savez(path, **record) | 460 | np.savez(path, **record) |
| 466 | return path | 461 | return path |
| 467 | 462 | ||
| 468 | 463 | ||
| 469 | def test_geoshift_candidate_paths_cover_the_three_conventions(tmp_path: Path) -> None: | 464 | def test_geoshift_candidate_paths_cover_the_three_conventions(tmp_path: pathlib.Path) -> None: |
| 470 | candidates = geoshift_candidate_paths(tmp_path / "dataset" / "lane_points") | 465 | candidates = segment_points_io.geoshift_candidate_paths(tmp_path / "dataset" / "lane_points") |
| 471 | assert candidates == [ | 466 | assert candidates == [ |
| 472 | tmp_path / "dataset" / "lane_points" / GEOSHIFT_NAME, | 467 | tmp_path / "dataset" / "lane_points" / segment_points_io.GEOSHIFT_NAME, |
| 473 | tmp_path / "dataset" / "lane_points" / "lane_points" / GEOSHIFT_NAME, | 468 | tmp_path / "dataset" / "lane_points" / "lane_points" / segment_points_io.GEOSHIFT_NAME, |
| 474 | tmp_path / "dataset" / GEOSHIFT_NAME, | 469 | tmp_path / "dataset" / segment_points_io.GEOSHIFT_NAME, |
| 475 | ] | 470 | ] |
| 476 | 471 | ||
| 477 | 472 | ||
| 478 | def test_find_geoshift_lane_points_dir_convention(tmp_path: Path) -> None: | 473 | def test_find_geoshift_lane_points_dir_convention(tmp_path: pathlib.Path) -> None: |
| 479 | lane_points = tmp_path / "lane_points" | 474 | lane_points = tmp_path / "lane_points" |
| 480 | _write_geoshift(lane_points / GEOSHIFT_NAME, (1.0, 2.0, 3.0)) | 475 | _write_geoshift(lane_points / segment_points_io.GEOSHIFT_NAME, (1.0, 2.0, 3.0)) |
| 481 | np.testing.assert_allclose(find_geoshift(lane_points), [1.0, 2.0, 3.0]) | 476 | np.testing.assert_allclose(segment_points_io.find_geoshift(lane_points), [1.0, 2.0, 3.0]) |
| 482 | 477 | ||
| 483 | 478 | ||
| 484 | def test_find_geoshift_dataset_root_convention(tmp_path: Path) -> None: | 479 | def test_find_geoshift_dataset_root_convention(tmp_path: pathlib.Path) -> None: |
| 485 | _write_geoshift(tmp_path / "lane_points" / GEOSHIFT_NAME, (4.0, 5.0, 6.0)) | 480 | _write_geoshift(tmp_path / "lane_points" / segment_points_io.GEOSHIFT_NAME, (4.0, 5.0, 6.0)) |
| 486 | np.testing.assert_allclose(find_geoshift(tmp_path), [4.0, 5.0, 6.0]) | 481 | np.testing.assert_allclose(segment_points_io.find_geoshift(tmp_path), [4.0, 5.0, 6.0]) |
| 487 | 482 | ||
| 488 | 483 | ||
| 489 | def test_find_geoshift_segment_dir_parent_convention(tmp_path: Path) -> None: | 484 | def test_find_geoshift_segment_dir_parent_convention(tmp_path: pathlib.Path) -> None: |
| 490 | lane_points = tmp_path / "lane_points" | 485 | lane_points = tmp_path / "lane_points" |
| 491 | seg_dir = lane_points / "segment_032" | 486 | seg_dir = lane_points / "segment_032" |
| 492 | seg_dir.mkdir(parents=True) | 487 | seg_dir.mkdir(parents=True) |
| 493 | _write_geoshift(lane_points / GEOSHIFT_NAME, (7.0, 8.0, 9.0)) | 488 | _write_geoshift(lane_points / segment_points_io.GEOSHIFT_NAME, (7.0, 8.0, 9.0)) |
| 494 | np.testing.assert_allclose(find_geoshift(seg_dir), [7.0, 8.0, 9.0]) | 489 | np.testing.assert_allclose(segment_points_io.find_geoshift(seg_dir), [7.0, 8.0, 9.0]) |
| 495 | 490 | ||
| 496 | 491 | ||
| 497 | def test_find_geoshift_prefers_directory_over_parent(tmp_path: Path) -> None: | 492 | def test_find_geoshift_prefers_directory_over_parent(tmp_path: pathlib.Path) -> None: |
| 498 | lane_points = tmp_path / "lane_points" | 493 | lane_points = tmp_path / "lane_points" |
| 499 | seg_dir = lane_points / "segment_032" | 494 | seg_dir = lane_points / "segment_032" |
| 500 | seg_dir.mkdir(parents=True) | 495 | seg_dir.mkdir(parents=True) |
| 501 | _write_geoshift(seg_dir / GEOSHIFT_NAME, (1.0, 1.0, 1.0)) | 496 | _write_geoshift(seg_dir / segment_points_io.GEOSHIFT_NAME, (1.0, 1.0, 1.0)) |
| 502 | _write_geoshift(lane_points / GEOSHIFT_NAME, (2.0, 2.0, 2.0)) | 497 | _write_geoshift(lane_points / segment_points_io.GEOSHIFT_NAME, (2.0, 2.0, 2.0)) |
| 503 | np.testing.assert_allclose(find_geoshift(seg_dir), [1.0, 1.0, 1.0]) | 498 | np.testing.assert_allclose(segment_points_io.find_geoshift(seg_dir), [1.0, 1.0, 1.0]) |
| 504 | 499 | ||
| 505 | 500 | ||
| 506 | def test_find_geoshift_or_none_returns_none_when_absent(tmp_path: Path) -> None: | 501 | def test_find_geoshift_or_none_returns_none_when_absent(tmp_path: pathlib.Path) -> None: |
| 507 | seg_dir = tmp_path / "lane_points" / "segment_000" | 502 | seg_dir = tmp_path / "lane_points" / "segment_000" |
| 508 | seg_dir.mkdir(parents=True) | 503 | seg_dir.mkdir(parents=True) |
| 509 | assert find_geoshift_or_none(seg_dir) is None | 504 | assert segment_points_io.find_geoshift_or_none(seg_dir) is None |
| 510 | 505 | ||
| 511 | 506 | ||
| 512 | def test_find_geoshift_raises_and_lists_searched_paths(tmp_path: Path) -> None: | 507 | def test_find_geoshift_raises_and_lists_searched_paths(tmp_path: pathlib.Path) -> None: |
| 513 | with pytest.raises(FileNotFoundError, match="searched:"): | 508 | with pytest.raises(FileNotFoundError, match="searched:"): |
| 514 | find_geoshift(tmp_path) | 509 | segment_points_io.find_geoshift(tmp_path) |
| 515 | 510 | ||
| 516 | 511 | ||
| 517 | def test_find_geoshift_returns_shift_unsigned(tmp_path: Path) -> None: | 512 | def test_find_geoshift_returns_shift_unsigned(tmp_path: pathlib.Path) -> None: |
| 518 | """The recorded shift is returned verbatim; no sign is baked in.""" | 513 | """The recorded shift is returned verbatim; no sign is baked in.""" |
| 519 | _write_geoshift(tmp_path / GEOSHIFT_NAME, (725883.5, -5422097.8, 390.8)) | 514 | _write_geoshift(tmp_path / segment_points_io.GEOSHIFT_NAME, (725883.5, -5422097.8, 390.8)) |
| 520 | np.testing.assert_allclose( | 515 | np.testing.assert_allclose( |
| 521 | find_geoshift(tmp_path), [725883.5, -5422097.8, 390.8] | 516 | segment_points_io.find_geoshift(tmp_path), [725883.5, -5422097.8, 390.8] |
| 522 | ) | 517 | ) |
| 523 | 518 | ||
| 524 | 519 | ||
| 525 | def test_geoshift_from_mapping_bare_and_nested() -> None: | 520 | def test_geoshift_from_mapping_bare_and_nested() -> None: |
| 526 | bare = geoshift_from_mapping({"x": 1.0, "y": 2.0, "z": 3.0}) | 521 | bare = segment_points_io.geoshift_from_mapping({"x": 1.0, "y": 2.0, "z": 3.0}) |
| 527 | nested = geoshift_from_mapping( | 522 | nested = segment_points_io.geoshift_from_mapping( |
| 528 | {"pixels_per_meter": 10, "geoshift": {"x": 1.0, "y": 2.0, "z": 3.0}} | 523 | {"pixels_per_meter": 10, "geoshift": {"x": 1.0, "y": 2.0, "z": 3.0}} |
| 529 | ) | 524 | ) |
| 530 | assert bare.dtype == np.float64 | 525 | assert bare.dtype == np.float64 |
| 531 | np.testing.assert_allclose(bare, [1.0, 2.0, 3.0]) | 526 | np.testing.assert_allclose(bare, [1.0, 2.0, 3.0]) |
| 533 | 528 | ||
| 534 | 529 | ||
| 535 | def test_geoshift_from_mapping_rejects_bad_input() -> None: | 530 | def test_geoshift_from_mapping_rejects_bad_input() -> None: |
| 536 | with pytest.raises(ValueError, match="missing geoshift key"): | 531 | with pytest.raises(ValueError, match="missing geoshift key"): |
| 537 | geoshift_from_mapping({"x": 1.0, "y": 2.0}) | 532 | segment_points_io.geoshift_from_mapping({"x": 1.0, "y": 2.0}) |
| 538 | with pytest.raises(ValueError, match="must be an object"): | 533 | with pytest.raises(ValueError, match="must be an object"): |
| 539 | geoshift_from_mapping([1.0, 2.0, 3.0]) # type: ignore[arg-type] | 534 | segment_points_io.geoshift_from_mapping([1.0, 2.0, 3.0]) # type: ignore[arg-type] |
| 540 | with pytest.raises(ValueError, match="must be numbers"): | 535 | with pytest.raises(ValueError, match="must be numbers"): |
| 541 | geoshift_from_mapping({"x": 1.0, "y": "north", "z": 3.0}) | 536 | segment_points_io.geoshift_from_mapping({"x": 1.0, "y": "north", "z": 3.0}) |
| 542 | 537 | ||
| 543 | 538 | ||
| 544 | def test_read_points_header_reports_stored_and_compressed(tmp_path: Path) -> None: | 539 | def test_read_points_header_reports_stored_and_compressed(tmp_path: pathlib.Path) -> None: |
| 545 | record = _make_record(5, seed=20) | 540 | record = _make_record(5, seed=20) |
| 546 | stored = _save_stored_npz(tmp_path / "stored.npz", record) | 541 | stored = _save_stored_npz(tmp_path / "stored.npz", record) |
| 547 | compressed = tmp_path / "compressed.npz" | 542 | compressed = tmp_path / "compressed.npz" |
| 548 | np.savez_compressed(compressed, **record) | 543 | np.savez_compressed(compressed, **record) |
| 549 | 544 | ||
| 550 | assert read_points_header(stored) == (5, True) | 545 | assert segment_points_io.read_points_header(stored) == (5, True) |
| 551 | assert read_points_header(compressed) == (5, False) | 546 | assert segment_points_io.read_points_header(compressed) == (5, False) |
| 552 | 547 | ||
| 553 | 548 | ||
| 554 | def test_read_points_header_unreadable_file(tmp_path: Path) -> None: | 549 | def test_read_points_header_unreadable_file(tmp_path: pathlib.Path) -> None: |
| 555 | broken = tmp_path / "broken.npz" | 550 | broken = tmp_path / "broken.npz" |
| 556 | broken.write_bytes(b"not a zip archive") | 551 | broken.write_bytes(b"not a zip archive") |
| 557 | assert read_points_header(broken) == (-1, False) | 552 | assert segment_points_io.read_points_header(broken) == (-1, False) |
| 558 | 553 | ||
| 559 | 554 | ||
| 560 | def test_iter_points_chunks_streams_tail_chunk_and_is_writeable(tmp_path: Path) -> None: | 555 | def test_iter_points_chunks_streams_tail_chunk_and_is_writeable(tmp_path: pathlib.Path) -> None: |
| 561 | record = _make_record(7, seed=21) | 556 | record = _make_record(7, seed=21) |
| 562 | path = _save_stored_npz(tmp_path / "big_run3_points.npz", record) | 557 | path = _save_stored_npz(tmp_path / "big_run3_points.npz", record) |
| 563 | 558 | ||
| 564 | chunks = list(iter_points_chunks(path, 3)) | 559 | chunks = list(segment_points_io.iter_points_chunks(path, 3)) |
| 565 | 560 | ||
| 566 | assert [len(chunk) for chunk in chunks] == [3, 3, 1] | 561 | assert [len(chunk) for chunk in chunks] == [3, 3, 1] |
| 567 | for chunk in chunks: | 562 | for chunk in chunks: |
| 568 | assert chunk.flags.writeable | 563 | assert chunk.flags.writeable |
| 569 | assert chunk.flags.c_contiguous | 564 | assert chunk.flags.c_contiguous |
| 570 | chunk[:] = 0.0 # must not raise: chunks are owned copies | 565 | chunk[:] = 0.0 # must not raise: chunks are owned copies |
| 571 | np.testing.assert_array_equal( | 566 | np.testing.assert_array_equal( |
| 572 | np.concatenate(list(iter_points_chunks(path, 3)), axis=0), record["points"] | 567 | np.concatenate(list(segment_points_io.iter_points_chunks(path, 3)), axis=0), |
| 568 | record["points"], | ||
| 573 | ) | 569 | ) |
| 574 | 570 | ||
| 575 | 571 | ||
| 576 | def test_iter_points_chunks_yields_whole_record_when_not_oversized(tmp_path: Path) -> None: | 572 | def test_iter_points_chunks_yields_whole_record_when_not_oversized(tmp_path: pathlib.Path) -> None: |
| 577 | record = _make_record(4, seed=22) | 573 | record = _make_record(4, seed=22) |
| 578 | path = _save_stored_npz(tmp_path / "small_run3_points.npz", record) | 574 | path = _save_stored_npz(tmp_path / "small_run3_points.npz", record) |
| 579 | 575 | ||
| 580 | chunks = list(iter_points_chunks(path, 4)) | 576 | chunks = list(segment_points_io.iter_points_chunks(path, 4)) |
| 581 | 577 | ||
| 582 | assert len(chunks) == 1 | 578 | assert len(chunks) == 1 |
| 583 | assert chunks[0].flags.writeable | 579 | assert chunks[0].flags.writeable |
| 584 | np.testing.assert_array_equal(chunks[0], record["points"]) | 580 | np.testing.assert_array_equal(chunks[0], record["points"]) |
| 585 | 581 | ||
| 586 | 582 | ||
| 587 | def test_iter_points_chunks_falls_back_for_compressed_records(tmp_path: Path) -> None: | 583 | def test_iter_points_chunks_falls_back_for_compressed_records(tmp_path: pathlib.Path) -> None: |
| 588 | record = _make_record(9, seed=23) | 584 | record = _make_record(9, seed=23) |
| 589 | path = tmp_path / "compressed_run3_points.npz" | 585 | path = tmp_path / "compressed_run3_points.npz" |
| 590 | np.savez_compressed(path, **record) | 586 | np.savez_compressed(path, **record) |
| 591 | 587 | ||
| 592 | chunks = list(iter_points_chunks(path, 2)) | 588 | chunks = list(segment_points_io.iter_points_chunks(path, 2)) |
| 593 | 589 | ||
| 594 | assert len(chunks) == 1 | 590 | assert len(chunks) == 1 |
| 595 | np.testing.assert_array_equal(chunks[0], record["points"]) | 591 | np.testing.assert_array_equal(chunks[0], record["points"]) |
| 596 | 592 | ||
| 597 | 593 | ||
| 598 | def test_iter_points_chunks_disabled_by_non_positive_chunk_size(tmp_path: Path) -> None: | 594 | def test_iter_points_chunks_disabled_by_non_positive_chunk_size(tmp_path: pathlib.Path) -> None: |
| 599 | record = _make_record(6, seed=24) | 595 | record = _make_record(6, seed=24) |
| 600 | path = _save_stored_npz(tmp_path / "run3_points.npz", record) | 596 | path = _save_stored_npz(tmp_path / "run3_points.npz", record) |
| 601 | 597 | ||
| 602 | chunks = list(iter_points_chunks(path, 0)) | 598 | chunks = list(segment_points_io.iter_points_chunks(path, 0)) |
| 603 | 599 | ||
| 604 | assert len(chunks) == 1 | 600 | assert len(chunks) == 1 |
| 605 | np.testing.assert_array_equal(chunks[0], record["points"]) | 601 | np.testing.assert_array_equal(chunks[0], record["points"]) |
| 606 | 602 | ||
| 607 | 603 | ||
| 608 | def test_discover_run3_files_skips_part_files(tmp_path: Path) -> None: | 604 | def test_discover_run3_files_skips_part_files(tmp_path: pathlib.Path) -> None: |
| 609 | seg_dir = tmp_path / "segment_003" | 605 | seg_dir = tmp_path / "segment_003" |
| 610 | seg_dir.mkdir() | 606 | seg_dir.mkdir() |
| 611 | save_points_npz(seg_dir / "beta_run3_points.npz", _make_record(2, seed=30)) | 607 | segment_points_io.save_points_npz(seg_dir / "beta_run3_points.npz", _make_record(2, seed=30)) |
| 612 | save_points_npz(seg_dir / "alpha_run3_points.npz", _make_record(2, seed=31)) | 608 | segment_points_io.save_points_npz(seg_dir / "alpha_run3_points.npz", _make_record(2, seed=31)) |
| 613 | (seg_dir / "gamma_run3_points.npz.part").write_bytes(b"partial write") | 609 | (seg_dir / "gamma_run3_points.npz.part").write_bytes(b"partial write") |
| 614 | (seg_dir / "delta_run4_road_surface.npz").write_bytes(b"other artifact") | 610 | (seg_dir / "delta_run4_road_surface.npz").write_bytes(b"other artifact") |
| 615 | 611 | ||
| 616 | found = discover_run3_files(seg_dir) | 612 | found = segment_points_io.discover_run3_files(seg_dir) |
| 617 | 613 | ||
| 618 | assert [path.name for path in found] == [ | 614 | assert [path.name for path in found] == [ |
| 619 | "alpha_run3_points.npz", | 615 | "alpha_run3_points.npz", |
| 620 | "beta_run3_points.npz", | 616 | "beta_run3_points.npz", |
| 621 | ] | 617 | ] |
| 622 | 618 | ||
| 623 | 619 | ||
| 624 | def test_discover_run3_files_missing_dir_is_empty(tmp_path: Path) -> None: | 620 | def test_discover_run3_files_missing_dir_is_empty(tmp_path: pathlib.Path) -> None: |
| 625 | assert discover_run3_files(tmp_path / "segment_999") == [] | 621 | assert segment_points_io.discover_run3_files(tmp_path / "segment_999") == [] |
| 626 | 622 | ||
| 627 | 623 | ||
| 628 | def test_load_run3_segment_boundary_table(tmp_path: Path) -> None: | 624 | def test_load_run3_segment_boundary_table(tmp_path: pathlib.Path) -> None: |
| 629 | seg_dir = tmp_path / "segment_007" | 625 | seg_dir = tmp_path / "segment_007" |
| 630 | seg_dir.mkdir() | 626 | seg_dir.mkdir() |
| 631 | first = _make_record(3, seed=40) | 627 | first = _make_record(3, seed=40) |
| 632 | second = _make_record(5, seed=41) | 628 | second = _make_record(5, seed=41) |
| 633 | save_points_npz(seg_dir / "a_run3_points.npz", first) | 629 | segment_points_io.save_points_npz(seg_dir / "a_run3_points.npz", first) |
| 634 | save_points_npz(seg_dir / "b_run3_points.npz", second) | 630 | segment_points_io.save_points_npz(seg_dir / "b_run3_points.npz", second) |
| 635 | (seg_dir / "c_run3_points.npz.part").write_bytes(b"partial write") | 631 | (seg_dir / "c_run3_points.npz.part").write_bytes(b"partial write") |
| 636 | 632 | ||
| 637 | merged, spans = load_run3_segment(seg_dir) | 633 | merged, spans = segment_points_io.load_run3_segment(seg_dir) |
| 638 | 634 | ||
| 639 | assert spans == [ | 635 | assert spans == [ |
| 640 | RecordSpan(name="a_run3_points.npz", offset=0, count=3), | 636 | segment_points_io.RecordSpan(name="a_run3_points.npz", offset=0, count=3), |
| 641 | RecordSpan(name="b_run3_points.npz", offset=3, count=5), | 637 | segment_points_io.RecordSpan(name="b_run3_points.npz", offset=3, count=5), |
| 642 | ] | 638 | ] |
| 643 | assert spans[1].end == 8 | 639 | assert spans[1].end == 8 |
| 644 | assert merged["points"].shape == (8, 3) | 640 | assert merged["points"].shape == (8, 3) |
| 645 | np.testing.assert_array_equal( | 641 | np.testing.assert_array_equal( |
| 653 | np.concatenate([first["intensity"], second["intensity"]], axis=0), | 649 | np.concatenate([first["intensity"], second["intensity"]], axis=0), |
| 654 | ) | 650 | ) |
| 655 | 651 | ||
| 656 | 652 | ||
| 657 | def test_load_run3_segment_requires_records(tmp_path: Path) -> None: | 653 | def test_load_run3_segment_requires_records(tmp_path: pathlib.Path) -> None: |
| 658 | seg_dir = tmp_path / "segment_008" | 654 | seg_dir = tmp_path / "segment_008" |
| 659 | seg_dir.mkdir() | 655 | seg_dir.mkdir() |
| 660 | (seg_dir / "a_run3_points.npz.part").write_bytes(b"partial write") | 656 | (seg_dir / "a_run3_points.npz.part").write_bytes(b"partial write") |
| 661 | with pytest.raises(FileNotFoundError, match=r"No \*_run3_points.npz"): | 657 | with pytest.raises(FileNotFoundError, match=r"No \*_run3_points.npz"): |
| 662 | load_run3_segment(seg_dir) | 658 | segment_points_io.load_run3_segment(seg_dir) |
| 663 | 659 | ||
| 664 | 660 | ||
| 665 | def test_concat_points_npz_rejects_mixed_dtypes(tmp_path: Path) -> None: | 661 | def test_concat_points_npz_rejects_mixed_dtypes(tmp_path: pathlib.Path) -> None: |
| 666 | first = _make_record(3, seed=50) | 662 | first = _make_record(3, seed=50) |
| 667 | second = _make_record(3, seed=51) | 663 | second = _make_record(3, seed=51) |
| 668 | second["points"] = second["points"].astype(np.float64) | 664 | second["points"] = second["points"].astype(np.float64) |
| 669 | save_points_npz(tmp_path / "a_run3_points.npz", first) | 665 | segment_points_io.save_points_npz(tmp_path / "a_run3_points.npz", first) |
| 670 | save_points_npz(tmp_path / "b_run3_points.npz", second) | 666 | segment_points_io.save_points_npz(tmp_path / "b_run3_points.npz", second) |
| 671 | 667 | ||
| 672 | with pytest.raises(ValueError, match="mixed dtypes"): | 668 | with pytest.raises(ValueError, match="mixed dtypes"): |
| 673 | concat_points_npz( | 669 | segment_points_io.concat_points_npz( |
| 674 | [tmp_path / "a_run3_points.npz", tmp_path / "b_run3_points.npz"] | 670 | [tmp_path / "a_run3_points.npz", tmp_path / "b_run3_points.npz"] |
| 675 | ) | 671 | ) |
| 676 | 672 | ||
| 677 | 673 | ||
| 678 | def test_concat_points_npz_rejects_mixed_ancillary_dtypes(tmp_path: Path) -> None: | 674 | def test_concat_points_npz_rejects_mixed_ancillary_dtypes(tmp_path: pathlib.Path) -> None: |
| 679 | first = _make_record(2, seed=52) | 675 | first = _make_record(2, seed=52) |
| 680 | second = _make_record(2, seed=53) | 676 | second = _make_record(2, seed=53) |
| 681 | second["intensity"] = second["intensity"].astype(np.uint8) | 677 | second["intensity"] = second["intensity"].astype(np.uint8) |
| 682 | save_points_npz(tmp_path / "a_run3_points.npz", first) | 678 | segment_points_io.save_points_npz(tmp_path / "a_run3_points.npz", first) |
| 683 | save_points_npz(tmp_path / "b_run3_points.npz", second) | 679 | segment_points_io.save_points_npz(tmp_path / "b_run3_points.npz", second) |
| 684 | 680 | ||
| 685 | with pytest.raises(ValueError, match="'intensity' dtype uint8"): | 681 | with pytest.raises(ValueError, match="'intensity' dtype uint8"): |
| 686 | concat_points_npz( | 682 | segment_points_io.concat_points_npz( |
| 687 | [tmp_path / "a_run3_points.npz", tmp_path / "b_run3_points.npz"] | 683 | [tmp_path / "a_run3_points.npz", tmp_path / "b_run3_points.npz"] |
| 688 | ) | 684 | ) |
| 689 | 685 | ||
| 690 | 686 | ||
| 691 | def test_concat_points_npz_empty_raises() -> None: | 687 | def test_concat_points_npz_empty_raises() -> None: |
| 692 | with pytest.raises(FileNotFoundError, match="No \\*_points.npz"): | 688 | with pytest.raises(FileNotFoundError, match="No \\*_points.npz"): |
| 693 | concat_points_npz([]) | 689 | segment_points_io.concat_points_npz([]) |
| 694 | 690 | ||
| 695 | 691 | ||
| 696 | def test_concat_points_npz_target_dtypes_casts_mixed_records(tmp_path: Path) -> None: | 692 | def test_concat_points_npz_target_dtypes_casts_mixed_records(tmp_path: pathlib.Path) -> None: |
| 697 | # seg3d-style normalization: historical records with different storage | 693 | # seg3d-style normalization: historical records with different storage |
| 698 | # dtypes (uint8 vs uint16 intensity) are cast to the target instead of | 694 | # dtypes (uint8 vs uint16 intensity) are cast to the target instead of |
| 699 | # rejected. | 695 | # rejected. |
| 700 | first = _make_record(2, seed=54) | 696 | first = _make_record(2, seed=54) |
| 701 | second = _make_record(2, seed=55) | 697 | second = _make_record(2, seed=55) |
| 702 | second["intensity"] = second["intensity"].astype(np.uint8) | 698 | second["intensity"] = second["intensity"].astype(np.uint8) |
| 703 | save_points_npz(tmp_path / "a_run3_points.npz", first) | 699 | segment_points_io.save_points_npz(tmp_path / "a_run3_points.npz", first) |
| 704 | save_points_npz(tmp_path / "b_run3_points.npz", second) | 700 | segment_points_io.save_points_npz(tmp_path / "b_run3_points.npz", second) |
| 705 | 701 | ||
| 706 | merged, spans = concat_points_npz( | 702 | merged, spans = segment_points_io.concat_points_npz( |
| 707 | [tmp_path / "a_run3_points.npz", tmp_path / "b_run3_points.npz"], | 703 | [tmp_path / "a_run3_points.npz", tmp_path / "b_run3_points.npz"], |
| 708 | target_dtypes={"points": np.dtype(np.float64), "intensity": np.dtype(np.uint16)}, | 704 | target_dtypes={"points": np.dtype(np.float64), "intensity": np.dtype(np.uint16)}, |
| 709 | ) | 705 | ) |
| 710 | assert merged["points"].dtype == np.float64 | 706 | assert merged["points"].dtype == np.float64 |
| 718 | ) | 714 | ) |
| 719 | 715 | ||
| 720 | 716 | ||
| 721 | def test_find_geoshift_warns_on_multiple_candidates( | 717 | def test_find_geoshift_warns_on_multiple_candidates( |
| 722 | tmp_path: Path, caplog: pytest.LogCaptureFixture | 718 | tmp_path: pathlib.Path, caplog: pytest.LogCaptureFixture |
| 723 | ) -> None: | 719 | ) -> None: |
| 724 | # A stale segment-local file shadowing the dataset-level one is legal but | 720 | # A stale segment-local file shadowing the dataset-level one is legal but |
| 725 | # suspicious; the lookup must say so out loud. | 721 | # suspicious; the lookup must say so out loud. |
| 726 | seg_dir = tmp_path / "lane_points" / "segment_032" | 722 | seg_dir = tmp_path / "lane_points" / "segment_032" |
| 728 | _write_geoshift(seg_dir / "run3_geoshift.json", (1.0, 1.0, 1.0)) | 724 | _write_geoshift(seg_dir / "run3_geoshift.json", (1.0, 1.0, 1.0)) |
| 729 | _write_geoshift(tmp_path / "lane_points" / "run3_geoshift.json", (2.0, 2.0, 2.0)) | 725 | _write_geoshift(tmp_path / "lane_points" / "run3_geoshift.json", (2.0, 2.0, 2.0)) |
| 730 | 726 | ||
| 731 | with caplog.at_level(logging.WARNING, logger="iolabs.common.segment_points_io"): | 727 | with caplog.at_level(logging.WARNING, logger="iolabs.common.segment_points_io"): |
| 732 | shift = find_geoshift(seg_dir) | 728 | shift = segment_points_io.find_geoshift(seg_dir) |
| 733 | 729 | ||
| 734 | np.testing.assert_array_equal(shift, np.array([1.0, 1.0, 1.0])) | 730 | np.testing.assert_array_equal(shift, np.array([1.0, 1.0, 1.0])) |
| 735 | assert any("Multiple" in message for message in caplog.messages) | 731 | assert any("Multiple" in message for message in caplog.messages) |
| 736 | 732 | ||
| 737 | 733 | ||
| 738 | def test_public_key_tuples_are_derived_from_the_schema() -> None: | 734 | def test_public_key_tuples_are_derived_from_the_schema() -> None: |
| 739 | assert POINT_RECORD_KEYS == tuple(POINT_RECORD_SCHEMA) | 735 | assert segment_points_io.POINT_RECORD_KEYS == tuple(segment_points_io.POINT_RECORD_SCHEMA) |
| 740 | assert REQUIRED_POINT_RECORD_KEYS == ( | 736 | assert segment_points_io.REQUIRED_POINT_RECORD_KEYS == ( |
| 741 | "points", | 737 | "points", |
| 742 | "red", | 738 | "red", |
| 743 | "green", | 739 | "green", |
| 744 | "blue", | 740 | "blue", |
| 745 | "intensity", | 741 | "intensity", |
| 746 | "scan_angle", | 742 | "scan_angle", |
| 747 | ) | 743 | ) |
| 748 | assert OPTIONAL_POINT_RECORD_KEYS == (NUMBER_OF_RETURNS_KEY,) | 744 | assert segment_points_io.OPTIONAL_POINT_RECORD_KEYS == ( |
| 749 | assert all(POINT_RECORD_SCHEMA[key].required for key in REQUIRED_POINT_RECORD_KEYS) | 745 | segment_points_io.NUMBER_OF_RETURNS_KEY, |
| 750 | assert POINT_RECORD_SCHEMA[NUMBER_OF_RETURNS_KEY].storage_dtype == np.dtype(np.uint8) | 746 | ) |
| 751 | assert POINT_RECORD_SCHEMA["points"].columns == 3 | 747 | assert all( |
| 748 | segment_points_io.POINT_RECORD_SCHEMA[key].required | ||
| 749 | for key in segment_points_io.REQUIRED_POINT_RECORD_KEYS | ||
| 750 | ) | ||
| 751 | assert segment_points_io.POINT_RECORD_SCHEMA[ | ||
| 752 | segment_points_io.NUMBER_OF_RETURNS_KEY | ||
| 753 | ].storage_dtype == np.dtype(np.uint8) | ||
| 754 | assert segment_points_io.POINT_RECORD_SCHEMA["points"].columns == 3 | ||
| 752 | 755 | ||
| 753 | 756 | ||
| 754 | def test_point_field_spec_rejects_an_optional_key_without_a_fill() -> None: | 757 | def test_point_field_spec_rejects_an_optional_key_without_a_fill() -> None: |
| 755 | """An optional key with no fill would make pre-existing records unloadable.""" | 758 | """An optional key with no fill would make pre-existing records unloadable.""" |
| 756 | with pytest.raises(ValueError, match="needs a fill factory"): | 759 | with pytest.raises(ValueError, match="needs a fill factory"): |
| 757 | PointFieldSpec(required=False) | 760 | segment_points_io.PointFieldSpec(required=False) |
| 758 | 761 | ||
| 759 | 762 | ||
| 760 | def test_mask_record_masks_points_rows_and_ancillary_elements() -> None: | 763 | def test_mask_record_masks_points_rows_and_ancillary_elements() -> None: |
| 761 | record = _make_record(5, seed=80) | 764 | record = _make_record(5, seed=80) |
| 762 | mask = np.array([True, False, True, False, True]) | 765 | mask = np.array([True, False, True, False, True]) |
| 763 | 766 | ||
| 764 | masked = mask_record(record, mask) | 767 | masked = segment_points_io.mask_record(record, mask) |
| 765 | 768 | ||
| 766 | assert list(masked) == list(record) | 769 | assert list(masked) == list(record) |
| 767 | assert masked["points"].shape == (3, 3) | 770 | assert masked["points"].shape == (3, 3) |
| 768 | np.testing.assert_array_equal(masked["points"], record["points"][mask]) | 771 | np.testing.assert_array_equal(masked["points"], record["points"][mask]) |
| 769 | for key in POINT_RECORD_KEYS: | 772 | for key in segment_points_io.POINT_RECORD_KEYS: |
| 770 | np.testing.assert_array_equal(masked[key], record[key][mask]) | 773 | np.testing.assert_array_equal(masked[key], record[key][mask]) |
| 771 | 774 | ||
| 772 | 775 | ||
| 773 | def test_mask_record_accepts_an_integer_index_array() -> None: | 776 | def test_mask_record_accepts_an_integer_index_array() -> None: |
| 774 | record = _make_record(4, seed=81) | 777 | record = _make_record(4, seed=81) |
| 775 | index = np.array([3, 0]) | 778 | index = np.array([3, 0]) |
| 776 | 779 | ||
| 777 | masked = mask_record(record, index) | 780 | masked = segment_points_io.mask_record(record, index) |
| 778 | 781 | ||
| 779 | np.testing.assert_array_equal(masked["points"], record["points"][index]) | 782 | np.testing.assert_array_equal(masked["points"], record["points"][index]) |
| 780 | np.testing.assert_array_equal(masked["intensity"], record["intensity"][index]) | 783 | np.testing.assert_array_equal(masked["intensity"], record["intensity"][index]) |
| 781 | 784 |
| 783 | def test_mask_record_rejects_misaligned_members() -> None: | 786 | def test_mask_record_rejects_misaligned_members() -> None: |
| 784 | record = _make_record(4, seed=82) | 787 | record = _make_record(4, seed=82) |
| 785 | record["red"] = record["red"][:2] | 788 | record["red"] = record["red"][:2] |
| 786 | with pytest.raises(ValueError, match="disagree on point count"): | 789 | with pytest.raises(ValueError, match="disagree on point count"): |
| 787 | mask_record(record, np.array([True, False, True, False])) | 790 | segment_points_io.mask_record(record, np.array([True, False, True, False])) |
| 788 | 791 | ||
| 789 | 792 | ||
| 790 | def test_mask_record_rejects_an_empty_record() -> None: | 793 | def test_mask_record_rejects_an_empty_record() -> None: |
| 791 | with pytest.raises(ValueError, match="empty point record"): | 794 | with pytest.raises(ValueError, match="empty point record"): |
| 792 | mask_record({}, np.array([True])) | 795 | segment_points_io.mask_record({}, np.array([True])) |
| 793 | 796 | ||
| 794 | 797 | ||
| 795 | def test_concat_records_joins_every_member_on_axis_zero() -> None: | 798 | def test_concat_records_joins_every_member_on_axis_zero() -> None: |
| 796 | first = _make_record(3, seed=83) | 799 | first = _make_record(3, seed=83) |
| 797 | second = _make_record(2, seed=84) | 800 | second = _make_record(2, seed=84) |
| 798 | 801 | ||
| 799 | merged = concat_records([first, second]) | 802 | merged = segment_points_io.concat_records([first, second]) |
| 800 | 803 | ||
| 801 | assert list(merged) == list(POINT_RECORD_KEYS) | 804 | assert list(merged) == list(segment_points_io.POINT_RECORD_KEYS) |
| 802 | assert merged["points"].shape == (5, 3) | 805 | assert merged["points"].shape == (5, 3) |
| 803 | for key in POINT_RECORD_KEYS: | 806 | for key in segment_points_io.POINT_RECORD_KEYS: |
| 804 | np.testing.assert_array_equal( | 807 | np.testing.assert_array_equal( |
| 805 | merged[key], np.concatenate([first[key], second[key]], axis=0) | 808 | merged[key], np.concatenate([first[key], second[key]], axis=0) |
| 806 | ) | 809 | ) |
| 807 | 810 | ||
| 808 | 811 | ||
| 809 | def test_concat_records_rejects_mismatched_key_sets() -> None: | 812 | def test_concat_records_rejects_mismatched_key_sets() -> None: |
| 810 | with pytest.raises(ValueError, match="different keys"): | 813 | with pytest.raises(ValueError, match="different keys"): |
| 811 | concat_records([_make_record(2, seed=85), _make_legacy_record(2, seed=86)]) | 814 | segment_points_io.concat_records( |
| 815 | [_make_record(2, seed=85), _make_legacy_record(2, seed=86)] | ||
| 816 | ) | ||
| 812 | 817 | ||
| 813 | 818 | ||
| 814 | def test_concat_records_rejects_an_empty_sequence() -> None: | 819 | def test_concat_records_rejects_an_empty_sequence() -> None: |
| 815 | with pytest.raises(ValueError, match="empty sequence"): | 820 | with pytest.raises(ValueError, match="empty sequence"): |
| 816 | concat_records([]) | 821 | segment_points_io.concat_records([]) |
| 817 | 822 | ||
| 818 | 823 | ||
| 819 | def test_concat_records_rejects_records_without_members() -> None: | 824 | def test_concat_records_rejects_records_without_members() -> None: |
| 820 | """Memberless records must raise, not silently concatenate to ``{}``.""" | 825 | """Memberless records must raise, not silently concatenate to ``{}``.""" |
| 821 | with pytest.raises(ValueError, match="no members"): | 826 | with pytest.raises(ValueError, match="no members"): |
| 822 | concat_records([{}, {}]) | 827 | segment_points_io.concat_records([{}, {}]) |
| 823 | 828 | ||
| 824 | 829 | ||
| 825 | def test_mask_record_rejects_zero_dimensional_members() -> None: | 830 | def test_mask_record_rejects_zero_dimensional_members() -> None: |
| 826 | """A 0-D member is not row-indexable: raise ValueError, not a raw IndexError.""" | 831 | """A 0-D member is not row-indexable: raise ValueError, not a raw IndexError.""" |
| 827 | record = {key: np.asarray(1, dtype=np.uint8) for key in POINT_RECORD_KEYS} | 832 | record = {key: np.asarray(1, dtype=np.uint8) for key in segment_points_io.POINT_RECORD_KEYS} |
| 828 | 833 | ||
| 829 | with pytest.raises(ValueError, match="not 1-D or 2-D"): | 834 | with pytest.raises(ValueError, match="not 1-D or 2-D"): |
| 830 | mask_record(record, np.array([True])) | 835 | segment_points_io.mask_record(record, np.array([True])) |
| 831 | 836 | ||
| 832 | 837 | ||
| 833 | # --- Adding a field to the contract must be a registry entry and nothing else --- | 838 | # --- Adding a field to the contract must be a registry entry and nothing else --- |
| 834 | 839 | ||
| 835 | EXTRA_KEY = "point_source_id" | 840 | EXTRA_KEY = "point_source_id" |
| 836 | 841 | ||
| 837 | EXTRA_SPEC = PointFieldSpec( | 842 | EXTRA_SPEC = segment_points_io.PointFieldSpec( |
| 838 | required=False, | 843 | required=False, |
| 839 | storage_dtype=np.dtype(np.int16), | 844 | storage_dtype=np.dtype(np.int16), |
| 840 | fill=functools.partial(np.full, fill_value=-1, dtype=np.int16), | 845 | fill=functools.partial(np.full, fill_value=-1, dtype=np.int16), |
| 841 | noun="source id", | 846 | noun="source id", |
| 847 | """Register one extra optional key, exactly as a real schema addition would.""" | 852 | """Register one extra optional key, exactly as a real schema addition would.""" |
| 848 | monkeypatch.setattr( | 853 | monkeypatch.setattr( |
| 849 | segment_points_io, | 854 | segment_points_io, |
| 850 | "POINT_RECORD_SCHEMA", | 855 | "POINT_RECORD_SCHEMA", |
| 851 | MappingProxyType({**POINT_RECORD_SCHEMA, EXTRA_KEY: EXTRA_SPEC}), | 856 | types.MappingProxyType({**segment_points_io.POINT_RECORD_SCHEMA, EXTRA_KEY: EXTRA_SPEC}), |
| 852 | ) | 857 | ) |
| 853 | return (*POINT_RECORD_KEYS, EXTRA_KEY) | 858 | return (*segment_points_io.POINT_RECORD_KEYS, EXTRA_KEY) |
| 854 | 859 | ||
| 855 | 860 | ||
| 856 | def _make_extended_record(n: int, *, seed: int) -> dict[str, np.ndarray]: | 861 | def _make_extended_record(n: int, *, seed: int) -> dict[str, np.ndarray]: |
| 857 | record = _make_record(n, seed=seed) | 862 | record = _make_record(n, seed=seed) |
| 859 | return record | 864 | return record |
| 860 | 865 | ||
| 861 | 866 | ||
| 862 | def test_registry_entry_alone_round_trips_a_new_field( | 867 | def test_registry_entry_alone_round_trips_a_new_field( |
| 863 | tmp_path: Path, extended_schema: tuple[str, ...] | 868 | tmp_path: pathlib.Path, extended_schema: tuple[str, ...] |
| 864 | ) -> None: | 869 | ) -> None: |
| 865 | record = _make_extended_record(4, seed=90) | 870 | record = _make_extended_record(4, seed=90) |
| 866 | 871 | ||
| 867 | path = save_points_npz(tmp_path / "extended_run3_points.npz", record) | 872 | path = segment_points_io.save_points_npz(tmp_path / "extended_run3_points.npz", record) |
| 868 | 873 | ||
| 869 | with np.load(path) as data: | 874 | with np.load(path) as data: |
| 870 | assert EXTRA_KEY in data.files | 875 | assert EXTRA_KEY in data.files |
| 871 | loaded = load_points_npz(path) | 876 | loaded = segment_points_io.load_points_npz(path) |
| 872 | assert list(loaded) == list(extended_schema) | 877 | assert list(loaded) == list(extended_schema) |
| 873 | assert loaded[EXTRA_KEY].dtype == np.int16 | 878 | assert loaded[EXTRA_KEY].dtype == np.int16 |
| 874 | np.testing.assert_array_equal(loaded[EXTRA_KEY], record[EXTRA_KEY]) | 879 | np.testing.assert_array_equal(loaded[EXTRA_KEY], record[EXTRA_KEY]) |
| 875 | 880 | ||
| 876 | 881 | ||
| 877 | def test_registry_entry_alone_coerces_and_validates_a_new_field( | 882 | def test_registry_entry_alone_coerces_and_validates_a_new_field( |
| 878 | tmp_path: Path, extended_schema: tuple[str, ...] | 883 | tmp_path: pathlib.Path, extended_schema: tuple[str, ...] |
| 879 | ) -> None: | 884 | ) -> None: |
| 880 | record = _make_extended_record(3, seed=91) | 885 | record = _make_extended_record(3, seed=91) |
| 881 | record[EXTRA_KEY] = record[EXTRA_KEY].astype(np.int64) | 886 | record[EXTRA_KEY] = record[EXTRA_KEY].astype(np.int64) |
| 882 | loaded = load_points_npz(save_points_npz(tmp_path / "cast_run3_points.npz", record)) | 887 | loaded = segment_points_io.load_points_npz( |
| 888 | segment_points_io.save_points_npz(tmp_path / "cast_run3_points.npz", record) | ||
| 889 | ) | ||
| 883 | assert loaded[EXTRA_KEY].dtype == np.int16 | 890 | assert loaded[EXTRA_KEY].dtype == np.int16 |
| 884 | 891 | ||
| 885 | out_of_range = _make_extended_record(3, seed=92) | 892 | out_of_range = _make_extended_record(3, seed=92) |
| 886 | out_of_range[EXTRA_KEY] = np.array([1, 40_000, 3], dtype=np.int32) | 893 | out_of_range[EXTRA_KEY] = np.array([1, 40_000, 3], dtype=np.int32) |
| 887 | with pytest.raises(ValueError, match=f"{EXTRA_KEY}.*fit in int16"): | 894 | with pytest.raises(ValueError, match=f"{EXTRA_KEY}.*fit in int16"): |
| 888 | save_points_npz(tmp_path / "range_run3_points.npz", out_of_range) | 895 | segment_points_io.save_points_npz(tmp_path / "range_run3_points.npz", out_of_range) |
| 889 | 896 | ||
| 890 | boolean = _make_extended_record(3, seed=93) | 897 | boolean = _make_extended_record(3, seed=93) |
| 891 | boolean[EXTRA_KEY] = np.array([True, False, True]) | 898 | boolean[EXTRA_KEY] = np.array([True, False, True]) |
| 892 | with pytest.raises(ValueError, match=f"{EXTRA_KEY}.*bool"): | 899 | with pytest.raises(ValueError, match=f"{EXTRA_KEY}.*bool"): |
| 893 | save_points_npz(tmp_path / "bool_run3_points.npz", boolean) | 900 | segment_points_io.save_points_npz(tmp_path / "bool_run3_points.npz", boolean) |
| 894 | 901 | ||
| 895 | misshaped = _make_extended_record(3, seed=94) | 902 | misshaped = _make_extended_record(3, seed=94) |
| 896 | misshaped[EXTRA_KEY] = misshaped[EXTRA_KEY].reshape(3, 1) | 903 | misshaped[EXTRA_KEY] = misshaped[EXTRA_KEY].reshape(3, 1) |
| 897 | with pytest.raises(ValueError, match=rf"'{EXTRA_KEY}' must have shape \(N,\)"): | 904 | with pytest.raises(ValueError, match=rf"'{EXTRA_KEY}' must have shape \(N,\)"): |
| 898 | save_points_npz(tmp_path / "shape_run3_points.npz", misshaped) | 905 | segment_points_io.save_points_npz(tmp_path / "shape_run3_points.npz", misshaped) |
| 899 | 906 | ||
| 900 | 907 | ||
| 901 | def test_registry_entry_alone_fills_and_requires_a_new_field( | 908 | def test_registry_entry_alone_fills_and_requires_a_new_field( |
| 902 | tmp_path: Path, extended_schema: tuple[str, ...] | 909 | tmp_path: pathlib.Path, extended_schema: tuple[str, ...] |
| 903 | ) -> None: | 910 | ) -> None: |
| 904 | older = _make_record(5, seed=95) | 911 | older = _make_record(5, seed=95) |
| 905 | path = tmp_path / "older_run3_points.npz" | 912 | path = tmp_path / "older_run3_points.npz" |
| 906 | np.savez_compressed(path, **older) | 913 | np.savez_compressed(path, **older) |
| 907 | 914 | ||
| 908 | loaded = load_points_npz(path) | 915 | loaded = segment_points_io.load_points_npz(path) |
| 909 | 916 | ||
| 910 | assert list(loaded) == list(extended_schema) | 917 | assert list(loaded) == list(extended_schema) |
| 911 | np.testing.assert_array_equal(loaded[EXTRA_KEY], np.full(5, -1, dtype=np.int16)) | 918 | np.testing.assert_array_equal(loaded[EXTRA_KEY], np.full(5, -1, dtype=np.int16)) |
| 912 | with pytest.raises(ValueError, match=EXTRA_KEY): | 919 | with pytest.raises(ValueError, match=EXTRA_KEY): |
| 913 | save_points_npz(tmp_path / "incomplete_run3_points.npz", older) | 920 | segment_points_io.save_points_npz(tmp_path / "incomplete_run3_points.npz", older) |
| 914 | 921 | ||
| 915 | 922 | ||
| 916 | def test_registry_entry_alone_flows_through_merge_mask_and_concat( | 923 | def test_registry_entry_alone_flows_through_merge_mask_and_concat( |
| 917 | tmp_path: Path, extended_schema: tuple[str, ...] | 924 | tmp_path: pathlib.Path, extended_schema: tuple[str, ...] |
| 918 | ) -> None: | 925 | ) -> None: |
| 919 | seg_dir = tmp_path / "segment_042" | 926 | seg_dir = tmp_path / "segment_042" |
| 920 | seg_dir.mkdir() | 927 | seg_dir.mkdir() |
| 921 | first = _make_extended_record(3, seed=96) | 928 | first = _make_extended_record(3, seed=96) |
| 922 | second = _make_extended_record(2, seed=97) | 929 | second = _make_extended_record(2, seed=97) |
| 923 | save_points_npz(seg_dir / "a_run3_points.npz", first) | 930 | segment_points_io.save_points_npz(seg_dir / "a_run3_points.npz", first) |
| 924 | save_points_npz(seg_dir / "b_run3_points.npz", second) | 931 | segment_points_io.save_points_npz(seg_dir / "b_run3_points.npz", second) |
| 925 | 932 | ||
| 926 | merged, spans = load_run3_segment(seg_dir) | 933 | merged, spans = segment_points_io.load_run3_segment(seg_dir) |
| 927 | segment_merged, point_file_ids, _ = load_segment_points( | 934 | segment_merged, point_file_ids, _ = segment_points_io.load_segment_points( |
| 928 | [seg_dir / "a_run3_points.npz", seg_dir / "b_run3_points.npz"] | 935 | [seg_dir / "a_run3_points.npz", seg_dir / "b_run3_points.npz"] |
| 929 | ) | 936 | ) |
| 930 | masked = mask_record(merged, np.array([True, False, True, False, True])) | 937 | masked = segment_points_io.mask_record(merged, np.array([True, False, True, False, True])) |
| 931 | joined = concat_records([first, second]) | 938 | joined = segment_points_io.concat_records([first, second]) |
| 932 | 939 | ||
| 933 | expected = np.concatenate([first[EXTRA_KEY], second[EXTRA_KEY]], axis=0) | 940 | expected = np.concatenate([first[EXTRA_KEY], second[EXTRA_KEY]], axis=0) |
| 934 | assert [span.count for span in spans] == [3, 2] | 941 | assert [span.count for span in spans] == [3, 2] |
| 935 | assert point_file_ids.shape == (5,) | 942 | assert point_file_ids.shape == (5,) |
| 940 | 947 | ||
| 941 | 948 | ||
| 942 | def test_extended_schema_does_not_leak_into_the_real_contract() -> None: | 949 | def test_extended_schema_does_not_leak_into_the_real_contract() -> None: |
| 943 | """The monkeypatched registry above must not pollute the shipped schema.""" | 950 | """The monkeypatched registry above must not pollute the shipped schema.""" |
| 944 | assert EXTRA_KEY not in POINT_RECORD_SCHEMA | 951 | assert EXTRA_KEY not in _POINT_RECORD_SCHEMA_AT_IMPORT |
| 945 | assert EXTRA_KEY not in segment_points_io.POINT_RECORD_SCHEMA | 952 | assert EXTRA_KEY not in segment_points_io.POINT_RECORD_SCHEMA |
| 946 | assert tuple(segment_points_io.POINT_RECORD_SCHEMA) == POINT_RECORD_KEYS | 953 | assert tuple(segment_points_io.POINT_RECORD_SCHEMA) == _POINT_RECORD_KEYS_AT_IMPORT |
| 34 | and consumers working in the local frame ignore it. This module never | 34 | and consumers working in the local frame ignore it. This module never |
| 35 | applies, negates, or bakes in a sign. | 35 | applies, negates, or bakes in a sign. |
| 36 | """ | 36 | """ |
| 37 | 37 | ||
| 38 | import dataclasses | ||
| 39 | import fnmatch | ||
| 38 | import functools | 40 | import functools |
| 39 | import json | 41 | import json |
| 40 | import logging | 42 | import logging |
| 43 | import pathlib | ||
| 44 | import types | ||
| 41 | import zipfile | 45 | import zipfile |
| 42 | from collections.abc import Callable, Iterator, Mapping, Sequence | 46 | from collections import abc |
| 43 | from dataclasses import dataclass | ||
| 44 | from fnmatch import fnmatch | ||
| 45 | from pathlib import Path | ||
| 46 | from types import MappingProxyType | ||
| 47 | 47 | ||
| 48 | import numpy as np | 48 | import numpy as np |
| 49 | 49 | ||
| 50 | from . import _dtype_coercion | 50 | from . import _dtype_coercion |
| 51 | 51 | ||
| 52 | logger = logging.getLogger(__name__) | 52 | logger = logging.getLogger(__name__) |
| 53 | 53 | ||
| 54 | #: Builds a member for a record that predates its key, given the point count. | 54 | #: Builds a member for a record that predates its key, given the point count. |
| 55 | FillFactory = Callable[[int], np.ndarray] | 55 | FillFactory = abc.Callable[[int], np.ndarray] |
| 56 | 56 | ||
| 57 | #: The point-record key whose row count defines ``N`` for every other member. | 57 | #: The point-record key whose row count defines ``N`` for every other member. |
| 58 | POINTS_KEY = "points" | 58 | POINTS_KEY = "points" |
| 59 | 59 |
| 63 | #: Storage dtype of :data:`NUMBER_OF_RETURNS_KEY` (LAS carries 3 bits, values 1-7). | 63 | #: Storage dtype of :data:`NUMBER_OF_RETURNS_KEY` (LAS carries 3 bits, values 1-7). |
| 64 | NUMBER_OF_RETURNS_DTYPE = np.uint8 | 64 | NUMBER_OF_RETURNS_DTYPE = np.uint8 |
| 65 | 65 | ||
| 66 | 66 | ||
| 67 | @dataclass(frozen=True) | 67 | @dataclasses.dataclass(frozen=True) |
| 68 | class PointFieldSpec: | 68 | class PointFieldSpec: |
| 69 | """How one member of the point record is validated, stored and synthesised. | 69 | """How one member of the point record is validated, stored and synthesised. |
| 70 | 70 | ||
| 71 | Attributes: | 71 | Attributes: |
| 108 | 108 | ||
| 109 | #: The point-record contract: ordered key -> spec. Adding a field to the NPZ | 109 | #: The point-record contract: ordered key -> spec. Adding a field to the NPZ |
| 110 | #: schema means adding an entry here (plus teaching producers to emit it); the | 110 | #: schema means adding an entry here (plus teaching producers to emit it); the |
| 111 | #: load/save/merge machinery below is entirely registry-driven. | 111 | #: load/save/merge machinery below is entirely registry-driven. |
| 112 | POINT_RECORD_SCHEMA: Mapping[str, PointFieldSpec] = MappingProxyType( | 112 | POINT_RECORD_SCHEMA: abc.Mapping[str, PointFieldSpec] = types.MappingProxyType( |
| 113 | { | 113 | { |
| 114 | POINTS_KEY: PointFieldSpec(required=True, columns=3, noun="coordinate"), | 114 | POINTS_KEY: PointFieldSpec(required=True, columns=3, noun="coordinate"), |
| 115 | "red": PointFieldSpec(required=True), | 115 | "red": PointFieldSpec(required=True), |
| 116 | "green": PointFieldSpec(required=True), | 116 | "green": PointFieldSpec(required=True), |
| 127 | ) | 127 | ) |
| 128 | 128 | ||
| 129 | 129 | ||
| 130 | def _schema_keys( | 130 | def _schema_keys( |
| 131 | schema: Mapping[str, PointFieldSpec], *, required: bool | None = None | 131 | schema: abc.Mapping[str, PointFieldSpec], *, required: bool | None = None |
| 132 | ) -> tuple[str, ...]: | 132 | ) -> tuple[str, ...]: |
| 133 | """List the schema's keys in registry order, optionally by required-ness.""" | 133 | """List the schema's keys in registry order, optionally by required-ness.""" |
| 134 | return tuple( | 134 | return tuple( |
| 135 | key for key, spec in schema.items() if required is None or spec.required is required | 135 | key for key, spec in schema.items() if required is None or spec.required is required |
| 190 | return "(N,)" if spec.columns is None else f"(N, {spec.columns})" | 190 | return "(N,)" if spec.columns is None else f"(N, {spec.columns})" |
| 191 | 191 | ||
| 192 | 192 | ||
| 193 | def _validate_record_shapes( | 193 | def _validate_record_shapes( |
| 194 | record: Mapping[str, np.ndarray], | 194 | record: abc.Mapping[str, np.ndarray], |
| 195 | schema: Mapping[str, PointFieldSpec], | 195 | schema: abc.Mapping[str, PointFieldSpec], |
| 196 | *, | 196 | *, |
| 197 | source: str, | 197 | source: str, |
| 198 | ) -> int: | 198 | ) -> int: |
| 199 | """Check every present member against its spec's shape and return ``N``. | 199 | """Check every present member against its spec's shape and return ``N``. |
| 227 | ) | 227 | ) |
| 228 | return n_points | 228 | return n_points |
| 229 | 229 | ||
| 230 | 230 | ||
| 231 | def load_points_npz(path: str | Path) -> PointRecord: | 231 | def load_points_npz(path: str | pathlib.Path) -> PointRecord: |
| 232 | """Load one ``*_run3_points.npz``-style file and validate the schema. | 232 | """Load one ``*_run3_points.npz``-style file and validate the schema. |
| 233 | 233 | ||
| 234 | Every required key of :data:`POINT_RECORD_SCHEMA` must be stored; ``points`` | 234 | Every required key of :data:`POINT_RECORD_SCHEMA` must be stored; ``points`` |
| 235 | must be shape ``(N, 3)`` and every ancillary array (``red``, ``green``, | 235 | must be shape ``(N, 3)`` and every ancillary array (``red``, ``green``, |
| 254 | ValueError: A required key is missing, an array has the wrong shape, or | 254 | ValueError: A required key is missing, an array has the wrong shape, or |
| 255 | a member with a declared storage dtype is stored with a | 255 | a member with a declared storage dtype is stored with a |
| 256 | non-integer/bool dtype or values outside that dtype's range. | 256 | non-integer/bool dtype or values outside that dtype's range. |
| 257 | """ | 257 | """ |
| 258 | npz_path = Path(path) | 258 | npz_path = pathlib.Path(path) |
| 259 | schema = POINT_RECORD_SCHEMA | 259 | schema = POINT_RECORD_SCHEMA |
| 260 | required_keys = _schema_keys(schema, required=True) | 260 | required_keys = _schema_keys(schema, required=True) |
| 261 | with np.load(npz_path) as data: | 261 | with np.load(npz_path) as data: |
| 262 | missing = [key for key in required_keys if key not in data.files] | 262 | missing = [key for key in required_keys if key not in data.files] |
| 285 | record[key] = spec.fill(n_points) | 285 | record[key] = spec.fill(n_points) |
| 286 | return {key: record[key] for key in schema} | 286 | return {key: record[key] for key in schema} |
| 287 | 287 | ||
| 288 | 288 | ||
| 289 | def save_points_npz(path: str | Path, record: Mapping[str, np.ndarray]) -> Path: | 289 | def save_points_npz(path: str | pathlib.Path, record: abc.Mapping[str, np.ndarray]) -> pathlib.Path: |
| 290 | """Write a point record as a compressed NPZ (inverse of :func:`load_points_npz`). | 290 | """Write a point record as a compressed NPZ (inverse of :func:`load_points_npz`). |
| 291 | 291 | ||
| 292 | Every key in :data:`POINT_RECORD_KEYS` must be present, ``number_of_returns`` | 292 | Every key in :data:`POINT_RECORD_KEYS` must be present, ``number_of_returns`` |
| 293 | included: it is optional on load only, to read datasets that predate it. | 293 | included: it is optional on load only, to read datasets that predate it. |
| 303 | Raises: | 303 | Raises: |
| 304 | ValueError: A key is missing, an array has the wrong shape, or a member | 304 | ValueError: A key is missing, an array has the wrong shape, or a member |
| 305 | with a declared storage dtype cannot be cast to it losslessly. | 305 | with a declared storage dtype cannot be cast to it losslessly. |
| 306 | """ | 306 | """ |
| 307 | npz_path = Path(path) | 307 | npz_path = pathlib.Path(path) |
| 308 | schema = POINT_RECORD_SCHEMA | 308 | schema = POINT_RECORD_SCHEMA |
| 309 | keys = _schema_keys(schema) | 309 | keys = _schema_keys(schema) |
| 310 | missing = [key for key in keys if key not in record] | 310 | missing = [key for key in keys if key not in record] |
| 311 | if missing: | 311 | if missing: |
| 323 | np.savez_compressed(npz_path, **payload) | 323 | np.savez_compressed(npz_path, **payload) |
| 324 | return npz_path | 324 | return npz_path |
| 325 | 325 | ||
| 326 | 326 | ||
| 327 | def mask_record(record: Mapping[str, np.ndarray], mask: np.ndarray) -> PointRecord: | 327 | def mask_record(record: abc.Mapping[str, np.ndarray], mask: np.ndarray) -> PointRecord: |
| 328 | """Select the same points from every member of a point record. | 328 | """Select the same points from every member of a point record. |
| 329 | 329 | ||
| 330 | Schema-agnostic: whatever keys the record carries are all indexed with | 330 | Schema-agnostic: whatever keys the record carries are all indexed with |
| 331 | *mask* along axis 0, so ``points`` keeps its ``(n, 3)`` rows while the | 331 | *mask* along axis 0, so ``points`` keeps its ``(n, 3)`` rows while the |
| 367 | ) | 367 | ) |
| 368 | return {key: array[mask] for key, array in arrays.items()} | 368 | return {key: array[mask] for key, array in arrays.items()} |
| 369 | 369 | ||
| 370 | 370 | ||
| 371 | def concat_records(records: Sequence[Mapping[str, np.ndarray]]) -> PointRecord: | 371 | def concat_records(records: abc.Sequence[abc.Mapping[str, np.ndarray]]) -> PointRecord: |
| 372 | """Concatenate point records member-by-member along axis 0. | 372 | """Concatenate point records member-by-member along axis 0. |
| 373 | 373 | ||
| 374 | Schema-agnostic: every key of the first record is concatenated across all | 374 | Schema-agnostic: every key of the first record is concatenated across all |
| 375 | records, so ``points`` grows by rows and the ancillary members by elements. | 375 | records, so ``points`` grows by rows and the ancillary members by elements. |
| 408 | } | 408 | } |
| 409 | 409 | ||
| 410 | 410 | ||
| 411 | def load_segment_points( | 411 | def load_segment_points( |
| 412 | files: Sequence[str | Path], | 412 | files: abc.Sequence[str | pathlib.Path], |
| 413 | ) -> tuple[PointRecord, np.ndarray, list[str]]: | 413 | ) -> tuple[PointRecord, np.ndarray, list[str]]: |
| 414 | """Load and concatenate multiple point NPZs for one segment. | 414 | """Load and concatenate multiple point NPZs for one segment. |
| 415 | 415 | ||
| 416 | Returns ``(merged_record, point_file_ids, file_stems)`` where | 416 | Returns ``(merged_record, point_file_ids, file_stems)`` where |
| 417 | ``point_file_ids[i]`` is the index into *files* for merged point ``i``, | 417 | ``point_file_ids[i]`` is the index into *files* for merged point ``i``, |
| 418 | and ``file_stems`` holds ``Path(f).stem`` for each input in order. | 418 | and ``file_stems`` holds ``Path(f).stem`` for each input in order. |
| 419 | """ | 419 | """ |
| 420 | npz_files = [Path(path) for path in files] | 420 | npz_files = [pathlib.Path(path) for path in files] |
| 421 | if not npz_files: | 421 | if not npz_files: |
| 422 | raise FileNotFoundError("No *_points.npz files provided for segment load") | 422 | raise FileNotFoundError("No *_points.npz files provided for segment load") |
| 423 | 423 | ||
| 424 | records: list[PointRecord] = [] | 424 | records: list[PointRecord] = [] |
| 441 | return merged, point_file_ids, file_stems | 441 | return merged, point_file_ids, file_stems |
| 442 | 442 | ||
| 443 | 443 | ||
| 444 | def geoshift_from_mapping( | 444 | def geoshift_from_mapping( |
| 445 | mapping: Mapping[str, object], | 445 | mapping: abc.Mapping[str, object], |
| 446 | *, | 446 | *, |
| 447 | source: str = "geoshift mapping", | 447 | source: str = "geoshift mapping", |
| 448 | ) -> np.ndarray: | 448 | ) -> np.ndarray: |
| 449 | """Convert an already-parsed geoshift mapping to a shape ``(3,)`` float64 array. | 449 | """Convert an already-parsed geoshift mapping to a shape ``(3,)`` float64 array. |
| 464 | Raises: | 464 | Raises: |
| 465 | ValueError: If *mapping* is not an object, a key is missing, or a | 465 | ValueError: If *mapping* is not an object, a key is missing, or a |
| 466 | value is not numeric. | 466 | value is not numeric. |
| 467 | """ | 467 | """ |
| 468 | if not isinstance(mapping, Mapping): | 468 | if not isinstance(mapping, abc.Mapping): |
| 469 | raise ValueError( | 469 | raise ValueError( |
| 470 | f"{source}: geoshift JSON must be an object with keys x/y/z, " | 470 | f"{source}: geoshift JSON must be an object with keys x/y/z, " |
| 471 | f"got {type(mapping).__name__}" | 471 | f"got {type(mapping).__name__}" |
| 472 | ) | 472 | ) |
| 473 | values: Mapping[str, object] = mapping | 473 | values: abc.Mapping[str, object] = mapping |
| 474 | nested = mapping.get("geoshift") | 474 | nested = mapping.get("geoshift") |
| 475 | if isinstance(nested, Mapping): | 475 | if isinstance(nested, abc.Mapping): |
| 476 | values = nested | 476 | values = nested |
| 477 | 477 | ||
| 478 | missing = [key for key in ("x", "y", "z") if key not in values] | 478 | missing = [key for key in ("x", "y", "z") if key not in values] |
| 479 | if missing: | 479 | if missing: |
| 489 | f"{[values['x'], values['y'], values['z']]!r}" | 489 | f"{[values['x'], values['y'], values['z']]!r}" |
| 490 | ) from exc | 490 | ) from exc |
| 491 | 491 | ||
| 492 | 492 | ||
| 493 | def load_geoshift(path: str | Path) -> np.ndarray: | 493 | def load_geoshift(path: str | pathlib.Path) -> np.ndarray: |
| 494 | """Load a Step-3 ``run3_geoshift.json`` as a shape ``(3,)`` float64 array. | 494 | """Load a Step-3 ``run3_geoshift.json`` as a shape ``(3,)`` float64 array. |
| 495 | 495 | ||
| 496 | Expected JSON layout (written by segmentation-trajectory SegmentMapper):: | 496 | Expected JSON layout (written by segmentation-trajectory SegmentMapper):: |
| 497 | 497 | ||
| 498 | {"x": float, "y": float, "z": float} | 498 | {"x": float, "y": float, "z": float} |
| 499 | 499 | ||
| 500 | The value is returned as recorded; no sign is applied (see module docstring). | 500 | The value is returned as recorded; no sign is applied (see module docstring). |
| 501 | """ | 501 | """ |
| 502 | geoshift_path = Path(path) | 502 | geoshift_path = pathlib.Path(path) |
| 503 | with geoshift_path.open(encoding="utf-8") as handle: | 503 | with geoshift_path.open(encoding="utf-8") as handle: |
| 504 | data = json.load(handle) | 504 | data = json.load(handle) |
| 505 | return geoshift_from_mapping(data, source=str(geoshift_path)) | 505 | return geoshift_from_mapping(data, source=str(geoshift_path)) |
| 506 | 506 | ||
| 507 | 507 | ||
| 508 | def geoshift_candidate_paths(directory: str | Path) -> list[Path]: | 508 | def geoshift_candidate_paths(directory: str | pathlib.Path) -> list[pathlib.Path]: |
| 509 | """List the ``run3_geoshift.json`` paths the fleet's three conventions use. | 509 | """List the ``run3_geoshift.json`` paths the fleet's three conventions use. |
| 510 | 510 | ||
| 511 | In lookup order: | 511 | In lookup order: |
| 512 | 512 |
| 523 | 523 | ||
| 524 | Returns: | 524 | Returns: |
| 525 | The candidate paths in lookup order (existence is not checked). | 525 | The candidate paths in lookup order (existence is not checked). |
| 526 | """ | 526 | """ |
| 527 | base = Path(directory) | 527 | base = pathlib.Path(directory) |
| 528 | return [ | 528 | return [ |
| 529 | base / GEOSHIFT_NAME, | 529 | base / GEOSHIFT_NAME, |
| 530 | base / LANE_POINTS_DIR_NAME / GEOSHIFT_NAME, | 530 | base / LANE_POINTS_DIR_NAME / GEOSHIFT_NAME, |
| 531 | base.parent / GEOSHIFT_NAME, | 531 | base.parent / GEOSHIFT_NAME, |
| 532 | ] | 532 | ] |
| 533 | 533 | ||
| 534 | 534 | ||
| 535 | def find_geoshift_or_none(directory: str | Path) -> np.ndarray | None: | 535 | def find_geoshift_or_none(directory: str | pathlib.Path) -> np.ndarray | None: |
| 536 | """Find and load a geoshift near *directory*, or return ``None``. | 536 | """Find and load a geoshift near *directory*, or return ``None``. |
| 537 | 537 | ||
| 538 | Searches :func:`geoshift_candidate_paths` in order and loads the first | 538 | Searches :func:`geoshift_candidate_paths` in order and loads the first |
| 539 | existing file. Datasets processed without a geoshift have no such file and | 539 | existing file. Datasets processed without a geoshift have no such file and |
| 561 | logger.debug("Using geoshift %s for %s", existing[0], directory) | 561 | logger.debug("Using geoshift %s for %s", existing[0], directory) |
| 562 | return load_geoshift(existing[0]) | 562 | return load_geoshift(existing[0]) |
| 563 | 563 | ||
| 564 | 564 | ||
| 565 | def find_geoshift(directory: str | Path) -> np.ndarray: | 565 | def find_geoshift(directory: str | pathlib.Path) -> np.ndarray: |
| 566 | """Find and load a geoshift near *directory*, raising when absent. | 566 | """Find and load a geoshift near *directory*, raising when absent. |
| 567 | 567 | ||
| 568 | Args: | 568 | Args: |
| 569 | directory: Dataset root, ``lane_points`` directory, or segment directory. | 569 | directory: Dataset root, ``lane_points`` directory, or segment directory. |
| 603 | ) from exc | 603 | ) from exc |
| 604 | 604 | ||
| 605 | 605 | ||
| 606 | def _normalize_file_patterns(file_patterns: object, *, source: str) -> list[str]: | 606 | def _normalize_file_patterns(file_patterns: object, *, source: str) -> list[str]: |
| 607 | raw_patterns: Sequence[object] | 607 | raw_patterns: abc.Sequence[object] |
| 608 | if isinstance(file_patterns, str): | 608 | if isinstance(file_patterns, str): |
| 609 | raw_patterns = [file_patterns] | 609 | raw_patterns = [file_patterns] |
| 610 | elif isinstance(file_patterns, Sequence): | 610 | elif isinstance(file_patterns, abc.Sequence): |
| 611 | raw_patterns = file_patterns | 611 | raw_patterns = file_patterns |
| 612 | else: | 612 | else: |
| 613 | raise ValueError( | 613 | raise ValueError( |
| 614 | f"{source} must be a string or a sequence of strings, got " | 614 | f"{source} must be a string or a sequence of strings, got " |
| 623 | return normalized | 623 | return normalized |
| 624 | 624 | ||
| 625 | 625 | ||
| 626 | def normalize_segment_file_blacklist( | 626 | def normalize_segment_file_blacklist( |
| 627 | mapping: Mapping[object, object] | None, | 627 | mapping: abc.Mapping[object, object] | None, |
| 628 | ) -> dict[int, list[str]]: | 628 | ) -> dict[int, list[str]]: |
| 629 | """Normalize a per-segment fnmatch blacklist mapping. | 629 | """Normalize a per-segment fnmatch blacklist mapping. |
| 630 | 630 | ||
| 631 | Keys may be ints or ``segment_<idx>`` strings. Values are a string or a | 631 | Keys may be ints or ``segment_<idx>`` strings. Values are a string or a |
| 633 | the same segment are merged and sorted. | 633 | the same segment are merged and sorted. |
| 634 | """ | 634 | """ |
| 635 | if mapping is None: | 635 | if mapping is None: |
| 636 | return {} | 636 | return {} |
| 637 | if not isinstance(mapping, Mapping): | 637 | if not isinstance(mapping, abc.Mapping): |
| 638 | raise ValueError("npz_blacklist_by_segment must be a mapping of segment to files") | 638 | raise ValueError("npz_blacklist_by_segment must be a mapping of segment to files") |
| 639 | 639 | ||
| 640 | normalized: dict[int, set[str]] = {} | 640 | normalized: dict[int, set[str]] = {} |
| 641 | for segment_key, file_patterns in mapping.items(): | 641 | for segment_key, file_patterns in mapping.items(): |
| 653 | } | 653 | } |
| 654 | 654 | ||
| 655 | 655 | ||
| 656 | def filter_segment_files( | 656 | def filter_segment_files( |
| 657 | files: Sequence[str | Path], | 657 | files: abc.Sequence[str | pathlib.Path], |
| 658 | segment_index: int, | 658 | segment_index: int, |
| 659 | blacklist: Mapping[int, Sequence[str]], | 659 | blacklist: abc.Mapping[int, abc.Sequence[str]], |
| 660 | ) -> tuple[list[Path], list[Path]]: | 660 | ) -> tuple[list[pathlib.Path], list[pathlib.Path]]: |
| 661 | """Split input files into kept and blacklisted buckets via fnmatch. | 661 | """Split input files into kept and blacklisted buckets via fnmatch. |
| 662 | 662 | ||
| 663 | Patterns from ``blacklist[segment_index]`` are matched against both the | 663 | Patterns from ``blacklist[segment_index]`` are matched against both the |
| 664 | file name and the full POSIX path. Returns ``(kept_files, excluded_files)``. | 664 | file name and the full POSIX path. Returns ``(kept_files, excluded_files)``. |
| 665 | """ | 665 | """ |
| 666 | kept_files: list[Path] = [] | 666 | kept_files: list[pathlib.Path] = [] |
| 667 | excluded_files: list[Path] = [] | 667 | excluded_files: list[pathlib.Path] = [] |
| 668 | file_patterns = list(blacklist.get(segment_index, [])) | 668 | file_patterns = list(blacklist.get(segment_index, [])) |
| 669 | 669 | ||
| 670 | for raw_path in files: | 670 | for raw_path in files: |
| 671 | npz_path = Path(raw_path) | 671 | npz_path = pathlib.Path(raw_path) |
| 672 | path_text = npz_path.as_posix() | 672 | path_text = npz_path.as_posix() |
| 673 | file_name = npz_path.name | 673 | file_name = npz_path.name |
| 674 | is_blacklisted = any( | 674 | is_blacklisted = any( |
| 675 | fnmatch(file_name, pattern) or fnmatch(path_text, pattern) | 675 | fnmatch.fnmatch(file_name, pattern) or fnmatch.fnmatch(path_text, pattern) |
| 676 | for pattern in file_patterns | 676 | for pattern in file_patterns |
| 677 | ) | 677 | ) |
| 678 | if is_blacklisted: | 678 | if is_blacklisted: |
| 679 | excluded_files.append(npz_path) | 679 | excluded_files.append(npz_path) |
| 695 | ) | 695 | ) |
| 696 | return kept_files, excluded_files | 696 | return kept_files, excluded_files |
| 697 | 697 | ||
| 698 | 698 | ||
| 699 | def read_points_header(path: str | Path) -> tuple[int, bool]: | 699 | def read_points_header(path: str | pathlib.Path) -> tuple[int, bool]: |
| 700 | """Return ``(row_count, chunk_streamable)`` for a record's ``points.npy`` member. | 700 | """Return ``(row_count, chunk_streamable)`` for a record's ``points.npy`` member. |
| 701 | 701 | ||
| 702 | Reads only the NPY header inside the NPZ ZIP container, so the point data | 702 | Reads only the NPY header inside the NPZ ZIP container, so the point data |
| 703 | is never materialised. Streaming needs an uncompressed (``ZIP_STORED``) | 703 | is never materialised. Streaming needs an uncompressed (``ZIP_STORED``) |
| 710 | Returns: | 710 | Returns: |
| 711 | ``(rows, streamable)``. ``rows`` is ``-1`` when the header is | 711 | ``(rows, streamable)``. ``rows`` is ``-1`` when the header is |
| 712 | unreadable (missing member, corrupt archive, unsupported NPY version). | 712 | unreadable (missing member, corrupt archive, unsupported NPY version). |
| 713 | """ | 713 | """ |
| 714 | npz_path = Path(path) | 714 | npz_path = pathlib.Path(path) |
| 715 | try: | 715 | try: |
| 716 | with zipfile.ZipFile(npz_path) as archive: | 716 | with zipfile.ZipFile(npz_path) as archive: |
| 717 | info = archive.getinfo("points.npy") | 717 | info = archive.getinfo("points.npy") |
| 718 | stored = info.compress_type == zipfile.ZIP_STORED | 718 | stored = info.compress_type == zipfile.ZIP_STORED |
| 733 | # caught here so a partial write degrades to "unreadable" instead. | 733 | # caught here so a partial write degrades to "unreadable" instead. |
| 734 | return -1, False | 734 | return -1, False |
| 735 | 735 | ||
| 736 | 736 | ||
| 737 | def _iter_points_chunks_streamed(path: Path, chunk_points: int) -> Iterator[np.ndarray]: | 737 | def _iter_points_chunks_streamed(path: pathlib.Path, chunk_points: int) -> abc.Iterator[np.ndarray]: |
| 738 | """Stream the ``points.npy`` member of a ZIP_STORED npz in row chunks.""" | 738 | """Stream the ``points.npy`` member of a ZIP_STORED npz in row chunks.""" |
| 739 | with zipfile.ZipFile(path) as archive, archive.open("points.npy") as handle: | 739 | with zipfile.ZipFile(path) as archive, archive.open("points.npy") as handle: |
| 740 | version = np.lib.format.read_magic(handle) | 740 | version = np.lib.format.read_magic(handle) |
| 741 | if version == (1, 0): | 741 | if version == (1, 0): |
| 753 | buffer = handle.read(count * row_bytes) | 753 | buffer = handle.read(count * row_bytes) |
| 754 | yield np.frombuffer(buffer, dtype=dtype).reshape(count, cols).copy() | 754 | yield np.frombuffer(buffer, dtype=dtype).reshape(count, cols).copy() |
| 755 | 755 | ||
| 756 | 756 | ||
| 757 | def iter_points_chunks(path: str | Path, chunk_points: int) -> Iterator[np.ndarray]: | 757 | def iter_points_chunks(path: str | pathlib.Path, chunk_points: int) -> abc.Iterator[np.ndarray]: |
| 758 | """Yield a record's ``points`` rows, chunked when the member is streamable. | 758 | """Yield a record's ``points`` rows, chunked when the member is streamable. |
| 759 | 759 | ||
| 760 | Records at or below *chunk_points* rows -- and any record whose | 760 | Records at or below *chunk_points* rows -- and any record whose |
| 761 | ``points.npy`` member is compressed, Fortran-ordered or otherwise | 761 | ``points.npy`` member is compressed, Fortran-ordered or otherwise |
| 774 | Yields: | 774 | Yields: |
| 775 | ``(rows_i, 3)`` point chunks in file order; the final chunk holds the | 775 | ``(rows_i, 3)`` point chunks in file order; the final chunk holds the |
| 776 | remainder and may be shorter. | 776 | remainder and may be shorter. |
| 777 | """ | 777 | """ |
| 778 | npz_path = Path(path) | 778 | npz_path = pathlib.Path(path) |
| 779 | rows, streamable = read_points_header(npz_path) | 779 | rows, streamable = read_points_header(npz_path) |
| 780 | if streamable and chunk_points > 0 and rows > chunk_points: | 780 | if streamable and chunk_points > 0 and rows > chunk_points: |
| 781 | logger.info( | 781 | logger.info( |
| 782 | "%s: streaming %d points in chunks of %d", | 782 | "%s: streaming %d points in chunks of %d", |
| 789 | with np.load(npz_path) as data: | 789 | with np.load(npz_path) as data: |
| 790 | yield np.asarray(data["points"]) | 790 | yield np.asarray(data["points"]) |
| 791 | 791 | ||
| 792 | 792 | ||
| 793 | @dataclass(frozen=True) | 793 | @dataclasses.dataclass(frozen=True) |
| 794 | class RecordSpan: | 794 | class RecordSpan: |
| 795 | """One input record's identity and row range within a concatenated cloud. | 795 | """One input record's identity and row range within a concatenated cloud. |
| 796 | 796 | ||
| 797 | Attributes: | 797 | Attributes: |
| 809 | """One past the last global row index contributed by this file.""" | 809 | """One past the last global row index contributed by this file.""" |
| 810 | return self.offset + self.count | 810 | return self.offset + self.count |
| 811 | 811 | ||
| 812 | 812 | ||
| 813 | def discover_run3_files(segment_dir: str | Path) -> list[Path]: | 813 | def discover_run3_files(segment_dir: str | pathlib.Path) -> list[pathlib.Path]: |
| 814 | """List a segment directory's ``*_run3_points.npz`` records in load order. | 814 | """List a segment directory's ``*_run3_points.npz`` records in load order. |
| 815 | 815 | ||
| 816 | Partial writes carry a trailing marker suffix (``*.npz.part`` and friends); | 816 | Partial writes carry a trailing marker suffix (``*.npz.part`` and friends); |
| 817 | those are skipped rather than handed to a loader that would crash on them. | 817 | those are skipped rather than handed to a loader that would crash on them. |
| 822 | Returns: | 822 | Returns: |
| 823 | Sorted, complete record paths. Empty when the directory is missing or | 823 | Sorted, complete record paths. Empty when the directory is missing or |
| 824 | holds no records. | 824 | holds no records. |
| 825 | """ | 825 | """ |
| 826 | seg_dir = Path(segment_dir) | 826 | seg_dir = pathlib.Path(segment_dir) |
| 827 | if not seg_dir.is_dir(): | 827 | if not seg_dir.is_dir(): |
| 828 | logger.warning("Segment directory missing: %s", seg_dir) | 828 | logger.warning("Segment directory missing: %s", seg_dir) |
| 829 | return [] | 829 | return [] |
| 830 | 830 | ||
| 831 | found: list[Path] = [] | 831 | found: list[pathlib.Path] = [] |
| 832 | for path in sorted(seg_dir.glob(f"{RUN3_POINTS_GLOB}*")): | 832 | for path in sorted(seg_dir.glob(f"{RUN3_POINTS_GLOB}*")): |
| 833 | if not path.name.endswith(RUN3_POINTS_SUFFIX): | 833 | if not path.name.endswith(RUN3_POINTS_SUFFIX): |
| 834 | logger.info("Skipping incomplete/partial run3 record: %s", path.name) | 834 | logger.info("Skipping incomplete/partial run3 record: %s", path.name) |
| 835 | continue | 835 | continue |
| 837 | return found | 837 | return found |
| 838 | 838 | ||
| 839 | 839 | ||
| 840 | def concat_points_npz( | 840 | def concat_points_npz( |
| 841 | files: Sequence[str | Path], | 841 | files: abc.Sequence[str | pathlib.Path], |
| 842 | *, | 842 | *, |
| 843 | target_dtypes: Mapping[str, np.dtype] | None = None, | 843 | target_dtypes: abc.Mapping[str, np.dtype] | None = None, |
| 844 | ) -> tuple[PointRecord, list[RecordSpan]]: | 844 | ) -> tuple[PointRecord, list[RecordSpan]]: |
| 845 | """Concatenate point records after validating a consistent schema. | 845 | """Concatenate point records after validating a consistent schema. |
| 846 | 846 | ||
| 847 | Each file must satisfy the :func:`load_points_npz` contract. By default | 847 | Each file must satisfy the :func:`load_points_npz` contract. By default |
| 872 | FileNotFoundError: If *files* is empty. | 872 | FileNotFoundError: If *files* is empty. |
| 873 | ValueError: If a file violates the point-record contract or its dtypes | 873 | ValueError: If a file violates the point-record contract or its dtypes |
| 874 | disagree with the first file's (for keys without a target dtype). | 874 | disagree with the first file's (for keys without a target dtype). |
| 875 | """ | 875 | """ |
| 876 | npz_files = [Path(path) for path in files] | 876 | npz_files = [pathlib.Path(path) for path in files] |
| 877 | if not npz_files: | 877 | if not npz_files: |
| 878 | raise FileNotFoundError("No *_points.npz files provided for concatenation") | 878 | raise FileNotFoundError("No *_points.npz files provided for concatenation") |
| 879 | 879 | ||
| 880 | casts: dict[str, np.dtype] = { | 880 | casts: dict[str, np.dtype] = { |
| 911 | return concat_records(records), spans | 911 | return concat_records(records), spans |
| 912 | 912 | ||
| 913 | 913 | ||
| 914 | def load_run3_segment( | 914 | def load_run3_segment( |
| 915 | segment_dir: str | Path, | 915 | segment_dir: str | pathlib.Path, |
| 916 | *, | 916 | *, |
| 917 | target_dtypes: Mapping[str, np.dtype] | None = None, | 917 | target_dtypes: abc.Mapping[str, np.dtype] | None = None, |
| 918 | ) -> tuple[PointRecord, list[RecordSpan]]: | 918 | ) -> tuple[PointRecord, list[RecordSpan]]: |
| 919 | """Discover and concatenate one segment's run3 records. | 919 | """Discover and concatenate one segment's run3 records. |
| 920 | 920 | ||
| 921 | Combines :func:`discover_run3_files` (sorted glob, ``.part`` skipped) with | 921 | Combines :func:`discover_run3_files` (sorted glob, ``.part`` skipped) with |
| 931 | Raises: | 931 | Raises: |
| 932 | FileNotFoundError: If the directory holds no complete run3 record. | 932 | FileNotFoundError: If the directory holds no complete run3 record. |
| 933 | ValueError: If the records' schema or dtypes disagree. | 933 | ValueError: If the records' schema or dtypes disagree. |
| 934 | """ | 934 | """ |
| 935 | seg_dir = Path(segment_dir) | 935 | seg_dir = pathlib.Path(segment_dir) |
| 936 | files = discover_run3_files(seg_dir) | 936 | files = discover_run3_files(seg_dir) |
| 937 | if not files: | 937 | if not files: |
| 938 | raise FileNotFoundError(f"No {RUN3_POINTS_GLOB} in {seg_dir}") | 938 | raise FileNotFoundError(f"No {RUN3_POINTS_GLOB} in {seg_dir}") |
| 939 | return concat_points_npz(files, target_dtypes=target_dtypes) | 939 | return concat_points_npz(files, target_dtypes=target_dtypes) |
| 1 | """Tests for ColorIntensityData selection and concatenation operations.""" | 1 | """Tests for ColorIntensityData selection and concatenation operations.""" |
| 2 | from dataclasses import dataclass, field | 2 | import dataclasses |
| 3 | 3 | ||
| 4 | import numpy as np | 4 | import numpy as np |
| 5 | import pytest | 5 | import pytest |
| 6 | 6 | ||
| 7 | from iolabs.common.color_intensity_data import ColorIntensityData | 7 | from iolabs.common import color_intensity_data |
| 8 | 8 | ||
| 9 | 9 | ||
| 10 | def _make_sample(n: int = 5, offset: int = 0) -> ColorIntensityData: | 10 | def _make_sample(n: int = 5, offset: int = 0) -> color_intensity_data.ColorIntensityData: |
| 11 | """Create a sample ColorIntensityData with n points and optional array offset.""" | 11 | """Create a sample ColorIntensityData with n points and optional array offset.""" |
| 12 | return ColorIntensityData( | 12 | return color_intensity_data.ColorIntensityData( |
| 13 | red=np.arange(offset, offset + n, dtype=np.uint8), | 13 | red=np.arange(offset, offset + n, dtype=np.uint8), |
| 14 | green=np.arange(offset + 10, offset + 10 + n, dtype=np.uint8), | 14 | green=np.arange(offset + 10, offset + 10 + n, dtype=np.uint8), |
| 15 | blue=np.arange(offset + 20, offset + 20 + n, dtype=np.uint8), | 15 | blue=np.arange(offset + 20, offset + 20 + n, dtype=np.uint8), |
| 16 | intensity=np.arange(offset + 100, offset + 100 + n, dtype=np.float64), | 16 | intensity=np.arange(offset + 100, offset + 100 + n, dtype=np.float64), |
| 85 | """Tests for the AI3D-382 number_of_returns field.""" | 85 | """Tests for the AI3D-382 number_of_returns field.""" |
| 86 | 86 | ||
| 87 | def test_defaults_to_zeros_when_omitted(self): | 87 | def test_defaults_to_zeros_when_omitted(self): |
| 88 | """Callers predating AI3D-382 get 0 (unknown) per point, never 1.""" | 88 | """Callers predating AI3D-382 get 0 (unknown) per point, never 1.""" |
| 89 | data = ColorIntensityData( | 89 | data = color_intensity_data.ColorIntensityData( |
| 90 | red=np.zeros(4, dtype=np.uint8), | 90 | red=np.zeros(4, dtype=np.uint8), |
| 91 | green=np.zeros(4, dtype=np.uint8), | 91 | green=np.zeros(4, dtype=np.uint8), |
| 92 | blue=np.zeros(4, dtype=np.uint8), | 92 | blue=np.zeros(4, dtype=np.uint8), |
| 93 | intensity=np.zeros(4, dtype=np.float64), | 93 | intensity=np.zeros(4, dtype=np.float64), |
| 99 | np.testing.assert_array_equal(data.number_of_returns, np.zeros(4, dtype=np.uint8)) | 99 | np.testing.assert_array_equal(data.number_of_returns, np.zeros(4, dtype=np.uint8)) |
| 100 | 100 | ||
| 101 | def test_zero_fill_survives_mask_and_append(self): | 101 | def test_zero_fill_survives_mask_and_append(self): |
| 102 | """A defaulted field stays aligned through the mask/concat paths.""" | 102 | """A defaulted field stays aligned through the mask/concat paths.""" |
| 103 | legacy = ColorIntensityData( | 103 | legacy = color_intensity_data.ColorIntensityData( |
| 104 | red=np.zeros(3, dtype=np.uint8), | 104 | red=np.zeros(3, dtype=np.uint8), |
| 105 | green=np.zeros(3, dtype=np.uint8), | 105 | green=np.zeros(3, dtype=np.uint8), |
| 106 | blue=np.zeros(3, dtype=np.uint8), | 106 | blue=np.zeros(3, dtype=np.uint8), |
| 107 | intensity=np.zeros(3, dtype=np.float64), | 107 | intensity=np.zeros(3, dtype=np.float64), |
| 123 | 123 | ||
| 124 | def test_extra_field_flows_through_mask_and_append(self): | 124 | def test_extra_field_flows_through_mask_and_append(self): |
| 125 | """A subclass field is masked and concatenated by the generic transforms.""" | 125 | """A subclass field is masked and concatenated by the generic transforms.""" |
| 126 | 126 | ||
| 127 | @dataclass | 127 | @dataclasses.dataclass |
| 128 | class WithClassification(ColorIntensityData): | 128 | class WithClassification(color_intensity_data.ColorIntensityData): |
| 129 | classification: np.ndarray | None = None | 129 | classification: np.ndarray | None = None |
| 130 | 130 | ||
| 131 | def _make(n: int, offset: int) -> WithClassification: | 131 | def _make(n: int, offset: int) -> WithClassification: |
| 132 | base = _make_sample(n, offset) | 132 | base = _make_sample(n, offset) |
| 160 | 160 | ||
| 161 | def test_non_init_subclass_field_is_not_passed_to_the_constructor(self): | 161 | def test_non_init_subclass_field_is_not_passed_to_the_constructor(self): |
| 162 | """A derived ``init=False`` field must not break the generic transforms.""" | 162 | """A derived ``init=False`` field must not break the generic transforms.""" |
| 163 | 163 | ||
| 164 | @dataclass | 164 | @dataclasses.dataclass |
| 165 | class WithPointCount(ColorIntensityData): | 165 | class WithPointCount(color_intensity_data.ColorIntensityData): |
| 166 | point_count: int = field(init=False, default=0) | 166 | point_count: int = dataclasses.field(init=False, default=0) |
| 167 | 167 | ||
| 168 | def __post_init__(self) -> None: | 168 | def __post_init__(self) -> None: |
| 169 | super().__post_init__() | 169 | super().__post_init__() |
| 170 | self.point_count = len(self.red) | 170 | self.point_count = len(self.red) |
| 195 | 195 | ||
| 196 | def test_subclass_may_declare_a_required_field(self): | 196 | def test_subclass_may_declare_a_required_field(self): |
| 197 | """A trailing defaulted field would make a required subclass field a TypeError.""" | 197 | """A trailing defaulted field would make a required subclass field a TypeError.""" |
| 198 | 198 | ||
| 199 | @dataclass | 199 | @dataclasses.dataclass |
| 200 | class WithRequiredClassification(ColorIntensityData): | 200 | class WithRequiredClassification(color_intensity_data.ColorIntensityData): |
| 201 | classification: np.ndarray | 201 | classification: np.ndarray |
| 202 | 202 | ||
| 203 | data = WithRequiredClassification( | 203 | data = WithRequiredClassification( |
| 204 | red=np.zeros(3, dtype=np.uint8), | 204 | red=np.zeros(3, dtype=np.uint8), |
| 217 | assert len(merged.classification) == 5 | 217 | assert len(merged.classification) == 5 |
| 218 | 218 | ||
| 219 | def test_pre_ai3d_382_fields_still_take_positional_args(self): | 219 | def test_pre_ai3d_382_fields_still_take_positional_args(self): |
| 220 | """The five original fields keep their positional order for old call sites.""" | 220 | """The five original fields keep their positional order for old call sites.""" |
| 221 | data = ColorIntensityData( | 221 | data = color_intensity_data.ColorIntensityData( |
| 222 | np.zeros(2, dtype=np.uint8), | 222 | np.zeros(2, dtype=np.uint8), |
| 223 | np.zeros(2, dtype=np.uint8), | 223 | np.zeros(2, dtype=np.uint8), |
| 224 | np.zeros(2, dtype=np.uint8), | 224 | np.zeros(2, dtype=np.uint8), |
| 225 | np.zeros(2, dtype=np.float64), | 225 | np.zeros(2, dtype=np.float64), |
| 230 | 230 | ||
| 231 | def test_number_of_returns_is_not_positional(self): | 231 | def test_number_of_returns_is_not_positional(self): |
| 232 | """Passing it as a sixth positional arg is a TypeError, not a silent mismatch.""" | 232 | """Passing it as a sixth positional arg is a TypeError, not a silent mismatch.""" |
| 233 | with pytest.raises(TypeError): | 233 | with pytest.raises(TypeError): |
| 234 | ColorIntensityData( | 234 | color_intensity_data.ColorIntensityData( |
| 235 | np.zeros(2, dtype=np.uint8), | 235 | np.zeros(2, dtype=np.uint8), |
| 236 | np.zeros(2, dtype=np.uint8), | 236 | np.zeros(2, dtype=np.uint8), |
| 237 | np.zeros(2, dtype=np.uint8), | 237 | np.zeros(2, dtype=np.uint8), |
| 238 | np.zeros(2, dtype=np.float64), | 238 | np.zeros(2, dtype=np.float64), |
| 244 | class TestNumberOfReturnsDtypeContract: | 244 | class TestNumberOfReturnsDtypeContract: |
| 245 | """The constructor enforces the same uint8 contract as segment_points_io.""" | 245 | """The constructor enforces the same uint8 contract as segment_points_io.""" |
| 246 | 246 | ||
| 247 | @staticmethod | 247 | @staticmethod |
| 248 | def _make(number_of_returns: np.ndarray) -> ColorIntensityData: | 248 | def _make(number_of_returns: np.ndarray) -> color_intensity_data.ColorIntensityData: |
| 249 | n = len(number_of_returns) | 249 | n = len(number_of_returns) |
| 250 | return ColorIntensityData( | 250 | return color_intensity_data.ColorIntensityData( |
| 251 | red=np.zeros(n, dtype=np.uint8), | 251 | red=np.zeros(n, dtype=np.uint8), |
| 252 | green=np.zeros(n, dtype=np.uint8), | 252 | green=np.zeros(n, dtype=np.uint8), |
| 253 | blue=np.zeros(n, dtype=np.uint8), | 253 | blue=np.zeros(n, dtype=np.uint8), |
| 254 | intensity=np.zeros(n, dtype=np.float64), | 254 | intensity=np.zeros(n, dtype=np.float64), |
| 2 | 2 | ||
| 3 | import functools | 3 | import functools |
| 4 | import json | 4 | import json |
| 5 | import logging | 5 | import logging |
| 6 | from pathlib import Path | 6 | import pathlib |
| 7 | from types import MappingProxyType | 7 | import types |
| 8 | 8 | ||
| 9 | import numpy as np | 9 | import numpy as np |
| 10 | import pytest | 10 | import pytest |
| 11 | 11 | ||
| 12 | from iolabs.common import segment_points_io | 12 | from iolabs.common import segment_points_io |
| 13 | from iolabs.common.segment_points_io import ( | 13 | |
| 14 | GEOSHIFT_NAME, | 14 | # Import-time snapshot of the shipped registry object. Monkeypatching the |
| 15 | NUMBER_OF_RETURNS_KEY, | 15 | # module attribute must not mutate this binding; the leak test below checks |
| 16 | OPTIONAL_POINT_RECORD_KEYS, | 16 | # both this snapshot and the live ``segment_points_io.POINT_RECORD_SCHEMA``. |
| 17 | POINT_RECORD_KEYS, | 17 | _POINT_RECORD_SCHEMA_AT_IMPORT = segment_points_io.POINT_RECORD_SCHEMA |
| 18 | POINT_RECORD_SCHEMA, | 18 | _POINT_RECORD_KEYS_AT_IMPORT = segment_points_io.POINT_RECORD_KEYS |
| 19 | REQUIRED_POINT_RECORD_KEYS, | ||
| 20 | PointFieldSpec, | ||
| 21 | RecordSpan, | ||
| 22 | concat_points_npz, | ||
| 23 | concat_records, | ||
| 24 | discover_run3_files, | ||
| 25 | filter_segment_files, | ||
| 26 | find_geoshift, | ||
| 27 | find_geoshift_or_none, | ||
| 28 | geoshift_candidate_paths, | ||
| 29 | geoshift_from_mapping, | ||
| 30 | iter_points_chunks, | ||
| 31 | load_geoshift, | ||
| 32 | load_points_npz, | ||
| 33 | load_run3_segment, | ||
| 34 | load_segment_points, | ||
| 35 | mask_record, | ||
| 36 | normalize_segment_file_blacklist, | ||
| 37 | parse_segment_key, | ||
| 38 | read_points_header, | ||
| 39 | save_points_npz, | ||
| 40 | ) | ||
| 41 | 19 | ||
| 42 | 20 | ||
| 43 | def _make_record(n: int, *, seed: int = 0) -> dict[str, np.ndarray]: | 21 | def _make_record(n: int, *, seed: int = 0) -> dict[str, np.ndarray]: |
| 44 | rng = np.random.default_rng(seed) | 22 | rng = np.random.default_rng(seed) |
| 55 | 33 | ||
| 56 | def _make_legacy_record(n: int, *, seed: int = 0) -> dict[str, np.ndarray]: | 34 | def _make_legacy_record(n: int, *, seed: int = 0) -> dict[str, np.ndarray]: |
| 57 | """Build a pre-AI3D-382 record: every key except ``number_of_returns``.""" | 35 | """Build a pre-AI3D-382 record: every key except ``number_of_returns``.""" |
| 58 | record = _make_record(n, seed=seed) | 36 | record = _make_record(n, seed=seed) |
| 59 | del record[NUMBER_OF_RETURNS_KEY] | 37 | del record[segment_points_io.NUMBER_OF_RETURNS_KEY] |
| 60 | return record | 38 | return record |
| 61 | 39 | ||
| 62 | 40 | ||
| 63 | def test_save_load_points_npz_round_trip(tmp_path: Path) -> None: | 41 | def test_save_load_points_npz_round_trip(tmp_path: pathlib.Path) -> None: |
| 64 | record = _make_record(7, seed=1) | 42 | record = _make_record(7, seed=1) |
| 65 | path = tmp_path / "scan_a_run3_points.npz" | 43 | path = tmp_path / "scan_a_run3_points.npz" |
| 66 | save_points_npz(path, record) | 44 | segment_points_io.save_points_npz(path, record) |
| 67 | 45 | ||
| 68 | loaded = load_points_npz(path) | 46 | loaded = segment_points_io.load_points_npz(path) |
| 69 | assert list(loaded.keys()) == list(POINT_RECORD_KEYS) | 47 | assert list(loaded.keys()) == list(segment_points_io.POINT_RECORD_KEYS) |
| 70 | for key in POINT_RECORD_KEYS: | 48 | for key in segment_points_io.POINT_RECORD_KEYS: |
| 71 | np.testing.assert_array_equal(loaded[key], record[key]) | 49 | np.testing.assert_array_equal(loaded[key], record[key]) |
| 72 | 50 | ||
| 73 | 51 | ||
| 74 | def test_load_points_npz_rejects_missing_keys(tmp_path: Path) -> None: | 52 | def test_load_points_npz_rejects_missing_keys(tmp_path: pathlib.Path) -> None: |
| 75 | path = tmp_path / "broken.npz" | 53 | path = tmp_path / "broken.npz" |
| 76 | np.savez_compressed(path, points=np.zeros((2, 3), dtype=np.float32), red=np.zeros(2)) | 54 | np.savez_compressed(path, points=np.zeros((2, 3), dtype=np.float32), red=np.zeros(2)) |
| 77 | with pytest.raises(ValueError, match="missing required key"): | 55 | with pytest.raises(ValueError, match="missing required key"): |
| 78 | load_points_npz(path) | 56 | segment_points_io.load_points_npz(path) |
| 79 | 57 | ||
| 80 | 58 | ||
| 81 | def test_load_points_npz_rejects_bad_points_shape(tmp_path: Path) -> None: | 59 | def test_load_points_npz_rejects_bad_points_shape(tmp_path: pathlib.Path) -> None: |
| 82 | record = _make_record(3, seed=2) | 60 | record = _make_record(3, seed=2) |
| 83 | record["points"] = np.zeros((3, 2), dtype=np.float32) | 61 | record["points"] = np.zeros((3, 2), dtype=np.float32) |
| 84 | path = tmp_path / "bad_shape.npz" | 62 | path = tmp_path / "bad_shape.npz" |
| 85 | np.savez_compressed(path, **record) | 63 | np.savez_compressed(path, **record) |
| 86 | with pytest.raises(ValueError, match=r"shape \(N, 3\)"): | 64 | with pytest.raises(ValueError, match=r"shape \(N, 3\)"): |
| 87 | load_points_npz(path) | 65 | segment_points_io.load_points_npz(path) |
| 88 | 66 | ||
| 89 | 67 | ||
| 90 | def test_load_points_npz_rejects_row_count_mismatch(tmp_path: Path) -> None: | 68 | def test_load_points_npz_rejects_row_count_mismatch(tmp_path: pathlib.Path) -> None: |
| 91 | record = _make_record(4, seed=3) | 69 | record = _make_record(4, seed=3) |
| 92 | record["intensity"] = record["intensity"][:2] | 70 | record["intensity"] = record["intensity"][:2] |
| 93 | path = tmp_path / "mismatch.npz" | 71 | path = tmp_path / "mismatch.npz" |
| 94 | np.savez_compressed(path, **record) | 72 | np.savez_compressed(path, **record) |
| 95 | with pytest.raises(ValueError, match="intensity"): | 73 | with pytest.raises(ValueError, match="intensity"): |
| 96 | load_points_npz(path) | 74 | segment_points_io.load_points_npz(path) |
| 97 | 75 | ||
| 98 | 76 | ||
| 99 | def test_load_points_npz_rejects_scalar_ancillary(tmp_path: Path) -> None: | 77 | def test_load_points_npz_rejects_scalar_ancillary(tmp_path: pathlib.Path) -> None: |
| 100 | record = _make_record(3, seed=4) | 78 | record = _make_record(3, seed=4) |
| 101 | record["red"] = np.array(42, dtype=np.uint16) | 79 | record["red"] = np.array(42, dtype=np.uint16) |
| 102 | path = tmp_path / "scalar_ancillary.npz" | 80 | path = tmp_path / "scalar_ancillary.npz" |
| 103 | np.savez_compressed(path, **record) | 81 | np.savez_compressed(path, **record) |
| 104 | with pytest.raises(ValueError, match=r"'red' must have shape \(N,\), got \(\)"): | 82 | with pytest.raises(ValueError, match=r"'red' must have shape \(N,\), got \(\)"): |
| 105 | load_points_npz(path) | 83 | segment_points_io.load_points_npz(path) |
| 106 | 84 | ||
| 107 | 85 | ||
| 108 | def test_load_points_npz_rejects_column_vector_ancillary(tmp_path: Path) -> None: | 86 | def test_load_points_npz_rejects_column_vector_ancillary(tmp_path: pathlib.Path) -> None: |
| 109 | record = _make_record(3, seed=5) | 87 | record = _make_record(3, seed=5) |
| 110 | record["intensity"] = record["intensity"].reshape(3, 1) | 88 | record["intensity"] = record["intensity"].reshape(3, 1) |
| 111 | path = tmp_path / "column_ancillary.npz" | 89 | path = tmp_path / "column_ancillary.npz" |
| 112 | np.savez_compressed(path, **record) | 90 | np.savez_compressed(path, **record) |
| 113 | with pytest.raises( | 91 | with pytest.raises( |
| 114 | ValueError, match=r"'intensity' must have shape \(N,\), got \(3, 1\)" | 92 | ValueError, match=r"'intensity' must have shape \(N,\), got \(3, 1\)" |
| 115 | ): | 93 | ): |
| 116 | load_points_npz(path) | 94 | segment_points_io.load_points_npz(path) |
| 117 | 95 | ||
| 118 | 96 | ||
| 119 | def test_save_points_npz_rejects_scalar_ancillary(tmp_path: Path) -> None: | 97 | def test_save_points_npz_rejects_scalar_ancillary(tmp_path: pathlib.Path) -> None: |
| 120 | record = _make_record(2, seed=6) | 98 | record = _make_record(2, seed=6) |
| 121 | record["green"] = np.array(7, dtype=np.uint16) | 99 | record["green"] = np.array(7, dtype=np.uint16) |
| 122 | with pytest.raises(ValueError, match=r"'green' must have shape \(N,\), got \(\)"): | 100 | with pytest.raises(ValueError, match=r"'green' must have shape \(N,\), got \(\)"): |
| 123 | save_points_npz(tmp_path / "out.npz", record) | 101 | segment_points_io.save_points_npz(tmp_path / "out.npz", record) |
| 124 | 102 | ||
| 125 | 103 | ||
| 126 | def test_save_points_npz_rejects_column_vector_ancillary(tmp_path: Path) -> None: | 104 | def test_save_points_npz_rejects_column_vector_ancillary(tmp_path: pathlib.Path) -> None: |
| 127 | record = _make_record(2, seed=7) | 105 | record = _make_record(2, seed=7) |
| 128 | record["scan_angle"] = record["scan_angle"].reshape(2, 1) | 106 | record["scan_angle"] = record["scan_angle"].reshape(2, 1) |
| 129 | with pytest.raises( | 107 | with pytest.raises( |
| 130 | ValueError, match=r"'scan_angle' must have shape \(N,\), got \(2, 1\)" | 108 | ValueError, match=r"'scan_angle' must have shape \(N,\), got \(2, 1\)" |
| 131 | ): | 109 | ): |
| 132 | save_points_npz(tmp_path / "out.npz", record) | 110 | segment_points_io.save_points_npz(tmp_path / "out.npz", record) |
| 133 | 111 | ||
| 134 | 112 | ||
| 135 | def test_save_points_npz_rejects_incomplete_record(tmp_path: Path) -> None: | 113 | def test_save_points_npz_rejects_incomplete_record(tmp_path: pathlib.Path) -> None: |
| 136 | with pytest.raises(ValueError, match="missing required key"): | 114 | with pytest.raises(ValueError, match="missing required key"): |
| 137 | save_points_npz(tmp_path / "out.npz", {"points": np.zeros((1, 3))}) | 115 | segment_points_io.save_points_npz(tmp_path / "out.npz", {"points": np.zeros((1, 3))}) |
| 138 | 116 | ||
| 139 | 117 | ||
| 140 | def test_number_of_returns_is_part_of_the_written_contract() -> None: | 118 | def test_number_of_returns_is_part_of_the_written_contract() -> None: |
| 141 | assert NUMBER_OF_RETURNS_KEY in POINT_RECORD_KEYS | 119 | assert segment_points_io.NUMBER_OF_RETURNS_KEY in segment_points_io.POINT_RECORD_KEYS |
| 142 | assert NUMBER_OF_RETURNS_KEY not in REQUIRED_POINT_RECORD_KEYS | 120 | assert ( |
| 143 | assert set(REQUIRED_POINT_RECORD_KEYS) < set(POINT_RECORD_KEYS) | 121 | segment_points_io.NUMBER_OF_RETURNS_KEY |
| 122 | not in segment_points_io.REQUIRED_POINT_RECORD_KEYS | ||
| 123 | ) | ||
| 124 | assert set(segment_points_io.REQUIRED_POINT_RECORD_KEYS) < set( | ||
| 125 | segment_points_io.POINT_RECORD_KEYS | ||
| 126 | ) | ||
| 144 | 127 | ||
| 145 | 128 | ||
| 146 | def test_save_points_npz_always_writes_number_of_returns(tmp_path: Path) -> None: | 129 | def test_save_points_npz_always_writes_number_of_returns(tmp_path: pathlib.Path) -> None: |
| 147 | record = _make_record(6, seed=60) | 130 | record = _make_record(6, seed=60) |
| 148 | path = save_points_npz(tmp_path / "with_returns_run3_points.npz", record) | 131 | path = segment_points_io.save_points_npz(tmp_path / "with_returns_run3_points.npz", record) |
| 149 | 132 | ||
| 150 | with np.load(path) as data: | 133 | with np.load(path) as data: |
| 151 | assert NUMBER_OF_RETURNS_KEY in data.files | 134 | assert segment_points_io.NUMBER_OF_RETURNS_KEY in data.files |
| 152 | stored = np.asarray(data[NUMBER_OF_RETURNS_KEY]) | 135 | stored = np.asarray(data[segment_points_io.NUMBER_OF_RETURNS_KEY]) |
| 153 | assert stored.dtype == np.uint8 | 136 | assert stored.dtype == np.uint8 |
| 154 | np.testing.assert_array_equal(stored, record[NUMBER_OF_RETURNS_KEY]) | 137 | np.testing.assert_array_equal(stored, record[segment_points_io.NUMBER_OF_RETURNS_KEY]) |
| 155 | 138 | ||
| 156 | 139 | ||
| 157 | def test_save_points_npz_casts_number_of_returns_to_uint8(tmp_path: Path) -> None: | 140 | def test_save_points_npz_casts_number_of_returns_to_uint8(tmp_path: pathlib.Path) -> None: |
| 158 | record = _make_record(4, seed=61) | 141 | record = _make_record(4, seed=61) |
| 159 | record[NUMBER_OF_RETURNS_KEY] = record[NUMBER_OF_RETURNS_KEY].astype(np.int64) | 142 | record[segment_points_io.NUMBER_OF_RETURNS_KEY] = record[ |
| 160 | path = save_points_npz(tmp_path / "cast_run3_points.npz", record) | 143 | segment_points_io.NUMBER_OF_RETURNS_KEY |
| 144 | ].astype(np.int64) | ||
| 145 | path = segment_points_io.save_points_npz(tmp_path / "cast_run3_points.npz", record) | ||
| 161 | 146 | ||
| 162 | loaded = load_points_npz(path) | 147 | loaded = segment_points_io.load_points_npz(path) |
| 163 | assert loaded[NUMBER_OF_RETURNS_KEY].dtype == np.uint8 | 148 | assert loaded[segment_points_io.NUMBER_OF_RETURNS_KEY].dtype == np.uint8 |
| 164 | 149 | ||
| 165 | 150 | ||
| 166 | def test_save_points_npz_rejects_out_of_range_number_of_returns(tmp_path: Path) -> None: | 151 | def test_save_points_npz_rejects_out_of_range_number_of_returns(tmp_path: pathlib.Path) -> None: |
| 167 | record = _make_record(3, seed=62) | 152 | record = _make_record(3, seed=62) |
| 168 | record[NUMBER_OF_RETURNS_KEY] = np.array([1, -1, 3], dtype=np.int16) | 153 | record[segment_points_io.NUMBER_OF_RETURNS_KEY] = np.array([1, -1, 3], dtype=np.int16) |
| 169 | with pytest.raises(ValueError, match="number_of_returns"): | 154 | with pytest.raises(ValueError, match="number_of_returns"): |
| 170 | save_points_npz(tmp_path / "out.npz", record) | 155 | segment_points_io.save_points_npz(tmp_path / "out.npz", record) |
| 171 | 156 | ||
| 172 | 157 | ||
| 173 | def test_save_points_npz_rejects_missing_number_of_returns(tmp_path: Path) -> None: | 158 | def test_save_points_npz_rejects_missing_number_of_returns(tmp_path: pathlib.Path) -> None: |
| 174 | with pytest.raises(ValueError, match="number_of_returns"): | 159 | with pytest.raises(ValueError, match="number_of_returns"): |
| 175 | save_points_npz(tmp_path / "out.npz", _make_legacy_record(3, seed=63)) | 160 | segment_points_io.save_points_npz(tmp_path / "out.npz", _make_legacy_record(3, seed=63)) |
| 176 | 161 | ||
| 177 | 162 | ||
| 178 | def test_load_points_npz_fills_zeros_for_legacy_records(tmp_path: Path) -> None: | 163 | def test_load_points_npz_fills_zeros_for_legacy_records(tmp_path: pathlib.Path) -> None: |
| 179 | """Datasets written before AI3D-382 lack the key; 0 means unknown, never 1.""" | 164 | """Datasets written before AI3D-382 lack the key; 0 means unknown, never 1.""" |
| 180 | legacy = _make_legacy_record(5, seed=64) | 165 | legacy = _make_legacy_record(5, seed=64) |
| 181 | path = tmp_path / "legacy_run3_points.npz" | 166 | path = tmp_path / "legacy_run3_points.npz" |
| 182 | np.savez_compressed(path, **legacy) | 167 | np.savez_compressed(path, **legacy) |
| 183 | 168 | ||
| 184 | loaded = load_points_npz(path) | 169 | loaded = segment_points_io.load_points_npz(path) |
| 185 | 170 | ||
| 186 | assert list(loaded.keys()) == list(POINT_RECORD_KEYS) | 171 | assert list(loaded.keys()) == list(segment_points_io.POINT_RECORD_KEYS) |
| 187 | returns = loaded[NUMBER_OF_RETURNS_KEY] | 172 | returns = loaded[segment_points_io.NUMBER_OF_RETURNS_KEY] |
| 188 | assert returns.shape == (5,) | 173 | assert returns.shape == (5,) |
| 189 | assert returns.dtype == np.uint8 | 174 | assert returns.dtype == np.uint8 |
| 190 | np.testing.assert_array_equal(returns, np.zeros(5, dtype=np.uint8)) | 175 | np.testing.assert_array_equal(returns, np.zeros(5, dtype=np.uint8)) |
| 191 | 176 | ||
| 192 | 177 | ||
| 193 | def test_load_points_npz_rejects_bad_number_of_returns_shape(tmp_path: Path) -> None: | 178 | def test_load_points_npz_rejects_bad_number_of_returns_shape(tmp_path: pathlib.Path) -> None: |
| 194 | record = _make_record(4, seed=65) | 179 | record = _make_record(4, seed=65) |
| 195 | record[NUMBER_OF_RETURNS_KEY] = record[NUMBER_OF_RETURNS_KEY][:2] | 180 | record[segment_points_io.NUMBER_OF_RETURNS_KEY] = record[ |
| 181 | segment_points_io.NUMBER_OF_RETURNS_KEY | ||
| 182 | ][:2] | ||
| 196 | path = tmp_path / "bad_returns.npz" | 183 | path = tmp_path / "bad_returns.npz" |
| 197 | np.savez_compressed(path, **record) | 184 | np.savez_compressed(path, **record) |
| 198 | with pytest.raises( | 185 | with pytest.raises( |
| 199 | ValueError, match=r"'number_of_returns' must have shape \(N,\), got \(2,\)" | 186 | ValueError, match=r"'number_of_returns' must have shape \(N,\), got \(2,\)" |
| 200 | ): | 187 | ): |
| 201 | load_points_npz(path) | 188 | segment_points_io.load_points_npz(path) |
| 202 | 189 | ||
| 203 | 190 | ||
| 204 | def test_save_points_npz_rejects_bool_number_of_returns(tmp_path: Path) -> None: | 191 | def test_save_points_npz_rejects_bool_number_of_returns(tmp_path: pathlib.Path) -> None: |
| 205 | """A bool mask is not a return count; casting it would fabricate 0/1 counts.""" | 192 | """A bool mask is not a return count; casting it would fabricate 0/1 counts.""" |
| 206 | record = _make_record(3, seed=68) | 193 | record = _make_record(3, seed=68) |
| 207 | record[NUMBER_OF_RETURNS_KEY] = np.array([True, False, True]) | 194 | record[segment_points_io.NUMBER_OF_RETURNS_KEY] = np.array([True, False, True]) |
| 208 | with pytest.raises(ValueError, match="number_of_returns.*bool"): | 195 | with pytest.raises(ValueError, match="number_of_returns.*bool"): |
| 209 | save_points_npz(tmp_path / "out.npz", record) | 196 | segment_points_io.save_points_npz(tmp_path / "out.npz", record) |
| 210 | 197 | ||
| 211 | 198 | ||
| 212 | def test_load_points_npz_rejects_bool_number_of_returns(tmp_path: Path) -> None: | 199 | def test_load_points_npz_rejects_bool_number_of_returns(tmp_path: pathlib.Path) -> None: |
| 213 | record = _make_record(3, seed=69) | 200 | record = _make_record(3, seed=69) |
| 214 | record[NUMBER_OF_RETURNS_KEY] = np.array([True, False, True]) | 201 | record[segment_points_io.NUMBER_OF_RETURNS_KEY] = np.array([True, False, True]) |
| 215 | path = tmp_path / "bool_returns_run3_points.npz" | 202 | path = tmp_path / "bool_returns_run3_points.npz" |
| 216 | np.savez_compressed(path, **record) | 203 | np.savez_compressed(path, **record) |
| 217 | with pytest.raises(ValueError, match="number_of_returns.*bool"): | 204 | with pytest.raises(ValueError, match="number_of_returns.*bool"): |
| 218 | load_points_npz(path) | 205 | segment_points_io.load_points_npz(path) |
| 219 | 206 | ||
| 220 | 207 | ||
| 221 | def test_load_points_npz_casts_stored_number_of_returns_to_uint8(tmp_path: Path) -> None: | 208 | def test_load_points_npz_casts_stored_number_of_returns_to_uint8(tmp_path: pathlib.Path) -> None: |
| 222 | """A hand-rolled producer's wider dtype must not leak into merges.""" | 209 | """A hand-rolled producer's wider dtype must not leak into merges.""" |
| 223 | record = _make_record(5, seed=70) | 210 | record = _make_record(5, seed=70) |
| 224 | stored = record[NUMBER_OF_RETURNS_KEY].astype(np.int32) | 211 | stored = record[segment_points_io.NUMBER_OF_RETURNS_KEY].astype(np.int32) |
| 225 | record[NUMBER_OF_RETURNS_KEY] = stored | 212 | record[segment_points_io.NUMBER_OF_RETURNS_KEY] = stored |
| 226 | path = tmp_path / "int32_returns_run3_points.npz" | 213 | path = tmp_path / "int32_returns_run3_points.npz" |
| 227 | np.savez_compressed(path, **record) | 214 | np.savez_compressed(path, **record) |
| 228 | 215 | ||
| 229 | loaded = load_points_npz(path) | 216 | loaded = segment_points_io.load_points_npz(path) |
| 230 | 217 | ||
| 231 | assert loaded[NUMBER_OF_RETURNS_KEY].dtype == np.uint8 | 218 | assert loaded[segment_points_io.NUMBER_OF_RETURNS_KEY].dtype == np.uint8 |
| 232 | np.testing.assert_array_equal(loaded[NUMBER_OF_RETURNS_KEY], stored) | 219 | np.testing.assert_array_equal(loaded[segment_points_io.NUMBER_OF_RETURNS_KEY], stored) |
| 233 | 220 | ||
| 234 | 221 | ||
| 235 | def test_load_points_npz_rejects_out_of_range_stored_number_of_returns( | 222 | def test_load_points_npz_rejects_out_of_range_stored_number_of_returns( |
| 236 | tmp_path: Path, | 223 | tmp_path: pathlib.Path, |
| 237 | ) -> None: | 224 | ) -> None: |
| 238 | record = _make_record(3, seed=71) | 225 | record = _make_record(3, seed=71) |
| 239 | record[NUMBER_OF_RETURNS_KEY] = np.array([1, 300, 3], dtype=np.int32) | 226 | record[segment_points_io.NUMBER_OF_RETURNS_KEY] = np.array([1, 300, 3], dtype=np.int32) |
| 240 | path = tmp_path / "out_of_range_returns_run3_points.npz" | 227 | path = tmp_path / "out_of_range_returns_run3_points.npz" |
| 241 | np.savez_compressed(path, **record) | 228 | np.savez_compressed(path, **record) |
| 242 | with pytest.raises(ValueError, match="number_of_returns.*fit in uint8"): | 229 | with pytest.raises(ValueError, match="number_of_returns.*fit in uint8"): |
| 243 | load_points_npz(path) | 230 | segment_points_io.load_points_npz(path) |
| 244 | 231 | ||
| 245 | 232 | ||
| 246 | def test_load_points_npz_rejects_float_stored_number_of_returns(tmp_path: Path) -> None: | 233 | def test_load_points_npz_rejects_float_stored_number_of_returns(tmp_path: pathlib.Path) -> None: |
| 247 | record = _make_record(3, seed=72) | 234 | record = _make_record(3, seed=72) |
| 248 | record[NUMBER_OF_RETURNS_KEY] = np.array([1.0, 2.0, 3.0], dtype=np.float32) | 235 | record[segment_points_io.NUMBER_OF_RETURNS_KEY] = np.array([1.0, 2.0, 3.0], dtype=np.float32) |
| 249 | path = tmp_path / "float_returns_run3_points.npz" | 236 | path = tmp_path / "float_returns_run3_points.npz" |
| 250 | np.savez_compressed(path, **record) | 237 | np.savez_compressed(path, **record) |
| 251 | with pytest.raises(ValueError, match="number_of_returns.*integer array"): | 238 | with pytest.raises(ValueError, match="number_of_returns.*integer array"): |
| 252 | load_points_npz(path) | 239 | segment_points_io.load_points_npz(path) |
| 253 | 240 | ||
| 254 | 241 | ||
| 255 | def test_concat_points_npz_merges_mixed_stored_return_dtypes(tmp_path: Path) -> None: | 242 | def test_concat_points_npz_merges_mixed_stored_return_dtypes(tmp_path: pathlib.Path) -> None: |
| 256 | """int32-stored and uint8-stored records concatenate after the load coercion.""" | 243 | """int32-stored and uint8-stored records concatenate after the load coercion.""" |
| 257 | wide = _make_record(3, seed=73) | 244 | wide = _make_record(3, seed=73) |
| 258 | wide[NUMBER_OF_RETURNS_KEY] = wide[NUMBER_OF_RETURNS_KEY].astype(np.int32) | 245 | wide[segment_points_io.NUMBER_OF_RETURNS_KEY] = wide[ |
| 246 | segment_points_io.NUMBER_OF_RETURNS_KEY | ||
| 247 | ].astype(np.int32) | ||
| 259 | wide_path = tmp_path / "alpha_run3_points.npz" | 248 | wide_path = tmp_path / "alpha_run3_points.npz" |
| 260 | np.savez_compressed(wide_path, **wide) | 249 | np.savez_compressed(wide_path, **wide) |
| 261 | narrow_path = save_points_npz(tmp_path / "beta_run3_points.npz", _make_record(2, seed=74)) | 250 | narrow_path = segment_points_io.save_points_npz( |
| 251 | tmp_path / "beta_run3_points.npz", _make_record(2, seed=74) | ||
| 252 | ) | ||
| 262 | 253 | ||
| 263 | merged, _ = concat_points_npz([wide_path, narrow_path]) | 254 | merged, _ = segment_points_io.concat_points_npz([wide_path, narrow_path]) |
| 264 | 255 | ||
| 265 | assert merged[NUMBER_OF_RETURNS_KEY].dtype == np.uint8 | 256 | assert merged[segment_points_io.NUMBER_OF_RETURNS_KEY].dtype == np.uint8 |
| 266 | assert merged[NUMBER_OF_RETURNS_KEY].shape == (5,) | 257 | assert merged[segment_points_io.NUMBER_OF_RETURNS_KEY].shape == (5,) |
| 267 | 258 | ||
| 268 | 259 | ||
| 269 | def test_load_segment_points_merges_legacy_and_new_records(tmp_path: Path) -> None: | 260 | def test_load_segment_points_merges_legacy_and_new_records(tmp_path: pathlib.Path) -> None: |
| 270 | legacy = _make_legacy_record(3, seed=66) | 261 | legacy = _make_legacy_record(3, seed=66) |
| 271 | modern = _make_record(4, seed=67) | 262 | modern = _make_record(4, seed=67) |
| 272 | legacy_path = tmp_path / "alpha_run3_points.npz" | 263 | legacy_path = tmp_path / "alpha_run3_points.npz" |
| 273 | np.savez_compressed(legacy_path, **legacy) | 264 | np.savez_compressed(legacy_path, **legacy) |
| 274 | modern_path = tmp_path / "beta_run3_points.npz" | 265 | modern_path = tmp_path / "beta_run3_points.npz" |
| 275 | save_points_npz(modern_path, modern) | 266 | segment_points_io.save_points_npz(modern_path, modern) |
| 276 | 267 | ||
| 277 | merged, _, _ = load_segment_points([legacy_path, modern_path]) | 268 | merged, _, _ = segment_points_io.load_segment_points([legacy_path, modern_path]) |
| 278 | 269 | ||
| 279 | np.testing.assert_array_equal( | 270 | np.testing.assert_array_equal( |
| 280 | merged[NUMBER_OF_RETURNS_KEY], | 271 | merged[segment_points_io.NUMBER_OF_RETURNS_KEY], |
| 281 | np.concatenate( | 272 | np.concatenate( |
| 282 | [np.zeros(3, dtype=np.uint8), modern[NUMBER_OF_RETURNS_KEY]], axis=0 | 273 | [np.zeros(3, dtype=np.uint8), modern[segment_points_io.NUMBER_OF_RETURNS_KEY]], axis=0 |
| 283 | ), | 274 | ), |
| 284 | ) | 275 | ) |
| 285 | 276 | ||
| 286 | 277 | ||
| 287 | def test_concat_points_npz_carries_number_of_returns(tmp_path: Path) -> None: | 278 | def test_concat_points_npz_carries_number_of_returns(tmp_path: pathlib.Path) -> None: |
| 288 | first = _make_record(3, seed=68) | 279 | first = _make_record(3, seed=68) |
| 289 | second = _make_legacy_record(2, seed=69) | 280 | second = _make_legacy_record(2, seed=69) |
| 290 | save_points_npz(tmp_path / "a_run3_points.npz", first) | 281 | segment_points_io.save_points_npz(tmp_path / "a_run3_points.npz", first) |
| 291 | np.savez_compressed(tmp_path / "b_run3_points.npz", **second) | 282 | np.savez_compressed(tmp_path / "b_run3_points.npz", **second) |
| 292 | 283 | ||
| 293 | merged, spans = concat_points_npz( | 284 | merged, spans = segment_points_io.concat_points_npz( |
| 294 | [tmp_path / "a_run3_points.npz", tmp_path / "b_run3_points.npz"] | 285 | [tmp_path / "a_run3_points.npz", tmp_path / "b_run3_points.npz"] |
| 295 | ) | 286 | ) |
| 296 | 287 | ||
| 297 | assert [span.count for span in spans] == [3, 2] | 288 | assert [span.count for span in spans] == [3, 2] |
| 298 | assert merged[NUMBER_OF_RETURNS_KEY].dtype == np.uint8 | 289 | assert merged[segment_points_io.NUMBER_OF_RETURNS_KEY].dtype == np.uint8 |
| 299 | np.testing.assert_array_equal( | 290 | np.testing.assert_array_equal( |
| 300 | merged[NUMBER_OF_RETURNS_KEY], | 291 | merged[segment_points_io.NUMBER_OF_RETURNS_KEY], |
| 301 | np.concatenate( | 292 | np.concatenate( |
| 302 | [first[NUMBER_OF_RETURNS_KEY], np.zeros(2, dtype=np.uint8)], axis=0 | 293 | [first[segment_points_io.NUMBER_OF_RETURNS_KEY], np.zeros(2, dtype=np.uint8)], axis=0 |
| 303 | ), | 294 | ), |
| 304 | ) | 295 | ) |
| 305 | 296 | ||
| 306 | 297 | ||
| 307 | def test_load_run3_segment_carries_number_of_returns(tmp_path: Path) -> None: | 298 | def test_load_run3_segment_carries_number_of_returns(tmp_path: pathlib.Path) -> None: |
| 308 | seg_dir = tmp_path / "segment_011" | 299 | seg_dir = tmp_path / "segment_011" |
| 309 | seg_dir.mkdir() | 300 | seg_dir.mkdir() |
| 310 | first = _make_record(2, seed=70) | 301 | first = _make_record(2, seed=70) |
| 311 | second = _make_record(3, seed=71) | 302 | second = _make_record(3, seed=71) |
| 312 | save_points_npz(seg_dir / "a_run3_points.npz", first) | 303 | segment_points_io.save_points_npz(seg_dir / "a_run3_points.npz", first) |
| 313 | save_points_npz(seg_dir / "b_run3_points.npz", second) | 304 | segment_points_io.save_points_npz(seg_dir / "b_run3_points.npz", second) |
| 314 | 305 | ||
| 315 | merged, _ = load_run3_segment(seg_dir) | 306 | merged, _ = segment_points_io.load_run3_segment(seg_dir) |
| 316 | 307 | ||
| 317 | np.testing.assert_array_equal( | 308 | np.testing.assert_array_equal( |
| 318 | merged[NUMBER_OF_RETURNS_KEY], | 309 | merged[segment_points_io.NUMBER_OF_RETURNS_KEY], |
| 319 | np.concatenate( | 310 | np.concatenate( |
| 320 | [first[NUMBER_OF_RETURNS_KEY], second[NUMBER_OF_RETURNS_KEY]], axis=0 | 311 | [ |
| 312 | first[segment_points_io.NUMBER_OF_RETURNS_KEY], | ||
| 313 | second[segment_points_io.NUMBER_OF_RETURNS_KEY], | ||
| 314 | ], | ||
| 315 | axis=0, | ||
| 321 | ), | 316 | ), |
| 322 | ) | 317 | ) |
| 323 | 318 | ||
| 324 | 319 | ||
| 325 | def test_load_segment_points_merges_and_tracks_file_ids(tmp_path: Path) -> None: | 320 | def test_load_segment_points_merges_and_tracks_file_ids(tmp_path: pathlib.Path) -> None: |
| 326 | records = [_make_record(3, seed=10), _make_record(5, seed=11)] | 321 | records = [_make_record(3, seed=10), _make_record(5, seed=11)] |
| 327 | paths = [ | 322 | paths = [ |
| 328 | tmp_path / "alpha_run3_points.npz", | 323 | tmp_path / "alpha_run3_points.npz", |
| 329 | tmp_path / "beta_run3_points.npz", | 324 | tmp_path / "beta_run3_points.npz", |
| 330 | ] | 325 | ] |
| 331 | for path, record in zip(paths, records, strict=True): | 326 | for path, record in zip(paths, records, strict=True): |
| 332 | save_points_npz(path, record) | 327 | segment_points_io.save_points_npz(path, record) |
| 333 | 328 | ||
| 334 | merged, point_file_ids, file_stems = load_segment_points(paths) | 329 | merged, point_file_ids, file_stems = segment_points_io.load_segment_points(paths) |
| 335 | 330 | ||
| 336 | assert file_stems == ["alpha_run3_points", "beta_run3_points"] | 331 | assert file_stems == ["alpha_run3_points", "beta_run3_points"] |
| 337 | assert merged["points"].shape == (8, 3) | 332 | assert merged["points"].shape == (8, 3) |
| 338 | assert point_file_ids.dtype == np.int32 | 333 | assert point_file_ids.dtype == np.int32 |
| 351 | 346 | ||
| 352 | 347 | ||
| 353 | def test_load_segment_points_empty_raises() -> None: | 348 | def test_load_segment_points_empty_raises() -> None: |
| 354 | with pytest.raises(FileNotFoundError, match="No \\*_points.npz"): | 349 | with pytest.raises(FileNotFoundError, match="No \\*_points.npz"): |
| 355 | load_segment_points([]) | 350 | segment_points_io.load_segment_points([]) |
| 356 | 351 | ||
| 357 | 352 | ||
| 358 | def test_load_geoshift_parses_xyz_json(tmp_path: Path) -> None: | 353 | def test_load_geoshift_parses_xyz_json(tmp_path: pathlib.Path) -> None: |
| 359 | path = tmp_path / "run3_geoshift.json" | 354 | path = tmp_path / "run3_geoshift.json" |
| 360 | path.write_text(json.dumps({"x": 725883.5, "y": 5422097.8, "z": 390.8}), encoding="utf-8") | 355 | path.write_text(json.dumps({"x": 725883.5, "y": 5422097.8, "z": 390.8}), encoding="utf-8") |
| 361 | geoshift = load_geoshift(path) | 356 | geoshift = segment_points_io.load_geoshift(path) |
| 362 | assert geoshift.shape == (3,) | 357 | assert geoshift.shape == (3,) |
| 363 | assert geoshift.dtype == np.float64 | 358 | assert geoshift.dtype == np.float64 |
| 364 | np.testing.assert_allclose(geoshift, [725883.5, 5422097.8, 390.8]) | 359 | np.testing.assert_allclose(geoshift, [725883.5, 5422097.8, 390.8]) |
| 365 | 360 | ||
| 366 | 361 | ||
| 367 | def test_load_geoshift_rejects_missing_keys(tmp_path: Path) -> None: | 362 | def test_load_geoshift_rejects_missing_keys(tmp_path: pathlib.Path) -> None: |
| 368 | path = tmp_path / "run3_geoshift.json" | 363 | path = tmp_path / "run3_geoshift.json" |
| 369 | path.write_text(json.dumps({"x": 1.0, "y": 2.0}), encoding="utf-8") | 364 | path.write_text(json.dumps({"x": 1.0, "y": 2.0}), encoding="utf-8") |
| 370 | with pytest.raises(ValueError, match="missing geoshift key"): | 365 | with pytest.raises(ValueError, match="missing geoshift key"): |
| 371 | load_geoshift(path) | 366 | segment_points_io.load_geoshift(path) |
| 372 | 367 | ||
| 373 | 368 | ||
| 374 | def test_parse_segment_key_accepts_int_and_prefixed_forms() -> None: | 369 | def test_parse_segment_key_accepts_int_and_prefixed_forms() -> None: |
| 375 | assert parse_segment_key(32) == 32 | 370 | assert segment_points_io.parse_segment_key(32) == 32 |
| 376 | assert parse_segment_key("32") == 32 | 371 | assert segment_points_io.parse_segment_key("32") == 32 |
| 377 | assert parse_segment_key("segment_33") == 33 | 372 | assert segment_points_io.parse_segment_key("segment_33") == 33 |
| 378 | assert parse_segment_key("segment_066") == 66 | 373 | assert segment_points_io.parse_segment_key("segment_066") == 66 |
| 379 | 374 | ||
| 380 | 375 | ||
| 381 | def test_parse_segment_key_rejects_invalid() -> None: | 376 | def test_parse_segment_key_rejects_invalid() -> None: |
| 382 | with pytest.raises(ValueError, match="Segment key must not be empty"): | 377 | with pytest.raises(ValueError, match="Segment key must not be empty"): |
| 383 | parse_segment_key("segment_") | 378 | segment_points_io.parse_segment_key("segment_") |
| 384 | with pytest.raises(ValueError, match="integer or segment_<idx>"): | 379 | with pytest.raises(ValueError, match="integer or segment_<idx>"): |
| 385 | parse_segment_key("lane_a") | 380 | segment_points_io.parse_segment_key("lane_a") |
| 386 | with pytest.raises(ValueError, match="integer or segment_<idx>"): | 381 | with pytest.raises(ValueError, match="integer or segment_<idx>"): |
| 387 | parse_segment_key(True) | 382 | segment_points_io.parse_segment_key(True) |
| 388 | 383 | ||
| 389 | 384 | ||
| 390 | def test_normalize_segment_file_blacklist_accepts_multiple_key_formats() -> None: | 385 | def test_normalize_segment_file_blacklist_accepts_multiple_key_formats() -> None: |
| 391 | blacklist = normalize_segment_file_blacklist( | 386 | blacklist = segment_points_io.normalize_segment_file_blacklist( |
| 392 | { | 387 | { |
| 393 | "32": ["scan_a_run3_points.npz", "nested/scan_b_run3_points.npz"], | 388 | "32": ["scan_a_run3_points.npz", "nested/scan_b_run3_points.npz"], |
| 394 | "segment_33": "scan_c_run3_points.npz", | 389 | "segment_33": "scan_c_run3_points.npz", |
| 395 | 33: ["scan_d_run3_points.npz"], | 390 | 33: ["scan_d_run3_points.npz"], |
| 402 | 397 | ||
| 403 | 398 | ||
| 404 | def test_normalize_segment_file_blacklist_rejects_invalid_values() -> None: | 399 | def test_normalize_segment_file_blacklist_rejects_invalid_values() -> None: |
| 405 | with pytest.raises(ValueError, match="mapping of segment to files"): | 400 | with pytest.raises(ValueError, match="mapping of segment to files"): |
| 406 | normalize_segment_file_blacklist(["32:scan_a_run3_points.npz"]) # type: ignore[arg-type] | 401 | segment_points_io.normalize_segment_file_blacklist(["32:scan_a_run3_points.npz"]) # type: ignore[arg-type] |
| 407 | with pytest.raises(ValueError, match="at least one non-empty"): | 402 | with pytest.raises(ValueError, match="at least one non-empty"): |
| 408 | normalize_segment_file_blacklist({"32": [" ", ""]}) | 403 | segment_points_io.normalize_segment_file_blacklist({"32": [" ", ""]}) |
| 409 | assert normalize_segment_file_blacklist(None) == {} | 404 | assert segment_points_io.normalize_segment_file_blacklist(None) == {} |
| 410 | 405 | ||
| 411 | 406 | ||
| 412 | def test_filter_segment_files_respects_segment_specific_exact_and_glob_rules( | 407 | def test_filter_segment_files_respects_segment_specific_exact_and_glob_rules( |
| 413 | caplog: pytest.LogCaptureFixture, | 408 | caplog: pytest.LogCaptureFixture, |
| 414 | ) -> None: | 409 | ) -> None: |
| 415 | segment_32_files = [ | 410 | segment_32_files = [ |
| 416 | Path("/tmp/lane_points/segment_32/scan_a_run3_points.npz"), | 411 | pathlib.Path("/tmp/lane_points/segment_32/scan_a_run3_points.npz"), |
| 417 | Path("/tmp/lane_points/segment_32/nested/scan_b_run3_points.npz"), | 412 | pathlib.Path("/tmp/lane_points/segment_32/nested/scan_b_run3_points.npz"), |
| 418 | Path("/tmp/lane_points/segment_32/scan_c_run3_points.npz"), | 413 | pathlib.Path("/tmp/lane_points/segment_32/scan_c_run3_points.npz"), |
| 419 | ] | 414 | ] |
| 420 | blacklist = normalize_segment_file_blacklist( | 415 | blacklist = segment_points_io.normalize_segment_file_blacklist( |
| 421 | { | 416 | { |
| 422 | "32": [ | 417 | "32": [ |
| 423 | "scan_a_run3_points.npz", | 418 | "scan_a_run3_points.npz", |
| 424 | "*/nested/scan_b_run3_points.npz", | 419 | "*/nested/scan_b_run3_points.npz", |
| 427 | } | 422 | } |
| 428 | ) | 423 | ) |
| 429 | 424 | ||
| 430 | with caplog.at_level(logging.INFO, logger="iolabs.common.segment_points_io"): | 425 | with caplog.at_level(logging.INFO, logger="iolabs.common.segment_points_io"): |
| 431 | kept_32, blacklisted_32 = filter_segment_files( | 426 | kept_32, blacklisted_32 = segment_points_io.filter_segment_files( |
| 432 | segment_32_files, | 427 | segment_32_files, |
| 433 | segment_index=32, | 428 | segment_index=32, |
| 434 | blacklist=blacklist, | 429 | blacklist=blacklist, |
| 435 | ) | 430 | ) |
| 436 | kept_33, blacklisted_33 = filter_segment_files( | 431 | kept_33, blacklisted_33 = segment_points_io.filter_segment_files( |
| 437 | segment_32_files, | 432 | segment_32_files, |
| 438 | segment_index=33, | 433 | segment_index=33, |
| 439 | blacklist=blacklist, | 434 | blacklist=blacklist, |
| 440 | ) | 435 | ) |
| 451 | assert [path.name for path in blacklisted_33] == ["scan_c_run3_points.npz"] | 446 | assert [path.name for path in blacklisted_33] == ["scan_c_run3_points.npz"] |
| 452 | assert any("excluding 2 input NPZ files" in record.message for record in caplog.records) | 447 | assert any("excluding 2 input NPZ files" in record.message for record in caplog.records) |
| 453 | 448 | ||
| 454 | 449 | ||
| 455 | def _write_geoshift(path: Path, xyz: tuple[float, float, float]) -> Path: | 450 | def _write_geoshift(path: pathlib.Path, xyz: tuple[float, float, float]) -> pathlib.Path: |
| 456 | path.parent.mkdir(parents=True, exist_ok=True) | 451 | path.parent.mkdir(parents=True, exist_ok=True) |
| 457 | path.write_text( | 452 | path.write_text( |
| 458 | json.dumps({"x": xyz[0], "y": xyz[1], "z": xyz[2]}), encoding="utf-8" | 453 | json.dumps({"x": xyz[0], "y": xyz[1], "z": xyz[2]}), encoding="utf-8" |
| 459 | ) | 454 | ) |
| 460 | return path | 455 | return path |
| 461 | 456 | ||
| 462 | 457 | ||
| 463 | def _save_stored_npz(path: Path, record: dict[str, np.ndarray]) -> Path: | 458 | def _save_stored_npz(path: pathlib.Path, record: dict[str, np.ndarray]) -> pathlib.Path: |
| 464 | """Write an uncompressed (ZIP_STORED) npz, the chunk-streamable layout.""" | 459 | """Write an uncompressed (ZIP_STORED) npz, the chunk-streamable layout.""" |
| 465 | np.savez(path, **record) | 460 | np.savez(path, **record) |
| 466 | return path | 461 | return path |
| 467 | 462 | ||
| 468 | 463 | ||
| 469 | def test_geoshift_candidate_paths_cover_the_three_conventions(tmp_path: Path) -> None: | 464 | def test_geoshift_candidate_paths_cover_the_three_conventions(tmp_path: pathlib.Path) -> None: |
| 470 | candidates = geoshift_candidate_paths(tmp_path / "dataset" / "lane_points") | 465 | candidates = segment_points_io.geoshift_candidate_paths(tmp_path / "dataset" / "lane_points") |
| 471 | assert candidates == [ | 466 | assert candidates == [ |
| 472 | tmp_path / "dataset" / "lane_points" / GEOSHIFT_NAME, | 467 | tmp_path / "dataset" / "lane_points" / segment_points_io.GEOSHIFT_NAME, |
| 473 | tmp_path / "dataset" / "lane_points" / "lane_points" / GEOSHIFT_NAME, | 468 | tmp_path / "dataset" / "lane_points" / "lane_points" / segment_points_io.GEOSHIFT_NAME, |
| 474 | tmp_path / "dataset" / GEOSHIFT_NAME, | 469 | tmp_path / "dataset" / segment_points_io.GEOSHIFT_NAME, |
| 475 | ] | 470 | ] |
| 476 | 471 | ||
| 477 | 472 | ||
| 478 | def test_find_geoshift_lane_points_dir_convention(tmp_path: Path) -> None: | 473 | def test_find_geoshift_lane_points_dir_convention(tmp_path: pathlib.Path) -> None: |
| 479 | lane_points = tmp_path / "lane_points" | 474 | lane_points = tmp_path / "lane_points" |
| 480 | _write_geoshift(lane_points / GEOSHIFT_NAME, (1.0, 2.0, 3.0)) | 475 | _write_geoshift(lane_points / segment_points_io.GEOSHIFT_NAME, (1.0, 2.0, 3.0)) |
| 481 | np.testing.assert_allclose(find_geoshift(lane_points), [1.0, 2.0, 3.0]) | 476 | np.testing.assert_allclose(segment_points_io.find_geoshift(lane_points), [1.0, 2.0, 3.0]) |
| 482 | 477 | ||
| 483 | 478 | ||
| 484 | def test_find_geoshift_dataset_root_convention(tmp_path: Path) -> None: | 479 | def test_find_geoshift_dataset_root_convention(tmp_path: pathlib.Path) -> None: |
| 485 | _write_geoshift(tmp_path / "lane_points" / GEOSHIFT_NAME, (4.0, 5.0, 6.0)) | 480 | _write_geoshift(tmp_path / "lane_points" / segment_points_io.GEOSHIFT_NAME, (4.0, 5.0, 6.0)) |
| 486 | np.testing.assert_allclose(find_geoshift(tmp_path), [4.0, 5.0, 6.0]) | 481 | np.testing.assert_allclose(segment_points_io.find_geoshift(tmp_path), [4.0, 5.0, 6.0]) |
| 487 | 482 | ||
| 488 | 483 | ||
| 489 | def test_find_geoshift_segment_dir_parent_convention(tmp_path: Path) -> None: | 484 | def test_find_geoshift_segment_dir_parent_convention(tmp_path: pathlib.Path) -> None: |
| 490 | lane_points = tmp_path / "lane_points" | 485 | lane_points = tmp_path / "lane_points" |
| 491 | seg_dir = lane_points / "segment_032" | 486 | seg_dir = lane_points / "segment_032" |
| 492 | seg_dir.mkdir(parents=True) | 487 | seg_dir.mkdir(parents=True) |
| 493 | _write_geoshift(lane_points / GEOSHIFT_NAME, (7.0, 8.0, 9.0)) | 488 | _write_geoshift(lane_points / segment_points_io.GEOSHIFT_NAME, (7.0, 8.0, 9.0)) |
| 494 | np.testing.assert_allclose(find_geoshift(seg_dir), [7.0, 8.0, 9.0]) | 489 | np.testing.assert_allclose(segment_points_io.find_geoshift(seg_dir), [7.0, 8.0, 9.0]) |
| 495 | 490 | ||
| 496 | 491 | ||
| 497 | def test_find_geoshift_prefers_directory_over_parent(tmp_path: Path) -> None: | 492 | def test_find_geoshift_prefers_directory_over_parent(tmp_path: pathlib.Path) -> None: |
| 498 | lane_points = tmp_path / "lane_points" | 493 | lane_points = tmp_path / "lane_points" |
| 499 | seg_dir = lane_points / "segment_032" | 494 | seg_dir = lane_points / "segment_032" |
| 500 | seg_dir.mkdir(parents=True) | 495 | seg_dir.mkdir(parents=True) |
| 501 | _write_geoshift(seg_dir / GEOSHIFT_NAME, (1.0, 1.0, 1.0)) | 496 | _write_geoshift(seg_dir / segment_points_io.GEOSHIFT_NAME, (1.0, 1.0, 1.0)) |
| 502 | _write_geoshift(lane_points / GEOSHIFT_NAME, (2.0, 2.0, 2.0)) | 497 | _write_geoshift(lane_points / segment_points_io.GEOSHIFT_NAME, (2.0, 2.0, 2.0)) |
| 503 | np.testing.assert_allclose(find_geoshift(seg_dir), [1.0, 1.0, 1.0]) | 498 | np.testing.assert_allclose(segment_points_io.find_geoshift(seg_dir), [1.0, 1.0, 1.0]) |
| 504 | 499 | ||
| 505 | 500 | ||
| 506 | def test_find_geoshift_or_none_returns_none_when_absent(tmp_path: Path) -> None: | 501 | def test_find_geoshift_or_none_returns_none_when_absent(tmp_path: pathlib.Path) -> None: |
| 507 | seg_dir = tmp_path / "lane_points" / "segment_000" | 502 | seg_dir = tmp_path / "lane_points" / "segment_000" |
| 508 | seg_dir.mkdir(parents=True) | 503 | seg_dir.mkdir(parents=True) |
| 509 | assert find_geoshift_or_none(seg_dir) is None | 504 | assert segment_points_io.find_geoshift_or_none(seg_dir) is None |
| 510 | 505 | ||
| 511 | 506 | ||
| 512 | def test_find_geoshift_raises_and_lists_searched_paths(tmp_path: Path) -> None: | 507 | def test_find_geoshift_raises_and_lists_searched_paths(tmp_path: pathlib.Path) -> None: |
| 513 | with pytest.raises(FileNotFoundError, match="searched:"): | 508 | with pytest.raises(FileNotFoundError, match="searched:"): |
| 514 | find_geoshift(tmp_path) | 509 | segment_points_io.find_geoshift(tmp_path) |
| 515 | 510 | ||
| 516 | 511 | ||
| 517 | def test_find_geoshift_returns_shift_unsigned(tmp_path: Path) -> None: | 512 | def test_find_geoshift_returns_shift_unsigned(tmp_path: pathlib.Path) -> None: |
| 518 | """The recorded shift is returned verbatim; no sign is baked in.""" | 513 | """The recorded shift is returned verbatim; no sign is baked in.""" |
| 519 | _write_geoshift(tmp_path / GEOSHIFT_NAME, (725883.5, -5422097.8, 390.8)) | 514 | _write_geoshift(tmp_path / segment_points_io.GEOSHIFT_NAME, (725883.5, -5422097.8, 390.8)) |
| 520 | np.testing.assert_allclose( | 515 | np.testing.assert_allclose( |
| 521 | find_geoshift(tmp_path), [725883.5, -5422097.8, 390.8] | 516 | segment_points_io.find_geoshift(tmp_path), [725883.5, -5422097.8, 390.8] |
| 522 | ) | 517 | ) |
| 523 | 518 | ||
| 524 | 519 | ||
| 525 | def test_geoshift_from_mapping_bare_and_nested() -> None: | 520 | def test_geoshift_from_mapping_bare_and_nested() -> None: |
| 526 | bare = geoshift_from_mapping({"x": 1.0, "y": 2.0, "z": 3.0}) | 521 | bare = segment_points_io.geoshift_from_mapping({"x": 1.0, "y": 2.0, "z": 3.0}) |
| 527 | nested = geoshift_from_mapping( | 522 | nested = segment_points_io.geoshift_from_mapping( |
| 528 | {"pixels_per_meter": 10, "geoshift": {"x": 1.0, "y": 2.0, "z": 3.0}} | 523 | {"pixels_per_meter": 10, "geoshift": {"x": 1.0, "y": 2.0, "z": 3.0}} |
| 529 | ) | 524 | ) |
| 530 | assert bare.dtype == np.float64 | 525 | assert bare.dtype == np.float64 |
| 531 | np.testing.assert_allclose(bare, [1.0, 2.0, 3.0]) | 526 | np.testing.assert_allclose(bare, [1.0, 2.0, 3.0]) |
| 533 | 528 | ||
| 534 | 529 | ||
| 535 | def test_geoshift_from_mapping_rejects_bad_input() -> None: | 530 | def test_geoshift_from_mapping_rejects_bad_input() -> None: |
| 536 | with pytest.raises(ValueError, match="missing geoshift key"): | 531 | with pytest.raises(ValueError, match="missing geoshift key"): |
| 537 | geoshift_from_mapping({"x": 1.0, "y": 2.0}) | 532 | segment_points_io.geoshift_from_mapping({"x": 1.0, "y": 2.0}) |
| 538 | with pytest.raises(ValueError, match="must be an object"): | 533 | with pytest.raises(ValueError, match="must be an object"): |
| 539 | geoshift_from_mapping([1.0, 2.0, 3.0]) # type: ignore[arg-type] | 534 | segment_points_io.geoshift_from_mapping([1.0, 2.0, 3.0]) # type: ignore[arg-type] |
| 540 | with pytest.raises(ValueError, match="must be numbers"): | 535 | with pytest.raises(ValueError, match="must be numbers"): |
| 541 | geoshift_from_mapping({"x": 1.0, "y": "north", "z": 3.0}) | 536 | segment_points_io.geoshift_from_mapping({"x": 1.0, "y": "north", "z": 3.0}) |
| 542 | 537 | ||
| 543 | 538 | ||
| 544 | def test_read_points_header_reports_stored_and_compressed(tmp_path: Path) -> None: | 539 | def test_read_points_header_reports_stored_and_compressed(tmp_path: pathlib.Path) -> None: |
| 545 | record = _make_record(5, seed=20) | 540 | record = _make_record(5, seed=20) |
| 546 | stored = _save_stored_npz(tmp_path / "stored.npz", record) | 541 | stored = _save_stored_npz(tmp_path / "stored.npz", record) |
| 547 | compressed = tmp_path / "compressed.npz" | 542 | compressed = tmp_path / "compressed.npz" |
| 548 | np.savez_compressed(compressed, **record) | 543 | np.savez_compressed(compressed, **record) |
| 549 | 544 | ||
| 550 | assert read_points_header(stored) == (5, True) | 545 | assert segment_points_io.read_points_header(stored) == (5, True) |
| 551 | assert read_points_header(compressed) == (5, False) | 546 | assert segment_points_io.read_points_header(compressed) == (5, False) |
| 552 | 547 | ||
| 553 | 548 | ||
| 554 | def test_read_points_header_unreadable_file(tmp_path: Path) -> None: | 549 | def test_read_points_header_unreadable_file(tmp_path: pathlib.Path) -> None: |
| 555 | broken = tmp_path / "broken.npz" | 550 | broken = tmp_path / "broken.npz" |
| 556 | broken.write_bytes(b"not a zip archive") | 551 | broken.write_bytes(b"not a zip archive") |
| 557 | assert read_points_header(broken) == (-1, False) | 552 | assert segment_points_io.read_points_header(broken) == (-1, False) |
| 558 | 553 | ||
| 559 | 554 | ||
| 560 | def test_iter_points_chunks_streams_tail_chunk_and_is_writeable(tmp_path: Path) -> None: | 555 | def test_iter_points_chunks_streams_tail_chunk_and_is_writeable(tmp_path: pathlib.Path) -> None: |
| 561 | record = _make_record(7, seed=21) | 556 | record = _make_record(7, seed=21) |
| 562 | path = _save_stored_npz(tmp_path / "big_run3_points.npz", record) | 557 | path = _save_stored_npz(tmp_path / "big_run3_points.npz", record) |
| 563 | 558 | ||
| 564 | chunks = list(iter_points_chunks(path, 3)) | 559 | chunks = list(segment_points_io.iter_points_chunks(path, 3)) |
| 565 | 560 | ||
| 566 | assert [len(chunk) for chunk in chunks] == [3, 3, 1] | 561 | assert [len(chunk) for chunk in chunks] == [3, 3, 1] |
| 567 | for chunk in chunks: | 562 | for chunk in chunks: |
| 568 | assert chunk.flags.writeable | 563 | assert chunk.flags.writeable |
| 569 | assert chunk.flags.c_contiguous | 564 | assert chunk.flags.c_contiguous |
| 570 | chunk[:] = 0.0 # must not raise: chunks are owned copies | 565 | chunk[:] = 0.0 # must not raise: chunks are owned copies |
| 571 | np.testing.assert_array_equal( | 566 | np.testing.assert_array_equal( |
| 572 | np.concatenate(list(iter_points_chunks(path, 3)), axis=0), record["points"] | 567 | np.concatenate(list(segment_points_io.iter_points_chunks(path, 3)), axis=0), |
| 568 | record["points"], | ||
| 573 | ) | 569 | ) |
| 574 | 570 | ||
| 575 | 571 | ||
| 576 | def test_iter_points_chunks_yields_whole_record_when_not_oversized(tmp_path: Path) -> None: | 572 | def test_iter_points_chunks_yields_whole_record_when_not_oversized(tmp_path: pathlib.Path) -> None: |
| 577 | record = _make_record(4, seed=22) | 573 | record = _make_record(4, seed=22) |
| 578 | path = _save_stored_npz(tmp_path / "small_run3_points.npz", record) | 574 | path = _save_stored_npz(tmp_path / "small_run3_points.npz", record) |
| 579 | 575 | ||
| 580 | chunks = list(iter_points_chunks(path, 4)) | 576 | chunks = list(segment_points_io.iter_points_chunks(path, 4)) |
| 581 | 577 | ||
| 582 | assert len(chunks) == 1 | 578 | assert len(chunks) == 1 |
| 583 | assert chunks[0].flags.writeable | 579 | assert chunks[0].flags.writeable |
| 584 | np.testing.assert_array_equal(chunks[0], record["points"]) | 580 | np.testing.assert_array_equal(chunks[0], record["points"]) |
| 585 | 581 | ||
| 586 | 582 | ||
| 587 | def test_iter_points_chunks_falls_back_for_compressed_records(tmp_path: Path) -> None: | 583 | def test_iter_points_chunks_falls_back_for_compressed_records(tmp_path: pathlib.Path) -> None: |
| 588 | record = _make_record(9, seed=23) | 584 | record = _make_record(9, seed=23) |
| 589 | path = tmp_path / "compressed_run3_points.npz" | 585 | path = tmp_path / "compressed_run3_points.npz" |
| 590 | np.savez_compressed(path, **record) | 586 | np.savez_compressed(path, **record) |
| 591 | 587 | ||
| 592 | chunks = list(iter_points_chunks(path, 2)) | 588 | chunks = list(segment_points_io.iter_points_chunks(path, 2)) |
| 593 | 589 | ||
| 594 | assert len(chunks) == 1 | 590 | assert len(chunks) == 1 |
| 595 | np.testing.assert_array_equal(chunks[0], record["points"]) | 591 | np.testing.assert_array_equal(chunks[0], record["points"]) |
| 596 | 592 | ||
| 597 | 593 | ||
| 598 | def test_iter_points_chunks_disabled_by_non_positive_chunk_size(tmp_path: Path) -> None: | 594 | def test_iter_points_chunks_disabled_by_non_positive_chunk_size(tmp_path: pathlib.Path) -> None: |
| 599 | record = _make_record(6, seed=24) | 595 | record = _make_record(6, seed=24) |
| 600 | path = _save_stored_npz(tmp_path / "run3_points.npz", record) | 596 | path = _save_stored_npz(tmp_path / "run3_points.npz", record) |
| 601 | 597 | ||
| 602 | chunks = list(iter_points_chunks(path, 0)) | 598 | chunks = list(segment_points_io.iter_points_chunks(path, 0)) |
| 603 | 599 | ||
| 604 | assert len(chunks) == 1 | 600 | assert len(chunks) == 1 |
| 605 | np.testing.assert_array_equal(chunks[0], record["points"]) | 601 | np.testing.assert_array_equal(chunks[0], record["points"]) |
| 606 | 602 | ||
| 607 | 603 | ||
| 608 | def test_discover_run3_files_skips_part_files(tmp_path: Path) -> None: | 604 | def test_discover_run3_files_skips_part_files(tmp_path: pathlib.Path) -> None: |
| 609 | seg_dir = tmp_path / "segment_003" | 605 | seg_dir = tmp_path / "segment_003" |
| 610 | seg_dir.mkdir() | 606 | seg_dir.mkdir() |
| 611 | save_points_npz(seg_dir / "beta_run3_points.npz", _make_record(2, seed=30)) | 607 | segment_points_io.save_points_npz(seg_dir / "beta_run3_points.npz", _make_record(2, seed=30)) |
| 612 | save_points_npz(seg_dir / "alpha_run3_points.npz", _make_record(2, seed=31)) | 608 | segment_points_io.save_points_npz(seg_dir / "alpha_run3_points.npz", _make_record(2, seed=31)) |
| 613 | (seg_dir / "gamma_run3_points.npz.part").write_bytes(b"partial write") | 609 | (seg_dir / "gamma_run3_points.npz.part").write_bytes(b"partial write") |
| 614 | (seg_dir / "delta_run4_road_surface.npz").write_bytes(b"other artifact") | 610 | (seg_dir / "delta_run4_road_surface.npz").write_bytes(b"other artifact") |
| 615 | 611 | ||
| 616 | found = discover_run3_files(seg_dir) | 612 | found = segment_points_io.discover_run3_files(seg_dir) |
| 617 | 613 | ||
| 618 | assert [path.name for path in found] == [ | 614 | assert [path.name for path in found] == [ |
| 619 | "alpha_run3_points.npz", | 615 | "alpha_run3_points.npz", |
| 620 | "beta_run3_points.npz", | 616 | "beta_run3_points.npz", |
| 621 | ] | 617 | ] |
| 622 | 618 | ||
| 623 | 619 | ||
| 624 | def test_discover_run3_files_missing_dir_is_empty(tmp_path: Path) -> None: | 620 | def test_discover_run3_files_missing_dir_is_empty(tmp_path: pathlib.Path) -> None: |
| 625 | assert discover_run3_files(tmp_path / "segment_999") == [] | 621 | assert segment_points_io.discover_run3_files(tmp_path / "segment_999") == [] |
| 626 | 622 | ||
| 627 | 623 | ||
| 628 | def test_load_run3_segment_boundary_table(tmp_path: Path) -> None: | 624 | def test_load_run3_segment_boundary_table(tmp_path: pathlib.Path) -> None: |
| 629 | seg_dir = tmp_path / "segment_007" | 625 | seg_dir = tmp_path / "segment_007" |
| 630 | seg_dir.mkdir() | 626 | seg_dir.mkdir() |
| 631 | first = _make_record(3, seed=40) | 627 | first = _make_record(3, seed=40) |
| 632 | second = _make_record(5, seed=41) | 628 | second = _make_record(5, seed=41) |
| 633 | save_points_npz(seg_dir / "a_run3_points.npz", first) | 629 | segment_points_io.save_points_npz(seg_dir / "a_run3_points.npz", first) |
| 634 | save_points_npz(seg_dir / "b_run3_points.npz", second) | 630 | segment_points_io.save_points_npz(seg_dir / "b_run3_points.npz", second) |
| 635 | (seg_dir / "c_run3_points.npz.part").write_bytes(b"partial write") | 631 | (seg_dir / "c_run3_points.npz.part").write_bytes(b"partial write") |
| 636 | 632 | ||
| 637 | merged, spans = load_run3_segment(seg_dir) | 633 | merged, spans = segment_points_io.load_run3_segment(seg_dir) |
| 638 | 634 | ||
| 639 | assert spans == [ | 635 | assert spans == [ |
| 640 | RecordSpan(name="a_run3_points.npz", offset=0, count=3), | 636 | segment_points_io.RecordSpan(name="a_run3_points.npz", offset=0, count=3), |
| 641 | RecordSpan(name="b_run3_points.npz", offset=3, count=5), | 637 | segment_points_io.RecordSpan(name="b_run3_points.npz", offset=3, count=5), |
| 642 | ] | 638 | ] |
| 643 | assert spans[1].end == 8 | 639 | assert spans[1].end == 8 |
| 644 | assert merged["points"].shape == (8, 3) | 640 | assert merged["points"].shape == (8, 3) |
| 645 | np.testing.assert_array_equal( | 641 | np.testing.assert_array_equal( |
| 653 | np.concatenate([first["intensity"], second["intensity"]], axis=0), | 649 | np.concatenate([first["intensity"], second["intensity"]], axis=0), |
| 654 | ) | 650 | ) |
| 655 | 651 | ||
| 656 | 652 | ||
| 657 | def test_load_run3_segment_requires_records(tmp_path: Path) -> None: | 653 | def test_load_run3_segment_requires_records(tmp_path: pathlib.Path) -> None: |
| 658 | seg_dir = tmp_path / "segment_008" | 654 | seg_dir = tmp_path / "segment_008" |
| 659 | seg_dir.mkdir() | 655 | seg_dir.mkdir() |
| 660 | (seg_dir / "a_run3_points.npz.part").write_bytes(b"partial write") | 656 | (seg_dir / "a_run3_points.npz.part").write_bytes(b"partial write") |
| 661 | with pytest.raises(FileNotFoundError, match=r"No \*_run3_points.npz"): | 657 | with pytest.raises(FileNotFoundError, match=r"No \*_run3_points.npz"): |
| 662 | load_run3_segment(seg_dir) | 658 | segment_points_io.load_run3_segment(seg_dir) |
| 663 | 659 | ||
| 664 | 660 | ||
| 665 | def test_concat_points_npz_rejects_mixed_dtypes(tmp_path: Path) -> None: | 661 | def test_concat_points_npz_rejects_mixed_dtypes(tmp_path: pathlib.Path) -> None: |
| 666 | first = _make_record(3, seed=50) | 662 | first = _make_record(3, seed=50) |
| 667 | second = _make_record(3, seed=51) | 663 | second = _make_record(3, seed=51) |
| 668 | second["points"] = second["points"].astype(np.float64) | 664 | second["points"] = second["points"].astype(np.float64) |
| 669 | save_points_npz(tmp_path / "a_run3_points.npz", first) | 665 | segment_points_io.save_points_npz(tmp_path / "a_run3_points.npz", first) |
| 670 | save_points_npz(tmp_path / "b_run3_points.npz", second) | 666 | segment_points_io.save_points_npz(tmp_path / "b_run3_points.npz", second) |
| 671 | 667 | ||
| 672 | with pytest.raises(ValueError, match="mixed dtypes"): | 668 | with pytest.raises(ValueError, match="mixed dtypes"): |
| 673 | concat_points_npz( | 669 | segment_points_io.concat_points_npz( |
| 674 | [tmp_path / "a_run3_points.npz", tmp_path / "b_run3_points.npz"] | 670 | [tmp_path / "a_run3_points.npz", tmp_path / "b_run3_points.npz"] |
| 675 | ) | 671 | ) |
| 676 | 672 | ||
| 677 | 673 | ||
| 678 | def test_concat_points_npz_rejects_mixed_ancillary_dtypes(tmp_path: Path) -> None: | 674 | def test_concat_points_npz_rejects_mixed_ancillary_dtypes(tmp_path: pathlib.Path) -> None: |
| 679 | first = _make_record(2, seed=52) | 675 | first = _make_record(2, seed=52) |
| 680 | second = _make_record(2, seed=53) | 676 | second = _make_record(2, seed=53) |
| 681 | second["intensity"] = second["intensity"].astype(np.uint8) | 677 | second["intensity"] = second["intensity"].astype(np.uint8) |
| 682 | save_points_npz(tmp_path / "a_run3_points.npz", first) | 678 | segment_points_io.save_points_npz(tmp_path / "a_run3_points.npz", first) |
| 683 | save_points_npz(tmp_path / "b_run3_points.npz", second) | 679 | segment_points_io.save_points_npz(tmp_path / "b_run3_points.npz", second) |
| 684 | 680 | ||
| 685 | with pytest.raises(ValueError, match="'intensity' dtype uint8"): | 681 | with pytest.raises(ValueError, match="'intensity' dtype uint8"): |
| 686 | concat_points_npz( | 682 | segment_points_io.concat_points_npz( |
| 687 | [tmp_path / "a_run3_points.npz", tmp_path / "b_run3_points.npz"] | 683 | [tmp_path / "a_run3_points.npz", tmp_path / "b_run3_points.npz"] |
| 688 | ) | 684 | ) |
| 689 | 685 | ||
| 690 | 686 | ||
| 691 | def test_concat_points_npz_empty_raises() -> None: | 687 | def test_concat_points_npz_empty_raises() -> None: |
| 692 | with pytest.raises(FileNotFoundError, match="No \\*_points.npz"): | 688 | with pytest.raises(FileNotFoundError, match="No \\*_points.npz"): |
| 693 | concat_points_npz([]) | 689 | segment_points_io.concat_points_npz([]) |
| 694 | 690 | ||
| 695 | 691 | ||
| 696 | def test_concat_points_npz_target_dtypes_casts_mixed_records(tmp_path: Path) -> None: | 692 | def test_concat_points_npz_target_dtypes_casts_mixed_records(tmp_path: pathlib.Path) -> None: |
| 697 | # seg3d-style normalization: historical records with different storage | 693 | # seg3d-style normalization: historical records with different storage |
| 698 | # dtypes (uint8 vs uint16 intensity) are cast to the target instead of | 694 | # dtypes (uint8 vs uint16 intensity) are cast to the target instead of |
| 699 | # rejected. | 695 | # rejected. |
| 700 | first = _make_record(2, seed=54) | 696 | first = _make_record(2, seed=54) |
| 701 | second = _make_record(2, seed=55) | 697 | second = _make_record(2, seed=55) |
| 702 | second["intensity"] = second["intensity"].astype(np.uint8) | 698 | second["intensity"] = second["intensity"].astype(np.uint8) |
| 703 | save_points_npz(tmp_path / "a_run3_points.npz", first) | 699 | segment_points_io.save_points_npz(tmp_path / "a_run3_points.npz", first) |
| 704 | save_points_npz(tmp_path / "b_run3_points.npz", second) | 700 | segment_points_io.save_points_npz(tmp_path / "b_run3_points.npz", second) |
| 705 | 701 | ||
| 706 | merged, spans = concat_points_npz( | 702 | merged, spans = segment_points_io.concat_points_npz( |
| 707 | [tmp_path / "a_run3_points.npz", tmp_path / "b_run3_points.npz"], | 703 | [tmp_path / "a_run3_points.npz", tmp_path / "b_run3_points.npz"], |
| 708 | target_dtypes={"points": np.dtype(np.float64), "intensity": np.dtype(np.uint16)}, | 704 | target_dtypes={"points": np.dtype(np.float64), "intensity": np.dtype(np.uint16)}, |
| 709 | ) | 705 | ) |
| 710 | assert merged["points"].dtype == np.float64 | 706 | assert merged["points"].dtype == np.float64 |
| 718 | ) | 714 | ) |
| 719 | 715 | ||
| 720 | 716 | ||
| 721 | def test_find_geoshift_warns_on_multiple_candidates( | 717 | def test_find_geoshift_warns_on_multiple_candidates( |
| 722 | tmp_path: Path, caplog: pytest.LogCaptureFixture | 718 | tmp_path: pathlib.Path, caplog: pytest.LogCaptureFixture |
| 723 | ) -> None: | 719 | ) -> None: |
| 724 | # A stale segment-local file shadowing the dataset-level one is legal but | 720 | # A stale segment-local file shadowing the dataset-level one is legal but |
| 725 | # suspicious; the lookup must say so out loud. | 721 | # suspicious; the lookup must say so out loud. |
| 726 | seg_dir = tmp_path / "lane_points" / "segment_032" | 722 | seg_dir = tmp_path / "lane_points" / "segment_032" |
| 728 | _write_geoshift(seg_dir / "run3_geoshift.json", (1.0, 1.0, 1.0)) | 724 | _write_geoshift(seg_dir / "run3_geoshift.json", (1.0, 1.0, 1.0)) |
| 729 | _write_geoshift(tmp_path / "lane_points" / "run3_geoshift.json", (2.0, 2.0, 2.0)) | 725 | _write_geoshift(tmp_path / "lane_points" / "run3_geoshift.json", (2.0, 2.0, 2.0)) |
| 730 | 726 | ||
| 731 | with caplog.at_level(logging.WARNING, logger="iolabs.common.segment_points_io"): | 727 | with caplog.at_level(logging.WARNING, logger="iolabs.common.segment_points_io"): |
| 732 | shift = find_geoshift(seg_dir) | 728 | shift = segment_points_io.find_geoshift(seg_dir) |
| 733 | 729 | ||
| 734 | np.testing.assert_array_equal(shift, np.array([1.0, 1.0, 1.0])) | 730 | np.testing.assert_array_equal(shift, np.array([1.0, 1.0, 1.0])) |
| 735 | assert any("Multiple" in message for message in caplog.messages) | 731 | assert any("Multiple" in message for message in caplog.messages) |
| 736 | 732 | ||
| 737 | 733 | ||
| 738 | def test_public_key_tuples_are_derived_from_the_schema() -> None: | 734 | def test_public_key_tuples_are_derived_from_the_schema() -> None: |
| 739 | assert POINT_RECORD_KEYS == tuple(POINT_RECORD_SCHEMA) | 735 | assert segment_points_io.POINT_RECORD_KEYS == tuple(segment_points_io.POINT_RECORD_SCHEMA) |
| 740 | assert REQUIRED_POINT_RECORD_KEYS == ( | 736 | assert segment_points_io.REQUIRED_POINT_RECORD_KEYS == ( |
| 741 | "points", | 737 | "points", |
| 742 | "red", | 738 | "red", |
| 743 | "green", | 739 | "green", |
| 744 | "blue", | 740 | "blue", |
| 745 | "intensity", | 741 | "intensity", |
| 746 | "scan_angle", | 742 | "scan_angle", |
| 747 | ) | 743 | ) |
| 748 | assert OPTIONAL_POINT_RECORD_KEYS == (NUMBER_OF_RETURNS_KEY,) | 744 | assert segment_points_io.OPTIONAL_POINT_RECORD_KEYS == ( |
| 749 | assert all(POINT_RECORD_SCHEMA[key].required for key in REQUIRED_POINT_RECORD_KEYS) | 745 | segment_points_io.NUMBER_OF_RETURNS_KEY, |
| 750 | assert POINT_RECORD_SCHEMA[NUMBER_OF_RETURNS_KEY].storage_dtype == np.dtype(np.uint8) | 746 | ) |
| 751 | assert POINT_RECORD_SCHEMA["points"].columns == 3 | 747 | assert all( |
| 748 | segment_points_io.POINT_RECORD_SCHEMA[key].required | ||
| 749 | for key in segment_points_io.REQUIRED_POINT_RECORD_KEYS | ||
| 750 | ) | ||
| 751 | assert segment_points_io.POINT_RECORD_SCHEMA[ | ||
| 752 | segment_points_io.NUMBER_OF_RETURNS_KEY | ||
| 753 | ].storage_dtype == np.dtype(np.uint8) | ||
| 754 | assert segment_points_io.POINT_RECORD_SCHEMA["points"].columns == 3 | ||
| 752 | 755 | ||
| 753 | 756 | ||
| 754 | def test_point_field_spec_rejects_an_optional_key_without_a_fill() -> None: | 757 | def test_point_field_spec_rejects_an_optional_key_without_a_fill() -> None: |
| 755 | """An optional key with no fill would make pre-existing records unloadable.""" | 758 | """An optional key with no fill would make pre-existing records unloadable.""" |
| 756 | with pytest.raises(ValueError, match="needs a fill factory"): | 759 | with pytest.raises(ValueError, match="needs a fill factory"): |
| 757 | PointFieldSpec(required=False) | 760 | segment_points_io.PointFieldSpec(required=False) |
| 758 | 761 | ||
| 759 | 762 | ||
| 760 | def test_mask_record_masks_points_rows_and_ancillary_elements() -> None: | 763 | def test_mask_record_masks_points_rows_and_ancillary_elements() -> None: |
| 761 | record = _make_record(5, seed=80) | 764 | record = _make_record(5, seed=80) |
| 762 | mask = np.array([True, False, True, False, True]) | 765 | mask = np.array([True, False, True, False, True]) |
| 763 | 766 | ||
| 764 | masked = mask_record(record, mask) | 767 | masked = segment_points_io.mask_record(record, mask) |
| 765 | 768 | ||
| 766 | assert list(masked) == list(record) | 769 | assert list(masked) == list(record) |
| 767 | assert masked["points"].shape == (3, 3) | 770 | assert masked["points"].shape == (3, 3) |
| 768 | np.testing.assert_array_equal(masked["points"], record["points"][mask]) | 771 | np.testing.assert_array_equal(masked["points"], record["points"][mask]) |
| 769 | for key in POINT_RECORD_KEYS: | 772 | for key in segment_points_io.POINT_RECORD_KEYS: |
| 770 | np.testing.assert_array_equal(masked[key], record[key][mask]) | 773 | np.testing.assert_array_equal(masked[key], record[key][mask]) |
| 771 | 774 | ||
| 772 | 775 | ||
| 773 | def test_mask_record_accepts_an_integer_index_array() -> None: | 776 | def test_mask_record_accepts_an_integer_index_array() -> None: |
| 774 | record = _make_record(4, seed=81) | 777 | record = _make_record(4, seed=81) |
| 775 | index = np.array([3, 0]) | 778 | index = np.array([3, 0]) |
| 776 | 779 | ||
| 777 | masked = mask_record(record, index) | 780 | masked = segment_points_io.mask_record(record, index) |
| 778 | 781 | ||
| 779 | np.testing.assert_array_equal(masked["points"], record["points"][index]) | 782 | np.testing.assert_array_equal(masked["points"], record["points"][index]) |
| 780 | np.testing.assert_array_equal(masked["intensity"], record["intensity"][index]) | 783 | np.testing.assert_array_equal(masked["intensity"], record["intensity"][index]) |
| 781 | 784 |
| 783 | def test_mask_record_rejects_misaligned_members() -> None: | 786 | def test_mask_record_rejects_misaligned_members() -> None: |
| 784 | record = _make_record(4, seed=82) | 787 | record = _make_record(4, seed=82) |
| 785 | record["red"] = record["red"][:2] | 788 | record["red"] = record["red"][:2] |
| 786 | with pytest.raises(ValueError, match="disagree on point count"): | 789 | with pytest.raises(ValueError, match="disagree on point count"): |
| 787 | mask_record(record, np.array([True, False, True, False])) | 790 | segment_points_io.mask_record(record, np.array([True, False, True, False])) |
| 788 | 791 | ||
| 789 | 792 | ||
| 790 | def test_mask_record_rejects_an_empty_record() -> None: | 793 | def test_mask_record_rejects_an_empty_record() -> None: |
| 791 | with pytest.raises(ValueError, match="empty point record"): | 794 | with pytest.raises(ValueError, match="empty point record"): |
| 792 | mask_record({}, np.array([True])) | 795 | segment_points_io.mask_record({}, np.array([True])) |
| 793 | 796 | ||
| 794 | 797 | ||
| 795 | def test_concat_records_joins_every_member_on_axis_zero() -> None: | 798 | def test_concat_records_joins_every_member_on_axis_zero() -> None: |
| 796 | first = _make_record(3, seed=83) | 799 | first = _make_record(3, seed=83) |
| 797 | second = _make_record(2, seed=84) | 800 | second = _make_record(2, seed=84) |
| 798 | 801 | ||
| 799 | merged = concat_records([first, second]) | 802 | merged = segment_points_io.concat_records([first, second]) |
| 800 | 803 | ||
| 801 | assert list(merged) == list(POINT_RECORD_KEYS) | 804 | assert list(merged) == list(segment_points_io.POINT_RECORD_KEYS) |
| 802 | assert merged["points"].shape == (5, 3) | 805 | assert merged["points"].shape == (5, 3) |
| 803 | for key in POINT_RECORD_KEYS: | 806 | for key in segment_points_io.POINT_RECORD_KEYS: |
| 804 | np.testing.assert_array_equal( | 807 | np.testing.assert_array_equal( |
| 805 | merged[key], np.concatenate([first[key], second[key]], axis=0) | 808 | merged[key], np.concatenate([first[key], second[key]], axis=0) |
| 806 | ) | 809 | ) |
| 807 | 810 | ||
| 808 | 811 | ||
| 809 | def test_concat_records_rejects_mismatched_key_sets() -> None: | 812 | def test_concat_records_rejects_mismatched_key_sets() -> None: |
| 810 | with pytest.raises(ValueError, match="different keys"): | 813 | with pytest.raises(ValueError, match="different keys"): |
| 811 | concat_records([_make_record(2, seed=85), _make_legacy_record(2, seed=86)]) | 814 | segment_points_io.concat_records( |
| 815 | [_make_record(2, seed=85), _make_legacy_record(2, seed=86)] | ||
| 816 | ) | ||
| 812 | 817 | ||
| 813 | 818 | ||
| 814 | def test_concat_records_rejects_an_empty_sequence() -> None: | 819 | def test_concat_records_rejects_an_empty_sequence() -> None: |
| 815 | with pytest.raises(ValueError, match="empty sequence"): | 820 | with pytest.raises(ValueError, match="empty sequence"): |
| 816 | concat_records([]) | 821 | segment_points_io.concat_records([]) |
| 817 | 822 | ||
| 818 | 823 | ||
| 819 | def test_concat_records_rejects_records_without_members() -> None: | 824 | def test_concat_records_rejects_records_without_members() -> None: |
| 820 | """Memberless records must raise, not silently concatenate to ``{}``.""" | 825 | """Memberless records must raise, not silently concatenate to ``{}``.""" |
| 821 | with pytest.raises(ValueError, match="no members"): | 826 | with pytest.raises(ValueError, match="no members"): |
| 822 | concat_records([{}, {}]) | 827 | segment_points_io.concat_records([{}, {}]) |
| 823 | 828 | ||
| 824 | 829 | ||
| 825 | def test_mask_record_rejects_zero_dimensional_members() -> None: | 830 | def test_mask_record_rejects_zero_dimensional_members() -> None: |
| 826 | """A 0-D member is not row-indexable: raise ValueError, not a raw IndexError.""" | 831 | """A 0-D member is not row-indexable: raise ValueError, not a raw IndexError.""" |
| 827 | record = {key: np.asarray(1, dtype=np.uint8) for key in POINT_RECORD_KEYS} | 832 | record = {key: np.asarray(1, dtype=np.uint8) for key in segment_points_io.POINT_RECORD_KEYS} |
| 828 | 833 | ||
| 829 | with pytest.raises(ValueError, match="not 1-D or 2-D"): | 834 | with pytest.raises(ValueError, match="not 1-D or 2-D"): |
| 830 | mask_record(record, np.array([True])) | 835 | segment_points_io.mask_record(record, np.array([True])) |
| 831 | 836 | ||
| 832 | 837 | ||
| 833 | # --- Adding a field to the contract must be a registry entry and nothing else --- | 838 | # --- Adding a field to the contract must be a registry entry and nothing else --- |
| 834 | 839 | ||
| 835 | EXTRA_KEY = "point_source_id" | 840 | EXTRA_KEY = "point_source_id" |
| 836 | 841 | ||
| 837 | EXTRA_SPEC = PointFieldSpec( | 842 | EXTRA_SPEC = segment_points_io.PointFieldSpec( |
| 838 | required=False, | 843 | required=False, |
| 839 | storage_dtype=np.dtype(np.int16), | 844 | storage_dtype=np.dtype(np.int16), |
| 840 | fill=functools.partial(np.full, fill_value=-1, dtype=np.int16), | 845 | fill=functools.partial(np.full, fill_value=-1, dtype=np.int16), |
| 841 | noun="source id", | 846 | noun="source id", |
| 847 | """Register one extra optional key, exactly as a real schema addition would.""" | 852 | """Register one extra optional key, exactly as a real schema addition would.""" |
| 848 | monkeypatch.setattr( | 853 | monkeypatch.setattr( |
| 849 | segment_points_io, | 854 | segment_points_io, |
| 850 | "POINT_RECORD_SCHEMA", | 855 | "POINT_RECORD_SCHEMA", |
| 851 | MappingProxyType({**POINT_RECORD_SCHEMA, EXTRA_KEY: EXTRA_SPEC}), | 856 | types.MappingProxyType({**segment_points_io.POINT_RECORD_SCHEMA, EXTRA_KEY: EXTRA_SPEC}), |
| 852 | ) | 857 | ) |
| 853 | return (*POINT_RECORD_KEYS, EXTRA_KEY) | 858 | return (*segment_points_io.POINT_RECORD_KEYS, EXTRA_KEY) |
| 854 | 859 | ||
| 855 | 860 | ||
| 856 | def _make_extended_record(n: int, *, seed: int) -> dict[str, np.ndarray]: | 861 | def _make_extended_record(n: int, *, seed: int) -> dict[str, np.ndarray]: |
| 857 | record = _make_record(n, seed=seed) | 862 | record = _make_record(n, seed=seed) |
| 859 | return record | 864 | return record |
| 860 | 865 | ||
| 861 | 866 | ||
| 862 | def test_registry_entry_alone_round_trips_a_new_field( | 867 | def test_registry_entry_alone_round_trips_a_new_field( |
| 863 | tmp_path: Path, extended_schema: tuple[str, ...] | 868 | tmp_path: pathlib.Path, extended_schema: tuple[str, ...] |
| 864 | ) -> None: | 869 | ) -> None: |
| 865 | record = _make_extended_record(4, seed=90) | 870 | record = _make_extended_record(4, seed=90) |
| 866 | 871 | ||
| 867 | path = save_points_npz(tmp_path / "extended_run3_points.npz", record) | 872 | path = segment_points_io.save_points_npz(tmp_path / "extended_run3_points.npz", record) |
| 868 | 873 | ||
| 869 | with np.load(path) as data: | 874 | with np.load(path) as data: |
| 870 | assert EXTRA_KEY in data.files | 875 | assert EXTRA_KEY in data.files |
| 871 | loaded = load_points_npz(path) | 876 | loaded = segment_points_io.load_points_npz(path) |
| 872 | assert list(loaded) == list(extended_schema) | 877 | assert list(loaded) == list(extended_schema) |
| 873 | assert loaded[EXTRA_KEY].dtype == np.int16 | 878 | assert loaded[EXTRA_KEY].dtype == np.int16 |
| 874 | np.testing.assert_array_equal(loaded[EXTRA_KEY], record[EXTRA_KEY]) | 879 | np.testing.assert_array_equal(loaded[EXTRA_KEY], record[EXTRA_KEY]) |
| 875 | 880 | ||
| 876 | 881 | ||
| 877 | def test_registry_entry_alone_coerces_and_validates_a_new_field( | 882 | def test_registry_entry_alone_coerces_and_validates_a_new_field( |
| 878 | tmp_path: Path, extended_schema: tuple[str, ...] | 883 | tmp_path: pathlib.Path, extended_schema: tuple[str, ...] |
| 879 | ) -> None: | 884 | ) -> None: |
| 880 | record = _make_extended_record(3, seed=91) | 885 | record = _make_extended_record(3, seed=91) |
| 881 | record[EXTRA_KEY] = record[EXTRA_KEY].astype(np.int64) | 886 | record[EXTRA_KEY] = record[EXTRA_KEY].astype(np.int64) |
| 882 | loaded = load_points_npz(save_points_npz(tmp_path / "cast_run3_points.npz", record)) | 887 | loaded = segment_points_io.load_points_npz( |
| 888 | segment_points_io.save_points_npz(tmp_path / "cast_run3_points.npz", record) | ||
| 889 | ) | ||
| 883 | assert loaded[EXTRA_KEY].dtype == np.int16 | 890 | assert loaded[EXTRA_KEY].dtype == np.int16 |
| 884 | 891 | ||
| 885 | out_of_range = _make_extended_record(3, seed=92) | 892 | out_of_range = _make_extended_record(3, seed=92) |
| 886 | out_of_range[EXTRA_KEY] = np.array([1, 40_000, 3], dtype=np.int32) | 893 | out_of_range[EXTRA_KEY] = np.array([1, 40_000, 3], dtype=np.int32) |
| 887 | with pytest.raises(ValueError, match=f"{EXTRA_KEY}.*fit in int16"): | 894 | with pytest.raises(ValueError, match=f"{EXTRA_KEY}.*fit in int16"): |
| 888 | save_points_npz(tmp_path / "range_run3_points.npz", out_of_range) | 895 | segment_points_io.save_points_npz(tmp_path / "range_run3_points.npz", out_of_range) |
| 889 | 896 | ||
| 890 | boolean = _make_extended_record(3, seed=93) | 897 | boolean = _make_extended_record(3, seed=93) |
| 891 | boolean[EXTRA_KEY] = np.array([True, False, True]) | 898 | boolean[EXTRA_KEY] = np.array([True, False, True]) |
| 892 | with pytest.raises(ValueError, match=f"{EXTRA_KEY}.*bool"): | 899 | with pytest.raises(ValueError, match=f"{EXTRA_KEY}.*bool"): |
| 893 | save_points_npz(tmp_path / "bool_run3_points.npz", boolean) | 900 | segment_points_io.save_points_npz(tmp_path / "bool_run3_points.npz", boolean) |
| 894 | 901 | ||
| 895 | misshaped = _make_extended_record(3, seed=94) | 902 | misshaped = _make_extended_record(3, seed=94) |
| 896 | misshaped[EXTRA_KEY] = misshaped[EXTRA_KEY].reshape(3, 1) | 903 | misshaped[EXTRA_KEY] = misshaped[EXTRA_KEY].reshape(3, 1) |
| 897 | with pytest.raises(ValueError, match=rf"'{EXTRA_KEY}' must have shape \(N,\)"): | 904 | with pytest.raises(ValueError, match=rf"'{EXTRA_KEY}' must have shape \(N,\)"): |
| 898 | save_points_npz(tmp_path / "shape_run3_points.npz", misshaped) | 905 | segment_points_io.save_points_npz(tmp_path / "shape_run3_points.npz", misshaped) |
| 899 | 906 | ||
| 900 | 907 | ||
| 901 | def test_registry_entry_alone_fills_and_requires_a_new_field( | 908 | def test_registry_entry_alone_fills_and_requires_a_new_field( |
| 902 | tmp_path: Path, extended_schema: tuple[str, ...] | 909 | tmp_path: pathlib.Path, extended_schema: tuple[str, ...] |
| 903 | ) -> None: | 910 | ) -> None: |
| 904 | older = _make_record(5, seed=95) | 911 | older = _make_record(5, seed=95) |
| 905 | path = tmp_path / "older_run3_points.npz" | 912 | path = tmp_path / "older_run3_points.npz" |
| 906 | np.savez_compressed(path, **older) | 913 | np.savez_compressed(path, **older) |
| 907 | 914 | ||
| 908 | loaded = load_points_npz(path) | 915 | loaded = segment_points_io.load_points_npz(path) |
| 909 | 916 | ||
| 910 | assert list(loaded) == list(extended_schema) | 917 | assert list(loaded) == list(extended_schema) |
| 911 | np.testing.assert_array_equal(loaded[EXTRA_KEY], np.full(5, -1, dtype=np.int16)) | 918 | np.testing.assert_array_equal(loaded[EXTRA_KEY], np.full(5, -1, dtype=np.int16)) |
| 912 | with pytest.raises(ValueError, match=EXTRA_KEY): | 919 | with pytest.raises(ValueError, match=EXTRA_KEY): |
| 913 | save_points_npz(tmp_path / "incomplete_run3_points.npz", older) | 920 | segment_points_io.save_points_npz(tmp_path / "incomplete_run3_points.npz", older) |
| 914 | 921 | ||
| 915 | 922 | ||
| 916 | def test_registry_entry_alone_flows_through_merge_mask_and_concat( | 923 | def test_registry_entry_alone_flows_through_merge_mask_and_concat( |
| 917 | tmp_path: Path, extended_schema: tuple[str, ...] | 924 | tmp_path: pathlib.Path, extended_schema: tuple[str, ...] |
| 918 | ) -> None: | 925 | ) -> None: |
| 919 | seg_dir = tmp_path / "segment_042" | 926 | seg_dir = tmp_path / "segment_042" |
| 920 | seg_dir.mkdir() | 927 | seg_dir.mkdir() |
| 921 | first = _make_extended_record(3, seed=96) | 928 | first = _make_extended_record(3, seed=96) |
| 922 | second = _make_extended_record(2, seed=97) | 929 | second = _make_extended_record(2, seed=97) |
| 923 | save_points_npz(seg_dir / "a_run3_points.npz", first) | 930 | segment_points_io.save_points_npz(seg_dir / "a_run3_points.npz", first) |
| 924 | save_points_npz(seg_dir / "b_run3_points.npz", second) | 931 | segment_points_io.save_points_npz(seg_dir / "b_run3_points.npz", second) |
| 925 | 932 | ||
| 926 | merged, spans = load_run3_segment(seg_dir) | 933 | merged, spans = segment_points_io.load_run3_segment(seg_dir) |
| 927 | segment_merged, point_file_ids, _ = load_segment_points( | 934 | segment_merged, point_file_ids, _ = segment_points_io.load_segment_points( |
| 928 | [seg_dir / "a_run3_points.npz", seg_dir / "b_run3_points.npz"] | 935 | [seg_dir / "a_run3_points.npz", seg_dir / "b_run3_points.npz"] |
| 929 | ) | 936 | ) |
| 930 | masked = mask_record(merged, np.array([True, False, True, False, True])) | 937 | masked = segment_points_io.mask_record(merged, np.array([True, False, True, False, True])) |
| 931 | joined = concat_records([first, second]) | 938 | joined = segment_points_io.concat_records([first, second]) |
| 932 | 939 | ||
| 933 | expected = np.concatenate([first[EXTRA_KEY], second[EXTRA_KEY]], axis=0) | 940 | expected = np.concatenate([first[EXTRA_KEY], second[EXTRA_KEY]], axis=0) |
| 934 | assert [span.count for span in spans] == [3, 2] | 941 | assert [span.count for span in spans] == [3, 2] |
| 935 | assert point_file_ids.shape == (5,) | 942 | assert point_file_ids.shape == (5,) |
| 940 | 947 | ||
| 941 | 948 | ||
| 942 | def test_extended_schema_does_not_leak_into_the_real_contract() -> None: | 949 | def test_extended_schema_does_not_leak_into_the_real_contract() -> None: |
| 943 | """The monkeypatched registry above must not pollute the shipped schema.""" | 950 | """The monkeypatched registry above must not pollute the shipped schema.""" |
| 944 | assert EXTRA_KEY not in POINT_RECORD_SCHEMA | 951 | assert EXTRA_KEY not in _POINT_RECORD_SCHEMA_AT_IMPORT |
| 945 | assert EXTRA_KEY not in segment_points_io.POINT_RECORD_SCHEMA | 952 | assert EXTRA_KEY not in segment_points_io.POINT_RECORD_SCHEMA |
| 946 | assert tuple(segment_points_io.POINT_RECORD_SCHEMA) == POINT_RECORD_KEYS | 953 | assert tuple(segment_points_io.POINT_RECORD_SCHEMA) == _POINT_RECORD_KEYS_AT_IMPORT |