fix: PSF sizing (radius cap, centering, normalization), add max_psf_radius_aus="auto"

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:

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

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)

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"
                     )

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