From 57950f8670780b21021669f69e543dd6ba16ca48 Mon Sep 17 00:00:00 2001 From: Talley Lambert Date: Tue, 13 Oct 2026 16:35:40 +0200 Subject: [PATCH 1/2] fix: compute confocal excitation at the excitation wavelengths; refactor PSF pipeline The confocal excitation spot was computed at the emission wavelength. It is now computed at the excitation wavelengths, each weighted by how strongly the fluorophore absorbs it, then saturated locally. To make this possible, the simulation -> modality -> PSF pipeline is restructured: - `Simulation._channel_fluorescence()` returns a `ChannelFluorescence` per (channel, fluorophore): absorption rate (photons/s by excitation wavelength) and detected fraction (photons detected per absorbed photon, by emission wavelength), always unsaturated. - Modalities combine an illumination (`WidefieldIllumination`, `FocusedIllumination`) and a detection (`CameraDetection`, `PinholeDetection`) into a system PSF (excitation spot x detection PSF), kept separate from the scalar brightness, so changing power doesn't recompute PSFs. - `microsim._psf` builds PSFs from two cached, single-wavelength primitives: `excitation_spot` and `detection_psf`. `make_psf` and `make_confocal_psf` keep their signatures. - The pinhole mask is anti-aliased (each pixel weighted by the fraction of its area inside the pinhole), so small pinholes don't snap to whole pixels. - The confocal pinhole has one physical size per channel: `pinhole_au` is in Airy units at the center of the emission filter band, or at the new `Confocal.airy_unit_wavelength_nm`. - New `Settings.excitation_bin_width_nm` (default 10). BREAKING CHANGES: - all confocal output changes (excitation wavelength, pinhole definition). - `filtered_emission_rates()` / `emission_flux()` apply saturation by default for every modality (new `saturate=` keyword). - removed `microsim.psf.cached_psf`, `Modality.psf()` and `point_scanning` (replaced by `detection_type.sums_over_pixel_area`); `Modality.render()` takes per-channel/fluorophore fluorescence instead of emission rates. - cached PSF arrays are read-only; `make_psf` returns a copy. - `Identity` output follows `settings.float_dtype` (float32 by default). - confocal warns for channels without a band-limited emission filter, and raises if a channel's pinhole detects no light. - PSF disk cache layout changed (old cached PSFs are orphaned, not reused). Co-Authored-By: Claude Opus 5.5 --- docs/concept.md | 10 +- docs/stages.md | 23 + src/microsim/_psf.py | 393 ++++++++-------- src/microsim/psf.py | 114 ++++- src/microsim/schema/modality/_optics.py | 140 ++++++ src/microsim/schema/modality/_simple_psf.py | 434 ++++++++++-------- .../schema/optical_config/_fluorescence.py | 104 +++++ src/microsim/schema/optical_config/config.py | 29 ++ src/microsim/schema/settings.py | 19 +- src/microsim/schema/simulation.py | 104 +++-- tests/test_cache.py | 41 +- tests/test_photometry.py | 215 ++++++++- tests/test_psf.py | 22 +- tests/test_simulation.py | 8 +- 14 files changed, 1185 insertions(+), 471 deletions(-) create mode 100644 src/microsim/schema/modality/_optics.py create mode 100644 src/microsim/schema/optical_config/_fluorescence.py diff --git a/docs/concept.md b/docs/concept.md index 6fd937fa..0c47d2a6 100644 --- a/docs/concept.md +++ b/docs/concept.md @@ -84,11 +84,11 @@ can be used to make the simulation reproducible, and other general parameters. ### Point Spread Functions -There is no top-level point spread function object in the simulation: most -modalities will be responsible for generating their own point spread function -based on the objective lens and other parameter that are passed into their -`render()` method. The `microsim.psf` module contains some common point spread -functions, with `make_psf()` being the most common entrypoint. +There is no top-level point spread function object in the simulation: each +modality builds a *system PSF* for every fluorophore in every channel (the +excitation spot times the detection PSF), based on the objective lens and on how bright each fluorophore is in each +channel (and at which wavelengths). The `microsim.psf` module contains +`make_psf()` and `make_confocal_psf()`, for making PSFs directly. ### FPbase diff --git a/docs/stages.md b/docs/stages.md index e971dcdf..6271f027 100644 --- a/docs/stages.md +++ b/docs/stages.md @@ -213,6 +213,12 @@ where: - $T_{\text{em}}(\lambda)$ is the transmission of the emission filters at each wavelength. - $\text{QE}(\lambda)$ is the quantum efficiency of the detector at each wavelength. +These rates are for a fluorophore at the *reference* illumination intensity (the +light source `power`): anywhere in a widefield field of view, or at the center +of the focused spot in a confocal. By default they include excitation saturation +at that intensity; use `filtered_emission_rates(saturate=False)` for the rates +without it. + The output of this stage (given by [`Simulation.filtered_emission_rates`][microsim.Simulation.filtered_emission_rates]) is a 3D array with dimensions `(C, F, W)`, but lacking spatial information: @@ -345,6 +351,23 @@ which retains the full `(C, F, Z, Y, X)` dimensions. In microsim, the PSF largely be determined by the [`ObjectiveLens`][microsim.schema.ObjectiveLens] and the [`Modality`](api.md#modality) of the simulation +!!! info "Widefield vs. confocal" + + For every modality, the signal from a fluorophore at some position (relative + to the focus) is *how strongly it is excited there* times *the fraction of its + emission that reaches the detector*: the excitation spot times the detection + PSF, or the **system PSF**. + + - **Widefield**: the whole field is evenly illuminated, so only the detection + PSF (the image of a point source on the camera) shapes the image. + - **Confocal**: the system PSF is the excitation spot times the detection PSF + through the pinhole. The excitation spot is computed at the excitation + wavelengths (with several lasers or a broadband source, each wavelength is + weighted by how strongly the fluorophore absorbs it). Excitation + saturation is applied locally, so it is strongest at the center of the + spot. `pinhole_au` is in Airy units at the center of the channel's + emission filter band (or at `airy_unit_wavelength_nm`, if set). + The output of this stage is a 4D array with dimensions `(C, Z, Y, X)` where the wavelength and fluorophore dimensions have been collapsed into the channel. Note that each channel dimension will contain the emission of *all* fluorophores that diff --git a/src/microsim/_psf.py b/src/microsim/_psf.py index 91c8ae26..e3850a3e 100644 --- a/src/microsim/_psf.py +++ b/src/microsim/_psf.py @@ -3,8 +3,9 @@ from __future__ import annotations import os +from dataclasses import dataclass from functools import cache, lru_cache -from typing import TYPE_CHECKING, Literal +from typing import TYPE_CHECKING, Any, Literal import numpy as np import numpy.typing as npt @@ -19,7 +20,7 @@ from ._logger import logger if TYPE_CHECKING: - from collections.abc import Sequence + from collections.abc import Callable, Sequence from pathlib import Path from microsim._data_array import ArrayProtocol @@ -249,44 +250,58 @@ def vectorial_psf_centered( ) -def make_confocal_psf( - nz: int, - ex_wvl_um: float = 0.475, - em_wvl_um: float = 0.525, - pinhole_au: float = 1.0, - dz: float = 0.05, - pz: float = 0, - nx: int = 31, - ny: int | None = None, - pos: tuple[float, float, float] = (0, 0, 0), - dxy: float = 0.05, - objective: ObjectiveKwargs | ObjectiveLens | None = None, - sf: int = 3, - normalize: bool | Literal["sum", "max"] = "sum", - saturation: float = 0, - xp: NumpyAPI | None = None, -) -> np.ndarray: - """Create a confocal PSF. +@dataclass(frozen=True) +class PSFGrid: + """Shape and voxel size (in um) of a 3D PSF array. - This function creates a confocal PSF by multiplying the excitation PSF with - the emission PSF convolved with a pinhole mask. + Sizes are rounded up to odd numbers, so that the PSF is centered on a voxel (an + even-sized PSF shifts the image by a fraction of a pixel when convolved). + """ + + nz: int + nx: int + dz: float + dx: float + + def __post_init__(self) -> None: + object.__setattr__(self, "nz", int(self.nz) | 1) + object.__setattr__(self, "nx", int(self.nx) | 1) - If `normalize` is False, the result is physically scaled: the relative excitation - intensity (1 at focus) times the fraction of collected emission that passes the - pinhole, for a point source at each position. - `saturation` is the excitation saturation parameter at the focus (`k * tau`, see - `OpticalConfig.saturation_parameter`). If non-zero, the relative excitation `r` - is replaced by the relative emission of a two-level system, `r / (1 + s * r)`. +def airy_radius_um(wvl_um: float, na: float) -> float: + """Radius of the Airy disk (to its first minimum), in um.""" + return 0.61 * wvl_um / na - All extra keyword arguments are passed to `vectorial_psf_centered`. + +def apply_saturation(spot: Any, saturation: float) -> Any: + """Apply excitation saturation to a relative excitation intensity (1 at focus). + + `saturation` is the fluorophore's saturation at the reference (focal) intensity + (see `ChannelFluorescence.saturation`). The returned array is the emission at + each position, relative to the *unsaturated* emission at the focus. """ - xp = NumpyAPI.create(xp) + if not saturation: + return spot + return spot / (1 + saturation * spot) - objective = _cast_objective(objective) - ex_psf = vectorial_psf_centered( + +def _excitation_spot( + nz: int, + wvl_um: float, + dz: float, + nx: int, + dxy: float, + objective: ObjectiveLens, + xp: NumpyAPI, + pz: float = 0, + ny: int | None = None, + pos: tuple[float, float, float] = (0, 0, 0), + sf: int = 3, +) -> npt.NDArray: + """Relative excitation intensity of a focused laser spot (1 at the focus).""" + return vectorial_psf_centered( nz=nz, - wvl=ex_wvl_um, + wvl=wvl_um, dz=dz, pz=pz, nx=nx, @@ -296,15 +311,34 @@ def make_confocal_psf( objective=objective, xp=xp, sf=sf, - normalize="max", # relative irradiance: 1 at focus + normalize="max", ) - if saturation: - # local steady-state saturation: emission relative to the unsaturated - # emission at focus (see `OpticalConfig.total_emission_rate`) - ex_psf = ex_psf / (1 + saturation * ex_psf) - em_psf = vectorial_psf_centered( + + +def _detection_psf( + nz: int, + wvl_um: float, + dz: float, + nx: int, + dxy: float, + objective: ObjectiveLens, + xp: NumpyAPI, + pinhole_radius_um: float | None = None, + pz: float = 0, + ny: int | None = None, + pos: tuple[float, float, float] = (0, 0, 0), + sf: int = 3, +) -> npt.NDArray: + """Fraction of the light collected from a point source that reaches the detector. + + Without a pinhole, each z-plane is the image of the point source on a camera + (see `_norm_to_collected_light`). With a pinhole (radius in sample space), each + voxel is the fraction of the collected light passing the pinhole, for a point + source at that position relative to the focus. + """ + psf = vectorial_psf_centered( nz=nz, - wvl=em_wvl_um, + wvl=wvl_um, dz=dz, pz=pz, nx=nx, @@ -316,32 +350,18 @@ def make_confocal_psf( sf=sf, normalize=False, ) - em_psf = _norm_to_collected_light(em_psf, em_wvl_um, dxy, objective, sf, xp) - - # The effective emission PSF is the regular emission PSF convolved with the - # pinhole mask. The pinhole mask is a disk with diameter equal to the pinhole - # size in AU, converted to pixels. - pinhole = _pinhole_mask( - nxy=ex_psf.shape[-1], - pinhole_au=pinhole_au, - wvl=em_wvl_um, - na=objective.numerical_aperture, - dxy=dxy, - xp=xp, - ) - pinhole = xp.asarray(pinhole) - - eff_em_psf = xp.empty_like(em_psf) - for i in tqdm.trange(len(em_psf), desc="convolving em_psf with pinhole..."): - plane = xp.fftconvolve(xp.asarray(em_psf[i]), pinhole, mode="same") - eff_em_psf = xp._array_assign(eff_em_psf, i, plane) + psf = _norm_to_collected_light(psf, wvl_um, dxy, objective, sf, xp) + if pinhole_radius_um is None: + return psf - # The final PSF is the excitation PSF multiplied by the effective emission PSF. - # Before normalization, this is the relative excitation times the fraction of - # collected emission that passes the pinhole, for a fluorophore at each position. - out = xp.asarray(ex_psf) * eff_em_psf - out = _norm_psf(out, normalize, xp) - return out + # light through the pinhole: the image of the point source, summed over the + # pinhole area (i.e. convolved with the pinhole mask) + pinhole = xp.asarray(_pinhole_mask(psf.shape[-1], pinhole_radius_um, dxy)) + out = xp.empty_like(psf) + for i in tqdm.trange(len(psf), desc="convolving em_psf with pinhole..."): + plane = xp.fftconvolve(xp.asarray(psf[i]), pinhole, mode="same") + out = xp._array_assign(out, i, plane) + return out # type: ignore [no-any-return] # lateral radius (in Airy units) of the in-focus plane used to normalize emission @@ -363,7 +383,7 @@ def _norm_to_collected_light( conserves energy across planes, so every plane holds the same total, and dividing by it makes each plane sum to ~1: a camera focused at any depth receives all the light collected from the point source. (Collection efficiency itself is applied - to the emission rates, in `Simulation.filtered_emission_rates`.) + to the detected fraction of emitted light, in `ChannelFluorescence.create`.) The total is the sum of the in-focus plane over at least `_NORM_RADIUS_AU`, rather than over `psf` itself: if `psf` is laterally truncated, its own planes miss some @@ -406,158 +426,131 @@ def _norm_psf( return psf -def _pinhole_mask( - nxy: int, - pinhole_au: float, - wvl: float, - na: float, - dxy: float, - xp: NumpyAPI | None = None, -) -> npt.NDArray: - """Create a 2D circular pinhole mask of specified `pinhole_au`.""" - xp = NumpyAPI.create(xp) +# minimum pinhole radius, in subpixels, used to compute the open area of pixels at +# the pinhole edge (area error < 0.5%) +_PINHOLE_RADIUS_SUBPX = 64 - pinhole_size = pinhole_au * 0.61 * wvl / na - pinhole_px = pinhole_size / dxy - x = xp.arange(nxy) - nxy // 2 - xx, yy = xp.meshgrid(x, x) - r = xp.sqrt(xx**2 + yy**2) - return (r <= pinhole_px).astype(int) # type: ignore +def _pinhole_mask(nxy: int, radius_um: float, dxy: float) -> npt.NDArray: + """Create a 2D circular pinhole mask of radius `radius_um` (in sample space). + Each pixel is the fraction of its area inside the pinhole, so the total open + area matches the pinhole's (a pinhole only a few pixels across would otherwise + snap to whole pixels, e.g. 9 px instead of 11.3 px for a radius of 1.9 px). + """ + radius_px = radius_um / dxy + mask = np.zeros((nxy, nxy)) + center = nxy // 2 + if radius_px < 0.5: # the pinhole lies within the central pixel + mask[center, center] = np.pi * radius_px**2 + return mask + + # only pixels near the pinhole can be (partly) open + half = min(int(np.ceil(radius_px)) + 1, (nxy - 1) // 2) + x = np.arange(-half, half + 1) + ss = max(4, int(np.ceil(_PINHOLE_RADIUS_SUBPX / radius_px))) + sub = (x[:, None] + (np.arange(ss) + 0.5) / ss - 0.5).ravel() + inside = np.hypot(sub[:, None], sub[None, :]) <= radius_px + local = inside.reshape(x.size, ss, x.size, ss).mean(axis=(1, 3)) + mask[center - half : center + half + 1, center - half : center + half + 1] = local + return mask -def make_psf( - nz: int, - nx: int, - dx: float, - dz: float, - objective: ObjectiveLens, - ex_wvl_nm: float | None = None, - em_wvl_nm: float | None = None, - pinhole_au: float | None = None, - max_au_relative: float | None = None, - saturation: float = 0, - xp: NumpyAPI | None = None, + +@lru_cache(maxsize=Settings().cache.in_mem_size.psf) +def excitation_spot( + grid: PSFGrid, wvl_nm: float, objective: ObjectiveLens, xp: NumpyAPI ) -> ArrayProtocol: - if ex_wvl_nm is None: - if em_wvl_nm is None: - raise ValueError( - "Either excitation or emission must be specified to generate a PSF." - ) - ex_wvl_nm = em_wvl_nm - if em_wvl_nm is None: - em_wvl_nm = ex_wvl_nm - - return cached_psf( - nz=nz, - nx=nx, - dx=dx, - dz=dz, - ex_wvl_um=ex_wvl_nm * 1e-3, # nm to um - em_wvl_um=em_wvl_nm * 1e-3, # nm to um - objective=_cast_objective(objective), - pinhole_au=pinhole_au, - max_au_relative=max_au_relative, - saturation=saturation, - xp=NumpyAPI.create(xp), + """Relative excitation intensity of a focused laser spot (1 at the focus). + + The result is cached (in memory, and on disk), so it must not be modified. + """ + return _disk_cached( + ("spot", wvl_nm), + grid, + objective, + xp, + lambda: _excitation_spot( + grid.nz, wvl_nm * 1e-3, grid.dz, grid.nx, grid.dx, objective, xp + ), ) -# variant of make_psf that only accepts hashable arguments @lru_cache(maxsize=Settings().cache.in_mem_size.psf) -def cached_psf( - nz: int, - nx: int, - dx: float, - dz: float, - ex_wvl_um: float, - em_wvl_um: float, +def detection_psf( + grid: PSFGrid, + wvl_nm: float, objective: ObjectiveLens, - pinhole_au: float | None, - max_au_relative: float | None, xp: NumpyAPI, - saturation: float = 0, + pinhole_radius_um: float | None = None, ) -> ArrayProtocol: - # now restrict nx to no more than max_au_relative - if max_au_relative is not None: - airy_radius = 0.61 * ex_wvl_um / objective.numerical_aperture - n_pix_per_airy_radius = airy_radius / dx - max_nx = int(n_pix_per_airy_radius * max_au_relative * 2) - nx = min(nx, max_nx) + """Fraction of the light collected from a point source that reaches the detector. - # odd sizes, so that the PSF is centered on a pixel (an even-sized PSF shifts - # the image by a fraction of a pixel when convolved in `mode="same"`) - nz, nx = nz | 1, nx | 1 + For a camera (no pinhole), each z-plane is the image of a point source at that + distance from the focal plane, as a fraction of the light collected from it. + With a pinhole (radius in um, in sample space), each voxel is the fraction of the + collected light that passes the pinhole, for a point source at that position. - use_cache = os.getenv("MICROSIM_CACHE", "").lower() not in {"0", "false", "no", "n"} - if use_cache: - cache_path = _psf_cache_path( - nz, nx, dz, dx, em_wvl_um, pinhole_au, ex_wvl_um, objective, saturation - ) - if cache_path.exists(): - logger.info( - f"Found cached PSF {nz=} {nx=} {dz=} {dx=} " - f"{ex_wvl_um=:.2f} {pinhole_au=}" - ) - # the file may have been written by a simulation of another precision - return xp.asarray(np.load(cache_path), dtype=xp.float_dtype) - - logger.info( - f"Creating new PSF {nz=} {nx=} {dz=} {dx=} {ex_wvl_um=:.2f} {pinhole_au=}" + The result is cached (in memory, and on disk), so it must not be modified. + """ + name: tuple[str | float, ...] = ("detection", wvl_nm) + if pinhole_radius_um is not None: + name = (*name, "pinhole", pinhole_radius_um) + return _disk_cached( + name, + grid, + objective, + xp, + lambda: _detection_psf( + grid.nz, + wvl_nm * 1e-3, + grid.dz, + grid.nx, + grid.dx, + objective, + xp, + pinhole_radius_um=pinhole_radius_um, + ), ) - if pinhole_au is None: - psf = vectorial_psf_centered( - wvl=em_wvl_um, - nz=nz, - nx=nx, - dz=dz, - dxy=dx, - objective=objective, - xp=xp, - normalize=False, - ) - psf = _norm_to_collected_light(psf, em_wvl_um, dx, objective, 3, xp) - else: - psf = make_confocal_psf( - nz=nz, - ex_wvl_um=ex_wvl_um, - em_wvl_um=em_wvl_um, - pinhole_au=pinhole_au, - nx=nx, - dz=dz, - dxy=dx, - objective=objective, - xp=xp, - normalize=False, # keep physical scale (pinhole throughput) - saturation=saturation, - ) + +def _disk_cached( + name: tuple[str | float, ...], + grid: PSFGrid, + objective: ObjectiveLens, + xp: NumpyAPI, + compute: Callable[[], ArrayProtocol], +) -> ArrayProtocol: + """Load a PSF array from the disk cache, or compute (and cache) it.""" + use_cache = os.getenv("MICROSIM_CACHE", "").lower() not in {"0", "false", "no", "n"} + path = _psf_cache_path(name, grid, objective) + if use_cache and path.exists(): + logger.info(f"Found cached PSF {path.stem}") + # the file may have been written by a simulation of another precision + return _read_only(xp.asarray(np.load(path), dtype=xp.float_dtype)) + + logger.info(f"Creating new PSF {path.stem}") + psf = compute() if use_cache: - cache_path.parent.mkdir(parents=True, exist_ok=True) - np.save(cache_path, psf) - return xp.asarray(psf) + path.parent.mkdir(parents=True, exist_ok=True) + np.save(path, psf) + return _read_only(xp.asarray(psf)) + + +def _read_only(psf: ArrayProtocol) -> ArrayProtocol: + """Protect a cached (shared) numpy array from in-place modification.""" + if isinstance(psf, np.ndarray): + psf.setflags(write=False) + return psf def _psf_cache_path( - nz: int, - nx: int, - dz: float, - dx: float, - em_wvl_um: float, - pinhole_au: float | None, - ex_wvl_um: float, - objective: ObjectiveLens, - saturation: float = 0, + name: tuple[str | float, ...], grid: PSFGrid, objective: ObjectiveLens ) -> Path: - """Return the cache location for these PSF parameters.""" - # "v2": PSFs are scaled to collected light, no longer sum-normalized - # "v3": PSF shape is (nz, nx, nx), rounded up to odd (widefield was nz + 1, ...), - # and normalization no longer depends on the lateral size (truncation) - cache_key: list[float | str] = [nz, nx, dz, dx, em_wvl_um, "v3"] - if pinhole_au is not None: - cache_key.extend([ex_wvl_um, pinhole_au]) - if saturation: - cache_key.append(f"sat{saturation:.6g}") - cache_path = microsim_cache("psf") / objective.cache_key() - cache_path = cache_path / "_".join([str(x).replace(".", "-") for x in cache_key]) - return cache_path.with_suffix(".npy") + """Return the cache location for a PSF array.""" + # "v4": PSFs are cached as separate excitation spots and detection PSFs + # "v5": anti-aliased pinhole mask + # e.g. 520 and 520.0 nm share a file + name = tuple(float(x) if isinstance(x, int | float) else x for x in name) + parts = [*name, grid.nz, grid.nx, grid.dz, grid.dx, "v5"] + stem = "_".join(str(x).replace(".", "-") for x in parts) + return (microsim_cache("psf") / objective.cache_key() / stem).with_suffix(".npy") diff --git a/src/microsim/psf.py b/src/microsim/psf.py index 56b67e49..4383fc39 100644 --- a/src/microsim/psf.py +++ b/src/microsim/psf.py @@ -3,10 +3,12 @@ The implementation lives in `microsim._psf`. """ +from __future__ import annotations + +from typing import TYPE_CHECKING, Any, Literal + +from microsim import _psf from microsim._psf import ( - cached_psf, - make_confocal_psf, - make_psf, radius_map, rz_to_xyz, simpson, @@ -14,9 +16,15 @@ vectorial_psf_centered, vectorial_rz, ) +from microsim.schema.backend import NumpyAPI + +if TYPE_CHECKING: + import numpy as np + + from microsim._data_array import ArrayProtocol + from microsim.schema.lens import ObjectiveKwargs, ObjectiveLens __all__ = [ - "cached_psf", "make_confocal_psf", "make_psf", "radius_map", @@ -26,3 +34,101 @@ "vectorial_psf_centered", "vectorial_rz", ] + + +def make_confocal_psf( + nz: int, + ex_wvl_um: float = 0.475, + em_wvl_um: float = 0.525, + pinhole_au: float = 1.0, + dz: float = 0.05, + pz: float = 0, + nx: int = 31, + ny: int | None = None, + pos: tuple[float, float, float] = (0, 0, 0), + dxy: float = 0.05, + objective: ObjectiveKwargs | ObjectiveLens | None = None, + sf: int = 3, + normalize: bool | Literal["sum", "max"] = "sum", + saturation: float = 0, + xp: NumpyAPI | None = None, +) -> np.ndarray: + """Create a confocal PSF. + + This function creates a confocal system PSF: the excitation spot times the + detection PSF (the emission PSF convolved with a pinhole mask). `pinhole_au` is + in Airy units at the emission wavelength. + + If `normalize` is False, the result is physically scaled: the relative excitation + intensity (1 at focus) times the fraction of collected emission that passes the + pinhole, for a point source at each position. + + `saturation` is the excitation saturation parameter at the focus (`k * tau`: the + absorption rate times the excited-state lifetime). If non-zero, the relative + excitation `r` is replaced by the relative emission of a two-level system, + `r / (1 + s * r)`. + + All extra keyword arguments are passed to `vectorial_psf_centered`. + """ + xp = NumpyAPI.create(xp) + objective = _psf._cast_objective(objective) + kw: dict[str, Any] = {"nz": nz, "dz": dz, "pz": pz, "nx": nx, "ny": ny} + kw.update(pos=pos, dxy=dxy, objective=objective, sf=sf, xp=xp) + spot = _psf.apply_saturation( + _psf._excitation_spot(wvl_um=ex_wvl_um, **kw), saturation + ) + pinhole_um = pinhole_au * _psf.airy_radius_um( + em_wvl_um, objective.numerical_aperture + ) + det = _psf._detection_psf(wvl_um=em_wvl_um, pinhole_radius_um=pinhole_um, **kw) + return _psf._norm_psf(xp.asarray(spot) * det, normalize, xp) + + +def make_psf( + nz: int, + nx: int, + dx: float, + dz: float, + objective: ObjectiveLens, + ex_wvl_nm: float | None = None, + em_wvl_nm: float | None = None, + pinhole_au: float | None = None, + max_au_relative: float | None = None, + saturation: float = 0, + xp: NumpyAPI | None = None, +) -> ArrayProtocol: + """Make a widefield PSF, or (if `pinhole_au` is given) a confocal PSF. + + Results are cached (in memory, and on disk). `pinhole_au` is + in Airy units at the emission wavelength, and `max_au_relative` (if given) caps + the lateral size of the PSF, in Airy units at the excitation wavelength. + """ + if ex_wvl_nm is None: + if em_wvl_nm is None: + raise ValueError( + "Either excitation or emission must be specified to generate a PSF." + ) + ex_wvl_nm = em_wvl_nm + if em_wvl_nm is None: + em_wvl_nm = ex_wvl_nm + ex_wvl_nm, em_wvl_nm = float(ex_wvl_nm), float(em_wvl_nm) + + objective = _psf._cast_objective(objective) + na = objective.numerical_aperture + if max_au_relative is not None: + airy_px = _psf.airy_radius_um(ex_wvl_nm * 1e-3, na) / dx + nx = min(nx, int(airy_px * max_au_relative * 2)) + grid = _psf.PSFGrid(nz=nz, nx=nx, dz=dz, dx=dx) + xp = NumpyAPI.create(xp) + + if pinhole_au is None: + # a copy: the cached PSF must not be modified + return _psf.detection_psf(grid, em_wvl_nm, objective, xp).copy() # type: ignore [attr-defined, no-any-return] + pinhole_um = pinhole_au * _psf.airy_radius_um(em_wvl_nm * 1e-3, na) + det = _psf.detection_psf( + grid, em_wvl_nm, objective, xp, pinhole_radius_um=pinhole_um + ) + spot = _psf.apply_saturation( + _psf.excitation_spot(grid, ex_wvl_nm, objective, xp), saturation + ) + return spot * det # type: ignore [no-any-return] diff --git a/src/microsim/schema/modality/_optics.py b/src/microsim/schema/modality/_optics.py new file mode 100644 index 00000000..ff3d2625 --- /dev/null +++ b/src/microsim/schema/modality/_optics.py @@ -0,0 +1,140 @@ +"""How a modality illuminates the sample, and how it detects the emitted light. + +Every modality in this package forms its image the same way: the signal from one +fluorophore, at some position relative to the focus, is + + (how strongly it is excited there) x (how much of its emission is detected) + +i.e. the excitation spot times the detection PSF: the *system PSF*. The +`Illumination` classes describe the first factor and the `Detection` classes the +second. A modality picks one of each. (In widefield, excitation is uniform, so the +system PSF is just the detection PSF.) + +Everything here is internal (not part of the user-facing schema): modalities build +these objects themselves. They are frozen dataclasses, so they can be used directly +in PSF cache keys. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import TYPE_CHECKING, Any, ClassVar + +import numpy as np + +from microsim._psf import detection_psf, excitation_spot + +if TYPE_CHECKING: + from microsim._data_array import ArrayProtocol + from microsim._psf import PSFGrid + from microsim.schema.backend import NumpyAPI + from microsim.schema.lens import ObjectiveLens + +# (wavelength_nm, weight) pairs: a spectrum discretized into a few wavelengths +SpectralLines = tuple[tuple[float, float], ...] + + +@dataclass(frozen=True) +class WidefieldIllumination: + """The whole field of view is evenly illuminated.""" + + def spot( + self, + excitation: SpectralLines, + grid: PSFGrid, + objective: ObjectiveLens, + xp: NumpyAPI, + ) -> float: + """Relative excitation intensity: the same (1) everywhere.""" + return 1.0 + + +@dataclass(frozen=True) +class FocusedIllumination: + """Excitation light is focused to a diffraction-limited spot (point scanning).""" + + def spot( + self, + excitation: SpectralLines, + grid: PSFGrid, + objective: ObjectiveLens, + xp: NumpyAPI, + ) -> ArrayProtocol: + """Relative excitation intensity at each position (1 at the focus). + + With several excitation wavelengths (e.g. two lasers, or a broadband source), + each wavelength's spot is weighted by how much of the fluorophore's total + absorption it accounts for (the `excitation` weights). + """ + if len(excitation) == 1: + return excitation_spot(grid, excitation[0][0], objective, xp) + total = sum(weight for _, weight in excitation) + spot: Any = None + for wvl_nm, weight in excitation: + line = excitation_spot(grid, wvl_nm, objective, xp) * (weight / total) + if spot is None: + spot = line + else: + spot += line + return spot # type: ignore [no-any-return] + + +@dataclass(frozen=True) +class CameraDetection: + """The sample is imaged onto a camera.""" + + # each camera pixel collects all of the light falling on its area + sums_over_pixel_area: ClassVar[bool] = True + + def psf( + self, wvl_nm: float, grid: PSFGrid, objective: ObjectiveLens, xp: NumpyAPI + ) -> ArrayProtocol: + """Fraction of the light collected from a point source, on each pixel.""" + return detection_psf(grid, wvl_nm, objective, xp) + + def psf_radius_aus( + self, depth_um: float, wvl_um: float, objective: ObjectiveLens + ) -> float: + """Lateral PSF radius (in Airy units) for `max_psf_radius_aus="auto"`. + + Calibrated to ~1% RMSE vs. an untruncated PSF: ~14 AU covers the tails of the + in-focus Airy pattern, plus half the radius of the cone of out-of-focus light + at the axial edge of the PSF (half the volume depth). + """ + airy_um = 0.61 * wvl_um / objective.numerical_aperture + cone_aus = depth_um / 2 * np.tan(objective.half_angle) / airy_um + return float(14 + cone_aus / 2) + + +@dataclass(frozen=True) +class PinholeDetection: + """A single detector behind a pinhole, which records one scan position at a time. + + `radius_um` is the radius of the pinhole, projected into the sample (i.e. its + physical radius divided by the magnification). It is the same for all + fluorophores and wavelengths in a channel. + """ + + radius_um: float + + # each pixel is the signal recorded at one scan position + sums_over_pixel_area: ClassVar[bool] = False + + def psf( + self, wvl_nm: float, grid: PSFGrid, objective: ObjectiveLens, xp: NumpyAPI + ) -> ArrayProtocol: + """Fraction of collected light passing the pinhole, at each source position.""" + return detection_psf( + grid, wvl_nm, objective, xp, pinhole_radius_um=self.radius_um + ) + + def psf_radius_aus( + self, depth_um: float, wvl_um: float, objective: ObjectiveLens + ) -> float: + # the pinhole rejects out-of-focus light, so a fixed radius suffices + # (<0.3% RMSE vs. an untruncated PSF, for volumes up to ~5 um deep) + return 8 + + +Illumination = WidefieldIllumination | FocusedIllumination +Detection = CameraDetection | PinholeDetection diff --git a/src/microsim/schema/modality/_simple_psf.py b/src/microsim/schema/modality/_simple_psf.py index fb706e06..6560cf16 100644 --- a/src/microsim/schema/modality/_simple_psf.py +++ b/src/microsim/schema/modality/_simple_psf.py @@ -1,99 +1,104 @@ +from __future__ import annotations + import warnings -from collections.abc import Hashable -from typing import Annotated, Any, ClassVar, Literal +from typing import TYPE_CHECKING, Annotated, Any, ClassVar, Literal import numpy as np from annotated_types import Ge +from pydantic import Field -from microsim._data_array import ArrayProtocol, DataArray, xrDataArray +from microsim._data_array import DataArray, xrDataArray from microsim._logger import logger, logging_indented -from microsim._psf import make_psf +from microsim._psf import PSFGrid, airy_radius_um, apply_saturation from microsim.fft.convolve import fftconvolve_same_pairs from microsim.schema._base_model import SimBaseModel -from microsim.schema.backend import NumpyAPI from microsim.schema.dimensions import Axis -from microsim.schema.lens import ObjectiveLens -from microsim.schema.settings import Settings -from microsim.schema.space import SpaceProtocol + +from ._optics import ( + CameraDetection, + Detection, + FocusedIllumination, + Illumination, + PinholeDetection, + SpectralLines, + WidefieldIllumination, +) + +if TYPE_CHECKING: + from collections.abc import Hashable, Sequence + + from microsim._data_array import ArrayProtocol + from microsim.schema.backend import NumpyAPI + from microsim.schema.lens import ObjectiveLens + from microsim.schema.optical_config._fluorescence import ChannelFluorescence + from microsim.schema.settings import Settings + from microsim.schema.space import SpaceProtocol class _PSFModality(SimBaseModel): - # whether the image is formed by scanning a focused spot (vs. a camera) - point_scanning: ClassVar[bool] = False + # how the image is detected (see `digital_image`, for how pixels are formed) + detection_type: ClassVar[type[Detection]] = CameraDetection - def psf( - self, - *, - nz: int, - nx: int, - dx: float, - dz: float, - objective_lens: ObjectiveLens, - xp: NumpyAPI, - ex_wvl_nm: float | None = None, - em_wvl_nm: float | None = None, - saturation: float = 0, - ) -> ArrayProtocol: - # default implementation is a widefield PSF (uniform illumination, so - # `saturation` is already applied to the emission rates and ignored here) - return make_psf( - nz=nz, - nx=nx, - dx=dx, - dz=dz, - objective=objective_lens, - ex_wvl_nm=ex_wvl_nm, - em_wvl_nm=em_wvl_nm, - xp=xp, - ) + def _illumination(self) -> Illumination: + return WidefieldIllumination() + + def _detection( + self, channel: Sequence[ChannelFluorescence], objective: ObjectiveLens + ) -> Detection: + """Detection for one channel (`channel` has one item per fluorophore).""" + return CameraDetection() def render( self, truth: xrDataArray, # (F, Z, Y, X) - em_rates: xrDataArray, # (C, F, W) + fluorescence: Sequence[Sequence[ChannelFluorescence]], # [channel][fluor] objective_lens: ObjectiveLens, settings: Settings, xp: NumpyAPI, ) -> xrDataArray: - """Render a 3D image of the truth for F fluorophores, in C channels.""" - space = truth.attrs["space"] + """Render a 3D image of the truth for F fluorophores, in C channels. - # summed (spectrally weighted) PSF for every channel/fluorophore pair - psfs: dict[tuple[int, int], tuple[Hashable, ArrayProtocol]] = {} - for c_idx, ch in enumerate(em_rates.coords[Axis.C].values): - logger.info(f"Rendering {type(self).__name__} channel {ch} ---------------") + `fluorescence[c]` describes each fluorophore in channel `c`. The result is + in photons/s (detected), with dimensions (C, F, Z, Y, X). + """ + space = truth.attrs["space"] + # system PSF (and brightness) for every channel/fluorophore pair + system_psfs: dict[tuple[int, int], tuple[Hashable, ArrayProtocol]] = {} + brightness: dict[tuple[int, int], float] = {} + for c_idx, channel in enumerate(fluorescence): + logger.info(f"Rendering {type(self).__name__} channel {channel[0].channel}") with logging_indented(): + detection: Detection | None = None + by_fluor = {fl.fluorophore: fl for fl in channel} for f_idx, fluor in enumerate(truth.coords[Axis.F].values): + fl = by_fluor[fluor] + if fl.detected_rate(saturate=False) <= 1e-12: + continue # this fluorophore is not seen in this channel + if detection is None: # (only needed if something is seen) + detection = self._detection(channel, objective_lens) logger.info(f"Fluor: {fluor}") with logging_indented(): - # discretize the em spectrum for this specific ch/fluor pair - em_spectrum = em_rates.sel({Axis.C: ch, Axis.F: fluor}) - # if we happen to have 2 spectra for the same fluorophore - # in the same channel, just take the first one (shouldn't happen - if Axis.F in em_spectrum.dims: # pragma: no cover - em_spectrum = em_spectrum.isel({Axis.F: 0}) - - if not (em_spectrum > 1e-12).any(): - # no emission at all for this fluorophore in this channel - continue - - psfs[(c_idx, f_idx)] = self._summed_weighted_psf( - em_spectrum, settings, space, objective_lens, xp + key, psf, bright = self._system_psf( + fl, detection, space, objective_lens, settings, xp ) + system_psfs[(c_idx, f_idx)] = (key, psf) + brightness[(c_idx, f_idx)] = bright # convolve each fluorophore's truth with its PSF in each channel -> [c][f] - channels = fftconvolve_same_pairs( + images = fftconvolve_same_pairs( [truth.data[f_idx] for f_idx in range(truth.sizes[Axis.F])], - psfs, - n_out=em_rates.sizes[Axis.C], + system_psfs, + n_out=len(fluorescence), xp=xp, cache_mb=settings.cache.in_mem_size.otf_mb, ) + for (c_idx, f_idx), bright in brightness.items(): + images[c_idx][f_idx] = images[c_idx][f_idx] * bright return DataArray( - [xp.stack(fluors, axis=0) for fluors in channels], + [xp.stack(fluors, axis=0) for fluors in images], dims=[Axis.C, Axis.F, Axis.Z, Axis.Y, Axis.X], coords={ - Axis.C: em_rates.coords[Axis.C], + Axis.C: [channel[0].channel for channel in fluorescence], Axis.F: truth.coords[Axis.F], Axis.Z: truth.coords[Axis.Z], Axis.Y: truth.coords[Axis.Y], @@ -106,93 +111,88 @@ def render( }, ) - def _summed_weighted_psf( + def _system_psf( self, - em_spectrum: xrDataArray, - settings: Settings, + fl: ChannelFluorescence, + detection: Detection, space: SpaceProtocol, objective_lens: ObjectiveLens, + settings: Settings, xp: NumpyAPI, - ) -> tuple[Hashable, ArrayProtocol]: - """Create a weighted sum of PSFs based on the emission spectrum. - - This takes advantage of the distributive property of convolution - (a * b) * c = a * (b * c) - We create a PSF for each emission wavelength, multiply it by the - emission rate at that wavelength, and sum them all together, prior - to convolving with the truth. - This creates a more realistic PSF for the fluorophore/channel, as it - accounts for the full emission spectrum. - - Returns a hashable key that uniquely identifies the summed PSF, and the PSF - itself. + ) -> tuple[Hashable, ArrayProtocol, float]: + """Return the system PSF of one fluorophore in one channel, and its brightness. + + The system PSF is the excitation spot (relative excitation, 1 at the focus) + times the detection PSF (fraction of collected light detected). The image of + the fluorophore is `brightness * (truth (*) system_psf)`, where `brightness` + is in detected photons/s per fluorophore. Also returns a key that uniquely + identifies the system PSF, for caching. + + Emission is spread over several wavelengths (see + `settings.spectral_bins_per_emission_channel`), so the detection PSF is a + weighted sum of PSFs at each. For focused illumination, the excitation spot + is likewise a weighted sum over the excitation wavelengths. """ - binned = bin_spectrum( - em_spectrum, + emission = emission_lines( + fl.detected_fraction, bins=settings.spectral_bins_per_emission_channel, threshold_percentage=settings.spectral_bin_threshold_percentage, ) - - # we need to pick a single nx size for all psfs we will sum, based on the - # maximum wavelength in the emission spectrum and `settings.max_psf_radius_aus` + excitation = excitation_lines( + fl.absorption_rate, + threshold_percentage=settings.spectral_bin_threshold_percentage, + bin_nm=settings.excitation_bin_width_nm, + ) + grid = self._psf_grid(space, emission, detection, objective_lens, settings) + # round, to reuse the PSF for (practically) identical saturation levels + saturation = float(f"{fl.saturation:.4g}") if fl.saturation > 1e-4 else 0.0 + + logger.info(f"Need PSF {grid} @ emission {emission} nm") + detected: Any = None + for wvl_nm, weight in emission: + psf = detection.psf(wvl_nm, grid, objective_lens, xp) * weight + if detected is None: + detected = psf # a new array (cached PSFs are read-only) + else: + detected += psf + + key: tuple = (detection, grid, emission, objective_lens.cache_key(), type(xp)) + spot = self._illumination().spot(excitation, grid, objective_lens, xp) + if isinstance(spot, float): + # uniform illumination: excitation (and saturation) is just a factor + brightness = fl.detected_rate(saturate=False) * apply_saturation( + spot, fl.saturation + ) + return key, detected, brightness + key = (*key, excitation, saturation) + # apply_saturation(spot, saturation) * detected, without temporary copies + detected *= spot + if saturation: + denominator: Any = spot * saturation + denominator += 1 + detected /= denominator + return key, detected, fl.detected_rate(saturate=False) + + def _psf_grid( + self, + space: SpaceProtocol, + emission: SpectralLines, + detection: Detection, + objective_lens: ObjectiveLens, + settings: Settings, + ) -> PSFGrid: + # one PSF size for all wavelengths, based on the longest emission wavelength + # and `settings.max_psf_radius_aus` nz, _ny, _nx = space.shape dz, _dy, dx = space.scale - max_wave_um = binned.coords[Axis.W].max().item() / 1000 + max_wave_um = max(wvl for wvl, _ in emission) / 1000 radius_aus = settings.max_psf_radius_aus if radius_aus == "auto": - radius_aus = self._auto_psf_radius_aus(nz * dz, max_wave_um, objective_lens) + radius_aus = detection.psf_radius_aus(nz * dz, max_wave_um, objective_lens) logger.info(f"Auto PSF radius: {radius_aus:.1f} AU") - nx = _pick_nx( - _nx, dx, radius_aus, max_wave_um, objective_lens.numerical_aperture - ) - - saturation = self._saturation_parameter(em_spectrum) - components: list[tuple[float, float]] = [] - for em_rate in binned: - em_wvl_nm = em_rate.w.item() - if em_rate.isnull().any() or em_rate == 0 or xp.isnan(em_wvl_nm): - continue - components.append((em_wvl_nm, em_rate.item())) - - key = ( - self.model_dump_json(), - objective_lens.cache_key(), - (nz, nx, dz, dx, saturation), - tuple(components), - type(xp), - ) - summed_psf: Any = 0 - for em_wvl_nm, weight in components: - logger.info(f"Need PSF ({nz},{nx}) @ {em_wvl_nm:.1f} nm ({weight=:.2f})") - psf = self.psf( - nz=nz, - nx=nx, - dx=dx, - dz=dz, - objective_lens=objective_lens, - em_wvl_nm=em_wvl_nm, - saturation=saturation, - xp=xp, - ) - summed_psf += psf * weight - return key, summed_psf - - def _saturation_parameter(self, em_spectrum: xrDataArray) -> float: - """Saturation parameter to apply in the PSF (0 if not applied spatially).""" - return 0.0 - - def _auto_psf_radius_aus( - self, depth_um: float, wvl_um: float, objective_lens: ObjectiveLens - ) -> float: - """Lateral PSF radius (in Airy units) for `max_psf_radius_aus="auto"`. - - Calibrated (for widefield) to ~1% RMSE vs. an untruncated PSF: ~14 AU covers - the tails of the in-focus Airy pattern, plus half the radius of the cone of - out-of-focus light at the axial edge of the PSF (half the volume depth). - """ - airy_um = 0.61 * wvl_um / objective_lens.numerical_aperture - cone_aus = depth_um / 2 * np.tan(objective_lens.half_angle) / airy_um - return float(14 + cone_aus / 2) + na = objective_lens.numerical_aperture + nx = _pick_nx(_nx, dx, radius_aus, max_wave_um, na) + return PSFGrid(nz=nz, nx=nx, dz=dz, dx=dx) class Confocal(_PSFModality): @@ -202,57 +202,47 @@ class Confocal(_PSFModality): position (relative to the scan spot) passes the pinhole, times the relative excitation intensity there. Light source `power` is therefore interpreted as the peak irradiance at the focus, and `exposure_ms` as the per-pixel integration - (dwell) time (it is not divided by the number of pixels). Excitation saturation - is applied locally (per position in the excitation PSF), rather than to the - emission rates. + (dwell) time (it is not divided by the number of pixels). The excitation spot is + computed at the excitation wavelengths, and excitation saturation is applied + locally (it is strongest at the center of the spot). """ type: Literal["confocal"] = "confocal" - point_scanning: ClassVar[bool] = True pinhole_au: Annotated[float, Ge(0)] = 1 - - def _saturation_parameter(self, em_spectrum: xrDataArray) -> float: - # emission rates are unsaturated for point-scanning modalities; - # saturation is applied locally to the excitation PSF instead. - oc = em_spectrum.coords[Axis.C].item() - fluor = em_spectrum.coords[Axis.F].item() - s = oc.saturation_parameter(fluor) - # ignore negligible saturation (reuses the unsaturated PSF), and round to - # keep PSF cache keys stable - return float(f"{s:.4g}") if s > 1e-4 else 0.0 - - def _auto_psf_radius_aus( - self, depth_um: float, wvl_um: float, objective_lens: ObjectiveLens - ) -> float: - # the pinhole rejects out-of-focus light, so a fixed radius suffices - # (<0.3% RMSE vs. an untruncated PSF, for volumes up to ~5 um deep) - return 8 - - def psf( - self, - *, - nz: int, - nx: int, - dx: float, - dz: float, - objective_lens: ObjectiveLens, - xp: NumpyAPI, - ex_wvl_nm: float | None = None, - em_wvl_nm: float | None = None, - saturation: float = 0, - ) -> ArrayProtocol: - return make_psf( - nz=nz, - nx=nx, - dx=dx, - dz=dz, - objective=objective_lens, - em_wvl_nm=em_wvl_nm, - ex_wvl_nm=ex_wvl_nm, - pinhole_au=self.pinhole_au, - saturation=saturation, - xp=xp, + airy_unit_wavelength_nm: float | None = Field( + default=None, + gt=0, + description="Wavelength (nm) at which `pinhole_au` is defined. By default, " + "1 AU is matched to the center of each channel's emission filter band.", + ) + + detection_type = PinholeDetection + + def _illumination(self) -> Illumination: + return FocusedIllumination() + + def _detection( + self, channel: Sequence[ChannelFluorescence], objective: ObjectiveLens + ) -> Detection: + if (wvl_nm := self.airy_unit_wavelength_nm) is None: + wvl_nm = channel[0].channel.emission_band_center() + if wvl_nm is None: + # no band-limited emission filter (e.g. a longpass): use the center of + # the light detected in this channel instead + wvl_nm = _detected_light_center(channel) + # (warnings are shown once per distinct message) + warnings.warn( + f"Channel {channel[0].channel.name!r} has no band-limited emission " + f"filter to match the pinhole size to; using 1 AU at {wvl_nm:.0f} nm " + "(the center of the detected emission). Set " + "`Confocal.airy_unit_wavelength_nm` to choose.", + stacklevel=2, + ) + radius_um = self.pinhole_au * airy_radius_um( + wvl_nm / 1000, objective.numerical_aperture ) + # round, to keep PSF cache keys stable + return PinholeDetection(radius_um=float(f"{radius_um:.6g}")) class Widefield(_PSFModality): @@ -272,7 +262,7 @@ class Identity(_PSFModality): def render( self, truth: xrDataArray, # (F, Z, Y, X) - em_rates: xrDataArray, # (C, F, W) + fluorescence: Sequence[Sequence[ChannelFluorescence]], # [channel][fluor] objective_lens: ObjectiveLens, settings: Settings, xp: NumpyAPI, @@ -283,8 +273,16 @@ def render( already convolved with the PSF. Therefore, we simply compute the emission flux for each fluorophore and each channel. """ - em_rates = em_rates.sum(Axis.W) - em_rates = em_rates.copy(data=xp.asarray(em_rates.data)) + fluors = truth.coords[Axis.F].values + rates = [] + for channel in fluorescence: + by_fluor = {fl.fluorophore: fl for fl in channel} + rates.append([by_fluor[fluor].detected_rate() for fluor in fluors]) + em_rates = DataArray( + xp.asarray(np.asarray(rates), dtype=truth.dtype), + dims=[Axis.C, Axis.F], + coords={Axis.C: [ch[0].channel for ch in fluorescence]}, + ) em_image = em_rates * truth return DataArray( em_image, @@ -304,6 +302,76 @@ def render( ) +def _detected_light_center(channel: Sequence[ChannelFluorescence]) -> float: + """Mean wavelength of all light detected in a channel (all fluorophores).""" + total: float = 0 + weighted: float = 0 + for fl in channel: + detected = fl.detected_spectrum(saturate=False) + total += float(detected.sum()) + weighted += float((detected * detected.coords[Axis.W]).sum()) + if total <= 0: + raise ValueError( + "Cannot determine the pinhole size: no light is detected in channel " + f"{channel[0].channel.name!r}. Set `Confocal.airy_unit_wavelength_nm`." + ) + return weighted / total + + +def emission_lines( + detection: xrDataArray, *, bins: int, threshold_percentage: float +) -> SpectralLines: + """Discretize an emission spectrum into `bins` `(wavelength, weight)` pairs. + + Each bin is represented by its centroid wavelength, and weights sum to 1. + """ + binned = bin_spectrum( + detection, bins=bins, threshold_percentage=threshold_percentage + ) + lines = [ + (float(wvl), float(weight)) + for wvl, weight in zip(binned.coords[Axis.W].values, binned.values, strict=True) + if weight > 0 and not np.isnan(wvl) + ] + total = sum(weight for _, weight in lines) + return tuple((_round_nm(wvl), float(f"{w / total:.4g}")) for wvl, w in lines) + + +def excitation_lines( + absorption: xrDataArray, *, threshold_percentage: float = 1, bin_nm: float = 10 +) -> SpectralLines: + """Discretize an absorption spectrum into `(wavelength, weight)` pairs. + + Wavelengths absorbing more than `threshold_percentage` of the peak are grouped + into bins at most `bin_nm` wide, each represented by its centroid wavelength. + The weights are each bin's fraction of the total (thresholded) absorption. + Laser lines more than `bin_nm` apart are never grouped together. + """ + wvl = np.asarray(absorption.coords[Axis.W].values, dtype=float) + rate = np.asarray(absorption.values, dtype=float) + keep = rate > rate.max() * threshold_percentage / 100 + wvl, rate = wvl[keep], rate[keep] + + idx = ((wvl - wvl.min()) // bin_nm).astype(int) + sums = np.bincount(idx, weights=rate) + centroids = np.bincount(idx, weights=rate * wvl)[sums > 0] / sums[sums > 0] + weights = sums[sums > 0] / sums.sum() + return tuple( + (_round_nm(w), float(f"{f:.4g}")) + for w, f in zip(centroids, weights, strict=True) + ) + + +# PSF wavelengths are rounded to this step, so that fluorophores with nearly the +# same (centroid) wavelength share PSFs. The PSF width scales with wavelength, so +# this changes it by < 0.1%. +_WAVELENGTH_STEP_NM = 0.5 + + +def _round_nm(wvl_nm: float) -> float: + return float(round(wvl_nm / _WAVELENGTH_STEP_NM) * _WAVELENGTH_STEP_NM) + + def bin_spectrum( spectrum: xrDataArray, bins: int | np.ndarray = 3, diff --git a/src/microsim/schema/optical_config/_fluorescence.py b/src/microsim/schema/optical_config/_fluorescence.py new file mode 100644 index 00000000..a71417d8 --- /dev/null +++ b/src/microsim/schema/optical_config/_fluorescence.py @@ -0,0 +1,104 @@ +from __future__ import annotations + +from dataclasses import dataclass +from typing import TYPE_CHECKING + +from microsim.schema.dimensions import Axis + +if TYPE_CHECKING: + import xarray as xr + + from microsim.schema.detectors import Detector + from microsim.schema.lens import ObjectiveLens + from microsim.schema.sample.fluorophore import Fluorophore + + from .config import OpticalConfig + + +@dataclass(frozen=True, eq=False) +class ChannelFluorescence: + """How bright one fluorophore is in one channel, and at which wavelengths. + + This has no spatial information: it describes a single fluorophore at the + *reference* illumination intensity (`OpticalConfig.power`), which is the intensity + everywhere in a widefield field of view, or at the center of a focused spot. It + is unsaturated: modalities apply saturation, since it depends on where the + fluorophore sits in the illumination. + + Attributes + ---------- + channel : OpticalConfig + The channel (light sources, filters). + fluorophore : Fluorophore + The fluorophore. + absorption_rate : xr.DataArray + Photons absorbed per second by one fluorophore, at each excitation wavelength. + detected_fraction : xr.DataArray + Photons detected per photon absorbed, at each emission wavelength: quantum + yield x emission spectrum x emission filters x detector QE x the fraction of + emitted light collected by the objective. + """ + + channel: OpticalConfig + fluorophore: Fluorophore + absorption_rate: xr.DataArray # (W,) + detected_fraction: xr.DataArray # (W,) + + @classmethod + def create( + cls, + channel: OpticalConfig, + fluorophore: Fluorophore, + objective: ObjectiveLens, + detector: Detector | None = None, + ) -> ChannelFluorescence: + """Describe `fluorophore` in `channel`.""" + absorption_rate = channel.absorption_rate(fluorophore).squeeze( + [Axis.F, Axis.C], drop=True + ) + qe = detector.qe if detector else None + detected = channel.detected_fraction(fluorophore, detector_qe=qe) + # fraction of isotropic emission entering the objective. PSFs are normalized + # to collected light, so this is not double-counted (see `detection_psf`) + detected = detected * objective.collection_efficiency + return cls(channel, fluorophore, absorption_rate, detected) + + @property + def total_absorption_rate(self) -> float: + """Photons absorbed per second by one fluorophore (all wavelengths).""" + return float(self.absorption_rate.sum()) + + @property + def saturation(self) -> float: + """How strongly the fluorophore is saturated at the reference intensity. + + This is the absorption rate times the excited-state lifetime (0 if the + lifetime is unknown). Emission is reduced by a factor of `1 + saturation` + relative to a linear response: at 1, a fluorophore emits half as much as it + would without saturation. + """ + if not (lifetime_ns := self.fluorophore.lifetime_ns): + return 0.0 + return self.total_absorption_rate * lifetime_ns * 1e-9 + + def detected_rate(self, *, saturate: bool = True) -> float: + """Photons detected per second from one fluorophore (all wavelengths). + + If `saturate` is True, saturation at the reference intensity is applied. + """ + rate = self.total_absorption_rate * float(self.detected_fraction.sum()) + return rate / (1 + self.saturation) if saturate else rate + + def detected_spectrum(self, *, saturate: bool = True) -> xr.DataArray: + """Photons detected per second from one fluorophore, at each wavelength. + + If `saturate` is True, saturation at the reference intensity is applied. + """ + rate = self.total_absorption_rate + if saturate: + rate /= 1 + self.saturation + out = self.detected_fraction * rate + out = out.expand_dims({Axis.F: [self.fluorophore], Axis.C: [self.channel]}) + out.name = "filtered_emission_rate" + out.attrs.update(long_name="Filtered Emission rate", units="photons/s") + return out diff --git a/src/microsim/schema/optical_config/config.py b/src/microsim/schema/optical_config/config.py index 268193ae..15a977be 100644 --- a/src/microsim/schema/optical_config/config.py +++ b/src/microsim/schema/optical_config/config.py @@ -2,6 +2,7 @@ from collections.abc import Sequence from typing import TYPE_CHECKING, Any, cast +import numpy as np import xarray as xr from pydantic import Field, model_validator from scipy.constants import c, h @@ -19,6 +20,11 @@ from matplotlib.axes import Axes +# emission bandpass filters are rarely wider than this; a wider band that runs off +# the end of the filter's spectrum is treated as open-ended (e.g. a longpass) +_MAX_EMISSION_BAND_NM = 150 + + class LightSource(SimBaseModel): name: str = "" spectrum: Spectrum @@ -193,6 +199,29 @@ def filtered_emission_rate( final.attrs["units"] = "photons/s" return final + def emission_band_center(self) -> float | None: + """Center (transmission-weighted mean wavelength) of the emission filter band. + + This combines all filters in the emission path. Returns None if there are no + emission filters, or if the band is open-ended: e.g. a longpass filter, which + is still transmitting at the end of its spectrum. (A spectrum that only + covers a narrow passband is fine: beyond its data, transmission is 0.) + """ + if (emission := self.emission) is None: + return None + spectrum = emission.spectrum + order = np.argsort(spectrum.wavelength) + wvl, trans = spectrum.wavelength[order], spectrum.intensity[order] + peak = trans.max() + if peak <= 0: + return None + passing = wvl[trans > peak / 2] + open_end = trans[0] > peak / 2 or trans[-1] > peak / 2 + if open_end and passing.max() - passing.min() > _MAX_EMISSION_BAND_NM: + return None + center = np.trapezoid(wvl * trans, wvl) / np.trapezoid(trans, wvl) + return float(center) + @property def excitation(self) -> Filter | None: """Combine all excitation filters into a single spectrum.""" diff --git a/src/microsim/schema/settings.py b/src/microsim/schema/settings.py index e397e7d7..d2a6113f 100644 --- a/src/microsim/schema/settings.py +++ b/src/microsim/schema/settings.py @@ -88,14 +88,23 @@ class Settings(SimBaseModel, BaseSettings): "realistic, by superposing multiple different PSFs, but will also increase " "memory usage and time.", ) + excitation_bin_width_nm: float = Field( + 10, + gt=0, + description="Maximum width (nm) of the wavelength bins used to represent the " + "excitation light absorbed by each fluorophore, for modalities that focus the " + "excitation (e.g. confocal). Each bin uses a separate excitation spot, so " + "laser lines further apart than this are always simulated separately.", + ) spectral_bin_threshold_percentage: float = Field( 1, - description="Percentage of the total intensity to use as a threshold when " - "binning the emission spectrum. This is used to determine the range of " - "wavelengths to consider when simulating the optical image. First, the " - "spectrum is trimmed to include only relevant wavelengths (those with " + description="Percentage of the peak intensity to use as a threshold when " + "binning the emission (and excitation) spectra. This is used to determine the " + "range of wavelengths to consider when simulating the optical image. First, " + "the spectrum is trimmed to include only relevant wavelengths (those with " "intensity above this threshold). Then, the spectrum is discretized into " - "spectral_bins_per_emission_channel bins.", + "spectral_bins_per_emission_channel bins (or, for excitation, bins of at most " + "excitation_bin_width_nm).", ) model_config: ClassVar[SettingsConfigDict] = SettingsConfigDict( diff --git a/src/microsim/schema/simulation.py b/src/microsim/schema/simulation.py index bec9de04..32e7c801 100644 --- a/src/microsim/schema/simulation.py +++ b/src/microsim/schema/simulation.py @@ -3,6 +3,7 @@ from collections import OrderedDict from collections.abc import Iterable, Iterator, Mapping from contextlib import suppress +from dataclasses import replace from pathlib import Path from typing import TYPE_CHECKING, Annotated, Any @@ -22,6 +23,7 @@ from .lens import ObjectiveLens from .modality import Modality, Widefield from .optical_config import OpticalConfig, Placement +from .optical_config._fluorescence import ChannelFluorescence from .optical_config.lib import FITC from .sample import Fluorophore, FluorophoreDistribution, Sample from .settings import Settings @@ -52,9 +54,9 @@ def _check_extensions(path: Path) -> Path: OutPath = Annotated[Path, AfterValidator(_check_extensions)] xr.set_options(keep_attrs=True) # type: ignore [no-untyped-call] -# LRU cache of `Simulation.filtered_emission_rates` results (small C x F x W arrays) -_RATES_CACHE: OrderedDict[tuple, xr.DataArray] = OrderedDict() -_RATES_CACHE_SIZE = 32 +# LRU cache of `Simulation._channel_fluorescence` results (small C x F x W arrays) +_FLUORESCENCE_CACHE: OrderedDict[tuple, list[list[ChannelFluorescence]]] = OrderedDict() +_FLUORESCENCE_CACHE_SIZE = 32 class Simulation(SimBaseModel): @@ -196,16 +198,56 @@ def ground_truth(self) -> xr.DataArray: ) return self._ground_truth - def filtered_emission_rates(self) -> xr.DataArray: + def _channel_fluorescence(self) -> list[list[ChannelFluorescence]]: + """Return how each fluorophore appears in each channel (`[channel][fluor]`). + + Each item describes one fluorophore in one channel, without spatial + information: how many photons it absorbs (and at which wavelengths) at the + reference illumination intensity, and what fraction of them are detected (and + at which wavelengths). Fluorophores are the unique fluorophores in + `sample.labels`. + """ + fluors = list({lbl.fluorophore: None for lbl in self.sample.labels}) + key = self._fluorescence_cache_key(fluors) + if key is not None and (cached := _FLUORESCENCE_CACHE.get(key)) is not None: + _FLUORESCENCE_CACHE.move_to_end(key) + # use this simulation's own (equal) channel and fluorophore objects + return [ + [ + replace(fl, channel=oc, fluorophore=fluor) + for fl, fluor in zip(row, fluors, strict=True) + ] + for row, oc in zip(cached, self.channels, strict=True) + ] + + fluorescence = [ + [ + ChannelFluorescence.create( + oc, fluor, self.objective_lens, self.detector + ) + for fluor in fluors + ] + for oc in self.channels + ] + if key is not None: + _FLUORESCENCE_CACHE[key] = fluorescence + if len(_FLUORESCENCE_CACHE) > _FLUORESCENCE_CACHE_SIZE: + _FLUORESCENCE_CACHE.popitem(last=False) + return fluorescence + + def filtered_emission_rates(self, *, saturate: bool = True) -> xr.DataArray: """Return the emission rates for each fluorophore in each channel. - Returns a (C, F, W) array of emission rates for each fluorophore in each - channel, as a function of wavelength. The units are photons/s. The range + Returns a (C, F, W) array of detected emission rates for each fluorophore in + each channel, as a function of wavelength. The units are photons/s. The range of wavelengths will encompass the union of all the fluorophores' emission spectra, and the rates will be zero where the fluorophore does not emit. - For point-scanning modalities (e.g. `Confocal`), these rates are *not* - saturated: excitation saturation is applied locally in the PSF instead. + These are the rates for a fluorophore at the reference illumination intensity + (`OpticalConfig.power`): anywhere in a widefield field of view, or at the + center of the excitation spot for point-scanning modalities (e.g. `Confocal`). + If `saturate` is True (the default), excitation saturation at that intensity + is applied. Examples -------- @@ -214,24 +256,10 @@ def filtered_emission_rates(self) -> xr.DataArray: >>> rates.isel(c=0).plot.line(x='w') # plot emission of all fluors in channel 0 >>> plt.show() """ - fluors = list({lbl.fluorophore: None for lbl in self.sample.labels}) - key = self._rates_cache_key(fluors) - if key is not None and (cached := _RATES_CACHE.get(key)) is not None: - _RATES_CACHE.move_to_end(key) - # use this simulation's own (equal) channel and fluorophore objects - return cached.copy().assign_coords({Axis.C: self.channels, Axis.F: fluors}) - - # point-scanning modalities apply saturation spatially, in the PSF - sat = not self.modality.point_scanning - # fraction of isotropic emission entering the objective. PSFs are normalized - # to collected light, so this is not double-counted (see `make_confocal_psf`) - collection = self.objective_lens.collection_efficiency - qe = self.detector.qe if self.detector else None - nested_rates: list[list[xr.DataArray]] = [ - [oc.filtered_emission_rate(f, detector_qe=qe, saturate=sat) for f in fluors] - for oc in self.channels + nested_rates = [ + [fl.detected_spectrum(saturate=saturate) for fl in row] + for row in self._channel_fluorescence() ] - # combine xarray objects along the C and F axes, with outer join on W rates: xr.DataArray = xr.combine_nested( # type: ignore [assignment] nested_rates, @@ -240,31 +268,25 @@ def filtered_emission_rates(self) -> xr.DataArray: join="outer", fill_value=0, ).transpose(Axis.C, Axis.F, Axis.W) - rates = rates * collection - if key is not None: - _RATES_CACHE[key] = rates.copy() - if len(_RATES_CACHE) > _RATES_CACHE_SIZE: - _RATES_CACHE.popitem(last=False) return rates - def _rates_cache_key(self, fluors: list[Fluorophore]) -> tuple | None: - """Key identifying `filtered_emission_rates` output, or None if unhashable. + def _fluorescence_cache_key(self, fluors: list[Fluorophore]) -> tuple | None: + """Key identifying `_channel_fluorescence` output, or None if unhashable. - The rates depend only on the optics and fluorophores (not on where the - fluorophores are), so they can be reused across simulations. + This depends only on the optics and fluorophores (not on where the + fluorophores are), so it can be reused across simulations. """ try: return ( tuple(oc.model_dump_json() for oc in self.channels), tuple(f.model_dump_json() for f in fluors), self.detector.model_dump_json(include={"qe"}) if self.detector else "", - self.modality.point_scanning, self.objective_lens.collection_efficiency, ) except PydanticSerializationError: return None - def emission_flux(self) -> xr.DataArray: + def emission_flux(self, *, saturate: bool = True) -> xr.DataArray: """Return the spatial emission in photons per second (after emission filters). This is the ground truth data multiplied by the filtered emission rates, *prior* @@ -291,7 +313,7 @@ def emission_flux(self) -> xr.DataArray: return truth # total photons/s emitted by each fluorophore in each channel - em_rates = self.filtered_emission_rates().sum(Axis.W) + em_rates = self.filtered_emission_rates(saturate=saturate).sum(Axis.W) em_rates = em_rates.copy(data=self._xp.asarray(em_rates.data)) total_flux = em_rates * truth total_flux.attrs.update(units="photon/sec", long_name="Emission Flux") @@ -310,7 +332,7 @@ def optical_image_per_fluor(self) -> xr.DataArray: # (C, F, Z, Y, X) return self.modality.render( self.ground_truth(), # (F, Z, Y, X) - self.filtered_emission_rates(), # (C, F, W) + self._channel_fluorescence(), # [C][F] objective_lens=self.objective_lens, settings=self.settings, xp=self._xp, @@ -358,9 +380,9 @@ def digital_image( # Each Z plane is a separate exposure at one focal position (it does not # collect light from neighboring focal planes), so always average over Z. mean_axes = ( - [str(d) for d in image.dims] - if self.modality.point_scanning - else [Axis.Z] + [Axis.Z] + if self.modality.detection_type.sums_over_pixel_area + else [str(d) for d in image.dims] ) image = self.output_space.rescale(image, mean_axes=mean_axes) diff --git a/tests/test_cache.py b/tests/test_cache.py index 64d743d4..1ab8e5bc 100644 --- a/tests/test_cache.py +++ b/tests/test_cache.py @@ -6,7 +6,8 @@ import microsim.schema as ms from microsim import _data_array -from microsim.psf import cached_psf +from microsim._psf import PSFGrid, detection_psf +from microsim.psf import make_psf def test_cache(sim1: ms.Simulation) -> None: @@ -28,22 +29,38 @@ def test_cache(sim1: ms.Simulation) -> None: def test_cache_psf(caplog: pytest.LogCaptureFixture) -> None: kwargs = { - "nz": 16, - "nx": 64, - "dx": 0.01, - "dz": 0.02, - "ex_wvl_um": 0.5, - "em_wvl_um": 0.6, + "grid": PSFGrid(nz=16, nx=64, dx=0.01, dz=0.02), + "wvl_nm": 600, "objective": ms.ObjectiveLens(numerical_aperture=1.4), - "pinhole_au": 1.0, - "max_au_relative": None, "xp": ms.NumpyAPI(), + "pinhole_radius_um": 0.26, } caplog.set_level(10) - psf1 = cached_psf(**kwargs) - cached_psf.cache_clear() # clear the functools cache + psf1 = detection_psf(**kwargs) + detection_psf.cache_clear() # clear the functools cache assert "Found cached PSF" not in caplog.text - psf2 = cached_psf(**kwargs) + psf2 = detection_psf(**kwargs) assert "Found cached PSF" in caplog.text np.testing.assert_array_equal(psf1, psf2) + + +def test_cached_psfs_are_protected() -> None: + """Modifying a PSF returned by `make_psf` doesn't corrupt the PSF cache.""" + kwargs = { + "nz": 5, + "nx": 15, + "dx": 0.05, + "dz": 0.1, + "objective": ms.ObjectiveLens(numerical_aperture=1.4), + "em_wvl_nm": 520, + } + psf = make_psf(**kwargs) + expected = psf.copy() + psf *= 0 + np.testing.assert_array_equal(make_psf(**kwargs), expected) + # the cached arrays themselves are read-only + grid = PSFGrid(nz=5, nx=15, dz=0.1, dx=0.05) + cached = detection_psf(grid, 520.0, kwargs["objective"], ms.NumpyAPI()) + with pytest.raises(ValueError, match="read-only"): + cached *= 0 diff --git a/tests/test_photometry.py b/tests/test_photometry.py index 2e260cbf..94a146c8 100644 --- a/tests/test_photometry.py +++ b/tests/test_photometry.py @@ -1,15 +1,33 @@ from __future__ import annotations +from typing import Any, cast + import numpy as np import pytest import xarray as xr from pydantic import ValidationError from microsim import schema as ms +from microsim._psf import ( + PSFGrid, + airy_radius_um, + apply_saturation, + detection_psf, + excitation_spot, +) from microsim.psf import make_confocal_psf from microsim.schema.dimensions import Axis from microsim.schema.lens import ObjectiveLens -from microsim.schema.modality._simple_psf import bin_spectrum +from microsim.schema.modality._optics import ( + FocusedIllumination, + PinholeDetection, + WidefieldIllumination, +) +from microsim.schema.modality._simple_psf import ( + bin_spectrum, + emission_lines, + excitation_lines, +) from microsim.schema.optical_config.lib import FITC from tests._util import skipif_no_internet @@ -128,7 +146,7 @@ def _green_fluor() -> ms.Fluorophore: ) -def test_confocal_saturation_in_psf_not_rates() -> None: +def test_emission_rates_saturation() -> None: fluor = _green_fluor() oc = FITC.model_copy(update={"power": 1e7}) s = oc.saturation_parameter(fluor) @@ -143,16 +161,15 @@ def test_confocal_saturation_in_psf_not_rates() -> None: wf = ms.Simulation(**kw, modality=ms.Widefield()) cf = ms.Simulation(**kw, modality=ms.Confocal()) - # confocal emission rates are unsaturated; widefield rates are saturated - cf_rates = cf.filtered_emission_rates() - np.testing.assert_allclose( - wf.filtered_emission_rates(), cf_rates / (1 + s), atol=1e-300 - ) - - # confocal recovers the same saturation parameter from the rate coords - em_spectrum = cf_rates.isel(c=0, f=0) - assert cf.modality._saturation_parameter(em_spectrum) == pytest.approx(s, rel=1e-3) - assert wf.modality._saturation_parameter(em_spectrum) == 0 + # rates are for a fluorophore at the reference intensity, for any modality + rates = cf.filtered_emission_rates() + np.testing.assert_allclose(wf.filtered_emission_rates(), rates) + unsaturated = cf.filtered_emission_rates(saturate=False) + np.testing.assert_allclose(rates, unsaturated / (1 + s), atol=1e-300) + [[fl]] = cf._channel_fluorescence() + assert fl.saturation == pytest.approx(s) + assert fl.detected_rate(saturate=False) == pytest.approx(float(unsaturated.sum())) + assert fl.detected_rate() == pytest.approx(float(rates.sum())) def test_off_grid_laser_wavelength() -> None: @@ -196,6 +213,180 @@ def test_bin_spectrum_keeps_all_samples() -> None: assert float(bin_spectrum(spectrum, bins=bins).sum()) == 10 +def _confocal_sim(oc: ms.OpticalConfig, **kwargs: Any) -> ms.Simulation: + return ms.Simulation( + truth_space={"shape": (9, 33, 33), "scale": (0.1, 0.03, 0.03)}, + sample=[{"distribution": ms.MatsLines(), "fluorophore": _green_fluor()}], + channels=[oc], + objective_lens=ObjectiveLens(numerical_aperture=1.4), + modality=ms.Confocal(**kwargs), + settings=ms.Settings(max_psf_radius_aus=None, cache=False), + ) + + +def _confocal_psf(sim: ms.Simulation) -> tuple[Any, float]: + modality = cast("ms.Confocal", sim.modality) + [channel] = sim._channel_fluorescence() + detection = modality._detection(channel, sim.objective_lens) + _key, psf, brightness = modality._system_psf( + channel[0], + detection, + sim.ground_truth().attrs["space"], + sim.objective_lens, + sim.settings, + sim._xp, + ) + return psf, brightness + + +@pytest.mark.parametrize("power", [1, 1e7], ids=["linear", "saturated"]) +def test_confocal_psf_uses_excitation_wavelength(power: float) -> None: + oc = ms.OpticalConfig( + lights=[ms.LightSource.laser(488)], filters=FITC.filters, power=power + ) + sim = _confocal_sim(oc, airy_unit_wavelength_nm=520) + psf, brightness = _confocal_psf(sim) + [[fl]] = sim._channel_fluorescence() + assert brightness == pytest.approx(fl.detected_rate(saturate=False)) + + # the excitation spot is computed at the laser wavelength, not the emission + xp, objective = sim._xp, sim.objective_lens + grid = PSFGrid(nz=9, nx=33, dz=0.1, dx=0.03) + emission = emission_lines(fl.detected_fraction, bins=1, threshold_percentage=1) + [(em_nm, _)] = emission + pinhole_um = airy_radius_um(0.52, 1.4) + expected = apply_saturation( + excitation_spot(grid, 488, objective, xp), fl.saturation + ) + expected = expected * detection_psf(grid, em_nm, objective, xp, pinhole_um) + np.testing.assert_allclose(psf, expected, rtol=1e-3) + assert psf.max() == pytest.approx( + detection_psf(grid, em_nm, objective, xp, pinhole_um).max() + / (1 + fl.saturation), + rel=1e-3, + ) + + +def test_confocal_polychromatic_excitation() -> None: + """Excitation spots are weighted by absorption, summed, then saturated.""" + fluor = _green_fluor() + lasers = [ms.LightSource.laser(w) for w in (405, 488, 491)] + oc = ms.OpticalConfig(lights=lasers) + lines = excitation_lines(oc.absorption_rate(fluor).squeeze()) + # 405 is excited too weakly to pass the threshold; 488 and 491 share a bin + assert len(lines) == 1 + assert 488 < lines[0][0] < 491 + assert lines[0][1] == 1 + + # broadband excitation is split into bins of at most `bin_nm` + broad = excitation_lines(FITC.absorption_rate(fluor).squeeze(), bin_nm=10) + assert len(broad) > 1 + assert np.all(np.diff([w for w, _ in broad]) < 2 * 10) + assert sum(f for _, f in broad) == pytest.approx(1, abs=1e-3) + + xp, objective = ms.NumpyAPI(), ObjectiveLens(numerical_aperture=1.4) + grid = PSFGrid(nz=9, nx=33, dz=0.1, dx=0.03) + blue = excitation_spot(grid, 405, objective, xp) + green = excitation_spot(grid, 488, objective, xp) + assert blue[4].sum() < green[4].sum() # shorter wavelength: tighter spot + mixed = FocusedIllumination().spot(((405, 1), (488, 3)), grid, objective, xp) + np.testing.assert_allclose(mixed, 0.25 * blue + 0.75 * green, rtol=1e-6) + assert mixed.max() == pytest.approx(1) + assert WidefieldIllumination().spot(((405, 1),), grid, objective, xp) == 1 + + +def test_confocal_pinhole_matched_to_emission_band() -> None: + """1 AU is defined at the center of the emission filter band, per channel.""" + objective = ObjectiveLens(numerical_aperture=1.4) + green, flat = _green_fluor(), FLAT + bp = ms.Bandpass(bandcenter=525, bandwidth=50, placement="EM") + oc = ms.OpticalConfig(lights=[ms.LightSource.laser(488)], filters=[bp]) + + def _pinhole(oc: ms.OpticalConfig, **kwargs: Any) -> PinholeDetection: + sim = ms.Simulation( + truth_space={"shape": (4, 4, 4), "scale": (1, 1, 1)}, + sample=[ + {"distribution": ms.MatsLines(), "fluorophore": green}, + {"distribution": ms.MatsLines(), "fluorophore": flat}, + ], + channels=[oc], + objective_lens=objective, + modality=ms.Confocal(pinhole_au=2, **kwargs), + ) + [channel] = sim._channel_fluorescence() + det = cast("ms.Confocal", sim.modality)._detection(channel, objective) + return cast("PinholeDetection", det) + + # one pinhole for both fluorophores, 2 AU at the band center + expected = 2 * airy_radius_um(0.525, 1.4) + assert _pinhole(oc).radius_um == pytest.approx(expected, rel=1e-3) + custom = _pinhole(oc, airy_unit_wavelength_nm=600).radius_um + assert custom == pytest.approx(2 * airy_radius_um(0.6, 1.4)) + + # a longpass emission filter has no band center: fall back to detected light + lp = ms.Longpass(cuton=500, placement="EM") + open_ended = oc.model_copy(update={"filters": [lp]}) + with pytest.warns(UserWarning, match="no band-limited emission filter"): + assert _pinhole(open_ended).radius_um > 0 + + # a filter spectrum covering only its passband is band-limited + passband = ms.SpectrumFilter( + transmission={"wavelength": [500, 525, 550], "intensity": [1, 1, 1]}, + placement="EM", + ) + narrow = oc.model_copy(update={"filters": [passband]}) + assert narrow.emission_band_center() == pytest.approx(525) + assert _pinhole(narrow).radius_um == pytest.approx(expected, rel=1e-3) + + with pytest.raises(ValidationError): + ms.Confocal(airy_unit_wavelength_nm=0) + + +def test_confocal_channel_without_signal() -> None: + """A channel that detects nothing is dark (no pinhole size is needed).""" + far_red = ms.OpticalConfig( + lights=[ms.LightSource.laser(700)], + filters=[ms.Longpass(cuton=720, placement="EM")], + ) + sim = ms.Simulation( + truth_space={"shape": (4, 16, 16), "scale": (0.1, 0.05, 0.05)}, + sample=[{"distribution": ms.MatsLines(), "fluorophore": _green_fluor()}], + channels=[FITC, far_red], + modality=ms.Confocal(), + settings=ms.Settings(cache=False), + ) + img = sim.optical_image() + assert float(img[0].sum()) > 0 + assert float(img[1].sum()) == 0 + + +def test_single_wavelength_light() -> None: + light = ms.LightSource(spectrum={"wavelength": [488], "intensity": [1]}) + sim = ms.Simulation( + truth_space={"shape": (4, 16, 16), "scale": (0.1, 0.05, 0.05)}, + sample=[{"distribution": ms.MatsLines(), "fluorophore": _green_fluor()}], + channels=[ms.OpticalConfig(lights=[light])], + settings=ms.Settings(cache=False), + ) + assert float(sim.optical_image().sum()) > 0 + + +def test_emission_binning_preserves_brightness() -> None: + """Splitting emission into more PSF wavelengths doesn't change total photons.""" + + def _total(bins: int) -> float: + sim = ms.Simulation( + truth_space={"shape": (8, 64, 64), "scale": (0.1, 0.05, 0.05)}, + sample=[{"distribution": ms.MatsLines(), "fluorophore": _green_fluor()}], + settings=ms.Settings( + spectral_bins_per_emission_channel=bins, random_seed=0, cache=False + ), + ) + return float(sim.optical_image().sum()) + + assert _total(4) == pytest.approx(_total(1), rel=1e-3) + + # flat spectra across the visible, so the fluorophore is excited and detected in FITC _WAVES = np.arange(350, 750) FLAT = ms.Fluorophore( diff --git a/tests/test_psf.py b/tests/test_psf.py index b783330e..f3c18d53 100644 --- a/tests/test_psf.py +++ b/tests/test_psf.py @@ -4,7 +4,9 @@ import pytest from microsim import schema as ms +from microsim._psf import _pinhole_mask from microsim.psf import make_psf +from microsim.schema.modality._optics import CameraDetection, PinholeDetection from microsim.schema.optical_config.lib import FITC NA1_4 = ms.ObjectiveLens(numerical_aperture=1.4) @@ -70,8 +72,18 @@ def _image(radius: float | str | None) -> np.ndarray: assert not np.allclose(_image(2), full, rtol=0.01) # 2 AU = ~9 px << 48 px np.testing.assert_allclose(_image(100), full) # never larger than the volume - wf, conf = ms.Widefield(), ms.Confocal() - assert wf._auto_psf_radius_aus(5, 0.52, NA1_4) > wf._auto_psf_radius_aus( - 1, 0.52, NA1_4 - ) - assert conf._auto_psf_radius_aus(5, 0.52, NA1_4) == 8 + camera, pinhole = CameraDetection(), PinholeDetection(radius_um=0.2) + assert camera.psf_radius_aus(5, 0.52, NA1_4) > camera.psf_radius_aus(1, 0.52, NA1_4) + assert pinhole.psf_radius_aus(5, 0.52, NA1_4) == 8 + + +@pytest.mark.parametrize("radius_px", [0, 1e-4, 0.4, 0.5, 1.9, 7.6]) +def test_pinhole_mask_area(radius_px: float) -> None: + """The open area of the pinhole mask matches the pinhole, even when tiny.""" + for nxy in (31, 30): + mask = _pinhole_mask(nxy, radius_um=radius_px * 0.03, dxy=0.03) + assert mask.shape == (nxy, nxy) + assert mask.sum() == pytest.approx(np.pi * radius_px**2, rel=0.01) + assert mask[nxy // 2, nxy // 2] == mask.max() # centered + if nxy % 2: # symmetric around the center pixel + np.testing.assert_allclose(mask, mask[::-1, ::-1]) diff --git a/tests/test_simulation.py b/tests/test_simulation.py index 3a75e2c6..eb93c48e 100644 --- a/tests/test_simulation.py +++ b/tests/test_simulation.py @@ -174,7 +174,7 @@ def test_simulation_write_from_ground_truth(tmp_path: Path) -> None: def test_reused_caches_across_simulations() -> None: """Simulations with the same optics reuse PSF transforms and emission rates.""" convolve._KERNEL_FT_CACHE.clear() - simulation._RATES_CACHE.clear() + simulation._FLUORESCENCE_CACHE.clear() def _sim(seed: int) -> ms.Simulation: return ms.Simulation( @@ -190,12 +190,12 @@ def _sim(seed: int) -> ms.Simulation: sim_a, sim_b = _sim(1), _sim(2) img_a = sim_a.optical_image() assert len(convolve._KERNEL_FT_CACHE) == 1 - assert len(simulation._RATES_CACHE) == 1 + assert len(simulation._FLUORESCENCE_CACHE) == 1 # the second simulation hits both caches... img_b = sim_b.optical_image() assert len(convolve._KERNEL_FT_CACHE) == 1 - assert len(simulation._RATES_CACHE) == 1 + assert len(simulation._FLUORESCENCE_CACHE) == 1 # ... but with its own coordinate objects, and its own (different) sample rates_b = sim_b.filtered_emission_rates() assert rates_b.coords[Axis.C].values[0] is sim_b.channels[0] @@ -206,7 +206,7 @@ def _sim(seed: int) -> ms.Simulation: # and the result is the same as without the caches convolve._KERNEL_FT_CACHE.clear() - simulation._RATES_CACHE.clear() + simulation._FLUORESCENCE_CACHE.clear() np.testing.assert_allclose(_sim(2).optical_image(), img_b, rtol=1e-6) From 9fc40dc5c90200e1f7a70b3846d02b758c8700ea Mon Sep 17 00:00:00 2001 From: Talley Lambert Date: Wed, 14 Oct 2026 12:36:33 +0200 Subject: [PATCH 2/2] fix: require Confocal.pinhole_au >= 0.05 A (nearly) closed pinhole passes (almost) no light, so pinhole_au=0 gave an all-zero image. On main, 0 meant a one-pixel pinhole. Co-Authored-By: Claude Opus 5.5 --- src/microsim/schema/modality/_simple_psf.py | 4 +++- tests/test_photometry.py | 3 +++ 2 files changed, 6 insertions(+), 1 deletion(-) diff --git a/src/microsim/schema/modality/_simple_psf.py b/src/microsim/schema/modality/_simple_psf.py index 6560cf16..1e6f1235 100644 --- a/src/microsim/schema/modality/_simple_psf.py +++ b/src/microsim/schema/modality/_simple_psf.py @@ -208,7 +208,9 @@ class Confocal(_PSFModality): """ type: Literal["confocal"] = "confocal" - pinhole_au: Annotated[float, Ge(0)] = 1 + # a (nearly) closed pinhole passes (almost) no light. Real scanners rarely go + # below ~0.2 AU. + pinhole_au: Annotated[float, Ge(0.05)] = 1 airy_unit_wavelength_nm: float | None = Field( default=None, gt=0, diff --git a/tests/test_photometry.py b/tests/test_photometry.py index 94a146c8..de892fac 100644 --- a/tests/test_photometry.py +++ b/tests/test_photometry.py @@ -340,6 +340,9 @@ def _pinhole(oc: ms.OpticalConfig, **kwargs: Any) -> PinholeDetection: with pytest.raises(ValidationError): ms.Confocal(airy_unit_wavelength_nm=0) + for closed in (0, 0.01): + with pytest.raises(ValidationError): + ms.Confocal(pinhole_au=closed) def test_confocal_channel_without_signal() -> None: