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/NOTES.md b/NOTES.md index b0190b7..ca88e24 100644 --- a/NOTES.md +++ b/NOTES.md @@ -1721,6 +1721,82 @@ 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`), 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, + 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 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 + 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 #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 uiPSF from before #17 as well. + ## 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..43a6ad2 --- /dev/null +++ b/src/smappy/plugins/docs/Localize/uiPSF calibration.md @@ -0,0 +1,392 @@ +--- +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.zernike_phases, smappy.uipsf.plots.transformation_residuals] +--- + +## 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 +(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. + +## 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 https://github.com/ries-lab/uiPSF +pip install -e uiPSF +``` + +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: + +``` +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 +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. + +**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 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, 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 + 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. +* **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 + 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 +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..2c104c9 --- /dev/null +++ b/src/smappy/plugins/uipsf.py @@ -0,0 +1,295 @@ +"""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 "") + 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, + 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)) + _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) + 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}) + + +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, 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] + 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", "")) + 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) diff --git a/src/smappy/uipsf/__init__.py b/src/smappy/uipsf/__init__.py new file mode 100644 index 0000000..693fc89 --- /dev/null +++ b/src/smappy/uipsf/__init__.py @@ -0,0 +1,107 @@ +"""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, 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, + channel_shift) + 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/plots.py b/src/smappy/uipsf/plots.py new file mode 100644 index 0000000..b957309 --- /dev/null +++ b/src/smappy/uipsf/plots.py @@ -0,0 +1,363 @@ +"""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_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 + +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"} +# 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): + """``[(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 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() + # 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"]) + 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, 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"]: + out["channel transformation"] = (lambda fig: draw_transformation(fig, data), 2, (10, 4.5)) + return out diff --git a/src/smappy/uipsf/prepare.py b/src/smappy/uipsf/prepare.py new file mode 100644 index 0000000..8fa9360 --- /dev/null +++ b/src/smappy/uipsf/prepare.py @@ -0,0 +1,279 @@ +"""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. 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:] + 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..201301b --- /dev/null +++ b/src/smappy/uipsf/runner.py @@ -0,0 +1,145 @@ +"""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, 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, + "channel_shift": channel_shift}, 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..4c6eec1 --- /dev/null +++ b/src/smappy/uipsf/worker.py @@ -0,0 +1,116 @@ +"""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 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 + 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)] + 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") + 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..fc54444 --- /dev/null +++ b/tests/test_uipsf.py @@ -0,0 +1,426 @@ +"""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, rois=None, locres=None): + """A file laid out as uiPSF's `writeh5file` lays one out.""" + def put(group, items): + for key, value in items.items(): + if isinstance(value, dict): + put(group.create_group(key), value) + else: + 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 + + +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 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) + 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") + plot.draw_into(figure) + figure.savefig(tmp_path / "tab.png") + page = Figure().subfigures(1, 2)[1] # the All page hands out subfigures + plot.draw_into(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) + # 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): + """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", "Zernike", "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 + (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 + assert [f.name for f in result.figures()] == ["data vs model", "localization bias", + "pupil", "Zernike", "beads"] + 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):