Added zonal fast basis generator (#8)

* Added zonal fast basis generator

* fixes for ci

Author: jacotay7

.github/workflows/ci.yml

@@ -11,7 +11,7 @@ jobs:
     runs-on: ubuntu-latest
     strategy:
       matrix:
-        python-version: ["3.8","3.9", "3.10", "3.11", "3.12"]
+        python-version: ["3.8","3.9", "3.10", "3.11", "3.12", "3.13"]
 
     steps:
     - uses: actions/checkout@v3

README.md

@@ -9,6 +9,7 @@ A Python package for generating various modal basis sets for Adaptive Optics (AO
 - **Zernike Polynomials**: Standard optical aberration modes (Noll indexing).
 - **Fourier Modes**: Sinusoidal basis sets.
 - **Zonal Basis**: Single actuator pokes (Identity).
+- **Zonal Fast Basis**: Distance-constrained grouped actuator pokes for faster calibration sweeps.
 - **Hadamard Basis**: Orthogonal binary patterns for calibration.
 - **Flexible Geometry**: Works with arbitrary actuator positions (defaulting to circular grids).
 - **Piston Removal**: Option to exclude piston/DC modes from generation.
@@ -107,6 +108,27 @@ kl_gen.plot(count=6, title_prefix="KL Mode")
 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.
+
+```python
+from aobasis import ZonalFastBasisGenerator, make_circular_actuator_grid
+
+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)
+
+# 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")
+```
+
+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.
+
 ## Performance
 
 Generation times for 100 modes benchmarked on the following system:
@@ -121,6 +143,7 @@ Generation times for 100 modes benchmarked on the following system:
 | **Zernike** | 0.001s | 0.002s | 0.005s |
 | **Fourier** | <0.001s | 0.001s | 0.003s |
 | **Zonal** | <0.001s | <0.001s | 0.003s |
+| **Zonal Fast** | depends on spacing threshold | depends on spacing threshold | depends on spacing threshold |
 | **Hadamard** | <0.001s | 0.001s | 0.031s |
 
 *Note: KL basis generation is computationally intensive ($O(N^3)$) due to the dense covariance matrix diagonalization. GPU acceleration provides significant speedup (8-15x) for larger grids.*
@@ -129,7 +152,7 @@ Generation times for 100 modes benchmarked on the following system:
 
 We provide Jupyter notebooks to help you get started.
 
-1.  **Getting Started**: `tutorials/getting_started.ipynb` covers all supported basis types and features.
+1.  **Getting Started**: `tutorials/getting_started.ipynb` covers all supported basis types, including zonal fast grouped pokes.
 
 To run the tutorials:
 ```bash
@@ -171,7 +194,7 @@ If you encounter any bugs or have feature requests, please file an issue on the
 For questions or support, please contact:
 
 **User Name**  
-Email: jtaylor@keck.hawaii.edu
+Email: jacobataylor7@gmail.com
 
 ## License
 

src/aobasis/__init__.py

@@ -9,7 +9,7 @@
 from .kl import KLBasisGenerator
 from .zernike import ZernikeBasisGenerator
 from .fourier import FourierBasisGenerator
-from .zonal import ZonalBasisGenerator
+from .zonal import ZonalBasisGenerator, ZonalFastBasisGenerator
 from .hadamard import HadamardBasisGenerator
 from .utils import make_circular_actuator_grid, make_concentric_actuator_grid, plot_basis_modes
 
@@ -19,6 +19,7 @@
     "ZernikeBasisGenerator",
     "FourierBasisGenerator",
     "ZonalBasisGenerator",
+    "ZonalFastBasisGenerator",
     "HadamardBasisGenerator",
     "make_circular_actuator_grid",
     "make_concentric_actuator_grid",

src/aobasis/zonal.py

@@ -1,4 +1,8 @@
+from typing import List, Optional, Set
+
 import numpy as np
+from scipy.spatial import cKDTree
+
 from .base import BasisGenerator
 
 class ZonalBasisGenerator(BasisGenerator):
@@ -24,3 +28,124 @@ def generate(self, n_modes: int, **kwargs) -> np.ndarray:
 
         self.modes = full_basis[:, :n_modes]
         return self.modes
+
+
+def _build_conflict_graph(positions: np.ndarray, min_distance: float) -> List[Set[int]]:
+    n_actuators = positions.shape[0]
+    adjacency = [set() for _ in range(n_actuators)]
+
+    if n_actuators == 0 or min_distance <= 0:
+        return adjacency
+
+    tree = cKDTree(positions)
+    search_radius = np.nextafter(min_distance, 0.0)
+    for first, second in tree.query_pairs(r=search_radius, output_type="ndarray"):
+        adjacency[int(first)].add(int(second))
+        adjacency[int(second)].add(int(first))
+
+    return adjacency
+
+
+def _dsatur_coloring(adjacency: List[Set[int]]) -> np.ndarray:
+    n_vertices = len(adjacency)
+    colors = np.full(n_vertices, -1, dtype=int)
+    neighbor_colors = [set() for _ in range(n_vertices)]
+    degrees = np.array([len(neighbors) for neighbors in adjacency], dtype=int)
+
+    for _ in range(n_vertices):
+        uncolored = np.flatnonzero(colors < 0)
+        if uncolored.size == 0:
+            break
+
+        saturation = np.array([len(neighbor_colors[index]) for index in uncolored], dtype=int)
+        candidate_order = np.lexsort((-degrees[uncolored], -saturation))
+        vertex = int(uncolored[candidate_order[0]])
+
+        used = neighbor_colors[vertex]
+        color = 0
+        while color in used:
+            color += 1
+
+        colors[vertex] = color
+        for neighbor in adjacency[vertex]:
+            if colors[neighbor] < 0:
+                neighbor_colors[neighbor].add(color)
+
+    return colors
+
+
+def compute_zonal_fast_basis(positions: np.ndarray, min_distance: float) -> np.ndarray:
+    """
+    Compute a distance-constrained zonal basis.
+
+    Each returned mode is a binary poke pattern. Actuators that are closer than
+    ``min_distance`` cannot appear in the same mode, so the basis is built by
+    coloring the actuator conflict graph and turning each color into one column
+    of the returned matrix.
+
+    Args:
+        positions: ``(n_actuators, 2)`` array of actuator coordinates.
+        min_distance: Minimum allowed pairwise distance within a mode.
+
+    Returns:
+        ``(n_actuators, n_modes)`` matrix of binary zonal-fast modes.
+    """
+    positions = np.asarray(positions, dtype=float)
+    if positions.ndim != 2 or positions.shape[1] != 2:
+        raise ValueError("positions must have shape (n_actuators, 2).")
+    if min_distance < 0:
+        raise ValueError("min_distance must be non-negative.")
+    if positions.shape[0] == 0:
+        return np.zeros((0, 0), dtype=float)
+
+    adjacency = _build_conflict_graph(positions, min_distance)
+    colors = _dsatur_coloring(adjacency)
+    n_modes = int(colors.max()) + 1
+
+    basis = np.zeros((positions.shape[0], n_modes), dtype=float)
+    basis[np.arange(positions.shape[0]), colors] = 1.0
+    return basis
+
+
+class ZonalFastBasisGenerator(BasisGenerator):
+    """
+    Generate grouped zonal poke patterns separated by a minimum distance.
+
+    A full zonal-fast basis covers every actuator exactly once while using a
+    compact coloring of the actuator conflict graph.
+    """
+
+    def __init__(self, positions: np.ndarray, min_distance: float):
+        super().__init__(positions)
+        if min_distance < 0:
+            raise ValueError("min_distance must be non-negative.")
+        self.min_distance = float(min_distance)
+        self.full_modes: Optional[np.ndarray] = None
+
+    def generate(self, n_modes: Optional[int] = None, **kwargs) -> np.ndarray:
+        """
+        Generate zonal-fast modes.
+
+        Args:
+            n_modes: Number of grouped poke modes to return. If omitted, return
+                the full distance-constrained basis.
+
+        Returns:
+            ``(n_actuators, n_modes)`` matrix of binary grouped poke patterns.
+        """
+        full_basis = compute_zonal_fast_basis(self.positions, self.min_distance)
+        self.full_modes = full_basis
+
+        if n_modes is None:
+            self.modes = full_basis
+            return self.modes
+
+        if n_modes < 0:
+            raise ValueError("n_modes must be non-negative.")
+        if n_modes > full_basis.shape[1]:
+            raise ValueError(
+                f"Cannot generate {n_modes} zonal-fast modes; full basis only contains {full_basis.shape[1]} modes."
+            )
+
+        self.modes = full_basis[:, :n_modes]
+        return self.modes

tests/test_generators.py

@@ -5,10 +5,22 @@
     ZernikeBasisGenerator, 
     FourierBasisGenerator,
     ZonalBasisGenerator,
+    ZonalFastBasisGenerator,
     HadamardBasisGenerator,
     make_circular_actuator_grid
 )
 
+
+def assert_min_spacing_within_modes(positions, modes, min_distance):
+    for mode_index in range(modes.shape[1]):
+        active = positions[modes[:, mode_index] > 0.5]
+        if active.shape[0] < 2:
+            continue
+        deltas = active[:, None, :] - active[None, :, :]
+        distances = np.linalg.norm(deltas, axis=-1)
+        upper_triangle = distances[np.triu_indices(active.shape[0], k=1)]
+        assert np.all(upper_triangle >= min_distance - 1e-12)
+
 @pytest.fixture
 def grid():
     return make_circular_actuator_grid(telescope_diameter=10.0, grid_size=10)
@@ -184,6 +196,65 @@ def test_zonal_all_modes(small_grid):
     assert modes.shape == (small_grid.shape[0], small_grid.shape[0])
     assert np.allclose(modes, np.eye(small_grid.shape[0]))
 
+def test_zonal_fast_generation_path_graph():
+    positions = np.array([
+        [0.0, 0.0],
+        [1.0, 0.0],
+        [2.0, 0.0],
+    ])
+
+    gen = ZonalFastBasisGenerator(positions, min_distance=1.1)
+    modes = gen.generate()
+
+    assert modes.shape == (3, 2)
+    assert np.allclose(modes.sum(axis=1), 1.0)
+    assert np.allclose(modes.T @ modes, np.diag(np.diag(modes.T @ modes)))
+    assert_min_spacing_within_modes(positions, modes, min_distance=1.1)
+
+def test_zonal_fast_generation_clique():
+    positions = np.array([
+        [0.0, 0.0],
+        [1.0, 0.0],
+        [0.5, np.sqrt(3.0) / 2.0],
+    ])
+
+    gen = ZonalFastBasisGenerator(positions, min_distance=1.01)
+    modes = gen.generate()
+
+    assert modes.shape == (3, 3)
+    assert np.allclose(modes, np.eye(3))
+
+def test_zonal_fast_single_mode_when_distance_is_small():
+    positions = np.array([
+        [0.0, 0.0],
+        [1.0, 0.0],
+        [2.0, 0.0],
+        [3.0, 0.0],
+    ])
+
+    gen = ZonalFastBasisGenerator(positions, min_distance=0.5)
+    modes = gen.generate()
+
+    assert modes.shape == (4, 1)
+    assert np.allclose(modes[:, 0], np.ones(4))
+
+def test_zonal_fast_mode_subset_and_errors():
+    positions = np.array([
+        [0.0, 0.0],
+        [1.0, 0.0],
+        [2.0, 0.0],
+    ])
+
+    gen = ZonalFastBasisGenerator(positions, min_distance=1.1)
+    subset = gen.generate(n_modes=1)
+    assert subset.shape == (3, 1)
+
+    with pytest.raises(ValueError, match="full basis only contains"):
+        gen.generate(n_modes=3)
+
+    with pytest.raises(ValueError, match="non-negative"):
+        ZonalFastBasisGenerator(positions, min_distance=-0.1)
+
 def test_hadamard_generation(grid):
     gen = HadamardBasisGenerator(grid)
     modes = gen.generate(n_modes=8)