KaneAI - World’s First E2E Software Testing Agent.Book a Demo 
How to use is_on_axis method in refurb
Best Python code snippet using refurb_python
fit_psf.py
Source:fit_psf.py 
1"""2Tools for working with point-spread functions and optical transfer functions.3Functions for estimating PSF's from images of fluorescent beads (z-stacks or single planes). Useful for generating4experimental PSF's from the average of many beads and fitting 2D and 3D PSF models to beads.5"""6import warnings7import numpy as np8import scipy.ndimage as ndi9import scipy.special as sp10import scipy.integrate11import scipy.interpolate12import scipy.signal13from scipy import fft14from localize_psf import affine, rois15# most of the functions don't require this module, and it does not easily pip install,16# so don't require it. Probably should enforce some reasonable behavior on the functions17# that require it...18# https://pypi.org/project/psfmodels/19try:20 import psfmodels as psfm21 psfmodels_available = True22except ImportError:23 psfmodels_available = False24def blur_img_otf(ground_truth, otf, apodization=1):25 """26 Blur image with OTF27 :param ground_truth:28 :param otf: optical transfer function evalated at the FFT frequencies (with f=0 near the center of the array)29 :return img_blurred:30 """31 gt_ft = fft.fftshift(fft.fftn(fft.ifftshift(ground_truth)))32 img_blurred = fft.fftshift(fft.ifftn(fft.ifftshift(gt_ft * otf * apodization)))33 return img_blurred34def blur_img_psf(ground_truth, psf, apodization=1):35 """36 Blur image with PSF37 :param ground_truth:38 :param psf: point-spread function. this array should be centered at ny//2, nx//239 :return blurred_img:40 """41 # todo: allow PSF of different size than image42 otf, _ = psf2otf(psf)43 return blur_img_otf(ground_truth, otf, apodization)44# tools for converting between different otf/psf representations45def otf2psf(otf, dfs=1):46 """47 Compute the point-spread function from the optical transfer function48 :param otf: otf, as a 1D, 2D or 3D array. Assumes that f=0 is near the center of the array, and frequency are49 arranged by the FFT convention50 :param dfs: (dfz, dfy, dfx), (dfy, dfx), or (dfx). If only a single number is provided, will assume these are the51 same52 :return psf, coords: where coords = (z, y, x)53 """54 if isinstance(dfs, (int, float)) and otf.ndim > 1:55 dfs = [dfs] * otf.ndim56 if len(dfs) != otf.ndim:57 raise ValueError("dfs length must be otf.ndim")58 shape = otf.shape59 drs = np.array([1 / (df * n) for df, n in zip(shape, dfs)])60 coords = [fft.fftshift(fft.fftfreq(n, 1 / (dr * n))) for n, dr in zip(shape, drs)]61 psf = fft.fftshift(fft.ifftn(fft.ifftshift(otf))).real62 return psf, coords63def psf2otf(psf, drs=1):64 """65 Compute the optical transfer function from the point-spread function66 :param psf: psf, as a 1D, 2D or 3D array. Assumes that r=0 is near the center of the array, and positions67 are arranged by the FFT convention68 :param drs: (dz, dy, dx), (dy, dx), or (dx). If only a single number is provided, will assume these are the69 same70 :return otf, coords: where coords = (fz, fy, fx)71 """72 if isinstance(drs, (int, float)) and psf.ndim > 1:73 drs = [drs] * psf.ndim74 if len(drs) != psf.ndim:75 raise ValueError("drs length must be psf.ndim")76 shape = psf.shape77 coords = [fft.fftshift(fft.fftfreq(n, dr)) for n, dr in zip(shape, drs)]78 otf = fft.fftshift(fft.fftn(fft.ifftshift(psf)))79 return otf, coords80def symm_fn_1d_to_2d(arr, fs, fmax, npts):81 """82 Convert a 1D function which is symmetric wrt to the radial variable to a 2D matrix.83 Useful helper function when computing PSFs from 2D OTFs84 :param arr:85 :param fs:86 :param fmax:87 :param npts:88 :return arr_out, fxs, fys:89 """90 ny = 2 * npts + 191 nx = 2 * npts + 192 # fmax = fs.max()93 dx = 1 / (2 * fmax)94 dy = dx95 not_nan = np.logical_not(np.isnan(arr))96 fxs = fft.fftshift(fft.fftfreq(nx, dx))97 fys = fft.fftshift(fft.fftfreq(ny, dy))98 fmag = np.sqrt(fxs[None, :]**2 + fys[:, None]**2)99 to_interp = np.logical_and(fmag >= fs[not_nan].min(), fmag <= fs[not_nan].max())100 arr_out = np.zeros((ny, nx), dtype=arr.dtype)101 arr_out[to_interp] = scipy.interpolate.interp1d(fs[not_nan], arr[not_nan])(fmag[to_interp])102 return arr_out, fxs, fys103def atf2otf(atf, dx=None, wavelength=0.5, ni=1.5, defocus_um=0, fx=None, fy=None):104 """105 Get incoherent transfer function (OTF) from autocorrelation of coherent transfer function (ATF)106 :param atf:107 :param dx:108 :param wavelength:109 :param ni:110 :param defocus_um:111 :param fx:112 :param fy:113 :return otf, atf_defocus:114 """115 ny, nx = atf.shape116 if defocus_um != 0:117 if fx is None:118 fx = fft.fftshift(fft.fftfreq(nx, dx))119 if fy is None:120 fy = fft.fftshift(fft.fftfreq(ny, dx))121 if dx is None or wavelength is None or ni is None:122 raise TypeError("if defocus != 0, dx, wavelength, ni must be provided")123 k = 2*np.pi / wavelength * ni124 kperp = np.sqrt(np.array(k**2 - (2 * np.pi)**2 * (fx[None, :]**2 + fy[:, None]**2), dtype=np.complex))125 defocus_fn = np.exp(1j * defocus_um * kperp)126 else:127 defocus_fn = 1128 atf_defocus = atf * defocus_fn129 # if even number of frequencies, we must translate otf_c by one so that f and -f match up130 otf_c_minus_conj = np.roll(np.roll(np.flip(atf_defocus, axis=(0, 1)), np.mod(ny + 1, 2), axis=0),131 np.mod(nx + 1, 2), axis=1).conj()132 otf = scipy.signal.fftconvolve(atf_defocus, otf_c_minus_conj, mode='same') / np.sum(np.abs(atf) ** 2)133 return otf, atf_defocus134# circular aperture functions135def circ_aperture_atf(fx, fy, na, wavelength):136 """137 Amplitude transfer function for circular aperture138 @param fx:139 @param fy:140 @param na:141 @param wavelength:142 @return atf:143 """144 fmax = 0.5 / (0.5 * wavelength / na)145 # ff = np.sqrt(fx[None, :]**2 + fy[:, None]**2)146 ff = np.sqrt(fx**2 + fy**2)147 atf = np.ones(ff.shape)148 atf[ff > fmax] = 0149 return atf150def circ_aperture_otf(fx, fy, na, wavelength):151 """152 OTF for roi_size circular aperture153 :param fx:154 :param fy:155 :param na: numerical aperture156 :param wavelength: in um157 :return otf:158 """159 # maximum frequency imaging system can pass160 fmax = 1 / (0.5 * wavelength / na)161 # freq data162 fx = np.asarray(fx)163 fy = np.asarray(fy)164 ff = np.asarray(np.sqrt(fx**2 + fy**2))165 with np.errstate(invalid='ignore'):166 # compute otf167 otf = np.asarray(2 / np.pi * (np.arccos(ff / fmax) - (ff / fmax) * np.sqrt(1 - (ff / fmax)**2)))168 otf[ff > fmax] = 0169 return otf170# helper functions for converting between NA and peak widths171def na2fwhm(na, wavelength):172 """173 Convert numerical aperture to full-width at half-maximum, assuming an Airy-function PSF174 FWHM ~ 0.51 * wavelength / na175 :param na: numerical aperture176 :param wavelength:177 :return fwhm: in same units as wavelength178 """179 fwhm = 1.6163399561827614 / np.pi * wavelength / na180 return fwhm181def fwhm2na(wavelength, fwhm):182 """183 Convert full-width half-maximum PSF value to the equivalent numerical aperture. Inverse function of na2fwhm184 @param wavelength:185 @param fwhm:186 @return na:187 """188 na = 1.6163399561827614 / np.pi * wavelength / fwhm189 return na190def na2sxy(na, wavelength):191 """192 Convert numerical aperture to standard deviation, assuming the numerical aperture and the sigma193 are related as in the Airy function PSF194 :param na:195 :param wavelength:196 :return sigma:197 """198 fwhm = na2fwhm(na, wavelength)199 sigma = 1.49686886 / 1.6163399561827614 / 2 * fwhm200 # 2 * sqrt{2*log(2)} * sigma = 0.5 * wavelength / NA201 # sigma = na2fwhm(na, wavelength) / (2*np.sqrt(2 * np.log(2)))202 return sigma203def sxy2na(wavelength, sigma_xy):204 """205 Convert sigma xy value to equivalent numerical aperture, assuming these are related as in the Airy function PSF206 @param wavelength:207 @param sigma_xy:208 @return fwhm:209 """210 fwhm = 2 * 1.6163399561827614 / 1.49686886 * sigma_xy211 # fwhm = na2fwhm(na, wavelength)212 # fwhm = sigma * (2*np.sqrt(2 * np.log(2)))213 return fwhm2na(wavelength, fwhm)214def na2sz(na, wavelength, ni):215 """216 Convert numerical aperture to equivalent sigma-z value,217 @param na: numerical aperture218 @param wavelength:219 @param ni: index of refraction220 @return sz:221 """222 # todo: believe this is a gaussian approx. Find reference223 return np.sqrt(6) / np.pi * ni * wavelength / na ** 2224def sz2na(sigma_z, wavelength, ni):225 """226 Convert sigma-z value to equivalent numerical aperture227 todo: believe this is a gaussian approx. Find reference228 @param wavelength:229 @param sigma_z:230 @param ni: index of refraction231 @ return na:232 """233 return np.sqrt(np.sqrt(6) / np.pi * ni * wavelength / sigma_z)234# PSF models235class psf_model:236 """237 TODO: should I include a fit method in this class?238 """239 def __init__(self,240 param_names: list[str],241 dc: float = None,242 sf: int = 1,243 angles: tuple[float] = (0., 0., 0.),244 has_jacobian: bool = False):245 """246 PSF functions, accounting for image pixelation along an arbitrary direction247 vectorized, i.e. can rely on obeying broadcasting rules for x,y,z248 @param param_names:249 @param dc: pixel size250 @param sf: factor to oversample pixels. The value of each pixel is determined by averaging sf**2 equally spaced251 @param angles: Euler angles describing orientation of pixel to resample252 @param has_jacobian:253 """254 if not isinstance(param_names, list):255 raise ValueError("param_names must be a list of strings")256 self.parameter_names = param_names257 self.nparams = len(param_names)258 if not isinstance(sf, int):259 raise ValueError("sf must be an integer")260 self.sf = sf261 if sf != 1 and dc is None:262 raise Exception("If sf != 1, then a value for dc must be provided")263 self.dc = dc264 self.angles = angles265 self.has_jacobian = has_jacobian266 def model(self, coords: tuple[np.ndarray], params: np.ndarray):267 pass268 def jacobian(self, coords: tuple[np.ndarray], params: np.ndarray):269 pass270 def test_jacobian(self, coords: tuple[np.ndarray], params: np.ndarray, dp: float = 1e-7):271 # numerical test for jacobian272 jac_calc = self.jacobian(coords, params)273 jac_numerical = []274 for ii in range(self.nparams):275 dp_now = np.zeros(self.nparams)276 dp_now[ii] = dp * 0.5277 jac_numerical.append((self.model(coords, params + dp_now) - self.model(coords, params - dp_now)) / dp)278 return jac_numerical, jac_calc279 def estimate_parameters(self, img: np.ndarray, coords: tuple[np.ndarray]):280 pass281 def normalize_parameters(self, params):282 """283 Return parameters in a standardized format, when there can be ambiguity. For example,284 a Gaussian model may include a standard deviation parameter. Since only the square of this quantity enters285 the model, a fit may return a negative value for standard deviation. In that case, this function286 would return the absolute value of the standard deviation287 @param params:288 @return:289 """290 return params291class gaussian3d_psf_model(psf_model):292 """293 Gaussian approximation to PSF. Matches well for equal peak intensity, but some deviations in area.294 See https://doi.org/10.1364/AO.46.001819 for more details.295 sigma_xy = 0.22 * lambda / NA.296 This comes from equating the FWHM of the Gaussian and the airy function.297 FWHM = 2 * sqrt{2*log(2)} * sigma ~ 0.51 * wavelength / NA298 sigma_z = np.sqrt(6) / np.pi * ni * wavelength / NA ** 2299 """300 def __init__(self, dc: float = None, sf=1, angles=(0., 0., 0.)):301 super().__init__(["A", "cx", "cy", "cz", "sxy", "sz", "bg"],302 dc=dc, sf=sf, angles=angles, has_jacobian=True)303 def model(self, coords: tuple[np.ndarray], p: np.ndarray):304 z, y, x, = coords305 # oversample points in pixel306 xx_s, yy_s, zz_s = oversample_pixel(x, y, z, self.dc, sf=self.sf, euler_angles=self.angles)307 # calculate psf at oversampled points308 psf_s = np.exp(-(xx_s - p[1]) ** 2 / 2 / p[4] ** 2309 -(yy_s - p[2]) ** 2 / 2 / p[4] ** 2310 -(zz_s - p[3]) ** 2 / 2 / p[5] ** 2311 )312 # normalization is such that the predicted peak value of the PSF ignoring binning would be p[0]313 psf = p[0] * np.mean(psf_s, axis=-1) + p[-1]314 return psf315 def jacobian(self, coords: tuple[np.ndarray], p: np.ndarray):316 z, y, x = coords317 # oversample points in pixel318 xx_s, yy_s, zz_s = oversample_pixel(x, y, z, self.dc, sf=self.sf, euler_angles=self.angles)319 # use sxy * |sxy| instead of sxy**2 to enforce sxy > 0320 psf_s = np.exp(-(xx_s - p[1]) ** 2 / 2 / p[4]**2321 -(yy_s - p[2]) ** 2 / 2 / p[4]**2322 -(zz_s - p[3]) ** 2 / 2 / p[5]**2323 )324 # normalization is such that the predicted peak value of the PSF ignoring binning would be p[0]325 # psf = p[0] * psf_sum + p[-1]326 bcast_shape = (x + y + z).shape327 # [A, cx, cy, cz, sxy, sz, bg]328 jac = [np.mean(psf_s, axis=-1),329 p[0] * np.mean(2 * (xx_s - p[1]) / 2 / p[4]**2 * psf_s, axis=-1),330 p[0] * np.mean(2 * (yy_s - p[2]) / 2 / p[4]**2 * psf_s, axis=-1),331 p[0] * np.mean(2 * (zz_s - p[3]) / 2 / p[5]**2 * psf_s, axis=-1),332 p[0] * np.mean((2 / p[4] ** 3 * (xx_s - p[1]) ** 2 / 2 +333 2 / p[4] ** 3 * (yy_s - p[2]) ** 2 / 2) * psf_s, axis=-1),334 p[0] * np.mean( 2 / p[5] ** 3 * (zz_s - p[3]) ** 2 / 2 * psf_s, axis=-1),335 np.ones(bcast_shape)]336 return jac337 def estimate_parameters(self, img: np.ndarray, coords: tuple[np.ndarray]):338 z, y, x = coords339 # subtract smallest value so positive340 # img_temp = img - np.nanmean(img)341 # to_use = np.logical_and(np.logical_not(np.isnan(img_temp)), img_temp > 0)342 img_temp = img - np.nanmin(img)343 to_use = np.logical_and(np.logical_not(np.isnan(img_temp)), img_temp > 0)344 if img.ndim != len(coords):345 raise ValueError("len(coords) != img.ndim")346 # compute moments347 c1s = np.zeros(img.ndim)348 c2s = np.zeros(img.ndim)349 for ii in range(img.ndim):350 c1s[ii] = np.sum(img_temp[to_use] * coords[ii][to_use]) / np.sum(img_temp[to_use])351 c2s[ii] = np.sum(img_temp[to_use] * coords[ii][to_use] ** 2) / np.sum(img_temp[to_use])352 sigmas = np.sqrt(c2s - c1s ** 2)353 sxy = np.mean(sigmas[:2])354 sz = sigmas[2]355 guess_params = np.concatenate((#np.array([np.nanmax(img) - np.nanmean(img)]),356 np.array([np.nanmax(img) - np.nanmin(img)]), # may be too large ... but typically have problems with being too small357 np.flip(c1s),358 np.array([sxy]),359 np.array([sz]),360 np.array([np.nanmean(img)])361 ),362 )363 return self.normalize_parameters(guess_params)364 def normalize_parameters(self, params):365 norm_params = np.array([params[0], params[1], params[2], params[3],366 np.abs(params[4]), np.abs(params[5]),367 params[6]])368 return norm_params369class asymmetric_gaussian3d(psf_model):370 def __init__(self, dc: float = None, sf=1, angles=(0., 0., 0.)):371 super().__init__(["A", "cx", "cy", "cz", "sx", "sy/sx", "sz", "theta_xy", "bg"],372 dc=dc, sf=sf, angles=angles, has_jacobian=True)373 def model(self, coords: tuple[np.ndarray], p: np.ndarray):374 z, y, x, = coords375 # oversample points in pixel376 xx_s, yy_s, zz_s = oversample_pixel(x, y, z, self.dc, sf=self.sf, euler_angles=self.angles)377 # rotated coordinates378 xx_rot = np.cos(p[7]) * (xx_s - p[1]) - np.sin(p[7]) * (yy_s - p[2])379 yy_rot = np.cos(p[7]) * (yy_s - p[2]) + np.sin(p[7]) * (xx_s - p[1])380 # calculate psf at oversampled points381 psf_s = np.exp(-xx_rot ** 2 / 2 / p[4] ** 2382 -yy_rot ** 2 / 2 / (p[4] * p[5]) ** 2383 -(zz_s - p[3]) ** 2 / 2 / p[6] ** 2)384 # normalization is such that the predicted peak value of the PSF ignoring binning would be p[0]385 psf = p[0] * np.mean(psf_s, axis=-1) + p[8]386 return psf387 def jacobian(self, coords: tuple[np.ndarray], p: np.ndarray):388 z, y, x, = coords389 # oversample points in pixel390 xx_s, yy_s, zz_s = oversample_pixel(x, y, z, self.dc, sf=self.sf, euler_angles=self.angles)391 # rotated coordinates392 xx_rot = np.cos(p[7]) * (xx_s - p[1]) - np.sin(p[7]) * (yy_s - p[2])393 yy_rot = np.cos(p[7]) * (yy_s - p[2]) + np.sin(p[7]) * (xx_s - p[1])394 dx_dphi = -np.sin(p[7]) * (xx_s - p[1]) - np.cos(p[7]) * (yy_s - p[2])395 dy_dphi = -np.sin(p[7]) * (yy_s - p[2]) + np.cos(p[7]) * (xx_s - p[1])396 # calculate psf at oversampled points397 psf_s = np.exp(-xx_rot ** 2 / 2 / p[4] ** 2398 -yy_rot ** 2 / 2 / (p[4] * p[5]) ** 2399 -(zz_s - p[3]) ** 2 / 2 / p[6] ** 2)400 dpsf_dcx = 2 * xx_rot * np.cos(p[7]) / 2 / p[4]**2 + \401 2 * yy_rot * np.sin(p[7]) / 2 / (p[4] * p[5])**2402 dpsf_dcy = -2 * xx_rot * np.sin(p[7]) / 2 / p[4] ** 2 + \403 2 * yy_rot * np.cos(p[7]) / 2 / (p[4] * p[5]) ** 2404 dpsf_dcz = 2 * (zz_s - p[3]) / 2 / p[6] ** 2405 dpsf_dsx = 2 / p[4] ** 3 * xx_rot ** 2 / 2 + \406 2 / p[4] ** 3 * yy_rot ** 2 / 2 / p[5] ** 2407 dpsf_dsrat = 2 / p[5] ** 3 * yy_rot**2 / 2 / p[4] ** 2408 dpsf_dsz = 2 / p[6] ** 3 * (zz_s - p[3]) ** 2 / 2409 dpsf_dtheta = 2 * xx_rot / 2 / p[4] ** 2 * dx_dphi + \410 2 * yy_rot / 2 / (p[4] * p[5]) ** 2 * dy_dphi411 bcast_shape = (x + y + z).shape412 jac = [np.mean(psf_s, axis=-1), # A413 p[0] * np.mean(dpsf_dcx * psf_s, axis=-1), # cx414 p[0] * np.mean(dpsf_dcy * psf_s, axis=-1), # cy415 p[0] * np.mean(dpsf_dcz * psf_s, axis=-1), # cz416 p[0] * np.mean(dpsf_dsx * psf_s, axis=-1), # sx417 p[0] * np.mean(dpsf_dsrat * psf_s, axis=-1), # sy/sx418 p[0] * np.mean(dpsf_dsz * psf_s, axis=-1), # sz419 p[0] * np.mean(dpsf_dtheta * psf_s, axis=-1), # theta420 np.ones(bcast_shape) # bg421 ]422 return jac423 def estimate_parameters(self, img: np.ndarray, coords: tuple[np.ndarray]):424 z, y, x = coords425 # subtract smallest value so positive426 # img_temp = img - np.nanmean(img)427 img_temp = img428 to_use = np.logical_and(np.logical_not(np.isnan(img_temp)), img_temp > 0)429 if img.ndim != len(coords):430 raise ValueError("len(coords) != img.ndim")431 # compute moments432 c1s = np.zeros(img.ndim)433 c2s = np.zeros(img.ndim)434 for ii in range(img.ndim):435 c1s[ii] = np.sum(img_temp[to_use] * coords[ii][to_use]) / np.sum(img_temp[to_use])436 c2s[ii] = np.sum(img_temp[to_use] * coords[ii][to_use] ** 2) / np.sum(img_temp[to_use])437 sigmas = np.sqrt(c2s - c1s ** 2)438 guess_params = np.concatenate((np.array([np.nanmax(img) - np.nanmean(img)]),439 np.flip(c1s),440 np.flip(sigmas),441 np.array([0.]),442 np.array([np.nanmean(img)])443 ),444 )445 return self.normalize_parameters(guess_params)446 def normalize_parameters(self, params):447 norm_params = np.array([params[0], params[1], params[2], params[3],448 np.abs(params[4]), np.abs(params[5]), np.abs(params[6]),449 params[7], params[8]])450 return norm_params451class gaussian2d_psf_model(psf_model):452 """453 Gaussian approximation to PSF. Matches well for equal peak intensity, but some deviations in area.454 See https://doi.org/10.1364/AO.46.001819 for more details.455 sigma_xy = 0.22 * lambda / NA.456 This comes from equating the FWHM of the Gaussian and the airy function.457 FWHM = 2 * sqrt{2*log(2)} * sigma ~ 0.51 * wavelength / NA458 """459 def __init__(self, dc: float = None, sf=1, angles=(0., 0., 0.)):460 super().__init__(["A", "cx", "cy", "sxy", "bg"],461 dc=dc, sf=sf, angles=angles, has_jacobian=True)462 self.gauss3d = gaussian3d_psf_model(dc=dc, sf=sf, angles=angles)463 def model(self, coords: tuple[np.ndarray], params: np.ndarray):464 y, x = coords465 bcast_shape = np.broadcast_shapes(y.shape, x.shape)466 z = np.zeros(bcast_shape)467 p3d = np.array([params[0], params[1], params[2], 0., params[3], 1., params[4]])468 return self.gauss3d.model((z, y, x), p3d)469 def jacobian(self, coords: tuple[np.ndarray], params: np.ndarray):470 y, x = coords471 bcast_shape = np.broadcast_shapes(y.shape, x.shape)472 z = np.zeros(bcast_shape)473 p3d = np.array([params[0], params[1], params[2], 0., params[3], 1., params[4]])474 j3d = self.gauss3d.jacobian((z, y, x), p3d)475 j2d = j3d[:3] + [j3d[4]] + [j3d[6]]476 return j2d477 def estimate_parameters(self, img: np.ndarray, coords: tuple[np.ndarray]):478 y, x = coords479 bcast_shape = np.broadcast_shapes(y.shape, x.shape)480 z = np.zeros(bcast_shape)481 p3d = self.gauss3d.estimate_parameters(img, (z, y, x))482 p2d = np.concatenate((p3d[:3],483 np.array([p3d[4]]),484 np.array([p3d[6]])485 )486 )487 return self.normalize_parameters(p2d)488 def normalize_parameters(self, params):489 norm_params = np.array([params[0], params[1], params[2], np.abs(params[3]), params[4]])490 return norm_params491class gaussian_lorentzian_psf_model(psf_model):492 """493 Gaussian-Lorentzian PSF model. One difficulty with the Gaussian PSF is the weight is not the same in every z-plane,494 as we expect it should be. The Gaussian-Lorentzian form remedies this, but the functional form495 is no longer separable496 """497 def __init__(self, dc=None, sf=1, angles=(0., 0., 0.)):498 super().__init__(["A", "cx", "cy", "cz", "sxy", "hwhm_z", "bg"],499 dc=dc, sf=sf, angles=angles, has_jacobian=True)500 def model(self, coords, p):501 (z, y, x) = coords502 # oversample points in pixel503 xx_s, yy_s, zz_s = oversample_pixel(x, y, z, self.dc, sf=self.sf, euler_angles=self.angles)504 # calculate psf at oversampled points505 lor_factor = 1 + (zz_s - p[3]) ** 2 / p[5] ** 2506 psf_s = np.exp(- ((xx_s - p[1]) ** 2 + (yy_s - p[2]) ** 2) / (2 * p[4] ** 2 * lor_factor)) / lor_factor507 # normalization is such that the predicted peak value of the PSF ignoring binning would be p[0]508 psf = p[0] * np.mean(psf_s, axis=-1) + p[6]509 return psf510 def jacobian(self, coords, p):511 z, y, x = coords512 # oversample points in pixel513 xx_s, yy_s, zz_s = oversample_pixel(x, y, z, self.dc, sf=self.sf, euler_angles=self.angles)514 lor = 1 + (zz_s - p[3]) ** 2 / p[5] ** 2515 r_sqr = (xx_s - p[1]) ** 2 + (yy_s - p[2]) ** 2516 factor = np.exp(- r_sqr / (2 * p[4] ** 2 * lor)) / lor517 bcast_shape = (x + y + z).shape518 jac = [np.mean(factor, axis=-1), # amp519 p[0] * np.mean((x - p[1]) / p[4] ** 2 / lor * factor, axis=-1), # cx520 p[0] * np.mean((y - p[2]) / p[4] ** 2 / lor * factor, axis=-1), # cy521 p[0] * np.mean(-(z - p[3]) / p[5] ** 2 / lor ** 2 * (r_sqr / 2 / p[4] ** 2) * factor, axis=-1) +522 p[0] * np.mean(factor / lor * 2 * (z - p[3]) / p[5] ** 2, axis=-1), # cz523 p[0] * np.mean(factor * r_sqr / lor / p[4] ** 3, axis=-1), # sxy524 p[0] * np.mean(factor * r_sqr / 2 / p[4] ** 2 / lor ** 2 * 2 * (z - p[3]) ** 2 / p[5] ** 3, axis=-1) +525 p[0] * np.mean(factor / lor * 2 * (z - p[3]) ** 2 / p[5] ** 3, axis=-1), # hwhm_z526 np.ones(bcast_shape) # bg527 ]528 return jac529 def estimate_parameters(self, img: np.ndarray, coords: tuple[np.ndarray]):530 z, y, x = coords531 # subtract smallest value so positive532 img_temp = img - np.nanmean(img)533 to_use = np.logical_and(np.logical_not(np.isnan(img_temp)), img_temp > 0)534 if img.ndim != len(coords):535 raise ValueError("len(coords) != img.ndim")536 # compute moments537 c1s = np.zeros(img.ndim)538 c2s = np.zeros(img.ndim)539 for ii in range(img.ndim):540 c1s[ii] = np.sum(img_temp[to_use] * coords[ii][to_use]) / np.sum(img_temp[to_use])541 c2s[ii] = np.sum(img_temp[to_use] * coords[ii][to_use] ** 2) / np.sum(img_temp[to_use])542 sigmas = np.sqrt(c2s - c1s ** 2)543 sxy = np.mean(sigmas[:2])544 sz = sigmas[2]545 guess_params = np.concatenate((np.array([np.nanmax(img) - np.nanmean(img)]),546 np.flip(c1s),547 np.array([sxy]),548 np.array([sz]),549 np.array([np.nanmean(img)])550 ),551 )552 return self.normalize_parameters(guess_params)553 def normalize_parameters(self, params):554 norm_params = np.array([params[0], params[1], params[2], params[3],555 np.abs(params[4]), np.abs(params[5]), params[6]])556 return norm_params557class born_wolf_psf_model(psf_model):558 """559 Born-wolf PSF function evaluated using Airy function if in-focus, and axial function if along the axis.560 Otherwise evaluated using numerical integration.561 """562 def __init__(self, wavelength: float, ni: float, dc: float = None, sf=1, angles=(0., 0., 0.)):563 """564 @param wavelength:565 @param ni: refractive index566 @param dc:567 @param sf:568 @param angles:569 """570 # TODO is it better to put wavelength and ni as arguments to model or as class members?571 # advantage to being model parameters is could conceivable fit572 self.wavelength = wavelength573 self.ni = ni574 if sf != 1:575 raise NotImplementedError("Only implemented for sf=1")576 super().__init__(["A", "cx", "cy", "cz", "na", "bg"],577 dc=dc, sf=sf, angles=angles, has_jacobian=False)578 def model(self, coords: tuple[np.ndarray], p: np.ndarray):579 """580 @param coords:581 @param p:582 @param wavelength:583 @param ni:584 @return:585 """586 z, y, x = coords587 x = np.atleast_1d(x)588 y = np.atleast_1d(y)589 z = np.atleast_1d(z)590 x, y, z = np.broadcast_arrays(x, y, z)591 k = 2 * np.pi / self.wavelength592 rr = np.sqrt((x - p[1]) ** 2 + (y - p[2]) ** 2)593 psfs = np.zeros(rr.shape) * np.nan594 is_in_focus = (z == p[3])595 is_on_axis = (rr == 0)596 # ################################597 # evaluate in-focus portion using airy function, which is much faster than integrating598 # ################################599 def airy_fn(rho):600 val = p[0] * 4 * np.abs(sp.j1(rho * k * p[4]) / (rho * k * p[4])) ** 2 + p[5]601 val[rho == 0] = p[0] + p[4]602 return val603 with np.errstate(invalid="ignore"):604 psfs[is_in_focus] = airy_fn(rr[is_in_focus])605 # ################################606 # evaluate on axis portion using exact expression607 # ################################608 def axial_fn(z):609 val = p[0] * 4 * (2 * self.ni ** 2) / (k ** 2 * p[4] ** 4 * (z - p[3]) ** 2) * \610 (1 - np.cos(0.5 * k * (z - p[3]) * p[4] ** 2 / self.ni)) + p[5]611 val[z == p[3]] = p[0] + p[5]612 return val613 with np.errstate(invalid="ignore"):614 psfs[is_on_axis] = axial_fn(z[is_on_axis])615 # ################################616 # evaluate out of focus portion using integral617 # ################################618 if not np.all(is_in_focus):619 def integrand(rho, r, z):620 return rho * sp.j0(k * r * p[4] * rho) * np.exp(-1j * k * (z - p[3]) * p[4] ** 2 * rho ** 2 / (2 * self.ni))621 # like this approach because allows rr, z, etc. to have arbitrary dimension622 already_evaluated = np.logical_or(is_in_focus, is_in_focus)623 for ii, (r, zc, already_eval) in enumerate(zip(rr.ravel(), z.ravel(), already_evaluated.ravel())):624 if already_eval:625 continue626 int_real = scipy.integrate.quad(lambda rho: integrand(rho, r, zc).real, 0, 1)[0]627 int_img = scipy.integrate.quad(lambda rho: integrand(rho, r, zc).imag, 0, 1)[0]628 coords = np.unravel_index(ii, rr.shape)629 psfs[coords] = p[0] * 4 * (int_real ** 2 + int_img ** 2) + p[5]630 return psfs631 def estimate_parameters(self, img: np.ndarray, coords: tuple[np.ndarray]):632 gauss3d = gaussian3d_psf_model(dc=self.dc, sf=self.sf, angles=self.angles)633 p3d_gauss = gauss3d.estimate_parameters(img, coords)634 na = sxy2na(self.wavelength, p3d_gauss[4])635 params_guess = np.concatenate((p3d_gauss[:4],636 np.array([na]),637 np.array([p3d_gauss[-1]])638 )639 )640 return params_guess641class model(psf_model):642 """643 Wrapper function for evaluating different PSF models where only the coordinate grid information is given644 (i.e. nx and dxy) and not the actual coordinates. The real coordinates can be obtained using get_psf_coords().645 This class primarily exists to wrap the psfmodel functions, but646 The size of these functions is parameterized by the numerical aperture, instead of the sigma or other size647 parameters that are only convenient for some specific models of the PSF648 For vectorial or gibson-lanni PSF's, this wraps the functions in the psfmodels package649 (https://pypi.org/project/psfmodels/) with added ability to shift the center of these functions away from the650 center of the ROI, which is useful when fitting them.651 For 'gaussian', it wraps the gaussian3d_pixelated_psf() function. More details about the relationship between652 the Gaussian sigma and the numerical aperture can be found here: https://doi.org/10.1364/AO.46.001819653 """654 def __init__(self, wavelength, ni, model_name="vectorial", dc: float = None, sf=1, angles=(0., 0., 0.)):655 """656 @param wavelength:657 @param ni: index of refraction658 @param model_name: 'gaussian', 'gibson-lanni', 'born-wolf', or 'vectorial'. 'gibson-lanni' relies on the659 psdmodels function scalar_psf(), while 'vectorial' relies on the psfmodels function vectorial_psf()660 @param dc:661 @param sf:662 @param angles:663 """664 self.wavelength = wavelength665 self.ni = ni666 allowed_models = ["gibson-lanni", "vectorial", "born-wolf", "gaussian"]667 if model_name not in allowed_models:668 raise ValueError(f"model={model_name:s} was not an allowed value. Allowed values are {allowed_models}")669 if not psfmodels_available and (model_name == "vectorial" or model_name == "gibson-lanni"):670 raise NotImplementedError(f"model={model_name:s} selected but psfmodels is not installed")671 self.model_name = model_name672 if sf != 1:673 raise NotImplementedError("not yet implemented for sf=/=1")674 super().__init__(["A", "cx", "cy", "cz", "na", "bg"],675 dc=dc, sf=sf, angles=angles, has_jacobian=False)676 def model(self, coords: tuple[np.ndarray], p: np.ndarray, **kwargs):677 """678 Unlike other model functions this ONLY works if coords are of the same style as obtained from679 get_psf_coords()680 @param coords: (z, y, x). Coordinates must be exactly as obtained from get_psf_coords() with681 nx=ny and dx=dy.682 @param p:683 @param kwargs: keywords passed through to vectorial_psf() or scalar_psf(). Note that in most cases684 these will shift the best focus PSF away from z=0685 @return:686 """687 zz, y, x = coords688 z = zz[:, 0, 0]689 dxy = x[0, 0, 1] - x[0, 0, 0]690 nx = x.shape[2]691 if 'NA' in kwargs.keys():692 raise ValueError("'NA' is not allowed to be passed as a named parameter. It is specified in p.")693 if self.model_name == 'vectorial':694 model_params = {'NA': p[4], 'sf': self.sf, 'ni': self.ni, 'ni0': self.ni}695 model_params.update(kwargs)696 psf_norm = psfm.vectorial_psf(0, 1, dxy, wvl=self.wavelength, params=model_params, normalize=False)697 val = psfm.vectorial_psf(z - p[3], nx, dxy, wvl=self.wavelength, params=model_params, normalize=False)698 val = p[0] / psf_norm * ndi.shift(val, [0, p[2] / dxy, p[1] / dxy], mode='nearest') + p[5]699 elif self.model_name == 'gibson-lanni':700 model_params = {'NA': p[4], 'sf': self.sf, 'ni': self.ni, 'ni0': self.ni}701 model_params.update(kwargs)702 psf_norm = psfm.scalar_psf(0, 1, dxy, wvl=self.wavelength, params=model_params, normalize=False)703 val = psfm.scalar_psf(z - p[3], nx, dxy, wvl=self.wavelength, params=model_params, normalize=False)704 val = p[0] / psf_norm * ndi.shift(val, [0, p[2] / dxy, p[1] / dxy], mode='nearest') + p[5]705 elif self.model_name == "born-wolf":706 bw_model = born_wolf_psf_model(wavelength=self.wavelength, ni=self.ni, dc=self.dc, sf=self.sf, angles=self.angles)707 val = bw_model.model(coords, p)708 elif self.model_name == 'gaussian':709 # transform NA to sigmas710 p_gauss = [p[0], p[1], p[2], p[3], 0.22 * self.wavelength / p[4],711 np.sqrt(6) / np.pi * self.ni * self.wavelength / p[4] ** 2,712 p[5]]713 gauss_model = gaussian3d_psf_model(dc=self.dc, sf=self.sf, angles=self.angles)714 # normalize so that amplitude parameter is actually amplitude715 psf_norm = gauss_model.model((p[3], p[2], p[1]), p_gauss) - p[5]716 val = p[0] / psf_norm * (gauss_model.model((z, y, x), p_gauss) - p[5]) + p[5]717 else:718 raise ValueError(f"model_name was '{self.model_name:s}',"719 f" but must be 'vectorial', 'gibson-lanni', 'born-wolf', or 'gaussian'")720 return val721 def estimate_parameters(self, img: np.ndarray, coords: tuple[np.ndarray]):722 gauss3d = gaussian3d_psf_model(dc=self.dc, sf=self.sf, angles=self.angles)723 p3d_gauss = gauss3d.estimate_parameters(img, coords)724 na = sxy2na(self.wavelength, p3d_gauss[4])725 params_guess = np.concatenate((p3d_gauss[:4],726 np.array([na]),727 np.array([p3d_gauss[-1]])728 )729 )730 return params_guess731# utility functions732def oversample_pixel(x, y, z, ds, sf, euler_angles=(0., 0., 0.)):733 """734 Generate coordinates to oversample a pixel on a 2D grid.735 Suppose we have a set of pixels centered at points given by x, y, z. Generate sf**2 points in this pixel equally736 spaced about the center. Allow the pixel to be orientated in an arbitrary direction with respect to the coordinate737 system. The pixel rotation is described by the Euler angles (psi, theta, phi), where the pixel "body" frame738 is a square with xy axis orientated along the legs of the square with z-normal to the square739 Compare with oversample_voxel() which works on a 3D grid with a fixed orientation. oversample_pixel() works740 on a 2D grid with an arbitrary orientation.741 :param x: x-coordinate with shape such that can be broadcast with y and z. e.g. z.shape = [nz, 1, 1];742 y.shape = [1, ny, 1]; x.shape = [1, 1, nx]743 :param y:744 :param z:745 :param ds: pixel size746 :param sf: sample factor747 :param euler_angles: [phi, theta, psi] where phi and theta are the polar angles describing the normal of the pixel,748 and psi describes the rotation of the pixel about its normal749 :return xx_s, yy_s, zz_s:750 """751 # generate new points in pixel, each of which is centered about an equal area of the pixel, so summing them is752 # giving an approximation of the integral753 if sf > 1:754 pts = np.arange(1 / (2*sf), 1 - 1 / (2*sf), 1 / sf) - 0.5755 xp, yp = np.meshgrid(ds * pts, ds * pts)756 zp = np.zeros(xp.shape)757 # rotate points to correct position using normal vector758 # for now we will fix x, but lose generality759 mat = affine.euler_mat(*euler_angles)760 result = mat.dot(np.concatenate((xp.ravel()[None, :],761 yp.ravel()[None, :],762 zp.ravel()[None, :]), axis=0))763 xs, ys, zs = result764 # now must add these to each point x, y, z765 xx_s = x[..., None] + xs[None, ...]766 yy_s = y[..., None] + ys[None, ...]767 zz_s = z[..., None] + zs[None, ...]768 else:769 xx_s = np.expand_dims(x, axis=-1)770 yy_s = np.expand_dims(y, axis=-1)771 zz_s = np.expand_dims(z, axis=-1)772 return xx_s, yy_s, zz_s773def oversample_voxel(coords, drs, sf=3):774 """775 Get coordinates to oversample a voxel on a 3D grid776 Compare with oversample_pixel(), which performs oversampling on a 2D grid.777 :param coords: tuple of coordinates, e.g. (z, y, x)778 :param drs: tuple giving voxel size (dz, dy, dx)779 :param sf: sampling factor. Assumed to be same for all directions780 :return coords_upsample: tuple of coordinates, e.g. (z_os, y_os, x_os). e.g. x_os has one more dimension than x781 with this extra dimension giving the oversampled coordinates782 """783 pts = np.arange(1 / (2 * sf), 1 - 1 / (2 * sf), 1 / sf) - 0.5784 pts_dims = np.meshgrid(*[pts * dr for dr in drs], indexing="ij")785 coords_upsample = [np.expand_dims(c, axis=-1) + np.expand_dims(np.ravel(r), axis=0)786 for c, r in zip(coords, pts_dims)]787 # now must add these to each point x, y, z788 return coords_upsample789def get_psf_coords(ns, drs, broadast=False):790 """791 Get centered coordinates for PSFmodels style PSF's from step size and number of coordinates792 :param ns: list of number of points793 :param drs: list of step sizes794 :return coords: list of coordinates [zs, ys, xs, ...]795 """796 ndims = len(drs)797 coords = [np.expand_dims(d * (np.arange(n) - n // 2),798 axis=tuple(range(ii)) + tuple(range(ii+1, ndims)))799 for ii, (n, d) in enumerate(zip(ns, drs))]800 if broadast:801 # return arrays instead of views coords802 coords = [np.array(c, copy=True) for c in np.broadcast_arrays(*coords)]803 return coords804def average_exp_psfs(imgs, coords, centers, roi_sizes, backgrounds=None):805 """806 Get experimental psf from imgs by averaging many localizations (after pixel shifting).807 :param imgs:z-stack of images808 :param coords: (z, y, x) of full image. Must be broadcastable to full image size809 :param centers: n x 3 array, (cz, cy, cx)810 :param roi_sizes: [sz, sy, sx]811 :param backgrounds: values to subtracted from each ROI812 :return psf_mean, psf_coords, otf_mean, otf_coords:813 """814 # if np.any(np.mod(np.array(roi_sizes), 2) == 0):815 # raise ValueError("roi_sizes must be odd")816 z, y, x, = coords817 dz = z[1, 0, 0] - z[0, 0, 0]818 dy = y[0, 1, 0] - y[0, 0, 0]819 dx = x[0, 0, 1] - x[0, 0, 0]820 # set up array to hold psfs821 nrois = len(centers)822 if backgrounds is None:823 backgrounds = np.zeros(nrois)824 psf_shifted = np.zeros((nrois, roi_sizes[0], roi_sizes[1], roi_sizes[2])) * np.nan825 # coordinates826 z_psf, y_psf, x_psf = get_psf_coords(roi_sizes, [dz, dy, dx], broadast=True)827 zc_pix_psf = np.argmin(np.abs(z_psf[:, 0, 0]))828 yc_pix_psf = np.argmin(np.abs(y_psf[0, :, 0]))829 xc_pix_psf = np.argmin(np.abs(x_psf[0, 0, :]))830 # loop over rois and shift psfs so they are centered831 for ii in range(nrois):832 # get closest pixels to center833 xc_pix = np.argmin(np.abs(x - centers[ii, 2]))834 yc_pix = np.argmin(np.abs(y - centers[ii, 1]))835 zc_pix = np.argmin(np.abs(z - centers[ii, 0]))836 # cut roi from image837 roi_unc = rois.get_centered_roi((zc_pix, yc_pix, xc_pix), roi_sizes)838 roi = rois.get_centered_roi((zc_pix, yc_pix, xc_pix), roi_sizes, min_vals=[0, 0, 0], max_vals=imgs.shape)839 img_roi = rois.cut_roi(roi, imgs)840 zroi = rois.cut_roi(roi, z)841 yroi = rois.cut_roi(roi, y)842 xroi = rois.cut_roi(roi, x)843 cx_pix_roi = (roi[5] - roi[4]) // 2844 cy_pix_roi = (roi[3] - roi[2]) // 2845 cz_pix_roi = (roi[1] - roi[0]) // 2846 xshift_pix = (xroi[0, 0, cx_pix_roi] - centers[ii, 2]) / dx847 yshift_pix = (yroi[0, cy_pix_roi, 0] - centers[ii, 1]) / dy848 zshift_pix = (zroi[cz_pix_roi, 0, 0] - centers[ii, 0]) / dz849 # get coordinates850 img_roi_shifted = ndi.shift(np.array(img_roi, dtype=float), [zshift_pix, yshift_pix, xshift_pix],851 mode="constant", cval=-1)852 img_roi_shifted[img_roi_shifted == -1] = np.nan853 # put into array in appropriate positions854 zstart = zc_pix_psf - cz_pix_roi855 zend = zstart + (roi[1] - roi[0])856 ystart = yc_pix_psf - cy_pix_roi857 yend = ystart + (roi[3] - roi[2])858 xstart = xc_pix_psf - cx_pix_roi859 xend = xstart + (roi[5] - roi[4])860 psf_shifted[ii, zstart:zend, ystart:yend, xstart:xend] = img_roi_shifted - backgrounds[ii]861 with warnings.catch_warnings():862 warnings.simplefilter("ignore", category=RuntimeWarning)863 with np.errstate(divide='ignore', invalid='ignore'):864 psf_mean = np.nanmean(psf_shifted, axis=0)865 # the above doesn't do a good enough job of normalizing PSF866 max_val = np.nanmax(psf_mean[psf_mean.shape[0]//2])867 psf_mean = psf_mean / max_val868 # get otf869 otf_mean, ks = psf2otf(psf_mean, drs=(dz, dy, dx))870 kz, ky, kx = np.meshgrid(*ks, indexing="ij")...
analysis_funcs.py
Source:analysis_funcs.py
1import json2import math3import datetime4import os5import fnmatch6def find(pattern, path):7 result = []8 for root, dirs, files in os.walk(path):9 for name in files:10 if fnmatch.fnmatch(name, pattern):11 result.append(os.path.join(root, name))12 return result13def relative_difference (a,b):14 try:15 return 100*(a-b)/(0.5*(a+b))16 except:17 return 018def group_results(resultsFilePath, nPeaks, dateSince = datetime.datetime(2001,1,1), dateBefore = datetime.datetime(2030,1,1)):19 rawResultsList = []20 with open(resultsFilePath, "r") as resultsFile:21 for line in resultsFile:22 rawResultsList.append(json.loads(line))23 results = []24 for i in rawResultsList:25 peak_list = i["peak_list"]26 precision = i["precision"]27 parameters = {"seed": i["seed"],\28 "period": i["period"],\29 "width": i["width"],\30 "snrpeak": i["snrpeak"],\31 "dm": i["dm"],\32 "nbits": i["nbits"],\33 "nchans": i["nchans"],\34 "tsamp": i["tsamp"],\35 "tobs": i["tobs"],\36 "fch1": i["fch1"],\37 "foff": i["foff"],\38 "binary": i["binary"],\39 "bper": i["bper"],\40 "bphase": i["bphase"],\41 "bpmass": i["bpmass"],\42 "bcmass": i["bcmass"],\43 }44 aaOutputDatFileName = i["aaOutputDatFileName"]45 date = i["date"]46 checkDate = datetime.datetime.strptime(date, "%d-%m-%Y-%H-%M-%S.%f")47 if (checkDate > dateSince) and (checkDate < dateBefore):48 results.append({"peak_list": peak_list, "precision":precision, "parameters":parameters, "date":date, "aaOutputDatFileName":aaOutputDatFileName})49 grouped_results = []50 for result in results:51 result_group = {}52 #print(result)53 if result["precision"] == "bfloat":54 result_group["bfloat_peaks"] = json.loads(result["peak_list"])[0:nPeaks]55 result_group["parameters"] = result["parameters"]56 result_group["aaOutputDatFileName"] = result["aaOutputDatFileName"]57 result_group["date"] = result["date"]58 for candidate in results:59 if candidate["parameters"] == result["parameters"]:60 if candidate["precision"] == "single":61 result_group["single_peaks"] = json.loads(candidate["peak_list"])[0:nPeaks]62 if candidate["precision"] == "double":63 result_group["double_peaks"] = json.loads(candidate["peak_list"])[0:nPeaks]64 if candidate["precision"] == "PRESTOsummed":65 result_group["PRESTOsummed_peaks"] = json.loads(candidate["peak_list"])[0:nPeaks]66 if candidate["precision"] == "PRESTOcoherent":67 result_group["PRESTOcoherent_peaks"] = json.loads(candidate["peak_list"])[0:nPeaks]68 if candidate["precision"] == "PRESTOsigma":69 result_group["PRESTOsigma_peaks"] = json.loads(candidate["peak_list"])[0:nPeaks]70 results.remove(result)71 grouped_results.append(result_group)72 #print("Returning grouped_results: " + str(grouped_results))73 return grouped_results74def extract_diff_hist_save_params(grouped_results, harmonic, index):75 b2s_hist = []76 s2d_hist = []77 b2d_hist = []78 for group in grouped_results:79 #print("group: " + str(group))80 if datetime.datetime.strptime(group["date"], '%d-%m-%Y-%H-%M-%S.%f') > datetime.datetime(2021,8,20):81 try:82 bfloat_peak = group["bfloat_peaks"][harmonic][index]83 single_peak = group["single_peaks"][harmonic][index]84 double_peak = group["double_peaks"][harmonic][index]85 if (index != 2) or ((index == 2) and (bfloat_peak != 0.0) and (single_peak != 0.0) and (double_peak != 0.0)):86 b2s_diff = relative_difference(bfloat_peak,single_peak)87 s2d_diff = relative_difference(single_peak,double_peak)88 b2d_diff = relative_difference(bfloat_peak,double_peak)89 b2s_hist.append({'diff': b2s_diff,'parameters': group["parameters"], "bfloat_peak" : bfloat_peak, "single_peak" : single_peak, "double_peak" : double_peak})90 s2d_hist.append({'diff': s2d_diff,'parameters': group["parameters"], "bfloat_peak" : bfloat_peak, "single_peak" : single_peak, "double_peak" : double_peak})91 b2d_hist.append({'diff': b2d_diff,'parameters': group["parameters"], "bfloat_peak" : bfloat_peak, "single_peak" : single_peak, "double_peak" : double_peak})92 except:93 print("nothing to add to histogram, dump group:")94 #print(group)95 #print("\n\n\n")96 return b2s_hist, s2d_hist, b2d_hist97def extract_diff_hist(grouped_results, harmonic, index):98 b2s_hist = []99 s2d_hist = []100 b2d_hist = []101 for group in grouped_results:102 try:103 bfloat_peak = group["bfloat_peaks"][harmonic][index]104 single_peak = group["single_peaks"][harmonic][index]105 double_peak = group["double_peaks"][harmonic][index]106 if (index != 2) or ((index == 2) and (bfloat_peak != 0.0) and (single_peak != 0.0) and (double_peak != 0.0)):107 b2s_diff = relative_difference(bfloat_peak,single_peak)108 s2d_diff = relative_difference(single_peak,double_peak)109 b2d_diff = relative_difference(bfloat_peak,double_peak)110 #print(str(b2s_diff) +"\t"+ str(s2d_diff) +"\t"+ str(b2d_diff))111 b2s_hist.append(b2s_diff)112 s2d_hist.append(s2d_diff)113 b2d_hist.append(b2d_diff)114 except:115 #pass116 print("nothing to add to histogram, dump group:")117 print(group)118 print("\n\n\n")119 return b2s_hist, s2d_hist, b2d_hist120def extract_diff_hist_freq_bin(grouped_results, harmonic, index):121 b2s_hist = []122 s2d_hist = []123 b2d_hist = []124 for group in grouped_results:125 try:126 bfloat_peak = group["bfloat_peaks"][harmonic][index]127 single_peak = group["single_peaks"][harmonic][index]128 double_peak = group["double_peaks"][harmonic][index]129 if (index != 2) or ((index == 2) and (bfloat_peak != 0.0) and (single_peak != 0.0) and (double_peak != 0.0)):130 b2s_diff = bfloat_peak-single_peak131 s2d_diff = single_peak-double_peak132 b2d_diff = bfloat_peak-double_peak133 #print(str(b2s_diff) +"\t"+ str(s2d_diff) +"\t"+ str(b2d_diff))134 b2s_hist.append(b2s_diff)135 s2d_hist.append(s2d_diff)136 b2d_hist.append(b2d_diff)137 except:138 #pass139 print("nothing to add to histogram, dump group:")140 print(group)141 print("\n\n\n")142 return b2s_hist, s2d_hist, b2d_hist143def extract_diff_hist_freq_bin_3d(grouped_results, harmonic, index):144 b2s_hist = [[],[],[]]145 s2d_hist = [[],[],[]]146 b2d_hist = [[],[],[]]147 for group in grouped_results:148 try:149 bfloat_peak = group["bfloat_peaks"][harmonic][index]150 single_peak = group["single_peaks"][harmonic][index]151 double_peak = group["double_peaks"][harmonic][index]152 if (index != 2) or ((index == 2) and (bfloat_peak != 0.0) and (single_peak != 0.0) and (double_peak != 0.0)):153 b2s_diff = bfloat_peak-single_peak154 s2d_diff = single_peak-double_peak155 b2d_diff = bfloat_peak-double_peak156 #if (b2s_diff != 0.0) and (s2d_peak != 0.0) and (b2d_peak != 0.0):157 if (b2s_diff != 0.0):158 #print(str(b2s_diff) +"\t"+ str(s2d_diff) +"\t"+ str(b2d_diff))159 b2s_hist[0].append(b2s_diff)160 s2d_hist[0].append(s2d_diff)161 b2d_hist[0].append(b2d_diff)162 b2s_hist[1].append(group["bfloat_peaks"][harmonic][2])163 s2d_hist[1].append(group["single_peaks"][harmonic][2])164 b2d_hist[1].append(group["double_peaks"][harmonic][2])165 b2s_hist[2].append(group['parameters']['period'])166 s2d_hist[2].append(group['parameters']['period'])167 b2d_hist[2].append(group['parameters']['period'])168 except:169 #pass170 print("nothing to add to histogram, dump group:")171 print(group)172 print("\n\n\n")173 return b2s_hist, s2d_hist, b2d_hist174def extract_zero_diff_hist_freq_bin_3d(grouped_results, harmonic, index):175 b2s_hist = [[],[],[]]176 s2d_hist = [[],[],[]]177 b2d_hist = [[],[],[]]178 for group in grouped_results:179 try:180 bfloat_peak = group["bfloat_peaks"][harmonic][index]181 single_peak = group["single_peaks"][harmonic][index]182 double_peak = group["double_peaks"][harmonic][index]183 if (index != 2) or ((index == 2) and (bfloat_peak != 0.0) and (single_peak != 0.0) and (double_peak != 0.0)):184 b2s_diff = bfloat_peak-single_peak185 s2d_diff = single_peak-double_peak186 b2d_diff = bfloat_peak-double_peak187 #if (b2s_diff != 0.0) and (s2d_peak != 0.0) and (b2d_peak != 0.0):188 if (b2s_diff == 0.0):189 #print(str(b2s_diff) +"\t"+ str(s2d_diff) +"\t"+ str(b2d_diff))190 b2s_hist[0].append(b2s_diff)191 s2d_hist[0].append(s2d_diff)192 b2d_hist[0].append(b2d_diff)193 b2s_hist[1].append(group["bfloat_peaks"][harmonic][2])194 s2d_hist[1].append(group["single_peaks"][harmonic][2])195 b2d_hist[1].append(group["double_peaks"][harmonic][2])196 b2s_hist[2].append(group["parameters"])197 s2d_hist[2].append(group["parameters"])198 b2d_hist[2].append(group["parameters"])199 except:200 #pass201 print("nothing to add to histogram, dump group:")202 print(group)203 print("\n\n\n")204 return b2s_hist, s2d_hist, b2d_hist205def extract_thresh_diff_hist_freq_bin_3d(grouped_results, harmonic, index,thresh = 2):206 b2s_hist = [[],[]]207 s2d_hist = [[],[]]208 b2d_hist = [[],[]]209 for group in grouped_results:210 try:211 bfloat_peak = group["bfloat_peaks"][harmonic][index]212 single_peak = group["single_peaks"][harmonic][index]213 double_peak = group["double_peaks"][harmonic][index]214 if (index != 2) or ((index == 2) and (bfloat_peak != 0.0) and (single_peak != 0.0) and (double_peak != 0.0)):215 b2s_diff = bfloat_peak-single_peak216 s2d_diff = single_peak-double_peak217 b2d_diff = bfloat_peak-double_peak218 #if (b2s_diff != 0.0) and (s2d_peak != 0.0) and (b2d_peak != 0.0):219 if (b2s_diff > 2):220 #print(str(b2s_diff) +"\t"+ str(s2d_diff) +"\t"+ str(b2d_diff))221 b2s_hist[0].append(b2s_diff)222 s2d_hist[0].append(s2d_diff)223 b2d_hist[0].append(b2d_diff)224 group["parameters"]["snrpeak"] = 0.0225 group["parameters"]["seed"] = int(group["parameters"]["seed"])226 b2s_hist[1].append(group["parameters"])227 s2d_hist[1].append(group["parameters"])228 b2d_hist[1].append(group["parameters"])229 except:230 #pass231 print("nothing to add to histogram, dump group:")232 print(group)233 print("\n\n\n")234 return b2s_hist, s2d_hist, b2d_hist235def extract_diff_hist_perc_diff_b2s(grouped_results, harmonic, index):236 b2s_hist = []237 for group in grouped_results:238 try:239 bfloat_peak = group["bfloat_peaks"][harmonic][index]240 single_peak = group["single_peaks"][harmonic][index]241 b2s_diff = abs(relative_difference(bfloat_peak,single_peak))242 b2s_hist.append([b2s_diff, bfloat_peak])243 except:244 print("nothing to add to histogram, dump group:")245 print(group)246 print("\n\n\n")247 return b2s_hist248def extract_bounded_diff_hist(grouped_results, harmonic, index, lb, ub):249 b2s_hist = []250 s2d_hist = []251 b2d_hist = []252 for group in grouped_results:253 try:254 bfloat_peak = group["bfloat_peaks"][harmonic][index]255 single_peak = group["single_peaks"][harmonic][index]256 double_peak = group["double_peaks"][harmonic][index]257 if (index != 2) or ((index == 2) and (bfloat_peak != 0.0) and (single_peak != 0.0) and (double_peak != 0.0)):258 b2s_diff = relative_difference(bfloat_peak,single_peak)259 s2d_diff = relative_difference(single_peak,double_peak)260 b2d_diff = relative_difference(bfloat_peak,double_peak)261 #print(str(b2s_diff) +"\t"+ str(s2d_diff) +"\t"+ str(b2d_diff))262 #print("Single peak: " + str(single_peak))263 if (single_peak > lb) and (single_peak < ub):264 b2s_hist.append(b2s_diff)265 s2d_hist.append(s2d_diff)266 b2d_hist.append(b2d_diff)267 except:268 print("nothing to add to histogram, dump group:")269 print(group)270 print("\n\n\n")271 return b2s_hist, s2d_hist, b2d_hist272def rounddown(number, nearest):273 if number < 0:274 return nearest*(-1)*math.ceil(abs(number/nearest)) - nearest275 elif number > 0:276 return nearest*math.floor(number/nearest) - nearest277 else:278 return (-1)*nearest279def roundup(number, nearest):280 if number < 0:281 return nearest*(-1)*math.floor(abs(number/nearest)) + nearest282 elif number > 0:283 return nearest*math.ceil(number/nearest) + nearest284 else:285 return nearest286def params_from_filename(filename, extract_date=True):287 #print("extracing params from " + filename)288 params = {}289 filename = filename.split("/")[-1]290 params_list = filename.split('_')291 i = 0292 for i in range(0, len(params_list), 2):293 try:294 params[params_list[i]] = float(params_list[i+1])295 except:296 pass297 #print("Reached end of string")298 if extract_date:299 try:300 date = datetime.datetime.strptime(params_list[i], '%d-%m-%Y-%H-%M-%S.%f')301 params["date"] = params_list[i]302 #print(date)303 except:304 #print("params_list: " + str(params_list) + " params_list[i]:" +params_list[i] + " is not datetime")305 pass306 for element in params_list:307 if element == "bfloat":308 params["precision"] = "bfloat"309 if element == "single":310 params["precision"] = "single"311 if element == "double":312 params["precision"] = "double"313 if element == "PRESTO":314 params["precision"] = "PRESTO"315 #print(element)316 #print("Returning params: " + str(params))317 return params318#add the bfloat, single, double peak if all parameters EXCEPT target are default values319def extract_1D_hist(grouped_results, harmonic, index, target_axis):320 b2s_hist = []321 s2d_hist = []322 b2d_hist = []323 for group in grouped_results:324 try:325 print(group)326 bfloat_peak = group["bfloat_peaks"][harmonic][index]327 single_peak = group["single_peaks"][harmonic][index]328 double_peak = group["double_peaks"][harmonic][index]329 b2s_diff = relative_difference(bfloat_peak,single_peak)330 s2d_diff = relative_difference(single_peak,double_peak)331 b2d_diff = relative_difference(bfloat_peak,double_peak)332 if check_1d(group["parameters"], target_axis):333 b2s_hist.append(b2s_diff)334 print("appending: " + str(b2s_diff))335 s2d_hist.append(s2d_diff)336 print("appending: " + str(s2d_diff))337 b2d_hist.append(b2d_diff)338 print("appending: " + str(b2d_diff))339 except:340 print("nothing to add to histogram, dump group:")341 print(group)342 #print("\n\n\n")343# pass344 return b2s_hist, s2d_hist, b2d_hist345#check all params are default, return true if the specified axis is different346def check_1d(params, target_axis):347 is_on_axis = True348 defaultParameters = {'period': [10],\349 'width':[25],\350 'snrpeak':[0.25],\351 'dm':[100],\352 'nbits':[8],\353 'nchans':[1024],\354 'tsamp':[128],\355 'tobs':[600],\356 'fch1':[1550],\357 'foff':[0.292968752],\358 'binary':[float('nan')],\359 'bper':[15],\360 'bphase':[0.2],\361 'bpmass':[1.25],\362 'bcmass':[2.0]}363 for param in params:364 if param != target_axis:365 if (not math.isnan(params[param])) and (params[param] != defaultParameters[param][0]):366 is_on_axis = False367 #print("FLAG1")368 #print(str(params[param]) + " is not equal to " + str(defaultParameters[param][0]))369 if param == target_axis:370 if params[param] == defaultParameters[param]:371 is_on_axis = False372 #print("FLAG2")373 print("is on axis: " + str(is_on_axis))...
simplify_return.py
Source:simplify_return.py
...20 if index >= len(nums):21 return default22 else:23 return nums[index]24 def is_on_axis(position: tuple[int, int]) -> bool:25 match position:26 case (0, _) | (_, 0):27 return True28 case _:29 return False30 ```31 Good:32 ```33 def index_or_default(nums: list[Any], index: int, default: Any):34 if index >= len(nums):35 return default36 return nums[index]37 def is_on_axis(position: tuple[int, int]) -> bool:38 match position:39 case (0, _) | (_, 0):40 return True41 return False42 ```43 """44 code = 12645def get_trailing_return(node: Statement) -> Statement | None:46 match node:47 case ReturnStmt(expr=Expression()):48 return node49 case MatchStmt(50 bodies=[*bodies, Block(body=[stmt])],51 patterns=[*_, AsPattern(pattern=None)],...
Automation Testing Tutorials
Learn to execute automation testing from scratch with LambdaTest Learning Hub. Right from setting up the prerequisites to run your first automation test, to following best practices and diving deeper into advanced test scenarios. LambdaTest Learning Hubs compile a list of step-by-step guides to help you be proficient with different test automation frameworks i.e. Selenium, Cypress, TestNG etc.
LambdaTest Learning Hubs:
- JUnit Tutorial
- TestNG Tutorial
- Webdriver Tutorial
- WebDriverIO Tutorial
- Protractor Tutorial
- Selenium 4 Tutorial
- Jenkins Tutorial
- NUnit Tutorial
- Jest Tutorial
- Playwright Tutorial
- Cypress Tutorial
- PyTest Tutorial
YouTube
You could also refer to video tutorials over LambdaTest YouTube channel to get step by step demonstration from industry experts.
Run refurb automation tests on LambdaTest cloud grid
Perform automation testing on 3000+ real desktop and mobile devices online.
Try LambdaTest Now !!
Get 100 minutes of automation test minutes FREE!!
Was this article helpful?
