optisketch 0.0.3__py3-none-any.whl

optisketch/__init__.py

@@ -0,0 +1,46 @@
+"""A light raytracing library."""
+
+from importlib.metadata import PackageNotFoundError, version
+
+try:
+    __version__ = version("trace-light")
+except PackageNotFoundError:
+    __version__ = "uninstalled"
+__author__ = "Kevin Yamauchi"
+__email__ = "kevin.yamauchi@gmail.com"
+
+from optisketch import analysis, backends, fluorescence, optimize, viz
+from optisketch.kernels import _trace_surfaces, trace
+from optisketch.rays import (
+    Rays,
+    Surface,
+    System,
+    _Params,
+    _Structure,
+    load_system,
+    save_system,
+)
+from optisketch.sources import Source, collimated_source, emit, point_source
+from optisketch.systems import SystemBuilder
+
+__all__ = [
+    "Rays",
+    "Source",
+    "Surface",
+    "System",
+    "SystemBuilder",
+    "_Params",
+    "_Structure",
+    "_trace_surfaces",
+    "analysis",
+    "backends",
+    "collimated_source",
+    "emit",
+    "fluorescence",
+    "load_system",
+    "optimize",
+    "point_source",
+    "save_system",
+    "trace",
+    "viz",
+]

optisketch/analysis/__init__.py

@@ -0,0 +1,22 @@
+"""Read-only analysis functions (Phases 4-5).
+
+* :func:`spot` — spot-diagram statistics for a ray bundle.
+* :func:`psf` — geometric point-spread-function kernel.
+* :func:`irradiance` — flux-density histogram at a plane.
+* :func:`image_sim` — incoherent image formation by PSF convolution.
+"""
+
+from __future__ import annotations
+
+from optisketch.analysis.image_sim import image_sim
+from optisketch.analysis.irradiance import irradiance
+from optisketch.analysis.psf import psf
+from optisketch.analysis.spot import SpotStats, spot
+
+__all__ = [
+    "SpotStats",
+    "image_sim",
+    "irradiance",
+    "psf",
+    "spot",
+]

optisketch/analysis/image_sim.py

@@ -0,0 +1,230 @@
+"""Incoherent image simulation by PSF convolution (Phase 5).
+
+:func:`image_sim` forms an image as the incoherent sum, over object depth
+slices, of each slice convolved with the system PSF appropriate to its
+``(field, depth, focus)``. Two modes are provided:
+
+* ``psf="single"`` — one lateral PSF per depth (shift-invariant); each slice is
+  FFT-convolved with that single kernel.
+* ``psf="varying"`` — PSFs sampled on a coarse ``grid`` of field points per
+  depth. Each field region is selected by a separable partition-of-unity weight
+  map; the object slice is multiplied by the weight, convolved with the local
+  PSF, and the contributions are summed. Because the weights sum to one
+  everywhere, a system whose PSF is field-independent reproduces the
+  ``psf="single"`` result exactly.
+
+A scalar ``focus`` yields a 2-D image; an array ``focus`` yields a 3-D focal
+stack.
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any
+
+import numpy as np
+
+from optisketch.analysis.psf import psf as _psf
+
+if TYPE_CHECKING:
+    from optisketch.rays import System
+
+
+def _axis_hats(n: int, g: int) -> tuple[np.ndarray, np.ndarray]:
+    """Return separable partition-of-unity hat weights along one axis.
+
+    For *g* nodes spread evenly over ``[0, n-1]``, each pixel is assigned
+    linear-interpolation weights to the nodes such that the weights sum to one
+    at every pixel (a partition of unity).
+
+    Parameters
+    ----------
+    n : int
+        Number of pixels along the axis.
+    g : int
+        Number of field-sample nodes along the axis.
+
+    Returns
+    -------
+    weights : numpy.ndarray
+        Array of shape ``(n, g)``; column *k* is the weight of node *k*.
+    nodes : numpy.ndarray
+        Node pixel positions, shape ``(g,)``.
+    """
+    coords = np.arange(n, dtype=np.float64)
+    nodes = np.linspace(0.0, n - 1, g)
+    weights = np.zeros((n, g), dtype=np.float64)
+    for k in range(g):
+        e = np.zeros(g, dtype=np.float64)
+        e[k] = 1.0
+        weights[:, k] = np.interp(coords, nodes, e)
+    return weights, nodes
+
+
+def _pixel_to_field(node_px: float, n: int, half_extent: float) -> float:
+    """Map a node pixel position to a lateral field coordinate (mm).
+
+    Parameters
+    ----------
+    node_px : float
+        Node position in pixel units (0 .. n-1).
+    n : int
+        Number of pixels along the axis.
+    half_extent : float
+        Half-width of the field of view along this axis (mm).
+
+    Returns
+    -------
+    float
+        Field coordinate in millimetres, centred on the array.
+    """
+    if n <= 1:
+        return 0.0
+    frac = (node_px - (n - 1) / 2.0) / ((n - 1) / 2.0)
+    return frac * half_extent
+
+
+def image_sim(
+    system: System,
+    obj: Any,
+    extent: float | tuple[float, float],
+    *,
+    psf: str = "varying",
+    field: tuple[float, float] = (0.0, 0.0),
+    grid: tuple[int, int] = (3, 3),
+    focus: float | Any = 0.0,
+    wavelength: float | None = None,
+    psf_grid: tuple[int, int] = (31, 31),
+    n_rays: int = 256,
+    z_object: float = -100.0,
+    depth_extent: float = 0.0,
+) -> Any:
+    """Simulate the incoherent image of *obj* formed by *system*.
+
+    Parameters
+    ----------
+    system : System
+        Imaging system.
+    obj : array
+        Object intensity. Shape ``(ny, nx)`` for a planar object or
+        ``(nz, ny, nx)`` for a volume (incoherently summed over depth).
+    extent : float or tuple of float
+        Lateral half-width of the field of view (mm). Scalar applies to both
+        axes; ``(ey, ex)`` sets them separately.
+    psf : str, optional
+        ``"varying"`` (field-dependent, default) or ``"single"``
+        (shift-invariant).
+    field : tuple of float, optional
+        Base lateral field offset (mm) added to every PSF evaluation.
+    grid : tuple of int, optional
+        ``(gy, gx)`` coarse field-sampling grid for ``psf="varying"``.
+        Ignored for ``psf="single"``. Default ``(3, 3)``.
+    focus : float or array, optional
+        Detector focus offset(s) from ``system.image_z`` (mm). A scalar yields
+        a 2-D image; a 1-D array yields a 3-D ``(nf, ny, nx)`` stack.
+    wavelength : float, optional
+        Imaging wavelength (µm). Defaults to the system's first wavelength.
+    psf_grid : tuple of int, optional
+        ``(ny, nx)`` grid of each evaluated PSF kernel. Default ``(31, 31)``.
+    n_rays : int, optional
+        Pupil samples per PSF evaluation. Default 256.
+    z_object : float, optional
+        Nominal object-plane z-position (mm). Default ``-100.0``.
+    depth_extent : float, optional
+        Half-range of object axial depth mapped across the ``nz`` slices (mm).
+        Default 0 (all slices at the nominal plane).
+
+    Returns
+    -------
+    array
+        Simulated image of shape ``(ny, nx)`` (scalar *focus*) or
+        ``(nf, ny, nx)`` (array *focus*).
+
+    Raises
+    ------
+    ValueError
+        If *psf* is not ``"single"`` or ``"varying"`` or *obj* is not 2-D/3-D.
+    """
+    if psf not in ("single", "varying"):
+        raise ValueError(f"Unknown psf mode {psf!r}. Choose 'single' or 'varying'.")
+
+    be = system.backend
+
+    shape = tuple(obj.shape)
+    if len(shape) == 2:
+        ny, nx = shape
+        slices = [obj]
+        nz = 1
+    elif len(shape) == 3:
+        nz, ny, nx = shape
+        slices = [obj[zi] for zi in range(nz)]
+    else:
+        raise ValueError("obj must be 2-D (ny, nx) or 3-D (nz, ny, nx).")
+
+    if isinstance(extent, (tuple, list)):
+        ext_y, ext_x = float(extent[0]), float(extent[1])
+    else:
+        ext_y = ext_x = float(extent)
+
+    # The PSF kernel must share the object's pixel pitch so the discrete
+    # convolution is physically meaningful. Pixel pitch = full width / pixels.
+    pitch = (2.0 * ext_x) / nx
+    psf_extent = pitch * psf_grid[1] / 2.0
+
+    # focus may be scalar or a 1-D array of focal planes
+    focus_arr = np.atleast_1d(np.asarray(be.to_numpy(focus), dtype=np.float64))
+    scalar_focus = np.ndim(be.to_numpy(focus)) == 0
+
+    # depth per slice
+    if nz == 1:
+        depths = [0.0]
+    else:
+        depths = list(np.linspace(-depth_extent, depth_extent, nz))
+
+    # field sampling weight maps (varying mode only)
+    if psf == "varying":
+        gy, gx = grid
+        wy, nodes_y = _axis_hats(ny, gy)
+        wx, nodes_x = _axis_hats(nx, gx)
+
+    images = []
+    for f in focus_arr:
+        acc = be.zeros((ny, nx))
+        for zi in range(nz):
+            depth = depths[zi]
+            slc = slices[zi]
+            if psf == "single":
+                k = _psf(
+                    system,
+                    field=field,
+                    depth=depth,
+                    focus=float(f),
+                    wavelength=wavelength,
+                    n_rays=n_rays,
+                    grid=psf_grid,
+                    extent=psf_extent,
+                    z_object=z_object,
+                )
+                acc = acc + be.fftconvolve(slc, k, mode="same")
+            else:
+                for gi in range(gy):
+                    fy = field[1] + _pixel_to_field(nodes_y[gi], ny, ext_y)
+                    for gj in range(gx):
+                        fx = field[0] + _pixel_to_field(nodes_x[gj], nx, ext_x)
+                        wmap = be.asarray(np.outer(wy[:, gi], wx[:, gj]))
+                        k = _psf(
+                            system,
+                            field=(fx, fy),
+                            depth=depth,
+                            focus=float(f),
+                            wavelength=wavelength,
+                            n_rays=n_rays,
+                            grid=psf_grid,
+                            extent=psf_extent,
+                            z_object=z_object,
+                        )
+                        acc = acc + be.fftconvolve(slc * wmap, k, mode="same")
+        images.append(acc)
+
+    if scalar_focus:
+        return images[0]
+    return be.stack(images, axis=0)

optisketch/analysis/irradiance.py

@@ -0,0 +1,79 @@
+"""Irradiance (flux density) estimation at a detector plane (Phase 4).
+
+:func:`irradiance` traces a :class:`~optisketch.sources.Source` through a
+system, propagates the exiting rays to a target plane, and accumulates a
+weighted 2-D histogram of the valid hits. It is a thin wrapper over
+:func:`~optisketch.kernels._trace_surfaces` and the backend ``histogram2d``.
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any
+
+from optisketch.kernels import _propagate_to_plane, _trace_surfaces
+from optisketch.sources import emit
+
+if TYPE_CHECKING:
+    from optisketch.rays import System
+    from optisketch.sources import Source
+
+
+def irradiance(
+    system: System,
+    source: Source,
+    z: float,
+    grid: tuple[int, int] = (64, 64),
+    *,
+    extent: float | tuple[float, float] | None = None,
+) -> Any:
+    """Compute the irradiance distribution of *source* at plane ``z``.
+
+    The source is emitted into the system, traced through every surface, and
+    propagated in free space to ``z``. Valid hits are binned into a weighted
+    2-D histogram (weights are the per-ray intensities). When *extent* is None
+    the window is sized to enclose all valid hits, so the histogram sum equals
+    the total weight of the valid rays.
+
+    Parameters
+    ----------
+    system : System
+        Optical system to trace through.
+    source : Source
+        Ray source to emit.
+    z : float
+        Detector-plane z-position (mm).
+    grid : tuple of int, optional
+        ``(ny, nx)`` histogram grid shape. Default ``(64, 64)``.
+    extent : float or tuple of float, optional
+        Half-width of the window about the origin (mm). Scalar applies to both
+        axes; ``(ey, ex)`` sets them separately. When None it is derived from
+        the data.
+
+    Returns
+    -------
+    array
+        2-D irradiance histogram of shape *grid*, indexed ``[iy, ix]``.
+    """
+    be = system.backend
+    rays = emit(source, system)
+    final, _ = _trace_surfaces(rays, system.structure, system.params, be)
+    final = _propagate_to_plane(final, float(z), be)
+
+    xs, ys, valid = final.x, final.y, final.valid
+    zeros = be.zeros_like(xs)
+    wv = be.where(valid, final.i, be.zeros_like(final.i))
+
+    if extent is None:
+        ax = float(be.to_numpy(be.max(be.where(valid, be.abs(xs), zeros))))
+        ay = float(be.to_numpy(be.max(be.where(valid, be.abs(ys), zeros))))
+        ex = ax * 1.0 + 1e-6
+        ey = ay * 1.0 + 1e-6
+    elif isinstance(extent, (tuple, list)):
+        ey, ex = float(extent[0]), float(extent[1])
+    else:
+        ex = ey = float(extent)
+
+    xs_safe = be.where(valid, xs, zeros)
+    ys_safe = be.where(valid, ys, zeros)
+    rng = ((-ey, ey), (-ex, ex))
+    return be.histogram2d(ys_safe, xs_safe, bins=grid, range=rng, weights=wv)

optisketch/analysis/psf.py

@@ -0,0 +1,144 @@
+"""Point-spread-function estimation by geometric ray histogramming (Phase 4).
+
+:func:`psf` traces a point emitter through a :class:`~optisketch.rays.System`,
+propagates the exiting rays to a (possibly defocused) detector plane, and
+histograms the valid image-plane hits into a normalised 2-D kernel centred on
+the spot centroid. Because the kernel is centred on the centroid, the result is
+a shift-invariant impulse response suitable for convolution in
+:func:`~optisketch.analysis.image_sim.image_sim`.
+
+Sweeping ``depth`` (object-side axial position) or ``focus`` (detector-side
+axial position) produces a through-focus stack.
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any
+
+from optisketch.kernels import _propagate_to_plane, _trace_surfaces
+from optisketch.sources import emit, point_source
+
+if TYPE_CHECKING:
+    from optisketch.rays import System
+
+
+def _masked_centroid(xs: Any, ys: Any, valid: Any, be: Any) -> tuple[float, float, int]:
+    """Return the NaN-safe centroid of valid samples as Python floats.
+
+    Parameters
+    ----------
+    xs : array
+        x-coordinates of the samples.
+    ys : array
+        y-coordinates of the samples.
+    valid : bool array
+        Mask selecting the samples that contribute.
+    be : Backend
+        Array-computation backend.
+
+    Returns
+    -------
+    cx : float
+        Centroid x-coordinate (0.0 if no valid samples).
+    cy : float
+        Centroid y-coordinate.
+    n : int
+        Number of valid samples.
+    """
+    zeros = be.zeros_like(xs)
+    vf = be.where(valid, be.ones_like(xs), zeros)
+    n = float(be.to_numpy(be.sum(vf)))
+    if n < 1.0:
+        return 0.0, 0.0, 0
+    cx = float(be.to_numpy(be.sum(be.where(valid, xs, zeros)))) / n
+    cy = float(be.to_numpy(be.sum(be.where(valid, ys, zeros)))) / n
+    return cx, cy, int(n)
+
+
+def psf(
+    system: System,
+    field: tuple[float, float] = (0.0, 0.0),
+    *,
+    depth: float = 0.0,
+    focus: float = 0.0,
+    wavelength: float | None = None,
+    n_rays: int = 256,
+    grid: tuple[int, int] = (64, 64),
+    extent: float | None = None,
+    z_object: float = -100.0,
+) -> Any:
+    """Estimate the geometric PSF kernel of *system* for a point emitter.
+
+    A point source at lateral position *field* and axial position
+    ``z_object + depth`` is traced through the system; the exiting rays are
+    propagated to the detector plane ``image_z + focus`` and histogrammed into
+    a grid centred on the spot centroid. The kernel is normalised to unit sum.
+
+    Parameters
+    ----------
+    system : System
+        Optical system to evaluate.
+    field : tuple of float, optional
+        ``(x, y)`` lateral position of the point emitter (mm). Default on-axis.
+    depth : float, optional
+        Axial offset of the emitter from the nominal object plane (mm).
+        Default 0.
+    focus : float, optional
+        Axial offset of the detector plane from ``system.image_z`` (mm).
+        Default 0.
+    wavelength : float, optional
+        Emission wavelength (µm). Defaults to the system's first wavelength.
+    n_rays : int, optional
+        Number of pupil samples (rays). Default 256.
+    grid : tuple of int, optional
+        ``(ny, nx)`` kernel grid shape. Default ``(64, 64)``.
+    extent : float, optional
+        Half-width of the kernel window about the centroid (mm). When None it
+        is derived from the geometric spread of the valid rays.
+    z_object : float, optional
+        Nominal object-plane z-position (mm). Default ``-100.0``.
+
+    Returns
+    -------
+    array
+        Normalised 2-D PSF kernel of shape *grid*, indexed ``[iy, ix]``.
+    """
+    be = system.backend
+    wl = float(system.wavelengths[0]) if wavelength is None else float(wavelength)
+
+    src = point_source(
+        (float(field[0]), float(field[1])),
+        z_object=float(z_object) + float(depth),
+        wavelength=wl,
+        pupil_pattern="disk",
+        n_samples=n_rays,
+    )
+    rays = emit(src, system)
+    final, _ = _trace_surfaces(rays, system.structure, system.params, be)
+
+    z_det = float(system.image_z) + float(focus)
+    final = _propagate_to_plane(final, z_det, be)
+
+    xs, ys, valid = final.x, final.y, final.valid
+    cx, cy, _n = _masked_centroid(xs, ys, valid, be)
+
+    if extent is None:
+        zeros = be.zeros_like(xs)
+        dx = be.where(valid, xs - cx, zeros)
+        dy = be.where(valid, ys - cy, zeros)
+        r2 = dx * dx + dy * dy
+        spread = float(be.to_numpy(be.sqrt(be.max(be.where(valid, r2, zeros)))))
+        extent = spread * 1.1 + 1e-6
+    e = float(extent)
+
+    # weights: valid-ray intensity, zero elsewhere; keep invalid coords in-range
+    wv = be.where(valid, final.i, be.zeros_like(final.i))
+    xs_safe = be.where(valid, xs, be.full_like(xs, cx))
+    ys_safe = be.where(valid, ys, be.full_like(ys, cy))
+
+    rng = ((cy - e, cy + e), (cx - e, cx + e))
+    h = be.histogram2d(ys_safe, xs_safe, bins=grid, range=rng, weights=wv)
+
+    total = be.sum(h)
+    total_safe = be.where(total > 0.0, total, be.ones_like(total))
+    return h / total_safe

optisketch/analysis/spot.py

@@ -0,0 +1,129 @@
+"""Spot-diagram statistics for a traced ray bundle (Phase 4).
+
+:func:`spot` reduces a :class:`~optisketch.rays.Rays` bundle to centroid,
+RMS, and geometric-radius statistics over the valid rays. All reductions are
+NaN-safe: invalid rays are masked out via ``backend.where`` before any sum or
+maximum, so a NaN in a missed/TIR ray never corrupts the result.
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, NamedTuple
+
+if TYPE_CHECKING:
+    from optisketch.backends._protocol import Backend
+    from optisketch.rays import Rays
+
+
+class SpotStats(NamedTuple):
+    """Summary statistics of a spot diagram.
+
+    Attributes
+    ----------
+    cx : float
+        x-coordinate of the centroid of valid ray hits (mm).
+    cy : float
+        y-coordinate of the centroid of valid ray hits (mm).
+    rms : float
+        Root-mean-square spot radius about the reference point (mm).
+    geo_radius : float
+        Maximum distance of any valid ray from the reference point (mm).
+    n_valid : int
+        Number of valid rays contributing to the statistics.
+    """
+
+    cx: float
+    cy: float
+    rms: float
+    geo_radius: float
+    n_valid: int
+
+
+def _default_backend() -> Backend:
+    """Return a NumpyBackend instance without importing it at module load.
+
+    Returns
+    -------
+    Backend
+        A fresh :class:`~optisketch.backends.NumpyBackend`.
+    """
+    from optisketch.backends._numpy import NumpyBackend
+
+    return NumpyBackend()
+
+
+def spot(
+    rays: Rays,
+    *,
+    reference: str = "centroid",
+    backend: Backend | None = None,
+) -> SpotStats:
+    """Compute spot-diagram statistics for a ray bundle.
+
+    Statistics are evaluated on the final ``(x, y)`` positions of the rays,
+    masked on ``rays.valid``. Two reference modes are supported: ``"centroid"``
+    measures RMS/geometric radius about the centroid of the valid hits, while
+    ``"chief"`` measures them about the chief ray (ray index 0).
+
+    Parameters
+    ----------
+    rays : Rays
+        Ray bundle, typically the output of a trace at the image plane.
+    reference : str, optional
+        ``"centroid"`` (default) or ``"chief"``. Selects the reference point
+        for the RMS and geometric-radius statistics.
+    backend : Backend, optional
+        Array-computation backend. Defaults to
+        :class:`~optisketch.backends.NumpyBackend`.
+
+    Returns
+    -------
+    SpotStats
+        Centroid, RMS radius, geometric radius, and valid-ray count.
+
+    Raises
+    ------
+    ValueError
+        If *reference* is not ``"centroid"`` or ``"chief"``.
+    """
+    if reference not in ("centroid", "chief"):
+        raise ValueError(
+            f"Unknown reference {reference!r}. Choose 'centroid' or 'chief'."
+        )
+
+    be = backend if backend is not None else _default_backend()
+
+    valid = rays.valid
+    zeros = be.zeros_like(rays.x)
+    # mask positions: invalid rays contribute 0 (never NaN) to the sums
+    xv = be.where(valid, rays.x, zeros)
+    yv = be.where(valid, rays.y, zeros)
+    vf = be.where(valid, be.ones_like(rays.x), zeros)
+
+    n = be.sum(vf)
+    n_safe = be.maximum(n, be.ones_like(n))
+    cx = be.sum(xv) / n_safe
+    cy = be.sum(yv) / n_safe
+
+    if reference == "chief":
+        rx = rays.x[0]
+        ry = rays.y[0]
+    else:
+        rx = cx
+        ry = cy
+
+    dx = be.where(valid, rays.x - rx, zeros)
+    dy = be.where(valid, rays.y - ry, zeros)
+    r2 = dx * dx + dy * dy
+
+    rms = be.sqrt(be.sum(r2) / n_safe)
+    geo = be.sqrt(be.max(be.where(valid, r2, zeros)))
+
+    to_np = be.to_numpy
+    return SpotStats(
+        cx=float(to_np(cx)),
+        cy=float(to_np(cy)),
+        rms=float(to_np(rms)),
+        geo_radius=float(to_np(geo)),
+        n_valid=int(to_np(n)),
+    )

optisketch/backends/__init__.py

@@ -0,0 +1,38 @@
+"""Backend constructors.
+
+Usage::
+
+    from optisketch.backends import numpy, jax
+
+    be = numpy()
+    be_jax = jax()  # raises ImportError if jax is not installed
+"""
+
+from __future__ import annotations
+
+from optisketch.backends._numpy import NumpyBackend
+from optisketch.backends._protocol import Backend, NotDifferentiable
+
+
+def numpy() -> NumpyBackend:
+    """Return a :class:`NumpyBackend` instance."""
+    return NumpyBackend()
+
+
+def jax():
+    """Return a :class:`JaxBackend` instance.
+
+    Raises :exc:`ImportError` if jax is not installed.
+    """
+    from optisketch.backends._jax import JaxBackend
+
+    return JaxBackend()
+
+
+__all__ = [
+    "Backend",
+    "NotDifferentiable",
+    "NumpyBackend",
+    "jax",
+    "numpy",
+]