improving zonal fast, adding robustness testing

Author: jacotay7

README.md

@@ -110,21 +110,27 @@ kl_gen.save("my_kl_basis.npz")
 
 ## Zonal Fast Basis
 
-`ZonalFastBasisGenerator` groups actuators into binary poke patterns such that no two actuators in the same mode are closer than a user-defined distance `D`. This is useful when you want a compact calibration basis that reduces the number of measurements compared with pure zonal pokes.
+`ZonalFastBasisGenerator` groups actuators into binary poke patterns such that no two actuators in the same mode are closer than a user-defined distance `D`. This is useful when you want a compact calibration basis that reduces the number of measurements compared with pure zonal pokes. For square-grid actuator layouts it uses a modulo lattice grouping directly, and for exotic layouts it falls back to a greedy graph-coloring approach.
 
 ```python
-from aobasis import ZonalFastBasisGenerator, make_circular_actuator_grid
+import numpy as np
 
-positions = make_circular_actuator_grid(telescope_diameter=10.0, grid_size=20)
-
-# Distance threshold in the same units as the actuator coordinates.
-zonal_fast_gen = ZonalFastBasisGenerator(positions, min_distance=0.8)
+from aobasis import ZonalFastBasisGenerator, make_circular_actuator_grid, make_concentric_actuator_grid
 
-# Omitting n_modes returns the full grouped basis.
-zonal_fast_modes = zonal_fast_gen.generate()
-print(zonal_fast_modes.shape)
-
-zonal_fast_gen.plot(count=min(12, zonal_fast_modes.shape[1]), title_prefix="Zonal Fast")
+# Example 1: grid-like actuator positions clipped by a circular pupil.
+positions = make_circular_actuator_grid(telescope_diameter=10.0, grid_size=20)
+grid_gen = ZonalFastBasisGenerator(positions, min_distance=0.8)
+grid_modes = grid_gen.generate()
+print("Grid layout:", grid_modes.shape)
+grid_gen.plot(count=min(12, grid_modes.shape[1]), title_prefix="Zonal Fast Grid")
+
+# Example 2: non-grid actuator positions.
+exotic_positions = make_concentric_actuator_grid(telescope_diameter=10.0, n_rings=5)
+exotic_positions = exotic_positions + 0.03 * np.sin(exotic_positions)
+exotic_gen = ZonalFastBasisGenerator(exotic_positions, min_distance=1.0)
+exotic_modes = exotic_gen.generate()
+print("Exotic layout:", exotic_modes.shape)
+exotic_gen.plot(count=min(12, exotic_modes.shape[1]), title_prefix="Zonal Fast Exotic")
 ```
 
 The returned matrix still has the standard `(n_actuators, n_modes)` layout, but each column is now a sparse binary pattern rather than a single-actuator poke. Every actuator appears in exactly one column of the full basis.

compare_zonal_fast.py

@@ -0,0 +1,184 @@
+import argparse
+from typing import Iterable, List, Sequence, Tuple
+
+import numpy as np
+
+from aobasis import ZonalFastBasisGenerator
+
+
+def make_integer_circular_grid(grid_size: int, radius: float) -> Tuple[np.ndarray, np.ndarray]:
+    """Build a unit-pitch square lattice clipped by a circular aperture."""
+    axis = np.arange(grid_size, dtype=float) - 0.5 * (grid_size - 1)
+    xx, yy = np.meshgrid(axis, axis, indexing="xy")
+    full_positions = np.column_stack((xx.ravel(), yy.ravel()))
+    full_indices = np.column_stack(np.unravel_index(np.arange(grid_size * grid_size), (grid_size, grid_size)))
+
+    mask = np.sum(full_positions**2, axis=1) <= radius**2 + 1e-12
+    return full_positions[mask], full_indices[mask]
+
+
+def build_corner_subgrid_basis(grid_indices: np.ndarray, spacing: int) -> np.ndarray:
+    """Group actuators by their row and column residue classes modulo spacing."""
+    if spacing <= 0:
+        raise ValueError("spacing must be a positive integer.")
+    if grid_indices.ndim != 2 or grid_indices.shape[1] != 2:
+        raise ValueError("grid_indices must have shape (n_actuators, 2).")
+
+    residues = np.mod(grid_indices, spacing)
+    groups = {}
+    for actuator_index, residue in enumerate(residues):
+        key = (int(residue[0]), int(residue[1]))
+        groups.setdefault(key, []).append(actuator_index)
+
+    ordered_groups = [groups[key] for key in sorted(groups)]
+    basis = np.zeros((grid_indices.shape[0], len(ordered_groups)), dtype=float)
+    for mode_index, actuator_group in enumerate(ordered_groups):
+        basis[actuator_group, mode_index] = 1.0
+    return basis
+
+
+def minimum_pairwise_distance(positions: np.ndarray, active_indices: np.ndarray) -> float:
+    if active_indices.size < 2:
+        return np.inf
+
+    active_positions = positions[active_indices]
+    deltas = active_positions[:, None, :] - active_positions[None, :, :]
+    distances = np.linalg.norm(deltas, axis=-1)
+    upper_triangle = distances[np.triu_indices(active_positions.shape[0], k=1)]
+    return float(upper_triangle.min())
+
+
+def validate_basis(positions: np.ndarray, basis: np.ndarray, min_distance: float) -> Tuple[bool, float]:
+    if basis.shape[0] != positions.shape[0]:
+        raise ValueError("basis row count must match the number of actuator positions.")
+
+    covered = np.allclose(basis.sum(axis=1), 1.0)
+    mode_min_distances: List[float] = []
+
+    for mode_index in range(basis.shape[1]):
+        active_indices = np.flatnonzero(basis[:, mode_index] > 0.5)
+        mode_min_distances.append(minimum_pairwise_distance(positions, active_indices))
+
+    worst_case_distance = min(mode_min_distances) if mode_min_distances else np.inf
+    spacing_ok = worst_case_distance >= min_distance - 1e-12
+    return covered and spacing_ok, worst_case_distance
+
+
+def compare_for_spacing(positions: np.ndarray, grid_indices: np.ndarray, spacing: int) -> dict:
+    naive_basis = build_corner_subgrid_basis(grid_indices, spacing)
+    fast_basis = ZonalFastBasisGenerator(positions, min_distance=float(spacing)).generate()
+
+    naive_valid, naive_min_distance = validate_basis(positions, naive_basis, float(spacing))
+    fast_valid, fast_min_distance = validate_basis(positions, fast_basis, float(spacing))
+
+    naive_modes = int(naive_basis.shape[1])
+    fast_modes = int(fast_basis.shape[1])
+    reduction = naive_modes - fast_modes
+    reduction_fraction = reduction / naive_modes if naive_modes else 0.0
+
+    return {
+        "spacing": spacing,
+        "n_actuators": int(positions.shape[0]),
+        "naive_modes": naive_modes,
+        "fast_modes": fast_modes,
+        "reduction": reduction,
+        "reduction_fraction": reduction_fraction,
+        "naive_valid": naive_valid,
+        "fast_valid": fast_valid,
+        "naive_min_distance": naive_min_distance,
+        "fast_min_distance": fast_min_distance,
+    }
+
+
+def parse_distances(values: Sequence[int]) -> List[int]:
+    unique_values = sorted({int(value) for value in values})
+    if not unique_values:
+        raise ValueError("At least one spacing value must be provided.")
+    if any(value <= 0 for value in unique_values):
+        raise ValueError("Spacing values must all be positive integers.")
+    return unique_values
+
+
+def print_report(results: Iterable[dict]) -> None:
+    header = (
+        f"{'D':>4} {'Actuators':>10} {'Naive':>8} {'Fast':>8} {'Saved':>8} {'Saved %':>9} "
+        f"{'Naive OK':>9} {'Fast OK':>8} {'Naive min d':>12} {'Fast min d':>11}"
+    )
+    print(header)
+    print("-" * len(header))
+
+    for result in results:
+        print(
+            f"{result['spacing']:>4d} "
+            f"{result['n_actuators']:>10d} "
+            f"{result['naive_modes']:>8d} "
+            f"{result['fast_modes']:>8d} "
+            f"{result['reduction']:>8d} "
+            f"{100.0 * result['reduction_fraction']:>8.2f}% "
+            f"{str(result['naive_valid']):>9} "
+            f"{str(result['fast_valid']):>8} "
+            f"{result['naive_min_distance']:>12.3f} "
+            f"{result['fast_min_distance']:>11.3f}"
+        )
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(
+        description=(
+            "Compare a naive corner-anchored D x D subgrid zonal-fast basis against "
+            "the graph-coloring zonal-fast basis on a circular aperture."
+        )
+    )
+    parser.add_argument("--grid-size", type=int, default=60, help="Number of points along each side of the square lattice.")
+    parser.add_argument("--radius", type=float, default=30.0, help="Circular aperture radius in lattice-pitch units.")
+    parser.add_argument(
+        "--distances",
+        type=int,
+        nargs="+",
+        default=[2, 3, 4, 5, 6, 8, 10],
+        help="Integer spacing values D to test, in lattice-pitch units.",
+    )
+    parser.add_argument(
+        "--fail-if-not-better",
+        action="store_true",
+        help="Exit with status 1 if the zonal-fast basis is not strictly better for every tested spacing.",
+    )
+    args = parser.parse_args()
+
+    distances = parse_distances(args.distances)
+    positions, grid_indices = make_integer_circular_grid(args.grid_size, args.radius)
+    results = [compare_for_spacing(positions, grid_indices, spacing) for spacing in distances]
+
+    print(
+        f"Circular aperture on a {args.grid_size}x{args.grid_size} unit-pitch lattice, "
+        f"radius={args.radius:.1f}, actuators inside pupil={positions.shape[0]}"
+    )
+    print_report(results)
+
+    wins = [result for result in results if result["fast_modes"] < result["naive_modes"]]
+    ties = [result for result in results if result["fast_modes"] == result["naive_modes"]]
+    losses = [result for result in results if result["fast_modes"] > result["naive_modes"]]
+
+    print()
+    print(f"Fast basis uses fewer modes for {len(wins)} / {len(results)} tested spacings.")
+    if ties:
+        tied_spacings = ", ".join(str(result["spacing"]) for result in ties)
+        print(f"Tied spacings: {tied_spacings}")
+    if losses:
+        loss_spacings = ", ".join(str(result["spacing"]) for result in losses)
+        print(f"Fast basis used more modes at: {loss_spacings}")
+
+    if not losses and not ties:
+        print("Verdict: zonal-fast is strictly better than the naive corner-subgrid basis for every tested spacing.")
+    elif losses:
+        print("Verdict: this zonal-fast implementation does not beat the naive corner-subgrid baseline on this geometry.")
+        print("The graph coloring is producing a valid basis, but not a minimal one for these spacings.")
+    else:
+        print("Verdict: zonal-fast matches the naive construction for some spacings and never improves on it in this run.")
+
+    if args.fail_if_not_better and (losses or ties):
+        raise SystemExit(1)
+
+
+if __name__ == "__main__":
+    main()

pyproject.toml

@@ -4,10 +4,10 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "aobasis"
-version = "1.0.1"
+version = "1.0.2"
 description = "A package for generating AO basis sets (KL, Zernike, Fourier)"
 readme = "README.md"
-authors = [{ name = "Jacob Taylor", email = "jtaylor@keck.hawaii.edu" }]
+authors = [{ name = "Jacob Taylor", email = "jacobataylor7@gmail.com" }]
 license = { file = "LICENSE" }
 classifiers = [
     "Programming Language :: Python :: 3",

src/aobasis/base.py

@@ -4,6 +4,20 @@
 from typing import Tuple, Optional, Union
 from .utils import plot_basis_modes
 
+
+def _validate_positions_array(positions: np.ndarray) -> np.ndarray:
+    try:
+        array = np.asarray(positions, dtype=float)
+    except (TypeError, ValueError) as exc:
+        raise ValueError("positions must be a finite numeric array with shape (n_actuators, 2).") from exc
+
+    if array.ndim != 2 or array.shape[1] != 2:
+        raise ValueError("positions must have shape (n_actuators, 2).")
+    if not np.all(np.isfinite(array)):
+        raise ValueError("positions must contain only finite values.")
+
+    return array
+
 class BasisGenerator(ABC):
     """
     Abstract base class for AO basis generators.
@@ -14,9 +28,21 @@ def __init__(self, positions: np.ndarray):
         Args:
             positions: (N, 2) array of actuator coordinates (x, y) in meters.
         """
-        self.positions = np.array(positions)
+        self.positions = _validate_positions_array(positions)
         self.n_actuators = self.positions.shape[0]
         self.modes: Optional[np.ndarray] = None
+
+    def _validate_n_modes(self, n_modes: int, max_modes: Optional[int] = None) -> int:
+        if isinstance(n_modes, bool) or not isinstance(n_modes, (int, np.integer)):
+            raise ValueError("n_modes must be an integer.")
+
+        n_modes = int(n_modes)
+        if n_modes < 0:
+            raise ValueError("n_modes must be non-negative.")
+        if max_modes is not None and n_modes > max_modes:
+            raise ValueError(f"Cannot generate {n_modes} modes; maximum available is {max_modes}.")
+
+        return n_modes
 
     @abstractmethod
     def generate(self, n_modes: int, **kwargs) -> np.ndarray:
@@ -73,4 +99,5 @@ class ConcreteBasis(BasisGenerator):
     def generate(self, n_modes: int, **kwargs) -> np.ndarray:
         if self.modes is None:
             raise NotImplementedError("This is a loaded basis container.")
+        n_modes = self._validate_n_modes(n_modes, max_modes=self.modes.shape[1])
         return self.modes[:, :n_modes]

src/aobasis/fourier.py

@@ -8,13 +8,20 @@ class FourierBasisGenerator(BasisGenerator):
 
     def __init__(self, positions: np.ndarray, pupil_diameter: float):
         super().__init__(positions)
+        if not np.isscalar(pupil_diameter) or not np.isfinite(pupil_diameter) or pupil_diameter <= 0:
+            raise ValueError("pupil_diameter must be a positive finite scalar.")
         self.pupil_diameter = pupil_diameter
 
     def generate(self, n_modes: int, ignore_piston: bool = False, **kwargs) -> np.ndarray:
         """
         Generate Fourier modes.
         We generate pairs of sin/cos for increasing spatial frequencies.
         """
+        n_modes = self._validate_n_modes(n_modes)
+        if n_modes == 0:
+            self.modes = np.zeros((self.n_actuators, 0), dtype=float)
+            return self.modes
+
         x = self.positions[:, 0]
         y = self.positions[:, 1]
 

src/aobasis/hadamard.py

@@ -16,6 +16,11 @@ def generate(self, n_modes: int, **kwargs) -> np.ndarray:
         we find the next power of 2 >= n_actuators, generate the Hadamard matrix,
         and truncate it to the number of actuators (rows) and requested modes (columns).
         """
+        n_modes = self._validate_n_modes(n_modes)
+        if n_modes == 0:
+            self.modes = np.zeros((self.n_actuators, 0), dtype=int)
+            return self.modes
+
         # Find next power of 2 covering the number of actuators
         # We need at least n_actuators rows to define the pattern on the grid
         # And we need enough columns for n_modes

src/aobasis/kl.py

@@ -99,6 +99,10 @@ class KLBasisGenerator(BasisGenerator):
 
     def __init__(self, positions: np.ndarray, fried_parameter: float = 0.16, outer_scale: float = 30.0, use_gpu: bool = False):
         super().__init__(positions)
+        if not np.isscalar(fried_parameter) or not np.isfinite(fried_parameter) or fried_parameter <= 0:
+            raise ValueError("fried_parameter must be a positive finite scalar.")
+        if not np.isscalar(outer_scale) or not np.isfinite(outer_scale) or outer_scale <= 0:
+            raise ValueError("outer_scale must be a positive finite scalar.")
         self.fried_parameter = fried_parameter
         self.outer_scale = outer_scale
         self.eigenvalues = None
@@ -177,6 +181,14 @@ def _von_karman_covariance_gpu(self):
         return cov
 
     def generate(self, n_modes: int, ignore_piston: bool = False, **kwargs) -> np.ndarray:
+        max_modes = self.n_actuators - (1 if ignore_piston else 0)
+        n_modes = self._validate_n_modes(n_modes, max_modes=max_modes)
+
+        if n_modes == 0:
+            self.eigenvalues = np.array([], dtype=float)
+            self.modes = np.zeros((self.n_actuators, 0), dtype=float)
+            return self.modes
+
         cov = self._von_karman_covariance()
 
         if self.use_gpu: