scisynth.testing.invariants

  1"""Reusable invariant checks for observations and pipelines."""
  2
  3from __future__ import annotations
  4
  5import numpy as np
  6
  7from ..core._coords import has_position
  8from ..core.observation import Observation
  9from ..core.rng import SeedLike
 10from ..latent.base import Latent
 11from ..observer.observer import Observer
 12from ..stages.base import Stage
 13
 14
 15class InvariantError(AssertionError):
 16    """An invariant of an Observation or pipeline was violated."""
 17
 18
 19def check_observation(obs: Observation) -> None:
 20    """Check that an observation is well formed.
 21
 22    ``values`` is ``(*spatial, d)``; ``mask`` must have the spatial shape and
 23    ``truth`` the shape of ``values``. ``ids`` must be unique integers with the
 24    spatial shape. Coordinates must match a grid or points layout, masked-out
 25    values must be NaN, valid locations must have finite coordinates, and the
 26    ``meta`` records must be well formed.
 27
 28    Parameters
 29    ----------
 30    obs : Observation
 31        The observation to check.
 32
 33    Raises
 34    ------
 35    InvariantError
 36        If any of these fails.
 37    """
 38    _check_arrays(obs)
 39    _check_ids(obs)
 40    _check_coords(obs)
 41    if not np.all(np.isnan(obs.values[~obs.mask])):
 42        raise InvariantError("masked-out values must be NaN")
 43    for entry in obs.meta:
 44        if set(entry) != {"name", "type", "params"}:
 45            raise InvariantError(f"malformed meta entry: {entry!r}")
 46
 47
 48def _check_arrays(obs: Observation) -> None:
 49    """Check that ``values``, ``mask`` and ``truth`` have matching shapes and types."""
 50    shape = obs.values.shape
 51    spatial = shape[:-1]
 52    if obs.mask.shape != spatial:
 53        raise InvariantError(
 54            f"mask shape {obs.mask.shape} != spatial shape {spatial} of values {shape}"
 55        )
 56    if obs.mask.dtype != np.bool_:
 57        raise InvariantError(f"mask dtype must be bool, got {obs.mask.dtype}")
 58    if obs.truth is not None and obs.truth.shape != shape:
 59        raise InvariantError(f"truth shape {obs.truth.shape} != values shape {shape}")
 60
 61
 62def _check_ids(obs: Observation) -> None:
 63    """Check that the ids are unique integers with the spatial shape."""
 64    spatial = obs.values.shape[:-1]
 65    if obs.ids.shape != spatial or not np.issubdtype(obs.ids.dtype, np.integer):
 66        raise InvariantError(f"ids must be integers of shape {spatial}")
 67    if np.unique(obs.ids).size != obs.ids.size:
 68        raise InvariantError("ids must be unique")
 69
 70
 71def _check_coords(obs: Observation) -> None:
 72    """Check that the coordinates fit a layout and valid locations have a position."""
 73    spatial = obs.values.shape[:-1]
 74    coords = list(obs.coords.values())
 75    if obs.layout == "points" and any(c.shape != spatial for c in coords):
 76        raise InvariantError(
 77            f"coords {[c.shape for c in coords]} match neither a grid nor points "
 78            f"layout for values of shape {obs.values.shape}"
 79        )
 80    if np.any(obs.mask & ~has_position(obs)):
 81        raise InvariantError("valid locations must have finite coordinates")
 82
 83
 84def observations_equal(a: Observation, b: Observation) -> bool:
 85    """Return whether two observations are exactly equal.
 86
 87    Compares values, mask, ids, coordinates, truth and meta, treating NaN as equal to
 88    NaN.
 89
 90    Parameters
 91    ----------
 92    a, b : Observation
 93        The observations to compare.
 94
 95    Returns
 96    -------
 97    bool
 98        True if they are exactly equal.
 99    """
100    if not (
101        np.array_equal(a.values, b.values, equal_nan=True)
102        and np.array_equal(a.mask, b.mask)
103        and np.array_equal(a.ids, b.ids)
104        and a.coords.keys() == b.coords.keys()
105        and all(np.array_equal(a.coords[k], b.coords[k]) for k in a.coords)
106        and a.meta == b.meta
107    ):
108        return False
109    if a.truth is None or b.truth is None:
110        return a.truth is b.truth
111    return np.array_equal(a.truth, b.truth, equal_nan=True)
112
113
114def assert_reproducible(observer: Observer, latent: Latent, seed: SeedLike) -> None:
115    """Check that two runs with the same seed give identical observations.
116
117    Parameters
118    ----------
119    observer : Observer
120        The pipeline to run.
121    latent : Latent
122        The ground truth to observe.
123    seed : int, SeedSequence or None
124        Seed for both runs.
125
126    Raises
127    ------
128    InvariantError
129        If the two runs differ.
130    """
131    if not observations_equal(observer.run(latent, seed), observer.run(latent, seed)):
132        raise InvariantError("same seed gave different observations")
133
134
135def assert_prefix_stable(observer: Observer, latent: Latent, seed: SeedLike) -> None:
136    """Check that removing trailing stages leaves earlier output unchanged.
137
138    Parameters
139    ----------
140    observer : Observer
141        The pipeline to run.
142    latent : Latent
143        The ground truth to observe.
144    seed : int, SeedSequence or None
145        Seed for all runs.
146
147    Raises
148    ------
149    InvariantError
150        If the observer built from the first ``k`` stages differs from the full run
151        at step ``k``.
152    """
153    full = observer.trace(latent, seed)
154    for k in range(len(observer.stages) + 1):
155        head = observer.with_stages(observer.stages[:k])
156        if not observations_equal(head.run(latent, seed), full[k]):
157            raise InvariantError(
158                f"the first {k} stage(s) differ from the full run at step {k}"
159            )
160
161
162def assert_append_stable(
163    observer: Observer, extra: Stage, latent: Latent, seed: SeedLike
164) -> None:
165    """Check that appending a stage leaves every existing step's output unchanged.
166
167    Parameters
168    ----------
169    observer : Observer
170        The pipeline to run.
171    extra : Stage
172        The stage to append.
173    latent : Latent
174        The ground truth to observe.
175    seed : int, SeedSequence or None
176        Seed for both runs.
177
178    Raises
179    ------
180    InvariantError
181        If any step of the run with ``extra`` differs from the run without it.
182    """
183    before = observer.trace(latent, seed)
184    taken = set(observer.names)
185    n = 1
186    while f"appended_{n}" in taken:
187        n += 1
188    after = observer.with_stages(
189        [*observer.stages.items(), (f"appended_{n}", extra)]
190    ).trace(latent, seed)
191    for k, obs in enumerate(before):
192        if not observations_equal(obs, after[k]):
193            raise InvariantError(f"appending {type(extra).__name__} changed step {k}")