optixstuff 1.0.2__tar.gz → 1.2.0__tar.gz

optixstuff-1.2.0/.pre-commit-config.yaml

@@ -0,0 +1,20 @@
+repos:
+  - repo: https://github.com/pre-commit/pre-commit-hooks
+    rev: "v6.0.0"
+    hooks:
+      - id: trailing-whitespace
+      - id: name-tests-test
+        args: [--pytest-test-first]
+      - id: end-of-file-fixer
+  - repo: https://github.com/astral-sh/ruff-pre-commit
+    rev: v0.15.2
+    hooks:
+      - id: ruff-check
+        args: [--fix]
+      - id: ruff-format
+  - repo: https://github.com/compilerla/conventional-pre-commit
+    rev: v4.3.0
+    hooks:
+      - id: conventional-pre-commit
+        stages: [commit-msg]
+        args: []

{optixstuff-1.0.2 → optixstuff-1.2.0}/CHANGELOG.md

@@ -1,5 +1,22 @@
 # Changelog
 
+## [1.2.0](https://github.com/CoreySpohn/optixstuff/compare/v1.1.0...v1.2.0) (2026-06-18)
+
+
+### Features
+
+* Add reference wavelength for the psflets ([a2fe241](https://github.com/CoreySpohn/optixstuff/commit/a2fe241ee778ef54c4832e68950fe5268f8a8b1c))
+* **disperser:** wavelength-dependent throughput via composed optical element ([c902c88](https://github.com/CoreySpohn/optixstuff/commit/c902c88ae08447e5dfbc868cbcc51a254098877e))
+
+## [1.1.0](https://github.com/CoreySpohn/optixstuff/compare/v1.0.2...v1.1.0) (2026-05-30)
+
+
+### Features
+
+* add AbstractDisperser and LensletDisperser descriptors ([6e9cd21](https://github.com/CoreySpohn/optixstuff/commit/6e9cd210c3b6f9b8af3b4b8a72cc85e412288c39))
+* **detector:** add per-pixel noise_variance (shot + dark + CIC + read) ([f983b0d](https://github.com/CoreySpohn/optixstuff/commit/f983b0dc628e6a2ef8f6489a396c18ff2fff173a))
+* **optical_path:** add optional disperser descriptor field ([06e486e](https://github.com/CoreySpohn/optixstuff/commit/06e486e31daaa15773d91f2b586ad7ec16359db5))
+
 ## [1.0.2](https://github.com/CoreySpohn/optixstuff/compare/v1.0.1...v1.0.2) (2026-05-25)
 
 

{optixstuff-1.0.2 → optixstuff-1.2.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: optixstuff
-Version: 1.0.2
+Version: 1.2.0
 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
@@ -68,11 +68,13 @@ hardware as composable JAX modules: primary aperture, throughput-affecting
 elements, coronagraph backend, detector. Its job is to be the single source of
 truth for the hardware configuration that downstream tools consume.
 
-Both [coronagraphoto](https://github.com/CoreySpohn/coronagraphoto)
-(2D image simulation) and [jaxEDITH](https://github.com/CoreySpohn/jaxedith)
+[coronagraphoto](https://github.com/CoreySpohn/coronagraphoto)
+(2D image simulation), [coronachrome](https://github.com/CoreySpohn/coronachrome)
+(lenslet-IFS dispersion + spectral extraction), and
+[jaxEDITH](https://github.com/CoreySpohn/jaxedith)
 (exposure-time and yield calculations) import the same `OpticalPath` class and
-the same detector / throughput / primary types from optixstuff. Change a value
-here and both downstream tools pick it up on the next import.
+the same detector / throughput / primary / disperser types from optixstuff.
+Change a value here and every downstream tool picks it up on the next import.
 
 ## What optixstuff is *not*
 
@@ -82,8 +84,8 @@ here and both downstream tools pick it up on the next import.
   optixstuff wraps a PSF backend via `YippyCoronagraph` but does not compute PSFs.
 - **Not a scene model.** Stars, planets, disks, and zodi live in
   [skyscapes](https://github.com/CoreySpohn/skyscapes).
-- **Not a simulator.** Downstream tools (coronagraphoto, jaxEDITH) consume an
-  `OpticalPath` to produce images or count rates.
+- **Not a simulator.** Downstream tools (coronagraphoto, coronachrome, jaxEDITH)
+  consume an `OpticalPath` to produce images, dispersed IFS frames, or count rates.
 
 ## Architecture
 
@@ -91,9 +93,9 @@ Built on [JAX](https://github.com/google/jax) and
 [Equinox](https://github.com/patrick-kidger/equinox), `optixstuff` provides:
 
 - **Abstract interfaces** — `AbstractPrimary`, `AbstractOpticalElement`,
-  `AbstractCoronagraph`, `AbstractDetector`
+  `AbstractCoronagraph`, `AbstractDetector`, `AbstractDisperser`
 - **Concrete implementations** — `SimplePrimary`, `ConstantThroughput`,
-  `IdealDetector`
+  `IdealDetector`, `Detector`, `LensletDisperser`
 - **Container** — `OpticalPath`, a composable hardware configuration passed to all
   simulators
 
@@ -111,11 +113,13 @@ flowchart TB
     optix["<b>optixstuff</b><br/>Telescope · Coronagraph · Detector · OpticalPath<br/>Throughput chains · QE · Noise rates"]
     jaxedith["<b>jaxEDITH</b><br/>Scalar count rates<br/>Exposure-time calculations"]
     corono["<b>coronagraphoto</b><br/>2D image simulation<br/>Multi-epoch scenes"]
+    coronachrome["<b>coronachrome</b><br/>Lenslet-IFS dispersion<br/>Spectral extraction"]
 
     physopt -- YIP --> yippy
     yippy -- flux patterns --> optix
     optix --> jaxedith
     optix --> corono
+    optix --> coronachrome
 ```
 
 ## Installation

{optixstuff-1.0.2 → optixstuff-1.2.0}/README.md

@@ -10,11 +10,13 @@ hardware as composable JAX modules: primary aperture, throughput-affecting
 elements, coronagraph backend, detector. Its job is to be the single source of
 truth for the hardware configuration that downstream tools consume.
 
-Both [coronagraphoto](https://github.com/CoreySpohn/coronagraphoto)
-(2D image simulation) and [jaxEDITH](https://github.com/CoreySpohn/jaxedith)
+[coronagraphoto](https://github.com/CoreySpohn/coronagraphoto)
+(2D image simulation), [coronachrome](https://github.com/CoreySpohn/coronachrome)
+(lenslet-IFS dispersion + spectral extraction), and
+[jaxEDITH](https://github.com/CoreySpohn/jaxedith)
 (exposure-time and yield calculations) import the same `OpticalPath` class and
-the same detector / throughput / primary types from optixstuff. Change a value
-here and both downstream tools pick it up on the next import.
+the same detector / throughput / primary / disperser types from optixstuff.
+Change a value here and every downstream tool picks it up on the next import.
 
 ## What optixstuff is *not*
 
@@ -24,8 +26,8 @@ here and both downstream tools pick it up on the next import.
   optixstuff wraps a PSF backend via `YippyCoronagraph` but does not compute PSFs.
 - **Not a scene model.** Stars, planets, disks, and zodi live in
   [skyscapes](https://github.com/CoreySpohn/skyscapes).
-- **Not a simulator.** Downstream tools (coronagraphoto, jaxEDITH) consume an
-  `OpticalPath` to produce images or count rates.
+- **Not a simulator.** Downstream tools (coronagraphoto, coronachrome, jaxEDITH)
+  consume an `OpticalPath` to produce images, dispersed IFS frames, or count rates.
 
 ## Architecture
 
@@ -33,9 +35,9 @@ Built on [JAX](https://github.com/google/jax) and
 [Equinox](https://github.com/patrick-kidger/equinox), `optixstuff` provides:
 
 - **Abstract interfaces** — `AbstractPrimary`, `AbstractOpticalElement`,
-  `AbstractCoronagraph`, `AbstractDetector`
+  `AbstractCoronagraph`, `AbstractDetector`, `AbstractDisperser`
 - **Concrete implementations** — `SimplePrimary`, `ConstantThroughput`,
-  `IdealDetector`
+  `IdealDetector`, `Detector`, `LensletDisperser`
 - **Container** — `OpticalPath`, a composable hardware configuration passed to all
   simulators
 
@@ -53,11 +55,13 @@ flowchart TB
     optix["<b>optixstuff</b><br/>Telescope · Coronagraph · Detector · OpticalPath<br/>Throughput chains · QE · Noise rates"]
     jaxedith["<b>jaxEDITH</b><br/>Scalar count rates<br/>Exposure-time calculations"]
     corono["<b>coronagraphoto</b><br/>2D image simulation<br/>Multi-epoch scenes"]
+    coronachrome["<b>coronachrome</b><br/>Lenslet-IFS dispersion<br/>Spectral extraction"]
 
     physopt -- YIP --> yippy
     yippy -- flux patterns --> optix
     optix --> jaxedith
     optix --> corono
+    optix --> coronachrome
 ```
 
 ## Installation

{optixstuff-1.0.2 → optixstuff-1.2.0}/src/optixstuff/__init__.py

@@ -10,6 +10,7 @@ from optixstuff.detector import (
     dark_current,
     read_noise,
 )
+from optixstuff.disperser import AbstractDisperser, LensletDisperser
 from optixstuff.exposure import ExposureConfig
 from optixstuff.optical_elements import (
     AbstractOpticalElement,
@@ -24,6 +25,7 @@ from optixstuff.yippy_coronagraph import YippyCoronagraph
 __all__ = [
     "AbstractCoronagraph",
     "AbstractDetector",
+    "AbstractDisperser",
     "AbstractOpticalElement",
     "AbstractPrimary",
     "AbstractScalarCoronagraph",
@@ -32,6 +34,7 @@ __all__ = [
     "Detector",
     "ExposureConfig",
     "IdealDetector",
+    "LensletDisperser",
     "OpticalPath",
     "SimplePrimary",
     "SpectralThroughput",

{optixstuff-1.0.2 → optixstuff-1.2.0}/src/optixstuff/_version.py

@@ -18,7 +18,7 @@ version_tuple: tuple[int | str, ...]
 commit_id: str | None
 __commit_id__: str | None
 
-__version__ = version = '1.0.2'
-__version_tuple__ = version_tuple = (1, 0, 2)
+__version__ = version = '1.2.0'
+__version_tuple__ = version_tuple = (1, 2, 0)
 
 __commit_id__ = commit_id = None

{optixstuff-1.0.2 → optixstuff-1.2.0}/src/optixstuff/detector.py

@@ -1,6 +1,5 @@
 """Detector abstractions and concrete implementations."""
 
-
 import abc
 from typing import final
 
@@ -154,6 +153,37 @@ class AbstractDetector(eqx.Module):
             prng_key, image_rate * exposure_time_s * self.quantum_efficiency
         )
 
+    def noise_variance(self, image_rate: Array, exposure_time_s: ArrayLike) -> Array:
+        """Deterministic per-pixel total noise variance in electrons^2.
+
+        The expected variance at each pixel for an incident photon-rate image
+        -- the deterministic companion to :meth:`readout`. Combines source shot
+        noise (Poisson on detected electrons) with the source-independent floor
+        (dark current, clock-induced charge, read noise)::
+
+            N = QE * rate * t
+                + dark_rate * t
+                + CIC_rate * n_frames
+                + read_noise^2 * n_frames
+
+        with ``n_frames = ceil(t / frame_time_s)``. Use ``1 / noise_variance(...)``
+        as inverse-variance weights for least-squares spectral extraction or its
+        GLS covariance, where the shot term carries the wavelength dependence.
+
+        Args:
+            image_rate: Incident photon rate [ph/s/pixel], any shape.
+            exposure_time_s: Exposure time in seconds.
+
+        Returns:
+            Per-pixel noise variance [electrons^2], same shape as image_rate.
+        """
+        n_frames = jnp.ceil(exposure_time_s / self.frame_time_s)
+        shot = self.quantum_efficiency * image_rate * exposure_time_s
+        dark = self.dark_current_rate_e_per_s * exposure_time_s
+        cic = self.clock_induced_charge_rate_e_per_frame * n_frames
+        read = self.read_noise_e**2 * n_frames
+        return shot + dark + cic + read
+
     @abc.abstractmethod
     def readout_noise_electrons(
         self,
@@ -196,7 +226,9 @@ def dark_current(
     Returns:
         Dark current electrons, shape (ny, nx).
     """
-    return jax.random.poisson(prng_key, dark_current_rate_e_per_s * exposure_time_s, shape=shape)
+    return jax.random.poisson(
+        prng_key, dark_current_rate_e_per_s * exposure_time_s, shape=shape
+    )
 
 
 def clock_induced_charge(
@@ -216,7 +248,9 @@ def clock_induced_charge(
     Returns:
         CIC electrons, shape (ny, nx).
     """
-    return jax.random.poisson(prng_key, clock_induced_charge_rate_e_per_frame * num_frames, shape=shape)
+    return jax.random.poisson(
+        prng_key, clock_induced_charge_rate_e_per_frame * num_frames, shape=shape
+    )
 
 
 def read_noise(
@@ -281,7 +315,9 @@ class IdealDetector(AbstractDetector):
         self.quantum_efficiency = quantum_efficiency
         self.dark_current_rate_e_per_s = dark_current_rate_e_per_s
         self.read_noise_e = read_noise_e
-        self.clock_induced_charge_rate_e_per_frame = clock_induced_charge_rate_e_per_frame
+        self.clock_induced_charge_rate_e_per_frame = (
+            clock_induced_charge_rate_e_per_frame
+        )
         self.frame_time_s = frame_time_s
         self.read_time_s = read_time_s
         self.dqe = dqe
@@ -297,7 +333,9 @@ class IdealDetector(AbstractDetector):
         Callers add (read_noise_e^2 * n_reads) / t_exp separately.
         """
         dark_variance_rate = self.dark_current_rate_e_per_s * n_pix
-        cic_variance_rate = self.clock_induced_charge_rate_e_per_frame * n_pix / t_photon
+        cic_variance_rate = (
+            self.clock_induced_charge_rate_e_per_frame * n_pix / t_photon
+        )
         return dark_variance_rate + cic_variance_rate
 
     def readout_noise_electrons(
@@ -373,7 +411,9 @@ class Detector(AbstractDetector):
         self.quantum_efficiency = quantum_efficiency
         self.dark_current_rate_e_per_s = dark_current_rate_e_per_s
         self.read_noise_e = read_noise_e
-        self.clock_induced_charge_rate_e_per_frame = clock_induced_charge_rate_e_per_frame
+        self.clock_induced_charge_rate_e_per_frame = (
+            clock_induced_charge_rate_e_per_frame
+        )
         self.frame_time_s = frame_time_s
         self.read_time_s = read_time_s
         self.dqe = dqe
@@ -385,7 +425,9 @@ class Detector(AbstractDetector):
     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_e_per_s * n_pix
-        cic_variance_rate = self.clock_induced_charge_rate_e_per_frame * n_pix / t_photon
+        cic_variance_rate = (
+            self.clock_induced_charge_rate_e_per_frame * n_pix / t_photon
+        )
         return dark_variance_rate + cic_variance_rate
 
     def readout_noise_electrons(
@@ -403,9 +445,7 @@ class Detector(AbstractDetector):
         cic_e = clock_induced_charge(
             self.clock_induced_charge_rate_e_per_frame, num_frames, self.shape, key_cic
         )
-        read_e = read_noise(
-            self.read_noise_e, num_frames, self.shape, key_read
-        )
+        read_e = read_noise(self.read_noise_e, num_frames, self.shape, key_read)
         return dark_e + cic_e + read_e
 
     def readout(

optixstuff-1.2.0/src/optixstuff/disperser.py

@@ -0,0 +1,114 @@
+"""Disperser hardware descriptors for integral field spectrographs.
+
+optixstuff owns the descriptor: the interface plus a cheap closed-form
+scalar/ETC face. The heavy render logic (building the forward operator) lives
+in coronachrome. This mirrors the coronagraph split, so jaxedith and yield
+tools can read IFS hardware info without importing the render engine.
+"""
+
+import abc
+
+import equinox as eqx
+import jax.numpy as jnp
+from jax import Array
+
+from optixstuff.optical_elements import AbstractOpticalElement, ConstantThroughput
+
+
+class AbstractDisperser(eqx.Module):
+    """Interface for a dispersing IFS element (lenslet array, slicer, MSA).
+
+    Only the scalar/ETC face is defined here. Render geometry lives in
+    coronachrome and dispatches on the concrete descriptor type.
+    """
+
+    @abc.abstractmethod
+    def spectral_resolution(self, wavelength_nm):
+        """Resolving power R = lambda / dlambda at the given wavelength."""
+
+    @abc.abstractmethod
+    def spectral_sampling(self):
+        """Detector pixels per resolution element."""
+
+    @abc.abstractmethod
+    def n_pix_spread(self, wavelength_min_nm, wavelength_max_nm):
+        """Detector pixels a single spaxel spectrum spans across a band."""
+
+    @abc.abstractmethod
+    def throughput(self, wavelength_nm):
+        """Disperser optical throughput in [0, 1] at the given wavelength."""
+
+
+def _polyval_deriv(coeffs, x):
+    """Evaluate the derivative of a descending-order polynomial at x."""
+    n = coeffs.shape[0]
+    if n <= 1:
+        return jnp.zeros_like(jnp.asarray(x, dtype=float))
+    powers = jnp.arange(n - 1, 0, -1)
+    return jnp.polyval(coeffs[:-1] * powers, x)
+
+
+class LensletDisperser(AbstractDisperser):
+    """Lenslet-array IFS disperser (CRISPY heritage).
+
+    Config only. The render geometry (IR build) is performed by coronachrome,
+    which reads these fields. Scalar/ETC methods derive from
+    ``dispersion_coeffs`` + ``pix_per_reselt`` so the dispersion model is the
+    single source of truth.
+
+    ``psflet_params[0]`` is the PSFlet core width in detector pixels (Gaussian
+    ``sigma`` or Moffat ``alpha``); any trailing entries are dimensionless shape
+    parameters (e.g. Moffat ``beta``). ``psflet_ref_nm`` is the wavelength at
+    which that core width is specified: a diffraction-limited spot scales as
+    ``lambda f / D``, so at fixed pixel scale coronachrome scales the core width
+    by ``lambda / psflet_ref_nm`` per wavelength (the shape parameters do not
+    scale).
+    """
+
+    pitch_m: float
+    pixsize_m: float
+    angle_rad: float
+    lam_ref_nm: float
+    pix_per_reselt: float
+    dispersion_coeffs: Array = eqx.field(converter=jnp.asarray)
+    psflet_params: Array = eqx.field(converter=jnp.asarray)
+    psflet_ref_nm: float
+    grid_kind: str = eqx.field(static=True)
+    n_lenslets: int = eqx.field(static=True)
+    psflet_kind: str = eqx.field(static=True)
+    detector_shape: tuple[int, int] = eqx.field(static=True)
+    throughput_element: AbstractOpticalElement = ConstantThroughput(1.0)
+
+    def _dispersion_px(self, wavelength_nm):
+        """Spectral-axis detector offset [px] for the wavelength(s)."""
+        u = jnp.log(jnp.asarray(wavelength_nm, dtype=float) / self.lam_ref_nm)
+        return jnp.polyval(self.dispersion_coeffs, u)
+
+    def spectral_resolution(self, wavelength_nm):
+        """R = (local px per unit log-lambda) / pixels-per-resolution-element."""
+        u = jnp.log(jnp.asarray(wavelength_nm, dtype=float) / self.lam_ref_nm)
+        local = jnp.abs(_polyval_deriv(self.dispersion_coeffs, u))
+        return local / self.pix_per_reselt
+
+    def spectral_sampling(self):
+        """Detector pixels per resolution element."""
+        return self.pix_per_reselt
+
+    def n_pix_spread(self, wavelength_min_nm, wavelength_max_nm):
+        """Spectral trace length [px] across a band, plus a PSFlet-width margin."""
+        span = jnp.abs(
+            self._dispersion_px(wavelength_max_nm)
+            - self._dispersion_px(wavelength_min_nm)
+        )
+        return span + self.psflet_params[0]
+
+    def throughput(self, wavelength_nm):
+        """Disperser optical throughput in [0, 1], shaped like wavelength_nm.
+
+        Delegates to the composed throughput element (ConstantThroughput by
+        default; SpectralThroughput for a tabulated blaze/transmission curve) and
+        broadcasts to the wavelength shape so the output shape is canonical
+        regardless of the backing element.
+        """
+        w = jnp.asarray(wavelength_nm, dtype=float)
+        return self.throughput_element.get_throughput(w) * jnp.ones_like(w)

{optixstuff-1.0.2 → optixstuff-1.2.0}/src/optixstuff/optical_path.py

@@ -10,6 +10,7 @@ import equinox as eqx
 from optixstuff._repr import indent
 from optixstuff.coronagraph import AbstractCoronagraph
 from optixstuff.detector import AbstractDetector, IdealDetector
+from optixstuff.disperser import AbstractDisperser
 from optixstuff.optical_elements import (
     AbstractOpticalElement,
     ConstantThroughput,
@@ -31,6 +32,7 @@ class OpticalPath(eqx.Module):
             between the primary and coronagraph (mirrors, filters, etc.).
         coronagraph: Coronagraph performance model.
         detector: Focal-plane detector model.
+        disperser: Optional IFS disperser descriptor; None for imaging mode.
         n_channels: Number of parallel identical optical-path copies
             (AYO shorthand, multiplicative factor on count rates;
             not a spectral channel count). Default 1.0.
@@ -42,6 +44,7 @@ class OpticalPath(eqx.Module):
     attenuating_elements: tuple[AbstractOpticalElement, ...]
     coronagraph: AbstractCoronagraph
     detector: AbstractDetector
+    disperser: AbstractDisperser | None = None
     n_channels: float = 1.0
     npix_multiplier: float = 1.0
 
@@ -110,9 +113,7 @@ class OpticalPath(eqx.Module):
         return cls(
             primary=SimplePrimary(diameter_m=diameter_m, obscuration=obscuration),
             attenuating_elements=(
-                ConstantThroughput(
-                    throughput=attenuating_throughput, name="optics"
-                ),
+                ConstantThroughput(throughput=attenuating_throughput, name="optics"),
             ),
             coronagraph=coro,
             detector=IdealDetector(