Adding code to the repo

Author: jacotay7

.github/workflows/ci.yml

@@ -0,0 +1,28 @@
+name: CI
+
+on:
+  push:
+    branches: [ main ]
+  pull_request:
+    branches: [ main ]
+
+jobs:
+  test:
+    runs-on: ubuntu-latest
+    strategy:
+      matrix:
+        python-version: ["3.8","3.9", "3.10", "3.11", "3.12"]
+
+    steps:
+    - uses: actions/checkout@v3
+    - name: Set up Python ${{ matrix.python-version }}
+      uses: actions/setup-python@v4
+      with:
+        python-version: ${{ matrix.python-version }}
+    - name: Install dependencies
+      run: |
+        python -m pip install --upgrade pip
+        pip install -e .[dev]
+    - name: Run tests
+      run: |
+        pytest

.gitignore

@@ -0,0 +1,10 @@
+__pycache__/
+*.pyc
+*.egg-info/
+dist/
+build/
+.pytest_cache/
+.DS_Store
+*.npz
+*.png
+*.gif

README.md

@@ -0,0 +1,123 @@
+# AO Basis (aobasis)
+
+A Python package for generating various modal basis sets for Adaptive Optics (AO) systems. This tool allows you to easily create, visualize, and save basis sets for any deformable mirror geometry.
+
+## Features
+
+- **Karhunen-Loève (KL) Modes**: Optimized for atmospheric turbulence (Von Kármán spectrum).
+- **Zernike Polynomials**: Standard optical aberration modes (Noll indexing).
+- **Fourier Modes**: Sinusoidal basis sets.
+- **Zonal Basis**: Single actuator pokes (Identity).
+- **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.
+- **Visualization**: Built-in plotting tools for quick inspection.
+- **Serialization**: Save and load basis sets to/from `.npz` files.
+
+## Installation
+
+### Prerequisites
+- Python 3.8 or higher
+
+### Install from Source
+Clone the repository and install using pip:
+
+```bash
+git clone https://github.com/jacotay7/aobasis.git
+cd aobasis
+pip install .
+```
+
+For development (editable install with test dependencies):
+```bash
+pip install -e ".[dev]"
+```
+
+## Quick Start
+
+Here is a simple example of generating and plotting KL modes for a 10-meter telescope:
+
+```python
+from aobasis import KLBasisGenerator, make_circular_actuator_grid
+
+# 1. Define the actuator geometry
+positions = make_circular_actuator_grid(telescope_diameter=10.0, grid_size=20)
+
+# 2. Initialize the generator
+kl_gen = KLBasisGenerator(positions, fried_parameter=0.16, outer_scale=30.0)
+
+# 3. Generate modes (excluding piston)
+modes = kl_gen.generate(n_modes=50, ignore_piston=True)
+
+# 4. Plot the first 6 modes
+kl_gen.plot(count=6, title_prefix="KL Mode")
+
+# 5. Save to disk
+kl_gen.save("my_kl_basis.npz")
+```
+
+## Performance
+
+Generation times for 100 modes on a standard laptop (M1/M2 class):
+
+| Basis | 16x16 Grid (~170 acts) | 32x32 Grid (~740 acts) | 64x64 Grid (~3100 acts) |
+|-------|------------------------|------------------------|-------------------------|
+| **KL** | 0.01s | 0.29s | 5.60s |
+| **Zernike** | 0.001s | 0.002s | 0.02s |
+| **Fourier** | 0.001s | 0.001s | 0.003s |
+| **Zonal** | <0.001s | <0.001s | 0.005s |
+| **Hadamard** | <0.001s | 0.004s | 0.09s |
+
+*Note: KL basis generation is computationally intensive ($O(N^3)$) due to the dense covariance matrix diagonalization.*
+
+## Tutorials
+
+We provide Jupyter notebooks to help you get started.
+
+1.  **Getting Started**: `tutorials/getting_started.ipynb` covers all supported basis types and features.
+
+To run the tutorials:
+```bash
+# Install Jupyter if you haven't already
+pip install jupyter
+
+# Launch the notebook server
+jupyter notebook tutorials/getting_started.ipynb
+```
+
+## Development & Testing
+
+This project uses `pytest` for testing. To run the test suite:
+
+```bash
+# Install dev dependencies
+pip install -e ".[dev]"
+
+# Run tests
+pytest
+```
+
+## Contributing
+
+Contributions are welcome! Please feel free to submit a Pull Request.
+
+1.  Fork the repository.
+2.  Create your feature branch (`git checkout -b feature/AmazingFeature`).
+3.  Commit your changes (`git commit -m 'Add some AmazingFeature'`).
+4.  Push to the branch (`git push origin feature/AmazingFeature`).
+5.  Open a Pull Request.
+
+## Issues
+
+If you encounter any bugs or have feature requests, please file an issue on the [GitHub Issues](https://github.com/jacotay7/aobasis/issues) page.
+
+## Contact
+
+For questions or support, please contact:
+
+**User Name**  
+Email: jtaylor@keck.hawaii.edu
+
+## License
+
+This project is licensed under the MIT License - see the [LICENSE](LICENSE) file for details.

benchmark.py

@@ -0,0 +1,52 @@
+import time
+import numpy as np
+from aobasis import (
+    KLBasisGenerator, 
+    ZernikeBasisGenerator, 
+    FourierBasisGenerator,
+    ZonalBasisGenerator,
+    HadamardBasisGenerator,
+    make_circular_actuator_grid
+)
+
+def benchmark():
+    # Grid sizes to test
+    grid_sizes = [16, 32, 64]
+    
+    generators = [
+        ("KL", KLBasisGenerator, {"fried_parameter": 0.16, "outer_scale": 30.0}),
+        ("Zernike", ZernikeBasisGenerator, {"pupil_radius": 5.0}),
+        ("Fourier", FourierBasisGenerator, {"pupil_diameter": 10.0}),
+        ("Zonal", ZonalBasisGenerator, {}),
+        ("Hadamard", HadamardBasisGenerator, {})
+    ]
+    
+    print(f"Benchmarking Basis Generation (n_modes=n_acts)")
+    print("=" * 80)
+    print(f"{'Basis':<10} | {'Grid Size':<10} | {'Actuators':<10} | {'Time (s)':<10}")
+    print("-" * 80)
+    
+    results = {}
+    
+    for name, GenClass, kwargs in generators:
+        results[name] = []
+        for size in grid_sizes:
+            # Setup
+            positions = make_circular_actuator_grid(10.0, size)
+            n_act = positions.shape[0]
+            
+            # Update kwargs based on size if needed (e.g. pupil radius is constant)
+            gen = GenClass(positions, **kwargs)
+            
+            # Timing
+            start_time = time.perf_counter()
+            gen.generate(n_modes=n_act)
+            end_time = time.perf_counter()
+            
+            duration = end_time - start_time
+            print(f"{name:<10} | {size:<10} | {n_act:<10} | {duration:<10.4f}")
+            results[name].append((size, n_act, duration))
+        print("-" * 80)
+
+if __name__ == "__main__":
+    benchmark()

pyproject.toml

@@ -0,0 +1,34 @@
+[build-system]
+requires = ["setuptools>=61.0"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "aobasis"
+version = "0.1.0"
+description = "A package for generating AO basis sets (KL, Zernike, Fourier)"
+readme = "README.md"
+authors = [{ name = "User", email = "user@example.com" }]
+license = { file = "LICENSE" }
+classifiers = [
+    "Programming Language :: Python :: 3",
+    "License :: OSI Approved :: MIT License",
+    "Operating System :: OS Independent",
+]
+dependencies = [
+    "numpy>=1.20",
+    "scipy>=1.7",
+    "matplotlib>=3.5",
+]
+requires-python = ">=3.8"
+
+[project.optional-dependencies]
+dev = [
+    "pytest",
+    "imageio",
+]
+
+[project.urls]
+"Homepage" = "https://github.com/example/aobasis"
+
+[tool.setuptools.packages.find]
+where = ["src"]

src/aobasis/__init__.py

@@ -0,0 +1,19 @@
+from .base import BasisGenerator
+from .kl import KLBasisGenerator
+from .zernike import ZernikeBasisGenerator
+from .fourier import FourierBasisGenerator
+from .zonal import ZonalBasisGenerator
+from .hadamard import HadamardBasisGenerator
+from .utils import make_circular_actuator_grid, make_concentric_actuator_grid, plot_basis_modes
+
+__all__ = [
+    "BasisGenerator",
+    "KLBasisGenerator",
+    "ZernikeBasisGenerator",
+    "FourierBasisGenerator",
+    "ZonalBasisGenerator",
+    "HadamardBasisGenerator",
+    "make_circular_actuator_grid",
+    "make_concentric_actuator_grid",
+    "plot_basis_modes",
+]

src/aobasis/base.py

@@ -0,0 +1,76 @@
+from abc import ABC, abstractmethod
+import numpy as np
+from pathlib import Path
+from typing import Tuple, Optional
+from .utils import plot_basis_modes
+
+class BasisGenerator(ABC):
+    """
+    Abstract base class for AO basis generators.
+    """
+    
+    def __init__(self, positions: np.ndarray):
+        """
+        Args:
+            positions: (N, 2) array of actuator coordinates (x, y) in meters.
+        """
+        self.positions = np.array(positions)
+        self.n_actuators = self.positions.shape[0]
+        self.modes: Optional[np.ndarray] = None
+        
+    @abstractmethod
+    def generate(self, n_modes: int, **kwargs) -> np.ndarray:
+        """
+        Generate the basis modes.
+        
+        Args:
+            n_modes: Number of modes to generate.
+            
+        Returns:
+            modes: (n_actuators, n_modes) matrix.
+        """
+        pass
+    
+    def save(self, filepath: str | Path) -> None:
+        """
+        Save the generated basis and actuator positions to a .npz file.
+        """
+        if self.modes is None:
+            raise ValueError("No modes generated yet. Call generate() first.")
+            
+        np.savez(
+            filepath,
+            modes=self.modes,
+            positions=self.positions,
+            basis_type=self.__class__.__name__
+        )
+        
+    @classmethod
+    def load(cls, filepath: str | Path) -> 'BasisGenerator':
+        """
+        Load a basis from a .npz file. 
+        Note: This returns a generic container or re-instantiates the specific class if possible.
+        For simplicity here, we might just return the data or a generic wrapper.
+        """
+        data = np.load(filepath)
+        positions = data['positions']
+        modes = data['modes']
+        
+        # Create a generic instance to hold the data
+        # In a more complex system, we might factory this based on basis_type
+        instance = ConcreteBasis(positions)
+        instance.modes = modes
+        return instance
+
+    def plot(self, count: int = 6, outfile: Optional[str | Path] = None, **kwargs):
+        """Plot the generated modes."""
+        if self.modes is None:
+            raise ValueError("No modes to plot.")
+        plot_basis_modes(self.modes, self.positions, count=count, outfile=outfile, **kwargs)
+
+class ConcreteBasis(BasisGenerator):
+    """Helper class for loading existing bases."""
+    def generate(self, n_modes: int, **kwargs) -> np.ndarray:
+        if self.modes is None:
+            raise NotImplementedError("This is a loaded basis container.")
+        return self.modes[:, :n_modes]