Module pyminflux.correct
Correction/restoration functions.
Expand source code
# Copyright (c) 2022 - 2024 D-BSSE, ETH Zurich.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
__doc__ = "Correction/restoration functions."
__all__ = [
"drift_correction_time_windows_2d",
"drift_correction_time_windows_3d",
]
from ._correct import drift_correction_time_windows_2d, drift_correction_time_windows_3d
Functions
def drift_correction_time_windows_2d(x: numpy.ndarray, y: numpy.ndarray, t: numpy.ndarray, sxy: float, rx: Optional[tuple] = None, ry: Optional[tuple] = None, T: Optional[float] = None, tid: Optional[numpy.ndarray] = None)
-
Estimate 2D drift correction based on auto-correlation.
Reimplemented (with modifications) from:
- [paper] Ostersehlt, L.M., Jans, D.C., Wittek, A. et al. DNA-PAINT MINFLUX nanoscopy. Nat Methods 19, 1072-1075 (2022). https://doi.org/10.1038/s41592-022-01577-1
- [code] https://zenodo.org/record/6563100
Parameters
x
:np.ndarray
- Array of localization x coordinates.
y
:np.ndarray
- Array of localization y coordinates.
t
:np.ndarray
- Array of localization time points.
sxy
:float (Default = 1.0)
- Resolution in nm in both the x and y direction.
rx
:tuple (Optional)
- (min, max) boundaries for the x coordinates. If omitted, it will default to (x.min(), x.max()).
ry
:float (Optional)
- (min, max) boundaries for the y coordinates. If omitted, it will default to (y.min(), y.max()).
T
:float (Optional)
- Time window for analysis.
tid
:np.ndarray (Optional)
- Only used if T is None. The unique trace IDs are used to calculate the time window for analysis using some heuristics.
Expand source code
def drift_correction_time_windows_2d( x: np.ndarray, y: np.ndarray, t: np.ndarray, sxy: float, rx: Optional[tuple] = None, ry: Optional[tuple] = None, T: Optional[float] = None, tid: Optional[np.ndarray] = None, ): """Estimate 2D drift correction based on auto-correlation. Reimplemented (with modifications) from: * [paper] Ostersehlt, L.M., Jans, D.C., Wittek, A. et al. DNA-PAINT MINFLUX nanoscopy. Nat Methods 19, 1072-1075 (2022). https://doi.org/10.1038/s41592-022-01577-1 * [code] https://zenodo.org/record/6563100 Parameters ---------- x: np.ndarray Array of localization x coordinates. y: np.ndarray Array of localization y coordinates. t: np.ndarray Array of localization time points. sxy: float (Default = 1.0) Resolution in nm in both the x and y direction. rx: tuple (Optional) (min, max) boundaries for the x coordinates. If omitted, it will default to (x.min(), x.max()). ry: float (Optional) (min, max) boundaries for the y coordinates. If omitted, it will default to (y.min(), y.max()). T: float (Optional) Time window for analysis. tid: np.ndarray (Optional) Only used if T is None. The unique trace IDs are used to calculate the time window for analysis using some heuristics. """ if T is None and tid is None: raise ValueError("If T is not defined, the array of TIDs must be provided.") # Make sure we are working with NumPy arrays x = np.array(x) y = np.array(y) t = np.array(t) if tid is not None: tid = np.array(tid) # Make sure we have valid ranges if rx is None: rx = (x.min(), x.max()) if ry is None: ry = (y.min(), y.max()) # Heuristics to define the length of the time window if T is None: rt = (t[0], t[-1]) T = len(np.unique(tid)) * np.diff(rx)[0] * np.diff(ry)[0] / 3e6 T = np.min([T, np.diff(rt)[0] / 2, 3600]) # At least two time windows T = max([T, 600]) # And at least 10 minutes long # Number of time windows to use Rt = [t[0], t[-1]] Ns = int(np.floor((Rt[1] - Rt[0]) / T)) # total number of time windows assert Ns > 1, "At least two time windows are needed, please reduce T" # Center of mass CR = 10 # Maximum number of frames in the cross-correlation D = Ns # Weight with distance w = np.linspace(1, 0.5, D) # Regularization term (roughness) l = 0.1 # get dimensions Nx = int(np.ceil((rx[1] - rx[0]) / sxy)) Ny = int(np.ceil((ry[1] - ry[0]) / sxy)) c = np.round(np.array([Ny, Nx]) / 2) # Create all the histograms h = [None] * Ns ti = np.zeros(Ns) # Average times of the snapshots for j in range(Ns): t0 = Rt[0] + j * T idx = (t >= t0) & (t < t0 + T) ti[j] = np.mean(t[idx]) hj, _, _, _ = render_xy( x[idx], y[idx], sx=sxy, sy=sxy, rx=rx, ry=ry, render_type="fixed_gaussian", fwhm=3 * sxy, ) h[j] = hj # Compute fourier transform of every histogram for j in range(Ns): h[j] = fftn(h[j]) # Compute cross-correlations dx = np.zeros((Ns, Ns)) dy = np.zeros((Ns, Ns)) dm = np.zeros((Ns, Ns), dtype=bool) dx0 = np.zeros(Ns - 1) dy0 = np.zeros(Ns - 1) for i in range(Ns - 1): hi = np.conj(h[i]) for j in range(i + 1, min(Ns, i + D)): # Either to Ns or maximally D more hj = ifftshift(np.real(ifftn(hi * h[j]))) yc = c[0] xc = c[1] # Centroid estimation gy = np.arange(yc - 2 * CR, yc + 2 * CR + 1).astype(int) gx = np.arange(xc - 2 * CR, xc + 2 * CR + 1).astype(int) gy, gx = np.meshgrid(gy, gx, indexing="ij") d = hj[gy, gx] gy = gy.flatten() gx = gx.flatten() d = d.flatten() - np.min(d) d = d / np.sum(d) for k in range(20): wc = np.exp( -4 * np.log(2) * ((xc - gx) ** 2 + (yc - gy) ** 2) / CR**2 ) n = np.sum(wc * d) xc = np.sum(gx * d * wc) / n yc = np.sum(gy * d * wc) / n sh = np.array([-1.0, 1.0]) * (np.array([yc, xc]) - c) dy[i, j] = sh[0] dx[i, j] = sh[1] dm[i, j] = True if j == i + 1: dy0[i] = sh[0] dx0[i] = sh[1] a, b = np.nonzero(dm) dx = dx[dm] dy = dy[dm] sx0 = np.cumsum(dx0) sy0 = np.cumsum(dy0) # Minimize cost function with some kind of regularization options = {"disp": False, "maxiter": 1e5} minimizer = lambda x: lse_distance(x, a, b, dx, w, l) res = minimize(minimizer, sx0, options=options, method="BFGS") sx = np.concatenate(([0], res.x)) minimizer = lambda x: lse_distance(x, a, b, dy, w, l) res = minimize(minimizer, sy0, options=options, method="BFGS") sy = np.concatenate(([0], res.x)) # Reduce by mean (so shift is minimal) sx = sx - np.mean(sx) sy = sy - np.mean(sy) # Multiply by pixel size sx = sx * sxy sy = sy * sxy # Create interpolants fx = interp1d(ti, sx, kind="slinear", fill_value="extrapolate") fy = interp1d(ti, sy, kind="slinear", fill_value="extrapolate") # Correct drift dx = fx(t) dy = fy(t) # Apply to the time points used to estimate the correction dxt = fx(ti) dyt = fy(ti) return dx, dy, dxt, dyt, ti, T
def drift_correction_time_windows_3d(x: numpy.ndarray, y: numpy.ndarray, z: numpy.ndarray, t: numpy.ndarray, sxyz: float, rx: Optional[tuple] = None, ry: Optional[tuple] = None, rz: Optional[tuple] = None, T: Optional[float] = None, tid: Optional[numpy.ndarray] = None)
-
Estimate 3D drift correction based on auto-correlation.
Reimplemented (with modifications) from:
- [paper] Ostersehlt, L.M., Jans, D.C., Wittek, A. et al. DNA-PAINT MINFLUX nanoscopy. Nat Methods 19, 1072-1075 (2022). https://doi.org/10.1038/s41592-022-01577-1
- [code] https://zenodo.org/record/6563100
Parameters
x
:np.ndarray
- Array of localization x coordinates.
y
:np.ndarray
- Array of localization y coordinates.
z
:np.ndarray
- Array of localization z coordinates.
t
:np.ndarray
- Array of localization time points.
sxyz
:float (Default = 1.0)
- Resolution in nm in x, y and z direction.
rx
:tuple (Optional)
- (min, max) boundaries for the x coordinates. If omitted, it will default to (x.min(), x.max()).
ry
:float (Optional)
- (min, max) boundaries for the y coordinates. If omitted, it will default to (y.min(), y.max()).
rz
:float (Optional)
- (min, max) boundaries for the z coordinates. If omitted, it will default to (z.min(), z.max()).
T
:float (Optional)
- Time window for analysis.
tid
:np.ndarray (Optional)
- Only used if T is None. The unique trace IDs are used to calculate the time window for analysis using some heuristics.
Expand source code
def drift_correction_time_windows_3d( x: np.ndarray, y: np.ndarray, z: np.ndarray, t: np.ndarray, sxyz: float, rx: Optional[tuple] = None, ry: Optional[tuple] = None, rz: Optional[tuple] = None, T: Optional[float] = None, tid: Optional[np.ndarray] = None, ): """Estimate 3D drift correction based on auto-correlation. Reimplemented (with modifications) from: * [paper] Ostersehlt, L.M., Jans, D.C., Wittek, A. et al. DNA-PAINT MINFLUX nanoscopy. Nat Methods 19, 1072-1075 (2022). https://doi.org/10.1038/s41592-022-01577-1 * [code] https://zenodo.org/record/6563100 Parameters ---------- x: np.ndarray Array of localization x coordinates. y: np.ndarray Array of localization y coordinates. z: np.ndarray Array of localization z coordinates. t: np.ndarray Array of localization time points. sxyz: float (Default = 1.0) Resolution in nm in x, y and z direction. rx: tuple (Optional) (min, max) boundaries for the x coordinates. If omitted, it will default to (x.min(), x.max()). ry: float (Optional) (min, max) boundaries for the y coordinates. If omitted, it will default to (y.min(), y.max()). rz: float (Optional) (min, max) boundaries for the z coordinates. If omitted, it will default to (z.min(), z.max()). T: float (Optional) Time window for analysis. tid: np.ndarray (Optional) Only used if T is None. The unique trace IDs are used to calculate the time window for analysis using some heuristics. """ if T is None and tid is None: raise ValueError("If T is not defined, the array of TIDs must be provided.") # Make sure we are working with NumPy arrays x = np.array(x) y = np.array(y) z = np.array(z) t = np.array(t) if tid is not None: tid = np.array(tid) # Make sure we have valid ranges if rx is None: rx = (x.min(), x.max()) if ry is None: ry = (y.min(), y.max()) if rz is None: rz = (z.min(), z.max()) # Heuristics to define the length of the time window if T is None: rt = (t[0], t[-1]) T = len(np.unique(tid)) * np.diff(rx)[0] * np.diff(ry)[0] / 3e6 T = np.min([T, np.diff(rt)[0] / 2, 3600]) # At least two time windows T = max([T, 600]) # And at least 10 minutes long # Number of time windows to use Rt = [t[0], t[-1]] Ns = int(np.floor((Rt[1] - Rt[0]) / T)) # total number of time windows assert Ns > 1, "At least two time windows are needed, please reduce T" # Center of mass CR = 8 # Maximum number of frames in the cross-correlation D = Ns # Weight with distance w = np.linspace(1, 0.2, D) # Regularization term (roughness) l = 0.01 # get dimensions Nx = int(np.ceil((rx[1] - rx[0]) / sxyz)) Ny = int(np.ceil((ry[1] - ry[0]) / sxyz)) Nz = int(np.ceil((rz[1] - rz[0]) / sxyz)) c = np.round(np.array([Nz, Ny, Nx]) / 2) # Create all the histograms h = [None] * Ns ti = np.zeros(Ns) # Average times of the snapshots for j in range(Ns): t0 = Rt[0] + j * T idx = (t >= t0) & (t < t0 + T) ti[j] = np.mean(t[idx]) hj, _, _, _, _ = render_xyz( x[idx], y[idx], z[idx], sx=sxyz, sy=sxyz, sz=sxyz, rx=rx, ry=ry, rz=rz, render_type="fixed_gaussian", fwhm=3 * sxyz, ) h[j] = hj # Compute fourier transform of every histogram for j in range(Ns): h[j] = fftn(h[j]) # Compute cross-correlations dx = np.zeros((Ns, Ns)) dy = np.zeros((Ns, Ns)) dz = np.zeros((Ns, Ns)) dm = np.zeros((Ns, Ns), dtype=bool) dx0 = np.zeros(Ns - 1) dy0 = np.zeros(Ns - 1) dz0 = np.zeros(Ns - 1) for i in range(Ns - 1): hi = np.conj(h[i]) for j in range(i + 1, min(Ns, i + D)): # either to Ns or maximally D more hj = ifftshift(np.real(ifftn(hi * h[j]))) zc = c[0] yc = c[1] xc = c[2] # Centroid estimation gz = np.arange(zc - 2 * CR, zc + 2 * CR + 1).astype(int) gy = np.arange(yc - 2 * CR, yc + 2 * CR + 1).astype(int) gx = np.arange(xc - 2 * CR, xc + 2 * CR + 1).astype(int) gz, gy, gx = np.meshgrid(gz, gy, gx, indexing="ij") d = hj[gz, gy, gx] gz = gz.flatten() gy = gy.flatten() gx = gx.flatten() d = d.flatten() - np.min(d) d = d / np.sum(d) for k in range(20): wc = np.exp( -4 * np.log(2) * ((xc - gx) ** 2 + (yc - gy) ** 2 + (zc - gz) ** 2) / CR**2 ) n = np.sum(wc * d) xc = np.sum(gx * d * wc) / n yc = np.sum(gy * d * wc) / n zc = np.sum(gz * d * wc) / n sh = np.array([1.0, -1.0, 1.0]) * (np.array([zc, yc, xc]) - c) dz[i, j] = sh[0] dy[i, j] = sh[1] dx[i, j] = sh[2] dm[i, j] = True if j == i + 1: dz0[i] = sh[0] dy0[i] = sh[1] dx0[i] = sh[2] a, b = np.nonzero(dm) dx = dx[dm] dy = dy[dm] dz = dz[dm] sx0 = np.cumsum(dx0) sy0 = np.cumsum(dy0) sz0 = np.cumsum(dz0) # Minimize cost function with some kind of regularization options = {"disp": False, "maxiter": 1e5} minimizer = lambda x: lse_distance(x, a, b, dx, w, l) res = minimize(minimizer, sx0, options=options, method="BFGS") sx = np.concatenate(([0], res.x)) minimizer = lambda x: lse_distance(x, a, b, dy, w, l) res = minimize(minimizer, sy0, options=options, method="BFGS") sy = np.concatenate(([0], res.x)) minimizer = lambda x: lse_distance(x, a, b, dz, w, l) res = minimize(minimizer, sz0, options=options, method="BFGS") sz = np.concatenate(([0], res.x)) # Reduce by mean (so shift is minimal) sx = sx - np.mean(sx) sy = sy - np.mean(sy) sz = sz - np.mean(sz) # Multiply by voxel size sx = sx * sxyz sy = sy * sxyz sz = sz * sxyz # Create interpolants fx = interp1d(ti, sx, kind="slinear", fill_value="extrapolate") fy = interp1d(ti, sy, kind="slinear", fill_value="extrapolate") fz = interp1d(ti, sz, kind="slinear", fill_value="extrapolate") # Correct drift dx = fx(t) dy = fy(t) dz = fz(t) # Apply to the time points used to estimate the correction dxt = fx(ti) dyt = fy(ti) dzt = fz(ti) return dx, dy, dz, dxt, dyt, dzt, ti, T