python-microscopy 25.11.20__cp313-cp313-win_amd64.whl

PYME/Analysis/PSFEst/extractImages.py

@@ -0,0 +1,316 @@
+#!/usr/bin/python
+
+##################
+# extractImages.py
+#
+# Copyright David Baddeley, 2009
+# d.baddeley@auckland.ac.nz
+#
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the GNU General Public License as published by
+# the Free Software Foundation, either version 3 of the License, or
+# (at your option) any later version.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# GNU General Public License for more details.
+#
+# You should have received a copy of the GNU General Public License
+# along with this program.  If not, see <http://www.gnu.org/licenses/>.
+#
+##################
+
+import numpy as np
+from PYME.IO import tabular
+import scipy.ndimage
+import logging
+logger = logging.getLogger(__name__)
+
+import PYME.pyme_warnings
+
+
+def getPSFSlice(datasource, resultsSource, metadata, zm=None):
+    f1 = tabular.ResultsFilter(resultsSource, error_x=[1, 30], A=[10, 500], sig=(150 / 2.35, 900 / 2.35))
+
+    ims, pts, zvals, zis = extractIms(datasource, f1, metadata, zm)
+    return getPSF(ims, pts, zvals, zis)
+
+#def getPointPoss(dataSource, zm, zmid):
+#    for i in range()
+
+def extractIms(dataSource, results, metadata, zm =None, roiSize=10, nmax = 1000):
+    ims = np.zeros((2*roiSize, 2*roiSize, len(results['x'])))
+    
+    vs = metadata.voxelsize_nm
+    points = (np.array([results['x']/vs.x,
+                        results['y']/vs.y,
+                        results['A']]).T)
+
+    pts = np.round(points[:,:2])
+    points[:,:2] = points[:,:2] - pts
+    ts = results['tIndex']
+    bs = results['fitResults_background']
+
+    ind = (pts[:,0] > roiSize)*(pts[:,1] > roiSize)*(pts[:,0] < (dataSource.shape[0] - roiSize))*(pts[:,1] < (dataSource.shape[1] - roiSize))
+
+    #print ind.sum()
+
+    points = points[ind,:]
+    pts = pts[ind,:]
+    ts = ts[ind]
+    bs = bs[ind]
+
+    if not zm is None:
+        zvals = np.array(list(set(zm.yvals)))
+        zvals.sort()
+
+        zv = zm(ts.astype('f'))
+        #print zvals
+        #print zv
+
+        zis = np.array([np.argmin(np.abs(zvals - z)) for z in zv])
+        #print zis
+    else:
+        zvals = np.array([0])
+        zis = 0.*ts
+
+    for i in range(len(ts)):
+        x = pts[i,0]
+        y = pts[i,1]
+
+        t = ts[i]
+        #print t
+
+        ims[:,:,i] = dataSource[(x-roiSize):(x+roiSize), (y-roiSize):(y+roiSize), t].squeeze() - bs[i]
+
+    return ims - metadata['Camera.ADOffset'], points, zvals, zis
+
+
+def getPSF(ims, points, zvals, zis):
+    height, width = ims.shape[0],ims.shape[1]
+    kx,ky = np.mgrid[:height,:width]#,:self.sliceShape[2]]
+
+    kx = np.fft.fftshift(kx - height/2.)/height
+    ky = np.fft.fftshift(ky - width/2.)/width
+    
+
+    d = np.zeros((height, width, len(zvals)))
+    print((d.shape))
+
+    for i in range(len(points)):
+        F = np.fft.fftn(ims[:,:,i])
+        p = points[i,:]
+        #print zis[i]
+        #print ifftn(F*exp(-2j*pi*(kx*-p[0] + ky*-p[1]))).real.shape
+        d[:,:,zis[i]] = d[:,:,zis[i]] + np.fft.ifftn(F*np.exp(-2j*np.pi*(kx*-p[0] + ky*-p[1]))).real
+
+    d = len(zvals)*d/(points[:,2].sum())
+    
+    #estimate background as a function of z by averaging rim pixels
+    bg = (d[0,:,:].squeeze().mean(0) + d[-1,:,:].squeeze().mean(0) + d[:,0,:].squeeze().mean(0) + d[:,-1,:].squeeze().mean(0))/4
+    d = d - bg[None,None,:]
+    d = d/d.sum(1).sum(0)[None,None,:]
+
+    return d
+
+def getIntCenter(im):
+    im = im.squeeze()
+    X, Y, Z = np.ogrid[0:im.shape[0], 0:im.shape[1], 0:im.shape[2]]
+
+    #from pylab import *
+    #imshow(im.max(2))
+
+    X = X.astype('f') - X.mean()
+    Y = Y.astype('f') - Y.mean()
+    Z = Z.astype('f')
+
+    im2 = im - im.min()
+    im2 = im2 - 0.5*im2.max()
+    im2 = im2*(im2 > 0)
+
+    ims = im2.sum()
+
+    x = (im2*X).sum()/ims
+    y = (im2*Y).sum()/ims
+    z = (im2*Z).sum()/ims
+
+    #print x, y, z
+
+    return x, y, z
+
+def _expand_z(ps_shape, im_shape, points):
+    """
+    Expand ROI size in z to maximum supported by image and point locations.
+    
+    Parameters
+    ----------
+    ps_shape : 3-tuple of psf half roi sizes
+    im_shape : 3-tuple giving
+    points : locations of points
+
+    Returns
+    -------
+
+    new_shape : 3-tuple of expanded PSF shape
+    """
+    
+    sx, sy, sz = ps_shape
+    logger.debug('Expanding PSF z size to maximum supported by data')
+    
+    pts = np.asarray(points)
+    dzl = np.min(pts[:,2]) - 1
+    dzh = np.min(im_shape[2] - pts[:,2]) -1
+    
+    szn = int(min(dzl, dzh))
+    
+    if szn < sz:
+        logger.warning('New PSF is SMALLER than requested (szn = %d, request = %d)' % (szn, sz))
+    
+    return sx, sy, szn
+
+
+def getPSF3D(im, points, PSshape=(30,30,30), blur=(0.5, 0.5, 1), normalize=True, centreZ=True, centreXY=True,
+             x_offset=0, y_offset=0, z_offset=0,expand_z=False, pad=(0,0,4)):
+    
+    """
+    Extract a 3D PSF by averaging the images of a set of point sources.
+
+    Parameters
+    ----------
+    im : ndarray
+        3D image stack
+    points : ndarray of shape (N, 3)
+        array of point locations
+    PSshape : tuple of 3 ints
+        half size of PSF in pixels
+    blur : tuple of 3 floats
+        standard deviation of Gaussian blur to apply to each dimension
+    normalize : bool, 'max', or 'sum'
+        If True, or max normalize the entire PSF to the maximum value. If 'sum', normalize to the sum of the in-focus plane of the PSF
+    centreZ : bool
+        If True, make sure the PSF is centred in z. If False, individual PSF images will still be aligned with respect to each other, but the
+        PSF as a whole will not be. Useful when taking multi-channel PSFs where you want to preserve chromatic shifts between the channels
+        in the PSF so that deconvolution will take care of chromatic shifts without an explicit separate shift correction step. Also used
+        when extracting multi-channel PSFs for Biplane localization.
+    centreXY : bool
+        Similar to centreZ, but for the xy plane. Do not use for biplane PSFs (it is generally not possible to assume a constant
+        lateral shift across the field of view, so this needs to be handled with shiftmaps).
+    expand_z : bool
+        Expand to the entire z range of the image. Was useful for deconvolution to ensure PSF sizes match the image size, but
+        largely unnecessary as deconvolution code now has PSF expansion built in.
+    pad : tuple of 3 ints
+        How much to pad the PSF in each dimension to avoid wrap-around when performing Fourier domain shifting. Default is (0,0,2)
+    """
+    
+    # pad the PSF shape to avoid wrap-around when performing Fourier domain shifting
+    PSshape = np.array(PSshape) + np.array(pad)
+    
+    if expand_z:
+        PSshape = _expand_z(PSshape, im.shape, points)
+        
+    sx, sy, sz = PSshape
+    height, width, depth = 2 * np.array(PSshape, dtype=int) + 1
+
+    kx, ky, kz = np.mgrid[:height, :width, :depth]  # ,:self.sliceShape[2]]
+
+    kx = np.fft.fftshift(kx - height / 2.) / height
+    ky = np.fft.fftshift(ky - width / 2.) / width
+    kz = np.fft.fftshift(kz - depth / 2.) / depth
+
+    logger.debug('Extracting PSF of size: %d, %d, %d [px]' % (height, width, depth))
+    d = np.zeros((height, width, depth))
+
+    imgs = []
+    dxs = []
+    dys = []
+    dzs = []
+
+    for px,py,pz in points:
+        print((px, py, pz))
+        px = int(px)
+        py = int(py)
+        pz = int(pz)
+        imi = im[(px-sx):(px+sx+1),(py-sy):(py+sy+1),(pz-sz):(pz+sz+1)]
+        print((imi.shape))
+        dx, dy, dz = getIntCenter(imi)
+        dz -= sz
+
+        imgs.append(imi)
+        dxs.append(dx)
+        dys.append(dy)
+        dzs.append(dz)
+
+    dxs = np.array(dxs)
+    dys = np.array(dys)
+    dzs = np.array(dzs)
+
+    dzm = dzs.mean()
+    dxm = dxs.mean()
+    dym = dys.mean()
+
+    if not centreZ:
+        #images will still be aligned with each other, but any shift in the channel as a whole will be maintained.
+        dzs = dzs - dzm + z_offset
+        dzm = 0
+
+    if not centreXY:
+        dxs = dxs - dxm + x_offset
+        dys = dys - dym + y_offset
+        dzm = 0
+
+    for imi, dx, dy, dz in zip(imgs, list(dxs), list(dys), list(dzs)):
+        F = np.fft.fftn(imi)
+        d = d + np.fft.ifftn(F*np.exp(-2j*np.pi*(kx*-dx + ky*-dy + kz*-dz))).real
+
+    print('dzs:', dzs)
+    if np.any(np.abs(dzs) > pad[2]):
+        PYME.pyme_warnings.warn('Axial shift of PSF is greater than padding. PSF is likely to exhibit wrap-around artifacts. This is \
+                           typically caused by insufficent z-extent in the requested PSF shape and/or the calibration z-stack.')
+    
+    # remove padding
+    px, py, pz = pad
+    d= d[px:(-px if (px > 0) else None), py:(-py if py > 0 else None), pz:(-pz if pz > 0 else None)]
+
+    d = scipy.ndimage.gaussian_filter(d, blur)
+    #estimate background as a function of z by averaging rim pixels
+    #bg = (d[0,:,:].squeeze().mean(0) + d[-1,:,:].squeeze().mean(0) + d[:,0,:].squeeze().mean(0) + d[:,-1,:].squeeze().mean(0))/4
+    d = d - d.min()
+
+    if (normalize == True) or (normalize == 'max'):
+        d = d/d.max()
+    elif normalize == 'sum':
+        d = d/d[:,:,d.shape[2]/2].sum()
+
+    return d, (dxm, dym, dzm)
+    
+def backgroundCorrectPSFWF(d):
+    import numpy as np
+    from scipy import linalg
+    
+    zf = int(d.shape[2]/2)
+        
+    #subtract a linear background in x
+    Ax = np.vstack([np.ones(d.shape[0]), np.arange(d.shape[0])]).T        
+    bgxf = (d[0,:,zf] + d[-1,:,zf])/2
+    gx = linalg.lstsq(Ax, bgxf)[0]
+    
+    d = d - np.dot(Ax, gx)[:,None,None]
+    
+    #do the same in y
+    Ay = np.vstack([np.ones(d.shape[1]), np.arange(d.shape[1])]).T        
+    bgyf = (d[:,0,zf] + d[:,-1,zf])/2
+    gy = linalg.lstsq(Ay, bgyf)[0]
+    
+    d = d - np.dot(Ay, gy)[None, :,None]
+    
+    
+    #estimate background on central slice as mean of rim pixels
+    #bgr = (d[0,:,zf].mean() + d[-1,:,zf].mean() + d[:,0,zf].mean() + d[:,-1,zf].mean())/4
+    
+    #sum over all pixels (and hence mean) should be preserved over z (for widefield psf)
+    dm = d.mean(1).mean(0)
+    
+    bg = dm - dm[zf]
+    
+    return np.maximum(d - bg[None, None, :], 0) +  1e-5

PYME/Analysis/PSFEst/fit_psf_zernikes.py

@@ -0,0 +1,194 @@
+"""Fits zernike polynomials to a PSF image (as an alternative to Gerchberg-Saxton pupil estimation)
+
+works by using the PSF simulation functions to generate as PSF image with a given set of zernike 
+coefficients and comparing to the measured PSF image.
+"""
+
+import numpy as np
+from PYME.Analysis.PSFGen import fourierHNA
+
+
+
+class ZernikePSFFitter(object):
+    def __init__(self, psf, voxelsize_nm, wavelength=700., n=1.51, NA=1.4, **kwargs):
+        self.psf = psf
+        self.voxelsize_nm = voxelsize_nm
+        self.wavelength = wavelength
+        self.n = n
+        self.NA = NA
+
+    def get_zernike_psf(self, zernike_coeffs=None):
+        """
+        Generate a PSF image with a given set of zernike coefficients
+
+        Parameters
+        ----------
+        zernike_coeffs : ndarray
+            array of zernike coefficients
+
+        Returns
+        -------
+        psf : ndarray
+            simulated PSF image
+        """
+
+        if zernike_coeffs is None:
+            zernike_coeffs = self.zernike_coefficients
+
+        zs = (np.arange(self.psf.shape[2]) - self.psf.shape[2]/2)*self.voxelsize_nm.z
+        
+        # pupil might be slightly undersampled at default PSF size,
+        # so we use a larger pupil to avoid out-of focus grid artifacts
+        # TODO - optimise the size here - doubling the pupil size
+        # may well be overkill (and slow)
+        X = np.arange(-64, 65)*self.voxelsize_nm.x
+        Y = np.arange(-64, 65)*self.voxelsize_nm.y
+
+        ps = fourierHNA.GenZernikePSF(zs, zernike_coeffs, X=X, Y=Y, dx=self.voxelsize_nm.x, lamb=self.wavelength, n=self.n, NA=self.NA, output_shape=self.psf.shape)
+        
+        # TODO - improve normalisation of PSF - maybe build into residual calculation
+        return ps/ps.max()
+    
+    def zernike_residuals(self, zernike_coeffs):
+        """
+        Calculate the residuals between a measured PSF and a simulated PSF with a given set of zernike coefficients
+
+        Parameters
+        ----------
+        zernike_coeffs : ndarray
+            array of zernike coefficients
+
+        Returns
+        -------
+        residuals : ndarray
+            array of residuals between the measured and simulated PSFs
+        """
+
+        sim_psf = self.get_zernike_psf(zernike_coeffs)
+        residuals = sim_psf - self.psf
+
+        # TODO - change noise model???
+        return np.sum(residuals**2)
+
+    def refine_zernike(self, zernike_coeffs, zernike_order, delta=2.0):
+        """
+        Refine a single zernike coefficient, assuming a quadratic dependence on the residuals
+        """
+        coeffs = np.copy(zernike_coeffs)
+        
+        res0 = self.zernike_residuals(coeffs)
+
+        coeffs[zernike_order] += delta
+        res1 = self.zernike_residuals(coeffs)
+
+        coeffs[zernike_order] -= 2*delta
+        res2 = self.zernike_residuals(coeffs)
+
+        # fit a quadratic to the residuals
+        a = (res1 + res2 - 2*res0)/(2*delta**2)
+        b = (res1 - res2)/(2*delta)
+        c = res0
+
+        return (-b/(2*a))
+    
+    def refine_zernikes(self, zernike_coeffs, delta):
+        """
+        Refine all zernike coefficients
+        """
+        refined_coeffs = np.copy(zernike_coeffs)
+        for i in range(1, len(zernike_coeffs)): # skip the piston term
+            d_i = self.refine_zernike(refined_coeffs, i, delta) # progressive update
+            #d_i = self.refine_zernike(zernike_coeffs, i, delta) # update all at once (after all have been calculated) - this doesn't seem to be as good
+            refined_coeffs[i] += np.clip(d_i, -delta, delta) # limit the step size to the range we are searching
+        
+        return refined_coeffs
+    
+    def approximate_zernikes(self, num_zernike_orders, num_iterations=5, starting_zernikes={}, plot=True, **kwargs):
+        """
+        Approximate the zernike coefficients of a PSF image using a least squares fit
+        """
+        zernikes = np.zeros(num_zernike_orders)
+        for k,v in starting_zernikes.items():
+            zernikes[k] = v
+
+        delta = 2.0
+
+        self._residuals = [] 
+        self._deltas = []
+
+        for i in range(num_iterations):
+            z0 = zernikes
+            zernikes = self.refine_zernikes(zernikes, delta)
+            self._residuals.append(self.zernike_residuals(zernikes))
+            self._deltas.append(delta)
+            print(self._residuals[-1])
+
+            # if we made a large move on any of the coefficients, hold step size, otherwise reduce it.
+            max_change = np.max(np.abs(zernikes - z0))
+            delta = np.max([delta/ 2.0, max_change])
+        
+        self.zernike_coefficients = zernikes
+
+        if plot:
+            self.plot_results()
+        
+        return zernikes
+    
+
+    def get_misfit_map(self, zernike_coeffs):
+        """
+        Get a map of the residuals between the measured and simulated PSFs
+        as a 5x5 tiled image of z planes
+        """
+        sim_psf = self.get_zernike_psf(zernike_coeffs)
+        res =  sim_psf - self.psf
+
+        step = res.shape[2]//25
+        out = np.zeros([5*self.psf.shape[0], 5*self.psf.shape[1]])
+
+        for i in range(5):
+            for j in range(5):
+                out[i*self.psf.shape[0]:(i+1)*self.psf.shape[0], j*self.psf.shape[1]:(j+1)*self.psf.shape[1]] = res[:,:,(i*5+j)*step]
+
+        return out
+    
+    def plot_results(self):
+        """
+        Plot the results of the zernike fitting
+        """
+        import matplotlib.pyplot as plt
+        plt.figure()
+        plt.imshow(100*self.get_misfit_map(self.zernike_coefficients), cmap=plt.cm.RdBu_r, clim=(-30, 30))
+        plt.colorbar()
+        plt.title('Fit residuals (% of max)')
+        plt.xticks([])
+        plt.yticks([])
+
+        plt.figure()
+        plt.subplot(311)
+        plt.plot(self._residuals)
+        plt.ylabel('Squared error')
+        plt.xlabel('Iteration')
+        plt.subplot(312)
+        plt.plot(self._deltas)
+        plt.ylabel('Step size')
+        plt.xlabel('Iteration')
+
+        plt.subplot(313)
+        b = plt.bar(np.arange(len(self.zernike_coefficients)), self.zernike_coefficients)
+        plt.bar_label(b, fmt='%.2f')
+        plt.ylabel('Zernike coefficient')
+        plt.xlabel('Zernike order')
+        plt.axhline(0, color='k', linestyle='--')
+
+        plt.show()
+
+
+
+
+
+
+
+    
+
+