From bb08929c634aec9890d87c2c1591fd4847ff7e53 Mon Sep 17 00:00:00 2001 From: Claude Date: Tue, 6 Oct 2026 10:01:33 +0000 Subject: [PATCH 1/7] Localize/uiPSF calibration: PSF models learnt by uiPSF, saved for Spline 3D A plugin of its own (not a mode of the bead calibration window) that has uiPSF learn a PSF from bead z-stacks or from the blinking molecules of an acquisition (in situ), in one channel or two, and saves it as smappy's calibration file, which Spline 3D and Spline 3D 2C read unchanged. - smappy.uipsf: prepare (camera to photons, the split of a two-channel camera recorded so uiPSF's channel transformation can be taken back to chip pixels, background subtraction and the channel shift for uiPSF's bead finder), microscopes (YAML profiles: NA, refractive indices, wavelength, splitter; $SMAPPY_MICROSCOPES, /microscopes, shipped example), runner + worker (uiPSF runs in its own Python, as a script importing nothing of smappy's; progress read from its output), convert (uiPSF's voxels splined and normalised as smappy's bead calibration does; uiPSF's T refitted as the secondary -> main projective map) - the plugin's page, and NOTES.md "uiPSF calibrations" with what was measured - tests/test_uipsf.py: conventions on hand-written uiPSF files (z sign and zero for bead and in situ models, split geometry for all four layouts, T -> transformation, joint normalisation), profiles, parameters, the runner protocol; and, when SMAPPY_UIPSF_PYTHON is set, uiPSF itself on simulated beads fitted back by Spline 3D (slow) Needs the claude/friendly-cerf-qt2sfj branch of ries-lab/uiPSF (numpy 2, current TensorFlow, channel_shift). On simulated data: single-channel beads fit with z slope 1.03, dual-channel beads with a transformation 0.02 px off and z slope 1.015. In situ learning runs but its z scale came out 0.76 on simulated data; the page says to check it against beads. --slow -n auto: 1342 passed, 4 skipped (with SMAPPY_UIPSF_PYTHON set) Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01GiNdULxyiqQEqWHGKTbVzV --- NOTES.md | 63 ++++ pyproject.toml | 2 +- src/smappy/data/microscopes/example.yaml | 21 ++ .../docs/Localize/uiPSF calibration.md | 299 ++++++++++++++++ src/smappy/plugins/uipsf.py | 316 +++++++++++++++++ src/smappy/uipsf/__init__.py | 102 ++++++ src/smappy/uipsf/convert.py | 264 ++++++++++++++ src/smappy/uipsf/microscopes.py | 138 ++++++++ src/smappy/uipsf/prepare.py | 278 +++++++++++++++ src/smappy/uipsf/runner.py | 144 ++++++++ src/smappy/uipsf/worker.py | 90 +++++ tests/test_uipsf.py | 327 ++++++++++++++++++ tests/test_workspace.py | 3 +- 13 files changed, 2045 insertions(+), 2 deletions(-) create mode 100644 src/smappy/data/microscopes/example.yaml create mode 100644 src/smappy/plugins/docs/Localize/uiPSF calibration.md create mode 100644 src/smappy/plugins/uipsf.py create mode 100644 src/smappy/uipsf/__init__.py create mode 100644 src/smappy/uipsf/convert.py create mode 100644 src/smappy/uipsf/microscopes.py create mode 100644 src/smappy/uipsf/prepare.py create mode 100644 src/smappy/uipsf/runner.py create mode 100644 src/smappy/uipsf/worker.py create mode 100644 tests/test_uipsf.py diff --git a/NOTES.md b/NOTES.md index b0190b7..5505866 100644 --- a/NOTES.md +++ b/NOTES.md @@ -1721,6 +1721,69 @@ Left out for now: a non-flat background, the precision from a spline PSF's CRLB (SMAP has it) rather than Mortensen's, and using a loaded table as the structure. +## uiPSF calibrations + +`Localize/uiPSF calibration` (`smappy.uipsf`) has uiPSF learn a PSF model +from bead stacks or blinking molecules and saves it as an ordinary spline +calibration. It is a plugin of its own, not a mode of the bead calibration +window: uiPSF is optional and external, and the duplication (reading stacks, +saving) is the price of keeping it out of that window. + +* **A process of its own.** uiPSF pins its own TensorFlow, so it runs under + the interpreter of its own environment (`runner.find_python`: the + plugin's setting, `uipsf_python` in the config, `$SMAPPY_UIPSF_PYTHON`, + else this Python if it has `psflearning`), as `uipsf/worker.py`, which + imports nothing of smappy's. The two talk through a job folder + (`images.npy`, `job.json`) and tagged lines on stdout; uiPSF's tqdm bars + are passed on, thinned to one per stage or five seconds. Nothing in + smappy imports TensorFlow. +* **smappy reads the data, not uiPSF.** Camera to photons, the stacks' z + from Micro-Manager, the split of a two-channel camera + (`prepare.Split`): all as smappy's own calibration and fitter do them, + and uiPSF's `prep_data` gets arrays. This is also what makes the channel + transformation recoverable: the cut is known exactly, so uiPSF's `T` + (main (y, x, 1) - imgcenter, right-multiplied, into the secondary's + flipped array) can be taken back to chip pixels and refitted as the + projective secondary -> main map `dualfit` reads. +* **Two things smappy does for uiPSF's bead finder.** It subtracts each + stack's median: uiPSF's threshold is relative to the brightest maximum + with the background in, and simulated two-colour beads (a background a + fifth of the peak) lost 15 of 18 beads to spurious maxima. And it + measures the shift between the channels by cross-correlating their + projections (`prepare.channel_shift`, passed as `channel_shift`, which + the uiPSF branch reads): uiPSF's own guess, from bead coordinates, took + the neighbour on a 30 px grid of beads and paired one bead. +* **Voxels, not uiPSF's spline.** uiPSF's `locres/coeff` is normalised to + the median plane sum after subtracting the minimum; the fitter's photons + assume the brightest plane sums to one. `res/I_model` is clipped, + normalised and splined exactly as `_single_model` / + `positive_pair_models` do (two channels jointly, the split from uiPSF's + per-channel intensities), so a fit cannot tell the calibrations apart. +* **z.** A bead model is uiPSF's `Zrange`, centred on the pupil focus, in + stack order: smappy's convention as is (checked: smappy's own calibration + of the same simulated beads has the same astigmatism against index). An + in situ model's planes are heights above the coverslip, the other way up + (uiPSF itself seeds an in situ run from a bead file's `I_model_reverse`): + reversed, with z = 0 at `stagepos * n_med / n_imm`. +* **Measured** on simulated astigmatic beads (2 stacks, 41 planes, 18 beads, + 4 CPU cores, uiPSF 100 iterations): 3 minutes; Spline 3D with the result + fits simulated frames with z slope 1.03, intercept 2 nm (smappy's own + calibration: 1.00). Frames drawn with that uiPSF model and fitted with + it: slope 1.00, photons 1.00. +* **In situ is not validated.** On frames drawn with a physical (uiPSF + bead) PSF, 3000 frames, two rounds (12 minutes), the in situ model had + the right sign and zero but its z ran 0.76 of the truth, with photons 10% + high; plane by plane it matches the true PSF at about 1.3 times the + height. Started from Noll 6 instead of 5 it converged to a model whose z + ran backwards. The conversion is not the cause (a hand-written in situ + file converts exactly, tests/test_uipsf.py); whether this is uiPSF on + real data or the simulation (a demo structure only +-300 nm deep) needs + real data with a bead calibration beside it. +* **uiPSF's branch.** uiPSF's main pins TensorFlow 2.9 / Python 3.7, which + has no Apple-silicon build. The `claude/friendly-cerf-qt2sfj` branch of + ries-lab/uiPSF runs on numpy 2 and TensorFlow 2.21 (hdfdict replaced, + tkinter dropped, NEP 50 dtype casts), tried with Python 3.12 on Linux. + ## Open questions * **Pages for what is not a plugin.** Every plugin has a page behind its ? diff --git a/pyproject.toml b/pyproject.toml index 34a9b65..03c7e76 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -122,7 +122,7 @@ where = ["src"] # the camera database: a file that is read at run time, so it has to be in the # wheel and not only in the checkout # and the simulator's built-in structures, likewise -"smappy" = ["data/*.json", "data/structures/*.yaml"] +"smappy" = ["data/*.json", "data/structures/*.yaml", "data/microscopes/*.yaml"] # the plugins' pages, read by the Help window (`smappy.docs`), one folder per # level of the menu tree "smappy.plugins" = ["*.chain.yaml", "docs/*.md", "docs/*/*.md", "docs/*/*/*.md", "docs/*/*/*/*.md"] diff --git a/src/smappy/data/microscopes/example.yaml b/src/smappy/data/microscopes/example.yaml new file mode 100644 index 0000000..3daed3a --- /dev/null +++ b/src/smappy/data/microscopes/example.yaml @@ -0,0 +1,21 @@ +# An example microscope profile for Localize > uiPSF calibration. +# +# Copy it to /microscopes/.yaml (or to a shared folder +# named in $SMAPPY_MICROSCOPES), change the values to your instrument's, and it +# appears in the plugin's "microscope" list. The numbers here are those of +# uiPSF's own example system "M2" (config/systemtype/M2.yaml): a 1.43 NA oil +# objective with an astigmatic lens and a mirrored two-colour splitter. +name: Example (uiPSF's M2) +NA: 1.43 +refractive_index: {immersion: 1.516, medium: 1.335, coverslip: 1.516} +emission_wavelength_nm: 680 +# pixelsize_um: [0.127, 0.116] # x, y; leave out to take it from the camera +roi_size_px: 25 +dual: + layout: up-down mirrored + main_channel: upper + split_position: null +insitu: + zernike_index: [5] + zernike_coeff: [0.5] +uipsf: {} diff --git a/src/smappy/plugins/docs/Localize/uiPSF calibration.md b/src/smappy/plugins/docs/Localize/uiPSF calibration.md new file mode 100644 index 0000000..8586d91 --- /dev/null +++ b/src/smappy/plugins/docs/Localize/uiPSF calibration.md @@ -0,0 +1,299 @@ +--- +version: "1" +covers: [smappy.uipsf.prepare.Split, smappy.uipsf.prepare.channel_shift, smappy.uipsf.prepare.subtract_background, smappy.uipsf.prepare.read_beads, smappy.uipsf.prepare.read_movie, smappy.uipsf.uipsf_parameters, smappy.uipsf.runner.run, smappy.uipsf.convert.single_calibration, smappy.uipsf.convert.dual_calibration, smappy.uipsf.convert.channel_transformation, smappy.uipsf.microscopes.Microscope] +--- + +## What it does + +A 3D fit needs to know what a molecule looks like at every height: the +**point spread function** (PSF). The [bead calibration](panel:Bead calibration) +measures it by averaging beads imaged at a series of focus positions. This +plugin hands the same job to **uiPSF** +([Liu et al. 2024](https://doi.org/10.1038/s41592-024-02282-x)), which instead +fits a *model of the microscope* to the images. The model is the light's +wavefront in the objective's pupil, written as a sum of Zernike polynomials +-- defocus, astigmatism, coma, spherical aberration and so on -- and +propagated to the camera with the physics of a high-NA objective. + +That gives two things the bead average cannot: + +* **A model from beads that is smooth without smoothing**, defined beyond the + beads' noise, and that names the microscope's aberrations in nanometres of + wavefront. +* **A model from the blinking molecules themselves** (*in situ*). A bead on + the coverslip does not see what a molecule a micrometre deep in a cell sees: + the refractive-index mismatch between oil and water adds spherical + aberration that grows with depth, and the sample adds its own. uiPSF can + learn the PSF from the raw frames of the acquisition you want to fit, so the + model is the one the data were taken with. + +The result is saved as an ordinary smappy calibration file, which +[Spline 3D](plugin:Localize/Spline 3D) and +[Spline 3D 2C](plugin:Localize/Spline 3D 2C) read exactly as they read a bead +calibration. For a split camera it holds both channels' models and the +transformation between the halves, which uiPSF learns together with the PSF. + +**What it needs.** uiPSF itself, installed in a Python environment of its own +(it brings TensorFlow; see *Installing uiPSF* below), and a **microscope +profile**: the objective's numerical aperture, the refractive indices, the +emission wavelength and, for a split camera, how the halves are arranged. +From beads: one or more z-stacks, as for the bead calibration. From blinking +molecules: a few thousand frames with molecules spread over the depth you want +to fit, and the depth the focus was at. + +**Learning from blinking molecules needs checking.** Unlike beads, the +molecules' heights are unknown, so uiPSF has to find them and the PSF +together. Where that ends up depends on where it starts: the starting pupil +from the profile (`insitu:` in the profile) and the *depth*. In tests on +simulated data the z scale of the learnt model came out up to a quarter +short, and a poor starting pupil gave a model whose z ran backwards. Before +relying on an in situ calibration, compare it with a bead calibration of +the same microscope: fit both to the same data near the coverslip, where +they should agree. Starting from that bead model (*start from*) helps. + +**When not to use it.** For a PSF no pupil describes -- a strongly +non-uniform illumination of the pupil, a damaged optic -- the voxel model or +the bead calibration is the safer choice. Field-dependent and 4Pi models, +which uiPSF also learns, are not offered: the fitters here use one PSF for the +whole field. + +## How it works + +**1. Reading the data.** smappy reads the frames, not uiPSF: the bead +stacks' z positions come from the microscope's metadata as for the bead +calibration, and every pixel is converted to photons with the camera's gain +and offset, $N = (c - c_0)\,g$, as the fitter converts them. Bead stacks of +different lengths are cut to the shortest about their centres. Each stack's +median is subtracted. uiPSF fits a background per bead anyway, but it finds +beads above a fraction (*threshold*) of the brightest pixel, and with the +background left in that fraction can fall into the noise. + +```figure-setup +from smappy.simulate import dual_bead_stacks +from smappy.uipsf.prepare import make_split +stacks, z = dual_bead_stacks(1, seed=0, size_px=60) +frame = stacks[0].max(axis=0).astype(float) +``` + +**2. Splitting a two-channel camera.** The frame is cut at the split into +two halves of equal size, the main channel first, and a mirrored half is +flipped back, so that uiPSF sees two images of the same molecules. Exactly +how they were cut is recorded, because the transformation uiPSF learns +between the halves has to be put back into camera pixels afterwards. For beads, +the shift between the halves is measured first, from the cross-correlation +of their projections, and handed to uiPSF as its starting point for pairing +the beads. + +```figure Each bead appears in both halves of the split camera; cut and stacked, the two channels show the same beads, a little shifted and turned -- the transformation uiPSF learns along with the PSF. +fig.set_size_inches(7.5, 3.2) +split = make_split(frame.shape, {"layout": "up-down", "main_channel": "upper"}) +channels = split.cut(frame[None])[:, 0] +axes = fig.subplots(1, 3) +axes[0].imshow(frame, cmap="magma") +axes[0].axhline(split.split_position - 0.5, color="w", lw=1) +axes[0].set_title("camera frame") +for ax, image, title in zip(axes[1:], channels, ("main", "secondary")): + ax.imshow(image, cmap="magma") + ax.set_title(title) +for ax in axes: + ax.set_xticks([]) + ax.set_yticks([]) +``` + +**3. Learning the PSF.** uiPSF finds the beads or molecules, cuts a small +image around each, and fits all of them at once with one PSF model: every +emitter gets its own position, brightness and background, and all share the +pupil. The pupil's Zernike coefficients are adjusted until the modelled +images match the measured ones. From blinking molecules the heights are not +known beforehand. They are estimated with a starting pupil (astigmatism +from the profile), then refined together with the pupil, over a few +*rounds*. This is the slow step: minutes for beads, longer for blinking +molecules (see *Running time*). + +**4. Making the calibration.** The learnt PSF is taken as voxels, a stack of +images at the model's z planes, and turned into the cubic spline the fitters +use. This is the same step as the bead calibration's: negative values set to +zero, the brightest plane normalised to one photon, and the spline computed +through every voxel. The fitter cannot tell which way a calibration was made. +The Zernike coefficients, the refined depth and uiPSF's parameters are stored +with it, and uiPSF's own result file is saved beside it. + +## In detail + +**Where z = 0 is.** For beads, uiPSF's model is centred on the pupil's +focus, and that plane is z = 0 -- the same convention as the bead calibration, +with the fitter's $z = -(i - z_0)\,\Delta z$ for plane $i$. For blinking +molecules uiPSF's planes are heights above the coverslip, from the bottom up. +They are reversed to run as a bead stack does. z = 0 is put at the nominal +focal plane: the focus moved $d$ into the sample sits at height +$d\, n_{\mathrm{med}} / n_{\mathrm{imm}}$, and $d$ is the *depth* as uiPSF +refined it. A fitted z is then the molecule's height relative to that plane, +and its height above the coverslip is $z + d\, n_{\mathrm{med}} / n_{\mathrm{imm}}$. + +**The scale of z from blinking molecules.** From beads the z steps are +known, so the model's z scale is the stage's. From blinking molecules it is +not: it comes only from the physics of the pupil model, from how defocus +changes the spot for this NA, wavelength and refractive index. The profile's +numbers therefore set the z scale of an in situ calibration, and a wrong +refractive index scales every z. + +**Two channels.** uiPSF stores, for each pair of channels, an affine map +$T$ that takes a main-channel position $(y, x, 1) - c$ to the other +channel's, both in the arrays it was given and $c$ the image centre. The +plugin pushes a grid of points through it, takes both sets back to camera +pixels through the recorded split, and fits the projective map, secondary to +main, that the two-colour fitter reads. The two models are normalised +together, as the dual-colour bead calibration does: each channel's light +around focus is its model's signal times the photons uiPSF fitted to that +channel, and together they sum to one, so the photon split between the +channels is kept. The secondary model is flipped back to the camera's +orientation when the splitter mirrors it. + +**Running uiPSF.** uiPSF is not imported by smappy. It runs as a separate +program in its own Python, so that smappy needs neither TensorFlow nor its +version. The plugin writes the photons and the configuration into a job +folder, starts uiPSF there, passes its progress on and reads back the file +uiPSF writes. If it fails, the message names the log with everything uiPSF +printed (*keep job folder* keeps the folder). + +**What uiPSF is told.** The profile's NA, refractive indices and wavelength; +the pixel size, from the camera's metadata or else from the profile; the z +step; and the settings below. Anything else uiPSF can be told goes into the +profile's `uipsf:` section, nested as in uiPSF's own YAML configuration, and +has the last word. + +**Running time.** On four CPU cores, 18 simulated beads in 41 planes took +three minutes. uiPSF's README reports, on a graphics card, half a minute for +beads in one channel, five for two channels, and five to fifteen minutes in +situ. On a CPU expect several times that. On a Mac uiPSF runs on the CPU; +TensorFlow's Apple-GPU plugin has not been tried with it, and the pupil model +is built on complex numbers and Fourier transforms, which that plugin has +supported only in part. + +**Installing uiPSF.** In a Python environment of its own (tried with Python +3.12 and TensorFlow 2.21): + + pip install tensorflow tensorflow-probability[tf] h5py omegaconf dotted_dict \ + czifile scikit-image matplotlib tqdm + pip install --no-deps -e /path/to/uiPSF + +The plugin needs uiPSF's branch that runs on current numpy and TensorFlow. +Then give that environment's python as *uiPSF python*, or once and for all as +`uipsf_python` in smappy's configuration file, or in `SMAPPY_UIPSF_PYTHON`. + +**Microscope profiles.** A profile is a YAML file, `example.yaml` ships as +a template: + + name: M2, astigmatic, two colour + NA: 1.43 + refractive_index: {immersion: 1.516, medium: 1.335, coverslip: 1.516} + emission_wavelength_nm: 680 + roi_size_px: 25 + dual: {layout: up-down mirrored, main_channel: upper} + insitu: {zernike_index: [5], zernike_coeff: [0.5]} + uipsf: {} + +Profiles are found in the folders named by `SMAPPY_MICROSCOPES` (a group's +shared folder, say), then in `microscopes/` in smappy's configuration folder, +then among those that ship. The first file of a name wins. + +## Parameters + +### data +Beads give the PSF on the coverslip, with the z scale of the stage. Blinking +molecules give the PSF where the molecules are, aberrations of the sample +included; prefer it for samples deeper than about a micrometre in water. + +### microscope +The profile sets the z scale of an in situ model (see *In detail*), so a +wrong NA or refractive index matters there more than for beads. + +### model +A voxel model has a free value for every voxel, so it needs many bright beads +and has no aberrations to report. + +### beads.z_step_nm +Only needed when the stack's metadata has no z positions, or has them wrong. + +### beads.bead_diameter_nm +Beads of 100 nm or more broaden the measured PSF. Given their size, uiPSF +models the bead and the PSF is learnt without it. + +### blinking.depth_um +A starting value: uiPSF refines it. A guess within a few hundred nanometres +is enough. + +### blinking.z_range_um +Cover the depth the molecules are spread over, with a few hundred nanometres +to spare. Wider costs time and leaves thinly sampled planes at the ends. + +### blinking.last_frame +Three thousand frames with ten or more molecules each give enough molecules +per z slice; sparser data need more frames. + +### blinking.start_from +A bead model of the same microscope is a good start, and the learning +converges in fewer rounds. Either uiPSF's own result file or the calibration +this plugin saved, whose uiPSF file lies beside it, will do. + +### blinking.min_photon +Raise it for data with many dim molecules, which carry little information +about the PSF. + +### learning.roi_size_px +The fitters cut their ROI from inside the model, so it must be at least three +pixels wider than the fit's ROI. Larger also captures the wings of a +defocused spot. + +### learning.iterations +uiPSF's own default is 200. 100 is usually converged for a Zernike model, +and the output shows the loss falling. + +### learning.zernike_order +Order 8 is 45 terms. Order 6 (28 terms) is enough for most systems, and +lower orders learn faster. + +## Output + +**The calibration** (*save as*, by default `_uipsf_3dcal.h5` +beside the data) is what Spline 3D or Spline 3D 2C take as their +*calibration*. Beside it is `_uipsf.h5`, uiPSF's own result: the +pupil, the per-emitter fits and the model, which SMAP and uiPSF's own +notebooks read. + +**The figures** show the model through focus, x–z and y–z, and its focal +plane; for two channels both, the secondary in the camera's orientation. An +astigmatic PSF is wide in x on one side of focus and in y on the other. The +second tab shows the learnt aberrations in nanometres of wavefront per +Zernike term: astigmatism (5, 6) for an astigmatic lens, spherical (11) +growing with depth in water. Large coma (7, 8) points at a tilted or +misaligned optic. + +**A bad result** shows as a model with structure that changes erratically +from plane to plane, or as fitted z that does not follow the stage when the +bead stack itself is fitted. The usual causes are too few emitters, a wrong +pixel size or wavelength, and, from blinking molecules, a depth guess or z +range far from the data. + +## Differences from SMAP + +SMAP has no uiPSF plugin: uiPSF's result file is loaded directly into SMAP's +fitter and peak finder ("load 3D cal", "load T"), using the spline uiPSF +computes. Here the spline is computed again from uiPSF's voxels, the way +smappy's bead calibration computes it. uiPSF normalises its spline by the +median plane sum after subtracting the minimum, while the fitter's photon +count assumes the brightest plane sums to one. The channel transformation is +converted to smappy's projective map in camera pixels, so a two-channel +calibration from uiPSF and one from beads are used in the same way. + +## References + +* Liu S, Chen J, Hellgoth J, *et al.* (2024) Universal inverse modeling of + point spread functions for SMLM localization and microscope + characterization. *Nat Methods* 21:1082–1093. + [doi:10.1038/s41592-024-02282-x](https://doi.org/10.1038/s41592-024-02282-x) + -- uiPSF. +* Li Y, Mund M, Hoess P, *et al.* (2018) Real-time 3D single-molecule + localization using experimental point spread functions. *Nat Methods* + 15:367–369. [doi:10.1038/nmeth.4661](https://doi.org/10.1038/nmeth.4661) + -- the cubic spline the fitters evaluate. diff --git a/src/smappy/plugins/uipsf.py b/src/smappy/plugins/uipsf.py new file mode 100644 index 0000000..a49ccd8 --- /dev/null +++ b/src/smappy/plugins/uipsf.py @@ -0,0 +1,316 @@ +"""uiPSF calibration: a PSF model learnt by uiPSF, saved as smappy's calibration. + +A separate plugin from the bead calibration window on purpose, even where the +two overlap (reading bead stacks, saving): uiPSF is an optional external +program with its own Python, and keeping it out of the window keeps that +window's code free of it. Everything is in `smappy.uipsf`; this is the form. + +What it adds to the bead calibration: a model of the optics (a pupil, its +Zernike aberrations) rather than an average of beads, so it is smooth without +smoothing and extends past the beads' noise; and in situ learning from the +blinking molecules of an acquisition, for a PSF the sample has changed. +""" +from __future__ import annotations + +import shutil +import tempfile +from dataclasses import dataclass, replace +from pathlib import Path +from typing import Optional + +import numpy as np + +from . import Context, Plot, Plugin, Result, param, register +from .fit import TIFF_FILTER, CameraSettings + +CALIBRATION_FILTER = "smappy calibration (*.h5);;All files (*)" +UIPSF_FILTER = "uiPSF or smappy uiPSF calibration (*.h5);;All files (*)" + + +def _microscopes(): + from ..uipsf.microscopes import choices + try: + return choices() + except Exception: # noqa: BLE001 -- the list must open + return [("", "uiPSF defaults")] + + +@dataclass +class BeadPart: + """Bead z-stacks: the objective stepped through focus.""" + z_step_nm: Optional[float] = param(None, label="z step", unit="nm", min=1, + help="auto: from the stack's z positions") + whole_folder: bool = param(True, label="all in folder", + help="every bead acquisition in the file's folder, " + "not only the one chosen") + max_beads: int = param(40, label="max beads", min=1, advanced=True, + help="at most this many beads are used") + bead_diameter_nm: float = param(0.0, label="bead size", unit="nm", min=0, + advanced=True, + help="the beads' diameter, deconvolved from " + "the model; 0 for point-like beads") + + +@dataclass +class BlinkingPart: + """Single molecules in the sample (in situ).""" + depth_um: float = param(1.0, label="depth", unit="um", min=0, + help="how far the objective's focus was moved into " + "the sample from the coverslip; uiPSF refines it") + z_range_um: float = param(2.0, label="z range", unit="um", min=0.2, + help="the depth of the model, around focus") + z_step_nm: float = param(50.0, label="z step", unit="nm", min=5, + help="the model's plane spacing") + first_frame: int = param(0, label="first frame", min=0, + help="frames before this one are not read") + last_frame: Optional[int] = param(3000, label="last frame", min=1, + help="auto: to the end; more frames, more " + "molecules, and more memory") + start_from: str = param("", label="start from", kind="open_file", + file_filter=UIPSF_FILTER, advanced=True, + help="a uiPSF result (a bead model, say) to start " + "the pupil from") + min_photon: float = param(0.4, label="photon quantile", min=0, max=1, + advanced=True, + help="molecules dimmer than this quantile are " + "not used") + z_bins: int = param(21, label="z bins", min=1, advanced=True, + help="molecules are chosen evenly over this many slices " + "in z") + per_bin: int = param(100, label="per bin", min=1, advanced=True, + help="at most this many molecules per slice") + repeat: int = param(2, label="rounds", min=1, advanced=True, + help="learning repeated, each round starting from the last") + + +@dataclass +class LearningPart: + """What uiPSF is told beyond the microscope.""" + roi_size_px: Optional[int] = param(None, label="model size", unit="px", min=11, + help="auto: the microscope profile's; at " + "least 3 px more than the fit's ROI") + emission_wavelength_nm: Optional[float] = param(None, label="wavelength", + unit="nm", min=300, max=1200, + help="auto: the profile's") + iterations: int = param(100, min=10, help="uiPSF's optimiser steps per round") + zernike_order: int = param(8, label="Zernike order", min=2, max=12, + advanced=True, + help="the highest radial order of the pupil's " + "Zernike terms") + peak_height: float = param(0.2, label="threshold", min=0, max=1, advanced=True, + help="relative brightness an emitter must reach " + "to be used") + python: str = param("", label="uiPSF python", kind="open_file", advanced=True, + help="the python of uiPSF's environment; auto: " + "uipsf_python in the config file, or " + "$SMAPPY_UIPSF_PYTHON") + keep_job: bool = param(False, label="keep job folder", advanced=True, + help="keep uiPSF's input, log and result in a " + "folder beside the calibration") + + +@dataclass +class UiPSFSettings: + data: str = param("beads", choices=(("beads", "bead z-stacks"), + ("blinking", "blinking molecules")), + help="beads: a calibration as the bead window makes one; " + "blinking: the PSF in the sample itself") + path: str = param("", label="file", kind="open_file", file_filter=TIFF_FILTER, + help="a file of the acquisition") + microscope: str = param("", choices=lambda: _microscopes(), + help="the optics: NA, refractive indices, " + "wavelength, the splitter") + channels: str = param("auto", choices=(("auto", "as the microscope"), + ("single", "one"), + ("dual", "two (split camera)")), + help="auto: two if the microscope has a splitter") + model: str = param("zernike", choices=(("zernike", "Zernike pupil"), + ("voxel", "voxels (beads)"), + ("pupil", "free pupil (blinking)")), + help="Zernike: smooth and physical; voxels: any shape, " + "needs many beads; free pupil: for unusual PSFs " + "such as a tetrapod") + output: str = param("", label="save as", kind="save_file", + file_filter=CALIBRATION_FILTER, + help="auto: _uipsf_3dcal.h5 beside the data") + beads: BeadPart = param(default_factory=BeadPart, label="beads") + blinking: BlinkingPart = param(default_factory=BlinkingPart, label="blinking", + collapsed=True) + learning: LearningPart = param(default_factory=LearningPart, label="learning", + collapsed=True) + camera: CameraSettings = param(default_factory=CameraSettings, label="camera", + collapsed=True) + + +def output_path(settings: UiPSFSettings) -> Path: + """Where the calibration goes.""" + if settings.output: + return Path(settings.output) + source = Path(settings.path) + if settings.data == "beads" and settings.beads.whole_folder: + base = source.parent if source.is_file() else source + return base / f"{base.name}_uipsf_3dcal.h5" + stem = source.name.split(".")[0] if source.is_file() else source.name + return (source.parent if source.is_file() else source) / f"{stem}_uipsf_3dcal.h5" + + +def check(settings: UiPSFSettings) -> None: + """Refuse what cannot work, before minutes are spent on it.""" + if not settings.path: + raise ValueError("choose a file of the acquisition") + if not Path(settings.path).exists(): + raise ValueError(f"{settings.path} does not exist") + if settings.data == "beads" and settings.model == "pupil": + raise ValueError("a free pupil is learnt from blinking molecules; for " + "beads choose Zernike or voxels") + if settings.data == "blinking" and settings.model == "voxel": + raise ValueError("a voxel model is learnt from beads; for blinking " + "molecules choose Zernike or a free pupil") + + +@register("Localize/uiPSF calibration") +class UiPSFCalibration(Plugin): + """A PSF model from beads or blinking molecules, learnt by uiPSF, for Spline 3D.""" + description = ("learn a PSF model with uiPSF from bead stacks or from the " + "blinking molecules themselves, and save it as a calibration " + "for Spline 3D and Spline 3D 2C") + Settings = UiPSFSettings + version = "1" + main = ("data", "path", "microscope", "channels", "model", "output") + + def run(self, ctx: Context, settings: UiPSFSettings) -> Result: + from ..uipsf import convert, learn, microscopes, prepare, uipsf_parameters + check(settings) + profile = microscopes.load(settings.microscope) + dual = (profile.dual is not None if settings.channels == "auto" + else settings.channels == "dual") + if dual and profile.dual is None: + profile = replace(profile, dual={}) # the bead window's default split + if settings.data == "beads": + paths = [Path(settings.path).parent if settings.beads.whole_folder + and Path(settings.path).is_file() else settings.path] + data = prepare.read_beads(paths, settings.camera, profile, dual, + settings.beads.z_step_nm, ctx.report) + dz, sources = data.dz_nm, data.sources + else: + b = settings.blinking + data = prepare.read_movie(settings.path, settings.camera, profile, dual, + b.first_frame, b.last_frame, ctx.report) + dz, sources = b.z_step_nm, [data.source] + ctx.report(f"{data.images.shape} photons, " + f"{'two channels' if dual else 'one channel'}") + lp, b = settings.learning, settings.blinking + params = uipsf_parameters( + profile, data=settings.data, pixelsize_xy_um=prepare.pixel_size_xy_um(data.camera), + dz_nm=dz, roi_size_px=lp.roi_size_px, + emission_wavelength_nm=lp.emission_wavelength_nm, peak_height=lp.peak_height, + max_beads=settings.beads.max_beads, + bead_diameter_nm=settings.beads.bead_diameter_nm, iterations=lp.iterations, + zernike_order=lp.zernike_order, depth_um=b.depth_um, z_range_um=b.z_range_um, + min_photon=b.min_photon, z_bins=b.z_bins, per_bin=b.per_bin, repeat=b.repeat, + start_from=_uipsf_file(b.start_from) if settings.data == "blinking" else "") + if dual and settings.data == "beads": + params["channel_shift"] = prepare.channel_shift(data.images) + out = output_path(settings) + job = (out.with_name(out.stem + "_uipsf_job") if lp.keep_job + else Path(tempfile.mkdtemp(prefix="smappy-uipsf-"))) + result = learn(data.images, data=settings.data, model=settings.model, dual=dual, + params=params, workdir=job, python=lp.python, progress=ctx.report) + em_on = getattr(data.camera, "em_on", None) + calibration = (convert.dual_calibration(result, data.split, sources, em_on=em_on) + if dual else convert.single_calibration(result, em_on=em_on)) + _note_sources(calibration, settings, sources) + convert.save(out, calibration) + raw = out.with_name(out.stem + "_uipsf.h5") + shutil.copyfile(result, raw) + if not lp.keep_job: + shutil.rmtree(job, ignore_errors=True) + text = summary(calibration, out, raw) + return Result(text=text, settings=settings, + plot=Plot(draw=lambda fig: draw_models(fig, calibration), + panels=6 if dual else 3, size=(10, 6 if dual else 3.5)), + plots={"Zernike": lambda ax: draw_zernike(ax, calibration)}, + data={"calibration": str(out), "uipsf_result": str(raw), + "dual": dual}) + + +def _uipsf_file(path: str) -> str: + """A smappy uiPSF calibration names the uiPSF file beside it; uiPSF wants that.""" + if not path: + return "" + p = Path(path) + beside = p.with_name(p.stem + "_uipsf.h5") + if not p.name.endswith("_uipsf.h5") and beside.exists(): + return str(beside) + return str(p) + + +def _note_sources(calibration, settings, sources) -> None: + models = ([calibration.main, calibration.secondary] if hasattr(calibration, "main") + else [calibration]) + for model in models: + model.parameters["data"] = settings.data + model.parameters["microscope"] = settings.microscope + model.parameters["sources"] = [s.get("path") for s in sources] + + +def _models(calibration): + if hasattr(calibration, "main"): + return [("main", calibration.main), ("secondary", calibration.secondary)] + return [("", calibration)] + + +def summary(calibration, out: Path, raw: Path) -> str: + """What the run made: the file, the model's extent and the main aberrations.""" + lines = [f"saved {out}", f"uiPSF's own result: {raw}"] + for name, model in _models(calibration): + nz = model.psf.shape[0] + lines.append(f"{name + ': ' if name else ''}{nz} planes, {model.dz:g} nm apart, " + f"z from {model.z_index_to_nm(nz - 1):.0f} to " + f"{model.z_index_to_nm(0):.0f} nm") + terms = [row for row in model.parameters.get("zernike", []) if row[0] > 4] + terms = sorted(terms, key=lambda row: -abs(row[3]))[:4] + if terms: + lines.append(" largest aberrations: " + ", ".join( + f"{label or 'Z' + str(j)} {phase:+.0f} nm" for j, label, _, phase in terms)) + if hasattr(calibration, "transformation"): + lines.append("channel transformation (secondary -> main, chip px): " + + np.array2string(calibration.transformation / calibration.transformation[2, 2], + precision=4, suppress_small=True).replace("\n", "")) + return "\n".join(lines) + + +def draw_models(figure, calibration) -> None: + """Each channel's model through focus: x-z, y-z, and the focal plane.""" + rows = _models(calibration) + axes = figure.subplots(len(rows), 3, squeeze=False) + for (name, model), row in zip(rows, axes): + psf = model.psf + nz, ny, nx = psf.shape + extent_z = (model.z_index_to_nm(nz - 0.5), model.z_index_to_nm(-0.5)) + focus = int(round(model.z0)) + row[0].imshow(psf[::-1, ny // 2, :], aspect="auto", cmap="magma", + extent=(-0.5, nx - 0.5) + extent_z) + row[0].set(xlabel="x (px)", ylabel="z (nm)", title=f"{name} x-z".strip()) + row[1].imshow(psf[::-1, :, nx // 2], aspect="auto", cmap="magma", + extent=(-0.5, ny - 0.5) + extent_z) + row[1].set(xlabel="y (px)", title=f"{name} y-z".strip()) + row[2].imshow(psf[min(max(focus, 0), nz - 1)], cmap="magma") + row[2].set(xlabel="x (px)", ylabel="y (px)", title=f"{name} z = 0".strip()) + + +def draw_zernike(ax, calibration) -> None: + """The pupil's aberrations, Noll index 5 up, in nm rms of wavefront.""" + width = 0.8 / len(_models(calibration)) + for k, (name, model) in enumerate(_models(calibration)): + rows = [row for row in model.parameters.get("zernike", []) if row[0] > 4] + if not rows: + continue + j = np.array([row[0] for row in rows]) + ax.bar(j + (k - (len(_models(calibration)) - 1) / 2) * width, + [row[3] for row in rows], width=width, label=name or None) + ax.axhline(0, color="0.5", lw=0.8) + ax.set(xlabel="Noll index", ylabel="phase (nm rms)", + title="aberrations (5, 6: astigmatism; 7, 8: coma; 11: spherical)") + if len(_models(calibration)) > 1: + ax.legend() diff --git a/src/smappy/uipsf/__init__.py b/src/smappy/uipsf/__init__.py new file mode 100644 index 0000000..274cf16 --- /dev/null +++ b/src/smappy/uipsf/__init__.py @@ -0,0 +1,102 @@ +"""PSF models learnt by uiPSF, made into smappy calibrations. + +uiPSF (Liu et al. 2024, https://github.com/ries-lab/uiPSF) learns a PSF by +inverse modelling: a pupil -- Zernike polynomials, or free -- is propagated +to the camera, and its parameters are fitted to the data with the emitters' +positions. From bead stacks that is a calibration like smappy's own, +with a physical model behind it; from the blinking molecules of the +acquisition itself (*in situ*) it is the PSF the data were actually taken +with, aberrations of the sample and all, which no bead on a coverslip shows. + +The work is split four ways, so that each part can be read and tested alone: + +* `prepare` reads the data, converts it to photons, and cuts a split camera + into its two channels, recording exactly how; +* `uipsf_parameters` turns a microscope profile (`microscopes`) and the + user's settings into uiPSF's configuration; +* `runner` runs uiPSF -- in its own Python, through `worker.py` -- and + follows its progress; +* `convert` makes its result into smappy's calibration file, which the + Spline 3D and Spline 3D 2C fitters read as they read a bead calibration. + +`learn` does all four. The plugin is ``Localize/uiPSF calibration`` +(`smappy.plugins.uipsf`). +""" +from __future__ import annotations + +import tempfile +from pathlib import Path +from typing import Callable, Optional + +from .microscopes import Microscope + +# uiPSF's configuration files by what is learnt from what +PSFTYPES = {("beads", "zernike"): ("zernike", None), + ("beads", "voxel"): ("voxel", None), + ("blinking", "zernike"): ("insitu", None), + ("blinking", "pupil"): ("insitu", "insitu_pupil")} + + +def uipsf_parameters(profile: Microscope, *, data: str, pixelsize_xy_um, dz_nm: float, + roi_size_px: Optional[int] = None, emission_wavelength_nm: Optional[float] = None, + peak_height: float = 0.2, max_beads: int = 40, bead_diameter_nm: float = 0.0, + iterations: int = 100, zernike_order: int = 8, + depth_um: float = 1.0, z_range_um: float = 2.0, min_photon: float = 0.4, + z_bins: int = 21, per_bin: int = 100, repeat: int = 2, + start_from: str = "") -> dict: + """uiPSF's parameters, nested as its YAML, for one run. + + What the profile does not say is uiPSF's default; what the profile's own + ``uipsf`` section says is applied last, so a microscope with an unusual + need (a deformable mirror's own Zernike terms) can say so there. + """ + ri = profile.refractive_index + roi = int(roi_size_px or profile.roi_size_px) + params = { + "pixel_size": {"x": float(pixelsize_xy_um[0]), "y": float(pixelsize_xy_um[1]), + "z": float(dz_nm) / 1000}, + "roi": {"roi_size": [roi, roi], "peak_height": float(peak_height), + "max_bead_number": int(max_beads), + "bead_radius": float(bead_diameter_nm) / 2000}, + "gain": 1.0, "ccd_offset": 0.0, # smappy hands over photons + "iteration": int(iterations), + "option": { + "imaging": {"emission_wavelength": float(emission_wavelength_nm + or profile.emission_wavelength_nm) / 1000, + "RI": {"imm": float(ri["immersion"]), "med": float(ri["medium"]), + "cov": float(ri["coverslip"])}, + "NA": float(profile.NA)}, + "model": {"n_max": int(zernike_order), "init_pupil_file": start_from or ""}, + }, + } + if data == "blinking": + insitu = dict(profile.insitu) + params["option"]["insitu"] = { + "stage_pos": float(depth_um), "z_range": float(z_range_um), + "min_photon": float(min_photon), "partition_size": [int(z_bins), int(per_bin)], + "repeat": int(repeat), + "zernike_index": list(insitu.get("zernike_index", [5])), + "zernike_coeff": list(insitu.get("zernike_coeff", [0.5]))} + _merge(params, profile.uipsf or {}) + return params + + +def _merge(target: dict, extra: dict) -> None: + for key, value in extra.items(): + if isinstance(value, dict) and isinstance(target.get(key), dict): + _merge(target[key], value) + else: + target[key] = value + + +def learn(images, *, data: str, model: str, dual: bool, params: dict, + workdir=None, python: str = "", + progress: Optional[Callable[[str], None]] = None) -> Path: + """Run uiPSF on ``images`` (`prepare`'s); the result file it wrote.""" + from . import runner + psftype, override = PSFTYPES[(data, model)] + if override: + params = dict(params, PSFtype=override) + workdir = Path(workdir) if workdir else Path(tempfile.mkdtemp(prefix="smappy-uipsf-")) + runner.write_job(workdir, images, psftype, "2ch" if dual else "1ch", params) + return runner.run(workdir, python=python, progress=progress) diff --git a/src/smappy/uipsf/convert.py b/src/smappy/uipsf/convert.py new file mode 100644 index 0000000..01db1b9 --- /dev/null +++ b/src/smappy/uipsf/convert.py @@ -0,0 +1,264 @@ +"""uiPSF's result file -> smappy's calibration file. + +uiPSF writes the PSF it learnt as voxels (``res/I_model``, z, y, x) beside +the pupil and its Zernike coefficients, and a cubic spline of it that SMAP +reads (``locres/coeff``). The spline is not used here: uiPSF normalises it by +the median plane sum after subtracting the minimum, which is not what the +fitter's photon number assumes, and its grid is uiPSF's. The voxels are +taken instead and made into a calibration the way smappy's own bead +calibration makes one (`calibrate.core._single_model`, +`positive_pair_models`): clipped at zero, normalised, splined by +`spline_coefficients` -- so the fitter cannot tell where a calibration came +from, and the pupil and coefficients go along in ``parameters`` as the +record of what was learnt. + +Conventions, each checked by fitting simulated data with the result +(tests/test_uipsf.py): + +* **z.** A bead model's planes are uiPSF's ``Zrange``, centred on the + pupil's focus and in the order of the stack (objective moving up), which is + smappy's: ``z0 = (nz - 1) / 2``, emitter ``z = -(index - z0) dz``. An in + situ model's planes run the other way -- the emitter's height above the + coverslip -- so they are reversed, and z = 0 is put at the height the + objective's focus sat at (uiPSF's learnt ``stagepos``, scaled by the + refractive indices), which is the z a bead calibration calls 0. +* **x, y.** numpy's (y, x) throughout; nothing is swapped, because smappy + hands uiPSF arrays rather than files (uiPSF's ``swapxy`` is for SMAP's + transposed TIFF reading). +* **two channels.** uiPSF's ``T`` maps the reference channel's (y, x, 1) row + vectors, about ``imgcenter``, to the other channel's, both in the arrays it + was handed; `prepare.Split.to_chip` puts both into chip pixels, and the + projective map smappy wants (secondary -> main, x, y) is fitted to points + pushed through it. The secondary PSF is flipped back to the camera's + orientation, as smappy's dual calibration stores it. +""" +from __future__ import annotations + +import json +from pathlib import Path +from typing import Dict, Optional + +import numpy as np + +NOLL_NAMES = {1: "piston", 2: "tilt x", 3: "tilt y", 4: "defocus", + 5: "astigmatism 45", 6: "astigmatism 0", 7: "coma y", 8: "coma x", + 9: "trefoil y", 10: "trefoil x", 11: "spherical", + 12: "2nd astigmatism 0", 13: "2nd astigmatism 45", + 22: "2nd spherical", 37: "3rd spherical"} + + +def read(path) -> Dict: + """uiPSF's result: ``{"res": ..., "locres": ..., "params": ...}`` as dicts.""" + import h5py + + def group(g): + return {k: group(v) if isinstance(v, h5py.Group) else v[()] for k, v in g.items()} + with h5py.File(path, "r") as f: + if "res" not in f or "params" not in f.attrs: + raise ValueError(f"{path} is not a uiPSF result (no res group or params)") + out = {k: group(f[k]) for k in f if isinstance(f[k], h5py.Group)} + out["params"] = json.loads(f.attrs["params"]) + return out + + +def _is_insitu(params) -> bool: + return "insitu" in str(params.get("PSFtype", "")) + + +def zernike_table(res: Dict, params: Dict) -> list: + """(Noll index, name, magnitude, phase in nm rms) for each learnt term.""" + coeff = res.get("zernike_coeff") + if coeff is None: + return [] + coeff = np.asarray(coeff).reshape(2, -1) + nl = params.get("option", {}).get("model", {}).get("zernike_nl") or [] + if nl: + noll = [_nl2noll(n, l) for n, l in nl] + else: + noll = list(range(1, coeff.shape[1] + 1)) + wavelength_nm = 1000 * float(params["option"]["imaging"]["emission_wavelength"]) + return [(int(j), NOLL_NAMES.get(int(j), ""), float(m), float(p) * wavelength_nm / (2 * np.pi)) + for j, m, p in zip(noll, coeff[0], coeff[1])] + + +def _nl2noll(n, l): + """uiPSF's `nl2noll`: Noll's index of radial order n, azimuthal l.""" + m = abs(l) + j = n * (n + 1) // 2 + 1 + max(0, m - 1) + if (l > 0 and n % 4 >= 2) or (l < 0 and n % 4 <= 1): + j += 1 + return int(j) + + +def _focus_index(res: Dict, params: Dict, nz: int) -> float: + """Where z = 0 lies in an in situ model, in its (reversed) planes. + + uiPSF's in situ planes are emitter heights ``zoffset + i`` (in z steps) + above the coverslip, the objective's focus sitting at ``stagepos`` (um) + of stage travel, which in the sample is the nominal focal plane + ``stagepos * n_medium / n_immersion`` above it. That plane is z = 0, as + a bead stack's centre is, and the planes are reversed to run as a bead + stack does; outside the model's range the middle is taken instead. + """ + dz_um = float(params["pixel_size"]["z"]) + ri = params["option"]["imaging"]["RI"] + stage = float(np.ravel(res.get("stagepos", [params["option"]["insitu"]["stage_pos"]]))[0]) + zoffset = float(np.real(np.ravel(res.get("zoffset", [0.0]))[0])) + height = stage * float(ri["med"]) / float(ri["imm"]) / dz_um # in planes + index = height - zoffset # unreversed + reversed_index = (nz - 1) - index + return float(reversed_index) if 0 <= reversed_index <= nz - 1 else (nz - 1) / 2 + + +def _model(res: Dict, params: Dict): + """The channel's PSF voxels in smappy's z order, and its z0.""" + psf = np.asarray(res["I_model"], np.float64) + if psf.ndim != 3: + raise ValueError(f"uiPSF's I_model has shape {psf.shape}; a 3D model was expected " + "(field-dependent and 4Pi models are not supported)") + if _is_insitu(params): + z0 = _focus_index(res, params, len(psf)) + psf = psf[::-1] + else: + z0 = (len(psf) - 1) / 2 + return np.real(psf), z0 + + +def _record(params: Dict, res: Dict, path, **extra) -> Dict: + imaging = params["option"]["imaging"] + out = {"method": "uipsf_" + str(params.get("PSFtype")), + "source": str(path), + "uipsf_psftype": params.get("PSFtype"), + "normalization": "maximum plane sum = 1", + "z_reference": "uiPSF focus" if not _is_insitu(params) + else "nominal focal plane (stagepos * n_med / n_imm)", + "z_convention": "emitter_nm = -(index-z0)*dz", + "lateral_unit": "pixel", + "NA": imaging["NA"], "refractive_index": imaging["RI"], + "emission_wavelength_um": imaging["emission_wavelength"], + "pixel_size_um": params["pixel_size"], + "zernike": [list(row) for row in zernike_table(res, params)]} + for key in ("stagepos", "zoffset", "sigma"): + if key in res: + out["uipsf_" + key] = np.real(np.ravel(res[key])).astype(float).tolist() + out.update(extra) + return out + + +def single_calibration(path, em_on: Optional[bool] = None): + """A `SplineCalibration` of a one-channel uiPSF result.""" + from ..calibrate.core import spline_coefficients + from ..io.calibration import SplineCalibration + data = read(path) + res, params = data["res"], data["params"] + if "channel0" in res: + raise ValueError(f"{path} is a multi-channel uiPSF result: use dual_calibration") + psf, z0 = _model(res, params) + psf = np.maximum(psf, 0) + norm = float(psf.sum(axis=(1, 2)).max()) + if not np.isfinite(norm) or norm <= 0: + raise ValueError("uiPSF's model has no positive signal") + psf = (psf / norm).astype(np.float32) + return SplineCalibration(spline_coefficients(psf).astype(np.float32), + 1000 * float(params["pixel_size"]["z"]), z0, + x0=(psf.shape[1] - 1) / 2, psf=psf, em_mirror=False, + source=Path(path), em_on=em_on, + parameters=_record(params, res, path)) + + +# the planes either side of focus the channels' light is compared over, as +# smappy's dual calibration does (calibrate.core.FOCAL_PLANES) +FOCAL_PLANES = 2 + + +def dual_calibration(path, split, sources=(), em_on: Optional[bool] = None): + """A `DualColorCalibration` of a two-channel uiPSF result. + + ``split`` is the `prepare.Split` the channels were cut with: without it + the transformation cannot be put back on the chip. + """ + from ..calibrate.core import spline_coefficients + from ..calibrate.dual import DualColorCalibration + from ..io.calibration import SplineCalibration + data = read(path) + res, params = data["res"], data["params"] + if "channel1" not in res or "T" not in res: + raise ValueError(f"{path} is not a two-channel uiPSF result") + models, intensities, light = [], [], [] + for k in (0, 1): + psf, z0 = _model(res[f"channel{k}"], params) + models.append((psf, z0)) + centre = int(round(z0)) + window = slice(max(centre - FOCAL_PLANES, 0), centre + FOCAL_PLANES + 1) + # the channel's light per emitter at focus: the model's focal signal + # times the photons uiPSF fitted to it -- uiPSF normalises each + # channel's model on its own, so the split between them is in the + # intensities + intensities.append(float(np.median(np.real(res[f"channel{k}"]["intensity"])))) + light.append(float(psf[window].sum(axis=(1, 2)).mean()) * intensities[-1]) + if models[0][0].shape != models[1][0].shape or models[0][1] != models[1][1]: + raise ValueError("uiPSF's two channel models differ in their grid") + total = sum(light) + if not np.isfinite(total) or total <= 0 or min(light) <= 0: + raise ValueError("uiPSF's channel intensities give no valid photon split") + mirror = split.geometry(sources)["mirror_axis_xy"] + out = [] + for k, ((psf, z0), share) in enumerate(zip(models, light)): + # one factor for the pair, as `positive_pair_models`: the channels' + # light around focus sums to one, each keeping its share + psf = np.maximum(psf, 0) * intensities[k] / total + if k == 1 and mirror is not None: + psf = np.flip(psf, axis=2 - mirror) # back to the camera's orientation + psf = np.ascontiguousarray(psf, np.float32) + out.append(SplineCalibration( + spline_coefficients(psf).astype(np.float32), 1000 * float(params["pixel_size"]["z"]), + z0, x0=(psf.shape[1] - 1) / 2, psf=psf, em_mirror=False, source=Path(path), + em_on=em_on, parameters=_record( + params, res[f"channel{k}"], path, + normalization="joint: both channels' light around focus sums to 1", + photon_normalization=share / total))) + transformation = channel_transformation(res["T"], res["imgcenter"], split) + calibration_parameters = { + "method": "uipsf_" + str(params.get("PSFtype")), + "source": str(path), + "secondary_main_brightness_ratio": light[1] / light[0], + "transform_method": "uiPSF's affine channel transformation, learnt with the PSF", + "transform_direction": "secondary to main, native camera coordinates", + "uipsf_T": np.asarray(res["T"]).tolist(), + "uipsf_imgcenter": np.asarray(res["imgcenter"]).tolist()} + return DualColorCalibration(out[0], out[1], transformation, + split.geometry(sources), calibration_parameters) + + +def uipsf_to_main(T, imgcenter, yx) -> np.ndarray: + """uiPSF's channel map: reference (y, x) -> the other channel's (y, x).""" + T = np.asarray(T, float).reshape(-1, 3, 3)[0] + c = np.asarray(imgcenter, float).ravel() + rows = np.c_[np.asarray(yx, float), np.ones(len(yx))] - c + return (rows @ T)[:, :2] + c[:2] + + +def channel_transformation(T, imgcenter, split) -> np.ndarray: + """smappy's 3x3 secondary -> main chip (x, y) map from uiPSF's ``T``.""" + from ..calibrate.dual import fit_projective + n = split.length + other = (split.image_shape[1 - split.axis]) + shape = [0, 0] + shape[split.axis], shape[1 - split.axis] = n, other + yy, xx = np.meshgrid(np.linspace(0, shape[0] - 1, 9), np.linspace(0, shape[1] - 1, 9), + indexing="ij") + main_yx = np.c_[yy.ravel(), xx.ravel()] + secondary_yx = uipsf_to_main(T, imgcenter, main_yx) + return fit_projective(split.to_chip(1, secondary_yx), split.to_chip(0, main_yx)) + + +def save(path, calibration, overwrite: bool = True) -> Path: + """Write either kind as smappy's calibration file.""" + from ..calibrate.dual import DualColorCalibration, save_dual_color_calibration + from ..io.calibration import save_spline_calibration + path = Path(path) + if isinstance(calibration, DualColorCalibration): + save_dual_color_calibration(path, calibration, overwrite=overwrite) + else: + save_spline_calibration(path, calibration, overwrite=overwrite) + return path diff --git a/src/smappy/uipsf/microscopes.py b/src/smappy/uipsf/microscopes.py new file mode 100644 index 0000000..0764441 --- /dev/null +++ b/src/smappy/uipsf/microscopes.py @@ -0,0 +1,138 @@ +"""Microscope profiles: what uiPSF needs to know about the instrument. + +A PSF model is learnt from the optics outward -- numerical aperture, +refractive indices, the emission wavelength -- and a user calibrating beads on +Tuesday should not type those in again. A profile is a small YAML file per +microscope, kept where the group shares it, with the values uiPSF's own +``config/systemtype/*.yaml`` files hold under names a user reads:: + + name: M2, astigmatic, two colour + NA: 1.43 + refractive_index: {immersion: 1.516, medium: 1.335, coverslip: 1.516} + emission_wavelength_nm: 680 + pixelsize_um: [0.127, 0.116] # optional: else from the camera + roi_size_px: 25 # the model's lateral size + dual: # a split camera; omit for one channel + layout: up-down mirrored # as the bead calibration: smappy.calibrate.dual.LAYOUTS + main_channel: upper + split_position: null # null: the middle + insitu: # the pupil learning from blinking starts at + zernike_index: [5] # Noll index (5: astigmatism) + zernike_coeff: [0.5] # its amplitude, in radians + uipsf: {} # any uiPSF parameter, nested as in its YAML + +The camera -- gain, offset, EM gain -- is not here: it comes from the +acquisition and the camera database, as it does for fitting. Pixel size is +both: the camera's file says it when the microscope software was told, and a +profile is the fallback for one that was not. + +Profiles are looked for in ``$SMAPPY_MICROSCOPES`` (one or more directories, +separated as ``PATH`` is), ``/microscopes`` and the shipped +``smappy/data/microscopes``; the first file of a name wins, so a group's +directory overrides what ships. +""" +from __future__ import annotations + +import os +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any, Dict, List, Optional, Tuple + +import yaml + +SHIPPED = Path(__file__).resolve().parents[1] / "data" / "microscopes" + + +@dataclass +class Microscope: + """One profile, with uiPSF's defaults for what it leaves out.""" + name: str = "default" + NA: float = 1.43 + refractive_index: Dict[str, float] = field(default_factory=lambda: dict( + immersion=1.516, medium=1.335, coverslip=1.516)) + emission_wavelength_nm: float = 680.0 + pixelsize_um: Optional[Any] = None + roi_size_px: int = 25 + dual: Optional[Dict[str, Any]] = None + insitu: Dict[str, Any] = field(default_factory=lambda: dict( + zernike_index=[5], zernike_coeff=[0.5])) + uipsf: Dict[str, Any] = field(default_factory=dict) + path: Optional[Path] = None + + @classmethod + def from_yaml(cls, path) -> "Microscope": + path = Path(path) + data = yaml.safe_load(path.read_text()) or {} + if not isinstance(data, dict): + raise ValueError(f"{path}: a microscope profile is a YAML mapping") + known = set(cls.__dataclass_fields__) - {"path"} + unknown = set(data) - known + if unknown: + raise ValueError(f"{path}: unknown keys {sorted(unknown)}; a profile " + f"has {sorted(known)} (uiPSF's own go under 'uipsf')") + data.setdefault("name", path.stem) + profile = cls(**data, path=path) + profile.validate() + return profile + + def validate(self) -> None: + if not 0 < self.NA < 2: + raise ValueError(f"{self.name}: NA {self.NA} is not a numerical aperture") + for key in ("immersion", "medium", "coverslip"): + if key not in self.refractive_index: + raise ValueError(f"{self.name}: refractive_index needs {key}") + if self.dual is not None: + from ..calibrate.dual import LAYOUTS + layout = self.dual.get("layout", "up-down mirrored") + if layout not in LAYOUTS: + raise ValueError(f"{self.name}: dual layout must be one of {LAYOUTS}") + + @property + def pixelsize_xy_um(self) -> Optional[Tuple[float, float]]: + p = self.pixelsize_um + if p is None: + return None + if isinstance(p, (list, tuple)): + return float(p[0]), float(p[-1]) + return float(p), float(p) + + +def directories() -> List[Path]: + from ..config import config_dir + found = [Path(p).expanduser() for p in + os.environ.get("SMAPPY_MICROSCOPES", "").split(os.pathsep) if p] + return found + [config_dir() / "microscopes", SHIPPED] + + +def profiles() -> Dict[str, Path]: + """Every profile there is, by file name, the first of a name winning.""" + out: Dict[str, Path] = {} + for d in directories(): + if d.is_dir(): + for p in sorted(d.glob("*.yaml")): + out.setdefault(p.stem, p) + return out + + +def load(name_or_path: str) -> Microscope: + """A profile by its name, or a YAML file by its path; '' is uiPSF's defaults.""" + if not name_or_path: + return Microscope() + found = profiles() + if name_or_path in found: + return Microscope.from_yaml(found[name_or_path]) + if Path(name_or_path).exists(): + return Microscope.from_yaml(name_or_path) + raise ValueError(f"no microscope profile {name_or_path!r}; there are " + f"{sorted(found)} in {[str(d) for d in directories()]}") + + +def choices() -> List[Tuple[str, str]]: + """For the plugin's drop-down: (name, the profile's own name).""" + out = [("", "uiPSF defaults")] + for stem, path in profiles().items(): + try: + out.append((stem, Microscope.from_yaml(path).name)) + except (OSError, ValueError, yaml.YAMLError): + out.append((stem, f"{stem} (unreadable)")) + return out diff --git a/src/smappy/uipsf/prepare.py b/src/smappy/uipsf/prepare.py new file mode 100644 index 0000000..58a9b54 --- /dev/null +++ b/src/smappy/uipsf/prepare.py @@ -0,0 +1,278 @@ +"""What uiPSF is handed: photons, in channels, cut as it wants them. + +smappy reads the data rather than uiPSF's loader, so that a bead stack's z +positions, the camera's gain and offset and the split of a two-colour camera +are found exactly as they are for smappy's own calibration and fit -- and the +transformation uiPSF learns can be put back into chip coordinates, which only +works if this side knows exactly how the halves were cut. + +uiPSF's arrays: + +* beads, one channel: ``(stacks, z, y, x)``; two: ``(2, stacks, z, y, x)`` +* blinking, one channel: ``(frames, y, x)``; two: ``(2, frames, y, x)`` + +with the main channel first (``ref_channel = 0``) and the secondary flipped +back to the main channel's orientation when the splitter mirrors it. +""" +from __future__ import annotations + +from dataclasses import dataclass +from typing import Optional, Sequence + +import numpy as np + + +@dataclass +class Split: + """How a split camera frame was cut into the two channels uiPSF sees. + + The halves are cut to the same length either side of ``split_position`` + -- ``[split-m, split)`` and ``[split, split+m)`` along the split axis -- + because uiPSF stacks the two channels into one array; the secondary is + flipped along that axis when the layout is mirrored. `to_chip` undoes + exactly this, which is what turns uiPSF's channel transformation back + into chip pixels. + """ + layout: str + main_channel: str + split_position: int + image_shape: tuple # (height, width) of the whole frame + length: int # m, each half's extent along the split axis + origin: tuple = (0, 0) # the frame's (x, y) on the chip, or (0, 0) + + @property + def axis(self) -> int: + """The array axis of a (y, x) frame that is split.""" + return 1 if "right-left" in self.layout else 0 + + @property + def mirrored(self) -> bool: + return "mirrored" in self.layout + + @property + def main_first(self) -> bool: + return self.main_channel in ("left", "upper") + + def starts(self): + """Where (main, secondary) begin along the split axis.""" + low, high = self.split_position - self.length, self.split_position + return (low, high) if self.main_first else (high, low) + + def cut(self, images: np.ndarray) -> np.ndarray: + """``(..., y, x)`` -> ``(2, ..., y', x')``, main first, secondary unmirrored.""" + out = [] + for channel, start in enumerate(self.starts()): + index = [slice(None)] * images.ndim + index[images.ndim - 2 + self.axis] = slice(start, start + self.length) + half = images[tuple(index)] + if channel == 1 and self.mirrored: + half = np.flip(half, axis=images.ndim - 2 + self.axis) + out.append(half) + return np.stack(out) + + def to_chip(self, channel: int, yx) -> np.ndarray: + """Pixel (y, x) positions in channel ``channel``'s array -> chip (x, y).""" + yx = np.array(yx, dtype=float, copy=True) + along = yx[:, self.axis] + if channel == 1 and self.mirrored: + along = self.length - 1 - along + yx[:, self.axis] = along + self.starts()[channel] + return yx[:, ::-1] + np.asarray(self.origin, float) + + def geometry(self, sources: Sequence[dict] = ()) -> dict: + """The split as `smappy.calibrate.dual` records it with a calibration.""" + coord_axis = 1 - self.axis + return {"layout": self.layout, "main_channel": self.main_channel, + "split_position": int(self.split_position), + "image_shape": list(self.image_shape), + "coordinate_system": "camera-chip" if any( + s.get("roi") is not None for s in sources) else "roi-local", + "sources": list(sources), + "mirror_axis_xy": coord_axis if self.mirrored else None, + "convention": "zero-based pixel centers; split is first index of second half"} + + +def make_split(shape, dual: dict, origin=(0, 0)) -> Split: + """The `Split` of a frame of ``shape`` for a profile's ``dual`` section.""" + from ..calibrate.dual import DualColorSettings + settings = DualColorSettings(layout=dual.get("layout", "up-down mirrored"), + main_channel=dual.get("main_channel", "upper"), + split_position=dual.get("split_position")) + settings.validate() + shape = tuple(int(n) for n in shape[-2:]) + axis = 1 if "right-left" in settings.layout else 0 + split = settings.split_position or shape[axis] // 2 + if not 0 < split < shape[axis]: + raise ValueError(f"the split at {split} px is outside the {shape[axis]} px frame") + return Split(settings.layout, settings.main_channel, int(split), shape, + min(split, shape[axis] - split), tuple(origin)) + + +def channel_shift(channels: np.ndarray) -> list: + """How far the secondary channel's image lies from the main's, [y, x] px. + + uiPSF pairs the two channels' beads from a first guess at this shift, + which it finds by matching bead coordinates -- and on a regular pattern + of beads that locked onto the neighbouring bead, 30 px off, leaving one + pair. The cross-correlation of the two channels' projections peaks where + *all* the beads overlap, which a neighbour's shift cannot match. Handed + to uiPSF as ``channel_shift``: per channel, target minus reference. + """ + from scipy import ndimage + ny, nx = channels.shape[-2:] + a, b = (np.asarray(c, float).reshape(-1, ny, nx).max(axis=0) for c in channels) + a, b = a - np.median(a), b - np.median(b) + # b(r + d) = a(r) where this peaks: the secondary is the main shifted by d + correlation = np.fft.ifft2(np.conj(np.fft.fft2(a)) * np.fft.fft2(b)).real + correlation = np.fft.fftshift(ndimage.gaussian_filter(np.fft.fftshift(correlation), 1)) + dy, dx = np.unravel_index(np.argmax(correlation), correlation.shape) + dy = dy - correlation.shape[0] if dy > correlation.shape[0] // 2 else dy + dx = dx - correlation.shape[1] if dx > correlation.shape[1] // 2 else dx + return [[0.0, 0.0], [float(dy), float(dx)]] + + +def to_photons(adu: np.ndarray, camera) -> np.ndarray: + """Camera counts -> photons, as the fitter converts them.""" + return ((np.asarray(adu, np.float32) - float(camera.offset)) + * float(camera.adu_to_photons)).astype(np.float32) + + +def resolve_camera(camera_settings, source, profile): + """The camera, with the profile's pixel size where nothing else gave one.""" + from ..io.tiff import camera_metadata + from ..metadata import CameraMetadata + overrides = camera_settings.overrides() if camera_settings is not None else {} + if camera_settings is not None and camera_settings.preset: + from dataclasses import asdict + preset = asdict(CameraMetadata.from_yaml(camera_settings.preset)) + overrides = {**{k: v for k, v in preset.items() if v is not None}, **overrides} + name = camera_settings.camera if camera_settings is not None else "" + if source is not None: + meta = camera_metadata(source, overrides=overrides, camera=name, require=False) + else: + meta = CameraMetadata.from_dict(overrides) + if meta.pixelsize_um is None and profile.pixelsize_xy_um is not None: + meta.pixelsize_um = list(profile.pixelsize_xy_um) + meta.require() + return meta + + +def pixel_size_xy_um(camera) -> tuple: + p = camera.pixelsize_um + if isinstance(p, (list, tuple, np.ndarray)): + return float(p[0]), float(p[-1]) + return float(p), float(p) + + +@dataclass +class BeadInput: + images: np.ndarray # photons, as uiPSF wants them + dz_nm: float + camera: object + sources: list # per stack: path, roi, shape, z + split: Optional[Split] = None + + +def read_beads(paths, camera_settings, profile, dual: bool, dz_nm: Optional[float] = None, + progress=None) -> BeadInput: + """Bead z-stacks, from files or `BeadStack`s, as uiPSF's photons. + + Stacks of different lengths are cut to the shortest, about their centres: + uiPSF stacks them into one array, and fits each bead's z itself. + """ + from ..calibrate.input import BeadStack, discover_acquisitions, read_bead_stacks + from ..io.tiff import open_stack + report = progress or (lambda text: None) + if isinstance(paths, (str, BeadStack)): + paths = [paths] + stacks, first_path = [], None + for item in paths: + if isinstance(item, BeadStack): + stacks.append(item) + continue + for acquisition in discover_acquisitions([item]): + first_path = first_path or acquisition + stacks.extend(read_bead_stacks(acquisition, dz_nm)) + if not stacks: + raise ValueError("no bead stacks: choose the files of a z-stack acquisition") + report(f"{len(stacks)} bead stacks") + shapes = {s.images.shape[-2:] for s in stacks} + if len(shapes) > 1: + raise ValueError(f"the bead stacks differ in frame size {sorted(shapes)}; " + "uiPSF needs them alike: calibrate them separately") + steps = np.array([s.dz for s in stacks]) + if dz_nm is None and np.ptp(steps) > 0.01 * np.median(steps): + raise ValueError(f"the stacks have different z steps {sorted(set(steps.round(1)))} nm") + dz = float(dz_nm or np.median(steps)) + nz = min(len(s.images) for s in stacks) + cut = [s.images[(len(s.images) - nz) // 2:(len(s.images) - nz) // 2 + nz] for s in stacks] + source = None + if first_path is not None: + try: + source = open_stack(first_path) + except Exception: # noqa: BLE001 -- the user's camera fields remain + source = None + camera = resolve_camera(camera_settings, source, profile) + if camera.roi is None and stacks[0].roi is not None: + camera.roi = stacks[0].roi + images = to_photons(np.stack(cut), camera) + sources = [{"path": str(s.source), "roi": list(s.roi) if s.roi is not None else None, + "shape": list(s.images.shape), "z_nm": np.asarray(s.z_nm).tolist(), + "axes": s.axes} for s in stacks] + split = None + if dual: + rois = {tuple(s.roi) for s in stacks if s.roi is not None} + if len(rois) > 1: + raise ValueError("the bead stacks were taken with different camera ROIs") + origin = tuple(next(iter(rois))[:2]) if rois else (0, 0) + split = make_split(images.shape, profile.dual or {}, origin) + images = split.cut(images) + return BeadInput(subtract_background(images), dz, camera, sources, split) + + +def subtract_background(images: np.ndarray) -> np.ndarray: + """Each bead stack less its median, per channel. + + uiPSF finds beads above a threshold relative to the brightest maximum of + the z projection *including* the background: with a background a fifth of + the peak -- two-colour beads, half the light each -- the noise of the + projection passed it everywhere and nine beads in ten were lost as too + close to a spurious neighbour. uiPSF fits a background per bead anyway, + so taking a constant off changes the model it learns nothing. + """ + images = np.asarray(images, np.float32) + return images - np.median(images, axis=(-3, -2, -1), keepdims=True) + + +@dataclass +class MovieInput: + images: np.ndarray + camera: object + source: dict + split: Optional[Split] = None + + +def read_movie(path, camera_settings, profile, dual: bool, start: int = 0, + stop: Optional[int] = None, progress=None) -> MovieInput: + """Frames of a blinking acquisition, as uiPSF's photons.""" + from ..io.tiff import open_stack + report = progress or (lambda text: None) + source = open_stack(path) + camera = resolve_camera(camera_settings, source, profile) + stop = len(source) if stop is None else min(stop, len(source)) + if stop - start < 100: + raise ValueError(f"frames {start} to {stop}: too few to learn a PSF from") + blocks = [] + for first, block in source.frames(chunk=500, start=start, stop=stop): + blocks.append(to_photons(block, camera)) + report(f"read {first + len(block) - start} of {stop - start} frames") + images = np.concatenate(blocks) + roi = list(camera.roi) if camera.roi is not None else None + info = {"path": str(path), "roi": roi, "shape": list(images.shape), + "frames": [int(start), int(stop)]} + split = None + if dual: + origin = tuple(roi[:2]) if roi is not None else (0, 0) + split = make_split(images.shape, profile.dual or {}, origin) + images = split.cut(images) + return MovieInput(images, camera, info, split) diff --git a/src/smappy/uipsf/runner.py b/src/smappy/uipsf/runner.py new file mode 100644 index 0000000..3be303c --- /dev/null +++ b/src/smappy/uipsf/runner.py @@ -0,0 +1,144 @@ +"""Starting uiPSF in its own Python and following it. + +uiPSF is not a dependency of smappy and is not imported here. It pins its own +TensorFlow, and a TensorFlow in the GUI's process would cost seconds of import +and gigabytes of memory for every user who never learns a PSF; it is run +instead as `worker.py` under the interpreter of an environment that has it, +the way the batch window runs `smappy-batch`. The same arrangement is what +lets the work move to another machine (a GPU) later without the plugin +knowing. + +Which interpreter: the plugin's own setting, else ``uipsf_python`` in the +config file (`smappy.config`), else ``$SMAPPY_UIPSF_PYTHON``, else this +Python if it can import ``psflearning``. +""" +from __future__ import annotations + +import json +import os +import re +import shutil +import subprocess +import sys +import time +from pathlib import Path +from typing import Callable, Optional + +import numpy as np + +TAG = "SMAPPY_UIPSF" +WORKER = Path(__file__).with_name("worker.py") +# uiPSF's progress bars: "3/6: learning: 99/150 [02:48s] ... current loss: 0.53" +_BAR = re.compile(r"^(\d/\d: [^:]+)(?::\s*(\d+)/(\d+))?") +_ESCAPE = re.compile(r"\x1b\[[0-9;]*[A-Za-z]") +# how often a progress bar that keeps moving is passed on, in seconds +REPORT_EVERY_S = 5.0 + + +class UiPSFError(RuntimeError): + """uiPSF could not be started, or stopped with an error.""" + + +def find_python(setting: str = "") -> str: + """The interpreter of an environment with uiPSF in it.""" + from .. import config + for candidate in (setting, config.get("uipsf_python") or "", + os.environ.get("SMAPPY_UIPSF_PYTHON", "")): + if candidate: + path = shutil.which(candidate) or candidate + if not Path(path).exists(): + raise UiPSFError(f"no Python at {candidate}: set the uiPSF " + "environment's python in the plugin, or " + "uipsf_python in the smappy config file") + return path + import importlib.util + if importlib.util.find_spec("psflearning") is not None: + return sys.executable + raise UiPSFError( + "uiPSF was not found: install it into an environment of its own " + "(see the plugin's page) and give that environment's python in the " + "plugin's 'uiPSF python', or as uipsf_python in the smappy config file") + + +def write_job(directory: Path, images: np.ndarray, psftype: str, channeltype: str, + params: dict) -> Path: + """The job directory the worker reads.""" + directory = Path(directory) + directory.mkdir(parents=True, exist_ok=True) + np.save(directory / "images.npy", np.ascontiguousarray(images, dtype=np.float32)) + (directory / "job.json").write_text(json.dumps( + {"psftype": psftype, "channeltype": channeltype, "params": params}, indent=1)) + return directory + + +def _lines(stream): + """Lines of a byte stream, a progress bar's carriage returns ending one too.""" + buffer = b"" + while True: + chunk = stream.read1(4096) if hasattr(stream, "read1") else stream.read(4096) + if not chunk: + break + buffer += chunk + parts = re.split(rb"[\r\n]", buffer) + buffer = parts.pop() + for part in parts: + text = _ESCAPE.sub("", part.decode(errors="replace")).strip() + if text: + yield text + text = _ESCAPE.sub("", buffer.decode(errors="replace")).strip() + if text: + yield text + + +def run(directory: Path, python: str = "", progress: Optional[Callable[[str], None]] = None, + log: Optional[Path] = None) -> Path: + """Run the job in ``directory``; the uiPSF result file it wrote. + + ``progress`` gets the worker's stages and uiPSF's progress bars, the + latter at most every `REPORT_EVERY_S`; everything the worker printed goes + to ``log`` (``directory/uipsf.log`` by default), which is where to look + when it fails. + """ + report = progress or (lambda text: None) + python = find_python(python) + directory = Path(directory) + log = Path(log) if log else directory / "uipsf.log" + env = dict(os.environ, MPLBACKEND="Agg", PYTHONUNBUFFERED="1", + TF_CPP_MIN_LOG_LEVEL="2") + command = [python, "-u", str(WORKER), str(directory)] + try: + process = subprocess.Popen(command, stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, env=env) + except OSError as error: + raise UiPSFError(f"could not start {python}: {error}") from error + result = error = None + last, shown = "", 0.0 + try: + with open(log, "w") as out: + for line in _lines(process.stdout): + out.write(line + "\n") + if line.startswith(TAG + "\t"): + _, kind, text = (line.split("\t", 2) + [""])[:3] + if kind == "result": + result = Path(text) + elif kind == "error": + error = text + else: + report(f"uiPSF: {text}") + continue + bar = _BAR.match(line) + if bar: + stage = bar.group(1) + now = time.monotonic() + if stage != last or now - shown > REPORT_EVERY_S: + report(f"uiPSF {line[:120]}") + last, shown = stage, now + process.wait() + finally: + if process.poll() is None: # interrupted: do not leave it running + process.kill() + process.wait() + if process.returncode != 0 or result is None: + raise UiPSFError(f"uiPSF failed: {error or 'exit code ' + str(process.returncode)} " + f"(the whole output is in {log})") + return result diff --git a/src/smappy/uipsf/worker.py b/src/smappy/uipsf/worker.py new file mode 100644 index 0000000..d0db51d --- /dev/null +++ b/src/smappy/uipsf/worker.py @@ -0,0 +1,90 @@ +"""Run one uiPSF learning job: the half of `smappy.uipsf` that runs *in uiPSF's +environment*, as a script. + +It imports nothing of smappy's, on purpose: uiPSF wants its own TensorFlow and +usually its own Python, so smappy starts this file with that interpreter +(`runner.run`) and talks to it through a job directory and its stdout. The +directory holds + + job.json {"psftype": ..., "channeltype": ..., "params": {...}} + images.npy photons, already split into channels and cut as uiPSF wants + +and the learnt model is written there as uiPSF writes it. Every line meant for +smappy starts with ``SMAPPY_UIPSF`` and a tab -- ``stage``, ``result`` or +``error`` -- and everything else (TensorFlow's chatter, uiPSF's progress bars, +which go to stderr) is passed on as it comes. + +The images bypass uiPSF's own loader (`psflearninglib.load_data`): smappy reads +the camera, the z positions, the split and the mirror itself, the same way its +own calibration does, so this only repeats what `load_data` does *after* +reading -- the channel axis first and the multi-channel defocus list. +""" +import json +import os +import sys +import time +import traceback + +TAG = "SMAPPY_UIPSF" + + +def say(kind, text=""): + print(f"{TAG}\t{kind}\t{text}", flush=True) + + +def merge(target, overrides): + """Write ``overrides`` (nested dicts) into an OmegaConf node, key by key.""" + for key, value in overrides.items(): + if isinstance(value, dict) and key in target and hasattr(target[key], "keys"): + merge(target[key], value) + else: + target[key] = value + + +def main(directory): + os.environ.setdefault("MPLBACKEND", "Agg") # uiPSF imports pyplot + import numpy as np + from psflearning import io + from psflearning.psflearninglib import psflearninglib + + with open(os.path.join(directory, "job.json")) as f: + job = json.load(f) + say("stage", "starting uiPSF") + L = psflearninglib() + L.param = io.param.combine("config_base", psftype=job["psftype"], + channeltype=job["channeltype"]) + merge(L.param, job.get("params", {})) + L.param.datapath = directory + os.sep + L.param.savename = os.path.join(directory, "psfmodel") + L.param.ref_channel = 0 # smappy puts its main channel first + L.param.swapxy = False # numpy (y, x), as smappy hands it over + L.param.plotall = False + images = np.load(os.path.join(directory, "images.npy")).astype(np.float32) + if L.param.channeltype == "multi": + # what load_data sets for a multi-channel acquisition: no defocus + # between the channels unless the configuration says so + n = images.shape[0] + multi = L.param.option.multi + L.param.option.multi.defocus = [multi.defocus_offset + i*multi.defocus_delay + for i in range(n)] + start = time.time() + L.getpsfclass() + say("stage", "finding the emitters") + data = L.prep_data(images) + say("stage", "learning the PSF") + if "insitu" in L.param.PSFtype: + result = L.iterlearn_psf(data, time=0) + else: + psf, fitter = L.learn_psf(data, time=0) + result = L.save_result(psf, data, fitter) + say("result", os.path.abspath(result)) + say("stage", f"done in {time.time() - start:.0f} s") + + +if __name__ == "__main__": + try: + main(sys.argv[1]) + except Exception as error: # noqa: BLE001 -- reported, not hidden + traceback.print_exc() + say("error", f"{type(error).__name__}: {error}".replace("\n", " ")) + sys.exit(1) diff --git a/tests/test_uipsf.py b/tests/test_uipsf.py new file mode 100644 index 0000000..404c79b --- /dev/null +++ b/tests/test_uipsf.py @@ -0,0 +1,327 @@ +"""uiPSF as a calibration: its model comes back in smappy's conventions. + +Most of this needs no uiPSF: a uiPSF result is an HDF5 file whose layout and +conventions are written here by hand, so what is tested is the conversion +and the geometry -- that z keeps its sign and its zero, that a split camera's +channels go back on the chip, that uiPSF's transformation becomes smappy's. +The last test runs uiPSF itself, when ``SMAPPY_UIPSF_PYTHON`` names a Python +that has it (slow tier: minutes on a CPU). +""" +import json +import os +import sys +import textwrap + +import h5py +import numpy as np +import pytest +from scipy.special import erf + +from smappy.calibrate.dual import map_points +from smappy.io.calibration import load_spline_calibration +from smappy.simulate import ASTIGMATISM, astigmatic_sigmas +from smappy.uipsf import convert, microscopes, prepare, runner, uipsf_parameters + +PIXEL_NM, DZ_NM, SIZE = 100.0, 50.0, 21 + + +def spot(z_nm, x=0.0, y=0.0, size=SIZE, photons=1.0): + """The simulator's astigmatic spot at emitter height ``z_nm``, pixel-integrated.""" + sx, sy = (w / PIXEL_NM * np.sqrt(2.0) for w in astigmatic_sigmas(z_nm, 130.0, *ASTIGMATISM)) + pix, c = np.arange(size), (size - 1) / 2 + px = 0.5 * (erf((pix + 0.5 - c - x) / sx) - erf((pix - 0.5 - c - x) / sx)) + py = 0.5 * (erf((pix + 0.5 - c - y) / sy) - erf((pix - 0.5 - c - y) / sy)) + return photons * np.outer(py, px) + + +def params(psftype="zernike_vector", stage_pos=1.0, med=1.335, imm=1.516): + return {"PSFtype": psftype, "pixel_size": {"x": 0.1, "y": 0.1, "z": DZ_NM / 1000}, + "option": {"imaging": {"emission_wavelength": 0.68, "NA": 1.43, + "RI": {"imm": imm, "med": med, "cov": 1.516}}, + "insitu": {"stage_pos": stage_pos}, + "model": {"zernike_nl": []}}} + + +def write_result(path, res, p): + """A file laid out as uiPSF's `writeh5file` lays one out.""" + with h5py.File(path, "w") as f: + f.attrs["params"] = json.dumps(p) + g = f.create_group("res") + for key, value in res.items(): + if isinstance(value, dict): + sub = g.create_group(key) + for k, v in value.items(): + sub[k] = v + else: + g[key] = value + f.create_group("locres")["coeff"] = np.zeros(1) + return path + + +def bead_model(nz=33): + """uiPSF's bead model: planes are the objective's steps, focus in the middle. + + Raised by dz, the objective leaves a bead on the coverslip dz *below* its + focus, as `simulate.bead_stacks` draws it. + """ + objective = (np.arange(nz) - (nz - 1) / 2) * DZ_NM + return np.stack([spot(-z) for z in objective]) + + +def fitted_z(calibration, z_true): + from smappy.psf import SplinePSF + rois = np.stack([spot(z, 0.3, -0.2, size=13, photons=5000) + 10 for z in z_true]) + fit = SplinePSF(calibration).fit(rois.astype(np.float32), iterations=100) + return np.array([calibration.z_index_to_nm(t) for t in fit.theta[:, 4]]), fit.theta + + +def test_a_bead_model_keeps_its_sign_and_its_zero(tmp_path): + zernike = np.zeros((2, 45)) + zernike[1, 5] = 0.5 # Noll 6, astigmatism 0, in radians + path = write_result(tmp_path / "beads_zernike_vector_single.h5", + {"I_model": bead_model(), "zernike_coeff": zernike}, params()) + calibration = convert.single_calibration(path) + assert calibration.z0 == 16 and calibration.dz == DZ_NM + assert calibration.psf.sum(axis=(1, 2)).max() == pytest.approx(1.0) + z_true = np.array([-400.0, -150.0, 0.0, 200.0, 450.0]) + z, theta = fitted_z(calibration, z_true) + np.testing.assert_allclose(z, z_true, atol=8) + np.testing.assert_allclose(theta[:, :2] - theta[0, :2], 0, atol=0.02) + # the aberration, read out in nm of wavefront: 0.5 rad of 680 nm + row = next(r for r in calibration.parameters["zernike"] if r[0] == 6) + assert row[1] == "astigmatism 0" and row[3] == pytest.approx(0.5 * 680 / (2 * np.pi)) + + +def test_an_in_situ_model_is_turned_over_and_zeroed_at_the_focal_plane(tmp_path): + """In situ planes are emitter heights above the coverslip, the other way up. + + uiPSF's planes run ``zoffset + i`` steps above the coverslip, the focus + sitting at ``stagepos`` um of stage travel, ``stagepos * n_med / n_imm`` + into the sample -- that plane is z = 0. + """ + nz, stagepos, zoffset, med, imm = 25, 1.137, 13.4, 1.335, 1.516 + focus_um = stagepos * med / imm + heights_um = (zoffset + np.arange(nz)) * DZ_NM / 1000 + model = np.stack([spot((h - focus_um) * 1000) for h in heights_um]) + path = write_result(tmp_path / "psfmodel0_insitu_zernike_single.h5", + {"I_model": model, "stagepos": stagepos, "zoffset": zoffset}, + params("insitu_zernike", stagepos, med, imm)) + calibration = convert.single_calibration(path) + z_true = np.array([-300.0, -100.0, 0.0, 150.0, 300.0]) + z, _ = fitted_z(calibration, z_true) + np.testing.assert_allclose(z, z_true, atol=8) + assert calibration.parameters["uipsf_stagepos"] == [stagepos] + + +@pytest.mark.parametrize("layout,main", [("up-down", "upper"), ("up-down mirrored", "lower"), + ("right-left mirrored", "left"), + ("right-left", "right")]) +def test_a_split_frame_goes_back_on_the_chip(layout, main): + shape, origin = (64, 80), (300, 120) + split = prepare.make_split(shape, {"layout": layout, "main_channel": main, + "split_position": 34 if "up" in layout else 38}, + origin=origin) + frame = np.zeros(shape) + marks = {(5, 7): 1.0, (50, 70): 2.0, (40, 3): 3.0, (20, 60): 4.0} + for (y, x), value in marks.items(): + frame[y, x] = value + channels = split.cut(frame[None])[:, 0] + assert channels.shape[0] == 2 and channels[0].shape == channels[1].shape + for channel in (0, 1): + for y, x in zip(*np.nonzero(channels[channel])): + chip_x, chip_y = split.to_chip(channel, [[y, x]])[0] + assert frame[int(chip_y) - origin[1], int(chip_x) - origin[0]] == \ + channels[channel][y, x] + + +def test_uipsfs_channel_transformation_becomes_smappys(): + """uiPSF's ``T``, made from a known chip map, gives that map back. + + uiPSF maps the main channel's (y, x, 1) rows about ``imgcenter`` into + the secondary channel's array, which a mirrored splitter has flipped. + """ + split = prepare.make_split((128, 100), {"layout": "up-down mirrored", + "main_channel": "upper"}, origin=(16, 8)) + truth = np.array([[1.003, 0.004, 1.7], [-0.006, 0.998, 130.2], [0.0, 0.0, 1.0]]) + truth = np.linalg.inv(truth) # secondary -> main + rng = np.random.default_rng(3) + main_yx = rng.uniform(5, 58, (40, 2)) + secondary_chip = map_points(np.linalg.inv(truth), split.to_chip(0, main_yx)) + # back into the secondary channel's array: unshift, unflip, (y, x) + local = secondary_chip - np.asarray(split.origin, float) + secondary_yx = np.c_[split.length - 1 - (local[:, 1] - split.starts()[1]), local[:, 0]] + centre = np.array([split.length / 2, 50.0, 0.0]) + rows = np.c_[main_yx, np.ones(len(main_yx))] - centre + target = np.c_[secondary_yx - centre[:2], np.ones(len(main_yx))] + T, *_ = np.linalg.lstsq(rows, target, rcond=None) + found = convert.channel_transformation(T, centre, split) + points = np.array([[20.0, 70.0], [90.0, 120.0], [50.0, 100.0]]) + np.testing.assert_allclose(map_points(found, points), map_points(truth, points), atol=1e-6) + + +def test_the_channel_shift_is_found_on_a_bead_grid_that_fools_nearest_neighbours(): + """uiPSF's own guess, from bead coordinates, took the neighbouring bead of + a 30 px grid for the partner and kept one pair of nine.""" + from scipy import ndimage + image = np.zeros((2, 3, 9, 100, 100)) + for gy in (20, 50, 80): + for gx in (20, 50, 80): + image[0, :, :, gy, gx] = 1000.0 + image[1] = ndimage.shift(image[0], (0, 0, -1.0, 2.0), order=1) + image += np.random.default_rng(1).poisson(5.0, image.shape) + assert prepare.channel_shift(image) == [[0.0, 0.0], [-1.0, 2.0]] + + +def test_two_channels_share_one_normalisation_and_the_secondary_is_unmirrored(tmp_path): + split = prepare.make_split((128, 100), {"layout": "up-down mirrored", + "main_channel": "upper"}) + model = bead_model() + lopsided = model * (1 + 0.5 * np.linspace(0, 1, SIZE))[None, :, None] # not y-symmetric + T = np.eye(3) + T[2, :2] = [0.5, -0.3] + res = {"channel0": {"I_model": model, "intensity": np.full(9, 3000.0)}, + "channel1": {"I_model": lopsided, "intensity": np.full(9, 1000.0)}, + "T": T, "imgcenter": np.array([32.0, 50.0, 0.0]), "xyshift": np.zeros(2)} + path = write_result(tmp_path / "psfmodel_zernike_vector_multi.h5", res, params()) + calibration = convert.dual_calibration(path, split) + main, secondary = calibration.main, calibration.secondary + window = slice(14, 19) + light = [m.psf[window].sum(axis=(1, 2)).mean() for m in (main, secondary)] + assert sum(light) == pytest.approx(1.0, rel=1e-3) + assert main.parameters["photon_normalization"] == pytest.approx(light[0], rel=1e-3) + assert light[0] / light[1] == pytest.approx(3000 / (1000 * lopsided[window].sum() + / model[window].sum()), rel=1e-3) + # flipped back along y, the mirror axis, to the camera's orientation + np.testing.assert_allclose(secondary.psf / secondary.psf.max(), + (lopsided[:, ::-1, :] / lopsided.max()).astype(np.float32), + atol=1e-6) + path_out = convert.save(tmp_path / "dual.h5", calibration) + from smappy.calibrate.dual import load_dual_color_calibration + again = load_dual_color_calibration(path_out) + np.testing.assert_allclose(again.transformation, calibration.transformation) + + +def test_a_profile_says_what_it_knows_and_refuses_what_it_does_not(tmp_path, monkeypatch): + shipped = microscopes.load("example") + assert shipped.dual["layout"] == "up-down mirrored" and shipped.NA == 1.43 + (tmp_path / "example.yaml").write_text("name: ours\nNA: 1.5\n") + (tmp_path / "typo.yaml").write_text("NA: 1.4\nrefractve_index: {}\n") + monkeypatch.setenv("SMAPPY_MICROSCOPES", str(tmp_path)) + assert microscopes.load("example").name == "ours" # the group's wins + with pytest.raises(ValueError, match="unknown keys"): + microscopes.load("typo") + assert ("typo", "typo (unreadable)") in microscopes.choices() + + +def test_the_profile_becomes_uipsfs_parameters_and_its_own_section_has_the_last_word(): + profile = microscopes.Microscope(NA=1.35, emission_wavelength_nm=600, + uipsf={"option": {"model": {"n_max": 6}}, + "batch_size": 400}) + p = uipsf_parameters(profile, data="blinking", pixelsize_xy_um=(0.11, 0.12), dz_nm=40, + depth_um=2.5, zernike_order=8) + assert p["pixel_size"] == {"x": 0.11, "y": 0.12, "z": 0.04} + assert p["option"]["imaging"]["NA"] == 1.35 + assert p["option"]["imaging"]["emission_wavelength"] == pytest.approx(0.6) + assert p["option"]["insitu"]["stage_pos"] == 2.5 + assert p["option"]["model"]["n_max"] == 6 and p["batch_size"] == 400 + assert "insitu" not in uipsf_parameters(profile, data="beads", + pixelsize_xy_um=(0.1, 0.1), dz_nm=50)["option"] + + +def test_the_runner_follows_the_worker_and_passes_its_failure_on(tmp_path, monkeypatch): + """The protocol, with a stand-in for the worker: stages, bars, result, error.""" + fake = tmp_path / "worker.py" + fake.write_text(textwrap.dedent(f""" + import sys + print("{runner.TAG}\\tstage\\tlearning the PSF", flush=True) + sys.stderr.write("3/6: learning: 1/150 [00:01s]\\r3/6: learning: 2/150 [00:02s]\\r") + sys.stderr.flush() + if "fail" in sys.argv[1]: + print("{runner.TAG}\\terror\\tValueError: no bead is found", flush=True) + sys.exit(1) + print("{runner.TAG}\\tresult\\t/somewhere/psfmodel.h5", flush=True) + """)) + monkeypatch.setattr(runner, "WORKER", fake) + said = [] + (tmp_path / "ok").mkdir() + assert str(runner.run(tmp_path / "ok", python=sys.executable, + progress=said.append)) == "/somewhere/psfmodel.h5" + assert said[0] == "uiPSF: learning the PSF" + assert any("3/6: learning" in s for s in said) + (tmp_path / "fail").mkdir() + with pytest.raises(runner.UiPSFError, match="no bead is found"): + runner.run(tmp_path / "fail", python=sys.executable) + assert "3/6: learning: 2/150" in (tmp_path / "fail" / "uipsf.log").read_text() + + +def test_without_uipsf_it_says_where_to_put_it(monkeypatch): + monkeypatch.delenv("SMAPPY_UIPSF_PYTHON", raising=False) + monkeypatch.setattr("smappy.config.get", lambda key, default=None: default) + import importlib.util + monkeypatch.setattr(importlib.util, "find_spec", lambda name: None) + with pytest.raises(runner.UiPSFError, match="uiPSF python"): + runner.find_python("") + with pytest.raises(runner.UiPSFError, match="no Python at"): + runner.find_python("/no/such/python") + + +@pytest.mark.slow +@pytest.mark.timeout(1200) # uiPSF on a CPU: two to five minutes alone +@pytest.mark.skipif(not os.environ.get("SMAPPY_UIPSF_PYTHON"), + reason="SMAPPY_UIPSF_PYTHON names no Python with uiPSF") +def test_uipsf_learns_simulated_beads_that_spline_3d_then_fits_in_z(tmp_path, monkeypatch): + """The whole path with uiPSF itself: beads -> uiPSF -> Spline 3D. + + The beads are the simulator's astigmatic Gaussians, which no pupil + produces exactly, so the slope is allowed a few per cent where smappy's + own calibration of the same beads gets 2 (test_camera_frames). + """ + tifffile = pytest.importorskip("tifffile") + from scipy.spatial import cKDTree + from smappy.plugins import Context + from smappy.plugins.fit import (CameraSettings, OutputSettings, SourceSettings, + SplineFit, SplineFitSettings, SplineModelSettings) + from smappy.plugins.uipsf import (BeadPart, LearningPart, UiPSFCalibration, + UiPSFSettings) + from smappy.session import Session + from smappy.simulate import ISOLATED_NM, LabellingSettings, bead_stacks, camera_frames + stacks, z = bead_stacks(2, seed=0, dz_nm=50, z_range_nm=(-1000, 1000)) + for i, stack in enumerate(stacks): + tifffile.imwrite(tmp_path / f"beads{i}.tif", stack) + camera = CameraSettings(conversion=0.5, offset=100.0, pixelsize_um=0.1) + (tmp_path / "scope").mkdir() + (tmp_path / "scope" / "sim.yaml").write_text( + "NA: 1.45\nemission_wavelength_nm: 670\nroi_size_px: 21\n") + monkeypatch.setenv("SMAPPY_MICROSCOPES", str(tmp_path / "scope")) + # a test worker's one BLAS thread (conftest.py) would be uiPSF's too + for name in ("OMP_NUM_THREADS", "OPENBLAS_NUM_THREADS", "MKL_NUM_THREADS"): + monkeypatch.delenv(name, raising=False) + result = UiPSFCalibration().run(Context(), UiPSFSettings( + path=str(tmp_path / "beads0.tif"), microscope="sim", channels="single", + beads=BeadPart(z_step_nm=50, max_beads=18), camera=camera, + learning=LearningPart(iterations=60))) + calibration = result.data["calibration"] + assert load_spline_calibration(calibration).psf.shape[1] == 21 + frames, truth = camera_frames(300, seed=3, astigmatism=ASTIGMATISM, + labelling=LabellingSettings(efficiency=0.1)) + tifffile.imwrite(tmp_path / "frames.tif", frames) + SplineFit().run(Context(), SplineFitSettings( + source=SourceSettings(path=str(tmp_path / "frames.tif")), camera=camera, + model=SplineModelSettings(calibration=calibration), + output=OutputSettings(path=str(tmp_path / "out.hdf5")))) + session = Session() + session.load(tmp_path / "out.hdf5") + locs = session.locs + clean = (truth["photons"] > 500) & (truth["neighbour_nm"] > ISOLATED_NM) + fitted, true = [], [] + for f in np.unique(locs["frame"]): + m, t = locs["frame"] == f, (truth["frame"] == f) & clean + if not t.any(): + continue + d, i = cKDTree(np.c_[truth["x_nm"][t], truth["y_nm"][t]]).query( + np.c_[locs["x_nm"][m], locs["y_nm"][m]]) + fitted += list(locs["z_nm"][m][d < 60]) + true += list(truth["z_nm"][t][i[d < 60]]) + slope, intercept = np.polyfit(true, fitted, 1) + assert slope == pytest.approx(1.0, abs=0.06) + assert abs(intercept) < 15 diff --git a/tests/test_workspace.py b/tests/test_workspace.py index 0d2a253..7f6ff17 100644 --- a/tests/test_workspace.py +++ b/tests/test_workspace.py @@ -168,7 +168,8 @@ def test_the_window_opens_the_shipped_tabs(window): ["File", "Localize", "Render", "Analysis", "ROI"] localize = tab_named(window, "Localize") assert [s.title for s in localize.sections] == \ - ["Gaussian 2D", "Gaussian 2D 2C", "Spline 3D", "Spline 3D 2C"] + ["Gaussian 2D", "Gaussian 2D 2C", "Spline 3D", "Spline 3D 2C", + "uiPSF calibration"] def test_opening_a_section_is_what_imports_the_plugin(window): From 3022f8a8ace17038148c0e7ba9030f2d309d942f Mon Sep 17 00:00:00 2001 From: Claude Date: Tue, 6 Oct 2026 10:34:41 +0000 Subject: [PATCH 2/7] uiPSF calibration: a section on installing uiPSF; CLAUDE.md brought along - the plugin's page: "Installing uiPSF", after "What it does" -- smappy does not install it; a conda environment with uiPSF's branch, a check, and the three ways to tell smappy where it is (config.yaml per platform, the plugin's setting, SMAPPY_UIPSF_PYTHON), and what a graphics card needs. Checked in a fresh Python 3.12 venv: pip install -e of the branch, then the real-uiPSF test passes. The page's code blocks fenced, which is what the page renderer turns into code (indented ones became paragraphs) - CLAUDE.md: the plugin-shape row for a plugin that runs an external program in its own Python; the Localize tab's titles listed in test_workspace; SMAPPY_UIPSF_PYTHON enabling the uiPSF test Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01GiNdULxyiqQEqWHGKTbVzV --- CLAUDE.md | 8 +- .../docs/Localize/uiPSF calibration.md | 91 +++++++++++++++---- 2 files changed, 78 insertions(+), 21 deletions(-) diff --git a/CLAUDE.md b/CLAUDE.md index 76d2ddc..4f92ede 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -72,6 +72,9 @@ Running the tests -- two tiers: `test_roi_session`), which is the environment and not your change. * `SMAPPY_TEST_CAL=/path/to/_3dcal.mat` enables the tests that want real calibration data; without it they skip. +* `SMAPPY_UIPSF_PYTHON=/path/to/python` (of an environment with uiPSF, as the + page of `Localize/uiPSF calibration` installs it) enables the slow test that + runs uiPSF itself, a few minutes on a CPU; without it it skips. ## Writing a plugin @@ -96,6 +99,7 @@ shape, so read the one closest to what you are writing: | finds a structure by its shape and judges it, then counts it | `npc.py` (a ring filter and circle fit in the segmenter, a per-site count, and a model of whole pores in `smappy/npc.py` fitted to the site table) | | summarises what evaluators measured | `roi.py`'s `SiteAnalysisPlugin` (an evaluator to read, a *results from* field, a preflight that asks which when there are several) | | runs the ROI manager's segment, evaluate and analyse as one | `npc_analysis.chain.yaml` (a shipped chain: an evaluator step runs over every ROI before the next; `docs/batch.md`, "ROIs: segment, evaluate, analyse") | +| runs an external program in its own Python | `uipsf.py` (`smappy.uipsf`: a worker script under that program's interpreter, progress read from its output, the result converted into smappy's own file) | A chain of plugins that runs as one, and running one over many files (`smappy-batch`, the batch window): `docs/batch.md`. @@ -355,7 +359,9 @@ a new plugin there fails four tests until it is listed in both (the workspace ones name the seeded Analysis tab, including the order its titles come back in). `tests/test_file_plugins.py::test_the_file_tab_ships_with_the_four_it_needs` asserts the File tab's exact contents, so a new File plugin needs a line there -too (it is a Qt test and skips without PySide6). +too (it is a Qt test and skips without PySide6). Likewise +`tests/test_workspace.py::test_the_window_opens_the_shipped_tabs` lists the +Localize tab's titles, so a new Localize plugin needs a line there. `tests/test_control_width.py` builds every shipped plugin's panel and fails if it is wider than the control column (`CONTROL_WIDTH`): a new plugin with a long diff --git a/src/smappy/plugins/docs/Localize/uiPSF calibration.md b/src/smappy/plugins/docs/Localize/uiPSF calibration.md index 8586d91..3c2a99f 100644 --- a/src/smappy/plugins/docs/Localize/uiPSF calibration.md +++ b/src/smappy/plugins/docs/Localize/uiPSF calibration.md @@ -34,7 +34,7 @@ calibration. For a split camera it holds both channels' models and the transformation between the halves, which uiPSF learns together with the PSF. **What it needs.** uiPSF itself, installed in a Python environment of its own -(it brings TensorFlow; see *Installing uiPSF* below), and a **microscope +(see *Installing uiPSF* below), and a **microscope profile**: the objective's numerical aperture, the refractive indices, the emission wavelength and, for a split camera, how the halves are arranged. From beads: one or more z-stacks, as for the bead calibration. From blinking @@ -57,6 +57,66 @@ the bead calibration is the safer choice. Field-dependent and 4Pi models, which uiPSF also learns, are not offered: the fitters here use one PSF for the whole field. +## Installing uiPSF + +smappy does **not** install uiPSF, and installing smappy does not bring it. +uiPSF needs TensorFlow, which is large and wants its own versions of numpy +and Python, so it lives in a Python environment of its own. smappy starts +uiPSF in that environment when the plugin runs, and needs to be told once +where it is. Everything below is done once per computer. + +**1. An environment with uiPSF.** With conda (Miniconda or Anaconda), in a +terminal (on Windows, the Anaconda Prompt): + +``` +conda create -n uipsf python=3.12 +conda activate uipsf +git clone --branch claude/friendly-cerf-qt2sfj https://github.com/ries-lab/uiPSF +pip install -e uiPSF +``` + +The last line installs uiPSF with TensorFlow and everything else it needs. +It has to be this branch of uiPSF: its main branch needs Python 3.7 and +TensorFlow 2.9, which do not run on current Macs or with current numpy. +Without conda, `python3.12 -m venv uipsf-env` and that environment's `pip` do +the same. + +**2. Check it.** Still in that environment: + +``` +python -c "import psflearning.psflearninglib; print('uiPSF works')" +python -c "import sys; print(sys.executable)" +``` + +The second line prints the environment's Python, for example +`/Users/me/miniconda3/envs/uipsf/bin/python` on a Mac, +`/home/me/miniconda3/envs/uipsf/bin/python` on Linux, or +`C:\Users\me\miniconda3\envs\uipsf\python.exe` on Windows. + +**3. Tell smappy.** Once for all plugins and sessions, add a line with that +path to smappy's configuration file, `config.yaml` in +`~/Library/Application Support/smappy/` (Mac), `~/.config/smappy/` (Linux) or +`%APPDATA%\smappy\` (Windows), creating the file if it is not there: + +``` +uipsf_python: /Users/me/miniconda3/envs/uipsf/bin/python +``` + +Or give the path in the plugin itself, as *uiPSF python* (under *learning*, +"more"), or in the environment variable `SMAPPY_UIPSF_PYTHON`. The plugin's +setting wins over the file, and the file over the variable. + +If uiPSF can be imported in the Python smappy itself runs in, nothing needs to +be set, but a separate environment is what keeps TensorFlow from fighting +smappy's own packages. + +**A graphics card.** uiPSF runs on the CPU by default, which works +everywhere and takes minutes for beads (see *Running time* below). On Linux +with an NVIDIA card, `pip install "tensorflow[and-cuda]"` in the uiPSF +environment lets it use the card. On Windows TensorFlow uses no graphics card +after version 2.10, except under WSL2. On a Mac uiPSF runs on the CPU, and +Apple-silicon Macs run it natively. + ## How it works **1. Reading the data.** smappy reads the frames, not uiPSF: the bead @@ -170,28 +230,19 @@ TensorFlow's Apple-GPU plugin has not been tried with it, and the pupil model is built on complex numbers and Fourier transforms, which that plugin has supported only in part. -**Installing uiPSF.** In a Python environment of its own (tried with Python -3.12 and TensorFlow 2.21): - - pip install tensorflow tensorflow-probability[tf] h5py omegaconf dotted_dict \ - czifile scikit-image matplotlib tqdm - pip install --no-deps -e /path/to/uiPSF - -The plugin needs uiPSF's branch that runs on current numpy and TensorFlow. -Then give that environment's python as *uiPSF python*, or once and for all as -`uipsf_python` in smappy's configuration file, or in `SMAPPY_UIPSF_PYTHON`. - **Microscope profiles.** A profile is a YAML file, `example.yaml` ships as a template: - name: M2, astigmatic, two colour - NA: 1.43 - refractive_index: {immersion: 1.516, medium: 1.335, coverslip: 1.516} - emission_wavelength_nm: 680 - roi_size_px: 25 - dual: {layout: up-down mirrored, main_channel: upper} - insitu: {zernike_index: [5], zernike_coeff: [0.5]} - uipsf: {} +``` +name: M2, astigmatic, two colour +NA: 1.43 +refractive_index: {immersion: 1.516, medium: 1.335, coverslip: 1.516} +emission_wavelength_nm: 680 +roi_size_px: 25 +dual: {layout: up-down mirrored, main_channel: upper} +insitu: {zernike_index: [5], zernike_coeff: [0.5]} +uipsf: {} +``` Profiles are found in the folders named by `SMAPPY_MICROSCOPES` (a group's shared folder, say), then in `microscopes/` in smappy's configuration folder, From fb60439597670a7cdde8faa0717018aa16073864 Mon Sep 17 00:00:00 2001 From: Claude Date: Tue, 6 Oct 2026 16:08:23 +0000 Subject: [PATCH 3/7] uiPSF calibration: uiPSF's own diagnostics as the result's tabs uiPSF draws its checks in its notebooks (psflearning.makeplots), which make their own figures and call plt.show(); smappy.uipsf.plots redraws them from the saved result into the figure a result window hands over, one tab each: - data vs model (showpsfvsdata / _insitu): a typical bead, or molecules averaged per model plane by fitted height, above the model, with x-z - localization bias (showlocalization): every bead refitted plane by plane, x, y, z bias in nm against stage z, the median in red; beads only - pupil (showzernike): magnitude, aberration phase from Noll 5 up in nm, and every Zernike term - beads / emitters (showcoord): found and used - channel transformation (showtransform): the pairs' residuals after T, in nm, as arrows over the field and scattered; two channels only They replace the plugin's own model and Zernike figures. The text adds the beads' median bias and the transformation's rms residual. The page's Output section describes each tab and what a good one looks like. --slow -n auto: 1346 passed, 4 skipped (with SMAPPY_UIPSF_PYTHON set) Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01GiNdULxyiqQEqWHGKTbVzV --- .../docs/Localize/uiPSF calibration.md | 51 ++- src/smappy/plugins/uipsf.py | 64 ++-- src/smappy/uipsf/plots.py | 311 ++++++++++++++++++ tests/test_uipsf.py | 114 ++++++- 4 files changed, 479 insertions(+), 61 deletions(-) create mode 100644 src/smappy/uipsf/plots.py diff --git a/src/smappy/plugins/docs/Localize/uiPSF calibration.md b/src/smappy/plugins/docs/Localize/uiPSF calibration.md index 3c2a99f..e577f96 100644 --- a/src/smappy/plugins/docs/Localize/uiPSF calibration.md +++ b/src/smappy/plugins/docs/Localize/uiPSF calibration.md @@ -1,6 +1,6 @@ --- version: "1" -covers: [smappy.uipsf.prepare.Split, smappy.uipsf.prepare.channel_shift, smappy.uipsf.prepare.subtract_background, smappy.uipsf.prepare.read_beads, smappy.uipsf.prepare.read_movie, smappy.uipsf.uipsf_parameters, smappy.uipsf.runner.run, smappy.uipsf.convert.single_calibration, smappy.uipsf.convert.dual_calibration, smappy.uipsf.convert.channel_transformation, smappy.uipsf.microscopes.Microscope] +covers: [smappy.uipsf.prepare.Split, smappy.uipsf.prepare.channel_shift, smappy.uipsf.prepare.subtract_background, smappy.uipsf.prepare.read_beads, smappy.uipsf.prepare.read_movie, smappy.uipsf.uipsf_parameters, smappy.uipsf.runner.run, smappy.uipsf.convert.single_calibration, smappy.uipsf.convert.dual_calibration, smappy.uipsf.convert.channel_transformation, smappy.uipsf.microscopes.Microscope, smappy.uipsf.plots.figures, smappy.uipsf.plots.data_and_model, smappy.uipsf.plots.localization_bias, smappy.uipsf.plots.pupil_images, smappy.uipsf.plots.transformation_residuals] --- ## What it does @@ -312,13 +312,48 @@ beside the data) is what Spline 3D or Spline 3D 2C take as their pupil, the per-emitter fits and the model, which SMAP and uiPSF's own notebooks read. -**The figures** show the model through focus, x–z and y–z, and its focal -plane; for two channels both, the secondary in the camera's orientation. An -astigmatic PSF is wide in x on one side of focus and in y on the other. The -second tab shows the learnt aberrations in nanometres of wavefront per -Zernike term: astigmatism (5, 6) for an astigmatic lens, spherical (11) -growing with depth in water. Large coma (7, 8) points at a tilted or -misaligned optic. +**The text** names both files, the model's z range, its largest +aberrations in nanometres of wavefront, and two numbers from uiPSF's own +check of the result. The first is the median bias of the beads, refitted +with the model. The second, for two channels, is how far the +transformation leaves each bead pair apart. The figures are uiPSF's own +diagnostics, the ones its notebooks show, one tab each: + +* **data vs model**: the measured spot above and uiPSF's model of it below, + at planes through the model and in an x–z section, per channel. For beads + it is one bead, the one the model fits as well as it fits most, and the + planes are labelled by stage position. From blinking molecules the + molecules are averaged per plane of the model by their fitted height, and + the planes are labelled by height above the coverslip. Planes no molecule + fell into stay black. Data and model should look alike in every plane; a + difference that grows away from focus means the model misses an + aberration. +* **localization bias** (beads): uiPSF refits every bead in every plane with + the learnt model. This is x, y and z of each fit minus where the bead + was, against the stage position, one grey line per bead and their median + in red. A few nanometres in x and y, and a z bias of a few tens of + nanometres at most over the range you will fit, are good. A slope in z means the model's z scale + is off. A bias that grows towards the ends of the stack marks where the + model stops being trustworthy. +* **pupil**: per channel, the pupil's magnitude, the aberration it adds to + the wavefront (from Zernike term 5 up, in nm), and every Zernike term: + phase as bars in nm rms, magnitude as dots. The box names the terms + uiPSF names. Astigmatism (5, 6) is expected with an astigmatic lens, and + spherical (11) grows with depth in water. Large coma (7, 8) points at a + tilted or misaligned optic. +* **beads** or **emitters**: where uiPSF found emitters (grey) and which it + used for the model (red circles), in the channel as uiPSF saw it. Beads + too close to another or to the edge are dropped before they count as + found. Of those found, uiPSF leaves out the ones beyond *max beads* and + those it rejects as outliers after a first round of learning. From + blinking molecules it also leaves out the dim and the poorly fitted, and + more than each z slice takes (*per bin*). +* **channel transformation** (two channels): for every bead pair, where the + secondary channel's bead lies minus where the transformation puts it, in + nanometres. On the left the residuals are drawn as arrows over the field, + with the longest given in the title; on the right they are scattered. + Residuals of a few nanometres are good. Arrows that all point one way in + part of the field mean the transformation cannot follow the splitter there. **A bad result** shows as a model with structure that changes erratically from plane to plane, or as fitted z that does not follow the stage when the diff --git a/src/smappy/plugins/uipsf.py b/src/smappy/plugins/uipsf.py index a49ccd8..9b514c5 100644 --- a/src/smappy/plugins/uipsf.py +++ b/src/smappy/plugins/uipsf.py @@ -225,11 +225,12 @@ def run(self, ctx: Context, settings: UiPSFSettings) -> Result: shutil.copyfile(result, raw) if not lp.keep_job: shutil.rmtree(job, ignore_errors=True) - text = summary(calibration, out, raw) - return Result(text=text, settings=settings, - plot=Plot(draw=lambda fig: draw_models(fig, calibration), - panels=6 if dual else 3, size=(10, 6 if dual else 3.5)), - plots={"Zernike": lambda ax: draw_zernike(ax, calibration)}, + from ..uipsf import plots + learnt = convert.read(raw) + figures = [Plot(draw=draw, name=name, panels=panels, size=size) + for name, (draw, panels, size) in plots.figures(learnt).items()] + return Result(text=summary(calibration, out, raw, learnt), settings=settings, + plot=figures[0], plots={f.name: f for f in figures[1:]}, data={"calibration": str(out), "uipsf_result": str(raw), "dual": dual}) @@ -260,8 +261,9 @@ def _models(calibration): return [("", calibration)] -def summary(calibration, out: Path, raw: Path) -> str: - """What the run made: the file, the model's extent and the main aberrations.""" +def summary(calibration, out: Path, raw: Path, learnt=None) -> str: + """What the run made: the file, the model's extent, the main aberrations, + and -- from uiPSF's result -- how well the beads and the channels fit.""" lines = [f"saved {out}", f"uiPSF's own result: {raw}"] for name, model in _models(calibration): nz = model.psf.shape[0] @@ -277,40 +279,16 @@ def summary(calibration, out: Path, raw: Path) -> str: lines.append("channel transformation (secondary -> main, chip px): " + np.array2string(calibration.transformation / calibration.transformation[2, 2], precision=4, suppress_small=True).replace("\n", "")) + if learnt is not None: + from ..uipsf import plots + if plots.has_localization_bias(learnt): + bias = [np.nanmedian(np.abs(b)) for b in plots.localization_bias(learnt)] + lines.append("beads refitted with the model, median |bias|: x {:.1f}, " + "y {:.1f}, z {:.1f} nm".format(*bias)) + if len(plots.channels(learnt)) > 1 and "T" in learnt["res"]: + _, residual = plots.transformation_residuals(learnt) + p = learnt["params"]["pixel_size"] + rms = np.sqrt(np.mean(residual ** 2, axis=0)) * 1000 * np.array( + [float(p["y"]), float(p["x"])]) + lines.append(f"transformation residual, rms: x {rms[1]:.1f}, y {rms[0]:.1f} nm") return "\n".join(lines) - - -def draw_models(figure, calibration) -> None: - """Each channel's model through focus: x-z, y-z, and the focal plane.""" - rows = _models(calibration) - axes = figure.subplots(len(rows), 3, squeeze=False) - for (name, model), row in zip(rows, axes): - psf = model.psf - nz, ny, nx = psf.shape - extent_z = (model.z_index_to_nm(nz - 0.5), model.z_index_to_nm(-0.5)) - focus = int(round(model.z0)) - row[0].imshow(psf[::-1, ny // 2, :], aspect="auto", cmap="magma", - extent=(-0.5, nx - 0.5) + extent_z) - row[0].set(xlabel="x (px)", ylabel="z (nm)", title=f"{name} x-z".strip()) - row[1].imshow(psf[::-1, :, nx // 2], aspect="auto", cmap="magma", - extent=(-0.5, ny - 0.5) + extent_z) - row[1].set(xlabel="y (px)", title=f"{name} y-z".strip()) - row[2].imshow(psf[min(max(focus, 0), nz - 1)], cmap="magma") - row[2].set(xlabel="x (px)", ylabel="y (px)", title=f"{name} z = 0".strip()) - - -def draw_zernike(ax, calibration) -> None: - """The pupil's aberrations, Noll index 5 up, in nm rms of wavefront.""" - width = 0.8 / len(_models(calibration)) - for k, (name, model) in enumerate(_models(calibration)): - rows = [row for row in model.parameters.get("zernike", []) if row[0] > 4] - if not rows: - continue - j = np.array([row[0] for row in rows]) - ax.bar(j + (k - (len(_models(calibration)) - 1) / 2) * width, - [row[3] for row in rows], width=width, label=name or None) - ax.axhline(0, color="0.5", lw=0.8) - ax.set(xlabel="Noll index", ylabel="phase (nm rms)", - title="aberrations (5, 6: astigmatism; 7, 8: coma; 11: spherical)") - if len(_models(calibration)) > 1: - ax.legend() diff --git a/src/smappy/uipsf/plots.py b/src/smappy/uipsf/plots.py new file mode 100644 index 0000000..458ca8c --- /dev/null +++ b/src/smappy/uipsf/plots.py @@ -0,0 +1,311 @@ +"""The figures uiPSF's notebooks show, drawn from its result file. + +uiPSF draws nothing useful while it learns; its demo notebooks call +``psflearning.makeplots`` afterwards on the saved result. Those functions +make their own figures and call ``plt.show()``, which a result window cannot +use -- a smappy plot is *handed* a figure, or a subfigure on the window's All +page -- so the ones worth keeping are redrawn here from the same arrays (the +dict `convert.read` gives), each into the figure it is handed: + +* `draw_data_vs_model`: ``showpsfvsdata`` / ``showpsfvsdata_insitu`` +* `draw_localization_bias`: ``showlocalization`` +* `draw_pupil`: ``showzernike`` (``showpupil`` for a model without Zernike + coefficients) +* `draw_emitters`: ``showcoord`` +* `draw_transformation`: ``showtransform``, as residuals rather than overlays + +Left out: ``showlearnedparam`` (per-emitter position, photons, background, +drift: diagnostics of the learning rather than of the model), ``showpsf`` +(the model alone: it is in `draw_data_vs_model`) and the field-dependent +maps, for models the plugin does not offer. + +All of it is in uiPSF's coordinates: a channel's array as uiPSF was handed it +(the secondary flipped back if the splitter mirrors it), y down, x across. +""" +from __future__ import annotations + +import numpy as np + +# how many planes the data-versus-model figure shows, plus the x-z section +PLANES = 6 +# the terms uiPSF's showzernike names, by Noll index +NAMED = {5: "astigmatism 45", 6: "astigmatism 0", 7: "coma y", 8: "coma x", + 11: "spherical", 22: "2nd spherical"} + + +def channels(data): + """``[(label, res)]``: one entry, or one per channel of a multi-channel result.""" + res = data["res"] + if "channel0" in res: + names = ("main", "secondary") + return [(names[i] if i < 2 else f"channel {i}", res[f"channel{i}"]) + for i in range(sum(k.startswith("channel") for k in res))] + return [("", res)] + + +def _insitu(data) -> bool: + return "insitu" in str(data["params"].get("PSFtype", "")) + + +def _rois(data, k): + """Channel ``k``'s measured and modelled ROIs.""" + rois = data["rois"] + multi = len(channels(data)) > 1 + return ((rois["psf_data"][k], rois["psf_fit"][k]) if multi + else (rois["psf_data"], rois["psf_fit"])) + + +def _binned(data, k, res): + """In situ: the molecules averaged per z plane of the model, as uiPSF does. + + Each molecule's fitted z picks the plane of the model it belongs to; the + planes no molecule fell into stay empty. + """ + measured, _ = _rois(data, k) + model = np.asarray(res["I_model"], float) + z = np.asarray(channels(data)[0][1]["pos"])[:, 0] + zoffset = float(np.real(np.ravel(channels(data)[0][1].get("zoffset", [0]))[0])) + index = np.digitize(z, zoffset + np.arange(len(model) + 1)) - 1 + average = np.zeros_like(model) + for plane in range(len(model)): + here = index == plane + if here.any(): + average[plane] = measured[here].mean(axis=0) + return average, model + + +def typical_bead(measured, modelled) -> int: + """The bead whose model fits it as well as most: the median relative error.""" + error = (np.sum((measured - modelled) ** 2, axis=tuple(range(1, measured.ndim))) + / np.sum(measured ** 2, axis=tuple(range(1, measured.ndim)))) + return int(np.argsort(error)[len(error) // 2]) + + +def data_and_model(data): + """``[(label, data stack, model stack, what)]`` per channel.""" + out = [] + for k, (label, res) in enumerate(channels(data)): + if _insitu(data): + measured, modelled = _binned(data, k, res) + what = "molecules averaged per plane" + else: + stacks, fits = _rois(data, k) + bead = typical_bead(stacks, fits) + measured, modelled = stacks[bead], fits[bead] + what = f"bead {bead}" + out.append((label, np.asarray(measured, float), np.asarray(modelled, float), what)) + return out + + +def draw_data_vs_model(figure, data) -> None: + """Measured above, uiPSF's model below, at planes through focus and in x-z.""" + rows = data_and_model(data) + dz_um = float(data["params"]["pixel_size"]["z"]) + axes = figure.subplots(2 * len(rows), PLANES + 1, squeeze=False) + insitu = _insitu(data) + for r, (label, measured, modelled, what) in enumerate(rows): + nz, ny, nx = measured.shape + planes = np.linspace(0, nz - 1, PLANES + 2)[1:-1].round().astype(int) + if insitu: # uiPSF's own axis: the height above the coverslip + zoffset = float(np.real(np.ravel(channels(data)[0][1].get("zoffset", [0]))[0])) + z_um, unit = (zoffset + np.arange(nz)) * dz_um, "um high" + else: # the stack's: where the objective was, focus at 0 + z_um, unit = (np.arange(nz) - (nz - 1) / 2) * dz_um, "um stage" + for row, stack, name in ((2 * r, measured, "data"), (2 * r + 1, modelled, "model")): + top = max(float(stack.max()), 1e-12) + for c, plane in enumerate(planes): + ax = axes[row, c] + ax.imshow(stack[plane], cmap="magma", vmin=0, vmax=top) + ax.set_xticks([]) + ax.set_yticks([]) + if row == 2 * r: + ax.set_title(f"{z_um[plane]:.2f} {unit}" if insitu + else f"{z_um[plane]:+.2f} {unit}", fontsize=8) + ax = axes[row, -1] + ax.imshow(stack[:, ny // 2, :], cmap="magma", aspect="auto", vmin=0, vmax=top) + ax.set_xticks([]) + ax.set_yticks([]) + if row == 2 * r: + ax.set_title("x-z", fontsize=8) + axes[2 * r + 1, 0].set_ylabel("model", fontsize=8) + axes[2 * r, 0].set_ylabel(f"{(label + ' ') if label else ''}data\n({what})", fontsize=8) + + +def has_localization_bias(data) -> bool: + """Bead results carry each bead refitted plane by plane; in situ ones do not.""" + loc = data.get("locres", {}).get("loc") + return bool(loc) and np.asarray(loc["z"]).ndim == 2 and np.asarray(loc["z"]).shape[1] > 1 + + +def localization_bias(data): + """``(x, y, z)`` bias in nm, beads by planes, as uiPSF's ``plotlocbias``. + + uiPSF refits every bead's every plane with the learnt model (spline MLE); + x and y are the deviation from the bead's position, z the fitted plane + less the plane it was taken at. + """ + loc, p = data["locres"]["loc"], data["params"]["pixel_size"] + x = np.asarray(loc["x"], float) * 1000 * float(p["x"]) + y = np.asarray(loc["y"], float) * 1000 * float(p["y"]) + z = np.asarray(loc["z"], float) + z = (z - np.arange(z.shape[1])) * 1000 * float(p["z"]) + return x, y, z + + +def draw_localization_bias(figure, data) -> None: + """Each bead's fitted x, y, z against where it was, through the stack.""" + x, y, z = localization_bias(data) + dz_nm = 1000 * float(data["params"]["pixel_size"]["z"]) + centre = (z.shape[1] - 1) / 2 + stage = (np.arange(z.shape[1]) - centre) * dz_nm + axes = figure.subplots(1, 3) + for ax, values, name in zip(axes, (x, y, z), ("x", "y", "z")): + ax.plot(stage, values.T, color="0.2", alpha=0.15, lw=0.8) + ax.plot(stage, np.nanmedian(values, axis=0), color="C3", lw=1.5, label="median") + ax.axhline(0, color="C0", lw=0.8) + ax.set(xlabel="stage z (nm)", ylabel=f"{name} bias (nm)") + inner = values[:, 2:-2] if values.shape[1] > 4 else values + lo, hi = np.nanquantile(inner, [0.005, 0.995]) + ax.set_ylim(max(lo, -300) - 5, min(hi, 300) + 5) + axes[0].legend(fontsize=8) + + +def _wavelength_nm(data) -> float: + return 1000 * float(data["params"]["option"]["imaging"]["emission_wavelength"]) + + +def pupil_images(res, wavelength_nm): + """Magnitude and aberration phase (nm of wavefront) over the pupil. + + The phase is rebuilt from the Zernike terms from Noll 5 up, as uiPSF's + ``showzernike`` does: piston, tilt and defocus only place the emitter. + Without Zernike coefficients (a free pupil) it is the learnt pupil's own. + """ + pupil = np.asarray(res["pupil"]) + pupil = pupil[0] if pupil.ndim > 2 else pupil + aperture = np.abs(pupil) > 0 + if "zernike_coeff" in res and "zernike_polynomial" in res: + coeff = np.asarray(res["zernike_coeff"], float).reshape(2, -1) + Zk = np.asarray(res["zernike_polynomial"], float) + n = min(len(Zk), coeff.shape[1]) + magnitude = np.tensordot(coeff[0, :n], Zk[:n], 1) + phase = np.tensordot(coeff[1, 4:n], Zk[4:n], 1) + else: + magnitude, phase = np.abs(pupil), np.angle(pupil) + phase = phase * wavelength_nm / (2 * np.pi) + return (np.where(aperture, magnitude, np.nan), np.where(aperture, phase, np.nan)) + + +def draw_pupil(figure, data) -> None: + """Per channel: the pupil's magnitude and aberration, and the Zernike terms.""" + rows = channels(data) + wavelength = _wavelength_nm(data) + grid = figure.add_gridspec(len(rows), 3, width_ratios=(1, 1, 2.2)) + for r, (label, res) in enumerate(rows): + magnitude, phase = pupil_images(res, wavelength) + ax = figure.add_subplot(grid[r, 0]) + image = ax.imshow(magnitude, cmap="viridis") + figure.colorbar(image, ax=ax, shrink=0.8) + ax.set_title(f"{label} magnitude".strip(), fontsize=9) + ax.set_axis_off() + ax = figure.add_subplot(grid[r, 1]) + limit = np.nanmax(np.abs(phase)) or 1.0 + image = ax.imshow(phase, cmap="RdBu_r", vmin=-limit, vmax=limit) + figure.colorbar(image, ax=ax, shrink=0.8, label="nm") + ax.set_title(f"{label} aberration".strip(), fontsize=9) + ax.set_axis_off() + ax = figure.add_subplot(grid[r, 2]) + if "zernike_coeff" not in res: + ax.text(0.5, 0.5, "free pupil: no Zernike terms", ha="center", + transform=ax.transAxes) + ax.set_axis_off() + continue + coeff = np.asarray(res["zernike_coeff"], float).reshape(2, -1) + noll = np.arange(1, coeff.shape[1] + 1) + nm = coeff[1] * wavelength / (2 * np.pi) + ax.bar(noll[4:], nm[4:], color="C0", width=0.8, label="phase (nm rms)") + twin = ax.twinx() + twin.plot(noll, coeff[0], ".", color="C1", ms=4, label="magnitude") + twin.set_ylabel("magnitude", fontsize=8, color="C1") + ax.axhline(0, color="0.5", lw=0.8) + ax.set_xlabel("Noll index", fontsize=8) + ax.set_ylabel("phase (nm rms)", fontsize=8) + named = [f"{name} {nm[j - 1]:+.0f}" for j, name in NAMED.items() if j <= len(nm)] + ax.text(0.99, 0.97, "\n".join(named), transform=ax.transAxes, ha="right", + va="top", fontsize=7, + bbox=dict(boxstyle="round", fc="white", ec="0.7", alpha=0.85)) + + +def draw_emitters(figure, data) -> None: + """Every emitter found, and those uiPSF kept for the model, per channel.""" + rows = channels(data) + size = np.asarray(data["rois"]["image_size"]).ravel() + axes = figure.subplots(1, len(rows), squeeze=False)[0] + beads = not _insitu(data) + for ax, (label, res) in zip(axes, rows): + found, kept = np.asarray(res["cor_all"]), np.asarray(res["cor"]) + ax.plot(found[:, -1], found[:, -2], ".", color="0.6", ms=3 if beads else 1, + label=f"found ({len(found)})") + ax.plot(kept[:, -1], kept[:, -2], "o", mfc="none", color="C3", + ms=7 if beads else 3, label=f"used ({len(kept)})") + ax.set_xlim(0, size[-1]) + ax.set_ylim(size[-2], 0) + ax.set_aspect("equal") + ax.set(xlabel="x (px)", ylabel="y (px)", title=label or None) + ax.legend(fontsize=7, loc="lower right") + + +def transformation_residuals(data): + """Main-channel positions, and the secondary's minus the main's mapped by ``T``. + + In uiPSF's channel arrays and pixels: ``T`` maps the main channel's + (y, x, 1) - imgcenter to the secondary's, as ``showtransform`` applies it. + """ + res = data["res"] + T = np.asarray(res["T"], float).reshape(-1, 3, 3)[0] + centre = np.asarray(res["imgcenter"], float).ravel() + main = np.asarray(res["channel0"]["pos"], float)[:, 1:] + secondary = np.asarray(res["channel1"]["pos"], float)[:, 1:] + mapped = (np.c_[main, np.ones(len(main))] - centre) @ T + mapped = mapped[:, :2] + centre[:2] + return main, secondary - mapped + + +def draw_transformation(figure, data) -> None: + """Where the transformation leaves each bead pair, in nm: arrows and spread.""" + main, residual = transformation_residuals(data) + p = data["params"]["pixel_size"] + nm = residual * 1000 * np.array([float(p["y"]), float(p["x"])]) + left, right = figure.subplots(1, 2) + scale = max(float(np.nanmax(np.hypot(*nm.T))), 1.0) + size = np.asarray(data["rois"]["image_size"]).ravel() + left.quiver(main[:, 1], main[:, 0], nm[:, 1], -nm[:, 0], angles="xy", + scale_units="xy", scale=scale / (0.08 * size[-1]), color="C3") + left.plot(main[:, 1], main[:, 0], ".", color="0.3", ms=3) + left.set_xlim(0, size[-1]) + left.set_ylim(size[-2], 0) + left.set_aspect("equal") + left.set(xlabel="x (px)", ylabel="y (px)", + title=f"residuals, longest arrow {scale:.1f} nm") + right.plot(nm[:, 1], nm[:, 0], "o", ms=4, color="C3") + right.axhline(0, color="0.5", lw=0.8) + right.axvline(0, color="0.5", lw=0.8) + rms = np.sqrt(np.mean(nm ** 2, axis=0)) + right.set(xlabel="x residual (nm)", ylabel="y residual (nm)", + title=f"rms x {rms[1]:.1f}, y {rms[0]:.1f} nm") + right.set_aspect("equal", adjustable="datalim") + + +def figures(data): + """``{tab name: (draw(figure), panels, size)}`` for what this result has.""" + n = len(channels(data)) + out = {"data vs model": (lambda fig: draw_data_vs_model(fig, data), + 2 * n * (PLANES + 1), (11, 2.6 * n))} + if has_localization_bias(data): + out["localization bias"] = (lambda fig: draw_localization_bias(fig, data), 3, (11, 3.5)) + out["pupil"] = (lambda fig: draw_pupil(fig, data), 3 * n, (11, 3.3 * n)) + out["emitters" if _insitu(data) else "beads"] = (lambda fig: draw_emitters(fig, data), + n, (4.5 * n, 4.5)) + if n > 1 and "T" in data["res"]: + out["channel transformation"] = (lambda fig: draw_transformation(fig, data), 2, (10, 4.5)) + return out diff --git a/tests/test_uipsf.py b/tests/test_uipsf.py index 404c79b..619e0ac 100644 --- a/tests/test_uipsf.py +++ b/tests/test_uipsf.py @@ -42,19 +42,20 @@ def params(psftype="zernike_vector", stage_pos=1.0, med=1.335, imm=1.516): "model": {"zernike_nl": []}}} -def write_result(path, res, p): +def write_result(path, res, p, rois=None, locres=None): """A file laid out as uiPSF's `writeh5file` lays one out.""" - with h5py.File(path, "w") as f: - f.attrs["params"] = json.dumps(p) - g = f.create_group("res") - for key, value in res.items(): + def put(group, items): + for key, value in items.items(): if isinstance(value, dict): - sub = g.create_group(key) - for k, v in value.items(): - sub[k] = v + put(group.create_group(key), value) else: - g[key] = value - f.create_group("locres")["coeff"] = np.zeros(1) + group[key] = value + with h5py.File(path, "w") as f: + f.attrs["params"] = json.dumps(p) + put(f.create_group("res"), res) + put(f.create_group("locres"), locres or {"coeff": np.zeros(1)}) + if rois is not None: + put(f.create_group("rois"), rois) return path @@ -201,6 +202,97 @@ def test_two_channels_share_one_normalisation_and_the_secondary_is_unmirrored(tm np.testing.assert_allclose(again.transformation, calibration.transformation) +def bead_result(tmp_path, dual=False): + """A complete uiPSF bead result, one or two channels, as the plots read it. + + The second channel sits 2 px right and 1 px down of the first, which is + what ``T`` says: the transformation's residuals are zero by construction. + The localization bias is a known 3 nm in x and one plane in z. + """ + rng = np.random.default_rng(5) + model = bead_model() + nz, n = len(model), 6 + zernike = np.zeros((2, 45)) + zernike[0, 0], zernike[1, 5] = 1.0, 0.5 + yy, xx = np.mgrid[:64, :64] + pupil = ((yy - 31.5) ** 2 + (xx - 31.5) ** 2 < 30 ** 2).astype(np.complex64) + cor = np.array([[20, 20], [20, 60], [60, 20], [60, 60], [40, 40], [30, 70]]) + pos = np.c_[np.full(n, nz / 2), cor + rng.uniform(-0.3, 0.3, (n, 2))] + + def channel(shift): + return {"I_model": model, "intensity": np.full((n, nz), 2000.0), + "zernike_coeff": zernike, "zernike_polynomial": np.ones((45, 64, 64)), + "pupil": pupil, "pos": pos + [0, *shift], "cor": cor + shift, + "cor_all": np.r_[cor + shift, [[5, 5]]]} + stacks = np.stack([model * 2000 + rng.normal(0, 1, model.shape) for _ in range(n)]) + loc = {"x": np.full((n, nz), 0.03), "y": np.zeros((n, nz)), + "z": np.tile(np.arange(nz) + 1.0, (n, 1))} + if not dual: + return write_result(tmp_path / "single.h5", channel(np.zeros(2)), params(), + {"psf_data": stacks, "psf_fit": stacks * 0.98, + "cor": cor, "image_size": [1, nz, 80, 80]}, + {"loc": loc}) + T = np.eye(3) + T[2, :2] = [1.0, 2.0] # rows (y, x, 1) - c: + (1, 2) + return write_result(tmp_path / "dual.h5", + {"channel0": channel(np.zeros(2)), "channel1": channel(np.array([1, 2])), + "T": T, "imgcenter": np.array([40.0, 40.0, 0.0])}, + params(), {"psf_data": np.stack([stacks, stacks]), + "psf_fit": np.stack([stacks, stacks]), + "cor": np.stack([cor, cor + [1, 2]]), + "image_size": [2, 1, nz, 80, 80]}, {"loc": loc}) + + +@pytest.mark.parametrize("dual", [False, True]) +def test_every_tab_of_a_result_draws_into_a_figure_and_into_a_page(tmp_path, dual): + import matplotlib + matplotlib.use("Agg") + from matplotlib.figure import Figure + from smappy.uipsf import plots + data = convert.read(bead_result(tmp_path, dual)) + tabs = plots.figures(data) + expected = ["data vs model", "localization bias", "pupil", "beads"] + assert list(tabs) == expected + (["channel transformation"] if dual else []) + for draw, panels, size in tabs.values(): + figure = Figure(figsize=size, layout="constrained") + draw(figure) + figure.savefig(tmp_path / "tab.png") + page = Figure().subfigures(1, 2)[1] # the All page hands out subfigures + draw(page) + assert page.axes + + +def test_the_bias_and_the_transformation_residuals_come_out_in_nanometres(tmp_path): + from smappy.uipsf import plots + data = convert.read(bead_result(tmp_path, dual=True)) + x, y, z = plots.localization_bias(data) + assert np.allclose(x, 3.0) and np.allclose(y, 0.0) and np.allclose(z, DZ_NM) + main, residual = plots.transformation_residuals(data) + np.testing.assert_allclose(residual, 0, atol=1e-6) + + +def test_in_situ_data_are_averaged_per_plane_of_the_model(tmp_path): + """Molecules fall into the model's planes by their fitted height, as uiPSF bins them.""" + from smappy.uipsf import plots + nz, zoffset = 5, 10.0 + model = np.ones((nz, 7, 7)) + z = np.array([10.2, 10.7, 12.5, 14.9]) # planes 0, 0, 2, 4 + rois = np.stack([np.full((7, 7), v) for v in (1.0, 3.0, 5.0, 7.0)]) + path = write_result(tmp_path / "insitu.h5", + {"I_model": model, "pos": np.c_[z, np.zeros((4, 2))], + "zoffset": np.array([[zoffset]]), "stagepos": np.array([1.0]), + "cor": np.zeros((4, 2)), "cor_all": np.zeros((9, 2)), + "pupil": np.ones((8, 8), np.complex64)}, + params("insitu_zernike"), + {"psf_data": rois, "psf_fit": rois, "image_size": [100, 64, 64]}, + {"loc": {"x": np.zeros((4, 1)), "y": np.zeros((4, 1)), + "z": np.zeros((4, 1))}}) + data = convert.read(path) + (label, measured, modelled, what), = plots.data_and_model(data) + np.testing.assert_allclose(measured[:, 0, 0], [2.0, 0.0, 5.0, 0.0, 7.0]) + assert list(plots.figures(data)) == ["data vs model", "pupil", "emitters"] + + def test_a_profile_says_what_it_knows_and_refuses_what_it_does_not(tmp_path, monkeypatch): shipped = microscopes.load("example") assert shipped.dual["layout"] == "up-down mirrored" and shipped.NA == 1.43 @@ -302,6 +394,8 @@ def test_uipsf_learns_simulated_beads_that_spline_3d_then_fits_in_z(tmp_path, mo learning=LearningPart(iterations=60))) calibration = result.data["calibration"] assert load_spline_calibration(calibration).psf.shape[1] == 21 + assert [f.name for f in result.figures()] == ["data vs model", "localization bias", + "pupil", "beads"] frames, truth = camera_frames(300, seed=3, astigmatism=ASTIGMATISM, labelling=LabellingSettings(efficiency=0.1)) tifffile.imwrite(tmp_path / "frames.tif", frames) From 02492374edc85b2f81cc19f364c17cf3f909ba84 Mon Sep 17 00:00:00 2001 From: Claude Date: Tue, 6 Oct 2026 16:14:21 +0000 Subject: [PATCH 4/7] uiPSF calibration: the Zernike bar chart back, as a tab; beads tab fixed - a "Zernike" tab after "pupil": the aberration phases from Noll 5 up in nm rms as bars, the channels side by side, the named terms labelled - the beads/emitters tab of a one-channel result failed to draw in the window: one panel is handed an axis, not a figure. The tests now draw every tab through Plot.draw_into, as the window does, which is what found it Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01GiNdULxyiqQEqWHGKTbVzV --- .../docs/Localize/uiPSF calibration.md | 9 ++- src/smappy/uipsf/plots.py | 56 ++++++++++++++++++- tests/test_uipsf.py | 15 +++-- 3 files changed, 71 insertions(+), 9 deletions(-) diff --git a/src/smappy/plugins/docs/Localize/uiPSF calibration.md b/src/smappy/plugins/docs/Localize/uiPSF calibration.md index e577f96..5e42e94 100644 --- a/src/smappy/plugins/docs/Localize/uiPSF calibration.md +++ b/src/smappy/plugins/docs/Localize/uiPSF calibration.md @@ -1,6 +1,6 @@ --- version: "1" -covers: [smappy.uipsf.prepare.Split, smappy.uipsf.prepare.channel_shift, smappy.uipsf.prepare.subtract_background, smappy.uipsf.prepare.read_beads, smappy.uipsf.prepare.read_movie, smappy.uipsf.uipsf_parameters, smappy.uipsf.runner.run, smappy.uipsf.convert.single_calibration, smappy.uipsf.convert.dual_calibration, smappy.uipsf.convert.channel_transformation, smappy.uipsf.microscopes.Microscope, smappy.uipsf.plots.figures, smappy.uipsf.plots.data_and_model, smappy.uipsf.plots.localization_bias, smappy.uipsf.plots.pupil_images, smappy.uipsf.plots.transformation_residuals] +covers: [smappy.uipsf.prepare.Split, smappy.uipsf.prepare.channel_shift, smappy.uipsf.prepare.subtract_background, smappy.uipsf.prepare.read_beads, smappy.uipsf.prepare.read_movie, smappy.uipsf.uipsf_parameters, smappy.uipsf.runner.run, smappy.uipsf.convert.single_calibration, smappy.uipsf.convert.dual_calibration, smappy.uipsf.convert.channel_transformation, smappy.uipsf.microscopes.Microscope, smappy.uipsf.plots.figures, smappy.uipsf.plots.data_and_model, smappy.uipsf.plots.localization_bias, smappy.uipsf.plots.pupil_images, smappy.uipsf.plots.zernike_phases, smappy.uipsf.plots.transformation_residuals] --- ## What it does @@ -317,7 +317,8 @@ aberrations in nanometres of wavefront, and two numbers from uiPSF's own check of the result. The first is the median bias of the beads, refitted with the model. The second, for two channels, is how far the transformation leaves each bead pair apart. The figures are uiPSF's own -diagnostics, the ones its notebooks show, one tab each: +diagnostics, the ones its notebooks show, one tab each, and a bar chart of +the aberrations: * **data vs model**: the measured spot above and uiPSF's model of it below, at planes through the model and in an x–z section, per channel. For beads @@ -341,6 +342,10 @@ diagnostics, the ones its notebooks show, one tab each: uiPSF names. Astigmatism (5, 6) is expected with an astigmatic lens, and spherical (11) grows with depth in water. Large coma (7, 8) points at a tilted or misaligned optic. +* **Zernike**: the same phases as bars, from term 5 up, the channels side by + side and the named terms labelled under the axis. Two channels of one + microscope should share most of their aberrations. A term that differs + between them belongs to the optics after the splitter. * **beads** or **emitters**: where uiPSF found emitters (grey) and which it used for the model (red circles), in the channel as uiPSF saw it. Beads too close to another or to the edge are dropped before they count as diff --git a/src/smappy/uipsf/plots.py b/src/smappy/uipsf/plots.py index 458ca8c..b957309 100644 --- a/src/smappy/uipsf/plots.py +++ b/src/smappy/uipsf/plots.py @@ -11,6 +11,8 @@ * `draw_localization_bias`: ``showlocalization`` * `draw_pupil`: ``showzernike`` (``showpupil`` for a model without Zernike coefficients) +* `draw_zernike`: the aberrations as bars, the channels side by side -- + ``showzernike``'s coefficients, read off more easily * `draw_emitters`: ``showcoord`` * `draw_transformation`: ``showtransform``, as residuals rather than overlays @@ -31,6 +33,9 @@ # the terms uiPSF's showzernike names, by Noll index NAMED = {5: "astigmatism 45", 6: "astigmatism 0", 7: "coma y", 8: "coma x", 11: "spherical", 22: "2nd spherical"} +# the same, short enough to sit side by side under the Zernike bars +SHORT = {5: "astig. 45", 6: "astig. 0", 7: "coma y", 8: "coma x", 11: "spherical", + 22: "2nd spher."} def channels(data): @@ -236,11 +241,52 @@ def draw_pupil(figure, data) -> None: bbox=dict(boxstyle="round", fc="white", ec="0.7", alpha=0.85)) +def zernike_phases(data): + """``[(label, Noll indices, phase in nm rms)]`` per channel, from Noll 5 up.""" + wavelength = _wavelength_nm(data) + out = [] + for label, res in channels(data): + if "zernike_coeff" not in res: + continue + phase = np.asarray(res["zernike_coeff"], float).reshape(2, -1)[1] + noll = np.arange(1, len(phase) + 1) + out.append((label, noll[4:], phase[4:] * wavelength / (2 * np.pi))) + return out + + +def draw_zernike(ax, data) -> None: + """The aberrations as bars, Noll index 5 up, the channels side by side. + + Piston, tilt and defocus are left out: they only place the emitter. The + named terms are labelled under the axis, so astigmatism, coma and + spherical can be read off without counting. + """ + rows = zernike_phases(data) + if not rows: + ax.text(0.5, 0.5, "free pupil: no Zernike terms", ha="center", transform=ax.transAxes) + ax.set_axis_off() + return + width = 0.8 / len(rows) + for k, (label, noll, nm) in enumerate(rows): + ax.bar(noll + (k - (len(rows) - 1) / 2) * width, nm, width=width, label=label or None) + ax.axhline(0, color="0.5", lw=0.8) + named = [j for j in NAMED if j <= rows[0][1][-1]] + ax.set_xticks(named, [f"{SHORT[j]} ({j})" for j in named], fontsize=7, + rotation=40, ha="right", rotation_mode="anchor") + ax.set_xticks(rows[0][1], minor=True) + ax.set_xlim(rows[0][1][0] - 1, rows[0][1][-1] + 1) + ax.set(xlabel="Noll index", ylabel="phase (nm rms)") + if len(rows) > 1: + ax.legend(fontsize=8) + + def draw_emitters(figure, data) -> None: """Every emitter found, and those uiPSF kept for the model, per channel.""" rows = channels(data) size = np.asarray(data["rois"]["image_size"]).ravel() - axes = figure.subplots(1, len(rows), squeeze=False)[0] + # one channel is one panel, and a one-panel plot is handed an axis + axes = (figure.subplots(1, len(rows), squeeze=False)[0] + if hasattr(figure, "subplots") else [figure]) beads = not _insitu(data) for ax, (label, res) in zip(axes, rows): found, kept = np.asarray(res["cor_all"]), np.asarray(res["cor"]) @@ -297,13 +343,19 @@ def draw_transformation(figure, data) -> None: def figures(data): - """``{tab name: (draw(figure), panels, size)}`` for what this result has.""" + """``{tab name: (draw, panels, size)}`` for what this result has. + + ``draw`` takes the figure when ``panels`` > 1 and an axis when it is 1, + as `smappy.plugins.Plot` hands them out. + """ n = len(channels(data)) out = {"data vs model": (lambda fig: draw_data_vs_model(fig, data), 2 * n * (PLANES + 1), (11, 2.6 * n))} if has_localization_bias(data): out["localization bias"] = (lambda fig: draw_localization_bias(fig, data), 3, (11, 3.5)) out["pupil"] = (lambda fig: draw_pupil(fig, data), 3 * n, (11, 3.3 * n)) + # one panel: handed an axis, as a plain plot(ax) is + out["Zernike"] = (lambda ax: draw_zernike(ax, data), 1, (9, 4)) out["emitters" if _insitu(data) else "beads"] = (lambda fig: draw_emitters(fig, data), n, (4.5 * n, 4.5)) if n > 1 and "T" in data["res"]: diff --git a/tests/test_uipsf.py b/tests/test_uipsf.py index 619e0ac..fc54444 100644 --- a/tests/test_uipsf.py +++ b/tests/test_uipsf.py @@ -251,14 +251,16 @@ def test_every_tab_of_a_result_draws_into_a_figure_and_into_a_page(tmp_path, dua from smappy.uipsf import plots data = convert.read(bead_result(tmp_path, dual)) tabs = plots.figures(data) - expected = ["data vs model", "localization bias", "pupil", "beads"] + from smappy.plugins import Plot + expected = ["data vs model", "localization bias", "pupil", "Zernike", "beads"] assert list(tabs) == expected + (["channel transformation"] if dual else []) for draw, panels, size in tabs.values(): + plot = Plot(draw=draw, panels=panels, size=size) figure = Figure(figsize=size, layout="constrained") - draw(figure) + plot.draw_into(figure) figure.savefig(tmp_path / "tab.png") page = Figure().subfigures(1, 2)[1] # the All page hands out subfigures - draw(page) + plot.draw_into(page) assert page.axes @@ -269,6 +271,9 @@ def test_the_bias_and_the_transformation_residuals_come_out_in_nanometres(tmp_pa assert np.allclose(x, 3.0) and np.allclose(y, 0.0) and np.allclose(z, DZ_NM) main, residual = plots.transformation_residuals(data) np.testing.assert_allclose(residual, 0, atol=1e-6) + # 0.5 rad of Noll 6 at 680 nm, in both channels + for label, noll, nm in plots.zernike_phases(data): + assert nm[list(noll).index(6)] == pytest.approx(0.5 * 680 / (2 * np.pi)) def test_in_situ_data_are_averaged_per_plane_of_the_model(tmp_path): @@ -290,7 +295,7 @@ def test_in_situ_data_are_averaged_per_plane_of_the_model(tmp_path): data = convert.read(path) (label, measured, modelled, what), = plots.data_and_model(data) np.testing.assert_allclose(measured[:, 0, 0], [2.0, 0.0, 5.0, 0.0, 7.0]) - assert list(plots.figures(data)) == ["data vs model", "pupil", "emitters"] + assert list(plots.figures(data)) == ["data vs model", "pupil", "Zernike", "emitters"] def test_a_profile_says_what_it_knows_and_refuses_what_it_does_not(tmp_path, monkeypatch): @@ -395,7 +400,7 @@ def test_uipsf_learns_simulated_beads_that_spline_3d_then_fits_in_z(tmp_path, mo calibration = result.data["calibration"] assert load_spline_calibration(calibration).psf.shape[1] == 21 assert [f.name for f in result.figures()] == ["data vs model", "localization bias", - "pupil", "beads"] + "pupil", "Zernike", "beads"] frames, truth = camera_frames(300, seed=3, astigmatism=ASTIGMATISM, labelling=LabellingSettings(efficiency=0.1)) tifffile.imwrite(tmp_path / "frames.tif", frames) From f268d802e42cf711d26bfef2cd456379d671a820 Mon Sep 17 00:00:00 2001 From: Claude Date: Tue, 6 Oct 2026 16:41:13 +0000 Subject: [PATCH 5/7] uiPSF: the channel shift set from smappy's worker, not a uiPSF parameter uiPSF's main stays without a channel_shift parameter: the worker sets the measured shift (prepare.channel_shift) as the multi-channel data's shiftxy before uiPSF pairs the beads, on uiPSF's class, since prep_data makes and processes the object in one call. The shift travels in job.json. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01GiNdULxyiqQEqWHGKTbVzV --- NOTES.md | 8 +++++--- src/smappy/plugins/uipsf.py | 7 ++++--- src/smappy/uipsf/__init__.py | 11 ++++++++--- src/smappy/uipsf/prepare.py | 5 +++-- src/smappy/uipsf/runner.py | 5 +++-- src/smappy/uipsf/worker.py | 26 ++++++++++++++++++++++++++ 6 files changed, 49 insertions(+), 13 deletions(-) diff --git a/NOTES.md b/NOTES.md index 5505866..77ef4fc 100644 --- a/NOTES.md +++ b/NOTES.md @@ -1750,9 +1750,11 @@ saving) is the price of keeping it out of that window. with the background in, and simulated two-colour beads (a background a fifth of the peak) lost 15 of 18 beads to spurious maxima. And it measures the shift between the channels by cross-correlating their - projections (`prepare.channel_shift`, passed as `channel_shift`, which - the uiPSF branch reads): uiPSF's own guess, from bead coordinates, took - the neighbour on a 30 px grid of beads and paired one bead. + projections (`prepare.channel_shift`), which `worker.py` sets as the + multi-channel data's `shiftxy` before uiPSF pairs the beads: uiPSF's own + guess, from bead coordinates, took the neighbour on a 30 px grid of beads + and paired one bead. It is set from the worker, on uiPSF's class, rather + than by a new uiPSF parameter, so that uiPSF's main stays as it is. * **Voxels, not uiPSF's spline.** uiPSF's `locres/coeff` is normalised to the median plane sum after subtracting the minimum; the fitter's photons assume the brightest plane sums to one. `res/I_model` is clipped, diff --git a/src/smappy/plugins/uipsf.py b/src/smappy/plugins/uipsf.py index 9b514c5..2c104c9 100644 --- a/src/smappy/plugins/uipsf.py +++ b/src/smappy/plugins/uipsf.py @@ -209,13 +209,14 @@ def run(self, ctx: Context, settings: UiPSFSettings) -> Result: zernike_order=lp.zernike_order, depth_um=b.depth_um, z_range_um=b.z_range_um, min_photon=b.min_photon, z_bins=b.z_bins, per_bin=b.per_bin, repeat=b.repeat, start_from=_uipsf_file(b.start_from) if settings.data == "blinking" else "") - if dual and settings.data == "beads": - params["channel_shift"] = prepare.channel_shift(data.images) + shift = (prepare.channel_shift(data.images) + if dual and settings.data == "beads" else None) out = output_path(settings) job = (out.with_name(out.stem + "_uipsf_job") if lp.keep_job else Path(tempfile.mkdtemp(prefix="smappy-uipsf-"))) result = learn(data.images, data=settings.data, model=settings.model, dual=dual, - params=params, workdir=job, python=lp.python, progress=ctx.report) + params=params, workdir=job, python=lp.python, progress=ctx.report, + channel_shift=shift) em_on = getattr(data.camera, "em_on", None) calibration = (convert.dual_calibration(result, data.split, sources, em_on=em_on) if dual else convert.single_calibration(result, em_on=em_on)) diff --git a/src/smappy/uipsf/__init__.py b/src/smappy/uipsf/__init__.py index 274cf16..693fc89 100644 --- a/src/smappy/uipsf/__init__.py +++ b/src/smappy/uipsf/__init__.py @@ -91,12 +91,17 @@ def _merge(target: dict, extra: dict) -> None: def learn(images, *, data: str, model: str, dual: bool, params: dict, workdir=None, python: str = "", - progress: Optional[Callable[[str], None]] = None) -> Path: - """Run uiPSF on ``images`` (`prepare`'s); the result file it wrote.""" + progress: Optional[Callable[[str], None]] = None, channel_shift=None) -> Path: + """Run uiPSF on ``images`` (`prepare`'s); the result file it wrote. + + ``channel_shift`` (`prepare.channel_shift`) is where the worker starts + uiPSF's pairing of two channels' beads; None leaves uiPSF to find it. + """ from . import runner psftype, override = PSFTYPES[(data, model)] if override: params = dict(params, PSFtype=override) workdir = Path(workdir) if workdir else Path(tempfile.mkdtemp(prefix="smappy-uipsf-")) - runner.write_job(workdir, images, psftype, "2ch" if dual else "1ch", params) + runner.write_job(workdir, images, psftype, "2ch" if dual else "1ch", params, + channel_shift) return runner.run(workdir, python=python, progress=progress) diff --git a/src/smappy/uipsf/prepare.py b/src/smappy/uipsf/prepare.py index 58a9b54..8fa9360 100644 --- a/src/smappy/uipsf/prepare.py +++ b/src/smappy/uipsf/prepare.py @@ -115,8 +115,9 @@ def channel_shift(channels: np.ndarray) -> list: which it finds by matching bead coordinates -- and on a regular pattern of beads that locked onto the neighbouring bead, 30 px off, leaving one pair. The cross-correlation of the two channels' projections peaks where - *all* the beads overlap, which a neighbour's shift cannot match. Handed - to uiPSF as ``channel_shift``: per channel, target minus reference. + *all* the beads overlap, which a neighbour's shift cannot match. Per + channel, target minus reference, [y, x]: what uiPSF's multi-channel data + calls ``shiftxy``, and what `worker.py` sets before uiPSF pairs the beads. """ from scipy import ndimage ny, nx = channels.shape[-2:] diff --git a/src/smappy/uipsf/runner.py b/src/smappy/uipsf/runner.py index 3be303c..201301b 100644 --- a/src/smappy/uipsf/runner.py +++ b/src/smappy/uipsf/runner.py @@ -61,13 +61,14 @@ def find_python(setting: str = "") -> str: def write_job(directory: Path, images: np.ndarray, psftype: str, channeltype: str, - params: dict) -> Path: + params: dict, channel_shift=None) -> Path: """The job directory the worker reads.""" directory = Path(directory) directory.mkdir(parents=True, exist_ok=True) np.save(directory / "images.npy", np.ascontiguousarray(images, dtype=np.float32)) (directory / "job.json").write_text(json.dumps( - {"psftype": psftype, "channeltype": channeltype, "params": params}, indent=1)) + {"psftype": psftype, "channeltype": channeltype, "params": params, + "channel_shift": channel_shift}, indent=1)) return directory diff --git a/src/smappy/uipsf/worker.py b/src/smappy/uipsf/worker.py index d0db51d..4c6eec1 100644 --- a/src/smappy/uipsf/worker.py +++ b/src/smappy/uipsf/worker.py @@ -41,6 +41,30 @@ def merge(target, overrides): target[key] = value +def preset_channel_shift(shift): + """Start uiPSF's pairing of two channels' beads from a known shift. + + uiPSF finds the shift between the channels from bead coordinates, unless + its multi-channel data already has one (``shiftxy``, target minus + reference, [y, x] per channel) when ``process`` runs -- which is set here, + on the class, because ``prep_data`` makes the object and processes it in + one call. smappy measures the shift by cross-correlation + (`prepare.channel_shift`): uiPSF's own guess paired each bead of a + regular grid with its neighbour. Done here rather than in uiPSF so that + uiPSF itself is unchanged. + """ + import numpy as np + from psflearning.learning.data_representation.PreprocessedImageDataMultiChannel_file \ + import PreprocessedImageDataMultiChannel as Multi + process = Multi.process + + def process_from_shift(self, *args, **kwargs): + if self.shiftxy is None: + self.shiftxy = np.float32(shift) + return process(self, *args, **kwargs) + Multi.process = process_from_shift + + def main(directory): os.environ.setdefault("MPLBACKEND", "Agg") # uiPSF imports pyplot import numpy as np @@ -67,6 +91,8 @@ def main(directory): multi = L.param.option.multi L.param.option.multi.defocus = [multi.defocus_offset + i*multi.defocus_delay for i in range(n)] + if job.get("channel_shift") and L.param.channeltype == "multi": + preset_channel_shift(job["channel_shift"]) start = time.time() L.getpsfclass() say("stage", "finding the emitters") From ff0baf8765c531c5302e9589fd61f6b3d598cf98 Mon Sep 17 00:00:00 2001 From: Claude Date: Tue, 6 Oct 2026 20:42:51 +0000 Subject: [PATCH 6/7] uiPSF: what main's users keep, and what the equivalence check showed - NOTES.md, "uiPSF's branch": the branch's setup.py keeps main's TensorFlow where it was pinned; on the old stack (TensorFlow 2.9.1, numpy 1.23.5) main and the branch learnt the same jobs bit-identically, and the worker drives an unchanged main as well - the page: an environment already made from uiPSF's main works too --slow -n auto: 1346 passed, 4 skipped (with SMAPPY_UIPSF_PYTHON set) Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01GiNdULxyiqQEqWHGKTbVzV --- NOTES.md | 15 +++++++++++++-- .../plugins/docs/Localize/uiPSF calibration.md | 10 ++++++---- 2 files changed, 19 insertions(+), 6 deletions(-) diff --git a/NOTES.md b/NOTES.md index 77ef4fc..62233bf 100644 --- a/NOTES.md +++ b/NOTES.md @@ -1783,8 +1783,19 @@ saving) is the price of keeping it out of that window. real data with a bead calibration beside it. * **uiPSF's branch.** uiPSF's main pins TensorFlow 2.9 / Python 3.7, which has no Apple-silicon build. The `claude/friendly-cerf-qt2sfj` branch of - ries-lab/uiPSF runs on numpy 2 and TensorFlow 2.21 (hdfdict replaced, - tkinter dropped, NEP 50 dtype casts), tried with Python 3.12 on Linux. + ries-lab/uiPSF (ries-lab/uiPSF#17) also runs on numpy 2 and TensorFlow + 2.21 (hdfdict replaced, tkinter dropped, NEP 50 dtype casts), tried with + Python 3.12 on Linux, and is meant to leave main's users where they were: + its `setup.py` still installs TensorFlow 2.9.1 / tensorflow-probability + 0.17 on Python 3.7 and keeps native Windows on Python 3.8-3.10 below + TensorFlow 2.11 (the last with its GPU support) and numpy below 1.24; + only elsewhere does it allow current TensorFlow. On the old stack + (TensorFlow 2.9.1, tensorflow-probability 0.17, numpy 1.23.5; Python 3.9, + as no 3.7 could be installed here) main and the branch learnt the same + jobs -- one-channel beads, two-channel beads, in situ, 15 iterations each, + driven by `worker.py` -- with every array of the result bit-identical, and + `io.h5.load` (hdfdict on main, h5py on the branch) read them back alike. + The worker drives an unchanged main as well. ## Open questions diff --git a/src/smappy/plugins/docs/Localize/uiPSF calibration.md b/src/smappy/plugins/docs/Localize/uiPSF calibration.md index 5e42e94..e3bbd69 100644 --- a/src/smappy/plugins/docs/Localize/uiPSF calibration.md +++ b/src/smappy/plugins/docs/Localize/uiPSF calibration.md @@ -76,10 +76,12 @@ pip install -e uiPSF ``` The last line installs uiPSF with TensorFlow and everything else it needs. -It has to be this branch of uiPSF: its main branch needs Python 3.7 and -TensorFlow 2.9, which do not run on current Macs or with current numpy. -Without conda, `python3.12 -m venv uipsf-env` and that environment's `pip` do -the same. +For a new environment it has to be this branch of uiPSF: its main branch +needs Python 3.7 and TensorFlow 2.9, which do not run on current Macs or with +current numpy. An environment already made from uiPSF's main branch, with +its old versions, works too (tried on Linux), and there the branch and main +learn exactly the same PSF. Without conda, +`python3.12 -m venv uipsf-env` and that environment's `pip` do the same. **2. Check it.** Still in that environment: From 779e529df6e7e26ccb673c2e12c7a61409847aee Mon Sep 17 00:00:00 2001 From: Claude Date: Wed, 7 Oct 2026 05:55:14 +0000 Subject: [PATCH 7/7] uiPSF calibration: install uiPSF from its main branch The page's install line clones uiPSF's main rather than the branch of ries-lab/uiPSF#17, which is to be merged first; NOTES.md speaks of #17 rather than of the branch. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01GiNdULxyiqQEqWHGKTbVzV --- NOTES.md | 10 +++++----- .../plugins/docs/Localize/uiPSF calibration.md | 16 ++++++++-------- 2 files changed, 13 insertions(+), 13 deletions(-) diff --git a/NOTES.md b/NOTES.md index 62233bf..ca88e24 100644 --- a/NOTES.md +++ b/NOTES.md @@ -1781,9 +1781,9 @@ saving) is the price of keeping it out of that window. file converts exactly, tests/test_uipsf.py); whether this is uiPSF on real data or the simulation (a demo structure only +-300 nm deep) needs real data with a bead calibration beside it. -* **uiPSF's branch.** uiPSF's main pins TensorFlow 2.9 / Python 3.7, which - has no Apple-silicon build. The `claude/friendly-cerf-qt2sfj` branch of - ries-lab/uiPSF (ries-lab/uiPSF#17) also runs on numpy 2 and TensorFlow +* **uiPSF on current numpy and TensorFlow.** uiPSF's main pinned + TensorFlow 2.9 / Python 3.7, which has no Apple-silicon build. + ries-lab/uiPSF#17 makes it run on numpy 2 and TensorFlow 2.21 (hdfdict replaced, tkinter dropped, NEP 50 dtype casts), tried with Python 3.12 on Linux, and is meant to leave main's users where they were: its `setup.py` still installs TensorFlow 2.9.1 / tensorflow-probability @@ -1791,11 +1791,11 @@ saving) is the price of keeping it out of that window. TensorFlow 2.11 (the last with its GPU support) and numpy below 1.24; only elsewhere does it allow current TensorFlow. On the old stack (TensorFlow 2.9.1, tensorflow-probability 0.17, numpy 1.23.5; Python 3.9, - as no 3.7 could be installed here) main and the branch learnt the same + as no 3.7 could be installed here) main and #17's branch learnt the same jobs -- one-channel beads, two-channel beads, in situ, 15 iterations each, driven by `worker.py` -- with every array of the result bit-identical, and `io.h5.load` (hdfdict on main, h5py on the branch) read them back alike. - The worker drives an unchanged main as well. + The worker drives uiPSF from before #17 as well. ## Open questions diff --git a/src/smappy/plugins/docs/Localize/uiPSF calibration.md b/src/smappy/plugins/docs/Localize/uiPSF calibration.md index e3bbd69..43a6ad2 100644 --- a/src/smappy/plugins/docs/Localize/uiPSF calibration.md +++ b/src/smappy/plugins/docs/Localize/uiPSF calibration.md @@ -71,17 +71,17 @@ terminal (on Windows, the Anaconda Prompt): ``` conda create -n uipsf python=3.12 conda activate uipsf -git clone --branch claude/friendly-cerf-qt2sfj https://github.com/ries-lab/uiPSF +git clone https://github.com/ries-lab/uiPSF pip install -e uiPSF ``` -The last line installs uiPSF with TensorFlow and everything else it needs. -For a new environment it has to be this branch of uiPSF: its main branch -needs Python 3.7 and TensorFlow 2.9, which do not run on current Macs or with -current numpy. An environment already made from uiPSF's main branch, with -its old versions, works too (tried on Linux), and there the branch and main -learn exactly the same PSF. Without conda, -`python3.12 -m venv uipsf-env` and that environment's `pip` do the same. +The last line installs uiPSF with TensorFlow and everything else it needs; +on Python 3.12 that is a current TensorFlow, which runs on current Macs too. +uiPSF from before October 2026 needed Python 3.7 and TensorFlow 2.9, which +do not install there; an environment already made from it works as it is +(tried on Linux), and the current uiPSF learns exactly the same PSF in it. +Without conda, `python3.12 -m venv uipsf-env` and that environment's `pip` +do the same. **2. Check it.** Still in that environment: