optixstuff 0.0.1__py3-none-any.whl

optixstuff/__init__.py

@@ -0,0 +1,43 @@
+"""optixstuff -- Hardware abstractions for the HWO simulation suite."""
+
+from optixstuff._version import __version__
+from optixstuff.coronagraph import AbstractCoronagraph, AbstractScalarOnlyCoronagraph
+from optixstuff.detector import (
+    AbstractDetector,
+    Detector,
+    SimpleDetector,
+    simulate_cic,
+    simulate_dark_current,
+    simulate_read_noise,
+)
+from optixstuff.optical_elements import (
+    AbstractOpticalElement,
+    AbstractUniformElement,
+    ConstantThroughputElement,
+    LinearThroughputElement,
+    OpticalFilter,
+)
+from optixstuff.optical_path import OpticalPath
+from optixstuff.primary import AbstractPrimary, SimplePrimary
+from optixstuff.yippy_coronagraph import YippyCoronagraph
+
+__all__ = [
+    "AbstractCoronagraph",
+    "AbstractDetector",
+    "AbstractOpticalElement",
+    "AbstractPrimary",
+    "AbstractScalarOnlyCoronagraph",
+    "AbstractUniformElement",
+    "ConstantThroughputElement",
+    "Detector",
+    "LinearThroughputElement",
+    "OpticalFilter",
+    "OpticalPath",
+    "SimpleDetector",
+    "SimplePrimary",
+    "YippyCoronagraph",
+    "__version__",
+    "simulate_cic",
+    "simulate_dark_current",
+    "simulate_read_noise",
+]

optixstuff/_version.py

@@ -0,0 +1,24 @@
+# file generated by vcs-versioning
+# don't change, don't track in version control
+from __future__ import annotations
+
+__all__ = [
+    "__version__",
+    "__version_tuple__",
+    "version",
+    "version_tuple",
+    "__commit_id__",
+    "commit_id",
+]
+
+version: str
+__version__: str
+__version_tuple__: tuple[int | str, ...]
+version_tuple: tuple[int | str, ...]
+commit_id: str | None
+__commit_id__: str | None
+
+__version__ = version = '0.0.1'
+__version_tuple__ = version_tuple = (0, 0, 1)
+
+__commit_id__ = commit_id = None

optixstuff/coronagraph.py

@@ -0,0 +1,193 @@
+"""Coronagraph abstractions."""
+
+import abc
+
+import equinox as eqx
+import jax.numpy as jnp
+from equinox import AbstractVar
+from jax.typing import ArrayLike
+from jaxtyping import Array
+
+
+class AbstractCoronagraph(eqx.Module):
+    """Abstract interface for coronagraph performance models.
+
+    Provides both scalar performance curves (for ETC use) and 2D PSF
+    generation (for image simulation). Implementations can be backed by
+    pre-computed interpolation tables (yippy), physical wavefront
+    propagation, or analytical models.
+
+    All wavelength arguments are in nanometres throughout.
+    All separations are in lambda/D units.
+    """
+
+    pixel_scale_lod: AbstractVar[float]
+    """Native pixel scale in lambda/D per pixel."""
+
+    IWA: AbstractVar[float]
+    """Inner working angle in lambda/D."""
+
+    OWA: AbstractVar[float]
+    """Outer working angle in lambda/D."""
+
+    # ------------------------------------------------------------------
+    # Scalar interface -- consumed by jaxEDITH and yield estimators
+    # ------------------------------------------------------------------
+
+    @abc.abstractmethod
+    def throughput(
+        self,
+        separation_lod: ArrayLike,
+        wavelength_nm: ArrayLike,
+        *,
+        time_s: ArrayLike = 0.0,
+    ) -> ArrayLike:
+        """Core (off-axis planet) throughput.
+
+        Args:
+            separation_lod: Angular separation in lambda/D.
+            wavelength_nm: Wavelength in nanometres.
+            time_s: Time since mission start in seconds.
+
+        Returns:
+            Fractional throughput in [0, 1].
+        """
+        ...
+
+    @abc.abstractmethod
+    def core_area(
+        self,
+        separation_lod: ArrayLike,
+        wavelength_nm: ArrayLike,
+        *,
+        time_s: ArrayLike = 0.0,
+    ) -> ArrayLike:
+        """Photometric aperture area in (lambda/D)^2.
+
+        Args:
+            separation_lod: Angular separation in lambda/D.
+            wavelength_nm: Wavelength in nanometres.
+            time_s: Time since mission start in seconds.
+
+        Returns:
+            Core area in (lambda/D)^2.
+        """
+        ...
+
+    @abc.abstractmethod
+    def core_mean_intensity(
+        self,
+        separation_lod: ArrayLike,
+        wavelength_nm: ArrayLike,
+        *,
+        time_s: ArrayLike = 0.0,
+    ) -> ArrayLike:
+        """Mean stellar intensity within the photometric aperture.
+
+        Args:
+            separation_lod: Angular separation in lambda/D.
+            wavelength_nm: Wavelength in nanometres.
+            time_s: Time since mission start in seconds.
+
+        Returns:
+            Mean stellar leakage intensity in (lambda/D)^-2.
+        """
+        ...
+
+    @abc.abstractmethod
+    def occulter_transmission(
+        self,
+        separation_lod: ArrayLike,
+        wavelength_nm: ArrayLike,
+        *,
+        time_s: ArrayLike = 0.0,
+    ) -> ArrayLike:
+        """Off-axis (sky/zodi) transmission through the occulter.
+
+        Args:
+            separation_lod: Angular separation in lambda/D.
+            wavelength_nm: Wavelength in nanometres.
+            time_s: Time since mission start in seconds.
+
+        Returns:
+            Fractional sky transmission in [0, 1].
+        """
+        ...
+
+    # ------------------------------------------------------------------
+    # Image interface -- consumed by coronagraphoto
+    # ------------------------------------------------------------------
+
+    @abc.abstractmethod
+    def on_axis_psf(
+        self,
+        wavelength_nm: ArrayLike,
+        pixel_scale_rad: float,
+        npixels: int,
+    ) -> Array:
+        """On-axis (stellar leakage) PSF.
+
+        Returns the coronagraphic PSF for an on-axis point source,
+        normalized to unit stellar flux before the coronagraph.
+
+        Args:
+            wavelength_nm: Wavelength in nanometres.
+            pixel_scale_rad: Output pixel scale in radians/pixel.
+            npixels: Output array side length in pixels. Must be a
+                Python int (not a JAX array) as it determines the
+                output shape at compile time.
+
+        Returns:
+            2D float array of shape (npixels, npixels).
+        """
+        ...
+
+    @abc.abstractmethod
+    def off_axis_psf(
+        self,
+        wavelength_nm: ArrayLike,
+        separation_lod: ArrayLike,
+        pixel_scale_rad: float,
+        npixels: int,
+    ) -> Array:
+        """Off-axis PSF at a given angular separation.
+
+        Args:
+            wavelength_nm: Wavelength in nanometres.
+            separation_lod: Source separation in lambda/D.
+            pixel_scale_rad: Output pixel scale in radians/pixel.
+            npixels: Output array side length in pixels. Must be a
+                Python int (not a JAX array) as it determines the
+                output shape at compile time.
+
+        Returns:
+            2D float array of shape (npixels, npixels).
+        """
+        ...
+
+
+class AbstractScalarOnlyCoronagraph(AbstractCoronagraph):
+    """Base for ETC-only coronagraph models that lack 2D PSF generation.
+
+    Stubs out the image interface with zero arrays so the class satisfies
+    AbstractCoronagraph without requiring a full optical model.
+    """
+
+    def on_axis_psf(
+        self,
+        wavelength_nm: ArrayLike,
+        pixel_scale_rad: float,
+        npixels: int,
+    ) -> Array:
+        """Return a zero PSF (not implemented for scalar-only models)."""
+        return jnp.zeros((npixels, npixels))
+
+    def off_axis_psf(
+        self,
+        wavelength_nm: ArrayLike,
+        separation_lod: ArrayLike,
+        pixel_scale_rad: float,
+        npixels: int,
+    ) -> Array:
+        """Return a zero PSF (not implemented for scalar-only models)."""
+        return jnp.zeros((npixels, npixels))

optixstuff/detector.py

@@ -0,0 +1,324 @@
+"""Detector abstractions and concrete implementations."""
+
+
+import abc
+from typing import final
+
+import equinox as eqx
+import jax
+import jax.numpy as jnp
+from equinox import AbstractVar
+from jax.typing import ArrayLike
+from jaxtyping import Array
+
+
+class AbstractDetector(eqx.Module):
+    """Abstract interface for a focal-plane detector.
+
+    Provides both scalar noise rates (for ETC use) and stochastic noise
+    realization (for image simulation).  All concrete implementations
+    must define the hardware parameters listed as ``AbstractVar`` fields.
+    """
+
+    pixel_scale: AbstractVar[float]
+    """Detector plate scale in arcsec/pixel."""
+
+    quantum_efficiency: AbstractVar[float]
+    """Baseline quantum efficiency as a fraction in [0, 1]."""
+
+    dark_current_rate: AbstractVar[float]
+    """Dark current rate in electrons/pixel/second."""
+
+    read_noise_electrons: AbstractVar[float]
+    """Read noise in electrons RMS per pixel per read."""
+
+    cic_rate: AbstractVar[float]
+    """Clock-induced charge in electrons/pixel/frame."""
+
+    frame_time: AbstractVar[float]
+    """Integration time per frame/read in seconds."""
+
+    read_time: AbstractVar[float]
+    """Time per read cycle in seconds (for RN^2/t_read in ETC)."""
+
+    dqe: AbstractVar[float]
+    """QE degradation factor (multiplicative correction over mission life)."""
+
+    shape: AbstractVar[tuple[int, int]]
+    """Detector dimensions (ny, nx) in pixels."""
+
+    @abc.abstractmethod
+    def get_qe(self, wavelength_nm: ArrayLike) -> ArrayLike:
+        """Quantum efficiency at a given wavelength.
+
+        Args:
+            wavelength_nm: Wavelength in nanometres.
+
+        Returns:
+            QE as a fraction in [0, 1].
+        """
+        ...
+
+    @abc.abstractmethod
+    def scalar_noise_rate(self, n_pix: ArrayLike, t_photon: ArrayLike) -> ArrayLike:
+        """Total scalar noise variance rate for the ETC.
+
+        Returns the combined noise variance per unit time (electrons^2/s)
+        for a photometric aperture of n_pix pixels.
+
+        Args:
+            n_pix: Number of pixels in the photometric aperture.
+            t_photon: Photon counting integration time in seconds.
+
+        Returns:
+            Noise variance rate in electrons^2/second.
+        """
+        ...
+
+    @abc.abstractmethod
+    def add_noise(
+        self,
+        image_rate: Array,
+        exposure_time: ArrayLike,
+        prng_key: Array,
+    ) -> Array:
+        """Apply stochastic noise realization to a photon rate image.
+
+        Converts photon rates to detected electrons including QE,
+        dark current, CIC, and read noise.
+
+        Args:
+            image_rate: Incident photon rate array in ph/s/pixel.
+            exposure_time: Exposure time in seconds.
+            prng_key: JAX PRNG key (required, no default).
+
+        Returns:
+            Detected electrons array, same shape as image_rate.
+        """
+        ...
+
+
+# -- Pure noise simulation functions -----------------------------------------
+
+
+def simulate_dark_current(
+    dark_current_rate: float,
+    exposure_time: ArrayLike,
+    shape: tuple[int, int],
+    prng_key: Array,
+) -> Array:
+    """Draw dark current electrons from a Poisson distribution.
+
+    Args:
+        dark_current_rate: Dark current rate in electrons/s/pixel.
+        exposure_time: Exposure time in seconds.
+        shape: Detector shape (ny, nx).
+        prng_key: PRNG key.
+
+    Returns:
+        Dark current electrons, shape (ny, nx).
+    """
+    return jax.random.poisson(prng_key, dark_current_rate * exposure_time, shape=shape)
+
+
+def simulate_cic(
+    cic_rate: float,
+    num_frames: ArrayLike,
+    shape: tuple[int, int],
+    prng_key: Array,
+) -> Array:
+    """Draw clock-induced charge electrons from a Poisson distribution.
+
+    Args:
+        cic_rate: CIC rate in electrons/pixel/frame.
+        num_frames: Number of frames (kept as float for JIT safety).
+        shape: Detector shape (ny, nx).
+        prng_key: PRNG key.
+
+    Returns:
+        CIC electrons, shape (ny, nx).
+    """
+    return jax.random.poisson(prng_key, cic_rate * num_frames, shape=shape)
+
+
+def simulate_read_noise(
+    read_noise: float,
+    num_frames: ArrayLike,
+    shape: tuple[int, int],
+    prng_key: Array,
+) -> Array:
+    """Draw read noise from a Gaussian distribution.
+
+    Total read noise sigma = sqrt(num_frames) * read_noise_per_read.
+
+    Args:
+        read_noise: Read noise in electrons/pixel/read.
+        num_frames: Number of frames.
+        shape: Detector shape (ny, nx).
+        prng_key: PRNG key.
+
+    Returns:
+        Read noise electrons, shape (ny, nx).
+    """
+    sigma = read_noise * jnp.sqrt(num_frames)
+    return sigma * jax.random.normal(prng_key, shape=shape)
+
+
+# -- Concrete implementations -----------------------------------------------
+
+
+@final
+class SimpleDetector(AbstractDetector):
+    """Detector with constant QE and minimal noise sources.
+
+    Suitable for broadband imager studies where wavelength-dependent
+    QE variation is not important and CIC/read noise are negligible.
+    """
+
+    pixel_scale: float
+    quantum_efficiency: float
+    dark_current_rate: float
+    read_noise_electrons: float
+    cic_rate: float
+    frame_time: float
+    read_time: float
+    dqe: float
+    shape: tuple[int, int] = eqx.field(static=True)
+
+    def __init__(
+        self,
+        pixel_scale: float,
+        shape: tuple[int, int],
+        quantum_efficiency: float = 1.0,
+        dark_current_rate: float = 0.0,
+        read_noise_electrons: float = 0.0,
+        cic_rate: float = 0.0,
+        frame_time: float = 1.0,
+        read_time: float = 0.05,
+        dqe: float = 0.0,
+    ) -> None:
+        """Create a simple constant-QE detector."""
+        self.pixel_scale = pixel_scale
+        self.shape = shape
+        self.quantum_efficiency = quantum_efficiency
+        self.dark_current_rate = dark_current_rate
+        self.read_noise_electrons = read_noise_electrons
+        self.cic_rate = cic_rate
+        self.frame_time = frame_time
+        self.read_time = read_time
+        self.dqe = dqe
+
+    def get_qe(self, wavelength_nm: ArrayLike) -> ArrayLike:
+        """Return constant QE, ignoring wavelength."""
+        return self.quantum_efficiency
+
+    def scalar_noise_rate(self, n_pix: ArrayLike, t_photon: ArrayLike) -> ArrayLike:
+        """Combined dark + CIC noise variance rate.
+
+        Read noise is not included here as it scales per-read, not per-second.
+        Callers add (read_noise^2 * n_reads) / t_exp separately.
+        """
+        dark_variance_rate = self.dark_current_rate * n_pix
+        cic_variance_rate = self.cic_rate * n_pix / t_photon
+        return dark_variance_rate + cic_variance_rate
+
+    def add_noise(
+        self,
+        image_rate: Array,
+        exposure_time: ArrayLike,
+        prng_key: Array,
+    ) -> Array:
+        """Apply dark current noise to a photon rate image."""
+        key_phot, key_qe, key_dark = jax.random.split(prng_key, 3)
+
+        inc_photons = jax.random.poisson(key_phot, image_rate * exposure_time)
+        qe = self.quantum_efficiency
+        photo_electrons = jax.random.binomial(key_qe, inc_photons, qe)
+
+        dark = simulate_dark_current(
+            self.dark_current_rate, exposure_time, self.shape, key_dark
+        )
+        return photo_electrons + dark
+
+
+@final
+class Detector(AbstractDetector):
+    """Full detector model with dark current, CIC, and read noise.
+
+    Suitable for detailed noise simulations where all detector noise
+    sources matter.  Uses Poisson statistics for dark/CIC and Gaussian
+    for read noise, matching the coronagraphoto convention.
+
+    Warning: ``num_frames = jnp.ceil(exposure_time / frame_time)`` is
+    kept as a float. Never cast it to int inside JIT -- that triggers a
+    ConcretizationTypeError when exposure_time is traced.
+    """
+
+    pixel_scale: float
+    quantum_efficiency: float
+    dark_current_rate: float
+    read_noise_electrons: float
+    cic_rate: float
+    frame_time: float
+    read_time: float
+    dqe: float
+    shape: tuple[int, int] = eqx.field(static=True)
+
+    def __init__(
+        self,
+        pixel_scale: float,
+        shape: tuple[int, int],
+        quantum_efficiency: float = 1.0,
+        dark_current_rate: float = 0.0,
+        read_noise_electrons: float = 0.0,
+        cic_rate: float = 0.0,
+        frame_time: float = 1.0,
+        read_time: float = 0.05,
+        dqe: float = 0.0,
+    ) -> None:
+        """Create a full detector with all noise sources."""
+        self.pixel_scale = pixel_scale
+        self.shape = shape
+        self.quantum_efficiency = quantum_efficiency
+        self.dark_current_rate = dark_current_rate
+        self.read_noise_electrons = read_noise_electrons
+        self.cic_rate = cic_rate
+        self.frame_time = frame_time
+        self.read_time = read_time
+        self.dqe = dqe
+
+    def get_qe(self, wavelength_nm: ArrayLike) -> ArrayLike:
+        """Return constant QE, ignoring wavelength."""
+        return self.quantum_efficiency
+
+    def scalar_noise_rate(self, n_pix: ArrayLike, t_photon: ArrayLike) -> ArrayLike:
+        """Combined dark + CIC noise variance rate."""
+        dark_variance_rate = self.dark_current_rate * n_pix
+        cic_variance_rate = self.cic_rate * n_pix / t_photon
+        return dark_variance_rate + cic_variance_rate
+
+    def add_noise(
+        self,
+        image_rate: Array,
+        exposure_time: ArrayLike,
+        prng_key: Array,
+    ) -> Array:
+        """Apply all noise sources: QE, dark current, CIC, and read noise."""
+        key_phot, key_qe, key_dark, key_cic, key_read = jax.random.split(prng_key, 5)
+
+        inc_photons = jax.random.poisson(key_phot, image_rate * exposure_time)
+        qe = self.quantum_efficiency
+        photo_electrons = jax.random.binomial(key_qe, inc_photons, qe)
+
+        dark = simulate_dark_current(
+            self.dark_current_rate, exposure_time, self.shape, key_dark
+        )
+
+        # num_frames stays as a traced float -- never cast to int
+        num_frames = jnp.ceil(exposure_time / self.frame_time)
+        cic = simulate_cic(self.cic_rate, num_frames, self.shape, key_cic)
+        read = simulate_read_noise(
+            self.read_noise_electrons, num_frames, self.shape, key_read
+        )
+
+        return photo_electrons + dark + cic + read

optixstuff/optical_elements.py

@@ -0,0 +1,132 @@
+"""Optical element abstractions (throughput, filters, field stops)."""
+
+
+import abc
+from typing import final
+
+import equinox as eqx
+import interpax
+import jax.numpy as jnp
+from jax.typing import ArrayLike
+from jaxtyping import Array
+
+
+class AbstractOpticalElement(eqx.Module):
+    """Abstract interface for an optical element in the beam path.
+
+    Elements reduce photon flux via wavelength-dependent throughput.
+    The ETC calls get_throughput() for scalar efficiency calculations.
+    The simulator calls apply() to attenuate 2D photon arrays.
+
+    Both methods are abstract: use AbstractUniformElement for elements
+    with spatially uniform throughput, which provides a default apply().
+    """
+
+    @abc.abstractmethod
+    def get_throughput(self, wavelength_nm: ArrayLike) -> ArrayLike:
+        """Fractional throughput at a given wavelength.
+
+        Args:
+            wavelength_nm: Wavelength in nanometres.
+
+        Returns:
+            Scalar throughput in [0, 1].
+        """
+        ...
+
+    @abc.abstractmethod
+    def apply(self, arr: Array, wavelength_nm: ArrayLike) -> Array:
+        """Apply this element to a 2D photon array.
+
+        Args:
+            arr: Input photon rate array [ph/s/pixel].
+            wavelength_nm: Wavelength in nanometres.
+
+        Returns:
+            Attenuated photon rate array, same shape as arr.
+        """
+        ...
+
+
+class AbstractUniformElement(AbstractOpticalElement):
+    """Base for elements with spatially uniform throughput.
+
+    Provides a default apply() that multiplies the array by the scalar
+    throughput. Override apply() for elements with spatially varying
+    transmission (e.g., field-dependent filter transmission maps).
+    """
+
+    def apply(self, arr: Array, wavelength_nm: ArrayLike) -> Array:
+        """Apply uniform throughput to a photon array."""
+        return arr * self.get_throughput(wavelength_nm)
+
+
+@final
+class ConstantThroughputElement(AbstractUniformElement):
+    """An optical element with wavelength-independent throughput.
+
+    Useful for modeling simple attenuators, beamsplitters, or as a
+    placeholder during instrument design studies.
+    """
+
+    throughput: float
+    name: str = eqx.field(default="element", static=True)
+
+    def get_throughput(self, wavelength_nm: ArrayLike) -> ArrayLike:
+        """Return constant throughput, ignoring wavelength."""
+        return self.throughput
+
+
+@final
+class LinearThroughputElement(AbstractUniformElement):
+    """An optical element with linearly interpolated wavelength-dependent throughput.
+
+    Throughput is specified at a set of wavelengths and linearly
+    interpolated between them. Extrapolation returns zero outside the
+    defined wavelength range.
+    """
+
+    wavelengths_nm: Array
+    throughputs: Array
+    interp: interpax.Interpolator1D
+
+    def __init__(self, wavelengths_nm: Array, throughputs: Array) -> None:
+        """Create a throughput element from sampled wavelength/throughput pairs."""
+        self.wavelengths_nm = wavelengths_nm
+        self.throughputs = throughputs
+        self.interp = interpax.Interpolator1D(
+            wavelengths_nm, throughputs, method="linear", extrap=jnp.array([0.0, 0.0])
+        )
+
+    def get_throughput(self, wavelength_nm: ArrayLike) -> ArrayLike:
+        """Interpolate throughput at the requested wavelength."""
+        return self.interp(wavelength_nm)
+
+
+@final
+class OpticalFilter(AbstractUniformElement):
+    """A bandpass filter with linearly interpolated transmittance.
+
+    Structurally identical to LinearThroughputElement but semantically
+    distinct -- represents a spectral bandpass selection rather than a
+    reflective coating or attenuator.
+    """
+
+    wavelengths_nm: Array
+    transmittances: Array
+    interp: interpax.Interpolator1D
+
+    def __init__(self, wavelengths_nm: Array, transmittances: Array) -> None:
+        """Create an optical filter from sampled wavelength/transmittance pairs."""
+        self.wavelengths_nm = wavelengths_nm
+        self.transmittances = transmittances
+        self.interp = interpax.Interpolator1D(
+            wavelengths_nm,
+            transmittances,
+            method="linear",
+            extrap=jnp.array([0.0, 0.0]),
+        )
+
+    def get_throughput(self, wavelength_nm: ArrayLike) -> ArrayLike:
+        """Interpolate filter transmittance at the requested wavelength."""
+        return self.interp(wavelength_nm)

optixstuff/optical_path.py

@@ -0,0 +1,47 @@
+"""OpticalPath container -- the universal hardware configuration."""
+
+import functools
+
+import equinox as eqx
+
+from optixstuff.coronagraph import AbstractCoronagraph
+from optixstuff.detector import AbstractDetector
+from optixstuff.optical_elements import AbstractOpticalElement
+from optixstuff.primary import AbstractPrimary
+
+
+class OpticalPath(eqx.Module):
+    """Universal hardware container for a coronagraphic telescope.
+
+    Bundles a primary mirror, ordered chain of attenuating elements,
+    a coronagraph, and a detector into a single configuration object.
+    This is the interface passed to simulators (coronagraphoto),
+    exposure time calculators (jaxEDITH), and IFS instruments (coronachrome).
+
+    Args:
+        primary: Primary mirror description.
+        attenuating_elements: Ordered tuple of throughput elements
+            between the primary and coronagraph (mirrors, filters, etc.).
+        coronagraph: Coronagraph performance model.
+        detector: Focal-plane detector model.
+    """
+
+    primary: AbstractPrimary
+    attenuating_elements: tuple[AbstractOpticalElement, ...]
+    coronagraph: AbstractCoronagraph
+    detector: AbstractDetector
+
+    def system_throughput(self, wavelength_nm: float) -> float:
+        """Total throughput of all attenuating elements.
+
+        Args:
+            wavelength_nm: Wavelength in nanometres.
+
+        Returns:
+            Combined fractional throughput in [0, 1].
+        """
+        return functools.reduce(
+            lambda acc, el: acc * el.get_throughput(wavelength_nm),
+            self.attenuating_elements,
+            1.0,
+        )

optixstuff/primary.py

@@ -0,0 +1,58 @@
+"""Primary mirror abstractions."""
+
+import equinox as eqx
+import jax.numpy as jnp
+from equinox import AbstractVar
+
+
+class AbstractPrimary(eqx.Module):
+    """Abstract interface for a primary aperture.
+
+    Any concrete implementation must provide the diameter and collecting
+    area of the primary mirror as scalar values in SI units. These are
+    consumed by exposure time calculators and simulation tools alike.
+    """
+
+    diameter_m: AbstractVar[float]
+    """Primary mirror diameter in metres."""
+
+    area_m2: AbstractVar[float]
+    """Effective collecting area in square metres."""
+
+
+class SimplePrimary(AbstractPrimary):
+    """A simple circular primary mirror with a central obscuration.
+
+    Args:
+        diameter_m: Primary mirror diameter in metres.
+        obscuration: Linear obscuration fraction (0 = no obscuration).
+        shape_factor: Fraction of unobscured area that is collecting
+            (accounts for struts, segment gaps, etc.). Default 1.0.
+    """
+
+    _diameter_m: float
+    obscuration: float
+    shape_factor: float
+
+    def __init__(
+        self,
+        diameter_m: float,
+        obscuration: float = 0.0,
+        shape_factor: float = 1.0,
+    ) -> None:
+        """Create a simple circular primary mirror."""
+        self._diameter_m = diameter_m
+        self.obscuration = obscuration
+        self.shape_factor = shape_factor
+
+    @property
+    def diameter_m(self) -> float:
+        """Primary mirror diameter in metres."""
+        return self._diameter_m
+
+    @property
+    def area_m2(self) -> float:
+        """Effective collecting area in square metres."""
+        r = self._diameter_m / 2.0
+        gross_area = jnp.pi * r**2
+        return gross_area * (1.0 - self.obscuration**2) * self.shape_factor

optixstuff/yippy_coronagraph.py

@@ -0,0 +1,167 @@
+"""YippyCoronagraph -- AbstractCoronagraph backed by a yippy EqxCoronagraph."""
+
+
+from typing import final
+
+from jax.typing import ArrayLike
+from jaxtyping import Array
+from yippy import EqxCoronagraph
+
+from optixstuff.coronagraph import AbstractCoronagraph
+
+
+@final
+class YippyCoronagraph(AbstractCoronagraph):
+    """Coronagraph performance model backed by a yippy YIP interpolation table.
+
+    Wraps a yippy ``EqxCoronagraph`` via composition, adapting its methods
+    to the ``AbstractCoronagraph`` interface.  The ``_backend`` field is
+    itself an ``eqx.Module``, so its internal JAX arrays flow through
+    ``filter_jit`` and ``filter_grad`` normally.
+
+    Construction mirrors ``EqxCoronagraph`` -- pass either a YIP path or
+    an existing ``EqxCoronagraph`` instance::
+
+        coro = YippyCoronagraph("/path/to/yip")
+        coro = YippyCoronagraph(backend=existing_eqx_coro)
+    """
+
+    _backend: EqxCoronagraph
+
+    def __init__(
+        self,
+        yip_path: str | None = None,
+        *,
+        backend: EqxCoronagraph | None = None,
+        **kwargs,
+    ) -> None:
+        """Create a YippyCoronagraph from a YIP path or existing backend.
+
+        Args:
+            yip_path: Path to a Yield Input Package directory.
+            backend: Pre-built EqxCoronagraph.  Takes precedence over yip_path.
+            **kwargs: Forwarded to ``EqxCoronagraph`` when building from yip_path.
+        """
+        if backend is not None:
+            self._backend = backend
+        elif yip_path is not None:
+            self._backend = EqxCoronagraph(yip_path, **kwargs)
+        else:
+            msg = "Provide either yip_path or backend"
+            raise ValueError(msg)
+
+    # -- AbstractVar fields satisfied via properties ----------------------
+
+    @property
+    def pixel_scale_lod(self) -> float:
+        """Native pixel scale in lambda/D per pixel."""
+        return self._backend.pixel_scale_lod
+
+    @property
+    def IWA(self) -> float:
+        """Inner working angle in lambda/D."""
+        return self._backend.IWA
+
+    @property
+    def OWA(self) -> float:
+        """Outer working angle in lambda/D."""
+        return self._backend.OWA
+
+    # -- Scalar interface -------------------------------------------------
+
+    def throughput(
+        self,
+        separation_lod: ArrayLike,
+        wavelength_nm: ArrayLike,
+        *,
+        time_s: ArrayLike = 0.0,
+    ) -> ArrayLike:
+        """Core throughput from the YIP interpolation table."""
+        return self._backend.throughput(separation_lod)
+
+    def core_area(
+        self,
+        separation_lod: ArrayLike,
+        wavelength_nm: ArrayLike,
+        *,
+        time_s: ArrayLike = 0.0,
+    ) -> ArrayLike:
+        """Photometric aperture area from the YIP interpolation table."""
+        return self._backend.core_area(separation_lod)
+
+    def core_mean_intensity(
+        self,
+        separation_lod: ArrayLike,
+        wavelength_nm: ArrayLike,
+        *,
+        time_s: ArrayLike = 0.0,
+    ) -> ArrayLike:
+        """Mean stellar leakage from the YIP interpolation table."""
+        return self._backend.core_mean_intensity(separation_lod)
+
+    def occulter_transmission(
+        self,
+        separation_lod: ArrayLike,
+        wavelength_nm: ArrayLike,
+        *,
+        time_s: ArrayLike = 0.0,
+    ) -> ArrayLike:
+        """Sky transmission from the YIP interpolation table."""
+        return self._backend.occulter_transmission(separation_lod)
+
+    # -- Image interface --------------------------------------------------
+
+    def on_axis_psf(
+        self,
+        wavelength_nm: ArrayLike,
+        pixel_scale_rad: float,
+        npixels: int,
+    ) -> Array:
+        """Stellar leakage PSF from the YIP stellar intensity model."""
+        return self._backend.stellar_intens(0.0)
+
+    def off_axis_psf(
+        self,
+        wavelength_nm: ArrayLike,
+        separation_lod: ArrayLike,
+        pixel_scale_rad: float,
+        npixels: int,
+    ) -> Array:
+        """Off-axis planet PSF from the YIP PSF interpolator.
+
+        Places the planet along the +x axis by convention.
+        """
+        return self._backend.create_psf(separation_lod, 0.0, npixels)
+
+    # -- Convenience methods (not on AbstractCoronagraph) ------------------
+
+    def noise_floor_ayo(
+        self,
+        separation_lod: ArrayLike,
+        ppf: float = 30.0,
+    ) -> ArrayLike:
+        """AYO noise floor: core_mean_intensity / ppf.
+
+        This is a convenience passthrough to the backend. Not part of the
+        AbstractCoronagraph contract -- downstream ETCs should compute
+        noise floors as pure functions.
+        """
+        return self._backend.noise_floor_ayo(separation_lod, ppf)
+
+    def raw_contrast(self, separation_lod: ArrayLike) -> ArrayLike:
+        """Raw contrast from the YIP interpolation table."""
+        return self._backend.raw_contrast(separation_lod)
+
+    def stellar_intens(self, stellar_diam_lod: float) -> Array:
+        """Stellar intensity map for a given stellar angular diameter."""
+        return self._backend.stellar_intens(stellar_diam_lod)
+
+    @property
+    def psf_shape(self) -> tuple[int, int]:
+        """Shape of the PSF arrays from the YIP file."""
+        return self._backend.psf_shape
+
+    @property
+    def sky_trans(self) -> Array:
+        """Full sky transmission map."""
+        return self._backend.sky_trans

optixstuff-0.0.1.dist-info/METADATA

@@ -0,0 +1,136 @@
+Metadata-Version: 2.4
+Name: optixstuff
+Version: 0.0.1
+Summary: Hardware abstractions for the HWO direct imaging simulation suite
+Project-URL: Homepage, https://github.com/CoreySpohn/optixstuff
+Project-URL: Issues, https://github.com/CoreySpohn/optixstuff/issues
+Author-email: Corey Spohn <corey.a.spohn@nasa.gov>
+License: MIT License
+        
+        Copyright (c) 2026 Corey Spohn
+        
+        Permission is hereby granted, free of charge, to any person obtaining a copy
+        of this software and associated documentation files (the "Software"), to deal
+        in the Software without restriction, including without limitation the rights
+        to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+        copies of the Software, and to permit persons to whom the Software is
+        furnished to do so, subject to the following conditions:
+        
+        The above copyright notice and this permission notice shall be included in all
+        copies or substantial portions of the Software.
+        
+        THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+        IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+        FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+        AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+        LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+        OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+        SOFTWARE.
+License-File: LICENSE
+Classifier: Development Status :: 2 - Pre-Alpha
+Classifier: Intended Audience :: Science/Research
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Programming Language :: Python :: 3
+Classifier: Topic :: Scientific/Engineering :: Astronomy
+Requires-Python: >=3.11
+Requires-Dist: equinox>=0.12.0
+Requires-Dist: interpax
+Requires-Dist: jax>=0.4.1
+Requires-Dist: jaxlib>=0.4.1
+Requires-Dist: yippy
+Provides-Extra: dev
+Requires-Dist: pre-commit; extra == 'dev'
+Provides-Extra: docs
+Requires-Dist: ipython; extra == 'docs'
+Requires-Dist: matplotlib; extra == 'docs'
+Requires-Dist: myst-nb; extra == 'docs'
+Requires-Dist: sphinx; extra == 'docs'
+Requires-Dist: sphinx-autoapi; extra == 'docs'
+Requires-Dist: sphinx-autodoc-typehints; extra == 'docs'
+Requires-Dist: sphinx-book-theme; extra == 'docs'
+Provides-Extra: test
+Requires-Dist: hypothesis; extra == 'test'
+Requires-Dist: nox; extra == 'test'
+Requires-Dist: pytest; extra == 'test'
+Requires-Dist: pytest-cov; extra == 'test'
+Description-Content-Type: text/markdown
+
+# optixstuff
+
+Flux-level optical system abstractions for the HWO direct imaging simulation suite.
+
+## What optixstuff is
+
+`optixstuff` is the **radiometric hardware layer** for HWO simulations. It operates on
+**flux** — photon rates, throughput fractions, and detector-level noise — providing the
+shared interfaces that simulation tools and exposure time calculators build on top of.
+
+The same `OpticalPath` object drives both scalar ETC calculations (`jaxEDITH`) and full
+2D image generation (`coronagraphoto`), ensuring that the hardware model is consistent
+across all downstream science products.
+
+## What optixstuff is *not*
+
+`optixstuff` does not model diffraction, wavefront propagation, or E-field interference.
+That level of physical optics belongs to tools like [dLux](https://github.com/LouisDesdoigts/dLux)
+and [HCIPy](https://github.com/ehpor/hcipy), which generate PSFs from first principles.
+
+`optixstuff` and these wavefront tools are **complementary**: dLux/HCIPy compute the PSFs,
+which are delivered as yield input packages (YIPs). `optixstuff` consumes those PSFs as
+flux patterns (via [yippy](https://github.com/CoreySpohn/yippy)) and composes them with
+the rest of the observatory — throughput chain, detector QE, noise — to produce
+science-level outputs.
+
+## Architecture
+
+Built on [JAX](https://github.com/google/jax) and
+[Equinox](https://github.com/patrick-kidger/equinox), `optixstuff` provides:
+
+- **Abstract interfaces** — `AbstractPrimary`, `AbstractOpticalElement`,
+  `AbstractCoronagraph`, `AbstractDetector`
+- **Concrete implementations** — `SimplePrimary`, `ConstantThroughputElement`,
+  `SimpleDetector`
+- **Container** — `OpticalPath`, a composable hardware configuration passed to all
+  simulators
+
+Every abstract method accepts three fidelity axes — **wavelength**, **position**, and
+**time** — with defaults so that simple implementations can ignore unused axes while
+future high-fidelity models (wavelength-dependent coatings, position-dependent vignetting,
+time-dependent detector degradation) can use them without breaking the interface.
+
+### Ecosystem position
+
+```
+                          ┌───────────────────────────┐
+                          │  Physical optics           │
+                          │  (dLux, HCIPy, PROPER)     │
+                          │  E-fields → PSFs           │
+                          └──────────┬────────────────┘
+                                     │ YIP
+                          ┌──────────▼────────────────┐
+                          │  yippy                     │
+                          │  PSF interpolation         │
+                          └──────────┬────────────────┘
+                                     │ flux patterns
+              ┌──────────────────────▼────────────────────────┐
+              │                  optixstuff                    │
+              │  Telescope • Coronagraph • Detector • OpticalPath │
+              │  Throughput chains • QE • Noise rates          │
+              └────────┬──────────────────────┬───────────────┘
+                       │                      │
+            ┌──────────▼──────────┐ ┌────────▼───────────────┐
+            │  jaxEDITH           │ │  coronagraphoto         │
+            │  Scalar count rates │ │  2D image simulation    │
+            │  Exposure times     │ │  Multi-epoch scenes     │
+            └─────────────────────┘ └────────────────────────┘
+```
+
+## Installation
+
+```bash
+pip install optixstuff
+```
+
+## Status
+
+This package is in early development (pre-v0.1.0).

optixstuff-0.0.1.dist-info/RECORD

@@ -0,0 +1,12 @@
+optixstuff/__init__.py,sha256=oAuCgjZZA2Ja0sYqLR06-Y5i9ZUzUAT5SYuPMCp_r2k,1180
+optixstuff/_version.py,sha256=8OsTLsIVB9D0HdPTmt5rVwyVUBe9xTVGkRslXicxzkM,520
+optixstuff/coronagraph.py,sha256=Y_L47A8tUnjthqv1riWwCSruy1jfpqOIKPqkBHQOCTc,5728
+optixstuff/detector.py,sha256=md-r0HWd4PKGC4F8GqfBwoVg2_7wPcFx-eLumzsyluQ,10275
+optixstuff/optical_elements.py,sha256=czPjlBb5bPWjUO0s6Tq_Mc6V3qIJprLnUudTpeqaWuc,4391
+optixstuff/optical_path.py,sha256=_5fE7saX3DXoJsqrToRkzGosCFaBvjKKQ9Wg6dO2ltg,1611
+optixstuff/primary.py,sha256=gqzR5tF37A2dghItQR4IsfkKPcmkp2CyhElcwtImo-s,1752
+optixstuff/yippy_coronagraph.py,sha256=RPGRvW1ECv7I5yeLK3dIzIzBTaxFr7dHfzk8mXTrfiE,5355
+optixstuff-0.0.1.dist-info/METADATA,sha256=DzsNtkE2KB0MLv3STb2ZDE-RjvhYNd7KwhRR-cdiRTY,6846
+optixstuff-0.0.1.dist-info/WHEEL,sha256=QccIxa26bgl1E6uMy58deGWi-0aeIkkangHcxk2kWfw,87
+optixstuff-0.0.1.dist-info/licenses/LICENSE,sha256=66ed08OMdjt-G3WoW0mfxmaAt-zqSwZ2A0kvPklkwwQ,1068
+optixstuff-0.0.1.dist-info/RECORD,,

optixstuff-0.0.1.dist-info/WHEEL

@@ -0,0 +1,4 @@
+Wheel-Version: 1.0
+Generator: hatchling 1.29.0
+Root-Is-Purelib: true
+Tag: py3-none-any

optixstuff-0.0.1.dist-info/licenses/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 Corey Spohn
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.