diff --git a/storm_analysis/multi_plane/fitting_mp.py b/storm_analysis/multi_plane/fitting_mp.py index 520771a2..abee7070 100644 --- a/storm_analysis/multi_plane/fitting_mp.py +++ b/storm_analysis/multi_plane/fitting_mp.py @@ -57,10 +57,11 @@ def __init__(self, parameters = None, **kwds): self.backgrounds = [] self.bg_filters = [] # Matched filters for background estimation. self.images = [] - self.mfilters = [] # Matched filters for foreground estimation. + self.mfilter_banks = [] # Foreground matched filter banks, one per channel. self.n_channels = None + self.n_zvals = 0 # The number of z values, i.e. the size of each bank. self.variances = [] # sCMOS variance data. - self.vfilters = [] # Matched filters for (background) variance estimation. + self.vfilter_banks = [] # (Background) variance matched filter banks, one per channel. self.xt = [] self.yt = [] self.z_values = [] @@ -77,10 +78,9 @@ def cleanUp(self): at.cleanup() # Clean up foreground filters and variance filters. - for i in range(len(self.mfilters)): - for j in range(len(self.mfilters[i])): - self.mfilters[i][j].cleanup() - self.vfilters[i][j].cleanup() + for i in range(len(self.mfilter_banks)): + self.mfilter_banks[i].cleanup() + self.vfilter_banks[i].cleanup() # Clean up background filters. for i in range(len(self.bg_filters)): @@ -197,19 +197,26 @@ def peakFinder(self, fit_peaks_images): for fi in fit_peaks_images: tf.write(fi.astype(numpy.float32)) + # + # Each channel's image is the same at every z value, so convolve it + # with all of that channel's variance filters in one pass rather than + # transforming it again for each z value. + # + # I believe that this is correct, the variance of the weighted average + # of independent processes is calculated using the square of the weights. + # + conv_vars = [] + for j in range(self.n_channels): + conv_vars.append(self.vfilter_banks[j].convolve(fit_peaks_images[j] + self.backgrounds[j])) + # Iterate over z values. - for i in range(len(self.vfilters)): + for i in range(self.n_zvals): bg_variance = numpy.zeros(fit_peaks_images[0].shape) # Iterate over channels / planes. - for j in range(len(self.vfilters[i])): + for j in range(self.n_channels): - # Convolve fit image + background with the appropriate variance filter. - # - # I believe that this is correct, the variance of the weighted average - # of independent processes is calculated using the square of the weights. - # - conv_var = self.vfilters[i][j].convolve(fit_peaks_images[j] + self.backgrounds[j]) + conv_var = conv_vars[j][i] # Transform variance to the channel 0 frame. if self.atrans[j] is None: @@ -240,16 +247,20 @@ def peakFinder(self, fit_peaks_images): fg_averages = [] # This is the average foreground across all the planes for each z value. foregrounds = [] # This is the foreground for each plane and z value. + # As above, one pass per channel rather than one per (z, channel). + conv_fgs = [] + for j in range(self.n_channels): + conv_fgs.append(self.mfilter_banks[j].convolve(self.images[j] - fit_peaks_images[j] - self.backgrounds[j])) + # Iterate over z values. - for i in range(len(self.mfilters)): + for i in range(self.n_zvals): foreground = numpy.zeros(fit_peaks_images[0].shape) foregrounds.append([]) # Iterate over channels / planes. - for j in range(len(self.mfilters[i])): + for j in range(self.n_channels): - # Convolve image / background with the appropriate PSF. - conv_fg = self.mfilters[i][j].convolve(self.images[j] - fit_peaks_images[j] - self.backgrounds[j]) + conv_fg = conv_fgs[j][i] # Store convolved image in foregrounds. foregrounds[i].append(conv_fg) @@ -295,7 +306,7 @@ def setVariances(self, variances): We initialize the following here because at __init__ we don't know how big the images are. This is called after the analysis specific version pads the variances and - creates the mfilters[] and the vfilters[] class members. + creates the mfilter_banks[] and the vfilter_banks[] class members. Note the assumption that every frame in all the movies is the same size. @@ -332,16 +343,21 @@ def setVariances(self, variances): # we are applying to the foreground. # + # One pass per channel, as in peakFinder(). + # + conv_vars = [] + for j in range(self.n_channels): + conv_vars.append(self.vfilter_banks[j].convolve(variances[j])) + # Iterate over z values. # - for i in range(len(self.mfilters)): + for i in range(self.n_zvals): variance = numpy.zeros(variances[0].shape) # Iterate over channels / planes. - for j in range(len(self.mfilters[i])): + for j in range(self.n_channels): - # Convolve variance with the appropriate variance filter. - conv_var = self.vfilters[i][j].convolve(variances[j]) + conv_var = conv_vars[j][i] # Transform variance to the channel 0 frame. if self.atrans[j] is None: @@ -451,13 +467,14 @@ def setVariances(self, variances): # Create "foreground" and "variance" filters. # - # These are stored in a list indexed by z value, then by - # channel / plane. So self.mfilters[1][2] is the filter - # for z value 1, plane 2. + # These are stored as one bank per channel / plane, each bank + # holding that channel's filter for every z value. So + # self.mfilter_banks[2] convolves plane 2 at all z values at once. # + mfilter_psfs = [[] for i in range(self.n_channels)] + vfilter_psfs = [[] for i in range(self.n_channels)] + for i, mfilter_z in enumerate(self.mfilters_z): - self.mfilters.append([]) - self.vfilters.append([]) for j, psf_object in enumerate(self.psf_objects): psf = psf_object.getPSF(mfilter_z, @@ -469,14 +486,8 @@ def setVariances(self, variances): # or if it does that they are very small. # psf_norm = psf/numpy.sum(psf) - self.mfilters[i].append(matchedFilterC.MatchedFilter(psf_norm, - fftw_estimate = self.parameters.getAttr("fftw_estimate"), - memoize = True, - max_diff = 1.0e-3)) - self.vfilters[i].append(matchedFilterC.MatchedFilter(psf_norm * psf_norm, - fftw_estimate = self.parameters.getAttr("fftw_estimate"), - memoize = True, - max_diff = 1.0e-3)) + mfilter_psfs[j].append(psf_norm) + vfilter_psfs[j].append(psf_norm * psf_norm) # Save a pictures of the PSFs for debugging purposes. if self.check_mode: @@ -484,6 +495,13 @@ def setVariances(self, variances): filename = "psf_z{0:.3f}_c{1:d}.tif".format(mfilter_z, j) tifffile.imwrite(filename, psf.astype(numpy.float32)) + self.n_zvals = len(self.mfilters_z) + for j in range(self.n_channels): + self.mfilter_banks.append(matchedFilterC.MatchedFilterBank(mfilter_psfs[j], + fftw_estimate = self.parameters.getAttr("fftw_estimate"))) + self.vfilter_banks.append(matchedFilterC.MatchedFilterBank(vfilter_psfs[j], + fftw_estimate = self.parameters.getAttr("fftw_estimate"))) + # This handles the rest of the initialization. # super(MPPeakFinderArb, self).setVariances(variances) @@ -578,25 +596,19 @@ def setVariances(self, variances): # Create "foreground" and "variance" filters. There is # only one z value here. # - # These are stored in a list indexed by z value, then by - # channel / plane. So self.mfilters[1][2] is the filter - # for z value 1, plane 2. + # These are stored as one bank per channel / plane, each bank + # holding that channel's filter for every z value. So + # self.mfilter_banks[2] convolves plane 2 at all z values at once. # - self.mfilters.append([]) - self.vfilters.append([]) - psf_norm = fitting.gaussianPSF(variances[0].shape, self.parameters.getAttr("foreground_sigma")) var_norm = psf_norm * psf_norm + self.n_zvals = 1 for i in range(self.n_channels): - self.mfilters[0].append(matchedFilterC.MatchedFilter(psf_norm, - fftw_estimate = self.parameters.getAttr("fftw_estimate"), - memoize = True, - max_diff = 1.0e-3)) - self.vfilters[0].append(matchedFilterC.MatchedFilter(var_norm, - fftw_estimate = self.parameters.getAttr("fftw_estimate"), - memoize = True, - max_diff = 1.0e-3)) + self.mfilter_banks.append(matchedFilterC.MatchedFilterBank([psf_norm], + fftw_estimate = self.parameters.getAttr("fftw_estimate"))) + self.vfilter_banks.append(matchedFilterC.MatchedFilterBank([var_norm], + fftw_estimate = self.parameters.getAttr("fftw_estimate"))) # Save a pictures of the PSFs for debugging purposes. if self.check_mode: diff --git a/storm_analysis/sa_library/fitting.py b/storm_analysis/sa_library/fitting.py index c30cf27a..72fc7653 100644 --- a/storm_analysis/sa_library/fitting.py +++ b/storm_analysis/sa_library/fitting.py @@ -619,9 +619,9 @@ def __init__(self, parameters = None, psf_object = None, **kwds): kwds["parameters"] = parameters super(PeakFinderArbitraryPSF, self).__init__(**kwds) - self.fg_mfilter = [] + self.fg_mfilter = None # Foreground MatchedFilterBank, one psf per z value. self.fg_mfilter_zval = [] - self.fg_vfilter = [] + self.fg_vfilter = None # Foreground variance MatchedFilterBank. self.psf_object = psf_object self.z_values = [] @@ -666,9 +666,9 @@ def __init__(self, parameters = None, psf_object = None, **kwds): def cleanUp(self): super(PeakFinderArbitraryPSF, self).cleanUp() - for i in range(len(self.fg_mfilter)): - self.fg_mfilter[i].cleanup() - self.fg_vfilter[i].cleanup() + if self.fg_mfilter is not None: + self.fg_mfilter.cleanup() + self.fg_vfilter.cleanup() def newImage(self, new_image): """ @@ -682,27 +682,32 @@ def newImage(self, new_image): # If does not already exist, create filter objects from # the PSF at different z values. # - if (len(self.fg_mfilter) == 0): + if self.fg_mfilter is None: + mfilter_psfs = [] + vfilter_psfs = [] for zval in self.fg_mfilter_zval: psf = self.psf_object.getPSF(zval, shape = new_image.shape, normalize = False) psf_norm = psf/numpy.sum(psf) - fg_mfilter = matchedFilterC.MatchedFilter(psf_norm, - fftw_estimate = self.parameters.getAttr("fftw_estimate"), - memoize = True, - max_diff = 1.0e-3) - self.fg_mfilter.append(fg_mfilter) - self.fg_vfilter.append(matchedFilterC.MatchedFilter(psf_norm * psf_norm, - fftw_estimate = self.parameters.getAttr("fftw_estimate"), - memoize = True, - max_diff = 1.0e-3)) + mfilter_psfs.append(psf_norm) + vfilter_psfs.append(psf_norm * psf_norm) # Save a picture of the PSF for debugging purposes. if self.check_mode: print("psf max", numpy.max(psf)) filename = "psf_{0:.3f}.tif".format(zval) tifffile.imwrite(filename, psf.astype(numpy.float32)) + + # + # These are banks rather than one filter per z value because + # peakFinder() applies all of them to the same image, so the + # forward FFT only needs to happen once per bank. + # + self.fg_mfilter = matchedFilterC.MatchedFilterBank(mfilter_psfs, + fftw_estimate = self.parameters.getAttr("fftw_estimate")) + self.fg_vfilter = matchedFilterC.MatchedFilterBank(vfilter_psfs, + fftw_estimate = self.parameters.getAttr("fftw_estimate")) def peakFinder(self, fit_peaks_image): """ @@ -736,11 +741,18 @@ def peakFinder(self, fit_peaks_image): fg_tif = tifffile.TiffWriter("foreground.tif") fg_bg_ratio_tif = tifffile.TiffWriter("fg_bg_ratio.tif") + # + # Both of these inputs are the same at every z value, so each bank + # transforms its input once instead of once per z value. + # + backgrounds = self.fg_vfilter.convolve(bg_var) + foregrounds = self.fg_mfilter.convolve(self.image - self.background - fit_peaks_image) + masked_images = [] - for i in range(len(self.fg_mfilter)): + for i in range(len(self.fg_mfilter_zval)): # Estimate background variance at this particular z value. - background = self.fg_vfilter[i].convolve(bg_var) + background = backgrounds[i] # Remove problematic values. # @@ -754,8 +766,7 @@ def peakFinder(self, fit_peaks_image): bg_std = numpy.sqrt(background) # Calculate foreground. - foreground = self.image - self.background - fit_peaks_image - foreground = self.fg_mfilter[i].convolve(foreground) + foreground = foregrounds[i] # Calculate foreground in units of signal to noise. fg_bg_ratio = foreground/bg_std diff --git a/storm_analysis/sa_library/matched_filter.c b/storm_analysis/sa_library/matched_filter.c index 8469c0a8..742c9226 100644 --- a/storm_analysis/sa_library/matched_filter.c +++ b/storm_analysis/sa_library/matched_filter.c @@ -49,11 +49,44 @@ struct filter_struct { }; typedef struct filter_struct filter; +/* + * A bank of filters that all get applied to the same image. + * + * The point of this is the forward FFT. Applying N filters with N separate + * 'filter' structures transforms the image N times, and every one of those + * transforms is identical. Here the image is transformed once and the result + * is reused, so the cost goes from 2N transforms to N+1. + * + * image_fft has to be kept separate from work_fft because convolve() does its + * complex multiply in place, which would destroy the transform we are trying + * to reuse. + */ +struct filter_bank_struct { + int fft_size; + int image_size; + int n_filters; + int x_size; + int y_size; + + double *fft_vector; + + fftw_plan fft_backward; + fftw_plan fft_forward; + + fftw_complex *image_fft; + fftw_complex *work_fft; + fftw_complex **psf_ffts; +}; +typedef struct filter_bank_struct filterBank; + /* Function Declarations */ void cleanup(filter *); +void cleanupBank(filterBank *); void convolve(filter *, double *, double *); +void convolveBank(filterBank *, double *, double *); void convolveMemo(filter *, double *, double *); filter *initialize(double *, double, int, int, int); +filterBank *initializeBank(double *, int, int, int, int); /* Functions */ @@ -214,6 +247,125 @@ filter *initialize(double *psf, double max_diff, int x_size, int y_size, int est return flt; } + +/* + * cleanupBank() + * + * bank - A pointer to a filterBank structure. + */ +void cleanupBank(filterBank *bank) +{ + int i; + + for(i=0;in_filters;i++){ + fftw_free(bank->psf_ffts[i]); + } + free(bank->psf_ffts); + + fftw_free(bank->fft_vector); + fftw_free(bank->image_fft); + fftw_free(bank->work_fft); + + fftw_destroy_plan(bank->fft_backward); + fftw_destroy_plan(bank->fft_forward); + + free(bank); +} + + +/* + * convolveBank() + * + * Convolve image with every psf in the bank. This is the same arithmetic + * that convolve() does, in the same order, just with the forward transform + * hoisted out of the loop. + * + * bank - A pointer to a filterBank structure. + * image - The image (must be the same size as the psfs the bank was made with). + * results - Pre-allocated storage for n_filters images, in psf order. + */ +void convolveBank(filterBank *bank, double *image, double *results) +{ + int i; + + /* Compute FFT of the image, once. */ + ftmDoubleCopy(image, bank->fft_vector, bank->image_size); + fftw_execute(bank->fft_forward); + + for(i=0;in_filters;i++){ + + /* Multiply by FFT of this psf and compute inverse FFT. */ + ftmComplexMultiply(bank->work_fft, bank->image_fft, bank->psf_ffts[i], bank->fft_size, 0); + fftw_execute(bank->fft_backward); + + /* Copy into the matching slice of result. */ + ftmDoubleCopy(bank->fft_vector, results + i*bank->image_size, bank->image_size); + } +} + + +/* + * initializeBank() + * + * Set things up for FFT convolution with several psfs at once. + * + * psfs - n_filters psfs, contiguous, each (x_size, y_size). + * n_filters - The number of psfs. + * x_size - the size of the psfs in x (slow dimension). + * y_size - the size of the psfs in y (fast dimension). + * estimate - 0/1 to just use an estimated FFT plan. + */ +filterBank *initializeBank(double *psfs, int n_filters, int x_size, int y_size, int estimate) +{ + int i; + double normalization; + filterBank *bank; + + bank = (filterBank *)malloc(sizeof(filterBank)); + + bank->fft_size = x_size * (y_size/2 + 1); + bank->image_size = x_size * y_size; + bank->n_filters = n_filters; + bank->x_size = x_size; + bank->y_size = y_size; + + normalization = 1.0/((double)(x_size * y_size)); + + /* Allocate storage. */ + bank->fft_vector = (double *)fftw_malloc(sizeof(double)*bank->image_size); + bank->image_fft = (fftw_complex *)fftw_malloc(sizeof(fftw_complex)*bank->fft_size); + bank->work_fft = (fftw_complex *)fftw_malloc(sizeof(fftw_complex)*bank->fft_size); + + bank->psf_ffts = (fftw_complex **)malloc(sizeof(fftw_complex *)*n_filters); + for(i=0;ipsf_ffts[i] = (fftw_complex *)fftw_malloc(sizeof(fftw_complex)*bank->fft_size); + } + + /* + * Create FFT plans. This has to happen before anything is written into + * fft_vector, as FFTW_MEASURE overwrites its arrays while planning. + * + * One pair of plans serves the whole bank, rather than a pair per psf. + */ + if (estimate){ + bank->fft_forward = fftw_plan_dft_r2c_2d(x_size, y_size, bank->fft_vector, bank->image_fft, FFTW_ESTIMATE); + bank->fft_backward = fftw_plan_dft_c2r_2d(x_size, y_size, bank->work_fft, bank->fft_vector, FFTW_ESTIMATE); + } + else { + bank->fft_forward = fftw_plan_dft_r2c_2d(x_size, y_size, bank->fft_vector, bank->image_fft, FFTW_MEASURE); + bank->fft_backward = fftw_plan_dft_c2r_2d(x_size, y_size, bank->work_fft, bank->fft_vector, FFTW_MEASURE); + } + + /* Compute FFT of each psf and save. */ + for(i=0;iimage_size, bank->fft_vector, bank->image_size); + fftw_execute(bank->fft_forward); + ftmComplexCopyNormalize(bank->image_fft, bank->psf_ffts[i], normalization, bank->fft_size); + } + + return bank; +} + /* * The MIT License * diff --git a/storm_analysis/sa_library/matched_filter_c.py b/storm_analysis/sa_library/matched_filter_c.py index 06bf5d12..7e3e829e 100644 --- a/storm_analysis/sa_library/matched_filter_c.py +++ b/storm_analysis/sa_library/matched_filter_c.py @@ -33,6 +33,20 @@ m_filter.initialize.restype = ctypes.c_void_p +m_filter.cleanupBank.argtypes = [ctypes.c_void_p] + +m_filter.convolveBank.argtypes = [ctypes.c_void_p, + ndpointer(dtype = numpy.float64), + ndpointer(dtype = numpy.float64)] + +m_filter.initializeBank.argtypes = [ndpointer(dtype = numpy.float64), + ctypes.c_int, + ctypes.c_int, + ctypes.c_int, + ctypes.c_int] + +m_filter.initializeBank.restype = ctypes.c_void_p + class MatchedFilterException(Exception): @@ -88,6 +102,60 @@ def convolve(self, image): return result +class MatchedFilterBank(object): + """ + Convolve one image with several psfs. + + This exists because applying N psfs to the same image with N MatchedFilter + objects computes the same forward FFT of that image N times. Here it is + computed once and reused, so the transform count per call goes from 2N to + N+1. + + The results are the same numbers that N MatchedFilter objects produce, in + psf order. There is no memoization, the caller is expected to be applying + the bank to an image that has changed. + """ + def __init__(self, psfs, fftw_estimate = False): + """ + psfs - A list of psfs, or an (n, x, y) array. All the same shape. + """ + psfs = numpy.asarray(psfs, dtype = numpy.float64) + if (psfs.ndim != 3): + raise MatchedFilterException("psfs must be (n, x, y), got " + str(psfs.shape)) + + self.n_filters = psfs.shape[0] + self.psf_shape = (psfs.shape[1], psfs.shape[2]) + + rc_psfs = numpy.zeros(psfs.shape) + for i in range(self.n_filters): + rc_psfs[i,:,:] = recenterPSF.recenterPSF(psfs[i,:,:]) + rc_psfs = numpy.ascontiguousarray(rc_psfs, dtype = numpy.float64) + + self.mfilter = m_filter.initializeBank(rc_psfs, + self.n_filters, + self.psf_shape[0], + self.psf_shape[1], + int(fftw_estimate)) + + def cleanup(self): + m_filter.cleanupBank(self.mfilter) + self.mfilter = None + + def convolve(self, image): + """ + Returns an (n_filters, x, y) array, one convolution per psf. + """ + if (image.shape[0] != self.psf_shape[0]) or (image.shape[1] != self.psf_shape[1]): + raise MatchedFilterException("Image shape must match psf shape! " + str(image.shape) + " != " + str(self.psf_shape)) + + image = numpy.ascontiguousarray(image, dtype = numpy.float64) + results = numpy.zeros((self.n_filters, self.psf_shape[0], self.psf_shape[1]), + dtype = numpy.float64) + m_filter.convolveBank(self.mfilter, image, results) + + return results + + if (__name__ == "__main__"): import tifffile diff --git a/storm_analysis/test/test_matched_filter.py b/storm_analysis/test/test_matched_filter.py index 260f4ec2..c782912d 100644 --- a/storm_analysis/test/test_matched_filter.py +++ b/storm_analysis/test/test_matched_filter.py @@ -179,6 +179,88 @@ def test_matched_filter5(): flt.cleanup() + +def test_matched_filter_bank1(): + """ + Verify that a bank gives exactly what the same psfs give one at a time. + + This is the property the bank exists to preserve: it only hoists the + forward FFT out of the loop, so every number should be identical, not + merely close. + """ + x_size = 80 + y_size = 90 + n_filters = 4 + + psfs = [] + objects = numpy.zeros((1, 5)) + for i in range(n_filters): + objects[0,:] = [x_size/2, y_size/2, 1.0, 1.0 + 0.4*i, 1.3 + 0.2*i] + psf = dg.drawGaussians((x_size, y_size), objects) + psfs.append(psf/numpy.sum(psf)) + + numpy.random.seed(42) + image = numpy.random.poisson(100, (x_size, y_size)).astype(numpy.float64) + + flts = [matchedFilterC.MatchedFilter(psf, fftw_estimate = True) for psf in psfs] + expected = numpy.array([flt.convolve(image) for flt in flts]) + + bank = matchedFilterC.MatchedFilterBank(psfs, fftw_estimate = True) + result = bank.convolve(image) + + assert (result.shape == (n_filters, x_size, y_size)) + assert numpy.array_equal(expected, result) + + for flt in flts: + flt.cleanup() + bank.cleanup() + + +def test_matched_filter_bank2(): + """ + Verify that a bank of one is just a matched filter, and that mismatched + shapes are rejected rather than read past the end of the array. + """ + x_size = 40 + y_size = 50 + + objects = numpy.zeros((1, 5)) + objects[0,:] = [x_size/2, y_size/2, 1.0, 2.0, 2.0] + psf = dg.drawGaussians((x_size, y_size), objects) + psf = psf/numpy.sum(psf) + + image = numpy.zeros((x_size, y_size)) + image[int(x_size/2), int(y_size/2)] = 3.0 + + flt = matchedFilterC.MatchedFilter(psf, fftw_estimate = True) + bank = matchedFilterC.MatchedFilterBank([psf], fftw_estimate = True) + + assert numpy.array_equal(flt.convolve(image), bank.convolve(image)[0]) + + okay = False + try: + bank.convolve(numpy.zeros((x_size + 2, y_size))) + except matchedFilterC.MatchedFilterException: + okay = True + assert okay, "A wrongly shaped image should raise." + + okay = False + try: + matchedFilterC.MatchedFilterBank(psf) + except matchedFilterC.MatchedFilterException: + okay = True + assert okay, "A 2D psf array should raise." + + flt.cleanup() + bank.cleanup() + + if (__name__ == "__main__"): + test_matched_filter1() + test_matched_filter2() + test_matched_filter3() + test_matched_filter4() test_matched_filter5() + test_matched_filter_bank1() + test_matched_filter_bank2()