From 45ef3adc029f59d130e25bd1d8e54c06624ec4ca Mon Sep 17 00:00:00 2001 From: Hazen Babcock Date: Tue, 8 Sep 2026 17:26:35 -0400 Subject: [PATCH 1/5] Normalize line endings, and add a .gitattributes Twenty-eight files were committed with CRLF endings while the other 579 text files in the repository used LF. Seven are C, twenty-one are Python, and they are spread across ten directories, so there was no rule to infer about which files were which - admm_lasso.c is CRLF and admm_lasso.py is not. Nothing depended on the difference. It only mattered when something wrote to one of these files, because an editor or script that rewrites a file whole converts the endings as a side effect, and a two line change then arrives as a diff touching all 765 lines. That is the kind of diff nobody can review, and it happened repeatedly while working on pupil_math.py, matched_filter.c and admm_lasso_c.py. The conversion is CRLF to LF and nothing else. Stripping every CR from each file before and after gives identical bytes, and git diff --ignore-cr-at-eol is empty across the whole change, which is the same statement made by git rather than by hand. The insertion and deletion counts are equal at 7140 each, as they must be when only line endings move. The .gitattributes then keeps it from coming back. text=auto lets git detect binaries itself; the explicit binary list below it is a backstop for the formats actually in the tree. The checked in Windows DLLs under storm_analysis/c_libraries are the ones worth being careful about, since rewriting those would break the wheel the Windows CI job builds. Rebuilt the C libraries and ran the suite afterwards: 280 tests, no failures. --- .gitattributes | 31 + storm_analysis/admm/admm_decon.py | 108 +- storm_analysis/admm/admm_lasso.c | 1010 +++++------ storm_analysis/admm/admm_lasso_c.py | 378 ++-- storm_analysis/dbscan/batch_analysis.py | 132 +- storm_analysis/dbscan/cluster_images.py | 338 ++-- storm_analysis/dbscan/dbscan.c | 1164 ++++++------- storm_analysis/dbscan/dbscan_analysis.py | 332 ++-- storm_analysis/dbscan/dbscan_c.py | 228 +-- storm_analysis/densestorm/densestorm.c | 1116 ++++++------ storm_analysis/densestorm/densestorm_c.py | 394 ++--- storm_analysis/densestorm/densestorm_decon.py | 114 +- storm_analysis/fista/fista_fft.c | 1122 ++++++------ storm_analysis/psf_fft/psf_fft.c | 840 ++++----- storm_analysis/sa_library/drift_utilities.py | 386 ++--- storm_analysis/sa_library/matched_filter.c | 478 ++--- storm_analysis/sa_library/rebin.py | 156 +- storm_analysis/sa_library/recenter_psf.py | 84 +- .../sa_utilities/xyz_drift_correction.py | 532 +++--- storm_analysis/simulator/background.py | 222 +-- storm_analysis/simulator/camera.py | 416 ++--- storm_analysis/simulator/draw_gaussians.c | 166 +- storm_analysis/simulator/draw_gaussians_c.py | 326 ++-- storm_analysis/simulator/photophysics.py | 480 +++--- storm_analysis/simulator/psf.py | 1244 +++++++------- storm_analysis/simulator/pupil_math.py | 1530 ++++++++--------- storm_analysis/simulator/simbase.py | 122 +- storm_analysis/simulator/simulate.py | 484 +++--- storm_analysis/voronoi/voronoi_analysis.py | 378 ++-- 29 files changed, 7171 insertions(+), 7140 deletions(-) create mode 100644 .gitattributes 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/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) + + From 955ab22ec403cf817207316ec395108b231e8612 Mon Sep 17 00:00:00 2001 From: Hazen Babcock Date: Tue, 8 Sep 2026 17:32:38 -0400 Subject: [PATCH 2/5] Add an AGENTS.md There was nothing in the repository describing how to work on it. README.md covers installing and using the package, and the documentation covers the analysis, but the things you need in order to change the code safely were only discoverable by doing it wrong once. The file collects what actually cost time rather than what a template would suggest. The build section says to re-run scons after editing C, because the Python side loads whatever is already in c_libraries/ and a stale library produces a plausible wrong answer instead of an error. The test section says to trust the exit status, because several C libraries flush their stdout at interpreter exit and push pytest's summary line off the end of the output. It also records that pytest does not run the __main__ blocks, which is how a renamed test stayed green across six pull requests. The conventions section is short and only lists things measured across the tree: all 242 modules that import numpy write 'import numpy', and no file anywhere uses 'as np'; class methods are camelCase 704 times against 6 in snake_case. Keyword argument spacing is explicitly called out as *not* a convention - it is close to evenly split and 108 of 254 files use both forms, so an agent tidying it in one direction would be inventing a rule. The traps section names decon_storm/, which has C sources that SConstruct never builds and nothing imports. That directory is why the scan behind 76080b6a first reported the wrong getPeaks: it kept the first definition matching each name and found the dead one. Nothing here changes code. Every path and count in the file was checked against the tree rather than recalled, and the claim that the diagnostics other than multiplane take no arguments was wrong when first written - fista_decon, spliner and spliner_2d all accept --no-splines. --- AGENTS.md | 149 ++++++++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 149 insertions(+) create mode 100644 AGENTS.md diff --git a/AGENTS.md b/AGENTS.md new file mode 100644 index 00000000..d468fc88 --- /dev/null +++ b/AGENTS.md @@ -0,0 +1,149 @@ +# 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. + + +## 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. From c9ffe74a65f03e290a441a46ef01e0d4f178ed79 Mon Sep 17 00:00:00 2001 From: Hazen Babcock Date: Tue, 8 Sep 2026 17:33:00 -0400 Subject: [PATCH 3/5] Say which diagnostics take arguments e644347e added a note that multiplane/configure.py needs --psf-model and said "the rest take no arguments". That is not true: fista_decon, spliner, spliner_2d and multiplane itself all accept --no-splines. The distinction that matters is required versus optional, since only the missing --psf-model stops a script from running, so say it that way. --- storm_analysis/diagnostics/README.txt | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/storm_analysis/diagnostics/README.txt b/storm_analysis/diagnostics/README.txt index 7d318ef5..25b4aa1d 100644 --- a/storm_analysis/diagnostics/README.txt +++ b/storm_analysis/diagnostics/README.txt @@ -33,7 +33,11 @@ does nothing anywhere in the package. Figures that are written to disk, such as mapping images from the multicolor diagnostic, are still written. Note that multiplane/configure.py needs --psf-model, one of psf_fft, pupilfn or -spline. The rest take no arguments. +spline. It is required and has no default, so configure.py exits without it. + +That is the only required argument in any of these. fista_decon, spliner, +spliner_2d and multiplane also accept an optional --no-splines. The rest run +with no arguments. (Linux) C profiling tools: From b2d2ee412681b51cea9dfcb4cd4050cdb2e29278 Mon Sep 17 00:00:00 2001 From: Hazen Babcock Date: Tue, 8 Sep 2026 17:46:37 -0400 Subject: [PATCH 4/5] Give the PSF measurement input files a format The Spliner section said to create "a text file containing the z value of each frame in the z stack movie". measure_psf_beads.py does not accept that. It reads the file with numpy.loadtxt and passes the result to makeZIndexArray(), which asserts a shape of (N,2): a valid flag in the first column and the z offset in microns in the second, with frames whose flag is 0 dropped from the measurement. Following the instructions as written therefore ends at AssertionError: Z offsets must have shape (N,2). with nothing to suggest what the second column should be. offset_to_z.py writes the correct format, but it was only mentioned in the Multiplane section a hundred lines further down, and it reads a storm-control .off file so it is not a general answer either. The bead location file had a smaller version of the same problem. It was described only as "a text file containing the locations of a few (isolated) beads", and the format - one space separated x y pair per line, in pixels - appeared nowhere outside a comment in the source. Both now have a stated format and a short example. Multiplane step 5 described its z offset file the same loose way and takes the same two columns, so it now refers to the Spliner description rather than repeating it. Also fixes beads_zoffset.text, which is spelled beads_zoffset.txt two lines below it and everywhere else. The documentation builds without warnings, as it did before. --- doc/source/analysis.rst | 33 ++++++++++++++++++++++++++++----- 1 file changed, 28 insertions(+), 5 deletions(-) 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. :: From 1631f6a0f4379ea299436817242b0658833e35ed Mon Sep 17 00:00:00 2001 From: Hazen Babcock Date: Tue, 8 Sep 2026 17:48:04 -0400 Subject: [PATCH 5/5] Say how to run an analysis in AGENTS.md The file covered how to change this code and said nothing about how to drive it, but "analyze this movie with 3D-DAOSTORM in 2d fixed mode" is a perfectly ordinary request and the answer was spread across analysis.rst, parameters.rst and the sample files. Running one of these turns out to work fine as documented. What is missing is not procedure, it is the two things that go wrong quietly: Parameter files are always copied from an existing one, and four of the parameters in them describe the microscope rather than the analysis - pixel_size, camera_gain, camera_offset and the initial sigma. Nothing compares them against the movie, so inheriting them from a sample file gives a clean run, a full set of localizations, and the wrong units. That is worth calling out where someone will read it before starting. The other is that parameter problems are reported at the top of the output. An unrecognized parameter prints a warning as the first line, which is then followed by several hundred lines of per-frame progress, so anything reading the tail of the output will conclude the run was clean. The warning already exists and works; it is just easy to miss. Also records that Spliner and Multiplane need a measured PSF first and that the text files for it have formats worth reading rather than guessing, and that multi_plane.py takes --basename instead of --movie since it loads one movie per channel. Entry points, the analyze() signature, the sample XML naming and the four parameter names were checked against the tree. --- AGENTS.md | 67 +++++++++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 67 insertions(+) diff --git a/AGENTS.md b/AGENTS.md index d468fc88..6f3c8bb0 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -66,6 +66,73 @@ Two things about reading the result: 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.