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..1e6f1235 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,49 @@ 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, + # 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, + 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 +264,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 +275,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 +304,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..de892fac 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,183 @@ 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) + for closed in (0, 0.01): + with pytest.raises(ValidationError): + ms.Confocal(pinhole_au=closed) + + +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)