OSCR

Aberration-aware 3D localization microscopy via self-supervised neural-physics learning.

Code ↔ Paper

16 matches between paragraphs of the paper and lines of its authors' code, computed by the harvester (lexical-v1). Click a colored paragraph or line to see its counterpart.

The 16 matches
  1. [1] § Methods › Training data generation ↔ ailoc/simulation/camera.py, lines 51–95 · score 0.74 · readout noise, electron multiplication, EMCCD camera, Poisson, simulation
  2. [2] § Methods › LUNAR architecture › Network encoder ↔ ailoc/lunar/transformer.py, lines 357–424 · score 0.69 · Transformer block, spatial dimension, aggregate, dropout, layer, masked
  3. [3] § Methods › Synchronized learning strategy ↔ ailoc/decode/loss.py, lines 35–76 · score 0.65 · Gaussian mixture model, log likelihood, GMM, loss
  4. [4] § Methods › Synchronized learning strategy ↔ ailoc/lunar/loss.py, lines 35–76 · score 0.65 · Gaussian mixture model, log likelihood, GMM, loss
  5. [5] § Results › Principle of LUNAR ↔ ailoc/lunar/network.py, lines 204–264 · score 0.62 · temporal attention module, feature extraction module, FEM, TAM, Transformer, LUNAR
  6. [6] § Methods › Microscope ↔ src/main/java/de/embl/rieslab/htsmlm/MainFramehtSMLM.java, lines 243–310 · score 0.58 · filter wheel, focus lock, SMART, beam, quadrant, TOPTICA
  7. [7] § Methods › Training data generation ↔ ailoc/common/notebook_gui.py, lines 727–804 · score 0.58 · readout noise, sCMOS, EMCCD, camera
  8. [8] § Methods › Sample preparation › Nuclear pore complex labeling ↔ usages/pyscripts/lunar_usage.py, lines 24–160 · score 0.57 · U2OS Nup96 SNAP, AF647, BG, blocked, cells
  9. [9] § Methods › LUNAR architecture › Network encoder ↔ ailoc/lunar/network.py, lines 204–264 · score 0.57 · Transformer block, dropout, FEM, TAM, layer, temporal
  10. [10] § Methods › Derivation of general objective ↔ ailoc/simulation/vectorpsf.py, lines 1514–1557 · score 0.56 · neural network, Zernike coefficients, PSF model, Derivation, optimize, objective
  11. [11] § Methods › Sample preparation › Cell culture and induced neuronal cells generation ↔ usages/pyscripts/lunar_usage.py, lines 24–160 · score 0.55 · Nup96 SNAP, U2OS, mm, cells
  12. [12] § Methods › Training data generation ↔ ailoc/simulation/mol_sampler.py, lines 230–282 · score 0.53 · predefined ranges, uniform distributions, batch, training, molecule, position
  13. [13] § Methods › Microscope ↔ ailoc/simulation/vectorpsf.py, lines 1514–1557 · score 0.52 · focal plane, microscope, NA, objective, camera, pixel
  14. [14] § Methods › PSF engineering ↔ ailoc/simulation/vectorpsf.py, lines 971–1076 · score 0.52 · optimized Zernike coefficients, pupil, wavelength, phase, nm, simulation
  15. [15] § Methods › Synchronized learning strategy ↔ ailoc/decode/loss.py, lines 35–76 · score 0.51 · negative log likelihood, GMM, loss, Gaussian, predicted, localization
  16. [16] § Methods › Synchronized learning strategy ↔ ailoc/lunar/loss.py, lines 35–76 · score 0.51 · negative log likelihood, GMM, loss, Gaussian, predicted, localization

Paper

Loaded from Europe PMC by your browser, not stored by OSCR: doi.org · Europe PMC

The paper is loaded when this pane is shown.

The authors' code

Python · 1,560 lines · 79 KB · GPL-3.0 · 3 matches

  1. import ctypes
  2. import sys
  3. from abc import ABC, abstractmethod # abstract class
  4. import torch
  5. import torch.nn as nn
  6. import numpy as np
  7. import os
  8. from deprecated import deprecated
  10. import ailoc.common
  13. class VectorPSF(ABC):
  14. """
  15. Abstract base class for vector psf simulators
  16. """
  18. @abstractmethod
  19. def simulate(self, *args, **kwargs):
  20. """
  21. run the simulation
  22. """
  24. raise NotImplementedError
  26. @staticmethod
  27. def _gauss2D_kernel(shape=(3, 3), sigmax=0.5, sigmay=0.5, data_type=torch.float32):
  28. """
  29. 2D gaussian mask for VectorPSF otf rescale
  30. """
  32. m, n = [(ss - 1.) / 2. for ss in shape]
  33. y, x = torch.meshgrid(ailoc.common.gpu(torch.arange(-m, m + 1, 1), data_type=data_type),
  34. ailoc.common.gpu(torch.arange(-n, n + 1, 1), data_type=data_type), indexing='ij')
  35. h = torch.exp(-(x * x) / (2. * sigmax * sigmax + 1e-6) - (y * y) / (2. * sigmay * sigmay + 1e-6))
  36. torch.clamp(h, min=0)
  37. # h[h < torch.finfo(h.dtype).eps * h.max()] = 0
  38. with torch.no_grad():
  39. sumh = h.sum()
  40. output = h / sumh if sumh != 0 else h
  41. return output
  43. @staticmethod
  44. def otf_rescale(psfdata, sigma_xy):
  45. """
  46. Convolve the psf with a gaussian kernel, namely otf rescale
  48. Args:
  49. psfdata (torch.Tensor): psf data to be blured
  50. sigma_xy (torch.Tensor): sigma x y of the gaussian kernel, unit pixel
  52. Returns:
  53. torch.Tensor: blured psf data
  54. """
  56. kernel = VectorPSF._gauss2D_kernel(shape=(5, 5),
  57. sigmax=sigma_xy[0],
  58. sigmay=sigma_xy[1],
  59. data_type=psfdata.dtype).reshape((1, 1, 5, 5))
  60. psf_size = psfdata.shape[1]
  61. psfdata = psfdata.view(-1, 1, psf_size, psf_size)
  62. tmp = nn.functional.conv2d(psfdata, kernel, padding=2, stride=1)
  63. outdata = tmp.view(-1, psf_size, psf_size)
  65. return outdata
  68. class VectorPSFCUDA(VectorPSF):
  69. """
  70. Vector psf simulator using cuda, used for training due to its speed
  71. """
  73. def __init__(self, psf_params):
  74. """
  76. Args:
  77. psf_params (dict): PSF parameters for simulation except emitter positions.
  78. na: numerical aperture;
  79. wavelength: wavelength, unit nm;
  80. refmed: refractive index of the sample medium;
  81. refcov: refractive index of the coverslip;
  82. refimm: refractive index of the immersion oil;
  83. zernike_mode: zernike orders, 2D array, first column is radial order,
  84. second column is azimuthal order;
  85. zernike_coef: amplitude for each zernike mode, unit nm;
  86. zernike_coef_map: spatially variant amplitude map for each zernike mode
  87. with shape (num zernike, row, column), unit nm;
  88. objstage0: initial objStage position,relative to focus at coverslip, unit nm;
  89. zemit0: initial emitter z position, distance relative to coverslip, unit nm;
  90. pixel_size_xy: pixel sizes for x and y direction, unit nm;
  91. otf_rescale_xy: sigma x y of the gaussian kernel for otf rescale, unit pixel;
  92. npupil: number of pupil sampling points for simulation, unit pixel;
  93. psf_size: number of image pixels;
  95. Returns:
  96. VectorPSFCUDA: an instance of VectorPSFCUDA
  97. """
  99. self.na = psf_params['na']
  100. self.wavelength = psf_params['wavelength']
  102. self.refmed = psf_params['refmed']
  103. self.refcov = psf_params['refcov']
  104. self.refimm = psf_params['refimm']
  106. self.zernike_mode = psf_params['zernike_mode']
  107. try:
  108. self.zernike_coef = psf_params['zernike_coef']
  109. except KeyError:
  110. self.zernike_coef = None
  111. try:
  112. self.zernike_coef_map = psf_params['zernike_coef_map']
  113. except KeyError:
  114. self.zernike_coef_map = None
  115. if (self.zernike_coef is None) != (self.zernike_coef_map is None):
  116. pass
  117. else:
  118. raise ValueError('you must define either zernike_coef xor zernike_coef_map')
  120. self.pixel_size_xy = psf_params['pixel_size_xy']
  121. self.otf_rescale_xy = psf_params['otf_rescale_xy']
  123. self.npupil = psf_params['npupil']
  124. self.psf_size = psf_params['psf_size']
  126. self.objstage0 = psf_params['objstage0']
  127. try:
  128. self.zemit0 = psf_params['zemit0']
  129. except KeyError:
  130. self.zemit0 = -self.objstage0 / self.refimm * self.refmed
  132. thispath = os.path.dirname(os.path.abspath(__file__))
  133. self.dll_path = thispath + '/../extensions/psf_simu_gpu.dll'
  135. def simulate(self, x, y, z, photons, zernike_coefs=None):
  136. """
  137. Run the simulation to generate the vector psfs with the given positions
  139. Args:
  140. x (torch.Tensor): x positions of the psfs, unit nm
  141. y (torch.Tensor): y positions of the psfs, unit nm
  142. z (torch.Tensor): z positions of the psfs, unit nm
  143. photons (torch.Tensor): photon counts of the psfs, unit photons
  144. zernike_coefs (torch.Tensor or None): if not None, each psf can be assigned a different zernike
  145. coefficients from this array with shape (npsf, 21), otherwise use the common class
  146. property self.zernike_coef, unit nm
  148. Returns:
  149. torch.Tensor: psfs, unit photons
  150. """
  152. psf_dll = ctypes.CDLL(self.dll_path, winmode=0)
  154. class _PSFParams(ctypes.Structure):
  155. _fields_ = [
  156. ('aberrations_', ctypes.POINTER(ctypes.c_float)),
  157. ('NA_', ctypes.c_float),
  158. ('refmed_', ctypes.c_float),
  159. ('refcov_', ctypes.c_float),
  160. ('refimm_', ctypes.c_float),
  161. ('lambdaX_', ctypes.c_float),
  162. ('objStage0_', ctypes.c_float),
  163. ('zemit0_', ctypes.c_float),
  164. ('pixelSizeX_', ctypes.c_float),
  165. ('pixelSizeY_', ctypes.c_float),
  166. ('sizeX_', ctypes.c_float),
  167. ('sizeY_', ctypes.c_float),
  168. ('PupilSize_', ctypes.c_float),
  169. ('Npupil_', ctypes.c_float),
  170. ('zernikeModesN_', ctypes.c_int),
  171. ('xemit_', ctypes.POINTER(ctypes.c_float)),
  172. ('yemit_', ctypes.POINTER(ctypes.c_float)),
  173. ('zemit_', ctypes.POINTER(ctypes.c_float)),
  174. ('objstage_', ctypes.POINTER(ctypes.c_float)),
  175. ('aberrationsParas_', ctypes.POINTER(ctypes.c_float)),
  176. ('psfOut_', ctypes.POINTER(ctypes.c_float)),
  177. ('aberrationOut_', ctypes.POINTER(ctypes.c_float)),
  178. ('Nmol_', ctypes.c_int),
  179. ('showAberrationNumber_', ctypes.c_int)
  180. ]
  182. param_struct = _PSFParams()
  184. # the x y positions and pixelsize should be inversed to ensure the
  185. # input X_os corresponds to the column
  186. param_struct.Npupil_ = self.npupil
  187. param_struct.Nmol_ = x.shape[0]
  188. param_struct.NA_ = self.na
  189. param_struct.refmed_ = self.refmed
  190. param_struct.refcov_ = self.refcov
  191. param_struct.refimm_ = self.refimm
  192. param_struct.lambdaX_ = self.wavelength
  193. param_struct.objStage0_ = self.objstage0
  194. param_struct.zemit0_ = self.zemit0
  195. param_struct.pixelSizeX_ = self.pixel_size_xy[1]
  196. param_struct.pixelSizeY_ = self.pixel_size_xy[0]
  197. param_struct.zernikeModesN_ = self.zernike_mode.shape[0]
  198. param_struct.sizeX_ = self.psf_size
  199. param_struct.sizeY_ = self.psf_size
  200. param_struct.PupilSize_ = 1.0
  201. param_struct.showAberrationNumber_ = 1
  203. np_xemit = ailoc.common.cpu(y).astype('float32') # nm
  204. np_yemit = ailoc.common.cpu(x).astype('float32') # nm
  205. np_zemit = ailoc.common.cpu(z).astype('float32') # nm
  206. np_objstage = 0 * np_zemit # nm
  207. np_aberrations = np.array(np.pad(self.zernike_mode, pad_width=((0, 0), (0, 1)), mode='constant',
  208. constant_values=((0, 0), (0, 0))).flatten('F'), dtype=np.float32)
  210. param_struct.xemit_ = np_xemit.ctypes.data_as(ctypes.POINTER(ctypes.c_float))
  211. param_struct.yemit_ = np_yemit.ctypes.data_as(ctypes.POINTER(ctypes.c_float))
  212. param_struct.zemit_ = np_zemit.ctypes.data_as(ctypes.POINTER(ctypes.c_float))
  213. param_struct.objstage_ = np_objstage.ctypes.data_as(ctypes.POINTER(ctypes.c_float))
  214. param_struct.aberrations_ = np_aberrations.ctypes.data_as(ctypes.POINTER(ctypes.c_float))
  216. n_mol = param_struct.Nmol_
  217. outdata_buffer = np.empty((n_mol, self.psf_size, self.psf_size), dtype=np.float32, order='C')
  218. outpupil_buffer = np.empty((int(param_struct.Npupil_), int(param_struct.Npupil_)), dtype=np.float32, order='C')
  219. input_zernikecoefs_buffer = np.empty((n_mol, param_struct.zernikeModesN_), dtype=np.float32, order='C')
  221. if zernike_coefs is None:
  222. assert self.zernike_coef is not None, \
  223. 'both self.zernike_coef and input zernike_coefs are None, please define either of them'
  224. zernike_coefs = np.tile(self.zernike_coef, reps=(n_mol, 1)).astype('float32')
  225. else:
  226. zernike_coefs = ailoc.common.cpu(zernike_coefs).astype('float32')
  228. for n in range(0, n_mol):
  229. input_zernikecoefs_buffer[n] = zernike_coefs[n]
  230. # input_zernikecoefs_buffer = input_zernikecoefs_buffer * self.wavelength # multiply with lambda
  232. param_struct.aberrationsParas_ = input_zernikecoefs_buffer.ctypes.data_as(ctypes.POINTER(ctypes.c_float))
  233. param_struct.psfOut_ = outdata_buffer.ctypes.data_as(ctypes.POINTER(ctypes.c_float))
  234. param_struct.aberrationOut_ = outpupil_buffer.ctypes.data_as(ctypes.POINTER(ctypes.c_float))
  236. # run simulation
  237. psf_dll.vectorPSFF1(param_struct)
  239. # check nan in the output PSFs
  240. assert not np.isnan(outdata_buffer).any(), "nan in the gpu psf, something wrong about the .dll call"
  242. psfs_out = ailoc.common.gpu(outdata_buffer)
  244. # otf rescale
  245. if self.otf_rescale_xy[0] or self.otf_rescale_xy[1] != 0:
  246. psfs_out = self.otf_rescale(psfdata=psfs_out, sigma_xy=ailoc.common.gpu(self.otf_rescale_xy))
  248. # normalize the psf to 1, then multiply with the photon number
  249. # psfs_out /= psfs_out.sum(-1).sum(-1)[:, None, None]
  250. psfs_out *= photons[:, None, None]
  252. return psfs_out
  255. class PartialOptimizeFunction(torch.autograd.Function):
  256. @staticmethod
  257. def forward(ctx, coef, indices_to_optimize):
  258. # save needed tensors and indices
  259. ctx.save_for_backward(coef, indices_to_optimize)
  260. return coef # directly return coef in the forward pass
  262. @staticmethod
  263. def backward(ctx, grad_output):
  264. # get the saved tensors and indices
  265. coef, indices_to_optimize = ctx.saved_tensors
  266. # create a zero gradient tensor
  267. grad_input = torch.zeros_like(coef)
  268. # only assign gradients to specified indices
  269. grad_input[indices_to_optimize] = grad_output[indices_to_optimize]
  270. return grad_input, None
  273. # 包装函数
  274. def partial_optimize(coef, indices_to_optimize):
  275. return PartialOptimizeFunction.apply(coef, indices_to_optimize)
  278. class VectorPSFTorch(VectorPSF):
  279. """
  280. Vector psf simulator using pytorch, thus psf parameters can be optimized by pytorch
  281. """
  283. def __init__(self, psf_params, req_grad=False, data_type=torch.float64, zernike_idx_learn=None):
  284. """
  286. Args:
  287. psf_params (dict): PSF parameters for simulation except emitter positions
  288. req_grad (bool): whether the PSF parameters are required gradient
  289. data_type (torch.dtype): data type for the PSF parameters
  290. zernike_idx_learn (list of bool): specified zernike coefficients to learn for LUNAR SL.
  292. Returns:
  293. VectorPSFTorch: an instance of VectorPSFTorch
  294. """
  296. self.data_type = data_type
  297. if data_type == torch.float64:
  298. self.complex_type = torch.complex128
  299. elif data_type == torch.float32:
  300. self.complex_type = torch.complex64
  301. elif data_type == torch.float16:
  302. self.complex_type = torch.complex32
  303. else:
  304. raise ValueError(f'unsupported data type {data_type}')
  306. self.na = torch.tensor(psf_params['na'], device='cuda', dtype=self.data_type)
  307. self.wavelength = torch.tensor(psf_params['wavelength'], device='cuda', dtype=self.data_type)
  308. self.refmed = torch.tensor(psf_params['refmed'], device='cuda', dtype=self.data_type,)
  309. self.refcov = torch.tensor(psf_params['refcov'], device='cuda', dtype=self.data_type,)
  310. self.refimm = torch.tensor(psf_params['refimm'], device='cuda', dtype=self.data_type,)
  311. self.zernike_mode = torch.tensor(psf_params['zernike_mode'], device='cuda', dtype=self.data_type)
  312. self.zernike_coef = torch.tensor(psf_params['zernike_coef'], device='cuda', dtype=self.data_type,
  313. requires_grad=req_grad)
  314. if zernike_idx_learn is None:
  315. self.zernike_idx_learn = torch.arange(self.zernike_coef.shape[0])
  316. else:
  317. self.zernike_idx_learn = torch.tensor(zernike_idx_learn)
  318. self.zernike_coef_map = None
  319. self.objstage0 = torch.tensor(psf_params['objstage0'], device='cuda', dtype=self.data_type,)
  320. try:
  321. self.zemit0 = torch.tensor(psf_params['zemit0'], device='cuda', dtype=self.data_type, )
  322. except KeyError:
  323. self.zemit0 = torch.tensor(-psf_params['objstage0']/psf_params['refimm']*psf_params['refmed'],
  324. device='cuda',
  325. dtype=self.data_type,)
  327. self.pixel_size_xy = torch.tensor(psf_params['pixel_size_xy'], device='cuda', dtype=self.data_type)
  328. self.otf_rescale_xy = torch.tensor(psf_params['otf_rescale_xy'], device='cuda', dtype=self.data_type,)
  329. self.npupil = psf_params['npupil']
  330. self.psf_size = psf_params['psf_size']
  332. self.focus_norm = psf_params.get('focus_norm', False)
  334. self._pre_compute()
  336. @staticmethod
  337. def get_zernike(orders, xpupil, ypupil):
  338. """
  339. Calculate zernike polynomials on pupil plane
  341. Args:
  342. orders (torch.Tensor): zernike orders, 2D array, first column is radial order,
  343. second column is azimuthal order
  344. xpupil (torch.Tensor): x coordinate of pupil plane
  345. ypupil (torch.Tensor): y coordinate of pupil plane
  347. Returns:
  348. torch.Tensor: zernike polynomials
  349. """
  351. xpupil = torch.real(xpupil)
  352. ypupil = torch.real(ypupil)
  353. zersize = orders.shape
  354. Nzer = zersize[0]
  355. radormax = int(max(orders[:, 0]))
  356. azormax = int(max(abs(orders[:, 1])))
  357. [Nx, Ny] = xpupil.shape
  359. # zerpol = np.zeros( [radormax+1,azormax+1,Nx,Ny] )
  360. zerpol = torch.zeros([21, 6, Nx, Ny], device='cuda')
  361. rhosq = xpupil ** 2 + ypupil ** 2
  362. rho = torch.sqrt(rhosq)
  363. zerpol[0, 0, :, :] = torch.ones_like(xpupil)
  365. for jm in range(1, azormax + 2 + 1):
  366. m = jm - 1
  367. if m > 0:
  368. zerpol[jm - 1, jm - 1, :, :] = rho * torch.squeeze(zerpol[jm - 1 - 1, jm - 1 - 1, :, :])
  370. zerpol[jm + 2 - 1, jm - 1, :, :] = ((m + 2) * rhosq - m - 1) * torch.squeeze(zerpol[jm - 1, jm - 1, :, :])
  371. for p in range(2, radormax - m + 2 + 1):
  372. n = m + 2 * p
  373. jn = n + 1
  374. zerpol[jn - 1, jm - 1, :, :] = (2 * (n - 1) * (n * (n - 2) * (2 * rhosq - 1) - m ** 2) * torch.squeeze(
  375. zerpol[jn - 2 - 1, jm - 1, :, :]) -
  376. n * (n + m - 2) * (n - m - 2) * torch.squeeze(
  377. zerpol[jn - 4 - 1, jm - 1, :, :])) / ((n - 2) * (n + m) * (n - m))
  379. phi = torch.atan2(ypupil, xpupil)
  380. allzernikes = torch.zeros([Nzer, Nx, Ny], device='cuda')
  381. for j in range(1, Nzer + 1):
  382. n = int(orders[j - 1, 0])
  383. m = int(orders[j - 1, 1])
  384. if m >= 0:
  385. allzernikes[j - 1, :, :] = torch.squeeze(zerpol[n + 1 - 1, m + 1 - 1, :, :]) * torch.cos(m * phi)
  386. else:
  387. allzernikes[j - 1, :, :] = torch.squeeze(zerpol[n + 1 - 1, -m + 1 - 1, :, :]) * torch.sin(-m * phi)
  389. # plt.figure(constrained_layout=True)
  390. # for i in range(21):
  391. # plt.subplot(3, 7, i + 1)
  392. # plt.imshow(cpu(allzernikes[i]))
  393. # plt.show()
  395. return allzernikes
  397. def czt(self, datain, A, B, D):
  398. """
  399. Execute the chirp-z transform
  401. Args:
  402. datain (torch.Tensor): input data, 2D array
  403. A (torch.Tensor): chirp parameter, 1D array
  404. B (torch.Tensor): chirp parameter, 1D array
  405. D (torch.Tensor): chirp parameter, 1D array
  407. Returns:
  408. torch.Tensor: output data, 2D array
  409. """
  411. N = A.shape[1]
  412. M = B.shape[1]
  413. L = D.shape[1]
  414. K = datain.shape[0]
  416. # torch.repeat_interleave is too slow
  417. # t0 = time.time()
  418. # Amt = torch.repeat_interleave(A, K, 0)
  419. # Bmt = torch.repeat_interleave(B, K, 0)
  420. # Dmt = torch.repeat_interleave(D, K, 0)
  421. # print('torch czt: ', time.time() - t0)
  422. Amt = A.expand(K, N)
  423. Bmt = B.expand(K, M)
  424. Dmt = D.expand(K, L)
  426. cztin = torch.zeros([K, L], dtype=self.complex_type, device='cuda')
  427. cztin[:, 0:N] = Amt * datain
  428. tmp = Dmt * torch.fft.fft(cztin)
  429. cztout = torch.fft.ifft(tmp)
  431. dataout = Bmt * cztout[:, 0:M]
  433. return dataout
  435. def czt_parallel(self, datain, A, B, D):
  436. """
  437. Execute the chirp-z transform
  439. Args:
  440. datain (torch.Tensor): input data, 2D array
  441. A (torch.Tensor): chirp parameter, 1D array
  442. B (torch.Tensor): chirp parameter, 1D array
  443. D (torch.Tensor): chirp parameter, 1D array
  445. Returns:
  446. torch.Tensor: output data, 2D array
  447. """
  449. N = A.shape[1]
  450. M = B.shape[1]
  451. L = D.shape[1]
  452. K = datain.shape[-2]
  453. n_mol = datain.shape[-3]
  455. # torch.repeat_interleave is too slow
  456. # t0 = time.time()
  457. # Amt = torch.repeat_interleave(A, K, 0)
  458. # Bmt = torch.repeat_interleave(B, K, 0)
  459. # Dmt = torch.repeat_interleave(D, K, 0)
  460. # print('torch czt: ', time.time() - t0)
  461. Amt = A.expand(K, N)
  462. Bmt = B.expand(K, M)
  463. Dmt = D.expand(K, L)
  465. cztin = torch.zeros([2, 3, n_mol, K, L], dtype=self.complex_type, device='cuda')
  466. cztin[:, :, :, :, 0:N] = Amt[None, None, None] * datain
  467. try:
  468. tmp = Dmt * torch.fft.fft(cztin, dim=-1)
  469. except:
  470. print('fft error')
  471. cztout = torch.fft.ifft(tmp, dim=-1)
  473. dataout = Bmt[None, None, None] * cztout[:, :, :, :, 0:M]
  475. return dataout
  477. def prechirpz(self, xsize, qsize, N, M):
  478. """
  479. Calculate the auxiliary vectors for chirp-z.
  481. Args:
  482. xsize (float): normalized pupil radius 1.0
  483. qsize (float): the original pixel number that could cover the region of interest
  484. in the image plane if using FFT
  485. N (int): the sampling number on the pupil
  486. M (int): the sampling number on the region of interest on the image plane
  488. Returns:
  489. (torch.Tensor,torch.Tensor,torch.Tensor): auxiliary vectors
  490. """
  492. L = N + M - 1
  493. sigma = 2 * np.pi * xsize * qsize / N / M
  494. Afac = torch.exp(2 * 1j * sigma * (1 - M))
  495. Bfac = torch.exp(2 * 1j * sigma * (1 - N))
  496. sqW = torch.exp(2 * 1j * sigma)
  497. W = sqW ** 2
  499. # fixed phase factor and amplitude factor
  500. Gfac = (2 * xsize / N) * torch.exp(1j * sigma * (1 - N) * (1 - M))
  502. # integration about n
  503. Utmp = torch.zeros([1, N], dtype=self.complex_type, device='cuda')
  504. A = torch.zeros([1, N], dtype=self.complex_type, device='cuda')
  505. Utmp[0, 0] = sqW * Afac
  506. A[0, 0] = 1.0
  507. for i in range(1, N):
  508. A[0, i] = Utmp[0, i - 1] * A[0, i - 1]
  509. Utmp[0, i] = Utmp[0, i - 1] * W
  511. # the factor before the summation
  512. Utmp = torch.zeros([1, M], dtype=self.complex_type, device='cuda')
  513. B = torch.ones([1, M], dtype=self.complex_type, device='cuda')
  514. Utmp[0, 0] = sqW * Bfac
  515. B[0, 0] = Gfac
  516. for i in range(1, M):
  517. B[0, i] = Utmp[0, i - 1] * B[0, i - 1]
  518. Utmp[0, i] = Utmp[0, i - 1] * W
  520. # for circular convolution
  521. Utmp = torch.zeros([1, max(N, M) + 1], dtype=self.complex_type, device='cuda')
  522. Vtmp = torch.zeros([1, max(N, M) + 1], dtype=self.complex_type, device='cuda')
  523. Utmp[0, 0] = sqW
  524. Vtmp[0, 0] = 1.0
  525. # Utmp_cp = Utmp.clone()
  526. # Vtmp_cp = Vtmp.clone()
  527. for i in range(1, max(N, M) + 1):
  528. Vtmp[0, i] = Utmp[0, i - 1] * Vtmp[0, i - 1]
  529. Utmp[0, i] = Utmp[0, i - 1] * W
  530. # Vtmp[0, i] = Utmp_cp[0, i - 1] * Vtmp_cp[0, i - 1]
  531. # Utmp[0, i] = Utmp_cp[0, i - 1] * W
  532. # Vtmp_cp[0, i] = Vtmp[0, i].clone()
  533. # Utmp_cp[0, i] = Utmp[0, i].clone()
  535. D = torch.ones([1, L], dtype=self.complex_type, device='cuda')
  536. for i in range(0, M):
  537. D[0, i] = torch.conj(Vtmp[0, i])
  538. for i in range(0, N):
  539. D[0, L - 1 - i] = torch.conj(Vtmp[0, i + 1])
  541. D = torch.fft.fft(D, axis=1)
  543. return A, B, D
  545. # old version
  546. @deprecated(reason="the same as matlab code, using for loop is slow")
  547. def _pre_compute_v1(self):
  548. """
  549. Compute the common intermediate variables in advance, this can save time for PSFs simulation
  550. """
  551. # pupil radius (in diffraction units) and pupil coordinate sampling
  552. pupil_size = 1.0
  553. dxypupil = 2 * pupil_size / self.npupil
  554. xypupil = torch.arange(-pupil_size + dxypupil / 2, pupil_size, dxypupil, device='cuda', dtype=self.data_type)
  555. [xpupil, ypupil] = torch.meshgrid(xypupil, xypupil, indexing='ij')
  556. ypupil = torch.complex(ypupil, torch.zeros_like(ypupil))
  557. xpupil = torch.complex(xpupil, torch.zeros_like(xpupil))
  559. # calculation of relevant Fresnel-coefficients for the interfaces
  560. costhetamed = torch.sqrt(1.0 - (xpupil ** 2 + ypupil ** 2) * (self.na ** 2) / (self.refmed ** 2))
  561. costhetacov = torch.sqrt(1.0 - (xpupil ** 2 + ypupil ** 2) * (self.na ** 2) / (self.refcov ** 2))
  562. costhetaimm = torch.sqrt(1.0 - (xpupil ** 2 + ypupil ** 2) * (self.na ** 2) / (self.refimm ** 2))
  563. fresnelpmedcov = 2 * self.refmed * costhetamed / (self.refmed * costhetacov + self.refcov * costhetamed)
  564. fresnelsmedcov = 2 * self.refmed * costhetamed / (self.refmed * costhetamed + self.refcov * costhetacov)
  565. fresnelpcovimm = 2 * self.refcov * costhetacov / (self.refcov * costhetaimm + self.refimm * costhetacov)
  566. fresnelscovimm = 2 * self.refcov * costhetacov / (self.refcov * costhetacov + self.refimm * costhetaimm)
  567. fresnelp = fresnelpmedcov * fresnelpcovimm
  568. fresnels = fresnelsmedcov * fresnelscovimm
  570. # apodization
  571. apod = 1 / torch.sqrt(costhetaimm)
  572. # define aperture
  573. aperturemask = torch.where((xpupil ** 2 + ypupil ** 2).real < 1.0, 1.0, 0.0)
  574. self.amplitude = aperturemask * apod
  576. # setting of vectorial functions
  577. phi = torch.atan2(torch.real(ypupil), torch.real(xpupil))
  578. cosphi = torch.cos(phi)
  579. sinphi = torch.sin(phi)
  580. costheta = costhetamed
  581. sintheta = torch.sqrt(1 - costheta ** 2)
  583. pvec = torch.empty([3, self.npupil, self.npupil], dtype=self.complex_type, device='cuda')
  584. pvec[0] = fresnelp * costheta * cosphi
  585. pvec[1] = fresnelp * costheta * sinphi
  586. pvec[2] = -fresnelp * sintheta
  587. svec = torch.empty([3, self.npupil, self.npupil], dtype=self.complex_type, device='cuda')
  588. svec[0] = -fresnels * sinphi
  589. svec[1] = fresnels * cosphi
  590. svec[2] = 0 * cosphi
  592. polarizationvector = torch.empty([2, 3, self.npupil, self.npupil], dtype=self.complex_type, device='cuda')
  593. for ipol in range(3):
  594. polarizationvector[0, ipol] = cosphi * pvec[ipol] - sinphi * svec[ipol]
  595. polarizationvector[1, ipol] = sinphi * pvec[ipol] + cosphi * svec[ipol]
  597. self.wavevector = torch.empty([2, self.npupil, self.npupil], dtype=self.complex_type, device='cuda')
  598. self.wavevector[0] = 2 * np.pi * self.na / self.wavelength * xpupil
  599. self.wavevector[1] = 2 * np.pi * self.na / self.wavelength * ypupil
  600. self.wavevectorzimm = 2 * np.pi * self.refimm / self.wavelength * costhetaimm
  601. self.wavevectorzmed = 2 * np.pi * self.refmed / self.wavelength * costhetamed
  603. # calculate aberration function
  604. waberration = torch.zeros_like(xpupil, dtype=self.complex_type, device='cuda')
  605. normfac = torch.sqrt(
  606. 2 * (self.zernike_mode[:, 0] + 1) / (1 + torch.where(self.zernike_mode[:, 1] == 0, 1.0, 0.0)))
  607. zernikecoefs_norm = self.zernike_coef * normfac
  608. allzernikes = self.get_zernike(self.zernike_mode, xpupil, ypupil)
  610. for izer in range(self.zernike_mode.shape[0]):
  611. waberration += zernikecoefs_norm[izer] * allzernikes[izer]
  612. waberration *= aperturemask
  613. self.zernike_phase = torch.exp(1j * 2 * np.pi * waberration / self.wavelength)
  615. self.pupilmatrix = torch.empty([2, 3, self.npupil, self.npupil], dtype=self.complex_type, device='cuda')
  616. for imat in range(2):
  617. for jmat in range(3):
  618. self.pupilmatrix[imat, jmat] = self.amplitude * self.zernike_phase * polarizationvector[imat, jmat]
  620. # czt transform(fft the pupil)
  621. xrange = self.pixel_size_xy[0] * self.psf_size / 2
  622. yrange = self.pixel_size_xy[1] * self.psf_size / 2
  623. imagesizex = xrange * self.na / self.wavelength
  624. imagesizey = yrange * self.na / self.wavelength
  626. # calculate the auxiliary vectors for chirp-z, pixelsize_xy should be inverse to match row and column
  627. self.ax, self.bx, self.dx = self.prechirpz(pupil_size, imagesizey, self.npupil, self.psf_size)
  628. self.ay, self.by, self.dy = self.prechirpz(pupil_size, imagesizex, self.npupil, self.psf_size)
  630. # calculate intensity normalization function using the PSF at focus
  631. fieldmatrix_norm = torch.empty([2, 3, self.psf_size, self.psf_size], dtype=self.complex_type, device='cuda')
  632. for itel in range(2):
  633. for jtel in range(3):
  634. Pupilfunction_norm = self.amplitude * polarizationvector[itel, jtel]
  635. inter_image_norm = torch.transpose(self.czt(Pupilfunction_norm, self.ax, self.bx, self.dx), 1, 0)
  636. fieldmatrix_norm[itel, jtel] = torch.transpose(self.czt(inter_image_norm, self.ay, self.by, self.dy), 1,
  637. 0)
  638. int_focus = torch.zeros([self.psf_size, self.psf_size], dtype=self.data_type, device='cuda')
  639. for jtel in range(3):
  640. for itel in range(2):
  641. int_focus += 1 / 3 * (torch.abs(fieldmatrix_norm[itel, jtel])) ** 2
  642. self.norm_intensity = torch.sum(int_focus)
  644. @deprecated(reason="parallel version of v1, but not compatible with simulate v3, where zernike phase "
  645. "is not computed in advance")
  646. def _pre_compute_v2(self):
  647. """
  648. Compute the common intermediate variables in advance, this can save time for PSFs simulation
  649. """
  650. # pupil radius (in diffraction units) and pupil coordinate sampling
  651. pupil_size = 1.0
  652. dxypupil = 2 * pupil_size / self.npupil
  653. xypupil = torch.arange(-pupil_size + dxypupil / 2, pupil_size, dxypupil, device='cuda', dtype=self.data_type)
  654. [xpupil, ypupil] = torch.meshgrid(xypupil, xypupil, indexing='ij')
  655. ypupil = torch.complex(ypupil, torch.zeros_like(ypupil))
  656. xpupil = torch.complex(xpupil, torch.zeros_like(xpupil))
  658. # calculation of relevant Fresnel-coefficients for the interfaces
  659. costhetamed = torch.sqrt(1.0 - (xpupil ** 2 + ypupil ** 2) * (self.na ** 2) / (self.refmed ** 2))
  660. costhetacov = torch.sqrt(1.0 - (xpupil ** 2 + ypupil ** 2) * (self.na ** 2) / (self.refcov ** 2))
  661. costhetaimm = torch.sqrt(1.0 - (xpupil ** 2 + ypupil ** 2) * (self.na ** 2) / (self.refimm ** 2))
  662. fresnelpmedcov = 2 * self.refmed * costhetamed / (self.refmed * costhetacov + self.refcov * costhetamed)
  663. fresnelsmedcov = 2 * self.refmed * costhetamed / (self.refmed * costhetamed + self.refcov * costhetacov)
  664. fresnelpcovimm = 2 * self.refcov * costhetacov / (self.refcov * costhetaimm + self.refimm * costhetacov)
  665. fresnelscovimm = 2 * self.refcov * costhetacov / (self.refcov * costhetacov + self.refimm * costhetaimm)
  666. fresnelp = fresnelpmedcov * fresnelpcovimm
  667. fresnels = fresnelsmedcov * fresnelscovimm
  669. # apodization
  670. # apod = 1 / torch.sqrt(costhetaimm) # previous version, for the simulated test dataset, should be deprecated
  671. # apod = 1 / torch.sqrt(costhetamed)
  672. apod = torch.sqrt(costhetaimm) / costhetamed # Sjoerd Stallinga version
  674. # define aperture
  675. aperturemask = torch.where((xpupil ** 2 + ypupil ** 2).real < 1.0, 1.0, 0.0)
  676. self.amplitude = aperturemask * apod
  678. # setting of vectorial functions
  679. phi = torch.atan2(torch.real(ypupil), torch.real(xpupil))
  680. cosphi = torch.cos(phi)
  681. sinphi = torch.sin(phi)
  682. costheta = costhetamed
  683. sintheta = torch.sqrt(1 - costheta ** 2)
  685. pvec = torch.empty([3, self.npupil, self.npupil], dtype=self.complex_type, device='cuda')
  686. pvec[0] = fresnelp * costheta * cosphi
  687. pvec[1] = fresnelp * costheta * sinphi
  688. pvec[2] = -fresnelp * sintheta
  689. svec = torch.empty([3, self.npupil, self.npupil], dtype=self.complex_type, device='cuda')
  690. svec[0] = -fresnels * sinphi
  691. svec[1] = fresnels * cosphi
  692. svec[2] = 0 * cosphi
  694. polarizationvector = torch.empty([2, 3, self.npupil, self.npupil], dtype=self.complex_type, device='cuda')
  695. polarizationvector[0,] = cosphi * pvec - sinphi * svec
  696. polarizationvector[1,] = sinphi * pvec + cosphi * svec
  698. self.wavevector = torch.empty([2, self.npupil, self.npupil], dtype=self.complex_type, device='cuda')
  699. self.wavevector[0] = 2 * np.pi * self.na / self.wavelength * xpupil
  700. self.wavevector[1] = 2 * np.pi * self.na / self.wavelength * ypupil
  701. self.wavevectorzimm = 2 * np.pi * self.refimm / self.wavelength * costhetaimm
  702. self.wavevectorzmed = 2 * np.pi * self.refmed / self.wavelength * costhetamed
  704. # calculate aberration function
  705. waberration = torch.zeros_like(xpupil, dtype=self.complex_type, device='cuda')
  706. normfac = torch.sqrt(
  707. 2 * (self.zernike_mode[:, 0] + 1) / (1 + torch.where(self.zernike_mode[:, 1] == 0, 1.0, 0.0)))
  708. zernikecoefs_norm = self.zernike_coef * normfac
  709. allzernikes = self.get_zernike(self.zernike_mode, xpupil, ypupil)
  711. waberration += torch.sum(zernikecoefs_norm[:, None, None]*allzernikes, dim=0)
  712. waberration *= aperturemask
  713. self.zernike_phase = torch.exp(1j * 2 * np.pi * waberration / self.wavelength)
  715. self.pupilmatrix = torch.empty([2, 3, self.npupil, self.npupil], dtype=self.complex_type, device='cuda')
  716. self.pupilmatrix = self.amplitude[None, None] * self.zernike_phase[None, None] * polarizationvector
  718. # czt transform(fft the pupil)
  719. xrange = self.pixel_size_xy[0] * self.psf_size / 2
  720. yrange = self.pixel_size_xy[1] * self.psf_size / 2
  721. imagesizex = xrange * self.na / self.wavelength
  722. imagesizey = yrange * self.na / self.wavelength
  724. # calculate the auxiliary vectors for chirp-z, pixelsize_xy should be inverse to match row and column
  725. self.ax, self.bx, self.dx = self.prechirpz(pupil_size, imagesizey, self.npupil, self.psf_size)
  726. self.ay, self.by, self.dy = self.prechirpz(pupil_size, imagesizex, self.npupil, self.psf_size)
  728. # calculate intensity normalization function using the PSF at focus
  729. fieldmatrix_norm = torch.empty([2, 3, self.psf_size, self.psf_size], dtype=self.complex_type, device='cuda')
  730. Pupilfunction_norm = self.amplitude[None, None, None] * polarizationvector[:, :, None]
  731. inter_image_norm = torch.transpose(self.czt_parallel(Pupilfunction_norm, self.ax, self.bx, self.dx), -1, -2)
  732. fieldmatrix_norm = torch.transpose(self.czt_parallel(inter_image_norm, self.ay, self.by, self.dy), -1, -2)
  734. int_focus = torch.zeros([self.psf_size, self.psf_size], dtype=self.data_type, device='cuda')
  735. int_focus += 1 / 3 * torch.sum(torch.abs(fieldmatrix_norm) ** 2, dim=(0, 1, 2))
  736. self.norm_intensity = torch.sum(int_focus)
  738. def _pre_compute(self):
  739. """
  740. Compute the common intermediate variables in advance, this can save time for PSFs simulation
  741. """
  742. # pupil radius (in diffraction units) and pupil coordinate sampling
  743. pupil_size = 1.0
  744. dxypupil = 2 * pupil_size / self.npupil
  745. xypupil = torch.arange(-pupil_size + dxypupil / 2, pupil_size, dxypupil, device='cuda', dtype=self.data_type)
  746. [xpupil, ypupil] = torch.meshgrid(xypupil, xypupil, indexing='ij')
  747. ypupil = torch.complex(ypupil, torch.zeros_like(ypupil))
  748. xpupil = torch.complex(xpupil, torch.zeros_like(xpupil))
  750. # calculation of relevant Fresnel-coefficients for the interfaces
  751. costhetamed = torch.sqrt(1.0 - (xpupil ** 2 + ypupil ** 2) * (self.na ** 2) / (self.refmed ** 2))
  752. costhetacov = torch.sqrt(1.0 - (xpupil ** 2 + ypupil ** 2) * (self.na ** 2) / (self.refcov ** 2))
  753. costhetaimm = torch.sqrt(1.0 - (xpupil ** 2 + ypupil ** 2) * (self.na ** 2) / (self.refimm ** 2))
  754. fresnelpmedcov = 2 * self.refmed * costhetamed / (self.refmed * costhetacov + self.refcov * costhetamed)
  755. fresnelsmedcov = 2 * self.refmed * costhetamed / (self.refmed * costhetamed + self.refcov * costhetacov)
  756. fresnelpcovimm = 2 * self.refcov * costhetacov / (self.refcov * costhetaimm + self.refimm * costhetacov)
  757. fresnelscovimm = 2 * self.refcov * costhetacov / (self.refcov * costhetacov + self.refimm * costhetaimm)
  758. fresnelp = fresnelpmedcov * fresnelpcovimm
  759. fresnels = fresnelsmedcov * fresnelscovimm
  761. # apodization
  762. # apod = 1 / torch.sqrt(costhetaimm) # previous version, for the simulated test dataset, should be deprecated
  763. # apod = 1 / torch.sqrt(costhetamed)
  764. apod = torch.sqrt(costhetaimm) / costhetamed # Sjoerd Stallinga version
  766. # define aperture
  767. aperturemask = torch.where((xpupil ** 2 + ypupil ** 2).real < 1.0, 1.0, 0.0)
  768. self.amplitude = aperturemask * apod
  770. # setting of vectorial functions
  771. phi = torch.atan2(torch.real(ypupil), torch.real(xpupil))
  772. cosphi = torch.cos(phi)
  773. sinphi = torch.sin(phi)
  774. costheta = costhetamed
  775. sintheta = torch.sqrt(1 - costheta ** 2)
  777. pvec = torch.empty([3, self.npupil, self.npupil], dtype=self.complex_type, device='cuda')
  778. pvec[0] = fresnelp * costheta * cosphi
  779. pvec[1] = fresnelp * costheta * sinphi
  780. pvec[2] = -fresnelp * sintheta
  781. svec = torch.empty([3, self.npupil, self.npupil], dtype=self.complex_type, device='cuda')
  782. svec[0] = -fresnels * sinphi
  783. svec[1] = fresnels * cosphi
  784. svec[2] = 0 * cosphi
  786. self.polarizationvector = torch.empty([2, 3, self.npupil, self.npupil], dtype=self.complex_type, device='cuda')
  787. self.polarizationvector[0,] = cosphi * pvec - sinphi * svec
  788. self.polarizationvector[1,] = sinphi * pvec + cosphi * svec
  790. self.wavevector = torch.empty([2, self.npupil, self.npupil], dtype=self.complex_type, device='cuda')
  791. self.wavevector[0] = 2 * np.pi * self.na / self.wavelength * xpupil
  792. self.wavevector[1] = 2 * np.pi * self.na / self.wavelength * ypupil
  793. self.wavevectorzimm = 2 * np.pi * self.refimm / self.wavelength * costhetaimm
  794. self.wavevectorzmed = 2 * np.pi * self.refmed / self.wavelength * costhetamed
  796. # calculate aberration function
  797. normfac = torch.sqrt(
  798. 2 * (self.zernike_mode[:, 0] + 1) / (1 + torch.where(self.zernike_mode[:, 1] == 0, 1.0, 0.0)))
  799. self.allzernikes = self.get_zernike(self.zernike_mode, xpupil, ypupil) * normfac[:, None, None] * aperturemask[None]
  801. # czt transform(fft the pupil)
  802. xrange = self.pixel_size_xy[0] * self.psf_size / 2
  803. yrange = self.pixel_size_xy[1] * self.psf_size / 2
  804. imagesizex = xrange * self.na / self.wavelength
  805. imagesizey = yrange * self.na / self.wavelength
  807. # calculate the auxiliary vectors for chirp-z, pixelsize_xy should be inverse to match row and column
  808. self.ax, self.bx, self.dx = self.prechirpz(pupil_size, imagesizey, self.npupil, self.psf_size)
  809. self.ay, self.by, self.dy = self.prechirpz(pupil_size, imagesizex, self.npupil, self.psf_size)
  811. # calculate intensity normalization function using the PSF at focus
  812. pupilfunction_norm = self.amplitude[None, None, None] * self.polarizationvector[:, :, None]
  813. inter_image_norm = torch.transpose(self.czt_parallel(pupilfunction_norm, self.ax, self.bx, self.dx), -1, -2)
  814. fieldmatrix_norm = torch.transpose(self.czt_parallel(inter_image_norm, self.ay, self.by, self.dy), -1, -2)
  816. int_focus = torch.zeros([self.psf_size, self.psf_size], dtype=self.data_type, device='cuda')
  817. int_focus += 1 / 3 * torch.sum(torch.abs(fieldmatrix_norm) ** 2, dim=(0, 1, 2))
  818. self.norm_intensity = torch.sum(int_focus)
  820. @deprecated(reason="the same as matlab code, using for loop is slow")
  821. def simulate_v1(self, x, y, z, photons, objstage=None):
  822. """
  823. Run the simulation to generate the vector PSFs with the given positions
  825. Args:
  826. x (torch.Tensor): x positions of the PSFs, unit nm
  827. y (torch.Tensor): y positions of the PSFs, unit nm
  828. z (torch.Tensor): z positions of the PSFs, unit nm
  829. photons (torch.Tensor): photon counts of the PSFs, unit photons
  830. objstage (torch.Tensor): objective stage positions relative to the cover-slip (0),
  831. the closer to the sample, the smaller this value is (-), unit nm
  833. Returns:
  834. torch.Tensor: PSFs, unit photons
  835. """
  837. n_mol = x.shape[0]
  838. objstage = torch.zeros(x.shape[0], dtype=self.data_type, device='cuda') if objstage is None else objstage
  840. field_matrix = torch.empty([2, 3, n_mol, self.psf_size, self.psf_size],
  841. dtype=self.complex_type, device='cuda')
  842. for jz in range(n_mol):
  843. # xyz induced phase, x,y should be inverse to match the column, row
  844. if z[jz] + self.zemit0 >= 0:
  845. phase_xyz = -y[jz] * self.wavevector[0] - x[jz] * self.wavevector[1] + \
  846. (z[jz] + self.zemit0) * self.wavevectorzmed
  847. position_phase = torch.exp(1j * (phase_xyz + (objstage[jz] + self.objstage0) *
  848. self.wavevectorzimm))
  849. else:
  850. # print("warning! the emitter's position may not have physical meaning")
  851. phase_xyz = -y[jz] * self.wavevector[0] - x[jz] * self.wavevector[1]
  852. position_phase = torch.exp(1j * (phase_xyz + (objstage[jz] + self.objstage0 + z[jz]
  853. + self.zemit0) * self.wavevectorzimm))
  854. for itel in range(2):
  855. for jtel in range(3):
  856. pupil_tmp = position_phase * self.pupilmatrix[itel, jtel]
  857. inter_image = torch.transpose(self.czt(pupil_tmp, self.ay, self.by, self.dy), 1, 0)
  858. field_matrix[itel, jtel, jz] = torch.transpose(self.czt(inter_image, self.ax, self.bx, self.dx), 1,
  859. 0)
  861. psfs_out = torch.zeros([n_mol, self.psf_size, self.psf_size], device='cuda', dtype=torch.float32)
  862. for jz in range(n_mol):
  863. for jtel in range(3):
  864. for itel in range(2):
  865. psfs_out[jz, :, :] += 1 / 3 * (torch.abs(field_matrix[itel, jtel, jz])) ** 2
  867. psfs_out /= self.norm_intensity
  869. # otf rescale
  870. if self.otf_rescale_xy[0] or self.otf_rescale_xy[1] != 0:
  871. psfs_out = self.otf_rescale(psfdata=psfs_out, sigma_xy=self.otf_rescale_xy)
  873. # normalize the psf to 1, then multiply with the photon number
  874. # psfs_out /= psfs_out.sum(-1).sum(-1)[:, None, None]
  875. psfs_out *= photons[:, None, None]
  877. return psfs_out
  879. @deprecated(reason="parallel version of v1, but not compatible with simulate v3, where zernike phase "
  880. "is not computed in advance")
  881. def simulate_v2(self, x, y, z, photons, objstage=None):
  882. """
  883. Run the simulation to generate the vector PSFs with the given positions
  885. Args:
  886. x (torch.Tensor): x positions of the PSFs, unit nm
  887. y (torch.Tensor): y positions of the PSFs, unit nm
  888. z (torch.Tensor): z positions of the PSFs, unit nm
  889. photons (torch.Tensor): photon counts of the PSFs, unit photons
  890. objstage (torch.Tensor): objective stage positions relative to the cover-slip (0),
  891. the closer to the sample, the smaller this value is (-), unit nm
  893. Returns:
  894. torch.Tensor: PSFs, unit photons
  895. """
  896. if not hasattr(self, "device"):
  897. self.device = x.device
  899. n_mol = x.shape[0]
  900. if n_mol == 0:
  901. return torch.zeros([0, self.psf_size, self.psf_size], device='cuda', dtype=torch.float32)
  903. objstage = torch.zeros(x.shape[0], dtype=self.data_type, device='cuda') if objstage is None else objstage
  905. # batch_size in parallel to save GPU memory
  906. slice_list = []
  907. batch_size = 100
  908. for i in np.arange(0, n_mol, batch_size):
  909. slice_list.append(slice(i, min(i + batch_size, n_mol)))
  911. psfs_out = torch.zeros([n_mol, self.psf_size, self.psf_size], device='cuda', dtype=torch.float32)
  912. for slice_tmp in slice_list:
  913. length_tmp = slice_tmp.stop - slice_tmp.start
  914. position_phase = torch.empty([length_tmp, self.npupil, self.npupil], dtype=self.complex_type, device='cuda')
  916. idx = torch.where(z[slice_tmp] + self.zemit0 >= 0)[0]
  917. phase_xyz_tmp = -y[slice_tmp][idx][:, None, None] * self.wavevector[0][None] - x[slice_tmp][idx][:, None, None] * \
  918. self.wavevector[1][None] + \
  919. (z[slice_tmp][idx] + self.zemit0)[:, None, None] * self.wavevectorzmed[None]
  920. position_phase[idx, :, :] = torch.exp(
  921. 1j * (phase_xyz_tmp + (objstage[slice_tmp][idx][:, None, None] + self.objstage0) *
  922. self.wavevectorzimm[None]))
  924. idx = torch.where(z[slice_tmp] + self.zemit0 < 0)[0]
  925. phase_xyz_tmp = -y[slice_tmp][idx][:, None, None] * self.wavevector[0][None] - x[slice_tmp][idx][:, None, None] * \
  926. self.wavevector[1][None]
  927. position_phase[idx, :, :] = torch.exp(
  928. 1j * (phase_xyz_tmp + (objstage[slice_tmp][idx][:, None, None] + self.objstage0 + z[slice_tmp][idx][:, None, None]
  929. + self.zemit0) * self.wavevectorzimm[None]))
  931. pupil_tmp = position_phase[None, None] * self.pupilmatrix[:, :, None]
  932. inter_image = torch.transpose(self.czt_parallel(pupil_tmp, self.ay, self.by, self.dy), -1, -2)
  933. field_matrix = torch.transpose(self.czt_parallel(inter_image, self.ax, self.bx, self.dx), -1, -2)
  935. psfs_out[slice_tmp] += 1 / 3 * torch.sum((torch.abs(field_matrix[:, :])) ** 2, dim=(0, 1))
  937. # # all in parallel, but may cause GPU memory overflow
  938. # position_phase = torch.empty([n_mol, self.npupil, self.npupil], dtype=self.complex_type, device='cuda')
  939. #
  940. # idx = torch.where(z + self.zemit0 >= 0)[0]
  941. # phase_xyz_tmp = -y[idx][:, None, None] * self.wavevector[0][None] - x[idx][:, None, None] * self.wavevector[1][None] + \
  942. # (z[idx] + self.zemit0)[:, None, None] * self.wavevectorzmed[None]
  943. # position_phase[idx, :, :] = torch.exp(1j * (phase_xyz_tmp + (objstage[idx][:, None, None] + self.objstage0) *
  944. # self.wavevectorzimm[None]))
  945. #
  946. # idx = torch.where(z + self.zemit0 < 0)[0]
  947. # phase_xyz_tmp = -y[idx][:, None, None] * self.wavevector[0][None] - x[idx][:, None, None] * self.wavevector[1][None]
  948. # position_phase[idx, :, :] = torch.exp(1j * (phase_xyz_tmp + (objstage[idx][:, None, None] + self.objstage0 + z[idx][:, None, None]
  949. # + self.zemit0) * self.wavevectorzimm[None]))
  950. #
  951. # pupil_tmp = position_phase[None, None] * self.pupilmatrix[:, :, None]
  952. # inter_image = torch.transpose(self.czt_parallel(pupil_tmp, self.ay, self.by, self.dy), -1, -2)
  953. # field_matrix = torch.transpose(self.czt_parallel(inter_image, self.ax, self.bx, self.dx), -1, -2)
  954. #
  955. # psfs_out = torch.zeros([n_mol, self.psf_size, self.psf_size], device='cuda', dtype=torch.float32)
  956. # psfs_out += 1 / 3 * torch.sum((torch.abs(field_matrix[:, :])) ** 2, dim=(0, 1))
  958. # intensity normalization
  959. psfs_out /= self.norm_intensity
  961. # otf rescale
  962. if self.otf_rescale_xy[0] or self.otf_rescale_xy[1] != 0:
  963. psfs_out = self.otf_rescale(psfdata=psfs_out, sigma_xy=self.otf_rescale_xy)
  965. # normalize the psf to 1, then multiply with the photon number
  966. # psfs_out /= psfs_out.sum(-1).sum(-1)[:, None, None]
  967. psfs_out *= photons[:, None, None]
  969. return psfs_out
  971. def simulate(self, x, y, z, photons, objstage=None, zernike_coefs=None):
  972. """
  973. Run the simulation to generate the vector PSFs with the given positions
  975. Args:
  976. x (torch.Tensor): x positions of the PSFs, unit nm
  977. y (torch.Tensor): y positions of the PSFs, unit nm
  978. z (torch.Tensor): z positions of the PSFs, unit nm
  979. photons (torch.Tensor): photon counts of the PSFs, unit photons
  980. objstage (torch.Tensor): objective stage positions relative to the cover-slip (0),
  981. the closer to the sample, the smaller this value is (-), unit nm
  982. zernike_coefs (torch.Tensor or None): if not None, each psf can be assigned a different zernike
  983. coefficients from this array with shape (npsf, 21), otherwise use the common class
  984. property self.zernike_coef, unit nm
  986. Returns:
  987. torch.Tensor: PSFs, unit photons
  988. """
  990. n_mol = x.shape[0]
  991. if n_mol == 0:
  992. return torch.zeros([0, self.psf_size, self.psf_size], device='cuda', dtype=torch.float32)
  994. objstage = torch.zeros(x.shape[0], dtype=self.data_type, device='cuda') if objstage is None else objstage
  996. # # temporally test, change the z to obj
  997. # objstage = z
  998. # z = torch.zeros_like(z, dtype=self.data_type, device='cuda')
  1001. if zernike_coefs is None:
  1002. try:
  1003. # for partially optimized zernike coefficients in LUNAR SL physics learning
  1004. zernike_coef_part_optm = partial_optimize(self.zernike_coef, self.zernike_idx_learn)
  1005. zernike_phase = torch.exp(1j * 2 * np.pi *
  1006. torch.sum(zernike_coef_part_optm[:, None, None] * self.allzernikes, dim=0)
  1007. / self.wavelength)
  1008. except:
  1009. zernike_phase = torch.exp(1j * 2 * np.pi *
  1010. torch.sum(self.zernike_coef[:, None, None] * self.allzernikes, dim=0)
  1011. / self.wavelength)
  1012. pupilmatrix = (self.amplitude[None, None] *
  1013. zernike_phase[None, None] *
  1014. self.polarizationvector)[:, :, None]
  1016. # batch_size in parallel to save GPU memory
  1017. slice_list = []
  1018. batch_size = 100
  1019. for i in np.arange(0, n_mol, batch_size):
  1020. slice_list.append(slice(i, min(i + batch_size, n_mol)))
  1022. psfs_out = torch.zeros([n_mol, self.psf_size, self.psf_size], device='cuda', dtype=torch.float32)
  1023. for slice_tmp in slice_list:
  1024. if zernike_coefs is not None:
  1025. zernike_phase = torch.exp(1j * 2 * np.pi *
  1026. torch.sum(zernike_coefs[slice_tmp, :, None, None] * self.allzernikes[None], dim=1)
  1027. / self.wavelength)
  1028. pupilmatrix = (zernike_phase[None, None] *
  1029. self.polarizationvector[:, :, None] *
  1030. self.amplitude[None, None, None])
  1032. # length_tmp = slice_tmp.stop - slice_tmp.start
  1033. # position_phase = torch.empty([length_tmp, self.npupil, self.npupil], dtype=self.complex_type, device='cuda')
  1035. # idx = torch.where(z[slice_tmp] + self.zemit0 >= 0)[0]
  1036. # phase_xyz_tmp = -y[slice_tmp][idx][:, None, None] * self.wavevector[0][None] - x[slice_tmp][idx][:, None, None] * \
  1037. # self.wavevector[1][None] + \
  1038. # (z[slice_tmp][idx] + self.zemit0)[:, None, None] * self.wavevectorzmed[None]
  1039. # position_phase[idx, :, :] = torch.exp(
  1040. # 1j * (phase_xyz_tmp + (objstage[slice_tmp][idx][:, None, None] + self.objstage0) *
  1041. # self.wavevectorzimm[None]))
  1042. #
  1043. # idx = torch.where(z[slice_tmp] + self.zemit0 < 0)[0]
  1044. # phase_xyz_tmp = -y[slice_tmp][idx][:, None, None] * self.wavevector[0][None] - x[slice_tmp][idx][:, None, None] * \
  1045. # self.wavevector[1][None]
  1046. # position_phase[idx, :, :] = torch.exp(
  1047. # 1j * (phase_xyz_tmp + (objstage[slice_tmp][idx][:, None, None] + self.objstage0 + z[slice_tmp][idx][:, None, None]
  1048. # + self.zemit0) * self.wavevectorzimm[None]))
  1050. phase_xyz_tmp = -y[slice_tmp][:, None, None] * self.wavevector[0][None] - x[slice_tmp][:, None, None] * \
  1051. self.wavevector[1][None] + (z[slice_tmp] + self.zemit0)[:, None, None] * self.wavevectorzmed[None]
  1052. position_phase = torch.exp(1j * (phase_xyz_tmp + (objstage[slice_tmp][:, None, None] + self.objstage0) *
  1053. self.wavevectorzimm[None]))
  1055. pupil_tmp = position_phase[None, None] * pupilmatrix
  1056. inter_image = torch.transpose(self.czt_parallel(pupil_tmp, self.ay, self.by, self.dy), -1, -2)
  1057. field_matrix = torch.transpose(self.czt_parallel(inter_image, self.ax, self.bx, self.dx), -1, -2)
  1059. psfs_out[slice_tmp] += 1 / 3 * torch.sum((torch.abs(field_matrix[:, :])) ** 2, dim=(0, 1))
  1061. if self.focus_norm:
  1062. # intensity normalization by focus
  1063. psfs_out /= self.norm_intensity
  1064. else:
  1065. # normalize by themselves
  1066. norm_factor = psfs_out.sum(dim=(-1, -2))
  1067. psfs_out /= norm_factor[:, None, None]
  1069. # otf rescale
  1070. if self.otf_rescale_xy[0] or self.otf_rescale_xy[1]:
  1071. psfs_out = self.otf_rescale(psfdata=psfs_out, sigma_xy=self.otf_rescale_xy)
  1073. # multiply with the photon number
  1074. psfs_out *= photons[:, None, None]
  1076. return psfs_out
  1078. def compute_crlb(self, x, y, z, photons, bgs):
  1079. """
  1080. Calculate the CRLB of this PSF model at give positions, photons and backgrounds.
  1082. Args:
  1083. x (torch.Tensor): x positions of the PSFs, unit nm
  1084. y (torch.Tensor): y positions of the PSFs, unit nm
  1085. z (torch.Tensor): z positions of the PSFs, unit nm
  1086. photons (torch.Tensor): photon counts of the PSFs, unit photons
  1087. bgs (torch.Tensor): background counts of the PSFs, unit photons
  1089. Returns:
  1090. (torch.Tensor,torch.Tensor): CRLB xyz (nPSFs, 3) and model PSFs, unit nm and photons
  1091. """
  1093. n_mol = x.shape[0]
  1095. # calculate the derivatives
  1096. [dudt, model] = self._compute_derivative(x, y, z, photons, bgs)
  1097. # dudt = self._torch_jacobian_derivative(x, y, z, photons, bgs)
  1099. # calculate hessian matrix, here only consider not shared parameters: x, y, z, photons, background
  1100. num_pars = n_mol * 5
  1101. t2 = 1 / model
  1102. hessian = torch.zeros([num_pars, num_pars], device='cuda', dtype=self.data_type)
  1103. for p1 in range(num_pars):
  1104. temp1_zind = int(np.floor(p1 / 5)) # the index of data
  1105. temp1_pind = int(p1 % 5) # the index of parameter type
  1106. temp1 = dudt[:, :, :, temp1_pind] # the derivative of data concerning this parameter type
  1107. for p2 in range(p1, num_pars):
  1108. temp2_zind = int(np.floor(p2 / 5))
  1109. temp2_pind = int(p2 % 5)
  1110. temp2 = dudt[:, :, :, temp2_pind]
  1112. # since all parameters are not shared, only the same molecule data makes sense
  1113. # when multiply gradients of two parameters
  1114. if temp1_zind == temp2_zind:
  1115. temp = t2[temp1_zind, :, :] * temp1[temp1_zind, :, :] * temp2[temp2_zind, :, :]
  1116. hessian[p1, p2] = torch.sum(temp)
  1117. hessian[p2, p1] = hessian[p1, p2]
  1119. # calculate local fisher matrix and crlb
  1120. xyz_crlb = torch.zeros([n_mol, 3], device='cuda')
  1121. for j in range(n_mol):
  1122. fisher_tmp = hessian[j * 5:j * 5 + 5, j * 5:j * 5 + 5]
  1123. sqrt_crlb_tmp = torch.sqrt(torch.diag(torch.inverse(fisher_tmp)))
  1124. xyz_crlb[j] = sqrt_crlb_tmp[0:3]
  1126. return xyz_crlb, model
  1128. def compute_crlb_mf(self, x, y, z, photons, bgs, attn_length):
  1129. #todo: need test
  1130. """
  1131. Calculate the CRLB of this PSF model at give positions, photons and backgrounds.
  1133. Args:
  1134. x (torch.Tensor): x positions of the PSFs, unit nm
  1135. y (torch.Tensor): y positions of the PSFs, unit nm
  1136. z (torch.Tensor): z positions of the PSFs, unit nm
  1137. photons (torch.Tensor): photon counts of the PSFs, unit photons
  1138. bgs (torch.Tensor): background counts of the PSFs, unit photons
  1139. attn_length (int): attention length of the network, used to multiply the Fisher matrix
  1141. Returns:
  1142. (torch.Tensor,torch.Tensor): CRLB xyz (nPSFs, 3) and model PSFs, unit nm and photons
  1143. """
  1145. n_mol = x.shape[0]
  1147. # calculate the derivatives
  1148. [dudt, model] = self._compute_derivative_parallel(x, y, z, photons, bgs)
  1150. # calculate hessian matrix, here only consider not shared parameters: x, y, z, photons, background
  1151. num_pars = n_mol * 5
  1152. t2 = 1 / model
  1153. hessian = torch.zeros([num_pars, num_pars], device='cuda', dtype=self.data_type)
  1154. for p1 in range(num_pars):
  1155. temp1_zind = int(np.floor(p1 / 5)) # the index of data
  1156. temp1_pind = int(p1 % 5) # the index of parameter type
  1157. temp1 = dudt[:, :, :, temp1_pind] # the derivative of data concerning this parameter type
  1158. for p2 in range(p1, num_pars):
  1159. temp2_zind = int(np.floor(p2 / 5))
  1160. temp2_pind = int(p2 % 5)
  1161. temp2 = dudt[:, :, :, temp2_pind]
  1163. # since all parameters are not shared, only the same molecule data makes sense
  1164. # when multiply gradients of two parameters
  1165. if temp1_zind == temp2_zind:
  1166. temp = t2[temp1_zind, :, :] * temp1[temp1_zind, :, :] * temp2[temp2_zind, :, :]
  1167. hessian[p1, p2] = torch.sum(temp)
  1168. hessian[p2, p1] = hessian[p1, p2]
  1170. # calculate local fisher matrix and crlb
  1171. xyz_crlb = torch.zeros([n_mol, 3], device='cuda')
  1172. for j in range(n_mol):
  1173. fisher_tmp = hessian[j * 5:j * 5 + 5, j * 5:j * 5 + 5]
  1174. fisher_tmp[:3, :3] *= attn_length
  1175. sqrt_crlb_tmp = torch.sqrt(torch.diag(torch.inverse(fisher_tmp)))
  1176. xyz_crlb[j] = sqrt_crlb_tmp[0:3]
  1178. return xyz_crlb, model
  1180. @deprecated(reason='the same as matlab code, using for loop is slow')
  1181. def _compute_derivative_v1(self, x, y, z, photons, bgs):
  1182. """
  1183. Calculate the analytical derivatives of the PSFs at given parameters with respect to x,y,z,photons,bg
  1184. """
  1186. n_mol = x.shape[0]
  1187. objstage = torch.zeros(x.shape[0], dtype=self.data_type, device='cuda')
  1188. field_matrix = torch.empty([2, 3, n_mol, self.psf_size, self.psf_size],
  1189. dtype=self.complex_type, device='cuda')
  1190. field_matrix_ders = torch.empty([2, 3, n_mol, 3, self.psf_size, self.psf_size],
  1191. dtype=self.complex_type, device='cuda')
  1192. for jz in range(n_mol):
  1193. # xyz induced phase
  1194. if z[jz] + self.zemit0 >= 0:
  1195. phase_xyz = -y[jz] * self.wavevector[0] - x[jz] * self.wavevector[1] + \
  1196. (z[jz] + self.zemit0) * self.wavevectorzmed
  1197. position_phase = torch.exp(1j * (phase_xyz + (objstage[jz] + self.objstage0) *
  1198. self.wavevectorzimm))
  1199. else:
  1200. # print("warning! the emitter's position may not have physical meaning")
  1201. phase_xyz = -y[jz] * self.wavevector[0] - x[jz] * self.wavevector[1]
  1202. position_phase = torch.exp(1j * (phase_xyz + (objstage[jz] + self.objstage0 + z[jz]
  1203. + self.zemit0) * self.wavevectorzimm))
  1204. for itel in range(2):
  1205. for jtel in range(3):
  1206. pupil_tmp = position_phase * self.pupilmatrix[itel, jtel]
  1207. inter_image = torch.transpose(self.czt(pupil_tmp, self.ay, self.by, self.dy), 1, 0)
  1208. field_matrix[itel, jtel, jz] = torch.transpose(self.czt(inter_image, self.ax, self.bx, self.dx), 1,
  1209. 0)
  1210. # derivatives with respect to x,y,z
  1211. pupilfunction_x = -1j * self.wavevector[1] * position_phase * self.pupilmatrix[itel, jtel]
  1212. inter_image_x = torch.transpose(self.czt(pupilfunction_x, self.ay, self.by, self.dy), 1, 0)
  1213. field_matrix_ders[itel, jtel, jz, 0] = torch.transpose(self.czt(inter_image_x, self.ax,
  1214. self.bx, self.dx), 1, 0)
  1216. pupilfunction_y = -1j * self.wavevector[0] * position_phase * self.pupilmatrix[itel, jtel]
  1217. inter_image_y = torch.transpose(self.czt(pupilfunction_y, self.ay, self.by, self.dy), 1, 0)
  1218. field_matrix_ders[itel, jtel, jz, 1] = torch.transpose(self.czt(inter_image_y, self.ax,
  1219. self.bx, self.dx), 1, 0)
  1221. if z[jz] + self.zemit0 >= 0:
  1222. pupilfunction_z = 1j * self.wavevectorzmed * position_phase * self.pupilmatrix[itel, jtel]
  1223. else:
  1224. pupilfunction_z = 1j * self.wavevectorzimm * position_phase * self.pupilmatrix[itel, jtel]
  1225. inter_image_z = torch.transpose(self.czt(pupilfunction_z, self.ay, self.by, self.dy), 1, 0)
  1226. field_matrix_ders[itel, jtel, jz, 2] = torch.transpose(self.czt(inter_image_z, self.ax,
  1227. self.bx, self.dx), 1, 0)
  1229. psfs = torch.zeros([n_mol, self.psf_size, self.psf_size], device='cuda', dtype=self.data_type)
  1230. psfs_ders = torch.zeros([n_mol, self.psf_size, self.psf_size, 3], device='cuda', dtype=self.data_type)
  1231. for jz in range(n_mol):
  1232. for jtel in range(3):
  1233. for itel in range(2):
  1234. psfs[jz, :, :] += 1 / 3 * (torch.abs(field_matrix[itel, jtel, jz])) ** 2
  1235. for jder in range(3):
  1236. psfs_ders[jz, :, :, jder] = psfs_ders[jz, :, :, jder] + 2 / 3 * \
  1237. torch.real(torch.conj(field_matrix[itel, jtel, jz]) *
  1238. field_matrix_ders[itel, jtel, jz, jder])
  1239. psfs /= self.norm_intensity
  1240. psfs_ders /= self.norm_intensity
  1242. # otf rescale
  1243. if self.otf_rescale_xy[0] or self.otf_rescale_xy[1] != 0:
  1244. psfs = self.otf_rescale(psfdata=psfs, sigma_xy=self.otf_rescale_xy)
  1245. psfs_ders[:, :, :, 0] = self.otf_rescale(psfdata=psfs_ders[:, :, :, 0], sigma_xy=self.otf_rescale_xy)
  1246. psfs_ders[:, :, :, 1] = self.otf_rescale(psfdata=psfs_ders[:, :, :, 1], sigma_xy=self.otf_rescale_xy)
  1247. psfs_ders[:, :, :, 2] = self.otf_rescale(psfdata=psfs_ders[:, :, :, 2], sigma_xy=self.otf_rescale_xy)
  1249. psfs_out = ailoc.common.gpu(psfs * photons[:, None, None] + bgs[:, None, None])
  1250. ders_out = torch.zeros([n_mol, self.psf_size, self.psf_size, 5], device='cuda', dtype=self.data_type)
  1251. ders_out[:, :, :, 0:3] = psfs_ders * photons[:, None, None, None]
  1252. ders_out[:, :, :, 3] = psfs
  1253. ders_out[:, :, :, 4] = torch.ones_like(psfs)
  1254. return ders_out, psfs_out
  1256. @deprecated(reason='parallel version of v1, but not compatible with the _pre_compute_v3 and simulate_v3')
  1257. def _compute_derivative_v2(self, x, y, z, photons, bgs):
  1258. """
  1259. Calculate the analytical derivatives of the PSFs at given parameters with respect to x,y,z,photons,bg
  1260. """
  1262. n_mol = x.shape[0]
  1263. objstage = torch.zeros(x.shape[0], dtype=self.data_type, device='cuda')
  1265. # batch_size in parallel to save GPU memory
  1266. slice_list = []
  1267. batch_size = 100
  1268. for i in np.arange(0, n_mol, batch_size):
  1269. slice_list.append(slice(i, min(i + batch_size, n_mol)))
  1271. psfs = torch.zeros([n_mol, self.psf_size, self.psf_size], device='cuda', dtype=self.data_type)
  1272. psfs_ders = torch.zeros([n_mol, self.psf_size, self.psf_size, 3], device='cuda', dtype=self.data_type)
  1273. for slice_tmp in slice_list:
  1274. length_tmp = slice_tmp.stop - slice_tmp.start
  1275. position_phase = torch.empty([length_tmp, self.npupil, self.npupil], dtype=self.complex_type, device='cuda')
  1277. idx_0 = torch.where(z[slice_tmp] + self.zemit0 >= 0)[0]
  1278. phase_xyz_tmp = -y[slice_tmp][idx_0][:, None, None] * self.wavevector[0][None] - x[slice_tmp][idx_0][:, None, None] * \
  1279. self.wavevector[1][None] + \
  1280. (z[slice_tmp][idx_0] + self.zemit0)[:, None, None] * self.wavevectorzmed[None]
  1281. position_phase[idx_0, :, :] = torch.exp(
  1282. 1j * (phase_xyz_tmp + (objstage[slice_tmp][idx_0][:, None, None] + self.objstage0) *
  1283. self.wavevectorzimm[None]))
  1285. idx_1 = torch.where(z[slice_tmp] + self.zemit0 < 0)[0]
  1286. phase_xyz_tmp = -y[slice_tmp][idx_1][:, None, None] * self.wavevector[0][None] - x[slice_tmp][idx_1][:, None, None] * \
  1287. self.wavevector[1][None]
  1288. position_phase[idx_1, :, :] = torch.exp(
  1289. 1j * (phase_xyz_tmp + (objstage[slice_tmp][idx_1][:, None, None] + self.objstage0 + z[slice_tmp][idx_1][:, None, None]
  1290. + self.zemit0) * self.wavevectorzimm[None]))
  1292. pupil_tmp = position_phase[None, None] * self.pupilmatrix[:, :, None]
  1293. inter_image = torch.transpose(self.czt_parallel(pupil_tmp, self.ay, self.by, self.dy), -1, -2)
  1294. field_matrix = torch.transpose(self.czt_parallel(inter_image, self.ax, self.bx, self.dx), -1, -2)
  1295. psfs[slice_tmp] += 1 / 3 * torch.sum((torch.abs(field_matrix[:, :])) ** 2, dim=(0, 1))
  1297. # derivatives with respect to x,y,z
  1298. pupil_tmp_x = -1j * self.wavevector[1] * position_phase[None, None] * self.pupilmatrix[:, :, None]
  1299. inter_image_x = torch.transpose(self.czt_parallel(pupil_tmp_x, self.ay, self.by, self.dy), -1, -2)
  1300. field_matrix_x = torch.transpose(self.czt_parallel(inter_image_x, self.ax, self.bx, self.dx), -1, -2)
  1301. psfs_ders[slice_tmp, :, :, 0] += 2 / 3 * torch.sum(torch.real(torch.conj(field_matrix) * field_matrix_x),
  1302. dim=(0, 1))
  1304. pupil_tmp_y = -1j * self.wavevector[0] * position_phase[None, None] * self.pupilmatrix[:, :, None]
  1305. inter_image_y = torch.transpose(self.czt_parallel(pupil_tmp_y, self.ay, self.by, self.dy), -1, -2)
  1306. field_matrix_y = torch.transpose(self.czt_parallel(inter_image_y, self.ax, self.bx, self.dx), -1, -2)
  1307. psfs_ders[slice_tmp, :, :, 1] += 2 / 3 * torch.sum(torch.real(torch.conj(field_matrix) * field_matrix_y),
  1308. dim=(0, 1))
  1310. pupil_tmp_z = torch.empty([2, 3, length_tmp, self.npupil, self.npupil], dtype=self.complex_type, device='cuda')
  1311. pupil_tmp_z[:, :, idx_0] = 1j * self.wavevectorzmed * position_phase[idx_0] * self.pupilmatrix[:, :, None]
  1312. pupil_tmp_z[:, :, idx_1] = 1j * self.wavevectorzimm * position_phase[idx_1] * self.pupilmatrix[:, :, None]
  1313. inter_image_z = torch.transpose(self.czt_parallel(pupil_tmp_z, self.ay, self.by, self.dy), -1, -2)
  1314. field_matrix_z = torch.transpose(self.czt_parallel(inter_image_z, self.ax, self.bx, self.dx), -1, -2)
  1315. psfs_ders[slice_tmp, :, :, 2] += 2 / 3 * torch.sum(torch.real(torch.conj(field_matrix) * field_matrix_z),
  1316. dim=(0, 1))
  1317. psfs /= self.norm_intensity
  1318. psfs_ders /= self.norm_intensity
  1320. # otf rescale
  1321. if self.otf_rescale_xy[0] or self.otf_rescale_xy[1] != 0:
  1322. psfs = self.otf_rescale(psfdata=psfs, sigma_xy=self.otf_rescale_xy)
  1323. psfs_ders[:, :, :, 0] = self.otf_rescale(psfdata=psfs_ders[:, :, :, 0], sigma_xy=self.otf_rescale_xy)
  1324. psfs_ders[:, :, :, 1] = self.otf_rescale(psfdata=psfs_ders[:, :, :, 1], sigma_xy=self.otf_rescale_xy)
  1325. psfs_ders[:, :, :, 2] = self.otf_rescale(psfdata=psfs_ders[:, :, :, 2], sigma_xy=self.otf_rescale_xy)
  1327. psfs_out = ailoc.common.gpu(psfs * photons[:, None, None] + bgs[:, None, None])
  1328. ders_out = torch.zeros([n_mol, self.psf_size, self.psf_size, 5], device='cuda', dtype=self.data_type)
  1329. ders_out[:, :, :, 0:3] = psfs_ders * photons[:, None, None, None]
  1330. ders_out[:, :, :, 3] = psfs
  1331. ders_out[:, :, :, 4] = torch.ones_like(psfs)
  1333. return ders_out, psfs_out
  1335. def _compute_derivative(self, x, y, z, photons, bgs):
  1336. """
  1337. Calculate the analytical derivatives of the PSFs at given parameters with respect to x,y,z,photons,bg
  1338. """
  1340. n_mol = x.shape[0]
  1341. objstage = torch.zeros(x.shape[0], dtype=self.data_type, device='cuda')
  1343. zernike_phase = torch.exp(1j * 2 * np.pi *
  1344. torch.sum(self.zernike_coef[:, None, None]*self.allzernikes, dim=0)
  1345. / self.wavelength)
  1346. pupilmatrix = (self.amplitude[None, None] *
  1347. zernike_phase[None, None] *
  1348. self.polarizationvector)[:, :, None]
  1350. # batch_size in parallel to save GPU memory
  1351. slice_list = []
  1352. batch_size = 100
  1353. for i in np.arange(0, n_mol, batch_size):
  1354. slice_list.append(slice(i, min(i + batch_size, n_mol)))
  1356. psfs = torch.zeros([n_mol, self.psf_size, self.psf_size], device='cuda', dtype=self.data_type)
  1357. psfs_ders = torch.zeros([n_mol, self.psf_size, self.psf_size, 3], device='cuda', dtype=self.data_type)
  1358. for slice_tmp in slice_list:
  1359. # length_tmp = slice_tmp.stop - slice_tmp.start
  1360. # position_phase = torch.empty([length_tmp, self.npupil, self.npupil], dtype=self.complex_type, device='cuda')
  1362. # idx_0 = torch.where(z[slice_tmp] + self.zemit0 >= 0)[0]
  1363. # phase_xyz_tmp = -y[slice_tmp][idx_0][:, None, None] * self.wavevector[0][None] - x[slice_tmp][idx_0][:, None, None] * \
  1364. # self.wavevector[1][None] + \
  1365. # (z[slice_tmp][idx_0] + self.zemit0)[:, None, None] * self.wavevectorzmed[None]
  1366. # position_phase[idx_0, :, :] = torch.exp(
  1367. # 1j * (phase_xyz_tmp + (objstage[slice_tmp][idx_0][:, None, None] + self.objstage0) *
  1368. # self.wavevectorzimm[None]))
  1369. #
  1370. # idx_1 = torch.where(z[slice_tmp] + self.zemit0 < 0)[0]
  1371. # phase_xyz_tmp = -y[slice_tmp][idx_1][:, None, None] * self.wavevector[0][None] - x[slice_tmp][idx_1][:, None, None] * \
  1372. # self.wavevector[1][None]
  1373. # position_phase[idx_1, :, :] = torch.exp(
  1374. # 1j * (phase_xyz_tmp + (objstage[slice_tmp][idx_1][:, None, None] + self.objstage0 + z[slice_tmp][idx_1][:, None, None]
  1375. # + self.zemit0) * self.wavevectorzimm[None]))
  1377. phase_xyz_tmp = -y[slice_tmp][:, None, None] * self.wavevector[0][None] - x[slice_tmp][:, None, None] * \
  1378. self.wavevector[1][None] + (z[slice_tmp] + self.zemit0)[:, None, None] * self.wavevectorzmed[None]
  1379. position_phase = torch.exp(1j * (phase_xyz_tmp + (objstage[slice_tmp][:, None, None] + self.objstage0) *
  1380. self.wavevectorzimm[None]))
  1382. pupil_tmp = position_phase[None, None] * pupilmatrix
  1383. inter_image = torch.transpose(self.czt_parallel(pupil_tmp, self.ay, self.by, self.dy), -1, -2)
  1384. field_matrix = torch.transpose(self.czt_parallel(inter_image, self.ax, self.bx, self.dx), -1, -2)
  1385. psfs[slice_tmp] += 1 / 3 * torch.sum((torch.abs(field_matrix[:, :])) ** 2, dim=(0, 1))
  1387. # derivatives with respect to x,y,z
  1388. pupil_tmp_x = -1j * self.wavevector[1] * position_phase[None, None] * pupilmatrix
  1389. inter_image_x = torch.transpose(self.czt_parallel(pupil_tmp_x, self.ay, self.by, self.dy), -1, -2)
  1390. field_matrix_x = torch.transpose(self.czt_parallel(inter_image_x, self.ax, self.bx, self.dx), -1, -2)
  1391. psfs_ders[slice_tmp, :, :, 0] += 2 / 3 * torch.sum(torch.real(torch.conj(field_matrix) * field_matrix_x),
  1392. dim=(0, 1))
  1394. pupil_tmp_y = -1j * self.wavevector[0] * position_phase[None, None] * pupilmatrix
  1395. inter_image_y = torch.transpose(self.czt_parallel(pupil_tmp_y, self.ay, self.by, self.dy), -1, -2)
  1396. field_matrix_y = torch.transpose(self.czt_parallel(inter_image_y, self.ax, self.bx, self.dx), -1, -2)
  1397. psfs_ders[slice_tmp, :, :, 1] += 2 / 3 * torch.sum(torch.real(torch.conj(field_matrix) * field_matrix_y),
  1398. dim=(0, 1))
  1400. # pupil_tmp_z = torch.empty([2, 3, length_tmp, self.npupil, self.npupil], dtype=self.complex_type, device='cuda')
  1401. # pupil_tmp_z[:, :, idx_0] = 1j * self.wavevectorzmed * position_phase[idx_0] * pupilmatrix
  1402. # pupil_tmp_z[:, :, idx_1] = 1j * self.wavevectorzimm * position_phase[idx_1] * pupilmatrix
  1403. pupil_tmp_z = 1j * self.wavevectorzmed * position_phase * pupilmatrix
  1404. inter_image_z = torch.transpose(self.czt_parallel(pupil_tmp_z, self.ay, self.by, self.dy), -1, -2)
  1405. field_matrix_z = torch.transpose(self.czt_parallel(inter_image_z, self.ax, self.bx, self.dx), -1, -2)
  1406. psfs_ders[slice_tmp, :, :, 2] += 2 / 3 * torch.sum(torch.real(torch.conj(field_matrix) * field_matrix_z),
  1407. dim=(0, 1))
  1408. if self.focus_norm:
  1409. # normalize by focus intensity
  1410. psfs /= self.norm_intensity
  1411. psfs_ders /= self.norm_intensity
  1412. else:
  1413. # normalize by themselves
  1414. norm_factor = psfs.sum(dim=(-1, -2))
  1415. psfs /= norm_factor[:, None, None]
  1416. psfs_ders /= norm_factor[:, None, None, None]
  1418. # otf rescale
  1419. if self.otf_rescale_xy[0] or self.otf_rescale_xy[1]:
  1420. psfs = self.otf_rescale(psfdata=psfs, sigma_xy=self.otf_rescale_xy)
  1421. psfs_ders[:, :, :, 0] = self.otf_rescale(psfdata=psfs_ders[:, :, :, 0], sigma_xy=self.otf_rescale_xy)
  1422. psfs_ders[:, :, :, 1] = self.otf_rescale(psfdata=psfs_ders[:, :, :, 1], sigma_xy=self.otf_rescale_xy)
  1423. psfs_ders[:, :, :, 2] = self.otf_rescale(psfdata=psfs_ders[:, :, :, 2], sigma_xy=self.otf_rescale_xy)
  1425. psfs_out = ailoc.common.gpu(psfs * photons[:, None, None] + bgs[:, None, None])
  1426. ders_out = torch.zeros([n_mol, self.psf_size, self.psf_size, 5], device='cuda', dtype=self.data_type)
  1427. ders_out[:, :, :, 0:3] = psfs_ders * photons[:, None, None, None]
  1428. ders_out[:, :, :, 3] = psfs
  1429. ders_out[:, :, :, 4] = torch.ones_like(psfs)
  1431. return ders_out, psfs_out
  1433. def _torch_jacobian_derivative(self, x, y, z, photons, bgs):
  1434. """Calculate the derivatives of the PSFs at given parameters with respect to x,y,z,photons,bg
  1435. using torch.autograd.functional.jacobian"""
  1436. x.requires_grad = True
  1437. y.requires_grad = True
  1438. z.requires_grad = True
  1439. photons.requires_grad = True
  1440. bgs.requires_grad = True
  1441. jacobian = torch.autograd.functional.jacobian(self.simulate_parallel, (x, y, z, photons))
  1443. ders_out = torch.zeros([x.shape[0], self.psf_size, self.psf_size, 5], device='cuda', dtype=self.data_type)
  1444. for i in range(len(jacobian)):
  1445. for j in range(x.shape[0]):
  1446. ders_out[j, :, :, i] = jacobian[i][j, :, :, j]
  1447. ders_out[:, :, :, 4] = torch.ones([x.shape[0], self.psf_size, self.psf_size])
  1449. return ders_out
  1451. def optimize_crlb(self, x, y, z, photons, bgs, tolerance):
  1452. """
  1453. Optimize the CRLB of this PSF model with respect to zernike coefficients
  1454. at give positions, photons and backgrounds. The instance should be initialized with req_grad=True
  1456. Args:
  1457. x (torch.Tensor): x positions of the PSFs, unit nm
  1458. y (torch.Tensor): y positions of the PSFs, unit nm
  1459. z (torch.Tensor): z positions of the PSFs, unit nm
  1460. photons (torch.Tensor): photon counts of the PSFs, unit photon
  1461. bgs (torch.Tensor): background counts of the PSFs, unit photon
  1462. tolerance (float): stop criteria, the difference of CRLB between two iterations
  1464. Returns:
  1466. """
  1468. # # ADAM optimizer
  1469. # n_mol = x.shape[0]
  1470. # crlb_optimizer = torch.optim.Adam([self.zernike_coef], lr=0.1)
  1471. # loss_1 = 1e6
  1472. # iter_num = 1
  1473. # # for iter_num in range(iterations):
  1474. # while True:
  1475. # xyz_crlb, model = self.compute_crlb(x, y, z, photons, bgs)
  1476. # crlb_3d_avg = torch.sum(xyz_crlb[:, 0]**2 + xyz_crlb[:, 1]**2 + xyz_crlb[:, 2]**2)/n_mol
  1477. # print(f"iter: {iter_num}, crlb_3d_avg: {crlb_3d_avg}")
  1478. # crlb_optimizer.zero_grad()
  1479. # crlb_3d_avg.backward()
  1480. # crlb_optimizer.step()
  1481. # self._pre_compute()
  1482. # if torch.abs((loss_1 - crlb_3d_avg) / loss_1) < tolerance:
  1483. # break
  1484. # loss_1 = crlb_3d_avg.detach()
  1485. # iter_num += 1
  1486. # print('CRLB optimization done')
  1488. # LBFGS optimizer
  1489. n_mol = x.shape[0]
  1490. crlb_optimizer = torch.optim.LBFGS([self.zernike_coef], lr=0.1, max_iter=20, tolerance_grad=1e-7,
  1491. tolerance_change=1e-9)
  1492. loss_1 = 1e6
  1493. iter_num = 1
  1495. def closure():
  1496. crlb_optimizer.zero_grad()
  1497. xyz_crlb, model = self.compute_crlb(x, y, z, photons, bgs)
  1498. crlb_3d_avg = torch.sum(xyz_crlb[:, 0] ** 2 + xyz_crlb[:, 1] ** 2 + xyz_crlb[:, 2] ** 2) / n_mol
  1499. crlb_3d_avg.backward()
  1500. return crlb_3d_avg
  1502. while True:
  1503. crlb_3d_avg = crlb_optimizer.step(closure)
  1504. print(f"iter: {iter_num}, crlb_3d_avg: {crlb_3d_avg}")
  1505. self._pre_compute()
  1506. if torch.abs((loss_1 - crlb_3d_avg) / loss_1) < tolerance:
  1507. break
  1508. loss_1 = crlb_3d_avg.detach()
  1509. iter_num += 1
  1511. print('CRLB optimization done')
  1514. class VectorPSFTorch_2channel(VectorPSFTorch):
  1515. """
  1516. Vectorial PSF model with two detection channels. The two channels have different focal planes.
  1518. Args:
  1519. wavelength (float): emission wavelength, unit nm
  1520. na (float): numerical aperture of the objective
  1521. nmed (float): refractive index of the medium
  1522. nimm (float): refractive index of the immersion medium
  1523. pixel_size (float): camera pixel size, unit nm
  1524. psf_size (int): size of the PSF image, should be an odd number
  1525. z_range (tuple of float): z range of the PSF, unit nm
  1526. z_step (float): z step of the PSF, unit nm
  1527. zemit0 (float, optional): initial axial position of the emitter, default 0, unit nm
  1528. objstage0 (float, optional): initial axial position of the objective stage, default 0, unit nm
  1529. n_zernike (int, optional): number of zernike modes used to model the aberrations,
  1530. default 15 (up to 5th order excluding piston, tip and tilt)
  1531. zernike_coef (torch.Tensor, optional): initial zernike coefficients in a torch tensor,
  1532. default None which means all zeros
  1533. focus_norm (bool, optional): whether to normalize the PSF intensity by the in-focus intensity,
  1534. default True. If False, the PSF is normalized by itself.
  1535. Note that focus_norm=True is more relevant to localization microscopy,
  1536. while focus_norm=False is more relevant to particle tracking microscopy.
  1537. otf_rescale_xy (tuple of float, optional): sigma_x and sigma_y for Gaussian rescaling of the OTF,
  1538. default (0,0) means no rescaling
  1539. req_grad (bool, optional): whether the zernike coefficients require gradients,
  1540. default False. Set True when optimizing the coefficients.
  1541. data_type (torch.dtype, optional): data type for real numbers, default torch.float32
  1542. device (str, optional): device to use, default 'cuda'
  1543. """
  1545. def __init__(self, psf_params, req_grad=False, data_type=torch.float64, zernike_idx_learn=None):
  1546. super(VectorPSFTorch_2channel, self).__init__(psf_params, req_grad, data_type, zernike_idx_learn)
  1548. self.z_offset = psf_params.get('z_offset', 300) # nm
  1549. self.reflection_ratio = psf_params.get('reflection_ratio', 0.5) # ratio of the main channel/sum of two channels
  1551. def simulate(self, x, y, z, photons, objstage=None, zernike_coefs=None):
  1552. main_psfs_out = super().simulate(x, y, z, photons*self.reflection_ratio, objstage, zernike_coefs)
  1553. auxi_psfs_out = super().simulate(x, y, z+self.z_offset, photons*(1-self.reflection_ratio), objstage, zernike_coefs)
  1554. psfs_out = torch.cat((main_psfs_out[None], auxi_psfs_out[None]), dim=0)
  1556. # todo: need to think how to calculate the derivatives for two channels and the CRLB optimization, further combined with neural network
  1557. return psfs_out

vectorpsf.py at commit a603000, under GPL-3.0 · at the source

Overview

Authors: Shuang Fu1, Wei Shi1, Eugene A Katrukha2, Xi Chen3, Yue Fei1, Ke Fang1,4, Ruixiong Wang1, Tianlun Zhang1, Donghan Ma3, Yiming Li1
ORCID iDs: Shuang Fu, Wei Shi, Yiming Li
  1. Department of Biomedical Engineering, Southern University of Science and Technology, Shenzhen, China
  2. Cell Biology, Neurobiology and Biophysics, Department of Biology, Faculty of Science, Utrecht University, Utrecht, the Netherlands
  3. School of Optoelectronic Engineering and Instrumentation Science, Dalian University of Technology, Dalian, China
  4. Center for Health Research, Guangzhou Institutes of Biomedicine and Health, Chinese Academy of Sciences, Guangzhou, China
Institutions: Southern University of Science and Technology (China); Utrecht University (Netherlands); Dalian University of Technology (China); Chinese Academy of Sciences (China); Guangzhou Institutes of Biomedicine and Health (China)
Journal: Nature communications, volume 17, issue 1, article 6504
Dates: received 7 October 2025; accepted 30 April 2026; published online 16 May 2026
Type: Research article · Language: English
License: CC BY-NC-ND
Identifiers: DOI 10.1038/s41467-026-73045-9 · PMID 42143028 · PMCID PMC13376372 · OpenAlex W7161408484
Open access: gold, a free copy (OpenAlex)
Status: code verified
Categories: histology / microscopy (modality), cellular / molecular (subfield)
Methods: Spectral & time-frequency, Physiology & signal measures, Machine learning
Keywords: Super-resolution microscopy, Fluorescence imaging
MeSH: Imaging, Three-Dimensional*, Neural Networks, Computer*, Single Molecule Imaging*, Animals, Deep Learning, Microscopy, Fluorescence, Mitochondria, Neurons (* major topic)
Topic: Advanced Fluorescence Microscopy Techniques (Biophysics, Biochemistry, Genetics and Molecular Biology), according to OpenAlex
Funding: National Natural Science Foundation of China (National Science Foundation of China) (62375116); China Postdoctoral Science Foundation (GZC20240651, 2025M772887, GZC20250546, 2025T180788)
Citations: cited by 2 papers (Europe PMC); 56 references in the paper

Abstract

The abstract is not reproduced here: the paper's license (CC BY-NC-ND) does not allow it. Read it in the paper, at the publisher or on Europe PMC.

Repositories

Its files are read in the Code ↔ Paper reader above, with 16 matches between paragraphs and lines of code.

Li-Lab-SUSTech/LUNAR

License: GPL-3.0
State: the link answers, verified on 28 September 2026
Evidence: files inventoried
Commit: a6030006892a2be8a862aa3e8d7bf03f7d3a0f12, 29 June 2026
Languages: Python (46), Jupyter (6)
Size: 65 files, 52 scripts
Software Heritage: not archived
Found in: “Code availability”
Holds: README, license file, environment (requirements.txt), documentation, 6 notebooks
Not found: CITATION.cff, tests, continuous integration
Tools: PyTorch (35 files), NumPy (34 files), Matplotlib (13 files), SciPy (12 files), tifffile (10 files), imageio (6 files), napari (2 files), OpenCV (2 files), Pillow (2 files), pandas (1 file), scikit-image (1 file)
Availability: 1 check, the latest on 28 September 2026: the link answers
  • 28 September 2026: the link answers
54 files
At the source: github.com/Li-Lab-SUSTech/LUNAR

jdeschamps/htSMLM

License: LGPL-2.1
State: the link answers, verified on 28 September 2026
Evidence: files inventoried
Commit: 30b6bfd923a13076c08011ee001cc3a0e83c5434, 8 January 2024
Languages: Java (88), Shell (3)
Size: 135 files, 91 scripts
Software Heritage: not archived
Found in: “Code availability”
Holds: README, license file, tests
Not found: CITATION.cff, environment file, continuous integration, documentation
Availability: 1 check, the latest on 28 September 2026: the link answers
  • 28 September 2026: the link answers
93 files
At the source: github.com/jdeschamps/htSMLM

Zenodo 19586991

License: GPL-3.0
State: the link answers, verified on 28 September 2026
Evidence: files inventoried
Size: 1 file
Software Heritage: not checked
Found in: “Code availability”
Not found: README, license file, CITATION.cff, environment file, tests, continuous integration, documentation
Availability: 1 check, the latest on 28 September 2026: the link answers (HTTP 200)
  • 28 September 2026: the link answers (HTTP 200)
At the source:

Code availability statement

The paper has a code availability statement. Its license (CC BY-NC-ND) does not allow reproducing it here; in short, from what the harvester recognized in it:

Read it in the paper: doi.org/10.1038/s41467-026-73045-9.

Tracing map

Proposed by the machine: these links were found in the paper and verified at the source, without human review. The map will receive a Zenodo DOI once one of the paper's authors has validated it with their ORCID.

What the map holds:

  • 3 repositories of the authors' code, each at its verified commit, with its license and how the link was found in the paper;
  • 143 scripts, each with its path and the digest of its content;
  • 16 matches between paragraphs of the paper and lines of the code (method lexical-v1);
  • neither the text of the paper nor the code itself.

Its JSON (tracing-map.json) is deposited on Zenodo with its DOI once the map is validated.

Data

Datasets cited

Data availability statement

The paper has a data availability statement. Its license (CC BY-NC-ND) does not allow reproducing it here; in short, from what the harvester recognized in it:

  • it points to a dataset: Zenodo 14709467
  • it says that the data are available on request

Read it in the paper: doi.org/10.1038/s41467-026-73045-9.

Versions

The history of this record: each version stored by the harvester or made by a correction of its authors or of the maintainers of its code, and what changed in its facts. The texts of the paper (its abstract, its availability statements) are not part of it; versions that changed only those are not listed.

Version 1, 28 September 2026: the first record

Recorded: type, language, journal, volume, issue, pages, dates, 10 authors, 2 keywords, 8 MeSH terms, 2 funders, 47 references.

Cite

This paper

Fu, S., Shi, W., Katrukha, E. A., Chen, X., Fei, Y., Fang, K., Wang, R., Zhang, T., Ma, D., & Li, Y. (2026). Aberration-aware 3D localization microscopy via self-supervised neural-physics learning. Nature communications, 17(1), 6504. https://doi.org/10.1038/s41467-026-73045-9

BibTeX

@article{fu2026aberration,
author = {Fu, Shuang and Shi, Wei and Katrukha, Eugene A and Chen, Xi and Fei, Yue and Fang, Ke and Wang, Ruixiong and Zhang, Tianlun and Ma, Donghan and Li, Yiming},
title = {{Aberration-aware 3D localization microscopy via self-supervised neural-physics learning}},
journal = {Nature communications},
year = {2026},
month = may,
volume = {17},
number = {1},
pages = {6504},
publisher = {Nature Publishing Group},
issn = {2041-1723},
doi = {10.1038/s41467-026-73045-9},
url = {https://doi.org/10.1038/s41467-026-73045-9},
pmid = {42143028},
pmcid = {PMC13376372}
}

RIS

TY - JOUR
AU - Fu, Shuang
AU - Shi, Wei
AU - Katrukha, Eugene A
AU - Chen, Xi
AU - Fei, Yue
AU - Fang, Ke
AU - Wang, Ruixiong
AU - Zhang, Tianlun
AU - Ma, Donghan
AU - Li, Yiming
TI - Aberration-aware 3D localization microscopy via self-supervised neural-physics learning
T2 - Nature communications
J2 - Nat Commun
PY - 2026
DA - 2026/05/16
VL - 17
IS - 1
SP - 6504
SN - 2041-1723
PB - Nature Publishing Group
DO - 10.1038/s41467-026-73045-9
UR - https://doi.org/10.1038/s41467-026-73045-9
LA - en
ER -

CSL-JSON

{
"id": "10.1038/s41467-026-73045-9",
"type": "article-journal",
"title": "Aberration-aware 3D localization microscopy via self-supervised neural-physics learning",
"container-title": "Nature communications",
"author": [
{
"family": "Fu",
"given": "Shuang"
},
{
"family": "Shi",
"given": "Wei"
},
{
"family": "Katrukha",
"given": "Eugene A"
},
{
"family": "Chen",
"given": "Xi"
},
{
"family": "Fei",
"given": "Yue"
},
{
"family": "Fang",
"given": "Ke"
},
{
"family": "Wang",
"given": "Ruixiong"
},
{
"family": "Zhang",
"given": "Tianlun"
},
{
"family": "Ma",
"given": "Donghan"
},
{
"family": "Li",
"given": "Yiming"
}
],
"container-title-short": "Nat Commun",
"volume": "17",
"issue": "1",
"page": "6504",
"DOI": "10.1038/s41467-026-73045-9",
"PMID": "42143028",
"PMCID": "PMC13376372",
"ISSN": "2041-1723",
"publisher": "Nature Publishing Group",
"URL": "https://doi.org/10.1038/s41467-026-73045-9",
"language": "en",
"issued": {
"date-parts": [
[
2026,
5,
16
]
]
}
}

The tracing map gets a citation of its own once an author has validated it and it has a DOI.

Similar papers

The papers with a page that share the most with this one: the tools found in their code, their categories, datasets, cited references and authors, the rarest counting most.

[1] doi:10.1371/journal.pcbi.1014571 [code]
SynAPSeg: A novel dataset and image analysis framework for deep learning-based synapse detection and quantification.
Journal: PLoS computational biology
In common: napari, imageio, tifffile, 8 other tools
[2] doi:10.1038/s41467-026-73389-2 [code]
Physics-informed multi-encoder adaptive optics enables rapid aberration correction for intravital microscopy of deep complex tissue.
Journal: Nature communications
In common: OpenCV, SciPy, Matplotlib, 1 other tool, histology / microscopy, 6 references
[3] doi:10.64898/2026.03.12.711239 [code]
Back-illumination Phase Imaging Enables Nanoscale Drift Stabilization in Non-transparent Biological Tissues
Journal: bioRxiv (preprint)
In common: napari, tifffile, OpenCV, 6 other tools, histology / microscopy, 1 reference
[4] doi: [code]
Real-time closed-loop feedback system for mouse mesoscale cortical signal and movement control
Journal: eLife
In common: napari, imageio, tifffile, 7 other tools
[5] doi:10.1038/s41467-026-72709-w [code]
An epifluorescence microscope design for naturalistic behavior and cellular activity in freely moving Caenorhabditis elegans.
Journal: Nature communications
In common: imageio, tifffile, OpenCV, 7 other tools, histology / microscopy, cellular / molecular
[6] doi:10.1016/j.isci.2026.116769 [code]
Graph-based modeling of optical system enables adaptive optics on dynamic samples with self-calibration.
Journal: iScience
In common: imageio, tifffile, scikit-image, 5 other tools, 2 references
[7] doi:10.1038/s41467-026-71614-6 [code]
Interferometric ultra-high resolution 3D imaging through brain sections.
Journal: Nature communications
In common: histology / microscopy, cellular / molecular, 7 references
[8] doi:10.1038/s41598-026-57519-w [code]
Automated segmentation of neurons and spinal cord structures in immunofluorescence images using SpineDL.
Journal: Scientific reports
In common: imageio, tifffile, OpenCV, 7 other tools, histology / microscopy
[9] doi:10.1038/s41467-026-75352-7 [code]
Mechanosensory encoding of surface mechanics optimizes locomotion.
Journal: Nature communications
In common: napari, imageio, tifffile, 6 other tools, cellular / molecular
[10] doi:10.1364/boe.605322 [code]
Generalized plaque digitization framework for multi-dimensional mesoscopic images.
Journal: Biomedical optics express
In common: imageio, tifffile, OpenCV, 7 other tools

Contribute

The authors of this paper can claim it, correct its record and validate its tracing map, and the maintainers of its code (its owner, or a public member of its organization) correct what it says of their repository; anyone signed in can ask for its removal. Every request goes to OSCR's own machine, which answers it; your account page follows them.

Sign in with ORCID to claim this paper as one of its authors, correct its record or validate its tracing map: when the paper's metadata lists your ORCID iD, you are recognized at once. Maintainers of its code: sign in with GitHub, then claim the repository on your account page.

Request its removal

To ask OSCR to remove this record, the copies of its authors' scripts or its tracing map, use the removal request page: signed in, you say who you are, what to remove and why, then review and confirm the request. Published rules decide every request (how).

Discussion, reproductions, activity

Discussion: questions and error reports about this paper and its code, from signed-in readers and its authors. It opens with sign-in.

Reproductions: reports from readers who ran the authors' code: what they reproduced, with which environment, commit and data. It opens with sign-in.

Activity: what happens around this paper: new versions of its record, its map's validation, discussions and reproductions. It opens with sign-in.