diff --git a/.gitattributes b/.gitattributes new file mode 100644 index 00000000..7b2b080f --- /dev/null +++ b/.gitattributes @@ -0,0 +1,31 @@ +# Line endings. +# +# Twenty-eight files in this repository were committed with CRLF endings +# while the other 579 text files used LF. Nothing depended on the +# difference, but it was invisible until something wrote to one of those +# files: an editor or script that rewrites a file whole would turn a two +# line change into a diff touching every line. Normalizing removed the +# inconsistency, and this file keeps it from coming back. +# +# text=auto lets git detect binary files rather than trusting the list +# below, which is a backstop for the formats this repository actually +# carries. The checked in Windows DLLs in storm_analysis/c_libraries are +# the ones that matter most - rewriting those would break the wheel that +# the Windows CI job builds. + +* text=auto eol=lf + +# Binary. Never touch these. +*.amat binary +*.bin binary +*.dax binary +*.dll binary +*.gif binary +*.hdf5 binary +*.jar binary +*.mat binary +*.npy binary +*.png binary +*.spline binary +*.tif binary +*.tiff binary diff --git a/AGENTS.md b/AGENTS.md new file mode 100644 index 00000000..6f3c8bb0 --- /dev/null +++ b/AGENTS.md @@ -0,0 +1,216 @@ +# Working on storm-analysis + +Notes for automated coding tools, and for anyone touching the code for the +first time. `README.md` is the entry point for *using* this package; this file +is about changing it. + +Everything below is checked against the tree rather than remembered, so if +something here disagrees with the code, the code is right and this file is +stale. Please fix it. + + +## Build + +The C libraries are built by SCons, not by setuptools, and `pip install -e .` +runs SCons for you through `_custom_build.py`: + +``` +pip install -e . +``` + +FFTW3 and LAPACK need to be present first. `doc/source/install.rst` has the +per-platform details, and the CI workflow in `.github/workflows` runs exactly +those steps on Linux, macOS and Windows. + +**After editing any `.c` or `.h` file, rebuild before running anything:** + +``` +scons -Q +``` + +Nothing reminds you. The Python side loads whatever shared library is already +sitting in `storm_analysis/c_libraries/`, so a test run against a stale library +looks like a normal result and quietly answers a question you did not ask. + +On Windows SCons has to be told which compiler to use; `_custom_build.py` +defaults to `compiler=mingw` and reads `SA_SCONS_COMPILER` if you want +something else. + + +## Tests + +From the repository root: + +``` +pytest +``` + +280 tests, around 45 seconds. Configuration lives in the `[tool.pytest]` table +of `pyproject.toml` — that is pytest's TOML mode, which is why `minversion` is +9.0. The older `[tool.pytest.ini_options]` spelling means ini mode; both are +valid and neither is a mistake to be corrected into the other. + +Two things about reading the result: + +- **Trust the exit status, not the last line printed.** Several of the C + libraries write progress to stdout through the C runtime. Those buffers are + flushed when the interpreter exits, which is after pytest has printed its + summary, so the final lines on screen are frequently library chatter and the + `N passed` line is nowhere in sight. Use `--junitxml=results.xml` when you + need counts rather than a yes or no. + +- **The `__main__` blocks are not covered.** Most test modules end with a block + that calls the tests so the file can be run directly. pytest never executes + it. Rename a test function without updating that block and the suite stays + green while the file is broken for anyone running it by hand — this has + happened, and went unnoticed across six pull requests. + + +## Running an analysis + +"Analyze this movie with 3D-DAOSTORM in 2d fixed mode" is a supported request, +and `doc/source/analysis.rst` is the reference for it. What follows is only the +part that is easy to get wrong. + +Each fitter has the same entry point shape: + +| Method | Entry point | +| --- | --- | +| 3D-DAOSTORM | `daostorm_3d/mufit_analysis.py` | +| sCMOS | `sCMOS/scmos_analysis.py` | +| Spliner | `spliner/spline_analysis.py` | +| Multiplane | `multi_plane/multi_plane.py` | +| Pupil function | `pupilfn/pupilfn_analysis.py` | +| PSF FFT | `psf_fft/psffft_analysis.py` | + +``` +python path/to/mufit_analysis.py --movie movie.dax --bin out.hdf5 --xml params.xml +``` + +Each also exposes `analyze(movie, output, parameters)` for calling from Python. +Multiplane is the exception in both forms: it takes `--basename` rather than +`--movie`, because it loads one movie per channel and builds each name by +appending the `channelX_ext` parameters to the basename. + +**Start from an existing parameters file rather than writing one.** +`storm_analysis/test/data/` holds working, heavily commented XML named +`test__.xml`: `test_3d_2d_fixed.xml` is 3D-DAOSTORM in 2dfixed +mode, `test_sc_Z.xml` is sCMOS with the Z model, `test_spliner_2D.xml` is +Spliner. The fitting model is the `model` parameter — one of `2dfixed`, `2d`, +`3d` or `Z`. `doc/source/parameters.rst` documents every parameter. + +**Four of them describe the instrument and have to be changed.** These are +properties of the microscope and camera rather than tuning choices, and the +values in the test files belong to some other setup: + +- `pixel_size` — camera pixel size in nm +- `camera_gain` — ADU per photo-electron +- `camera_offset` — what the camera reads with the shutter closed +- `sigma` — initial PSF width guess in pixels. For `2dfixed` this has to be + close to right or the fit produces spurious double and triple peaks + +Nothing checks these against the data. Wrong values give a clean run and +localizations in the wrong units, so ask for them rather than inheriting them +from a sample file. + +**Parameter errors are reported at the top of the output, not the bottom.** A +missing or misspelled required parameter raises `ParametersException` naming +it. An unrecognized parameter prints + +``` +Warning!! is not a relevant parameter!! +``` + +as the very first line, ahead of several hundred lines of per-frame progress. +Read the head of the output, not the tail. + +**Spliner and Multiplane need a measured PSF first**, which means a bead z +stack plus two hand-made text files. Their formats are in +`doc/source/analysis.rst` and are worth reading rather than guessing — the z +offset file takes two columns, a valid flag and the z position, not one. + +Output is HDF5, described in `doc/source/output_files.rst` and read with +`sa_library/sa_h5py.py`. + + +## Diagnostics + +`storm_analysis/diagnostics/` holds end-to-end runs against simulated data. +They are the check that matters after changing a fitter or anything in the +analysis pipeline: the unit tests confirm the code runs, the diagnostics +confirm it still recovers the right positions. Read +`storm_analysis/diagnostics/README.txt` first — each one is a +`settings.py` / `configure.py` / `make_data.py` / `analyze_data.py` / +`collate.py` sequence run from a working directory. + +Two things worth knowing before you start one: + +- Set `STORM_ANALYSIS_HEADLESS=1` unless you are sitting in front of the + machine. Without it several of these open a plot window and block until you + close it. Figures written to disk are unaffected. +- `multiplane/configure.py` requires `--psf-model`, one of `psf_fft`, + `pupilfn` or `spline`. It has no default and exits 2 without it. This is the + only required argument in any of them. `fista_decon`, `spliner`, + `spliner_2d` and `multiplane` also accept an optional `--no-splines`; + everything else runs bare. + + +## Layout + +| Path | What's in it | +| --- | --- | +| `sa_library/` | Shared machinery: fitters (`dao_fit.c`, `multi_fit.c`), readers and writers, the HDF5 localization format | +| `sa_utilities/` | Pipeline stages: drift correction, tracking, format conversion | +| `spliner/` | Cubic-spline PSF fitting, and the tools that measure a spline from bead data | +| `multi_plane/`, `pupilfn/`, `psf_fft/` | Multi-plane and pupil-function fitting | +| `simulator/` | Synthetic movie generation, used heavily by the tests and diagnostics | +| `diagnostics/` | End-to-end accuracy runs (above) | +| `test/` | The pytest suite | +| `c_libraries/` | Build output. Not source — SCons writes here | + +Python talks to C through `ctypes`, and the wrappers are named for what they +wrap: `sa_library/dao_fit_c.py` is the interface to `sa_library/dao_fit.c`. A +change to a C signature has to be made in both places, and see the ctypes note +under Traps. + + +## Conventions + +These are consistent across the tree and are deliberate. Matching them matters +more than matching any external style guide. + +- **`import numpy`, never `import numpy as np`.** All 242 modules that import + numpy do it this way, and `numpy.` appears in full at every use. +- **Class methods are camelCase** — `getPeaks()`, `cleanUp()`. 704 of them, + against 6 in snake_case. Module-level functions are more mixed, and test + functions are snake_case because pytest collects on the `test_` prefix. +- **Line endings are LF**, enforced by `.gitattributes`. + +Where the tree is genuinely inconsistent, follow the file you are editing +rather than imposing a rule. Spacing around keyword arguments is the main one: +`f(x = 1)` and `f(x=1)` are close to evenly split overall, and 108 of 254 files +contain both. + + +## Traps + +**`decon_storm/` is not built.** It contains `.c` files, but `SConstruct` has +no entry for them and nothing in the package imports it — it is mostly MATLAB +kept for reference. Two consequences: changes there have no effect, and a +function definition found there is not the one that runs. A search that keys C +definitions by name and stops at the first match will find the dead copy and +hide the live one. + +**Stale bytecode can survive an edit.** Python decides whether a `.pyc` is +current from the source file's size and modification time, and mtime has +one-second resolution. An edit that leaves the size unchanged and lands within +the same second as the previous one keeps the old bytecode. This is a real +problem when comparing two versions of a small change: clear `__pycache__` +between the two runs or you will measure the same code twice. + +**The ctypes boundary is checked on one side only.** A wrapper declaring +`ndpointer(dtype=numpy.int32)` guarantees what Python passes; it says nothing +about what the C function expects. If the two disagree the result is a wrong +stride through valid memory, not an exception. C functions on this boundary +take `int32_t` rather than `int` for that reason. Scalars passed as +`ctypes.c_int` are fine as they are. diff --git a/doc/source/analysis.rst b/doc/source/analysis.rst index 5ea20c95..d4be859d 100644 --- a/doc/source/analysis.rst +++ b/doc/source/analysis.rst @@ -74,10 +74,32 @@ Measuring the PSF 1. Take a z stack of beads, ``beads_zcal.tif`` 2. Create a text file containing the locations of a few (isolated) - beads, ``beads_locs.txt`` + beads, ``beads_locs.txt``. This is one bead per line, as a space + separated ``x y`` pair in pixels. :: -3. Create a text file containing the z value of each frame in the z - stack movie, ``beads_zoffset.text`` + 128.5 64.0 + 201.0 187.5 + + .. note:: One way to make this file is to open the bead movie with + ``visualizer/visualizer.py`` and record the center position + of several well separated beads. + +3. Create a text file containing the z offsets, ``beads_zoffset.txt``. + This has two space separated columns per line, a valid flag and the z + offset in microns, with one line per frame of the z stack. Frames + whose valid flag is 0 are left out of the measurement. :: + + 1 -0.600 + 1 -0.550 + 0 -0.500 + 1 -0.450 + + .. note:: A single column of z values is not enough. It fails with + "Z offsets must have shape (N,2)". + + .. note:: If the z stack was acquired with storm-control then + ``spliner/offset_to_z.py`` will convert the ``.off`` file it + writes into this format. 4. Use measure_psf_beads.py to measure the PSF. :: @@ -235,8 +257,9 @@ Measuring the PSFs .. note:: Drift can be estimated with the program ``zstack_xydrift.py``. You will need to have found localizations in the first and last frame of the PSF calibration movie. -5. Create a text file containing the z offset of each frame of the PSF calibration movie. One - possibility is to use ``spliner/offset_to_z.py``. +5. Create a text file containing the z offsets of the PSF calibration movie. This is the same + two column format described under ``Spliner`` above, a valid flag and the z offset in + microns with one line per frame. One possibility is to use ``spliner/offset_to_z.py``. 6. Measure the PSF for each plane. :: diff --git a/storm_analysis/admm/admm_decon.py b/storm_analysis/admm/admm_decon.py index ae971ad6..d001cb7e 100644 --- a/storm_analysis/admm/admm_decon.py +++ b/storm_analysis/admm/admm_decon.py @@ -1,54 +1,54 @@ -#!/usr/bin/env python -""" -Uses ADMM to perform image deconvolution. - -This minimizes 1/2*|| Ax - b ||_2^2 + lambda || x ||_1. - -Hazen 11/19 -""" -import numpy - -import storm_analysis.sa_library.cs_decon as csDecon - -import storm_analysis.admm.admm_3d as admm3D -import storm_analysis.admm.admm_lasso_c as admmLassoC - - -class ADMMDecon(csDecon.CSDecon): - - def __init__(self, image_size, psf_object, number_zvals, rho): - super(ADMMDecon, self).__init__(image_size, psf_object, number_zvals) - - psfs = self.createPSFs() - - if False: - # Python solver (useful for debugging). - self.cs_solver = admm3D.ADMM(psfs, rho) - else: - # C solver (faster). - self.cs_solver = admmLassoC.ADMMLasso(psfs, rho) - - -# -# The MIT License -# -# Copyright (c) 2018 Babcock Lab, Harvard University -# -# Permission is hereby granted, free of charge, to any person obtaining a copy -# of this software and associated documentation files (the "Software"), to deal -# in the Software without restriction, including without limitation the rights -# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -# copies of the Software, and to permit persons to whom the Software is -# furnished to do so, subject to the following conditions: -# -# The above copyright notice and this permission notice shall be included in -# all copies or substantial portions of the Software. -# -# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN -# THE SOFTWARE. -# +#!/usr/bin/env python +""" +Uses ADMM to perform image deconvolution. + +This minimizes 1/2*|| Ax - b ||_2^2 + lambda || x ||_1. + +Hazen 11/19 +""" +import numpy + +import storm_analysis.sa_library.cs_decon as csDecon + +import storm_analysis.admm.admm_3d as admm3D +import storm_analysis.admm.admm_lasso_c as admmLassoC + + +class ADMMDecon(csDecon.CSDecon): + + def __init__(self, image_size, psf_object, number_zvals, rho): + super(ADMMDecon, self).__init__(image_size, psf_object, number_zvals) + + psfs = self.createPSFs() + + if False: + # Python solver (useful for debugging). + self.cs_solver = admm3D.ADMM(psfs, rho) + else: + # C solver (faster). + self.cs_solver = admmLassoC.ADMMLasso(psfs, rho) + + +# +# The MIT License +# +# Copyright (c) 2018 Babcock Lab, Harvard University +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. +# diff --git a/storm_analysis/admm/admm_lasso.c b/storm_analysis/admm/admm_lasso.c index af22117d..40f2b6db 100644 --- a/storm_analysis/admm/admm_lasso.c +++ b/storm_analysis/admm/admm_lasso.c @@ -1,505 +1,505 @@ -/* - * Use the ADMM approach to solve the lasso problem: - * - * minimize 1/2*|| Ax - b ||_2^2 + \lambda || x ||_1 - * - * As described in: - * Boyd et al., "Distributed Optimization and Statistical Learning - * via the Alternating Direction Method of Multipliers", Foundations - * and Trends in Machine Learning, 2010. - * - * Notes: - * 1. This has been specialized for the problem of image - * deconvolution and makes heavy use of the FFT. This - * program cannot be used to solve the generic lasso problem. - * - * Hazen 11/19 - */ - -#include -#include -#include - -#include - -#include "../sa_library/ft_math.h" - - -/* ADMM Structure. */ -typedef struct{ - int fft_size; - int image_size; - int number_psfs; - int stale_Ax; - - double normalization; - double rho; - - double *Ax; - double *fft_vector; - double *image; - - double **Atb; - double **u_vector; - double **x_vector; - double **z_vector; - - fftw_plan fft_backward; - fftw_plan fft_forward; - - fftw_complex *fft_vector_fft; - - fftw_complex **A; - fftw_complex **G_inv; - fftw_complex **work1; - -} admmData; - - -/* Functions Declarations. */ -void calculateAx(admmData *); -void cleanup(admmData *); -void getAx(admmData *, double *); -void getXVector(admmData *, double *); -admmData* initialize2D(double, int, int); -admmData* initialize3D(double, int, int, int); -void initializeA(admmData *, fftw_complex *, int); -void initializeGInv(admmData *, fftw_complex *, int); -double l1Error(admmData *); -double l2Error(admmData *); -void iterate(admmData *, double, int); -void newImage(admmData *, double *); -void run(admmData *, double, int); - - -/* Functions */ - -/* - * calculateAx() - * - * admm_data - A pointer to a admmData structure. - */ -void calculateAx(admmData *admm_data) -{ - int i; - - if (admm_data->stale_Ax){ - - /* Compute Ax. */ - ftmComplexZero(admm_data->work1[0], admm_data->fft_size); - - for(i=0;inumber_psfs;i++){ - - // Compute FFT of x vector for each z plane. - ftmDoubleCopy(admm_data->x_vector[i], admm_data->fft_vector, admm_data->image_size); - fftw_execute(admm_data->fft_forward); - - // Multiply FFT of x vector by FFT of the PSF for this z plane. - ftmComplexMultiplyAccum(admm_data->work1[0], admm_data->fft_vector_fft, admm_data->A[i], admm_data->fft_size, 0); - } - - ftmComplexCopy(admm_data->work1[0], admm_data->fft_vector_fft, admm_data->fft_size); - fftw_execute(admm_data->fft_backward); - - /* Store Ax. */ - ftmDoubleCopyNormalize(admm_data->fft_vector, admm_data->Ax, admm_data->normalization, admm_data->image_size); - } - - admm_data->stale_Ax = 0; -} - -/* - * cleanup - * - * admm_data - A pointer to a admmData structure. - */ -void cleanup(admmData *admm_data) -{ - int i; - - free(admm_data->Ax); - free(admm_data->image); - - for(i=0;inumber_psfs;i++){ - free(admm_data->Atb[i]); - free(admm_data->x_vector[i]); - free(admm_data->z_vector[i]); - free(admm_data->u_vector[i]); - } - free(admm_data->Atb); - free(admm_data->x_vector); - free(admm_data->z_vector); - free(admm_data->u_vector); - - fftw_destroy_plan(admm_data->fft_backward); - fftw_destroy_plan(admm_data->fft_forward); - - fftw_free(admm_data->fft_vector); - fftw_free(admm_data->fft_vector_fft); - - for(i=0;inumber_psfs;i++){ - free(admm_data->A[i]); - free(admm_data->work1[i]); - } - for(i=0;i<(admm_data->number_psfs*admm_data->number_psfs);i++){ - free(admm_data->G_inv[i]); - } - - free(admm_data->A); - free(admm_data->G_inv); - free(admm_data->work1); - - free(admm_data); -} - -/* - * getAx() - * - * Copies the current Ax vector into user supplied storage. - * - * admm_data - A pointer to a admmData structure. - * data - Storage for the Ax vector. - */ -void getAx(admmData *admm_data, double *data) -{ - calculateAx(admm_data); - ftmDoubleCopy(admm_data->Ax, data, admm_data->image_size); -} - -/* - * getXVector() - * - * Copies the current x vector into user supplied storage. Also - * converts back from [(x1,y1), (x2,y2), ..] to (x,y,z). - * - * admm_data - A pointer to a admmData structure. - * data - Storage for the x vector. - */ -void getXVector(admmData *admm_data, double *data) -{ - int i,j; - double *t1; - - for(i=0;inumber_psfs;i++){ - t1 = admm_data->x_vector[i]; - for(j=0;jimage_size;j++){ - data[j*admm_data->number_psfs+i] = t1[j]; - } - } -} - -/* - * initialize2D() - * - * Set things up for 2D analysis. - * - * rho - The rho parameter. - * x_size - Size in x of data and psf. - * y_size - Size in y of data and psf. - */ -admmData* initialize2D(double rho, int x_size, int y_size) -{ - return initialize3D(rho, x_size, y_size, 1); -} - -/* - * initialize3D() - * - * Set things up for 3D analysis. Note that setting the values in the - * A and G_inv arrays is done in a separate step. - * - * rho - The rho parameter. - * x_size - Size in x of data and psf. - * y_size - Size in y of data and psf. - * z_size - The number of z planes. - */ -admmData* initialize3D(double rho, int x_size, int y_size, int z_size) -{ - int i; - admmData *admm_data; - - admm_data = (admmData *)malloc(sizeof(admmData)); - - admm_data->fft_size = x_size * (y_size/2+1); - admm_data->image_size = x_size * y_size; - admm_data->number_psfs = z_size; - - admm_data->normalization = 1.0/((double)(admm_data->image_size)); - admm_data->rho = rho; - - /* Allocate storage. */ - admm_data->Ax = (double *)malloc(sizeof(double) * admm_data->image_size); - admm_data->image = (double *)malloc(sizeof(double) * admm_data->image_size); - - admm_data->Atb = (double **)malloc(sizeof(double *) * z_size); - admm_data->x_vector = (double **)malloc(sizeof(double *) * z_size); - admm_data->z_vector = (double **)malloc(sizeof(double *) * z_size); - admm_data->u_vector = (double **)malloc(sizeof(double *) * z_size); - - for(i=0;iAtb[i] = (double *)malloc(sizeof(double) * admm_data->image_size); - admm_data->x_vector[i] = (double *)malloc(sizeof(double) * admm_data->image_size); - admm_data->z_vector[i] = (double *)malloc(sizeof(double) * admm_data->image_size); - admm_data->u_vector[i] = (double *)malloc(sizeof(double) * admm_data->image_size); - } - - admm_data->fft_vector = (double *)fftw_malloc(sizeof(double) * admm_data->image_size); - admm_data->fft_vector_fft = (fftw_complex *)fftw_malloc(sizeof(fftw_complex) * admm_data->fft_size); - - admm_data->A = (fftw_complex **)malloc(sizeof(fftw_complex *) * z_size); - admm_data->G_inv = (fftw_complex **)malloc(sizeof(fftw_complex *) * z_size * z_size); - admm_data->work1 = (fftw_complex **)malloc(sizeof(fftw_complex *) * z_size); - - for(i=0;iA[i] = (fftw_complex *)malloc(sizeof(fftw_complex) * admm_data->fft_size); - admm_data->work1[i] = (fftw_complex *)malloc(sizeof(fftw_complex) * admm_data->fft_size); - } - - for(i=0;i<(z_size*z_size);i++){ - admm_data->G_inv[i] = (fftw_complex *)malloc(sizeof(fftw_complex) * admm_data->fft_size); - } - - /* Create FFT plans. */ - admm_data->fft_forward = fftw_plan_dft_r2c_2d(x_size, y_size, admm_data->fft_vector, admm_data->fft_vector_fft, FFTW_MEASURE); - admm_data->fft_backward = fftw_plan_dft_c2r_2d(x_size, y_size, admm_data->fft_vector_fft, admm_data->fft_vector, FFTW_MEASURE); - - return admm_data; -} - -/* - * initializeA() - * - * Set initial values in the A array. - * - * admm_data - A pointer to an admmData structure. - * data - Complex values to use for A. - * index - Which A matrix to store data in. - */ -void initializeA(admmData *admm_data, fftw_complex *data, int index) -{ - ftmComplexCopy(data, admm_data->A[index], admm_data->fft_size); -} - -/* - * initializeGInv() - * - * Set initial values in the G_inv array. - * - * admm_data - A pointer to an admmData structure. - * data - Complex values to use for G. - * index - Which G matrix to store data in. - */ -void initializeGInv(admmData *admm_data, fftw_complex *data, int index) -{ - ftmComplexCopy(data, admm_data->G_inv[index], admm_data->fft_size); -} - -/* - * iterate() - * - * Performs one cycle of improvement. - * - * lambda - The lambda value to use. - * pos_only - If 1, all values of z < lambda/rho are set to 0.0. - */ -void iterate(admmData *admm_data, double lambda, int pos_only) -{ - int i,j,n_psfs; - double lr, t1, t2; - double *Atb, *uv, *xv, *zv; - - /* Flag that we need to recaculate Ax (if getAx is called). */ - admm_data->stale_Ax = 1; - - n_psfs = admm_data->number_psfs; - - /* Calculate fft of (Atb + rho * (z - u)), store in work1. */ - for(i=0;iAtb[i]; - uv = admm_data->u_vector[i]; - zv = admm_data->z_vector[i]; - for(j=0;jimage_size;j++){ - admm_data->fft_vector[j] = Atb[j] + admm_data->rho * (zv[j] - uv[j]); - } - fftw_execute(admm_data->fft_forward); - ftmComplexCopy(admm_data->fft_vector_fft, admm_data->work1[i], admm_data->fft_size); - } - - /* Calculate x(k+1) by multiplying by G_inv. */ - for(i=0;ifft_vector_fft, admm_data->fft_size); - for(j=0;jfft_vector_fft, admm_data->G_inv[i*n_psfs+j], admm_data->work1[j], admm_data->fft_size, 0); - } - fftw_execute(admm_data->fft_backward); - ftmDoubleCopyNormalize(admm_data->fft_vector, admm_data->x_vector[i], admm_data->normalization, admm_data->image_size); - } - - /* Calculate z(k+1). */ - lr = lambda/admm_data->rho; - for(i=0;iu_vector[i]; - xv = admm_data->x_vector[i]; - zv = admm_data->z_vector[i]; - if (pos_only){ - for(j=0;jimage_size;j++){ - t1 = xv[j] + uv[j] - lr; - zv[j] = (t1 > 0.0) ? t1 : 0.0; - } - } - else{ - for(j=0;jimage_size;j++){ - t1 = xv[j] + uv[j]; - t2 = fabs(t1) - lr; - if (t2 <= 0.0){ - zv[j] = 0.0; - } - else { - zv[j] = (t1 > 0.0) ? t2 : -t2; - } - } - } - } - - /* Calculate u(k+1). */ - for(i=0;iu_vector[i]; - xv = admm_data->x_vector[i]; - zv = admm_data->z_vector[i]; - for(j=0;jimage_size;j++){ - uv[j] += xv[j] - zv[j]; - } - } -} - -/* - * l1Error - * - * Calculate l1 error term (sum of the x vector). - * - * admm_data - A pointer to a fistaData structure. - * - * Return sum of the x vector. - */ -double l1Error(admmData *admm_data) -{ - int i,j; - double sum; - double *xv; - - sum = 0.0; - for(i=0;inumber_psfs;i++){ - xv = admm_data->x_vector[i]; - for(j=0;jimage_size;j++){ - sum += fabs(xv[j]); - } - } - - return sum; -} - -/* - * l2Error - * - * Calculate l2 error (Ax - b). - * - * admm_data - A pointer to a admmData structure. - * - * Return the error in the fit. - */ -double l2Error(admmData *admm_data) -{ - int i; - double l2_error,t1; - - /* Compute Ax. */ - calculateAx(admm_data); - - /* Compute (Ax - b)^2. */ - l2_error = 0.0; - for(i=0;iimage_size;i++){ - t1 = admm_data->Ax[i] - admm_data->image[i]; - l2_error += t1*t1; - } - - return sqrt(l2_error); -} - -/* - * newImage() - * - * Sets things up for the analysis of a new image. - * - * image - The image data (with estimated background subtracted). - */ -void newImage(admmData *admm_data, double *image) -{ - int i; - - /* Flag that we need to recaculate Ax (if getAx is called). */ - admm_data->stale_Ax = 1; - - /* Save a copy of the image. */ - ftmDoubleCopy(image, admm_data->image, admm_data->image_size); - - /* Calculate Atb vector. */ - ftmDoubleCopy(image, admm_data->fft_vector, admm_data->image_size); - fftw_execute(admm_data->fft_forward); - ftmComplexCopy(admm_data->fft_vector_fft, admm_data->work1[0], admm_data->fft_size); - - for(i=0;inumber_psfs;i++){ - ftmComplexMultiply(admm_data->fft_vector_fft, admm_data->work1[0], admm_data->A[i], admm_data->fft_size, 1); - fftw_execute(admm_data->fft_backward); - ftmDoubleCopyNormalize(admm_data->fft_vector, admm_data->Atb[i], admm_data->normalization, admm_data->image_size); - } - - /* zero u, x, z vectors. */ - for(i=0;inumber_psfs;i++){ - ftmDoubleZero(admm_data->u_vector[i], admm_data->image_size); - ftmDoubleZero(admm_data->x_vector[i], admm_data->image_size); - ftmDoubleZero(admm_data->z_vector[i], admm_data->image_size); - } -} - -/* - * run() - * - * admm_data - A pointer to a admmData structure. - * lambda - Lambda value. - * cycles - Number of iterations to perform. - * - * Performs a fixed number of cycles at a fixed lambda. - */ -void run(admmData *admm_data, double lambda, int cycles) -{ - int i; - - for(i=0;i +#include +#include + +#include + +#include "../sa_library/ft_math.h" + + +/* ADMM Structure. */ +typedef struct{ + int fft_size; + int image_size; + int number_psfs; + int stale_Ax; + + double normalization; + double rho; + + double *Ax; + double *fft_vector; + double *image; + + double **Atb; + double **u_vector; + double **x_vector; + double **z_vector; + + fftw_plan fft_backward; + fftw_plan fft_forward; + + fftw_complex *fft_vector_fft; + + fftw_complex **A; + fftw_complex **G_inv; + fftw_complex **work1; + +} admmData; + + +/* Functions Declarations. */ +void calculateAx(admmData *); +void cleanup(admmData *); +void getAx(admmData *, double *); +void getXVector(admmData *, double *); +admmData* initialize2D(double, int, int); +admmData* initialize3D(double, int, int, int); +void initializeA(admmData *, fftw_complex *, int); +void initializeGInv(admmData *, fftw_complex *, int); +double l1Error(admmData *); +double l2Error(admmData *); +void iterate(admmData *, double, int); +void newImage(admmData *, double *); +void run(admmData *, double, int); + + +/* Functions */ + +/* + * calculateAx() + * + * admm_data - A pointer to a admmData structure. + */ +void calculateAx(admmData *admm_data) +{ + int i; + + if (admm_data->stale_Ax){ + + /* Compute Ax. */ + ftmComplexZero(admm_data->work1[0], admm_data->fft_size); + + for(i=0;inumber_psfs;i++){ + + // Compute FFT of x vector for each z plane. + ftmDoubleCopy(admm_data->x_vector[i], admm_data->fft_vector, admm_data->image_size); + fftw_execute(admm_data->fft_forward); + + // Multiply FFT of x vector by FFT of the PSF for this z plane. + ftmComplexMultiplyAccum(admm_data->work1[0], admm_data->fft_vector_fft, admm_data->A[i], admm_data->fft_size, 0); + } + + ftmComplexCopy(admm_data->work1[0], admm_data->fft_vector_fft, admm_data->fft_size); + fftw_execute(admm_data->fft_backward); + + /* Store Ax. */ + ftmDoubleCopyNormalize(admm_data->fft_vector, admm_data->Ax, admm_data->normalization, admm_data->image_size); + } + + admm_data->stale_Ax = 0; +} + +/* + * cleanup + * + * admm_data - A pointer to a admmData structure. + */ +void cleanup(admmData *admm_data) +{ + int i; + + free(admm_data->Ax); + free(admm_data->image); + + for(i=0;inumber_psfs;i++){ + free(admm_data->Atb[i]); + free(admm_data->x_vector[i]); + free(admm_data->z_vector[i]); + free(admm_data->u_vector[i]); + } + free(admm_data->Atb); + free(admm_data->x_vector); + free(admm_data->z_vector); + free(admm_data->u_vector); + + fftw_destroy_plan(admm_data->fft_backward); + fftw_destroy_plan(admm_data->fft_forward); + + fftw_free(admm_data->fft_vector); + fftw_free(admm_data->fft_vector_fft); + + for(i=0;inumber_psfs;i++){ + free(admm_data->A[i]); + free(admm_data->work1[i]); + } + for(i=0;i<(admm_data->number_psfs*admm_data->number_psfs);i++){ + free(admm_data->G_inv[i]); + } + + free(admm_data->A); + free(admm_data->G_inv); + free(admm_data->work1); + + free(admm_data); +} + +/* + * getAx() + * + * Copies the current Ax vector into user supplied storage. + * + * admm_data - A pointer to a admmData structure. + * data - Storage for the Ax vector. + */ +void getAx(admmData *admm_data, double *data) +{ + calculateAx(admm_data); + ftmDoubleCopy(admm_data->Ax, data, admm_data->image_size); +} + +/* + * getXVector() + * + * Copies the current x vector into user supplied storage. Also + * converts back from [(x1,y1), (x2,y2), ..] to (x,y,z). + * + * admm_data - A pointer to a admmData structure. + * data - Storage for the x vector. + */ +void getXVector(admmData *admm_data, double *data) +{ + int i,j; + double *t1; + + for(i=0;inumber_psfs;i++){ + t1 = admm_data->x_vector[i]; + for(j=0;jimage_size;j++){ + data[j*admm_data->number_psfs+i] = t1[j]; + } + } +} + +/* + * initialize2D() + * + * Set things up for 2D analysis. + * + * rho - The rho parameter. + * x_size - Size in x of data and psf. + * y_size - Size in y of data and psf. + */ +admmData* initialize2D(double rho, int x_size, int y_size) +{ + return initialize3D(rho, x_size, y_size, 1); +} + +/* + * initialize3D() + * + * Set things up for 3D analysis. Note that setting the values in the + * A and G_inv arrays is done in a separate step. + * + * rho - The rho parameter. + * x_size - Size in x of data and psf. + * y_size - Size in y of data and psf. + * z_size - The number of z planes. + */ +admmData* initialize3D(double rho, int x_size, int y_size, int z_size) +{ + int i; + admmData *admm_data; + + admm_data = (admmData *)malloc(sizeof(admmData)); + + admm_data->fft_size = x_size * (y_size/2+1); + admm_data->image_size = x_size * y_size; + admm_data->number_psfs = z_size; + + admm_data->normalization = 1.0/((double)(admm_data->image_size)); + admm_data->rho = rho; + + /* Allocate storage. */ + admm_data->Ax = (double *)malloc(sizeof(double) * admm_data->image_size); + admm_data->image = (double *)malloc(sizeof(double) * admm_data->image_size); + + admm_data->Atb = (double **)malloc(sizeof(double *) * z_size); + admm_data->x_vector = (double **)malloc(sizeof(double *) * z_size); + admm_data->z_vector = (double **)malloc(sizeof(double *) * z_size); + admm_data->u_vector = (double **)malloc(sizeof(double *) * z_size); + + for(i=0;iAtb[i] = (double *)malloc(sizeof(double) * admm_data->image_size); + admm_data->x_vector[i] = (double *)malloc(sizeof(double) * admm_data->image_size); + admm_data->z_vector[i] = (double *)malloc(sizeof(double) * admm_data->image_size); + admm_data->u_vector[i] = (double *)malloc(sizeof(double) * admm_data->image_size); + } + + admm_data->fft_vector = (double *)fftw_malloc(sizeof(double) * admm_data->image_size); + admm_data->fft_vector_fft = (fftw_complex *)fftw_malloc(sizeof(fftw_complex) * admm_data->fft_size); + + admm_data->A = (fftw_complex **)malloc(sizeof(fftw_complex *) * z_size); + admm_data->G_inv = (fftw_complex **)malloc(sizeof(fftw_complex *) * z_size * z_size); + admm_data->work1 = (fftw_complex **)malloc(sizeof(fftw_complex *) * z_size); + + for(i=0;iA[i] = (fftw_complex *)malloc(sizeof(fftw_complex) * admm_data->fft_size); + admm_data->work1[i] = (fftw_complex *)malloc(sizeof(fftw_complex) * admm_data->fft_size); + } + + for(i=0;i<(z_size*z_size);i++){ + admm_data->G_inv[i] = (fftw_complex *)malloc(sizeof(fftw_complex) * admm_data->fft_size); + } + + /* Create FFT plans. */ + admm_data->fft_forward = fftw_plan_dft_r2c_2d(x_size, y_size, admm_data->fft_vector, admm_data->fft_vector_fft, FFTW_MEASURE); + admm_data->fft_backward = fftw_plan_dft_c2r_2d(x_size, y_size, admm_data->fft_vector_fft, admm_data->fft_vector, FFTW_MEASURE); + + return admm_data; +} + +/* + * initializeA() + * + * Set initial values in the A array. + * + * admm_data - A pointer to an admmData structure. + * data - Complex values to use for A. + * index - Which A matrix to store data in. + */ +void initializeA(admmData *admm_data, fftw_complex *data, int index) +{ + ftmComplexCopy(data, admm_data->A[index], admm_data->fft_size); +} + +/* + * initializeGInv() + * + * Set initial values in the G_inv array. + * + * admm_data - A pointer to an admmData structure. + * data - Complex values to use for G. + * index - Which G matrix to store data in. + */ +void initializeGInv(admmData *admm_data, fftw_complex *data, int index) +{ + ftmComplexCopy(data, admm_data->G_inv[index], admm_data->fft_size); +} + +/* + * iterate() + * + * Performs one cycle of improvement. + * + * lambda - The lambda value to use. + * pos_only - If 1, all values of z < lambda/rho are set to 0.0. + */ +void iterate(admmData *admm_data, double lambda, int pos_only) +{ + int i,j,n_psfs; + double lr, t1, t2; + double *Atb, *uv, *xv, *zv; + + /* Flag that we need to recaculate Ax (if getAx is called). */ + admm_data->stale_Ax = 1; + + n_psfs = admm_data->number_psfs; + + /* Calculate fft of (Atb + rho * (z - u)), store in work1. */ + for(i=0;iAtb[i]; + uv = admm_data->u_vector[i]; + zv = admm_data->z_vector[i]; + for(j=0;jimage_size;j++){ + admm_data->fft_vector[j] = Atb[j] + admm_data->rho * (zv[j] - uv[j]); + } + fftw_execute(admm_data->fft_forward); + ftmComplexCopy(admm_data->fft_vector_fft, admm_data->work1[i], admm_data->fft_size); + } + + /* Calculate x(k+1) by multiplying by G_inv. */ + for(i=0;ifft_vector_fft, admm_data->fft_size); + for(j=0;jfft_vector_fft, admm_data->G_inv[i*n_psfs+j], admm_data->work1[j], admm_data->fft_size, 0); + } + fftw_execute(admm_data->fft_backward); + ftmDoubleCopyNormalize(admm_data->fft_vector, admm_data->x_vector[i], admm_data->normalization, admm_data->image_size); + } + + /* Calculate z(k+1). */ + lr = lambda/admm_data->rho; + for(i=0;iu_vector[i]; + xv = admm_data->x_vector[i]; + zv = admm_data->z_vector[i]; + if (pos_only){ + for(j=0;jimage_size;j++){ + t1 = xv[j] + uv[j] - lr; + zv[j] = (t1 > 0.0) ? t1 : 0.0; + } + } + else{ + for(j=0;jimage_size;j++){ + t1 = xv[j] + uv[j]; + t2 = fabs(t1) - lr; + if (t2 <= 0.0){ + zv[j] = 0.0; + } + else { + zv[j] = (t1 > 0.0) ? t2 : -t2; + } + } + } + } + + /* Calculate u(k+1). */ + for(i=0;iu_vector[i]; + xv = admm_data->x_vector[i]; + zv = admm_data->z_vector[i]; + for(j=0;jimage_size;j++){ + uv[j] += xv[j] - zv[j]; + } + } +} + +/* + * l1Error + * + * Calculate l1 error term (sum of the x vector). + * + * admm_data - A pointer to a fistaData structure. + * + * Return sum of the x vector. + */ +double l1Error(admmData *admm_data) +{ + int i,j; + double sum; + double *xv; + + sum = 0.0; + for(i=0;inumber_psfs;i++){ + xv = admm_data->x_vector[i]; + for(j=0;jimage_size;j++){ + sum += fabs(xv[j]); + } + } + + return sum; +} + +/* + * l2Error + * + * Calculate l2 error (Ax - b). + * + * admm_data - A pointer to a admmData structure. + * + * Return the error in the fit. + */ +double l2Error(admmData *admm_data) +{ + int i; + double l2_error,t1; + + /* Compute Ax. */ + calculateAx(admm_data); + + /* Compute (Ax - b)^2. */ + l2_error = 0.0; + for(i=0;iimage_size;i++){ + t1 = admm_data->Ax[i] - admm_data->image[i]; + l2_error += t1*t1; + } + + return sqrt(l2_error); +} + +/* + * newImage() + * + * Sets things up for the analysis of a new image. + * + * image - The image data (with estimated background subtracted). + */ +void newImage(admmData *admm_data, double *image) +{ + int i; + + /* Flag that we need to recaculate Ax (if getAx is called). */ + admm_data->stale_Ax = 1; + + /* Save a copy of the image. */ + ftmDoubleCopy(image, admm_data->image, admm_data->image_size); + + /* Calculate Atb vector. */ + ftmDoubleCopy(image, admm_data->fft_vector, admm_data->image_size); + fftw_execute(admm_data->fft_forward); + ftmComplexCopy(admm_data->fft_vector_fft, admm_data->work1[0], admm_data->fft_size); + + for(i=0;inumber_psfs;i++){ + ftmComplexMultiply(admm_data->fft_vector_fft, admm_data->work1[0], admm_data->A[i], admm_data->fft_size, 1); + fftw_execute(admm_data->fft_backward); + ftmDoubleCopyNormalize(admm_data->fft_vector, admm_data->Atb[i], admm_data->normalization, admm_data->image_size); + } + + /* zero u, x, z vectors. */ + for(i=0;inumber_psfs;i++){ + ftmDoubleZero(admm_data->u_vector[i], admm_data->image_size); + ftmDoubleZero(admm_data->x_vector[i], admm_data->image_size); + ftmDoubleZero(admm_data->z_vector[i], admm_data->image_size); + } +} + +/* + * run() + * + * admm_data - A pointer to a admmData structure. + * lambda - Lambda value. + * cycles - Number of iterations to perform. + * + * Performs a fixed number of cycles at a fixed lambda. + */ +void run(admmData *admm_data, double lambda, int cycles) +{ + int i; + + for(i=0;i= 0) and (xi < image_size[0]) and (yi >= 0) and (yi < image_size[1]): - red_image[xi,yi] += r - grn_image[xi,yi] += g - blu_image[xi,yi] += b - sum_image[xi,yi] += 1 - - # Some wacky normalization scheme.. - mask = (sum_image > image_max) - - red_image[mask] = (red_image[mask]/sum_image[mask]) * image_max - grn_image[mask] = (grn_image[mask]/sum_image[mask]) * image_max - blu_image[mask] = (blu_image[mask]/sum_image[mask]) * image_max - sum_image[mask] = image_max - - rgb_image = numpy.zeros((image_size[0], image_size[1], 3)) - rgb_image[:,:,0] = red_image - rgb_image[:,:,1] = grn_image - rgb_image[:,:,2] = blu_image - - rgb_image = (255.0/image_max)*rgb_image - sum_image = (255.0/image_max)*sum_image - - rgb_image = rgb_image.astype(numpy.uint8) - sum_image = sum_image.astype(numpy.uint8) - - return [rgb_image, sum_image, num_clusters] - - -def clusterImagesPNG(h5_name, png_base, title, min_size, image_max, scale): - """ - Saves the results from clusterImages as a PNG file. - - FIXME: This was a convenience function that I used a lot in one paper - but now it seems pretty specialized. Not sure it is worth - keeping in this project. - """ - [rgb_image, sum_image, num_clusters] = clusterImages(h5_name, min_size, image_max, scale = scale) - - # Save RGB image. - img1 = Image.fromarray(rgb_image, "RGB") - try: - draw1 = ImageDraw.Draw(img1) - font1 = ImageFont.truetype("FreeMono.ttf", 24) - draw1.text((2,2), title, (255,255,255), font = font1) - except IOError: - print("Text drawing disabled, true type font file may be missing?") - - img1.save(png_base + "_01.png") - - # Save grayscale image. - img2 = Image.fromarray(sum_image, "L") - try: - img2 = img2.convert("RGB") - draw2 = ImageDraw.Draw(img2) - font2 = ImageFont.truetype("FreeMono.ttf", 24) - draw2.text((2,2), title, (255,255,255), font = font2) - except IOError: - print("Text drawing disabled, true type font file may be missing?") - - img2.save(png_base + "_02.png") - - - -if (__name__ == "__main__"): - - import argparse - - parser = argparse.ArgumentParser(description = 'Make images of the identified clusters in a cell') - - parser.add_argument('--bin', dest='clsah5', type=str, required=True, - help = "The name of the clustered HDF5 format localization file.") - parser.add_argument('--image', dest='image', type=str, required=True, - help = "Root name of image file. Two images are generated, one with the clusters colored and one without.") - parser.add_argument('--title', dest='title', type=str, required=True, - help = "Title to add to the output images.") - parser.add_argument('--min_size', dest='min_size', type=int, required=True, - help = "The minimum cluster size to render.") - parser.add_argument('--im_max', dest='image_max', type=int, required=True, - help = "Image maximum, a good value for this is (usually) something in the range 6-20.") - parser.add_argument('--scale', dest='scale', type=int, required=False, default = 4, - help = "The 'zoom' of the output image (an integer).") - - args = parser.parse_args() - - clusterImages(args.clsah5, args.image, args.title, args.min_size, args.image_max, args.scale) - +#!/usr/bin/env python +""" +Make images of the clusters. + +Hazen 08/18 +""" +import numpy +from PIL import Image +from PIL import ImageFont +from PIL import ImageDraw +import randomcolor + +import storm_analysis.simulator.draw_gaussians_c as dg + +import storm_analysis.dbscan.clusters_sa_h5py as clSAH5Py + + +def clusterImages(h5_name, min_size, image_max, scale = 4, show_unclustered = True): + """ + Creates a RGB and an intensity array of the clusters with more than + min_size elements. + + h5_name - The name of the HDF5 file with clustering information. + min_size - The minimum size cluster to color. + image_max - Maximum value for color normalization. + scale - (int) Image rendering scale, default is 4. + show_unclustered - Include unclustered localizations (in white). + """ + rand_color = randomcolor.RandomColor() + + with clSAH5Py.SAH5Clusters(h5_name) as cl_h5: + [movie_x, movie_y] = cl_h5.getMovieInformation()[:2] + + image_size = [movie_y * scale, movie_x * scale] + + red_image = numpy.zeros(image_size) + grn_image = numpy.zeros(image_size) + blu_image = numpy.zeros(image_size) + sum_image = numpy.zeros(image_size) + + num_clusters = 0 + for index, cluster in cl_h5.clustersIterator(skip_unclustered = False, fields = ["x", "y"]): + if ((index%200)==0): + print("Processing cluster", index) + + x = scale * cluster['x'] + y = scale * cluster['y'] + + if (index == 0) or (x.size < min_size): + + # Skip to next cluster if this cluster is too small and we are + # not supposed to draw the unclustered localizations. + if not show_unclustered: + continue + + else: + r = 1.0 + g = 1.0 + b = 1.0 + else: + num_clusters += 1 + color = rand_color.generate() + r = int(color[0][1:3], 16)/255.0 + g = int(color[0][3:5], 16)/255.0 + b = int(color[0][5:7], 16)/255.0 + + # FIXME: This would draw the image clusters as Gaussians. Not sure + # how well it actually works. Sigma should also probably be + # a user argument. + # + if False: + temp = numpy.zeros(image_size) + dg.drawGaussiansXYOnImage(temp, y, x, sigma = 1.5) + + red_image += r * temp + grn_image += g * temp + blu_image += b * temp + sum_image += temp + + else: + for i in range(x.size): + yi = int(x[i]) + xi = int(y[i]) + if (xi >= 0) and (xi < image_size[0]) and (yi >= 0) and (yi < image_size[1]): + red_image[xi,yi] += r + grn_image[xi,yi] += g + blu_image[xi,yi] += b + sum_image[xi,yi] += 1 + + # Some wacky normalization scheme.. + mask = (sum_image > image_max) + + red_image[mask] = (red_image[mask]/sum_image[mask]) * image_max + grn_image[mask] = (grn_image[mask]/sum_image[mask]) * image_max + blu_image[mask] = (blu_image[mask]/sum_image[mask]) * image_max + sum_image[mask] = image_max + + rgb_image = numpy.zeros((image_size[0], image_size[1], 3)) + rgb_image[:,:,0] = red_image + rgb_image[:,:,1] = grn_image + rgb_image[:,:,2] = blu_image + + rgb_image = (255.0/image_max)*rgb_image + sum_image = (255.0/image_max)*sum_image + + rgb_image = rgb_image.astype(numpy.uint8) + sum_image = sum_image.astype(numpy.uint8) + + return [rgb_image, sum_image, num_clusters] + + +def clusterImagesPNG(h5_name, png_base, title, min_size, image_max, scale): + """ + Saves the results from clusterImages as a PNG file. + + FIXME: This was a convenience function that I used a lot in one paper + but now it seems pretty specialized. Not sure it is worth + keeping in this project. + """ + [rgb_image, sum_image, num_clusters] = clusterImages(h5_name, min_size, image_max, scale = scale) + + # Save RGB image. + img1 = Image.fromarray(rgb_image, "RGB") + try: + draw1 = ImageDraw.Draw(img1) + font1 = ImageFont.truetype("FreeMono.ttf", 24) + draw1.text((2,2), title, (255,255,255), font = font1) + except IOError: + print("Text drawing disabled, true type font file may be missing?") + + img1.save(png_base + "_01.png") + + # Save grayscale image. + img2 = Image.fromarray(sum_image, "L") + try: + img2 = img2.convert("RGB") + draw2 = ImageDraw.Draw(img2) + font2 = ImageFont.truetype("FreeMono.ttf", 24) + draw2.text((2,2), title, (255,255,255), font = font2) + except IOError: + print("Text drawing disabled, true type font file may be missing?") + + img2.save(png_base + "_02.png") + + + +if (__name__ == "__main__"): + + import argparse + + parser = argparse.ArgumentParser(description = 'Make images of the identified clusters in a cell') + + parser.add_argument('--bin', dest='clsah5', type=str, required=True, + help = "The name of the clustered HDF5 format localization file.") + parser.add_argument('--image', dest='image', type=str, required=True, + help = "Root name of image file. Two images are generated, one with the clusters colored and one without.") + parser.add_argument('--title', dest='title', type=str, required=True, + help = "Title to add to the output images.") + parser.add_argument('--min_size', dest='min_size', type=int, required=True, + help = "The minimum cluster size to render.") + parser.add_argument('--im_max', dest='image_max', type=int, required=True, + help = "Image maximum, a good value for this is (usually) something in the range 6-20.") + parser.add_argument('--scale', dest='scale', type=int, required=False, default = 4, + help = "The 'zoom' of the output image (an integer).") + + args = parser.parse_args() + + clusterImages(args.clsah5, args.image, args.title, args.min_size, args.image_max, args.scale) + diff --git a/storm_analysis/dbscan/dbscan.c b/storm_analysis/dbscan/dbscan.c index 21fe3011..f2700a7f 100644 --- a/storm_analysis/dbscan/dbscan.c +++ b/storm_analysis/dbscan/dbscan.c @@ -1,582 +1,582 @@ -/* - * A C implementation of the DBSCAN algorithm as described here: - * http://en.wikipedia.org/wiki/DBSCAN - * - * 2011/11 - * - * - * Modified to use a kdtree for region queries. This is way way - * faster than the simple search algorithm I was using before. - * - * kdtree code from here: - * https://github.com/jtsiomb/kdtree - * - * 2011/11 - * - * - * Added some additional functions that are somewhat related - * to clustering. - * - * 2011/11 - * - * Why do we sometimes get clusters that are smaller than mc? - * Quoting Wikipedia, "DBSCAN is not entirely deterministic: - * border points that are reachable from more than one cluster - * can be part of either cluster, depending on the order the - * data are processed." So this can happen if there are enough - * points to create a cluster but the some of the points already - * belong to another cluster. - * - * 2018/09 - * - * Hazen - */ - -/* Includes */ -#include -#include -#include - -#include - -#include "kdtree.h" - -/* Defines */ -#define TESTING 0 -#define USEKD 1 - -#define NOISE -1 -#define UNVISITED 0 -#define VISITED 1 - -#define MAXSIZE 10000 // maximum cluster size. - - -/* Structures */ -typedef struct -{ - int end; // location of the last non-zero element in the array. - int size; // total size of the array. - int *data; // storage for array data. -} intArr; - - -/* Function Declarations */ - -// dbscan specific -void appendN(intArr *, intArr *); -intArr* createIntArr(int); -void dbscan(float *, float *, float *, int32_t *, int32_t *, int, float, int, int); -int expandCluster(intArr *, int, float, int, int); -void mergeN(intArr *, intArr *); -intArr* regionQuery(int *, int, int, int, float); -intArr* regionQueryKD(int *, int, int, int, float); - -// additional -void clusterSize(int *, int *, int, int); -void locClSize(int32_t *, int32_t *, int, int); -void recategorize(int32_t *, int32_t *, int, int, int); - - -/* Global Variables */ -int nsize; // size of c,l,x,y,z. -int *c; // localization category. -int *l; // localization label. -float *x; // localization x location (in nm). -float *y; // localization y location (in nm). -float *z; // localization z location (in nm). - -void *kd; // kdtree structure pointer. - - -/* Code */ - - -/* - * appendN - * - * Append the contents of the second intArr to the - * the end of the first intArr. - * - * N - first intArr. - * Np - second intArr. - */ -void appendN(intArr *N, intArr *Np) -{ - int i,n_end; - int *n_data, *np_data; - - // initializations - n_end = N->end; - n_data = N->data; - np_data = Np->data; - - // check size - if((N->end+Np->end)>N->size){ - printf("N reached maxsize! %d (appendN)\n", N->size); - return; - } - - // append data - for(i=0;i<(Np->end);i++){ - n_data[n_end] = np_data[i]; - n_end++; - } - N->end = n_end; -} - - -/* - * createIntArr - * - * Allocates and initializes a intArr structure. - * - * n - size of the array. - */ -intArr* createIntArr(int n) -{ - intArr *ia; - - if(TESTING){ - printf("createIntArr\n"); - } - - ia = (intArr *)malloc(sizeof(intArr)); - ia->end = 0; - ia->size = n; - ia->data = (int *)malloc(sizeof(int)*n); - - return ia; -} - - -/* - * dbscan - * - * Performs the DBSCAN algorithm to identify clusters. - * - * db_x - array of x locations. - * db_y - array of y locations. - * db_z - array of z locations. - * db_cat - array localization categories. - * db_l - pre-allocated storage for the cluster label. - * this should be the same size as x,y and z. - * db_nsize - size of x, y, z and l. - * eps - cluster size parameter (in nm). - * min_points - minimum number of points in a cluster. - * verbose - print cluster information as we go. - */ -void dbscan(float *db_x, float *db_y, float *db_z, int32_t *db_cat, int32_t *db_l, int db_nsize, float eps, int min_points, int verbose) -{ - int i,cluster_size,cn,counts,cur_cat; - intArr *N; - - if(TESTING){ - printf("dbscan\n"); - } - - // initializations. - nsize = db_nsize; - l = db_l; - x = db_x; - y = db_y; - z = db_z; - c = db_cat; - cn = 2; - if(!USEKD){ - eps = eps*eps; - } - - if(USEKD){ - kd = kd_create(3); - for(i=0;i=min_points){ - l[i] = cn; - cluster_size = expandCluster(N,cur_cat,eps,min_points,cn); - if(verbose){ - printf(" Cluster %d contains %d localizations\n", cn, cluster_size); - } - cn++; - } - else{ - l[i] = NOISE; - } - free(N->data); - free(N); - } - } - - if(USEKD){ - kd_free(kd); - } -} - - -/* - * expandCluster - * - * Searches the neighbors of the current cluster to - * to see if they also have enough neighbors to merit - * additional analysis. - * - * N - array of not yet visited points in the cluster. - * cur_cat - localization category. - * eps - cluster size parameter (in nm). - * min_points - minimum number of neighbors. - * cn - current cluster number. - * - * returns the cluster size. - */ -int expandCluster(intArr *N, int cur_cat, float eps, int min_points, int cn) -{ - int i,cluster_size,counts; - int *data; - intArr *Np; - - if(TESTING){ - printf("expandCluster\n"); - } - - cluster_size = 1; - data = N->data; - while((N->end)>0){ - - // Mark current position with current cluster number. - i = data[0]; - l[i] = cn; - - // Find the neighbors of the current position - // if it has not already been determined to - // have too few neighbors. - if(USEKD){ - Np = regionQueryKD(&counts,i,min_points,cur_cat,eps); - } - else{ - Np = regionQuery(&counts,i,min_points,cur_cat,eps); - } - - - // If there are enough neighbors, merge N and Np. - if(counts>=min_points){ - //mergeN(N,Np); - appendN(N,Np); - } - - // Free Np. - free(Np->data); - free(Np); - - // Remove current position from the N array. - for(i=0;i<(N->end);i++){ - data[i]=data[i+1]; - } - N->end--; - cluster_size++; - } - - return cluster_size; -} - - -/* - * mergeN - * - * Merge the contents of a second intArr into the first - * ignoring duplicates. - * - * N - first intArr. - * Np - second intArr. - */ -void mergeN(intArr *N, intArr *Np) -{ - int i,j,cur,n_end,found; - int *n_data, *np_data; - - if(TESTING){ - printf("mergeN\n"); - } - - n_end = N->end; - n_data = N->data; - np_data = Np->data; - - if(TESTING){ - printf(" N size %d\n", N->end); - for(i=0;i<(N->end);i++){ - printf(" N %d %d\n", i, n_data[i]); - } - printf("\n"); - printf(" Np size %d\n", Np->end); - for(i=0;i<(Np->end);i++){ - printf(" Np %d %d\n", i, np_data[i]); - } - printf("\n"); - } - - for(i=0;i<(Np->end);i++){ - cur = np_data[i]; - found = 0; - for(j=0;jend] = cur; - N->end++; - if((N->end)>(N->size)){ - printf("N reached maxsize! %d (mergeN)\n", N->size); - break; - } - } - } - - if(TESTING){ - printf(" N size %d\n", N->end); - for(i=0;i<(N->end);i++){ - printf(" N %d %d\n", i, n_data[i]); - } - printf("\n"); - } -} - - -/* - * regionQuery - * - * Counts how many neighbors surrounding the current point are - * within eps of the current point. Returns a intArr structure - * containing all the unvisited neighbors of the current point. - * - * counts - returns the total number of neighbors. - * index - index of the current point. - * min_points - minimum cluster size. - * cur_cat - localization category. - * eps - cluster size parameter (in nm). - */ -intArr *regionQuery(int *counts, int index, int min_points, int cur_cat, float eps) -{ - int i,cnt; - float cx,cy,cz,dx,dy,dz; - intArr *N; - - if(TESTING){ - printf("regionQuery\n"); - } - - // initialization. - cnt = 0; - cx = x[index]; - cy = y[index]; - cz = z[index]; - N = createIntArr(MAXSIZE); - - // loop over all the localizations. - for(i=0;idata[N->end] = i; - N->end++; - if((N->end)==(N->size)){ - printf("N reached maxsize! %d (regionQuery)\n", N->size); - break; - } - l[i] = VISITED; - } - } - } - } - - *counts = cnt; - return N; -} - - -/* - * regionQueryKD - * - * Counts how many neighbors surrounding the current point are - * within eps of the current point. Returns a intArr structure - * containing all the unvisited neighbors of the current point. - * Uses a kd-tree to quickly find the points neighbors. - * - * counts - returns the total number of neighbors. - * index - index of the current point. - * min_points - minimum cluster size. - * cur_cat - localization category. - * eps - cluster size parameter (in nm). - */ -intArr *regionQueryKD(int *counts, int index, int min_points, int cur_cat, float eps) -{ - int i,j,cnt; - void *set; - intArr *N; - - if(TESTING){ - printf("regionQuery\n"); - } - - cnt = 0; - N = createIntArr(MAXSIZE); - - // find the set of localizations within eps. - set = kd_nearest_range3f(kd, x[index], y[index], z[index], eps); - // printf("rqkd: %d %d %d %d\n", index, kd_res_size(set), c[index], cur_cat); - - // loop over the set to populate N. - for(i=0;idata[N->end] = j; - N->end++; - if((N->end)==(N->size)){ - printf("N reached maxsize! %d (regionQuery)\n", N->size); - break; - } - //l[j] = VISITED; - } - } - kd_res_next(set); - } - kd_res_free(set); - - // if cnts is larger than the min cluster size - // mark all members of N as visited. - // - // FIXME: Marks noise points as visited, which - // is not optimal since this means they - // will be visited twice. - if(cnt>=min_points){ - for(i=0;i<(N->end);i++){ - l[N->data[i]] = VISITED; - } - } - - *counts = cnt; - return N; -} - - -/********************************************************************/ - -/* - * clusterSize - * - * Returns the size of each cluster. - * - * The fiddly adding 2 to max_id & 1 to the cluster id - * is due to cluster id starting at -1 and going to max_id. - * - * counts - preallocated storage for cluster size. - * cl_id - array of molecule cluster ids. - * cl_size - size of cl_id. - * max_id - maximum cluster id number (and size of cl_cnts). - * - */ -void clusterSize(int *counts, int *cl_id, int cl_size, int max_id) -{ - int i; - - for(i=0;i +#include +#include + +#include + +#include "kdtree.h" + +/* Defines */ +#define TESTING 0 +#define USEKD 1 + +#define NOISE -1 +#define UNVISITED 0 +#define VISITED 1 + +#define MAXSIZE 10000 // maximum cluster size. + + +/* Structures */ +typedef struct +{ + int end; // location of the last non-zero element in the array. + int size; // total size of the array. + int *data; // storage for array data. +} intArr; + + +/* Function Declarations */ + +// dbscan specific +void appendN(intArr *, intArr *); +intArr* createIntArr(int); +void dbscan(float *, float *, float *, int32_t *, int32_t *, int, float, int, int); +int expandCluster(intArr *, int, float, int, int); +void mergeN(intArr *, intArr *); +intArr* regionQuery(int *, int, int, int, float); +intArr* regionQueryKD(int *, int, int, int, float); + +// additional +void clusterSize(int *, int *, int, int); +void locClSize(int32_t *, int32_t *, int, int); +void recategorize(int32_t *, int32_t *, int, int, int); + + +/* Global Variables */ +int nsize; // size of c,l,x,y,z. +int *c; // localization category. +int *l; // localization label. +float *x; // localization x location (in nm). +float *y; // localization y location (in nm). +float *z; // localization z location (in nm). + +void *kd; // kdtree structure pointer. + + +/* Code */ + + +/* + * appendN + * + * Append the contents of the second intArr to the + * the end of the first intArr. + * + * N - first intArr. + * Np - second intArr. + */ +void appendN(intArr *N, intArr *Np) +{ + int i,n_end; + int *n_data, *np_data; + + // initializations + n_end = N->end; + n_data = N->data; + np_data = Np->data; + + // check size + if((N->end+Np->end)>N->size){ + printf("N reached maxsize! %d (appendN)\n", N->size); + return; + } + + // append data + for(i=0;i<(Np->end);i++){ + n_data[n_end] = np_data[i]; + n_end++; + } + N->end = n_end; +} + + +/* + * createIntArr + * + * Allocates and initializes a intArr structure. + * + * n - size of the array. + */ +intArr* createIntArr(int n) +{ + intArr *ia; + + if(TESTING){ + printf("createIntArr\n"); + } + + ia = (intArr *)malloc(sizeof(intArr)); + ia->end = 0; + ia->size = n; + ia->data = (int *)malloc(sizeof(int)*n); + + return ia; +} + + +/* + * dbscan + * + * Performs the DBSCAN algorithm to identify clusters. + * + * db_x - array of x locations. + * db_y - array of y locations. + * db_z - array of z locations. + * db_cat - array localization categories. + * db_l - pre-allocated storage for the cluster label. + * this should be the same size as x,y and z. + * db_nsize - size of x, y, z and l. + * eps - cluster size parameter (in nm). + * min_points - minimum number of points in a cluster. + * verbose - print cluster information as we go. + */ +void dbscan(float *db_x, float *db_y, float *db_z, int32_t *db_cat, int32_t *db_l, int db_nsize, float eps, int min_points, int verbose) +{ + int i,cluster_size,cn,counts,cur_cat; + intArr *N; + + if(TESTING){ + printf("dbscan\n"); + } + + // initializations. + nsize = db_nsize; + l = db_l; + x = db_x; + y = db_y; + z = db_z; + c = db_cat; + cn = 2; + if(!USEKD){ + eps = eps*eps; + } + + if(USEKD){ + kd = kd_create(3); + for(i=0;i=min_points){ + l[i] = cn; + cluster_size = expandCluster(N,cur_cat,eps,min_points,cn); + if(verbose){ + printf(" Cluster %d contains %d localizations\n", cn, cluster_size); + } + cn++; + } + else{ + l[i] = NOISE; + } + free(N->data); + free(N); + } + } + + if(USEKD){ + kd_free(kd); + } +} + + +/* + * expandCluster + * + * Searches the neighbors of the current cluster to + * to see if they also have enough neighbors to merit + * additional analysis. + * + * N - array of not yet visited points in the cluster. + * cur_cat - localization category. + * eps - cluster size parameter (in nm). + * min_points - minimum number of neighbors. + * cn - current cluster number. + * + * returns the cluster size. + */ +int expandCluster(intArr *N, int cur_cat, float eps, int min_points, int cn) +{ + int i,cluster_size,counts; + int *data; + intArr *Np; + + if(TESTING){ + printf("expandCluster\n"); + } + + cluster_size = 1; + data = N->data; + while((N->end)>0){ + + // Mark current position with current cluster number. + i = data[0]; + l[i] = cn; + + // Find the neighbors of the current position + // if it has not already been determined to + // have too few neighbors. + if(USEKD){ + Np = regionQueryKD(&counts,i,min_points,cur_cat,eps); + } + else{ + Np = regionQuery(&counts,i,min_points,cur_cat,eps); + } + + + // If there are enough neighbors, merge N and Np. + if(counts>=min_points){ + //mergeN(N,Np); + appendN(N,Np); + } + + // Free Np. + free(Np->data); + free(Np); + + // Remove current position from the N array. + for(i=0;i<(N->end);i++){ + data[i]=data[i+1]; + } + N->end--; + cluster_size++; + } + + return cluster_size; +} + + +/* + * mergeN + * + * Merge the contents of a second intArr into the first + * ignoring duplicates. + * + * N - first intArr. + * Np - second intArr. + */ +void mergeN(intArr *N, intArr *Np) +{ + int i,j,cur,n_end,found; + int *n_data, *np_data; + + if(TESTING){ + printf("mergeN\n"); + } + + n_end = N->end; + n_data = N->data; + np_data = Np->data; + + if(TESTING){ + printf(" N size %d\n", N->end); + for(i=0;i<(N->end);i++){ + printf(" N %d %d\n", i, n_data[i]); + } + printf("\n"); + printf(" Np size %d\n", Np->end); + for(i=0;i<(Np->end);i++){ + printf(" Np %d %d\n", i, np_data[i]); + } + printf("\n"); + } + + for(i=0;i<(Np->end);i++){ + cur = np_data[i]; + found = 0; + for(j=0;jend] = cur; + N->end++; + if((N->end)>(N->size)){ + printf("N reached maxsize! %d (mergeN)\n", N->size); + break; + } + } + } + + if(TESTING){ + printf(" N size %d\n", N->end); + for(i=0;i<(N->end);i++){ + printf(" N %d %d\n", i, n_data[i]); + } + printf("\n"); + } +} + + +/* + * regionQuery + * + * Counts how many neighbors surrounding the current point are + * within eps of the current point. Returns a intArr structure + * containing all the unvisited neighbors of the current point. + * + * counts - returns the total number of neighbors. + * index - index of the current point. + * min_points - minimum cluster size. + * cur_cat - localization category. + * eps - cluster size parameter (in nm). + */ +intArr *regionQuery(int *counts, int index, int min_points, int cur_cat, float eps) +{ + int i,cnt; + float cx,cy,cz,dx,dy,dz; + intArr *N; + + if(TESTING){ + printf("regionQuery\n"); + } + + // initialization. + cnt = 0; + cx = x[index]; + cy = y[index]; + cz = z[index]; + N = createIntArr(MAXSIZE); + + // loop over all the localizations. + for(i=0;idata[N->end] = i; + N->end++; + if((N->end)==(N->size)){ + printf("N reached maxsize! %d (regionQuery)\n", N->size); + break; + } + l[i] = VISITED; + } + } + } + } + + *counts = cnt; + return N; +} + + +/* + * regionQueryKD + * + * Counts how many neighbors surrounding the current point are + * within eps of the current point. Returns a intArr structure + * containing all the unvisited neighbors of the current point. + * Uses a kd-tree to quickly find the points neighbors. + * + * counts - returns the total number of neighbors. + * index - index of the current point. + * min_points - minimum cluster size. + * cur_cat - localization category. + * eps - cluster size parameter (in nm). + */ +intArr *regionQueryKD(int *counts, int index, int min_points, int cur_cat, float eps) +{ + int i,j,cnt; + void *set; + intArr *N; + + if(TESTING){ + printf("regionQuery\n"); + } + + cnt = 0; + N = createIntArr(MAXSIZE); + + // find the set of localizations within eps. + set = kd_nearest_range3f(kd, x[index], y[index], z[index], eps); + // printf("rqkd: %d %d %d %d\n", index, kd_res_size(set), c[index], cur_cat); + + // loop over the set to populate N. + for(i=0;idata[N->end] = j; + N->end++; + if((N->end)==(N->size)){ + printf("N reached maxsize! %d (regionQuery)\n", N->size); + break; + } + //l[j] = VISITED; + } + } + kd_res_next(set); + } + kd_res_free(set); + + // if cnts is larger than the min cluster size + // mark all members of N as visited. + // + // FIXME: Marks noise points as visited, which + // is not optimal since this means they + // will be visited twice. + if(cnt>=min_points){ + for(i=0;i<(N->end);i++){ + l[N->data[i]] = VISITED; + } + } + + *counts = cnt; + return N; +} + + +/********************************************************************/ + +/* + * clusterSize + * + * Returns the size of each cluster. + * + * The fiddly adding 2 to max_id & 1 to the cluster id + * is due to cluster id starting at -1 and going to max_id. + * + * counts - preallocated storage for cluster size. + * cl_id - array of molecule cluster ids. + * cl_size - size of cl_id. + * max_id - maximum cluster id number (and size of cl_cnts). + * + */ +void clusterSize(int *counts, int *cl_id, int cl_size, int max_id) +{ + int i; + + for(i=0;i -#include -#include - -#include - -#include "../sa_library/ft_math.h" - - -/* 3denseSTORM Structure. */ -typedef struct{ - int fft_size; - int image_size; - int number_psfs; - int stale_Ax; - - double beta; - double eta; - double micro; - double normalization; - - double *Ax; - double *b_vec; - double *d_vec; - double *fft_vector; - double *image; - double *w_vec; - double *y_vec; - double *y_vec_h; - - double **e_vec; - double **work1; - double **x_vec_h; - double **x_vec_t; - - fftw_plan fft_backward; - fftw_plan fft_forward; - - fftw_complex *fft_vector_fft; - - fftw_complex **A; - fftw_complex **G_inv; - fftw_complex **work2; - -} dsData; - - -/* Functions Declarations. */ -void calculateAx(dsData *); -void cleanup(dsData *); -void getAx(dsData *, double *); -void getXVector(dsData *, double *, int); -dsData* initialize2D(double, double, double, int, int); -dsData* initialize3D(double, double, double, int, int, int); -void initializeA(dsData *, fftw_complex *, int); -void initializeGInv(dsData *, fftw_complex *, int); -void iterate(dsData *); -double l1Error(dsData *); -double l2Error(dsData *); -void newImage(dsData *, double *, double *); -void run(dsData *, int); - - -/* Functions */ - -/* - * calculateAx() - * - * ds_data - A pointer to a dsData structure. - */ -void calculateAx(dsData *ds_data) -{ - int i; - - if (ds_data->stale_Ax){ - - /* Compute Ax. */ - ftmComplexZero(ds_data->work2[0], ds_data->fft_size); - - for(i=0;inumber_psfs;i++){ - - // Compute FFT of x vector for each z plane. - ftmDoubleCopy(ds_data->x_vec_h[i], ds_data->fft_vector, ds_data->image_size); - fftw_execute(ds_data->fft_forward); - - // Multiply FFT of x vector by FFT of the PSF for this z plane. - ftmComplexMultiplyAccum(ds_data->work2[0], ds_data->fft_vector_fft, ds_data->A[i], ds_data->fft_size, 0); - } - - ftmComplexCopy(ds_data->work2[0], ds_data->fft_vector_fft, ds_data->fft_size); - fftw_execute(ds_data->fft_backward); - - /* Store Ax. */ - ftmDoubleCopyNormalize(ds_data->fft_vector, ds_data->Ax, ds_data->normalization, ds_data->image_size); - } - - ds_data->stale_Ax = 0; -} - -/* - * cleanup - * - * ds_data - A pointer to a dsData structure. - */ -void cleanup(dsData *ds_data) -{ - int i; - - free(ds_data->Ax); - free(ds_data->b_vec); - free(ds_data->d_vec); - free(ds_data->image); - free(ds_data->w_vec); - free(ds_data->y_vec); - free(ds_data->y_vec_h); - - for(i=0;inumber_psfs;i++){ - free(ds_data->e_vec[i]); - free(ds_data->work1[i]); - free(ds_data->x_vec_h[i]); - free(ds_data->x_vec_t[i]); - } - free(ds_data->e_vec); - free(ds_data->work1); - free(ds_data->x_vec_h); - free(ds_data->x_vec_t); - - fftw_destroy_plan(ds_data->fft_backward); - fftw_destroy_plan(ds_data->fft_forward); - - fftw_free(ds_data->fft_vector); - fftw_free(ds_data->fft_vector_fft); - - for(i=0;inumber_psfs;i++){ - free(ds_data->A[i]); - free(ds_data->work2[i]); - } - for(i=0;i<(ds_data->number_psfs*ds_data->number_psfs);i++){ - free(ds_data->G_inv[i]); - } - - free(ds_data->A); - free(ds_data->G_inv); - free(ds_data->work2); - - free(ds_data); -} - -/* - * getAx() - * - * Copies the current Ax vector into user supplied storage. - * - * ds_data - A pointer to a dsData structure. - * data - Storage for the Ax vector. - */ -void getAx(dsData *ds_data, double *data) -{ - calculateAx(ds_data); - ftmDoubleCopy(ds_data->Ax, data, ds_data->image_size); -} - -/* - * getXVector() - * - * Copies the current x vector into user supplied storage. Also - * converts back from [(x1,y1), (x2,y2), ..] to (x,y,z). - * - * ds_data - A pointer to a dsData structure. - * data - Storage for the x vector. - * compressed - Return x_vec_t (1) or x_vec_h (0). - */ -void getXVector(dsData *ds_data, double *data, int compressed) -{ - int i,j; - double *t1; - - for(i=0;inumber_psfs;i++){ - if (compressed){ - t1 = ds_data->x_vec_t[i]; - } - else{ - t1 = ds_data->x_vec_h[i]; - } - for(j=0;jimage_size;j++){ - data[j*ds_data->number_psfs+i] = t1[j]; - } - } -} - -/* - * initialize2D() - * - * Initialization for 2D analysis. - * - * beta - The beta parameter. - * eta - The eta parameter. - * micro - The micro parameter. - * x_size - Size in x of data and psf. - * y_size - Size in y of data and psf. - */ -dsData* initialize2D(double beta, double eta, double micro, int x_size, int y_size) -{ - return initialize3D(beta, eta, micro, x_size, y_size, 1); -} - -/* - * initialize3D() - * - * Initialization for 3D analysis. Note that setting the values - * in the A and G_inv arrays is done in a separate step. - * - * beta - The beta parameter. - * eta - The eta parameter. - * micro - The micro parameter. - * x_size - Size in x of data and psf. - * y_size - Size in y of data and psf. - * z_size - The number of z planes. - */ -dsData* initialize3D(double beta, double eta, double micro, int x_size, int y_size, int z_size) -{ - int i; - dsData *ds_data; - - ds_data = (dsData *)malloc(sizeof(dsData)); - - ds_data->fft_size = x_size * (y_size/2+1); - ds_data->image_size = x_size * y_size; - ds_data->number_psfs = z_size; - - ds_data->normalization = 1.0/((double)(ds_data->image_size)); - ds_data->beta = beta; - ds_data->eta = eta; - ds_data->micro = micro; - - /* Allocate storage. */ - ds_data->Ax = (double *)malloc(sizeof(double) * ds_data->image_size); - ds_data->b_vec = (double *)malloc(sizeof(double) * ds_data->image_size); - ds_data->d_vec = (double *)malloc(sizeof(double) * ds_data->image_size); - ds_data->image = (double *)malloc(sizeof(double) * ds_data->image_size); - ds_data->w_vec = (double *)malloc(sizeof(double) * ds_data->image_size); - ds_data->y_vec = (double *)malloc(sizeof(double) * ds_data->image_size); - ds_data->y_vec_h = (double *)malloc(sizeof(double) * ds_data->image_size); - - ds_data->e_vec = (double **)malloc(sizeof(double *) * z_size); - ds_data->work1 = (double **)malloc(sizeof(double *) * z_size); - ds_data->x_vec_h = (double **)malloc(sizeof(double *) * z_size); - ds_data->x_vec_t = (double **)malloc(sizeof(double *) * z_size); - - for(i=0;ie_vec[i] = (double *)malloc(sizeof(double) * ds_data->image_size); - ds_data->work1[i] = (double *)malloc(sizeof(double) * ds_data->image_size); - ds_data->x_vec_h[i] = (double *)malloc(sizeof(double) * ds_data->image_size); - ds_data->x_vec_t[i] = (double *)malloc(sizeof(double) * ds_data->image_size); - } - - ds_data->fft_vector = (double *)fftw_malloc(sizeof(double) * ds_data->image_size); - ds_data->fft_vector_fft = (fftw_complex *)fftw_malloc(sizeof(fftw_complex) * ds_data->fft_size); - - ds_data->A = (fftw_complex **)malloc(sizeof(fftw_complex *) * z_size); - ds_data->G_inv = (fftw_complex **)malloc(sizeof(fftw_complex *) * z_size * z_size); - ds_data->work2 = (fftw_complex **)malloc(sizeof(fftw_complex *) * z_size); - - for(i=0;iA[i] = (fftw_complex *)malloc(sizeof(fftw_complex) * ds_data->fft_size); - ds_data->work2[i] = (fftw_complex *)malloc(sizeof(fftw_complex) * ds_data->fft_size); - } - - for(i=0;i<(z_size*z_size);i++){ - ds_data->G_inv[i] = (fftw_complex *)malloc(sizeof(fftw_complex) * ds_data->fft_size); - } - - /* Create FFT plans. */ - ds_data->fft_forward = fftw_plan_dft_r2c_2d(x_size, y_size, ds_data->fft_vector, ds_data->fft_vector_fft, FFTW_MEASURE); - ds_data->fft_backward = fftw_plan_dft_c2r_2d(x_size, y_size, ds_data->fft_vector_fft, ds_data->fft_vector, FFTW_MEASURE); - - return ds_data; -} - -/* - * initializeA() - * - * Set initial values in the A array. - * - * ds_data - A pointer to a dsData structure. - * data - Complex values to use for A. - * index - Which A matrix to store data in. - */ -void initializeA(dsData *ds_data, fftw_complex *data, int index) -{ - ftmComplexCopy(data, ds_data->A[index], ds_data->fft_size); -} - -/* - * initializeGInv() - * - * Set initial values in the G_inv array. - * - * ds_data - A pointer to an admmData structure. - * data - Complex values to use for G. - * index - Which G matrix to store data in. - */ -void initializeGInv(dsData *ds_data, fftw_complex *data, int index) -{ - ftmComplexCopy(data, ds_data->G_inv[index], ds_data->fft_size); -} - -/* - * iterate() - * - * Performs one cycle of improvement. - * - * ds_data - A pointer to a dsData structure. - */ - -void iterate(dsData *ds_data) -{ - int i,j,n_psfs; - double t1, t2, t3; - double *ev,*w1,*xvh,*xvt; - - n_psfs = ds_data->number_psfs; - - /* - * Equation 6. - */ - /* Calculate At_ydb. */ - for(i=0;iimage_size;i++){ - ds_data->fft_vector[i] = ds_data->y_vec_h[i] + ds_data->d_vec[i] - ds_data->b_vec[i]; - } - fftw_execute(ds_data->fft_forward); - ftmComplexCopy(ds_data->fft_vector_fft, ds_data->work2[0], ds_data->fft_size); - - for(i=0;ifft_vector_fft, ds_data->work2[0], ds_data->A[i], ds_data->fft_size, 1); - fftw_execute(ds_data->fft_backward); - ftmDoubleCopyNormalize(ds_data->fft_vector, ds_data->work1[i], ds_data->normalization, ds_data->image_size); - } - - /* Calculate At_ydb + micro * (x-tilde + e). */ - for(i=0;ie_vec[i]; - w1 = ds_data->work1[i]; - xvt = ds_data->x_vec_t[i]; - for(j=0;jimage_size;j++){ - ds_data->fft_vector[j] = w1[j] + ds_data->micro * (xvt[j] + ev[j]); - } - fftw_execute(ds_data->fft_forward); - ftmComplexCopy(ds_data->fft_vector_fft, ds_data->work2[i], ds_data->fft_size); - } - - /* Calculate updated x-hat. */ - for(i=0;ifft_vector_fft, ds_data->fft_size); - for(j=0;jfft_vector_fft, ds_data->G_inv[i*n_psfs+j], ds_data->work2[j], ds_data->fft_size, 0); - } - fftw_execute(ds_data->fft_backward); - ftmDoubleCopyNormalize(ds_data->fft_vector, ds_data->x_vec_h[i], ds_data->normalization, ds_data->image_size); - } - - /* - * Equation 7. - */ - ds_data->stale_Ax = 1; - calculateAx(ds_data); - t1 = 1.0/(2.0*ds_data->eta); - for(i=0;iimage_size;i++){ - t2 = 1.0 - ds_data->eta*(ds_data->Ax[i] + ds_data->b_vec[i] - ds_data->d_vec[i]); - t3 = sqrt(t2*t2 + 4.0*ds_data->eta*ds_data->y_vec[i]); - ds_data->y_vec_h[i] = t1*(-t2 + t3); - } - - /* - * Equation 8. - */ - for(i=0;ix_vec_h[i]; - xvt = ds_data->x_vec_t[i]; - ev = ds_data->e_vec[i]; - for(j=0;jimage_size;j++){ - t1 = xvh[j] - ev[j] - ds_data->w_vec[j]; - if (t1 <= 0.0){ - xvt[j] = 0.0; - } - else { - xvt[j] = t1; - } - } - } - - /* - * Equation 9. - */ - for(i=0;iimage_size;i++){ - ds_data->d_vec[i] = ds_data->d_vec[i] - (ds_data->Ax[i] + ds_data->b_vec[i] - ds_data->y_vec_h[i]); - } - - /* - * Equation 10. - */ - for(i=0;ix_vec_h[i]; - xvt = ds_data->x_vec_t[i]; - ev = ds_data->e_vec[i]; - for(j=0;jimage_size;j++){ - ev[j] = ev[j] - (xvh[j] - xvt[j]); - } - } - - /* Flag that we need to recaculate Ax (if getAx is called). */ - /* ds_data->stale_Ax = 1; */ -} - -/* - * l1Error - * - * Calculate l1 error term (sum of the x vector). - * - * ds_data - A pointer to a dsData structure. - * - * Return sum of the x vector. - */ -double l1Error(dsData *ds_data) -{ - int i,j; - double sum; - double *xv; - - sum = 0.0; - for(i=0;inumber_psfs;i++){ - xv = ds_data->x_vec_h[i]; - for(j=0;jimage_size;j++){ - sum += fabs(xv[j]); - } - } - - return sum; -} - -/* - * l2Error - * - * Calculate l2 error (Ax - b). - * - * ds_data - A pointer to a dsData structure. - * - * Return the error in the fit. - */ -double l2Error(dsData *ds_data) -{ - int i; - double l2_error,t1; - - /* Compute Ax. */ - calculateAx(ds_data); - - /* Compute (Ax - b)^2. */ - l2_error = 0.0; - for(i=0;iimage_size;i++){ - t1 = ds_data->Ax[i] - ds_data->image[i]; - l2_error += t1*t1; - } - - return sqrt(l2_error); -} - -/* - * newImage() - * - * Sets things up for the analysis of a new image. - * - * ds_data - A pointer to a dsData structure. - * image - The image data. - * background - Image background estimate. - */ -void newImage(dsData *ds_data, double *image, double *background) -{ - int i; - - /* Flag that we need to recaculate Ax (if getAx is called). */ - ds_data->stale_Ax = 1; - - /* Store image - background. */ - for(i=0;iimage_size;i++){ - ds_data->image[i] = image[i] - background[i]; - } - - ftmDoubleCopy(image, ds_data->y_vec, ds_data->image_size); - ftmDoubleCopy(background, ds_data->b_vec, ds_data->image_size); - - for(i=0;iimage_size;i++){ - ds_data->w_vec[i] = ds_data->beta * sqrt(background[i]) / ds_data->micro; - } - - /* d, y-hat vectors. */ - ftmDoubleCopy(image, ds_data->y_vec_h, ds_data->image_size); - ftmDoubleZero(ds_data->d_vec, ds_data->image_size); - - /* zero e, x-hat, x-tidle vectors. */ - for(i=0;inumber_psfs;i++){ - ftmDoubleZero(ds_data->e_vec[i], ds_data->image_size); - ftmDoubleZero(ds_data->x_vec_h[i], ds_data->image_size); - ftmDoubleZero(ds_data->x_vec_t[i], ds_data->image_size); - } -} - -/* - * run() - * - * ds_data - A pointer to a dsData structure. - * cycles - Number of iterations to perform. - * - * Performs a fixed number of cycles. - */ -void run(dsData *ds_data, int cycles) -{ - int i; - - for(i=0;i +#include +#include + +#include + +#include "../sa_library/ft_math.h" + + +/* 3denseSTORM Structure. */ +typedef struct{ + int fft_size; + int image_size; + int number_psfs; + int stale_Ax; + + double beta; + double eta; + double micro; + double normalization; + + double *Ax; + double *b_vec; + double *d_vec; + double *fft_vector; + double *image; + double *w_vec; + double *y_vec; + double *y_vec_h; + + double **e_vec; + double **work1; + double **x_vec_h; + double **x_vec_t; + + fftw_plan fft_backward; + fftw_plan fft_forward; + + fftw_complex *fft_vector_fft; + + fftw_complex **A; + fftw_complex **G_inv; + fftw_complex **work2; + +} dsData; + + +/* Functions Declarations. */ +void calculateAx(dsData *); +void cleanup(dsData *); +void getAx(dsData *, double *); +void getXVector(dsData *, double *, int); +dsData* initialize2D(double, double, double, int, int); +dsData* initialize3D(double, double, double, int, int, int); +void initializeA(dsData *, fftw_complex *, int); +void initializeGInv(dsData *, fftw_complex *, int); +void iterate(dsData *); +double l1Error(dsData *); +double l2Error(dsData *); +void newImage(dsData *, double *, double *); +void run(dsData *, int); + + +/* Functions */ + +/* + * calculateAx() + * + * ds_data - A pointer to a dsData structure. + */ +void calculateAx(dsData *ds_data) +{ + int i; + + if (ds_data->stale_Ax){ + + /* Compute Ax. */ + ftmComplexZero(ds_data->work2[0], ds_data->fft_size); + + for(i=0;inumber_psfs;i++){ + + // Compute FFT of x vector for each z plane. + ftmDoubleCopy(ds_data->x_vec_h[i], ds_data->fft_vector, ds_data->image_size); + fftw_execute(ds_data->fft_forward); + + // Multiply FFT of x vector by FFT of the PSF for this z plane. + ftmComplexMultiplyAccum(ds_data->work2[0], ds_data->fft_vector_fft, ds_data->A[i], ds_data->fft_size, 0); + } + + ftmComplexCopy(ds_data->work2[0], ds_data->fft_vector_fft, ds_data->fft_size); + fftw_execute(ds_data->fft_backward); + + /* Store Ax. */ + ftmDoubleCopyNormalize(ds_data->fft_vector, ds_data->Ax, ds_data->normalization, ds_data->image_size); + } + + ds_data->stale_Ax = 0; +} + +/* + * cleanup + * + * ds_data - A pointer to a dsData structure. + */ +void cleanup(dsData *ds_data) +{ + int i; + + free(ds_data->Ax); + free(ds_data->b_vec); + free(ds_data->d_vec); + free(ds_data->image); + free(ds_data->w_vec); + free(ds_data->y_vec); + free(ds_data->y_vec_h); + + for(i=0;inumber_psfs;i++){ + free(ds_data->e_vec[i]); + free(ds_data->work1[i]); + free(ds_data->x_vec_h[i]); + free(ds_data->x_vec_t[i]); + } + free(ds_data->e_vec); + free(ds_data->work1); + free(ds_data->x_vec_h); + free(ds_data->x_vec_t); + + fftw_destroy_plan(ds_data->fft_backward); + fftw_destroy_plan(ds_data->fft_forward); + + fftw_free(ds_data->fft_vector); + fftw_free(ds_data->fft_vector_fft); + + for(i=0;inumber_psfs;i++){ + free(ds_data->A[i]); + free(ds_data->work2[i]); + } + for(i=0;i<(ds_data->number_psfs*ds_data->number_psfs);i++){ + free(ds_data->G_inv[i]); + } + + free(ds_data->A); + free(ds_data->G_inv); + free(ds_data->work2); + + free(ds_data); +} + +/* + * getAx() + * + * Copies the current Ax vector into user supplied storage. + * + * ds_data - A pointer to a dsData structure. + * data - Storage for the Ax vector. + */ +void getAx(dsData *ds_data, double *data) +{ + calculateAx(ds_data); + ftmDoubleCopy(ds_data->Ax, data, ds_data->image_size); +} + +/* + * getXVector() + * + * Copies the current x vector into user supplied storage. Also + * converts back from [(x1,y1), (x2,y2), ..] to (x,y,z). + * + * ds_data - A pointer to a dsData structure. + * data - Storage for the x vector. + * compressed - Return x_vec_t (1) or x_vec_h (0). + */ +void getXVector(dsData *ds_data, double *data, int compressed) +{ + int i,j; + double *t1; + + for(i=0;inumber_psfs;i++){ + if (compressed){ + t1 = ds_data->x_vec_t[i]; + } + else{ + t1 = ds_data->x_vec_h[i]; + } + for(j=0;jimage_size;j++){ + data[j*ds_data->number_psfs+i] = t1[j]; + } + } +} + +/* + * initialize2D() + * + * Initialization for 2D analysis. + * + * beta - The beta parameter. + * eta - The eta parameter. + * micro - The micro parameter. + * x_size - Size in x of data and psf. + * y_size - Size in y of data and psf. + */ +dsData* initialize2D(double beta, double eta, double micro, int x_size, int y_size) +{ + return initialize3D(beta, eta, micro, x_size, y_size, 1); +} + +/* + * initialize3D() + * + * Initialization for 3D analysis. Note that setting the values + * in the A and G_inv arrays is done in a separate step. + * + * beta - The beta parameter. + * eta - The eta parameter. + * micro - The micro parameter. + * x_size - Size in x of data and psf. + * y_size - Size in y of data and psf. + * z_size - The number of z planes. + */ +dsData* initialize3D(double beta, double eta, double micro, int x_size, int y_size, int z_size) +{ + int i; + dsData *ds_data; + + ds_data = (dsData *)malloc(sizeof(dsData)); + + ds_data->fft_size = x_size * (y_size/2+1); + ds_data->image_size = x_size * y_size; + ds_data->number_psfs = z_size; + + ds_data->normalization = 1.0/((double)(ds_data->image_size)); + ds_data->beta = beta; + ds_data->eta = eta; + ds_data->micro = micro; + + /* Allocate storage. */ + ds_data->Ax = (double *)malloc(sizeof(double) * ds_data->image_size); + ds_data->b_vec = (double *)malloc(sizeof(double) * ds_data->image_size); + ds_data->d_vec = (double *)malloc(sizeof(double) * ds_data->image_size); + ds_data->image = (double *)malloc(sizeof(double) * ds_data->image_size); + ds_data->w_vec = (double *)malloc(sizeof(double) * ds_data->image_size); + ds_data->y_vec = (double *)malloc(sizeof(double) * ds_data->image_size); + ds_data->y_vec_h = (double *)malloc(sizeof(double) * ds_data->image_size); + + ds_data->e_vec = (double **)malloc(sizeof(double *) * z_size); + ds_data->work1 = (double **)malloc(sizeof(double *) * z_size); + ds_data->x_vec_h = (double **)malloc(sizeof(double *) * z_size); + ds_data->x_vec_t = (double **)malloc(sizeof(double *) * z_size); + + for(i=0;ie_vec[i] = (double *)malloc(sizeof(double) * ds_data->image_size); + ds_data->work1[i] = (double *)malloc(sizeof(double) * ds_data->image_size); + ds_data->x_vec_h[i] = (double *)malloc(sizeof(double) * ds_data->image_size); + ds_data->x_vec_t[i] = (double *)malloc(sizeof(double) * ds_data->image_size); + } + + ds_data->fft_vector = (double *)fftw_malloc(sizeof(double) * ds_data->image_size); + ds_data->fft_vector_fft = (fftw_complex *)fftw_malloc(sizeof(fftw_complex) * ds_data->fft_size); + + ds_data->A = (fftw_complex **)malloc(sizeof(fftw_complex *) * z_size); + ds_data->G_inv = (fftw_complex **)malloc(sizeof(fftw_complex *) * z_size * z_size); + ds_data->work2 = (fftw_complex **)malloc(sizeof(fftw_complex *) * z_size); + + for(i=0;iA[i] = (fftw_complex *)malloc(sizeof(fftw_complex) * ds_data->fft_size); + ds_data->work2[i] = (fftw_complex *)malloc(sizeof(fftw_complex) * ds_data->fft_size); + } + + for(i=0;i<(z_size*z_size);i++){ + ds_data->G_inv[i] = (fftw_complex *)malloc(sizeof(fftw_complex) * ds_data->fft_size); + } + + /* Create FFT plans. */ + ds_data->fft_forward = fftw_plan_dft_r2c_2d(x_size, y_size, ds_data->fft_vector, ds_data->fft_vector_fft, FFTW_MEASURE); + ds_data->fft_backward = fftw_plan_dft_c2r_2d(x_size, y_size, ds_data->fft_vector_fft, ds_data->fft_vector, FFTW_MEASURE); + + return ds_data; +} + +/* + * initializeA() + * + * Set initial values in the A array. + * + * ds_data - A pointer to a dsData structure. + * data - Complex values to use for A. + * index - Which A matrix to store data in. + */ +void initializeA(dsData *ds_data, fftw_complex *data, int index) +{ + ftmComplexCopy(data, ds_data->A[index], ds_data->fft_size); +} + +/* + * initializeGInv() + * + * Set initial values in the G_inv array. + * + * ds_data - A pointer to an admmData structure. + * data - Complex values to use for G. + * index - Which G matrix to store data in. + */ +void initializeGInv(dsData *ds_data, fftw_complex *data, int index) +{ + ftmComplexCopy(data, ds_data->G_inv[index], ds_data->fft_size); +} + +/* + * iterate() + * + * Performs one cycle of improvement. + * + * ds_data - A pointer to a dsData structure. + */ + +void iterate(dsData *ds_data) +{ + int i,j,n_psfs; + double t1, t2, t3; + double *ev,*w1,*xvh,*xvt; + + n_psfs = ds_data->number_psfs; + + /* + * Equation 6. + */ + /* Calculate At_ydb. */ + for(i=0;iimage_size;i++){ + ds_data->fft_vector[i] = ds_data->y_vec_h[i] + ds_data->d_vec[i] - ds_data->b_vec[i]; + } + fftw_execute(ds_data->fft_forward); + ftmComplexCopy(ds_data->fft_vector_fft, ds_data->work2[0], ds_data->fft_size); + + for(i=0;ifft_vector_fft, ds_data->work2[0], ds_data->A[i], ds_data->fft_size, 1); + fftw_execute(ds_data->fft_backward); + ftmDoubleCopyNormalize(ds_data->fft_vector, ds_data->work1[i], ds_data->normalization, ds_data->image_size); + } + + /* Calculate At_ydb + micro * (x-tilde + e). */ + for(i=0;ie_vec[i]; + w1 = ds_data->work1[i]; + xvt = ds_data->x_vec_t[i]; + for(j=0;jimage_size;j++){ + ds_data->fft_vector[j] = w1[j] + ds_data->micro * (xvt[j] + ev[j]); + } + fftw_execute(ds_data->fft_forward); + ftmComplexCopy(ds_data->fft_vector_fft, ds_data->work2[i], ds_data->fft_size); + } + + /* Calculate updated x-hat. */ + for(i=0;ifft_vector_fft, ds_data->fft_size); + for(j=0;jfft_vector_fft, ds_data->G_inv[i*n_psfs+j], ds_data->work2[j], ds_data->fft_size, 0); + } + fftw_execute(ds_data->fft_backward); + ftmDoubleCopyNormalize(ds_data->fft_vector, ds_data->x_vec_h[i], ds_data->normalization, ds_data->image_size); + } + + /* + * Equation 7. + */ + ds_data->stale_Ax = 1; + calculateAx(ds_data); + t1 = 1.0/(2.0*ds_data->eta); + for(i=0;iimage_size;i++){ + t2 = 1.0 - ds_data->eta*(ds_data->Ax[i] + ds_data->b_vec[i] - ds_data->d_vec[i]); + t3 = sqrt(t2*t2 + 4.0*ds_data->eta*ds_data->y_vec[i]); + ds_data->y_vec_h[i] = t1*(-t2 + t3); + } + + /* + * Equation 8. + */ + for(i=0;ix_vec_h[i]; + xvt = ds_data->x_vec_t[i]; + ev = ds_data->e_vec[i]; + for(j=0;jimage_size;j++){ + t1 = xvh[j] - ev[j] - ds_data->w_vec[j]; + if (t1 <= 0.0){ + xvt[j] = 0.0; + } + else { + xvt[j] = t1; + } + } + } + + /* + * Equation 9. + */ + for(i=0;iimage_size;i++){ + ds_data->d_vec[i] = ds_data->d_vec[i] - (ds_data->Ax[i] + ds_data->b_vec[i] - ds_data->y_vec_h[i]); + } + + /* + * Equation 10. + */ + for(i=0;ix_vec_h[i]; + xvt = ds_data->x_vec_t[i]; + ev = ds_data->e_vec[i]; + for(j=0;jimage_size;j++){ + ev[j] = ev[j] - (xvh[j] - xvt[j]); + } + } + + /* Flag that we need to recaculate Ax (if getAx is called). */ + /* ds_data->stale_Ax = 1; */ +} + +/* + * l1Error + * + * Calculate l1 error term (sum of the x vector). + * + * ds_data - A pointer to a dsData structure. + * + * Return sum of the x vector. + */ +double l1Error(dsData *ds_data) +{ + int i,j; + double sum; + double *xv; + + sum = 0.0; + for(i=0;inumber_psfs;i++){ + xv = ds_data->x_vec_h[i]; + for(j=0;jimage_size;j++){ + sum += fabs(xv[j]); + } + } + + return sum; +} + +/* + * l2Error + * + * Calculate l2 error (Ax - b). + * + * ds_data - A pointer to a dsData structure. + * + * Return the error in the fit. + */ +double l2Error(dsData *ds_data) +{ + int i; + double l2_error,t1; + + /* Compute Ax. */ + calculateAx(ds_data); + + /* Compute (Ax - b)^2. */ + l2_error = 0.0; + for(i=0;iimage_size;i++){ + t1 = ds_data->Ax[i] - ds_data->image[i]; + l2_error += t1*t1; + } + + return sqrt(l2_error); +} + +/* + * newImage() + * + * Sets things up for the analysis of a new image. + * + * ds_data - A pointer to a dsData structure. + * image - The image data. + * background - Image background estimate. + */ +void newImage(dsData *ds_data, double *image, double *background) +{ + int i; + + /* Flag that we need to recaculate Ax (if getAx is called). */ + ds_data->stale_Ax = 1; + + /* Store image - background. */ + for(i=0;iimage_size;i++){ + ds_data->image[i] = image[i] - background[i]; + } + + ftmDoubleCopy(image, ds_data->y_vec, ds_data->image_size); + ftmDoubleCopy(background, ds_data->b_vec, ds_data->image_size); + + for(i=0;iimage_size;i++){ + ds_data->w_vec[i] = ds_data->beta * sqrt(background[i]) / ds_data->micro; + } + + /* d, y-hat vectors. */ + ftmDoubleCopy(image, ds_data->y_vec_h, ds_data->image_size); + ftmDoubleZero(ds_data->d_vec, ds_data->image_size); + + /* zero e, x-hat, x-tidle vectors. */ + for(i=0;inumber_psfs;i++){ + ftmDoubleZero(ds_data->e_vec[i], ds_data->image_size); + ftmDoubleZero(ds_data->x_vec_h[i], ds_data->image_size); + ftmDoubleZero(ds_data->x_vec_t[i], ds_data->image_size); + } +} + +/* + * run() + * + * ds_data - A pointer to a dsData structure. + * cycles - Number of iterations to perform. + * + * Performs a fixed number of cycles. + */ +void run(dsData *ds_data, int cycles) +{ + int i; + + for(i=0;i -#include -#include - -#include - -#include "../sa_library/ft_math.h" - - -/* FISTA Structure */ -typedef struct{ - int dwls; - int fft_size; - int image_size; - int number_psfs; - int stale_Ax; - - double normalization; - double time_step; - double tk; - double *Ax; - double *fft_vector; - double *image; - double *weights; - double **x_vector; - double **x_vector_old; - double **y_vector; - - fftw_plan fft_forward; - fftw_plan fft_backward; - - fftw_complex *Ax_fft; - fftw_complex *fft_vector_fft; - fftw_complex *image_fft; - fftw_complex **psf_fft; -} fistaData; - - -/* Function Declarations */ -void calculateAx(fistaData *); -void cleanup(fistaData *); -void getAx(fistaData *, double *); -void getXVector(fistaData *, double *); -fistaData* initialize2D(double *, double, int, int, int); -fistaData* initialize3D(double *, double, int, int, int, int); -void iterate(fistaData *, double); -double l1Error(fistaData *); -double l2Error(fistaData *); -void newImage(fistaData *, double *, double *); -void run(fistaData *, double, int); - - -/* Functions */ - -/* - * calculateAx() - * - * fista_data - A pointer to a fistaData structure. - */ -void calculateAx(fistaData *fista_data) -{ - int i; - - if (fista_data->stale_Ax){ - - /* Compute Ax. */ - ftmComplexZero(fista_data->Ax_fft, fista_data->fft_size); - - for(i=0;inumber_psfs;i++){ - - // Compute FFT of x vector for each z plane. - ftmDoubleCopy(fista_data->x_vector[i], fista_data->fft_vector, fista_data->image_size); - fftw_execute(fista_data->fft_forward); - - // Multiply FFT of x vector by FFT of the PSF for this z plane. - ftmComplexMultiplyAccum(fista_data->Ax_fft, fista_data->fft_vector_fft, fista_data->psf_fft[i], fista_data->fft_size, 0); - } - - ftmComplexCopy(fista_data->Ax_fft, fista_data->fft_vector_fft, fista_data->fft_size); - fftw_execute(fista_data->fft_backward); - - /* Store Ax. */ - ftmDoubleCopyNormalize(fista_data->fft_vector, fista_data->Ax, fista_data->normalization, fista_data->image_size); - } - - fista_data->stale_Ax = 0; -} - -/* - * cleanup() - * - * fista_data - A pointer to a fistaData structure. - */ -void cleanup(fistaData *fista_data) -{ - int i; - - free(fista_data->Ax); - free(fista_data->image); - free(fista_data->weights); - - for(i=0;inumber_psfs;i++){ - free(fista_data->x_vector[i]); - free(fista_data->x_vector_old[i]); - free(fista_data->y_vector[i]); - } - free(fista_data->x_vector); - free(fista_data->x_vector_old); - free(fista_data->y_vector); - - fftw_destroy_plan(fista_data->fft_forward); - fftw_destroy_plan(fista_data->fft_backward); - - fftw_free(fista_data->Ax_fft); - fftw_free(fista_data->fft_vector); - fftw_free(fista_data->fft_vector_fft); - fftw_free(fista_data->image_fft); - - for(i=0;inumber_psfs;i++){ - fftw_free(fista_data->psf_fft[i]); - } - free(fista_data->psf_fft); - - free(fista_data); -} - -/* - * dwlsError - * - * Calculate DWLS error (Ax - b)^2 / image. - * Note that b is the background corrected image. - * - * fista_data - A pointer to a fistaData structure. - * - * Return the error in the fit. - */ -double dwlsError(fistaData *fista_data) -{ - int i; - double dwls_error,t1; - - /* Compute Ax. */ - calculateAx(fista_data); - - /* Compute (Ax - b)^2 / image. */ - dwls_error = 0.0; - for(i=0;iimage_size;i++){ - t1 = fista_data->Ax[i] - fista_data->image[i]; - dwls_error += t1*t1*fista_data->weights[i]; - } - - return dwls_error; -} - -/* - * getAx() - * - * Copies the current Ax vector into user supplied storage. - * - * fista_data - A pointer to a fistaData structure. - * data - Storage for the Ax vector. - */ -void getAx(fistaData *fista_data, double *data) -{ - calculateAx(fista_data); - ftmDoubleCopy(fista_data->Ax, data, fista_data->image_size); -} - -/* - * getXVector() - * - * Copies the current x vector into user supplied storage. Also - * converts back from [(x1,y1), (x2,y2), ..] to (x,y,z). - * - * fista_data - A pointer to a fistaData structure. - * data - Storage for the x vector. - */ -void getXVector(fistaData *fista_data, double *data) -{ - int i,j; - double *t1; - - for(i=0;inumber_psfs;i++){ - t1 = fista_data->x_vector[i]; - for(j=0;jimage_size;j++){ - data[j*fista_data->number_psfs+i] = t1[j]; - } - } -} - -/* - * initialize2D() - * - * Set things up for 2D analysis. - * - * psf - The psf (also x_size x y_size) - * t_step - The time step to use. - * x_size - Size in x of data and psf. - * y_size - Size in y of data and psf. - * dwls - Data weighted least squares optimization. - */ -fistaData* initialize2D(double *psf, double t_step, int x_size, int y_size, int dwls) -{ - return initialize3D(psf, t_step, x_size, y_size, 1, dwls); -} - -/* - * initialize3D() - * - * Set things up for 3D analysis. - * - * psf - The psf (x_size, y_size, z_size) - * t_step - the time step to use. - * x_size - Size in x of data and psf. - * y_size - Size in y of data and psf. - * z_size - The z dimension of the psf. - * dwls - Data weighted least squares optimization. - */ -fistaData* initialize3D(double *psf, double t_step, int x_size, int y_size, int z_size, int dwls) -{ - int i,j; - double t1; - fftw_complex *t2; - fistaData *fista_data; - - fista_data = (fistaData *)malloc(sizeof(fistaData)); - - fista_data->dwls = dwls; - fista_data->stale_Ax = 1; - - /* Initialize some variables. */ - fista_data->fft_size = x_size * (y_size/2 + 1); - fista_data->image_size = x_size * y_size; - fista_data->normalization = 1.0/((double)(fista_data->image_size)); - fista_data->number_psfs = z_size; - - /* Allocate storage. */ - fista_data->Ax = (double *)malloc(sizeof(double) * fista_data->image_size); - fista_data->image = (double *)malloc(sizeof(double) * fista_data->image_size); - fista_data->weights = (double *)malloc(sizeof(double) * fista_data->image_size); - - fista_data->x_vector = (double **)malloc(sizeof(double *) * z_size); - fista_data->x_vector_old = (double **)malloc(sizeof(double *) * z_size); - fista_data->y_vector = (double **)malloc(sizeof(double *) * z_size); - for(i=0;ix_vector[i] = (double *)malloc(sizeof(double) * fista_data->image_size); - fista_data->x_vector_old[i] = (double *)malloc(sizeof(double) * fista_data->image_size); - fista_data->y_vector[i] = (double *)malloc(sizeof(double) * fista_data->image_size); - } - - fista_data->fft_vector = (double *)fftw_malloc(sizeof(double) * fista_data->image_size); - fista_data->Ax_fft = (fftw_complex *)fftw_malloc(sizeof(fftw_complex) * fista_data->fft_size); - fista_data->fft_vector_fft = (fftw_complex *)fftw_malloc(sizeof(fftw_complex) * fista_data->fft_size); - fista_data->image_fft = (fftw_complex *)fftw_malloc(sizeof(fftw_complex) * fista_data->fft_size); - - fista_data->psf_fft = (fftw_complex **)malloc(sizeof(fftw_complex *) * z_size); - for(i=0;ipsf_fft[i] = (fftw_complex *)fftw_malloc(sizeof(fftw_complex) * fista_data->fft_size); - } - - /* Create FFT plans. */ - fista_data->fft_forward = fftw_plan_dft_r2c_2d(x_size, y_size, fista_data->fft_vector, fista_data->fft_vector_fft, FFTW_MEASURE); - fista_data->fft_backward = fftw_plan_dft_c2r_2d(x_size, y_size, fista_data->fft_vector_fft, fista_data->fft_vector, FFTW_MEASURE); - - /* - Compute FFTs of the psfs and save in psf_fft. - Note: The input psfs are indexed (x,y,z) but internally we - use [(x1,y1), (x2,y2), ..] as this is hopefully clearer. - */ - for(i=0;iimage_size;j++){ - fista_data->fft_vector[j] = psf[j*z_size+i]; - } - fftw_execute(fista_data->fft_forward); - ftmComplexCopy(fista_data->fft_vector_fft, fista_data->psf_fft[i], fista_data->fft_size); - } - - /* Calculate optimal time step. */ - fista_data->time_step = 0.0; - for(i=0;ipsf_fft[i]; - for(j=0;jfft_size;j++){ - t1 = t2[j][0] * t2[j][0] + t2[j][1] * t2[j][1]; - if(t1>fista_data->time_step){ - fista_data->time_step = t1; - } - } - } - - fista_data->time_step = 1.0/(2.0*fista_data->time_step); - printf("Optimal time step is: %f\n", fista_data->time_step); - - fista_data->time_step = t_step; - - return fista_data; -} - -/* - * iterate() - * - * Performs a single cycle of optimization. - * - * fista_data - A pointer to a fistaData structure. - * lambda - l1 error term weigth. - */ -void iterate(fistaData *fista_data, double lambda) -{ - int i,j; - double lt,new_tk,t1,t2; - double *xv,*xv_old,*yv; - - /* Flag that we need to recaculate Ax (if getAx is called). */ - fista_data->stale_Ax = 1; - - /* Copy current x vector into old x vector. */ - for(i=0;inumber_psfs;i++){ - ftmDoubleCopy(fista_data->x_vector[i], fista_data->x_vector_old[i], fista_data->image_size); - } - - /* - * This is the plk(y) step in the FISTA algorithm. - */ - - /* Compute Ax_fft (n,n). x is generic here, and does not - implicitly refer to x_vector. */ - ftmComplexZero(fista_data->Ax_fft, fista_data->fft_size); - - for(i=0;inumber_psfs;i++){ - - /* Compute FFT of y vector for each z plane. */ - ftmDoubleCopy(fista_data->y_vector[i], fista_data->fft_vector, fista_data->image_size); - fftw_execute(fista_data->fft_forward); - - /* Multiply FFT of y vector by FFT of the PSF for this z plane. */ - ftmComplexMultiplyAccum(fista_data->Ax_fft, fista_data->fft_vector_fft, fista_data->psf_fft[i], fista_data->fft_size, 0); - } - - /* Compute Ax_fft - b_fft (image_fft) (n,n). */ - for(i=0;ifft_size;i++){ - fista_data->Ax_fft[i][0] -= fista_data->image_fft[i][0]; - fista_data->Ax_fft[i][1] -= fista_data->image_fft[i][1]; - } - - /* Compute x = y - At(Ax-b) (z,n,n). */ - for(i=0;inumber_psfs;i++){ - - // Compute inverse FFT of At(Ax-b) for each image plane. - ftmComplexMultiply(fista_data->fft_vector_fft, fista_data->Ax_fft, fista_data->psf_fft[i], fista_data->fft_size, 1); - fftw_execute(fista_data->fft_backward); - - // Update x vector. - t1 = 2.0*fista_data->time_step*fista_data->normalization; - xv = fista_data->x_vector[i]; - yv = fista_data->y_vector[i]; - - // Data Weighted Least Squares update. - if (fista_data->dwls){ - for(j=0;jimage_size;j++){ - xv[j] = yv[j] - t1*fista_data->fft_vector[j]*fista_data->weights[j]; - } - } - - // Least Squares update. - else{ - for(j=0;jimage_size;j++){ - xv[j] = yv[j] - t1*fista_data->fft_vector[j]; - } - } - } - - /* Shrink x vector. */ - lt = fista_data->time_step*lambda; - for(i=0;inumber_psfs;i++){ - xv = fista_data->x_vector[i]; - for(j=0;jimage_size;j++){ - t1 = xv[j]; - t2 = fabs(xv[j]) - lt; - if (t2 <= 0.0){ - xv[j] = 0.0; - } - else{ - xv[j] = (t1 > 0.0) ? t2 : -t2; - } - } - } - - /* - * Update tk term step. - */ - new_tk = 0.5*(1.0 + sqrt(1.0 + 4.0 * fista_data->tk * fista_data->tk)); - - /* - * Compute new y vector step. - */ - t1 = (fista_data->tk - 1.0)/new_tk; - for(i=0;inumber_psfs;i++){ - xv = fista_data->x_vector[i]; - xv_old = fista_data->x_vector_old[i]; - yv = fista_data->y_vector[i]; - for(j=0;jimage_size;j++){ - yv[j] = xv[j] + t1*(xv[j] - xv_old[j]); - } - } - - /* - * Update tk term step. - */ - fista_data->tk = new_tk; -} - -/* - * l1Error - * - * Calculate l1 error term (sum of the x vector). - * - * fista_data - A pointer to a fistaData structure. - * - * Return sum of the x vector. - */ -double l1Error(fistaData *fista_data) -{ - int i,j; - double sum; - double *xv; - - sum = 0.0; - for(i=0;inumber_psfs;i++){ - xv = fista_data->x_vector[i]; - for(j=0;jimage_size;j++){ - sum += fabs(xv[j]); - } - } - - return sum; -} - -/* - * l2Error - * - * Calculate l2 error (Ax - b). - * - * fista_data - A pointer to a fistaData structure. - * - * Return the error in the fit. - */ -double l2Error(fistaData *fista_data) -{ - int i; - double l2_error,t1; - - /* Compute Ax. */ - calculateAx(fista_data); - - /* Compute (Ax - b)^2. */ - l2_error = 0.0; - for(i=0;iimage_size;i++){ - t1 = fista_data->Ax[i] - fista_data->image[i]; - l2_error += t1*t1; - } - - return sqrt(l2_error); -} - -/* - * newImage() - * - * Initialize with a new image. - * - * fista_data - A pointer to a fistaData structure. - * image - The image data (including the background estimate). - * background - The background estimate. - */ -void newImage(fistaData *fista_data, double *image, double *background) -{ - int i; - - /* Flag that we need to recaculate Ax (if getAx is called). */ - fista_data->stale_Ax = 1; - - fista_data->tk = 1.0; - - /* Calculate background corrected image. */ - for(i=0;iimage_size;i++){ - fista_data->image[i] = image[i] - background[i]; - } - - /* Calculate weights. */ - for(i=0;iimage_size;i++){ - fista_data->weights[i] = 1.0/image[i]; - } - - /* Compute FFT of image. */ - ftmDoubleCopy(fista_data->image, fista_data->fft_vector, fista_data->image_size); - fftw_execute(fista_data->fft_forward); - ftmComplexCopy(fista_data->fft_vector_fft, fista_data->image_fft, fista_data->fft_size); - - /* Initialize x_vector, y_vectors */ - for(i=0;inumber_psfs;i++){ - ftmDoubleZero(fista_data->x_vector[i], fista_data->image_size); - ftmDoubleZero(fista_data->x_vector_old[i], fista_data->image_size); - ftmDoubleZero(fista_data->y_vector[i], fista_data->image_size); - } - -} - -/* - * run() - * - * fista_data - A pointer to a fistaData structure. - * lambda - Lambda value. - * cycles - Number of iterations to perform. - * - * Performs a fixed number of cycles at a fixed lambda. - */ -void run(fistaData *fista_data, double lambda, int cycles) -{ - int i; - - for(i=0;i +#include +#include + +#include + +#include "../sa_library/ft_math.h" + + +/* FISTA Structure */ +typedef struct{ + int dwls; + int fft_size; + int image_size; + int number_psfs; + int stale_Ax; + + double normalization; + double time_step; + double tk; + double *Ax; + double *fft_vector; + double *image; + double *weights; + double **x_vector; + double **x_vector_old; + double **y_vector; + + fftw_plan fft_forward; + fftw_plan fft_backward; + + fftw_complex *Ax_fft; + fftw_complex *fft_vector_fft; + fftw_complex *image_fft; + fftw_complex **psf_fft; +} fistaData; + + +/* Function Declarations */ +void calculateAx(fistaData *); +void cleanup(fistaData *); +void getAx(fistaData *, double *); +void getXVector(fistaData *, double *); +fistaData* initialize2D(double *, double, int, int, int); +fistaData* initialize3D(double *, double, int, int, int, int); +void iterate(fistaData *, double); +double l1Error(fistaData *); +double l2Error(fistaData *); +void newImage(fistaData *, double *, double *); +void run(fistaData *, double, int); + + +/* Functions */ + +/* + * calculateAx() + * + * fista_data - A pointer to a fistaData structure. + */ +void calculateAx(fistaData *fista_data) +{ + int i; + + if (fista_data->stale_Ax){ + + /* Compute Ax. */ + ftmComplexZero(fista_data->Ax_fft, fista_data->fft_size); + + for(i=0;inumber_psfs;i++){ + + // Compute FFT of x vector for each z plane. + ftmDoubleCopy(fista_data->x_vector[i], fista_data->fft_vector, fista_data->image_size); + fftw_execute(fista_data->fft_forward); + + // Multiply FFT of x vector by FFT of the PSF for this z plane. + ftmComplexMultiplyAccum(fista_data->Ax_fft, fista_data->fft_vector_fft, fista_data->psf_fft[i], fista_data->fft_size, 0); + } + + ftmComplexCopy(fista_data->Ax_fft, fista_data->fft_vector_fft, fista_data->fft_size); + fftw_execute(fista_data->fft_backward); + + /* Store Ax. */ + ftmDoubleCopyNormalize(fista_data->fft_vector, fista_data->Ax, fista_data->normalization, fista_data->image_size); + } + + fista_data->stale_Ax = 0; +} + +/* + * cleanup() + * + * fista_data - A pointer to a fistaData structure. + */ +void cleanup(fistaData *fista_data) +{ + int i; + + free(fista_data->Ax); + free(fista_data->image); + free(fista_data->weights); + + for(i=0;inumber_psfs;i++){ + free(fista_data->x_vector[i]); + free(fista_data->x_vector_old[i]); + free(fista_data->y_vector[i]); + } + free(fista_data->x_vector); + free(fista_data->x_vector_old); + free(fista_data->y_vector); + + fftw_destroy_plan(fista_data->fft_forward); + fftw_destroy_plan(fista_data->fft_backward); + + fftw_free(fista_data->Ax_fft); + fftw_free(fista_data->fft_vector); + fftw_free(fista_data->fft_vector_fft); + fftw_free(fista_data->image_fft); + + for(i=0;inumber_psfs;i++){ + fftw_free(fista_data->psf_fft[i]); + } + free(fista_data->psf_fft); + + free(fista_data); +} + +/* + * dwlsError + * + * Calculate DWLS error (Ax - b)^2 / image. + * Note that b is the background corrected image. + * + * fista_data - A pointer to a fistaData structure. + * + * Return the error in the fit. + */ +double dwlsError(fistaData *fista_data) +{ + int i; + double dwls_error,t1; + + /* Compute Ax. */ + calculateAx(fista_data); + + /* Compute (Ax - b)^2 / image. */ + dwls_error = 0.0; + for(i=0;iimage_size;i++){ + t1 = fista_data->Ax[i] - fista_data->image[i]; + dwls_error += t1*t1*fista_data->weights[i]; + } + + return dwls_error; +} + +/* + * getAx() + * + * Copies the current Ax vector into user supplied storage. + * + * fista_data - A pointer to a fistaData structure. + * data - Storage for the Ax vector. + */ +void getAx(fistaData *fista_data, double *data) +{ + calculateAx(fista_data); + ftmDoubleCopy(fista_data->Ax, data, fista_data->image_size); +} + +/* + * getXVector() + * + * Copies the current x vector into user supplied storage. Also + * converts back from [(x1,y1), (x2,y2), ..] to (x,y,z). + * + * fista_data - A pointer to a fistaData structure. + * data - Storage for the x vector. + */ +void getXVector(fistaData *fista_data, double *data) +{ + int i,j; + double *t1; + + for(i=0;inumber_psfs;i++){ + t1 = fista_data->x_vector[i]; + for(j=0;jimage_size;j++){ + data[j*fista_data->number_psfs+i] = t1[j]; + } + } +} + +/* + * initialize2D() + * + * Set things up for 2D analysis. + * + * psf - The psf (also x_size x y_size) + * t_step - The time step to use. + * x_size - Size in x of data and psf. + * y_size - Size in y of data and psf. + * dwls - Data weighted least squares optimization. + */ +fistaData* initialize2D(double *psf, double t_step, int x_size, int y_size, int dwls) +{ + return initialize3D(psf, t_step, x_size, y_size, 1, dwls); +} + +/* + * initialize3D() + * + * Set things up for 3D analysis. + * + * psf - The psf (x_size, y_size, z_size) + * t_step - the time step to use. + * x_size - Size in x of data and psf. + * y_size - Size in y of data and psf. + * z_size - The z dimension of the psf. + * dwls - Data weighted least squares optimization. + */ +fistaData* initialize3D(double *psf, double t_step, int x_size, int y_size, int z_size, int dwls) +{ + int i,j; + double t1; + fftw_complex *t2; + fistaData *fista_data; + + fista_data = (fistaData *)malloc(sizeof(fistaData)); + + fista_data->dwls = dwls; + fista_data->stale_Ax = 1; + + /* Initialize some variables. */ + fista_data->fft_size = x_size * (y_size/2 + 1); + fista_data->image_size = x_size * y_size; + fista_data->normalization = 1.0/((double)(fista_data->image_size)); + fista_data->number_psfs = z_size; + + /* Allocate storage. */ + fista_data->Ax = (double *)malloc(sizeof(double) * fista_data->image_size); + fista_data->image = (double *)malloc(sizeof(double) * fista_data->image_size); + fista_data->weights = (double *)malloc(sizeof(double) * fista_data->image_size); + + fista_data->x_vector = (double **)malloc(sizeof(double *) * z_size); + fista_data->x_vector_old = (double **)malloc(sizeof(double *) * z_size); + fista_data->y_vector = (double **)malloc(sizeof(double *) * z_size); + for(i=0;ix_vector[i] = (double *)malloc(sizeof(double) * fista_data->image_size); + fista_data->x_vector_old[i] = (double *)malloc(sizeof(double) * fista_data->image_size); + fista_data->y_vector[i] = (double *)malloc(sizeof(double) * fista_data->image_size); + } + + fista_data->fft_vector = (double *)fftw_malloc(sizeof(double) * fista_data->image_size); + fista_data->Ax_fft = (fftw_complex *)fftw_malloc(sizeof(fftw_complex) * fista_data->fft_size); + fista_data->fft_vector_fft = (fftw_complex *)fftw_malloc(sizeof(fftw_complex) * fista_data->fft_size); + fista_data->image_fft = (fftw_complex *)fftw_malloc(sizeof(fftw_complex) * fista_data->fft_size); + + fista_data->psf_fft = (fftw_complex **)malloc(sizeof(fftw_complex *) * z_size); + for(i=0;ipsf_fft[i] = (fftw_complex *)fftw_malloc(sizeof(fftw_complex) * fista_data->fft_size); + } + + /* Create FFT plans. */ + fista_data->fft_forward = fftw_plan_dft_r2c_2d(x_size, y_size, fista_data->fft_vector, fista_data->fft_vector_fft, FFTW_MEASURE); + fista_data->fft_backward = fftw_plan_dft_c2r_2d(x_size, y_size, fista_data->fft_vector_fft, fista_data->fft_vector, FFTW_MEASURE); + + /* + Compute FFTs of the psfs and save in psf_fft. + Note: The input psfs are indexed (x,y,z) but internally we + use [(x1,y1), (x2,y2), ..] as this is hopefully clearer. + */ + for(i=0;iimage_size;j++){ + fista_data->fft_vector[j] = psf[j*z_size+i]; + } + fftw_execute(fista_data->fft_forward); + ftmComplexCopy(fista_data->fft_vector_fft, fista_data->psf_fft[i], fista_data->fft_size); + } + + /* Calculate optimal time step. */ + fista_data->time_step = 0.0; + for(i=0;ipsf_fft[i]; + for(j=0;jfft_size;j++){ + t1 = t2[j][0] * t2[j][0] + t2[j][1] * t2[j][1]; + if(t1>fista_data->time_step){ + fista_data->time_step = t1; + } + } + } + + fista_data->time_step = 1.0/(2.0*fista_data->time_step); + printf("Optimal time step is: %f\n", fista_data->time_step); + + fista_data->time_step = t_step; + + return fista_data; +} + +/* + * iterate() + * + * Performs a single cycle of optimization. + * + * fista_data - A pointer to a fistaData structure. + * lambda - l1 error term weigth. + */ +void iterate(fistaData *fista_data, double lambda) +{ + int i,j; + double lt,new_tk,t1,t2; + double *xv,*xv_old,*yv; + + /* Flag that we need to recaculate Ax (if getAx is called). */ + fista_data->stale_Ax = 1; + + /* Copy current x vector into old x vector. */ + for(i=0;inumber_psfs;i++){ + ftmDoubleCopy(fista_data->x_vector[i], fista_data->x_vector_old[i], fista_data->image_size); + } + + /* + * This is the plk(y) step in the FISTA algorithm. + */ + + /* Compute Ax_fft (n,n). x is generic here, and does not + implicitly refer to x_vector. */ + ftmComplexZero(fista_data->Ax_fft, fista_data->fft_size); + + for(i=0;inumber_psfs;i++){ + + /* Compute FFT of y vector for each z plane. */ + ftmDoubleCopy(fista_data->y_vector[i], fista_data->fft_vector, fista_data->image_size); + fftw_execute(fista_data->fft_forward); + + /* Multiply FFT of y vector by FFT of the PSF for this z plane. */ + ftmComplexMultiplyAccum(fista_data->Ax_fft, fista_data->fft_vector_fft, fista_data->psf_fft[i], fista_data->fft_size, 0); + } + + /* Compute Ax_fft - b_fft (image_fft) (n,n). */ + for(i=0;ifft_size;i++){ + fista_data->Ax_fft[i][0] -= fista_data->image_fft[i][0]; + fista_data->Ax_fft[i][1] -= fista_data->image_fft[i][1]; + } + + /* Compute x = y - At(Ax-b) (z,n,n). */ + for(i=0;inumber_psfs;i++){ + + // Compute inverse FFT of At(Ax-b) for each image plane. + ftmComplexMultiply(fista_data->fft_vector_fft, fista_data->Ax_fft, fista_data->psf_fft[i], fista_data->fft_size, 1); + fftw_execute(fista_data->fft_backward); + + // Update x vector. + t1 = 2.0*fista_data->time_step*fista_data->normalization; + xv = fista_data->x_vector[i]; + yv = fista_data->y_vector[i]; + + // Data Weighted Least Squares update. + if (fista_data->dwls){ + for(j=0;jimage_size;j++){ + xv[j] = yv[j] - t1*fista_data->fft_vector[j]*fista_data->weights[j]; + } + } + + // Least Squares update. + else{ + for(j=0;jimage_size;j++){ + xv[j] = yv[j] - t1*fista_data->fft_vector[j]; + } + } + } + + /* Shrink x vector. */ + lt = fista_data->time_step*lambda; + for(i=0;inumber_psfs;i++){ + xv = fista_data->x_vector[i]; + for(j=0;jimage_size;j++){ + t1 = xv[j]; + t2 = fabs(xv[j]) - lt; + if (t2 <= 0.0){ + xv[j] = 0.0; + } + else{ + xv[j] = (t1 > 0.0) ? t2 : -t2; + } + } + } + + /* + * Update tk term step. + */ + new_tk = 0.5*(1.0 + sqrt(1.0 + 4.0 * fista_data->tk * fista_data->tk)); + + /* + * Compute new y vector step. + */ + t1 = (fista_data->tk - 1.0)/new_tk; + for(i=0;inumber_psfs;i++){ + xv = fista_data->x_vector[i]; + xv_old = fista_data->x_vector_old[i]; + yv = fista_data->y_vector[i]; + for(j=0;jimage_size;j++){ + yv[j] = xv[j] + t1*(xv[j] - xv_old[j]); + } + } + + /* + * Update tk term step. + */ + fista_data->tk = new_tk; +} + +/* + * l1Error + * + * Calculate l1 error term (sum of the x vector). + * + * fista_data - A pointer to a fistaData structure. + * + * Return sum of the x vector. + */ +double l1Error(fistaData *fista_data) +{ + int i,j; + double sum; + double *xv; + + sum = 0.0; + for(i=0;inumber_psfs;i++){ + xv = fista_data->x_vector[i]; + for(j=0;jimage_size;j++){ + sum += fabs(xv[j]); + } + } + + return sum; +} + +/* + * l2Error + * + * Calculate l2 error (Ax - b). + * + * fista_data - A pointer to a fistaData structure. + * + * Return the error in the fit. + */ +double l2Error(fistaData *fista_data) +{ + int i; + double l2_error,t1; + + /* Compute Ax. */ + calculateAx(fista_data); + + /* Compute (Ax - b)^2. */ + l2_error = 0.0; + for(i=0;iimage_size;i++){ + t1 = fista_data->Ax[i] - fista_data->image[i]; + l2_error += t1*t1; + } + + return sqrt(l2_error); +} + +/* + * newImage() + * + * Initialize with a new image. + * + * fista_data - A pointer to a fistaData structure. + * image - The image data (including the background estimate). + * background - The background estimate. + */ +void newImage(fistaData *fista_data, double *image, double *background) +{ + int i; + + /* Flag that we need to recaculate Ax (if getAx is called). */ + fista_data->stale_Ax = 1; + + fista_data->tk = 1.0; + + /* Calculate background corrected image. */ + for(i=0;iimage_size;i++){ + fista_data->image[i] = image[i] - background[i]; + } + + /* Calculate weights. */ + for(i=0;iimage_size;i++){ + fista_data->weights[i] = 1.0/image[i]; + } + + /* Compute FFT of image. */ + ftmDoubleCopy(fista_data->image, fista_data->fft_vector, fista_data->image_size); + fftw_execute(fista_data->fft_forward); + ftmComplexCopy(fista_data->fft_vector_fft, fista_data->image_fft, fista_data->fft_size); + + /* Initialize x_vector, y_vectors */ + for(i=0;inumber_psfs;i++){ + ftmDoubleZero(fista_data->x_vector[i], fista_data->image_size); + ftmDoubleZero(fista_data->x_vector_old[i], fista_data->image_size); + ftmDoubleZero(fista_data->y_vector[i], fista_data->image_size); + } + +} + +/* + * run() + * + * fista_data - A pointer to a fistaData structure. + * lambda - Lambda value. + * cycles - Number of iterations to perform. + * + * Performs a fixed number of cycles at a fixed lambda. + */ +void run(fistaData *fista_data, double lambda, int cycles) +{ + int i; + + for(i=0;i -#include -#include - -#include "psf_fft.h" - -/* Functions */ - -/* - * pFTCalcShiftVector() - * - * Calculate the FFT shift vector. - * - * sr - The real part of the shift vector. - * sc - The complex part of the shift vector. - * dx - Shift delta (in pixels). - * size - The size of the vector. - */ -void pFTCalcShiftVector(double *sr, double *sc, double dx, int size) -{ - int i; - double t1, t2, tc, tr; - - sr[0] = 1.0; - sc[0] = 0.0; - t1 = -M_PI * 2.0 * dx/((double)size); - for(i=1;i<(size/2+1);i++){ - t2 = t1 * (double)i; - tr = cos(t2); - tc = sin(t2); - sr[size-i] = tr; - sr[i] = tr; - sc[size-i] = -tc; - sc[i] = tc; - } -} - -/* - * pFTCalcShiftVectorDerivative() - * - * Calculate the derivative of the FFT shift vector. - * - * sc - The complex part of the derivative of the shift vector. - * size - The size of the vector. - */ -void pFTCalcShiftVectorDerivative(double *sc, int size) -{ - int i; - double t1,t2; - - sc[0] = 0.0; - t1 = -M_PI * 2.0/((double)size); - for(i=1;i<(size/2+1);i++){ - t2 = t1 * (double)i; - sc[size-i] = t2; - sc[i] = -t2; - } -} - -/* - * pFTCleanup() - * - * pfft - A pointer to a psfFFT structure. - */ -void pFTCleanup(psfFFT *pfft) -{ - free(pfft->kx_c); - free(pfft->kx_r); - free(pfft->ky_c); - free(pfft->ky_r); - free(pfft->kz_c); - free(pfft->kz_r); - - fftw_free(pfft->fftw_real); - fftw_free(pfft->fftw_fft); - fftw_free(pfft->ws); - fftw_free(pfft->psf); - - fftw_destroy_plan(pfft->fft_backward); - - free(pfft); -} - -/* - * pFTGetPSF() - * - * Return the current PSF. - * - * pfft - A pointer to a psfFFT structure. - * psf - Pre-allocated storage for the result. - */ -void pFTGetPSF(psfFFT *pfft, double *psf) -{ - int i; - int mid_z, size_xy; - - /* Copy current FFT of PSF into FFTW input, */ - for(i=0;ifft_size;i++){ - pfft->fftw_fft[i][0] = pfft->ws[i][0]; - pfft->fftw_fft[i][1] = pfft->ws[i][1]; - } - - /* Do reverse transform. */ - fftw_execute(pfft->fft_backward); - - /* The 2D PSF is the middle plane of the 3D PSF. */ - size_xy = pfft->x_size*pfft->y_size; - mid_z = size_xy * (pfft->z_size/2); - - for(i=0;ifftw_real[mid_z+i]; - } -} - -/* - * pFTGetPSFdx() - * - * Return the derivative with respect to x of the current PSF. - * - * pfft - A pointer to a psfFFT structure. - * dx - Pre-allocated storage for the result. - */ -void pFTGetPSFdx(psfFFT *pfft, double *dx) -{ - int i,j,k,t1,t2,t3; - int mid_z,size_xy; - double *dd; - - /* Calculate FFT derivative vector. */ - dd = pfft->kx_c; - pFTCalcShiftVectorDerivative(dd, pfft->x_size); - - /* Copy current PF multiplied by kx into FFTW input. */ - for(i=0;iz_size;i++){ - t1 = i * (pfft->y_size * pfft->fft_x_size); - for(j=0;jy_size;j++){ - t2 = j * pfft->fft_x_size; - for(k=0;kfft_x_size;k++){ - t3 = t1 + t2 + k; - - pfft->fftw_fft[t3][0] = pfft->ws[t3][1]*dd[k]; - pfft->fftw_fft[t3][1] = -1.0*pfft->ws[t3][0]*dd[k]; - } - } - } - - /* Do reverse transform. */ - fftw_execute(pfft->fft_backward); - - /* The 2D PSF is the middle plane of the 3D PSF. */ - size_xy = pfft->x_size*pfft->y_size; - mid_z = size_xy * (pfft->z_size/2); - - for(i=0;ifftw_real[mid_z+i]; - } -} - -/* - * pFTGetPSFdy() - * - * Return the derivative with respect to y of the current PSF. - * - * pfft - A pointer to a psfFFT structure. - * dy - Pre-allocated storage for the result. - */ -void pFTGetPSFdy(psfFFT *pfft, double *dy) -{ - int i,j,k,t1,t2,t3; - int mid_z,size_xy; - double *dd; - - /* Calculate FFT derivative vector. */ - dd = pfft->ky_c; - pFTCalcShiftVectorDerivative(dd, pfft->y_size); - - /* Copy current PF multiplied by ky into FFTW input. */ - for(i=0;iz_size;i++){ - t1 = i * (pfft->y_size * pfft->fft_x_size); - for(j=0;jy_size;j++){ - t2 = j * pfft->fft_x_size; - for(k=0;kfft_x_size;k++){ - t3 = t1 + t2 + k; - - pfft->fftw_fft[t3][0] = pfft->ws[t3][1]*dd[j]; - pfft->fftw_fft[t3][1] = -1.0*pfft->ws[t3][0]*dd[j]; - } - } - } - - /* Do reverse transform. */ - fftw_execute(pfft->fft_backward); - - /* The 2D PSF is the middle plane of the 3D PSF. */ - size_xy = pfft->x_size*pfft->y_size; - mid_z = size_xy * (pfft->z_size/2); - - for(i=0;ifftw_real[mid_z+i]; - } -} - -/* - * pFTGetPSFdz() - * - * Return the derivative with respect to z of the current PSF. - * - * pfft - A pointer to a psfFFT structure. - * dz - Pre-allocated storage for the result. - */ -void pFTGetPSFdz(psfFFT *pfft, double *dz) -{ - int i,j,k,t1,t2,t3; - int mid_z,size_xy; - double *dd; - - /* Calculate FFT derivative vector. */ - dd = pfft->kz_c; - pFTCalcShiftVectorDerivative(dd, pfft->z_size); - - /* Copy current PF multiplied by ky into FFTW input. */ - for(i=0;iz_size;i++){ - t1 = i * (pfft->y_size * pfft->fft_x_size); - for(j=0;jy_size;j++){ - t2 = j * pfft->fft_x_size; - for(k=0;kfft_x_size;k++){ - t3 = t1 + t2 + k; - - pfft->fftw_fft[t3][0] = pfft->ws[t3][1]*dd[i]; - pfft->fftw_fft[t3][1] = -1.0*pfft->ws[t3][0]*dd[i]; - } - } - } - - /* Do reverse transform. */ - fftw_execute(pfft->fft_backward); - - /* The 2D PSF is the middle plane of the 3D PSF. */ - size_xy = pfft->x_size*pfft->y_size; - mid_z = size_xy * (pfft->z_size/2); - - for(i=0;ifftw_real[mid_z+i]; - } -} - -/* - * pFTGetXSize() - */ -int pFTGetXSize(psfFFT *pfft) -{ - return pfft->x_size; -} - -/* - * pFTGetYSize() - */ -int pFTGetYSize(psfFFT *pfft) -{ - return pfft->y_size; -} - -/* - * pFTGetZSize() - */ -int pFTGetZSize(psfFFT *pfft) -{ - return pfft->z_size; -} - -/* - * pFTInitialize() - * - * Set things up for FFT based PSF calculations. - * - * Note: In the psf array, the z axis is the slowest followed - * by the y axis and then the x axis. - * - * psf - The psf (z_size, y_size, x_size). - * z_size - The size of the psf in z (slowest dimension). - * y_size - The size of the psf in y. - * x_size - The size of the psf in x (fastest dimension). - * - * return - A psf_fft structure. - */ -psfFFT *pFTInitialize(double *psf, int z_size, int y_size, int x_size) -{ - int i; - double normalization; - fftw_plan fft_forward; - psfFFT *pfft; - - pfft = (psfFFT *)malloc(sizeof(psfFFT)); - - /* Initialize some variables. */ - pfft->x_size = x_size; - pfft->y_size = y_size; - pfft->z_size = z_size; - pfft->psf_size = z_size * y_size * x_size; - - pfft->fft_x_size = (x_size/2 + 1); - pfft->fft_size = z_size * y_size * (x_size/2 + 1); - - normalization = 1.0/((double)(pfft->psf_size)); - - /* Allocate storage. */ - pfft->kx_c = (double *)malloc(sizeof(double)*x_size); - pfft->kx_r = (double *)malloc(sizeof(double)*x_size); - pfft->ky_c = (double *)malloc(sizeof(double)*y_size); - pfft->ky_r = (double *)malloc(sizeof(double)*y_size); - pfft->kz_c = (double *)malloc(sizeof(double)*z_size); - pfft->kz_r = (double *)malloc(sizeof(double)*z_size); - - pfft->fftw_real = (double *)fftw_malloc(sizeof(double)*pfft->psf_size); - pfft->fftw_fft = (fftw_complex *)fftw_malloc(sizeof(fftw_complex)*pfft->fft_size); - pfft->ws = (fftw_complex *)fftw_malloc(sizeof(fftw_complex)*pfft->fft_size); - pfft->psf = (fftw_complex *)fftw_malloc(sizeof(fftw_complex)*pfft->fft_size); - - /* Create FFT plan. */ - pfft->fft_backward = fftw_plan_dft_c2r_3d(z_size, y_size, x_size, pfft->fftw_fft, pfft->fftw_real, FFTW_MEASURE); - - /* Compute FFT of psf and save. */ - fft_forward = fftw_plan_dft_r2c_3d(z_size, y_size, x_size, pfft->fftw_real, pfft->fftw_fft, FFTW_ESTIMATE); - - for(i=0;ipsf_size;i++){ - pfft->fftw_real[i] = psf[i] * normalization; - } - - fftw_execute(fft_forward); - - for(i=0;ifft_size;i++){ - pfft->psf[i][0] = pfft->fftw_fft[i][0]; - pfft->psf[i][1] = pfft->fftw_fft[i][1]; - pfft->ws[i][0] = pfft->fftw_fft[i][0]; - pfft->ws[i][1] = pfft->fftw_fft[i][1]; - } - - fftw_destroy_plan(fft_forward); - - return pfft; -} - -/* - * pFTTranslate() - * - * Translate the psf by dx, dy, dz. - */ -void pFTTranslate(psfFFT *pfft, double dx, double dy, double dz) -{ - int i,j,k,t1,t2,t3; - double c1,c2,r1,r2; - double *sxc,*sxr,*syc,*syr,*szc,*szr; - - /* Calculate FFT translation vectors. */ - sxc = pfft->kx_c; - sxr = pfft->kx_r; - pFTCalcShiftVector(sxr, sxc, dx, pfft->x_size); - - syc = pfft->ky_c; - syr = pfft->ky_r; - pFTCalcShiftVector(syr, syc, dy, pfft->y_size); - - /* We use -dz to match the convention of Z convention of simulator.pupil_math. */ - szc = pfft->kz_c; - szr = pfft->kz_r; - pFTCalcShiftVector(szr, szc, -dz, pfft->z_size); - - /* Translate */ - for(i=0;iz_size;i++){ - t1 = i * (pfft->y_size * pfft->fft_x_size); - for(j=0;jy_size;j++){ - t2 = j * pfft->fft_x_size; - for(k=0;kfft_x_size;k++){ - t3 = t1 + t2 + k; - - r1 = pfft->psf[t3][0]; - c1 = pfft->psf[t3][1]; - - /* z shift. */ - r2 = r1 * szr[i] - c1 * szc[i]; - c2 = r1 * szc[i] + c1 * szr[i]; - - /* y shift. */ - r1 = r2 * syr[j] - c2 * syc[j]; - c1 = r2 * syc[j] + c2 * syr[j]; - - /* x shift. */ - r2 = r1 * sxr[k] - c1 * sxc[k]; - c2 = r1 * sxc[k] + c1 * sxr[k]; - - pfft->ws[t3][0] = r2; - pfft->ws[t3][1] = c2; - } - } - } -} +/* + * C library for PSF manipulation using the FFT. + * + * This is designed to work in the context of a fitting algorithm, + * so basically the steps are: + * + * 1. pFTInitialize() to set things up. + * 2. A. pFTTranslate() to adjust position in x,y,z. + * B. pFTGetPSF() to get the PSF at the adjusted position. + * C. pFTGetPSFdx() to get the derivative of the PSF at the adjusted + * position. + * D. ... + * + * So the PSF will stay at the adjusted position until pFTTranslate() + * is called again. + * + * Note: The boundary conditions are periodic, so the size of the + * PSF should be large enough that it goes to zero at the edges. + * + * Hazen 10/17. + */ + + +/* Include */ +#include +#include +#include + +#include "psf_fft.h" + +/* Functions */ + +/* + * pFTCalcShiftVector() + * + * Calculate the FFT shift vector. + * + * sr - The real part of the shift vector. + * sc - The complex part of the shift vector. + * dx - Shift delta (in pixels). + * size - The size of the vector. + */ +void pFTCalcShiftVector(double *sr, double *sc, double dx, int size) +{ + int i; + double t1, t2, tc, tr; + + sr[0] = 1.0; + sc[0] = 0.0; + t1 = -M_PI * 2.0 * dx/((double)size); + for(i=1;i<(size/2+1);i++){ + t2 = t1 * (double)i; + tr = cos(t2); + tc = sin(t2); + sr[size-i] = tr; + sr[i] = tr; + sc[size-i] = -tc; + sc[i] = tc; + } +} + +/* + * pFTCalcShiftVectorDerivative() + * + * Calculate the derivative of the FFT shift vector. + * + * sc - The complex part of the derivative of the shift vector. + * size - The size of the vector. + */ +void pFTCalcShiftVectorDerivative(double *sc, int size) +{ + int i; + double t1,t2; + + sc[0] = 0.0; + t1 = -M_PI * 2.0/((double)size); + for(i=1;i<(size/2+1);i++){ + t2 = t1 * (double)i; + sc[size-i] = t2; + sc[i] = -t2; + } +} + +/* + * pFTCleanup() + * + * pfft - A pointer to a psfFFT structure. + */ +void pFTCleanup(psfFFT *pfft) +{ + free(pfft->kx_c); + free(pfft->kx_r); + free(pfft->ky_c); + free(pfft->ky_r); + free(pfft->kz_c); + free(pfft->kz_r); + + fftw_free(pfft->fftw_real); + fftw_free(pfft->fftw_fft); + fftw_free(pfft->ws); + fftw_free(pfft->psf); + + fftw_destroy_plan(pfft->fft_backward); + + free(pfft); +} + +/* + * pFTGetPSF() + * + * Return the current PSF. + * + * pfft - A pointer to a psfFFT structure. + * psf - Pre-allocated storage for the result. + */ +void pFTGetPSF(psfFFT *pfft, double *psf) +{ + int i; + int mid_z, size_xy; + + /* Copy current FFT of PSF into FFTW input, */ + for(i=0;ifft_size;i++){ + pfft->fftw_fft[i][0] = pfft->ws[i][0]; + pfft->fftw_fft[i][1] = pfft->ws[i][1]; + } + + /* Do reverse transform. */ + fftw_execute(pfft->fft_backward); + + /* The 2D PSF is the middle plane of the 3D PSF. */ + size_xy = pfft->x_size*pfft->y_size; + mid_z = size_xy * (pfft->z_size/2); + + for(i=0;ifftw_real[mid_z+i]; + } +} + +/* + * pFTGetPSFdx() + * + * Return the derivative with respect to x of the current PSF. + * + * pfft - A pointer to a psfFFT structure. + * dx - Pre-allocated storage for the result. + */ +void pFTGetPSFdx(psfFFT *pfft, double *dx) +{ + int i,j,k,t1,t2,t3; + int mid_z,size_xy; + double *dd; + + /* Calculate FFT derivative vector. */ + dd = pfft->kx_c; + pFTCalcShiftVectorDerivative(dd, pfft->x_size); + + /* Copy current PF multiplied by kx into FFTW input. */ + for(i=0;iz_size;i++){ + t1 = i * (pfft->y_size * pfft->fft_x_size); + for(j=0;jy_size;j++){ + t2 = j * pfft->fft_x_size; + for(k=0;kfft_x_size;k++){ + t3 = t1 + t2 + k; + + pfft->fftw_fft[t3][0] = pfft->ws[t3][1]*dd[k]; + pfft->fftw_fft[t3][1] = -1.0*pfft->ws[t3][0]*dd[k]; + } + } + } + + /* Do reverse transform. */ + fftw_execute(pfft->fft_backward); + + /* The 2D PSF is the middle plane of the 3D PSF. */ + size_xy = pfft->x_size*pfft->y_size; + mid_z = size_xy * (pfft->z_size/2); + + for(i=0;ifftw_real[mid_z+i]; + } +} + +/* + * pFTGetPSFdy() + * + * Return the derivative with respect to y of the current PSF. + * + * pfft - A pointer to a psfFFT structure. + * dy - Pre-allocated storage for the result. + */ +void pFTGetPSFdy(psfFFT *pfft, double *dy) +{ + int i,j,k,t1,t2,t3; + int mid_z,size_xy; + double *dd; + + /* Calculate FFT derivative vector. */ + dd = pfft->ky_c; + pFTCalcShiftVectorDerivative(dd, pfft->y_size); + + /* Copy current PF multiplied by ky into FFTW input. */ + for(i=0;iz_size;i++){ + t1 = i * (pfft->y_size * pfft->fft_x_size); + for(j=0;jy_size;j++){ + t2 = j * pfft->fft_x_size; + for(k=0;kfft_x_size;k++){ + t3 = t1 + t2 + k; + + pfft->fftw_fft[t3][0] = pfft->ws[t3][1]*dd[j]; + pfft->fftw_fft[t3][1] = -1.0*pfft->ws[t3][0]*dd[j]; + } + } + } + + /* Do reverse transform. */ + fftw_execute(pfft->fft_backward); + + /* The 2D PSF is the middle plane of the 3D PSF. */ + size_xy = pfft->x_size*pfft->y_size; + mid_z = size_xy * (pfft->z_size/2); + + for(i=0;ifftw_real[mid_z+i]; + } +} + +/* + * pFTGetPSFdz() + * + * Return the derivative with respect to z of the current PSF. + * + * pfft - A pointer to a psfFFT structure. + * dz - Pre-allocated storage for the result. + */ +void pFTGetPSFdz(psfFFT *pfft, double *dz) +{ + int i,j,k,t1,t2,t3; + int mid_z,size_xy; + double *dd; + + /* Calculate FFT derivative vector. */ + dd = pfft->kz_c; + pFTCalcShiftVectorDerivative(dd, pfft->z_size); + + /* Copy current PF multiplied by ky into FFTW input. */ + for(i=0;iz_size;i++){ + t1 = i * (pfft->y_size * pfft->fft_x_size); + for(j=0;jy_size;j++){ + t2 = j * pfft->fft_x_size; + for(k=0;kfft_x_size;k++){ + t3 = t1 + t2 + k; + + pfft->fftw_fft[t3][0] = pfft->ws[t3][1]*dd[i]; + pfft->fftw_fft[t3][1] = -1.0*pfft->ws[t3][0]*dd[i]; + } + } + } + + /* Do reverse transform. */ + fftw_execute(pfft->fft_backward); + + /* The 2D PSF is the middle plane of the 3D PSF. */ + size_xy = pfft->x_size*pfft->y_size; + mid_z = size_xy * (pfft->z_size/2); + + for(i=0;ifftw_real[mid_z+i]; + } +} + +/* + * pFTGetXSize() + */ +int pFTGetXSize(psfFFT *pfft) +{ + return pfft->x_size; +} + +/* + * pFTGetYSize() + */ +int pFTGetYSize(psfFFT *pfft) +{ + return pfft->y_size; +} + +/* + * pFTGetZSize() + */ +int pFTGetZSize(psfFFT *pfft) +{ + return pfft->z_size; +} + +/* + * pFTInitialize() + * + * Set things up for FFT based PSF calculations. + * + * Note: In the psf array, the z axis is the slowest followed + * by the y axis and then the x axis. + * + * psf - The psf (z_size, y_size, x_size). + * z_size - The size of the psf in z (slowest dimension). + * y_size - The size of the psf in y. + * x_size - The size of the psf in x (fastest dimension). + * + * return - A psf_fft structure. + */ +psfFFT *pFTInitialize(double *psf, int z_size, int y_size, int x_size) +{ + int i; + double normalization; + fftw_plan fft_forward; + psfFFT *pfft; + + pfft = (psfFFT *)malloc(sizeof(psfFFT)); + + /* Initialize some variables. */ + pfft->x_size = x_size; + pfft->y_size = y_size; + pfft->z_size = z_size; + pfft->psf_size = z_size * y_size * x_size; + + pfft->fft_x_size = (x_size/2 + 1); + pfft->fft_size = z_size * y_size * (x_size/2 + 1); + + normalization = 1.0/((double)(pfft->psf_size)); + + /* Allocate storage. */ + pfft->kx_c = (double *)malloc(sizeof(double)*x_size); + pfft->kx_r = (double *)malloc(sizeof(double)*x_size); + pfft->ky_c = (double *)malloc(sizeof(double)*y_size); + pfft->ky_r = (double *)malloc(sizeof(double)*y_size); + pfft->kz_c = (double *)malloc(sizeof(double)*z_size); + pfft->kz_r = (double *)malloc(sizeof(double)*z_size); + + pfft->fftw_real = (double *)fftw_malloc(sizeof(double)*pfft->psf_size); + pfft->fftw_fft = (fftw_complex *)fftw_malloc(sizeof(fftw_complex)*pfft->fft_size); + pfft->ws = (fftw_complex *)fftw_malloc(sizeof(fftw_complex)*pfft->fft_size); + pfft->psf = (fftw_complex *)fftw_malloc(sizeof(fftw_complex)*pfft->fft_size); + + /* Create FFT plan. */ + pfft->fft_backward = fftw_plan_dft_c2r_3d(z_size, y_size, x_size, pfft->fftw_fft, pfft->fftw_real, FFTW_MEASURE); + + /* Compute FFT of psf and save. */ + fft_forward = fftw_plan_dft_r2c_3d(z_size, y_size, x_size, pfft->fftw_real, pfft->fftw_fft, FFTW_ESTIMATE); + + for(i=0;ipsf_size;i++){ + pfft->fftw_real[i] = psf[i] * normalization; + } + + fftw_execute(fft_forward); + + for(i=0;ifft_size;i++){ + pfft->psf[i][0] = pfft->fftw_fft[i][0]; + pfft->psf[i][1] = pfft->fftw_fft[i][1]; + pfft->ws[i][0] = pfft->fftw_fft[i][0]; + pfft->ws[i][1] = pfft->fftw_fft[i][1]; + } + + fftw_destroy_plan(fft_forward); + + return pfft; +} + +/* + * pFTTranslate() + * + * Translate the psf by dx, dy, dz. + */ +void pFTTranslate(psfFFT *pfft, double dx, double dy, double dz) +{ + int i,j,k,t1,t2,t3; + double c1,c2,r1,r2; + double *sxc,*sxr,*syc,*syr,*szc,*szr; + + /* Calculate FFT translation vectors. */ + sxc = pfft->kx_c; + sxr = pfft->kx_r; + pFTCalcShiftVector(sxr, sxc, dx, pfft->x_size); + + syc = pfft->ky_c; + syr = pfft->ky_r; + pFTCalcShiftVector(syr, syc, dy, pfft->y_size); + + /* We use -dz to match the convention of Z convention of simulator.pupil_math. */ + szc = pfft->kz_c; + szr = pfft->kz_r; + pFTCalcShiftVector(szr, szc, -dz, pfft->z_size); + + /* Translate */ + for(i=0;iz_size;i++){ + t1 = i * (pfft->y_size * pfft->fft_x_size); + for(j=0;jy_size;j++){ + t2 = j * pfft->fft_x_size; + for(k=0;kfft_x_size;k++){ + t3 = t1 + t2 + k; + + r1 = pfft->psf[t3][0]; + c1 = pfft->psf[t3][1]; + + /* z shift. */ + r2 = r1 * szr[i] - c1 * szc[i]; + c2 = r1 * szc[i] + c1 * szr[i]; + + /* y shift. */ + r1 = r2 * syr[j] - c2 * syc[j]; + c1 = r2 * syc[j] + c2 * syr[j]; + + /* x shift. */ + r2 = r1 * sxr[k] - c1 * sxc[k]; + c2 = r1 * sxc[k] + c1 * sxr[k]; + + pfft->ws[t3][0] = r2; + pfft->ws[t3][1] = c2; + } + } + } +} diff --git a/storm_analysis/sa_library/drift_utilities.py b/storm_analysis/sa_library/drift_utilities.py index 26b36d47..121a63f4 100644 --- a/storm_analysis/sa_library/drift_utilities.py +++ b/storm_analysis/sa_library/drift_utilities.py @@ -1,193 +1,193 @@ -#!/usr/bin/env python -""" -Utility classes functions that are used for drift correction. - -Hazen 02/17 -""" - -import numpy -import scipy - -import storm_analysis.sa_library.grid_c as gridC -import storm_analysis.sa_library.sa_h5py as saH5Py - - -class SAH5DriftCorrection(saH5Py.SAH5Py): - """ - A sub-class of SAH5Py designed for use in drift correction. - """ - def __init__(self, scale = None, z_bins = 1, **kwds): - super(SAH5DriftCorrection, self).__init__(**kwds) - - self.dx = 0.0 - self.dy = 0.0 - self.dz = 0.0 - self.fmin = None - self.fmax = None - self.im_shape_2D = (self.hdf5.attrs['movie_x']*scale, - self.hdf5.attrs['movie_y']*scale) - self.im_shape_3D = (self.hdf5.attrs['movie_x']*scale, - self.hdf5.attrs['movie_y']*scale, - z_bins) - self.scale = scale - self.z_bins = z_bins - - def grid2D(self, drift_corrected = False): - image = numpy.zeros(self.im_shape_2D, dtype = numpy.int32) - for locs in self.locsInFrameRangeIterator(self.fmin, self.fmax, ["x", "y"]): - if drift_corrected: - locs["x"] += self.dx - locs["y"] += self.dy - i_x = numpy.floor(locs["x"]*self.scale).astype(numpy.int32) - i_y = numpy.floor(locs["y"]*self.scale).astype(numpy.int32) - gridC.grid2D(i_x, i_y, image) - return image - - def grid3D(self, z_min, z_max, drift_corrected = False): - z_scale = float(self.z_bins)/(z_max - z_min) - image = numpy.zeros(self.im_shape_3D, dtype = numpy.int32) - for locs in self.locsInFrameRangeIterator(self.fmin, self.fmax, ["x", "y", "z"]): - - # Create z value filter. - # - # We filter here rather than just relying on gridC.grid3D as application - # of z drift correction could move out of z range peaks into the acceptable - # range. - # - mask = (locs["z"] > z_min) & (locs["z"] < z_max) - if (numpy.count_nonzero(mask) == 0): - continue - - # Remove localizations that are out of range. - locs["x"] = locs["x"][mask] - locs["y"] = locs["y"][mask] - locs["z"] = locs["z"][mask] - - # Apply drift correction if requested. - if drift_corrected: - locs["x"] += self.dx - locs["y"] += self.dy - locs["z"] += self.dz - - # Add to image. - i_x = numpy.floor(locs["x"]*self.scale).astype(numpy.int32) - i_y = numpy.floor(locs["y"]*self.scale).astype(numpy.int32) - i_z = numpy.floor((locs["z"] - z_min)*z_scale).astype(numpy.int32) - gridC.grid3D(i_x, i_y, i_z, image) - return image - - def locsInFrameRangeIterator(self, start, stop, fields): - for i in range(start, stop): - locs = self.getLocalizationsInFrame(i, - drift_corrected = False, - fields = fields) - - # Skip empty frames. - if not bool(locs): - continue - - yield locs - - def saveDriftData(self, all_dx, all_dy, all_dz): - """ - Store drift correction data in the HDF5 file. The all_** arrays - contain the drift corrections for every frame in the movie in - units of pixels (X,Y) or microns (Z). - """ - assert(len(all_dx) == self.getMovieLength()) - for i in range(self.getMovieLength()): - try: - self.setDriftCorrection(i, - dx = all_dx[i], - dy = all_dy[i], - dz = all_dz[i]) - except saH5Py.SAH5PyException: - pass - - def setDriftCorrectionXY(self, dx, dy): - self.dx = dx - self.dy = dy - - def setDriftCorrectionZ(self, dz): - self.dz = dz - - def setFrameRange(self, fmin, fmax): - self.fmin = fmin - self.fmax = fmax - - -class SAH5DriftCorrectionTest(SAH5DriftCorrection): - """ - A sub-class of SAH5PyDriftCorrection for testing purposes. - """ - def locsInFrameRangeIterator(self, start, stop, fields): - for i in range(start, stop): - locs = self.getLocalizationsInFrame(i, - drift_corrected = True, - fields = fields) - yield locs - - -def interpolateData(xvals, yvals, film_l): - """ - Interpolate drift data to the length of the film. - """ - - final_drift = numpy.zeros(film_l) - - # Use polyfit for extrapolation at the end points. - pe = numpy.poly1d(numpy.polyfit(xvals[0:2], yvals[0:2], 1)) - for i in range(int(xvals[0])): - final_drift[i] = pe(i) - - pe = numpy.poly1d(numpy.polyfit(xvals[-2:], yvals[-2:], 1)) - for i in range(int(xvals[-1]), film_l): - final_drift[i] = pe(i) - - # Create linear spline for interpolation. - sp = scipy.interpolate.interp1d(xvals, yvals, kind = "linear") - - # Interpolate. - i = int(xvals[0]) - while (i <= int(xvals[-1])): - final_drift[i] = sp(i) - i += 1 - - return final_drift - -def saveDriftData(filename, fdx, fdy, fdz): - """ - Save the x,y and z drift data to a file. - """ - frames = numpy.arange(fdx.size) + 1 - numpy.savetxt(filename, - numpy.column_stack((frames, - -fdx, - -fdy, - -fdz)), - fmt = "%d\t%.3f\t%.3f\t%.3f") - - -# -# The MIT License -# -# Copyright (c) 2017 Zhuang Lab, Harvard University -# -# Permission is hereby granted, free of charge, to any person obtaining a copy -# of this software and associated documentation files (the "Software"), to deal -# in the Software without restriction, including without limitation the rights -# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -# copies of the Software, and to permit persons to whom the Software is -# furnished to do so, subject to the following conditions: -# -# The above copyright notice and this permission notice shall be included in -# all copies or substantial portions of the Software. -# -# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN -# THE SOFTWARE. -# +#!/usr/bin/env python +""" +Utility classes functions that are used for drift correction. + +Hazen 02/17 +""" + +import numpy +import scipy + +import storm_analysis.sa_library.grid_c as gridC +import storm_analysis.sa_library.sa_h5py as saH5Py + + +class SAH5DriftCorrection(saH5Py.SAH5Py): + """ + A sub-class of SAH5Py designed for use in drift correction. + """ + def __init__(self, scale = None, z_bins = 1, **kwds): + super(SAH5DriftCorrection, self).__init__(**kwds) + + self.dx = 0.0 + self.dy = 0.0 + self.dz = 0.0 + self.fmin = None + self.fmax = None + self.im_shape_2D = (self.hdf5.attrs['movie_x']*scale, + self.hdf5.attrs['movie_y']*scale) + self.im_shape_3D = (self.hdf5.attrs['movie_x']*scale, + self.hdf5.attrs['movie_y']*scale, + z_bins) + self.scale = scale + self.z_bins = z_bins + + def grid2D(self, drift_corrected = False): + image = numpy.zeros(self.im_shape_2D, dtype = numpy.int32) + for locs in self.locsInFrameRangeIterator(self.fmin, self.fmax, ["x", "y"]): + if drift_corrected: + locs["x"] += self.dx + locs["y"] += self.dy + i_x = numpy.floor(locs["x"]*self.scale).astype(numpy.int32) + i_y = numpy.floor(locs["y"]*self.scale).astype(numpy.int32) + gridC.grid2D(i_x, i_y, image) + return image + + def grid3D(self, z_min, z_max, drift_corrected = False): + z_scale = float(self.z_bins)/(z_max - z_min) + image = numpy.zeros(self.im_shape_3D, dtype = numpy.int32) + for locs in self.locsInFrameRangeIterator(self.fmin, self.fmax, ["x", "y", "z"]): + + # Create z value filter. + # + # We filter here rather than just relying on gridC.grid3D as application + # of z drift correction could move out of z range peaks into the acceptable + # range. + # + mask = (locs["z"] > z_min) & (locs["z"] < z_max) + if (numpy.count_nonzero(mask) == 0): + continue + + # Remove localizations that are out of range. + locs["x"] = locs["x"][mask] + locs["y"] = locs["y"][mask] + locs["z"] = locs["z"][mask] + + # Apply drift correction if requested. + if drift_corrected: + locs["x"] += self.dx + locs["y"] += self.dy + locs["z"] += self.dz + + # Add to image. + i_x = numpy.floor(locs["x"]*self.scale).astype(numpy.int32) + i_y = numpy.floor(locs["y"]*self.scale).astype(numpy.int32) + i_z = numpy.floor((locs["z"] - z_min)*z_scale).astype(numpy.int32) + gridC.grid3D(i_x, i_y, i_z, image) + return image + + def locsInFrameRangeIterator(self, start, stop, fields): + for i in range(start, stop): + locs = self.getLocalizationsInFrame(i, + drift_corrected = False, + fields = fields) + + # Skip empty frames. + if not bool(locs): + continue + + yield locs + + def saveDriftData(self, all_dx, all_dy, all_dz): + """ + Store drift correction data in the HDF5 file. The all_** arrays + contain the drift corrections for every frame in the movie in + units of pixels (X,Y) or microns (Z). + """ + assert(len(all_dx) == self.getMovieLength()) + for i in range(self.getMovieLength()): + try: + self.setDriftCorrection(i, + dx = all_dx[i], + dy = all_dy[i], + dz = all_dz[i]) + except saH5Py.SAH5PyException: + pass + + def setDriftCorrectionXY(self, dx, dy): + self.dx = dx + self.dy = dy + + def setDriftCorrectionZ(self, dz): + self.dz = dz + + def setFrameRange(self, fmin, fmax): + self.fmin = fmin + self.fmax = fmax + + +class SAH5DriftCorrectionTest(SAH5DriftCorrection): + """ + A sub-class of SAH5PyDriftCorrection for testing purposes. + """ + def locsInFrameRangeIterator(self, start, stop, fields): + for i in range(start, stop): + locs = self.getLocalizationsInFrame(i, + drift_corrected = True, + fields = fields) + yield locs + + +def interpolateData(xvals, yvals, film_l): + """ + Interpolate drift data to the length of the film. + """ + + final_drift = numpy.zeros(film_l) + + # Use polyfit for extrapolation at the end points. + pe = numpy.poly1d(numpy.polyfit(xvals[0:2], yvals[0:2], 1)) + for i in range(int(xvals[0])): + final_drift[i] = pe(i) + + pe = numpy.poly1d(numpy.polyfit(xvals[-2:], yvals[-2:], 1)) + for i in range(int(xvals[-1]), film_l): + final_drift[i] = pe(i) + + # Create linear spline for interpolation. + sp = scipy.interpolate.interp1d(xvals, yvals, kind = "linear") + + # Interpolate. + i = int(xvals[0]) + while (i <= int(xvals[-1])): + final_drift[i] = sp(i) + i += 1 + + return final_drift + +def saveDriftData(filename, fdx, fdy, fdz): + """ + Save the x,y and z drift data to a file. + """ + frames = numpy.arange(fdx.size) + 1 + numpy.savetxt(filename, + numpy.column_stack((frames, + -fdx, + -fdy, + -fdz)), + fmt = "%d\t%.3f\t%.3f\t%.3f") + + +# +# The MIT License +# +# Copyright (c) 2017 Zhuang Lab, Harvard University +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. +# diff --git a/storm_analysis/sa_library/matched_filter.c b/storm_analysis/sa_library/matched_filter.c index d4c827fc..8469c0a8 100644 --- a/storm_analysis/sa_library/matched_filter.c +++ b/storm_analysis/sa_library/matched_filter.c @@ -1,239 +1,239 @@ -/* - * C library for the matched filter approach. Basically - * this is just convolving the image with the PSF. - * - * This uses an FFT to do the convolution so there will - * be edge effects. - * - * Hazen 3/16 - * - * Compilation instructions: - * - * Linux: - * gcc -fPIC -g -c -Wall matched_filter.c - * gcc -shared -Wl,-soname,matched_filter.so.1 -o matched_filter.so.1.0.1 matched_filter.o -lc -lfftw3 - * ln -s matched_filter.so.1.0.1 matched_filter.so - * - * Windows: - * gcc -c -O3 matched_filter.c - * gcc -shared -o matched_filter.dll matched_filter.o -lfftw3-3 c:\path\to\libfftw3-3.dll - */ - -/* Include */ -#include -#include -#include - -#include - -#include "ft_math.h" - - -/* Structures & Types */ -struct filter_struct { - int fft_size; - int image_size; - int x_size; - int y_size; - - double max_diff; - - double *fft_vector; - double *old_image; - - fftw_plan fft_backward; - fftw_plan fft_forward; - - fftw_complex *fft_vector_fft; - fftw_complex *psf_fft; -}; -typedef struct filter_struct filter; - -/* Function Declarations */ -void cleanup(filter *); -void convolve(filter *, double *, double *); -void convolveMemo(filter *, double *, double *); -filter *initialize(double *, double, int, int, int); - -/* Functions */ - - -/* - * cleanup() - * - * flt - A pointer to a filter structure. - */ -void cleanup(filter *flt) -{ - if (flt->old_image != NULL){ - free(flt->old_image); - } - - fftw_free(flt->fft_vector); - - fftw_destroy_plan(flt->fft_backward); - fftw_destroy_plan(flt->fft_forward); - - fftw_free(flt->fft_vector_fft); - fftw_free(flt->psf_fft); - - free(flt); -} - - -/* - * convolve() - * - * Convolve image with psf. - * - * flt - A pointer to a filter structure. - * image - The image (must be the same size as the original psf image). - * result - Pre-allocated storage for the result of the convolution. - */ -void convolve(filter *flt, double *image, double *result) -{ - /* Compute FFT of the image. */ - ftmDoubleCopy(image, flt->fft_vector, flt->image_size); - fftw_execute(flt->fft_forward); - - /* Multiple by FFT of the PSF and compute inverse FFT. */ - ftmComplexMultiply(flt->fft_vector_fft, flt->fft_vector_fft, flt->psf_fft, flt->fft_size, 0); - fftw_execute(flt->fft_backward); - - /* Copy into result, */ - ftmDoubleCopy(flt->fft_vector, result, flt->image_size); -} - - -/* - * convolveMemo() - * - * Convolve image with psf, but only if image is different enough from - * the previous image, otherwise just return the previous result. - * - * flt - A pointer to a filter structure. - * image - The image (must be the same size as the original psf image). - * result - Pre-allocated storage for the result of the convolution. - */ -void convolveMemo(filter *flt, double *image, double *result) -{ - int i,different; - - different = 0; - - /* Just check, don't copy so that we don't slowly drift.. */ - for(i=0;iimage_size;i++){ - if(fabs(image[i] - flt->old_image[i]) > flt->max_diff){ - different = 1; - break; - } - } - - if (different){ - - /* Copy into old image. */ - for(i=0;iimage_size;i++){ - flt->old_image[i] = image[i]; - } - //ftmDoubleCopy(image, flt->old_image, flt->image_size); - - /* Do the convolution. */ - convolve(flt, image, result); - } - else { - - /* Otherwise just return the previous result. */ - ftmDoubleCopy(flt->fft_vector, result, flt->image_size); - } -} - - -/* - * initialize() - * - * Set things up for FFT convolution. - * - * psf - the psf (x_size, y_size). - * max_diff - if this it not zero configure for memoization of results. - * x_size - the size of the psf in x (slow dimension). - * y_size - the size of the psf in y (fast dimension). - * estimate - 0/1 to just use an estimated FFT plan. If you are only going to - * to use the FFT a few times this can be much faster. - */ -filter *initialize(double *psf, double max_diff, int x_size, int y_size, int estimate) -{ - int i; - double normalization; - filter *flt; - - flt = (filter *)malloc(sizeof(filter)); - - /* Initialize some variables. */ - flt->fft_size = x_size * (y_size/2 + 1); - flt->image_size = x_size * y_size; - - flt->x_size = x_size; - flt->y_size = y_size; - normalization = 1.0/((double)(x_size * y_size)); - - /* Check whether we are memoizing. */ - if(max_diff > 0.0){ - flt->max_diff = max_diff; - flt->old_image = (double *)malloc(sizeof(double)*flt->image_size); - - for(i=0;iimage_size;i++){ - flt->old_image[i] = 0; - } - } - else{ - flt->max_diff = 0.0; - flt->old_image = NULL; - } - - /* Allocate storage. */ - flt->fft_vector = (double *)fftw_malloc(sizeof(double)*flt->image_size); - flt->fft_vector_fft = (fftw_complex *)fftw_malloc(sizeof(fftw_complex)*flt->fft_size); - flt->psf_fft = (fftw_complex *)fftw_malloc(sizeof(fftw_complex)*flt->fft_size); - - /* Create FFT plans. */ - if (estimate){ - flt->fft_forward = fftw_plan_dft_r2c_2d(x_size, y_size, flt->fft_vector, flt->fft_vector_fft, FFTW_ESTIMATE); - flt->fft_backward = fftw_plan_dft_c2r_2d(x_size, y_size, flt->fft_vector_fft, flt->fft_vector, FFTW_ESTIMATE); - } - else { - flt->fft_forward = fftw_plan_dft_r2c_2d(x_size, y_size, flt->fft_vector, flt->fft_vector_fft, FFTW_MEASURE); - flt->fft_backward = fftw_plan_dft_c2r_2d(x_size, y_size, flt->fft_vector_fft, flt->fft_vector, FFTW_MEASURE); - } - - /* Compute FFT of psf and save. */ - - ftmDoubleCopy(psf, flt->fft_vector, flt->image_size); - fftw_execute(flt->fft_forward); - ftmComplexCopyNormalize(flt->fft_vector_fft, flt->psf_fft, normalization, flt->fft_size); - - return flt; -} - -/* - * The MIT License - * - * Copyright (c) 2016 Zhuang Lab, Harvard University - * - * Permission is hereby granted, free of charge, to any person obtaining a copy - * of this software and associated documentation files (the "Software"), to deal - * in the Software without restriction, including without limitation the rights - * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell - * copies of the Software, and to permit persons to whom the Software is - * furnished to do so, subject to the following conditions: - * - * The above copyright notice and this permission notice shall be included in - * all copies or substantial portions of the Software. - * - * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR - * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, - * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE - * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER - * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, - * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN - * THE SOFTWARE. - */ +/* + * C library for the matched filter approach. Basically + * this is just convolving the image with the PSF. + * + * This uses an FFT to do the convolution so there will + * be edge effects. + * + * Hazen 3/16 + * + * Compilation instructions: + * + * Linux: + * gcc -fPIC -g -c -Wall matched_filter.c + * gcc -shared -Wl,-soname,matched_filter.so.1 -o matched_filter.so.1.0.1 matched_filter.o -lc -lfftw3 + * ln -s matched_filter.so.1.0.1 matched_filter.so + * + * Windows: + * gcc -c -O3 matched_filter.c + * gcc -shared -o matched_filter.dll matched_filter.o -lfftw3-3 c:\path\to\libfftw3-3.dll + */ + +/* Include */ +#include +#include +#include + +#include + +#include "ft_math.h" + + +/* Structures & Types */ +struct filter_struct { + int fft_size; + int image_size; + int x_size; + int y_size; + + double max_diff; + + double *fft_vector; + double *old_image; + + fftw_plan fft_backward; + fftw_plan fft_forward; + + fftw_complex *fft_vector_fft; + fftw_complex *psf_fft; +}; +typedef struct filter_struct filter; + +/* Function Declarations */ +void cleanup(filter *); +void convolve(filter *, double *, double *); +void convolveMemo(filter *, double *, double *); +filter *initialize(double *, double, int, int, int); + +/* Functions */ + + +/* + * cleanup() + * + * flt - A pointer to a filter structure. + */ +void cleanup(filter *flt) +{ + if (flt->old_image != NULL){ + free(flt->old_image); + } + + fftw_free(flt->fft_vector); + + fftw_destroy_plan(flt->fft_backward); + fftw_destroy_plan(flt->fft_forward); + + fftw_free(flt->fft_vector_fft); + fftw_free(flt->psf_fft); + + free(flt); +} + + +/* + * convolve() + * + * Convolve image with psf. + * + * flt - A pointer to a filter structure. + * image - The image (must be the same size as the original psf image). + * result - Pre-allocated storage for the result of the convolution. + */ +void convolve(filter *flt, double *image, double *result) +{ + /* Compute FFT of the image. */ + ftmDoubleCopy(image, flt->fft_vector, flt->image_size); + fftw_execute(flt->fft_forward); + + /* Multiple by FFT of the PSF and compute inverse FFT. */ + ftmComplexMultiply(flt->fft_vector_fft, flt->fft_vector_fft, flt->psf_fft, flt->fft_size, 0); + fftw_execute(flt->fft_backward); + + /* Copy into result, */ + ftmDoubleCopy(flt->fft_vector, result, flt->image_size); +} + + +/* + * convolveMemo() + * + * Convolve image with psf, but only if image is different enough from + * the previous image, otherwise just return the previous result. + * + * flt - A pointer to a filter structure. + * image - The image (must be the same size as the original psf image). + * result - Pre-allocated storage for the result of the convolution. + */ +void convolveMemo(filter *flt, double *image, double *result) +{ + int i,different; + + different = 0; + + /* Just check, don't copy so that we don't slowly drift.. */ + for(i=0;iimage_size;i++){ + if(fabs(image[i] - flt->old_image[i]) > flt->max_diff){ + different = 1; + break; + } + } + + if (different){ + + /* Copy into old image. */ + for(i=0;iimage_size;i++){ + flt->old_image[i] = image[i]; + } + //ftmDoubleCopy(image, flt->old_image, flt->image_size); + + /* Do the convolution. */ + convolve(flt, image, result); + } + else { + + /* Otherwise just return the previous result. */ + ftmDoubleCopy(flt->fft_vector, result, flt->image_size); + } +} + + +/* + * initialize() + * + * Set things up for FFT convolution. + * + * psf - the psf (x_size, y_size). + * max_diff - if this it not zero configure for memoization of results. + * x_size - the size of the psf in x (slow dimension). + * y_size - the size of the psf in y (fast dimension). + * estimate - 0/1 to just use an estimated FFT plan. If you are only going to + * to use the FFT a few times this can be much faster. + */ +filter *initialize(double *psf, double max_diff, int x_size, int y_size, int estimate) +{ + int i; + double normalization; + filter *flt; + + flt = (filter *)malloc(sizeof(filter)); + + /* Initialize some variables. */ + flt->fft_size = x_size * (y_size/2 + 1); + flt->image_size = x_size * y_size; + + flt->x_size = x_size; + flt->y_size = y_size; + normalization = 1.0/((double)(x_size * y_size)); + + /* Check whether we are memoizing. */ + if(max_diff > 0.0){ + flt->max_diff = max_diff; + flt->old_image = (double *)malloc(sizeof(double)*flt->image_size); + + for(i=0;iimage_size;i++){ + flt->old_image[i] = 0; + } + } + else{ + flt->max_diff = 0.0; + flt->old_image = NULL; + } + + /* Allocate storage. */ + flt->fft_vector = (double *)fftw_malloc(sizeof(double)*flt->image_size); + flt->fft_vector_fft = (fftw_complex *)fftw_malloc(sizeof(fftw_complex)*flt->fft_size); + flt->psf_fft = (fftw_complex *)fftw_malloc(sizeof(fftw_complex)*flt->fft_size); + + /* Create FFT plans. */ + if (estimate){ + flt->fft_forward = fftw_plan_dft_r2c_2d(x_size, y_size, flt->fft_vector, flt->fft_vector_fft, FFTW_ESTIMATE); + flt->fft_backward = fftw_plan_dft_c2r_2d(x_size, y_size, flt->fft_vector_fft, flt->fft_vector, FFTW_ESTIMATE); + } + else { + flt->fft_forward = fftw_plan_dft_r2c_2d(x_size, y_size, flt->fft_vector, flt->fft_vector_fft, FFTW_MEASURE); + flt->fft_backward = fftw_plan_dft_c2r_2d(x_size, y_size, flt->fft_vector_fft, flt->fft_vector, FFTW_MEASURE); + } + + /* Compute FFT of psf and save. */ + + ftmDoubleCopy(psf, flt->fft_vector, flt->image_size); + fftw_execute(flt->fft_forward); + ftmComplexCopyNormalize(flt->fft_vector_fft, flt->psf_fft, normalization, flt->fft_size); + + return flt; +} + +/* + * The MIT License + * + * Copyright (c) 2016 Zhuang Lab, Harvard University + * + * Permission is hereby granted, free of charge, to any person obtaining a copy + * of this software and associated documentation files (the "Software"), to deal + * in the Software without restriction, including without limitation the rights + * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell + * copies of the Software, and to permit persons to whom the Software is + * furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + * all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, + * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN + * THE SOFTWARE. + */ diff --git a/storm_analysis/sa_library/rebin.py b/storm_analysis/sa_library/rebin.py index b8231aa4..c0faacce 100644 --- a/storm_analysis/sa_library/rebin.py +++ b/storm_analysis/sa_library/rebin.py @@ -1,78 +1,78 @@ -#!/usr/bin/env python -""" -Rebinning of arrays. - -Hazen 2/15 -""" - -import numpy -import scipy -import scipy.fftpack - -# -# This comes directly from here: -# -# http://scipy-cookbook.readthedocs.io/items/Rebinning.html -# -# Copyright 2015, Various authors. Revision 28f8936a. -# -def downSample(a, *args): - '''rebin ndarray data into a smaller ndarray of the same rank whose dimensions - are factors of the original dimensions. eg. An array with 6 columns and 4 rows - can be reduced to have 6,3,2 or 1 columns and 4,2 or 1 rows. - example usages: - >>> a=rand(6,4); b=rebin(a,3,2) - >>> a=rand(6); b=rebin(a,2) - ''' - shape = a.shape - lenShape = len(shape) - factor = numpy.asarray(shape)/numpy.asarray(args) - evList = ['a.reshape('] + \ - ['args[%d],int(factor[%d]),'%(i,i) for i in range(lenShape)] + \ - [')'] + ['.sum(%d)'%(i+1) for i in range(lenShape)] + \ - ['/factor[%d]'%i for i in range(lenShape)] - return eval(''.join(evList)) - -def upSampleFFT(image, factor): - """ - Upsample using a FFT (high frequencies are set to zero). - """ - xsize = image.shape[0]*factor - ysize = image.shape[1]*factor - - new_image_fft = numpy.zeros((xsize, ysize),dtype=numpy.complex) - image_fft = scipy.fftpack.fft2(image) - half_x = int(image.shape[0]/2) - half_y = int(image.shape[1]/2) - new_image_fft[:half_x,:half_y] = image_fft[:half_x,:half_y] - new_image_fft[-half_x:,:half_y] = image_fft[half_x:,:half_y] - new_image_fft[:half_x,-half_y:] = image_fft[:half_x,half_y:] - new_image_fft[-half_x:,-half_y:] = image_fft[half_x:,half_y:] - - new_image = numpy.real(scipy.fftpack.ifft2(new_image_fft)) - new_image[(new_image<0.0)] = 0.0 - return new_image - -# -# The MIT License -# -# Copyright (c) 2015 Zhuang Lab, Harvard University -# -# Permission is hereby granted, free of charge, to any person obtaining a copy -# of this software and associated documentation files (the "Software"), to deal -# in the Software without restriction, including without limitation the rights -# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -# copies of the Software, and to permit persons to whom the Software is -# furnished to do so, subject to the following conditions: -# -# The above copyright notice and this permission notice shall be included in -# all copies or substantial portions of the Software. -# -# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN -# THE SOFTWARE. -# +#!/usr/bin/env python +""" +Rebinning of arrays. + +Hazen 2/15 +""" + +import numpy +import scipy +import scipy.fftpack + +# +# This comes directly from here: +# +# http://scipy-cookbook.readthedocs.io/items/Rebinning.html +# +# Copyright 2015, Various authors. Revision 28f8936a. +# +def downSample(a, *args): + '''rebin ndarray data into a smaller ndarray of the same rank whose dimensions + are factors of the original dimensions. eg. An array with 6 columns and 4 rows + can be reduced to have 6,3,2 or 1 columns and 4,2 or 1 rows. + example usages: + >>> a=rand(6,4); b=rebin(a,3,2) + >>> a=rand(6); b=rebin(a,2) + ''' + shape = a.shape + lenShape = len(shape) + factor = numpy.asarray(shape)/numpy.asarray(args) + evList = ['a.reshape('] + \ + ['args[%d],int(factor[%d]),'%(i,i) for i in range(lenShape)] + \ + [')'] + ['.sum(%d)'%(i+1) for i in range(lenShape)] + \ + ['/factor[%d]'%i for i in range(lenShape)] + return eval(''.join(evList)) + +def upSampleFFT(image, factor): + """ + Upsample using a FFT (high frequencies are set to zero). + """ + xsize = image.shape[0]*factor + ysize = image.shape[1]*factor + + new_image_fft = numpy.zeros((xsize, ysize),dtype=numpy.complex) + image_fft = scipy.fftpack.fft2(image) + half_x = int(image.shape[0]/2) + half_y = int(image.shape[1]/2) + new_image_fft[:half_x,:half_y] = image_fft[:half_x,:half_y] + new_image_fft[-half_x:,:half_y] = image_fft[half_x:,:half_y] + new_image_fft[:half_x,-half_y:] = image_fft[:half_x,half_y:] + new_image_fft[-half_x:,-half_y:] = image_fft[half_x:,half_y:] + + new_image = numpy.real(scipy.fftpack.ifft2(new_image_fft)) + new_image[(new_image<0.0)] = 0.0 + return new_image + +# +# The MIT License +# +# Copyright (c) 2015 Zhuang Lab, Harvard University +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. +# diff --git a/storm_analysis/sa_library/recenter_psf.py b/storm_analysis/sa_library/recenter_psf.py index 81453dab..a177ee04 100644 --- a/storm_analysis/sa_library/recenter_psf.py +++ b/storm_analysis/sa_library/recenter_psf.py @@ -1,42 +1,42 @@ -#!/usr/bin/env python -""" -Recenter a PSF for FFT image convolution. - -Hazen 3/16 -""" -import numpy - -def recenterPSF(psf): - """ - Conceptually anyway it is easier to draw the PSF in the center of a - array, but this does not work well when combined with FFT convolution. - """ - recentered = numpy.roll(psf, int(psf.shape[0]/2), axis = 0) - recentered = numpy.roll(recentered, int(psf.shape[1]/2), axis = 1) - - return recentered - - -# -# The MIT License -# -# Copyright (c) 2016 Zhuang Lab, Harvard University -# -# Permission is hereby granted, free of charge, to any person obtaining a copy -# of this software and associated documentation files (the "Software"), to deal -# in the Software without restriction, including without limitation the rights -# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -# copies of the Software, and to permit persons to whom the Software is -# furnished to do so, subject to the following conditions: -# -# The above copyright notice and this permission notice shall be included in -# all copies or substantial portions of the Software. -# -# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN -# THE SOFTWARE. -# +#!/usr/bin/env python +""" +Recenter a PSF for FFT image convolution. + +Hazen 3/16 +""" +import numpy + +def recenterPSF(psf): + """ + Conceptually anyway it is easier to draw the PSF in the center of a + array, but this does not work well when combined with FFT convolution. + """ + recentered = numpy.roll(psf, int(psf.shape[0]/2), axis = 0) + recentered = numpy.roll(recentered, int(psf.shape[1]/2), axis = 1) + + return recentered + + +# +# The MIT License +# +# Copyright (c) 2016 Zhuang Lab, Harvard University +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. +# diff --git a/storm_analysis/sa_utilities/xyz_drift_correction.py b/storm_analysis/sa_utilities/xyz_drift_correction.py index a8a4edd9..d5bbf1e2 100644 --- a/storm_analysis/sa_utilities/xyz_drift_correction.py +++ b/storm_analysis/sa_utilities/xyz_drift_correction.py @@ -1,266 +1,266 @@ -#!/usr/bin/env python -""" -Image based XYZ drift correction for STORM movies. - -Hazen 12/17 -""" - -import numpy -import os -import scipy.signal -import sys - -import storm_analysis.sa_library.drift_utilities as driftUtils -import storm_analysis.sa_library.imagecorrelation as imagecorrelation - - -def xyzDriftCorrection(hdf5_filename, drift_filename, step, scale, z_min, z_max, correct_z): - """ - hdf5_filename - The localizations file for drift estimation. - drift_filename - A text file to save the estimated drift in. - step - Number of frames to group together to create a single image. - scale - Image upsampling factor, 2.0 = 2x upsampling. - z_min - Minimum localization z value in microns. - z_max - Maximum localization z value in microns. - correct_z - Estimate drift in z as well as in x/y. - """ - # - # FIXME? This assumes that we also analyzed all the frames in the - # movie. If the user set the 'max_frame' parameter this - # might not actually be true. For now we're just skipping - # over all the empty frames, but it might make more sense - # not to do anything at all. - # - z_bins = int((z_max - z_min)/0.05) - h5_dc = driftUtils.SAH5DriftCorrection(filename = hdf5_filename, - scale = scale, - z_bins = z_bins) - film_l = h5_dc.getMovieLength() - - # Check if we have z data for z drift correction. - if correct_z: - assert h5_dc.hasLocalizationsField("z"), "Cannot do z drift correction without 'z' position information. Set 'z_correction' parameter to 0." - - # Sub-routines. - def saveDriftData(fdx, fdy, fdz): - driftUtils.saveDriftData(drift_filename, fdx, fdy, fdz) - h5_dc.saveDriftData(fdx, fdy, fdz) - - def interpolateData(xvals, yvals): - return driftUtils.interpolateData(xvals, yvals, film_l) - - # Don't analyze films that are empty. - if (h5_dc.getNLocalizations() == 0): - saveDriftData(numpy.zeros(film_l), - numpy.zeros(film_l), - numpy.zeros(film_l)) - return() - - # Don't analyze films that are too short. - if ((4*step) >= film_l): - saveDriftData(numpy.zeros(film_l), - numpy.zeros(film_l), - numpy.zeros(film_l)) - return() - - # - # Drift correction (XY and Z are all done at the same time) - # - # Note that drift corrected localizations are added back into - # the reference image in the hopes of improving the correction - # for subsequent localizations. - # - - # - # Figure out how to bin the movie. It seemed easier to do - # this at the beginning rather than dynamically as we - # went through the movie. - # - frame = 0 - bin_edges = [0] - while(frame < film_l): - if ((frame + 2*step) > film_l): - frame = film_l - else: - frame += step - bin_edges.append(frame) - - xy_master = None - xyz_master = None - t = [] - x = [] - y = [] - z = [] - old_dx = 0.0 - old_dy = 0.0 - old_dz = 0.0 - for i in range(len(bin_edges)-1): - - # Load correct frame range. - h5_dc.setFrameRange(bin_edges[i], bin_edges[i+1]) - - midp = (bin_edges[i+1] + bin_edges[i])/2 - - xy_curr = h5_dc.grid2D() - - # - # This is to handle analysis that did not start at frame 0 - # of the movie. Basically we keep skipping ahead until we - # find a group of frames that have some localizations. - # - # FIXME: There could still be problems if the movie does not - # start on a multiple of the step size. - # - if xy_master is None: - if (numpy.sum(xy_curr) > 0): - xy_master = xy_curr - if correct_z: - xyz_master = h5_dc.grid3D(z_min, z_max) - t.append(midp) - x.append(0.0) - y.append(0.0) - z.append(0.0) - print(bin_edges[i], bin_edges[i+1], numpy.sum(xy_curr), 0.0, 0.0, 0.0) - continue - - # Correlate to master image, skipping empty images. - if (numpy.sum(xy_curr) > 0): - [corr, dx, dy, xy_success] = imagecorrelation.xyOffset(xy_master, xy_curr, scale, - center = [x[i-1] * scale, - y[i-1] * scale]) - else: - [corr, dx, dy, xy_success] = [0.0, 0.0, 0.0, False] - - # - # Update values. If we failed, we just use the last successful - # offset measurement and hope this is close enough. - # - if xy_success: - old_dx = dx - old_dy = dy - else: - dx = old_dx - dy = old_dy - - dx = dx/float(scale) - dy = dy/float(scale) - - t.append(midp) - x.append(dx) - y.append(dy) - - # - # Apply the x/y drift correction to the current 'test' - # localizations and add them into the master, but only - # if the offset was measured successfully. - # - h5_dc.setDriftCorrectionXY(dx,dy) - if xy_success: - # Add current to master - xy_master += h5_dc.grid2D(drift_corrected = True) - - # - # Do Z correlation if requested. - # - # Note that dz and old_dz are in microns. zOffset() returns an offset - # in z bins, so convert as soon as we have it and not later, otherwise - # the paths that fall back on old_dz record a value in the wrong units. - # - dz = old_dz - if correct_z and xy_success: - - # Create 3D image with XY corrections only. We set the z offset to 0.0 to - # reset stale values from the previous cycle, if any. - h5_dc.setDriftCorrectionZ(0.0) - xyz_curr = h5_dc.grid3D(z_min, z_max, drift_corrected = True) - - # Do z correlation, skipping empty images. - if (numpy.sum(xyz_curr) > 0): - [corr, fit, dz_bins, z_success] = imagecorrelation.zOffset(xyz_master, xyz_curr) - else: - [corr, fit, dz_bins, z_success] = [0.0, 0.0, 0.0, False] - - # Update values. If we failed, use the last successful offset - # measurement, as we do for x and y. - if z_success: - dz = dz_bins * (z_max - z_min)/float(z_bins) - old_dz = dz - - h5_dc.setDriftCorrectionZ(-dz) - xyz_master += h5_dc.grid3D(z_min, z_max, drift_corrected = True) - else: - dz = old_dz - - z.append(-dz) - - print("{0:d} {1:d} {2:d} {3:0.3f} {4:0.3f} {5:0.3f}".format(bin_edges[i], - bin_edges[i+1], - numpy.sum(xy_curr), - dx, dy, dz)) - - # - # Create numpy versions of the drift arrays. We estimated the drift - # for groups of frames. We use interpolation to create an estimation - # for each individual frame. - # - nt = numpy.array(t) - final_driftx = interpolateData(nt, numpy.array(x)) - final_drifty = interpolateData(nt, numpy.array(y)) - final_driftz = interpolateData(nt, numpy.array(z)) - - saveDriftData(final_driftx, - final_drifty, - final_driftz) - - h5_dc.close(verbose = False) - - -if (__name__ == "__main__"): - - import argparse - - parser = argparse.ArgumentParser(description='Calculate drift correction using image correlation') - - parser.add_argument('--bin', dest='mlist', type=str, required=True, - help = "Localizations binary file to calculate drift correction from.") - parser.add_argument('--drift', dest='drift', type=str, required=True, - help = "Text file to save drift correction results in.") - parser.add_argument('--step', dest='step', type=int, required=True, - help = "Step size in frames.") - parser.add_argument('--scale', dest='scale', type=int, required=True, - help = "Scale for up-sampled images to use for correlation. 2 is usually a good value.") - parser.add_argument('--zmin', dest='zmin', type=float, required=False, default=-0.5, - help = "Minimum z value in microns.") - parser.add_argument('--zmax', dest='zmax', type=float, required=False, default=0.5, - help = "Maximum z value in microns.") - parser.add_argument('--zcorrect', dest='correct_z', type=bool, required=False, default=True, - help = "Also perform drift correction in Z.") - - args = parser.parse_args() - - xyzDriftCorrection(args.mlist, args.drift, args.step, args.scale, args.zmin, args.zmax, args.correct_z) - - -# -# The MIT License -# -# Copyright (c) 2014 Zhuang Lab, Harvard University -# -# Permission is hereby granted, free of charge, to any person obtaining a copy -# of this software and associated documentation files (the "Software"), to deal -# in the Software without restriction, including without limitation the rights -# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -# copies of the Software, and to permit persons to whom the Software is -# furnished to do so, subject to the following conditions: -# -# The above copyright notice and this permission notice shall be included in -# all copies or substantial portions of the Software. -# -# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN -# THE SOFTWARE. -# +#!/usr/bin/env python +""" +Image based XYZ drift correction for STORM movies. + +Hazen 12/17 +""" + +import numpy +import os +import scipy.signal +import sys + +import storm_analysis.sa_library.drift_utilities as driftUtils +import storm_analysis.sa_library.imagecorrelation as imagecorrelation + + +def xyzDriftCorrection(hdf5_filename, drift_filename, step, scale, z_min, z_max, correct_z): + """ + hdf5_filename - The localizations file for drift estimation. + drift_filename - A text file to save the estimated drift in. + step - Number of frames to group together to create a single image. + scale - Image upsampling factor, 2.0 = 2x upsampling. + z_min - Minimum localization z value in microns. + z_max - Maximum localization z value in microns. + correct_z - Estimate drift in z as well as in x/y. + """ + # + # FIXME? This assumes that we also analyzed all the frames in the + # movie. If the user set the 'max_frame' parameter this + # might not actually be true. For now we're just skipping + # over all the empty frames, but it might make more sense + # not to do anything at all. + # + z_bins = int((z_max - z_min)/0.05) + h5_dc = driftUtils.SAH5DriftCorrection(filename = hdf5_filename, + scale = scale, + z_bins = z_bins) + film_l = h5_dc.getMovieLength() + + # Check if we have z data for z drift correction. + if correct_z: + assert h5_dc.hasLocalizationsField("z"), "Cannot do z drift correction without 'z' position information. Set 'z_correction' parameter to 0." + + # Sub-routines. + def saveDriftData(fdx, fdy, fdz): + driftUtils.saveDriftData(drift_filename, fdx, fdy, fdz) + h5_dc.saveDriftData(fdx, fdy, fdz) + + def interpolateData(xvals, yvals): + return driftUtils.interpolateData(xvals, yvals, film_l) + + # Don't analyze films that are empty. + if (h5_dc.getNLocalizations() == 0): + saveDriftData(numpy.zeros(film_l), + numpy.zeros(film_l), + numpy.zeros(film_l)) + return() + + # Don't analyze films that are too short. + if ((4*step) >= film_l): + saveDriftData(numpy.zeros(film_l), + numpy.zeros(film_l), + numpy.zeros(film_l)) + return() + + # + # Drift correction (XY and Z are all done at the same time) + # + # Note that drift corrected localizations are added back into + # the reference image in the hopes of improving the correction + # for subsequent localizations. + # + + # + # Figure out how to bin the movie. It seemed easier to do + # this at the beginning rather than dynamically as we + # went through the movie. + # + frame = 0 + bin_edges = [0] + while(frame < film_l): + if ((frame + 2*step) > film_l): + frame = film_l + else: + frame += step + bin_edges.append(frame) + + xy_master = None + xyz_master = None + t = [] + x = [] + y = [] + z = [] + old_dx = 0.0 + old_dy = 0.0 + old_dz = 0.0 + for i in range(len(bin_edges)-1): + + # Load correct frame range. + h5_dc.setFrameRange(bin_edges[i], bin_edges[i+1]) + + midp = (bin_edges[i+1] + bin_edges[i])/2 + + xy_curr = h5_dc.grid2D() + + # + # This is to handle analysis that did not start at frame 0 + # of the movie. Basically we keep skipping ahead until we + # find a group of frames that have some localizations. + # + # FIXME: There could still be problems if the movie does not + # start on a multiple of the step size. + # + if xy_master is None: + if (numpy.sum(xy_curr) > 0): + xy_master = xy_curr + if correct_z: + xyz_master = h5_dc.grid3D(z_min, z_max) + t.append(midp) + x.append(0.0) + y.append(0.0) + z.append(0.0) + print(bin_edges[i], bin_edges[i+1], numpy.sum(xy_curr), 0.0, 0.0, 0.0) + continue + + # Correlate to master image, skipping empty images. + if (numpy.sum(xy_curr) > 0): + [corr, dx, dy, xy_success] = imagecorrelation.xyOffset(xy_master, xy_curr, scale, + center = [x[i-1] * scale, + y[i-1] * scale]) + else: + [corr, dx, dy, xy_success] = [0.0, 0.0, 0.0, False] + + # + # Update values. If we failed, we just use the last successful + # offset measurement and hope this is close enough. + # + if xy_success: + old_dx = dx + old_dy = dy + else: + dx = old_dx + dy = old_dy + + dx = dx/float(scale) + dy = dy/float(scale) + + t.append(midp) + x.append(dx) + y.append(dy) + + # + # Apply the x/y drift correction to the current 'test' + # localizations and add them into the master, but only + # if the offset was measured successfully. + # + h5_dc.setDriftCorrectionXY(dx,dy) + if xy_success: + # Add current to master + xy_master += h5_dc.grid2D(drift_corrected = True) + + # + # Do Z correlation if requested. + # + # Note that dz and old_dz are in microns. zOffset() returns an offset + # in z bins, so convert as soon as we have it and not later, otherwise + # the paths that fall back on old_dz record a value in the wrong units. + # + dz = old_dz + if correct_z and xy_success: + + # Create 3D image with XY corrections only. We set the z offset to 0.0 to + # reset stale values from the previous cycle, if any. + h5_dc.setDriftCorrectionZ(0.0) + xyz_curr = h5_dc.grid3D(z_min, z_max, drift_corrected = True) + + # Do z correlation, skipping empty images. + if (numpy.sum(xyz_curr) > 0): + [corr, fit, dz_bins, z_success] = imagecorrelation.zOffset(xyz_master, xyz_curr) + else: + [corr, fit, dz_bins, z_success] = [0.0, 0.0, 0.0, False] + + # Update values. If we failed, use the last successful offset + # measurement, as we do for x and y. + if z_success: + dz = dz_bins * (z_max - z_min)/float(z_bins) + old_dz = dz + + h5_dc.setDriftCorrectionZ(-dz) + xyz_master += h5_dc.grid3D(z_min, z_max, drift_corrected = True) + else: + dz = old_dz + + z.append(-dz) + + print("{0:d} {1:d} {2:d} {3:0.3f} {4:0.3f} {5:0.3f}".format(bin_edges[i], + bin_edges[i+1], + numpy.sum(xy_curr), + dx, dy, dz)) + + # + # Create numpy versions of the drift arrays. We estimated the drift + # for groups of frames. We use interpolation to create an estimation + # for each individual frame. + # + nt = numpy.array(t) + final_driftx = interpolateData(nt, numpy.array(x)) + final_drifty = interpolateData(nt, numpy.array(y)) + final_driftz = interpolateData(nt, numpy.array(z)) + + saveDriftData(final_driftx, + final_drifty, + final_driftz) + + h5_dc.close(verbose = False) + + +if (__name__ == "__main__"): + + import argparse + + parser = argparse.ArgumentParser(description='Calculate drift correction using image correlation') + + parser.add_argument('--bin', dest='mlist', type=str, required=True, + help = "Localizations binary file to calculate drift correction from.") + parser.add_argument('--drift', dest='drift', type=str, required=True, + help = "Text file to save drift correction results in.") + parser.add_argument('--step', dest='step', type=int, required=True, + help = "Step size in frames.") + parser.add_argument('--scale', dest='scale', type=int, required=True, + help = "Scale for up-sampled images to use for correlation. 2 is usually a good value.") + parser.add_argument('--zmin', dest='zmin', type=float, required=False, default=-0.5, + help = "Minimum z value in microns.") + parser.add_argument('--zmax', dest='zmax', type=float, required=False, default=0.5, + help = "Maximum z value in microns.") + parser.add_argument('--zcorrect', dest='correct_z', type=bool, required=False, default=True, + help = "Also perform drift correction in Z.") + + args = parser.parse_args() + + xyzDriftCorrection(args.mlist, args.drift, args.step, args.scale, args.zmin, args.zmax, args.correct_z) + + +# +# The MIT License +# +# Copyright (c) 2014 Zhuang Lab, Harvard University +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. +# diff --git a/storm_analysis/simulator/background.py b/storm_analysis/simulator/background.py index fed05856..a62ae0fa 100644 --- a/storm_analysis/simulator/background.py +++ b/storm_analysis/simulator/background.py @@ -1,111 +1,111 @@ -#!/usr/bin/env python -""" -Classes for creating different kinds of backgrounds. - -Hazen 11/16 -""" - -import numpy -import random - -import storm_analysis.simulator.draw_gaussians_c as dg -import storm_analysis.simulator.simbase as simbase - - -class Background(simbase.SimBase): - """ - Generate a background image (in photons). - """ - def __init__(self, sim_fp, x_size, y_size, h5_data): - super(Background, self).__init__(sim_fp, x_size, y_size, h5_data) - - def getBackground(self, frame): - return self.bg_image - - def getEmitterBackground(self, h5_data): - h5_data['background'] = numpy.zeros(h5_data["x"].size) - for i in range(h5_data['x'].size): - x = int(round(h5_data['x'][i])) - y = int(round(h5_data['y'][i])) - if((x >= 0) and (x < self.bg_image.shape[0]) and (y >= 0) and (y < self.bg_image.shape[1])): - h5_data['background'][i] = self.bg_image[x,y] - else: - h5_data['background'][i] = 0.0 - return h5_data - - -class GaussianBackground(Background): - - def __init__(self, sim_fp, x_size, y_size, h5_data, photons = 100, sigma = 100.0): - super(GaussianBackground, self).__init__(sim_fp, x_size, y_size, h5_data) - self.saveJSON({"background" : {"class" : "GaussianBackground", - "photons" : str(photons), - "sigma" : str(sigma)}}) - self.bg_image = dg.drawGaussiansXY((x_size, y_size), - numpy.array([0.5*x_size]), - numpy.array([0.5*y_size]), - sigma = sigma) - self.bg_image = photons * self.bg_image/numpy.max(self.bg_image) - - -class SineBackground(Background): - - def __init__(self, sim_fp, x_size, y_size, h5_data, photons = 100, period = 20.0): - super(SineBackground, self).__init__(sim_fp, x_size, y_size, h5_data) - self.saveJSON({"background" : {"class" : "SineBackground", - "photons" : str(photons), - "period" : str(period)}}) - self.bg_image = numpy.zeros((x_size, y_size)) - - sine_arr = photons * (0.5 + 0.5*numpy.sin(numpy.arange(x_size) * 2.0 * numpy.pi / period)) + 1.0 - for i in range(y_size): - self.bg_image[:,i] += sine_arr - - -class SlopedBackground(Background): - - def __init__(self, sim_fp, x_size, y_size, h5_data, slope = 0.1, offset = 0.0): - super(SlopedBackground, self).__init__(sim_fp, x_size, y_size, h5_data) - self.saveJSON({"background" : {"class" : "SlopedBackground", - "offset" : str(offset), - "slope" : str(slope)}}) - self.bg_image = numpy.zeros((x_size, y_size)) + offset - - slope_arr = numpy.arange(x_size) * slope - for i in range(y_size): - self.bg_image[:,i] += slope_arr - - -class UniformBackground(Background): - - def __init__(self, sim_fp, x_size, y_size, h5_data, photons = 100): - super(UniformBackground, self).__init__(sim_fp, x_size, y_size, h5_data) - self.saveJSON({"background" : {"class" : "UniformBackground", - "photons" : str(photons)}}) - self.bg_image = numpy.ones((x_size, y_size)) * photons - - - -# -# The MIT License -# -# Copyright (c) 2016 Zhuang Lab, Harvard University -# -# Permission is hereby granted, free of charge, to any person obtaining a copy -# of this software and associated documentation files (the "Software"), to deal -# in the Software without restriction, including without limitation the rights -# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -# copies of the Software, and to permit persons to whom the Software is -# furnished to do so, subject to the following conditions: -# -# The above copyright notice and this permission notice shall be included in -# all copies or substantial portions of the Software. -# -# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN -# THE SOFTWARE. -# +#!/usr/bin/env python +""" +Classes for creating different kinds of backgrounds. + +Hazen 11/16 +""" + +import numpy +import random + +import storm_analysis.simulator.draw_gaussians_c as dg +import storm_analysis.simulator.simbase as simbase + + +class Background(simbase.SimBase): + """ + Generate a background image (in photons). + """ + def __init__(self, sim_fp, x_size, y_size, h5_data): + super(Background, self).__init__(sim_fp, x_size, y_size, h5_data) + + def getBackground(self, frame): + return self.bg_image + + def getEmitterBackground(self, h5_data): + h5_data['background'] = numpy.zeros(h5_data["x"].size) + for i in range(h5_data['x'].size): + x = int(round(h5_data['x'][i])) + y = int(round(h5_data['y'][i])) + if((x >= 0) and (x < self.bg_image.shape[0]) and (y >= 0) and (y < self.bg_image.shape[1])): + h5_data['background'][i] = self.bg_image[x,y] + else: + h5_data['background'][i] = 0.0 + return h5_data + + +class GaussianBackground(Background): + + def __init__(self, sim_fp, x_size, y_size, h5_data, photons = 100, sigma = 100.0): + super(GaussianBackground, self).__init__(sim_fp, x_size, y_size, h5_data) + self.saveJSON({"background" : {"class" : "GaussianBackground", + "photons" : str(photons), + "sigma" : str(sigma)}}) + self.bg_image = dg.drawGaussiansXY((x_size, y_size), + numpy.array([0.5*x_size]), + numpy.array([0.5*y_size]), + sigma = sigma) + self.bg_image = photons * self.bg_image/numpy.max(self.bg_image) + + +class SineBackground(Background): + + def __init__(self, sim_fp, x_size, y_size, h5_data, photons = 100, period = 20.0): + super(SineBackground, self).__init__(sim_fp, x_size, y_size, h5_data) + self.saveJSON({"background" : {"class" : "SineBackground", + "photons" : str(photons), + "period" : str(period)}}) + self.bg_image = numpy.zeros((x_size, y_size)) + + sine_arr = photons * (0.5 + 0.5*numpy.sin(numpy.arange(x_size) * 2.0 * numpy.pi / period)) + 1.0 + for i in range(y_size): + self.bg_image[:,i] += sine_arr + + +class SlopedBackground(Background): + + def __init__(self, sim_fp, x_size, y_size, h5_data, slope = 0.1, offset = 0.0): + super(SlopedBackground, self).__init__(sim_fp, x_size, y_size, h5_data) + self.saveJSON({"background" : {"class" : "SlopedBackground", + "offset" : str(offset), + "slope" : str(slope)}}) + self.bg_image = numpy.zeros((x_size, y_size)) + offset + + slope_arr = numpy.arange(x_size) * slope + for i in range(y_size): + self.bg_image[:,i] += slope_arr + + +class UniformBackground(Background): + + def __init__(self, sim_fp, x_size, y_size, h5_data, photons = 100): + super(UniformBackground, self).__init__(sim_fp, x_size, y_size, h5_data) + self.saveJSON({"background" : {"class" : "UniformBackground", + "photons" : str(photons)}}) + self.bg_image = numpy.ones((x_size, y_size)) * photons + + + +# +# The MIT License +# +# Copyright (c) 2016 Zhuang Lab, Harvard University +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. +# diff --git a/storm_analysis/simulator/camera.py b/storm_analysis/simulator/camera.py index 2cfe7d30..ae8d5347 100644 --- a/storm_analysis/simulator/camera.py +++ b/storm_analysis/simulator/camera.py @@ -1,208 +1,208 @@ -#!/usr/bin/env python -""" -Classes simulating different kinds of cameras. - -These are also responsible for adding the Poisson noise -to the image (if desired). - -Hazen 11/16 -""" - -import numpy -import pickle -import random - -import storm_analysis.sa_library.analysis_io as analysisIO -import storm_analysis.simulator.simbase as simbase - -class Camera(simbase.SimBase): - """ - Converts the image from photons to counts and adds camera - noise, baseline, etc. - """ - def __init__(self, sim_fp, x_size, y_size, i3_data): - simbase.SimBase.__init__(self, sim_fp, x_size, y_size, i3_data) - - -class Ideal(Camera): - """ - Perfect camera with only shot noise. - """ - def __init__(self, sim_fp, x_size, y_size, i3_data, baseline, gain = 1.0): - super(Ideal, self).__init__(sim_fp, x_size, y_size, i3_data) - - self.baseline = baseline - self.gain = gain - self.saveJSON({"camera" : {"class" : "Ideal", - "baseline" : str(baseline), - "gain" : str(gain)}}) - - def readImage(self, image): - return self.gain*numpy.random.poisson(image) + self.baseline - - -class IdealNoNoise(Camera): - """ - Perfect camera without shot noise. - """ - def __init__(self, sim_fp, x_size, y_size, i3_data, baseline, gain = 1.0): - super(IdealNoNoise, self).__init__(sim_fp, x_size, y_size, i3_data) - - self.baseline = baseline - self.gain = gain - self.saveJSON({"camera" : {"class" : "IdealNoNoise", - "baseline" : str(baseline), - "gain" : str(gain)}}) - - def readImage(self, image): - return self.gain*image + self.baseline - - -class EMCCD(Camera): - """ - A camera with EMCCD gain. - """ - def __init__(self, sim_fp, x_size, y_size, i3_data, baseline, emccd_gain = 30.0, preamp_gain = 1.0/5.0, read_noise = 20.0): - super(EMCCD, self).__init__(sim_fp, x_size, y_size, i3_data) - - self.baseline = baseline - self.emccd_gain = emccd_gain - self.preamp_gain = preamp_gain - self.read_noise = read_noise - self.saveJSON({"camera" : {"class" : "EMCCD", - "baseline" : str(baseline), - "emccd_gain" : str(emccd_gain), - "preamp_gain" : str(preamp_gain), - "read_noise" : str(read_noise)}}) - - def readImage(self, image): - - # Detected image. - image = numpy.random.poisson(image) - - # EMCCD part. - for i in range(image.shape[0]): - for j in range(image.shape[1]): - image[i,j] = numpy.sum(numpy.random.exponential(self.emccd_gain, image[i,j])) - image = image.astype(numpy.float64) - - # Add read-noise. - image += numpy.random.normal(scale = self.read_noise, size = image.shape) - - # Pre-amp. - image = image * self.preamp_gain - - return image + self.baseline - - -class SCMOS(Camera): - """ - A sCMOS camera. The sCMOS calibration data needs to be the same size as - the simulated images. - """ - def __init__(self, sim_fp, x_size, y_size, i3_data, scmos_cal): - super(SCMOS, self).__init__(sim_fp, x_size, y_size, i3_data) - - # We transpose the calibration data here because the images we process - # are the transpose of the images that are saved for the analysis. - # - [self.offset, variance, self.gain, self.rqe] = map(numpy.transpose, - analysisIO.loadCMOSCalibration(scmos_cal)) - self.std_dev = numpy.sqrt(variance) - - if (self.offset.shape[0] != x_size) or (self.offset.shape[1] != y_size): - raise simbase.SimException("sCMOS calibration data size does not match the image size.") - - self.saveJSON({"camera" : {"class" : "SCMOS", - "scmos_cal" : scmos_cal}}) - - def readImage(self, image): - - # Multiply by relative QE. - image = image * self.rqe - - # Detected image. - image = numpy.random.poisson(image) - - # Multiply by pixel dependent gain. - image = self.gain * image - - # Add pixel dependent noise. - image += numpy.random.normal(scale = self.std_dev) - - # Add pixel dependent offset. - image += self.offset - - return image - - -def createSCMOSCalibration(x_size, y_size, gain, read_noise, hot_fraction = 0.05, hot_lambda = None, offset = 100.0): - """ - Create a simulated calibration file for a sCMOS camera. - - The camera is modeled as having a gaussian distribution of read noise values - with mean = 'read_noise' and sigma = 0.1 * read_noise. In addition a small - fraction of the pixels are 'hot' and have an additional exponential distributed - value added to their read noise. - - x_size - Size in x in pixels. - y_size - Size in y in pixels. - gain - Camera gain in ADU/electrons. - read_noise - Read noise in electrons. - hot_fraction - The fraction of the pixels that are 'hot'. - hot_lambda - The lambda for the exponential distribution of the hot pixels. - offset - Camera baseline in ADU. - - returns [offset, variance, gain, rqe] - """ - offset = numpy.zeros((x_size, y_size)) + offset - gain = gain * numpy.ones((x_size, y_size)) - rqe = numpy.ones((x_size, y_size)) - - # We're squaring everything as calibration file contains the variance, not the - # standard deviation. - # - r_noise = numpy.random.normal(loc = read_noise, scale = 0.1 * read_noise, size = (x_size, y_size)) - r_noise = r_noise * r_noise - - # Use 10x the read noise if hot_lambda is not specified. - if hot_lambda is None: - hot_lambda = 10.0 * read_noise - - e_noise = numpy.random.exponential(scale = hot_lambda, size = (x_size, y_size)) - e_noise = e_noise * e_noise - - mask = (numpy.random.uniform(size = (x_size, y_size)) > hot_fraction) - e_noise[mask] = 0.0 - - variance = r_noise + e_noise - - # Convert variance to ADU^2 - variance = variance * gain * gain - - return [offset, variance, gain, rqe] - - -# -# The MIT License -# -# Copyright (c) 2016 Zhuang Lab, Harvard University -# -# Permission is hereby granted, free of charge, to any person obtaining a copy -# of this software and associated documentation files (the "Software"), to deal -# in the Software without restriction, including without limitation the rights -# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -# copies of the Software, and to permit persons to whom the Software is -# furnished to do so, subject to the following conditions: -# -# The above copyright notice and this permission notice shall be included in -# all copies or substantial portions of the Software. -# -# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN -# THE SOFTWARE. -# +#!/usr/bin/env python +""" +Classes simulating different kinds of cameras. + +These are also responsible for adding the Poisson noise +to the image (if desired). + +Hazen 11/16 +""" + +import numpy +import pickle +import random + +import storm_analysis.sa_library.analysis_io as analysisIO +import storm_analysis.simulator.simbase as simbase + +class Camera(simbase.SimBase): + """ + Converts the image from photons to counts and adds camera + noise, baseline, etc. + """ + def __init__(self, sim_fp, x_size, y_size, i3_data): + simbase.SimBase.__init__(self, sim_fp, x_size, y_size, i3_data) + + +class Ideal(Camera): + """ + Perfect camera with only shot noise. + """ + def __init__(self, sim_fp, x_size, y_size, i3_data, baseline, gain = 1.0): + super(Ideal, self).__init__(sim_fp, x_size, y_size, i3_data) + + self.baseline = baseline + self.gain = gain + self.saveJSON({"camera" : {"class" : "Ideal", + "baseline" : str(baseline), + "gain" : str(gain)}}) + + def readImage(self, image): + return self.gain*numpy.random.poisson(image) + self.baseline + + +class IdealNoNoise(Camera): + """ + Perfect camera without shot noise. + """ + def __init__(self, sim_fp, x_size, y_size, i3_data, baseline, gain = 1.0): + super(IdealNoNoise, self).__init__(sim_fp, x_size, y_size, i3_data) + + self.baseline = baseline + self.gain = gain + self.saveJSON({"camera" : {"class" : "IdealNoNoise", + "baseline" : str(baseline), + "gain" : str(gain)}}) + + def readImage(self, image): + return self.gain*image + self.baseline + + +class EMCCD(Camera): + """ + A camera with EMCCD gain. + """ + def __init__(self, sim_fp, x_size, y_size, i3_data, baseline, emccd_gain = 30.0, preamp_gain = 1.0/5.0, read_noise = 20.0): + super(EMCCD, self).__init__(sim_fp, x_size, y_size, i3_data) + + self.baseline = baseline + self.emccd_gain = emccd_gain + self.preamp_gain = preamp_gain + self.read_noise = read_noise + self.saveJSON({"camera" : {"class" : "EMCCD", + "baseline" : str(baseline), + "emccd_gain" : str(emccd_gain), + "preamp_gain" : str(preamp_gain), + "read_noise" : str(read_noise)}}) + + def readImage(self, image): + + # Detected image. + image = numpy.random.poisson(image) + + # EMCCD part. + for i in range(image.shape[0]): + for j in range(image.shape[1]): + image[i,j] = numpy.sum(numpy.random.exponential(self.emccd_gain, image[i,j])) + image = image.astype(numpy.float64) + + # Add read-noise. + image += numpy.random.normal(scale = self.read_noise, size = image.shape) + + # Pre-amp. + image = image * self.preamp_gain + + return image + self.baseline + + +class SCMOS(Camera): + """ + A sCMOS camera. The sCMOS calibration data needs to be the same size as + the simulated images. + """ + def __init__(self, sim_fp, x_size, y_size, i3_data, scmos_cal): + super(SCMOS, self).__init__(sim_fp, x_size, y_size, i3_data) + + # We transpose the calibration data here because the images we process + # are the transpose of the images that are saved for the analysis. + # + [self.offset, variance, self.gain, self.rqe] = map(numpy.transpose, + analysisIO.loadCMOSCalibration(scmos_cal)) + self.std_dev = numpy.sqrt(variance) + + if (self.offset.shape[0] != x_size) or (self.offset.shape[1] != y_size): + raise simbase.SimException("sCMOS calibration data size does not match the image size.") + + self.saveJSON({"camera" : {"class" : "SCMOS", + "scmos_cal" : scmos_cal}}) + + def readImage(self, image): + + # Multiply by relative QE. + image = image * self.rqe + + # Detected image. + image = numpy.random.poisson(image) + + # Multiply by pixel dependent gain. + image = self.gain * image + + # Add pixel dependent noise. + image += numpy.random.normal(scale = self.std_dev) + + # Add pixel dependent offset. + image += self.offset + + return image + + +def createSCMOSCalibration(x_size, y_size, gain, read_noise, hot_fraction = 0.05, hot_lambda = None, offset = 100.0): + """ + Create a simulated calibration file for a sCMOS camera. + + The camera is modeled as having a gaussian distribution of read noise values + with mean = 'read_noise' and sigma = 0.1 * read_noise. In addition a small + fraction of the pixels are 'hot' and have an additional exponential distributed + value added to their read noise. + + x_size - Size in x in pixels. + y_size - Size in y in pixels. + gain - Camera gain in ADU/electrons. + read_noise - Read noise in electrons. + hot_fraction - The fraction of the pixels that are 'hot'. + hot_lambda - The lambda for the exponential distribution of the hot pixels. + offset - Camera baseline in ADU. + + returns [offset, variance, gain, rqe] + """ + offset = numpy.zeros((x_size, y_size)) + offset + gain = gain * numpy.ones((x_size, y_size)) + rqe = numpy.ones((x_size, y_size)) + + # We're squaring everything as calibration file contains the variance, not the + # standard deviation. + # + r_noise = numpy.random.normal(loc = read_noise, scale = 0.1 * read_noise, size = (x_size, y_size)) + r_noise = r_noise * r_noise + + # Use 10x the read noise if hot_lambda is not specified. + if hot_lambda is None: + hot_lambda = 10.0 * read_noise + + e_noise = numpy.random.exponential(scale = hot_lambda, size = (x_size, y_size)) + e_noise = e_noise * e_noise + + mask = (numpy.random.uniform(size = (x_size, y_size)) > hot_fraction) + e_noise[mask] = 0.0 + + variance = r_noise + e_noise + + # Convert variance to ADU^2 + variance = variance * gain * gain + + return [offset, variance, gain, rqe] + + +# +# The MIT License +# +# Copyright (c) 2016 Zhuang Lab, Harvard University +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. +# diff --git a/storm_analysis/simulator/draw_gaussians.c b/storm_analysis/simulator/draw_gaussians.c index 75663113..18a3d36b 100644 --- a/storm_analysis/simulator/draw_gaussians.c +++ b/storm_analysis/simulator/draw_gaussians.c @@ -1,83 +1,83 @@ -/* - * Draw gaussians on an array. - * - * Hazen 03/13 - */ - -/* Include */ -#include -#include -#include - -/* Define */ -#define NPARAMS 5 -#define XC 0 -#define YC 1 -#define HEIGHT 2 -#define XW 3 -#define YW 4 - -/* Function definitions */ -void drawGaussians(double *, double *, int, int, int, int); - -/* Functions */ -void drawGaussians(double *image, double *gaussian_params, int image_x, int image_y, int number_gaussians, int resolution) -{ - int i, j, k; - int awidth, fx, fy, sx, sy; - double x, y, xw, yw; - double dx, dy, px, py, sgx, sgy; - double intens, sum; - double norm, step_size, start; - - norm = 1.0/((double)(resolution*resolution)); - step_size = 1.0/((double)resolution); - start = -0.5 + 0.5/((double)resolution); - - for(i=0; i < number_gaussians; i++){ - intens = gaussian_params[i*NPARAMS+HEIGHT]; - if(intens > 0.0){ - px = gaussian_params[i*NPARAMS+XC]; - py = gaussian_params[i*NPARAMS+YC]; - xw = gaussian_params[i*NPARAMS+XW]; - yw = gaussian_params[i*NPARAMS+YW]; - - if(xw > yw){ - awidth = (int)(5.0 * xw); - } else { - awidth = (int)(5.0 * yw); - } - - if(awidth < 1){ - awidth = 1; - } - - sgx = 1.0/(2.0 * xw * xw); - sgy = 1.0/(2.0 * yw * yw); - - sx = (int)px - awidth; - sy = (int)py - awidth; - fx = (int)px + awidth + 1; - fy = (int)py + awidth + 1; - if(sx < 0) { sx = 0; } - if(fx > image_y) { fx = image_y; } - if(sy < 0) { sy = 0; } - if(fy > image_x) { fy = image_x; } - - for(j=sx;j +#include +#include + +/* Define */ +#define NPARAMS 5 +#define XC 0 +#define YC 1 +#define HEIGHT 2 +#define XW 3 +#define YW 4 + +/* Function definitions */ +void drawGaussians(double *, double *, int, int, int, int); + +/* Functions */ +void drawGaussians(double *image, double *gaussian_params, int image_x, int image_y, int number_gaussians, int resolution) +{ + int i, j, k; + int awidth, fx, fy, sx, sy; + double x, y, xw, yw; + double dx, dy, px, py, sgx, sgy; + double intens, sum; + double norm, step_size, start; + + norm = 1.0/((double)(resolution*resolution)); + step_size = 1.0/((double)resolution); + start = -0.5 + 0.5/((double)resolution); + + for(i=0; i < number_gaussians; i++){ + intens = gaussian_params[i*NPARAMS+HEIGHT]; + if(intens > 0.0){ + px = gaussian_params[i*NPARAMS+XC]; + py = gaussian_params[i*NPARAMS+YC]; + xw = gaussian_params[i*NPARAMS+XW]; + yw = gaussian_params[i*NPARAMS+YW]; + + if(xw > yw){ + awidth = (int)(5.0 * xw); + } else { + awidth = (int)(5.0 * yw); + } + + if(awidth < 1){ + awidth = 1; + } + + sgx = 1.0/(2.0 * xw * xw); + sgy = 1.0/(2.0 * yw * yw); + + sx = (int)px - awidth; + sy = (int)py - awidth; + fx = (int)px + awidth + 1; + fy = (int)py + awidth + 1; + if(sx < 0) { sx = 0; } + if(fx > image_y) { fx = image_y; } + if(sy < 0) { sy = 0; } + if(fy > image_x) { fy = image_x; } + + for(j=sx;j= (frame + 1.0)): - if self.am_on[i]: - integrated_on[i] = 1.0 - - else: - last_transistion = frame - while (self.next_transistion[i] < (frame + 1.0)): - if self.am_on[i]: - integrated_on[i] += self.next_transistion[i] - last_transistion - - self.am_on[i] = not self.am_on[i] - last_transistion = self.next_transistion[i] - if self.am_on[i]: - self.next_transistion[i] += numpy.random.exponential(self.on_time[i]) - else: - self.next_transistion[i] += numpy.random.exponential(self.off_time[i]) - - # Turned on and not off again in the current frame. - if self.am_on[i]: - integrated_on[i] += (frame + 1.0) - last_transistion - - # Set sum and return only those emitters that have sum > 0. - self.h5_data['sum'] = integrated_on * self.photons - mask = (self.h5_data['sum'] > 0.0) - - temp = {} - for key in self.h5_data: - temp[key] = self.h5_data[key][mask].copy() - - return temp - - -class ComplexSTORM(STORM): - """ - Each emitter has different on & off kinetics. - """ - def __init__(self, sim_fp, x_size, y_size, h5_data): - super(ComplexSTORM, self).__init__(sim_fp, x_size, y_size, h5_data) - - # The on/off time and intensity come from the H5 file. We'll crash - # here if these fields are not available. - self.photons = self.h5_data['photons'] - self.off_time = self.h5_data['off_time'] - self.on_time = self.h5_data['on_time'] - - self.n_emitters = self.h5_data['x'].size - self.saveJSON({"photophysics" : {"class" : "ComplexSTORM", - "photons" : str(numpy.mean(self.photons)), - "on_time" : str(numpy.mean(self.on_time)), - "off_time" : str(numpy.mean(self.off_time))}}) - - # Initially all the emitters are off. - self.am_on = numpy.zeros(self.n_emitters, dtype = numpy.bool_) - self.next_transistion = numpy.random.exponential(self.off_time, self.n_emitters) - - -class SimpleSTORM(STORM): - """ - Each emitter on for 1 frame out of every 1000 frames - on average, both are exponentially distributed. - - Args: - on_time : Average on time in frames. - off_time : Average off time in frames. - - """ - def __init__(self, sim_fp, x_size, y_size, h5_data, photons = 2000, on_time = 1.0, off_time = 1000.0): - super(SimpleSTORM, self).__init__(sim_fp, x_size, y_size, h5_data) - - self.photons = numpy.ones(self.n_emitters)*photons - self.off_time = numpy.ones(self.n_emitters)*off_time - self.on_time = numpy.ones(self.n_emitters)*on_time - - self.n_emitters = self.h5_data['x'].size - self.saveJSON({"photophysics" : {"class" : "SimpleSTORM", - "photons" : str(photons), - "on_time" : str(on_time), - "off_time" : str(off_time)}}) - - # Initially all the emitters are off. - self.am_on = numpy.zeros(self.n_emitters, dtype = numpy.bool_) - self.next_transistion = numpy.random.exponential(self.off_time, self.n_emitters) - - -# -# The MIT License -# -# Copyright (c) 2016 Zhuang Lab, Harvard University -# -# Permission is hereby granted, free of charge, to any person obtaining a copy -# of this software and associated documentation files (the "Software"), to deal -# in the Software without restriction, including without limitation the rights -# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -# copies of the Software, and to permit persons to whom the Software is -# furnished to do so, subject to the following conditions: -# -# The above copyright notice and this permission notice shall be included in -# all copies or substantial portions of the Software. -# -# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN -# THE SOFTWARE. -# +#!/usr/bin/env python +""" +Classes for simulating different dye photophysics. + +Hazen 11/16 +""" + +import numpy +import random + +import storm_analysis.sa_library.sa_h5py as saH5Py + +import storm_analysis.simulator.simbase as simbase + + +class PhotoPhysics(simbase.SimBase): + """ + Returns location and intensity (peak height in photons) of + the emitters that are on in the current frame. + """ + def __init__(self, sim_fp, x_size, y_size, h5_data): + super(PhotoPhysics, self).__init__(sim_fp, x_size, y_size, h5_data) + + +class AlwaysOn(PhotoPhysics): + """ + All the emitters are on all the time. + """ + def __init__(self, sim_fp, x_size, y_size, h5_data, photons = 2000): + super(AlwaysOn, self).__init__(sim_fp, x_size, y_size, h5_data) + self.saveJSON({"photophysics" : {"class" : "AlwaysOn", + "photons" : str(photons)}}) + self.h5_data['sum'] = photons * numpy.ones(self.h5_data['x'].size) + + def getEmitters(self, frame): + temp = {} + for key in self.h5_data: + temp[key] = self.h5_data[key].copy() + + return temp + + +class AlwaysOnMC(PhotoPhysics): + """ + All the emitters are on all the time, adapted for multicolor simulations. + """ + def __init__(self, sim_fp, x_size, y_size, h5_data, color = 0, photons = 2000): + super(AlwaysOnMC, self).__init__(sim_fp, x_size, y_size, h5_data) + self.saveJSON({"photophysics" : {"class" : "AlwaysOnMC", + "color" : str(color), + "photons" : str(photons)}}) + mask = (self.h5_data['color'] != color) + self.h5_data['sum'] = photons * numpy.ones(self.h5_data['x'].size) + self.h5_data['sum'][mask] = 0.1 * self.h5_data['sum'][mask] + + def getEmitters(self, frame): + temp = {} + for key in self.h5_data: + temp[key] = self.h5_data[key].copy() + + return temp + + +class Duplicate(PhotoPhysics): + """ + This just returns the localizations in an existing HDF5 file. It is perhaps + most useful for simulating multiplane data. + """ + def __init__(self, sim_fp, x_size, y_size, h5_data, h5_name, cx = None, cy = None, z_offset = None): + super(Duplicate, self).__init__(sim_fp, x_size, y_size, h5_data) + self.saveJSON({"photophysics" : {"class" : "Duplicate", + "h5_name" : h5_name, + "cx" : str(cx), + "cy" : str(cy), + "z_offset" : str(z_offset)}}) + + self.h5_data = saH5Py.SAH5Py(h5_name) + self.cx = cx + self.cy = cy + self.z_offset = z_offset + + def getEmitters(self, frame): + locs = self.h5_data.getLocalizationsInFrame(frame) + + temp = {} + + # Check for no localization data for this frame. If there are none add the minimum + # needed fields. + if not locs: + for elt in ['x', 'y', 'z', 'sum']: + temp[elt] = numpy.array([]) + + else: + for elt in locs: + temp[elt] = locs[elt] + + # Apply transforms if requested. + if self.cx is not None: + xf = self.cx[0] + self.cx[1] * temp['x'] + self.cx[2] * temp['y'] + yf = self.cy[0] + self.cy[1] * temp['x'] + self.cy[2] * temp['y'] + temp['x'] = xf + temp['y'] = yf + + if self.z_offset is not None: + temp['z'] += self.z_offset + + return temp + + +class STORM(PhotoPhysics): + """ + Emitters go on and off with an exponentially distributed waiting times. + Emitters can have a different brightness when on. + """ + def __init__(self, sim_fp, x_size, y_size, h5_data): + super(STORM, self).__init__(sim_fp, x_size, y_size, h5_data) + self.n_emitters = self.h5_data['x'].size + + # Sub-classes need to set these as appropriate. + self.photons = None + self.off_time = None + self.on_time = None + + def getEmitters(self, frame): + integrated_on = numpy.zeros(self.n_emitters) + + # + # This is a little complicated because we are trying to include accurate + # modeling of emitters that turned on/off more than once in a single frame. + # + for i in range(self.n_emitters): + + # The easy case, no change in state in the current frame. + if (self.next_transistion[i] >= (frame + 1.0)): + if self.am_on[i]: + integrated_on[i] = 1.0 + + else: + last_transistion = frame + while (self.next_transistion[i] < (frame + 1.0)): + if self.am_on[i]: + integrated_on[i] += self.next_transistion[i] - last_transistion + + self.am_on[i] = not self.am_on[i] + last_transistion = self.next_transistion[i] + if self.am_on[i]: + self.next_transistion[i] += numpy.random.exponential(self.on_time[i]) + else: + self.next_transistion[i] += numpy.random.exponential(self.off_time[i]) + + # Turned on and not off again in the current frame. + if self.am_on[i]: + integrated_on[i] += (frame + 1.0) - last_transistion + + # Set sum and return only those emitters that have sum > 0. + self.h5_data['sum'] = integrated_on * self.photons + mask = (self.h5_data['sum'] > 0.0) + + temp = {} + for key in self.h5_data: + temp[key] = self.h5_data[key][mask].copy() + + return temp + + +class ComplexSTORM(STORM): + """ + Each emitter has different on & off kinetics. + """ + def __init__(self, sim_fp, x_size, y_size, h5_data): + super(ComplexSTORM, self).__init__(sim_fp, x_size, y_size, h5_data) + + # The on/off time and intensity come from the H5 file. We'll crash + # here if these fields are not available. + self.photons = self.h5_data['photons'] + self.off_time = self.h5_data['off_time'] + self.on_time = self.h5_data['on_time'] + + self.n_emitters = self.h5_data['x'].size + self.saveJSON({"photophysics" : {"class" : "ComplexSTORM", + "photons" : str(numpy.mean(self.photons)), + "on_time" : str(numpy.mean(self.on_time)), + "off_time" : str(numpy.mean(self.off_time))}}) + + # Initially all the emitters are off. + self.am_on = numpy.zeros(self.n_emitters, dtype = numpy.bool_) + self.next_transistion = numpy.random.exponential(self.off_time, self.n_emitters) + + +class SimpleSTORM(STORM): + """ + Each emitter on for 1 frame out of every 1000 frames + on average, both are exponentially distributed. + + Args: + on_time : Average on time in frames. + off_time : Average off time in frames. + + """ + def __init__(self, sim_fp, x_size, y_size, h5_data, photons = 2000, on_time = 1.0, off_time = 1000.0): + super(SimpleSTORM, self).__init__(sim_fp, x_size, y_size, h5_data) + + self.photons = numpy.ones(self.n_emitters)*photons + self.off_time = numpy.ones(self.n_emitters)*off_time + self.on_time = numpy.ones(self.n_emitters)*on_time + + self.n_emitters = self.h5_data['x'].size + self.saveJSON({"photophysics" : {"class" : "SimpleSTORM", + "photons" : str(photons), + "on_time" : str(on_time), + "off_time" : str(off_time)}}) + + # Initially all the emitters are off. + self.am_on = numpy.zeros(self.n_emitters, dtype = numpy.bool_) + self.next_transistion = numpy.random.exponential(self.off_time, self.n_emitters) + + +# +# The MIT License +# +# Copyright (c) 2016 Zhuang Lab, Harvard University +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. +# diff --git a/storm_analysis/simulator/psf.py b/storm_analysis/simulator/psf.py index 38c70de3..7eb53571 100644 --- a/storm_analysis/simulator/psf.py +++ b/storm_analysis/simulator/psf.py @@ -1,622 +1,622 @@ -#!/usr/bin/env python -""" -Classes for creating different kinds of PSFs. - -Note that as a side-effect the getPSFs() method is expected -to also set the correct value of the 'h' field based on the -PSF and the 'a' field. - -Hazen 11/16 -""" - -import numpy -import pickle -import random -import scipy -import scipy.fftpack - -import storm_analysis.spliner.spline_to_psf as splineToPSF - -import storm_analysis.simulator.draw_gaussians_c as dg -import storm_analysis.simulator.pupil_math as pupilMath -import storm_analysis.simulator.simbase as simbase - -# These are here to make it easier to create pupil functions that match -# those used in the simulations. -pf_size = 128 -pf_wavelength = 600 -pf_refractive_index = 1.5 -pf_numerical_aperture = 1.4 - -class PSF(simbase.SimBase): - """ - Draws the emitter PSFs on an image. - """ - def __init__(self, sim_fp, x_size, y_size, h5_data, nm_per_pixel): - super(PSF, self).__init__(sim_fp, x_size, y_size, h5_data) - self.nm_per_pixel = nm_per_pixel - - -class DHPSF(PSF): - """ - A very simplistic approximation of the double helix PSF. - """ - def __init__(self, sim_fp, x_size, y_size, h5_data, nm_per_pixel, z_range = 0.75): - """ - z_range - (half) Z range in microns. - """ - super(DHPSF, self).__init__(sim_fp, x_size, y_size, h5_data, nm_per_pixel) - self.z_max = z_range - self.z_min = -z_range - - self.saveJSON({"psf" : {"class" : "DHPSF", - "z_range" : str(z_range)}}) - - def getPSFs(self, h5_data): - x = h5_data['x'] - y = h5_data['y'] - a = h5_data['sum'] - - sx = h5_data['ysigma'] - sy = h5_data['xsigma'] - - h = 0.5*a/(2.0 * numpy.pi * sx * sy) - h5_data['height'] = h - - angle = numpy.pi * 0.9 * ((h5_data['z'] - self.z_min)/(self.z_max - self.z_min)) - dx = 2.0 * sx * numpy.cos(angle) - dy = 2.0 * sy * numpy.sin(angle) - - x1 = x + dx - y1 = y + dy - x2 = x - dx - y2 = y - dy - - x = numpy.concatenate((x1, x2)) - y = numpy.concatenate((y1, y2)) - h = numpy.concatenate((h, h)) - sx = numpy.concatenate((sx, sx)) - sy = numpy.concatenate((sy, sy)) - - return dg.drawGaussians((self.x_size, self.y_size), - numpy.stack((x, y, h, sx, sy), axis = 1), - res = 5) - - -class GaussianPSF(PSF): - """ - Gaussian PSF. - """ - def __init__(self, sim_fp, x_size, y_size, h5_data, nm_per_pixel): - super(GaussianPSF, self).__init__(sim_fp, x_size, y_size, h5_data, nm_per_pixel) - self.saveJSON({"psf" : {"class" : "GaussianPSF", - "nm_per_pixel" : str(nm_per_pixel)}}) - - def getPSFs(self, h5_data): - x = h5_data['x'] - y = h5_data['y'] - a = h5_data['sum'] - - sx = h5_data['ysigma'] - sy = h5_data['xsigma'] - - h = a/(2.0 * numpy.pi * sx * sy) - h5_data['height'] = h - - return dg.drawGaussians((self.x_size, self.y_size), - numpy.concatenate((x[:,None], - y[:,None], - h[:,None], - sx[:,None], - sy[:,None]), - axis = 1), - res = 5) - - -class PupilFunctionBase(PSF): - """ - PSF using the pupil function approach. - """ - def __init__(self, sim_fp, x_size, y_size, h5_data, nm_per_pixel, pf_size): - super(PupilFunctionBase, self).__init__(sim_fp, x_size, y_size, h5_data, nm_per_pixel) - - self.otf_scaler = None - self.psf_size = pf_size - self.sample_index = None - self.z_center = None - - self.margin = int(self.psf_size/2) + 1 - - self.im_size_x = self.x_size + 2 * self.margin - self.im_size_y = self.y_size + 2 * self.margin - - def getAberrationFn(self, z_value): - return None - - def getPSFs(self, h5_data): - """ - The expected form for the h5 data fields are x,y in pixels and z in microns. - """ - image = numpy.zeros((self.im_size_x, self.im_size_y)) - x = h5_data['x']+1 # Pixels - y = h5_data['y']+1 # Pixels - z = h5_data['z'] # Expected to be in microns. - a = h5_data['sum'] - - h5_data['height'] = numpy.zeros(a.size) - - dx = x - numpy.floor(x) - dy = y - numpy.floor(y) - - for i in range(x.size): - - ix = int(x[i]) - iy = int(y[i]) - - if (ix >= 0.0) and (ix < self.x_size) and (iy >= 0.0) and (iy < self.y_size): - - pf = self.pf - - # Apply aberration function if available. - ab_fn = self.getAberrationFn(z[i]) - if ab_fn is not None: - pf = pf*ab_fn - - # Shift to the desired z value. - defocused = self.geo.changeFocus(pf, z[i]) - - # Translate to the correct sub-pixel position. - #translated = defocused - translated = self.geo.translatePf(defocused, dx[i], dy[i]) - - # Get real-space intensity. - psf = pupilMath.intensity(pupilMath.toRealSpace(translated)) * a[i] - - # Apply OTF scaling if requested. - if self.otf_scaler is not None: - otf = scipy.fftpack.fftshift(scipy.fftpack.fft2(psf)) - otf_scaled = otf * self.otf_scaler - psf = numpy.abs(scipy.fftpack.ifft2(otf_scaled)) - - h5_data['height'][i] = numpy.max(psf) - - image[ix:ix+self.psf_size,iy:iy+self.psf_size] += psf - - return image[self.margin:self.margin+self.x_size,self.margin:self.margin+self.y_size] - - -class PupilFunction(PupilFunctionBase): - """ - PSF using the pupil function approach. This is the most basic version. - - Note that by default this uses pupilMath.GeometrySim class which (more - or less) emulates OTF scaling by using a pixel size that is half the - specified pixel size to create the PSF. - """ - def __init__(self, sim_fp, x_size, y_size, h5_data, nm_per_pixel, zmn, - wavelength = pf_wavelength, - refractive_index = pf_refractive_index, - numerical_aperture = pf_numerical_aperture, - pf_size = pf_size, - geo_sim_pf = True): - """ - zmn is a list of lists containing the zernike mode terms, e.g. - [[1.3, 2, 2]] for pure astigmatism. - wavelength is the mean emission wavelength in nm. - """ - super(PupilFunction, self).__init__(sim_fp, x_size, y_size, h5_data, nm_per_pixel, pf_size) - self.saveJSON({"psf" : {"class" : "PupilFunction", - "geo_sim_pf" : str(geo_sim_pf), - "nm_per_pixel" : str(nm_per_pixel), - "numerical_aperture" : str(numerical_aperture), - "pf_size" : str(pf_size), - "refactrive_index" : str(refractive_index), - "wavelength" : str(wavelength), - "zmn" : str(zmn)}}) - - if geo_sim_pf: - self.geo = pupilMath.GeometrySim(pf_size, - nm_per_pixel * 0.001, - wavelength * 0.001, - refractive_index, - numerical_aperture) - else: - self.geo = pupilMath.Geometry(pf_size, - nm_per_pixel * 0.001, - wavelength * 0.001, - refractive_index, - numerical_aperture) - - self.pf = self.geo.createFromZernike(1.0, zmn) - - -class PupilFunctionScalar(PupilFunctionBase): - """ - This version supports OTF scaling and sample aberrations. - """ - def __init__(self, sim_fp, x_size, y_size, h5_data, nm_per_pixel, zmn, - wavelength = pf_wavelength, - refractive_index = pf_refractive_index, - numerical_aperture = pf_numerical_aperture, - pf_size = pf_size, - otf_sigma = None, - sample_index = None, - z_center = None): - """ - zmn is a list of lists containing the zernike mode terms, e.g. - [[1.3, 2, 2]] for pure astigmatism. - wavelength is the mean emission wavelength in nm. - - otf_sigma is the sigma to use for Gaussian OTF scaling. - sample_index is the index of the sample media. - z_center is the focal plane position in microns. - - If you specify sample_index you also need to specify z_center. - """ - super(PupilFunctionScalar, self).__init__(sim_fp, x_size, y_size, h5_data, nm_per_pixel, pf_size) - self.saveJSON({"psf" : {"class" : "PupilFunctionScalar", - "nm_per_pixel" : str(nm_per_pixel), - "numerical_aperture" : str(numerical_aperture), - "otf_sigma" : str(otf_sigma), - "pf_size" : str(pf_size), - "refactrive_index" : str(refractive_index), - "sample_index" : str(sample_index), - "wavelength" : str(wavelength), - "z_center" : str(z_center), - "zmn" : str(zmn)}}) - - self.geo = pupilMath.Geometry(pf_size, - nm_per_pixel * 0.001, - wavelength * 0.001, - refractive_index, - numerical_aperture) - - self.pf = self.geo.createFromZernike(1.0, zmn) - - if otf_sigma is not None: - self.otf_scaler = self.geo.gaussianScalingFactor(otf_sigma) - - self.sample_index = sample_index - self.z_center = z_center - - def getAberrationFn(self, z_value): - - # Apply sample index aberration if requested. - if self.sample_index is not None: - - # Distance 0.0 is the coverslip. Z values are relative to the focal - # plane which is at z_center. - # - depth = self.z_center + z_value - - if (depth < 0.0): - depth = 0.0 - - return self.geo.aberration(depth, self.sample_index) - - else: - return None - - -class PupilFunctionScalarCalibration(PupilFunctionBase): - """ - This version supports OTF scaling and a fixed sample aberration. It - simulates the PSF you'd expect from a calibration movie of a bead - on a coverslip. - """ - def __init__(self, sim_fp, x_size, y_size, h5_data, nm_per_pixel, zmn, - wavelength = pf_wavelength, - refractive_index = pf_refractive_index, - numerical_aperture = pf_numerical_aperture, - pf_size = pf_size, - bead_z_center = None, - otf_sigma = None, - sample_index = None): - """ - zmn is a list of lists containing the zernike mode terms, e.g. - [[1.3, 2, 2]] for pure astigmatism. - wavelength is the mean emission wavelength in nm. - - bead_z_center is the height of the bead above the coverslip in microns. - otf_sigma is the sigma to use for Gaussian OTF scaling. - sample_index is the index of the sample media. - - If you specify sample_index you also need to specify z_center. - """ - super(PupilFunctionScalarCalibration, self).__init__(sim_fp, x_size, y_size, h5_data, nm_per_pixel, pf_size) - self.saveJSON({"psf" : {"class" : "PupilFunctionScalarCalibration", - "bead_z_center" : str(bead_z_center), - "nm_per_pixel" : str(nm_per_pixel), - "numerical_aperture" : str(numerical_aperture), - "otf_sigma" : str(otf_sigma), - "pf_size" : str(pf_size), - "refactrive_index" : str(refractive_index), - "sample_index" : str(sample_index), - "wavelength" : str(wavelength), - "zmn" : str(zmn)}}) - - self.geo = pupilMath.Geometry(pf_size, - nm_per_pixel * 0.001, - wavelength * 0.001, - refractive_index, - numerical_aperture) - - self.pf = self.geo.createFromZernike(1.0, zmn) - - if otf_sigma is not None: - self.otf_scaler = self.geo.gaussianScalingFactor(otf_sigma) - - self.ab_fn = None - if sample_index is not None: - self.ab_fn = self.geo.aberration(bead_z_center, sample_index) - - def getAberrationFn(self, z_value): - return self.ab_fn - - -class PupilFunctionVectorial(PupilFunctionBase): - """ - This is a vectorial PSF version and supports bead scaling and sample aberrations. - """ - def __init__(self, sim_fp, x_size, y_size, h5_data, nm_per_pixel, zmn, - wavelength = pf_wavelength, - refractive_index = pf_refractive_index, - numerical_aperture = pf_numerical_aperture, - pf_size = pf_size, - bead_diameter = None, - sample_index = None, - z_center = None): - """ - zmn is a list of lists containing the zernike mode terms, e.g. - [[1.3, 2, 2]] for pure astigmatism. - wavelength is the mean emission wavelength in nm. - - bead_diameter is the bead diameter in microns. - sample_index is the index of the sample media. - z_center is the focal plane position in microns. - - If you specify sample_index you also need to specify z_center. - """ - super(PupilFunctionVectorial, self).__init__(sim_fp, x_size, y_size, h5_data, nm_per_pixel, pf_size) - self.saveJSON({"psf" : {"class" : "PupilFunctionVectorial", - "nm_per_pixel" : str(nm_per_pixel), - "numerical_aperture" : str(numerical_aperture), - "bead_diameter" : str(bead_diameter), - "pf_size" : str(pf_size), - "refactrive_index" : str(refractive_index), - "sample_index" : str(sample_index), - "wavelength" : str(wavelength), - "z_center" : str(z_center), - "zmn" : str(zmn)}}) - - self.geo = pupilMath.GeometryVectorial(pf_size, - nm_per_pixel * 0.001, - wavelength * 0.001, - refractive_index, - numerical_aperture) - - self.pf = self.geo.createFromZernike(1.0, zmn) - - if bead_diameter is not None: - self.otf_scaler = self.geo.beadScalingFactor(bead_diameter) - - self.sample_index = sample_index - self.z_center = z_center - - def getAberrationFn(self, z_value): - - # Apply sample index aberration if requested. - if self.sample_index is not None: - - # Distance 0.0 is the coverslip. Z values are relative to the focal - # plane which is at z_center. - # - depth = self.z_center + z_value - - if (depth < 0.0): - depth = 0.0 - - return self.geo.aberration(depth, self.sample_index) - - else: - return None - - -class Spline2D(splineToPSF.SplineToPSF2D): - """ - 2D spline with non-zero offsets in x, y. - """ - def __init__(self, spline_file): - super(Spline2D, self).__init__(spline_file) - self.psf_size = self.spline_size - 1 - - def getPSF(self, z_value, dx, dy): - """ - Use C library for better performance. - - z_value is ignored, it is only a parameter so that - the signature matches Spline3D. - - dx, dy are in the range 0.0 - 1.0. - """ - if((self.psf_size%2) == 0): - return self.c_spline.getPSF(dy + 0.5, dx + 0.5) - else: - return self.c_spline.getPSF(dy + 1.0, dx + 1.0) - - def getPSFPy(self, z_value, dx, dy): - """ - Python reference version. - - z_value is ignored, it is only a parameter so that - the signature matches Spline3D. - - dx, dy are in the range 0.0 - 1.0. - """ - psf = numpy.zeros((self.psf_size, self.psf_size)) - if((self.psf_size%2) == 0): - for x in range(self.psf_size): - for y in range(self.psf_size): - psf[y,x] = self.spline.f(float(y+dy) + 0.5, - float(x+dx) + 0.5) - else: - for x in range(self.psf_size): - for y in range(self.psf_size): - psf[y,x] = self.spline.f(float(y+dy) + 1.0, - float(x+dx) + 1.0) - - return psf - - -class Spline3D(splineToPSF.SplineToPSF3D): - """ - 3D spline with non-zero offsets in x, y. - """ - def __init__(self, spline_file): - super(Spline3D, self).__init__(spline_file) - self.psf_size = self.spline_size - 1 - - def getPSF(self, z_value, dx, dy): - """ - Use C library for better performance. - - z_value is ignored, it is only a parameter so that - the signature matches Spline3D. - - dx, dy are in the range 0.0 - 1.0. - """ - scaled_z = self.getScaledZ(z_value) - if((self.psf_size%2) == 0): - return self.c_spline.getPSF(scaled_z, dy + 0.5, dx + 0.5) - else: - return self.c_spline.getPSF(scaled_z, dy + 1.0, dx + 1.0) - - def getPSFPy(self, z_value, dx, dy): - """ - Python reference version. - - z_value needs to be inside the z range covered by the spline. - z_value should be in nanometers. - - dx, dy are in the range 0.0 - 1.0. - """ - scaled_z = self.getScaledZ(z_value) - - psf = numpy.zeros((self.psf_size, self.psf_size)) - if((self.psf_size%2) == 0): - for x in range(self.psf_size): - for y in range(self.psf_size): - psf[y,x] = self.spline.f(scaled_z, - float(y+dy) + 0.5, - float(x+dx) + 0.5) - else: - for x in range(self.psf_size): - for y in range(self.psf_size): - psf[y,x] = self.spline.f(scaled_z, - float(y+dy) + 1.0, - float(x+dx) + 1.0) - - return psf - - -class Spline(PSF): - """ - PSF from a (cubic) spline. - """ - def __init__(self, sim_fp, x_size, y_size, h5_data, nm_per_pixel, spline_file): - """ - spline_file is the name of a .spline file as generated by spliner/psf_to_spline.py. - - Splines are always pixel based, so the spline pixel size should match what you - want for the simulation. The pixel size of the spline is recorded in the spline - file and in theory that could be used to adjust appropriately depending on the - desired pixel size for the simulation, but that is not currently being done. - """ - super(Spline, self).__init__(sim_fp, x_size, y_size, h5_data, nm_per_pixel) - self.saveJSON({"psf" : {"class" : "Spline", - "spline_file" : spline_file}}) - with open(spline_file, 'rb') as fp: - spline_data = pickle.load(fp) - - if (spline_data["type"] == "3D"): - self.spline = Spline3D(spline_data) - else: - self.spline = Spline2D(spline_data) - - self.psf_size = self.spline.psf_size - - if ((self.psf_size%2)==0): - self.margin = int(self.psf_size/2) - else: - self.margin = int(self.psf_size/2) + 1 - - self.im_size_x = self.x_size + 2*self.margin - self.im_size_y = self.y_size + 2*self.margin - - def cleanup(self): - """ - Clean up the C library. - """ - self.spline.c_spline.cleanup() - - def getPSFs(self, h5_data): - """ - The expected form for the h5 data fields are x,y in pixels and z in microns. - """ - image = numpy.zeros((self.im_size_x, self.im_size_y)) - x = h5_data['x']+1 # Pixels - y = h5_data['y']+1 # Pixels - z = 1.0*h5_data['z']*1000.0 - a = h5_data['sum'] - - h5_data['height'] = numpy.zeros(a.size) - - dx = 1.0 - (x - numpy.floor(x)) - dy = 1.0 - (y - numpy.floor(y)) - - for i in range(x.size): - - ix = int(x[i]) - iy = int(y[i]) - - if (ix >= 0.0) and (ix < self.x_size) and (iy >= 0.0) and (iy < self.y_size): - - # Calculate psf, dx and dy are transposed to match the spline coordinate system. - psf = self.spline.getPSF(z[i], dx[i], dy[i]).transpose() - - # Remove negative values. - psf[(psf < 0.0)] = 0.0 - - # Scale to correct number of photons. - psf = a[i] * psf/numpy.sum(psf) - h5_data['height'][i] = numpy.max(psf) - - # Add to image. - image[ix:ix+self.psf_size,iy:iy+self.psf_size] += psf - - return image[self.margin:self.margin+self.x_size,self.margin:self.margin+self.y_size] - - -# -# The MIT License -# -# Copyright (c) 2016 Zhuang Lab, Harvard University -# -# Permission is hereby granted, free of charge, to any person obtaining a copy -# of this software and associated documentation files (the "Software"), to deal -# in the Software without restriction, including without limitation the rights -# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -# copies of the Software, and to permit persons to whom the Software is -# furnished to do so, subject to the following conditions: -# -# The above copyright notice and this permission notice shall be included in -# all copies or substantial portions of the Software. -# -# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN -# THE SOFTWARE. -# +#!/usr/bin/env python +""" +Classes for creating different kinds of PSFs. + +Note that as a side-effect the getPSFs() method is expected +to also set the correct value of the 'h' field based on the +PSF and the 'a' field. + +Hazen 11/16 +""" + +import numpy +import pickle +import random +import scipy +import scipy.fftpack + +import storm_analysis.spliner.spline_to_psf as splineToPSF + +import storm_analysis.simulator.draw_gaussians_c as dg +import storm_analysis.simulator.pupil_math as pupilMath +import storm_analysis.simulator.simbase as simbase + +# These are here to make it easier to create pupil functions that match +# those used in the simulations. +pf_size = 128 +pf_wavelength = 600 +pf_refractive_index = 1.5 +pf_numerical_aperture = 1.4 + +class PSF(simbase.SimBase): + """ + Draws the emitter PSFs on an image. + """ + def __init__(self, sim_fp, x_size, y_size, h5_data, nm_per_pixel): + super(PSF, self).__init__(sim_fp, x_size, y_size, h5_data) + self.nm_per_pixel = nm_per_pixel + + +class DHPSF(PSF): + """ + A very simplistic approximation of the double helix PSF. + """ + def __init__(self, sim_fp, x_size, y_size, h5_data, nm_per_pixel, z_range = 0.75): + """ + z_range - (half) Z range in microns. + """ + super(DHPSF, self).__init__(sim_fp, x_size, y_size, h5_data, nm_per_pixel) + self.z_max = z_range + self.z_min = -z_range + + self.saveJSON({"psf" : {"class" : "DHPSF", + "z_range" : str(z_range)}}) + + def getPSFs(self, h5_data): + x = h5_data['x'] + y = h5_data['y'] + a = h5_data['sum'] + + sx = h5_data['ysigma'] + sy = h5_data['xsigma'] + + h = 0.5*a/(2.0 * numpy.pi * sx * sy) + h5_data['height'] = h + + angle = numpy.pi * 0.9 * ((h5_data['z'] - self.z_min)/(self.z_max - self.z_min)) + dx = 2.0 * sx * numpy.cos(angle) + dy = 2.0 * sy * numpy.sin(angle) + + x1 = x + dx + y1 = y + dy + x2 = x - dx + y2 = y - dy + + x = numpy.concatenate((x1, x2)) + y = numpy.concatenate((y1, y2)) + h = numpy.concatenate((h, h)) + sx = numpy.concatenate((sx, sx)) + sy = numpy.concatenate((sy, sy)) + + return dg.drawGaussians((self.x_size, self.y_size), + numpy.stack((x, y, h, sx, sy), axis = 1), + res = 5) + + +class GaussianPSF(PSF): + """ + Gaussian PSF. + """ + def __init__(self, sim_fp, x_size, y_size, h5_data, nm_per_pixel): + super(GaussianPSF, self).__init__(sim_fp, x_size, y_size, h5_data, nm_per_pixel) + self.saveJSON({"psf" : {"class" : "GaussianPSF", + "nm_per_pixel" : str(nm_per_pixel)}}) + + def getPSFs(self, h5_data): + x = h5_data['x'] + y = h5_data['y'] + a = h5_data['sum'] + + sx = h5_data['ysigma'] + sy = h5_data['xsigma'] + + h = a/(2.0 * numpy.pi * sx * sy) + h5_data['height'] = h + + return dg.drawGaussians((self.x_size, self.y_size), + numpy.concatenate((x[:,None], + y[:,None], + h[:,None], + sx[:,None], + sy[:,None]), + axis = 1), + res = 5) + + +class PupilFunctionBase(PSF): + """ + PSF using the pupil function approach. + """ + def __init__(self, sim_fp, x_size, y_size, h5_data, nm_per_pixel, pf_size): + super(PupilFunctionBase, self).__init__(sim_fp, x_size, y_size, h5_data, nm_per_pixel) + + self.otf_scaler = None + self.psf_size = pf_size + self.sample_index = None + self.z_center = None + + self.margin = int(self.psf_size/2) + 1 + + self.im_size_x = self.x_size + 2 * self.margin + self.im_size_y = self.y_size + 2 * self.margin + + def getAberrationFn(self, z_value): + return None + + def getPSFs(self, h5_data): + """ + The expected form for the h5 data fields are x,y in pixels and z in microns. + """ + image = numpy.zeros((self.im_size_x, self.im_size_y)) + x = h5_data['x']+1 # Pixels + y = h5_data['y']+1 # Pixels + z = h5_data['z'] # Expected to be in microns. + a = h5_data['sum'] + + h5_data['height'] = numpy.zeros(a.size) + + dx = x - numpy.floor(x) + dy = y - numpy.floor(y) + + for i in range(x.size): + + ix = int(x[i]) + iy = int(y[i]) + + if (ix >= 0.0) and (ix < self.x_size) and (iy >= 0.0) and (iy < self.y_size): + + pf = self.pf + + # Apply aberration function if available. + ab_fn = self.getAberrationFn(z[i]) + if ab_fn is not None: + pf = pf*ab_fn + + # Shift to the desired z value. + defocused = self.geo.changeFocus(pf, z[i]) + + # Translate to the correct sub-pixel position. + #translated = defocused + translated = self.geo.translatePf(defocused, dx[i], dy[i]) + + # Get real-space intensity. + psf = pupilMath.intensity(pupilMath.toRealSpace(translated)) * a[i] + + # Apply OTF scaling if requested. + if self.otf_scaler is not None: + otf = scipy.fftpack.fftshift(scipy.fftpack.fft2(psf)) + otf_scaled = otf * self.otf_scaler + psf = numpy.abs(scipy.fftpack.ifft2(otf_scaled)) + + h5_data['height'][i] = numpy.max(psf) + + image[ix:ix+self.psf_size,iy:iy+self.psf_size] += psf + + return image[self.margin:self.margin+self.x_size,self.margin:self.margin+self.y_size] + + +class PupilFunction(PupilFunctionBase): + """ + PSF using the pupil function approach. This is the most basic version. + + Note that by default this uses pupilMath.GeometrySim class which (more + or less) emulates OTF scaling by using a pixel size that is half the + specified pixel size to create the PSF. + """ + def __init__(self, sim_fp, x_size, y_size, h5_data, nm_per_pixel, zmn, + wavelength = pf_wavelength, + refractive_index = pf_refractive_index, + numerical_aperture = pf_numerical_aperture, + pf_size = pf_size, + geo_sim_pf = True): + """ + zmn is a list of lists containing the zernike mode terms, e.g. + [[1.3, 2, 2]] for pure astigmatism. + wavelength is the mean emission wavelength in nm. + """ + super(PupilFunction, self).__init__(sim_fp, x_size, y_size, h5_data, nm_per_pixel, pf_size) + self.saveJSON({"psf" : {"class" : "PupilFunction", + "geo_sim_pf" : str(geo_sim_pf), + "nm_per_pixel" : str(nm_per_pixel), + "numerical_aperture" : str(numerical_aperture), + "pf_size" : str(pf_size), + "refactrive_index" : str(refractive_index), + "wavelength" : str(wavelength), + "zmn" : str(zmn)}}) + + if geo_sim_pf: + self.geo = pupilMath.GeometrySim(pf_size, + nm_per_pixel * 0.001, + wavelength * 0.001, + refractive_index, + numerical_aperture) + else: + self.geo = pupilMath.Geometry(pf_size, + nm_per_pixel * 0.001, + wavelength * 0.001, + refractive_index, + numerical_aperture) + + self.pf = self.geo.createFromZernike(1.0, zmn) + + +class PupilFunctionScalar(PupilFunctionBase): + """ + This version supports OTF scaling and sample aberrations. + """ + def __init__(self, sim_fp, x_size, y_size, h5_data, nm_per_pixel, zmn, + wavelength = pf_wavelength, + refractive_index = pf_refractive_index, + numerical_aperture = pf_numerical_aperture, + pf_size = pf_size, + otf_sigma = None, + sample_index = None, + z_center = None): + """ + zmn is a list of lists containing the zernike mode terms, e.g. + [[1.3, 2, 2]] for pure astigmatism. + wavelength is the mean emission wavelength in nm. + + otf_sigma is the sigma to use for Gaussian OTF scaling. + sample_index is the index of the sample media. + z_center is the focal plane position in microns. + + If you specify sample_index you also need to specify z_center. + """ + super(PupilFunctionScalar, self).__init__(sim_fp, x_size, y_size, h5_data, nm_per_pixel, pf_size) + self.saveJSON({"psf" : {"class" : "PupilFunctionScalar", + "nm_per_pixel" : str(nm_per_pixel), + "numerical_aperture" : str(numerical_aperture), + "otf_sigma" : str(otf_sigma), + "pf_size" : str(pf_size), + "refactrive_index" : str(refractive_index), + "sample_index" : str(sample_index), + "wavelength" : str(wavelength), + "z_center" : str(z_center), + "zmn" : str(zmn)}}) + + self.geo = pupilMath.Geometry(pf_size, + nm_per_pixel * 0.001, + wavelength * 0.001, + refractive_index, + numerical_aperture) + + self.pf = self.geo.createFromZernike(1.0, zmn) + + if otf_sigma is not None: + self.otf_scaler = self.geo.gaussianScalingFactor(otf_sigma) + + self.sample_index = sample_index + self.z_center = z_center + + def getAberrationFn(self, z_value): + + # Apply sample index aberration if requested. + if self.sample_index is not None: + + # Distance 0.0 is the coverslip. Z values are relative to the focal + # plane which is at z_center. + # + depth = self.z_center + z_value + + if (depth < 0.0): + depth = 0.0 + + return self.geo.aberration(depth, self.sample_index) + + else: + return None + + +class PupilFunctionScalarCalibration(PupilFunctionBase): + """ + This version supports OTF scaling and a fixed sample aberration. It + simulates the PSF you'd expect from a calibration movie of a bead + on a coverslip. + """ + def __init__(self, sim_fp, x_size, y_size, h5_data, nm_per_pixel, zmn, + wavelength = pf_wavelength, + refractive_index = pf_refractive_index, + numerical_aperture = pf_numerical_aperture, + pf_size = pf_size, + bead_z_center = None, + otf_sigma = None, + sample_index = None): + """ + zmn is a list of lists containing the zernike mode terms, e.g. + [[1.3, 2, 2]] for pure astigmatism. + wavelength is the mean emission wavelength in nm. + + bead_z_center is the height of the bead above the coverslip in microns. + otf_sigma is the sigma to use for Gaussian OTF scaling. + sample_index is the index of the sample media. + + If you specify sample_index you also need to specify z_center. + """ + super(PupilFunctionScalarCalibration, self).__init__(sim_fp, x_size, y_size, h5_data, nm_per_pixel, pf_size) + self.saveJSON({"psf" : {"class" : "PupilFunctionScalarCalibration", + "bead_z_center" : str(bead_z_center), + "nm_per_pixel" : str(nm_per_pixel), + "numerical_aperture" : str(numerical_aperture), + "otf_sigma" : str(otf_sigma), + "pf_size" : str(pf_size), + "refactrive_index" : str(refractive_index), + "sample_index" : str(sample_index), + "wavelength" : str(wavelength), + "zmn" : str(zmn)}}) + + self.geo = pupilMath.Geometry(pf_size, + nm_per_pixel * 0.001, + wavelength * 0.001, + refractive_index, + numerical_aperture) + + self.pf = self.geo.createFromZernike(1.0, zmn) + + if otf_sigma is not None: + self.otf_scaler = self.geo.gaussianScalingFactor(otf_sigma) + + self.ab_fn = None + if sample_index is not None: + self.ab_fn = self.geo.aberration(bead_z_center, sample_index) + + def getAberrationFn(self, z_value): + return self.ab_fn + + +class PupilFunctionVectorial(PupilFunctionBase): + """ + This is a vectorial PSF version and supports bead scaling and sample aberrations. + """ + def __init__(self, sim_fp, x_size, y_size, h5_data, nm_per_pixel, zmn, + wavelength = pf_wavelength, + refractive_index = pf_refractive_index, + numerical_aperture = pf_numerical_aperture, + pf_size = pf_size, + bead_diameter = None, + sample_index = None, + z_center = None): + """ + zmn is a list of lists containing the zernike mode terms, e.g. + [[1.3, 2, 2]] for pure astigmatism. + wavelength is the mean emission wavelength in nm. + + bead_diameter is the bead diameter in microns. + sample_index is the index of the sample media. + z_center is the focal plane position in microns. + + If you specify sample_index you also need to specify z_center. + """ + super(PupilFunctionVectorial, self).__init__(sim_fp, x_size, y_size, h5_data, nm_per_pixel, pf_size) + self.saveJSON({"psf" : {"class" : "PupilFunctionVectorial", + "nm_per_pixel" : str(nm_per_pixel), + "numerical_aperture" : str(numerical_aperture), + "bead_diameter" : str(bead_diameter), + "pf_size" : str(pf_size), + "refactrive_index" : str(refractive_index), + "sample_index" : str(sample_index), + "wavelength" : str(wavelength), + "z_center" : str(z_center), + "zmn" : str(zmn)}}) + + self.geo = pupilMath.GeometryVectorial(pf_size, + nm_per_pixel * 0.001, + wavelength * 0.001, + refractive_index, + numerical_aperture) + + self.pf = self.geo.createFromZernike(1.0, zmn) + + if bead_diameter is not None: + self.otf_scaler = self.geo.beadScalingFactor(bead_diameter) + + self.sample_index = sample_index + self.z_center = z_center + + def getAberrationFn(self, z_value): + + # Apply sample index aberration if requested. + if self.sample_index is not None: + + # Distance 0.0 is the coverslip. Z values are relative to the focal + # plane which is at z_center. + # + depth = self.z_center + z_value + + if (depth < 0.0): + depth = 0.0 + + return self.geo.aberration(depth, self.sample_index) + + else: + return None + + +class Spline2D(splineToPSF.SplineToPSF2D): + """ + 2D spline with non-zero offsets in x, y. + """ + def __init__(self, spline_file): + super(Spline2D, self).__init__(spline_file) + self.psf_size = self.spline_size - 1 + + def getPSF(self, z_value, dx, dy): + """ + Use C library for better performance. + + z_value is ignored, it is only a parameter so that + the signature matches Spline3D. + + dx, dy are in the range 0.0 - 1.0. + """ + if((self.psf_size%2) == 0): + return self.c_spline.getPSF(dy + 0.5, dx + 0.5) + else: + return self.c_spline.getPSF(dy + 1.0, dx + 1.0) + + def getPSFPy(self, z_value, dx, dy): + """ + Python reference version. + + z_value is ignored, it is only a parameter so that + the signature matches Spline3D. + + dx, dy are in the range 0.0 - 1.0. + """ + psf = numpy.zeros((self.psf_size, self.psf_size)) + if((self.psf_size%2) == 0): + for x in range(self.psf_size): + for y in range(self.psf_size): + psf[y,x] = self.spline.f(float(y+dy) + 0.5, + float(x+dx) + 0.5) + else: + for x in range(self.psf_size): + for y in range(self.psf_size): + psf[y,x] = self.spline.f(float(y+dy) + 1.0, + float(x+dx) + 1.0) + + return psf + + +class Spline3D(splineToPSF.SplineToPSF3D): + """ + 3D spline with non-zero offsets in x, y. + """ + def __init__(self, spline_file): + super(Spline3D, self).__init__(spline_file) + self.psf_size = self.spline_size - 1 + + def getPSF(self, z_value, dx, dy): + """ + Use C library for better performance. + + z_value is ignored, it is only a parameter so that + the signature matches Spline3D. + + dx, dy are in the range 0.0 - 1.0. + """ + scaled_z = self.getScaledZ(z_value) + if((self.psf_size%2) == 0): + return self.c_spline.getPSF(scaled_z, dy + 0.5, dx + 0.5) + else: + return self.c_spline.getPSF(scaled_z, dy + 1.0, dx + 1.0) + + def getPSFPy(self, z_value, dx, dy): + """ + Python reference version. + + z_value needs to be inside the z range covered by the spline. + z_value should be in nanometers. + + dx, dy are in the range 0.0 - 1.0. + """ + scaled_z = self.getScaledZ(z_value) + + psf = numpy.zeros((self.psf_size, self.psf_size)) + if((self.psf_size%2) == 0): + for x in range(self.psf_size): + for y in range(self.psf_size): + psf[y,x] = self.spline.f(scaled_z, + float(y+dy) + 0.5, + float(x+dx) + 0.5) + else: + for x in range(self.psf_size): + for y in range(self.psf_size): + psf[y,x] = self.spline.f(scaled_z, + float(y+dy) + 1.0, + float(x+dx) + 1.0) + + return psf + + +class Spline(PSF): + """ + PSF from a (cubic) spline. + """ + def __init__(self, sim_fp, x_size, y_size, h5_data, nm_per_pixel, spline_file): + """ + spline_file is the name of a .spline file as generated by spliner/psf_to_spline.py. + + Splines are always pixel based, so the spline pixel size should match what you + want for the simulation. The pixel size of the spline is recorded in the spline + file and in theory that could be used to adjust appropriately depending on the + desired pixel size for the simulation, but that is not currently being done. + """ + super(Spline, self).__init__(sim_fp, x_size, y_size, h5_data, nm_per_pixel) + self.saveJSON({"psf" : {"class" : "Spline", + "spline_file" : spline_file}}) + with open(spline_file, 'rb') as fp: + spline_data = pickle.load(fp) + + if (spline_data["type"] == "3D"): + self.spline = Spline3D(spline_data) + else: + self.spline = Spline2D(spline_data) + + self.psf_size = self.spline.psf_size + + if ((self.psf_size%2)==0): + self.margin = int(self.psf_size/2) + else: + self.margin = int(self.psf_size/2) + 1 + + self.im_size_x = self.x_size + 2*self.margin + self.im_size_y = self.y_size + 2*self.margin + + def cleanup(self): + """ + Clean up the C library. + """ + self.spline.c_spline.cleanup() + + def getPSFs(self, h5_data): + """ + The expected form for the h5 data fields are x,y in pixels and z in microns. + """ + image = numpy.zeros((self.im_size_x, self.im_size_y)) + x = h5_data['x']+1 # Pixels + y = h5_data['y']+1 # Pixels + z = 1.0*h5_data['z']*1000.0 + a = h5_data['sum'] + + h5_data['height'] = numpy.zeros(a.size) + + dx = 1.0 - (x - numpy.floor(x)) + dy = 1.0 - (y - numpy.floor(y)) + + for i in range(x.size): + + ix = int(x[i]) + iy = int(y[i]) + + if (ix >= 0.0) and (ix < self.x_size) and (iy >= 0.0) and (iy < self.y_size): + + # Calculate psf, dx and dy are transposed to match the spline coordinate system. + psf = self.spline.getPSF(z[i], dx[i], dy[i]).transpose() + + # Remove negative values. + psf[(psf < 0.0)] = 0.0 + + # Scale to correct number of photons. + psf = a[i] * psf/numpy.sum(psf) + h5_data['height'][i] = numpy.max(psf) + + # Add to image. + image[ix:ix+self.psf_size,iy:iy+self.psf_size] += psf + + return image[self.margin:self.margin+self.x_size,self.margin:self.margin+self.y_size] + + +# +# The MIT License +# +# Copyright (c) 2016 Zhuang Lab, Harvard University +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. +# diff --git a/storm_analysis/simulator/pupil_math.py b/storm_analysis/simulator/pupil_math.py index cb03e712..3b205cd3 100644 --- a/storm_analysis/simulator/pupil_math.py +++ b/storm_analysis/simulator/pupil_math.py @@ -1,765 +1,765 @@ -#!/usr/bin/env python -""" -Some math for calculating PSFs from pupil functions. All units are in -microns. - -Important note - The default for the simulator, and what is also used -in the diagnostics, is a pupil function with a pixel size of 1/2 the -actual pixel size. This was done as it has a more realistic width. If -you use the correct pixel size the PSF will be too narrow. This can -also be handled using OTF scaling as described in the Hanser paper, -but as this is more complicated. Also the pupil function localization -software would need to be updated to include OTF scaling, which it -currently does not. - - -This is based on code provided by the Huang Lab at UCSF. - -McGorty et al. "Correction of depth-dependent aberrations in 3D -single-molecule localization and super-resolution microscopy", -Optics Letters, 2014. - - -Another reference for pupil functions is: - -Hanser et al. "Phase-retrieved pupil functions in wide-field fluorescence -microscopy", Journal of Microscopy, 2004. - - -Reference for sample index mismatch aberration: - -Liu et al., "Three dimensional single molecule localization using a phase -retrieved pupil function", Optics Express, 2013 - - -Also, thanks to Sjoerd Stallinga for providing his MATLAB code for calculating -vectorial PSFs. - -M. Siemons, C. N. Hulleman, R. O. Thorsen, C. S. Smith, and S. Stallinga, -"High precision wavefront control in point spread function engineering -for single emitter localization", Optics Express, 26, pp. 8397-8416, 2018. - -Hazen 05/19 -""" - -import math -import numpy -import scipy -import scipy.fftpack -import tifffile - -import storm_analysis -import storm_analysis.pupilfn.otf_scaling_c as otfSC -import storm_analysis.pupilfn.pupil_function_c as pfFnC -import storm_analysis.simulator.pf_math_c as pfMathC - - -class PupilMathException(storm_analysis.SAException): - pass - - -class Geometry(object): - - def __init__(self, size, pixel_size, wavelength, imm_index, NA): - """ - size - The number of pixels in the PSF image, assumed square - and a multiple of 2. - pixel_size - The size of the camera pixel in um. - wavelength - The wavelength of the flourescence in um. - imm_index - The index of the immersion media. - NA - The numerical aperature of the objective. - """ - super(Geometry, self).__init__() - - if not ((size%2)==0): - raise PupilMathException("PF size must be a multiple of 2!") - - # imm_index must be larger than the objective NA. - if (imm_index <= NA): - raise PupilMathException("Immersion media index must be larger than objective NA!") - - self.imm_index = float(imm_index) - self.NA = float(NA) - self.pixel_size = float(pixel_size) - self.size = int(size) - self.wavelength = float(wavelength) - - # Hanser, 2004, page 35. - self.k_max = NA/wavelength - - dk = 1.0/(size * pixel_size) - self.r_max = self.k_max/dk - - [x,y] = numpy.mgrid[ -self.size/2.0 : self.size/2.0, -self.size/2.0 : self.size/2.0] - - # Vectors to use for X/Y translation. - self.kx = x/size - self.ky = y/size - - kx = dk * x - ky = dk * y - self.k = numpy.sqrt(kx * kx + ky * ky) - - # Hanser, 2004, page 34. - tmp = imm_index/wavelength - - # Vector to use for Z translation. - self.kz = numpy.lib.scimath.sqrt(tmp * tmp - self.k * self.k) - - self.r = self.k/self.k_max - - self.kz[(self.r > 1.0)] = 0.0 - self.n_pixels = numpy.sum(self.r <= 1) - self.norm = math.sqrt(self.r.size) - - if False: - with tifffile.TiffWriter("kz.tif") as tf: - tf.write(numpy.abs(self.kz).astype(numpy.float32)) - tf.write(numpy.angle(self.kz).astype(numpy.float32)) - - def aberration(self, depth, smp_index): - """ - Models the effect of a refractive index difference between the sample - and the immersion media. See Hanser 2004, equations 4-9. - - depth - Point source depth in microns. - smp_index - Refractive index of the sample media. - - Returns total aberration function (Hanser 2004, equation 8). Multiply the PF by - this numpy array to include this aberration. - - This approach appears to have the problem that at the d = 0 limit it does not - converge to the no aberration PF. - """ - # Use complex numbers to include super critical angle (near field) fluorescence - # effects. - # - sin_theta_1 = (self.wavelength/self.imm_index)*self.k + 0j - - # Special handling of the center point where self.k = 0.0, this - # will cause problems because then theta_1 will also be 0.0 and - # we'll end up with 0.0/0.0 when we calculate amp_comp. So instead - # we just use a really small number. - # - cp = int(sin_theta_1.shape[0]/2) - sin_theta_1[cp,cp] = 1.0e-6 - - sin_theta_2 = (self.imm_index/smp_index)*sin_theta_1 + 0j - - theta_1 = numpy.arcsin(sin_theta_1) - theta_2 = numpy.arcsin(sin_theta_2) - - amp_trans = (sin_theta_1 * numpy.cos(theta_2)/numpy.sin(theta_1 + theta_2))*(1.0 + (1.0/numpy.cos(theta_2 - theta_1))) - amp_comp = (self.imm_index * numpy.tan(theta_2))/(smp_index * numpy.tan(theta_1)) - new_phase = numpy.pi * 2.0 * depth * (smp_index * numpy.cos(theta_2) - self.imm_index * numpy.cos(theta_1))/self.wavelength - - # Minus sign here because this gives what I believe to be the expected result. - # The amplitude of the aberration in the region where theta_2 is imaginary will - # decrease exponentially with increasing distance from coverslip. - # - ab_fn = amp_trans * amp_comp * numpy.exp(-1j * new_phase) - return self.applyNARestriction(ab_fn) - - def aberrationOPD(self, z_stage, z_p, smp_index): - """ - Calculates an optical path difference aberration like the ones used in the - Gibson-Lanni PSF model. - - Reference: - Liu et al., "Three dimensional single molecule localization using a phase - retrieved pupil function", Optics Express, 2013 - - Equations 7,8,9. - - z_stage - Stage offset above 0 (microns, positive value). - z_p - Relative position (microns). - smp_index - Refractive index of the sample. - - Returns the aberration function, multiply the PF by this numpy array to include - this aberration. - """ - # Calculate theta1 and theta2. Use complex numbers to include near - # field effects. - # - sin_theta_1 = (self.wavelength/self.imm_index)*self.k - theta_1 = numpy.arcsin(sin_theta_1 + 0j) - - sin_theta_2 = (self.imm_index/smp_index)*sin_theta_1 - theta_2 = numpy.arcsin(sin_theta_2 + 0j) - - z_o = smp_index/self.imm_index * z_stage - - # In the coverslip warning. - # - # The distance of the emitter above the coverslip is z_p + z_o. - # - if ((z_p + z_o) < 0.0): - print("Warning! Negative Z value detected in aberrationOPD()!") - - t1 = z_o * smp_index * numpy.cos(theta_2) - t2 = z_o * self.imm_index * (self.imm_index/smp_index) * numpy.cos(theta_1) - t3 = z_p * smp_index * numpy.cos(theta_2) - t4 = 2.0 * (numpy.pi/self.wavelength) * (t1 - t2 + t3) - - # Minus sign here because this gives what I believe to be the expected result. - # The amplitude of the aberration in the region where theta_2 is imaginary will - # decrease exponentially with increasing distance from coverslip. - # - pf_ab = numpy.exp(-1j * t4) - return self.applyNARestriction(pf_ab) - - def applyNARestriction(self, pupil_fn): - """ - pupil_fn - The pupil function to restrict the NA of. - - return - The NA restricted pupil function. - """ - pupil_fn[(self.r > 1.0)] = 0.0 - return pupil_fn - - def beadScalingFactor(self, diameter): - """ - diameter - Bead diameter in microns. - - Return a bead function to use for OTF scaling. - - Hanser, 2004, page 36. - """ - x = numpy.pi * self.k * diameter - - # Special handling of the center point where k = 0.0, at - # this point the correct value is 1.0. - # - cp = int(self.size/2) - x[cp,cp] = 1.0e-12 - - b_k = 3.0*(numpy.sin(x)/(x*x*x) - numpy.cos(x)/(x*x)) - b_k[cp,cp] = 1.0 - - return b_k - - def changeFocus(self, pupil_fn, z_dist): - """ - pupil_fn - The pupil function. - z_dist - The distance to the new focal plane. - - return - The pupil function at the new focal plane. - """ - return numpy.exp(1j * 2.0 * numpy.pi * self.kz * z_dist) * pupil_fn - - def createPlaneWave(self, n_photons): - """ - n_photons - The intensity of the pupil function. - - return - The pupil function for a plane wave. - """ - plane = numpy.sqrt(n_photons/self.n_pixels) * numpy.exp(1j * numpy.zeros(self.r.shape)) - if False: - tifffile.imwrite("plane.tif", numpy.angle(self.applyNARestriction(plane)).astype(numpy.float32)) - return self.applyNARestriction(plane) - - def createFromZernike(self, n_photons, zernike_modes): - """ - n_photons - The intensity of the pupil function - zernike_modes - List of lists, [[magnitude (in radians), m, n], [..]] - - return - The pupil function for this combination of zernike modes. - """ - if (len(zernike_modes) == 0): - return self.createPlaneWave(n_photons) - else: - phases = numpy.zeros(self.r.shape) - for zmn in zernike_modes: - phases = pfMathC.zernikeGrid(phases, zmn[0], zmn[1], zmn[2], radius = self.r_max) - zmnpf = numpy.sqrt(n_photons/self.n_pixels) * numpy.exp(1j * phases) - if False: - tifffile.imwrite("zmnpf.tif", numpy.angle(self.applyNARestriction(zmnpf)).astype(numpy.float32)) - return self.applyNARestriction(zmnpf) - - def dx(self, pupil_fn): - """ - Returns the derivative of the pupil function in x. - """ - return -1j * 2.0 * numpy.pi * self.kx * pupil_fn - - def gaussianScalingFactor(self, sigma): - """ - Returns a gaussian function to use for OTF rescaling. - """ - return numpy.exp(-1 * self.k * self.k / (2.0 * sigma * sigma)) - - def pfToPSF(self, pf, z_vals, want_intensity = True, scaling_factor = None): - """ - pf - A pupil function. - z_vals - The z values (focal planes) of the desired PSF. - want_intensity - (Optional) Return intensity, default is True. - scaling_factor - (Optional) The OTF rescaling factor, default is None. - - return - The PSF that corresponds to pf at the requested z_vals. - """ - if want_intensity: - psf = numpy.zeros((len(z_vals), pf.shape[0], pf.shape[1])) - for i, z in enumerate(z_vals): - defocused = toRealSpace(self.changeFocus(pf, z)) - if scaling_factor is not None: - otf = scipy.fftpack.fftshift(scipy.fftpack.fft2(intensity(defocused))) - otf_scaled = otf * scaling_factor - psf[i,:,:] = numpy.abs(scipy.fftpack.ifft2(otf_scaled)) - else: - psf[i,:,:] = intensity(defocused) - return psf - else: - if scaling_factor is not None: - raise PupilMathException("OTF scaling of a complex valued PSF is not supported!") - - psf = numpy.zeros((len(z_vals), pf.shape[0], pf.shape[1]), - dtype = numpy.complex128) - for i, z in enumerate(z_vals): - psf[i,:,:] = toRealSpace(self.changeFocus(pf, z)) - return psf - - def theoreticalOTF(self): - """ - Returns a theoretical OTF. - - OTF = 2 * (psi - cos(psi)sin(psi)) / pi - psi = inv_cos(lambda * wavelength / 2 * NA) - - Reference: - https://www.microscopyu.com/microscopy-basics/modulation-transfer-function - - I'm assuming that the formula in this reference is for the FT of the PSF and - not the FT of the square root of the PSF. - """ - tmp = (1.5 * self.wavelength * self.k / (2.0 * self.NA)) - tmp[(tmp > 1.0)] = 1.0 - tmp[(tmp < -1.0)] = -1.0 - - psi = numpy.arccos(tmp) - otf = 2.0 * (psi - numpy.cos(psi)*numpy.sin(psi)) / numpy.pi - - otf = otf/numpy.sum(otf) - - if False: - tifffile.imwrite("otf.tif", otf.astype(numpy.float32)) - - return otf - - def translatePf(self, pupil_fn, dx, dy): - """ - Translate the Pf using Fourier translation. - - pupil_fn - A pupil function. - dx - Translation in x in pixels. - dy - Translation in y in pixels. - - return - The PF translated by dx, dy. - """ - return numpy.exp(-1j * 2.0 * numpy.pi * (self.kx * dx + self.ky * dy)) * pupil_fn - - -class GeometryVectorial(Geometry): - """ - PF model including the vectorial nature of light. - """ - def __init__(self, size, pixel_size, wavelength, imm_index, NA): - """ - size - The number of pixels in the PSF image, assumed square. - pixel_size - The size of the camera pixel in um. - wavelength - The wavelength of the flourescence in um. - imm_index - The index of the immersion media. - NA - The numerical aperature of the objective. - """ - super(GeometryVectorial, self).__init__(size, pixel_size, wavelength, imm_index, NA) - - t1 = self.k*self.wavelength/self.imm_index - t1[(t1 > 1.0)] = 1.0 - - self.theta = numpy.arcsin(t1) - self.theta[(self.r > 1.0)] = 0.0 - - self.phi = numpy.arctan2(self.ky, self.kx) - self.phi[(self.r > 1.0)] = 0.0 - - c_theta = numpy.cos(self.theta) - c_theta[(self.r > 1.0)] = 0.0 - - s_theta = numpy.sin(self.theta) - s_theta[(self.r > 1.0)] = 0.0 - - c_phi = numpy.cos(self.phi) - c_phi[(self.r > 1.0)] = 0.0 - - s_phi = numpy.sin(self.phi) - s_phi[(self.r > 1.0)] = 0.0 - - self.px_ex = c_theta*c_phi*c_phi + s_phi*s_phi - self.px_ey = (1.0 - c_theta)*s_phi*c_phi - self.py_ex = (c_theta - 1.0)*s_phi*c_phi - self.py_ey = c_theta*s_phi*s_phi + c_phi*c_phi - self.pz_ex = s_theta*c_phi - self.pz_ey = s_theta*s_phi - - def pfToPSF(self, pf, z_vals, scaling_factor = None): - """ - This matches the super-class function of the same name, but there is no - option of returning the (complex) real space PF. Also the PSF that is - returned is that for a freely rotating dipole. - - pf - A pupil function. - z_vals - The z values (focal planes) of the desired PSF. - scaling_factor - (Optional) The OTF rescaling factor, default is None. - - return - The PSF that corresponds to pf at the requested z_vals. - """ - rs_pf = self.pfToRS(pf, z_vals) - psf = self.rsToPSF(rs_pf) - - if scaling_factor is not None: - for i in range(psf.shape[0]): - otf = scipy.fftpack.fftshift(scipy.fftpack.fft2(psf[i,:,:])) - otf_scaled = otf * scaling_factor - psf[i,:,:] = numpy.abs(scipy.fftpack.ifft2(otf_scaled)) - - return psf - - def pfToRS(self, pf, z_vals): - """ - pf - A pupil function. - z_vals - The z values (focal planes) of the desired PSF. - - return - The real space counterpart of the PF at z_vals as a list. - [[rs_px_ex, rs_px_ey], [rs_py_ex, rs_py_ey], [rs_pz_ex, rs_pz_ey]] - """ - rs_px_ex = numpy.zeros((len(z_vals), pf.shape[0], pf.shape[1]), - dtype = numpy.complex128) - rs_px_ey = numpy.zeros((len(z_vals), pf.shape[0], pf.shape[1]), - dtype = numpy.complex128) - - rs_py_ex = numpy.zeros((len(z_vals), pf.shape[0], pf.shape[1]), - dtype = numpy.complex128) - rs_py_ey = numpy.zeros((len(z_vals), pf.shape[0], pf.shape[1]), - dtype = numpy.complex128) - - rs_pz_ex = numpy.zeros((len(z_vals), pf.shape[0], pf.shape[1]), - dtype = numpy.complex128) - rs_pz_ey = numpy.zeros((len(z_vals), pf.shape[0], pf.shape[1]), - dtype = numpy.complex128) - - for i, z in enumerate(z_vals): - pf_at_z = self.changeFocus(pf, z) - - rs_px_ex[i,:,:] = toRealSpace(self.px_ex * pf_at_z) - rs_px_ey[i,:,:] = toRealSpace(self.px_ey * pf_at_z) - - rs_py_ex[i,:,:] = toRealSpace(self.py_ex * pf_at_z) - rs_py_ey[i,:,:] = toRealSpace(self.py_ey * pf_at_z) - - rs_pz_ex[i,:,:] = toRealSpace(self.pz_ex * pf_at_z) - rs_pz_ey[i,:,:] = toRealSpace(self.pz_ey * pf_at_z) - - return [[rs_px_ex, rs_px_ey], [rs_py_ex, rs_py_ey], [rs_pz_ex, rs_pz_ey]] - - def rsToPSF(self, rs_pf, ratios = [1.0/3.0, 1.0/3.0, 1.0/3.0]): - """ - rs_pf - List of real space PFs from the pfToRS() method. - ratios - Ratios of dipole contribution, should sum to 1.0. - """ - psf = numpy.zeros((rs_pf[0][0].shape[0], self.size, self.size)) - for j in range(3): - psf += ratios[j] * intensity(rs_pf[j][0]) - psf += ratios[j] * intensity(rs_pf[j][1]) - - return psf - - -class GeometryC(Geometry): - """ - This class uses some of the C libraries in pupilfn to do the heavy lifting. It assumes - that the OTF scaling array is symmetric in X/Y. - - Based on profiling this is 3-4x faster than the pure Python version. - """ - def __init__(self, size, pixel_size, wavelength, imm_index, NA): - """ - size - The number of pixels in the PSF image, assumed square. - pixel_size - The size of the camera pixel in um. - wavelength - The wavelength of the flourescence in um. - imm_index - The index of the immersion media. - NA - The numerical aperature of the objective. - """ - super(GeometryC, self).__init__(size, pixel_size, wavelength, imm_index, NA) - - self.otf_sc = otfSC.OTFScaler(size = size) - self.pf_c = pfFnC.PupilFunction(geometry = self) - - def __del__(self): - self.otf_sc.cleanup() - self.pf_c.cleanup() - - def pfToPSF(self, pf, z_vals, want_intensity = True, scaling_factor = None): - """ - pf - A pupil function. - z_vals - The z values (focal planes) of the desired PSF. - want_intensity - (Optional) Return intensity, default is True. - scaling_factor - (Optional) The OTF rescaling factor, default is None. - - return - The PSF that corresponds to pf at the requested z_vals. - """ - self.pf_c.setPF(pf) - - if want_intensity: - if scaling_factor is not None: - self.otf_sc.setScale(scaling_factor) - - psf = numpy.zeros((len(z_vals), pf.shape[0], pf.shape[1])) - for i, z in enumerate(z_vals): - self.pf_c.translateZ(z) - temp = self.pf_c.getPSFIntensity() - if scaling_factor is not None: - psf[i,:,:] = self.otf_sc.scale(temp) - else: - psf[i,:,:] = temp - - else: - if scaling_factor is not None: - raise PupilMathException("OTF scaling of a complex valued PSF is not supported!") - - psf = numpy.zeros((len(z_vals), pf.shape[0], pf.shape[1]), dtype = numpy.complex128) - for i, z in enumerate(z_vals): - self.pf_c.translateZ(z) - psf[i,:,:] = self.pf_c.getPSF() - return psf - - return psf - - -class GeometryCVectorial(GeometryVectorial): - """ - The vectorial PF version of GeometryC. - """ - def __init__(self, size, pixel_size, wavelength, imm_index, NA): - """ - size - The number of pixels in the PSF image, assumed square. - pixel_size - The size of the camera pixel in um. - wavelength - The wavelength of the flourescence in um. - imm_index - The index of the immersion media. - NA - The numerical aperature of the objective. - """ - super(GeometryCVectorial, self).__init__(size, pixel_size, wavelength, imm_index, NA) - - self.otf_sc = otfSC.OTFScaler(size = size) - self.pf_c = pfFnC.PupilFunction(geometry = self) - - def __del__(self): - self.otf_sc.cleanup() - self.pf_c.cleanup() - - def pfToPSF(self, pf, z_vals, scaling_factor = None): - """ - This matches the super-class function of the same name, but there is no - option of returning the (complex) real space PF. Also the PSF that is - returned is that for a freely rotating dipole. - - pf - A pupil function. - z_vals - The z values (focal planes) of the desired PSF. - scaling_factor - (Optional) The OTF rescaling factor, default is None. - - return - The PSF that corresponds to pf at the requested z_vals. - """ - rs_pf = self.pfToRS(pf, z_vals) - psf = self.rsToPSF(rs_pf) - - if scaling_factor is not None: - self.otf_sc.setScale(scaling_factor) - for i in range(psf.shape[0]): - psf[i,:,:] = self.otf_sc.scale(psf[i,:,:]) - - return psf - - def pfToRS(self, pf, z_vals): - """ - pf - A pupil function. - z_vals - The z values (focal planes) of the desired PSF. - - return - The real space counterpart of the PF at z_vals as a list. - [[rs_px_ex, rs_px_ey], [rs_py_ex, rs_py_ey], [rs_pz_ex, rs_pz_ey]] - """ - rs_px_ex = numpy.zeros((len(z_vals), pf.shape[0], pf.shape[1]), - dtype = numpy.complex128) - rs_px_ey = numpy.zeros((len(z_vals), pf.shape[0], pf.shape[1]), - dtype = numpy.complex128) - - rs_py_ex = numpy.zeros((len(z_vals), pf.shape[0], pf.shape[1]), - dtype = numpy.complex128) - rs_py_ey = numpy.zeros((len(z_vals), pf.shape[0], pf.shape[1]), - dtype = numpy.complex128) - - rs_pz_ex = numpy.zeros((len(z_vals), pf.shape[0], pf.shape[1]), - dtype = numpy.complex128) - rs_pz_ey = numpy.zeros((len(z_vals), pf.shape[0], pf.shape[1]), - dtype = numpy.complex128) - - self.pf_c.setPF(pf) - for i, z in enumerate(z_vals): - self.pf_c.translateZ(z) - - rs_px_ex[i,:,:] = self.pf_c.getPXEX() - rs_px_ey[i,:,:] = self.pf_c.getPXEY() - - rs_py_ex[i,:,:] = self.pf_c.getPYEX() - rs_py_ey[i,:,:] = self.pf_c.getPYEY() - - rs_pz_ex[i,:,:] = self.pf_c.getPZEX() - rs_pz_ey[i,:,:] = self.pf_c.getPZEY() - - return [[rs_px_ex, rs_px_ey], [rs_py_ex, rs_py_ey], [rs_pz_ex, rs_pz_ey]] - - def rsToPSF(self, rs_pf, ratios = [1.0/3.0, 1.0/3.0, 1.0/3.0]): - """ - rs_pf - List of real space PFs from the pfToRS() method. - ratios - Ratios of dipole contribution, should sum to 1.0. - """ - psf = numpy.zeros((rs_pf[0][0].shape[0], self.size, self.size)) - for j in range(3): - psf += ratios[j] * intensity(rs_pf[j][0]) - psf += ratios[j] * intensity(rs_pf[j][1]) - - return psf - - -class GeometrySim(Geometry): - """ - This class is used in the simulations. It divides the pixel size - by two so that simulations look more realistic without having - to add the overhead of OTF scaling. - """ - def __init__(self, size, pixel_size, wavelength, imm_index, NA): - """ - size - The number of pixels in the PSF image, assumed square. - pixel_size - The size of the camera pixel in um. - wavelength - The wavelength of the flourescence in um. - imm_index - The index of the immersion media. - NA - The numerical aperature of the objective. - """ - super(GeometrySim, self).__init__(size, 0.5*pixel_size, wavelength, imm_index, NA) - - -def intensity(x): - """ - x - The (numpy array) to convert to intensity. - - return - The product of x and the complex conjugate of x. - """ - return numpy.abs(x * numpy.conj(x)) - - -def toRealSpace(pupil_fn): - """ - pupil_fn - A pupil function. - - return - The pupil function in real space (as opposed to fourier space). - """ - return scipy.fftpack.ifftshift(math.sqrt(pupil_fn.size) * scipy.fftpack.ifft2(pupil_fn)) - - -if (__name__ == "__main__"): - - import pickle - import sys - - if (len(sys.argv) < 2): - print("usage: ") - exit() - - pixel_size = 0.10 - #pixel_size = 0.020 - wavelength = 0.6 - refractive_index = 1.5 - numerical_aperture = 1.4 - z_range = 1.0 - z_pixel_size = 0.010 - - geo = Geometry(int(20.0/pixel_size), - pixel_size, - wavelength, - refractive_index, - numerical_aperture) - - if (len(sys.argv) == 4): - zmn = [] - amp = float(sys.argv[3]) - with open(sys.argv[2]) as fp: - for line in fp: - data = line.strip().split(" ") - if (len(data) == 3): - zmn.append([amp * float(data[2]), int(data[0]), int(data[1])]) - else: - zmn = [[1.3, 2, 2]] - #zmn = [[1, 0, 4]] - #zmn = [] - - pf = geo.createFromZernike(1.0, zmn) - z_values = numpy.arange(-z_range, z_range + 0.5 * z_pixel_size, z_pixel_size) - psfs = geo.pfToPSF(pf, z_values) - - #xy_size = 2.0*psfs.shape[0] - xy_size = 100 - xy_start = int(0.5 * (psfs.shape[1] - xy_size) + 1) - xy_end = int(xy_start + xy_size) - psfs = psfs[:,xy_start:xy_end,xy_start:xy_end] - - if True: - tifffile.imwrite("kz.tif", numpy.real(geo.kz).astype(numpy.float32)) - - if False: - tifffile.imwrite("pf_abs.tif", numpy.abs(pf).astype(numpy.float32)) - tifffile.imwrite("pf_angle.tif", (180.0 * numpy.angle(pf)/numpy.pi + 180).astype(numpy.float32)) - - if False: - with tifffile.TiffWriter(sys.argv[1]) as psf_tif: - temp = (psfs/numpy.max(psfs)).astype(numpy.float32) - psf_tif.write(temp) - - if False: - with open("z_offset.txt", "w") as fp: - for i in range(z_values.size): - fp.write("1 {0:.6f}\n".format(1000.0 * z_values[i])) - - if False: - psfs = (65000.0 * (psfs/numpy.max(psfs))).astype(numpy.uint16) - psf_dict = {"pixel_size" : pixel_size, - "wavelength" : wavelength, - "refractive_index" : refractive_index, - "numerical_aperture" : numerical_aperture, - "z_range" : z_range, - "zmm" : zmn, - "psf" : psfs} - pickle.dump(psf_dict, open(sys.argv[1], "wb"), protocol = 2) - - -# -# The MIT License -# -# Copyright (c) 2016 Zhuang Lab, Harvard University -# -# Permission is hereby granted, free of charge, to any person obtaining a copy -# of this software and associated documentation files (the "Software"), to deal -# in the Software without restriction, including without limitation the rights -# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -# copies of the Software, and to permit persons to whom the Software is -# furnished to do so, subject to the following conditions: -# -# The above copyright notice and this permission notice shall be included in -# all copies or substantial portions of the Software. -# -# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN -# THE SOFTWARE. -# - - +#!/usr/bin/env python +""" +Some math for calculating PSFs from pupil functions. All units are in +microns. + +Important note - The default for the simulator, and what is also used +in the diagnostics, is a pupil function with a pixel size of 1/2 the +actual pixel size. This was done as it has a more realistic width. If +you use the correct pixel size the PSF will be too narrow. This can +also be handled using OTF scaling as described in the Hanser paper, +but as this is more complicated. Also the pupil function localization +software would need to be updated to include OTF scaling, which it +currently does not. + + +This is based on code provided by the Huang Lab at UCSF. + +McGorty et al. "Correction of depth-dependent aberrations in 3D +single-molecule localization and super-resolution microscopy", +Optics Letters, 2014. + + +Another reference for pupil functions is: + +Hanser et al. "Phase-retrieved pupil functions in wide-field fluorescence +microscopy", Journal of Microscopy, 2004. + + +Reference for sample index mismatch aberration: + +Liu et al., "Three dimensional single molecule localization using a phase +retrieved pupil function", Optics Express, 2013 + + +Also, thanks to Sjoerd Stallinga for providing his MATLAB code for calculating +vectorial PSFs. + +M. Siemons, C. N. Hulleman, R. O. Thorsen, C. S. Smith, and S. Stallinga, +"High precision wavefront control in point spread function engineering +for single emitter localization", Optics Express, 26, pp. 8397-8416, 2018. + +Hazen 05/19 +""" + +import math +import numpy +import scipy +import scipy.fftpack +import tifffile + +import storm_analysis +import storm_analysis.pupilfn.otf_scaling_c as otfSC +import storm_analysis.pupilfn.pupil_function_c as pfFnC +import storm_analysis.simulator.pf_math_c as pfMathC + + +class PupilMathException(storm_analysis.SAException): + pass + + +class Geometry(object): + + def __init__(self, size, pixel_size, wavelength, imm_index, NA): + """ + size - The number of pixels in the PSF image, assumed square + and a multiple of 2. + pixel_size - The size of the camera pixel in um. + wavelength - The wavelength of the flourescence in um. + imm_index - The index of the immersion media. + NA - The numerical aperature of the objective. + """ + super(Geometry, self).__init__() + + if not ((size%2)==0): + raise PupilMathException("PF size must be a multiple of 2!") + + # imm_index must be larger than the objective NA. + if (imm_index <= NA): + raise PupilMathException("Immersion media index must be larger than objective NA!") + + self.imm_index = float(imm_index) + self.NA = float(NA) + self.pixel_size = float(pixel_size) + self.size = int(size) + self.wavelength = float(wavelength) + + # Hanser, 2004, page 35. + self.k_max = NA/wavelength + + dk = 1.0/(size * pixel_size) + self.r_max = self.k_max/dk + + [x,y] = numpy.mgrid[ -self.size/2.0 : self.size/2.0, -self.size/2.0 : self.size/2.0] + + # Vectors to use for X/Y translation. + self.kx = x/size + self.ky = y/size + + kx = dk * x + ky = dk * y + self.k = numpy.sqrt(kx * kx + ky * ky) + + # Hanser, 2004, page 34. + tmp = imm_index/wavelength + + # Vector to use for Z translation. + self.kz = numpy.lib.scimath.sqrt(tmp * tmp - self.k * self.k) + + self.r = self.k/self.k_max + + self.kz[(self.r > 1.0)] = 0.0 + self.n_pixels = numpy.sum(self.r <= 1) + self.norm = math.sqrt(self.r.size) + + if False: + with tifffile.TiffWriter("kz.tif") as tf: + tf.write(numpy.abs(self.kz).astype(numpy.float32)) + tf.write(numpy.angle(self.kz).astype(numpy.float32)) + + def aberration(self, depth, smp_index): + """ + Models the effect of a refractive index difference between the sample + and the immersion media. See Hanser 2004, equations 4-9. + + depth - Point source depth in microns. + smp_index - Refractive index of the sample media. + + Returns total aberration function (Hanser 2004, equation 8). Multiply the PF by + this numpy array to include this aberration. + + This approach appears to have the problem that at the d = 0 limit it does not + converge to the no aberration PF. + """ + # Use complex numbers to include super critical angle (near field) fluorescence + # effects. + # + sin_theta_1 = (self.wavelength/self.imm_index)*self.k + 0j + + # Special handling of the center point where self.k = 0.0, this + # will cause problems because then theta_1 will also be 0.0 and + # we'll end up with 0.0/0.0 when we calculate amp_comp. So instead + # we just use a really small number. + # + cp = int(sin_theta_1.shape[0]/2) + sin_theta_1[cp,cp] = 1.0e-6 + + sin_theta_2 = (self.imm_index/smp_index)*sin_theta_1 + 0j + + theta_1 = numpy.arcsin(sin_theta_1) + theta_2 = numpy.arcsin(sin_theta_2) + + amp_trans = (sin_theta_1 * numpy.cos(theta_2)/numpy.sin(theta_1 + theta_2))*(1.0 + (1.0/numpy.cos(theta_2 - theta_1))) + amp_comp = (self.imm_index * numpy.tan(theta_2))/(smp_index * numpy.tan(theta_1)) + new_phase = numpy.pi * 2.0 * depth * (smp_index * numpy.cos(theta_2) - self.imm_index * numpy.cos(theta_1))/self.wavelength + + # Minus sign here because this gives what I believe to be the expected result. + # The amplitude of the aberration in the region where theta_2 is imaginary will + # decrease exponentially with increasing distance from coverslip. + # + ab_fn = amp_trans * amp_comp * numpy.exp(-1j * new_phase) + return self.applyNARestriction(ab_fn) + + def aberrationOPD(self, z_stage, z_p, smp_index): + """ + Calculates an optical path difference aberration like the ones used in the + Gibson-Lanni PSF model. + + Reference: + Liu et al., "Three dimensional single molecule localization using a phase + retrieved pupil function", Optics Express, 2013 + + Equations 7,8,9. + + z_stage - Stage offset above 0 (microns, positive value). + z_p - Relative position (microns). + smp_index - Refractive index of the sample. + + Returns the aberration function, multiply the PF by this numpy array to include + this aberration. + """ + # Calculate theta1 and theta2. Use complex numbers to include near + # field effects. + # + sin_theta_1 = (self.wavelength/self.imm_index)*self.k + theta_1 = numpy.arcsin(sin_theta_1 + 0j) + + sin_theta_2 = (self.imm_index/smp_index)*sin_theta_1 + theta_2 = numpy.arcsin(sin_theta_2 + 0j) + + z_o = smp_index/self.imm_index * z_stage + + # In the coverslip warning. + # + # The distance of the emitter above the coverslip is z_p + z_o. + # + if ((z_p + z_o) < 0.0): + print("Warning! Negative Z value detected in aberrationOPD()!") + + t1 = z_o * smp_index * numpy.cos(theta_2) + t2 = z_o * self.imm_index * (self.imm_index/smp_index) * numpy.cos(theta_1) + t3 = z_p * smp_index * numpy.cos(theta_2) + t4 = 2.0 * (numpy.pi/self.wavelength) * (t1 - t2 + t3) + + # Minus sign here because this gives what I believe to be the expected result. + # The amplitude of the aberration in the region where theta_2 is imaginary will + # decrease exponentially with increasing distance from coverslip. + # + pf_ab = numpy.exp(-1j * t4) + return self.applyNARestriction(pf_ab) + + def applyNARestriction(self, pupil_fn): + """ + pupil_fn - The pupil function to restrict the NA of. + + return - The NA restricted pupil function. + """ + pupil_fn[(self.r > 1.0)] = 0.0 + return pupil_fn + + def beadScalingFactor(self, diameter): + """ + diameter - Bead diameter in microns. + + Return a bead function to use for OTF scaling. + + Hanser, 2004, page 36. + """ + x = numpy.pi * self.k * diameter + + # Special handling of the center point where k = 0.0, at + # this point the correct value is 1.0. + # + cp = int(self.size/2) + x[cp,cp] = 1.0e-12 + + b_k = 3.0*(numpy.sin(x)/(x*x*x) - numpy.cos(x)/(x*x)) + b_k[cp,cp] = 1.0 + + return b_k + + def changeFocus(self, pupil_fn, z_dist): + """ + pupil_fn - The pupil function. + z_dist - The distance to the new focal plane. + + return - The pupil function at the new focal plane. + """ + return numpy.exp(1j * 2.0 * numpy.pi * self.kz * z_dist) * pupil_fn + + def createPlaneWave(self, n_photons): + """ + n_photons - The intensity of the pupil function. + + return - The pupil function for a plane wave. + """ + plane = numpy.sqrt(n_photons/self.n_pixels) * numpy.exp(1j * numpy.zeros(self.r.shape)) + if False: + tifffile.imwrite("plane.tif", numpy.angle(self.applyNARestriction(plane)).astype(numpy.float32)) + return self.applyNARestriction(plane) + + def createFromZernike(self, n_photons, zernike_modes): + """ + n_photons - The intensity of the pupil function + zernike_modes - List of lists, [[magnitude (in radians), m, n], [..]] + + return - The pupil function for this combination of zernike modes. + """ + if (len(zernike_modes) == 0): + return self.createPlaneWave(n_photons) + else: + phases = numpy.zeros(self.r.shape) + for zmn in zernike_modes: + phases = pfMathC.zernikeGrid(phases, zmn[0], zmn[1], zmn[2], radius = self.r_max) + zmnpf = numpy.sqrt(n_photons/self.n_pixels) * numpy.exp(1j * phases) + if False: + tifffile.imwrite("zmnpf.tif", numpy.angle(self.applyNARestriction(zmnpf)).astype(numpy.float32)) + return self.applyNARestriction(zmnpf) + + def dx(self, pupil_fn): + """ + Returns the derivative of the pupil function in x. + """ + return -1j * 2.0 * numpy.pi * self.kx * pupil_fn + + def gaussianScalingFactor(self, sigma): + """ + Returns a gaussian function to use for OTF rescaling. + """ + return numpy.exp(-1 * self.k * self.k / (2.0 * sigma * sigma)) + + def pfToPSF(self, pf, z_vals, want_intensity = True, scaling_factor = None): + """ + pf - A pupil function. + z_vals - The z values (focal planes) of the desired PSF. + want_intensity - (Optional) Return intensity, default is True. + scaling_factor - (Optional) The OTF rescaling factor, default is None. + + return - The PSF that corresponds to pf at the requested z_vals. + """ + if want_intensity: + psf = numpy.zeros((len(z_vals), pf.shape[0], pf.shape[1])) + for i, z in enumerate(z_vals): + defocused = toRealSpace(self.changeFocus(pf, z)) + if scaling_factor is not None: + otf = scipy.fftpack.fftshift(scipy.fftpack.fft2(intensity(defocused))) + otf_scaled = otf * scaling_factor + psf[i,:,:] = numpy.abs(scipy.fftpack.ifft2(otf_scaled)) + else: + psf[i,:,:] = intensity(defocused) + return psf + else: + if scaling_factor is not None: + raise PupilMathException("OTF scaling of a complex valued PSF is not supported!") + + psf = numpy.zeros((len(z_vals), pf.shape[0], pf.shape[1]), + dtype = numpy.complex128) + for i, z in enumerate(z_vals): + psf[i,:,:] = toRealSpace(self.changeFocus(pf, z)) + return psf + + def theoreticalOTF(self): + """ + Returns a theoretical OTF. + + OTF = 2 * (psi - cos(psi)sin(psi)) / pi + psi = inv_cos(lambda * wavelength / 2 * NA) + + Reference: + https://www.microscopyu.com/microscopy-basics/modulation-transfer-function + + I'm assuming that the formula in this reference is for the FT of the PSF and + not the FT of the square root of the PSF. + """ + tmp = (1.5 * self.wavelength * self.k / (2.0 * self.NA)) + tmp[(tmp > 1.0)] = 1.0 + tmp[(tmp < -1.0)] = -1.0 + + psi = numpy.arccos(tmp) + otf = 2.0 * (psi - numpy.cos(psi)*numpy.sin(psi)) / numpy.pi + + otf = otf/numpy.sum(otf) + + if False: + tifffile.imwrite("otf.tif", otf.astype(numpy.float32)) + + return otf + + def translatePf(self, pupil_fn, dx, dy): + """ + Translate the Pf using Fourier translation. + + pupil_fn - A pupil function. + dx - Translation in x in pixels. + dy - Translation in y in pixels. + + return - The PF translated by dx, dy. + """ + return numpy.exp(-1j * 2.0 * numpy.pi * (self.kx * dx + self.ky * dy)) * pupil_fn + + +class GeometryVectorial(Geometry): + """ + PF model including the vectorial nature of light. + """ + def __init__(self, size, pixel_size, wavelength, imm_index, NA): + """ + size - The number of pixels in the PSF image, assumed square. + pixel_size - The size of the camera pixel in um. + wavelength - The wavelength of the flourescence in um. + imm_index - The index of the immersion media. + NA - The numerical aperature of the objective. + """ + super(GeometryVectorial, self).__init__(size, pixel_size, wavelength, imm_index, NA) + + t1 = self.k*self.wavelength/self.imm_index + t1[(t1 > 1.0)] = 1.0 + + self.theta = numpy.arcsin(t1) + self.theta[(self.r > 1.0)] = 0.0 + + self.phi = numpy.arctan2(self.ky, self.kx) + self.phi[(self.r > 1.0)] = 0.0 + + c_theta = numpy.cos(self.theta) + c_theta[(self.r > 1.0)] = 0.0 + + s_theta = numpy.sin(self.theta) + s_theta[(self.r > 1.0)] = 0.0 + + c_phi = numpy.cos(self.phi) + c_phi[(self.r > 1.0)] = 0.0 + + s_phi = numpy.sin(self.phi) + s_phi[(self.r > 1.0)] = 0.0 + + self.px_ex = c_theta*c_phi*c_phi + s_phi*s_phi + self.px_ey = (1.0 - c_theta)*s_phi*c_phi + self.py_ex = (c_theta - 1.0)*s_phi*c_phi + self.py_ey = c_theta*s_phi*s_phi + c_phi*c_phi + self.pz_ex = s_theta*c_phi + self.pz_ey = s_theta*s_phi + + def pfToPSF(self, pf, z_vals, scaling_factor = None): + """ + This matches the super-class function of the same name, but there is no + option of returning the (complex) real space PF. Also the PSF that is + returned is that for a freely rotating dipole. + + pf - A pupil function. + z_vals - The z values (focal planes) of the desired PSF. + scaling_factor - (Optional) The OTF rescaling factor, default is None. + + return - The PSF that corresponds to pf at the requested z_vals. + """ + rs_pf = self.pfToRS(pf, z_vals) + psf = self.rsToPSF(rs_pf) + + if scaling_factor is not None: + for i in range(psf.shape[0]): + otf = scipy.fftpack.fftshift(scipy.fftpack.fft2(psf[i,:,:])) + otf_scaled = otf * scaling_factor + psf[i,:,:] = numpy.abs(scipy.fftpack.ifft2(otf_scaled)) + + return psf + + def pfToRS(self, pf, z_vals): + """ + pf - A pupil function. + z_vals - The z values (focal planes) of the desired PSF. + + return - The real space counterpart of the PF at z_vals as a list. + [[rs_px_ex, rs_px_ey], [rs_py_ex, rs_py_ey], [rs_pz_ex, rs_pz_ey]] + """ + rs_px_ex = numpy.zeros((len(z_vals), pf.shape[0], pf.shape[1]), + dtype = numpy.complex128) + rs_px_ey = numpy.zeros((len(z_vals), pf.shape[0], pf.shape[1]), + dtype = numpy.complex128) + + rs_py_ex = numpy.zeros((len(z_vals), pf.shape[0], pf.shape[1]), + dtype = numpy.complex128) + rs_py_ey = numpy.zeros((len(z_vals), pf.shape[0], pf.shape[1]), + dtype = numpy.complex128) + + rs_pz_ex = numpy.zeros((len(z_vals), pf.shape[0], pf.shape[1]), + dtype = numpy.complex128) + rs_pz_ey = numpy.zeros((len(z_vals), pf.shape[0], pf.shape[1]), + dtype = numpy.complex128) + + for i, z in enumerate(z_vals): + pf_at_z = self.changeFocus(pf, z) + + rs_px_ex[i,:,:] = toRealSpace(self.px_ex * pf_at_z) + rs_px_ey[i,:,:] = toRealSpace(self.px_ey * pf_at_z) + + rs_py_ex[i,:,:] = toRealSpace(self.py_ex * pf_at_z) + rs_py_ey[i,:,:] = toRealSpace(self.py_ey * pf_at_z) + + rs_pz_ex[i,:,:] = toRealSpace(self.pz_ex * pf_at_z) + rs_pz_ey[i,:,:] = toRealSpace(self.pz_ey * pf_at_z) + + return [[rs_px_ex, rs_px_ey], [rs_py_ex, rs_py_ey], [rs_pz_ex, rs_pz_ey]] + + def rsToPSF(self, rs_pf, ratios = [1.0/3.0, 1.0/3.0, 1.0/3.0]): + """ + rs_pf - List of real space PFs from the pfToRS() method. + ratios - Ratios of dipole contribution, should sum to 1.0. + """ + psf = numpy.zeros((rs_pf[0][0].shape[0], self.size, self.size)) + for j in range(3): + psf += ratios[j] * intensity(rs_pf[j][0]) + psf += ratios[j] * intensity(rs_pf[j][1]) + + return psf + + +class GeometryC(Geometry): + """ + This class uses some of the C libraries in pupilfn to do the heavy lifting. It assumes + that the OTF scaling array is symmetric in X/Y. + + Based on profiling this is 3-4x faster than the pure Python version. + """ + def __init__(self, size, pixel_size, wavelength, imm_index, NA): + """ + size - The number of pixels in the PSF image, assumed square. + pixel_size - The size of the camera pixel in um. + wavelength - The wavelength of the flourescence in um. + imm_index - The index of the immersion media. + NA - The numerical aperature of the objective. + """ + super(GeometryC, self).__init__(size, pixel_size, wavelength, imm_index, NA) + + self.otf_sc = otfSC.OTFScaler(size = size) + self.pf_c = pfFnC.PupilFunction(geometry = self) + + def __del__(self): + self.otf_sc.cleanup() + self.pf_c.cleanup() + + def pfToPSF(self, pf, z_vals, want_intensity = True, scaling_factor = None): + """ + pf - A pupil function. + z_vals - The z values (focal planes) of the desired PSF. + want_intensity - (Optional) Return intensity, default is True. + scaling_factor - (Optional) The OTF rescaling factor, default is None. + + return - The PSF that corresponds to pf at the requested z_vals. + """ + self.pf_c.setPF(pf) + + if want_intensity: + if scaling_factor is not None: + self.otf_sc.setScale(scaling_factor) + + psf = numpy.zeros((len(z_vals), pf.shape[0], pf.shape[1])) + for i, z in enumerate(z_vals): + self.pf_c.translateZ(z) + temp = self.pf_c.getPSFIntensity() + if scaling_factor is not None: + psf[i,:,:] = self.otf_sc.scale(temp) + else: + psf[i,:,:] = temp + + else: + if scaling_factor is not None: + raise PupilMathException("OTF scaling of a complex valued PSF is not supported!") + + psf = numpy.zeros((len(z_vals), pf.shape[0], pf.shape[1]), dtype = numpy.complex128) + for i, z in enumerate(z_vals): + self.pf_c.translateZ(z) + psf[i,:,:] = self.pf_c.getPSF() + return psf + + return psf + + +class GeometryCVectorial(GeometryVectorial): + """ + The vectorial PF version of GeometryC. + """ + def __init__(self, size, pixel_size, wavelength, imm_index, NA): + """ + size - The number of pixels in the PSF image, assumed square. + pixel_size - The size of the camera pixel in um. + wavelength - The wavelength of the flourescence in um. + imm_index - The index of the immersion media. + NA - The numerical aperature of the objective. + """ + super(GeometryCVectorial, self).__init__(size, pixel_size, wavelength, imm_index, NA) + + self.otf_sc = otfSC.OTFScaler(size = size) + self.pf_c = pfFnC.PupilFunction(geometry = self) + + def __del__(self): + self.otf_sc.cleanup() + self.pf_c.cleanup() + + def pfToPSF(self, pf, z_vals, scaling_factor = None): + """ + This matches the super-class function of the same name, but there is no + option of returning the (complex) real space PF. Also the PSF that is + returned is that for a freely rotating dipole. + + pf - A pupil function. + z_vals - The z values (focal planes) of the desired PSF. + scaling_factor - (Optional) The OTF rescaling factor, default is None. + + return - The PSF that corresponds to pf at the requested z_vals. + """ + rs_pf = self.pfToRS(pf, z_vals) + psf = self.rsToPSF(rs_pf) + + if scaling_factor is not None: + self.otf_sc.setScale(scaling_factor) + for i in range(psf.shape[0]): + psf[i,:,:] = self.otf_sc.scale(psf[i,:,:]) + + return psf + + def pfToRS(self, pf, z_vals): + """ + pf - A pupil function. + z_vals - The z values (focal planes) of the desired PSF. + + return - The real space counterpart of the PF at z_vals as a list. + [[rs_px_ex, rs_px_ey], [rs_py_ex, rs_py_ey], [rs_pz_ex, rs_pz_ey]] + """ + rs_px_ex = numpy.zeros((len(z_vals), pf.shape[0], pf.shape[1]), + dtype = numpy.complex128) + rs_px_ey = numpy.zeros((len(z_vals), pf.shape[0], pf.shape[1]), + dtype = numpy.complex128) + + rs_py_ex = numpy.zeros((len(z_vals), pf.shape[0], pf.shape[1]), + dtype = numpy.complex128) + rs_py_ey = numpy.zeros((len(z_vals), pf.shape[0], pf.shape[1]), + dtype = numpy.complex128) + + rs_pz_ex = numpy.zeros((len(z_vals), pf.shape[0], pf.shape[1]), + dtype = numpy.complex128) + rs_pz_ey = numpy.zeros((len(z_vals), pf.shape[0], pf.shape[1]), + dtype = numpy.complex128) + + self.pf_c.setPF(pf) + for i, z in enumerate(z_vals): + self.pf_c.translateZ(z) + + rs_px_ex[i,:,:] = self.pf_c.getPXEX() + rs_px_ey[i,:,:] = self.pf_c.getPXEY() + + rs_py_ex[i,:,:] = self.pf_c.getPYEX() + rs_py_ey[i,:,:] = self.pf_c.getPYEY() + + rs_pz_ex[i,:,:] = self.pf_c.getPZEX() + rs_pz_ey[i,:,:] = self.pf_c.getPZEY() + + return [[rs_px_ex, rs_px_ey], [rs_py_ex, rs_py_ey], [rs_pz_ex, rs_pz_ey]] + + def rsToPSF(self, rs_pf, ratios = [1.0/3.0, 1.0/3.0, 1.0/3.0]): + """ + rs_pf - List of real space PFs from the pfToRS() method. + ratios - Ratios of dipole contribution, should sum to 1.0. + """ + psf = numpy.zeros((rs_pf[0][0].shape[0], self.size, self.size)) + for j in range(3): + psf += ratios[j] * intensity(rs_pf[j][0]) + psf += ratios[j] * intensity(rs_pf[j][1]) + + return psf + + +class GeometrySim(Geometry): + """ + This class is used in the simulations. It divides the pixel size + by two so that simulations look more realistic without having + to add the overhead of OTF scaling. + """ + def __init__(self, size, pixel_size, wavelength, imm_index, NA): + """ + size - The number of pixels in the PSF image, assumed square. + pixel_size - The size of the camera pixel in um. + wavelength - The wavelength of the flourescence in um. + imm_index - The index of the immersion media. + NA - The numerical aperature of the objective. + """ + super(GeometrySim, self).__init__(size, 0.5*pixel_size, wavelength, imm_index, NA) + + +def intensity(x): + """ + x - The (numpy array) to convert to intensity. + + return - The product of x and the complex conjugate of x. + """ + return numpy.abs(x * numpy.conj(x)) + + +def toRealSpace(pupil_fn): + """ + pupil_fn - A pupil function. + + return - The pupil function in real space (as opposed to fourier space). + """ + return scipy.fftpack.ifftshift(math.sqrt(pupil_fn.size) * scipy.fftpack.ifft2(pupil_fn)) + + +if (__name__ == "__main__"): + + import pickle + import sys + + if (len(sys.argv) < 2): + print("usage: ") + exit() + + pixel_size = 0.10 + #pixel_size = 0.020 + wavelength = 0.6 + refractive_index = 1.5 + numerical_aperture = 1.4 + z_range = 1.0 + z_pixel_size = 0.010 + + geo = Geometry(int(20.0/pixel_size), + pixel_size, + wavelength, + refractive_index, + numerical_aperture) + + if (len(sys.argv) == 4): + zmn = [] + amp = float(sys.argv[3]) + with open(sys.argv[2]) as fp: + for line in fp: + data = line.strip().split(" ") + if (len(data) == 3): + zmn.append([amp * float(data[2]), int(data[0]), int(data[1])]) + else: + zmn = [[1.3, 2, 2]] + #zmn = [[1, 0, 4]] + #zmn = [] + + pf = geo.createFromZernike(1.0, zmn) + z_values = numpy.arange(-z_range, z_range + 0.5 * z_pixel_size, z_pixel_size) + psfs = geo.pfToPSF(pf, z_values) + + #xy_size = 2.0*psfs.shape[0] + xy_size = 100 + xy_start = int(0.5 * (psfs.shape[1] - xy_size) + 1) + xy_end = int(xy_start + xy_size) + psfs = psfs[:,xy_start:xy_end,xy_start:xy_end] + + if True: + tifffile.imwrite("kz.tif", numpy.real(geo.kz).astype(numpy.float32)) + + if False: + tifffile.imwrite("pf_abs.tif", numpy.abs(pf).astype(numpy.float32)) + tifffile.imwrite("pf_angle.tif", (180.0 * numpy.angle(pf)/numpy.pi + 180).astype(numpy.float32)) + + if False: + with tifffile.TiffWriter(sys.argv[1]) as psf_tif: + temp = (psfs/numpy.max(psfs)).astype(numpy.float32) + psf_tif.write(temp) + + if False: + with open("z_offset.txt", "w") as fp: + for i in range(z_values.size): + fp.write("1 {0:.6f}\n".format(1000.0 * z_values[i])) + + if False: + psfs = (65000.0 * (psfs/numpy.max(psfs))).astype(numpy.uint16) + psf_dict = {"pixel_size" : pixel_size, + "wavelength" : wavelength, + "refractive_index" : refractive_index, + "numerical_aperture" : numerical_aperture, + "z_range" : z_range, + "zmm" : zmn, + "psf" : psfs} + pickle.dump(psf_dict, open(sys.argv[1], "wb"), protocol = 2) + + +# +# The MIT License +# +# Copyright (c) 2016 Zhuang Lab, Harvard University +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. +# + + diff --git a/storm_analysis/simulator/simbase.py b/storm_analysis/simulator/simbase.py index 64b6cd3c..9fa3774e 100644 --- a/storm_analysis/simulator/simbase.py +++ b/storm_analysis/simulator/simbase.py @@ -1,61 +1,61 @@ -#!/usr/bin/env python -""" -Base class for the various simulation classes. - -Hazen 12/16 -""" - -import json - - -class SimException(Exception): - pass - -class SimBase(object): - """ - Base class for all the simulation classes. - """ - def __init__(self, sim_fp, x_size, y_size, h5_data): - self.h5_data = h5_data - self.sim_fp = sim_fp - self.x_size = x_size - self.y_size = y_size - - def checkSize(self, x_size, y_size): - if (self.x_size != x_size) or (self.y_size != y_size): - raise SimException("Size mismatch (X) " + str(self.x_size) + " != " + x_size + " or (Y) " + self.y_size + " != " + y_size) - - def cleanup(self): - pass - - def saveJSON(self, elt): - self.sim_fp.write(json.dumps(elt) + "\n") - - def setSize(self, x_size, y_size): - self.x_size = x_size - self.y_size = y_size - - -# -# The MIT License -# -# Copyright (c) 2016 Zhuang Lab, Harvard University -# -# Permission is hereby granted, free of charge, to any person obtaining a copy -# of this software and associated documentation files (the "Software"), to deal -# in the Software without restriction, including without limitation the rights -# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -# copies of the Software, and to permit persons to whom the Software is -# furnished to do so, subject to the following conditions: -# -# The above copyright notice and this permission notice shall be included in -# all copies or substantial portions of the Software. -# -# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN -# THE SOFTWARE. -# +#!/usr/bin/env python +""" +Base class for the various simulation classes. + +Hazen 12/16 +""" + +import json + + +class SimException(Exception): + pass + +class SimBase(object): + """ + Base class for all the simulation classes. + """ + def __init__(self, sim_fp, x_size, y_size, h5_data): + self.h5_data = h5_data + self.sim_fp = sim_fp + self.x_size = x_size + self.y_size = y_size + + def checkSize(self, x_size, y_size): + if (self.x_size != x_size) or (self.y_size != y_size): + raise SimException("Size mismatch (X) " + str(self.x_size) + " != " + x_size + " or (Y) " + self.y_size + " != " + y_size) + + def cleanup(self): + pass + + def saveJSON(self, elt): + self.sim_fp.write(json.dumps(elt) + "\n") + + def setSize(self, x_size, y_size): + self.x_size = x_size + self.y_size = y_size + + +# +# The MIT License +# +# Copyright (c) 2016 Zhuang Lab, Harvard University +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. +# diff --git a/storm_analysis/simulator/simulate.py b/storm_analysis/simulator/simulate.py index 2b9fd034..62281b0c 100644 --- a/storm_analysis/simulator/simulate.py +++ b/storm_analysis/simulator/simulate.py @@ -1,242 +1,242 @@ -#!/usr/bin/env python -""" -Generate simulated. The basic idea is that you provide a list of -localizations in Insight3 .bin format and these are used to -generate a series of images using the following steps: - -Initialization: - 1. locs = h5.getLocalizations() - 2. bg = background.Background() - 3. camera = camera.Camera() - 4. drift = drift.Drift() - 5. pp = photophysics.Photophysics() - 6. psf = psf.PSF() - - -Generation: - 1. image = numpy.zeros() - 2. image += bg.getBackground() - 3. cur_locs = pp.getEmitters() - 4. drift.drift(cur_locs, frame_number) - 5. image += psf.getPSFs(cur_locs) - 6. image = camera.readImage(image) - 7. saveimage() - 8. savelocs() - -Note: This is expected to set the 'height','sum' and 'background' fields - in the output to the correct values. The values 'x', 'y', 'z', 'xsigma' - and 'ysigma' are just passed through. - -Hazen 01/18 -""" - -import json -import numpy - -import storm_analysis.sa_library.datawriter as datawriter -import storm_analysis.sa_library.sa_h5py as saH5Py - - -class Simulate(object): - - def __init__(self, - background_factory = None, - camera_factory = None, - drift_factory = None, - photophysics_factory = None, - psf_factory = None, - dither = False, - x_size = 256, - y_size = 256, - **kwds): - """ - The factory variables should be functions that return the correct class - to run a simulation with the following signature: - - factory_fn(sim_settings, x_size, y_size, h5_data_in) - - """ - super(Simulate, self).__init__(**kwds) - - self.bg_factory = background_factory - self.cam_factory = camera_factory - self.drift_factory = drift_factory - self.pphys_factory = photophysics_factory - self.psf_factory = psf_factory - - self.dither = dither - self.x_size = x_size - self.y_size = y_size - - def setBackgroundFactory(self, new_factory): - self.bg_factory = new_factory - - def setCameraFactory(self, new_factory): - self.cam_factory = new_factory - - def setDriftFactory(self, new_drift_factory): - self.drift_factory = new_drift_factory - - def setPhotoPhysicsFactory(self, new_factory): - self.pphys_factory = new_factory - - def setPSFFactory(self, new_factory): - self.psf_factory = new_factory - - def simulate(self, dax_file, bin_file, n_frames, verbosity = 1): - - assert (verbosity > 0), "Verbosity must be greater than 1." - - # - # Initialization. - # - movie_data = datawriter.inferWriter(dax_file, - width = self.x_size, - height = self.y_size) - with saH5Py.SAH5Py(bin_file) as h5: - h5_data_in = h5.getLocalizations() - - if not h5_data_in: - print("Warning! No localizations input HDF5 file!") - - out_fname_base = dax_file[:-4] - h5_data_out = saH5Py.SAH5Py(filename = out_fname_base + "_ref.hdf5", - is_existing = False, - overwrite = True) - h5_data_out.setMovieInformation(self.x_size, self.y_size, n_frames, "") - - sim_settings = open(out_fname_base + "_sim_params.txt", "w") - - # - # Create the user-specified class instances that will do - # most of the actual work of the simulation. - # - bg = self.bg_factory(sim_settings, self.x_size, self.y_size, h5_data_in) - cam = self.cam_factory(sim_settings, self.x_size, self.y_size, h5_data_in) - drift = None - if self.drift_factory is not None: - drift = self.drift_factory(sim_settings, self.x_size, self.y_size, h5_data_in) - - if h5_data_in: - pp = self.pphys_factory(sim_settings, self.x_size, self.y_size, h5_data_in) - - psf = self.psf_factory(sim_settings, self.x_size, self.y_size, h5_data_in) - - sim_settings.write(json.dumps({"simulation" : {"bin_file" : bin_file, - "x_size" : str(self.x_size), - "y_size" : str(self.y_size)}}) + "\n") - - h5_data_out.setPixelSize(psf.nm_per_pixel) - - # - # Generate the simulated movie. - # - for i in range(n_frames): - - # Generate the new image. - image = numpy.zeros((self.x_size, self.y_size)) - - if h5_data_in: - # Get the emitters that are on in the current frame. - cur_h5 = pp.getEmitters(i) - - if ((i % verbosity) == 0): - print("Frame", i, cur_h5['x'].size, "emitters") - - # Dither points x,y values if requested. This is useful for things - # like looking for pixel level biases in simulated data with gridded - # localizations. - # - if self.dither: - cur_h5['x'] += numpy.random.uniform(size = cur_h5['x'].size) - 0.5 - cur_h5['y'] += numpy.random.uniform(size = cur_h5['y'].size) - 0.5 - - # Add background to image. - image += bg.getBackground(i) - - if h5_data_in: - # Set 'bg' parameter of the emitters. - cur_h5 = bg.getEmitterBackground(cur_h5) - - # Apply drift to the localizations. - if drift is not None: - drift.drift(i, cur_h5) - - # Foreground - image += psf.getPSFs(cur_h5) - - # Camera - image = cam.readImage(image) - - # Save the image. - movie_data.addFrame(numpy.transpose(image)) - - if h5_data_in: - # Save the molecule locations. - h5_data_out.addLocalizations(cur_h5, i) - - # Clean up user-specified class instances. - for ci in [bg, cam, drift, pp, psf]: - if ci is not None: - ci.cleanup() - - movie_data.close() - h5_data_out.close() - sim_settings.close() - - -# The simplest simulation. - -if (__name__ == "__main__"): - - import argparse - - import storm_analysis.simulator.background as background - import storm_analysis.simulator.camera as camera - import storm_analysis.simulator.photophysics as photophysics - import storm_analysis.simulator.psf as psf - - parser = argparse.ArgumentParser(description = "A simple simulation example.") - - parser.add_argument('--dax', dest='dax_file', type=str, required=True, - help = "The name of the dax file to save the simulated STORM movie.") - parser.add_argument('--bin', dest='hdf5', type=str, required=True, - help = "The name of the HDF5 file containing the emitter locations.") - parser.add_argument('--frames', dest='frames', type=int, required=True, - help = "The length of the movie in frames.") - parser.add_argument('--photons', dest='photons', type=float, required=True, - help = "The integral of a single emitter in photons.") - - args = parser.parse_args() - - sim = Simulate(lambda settings, xs, ys, h5data : background.UniformBackground(settings, xs, ys, h5data), - lambda settings, xs, ys, h5data : camera.Ideal(settings, xs, ys, h5data, 100.0), - lambda settings, xs, ys, h5data : photophysics.AlwaysOn(settings, xs, ys, h5data, args.photons), - lambda settings, xs, ys, h5data : psf.GaussianPSF(settings, xs, ys, h5data, 160.0)) - - sim.simulate(args.dax_file, args.hdf5, args.frames) - - -# -# The MIT License -# -# Copyright (c) 2016 Zhuang Lab, Harvard University -# -# Permission is hereby granted, free of charge, to any person obtaining a copy -# of this software and associated documentation files (the "Software"), to deal -# in the Software without restriction, including without limitation the rights -# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -# copies of the Software, and to permit persons to whom the Software is -# furnished to do so, subject to the following conditions: -# -# The above copyright notice and this permission notice shall be included in -# all copies or substantial portions of the Software. -# -# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN -# THE SOFTWARE. -# +#!/usr/bin/env python +""" +Generate simulated. The basic idea is that you provide a list of +localizations in Insight3 .bin format and these are used to +generate a series of images using the following steps: + +Initialization: + 1. locs = h5.getLocalizations() + 2. bg = background.Background() + 3. camera = camera.Camera() + 4. drift = drift.Drift() + 5. pp = photophysics.Photophysics() + 6. psf = psf.PSF() + + +Generation: + 1. image = numpy.zeros() + 2. image += bg.getBackground() + 3. cur_locs = pp.getEmitters() + 4. drift.drift(cur_locs, frame_number) + 5. image += psf.getPSFs(cur_locs) + 6. image = camera.readImage(image) + 7. saveimage() + 8. savelocs() + +Note: This is expected to set the 'height','sum' and 'background' fields + in the output to the correct values. The values 'x', 'y', 'z', 'xsigma' + and 'ysigma' are just passed through. + +Hazen 01/18 +""" + +import json +import numpy + +import storm_analysis.sa_library.datawriter as datawriter +import storm_analysis.sa_library.sa_h5py as saH5Py + + +class Simulate(object): + + def __init__(self, + background_factory = None, + camera_factory = None, + drift_factory = None, + photophysics_factory = None, + psf_factory = None, + dither = False, + x_size = 256, + y_size = 256, + **kwds): + """ + The factory variables should be functions that return the correct class + to run a simulation with the following signature: + + factory_fn(sim_settings, x_size, y_size, h5_data_in) + + """ + super(Simulate, self).__init__(**kwds) + + self.bg_factory = background_factory + self.cam_factory = camera_factory + self.drift_factory = drift_factory + self.pphys_factory = photophysics_factory + self.psf_factory = psf_factory + + self.dither = dither + self.x_size = x_size + self.y_size = y_size + + def setBackgroundFactory(self, new_factory): + self.bg_factory = new_factory + + def setCameraFactory(self, new_factory): + self.cam_factory = new_factory + + def setDriftFactory(self, new_drift_factory): + self.drift_factory = new_drift_factory + + def setPhotoPhysicsFactory(self, new_factory): + self.pphys_factory = new_factory + + def setPSFFactory(self, new_factory): + self.psf_factory = new_factory + + def simulate(self, dax_file, bin_file, n_frames, verbosity = 1): + + assert (verbosity > 0), "Verbosity must be greater than 1." + + # + # Initialization. + # + movie_data = datawriter.inferWriter(dax_file, + width = self.x_size, + height = self.y_size) + with saH5Py.SAH5Py(bin_file) as h5: + h5_data_in = h5.getLocalizations() + + if not h5_data_in: + print("Warning! No localizations input HDF5 file!") + + out_fname_base = dax_file[:-4] + h5_data_out = saH5Py.SAH5Py(filename = out_fname_base + "_ref.hdf5", + is_existing = False, + overwrite = True) + h5_data_out.setMovieInformation(self.x_size, self.y_size, n_frames, "") + + sim_settings = open(out_fname_base + "_sim_params.txt", "w") + + # + # Create the user-specified class instances that will do + # most of the actual work of the simulation. + # + bg = self.bg_factory(sim_settings, self.x_size, self.y_size, h5_data_in) + cam = self.cam_factory(sim_settings, self.x_size, self.y_size, h5_data_in) + drift = None + if self.drift_factory is not None: + drift = self.drift_factory(sim_settings, self.x_size, self.y_size, h5_data_in) + + if h5_data_in: + pp = self.pphys_factory(sim_settings, self.x_size, self.y_size, h5_data_in) + + psf = self.psf_factory(sim_settings, self.x_size, self.y_size, h5_data_in) + + sim_settings.write(json.dumps({"simulation" : {"bin_file" : bin_file, + "x_size" : str(self.x_size), + "y_size" : str(self.y_size)}}) + "\n") + + h5_data_out.setPixelSize(psf.nm_per_pixel) + + # + # Generate the simulated movie. + # + for i in range(n_frames): + + # Generate the new image. + image = numpy.zeros((self.x_size, self.y_size)) + + if h5_data_in: + # Get the emitters that are on in the current frame. + cur_h5 = pp.getEmitters(i) + + if ((i % verbosity) == 0): + print("Frame", i, cur_h5['x'].size, "emitters") + + # Dither points x,y values if requested. This is useful for things + # like looking for pixel level biases in simulated data with gridded + # localizations. + # + if self.dither: + cur_h5['x'] += numpy.random.uniform(size = cur_h5['x'].size) - 0.5 + cur_h5['y'] += numpy.random.uniform(size = cur_h5['y'].size) - 0.5 + + # Add background to image. + image += bg.getBackground(i) + + if h5_data_in: + # Set 'bg' parameter of the emitters. + cur_h5 = bg.getEmitterBackground(cur_h5) + + # Apply drift to the localizations. + if drift is not None: + drift.drift(i, cur_h5) + + # Foreground + image += psf.getPSFs(cur_h5) + + # Camera + image = cam.readImage(image) + + # Save the image. + movie_data.addFrame(numpy.transpose(image)) + + if h5_data_in: + # Save the molecule locations. + h5_data_out.addLocalizations(cur_h5, i) + + # Clean up user-specified class instances. + for ci in [bg, cam, drift, pp, psf]: + if ci is not None: + ci.cleanup() + + movie_data.close() + h5_data_out.close() + sim_settings.close() + + +# The simplest simulation. + +if (__name__ == "__main__"): + + import argparse + + import storm_analysis.simulator.background as background + import storm_analysis.simulator.camera as camera + import storm_analysis.simulator.photophysics as photophysics + import storm_analysis.simulator.psf as psf + + parser = argparse.ArgumentParser(description = "A simple simulation example.") + + parser.add_argument('--dax', dest='dax_file', type=str, required=True, + help = "The name of the dax file to save the simulated STORM movie.") + parser.add_argument('--bin', dest='hdf5', type=str, required=True, + help = "The name of the HDF5 file containing the emitter locations.") + parser.add_argument('--frames', dest='frames', type=int, required=True, + help = "The length of the movie in frames.") + parser.add_argument('--photons', dest='photons', type=float, required=True, + help = "The integral of a single emitter in photons.") + + args = parser.parse_args() + + sim = Simulate(lambda settings, xs, ys, h5data : background.UniformBackground(settings, xs, ys, h5data), + lambda settings, xs, ys, h5data : camera.Ideal(settings, xs, ys, h5data, 100.0), + lambda settings, xs, ys, h5data : photophysics.AlwaysOn(settings, xs, ys, h5data, args.photons), + lambda settings, xs, ys, h5data : psf.GaussianPSF(settings, xs, ys, h5data, 160.0)) + + sim.simulate(args.dax_file, args.hdf5, args.frames) + + +# +# The MIT License +# +# Copyright (c) 2016 Zhuang Lab, Harvard University +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. +# diff --git a/storm_analysis/voronoi/voronoi_analysis.py b/storm_analysis/voronoi/voronoi_analysis.py index e1ea1aa1..1aef82bb 100644 --- a/storm_analysis/voronoi/voronoi_analysis.py +++ b/storm_analysis/voronoi/voronoi_analysis.py @@ -1,189 +1,189 @@ -#!/usr/bin/env python -""" -Does 2D Voroni cluster analysis. - -Hazen 08/18 -""" - -import numpy -from scipy.spatial import Voronoi, voronoi_plot_2d -from shapely.geometry import Polygon - -import storm_analysis.dbscan.clusters_sa_h5py as clSAH5Py -import storm_analysis.dbscan.dbscan_analysis as dbscanAnalysis - - -def findClusters(h5_name, density_factor, min_size, verbose = True): - """ - h5_name - The localizations HDF5 file. - density_factor - Multiple of the median density to be a cluster member. - min_size - The minimum number of localizations a cluster can have. - """ - - with clSAH5Py.SAH5Clusters(h5_name) as cl_h5: - [x, y, z, c, cl_dict] = cl_h5.getDataForClustering() - - n_locs = x.size - labels = numpy.zeros(n_locs, dtype = numpy.int32) - 1 - density = numpy.zeros(n_locs) - - # Convert data to nanometers - pix_to_nm = cl_h5.getPixelSize() - x_nm = x * pix_to_nm - y_nm = y * pix_to_nm - points = numpy.column_stack((x_nm, y_nm)) - - if verbose: - print("Creating Voronoi object.") - vor = Voronoi(points) - - if verbose: - print("Calculating 2D region sizes.") - for i, region_index in enumerate(vor.point_region): - if ((i%10000) == 0) and verbose: - print("Processing point", i) - - vertices = [] - for vertex in vor.regions[region_index]: - - # I think these are edge regions? - if (vertex == -1): - vertices = [] - break - - vertices.append(vor.vertices[vertex]) - - if (len(vertices) > 0): - area = Polygon(vertices).area - density[i] = 1.0/area - - # Used median density based threshold. - ave_density = numpy.median(density) - if verbose: - print("Min density", numpy.amin(density)) - print("Max density", numpy.amax(density)) - print("Median density", ave_density) - - # Record the neighbors of each point. These are polygons so there shouldn't - # be that many neighbors (sides). 40 is more than safe? - # - max_neighbors = 40 - neighbors = numpy.zeros((n_locs, max_neighbors), dtype = numpy.int32) - 1 - neighbors_counts = numpy.zeros((n_locs), dtype = numpy.int32) - - if verbose: - print("Calculating neighbors") - for ridge_p in vor.ridge_points: - - p1 = ridge_p[0] - p2 = ridge_p[1] - - # Add p2 to the list for p1 - neighbors[p1,neighbors_counts[p1]] = p2 - neighbors_counts[p1] += 1 - - # Add p1 to the list for p2 - neighbors[p2,neighbors_counts[p2]] = p1 - neighbors_counts[p2] += 1 - - if False: - n1 = neighbors[0,:] - print(n1) - print(neighbors[n1[0],:]) - - # Mark connected points that meet the minimum density criteria. - if verbose: - print("Marking connected regions") - min_density = density_factor * ave_density - visited = numpy.zeros(n_locs, dtype = numpy.int32) - - def neighborsList(index): - nlist = [] - for i in range(neighbors_counts[index]): - loc_index = neighbors[index,i] - if (visited[loc_index] == 0): - nlist.append(neighbors[index,i]) - visited[loc_index] = 1 - return nlist - - cluster_id = 0 - for i in range(n_locs): - if (visited[i] == 0): - visited[i] = 1 - if (density[i] > min_density): - cluster_elt = [i] - c_size = 1 - to_check = neighborsList(i) - while (len(to_check) > 0): - - # Remove last localization from the list. - loc_index = to_check[-1] - to_check = to_check[:-1] - - # If the localization has sufficient density add to cluster and - # check neighbors. - if (density[loc_index] > min_density): - to_check += neighborsList(loc_index) - cluster_elt.append(loc_index) - c_size += 1 - - # Mark as visited. - visited[loc_index] = 1 - - # Mark the cluster if there are enough localizations in the cluster. - if (c_size > min_size): - if verbose: - print("cluster", cluster_id, "size", c_size) - for elt in cluster_elt: - labels[elt] = cluster_id - cluster_id += 1 - - if verbose: - print(cluster_id, "clusters") - - # Save the clustering results. - cl_dict["x"] = x - cl_dict["y"] = y - cl_dict["z"] = z - cl_dict["density"] = density - cl_dict["category"] = c - cl_h5.addClusters(labels, cl_dict) - - # Save clustering info. - info = "voronoi,df,{0:0.3f},ms,{1:d}".format(density_factor,min_size) - cl_h5.setClusteringInfo(info) - - -def voronoiAnalysis(h5_name, density_factor, min_size = 30): - """ - Runs voronoi() and clusterStats() on a storm-analysis format HDF5 - format file. - - h5_name - The name of the HDF5 file. - density_factor - Voronoi clustering density factor. - min_size - Minimum size cluster for cluster statistics. - - Note: This ignores z values and category. - """ - findClusters(h5_name, density_factor, min_size) - dbscanAnalysis.clusterStats(h5_name, min_size) - - -if (__name__ == "__main__"): - - import argparse - - parser = argparse.ArgumentParser(description = 'Voronoi based clustering following Levet, Nature Methods, 2015') - - parser.add_argument('--bin', dest='hdf5', type=str, required=True, - help = "The name of the HDF5 localizations file.") - parser.add_argument('--density', dest='density', type=float, required=True, - help = "The density multiplier to be in a cluster. The median polygon size is multiplied by this value to give a threshold for polygon size to be in a cluster.") - parser.add_argument('--min_size', dest='min_size', type=int, required=False, default=30, - help = "The minimum cluster size to include when calculating cluster statistics.") - - args = parser.parse_args() - - voronoiAnalysis(args.hdf5, args.density, args.min_size) - - +#!/usr/bin/env python +""" +Does 2D Voroni cluster analysis. + +Hazen 08/18 +""" + +import numpy +from scipy.spatial import Voronoi, voronoi_plot_2d +from shapely.geometry import Polygon + +import storm_analysis.dbscan.clusters_sa_h5py as clSAH5Py +import storm_analysis.dbscan.dbscan_analysis as dbscanAnalysis + + +def findClusters(h5_name, density_factor, min_size, verbose = True): + """ + h5_name - The localizations HDF5 file. + density_factor - Multiple of the median density to be a cluster member. + min_size - The minimum number of localizations a cluster can have. + """ + + with clSAH5Py.SAH5Clusters(h5_name) as cl_h5: + [x, y, z, c, cl_dict] = cl_h5.getDataForClustering() + + n_locs = x.size + labels = numpy.zeros(n_locs, dtype = numpy.int32) - 1 + density = numpy.zeros(n_locs) + + # Convert data to nanometers + pix_to_nm = cl_h5.getPixelSize() + x_nm = x * pix_to_nm + y_nm = y * pix_to_nm + points = numpy.column_stack((x_nm, y_nm)) + + if verbose: + print("Creating Voronoi object.") + vor = Voronoi(points) + + if verbose: + print("Calculating 2D region sizes.") + for i, region_index in enumerate(vor.point_region): + if ((i%10000) == 0) and verbose: + print("Processing point", i) + + vertices = [] + for vertex in vor.regions[region_index]: + + # I think these are edge regions? + if (vertex == -1): + vertices = [] + break + + vertices.append(vor.vertices[vertex]) + + if (len(vertices) > 0): + area = Polygon(vertices).area + density[i] = 1.0/area + + # Used median density based threshold. + ave_density = numpy.median(density) + if verbose: + print("Min density", numpy.amin(density)) + print("Max density", numpy.amax(density)) + print("Median density", ave_density) + + # Record the neighbors of each point. These are polygons so there shouldn't + # be that many neighbors (sides). 40 is more than safe? + # + max_neighbors = 40 + neighbors = numpy.zeros((n_locs, max_neighbors), dtype = numpy.int32) - 1 + neighbors_counts = numpy.zeros((n_locs), dtype = numpy.int32) + + if verbose: + print("Calculating neighbors") + for ridge_p in vor.ridge_points: + + p1 = ridge_p[0] + p2 = ridge_p[1] + + # Add p2 to the list for p1 + neighbors[p1,neighbors_counts[p1]] = p2 + neighbors_counts[p1] += 1 + + # Add p1 to the list for p2 + neighbors[p2,neighbors_counts[p2]] = p1 + neighbors_counts[p2] += 1 + + if False: + n1 = neighbors[0,:] + print(n1) + print(neighbors[n1[0],:]) + + # Mark connected points that meet the minimum density criteria. + if verbose: + print("Marking connected regions") + min_density = density_factor * ave_density + visited = numpy.zeros(n_locs, dtype = numpy.int32) + + def neighborsList(index): + nlist = [] + for i in range(neighbors_counts[index]): + loc_index = neighbors[index,i] + if (visited[loc_index] == 0): + nlist.append(neighbors[index,i]) + visited[loc_index] = 1 + return nlist + + cluster_id = 0 + for i in range(n_locs): + if (visited[i] == 0): + visited[i] = 1 + if (density[i] > min_density): + cluster_elt = [i] + c_size = 1 + to_check = neighborsList(i) + while (len(to_check) > 0): + + # Remove last localization from the list. + loc_index = to_check[-1] + to_check = to_check[:-1] + + # If the localization has sufficient density add to cluster and + # check neighbors. + if (density[loc_index] > min_density): + to_check += neighborsList(loc_index) + cluster_elt.append(loc_index) + c_size += 1 + + # Mark as visited. + visited[loc_index] = 1 + + # Mark the cluster if there are enough localizations in the cluster. + if (c_size > min_size): + if verbose: + print("cluster", cluster_id, "size", c_size) + for elt in cluster_elt: + labels[elt] = cluster_id + cluster_id += 1 + + if verbose: + print(cluster_id, "clusters") + + # Save the clustering results. + cl_dict["x"] = x + cl_dict["y"] = y + cl_dict["z"] = z + cl_dict["density"] = density + cl_dict["category"] = c + cl_h5.addClusters(labels, cl_dict) + + # Save clustering info. + info = "voronoi,df,{0:0.3f},ms,{1:d}".format(density_factor,min_size) + cl_h5.setClusteringInfo(info) + + +def voronoiAnalysis(h5_name, density_factor, min_size = 30): + """ + Runs voronoi() and clusterStats() on a storm-analysis format HDF5 + format file. + + h5_name - The name of the HDF5 file. + density_factor - Voronoi clustering density factor. + min_size - Minimum size cluster for cluster statistics. + + Note: This ignores z values and category. + """ + findClusters(h5_name, density_factor, min_size) + dbscanAnalysis.clusterStats(h5_name, min_size) + + +if (__name__ == "__main__"): + + import argparse + + parser = argparse.ArgumentParser(description = 'Voronoi based clustering following Levet, Nature Methods, 2015') + + parser.add_argument('--bin', dest='hdf5', type=str, required=True, + help = "The name of the HDF5 localizations file.") + parser.add_argument('--density', dest='density', type=float, required=True, + help = "The density multiplier to be in a cluster. The median polygon size is multiplied by this value to give a threshold for polygon size to be in a cluster.") + parser.add_argument('--min_size', dest='min_size', type=int, required=False, default=30, + help = "The minimum cluster size to include when calculating cluster statistics.") + + args = parser.parse_args() + + voronoiAnalysis(args.hdf5, args.density, args.min_size) + +