diff --git a/src/microsim/psf.py b/src/microsim/psf.py index f0d42e46..89ffe554 100644 --- a/src/microsim/psf.py +++ b/src/microsim/psf.py @@ -314,7 +314,7 @@ def make_confocal_psf( sf=sf, normalize=False, ) - em_psf = _norm_to_collected_light(em_psf, xp) + 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 @@ -342,17 +342,52 @@ def make_confocal_psf( return out -def _norm_to_collected_light(psf: np.ndarray, xp: NumpyAPI) -> np.ndarray: +# lateral radius (in Airy units) of the in-focus plane used to normalize emission +# PSFs (see `_norm_to_collected_light`). Light beyond it is ~0.17 / 64 = ~0.3%. +_NORM_RADIUS_AU = 64 + + +def _norm_to_collected_light( + psf: np.ndarray, + wvl_um: float, + dxy: float, + objective: ObjectiveLens, + sf: int, + xp: NumpyAPI, +) -> np.ndarray: """Scale an emission PSF to the fraction of collected light on each pixel. Each z-plane is the image of a point source at that defocus. The vectorial model - conserves energy across planes, so every plane holds the same total (apart from - light falling outside the window). Dividing by the largest plane sum 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`.) + 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`.) + + 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 + light, and normalizing by them would brighten the PSF to compensate. """ - return psf / xp.max(xp.sum(psf, axis=(-2, -1))) # type: ignore[no-any-return] + airy_px = 0.61 * wvl_um / objective.numerical_aperture / dxy + nx = max(psf.shape[-1], int(2 * _NORM_RADIUS_AU * airy_px) | 1) + return psf / _in_focus_sum(nx, dxy, wvl_um, objective, sf, xp) + + +@cache +def _in_focus_sum( + nx: int, dxy: float, wvl_um: float, objective: ObjectiveLens, sf: int, xp: NumpyAPI +) -> float: + """Sum of the (unnormalized) in-focus plane of an emission PSF of width `nx`.""" + plane = vectorial_psf_centered( + nz=1, + nx=nx, + dxy=dxy, + wvl=wvl_um, + objective=objective, + sf=sf, + normalize=False, + xp=xp, + ) + return float(xp.sum(plane)) def _norm_psf( @@ -447,9 +482,10 @@ def cached_psf( 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) - # if even make odd - if nx % 2 == 0: - nx += 1 + + # 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 use_cache = os.getenv("MICROSIM_CACHE", "").lower() not in {"0", "false", "no", "n"} if use_cache: @@ -470,15 +506,15 @@ def cached_psf( if pinhole_au is None: psf = vectorial_psf_centered( wvl=em_wvl_um, - nz=nz + 1, - nx=nx + 1, + nz=nz, + nx=nx, dz=dz, dxy=dx, objective=objective, xp=xp, normalize=False, ) - psf = _norm_to_collected_light(psf, xp) + psf = _norm_to_collected_light(psf, em_wvl_um, dx, objective, 3, xp) else: psf = make_confocal_psf( nz=nz, @@ -513,7 +549,9 @@ def _psf_cache_path( ) -> Path: """Return the cache location for these PSF parameters.""" # "v2": PSFs are scaled to collected light, no longer sum-normalized - cache_key: list[float | str] = [nz, nx, dz, dx, em_wvl_um, "v2"] + # "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: diff --git a/src/microsim/schema/modality/_simple_psf.py b/src/microsim/schema/modality/_simple_psf.py index 3c29f85c..701e408c 100644 --- a/src/microsim/schema/modality/_simple_psf.py +++ b/src/microsim/schema/modality/_simple_psf.py @@ -137,13 +137,13 @@ def _summed_weighted_psf( # maximum wavelength in the emission spectrum and `settings.max_psf_radius_aus` nz, _ny, _nx = space.shape dz, _dy, dx = space.scale - max_wave = binned.coords[Axis.W].max().item() + max_wave_um = binned.coords[Axis.W].max().item() / 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) + logger.info(f"Auto PSF radius: {radius_aus:.1f} AU") nx = _pick_nx( - _nx, - dx, - settings.max_psf_radius_aus, - max_wave, - objective_lens.numerical_aperture, + _nx, dx, radius_aus, max_wave_um, objective_lens.numerical_aperture ) saturation = self._saturation_parameter(em_spectrum) @@ -181,6 +181,19 @@ 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) + class Confocal(_PSFModality): """Point-scanning confocal. @@ -208,6 +221,13 @@ def _saturation_parameter(self, em_spectrum: xrDataArray) -> float: # 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, *, @@ -350,7 +370,4 @@ def _pick_nx( 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) - # if even make odd - if nx % 2 == 0: - nx += 1 - return nx + return nx | 1 # odd, so that the PSF is centered on a pixel diff --git a/src/microsim/schema/settings.py b/src/microsim/schema/settings.py index b2f2330d..e397e7d7 100644 --- a/src/microsim/schema/settings.py +++ b/src/microsim/schema/settings.py @@ -1,5 +1,5 @@ import random -from typing import Any, ClassVar +from typing import Any, ClassVar, Literal from pydantic import Field, model_validator from pydantic_settings import BaseSettings, SettingsConfigDict @@ -59,13 +59,24 @@ class Settings(SimBaseModel, BaseSettings): random_seed: int | None = Field( default_factory=lambda: random.randint(0, 2**32 - 1) ) - max_psf_radius_aus: float | None = Field( - 6, + max_psf_radius_aus: float | Literal["auto"] | None = Field( + "auto", description=( - "When simulating, restrict generated lateral PSF size to no more than this " - "many Airy units. Decreasing this can *dramatically* speed up simulations, " - "but will decrease accuracy. If `None`, no restriction is applied, and the " - "psf will be generated to the full extent of the simulated space." + "Restrict the lateral radius of simulated PSFs to this many Airy units. " + "`'auto'` (default) is our opinionated balance between speed and " + "accuracy, aiming for ~1% RMSE (and <2% max error) relative to an " + "untruncated PSF, which is below shot noise unless pixels collect well " + "over ~10^4 photons. For widefield, the radius grows with the depth of the " + "simulated volume, since out-of-focus light spreads in a cone of " + "half-angle asin(NA/n): it is ~14 AU (for the tails of the in-focus Airy " + "pattern) plus half the radius of that cone at half the volume depth " + "(e.g. ~17 AU for a 1 um deep volume, ~27 AU for 5 um, at NA 1.4). For " + "confocal, where the pinhole rejects out-of-focus light, it is 8 AU. A " + "number sets the radius directly; a too-small radius loses out-of-focus " + "light (in a 5 um deep widefield volume, 8 AU loses ~20% of it). `None` " + "uses the full extent of the simulated space (most accurate, but much " + "slower for large volumes: ~2-4x slower per sample than 'auto' at " + "2048x2048 pixels). The PSF is never larger than the simulated space." ), ) cache: CacheSettings = Field(default_factory=CacheSettings) diff --git a/src/microsim/schema/spectrum.py b/src/microsim/schema/spectrum.py index 445ab6f3..7d7ebc47 100644 --- a/src/microsim/schema/spectrum.py +++ b/src/microsim/schema/spectrum.py @@ -55,7 +55,6 @@ def _cast_spectrum(cls, value: Any) -> Any: if isinstance(value, dict): if "wavelength" in value and "intensity" in value: if not len(value["wavelength"]) == len(value["intensity"]): - breakpoint() raise ValueError( "Wavelength and intensity must have the same length" ) diff --git a/tests/test_psf.py b/tests/test_psf.py new file mode 100644 index 00000000..b783330e --- /dev/null +++ b/tests/test_psf.py @@ -0,0 +1,77 @@ +from __future__ import annotations + +import numpy as np +import pytest + +from microsim import schema as ms +from microsim.psf import make_psf +from microsim.schema.optical_config.lib import FITC + +NA1_4 = ms.ObjectiveLens(numerical_aperture=1.4) +PSF_KW = {"dx": 0.05, "dz": 0.1, "objective": NA1_4, "em_wvl_nm": 520} + + +@pytest.mark.parametrize("pinhole_au", [None, 1], ids=["widefield", "confocal"]) +def test_make_psf_odd_and_centered(pinhole_au: float | None) -> None: + """PSFs have odd sizes (rounded up), with the peak on the center pixel.""" + for n, expect in [(17, 17), (16, 17)]: + psf = make_psf(nz=n, nx=n, pinhole_au=pinhole_au, **PSF_KW) + assert psf.shape == (expect, expect, expect) + assert np.unravel_index(np.argmax(psf), psf.shape) == (expect // 2,) * 3 + + +@pytest.mark.parametrize("pinhole_au", [None, 1], ids=["widefield", "confocal"]) +def test_make_psf_normalization_independent_of_size(pinhole_au: float | None) -> None: + """A laterally truncated PSF is not brightened to make up for the lost light.""" + small = make_psf(nz=9, nx=15, pinhole_au=pinhole_au, **PSF_KW) + large = make_psf(nz=9, nx=101, pinhole_au=pinhole_au, **PSF_KW) + np.testing.assert_allclose(small.max(), large.max(), rtol=1e-6) + assert small[4].sum() < large[4].sum() + if pinhole_au is None: # widefield planes are fractions of the collected light + assert large[4].sum() <= 1 + + +@pytest.mark.parametrize( + "modality", [ms.Widefield(), ms.Confocal()], ids=lambda m: m.type +) +def test_point_source_is_not_shifted(modality: ms.Modality) -> None: + truth = np.zeros((16, 48, 48)) + truth[8, 24, 24] = 1000 + for radius in (3, None): # odd sizes, whether or not the radius is capped + sim = ms.Simulation.from_ground_truth( + truth, + scale=(0.05, 0.02, 0.02), + modality=modality, + channels=[FITC], + settings=ms.Settings(max_psf_radius_aus=radius, cache=False), + ) + img = np.asarray(sim.optical_image().data[0]) + assert np.unravel_index(np.argmax(img), img.shape) == (8, 24, 24) + # laterally symmetric around the source + profile = img[8, 24, 1:] # pixels 1..47, centered on 24 + np.testing.assert_allclose(profile, profile[::-1], atol=1e-5 * profile.max()) + + +def test_max_psf_radius_aus() -> None: + """A smaller radius truncates the PSF (and auto grows with depth for widefield).""" + + def _image(radius: float | str | None) -> np.ndarray: + sim = ms.Simulation( + truth_space=ms.ShapeScaleSpace(shape=(8, 96, 96), scale=(0.05, 0.02, 0.02)), + sample=ms.Sample( + labels=[ms.FluorophoreDistribution(distribution=ms.MatsLines())] + ), + objective_lens=NA1_4, + settings=ms.Settings(random_seed=1, max_psf_radius_aus=radius, cache=False), + ) + return np.asarray(sim.optical_image().data) + + full = _image(None) + 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