From c221ee72a514b47cdd4f2e8322801a5a885b367b Mon Sep 17 00:00:00 2001 From: LREIN663 Date: Fri, 27 Mar 2026 15:55:29 +0100 Subject: [PATCH 01/19] Integrated COMET into Picasso render and made it the default drift correction. --- picasso/comet.py | 852 ++++++++++++++++++++++++++++++++++++++++++ picasso/gui/render.py | 115 +++++- 2 files changed, 965 insertions(+), 2 deletions(-) create mode 100644 picasso/comet.py diff --git a/picasso/comet.py b/picasso/comet.py new file mode 100644 index 00000000..11321927 --- /dev/null +++ b/picasso/comet.py @@ -0,0 +1,852 @@ +import numpy as np +import pandas as pd +import matplotlib.pyplot as plt +from numba import cuda +from scipy.ndimage import convolve +from scipy.optimize import minimize +import time +import math +from scipy.spatial import cKDTree + +from . import lib, __version__ + + + + +def comet( + locs: pd.DataFrame, + info: list[dict], + locs_per_segment: int, + max_drift_nm: float, + max_locs_per_segment: int | None = None, + *, + initial_sigma_nm: float | None = None, + target_sigma_nm: float = 10.0, + boxcar_width: int = 3, + drift_max_bound_factor: float = 2.0, + interpolation_method: str = "cubic", + mode: str = "cuda", + display: bool = False, + progress=None, # currently unused +) -> tuple[pd.DataFrame, list[dict], pd.DataFrame]: + """Apply COMET undrifting to the localizations. + + Parameters + ---------- + locs : pd.DataFrame + Localization table. Must contain columns 'x', 'y', 'frame'. + Optional: 'z'. + x/y are expected in camera pixels, z in the native z-unit used by Picasso. + info : list of dict + Localization metadata. + locs_per_segment : int + Target number of localizations per segment for COMET segmentation. + max_drift_nm : float + Maximum expected drift in nm. + max_locs_per_segment : int or None, optional + Optional cap for localizations per segment. If negative, treated as None. + initial_sigma_nm : float or None, optional + Initial Gaussian sigma for COMET. If None, defaults to max_drift_nm / 3. + target_sigma_nm : float, optional + Final Gaussian sigma for refinement. + boxcar_width : int, optional + Smoothing width over segments. + drift_max_bound_factor : float, optional + Bound factor for optimizer. + interpolation_method : str, optional + Drift interpolation method. + mode : str, optional + Backend mode, only "cuda". + display : bool, optional + Whether to show diagnostic plots. + progress : optional + Placeholder for API compatibility with AIM. + + Returns + ------- + locs : pd.DataFrame + Undrift-corrected localizations. + new_info : list[dict] + Updated metadata. + drift : pd.DataFrame + Per-frame drift table with columns x, y, and optionally z. + Drift is returned in the same units as locs: + x/y in pixels, z in the original z-unit. + """ + locs = locs.copy() + + if max_locs_per_segment is not None and max_locs_per_segment < 0: + max_locs_per_segment = None + + pixelsize_nm = lib.get_from_metadata(info, "Pixelsize", raise_error=True) + + required_cols = {"x", "y", "frame"} + missing = required_cols - set(locs.columns) + if missing: + raise KeyError(f"Missing required columns for COMET: {sorted(missing)}") + + has_z = "z" in locs.columns + + # Normalize frames so they start at 0 for internal indexing. + # This is important because COMET later indexes drift by frame number. + frame = locs["frame"].to_numpy(dtype=np.int64) + frame0 = frame - frame.min() + + # Build COMET input array in nm + dataset = np.zeros((len(locs), 4), dtype=np.float64) + dataset[:, 0] = locs["x"].to_numpy(dtype=np.float64) * pixelsize_nm + dataset[:, 1] = locs["y"].to_numpy(dtype=np.float64) * pixelsize_nm + dataset[:, 2] = ( + locs["z"].to_numpy(dtype=np.float64) if has_z else 0.0 + ) + dataset[:, 3] = frame0 + + if initial_sigma_nm is None: + initial_sigma_nm = max_drift_nm / 3 + + drift_nm_with_frame = comet_run_kd( + dataset=dataset, + segmentation_mode=1, # fixed locs per segment + segmentation_var=locs_per_segment, + max_locs_per_segment=max_locs_per_segment, + initial_sigma_nm=initial_sigma_nm, + gt_drift=None, + display=display, + return_corrected_locs=False, + max_drift_nm=max_drift_nm, + target_sigma_nm=target_sigma_nm, + boxcar_width=boxcar_width, + drift_max_bound_factor=drift_max_bound_factor, + interpolation_method=interpolation_method, + mode=mode, + min_max_frames=(int(frame0.min()), int(frame0.max())), + ) + + # drift_nm_with_frame columns: dx_nm, dy_nm, dz_nm, frame0 + drift_nm = drift_nm_with_frame[:, :3] + + # Apply drift back to dataframe in original units + frame_idx = frame0.astype(np.int64) + locs["x"] = locs["x"].to_numpy(dtype=np.float64) - drift_nm[frame_idx, 0] / pixelsize_nm + locs["y"] = locs["y"].to_numpy(dtype=np.float64) - drift_nm[frame_idx, 1] / pixelsize_nm + if has_z: + locs["z"] = locs["z"].to_numpy(dtype=np.float64) - drift_nm[frame_idx, 2] + + # Build drift dataframe in Picasso-style units + drift_dict = { + "x": (drift_nm[:, 0] / pixelsize_nm).astype("float32"), + "y": (drift_nm[:, 1] / pixelsize_nm).astype("float32"), + } + if has_z: + drift_dict["z"] = drift_nm[:, 2].astype("float32") + + drift = pd.DataFrame(drift_dict) + + new_info_entry = { + "Generated by": f"Picasso v{__version__} COMET", + "Segmentation mode": "localizations per segment", + "Localizations per segment": locs_per_segment, + "Maximum drift (nm)": max_drift_nm, + "Initial sigma (nm)": float(initial_sigma_nm), + "Target sigma (nm)": float(target_sigma_nm), + "Boxcar width": int(boxcar_width), + "Max localizations per segment": ( + None if max_locs_per_segment is None else int(max_locs_per_segment) + ), + "Interpolation method": interpolation_method, + "Backend": mode, + } + new_info = info + [new_info_entry] + + return locs, new_info, drift + +def comet_run_kd(dataset, segmentation_mode, segmentation_var, max_locs_per_segment=None, + initial_sigma_nm=None, gt_drift=None, display=False, return_corrected_locs=False, + max_drift_nm=300, target_sigma_nm=1, boxcar_width=1, drift_max_bound_factor=2, + save_intermediate_results=False, + interpolation_method='cubic', mode="cuda", min_max_frames=None, + pair_indices_safety_check=False): + """ + Run COMET drift correction end-to-end. + + Pipeline: temporal segmentation -> KD-tree neighbor pairs -> cost optimization (L-BFGS-B) + -> optional temporal smoothing -> spline interpolation to per-frame drift -> (optional) subtract drift. + + Parameters + ---------- + dataset : ndarray of shape (N, 4) + Localization array with columns [x_nm, y_nm, z_nm, frame]. Units in nm; frame is int. + For 2D CSVs, insert a zero z column to get (N, 4). + segmentation_mode : {0, 1, 2} + Temporal segmentation mode: + 0 = number of windows (choose S directly), + 1 = localizations per window (accumulate frames until >= X locs), + 2 = fixed frame window size (default). + segmentation_var : int + Mode-dependent value (S, locs per window, or frames per window). + initial_sigma_nm : float, default=100 + Initial Gaussian length scale for the overlap kernel (coarse scale). + target_sigma_nm : float, default=1 + Target (final) Gaussian length scale for fine refinement. + max_drift_nm : float or None, default=None + Pair radius in nm used for neighbor search. If None, uses 3 * initial_sigma_nm. + drift_max_bound_factor : float, default=1.0 + Multiplicative factor for L-BFGS-B box bounds around +-max_drift. + boxcar_width : int, default=1 + Temporal smoothing width (segments) applied to the estimated drift between optimizer steps. + interpolation_method : {"cubic", "catmull-rom"}, default="cubic" + Spline used to convert per-segment drift to per-frame drift. + max_locs_per_segment : int or None, default=None + Optional downsampling cap per segment (to control memory/time). + mode : str, "cuda" + only Cuda in this version + return_corrected_locs : bool, default=False + If True, also return drift-corrected localizations. + + Returns + ------- + drift_interp_with_frames : ndarray of shape (F, 4) + Per-frame drift with columns [dx_nm, dy_nm, dz_nm, frame]. + corrected_locs : ndarray of shape (N, 4), optional + Only if return_corrected_locs=True. Columns are [x_nm, y_nm, z_nm, segment_id]. + """ + + loc_frames = dataset[:, -1] + if min_max_frames is None: + min_max_frames = (loc_frames.min(), loc_frames.max()) + + # Segment the dataset based on frame numbers into time windows + + result, sorted_dataset, idx_i, idx_j = segmentation_and_pair_indices_wrapper( + dataset, segmentation_var, segmentation_mode, max_drift_nm, max_locs_per_segment, + pair_indices_safety_check=pair_indices_safety_check) + + # Set default initial sigma if not provided + if initial_sigma_nm is None: + initial_sigma_nm = max_drift_nm // 3 + + # Run drift optimization + t0 = time.time() + drift_est = optimize_3d_chunked_better_moving_avg_kd( + result.n_segments, sorted_dataset, + idx_i, idx_j, + sigma_nm=initial_sigma_nm, + target_sigma_nm=target_sigma_nm, + drift_max_nm=max_drift_nm, + drift_max_bound_factor=drift_max_bound_factor, + display_steps=display, + boxcar_width=boxcar_width, + segmentation_result=result, + mode=mode + ) + elapsed = time.time() - t0 + + # Reshape and interpolate drift across all frames + drift_est = drift_est.reshape((result.n_segments, 3)) + vld_tp = np.where(~np.isnan(drift_est[:, 0])) + + frame_interp = np.arange(0, min_max_frames[1] + 1, dtype=int) + drift_interp = interpolate_drift(result.center_frames[vld_tp], drift_est[vld_tp], frame_interp, + method=interpolation_method) + drift_interp_with_frames = np.hstack((drift_interp, frame_interp[:, np.newaxis])) + + # Apply drift correction to original localizations + for i in range(3): + dataset[:, i] = dataset[:, i] - drift_interp[dataset[:, -1].astype(int), i] + + # Optionally show estimated drift curve + if display: + print(f"Drift estimation completed in {elapsed:.2f} seconds.") + plt.figure() + plt.plot(frame_interp, drift_interp) + plt.title("Estimated Drift") + plt.xlabel("Frames") + plt.ylabel("Drift (nm)") + plt.legend(['X', 'Y', 'Z']) + plt.show() + + + # Optional GT comparison plot + if display and gt_drift is not None: + fig, ax = plt.subplots(3, 1) + ax[0].plot(gt_drift[:, 3], gt_drift[:, 0], label='GT Drift X') + ax[1].plot(gt_drift[:, 3], gt_drift[:, 1], label='GT Drift Y') + ax[2].plot(gt_drift[:, 3], gt_drift[:, 2], label='GT Drift Z') + ax[0].plot(frame_interp, drift_interp[:, 0], label='Estimated Drift X', linestyle='--') + ax[1].plot(frame_interp, drift_interp[:, 1], label='Estimated Drift Y', linestyle='--') + ax[2].plot(frame_interp, drift_interp[:, 2], label='Estimated Drift Z', linestyle='--') + ax[1].set_title("Ground Truth vs Estimated Drift") + ax[1].set_xlabel("Frames") + ax[0].set_ylabel("Drift (nm)") + plt.legend() + plt.show() + + # Return corrected locs + drift + if return_corrected_locs: + return drift_interp_with_frames, dataset + else: + return drift_interp_with_frames + + +def optimize_3d_chunked_better_moving_avg_kd(n_segments, locs_nm, idx_i, idx_j, sigma_nm=30, drift_max_nm=300, + target_sigma_nm=30, display_steps=False, + boxcar_width=3, drift_max_bound_factor=2, + segmentation_result=None, + mode="cuda", return_calc_time=False): + """ + Estimate per-segment drift (mu) by minimizing the negative Gaussian-overlap cost + with an L-BFGS-B optimizer and a coarse-to-fine schedule on sigma. + + The routine operates on temporally segmented localizations and reuses a static + neighbor graph (pairs within drift_max_nm). Between optimizer steps, a moving + average (boxcar) can be applied to mu as temporal regularization. Sigma is + reduced iteratively from `sigma_nm` toward `target_sigma_nm` for robust + convergence. + + Parameters + ---------- + n_segments : int + Number of temporal segments S (0..S-1). One 3D drift vector is estimated per segment. + locs_nm : ndarray of shape (M, 3) + localizations in nanometers, columns [x, y, z]; + sigma_nm : float, default=30 + Initial Gaussian width (coarse scale) for the overlap kernel. + drift_max_nm : float, default=300 + Maximum expected drift (nm). Also used as the radius for neighbor pairs and + as the L-BFGS-B bound scale (see `drift_max_bound_factor`). + target_sigma_nm : float, default=30 + Target / final Gaussian width (fine scale). The optimizer reduces sigma toward + this value over iterations. + display_steps : bool, default=False + If True, print or log intermediate progress per iteration/scale. + + boxcar_width : int, default=3 + Temporal smoothing width (in segments) for a moving average applied to mu between steps. + Use 0 or 1 to disable smoothing. + drift_max_bound_factor : float, default=2 + Multiplier for L-BFGS-B bounds around +- drift_max_nm to keep updates physically reasonable. + segmentation_result : object or None, default=None + Segmentation info and metadata. Expected to provide, at minimum: + - segment IDs per localization + - center frames per segment + - any additional structures required by backend (e.g., pair indices) + If None, pairs/ids may be built internally depending on implementation. + mode : str, default=cuda + If True, use the CPU backend; otherwise try GPU (CUDA) and fall back to CPU if unavailable. + return_calc_time : bool, default=False + If True, also return the total computation time in seconds. + + Returns + ------- + mu : ndarray of shape (S, 3) + Estimated per-segment drift (dx, dy, dz) in nanometers. + calc_time_s : float, optional + Only when `return_calc_time=True`. Wall-clock time for the optimization. + """ + + if segmentation_result is None: + segmentation_result = {} + intermediate_results_filehandle = None + sigma_factor = 1.0 + + # Extract coordinate + time arrays, convert to device if CUDA + coords = locs_nm[:, :3].astype(np.float32).copy() + times = locs_nm[:, 3].astype(np.int32).copy() + + chunk_size = int(1E8) # 1E7 + + quality_control = mode == "torch_qc" or mode == "cuda_qc" + + d_coords = cuda.to_device(coords) + d_times = cuda.to_device(times) + if len(idx_i) * 4 > 2e9: + # Use mapped memory if index arrays are large + print("Large index arrays — using mapped memory.") + d_idx_i = cuda.mapped_array_like(idx_i.astype(np.int32), wc=True) + d_idx_j = cuda.mapped_array_like(idx_j.astype(np.int32), wc=True) + d_idx_i[:] = idx_i + d_idx_j[:] = idx_j + else: + d_idx_i = cuda.to_device(idx_i.astype(np.int32)) + d_idx_j = cuda.to_device(idx_j.astype(np.int32)) + # Preallocate device arrays + d_sigma = np.float64(sigma_nm) + d_val = cuda.to_device(np.zeros(chunk_size)) + d_deri = cuda.to_device(np.zeros((n_segments, 3), dtype=np.float64)) + wrapper = cuda_wrapper_chunked + + # Initial drift estimate + bounds + drift_est = np.zeros(n_segments * 3) + bounds = [(-drift_max_nm * drift_max_bound_factor, drift_max_nm * drift_max_bound_factor)] * (3 * n_segments) + + drift_est_gradient = np.inf + fails = 0 + done = False + itr_counter = 0 + start_time = time.time() + + while not done: + # Apply boxcar smoothing to current estimate + tmp = drift_est.reshape((-1, 3)) + for i in range(3): + tmp[:, i] = convolve(tmp[:, i], np.ones(boxcar_width) / boxcar_width) + drift_est = tmp.flatten() + + # Run L-BFGS-B optimization step + result = minimize(wrapper, drift_est, method='L-BFGS-B', + args=( + d_coords, d_times, d_idx_i, d_idx_j, d_sigma, sigma_factor, d_val, d_deri, chunk_size), + jac=True, bounds=bounds, + options={'disp': display_steps, 'gtol': 1E-5, 'ftol': 1E3 * np.finfo(float).eps, + 'maxls': 40}) + itr_counter += 1 + print(f"Iteration {itr_counter}: status = {result.status}, success = {result.success}") + print(f" current sigma: {np.round(sigma_nm * sigma_factor, 2)} nm") + + # Update if successful + if result.success: + delta = np.median((result.x - drift_est) ** 2) + print(f" drift estimate gradient: {delta}") + print(f" previous gradient: {drift_est_gradient}") + # Check convergence + if (delta > drift_est_gradient or sigma_nm * sigma_factor <= 1.0) and sigma_nm * sigma_factor <= target_sigma_nm: + done = True + calc_time = time.time() - start_time + print(f"Optimization completed in {calc_time:.2f} s") + else: + sigma_factor /= 1.5 + drift_est_gradient = delta + drift_est = result.x + else: + fails += 1 + if fails > 2: + sigma_factor *= 2 + print("Restarting with larger sigma_factor") + if fails > 5: + raise RuntimeError("L-BFGS-B Optimization failed after multiple retries") + + if return_calc_time: + return drift_est, time.time() - start_time, itr_counter + else: + return drift_est + + +def segmentation_and_pair_indices_wrapper(dataset, segmentation_var, segmentation_mode, max_drift_nm, + max_locs_per_segment, pair_indices_safety_check=False, hard_limit_pairs=None): + if not segmentation_mode == -1: # -1 is for pre-segmented data + result = segmentation_wrapper(dataset[:, -1], segmentation_var, segmentation_mode, + max_locs_per_segment, return_param_dict=True) + else: + # pre segmented data, anyway we set these values in case auto downsampling is needed + segmentation_mode = 2 # dummy --> segment per frame ... + segmentation_var = 1 # dummy --> ... using 1 frame per segment + if pair_indices_safety_check: + n_pairs_est = estimate_pairs(dataset[result.loc_valid, :3], max_drift_nm) + print(f"Estimated number of pairs within {max_drift_nm} nm: {n_pairs_est:,}") + if hard_limit_pairs is not None and n_pairs_est > hard_limit_pairs: + raise RuntimeError(f"Estimated number of pairs {n_pairs_est} exceeds hard limit of {hard_limit_pairs}. " + f"Aborting to avoid crash.") + if n_pairs_est > 5e8: + print(f"Estimated number of pairs is very large {n_pairs_est}. " + f"Automatic down-sampling is usually required above 500 mil." + f"Billions of pairs can lead to crash.") + ans = input("Continue anyway? (y/n): ") + if ans.lower() != 'y': + raise RuntimeError("Aborted by user due to large estimated number of pairs.") + idx_i, idx_j, successful = pair_indices_kdtree(dataset[result.loc_valid, :3], max_drift_nm) + if not successful: + if max_locs_per_segment is None: + max_locs_per_segment = int(result.out_dict['locs_per_segment'].max()) + while not successful: + max_locs_per_segment = int(max_locs_per_segment * 0.9) + print(f"Segmentation and Pairing attempt failed, automatic down-sampling active...") + print(f"Retrying segmentation with max_locs_per_segment={max_locs_per_segment}...") + result = segmentation_wrapper(dataset[:, -1], segmentation_var, segmentation_mode, + max_locs_per_segment, return_param_dict=True) + sorted_dataset = dataset.copy() + sorted_dataset[:, -1] = result.loc_segments + sorted_dataset = sorted_dataset[result.loc_valid] + idx_i, idx_j, successful = pair_indices_kdtree(sorted_dataset[:, :3], max_drift_nm) + print(f"Segmentation and Pairing successful resulting in {result.n_segments:,} time windows with on average " + f"{int(np.median(result.out_dict['locs_per_segment']))} locs per time window. " + f"{len(idx_i):,} Pairs where found.") + sorted_dataset = dataset.copy() + sorted_dataset[:, -1] = result.loc_segments + sorted_dataset = sorted_dataset[result.loc_valid] + return result, sorted_dataset, idx_i, idx_j + + + +def estimate_pairs(coordinates, distance): + for i in range(len(coordinates[0])): + coordinates[:, i] -= np.min(coordinates[:, i]) + coordinates = np.array(np.floor(coordinates / distance), dtype=int) + coordinates = np.array(list(map(tuple, coordinates))) + sort_indices = np.lexsort(coordinates.T)# get the unique tuples and their counts + unique_tuples, counts = np.unique(coordinates[sort_indices], axis=0, return_counts=True) + # get the indices of the similar tuples + similar_indices = np.split(sort_indices, np.cumsum(counts[:-1])) + idx_i = [] + idx_j = [] + pair_idc_estimate = 0 + for i in range(len(similar_indices)): + n_elements = len(similar_indices[i]) + pair_idc_estimate += n_elements * (n_elements - 1) + rounded = round(pair_idc_estimate,-4) + return rounded + + +def interpolate_drift(center_frames, drift_est, frame_range, method='cubic'): + """ + Interpolates drift estimates to all frames using specified method. + Parameters: + - center_frames: np.ndarray of shape (M,), frames corresponding to drift estimates. + - drift_est: np.ndarray of shape (M, 3), drift estimates at center frames. + - frame_range: array-like, frames to interpolate drift estimates to. + - method: str, interpolation method ('cubic' or 'catmull-rom'). + Returns: + - drift_interp: np.ndarray of shape (len(frame_range), 3), interpolated drift estimates. + """ + if method == 'cubic': + return _interpolate_cubic(center_frames, drift_est, frame_range) + elif method == 'catmull-rom': + return _interpolate_catmull_rom(center_frames, drift_est, frame_range) + elif method == 'linear': + drift_x = np.interp(frame_range, center_frames, drift_est[:, 0]) + drift_y = np.interp(frame_range, center_frames, drift_est[:, 1]) + drift_z = np.interp(frame_range, center_frames, drift_est[:, 2]) + return np.vstack([drift_x, drift_y, drift_z]).T + else: + raise ValueError(f"Unknown interpolation method: {method}") + + +def _interpolate_cubic(center_frames, drift_est, frame_range): + from scipy.interpolate import CubicSpline + drift_x = CubicSpline(center_frames, drift_est[:, 0])(frame_range) + drift_y = CubicSpline(center_frames, drift_est[:, 1])(frame_range) + drift_z = CubicSpline(center_frames, drift_est[:, 2])(frame_range) + return np.vstack([drift_x, drift_y, drift_z]).T + + +def _interpolate_catmull_rom(center_frames, drift_est, frame_range): + def catmull_rom_1d(x, y, x_interp): + result = np.zeros_like(x_interp) + for i in range(1, len(x) - 2): + x0, x1, x2, x3 = x[i-1], x[i], x[i+1], x[i+2] + y0, y1, y2, y3 = y[i-1], y[i], y[i+1], y[i+2] + + mask = (x_interp >= x1) & (x_interp <= x2) + t = (x_interp[mask] - x1) / (x2 - x1) + + result[mask] = ( + 0.5 * ( + (2 * y1) + + (-y0 + y2) * t + + (2*y0 - 5*y1 + 4*y2 - y3) * t**2 + + (-y0 + 3*y1 - 3*y2 + y3) * t**3 + ) + ) + return result + + x_interp = np.asarray(frame_range) + x = np.asarray(center_frames) + + if len(x) < 4: + raise ValueError("Catmull-Rom interpolation requires at least 4 points.") + + drift_x = catmull_rom_1d(x, drift_est[:, 0], x_interp) + drift_y = catmull_rom_1d(x, drift_est[:, 1], x_interp) + drift_z = catmull_rom_1d(x, drift_est[:, 2], x_interp) + + return np.vstack([drift_x, drift_y, drift_z]).T + + +# Numba cuda code of cost-function --> quantized overlap of pairs of localizations, gaussian mixture model +@cuda.jit +def cost_function_full_3d_chunked(d_locs_time, start_idx, chunk_size, d_idx_i, d_idx_j, d_sigma, d_sigma_factor, d_val, + d_val_sum, d_deri, d_locs_coords, mu): + """Compute negative-overlap cost and gradient for 3D localizations (GPU/CPU backend; internal use).""" + tx = cuda.threadIdx.x + ty = cuda.blockIdx.x + bw = cuda.blockDim.x + pos = tx + ty * bw + + if pos < chunk_size: + i = d_idx_i[pos + start_idx] + j = d_idx_j[pos + start_idx] + + ti = d_locs_time[i] + tj = d_locs_time[j] + + dx = (d_locs_coords[i, 0] - mu[ti, 0]) - (d_locs_coords[j, 0] - mu[tj, 0]) + dy = (d_locs_coords[i, 1] - mu[ti, 1]) - (d_locs_coords[j, 1] - mu[tj, 1]) + dz = (d_locs_coords[i, 2] - mu[ti, 2]) - (d_locs_coords[j, 2] - mu[tj, 2]) + sigma_sq = (2 * d_sigma * d_sigma_factor) ** 2 + + diff_sq = dx * dx + dy * dy + dz * dz + val = 1 / (d_sigma * d_sigma_factor) * math.exp(-diff_sq / sigma_sq) + d_val[pos] = val + + # Update derivatives + cuda.atomic.add(d_deri, (tj, 0), 2 * val * (d_locs_coords[j, 0] - d_locs_coords[i, 0] + mu[ti, 0] - mu[tj, 0]) / sigma_sq) + cuda.atomic.add(d_deri, (tj, 1), 2 * val * (d_locs_coords[j, 1] - d_locs_coords[i, 1] + mu[ti, 1] - mu[tj, 1]) / sigma_sq) + cuda.atomic.add(d_deri, (tj, 2), 2 * val * (d_locs_coords[j, 2] - d_locs_coords[i, 2] + mu[ti, 2] - mu[tj, 2]) / sigma_sq) + + cuda.atomic.add(d_deri, (ti, 0), 2 * val * (d_locs_coords[i, 0] - d_locs_coords[j, 0] + mu[tj, 0] - mu[ti, 0]) / sigma_sq) + cuda.atomic.add(d_deri, (ti, 1), 2 * val * (d_locs_coords[i, 1] - d_locs_coords[j, 1] + mu[tj, 1] - mu[ti, 1]) / sigma_sq) + cuda.atomic.add(d_deri, (ti, 2), 2 * val * (d_locs_coords[i, 2] - d_locs_coords[j, 2] + mu[tj, 2] - mu[ti, 2]) / sigma_sq) + + cuda.atomic.add(d_val_sum, 0, val) + d_val[pos] = 0 + + +# Interface between the Python code and the CUDA kernel, mainly for chunking the data to avoid memory issues +def cuda_wrapper_chunked(mu, d_locs_coords, d_locs_time, d_idx_i, d_idx_j, d_sigma, d_sigma_factor, d_val, d_deri, + chunk_size): + val_total = 0 + d_val_sum = cuda.to_device(np.zeros(1, dtype=np.float64)) + mu_dev = cuda.to_device(np.asarray(mu.reshape(int(mu.size / 3), 3), dtype=np.float64)) + + n_chunks = int(np.ceil(d_idx_i.size / chunk_size)) + threadsperblock = 128 + + for i in range(n_chunks - 1): + idc_start = i*chunk_size + blockspergrid = (chunk_size + (threadsperblock - 1)) // threadsperblock + cost_function_full_3d_chunked[blockspergrid, threadsperblock]( + d_locs_time, idc_start, chunk_size, d_idx_i, d_idx_j, + d_sigma, d_sigma_factor, d_val, d_val_sum, d_deri, d_locs_coords, mu_dev + ) + val_total += d_val_sum.copy_to_host() + + # Final chunk + n_remaining = d_idx_i.size - (n_chunks - 1) * chunk_size + idc_start = (n_chunks - 1) * chunk_size + blockspergrid = (n_remaining + (threadsperblock - 1)) // threadsperblock + cost_function_full_3d_chunked[blockspergrid, threadsperblock]( + d_locs_time, idc_start, n_remaining, d_idx_i, d_idx_j, + d_sigma, d_sigma_factor, d_val, d_val_sum, d_deri, d_locs_coords, mu_dev + ) + val_total += d_val_sum.copy_to_host() + deri = d_deri.copy_to_host() + d_deri[:] = 0 + + return -np.nansum(val_total), -deri.flatten() + +from dataclasses import dataclass +from typing import Optional, Dict + + +@dataclass +class SegmentationResult: + """Container for segmentation results.""" + loc_segments: np.ndarray + loc_valid: np.ndarray + center_frames: np.ndarray + n_segments: int + out_dict: Optional[Dict] = None + + +def _group_by_frame(loc_frames: np.ndarray): + """Returns a dict {frame_number: indices_in_loc_frames} efficiently.""" + sort_idx = np.argsort(loc_frames) + sorted_frames = loc_frames[sort_idx] + unique_frames, start_idx, counts = np.unique(sorted_frames, return_index=True, return_counts=True) + frame_to_indices = { + frame: sort_idx[start:start + count] + for frame, start, count in zip(unique_frames, start_idx, counts) + } + return unique_frames, frame_to_indices + + +def segment_by_num_locs_per_window(loc_frames: np.ndarray, min_n_locs_per_window: int, + max_locs_per_segment = None, + return_param_dict: bool = False) -> SegmentationResult: + """ + Segments by collecting a minimum number of localizations per window. + Once the threshold is met and enough locs remain, a new segment is created. + This method ensures that each segment has at least `min_n_locs_per_window` localizations, + while also trying to avoid creating segments that are too small at the end of the dataset. + If `max_locs_per_segment` is set, a random subset of that size is chosen from each segment. + This method is particularly useful for datasets with varying localization densities over time. + Parameters: + loc_frames (np.ndarray): Array of frame numbers for each localization. + min_n_locs_per_window (int): Minimum number of localizations per segment. + max_locs_per_segment (Optional[int]): Maximum number of localizations per segment. If None, all locs are used. + return_param_dict (bool): Whether to return a dictionary of segmentation parameters. + Returns: + SegmentationResult: A dataclass containing segmentation results and parameters. + """ + loc_frames = np.asarray(loc_frames, dtype=int) + n_locs = len(loc_frames) + + if max_locs_per_segment is not None and max_locs_per_segment < 1: # downsampling in percentage + max_locs_per_segment = int(min_n_locs_per_window * max_locs_per_segment) + + unique_frames, frame_to_indices = _group_by_frame(loc_frames) + loc_segments = np.full(n_locs, -1, dtype=int) # Default to -1 for safety + segment_counter = 0 + n_locs_in_current_segment = 0 + current_segment_indices = [] + start_frames, end_frames, locs_per_segment = [], [], [] + + for i, frame in enumerate(unique_frames): + indices = frame_to_indices[frame] + n_locs_this_frame = len(indices) + remaining_locs = n_locs - (len(current_segment_indices) + n_locs_this_frame + np.sum(locs_per_segment)) + + # Add frame to current segment if: + # - It fills the current segment to threshold + # - AND there are enough locs left for another segment (or it's the last frame) + if (n_locs_in_current_segment + n_locs_this_frame >= min_n_locs_per_window) and \ + (remaining_locs >= min_n_locs_per_window or i == len(unique_frames) - 1): + current_segment_indices.extend(indices) + n_locs_in_current_segment += n_locs_this_frame + + loc_segments[current_segment_indices] = segment_counter + start_frames.append(loc_frames[current_segment_indices[0]]) + end_frames.append(loc_frames[current_segment_indices[-1]]) + locs_per_segment.append(len(current_segment_indices)) + + segment_counter += 1 + current_segment_indices = [] + n_locs_in_current_segment = 0 + else: + # Defer frame to current segment + current_segment_indices.extend(indices) + n_locs_in_current_segment += n_locs_this_frame + + n_segments = segment_counter + center_frames = np.zeros(n_segments) + loc_valid = np.zeros(n_locs, dtype=bool) + + for i in range(n_segments): + segment_indices = np.where(loc_segments == i)[0] + if max_locs_per_segment and len(segment_indices) > max_locs_per_segment: + selected = np.random.choice(segment_indices, max_locs_per_segment, replace=False) + else: + selected = segment_indices + loc_valid[selected] = True + locs_per_segment[i] = len(selected) + center_frames[i] = np.mean(loc_frames[selected]) + + out_dict = None + if return_param_dict: + n_locs_valid = loc_valid.sum() + out_dict = { + "n_segments": n_segments, + "min_n_locs_per_window": min_n_locs_per_window, + "frames_per_window": -1, + "start_frames": np.array(start_frames), + "end_frames": np.array(end_frames), + "locs_per_segment": np.array(locs_per_segment), + "n_locs": n_locs, + "n_locs_valid": n_locs_valid, + "n_locs_invalid": n_locs - n_locs_valid, + "center_frames": center_frames + } + + return SegmentationResult(loc_segments, loc_valid, center_frames, n_segments, out_dict) + + +def segment_by_frame_windows(loc_frames: np.ndarray, n_frames_per_window: int, + max_locs_per_segment = None, + return_param_dict: bool = False) -> SegmentationResult: + """ + Splits localization data into fixed-size windows of N frames. + All localizations in those frames are grouped into one segment. + """ + loc_frames = np.asarray(loc_frames, dtype=int) + frames, frame_to_indices = _group_by_frame(loc_frames) + n_locs = len(loc_frames) + n_segments = int(np.ceil(len(frames) / n_frames_per_window)) + + if max_locs_per_segment is not None and max_locs_per_segment < 1: # downsampling in percentage + max_locs_per_segment = n_locs / n_segments * max_locs_per_segment + + loc_segments = np.zeros(n_locs, dtype=int) + center_frames = np.zeros(n_segments) + loc_valid = np.ones(n_locs, dtype=bool) + start_frames, end_frames, locs_per_segment = [], [], [] + + for i in range(n_segments): + frame_window = frames[i * n_frames_per_window:(i + 1) * n_frames_per_window] + indices = np.concatenate([frame_to_indices[frame] for frame in frame_window if frame in frame_to_indices]) + if len(indices) == 0: + continue + loc_segments[indices] = i + start_frames.append(frame_window[0]) + end_frames.append(frame_window[-1]) + center_frames[i] = np.mean(loc_frames[indices]) + locs_per_segment.append(len(indices)) + if max_locs_per_segment and len(indices) > max_locs_per_segment: + mask = np.ones(len(indices), dtype=bool) + mask[np.random.choice(len(indices), len(indices) - max_locs_per_segment, replace=False)] = False + loc_valid[indices[~mask]] = False + + out_dict = None + if return_param_dict: + n_locs_valid = loc_valid.sum() + out_dict = { + "n_segments": n_segments, + "min_n_locs_per_window": -1, + "frames_per_window": n_frames_per_window, + "start_frames": np.array(start_frames), + "end_frames": np.array(end_frames), + "locs_per_segment": np.array(locs_per_segment), + "n_locs": n_locs, + "n_locs_valid": n_locs_valid, + "n_locs_invalid": n_locs - n_locs_valid, + "center_frames": center_frames + } + + return SegmentationResult(loc_segments, loc_valid, center_frames, n_segments, out_dict) + + +def segment_by_num_windows(loc_frames: np.ndarray, n_windows: int, max_locs_per_segment = None, + return_param_dict: bool = False) -> SegmentationResult: + """ + Converts number of windows into an equivalent minimum locs per window, + then calls `segment_by_num_locs_per_window`. + """ + n_locs = len(loc_frames) + n_locs_per_window = int(np.ceil(n_locs / n_windows)) + if max_locs_per_segment is not None and max_locs_per_segment < 1: # downsampling in percentage + max_locs_per_segment = int(n_locs_per_window * max_locs_per_segment) + return segment_by_num_locs_per_window(loc_frames, n_locs_per_window, max_locs_per_segment, return_param_dict) + + +def segmentation_wrapper(loc_frames: np.ndarray, segmentation_var: int, segmentation_mode: int = 2, + max_locs_per_segment = None, + return_param_dict: bool = False) -> SegmentationResult: + """ + Dispatch function that selects segmentation method: + 0 → fixed number of windows + 1 → fixed number of localizations per window + 2 → fixed number of frames per window (default) + """ + if segmentation_mode == 0: + return segment_by_num_windows(loc_frames, segmentation_var, max_locs_per_segment, return_param_dict) + elif segmentation_mode == 1: + return segment_by_num_locs_per_window(loc_frames, segmentation_var, max_locs_per_segment, return_param_dict) + else: + return segment_by_frame_windows(loc_frames, segmentation_var, max_locs_per_segment, return_param_dict) + + +def pair_indices_kdtree(coordinates, distance): + """ + Find all pairs of points within a certain distance using a KD-tree. + Parameters: + - coordinates: np.ndarray of shape (N, D) where N is the number of points and D is the dimensionality. + - distance: float, the maximum distance to consider points as a pair. + Returns: + - idx1: np.ndarray of shape (M,), indices of the first point in each pair. + - idx2: np.ndarray of shape (M,), indices of the second point in each pair. + """ + tree = cKDTree(coordinates) + try: + pairs = tree.query_pairs(r=distance, output_type='ndarray') + except MemoryError: + print("[pair_indices_kdtree] MemoryError encountered") + return [], [], False + return np.ascontiguousarray(pairs[:, 0], dtype=np.int32), np.ascontiguousarray(pairs[:, 1], dtype=np.int32), True \ No newline at end of file diff --git a/picasso/gui/render.py b/picasso/gui/render.py index 21e0b4b9..4952a21b 100644 --- a/picasso/gui/render.py +++ b/picasso/gui/render.py @@ -38,9 +38,9 @@ from sklearn.metrics.pairwise import euclidean_distances from sklearn.cluster import KMeans from PyQt5 import QtCore, QtGui, QtWidgets - from .. import ( aim, + comet, clusterer, g5m, imageprocess, @@ -1685,6 +1685,87 @@ def getParams( clustered_locs, ) +class COMETDialog(QtWidgets.QDialog): + """Dialog to choose parameters for COMET undrifting. + + Attributes + ---------- + locs_per_segment : QSpinBox + Target number of localizations per temporal segment. + max_drift_nm : QDoubleSpinBox + Maximum expected drift in nm. + max_locs_per_segment : QSpinBox + Optional cap for downsampling localizations per segment. + Use -1 to disable downsampling. + """ + + def __init__(self, window: QtWidgets.QMainWindow) -> None: + super().__init__(window) + self.window = window + self.setWindowTitle("COMET undrifting") + + vbox = QtWidgets.QVBoxLayout(self) + grid = QtWidgets.QGridLayout() + + locs_label = QtWidgets.QLabel("Localizations per segment:") + locs_label.setToolTip( + "Target number of localizations used to form each temporal segment." + ) + grid.addWidget(locs_label, 0, 0) + self.locs_per_segment = QtWidgets.QSpinBox() + self.locs_per_segment.setRange(1, int(1e6)) + self.locs_per_segment.setValue(500) + grid.addWidget(self.locs_per_segment, 0, 1) + + max_drift_label = QtWidgets.QLabel("Maximum drift (nm):") + max_drift_label.setToolTip( + "Maximum expected drift over the dataset. " + "Used for pairing and optimization bounds." + ) + grid.addWidget(max_drift_label, 1, 0) + self.max_drift_nm = QtWidgets.QDoubleSpinBox() + self.max_drift_nm.setRange(0.1, 1e6) + self.max_drift_nm.setValue(60.0) + self.max_drift_nm.setDecimals(1) + self.max_drift_nm.setSingleStep(1.0) + grid.addWidget(self.max_drift_nm, 1, 1) + + downsample_label = QtWidgets.QLabel("Max. localizations per segment:") + downsample_label.setToolTip( + "Optional downsampling cap per segment. " + "Set to -1 to keep all localizations." + ) + grid.addWidget(downsample_label, 2, 0) + self.max_locs_per_segment = QtWidgets.QSpinBox() + self.max_locs_per_segment.setRange(-1, int(1e6)) + self.max_locs_per_segment.setValue(-1) + grid.addWidget(self.max_locs_per_segment, 2, 1) + + vbox.addLayout(grid) + + self.buttons = QtWidgets.QDialogButtonBox( + QtWidgets.QDialogButtonBox.Ok | QtWidgets.QDialogButtonBox.Cancel, + QtCore.Qt.Horizontal, + self, + ) + vbox.addWidget(self.buttons) + self.buttons.accepted.connect(self.accept) + self.buttons.rejected.connect(self.reject) + + @staticmethod + def getParams( + parent: QtWidgets.QMainWindow | None = None, + ) -> tuple[dict, bool]: + """Create the dialog and return the requested COMET parameters.""" + dialog = COMETDialog(parent) + result = dialog.exec_() + params = { + "locs_per_segment": dialog.locs_per_segment.value(), + "max_drift_nm": dialog.max_drift_nm.value(), + "max_locs_per_segment": dialog.max_locs_per_segment.value(), + } + return params, result == QtWidgets.QDialog.Accepted + class AIMDialog(QtWidgets.QDialog): """Choose parameters for AIM undrifting. @@ -10592,6 +10673,33 @@ def show_drift(self) -> None: self.plot_window.show() + def undrift_comet(self) -> None: + """Undrift with COMET. + + See https://comet.smlm.tools + """ + channel = self.get_channel("Undrift by COMET") + if channel is not None: + locs = self.all_locs[channel] + info = self.infos[channel] + + params, ok = COMETDialog.getParams(self.window) + if ok: + locs, new_info, drift = comet.comet( + locs, + info, + **params, + ) + + locs = lib.ensure_sanity(locs, info) + self.all_locs[channel] = locs + self.locs[channel] = copy.copy(locs) + self.infos[channel] = new_info + self.index_blocks[channel] = None + self.add_drift(channel, drift) + self.update_scene() + self.show_drift() + def undrift_aim(self) -> None: """Undrift with Adaptive Intersection Maximization (AIM). @@ -11746,8 +11854,11 @@ def initUI(self, plugins_loaded: bool) -> None: # menu bar - Postprocess postprocess_menu = self.menu_bar.addMenu("Postprocess") + undrift_comet_action = postprocess_menu.addAction("Undrift by COMET") + undrift_comet_action.setShortcut("Ctrl+U") + undrift_comet_action.triggered.connect(self.view.undrift_comet) + undrift_aim_action = postprocess_menu.addAction("Undrift by AIM") - undrift_aim_action.setShortcut("Ctrl+U") undrift_aim_action.triggered.connect(self.view.undrift_aim) undrift_from_picked_action = postprocess_menu.addAction( "Undrift from picked" From 23c20e68576ad077a828a1a6738d357c6f1d8d5e Mon Sep 17 00:00:00 2001 From: Lenny Reinkensmeier Date: Thu, 28 May 2026 17:17:43 +0200 Subject: [PATCH 02/19] Handle COMET availability in drift correction and update menu shortcuts. Added proper testing for COMET. --- picasso/comet.py | 852 ------------------------------------------ picasso/gui/render.py | 28 +- tests/test_comet.py | 322 ++++++++++++++++ 3 files changed, 342 insertions(+), 860 deletions(-) delete mode 100644 picasso/comet.py create mode 100644 tests/test_comet.py diff --git a/picasso/comet.py b/picasso/comet.py deleted file mode 100644 index 11321927..00000000 --- a/picasso/comet.py +++ /dev/null @@ -1,852 +0,0 @@ -import numpy as np -import pandas as pd -import matplotlib.pyplot as plt -from numba import cuda -from scipy.ndimage import convolve -from scipy.optimize import minimize -import time -import math -from scipy.spatial import cKDTree - -from . import lib, __version__ - - - - -def comet( - locs: pd.DataFrame, - info: list[dict], - locs_per_segment: int, - max_drift_nm: float, - max_locs_per_segment: int | None = None, - *, - initial_sigma_nm: float | None = None, - target_sigma_nm: float = 10.0, - boxcar_width: int = 3, - drift_max_bound_factor: float = 2.0, - interpolation_method: str = "cubic", - mode: str = "cuda", - display: bool = False, - progress=None, # currently unused -) -> tuple[pd.DataFrame, list[dict], pd.DataFrame]: - """Apply COMET undrifting to the localizations. - - Parameters - ---------- - locs : pd.DataFrame - Localization table. Must contain columns 'x', 'y', 'frame'. - Optional: 'z'. - x/y are expected in camera pixels, z in the native z-unit used by Picasso. - info : list of dict - Localization metadata. - locs_per_segment : int - Target number of localizations per segment for COMET segmentation. - max_drift_nm : float - Maximum expected drift in nm. - max_locs_per_segment : int or None, optional - Optional cap for localizations per segment. If negative, treated as None. - initial_sigma_nm : float or None, optional - Initial Gaussian sigma for COMET. If None, defaults to max_drift_nm / 3. - target_sigma_nm : float, optional - Final Gaussian sigma for refinement. - boxcar_width : int, optional - Smoothing width over segments. - drift_max_bound_factor : float, optional - Bound factor for optimizer. - interpolation_method : str, optional - Drift interpolation method. - mode : str, optional - Backend mode, only "cuda". - display : bool, optional - Whether to show diagnostic plots. - progress : optional - Placeholder for API compatibility with AIM. - - Returns - ------- - locs : pd.DataFrame - Undrift-corrected localizations. - new_info : list[dict] - Updated metadata. - drift : pd.DataFrame - Per-frame drift table with columns x, y, and optionally z. - Drift is returned in the same units as locs: - x/y in pixels, z in the original z-unit. - """ - locs = locs.copy() - - if max_locs_per_segment is not None and max_locs_per_segment < 0: - max_locs_per_segment = None - - pixelsize_nm = lib.get_from_metadata(info, "Pixelsize", raise_error=True) - - required_cols = {"x", "y", "frame"} - missing = required_cols - set(locs.columns) - if missing: - raise KeyError(f"Missing required columns for COMET: {sorted(missing)}") - - has_z = "z" in locs.columns - - # Normalize frames so they start at 0 for internal indexing. - # This is important because COMET later indexes drift by frame number. - frame = locs["frame"].to_numpy(dtype=np.int64) - frame0 = frame - frame.min() - - # Build COMET input array in nm - dataset = np.zeros((len(locs), 4), dtype=np.float64) - dataset[:, 0] = locs["x"].to_numpy(dtype=np.float64) * pixelsize_nm - dataset[:, 1] = locs["y"].to_numpy(dtype=np.float64) * pixelsize_nm - dataset[:, 2] = ( - locs["z"].to_numpy(dtype=np.float64) if has_z else 0.0 - ) - dataset[:, 3] = frame0 - - if initial_sigma_nm is None: - initial_sigma_nm = max_drift_nm / 3 - - drift_nm_with_frame = comet_run_kd( - dataset=dataset, - segmentation_mode=1, # fixed locs per segment - segmentation_var=locs_per_segment, - max_locs_per_segment=max_locs_per_segment, - initial_sigma_nm=initial_sigma_nm, - gt_drift=None, - display=display, - return_corrected_locs=False, - max_drift_nm=max_drift_nm, - target_sigma_nm=target_sigma_nm, - boxcar_width=boxcar_width, - drift_max_bound_factor=drift_max_bound_factor, - interpolation_method=interpolation_method, - mode=mode, - min_max_frames=(int(frame0.min()), int(frame0.max())), - ) - - # drift_nm_with_frame columns: dx_nm, dy_nm, dz_nm, frame0 - drift_nm = drift_nm_with_frame[:, :3] - - # Apply drift back to dataframe in original units - frame_idx = frame0.astype(np.int64) - locs["x"] = locs["x"].to_numpy(dtype=np.float64) - drift_nm[frame_idx, 0] / pixelsize_nm - locs["y"] = locs["y"].to_numpy(dtype=np.float64) - drift_nm[frame_idx, 1] / pixelsize_nm - if has_z: - locs["z"] = locs["z"].to_numpy(dtype=np.float64) - drift_nm[frame_idx, 2] - - # Build drift dataframe in Picasso-style units - drift_dict = { - "x": (drift_nm[:, 0] / pixelsize_nm).astype("float32"), - "y": (drift_nm[:, 1] / pixelsize_nm).astype("float32"), - } - if has_z: - drift_dict["z"] = drift_nm[:, 2].astype("float32") - - drift = pd.DataFrame(drift_dict) - - new_info_entry = { - "Generated by": f"Picasso v{__version__} COMET", - "Segmentation mode": "localizations per segment", - "Localizations per segment": locs_per_segment, - "Maximum drift (nm)": max_drift_nm, - "Initial sigma (nm)": float(initial_sigma_nm), - "Target sigma (nm)": float(target_sigma_nm), - "Boxcar width": int(boxcar_width), - "Max localizations per segment": ( - None if max_locs_per_segment is None else int(max_locs_per_segment) - ), - "Interpolation method": interpolation_method, - "Backend": mode, - } - new_info = info + [new_info_entry] - - return locs, new_info, drift - -def comet_run_kd(dataset, segmentation_mode, segmentation_var, max_locs_per_segment=None, - initial_sigma_nm=None, gt_drift=None, display=False, return_corrected_locs=False, - max_drift_nm=300, target_sigma_nm=1, boxcar_width=1, drift_max_bound_factor=2, - save_intermediate_results=False, - interpolation_method='cubic', mode="cuda", min_max_frames=None, - pair_indices_safety_check=False): - """ - Run COMET drift correction end-to-end. - - Pipeline: temporal segmentation -> KD-tree neighbor pairs -> cost optimization (L-BFGS-B) - -> optional temporal smoothing -> spline interpolation to per-frame drift -> (optional) subtract drift. - - Parameters - ---------- - dataset : ndarray of shape (N, 4) - Localization array with columns [x_nm, y_nm, z_nm, frame]. Units in nm; frame is int. - For 2D CSVs, insert a zero z column to get (N, 4). - segmentation_mode : {0, 1, 2} - Temporal segmentation mode: - 0 = number of windows (choose S directly), - 1 = localizations per window (accumulate frames until >= X locs), - 2 = fixed frame window size (default). - segmentation_var : int - Mode-dependent value (S, locs per window, or frames per window). - initial_sigma_nm : float, default=100 - Initial Gaussian length scale for the overlap kernel (coarse scale). - target_sigma_nm : float, default=1 - Target (final) Gaussian length scale for fine refinement. - max_drift_nm : float or None, default=None - Pair radius in nm used for neighbor search. If None, uses 3 * initial_sigma_nm. - drift_max_bound_factor : float, default=1.0 - Multiplicative factor for L-BFGS-B box bounds around +-max_drift. - boxcar_width : int, default=1 - Temporal smoothing width (segments) applied to the estimated drift between optimizer steps. - interpolation_method : {"cubic", "catmull-rom"}, default="cubic" - Spline used to convert per-segment drift to per-frame drift. - max_locs_per_segment : int or None, default=None - Optional downsampling cap per segment (to control memory/time). - mode : str, "cuda" - only Cuda in this version - return_corrected_locs : bool, default=False - If True, also return drift-corrected localizations. - - Returns - ------- - drift_interp_with_frames : ndarray of shape (F, 4) - Per-frame drift with columns [dx_nm, dy_nm, dz_nm, frame]. - corrected_locs : ndarray of shape (N, 4), optional - Only if return_corrected_locs=True. Columns are [x_nm, y_nm, z_nm, segment_id]. - """ - - loc_frames = dataset[:, -1] - if min_max_frames is None: - min_max_frames = (loc_frames.min(), loc_frames.max()) - - # Segment the dataset based on frame numbers into time windows - - result, sorted_dataset, idx_i, idx_j = segmentation_and_pair_indices_wrapper( - dataset, segmentation_var, segmentation_mode, max_drift_nm, max_locs_per_segment, - pair_indices_safety_check=pair_indices_safety_check) - - # Set default initial sigma if not provided - if initial_sigma_nm is None: - initial_sigma_nm = max_drift_nm // 3 - - # Run drift optimization - t0 = time.time() - drift_est = optimize_3d_chunked_better_moving_avg_kd( - result.n_segments, sorted_dataset, - idx_i, idx_j, - sigma_nm=initial_sigma_nm, - target_sigma_nm=target_sigma_nm, - drift_max_nm=max_drift_nm, - drift_max_bound_factor=drift_max_bound_factor, - display_steps=display, - boxcar_width=boxcar_width, - segmentation_result=result, - mode=mode - ) - elapsed = time.time() - t0 - - # Reshape and interpolate drift across all frames - drift_est = drift_est.reshape((result.n_segments, 3)) - vld_tp = np.where(~np.isnan(drift_est[:, 0])) - - frame_interp = np.arange(0, min_max_frames[1] + 1, dtype=int) - drift_interp = interpolate_drift(result.center_frames[vld_tp], drift_est[vld_tp], frame_interp, - method=interpolation_method) - drift_interp_with_frames = np.hstack((drift_interp, frame_interp[:, np.newaxis])) - - # Apply drift correction to original localizations - for i in range(3): - dataset[:, i] = dataset[:, i] - drift_interp[dataset[:, -1].astype(int), i] - - # Optionally show estimated drift curve - if display: - print(f"Drift estimation completed in {elapsed:.2f} seconds.") - plt.figure() - plt.plot(frame_interp, drift_interp) - plt.title("Estimated Drift") - plt.xlabel("Frames") - plt.ylabel("Drift (nm)") - plt.legend(['X', 'Y', 'Z']) - plt.show() - - - # Optional GT comparison plot - if display and gt_drift is not None: - fig, ax = plt.subplots(3, 1) - ax[0].plot(gt_drift[:, 3], gt_drift[:, 0], label='GT Drift X') - ax[1].plot(gt_drift[:, 3], gt_drift[:, 1], label='GT Drift Y') - ax[2].plot(gt_drift[:, 3], gt_drift[:, 2], label='GT Drift Z') - ax[0].plot(frame_interp, drift_interp[:, 0], label='Estimated Drift X', linestyle='--') - ax[1].plot(frame_interp, drift_interp[:, 1], label='Estimated Drift Y', linestyle='--') - ax[2].plot(frame_interp, drift_interp[:, 2], label='Estimated Drift Z', linestyle='--') - ax[1].set_title("Ground Truth vs Estimated Drift") - ax[1].set_xlabel("Frames") - ax[0].set_ylabel("Drift (nm)") - plt.legend() - plt.show() - - # Return corrected locs + drift - if return_corrected_locs: - return drift_interp_with_frames, dataset - else: - return drift_interp_with_frames - - -def optimize_3d_chunked_better_moving_avg_kd(n_segments, locs_nm, idx_i, idx_j, sigma_nm=30, drift_max_nm=300, - target_sigma_nm=30, display_steps=False, - boxcar_width=3, drift_max_bound_factor=2, - segmentation_result=None, - mode="cuda", return_calc_time=False): - """ - Estimate per-segment drift (mu) by minimizing the negative Gaussian-overlap cost - with an L-BFGS-B optimizer and a coarse-to-fine schedule on sigma. - - The routine operates on temporally segmented localizations and reuses a static - neighbor graph (pairs within drift_max_nm). Between optimizer steps, a moving - average (boxcar) can be applied to mu as temporal regularization. Sigma is - reduced iteratively from `sigma_nm` toward `target_sigma_nm` for robust - convergence. - - Parameters - ---------- - n_segments : int - Number of temporal segments S (0..S-1). One 3D drift vector is estimated per segment. - locs_nm : ndarray of shape (M, 3) - localizations in nanometers, columns [x, y, z]; - sigma_nm : float, default=30 - Initial Gaussian width (coarse scale) for the overlap kernel. - drift_max_nm : float, default=300 - Maximum expected drift (nm). Also used as the radius for neighbor pairs and - as the L-BFGS-B bound scale (see `drift_max_bound_factor`). - target_sigma_nm : float, default=30 - Target / final Gaussian width (fine scale). The optimizer reduces sigma toward - this value over iterations. - display_steps : bool, default=False - If True, print or log intermediate progress per iteration/scale. - - boxcar_width : int, default=3 - Temporal smoothing width (in segments) for a moving average applied to mu between steps. - Use 0 or 1 to disable smoothing. - drift_max_bound_factor : float, default=2 - Multiplier for L-BFGS-B bounds around +- drift_max_nm to keep updates physically reasonable. - segmentation_result : object or None, default=None - Segmentation info and metadata. Expected to provide, at minimum: - - segment IDs per localization - - center frames per segment - - any additional structures required by backend (e.g., pair indices) - If None, pairs/ids may be built internally depending on implementation. - mode : str, default=cuda - If True, use the CPU backend; otherwise try GPU (CUDA) and fall back to CPU if unavailable. - return_calc_time : bool, default=False - If True, also return the total computation time in seconds. - - Returns - ------- - mu : ndarray of shape (S, 3) - Estimated per-segment drift (dx, dy, dz) in nanometers. - calc_time_s : float, optional - Only when `return_calc_time=True`. Wall-clock time for the optimization. - """ - - if segmentation_result is None: - segmentation_result = {} - intermediate_results_filehandle = None - sigma_factor = 1.0 - - # Extract coordinate + time arrays, convert to device if CUDA - coords = locs_nm[:, :3].astype(np.float32).copy() - times = locs_nm[:, 3].astype(np.int32).copy() - - chunk_size = int(1E8) # 1E7 - - quality_control = mode == "torch_qc" or mode == "cuda_qc" - - d_coords = cuda.to_device(coords) - d_times = cuda.to_device(times) - if len(idx_i) * 4 > 2e9: - # Use mapped memory if index arrays are large - print("Large index arrays — using mapped memory.") - d_idx_i = cuda.mapped_array_like(idx_i.astype(np.int32), wc=True) - d_idx_j = cuda.mapped_array_like(idx_j.astype(np.int32), wc=True) - d_idx_i[:] = idx_i - d_idx_j[:] = idx_j - else: - d_idx_i = cuda.to_device(idx_i.astype(np.int32)) - d_idx_j = cuda.to_device(idx_j.astype(np.int32)) - # Preallocate device arrays - d_sigma = np.float64(sigma_nm) - d_val = cuda.to_device(np.zeros(chunk_size)) - d_deri = cuda.to_device(np.zeros((n_segments, 3), dtype=np.float64)) - wrapper = cuda_wrapper_chunked - - # Initial drift estimate + bounds - drift_est = np.zeros(n_segments * 3) - bounds = [(-drift_max_nm * drift_max_bound_factor, drift_max_nm * drift_max_bound_factor)] * (3 * n_segments) - - drift_est_gradient = np.inf - fails = 0 - done = False - itr_counter = 0 - start_time = time.time() - - while not done: - # Apply boxcar smoothing to current estimate - tmp = drift_est.reshape((-1, 3)) - for i in range(3): - tmp[:, i] = convolve(tmp[:, i], np.ones(boxcar_width) / boxcar_width) - drift_est = tmp.flatten() - - # Run L-BFGS-B optimization step - result = minimize(wrapper, drift_est, method='L-BFGS-B', - args=( - d_coords, d_times, d_idx_i, d_idx_j, d_sigma, sigma_factor, d_val, d_deri, chunk_size), - jac=True, bounds=bounds, - options={'disp': display_steps, 'gtol': 1E-5, 'ftol': 1E3 * np.finfo(float).eps, - 'maxls': 40}) - itr_counter += 1 - print(f"Iteration {itr_counter}: status = {result.status}, success = {result.success}") - print(f" current sigma: {np.round(sigma_nm * sigma_factor, 2)} nm") - - # Update if successful - if result.success: - delta = np.median((result.x - drift_est) ** 2) - print(f" drift estimate gradient: {delta}") - print(f" previous gradient: {drift_est_gradient}") - # Check convergence - if (delta > drift_est_gradient or sigma_nm * sigma_factor <= 1.0) and sigma_nm * sigma_factor <= target_sigma_nm: - done = True - calc_time = time.time() - start_time - print(f"Optimization completed in {calc_time:.2f} s") - else: - sigma_factor /= 1.5 - drift_est_gradient = delta - drift_est = result.x - else: - fails += 1 - if fails > 2: - sigma_factor *= 2 - print("Restarting with larger sigma_factor") - if fails > 5: - raise RuntimeError("L-BFGS-B Optimization failed after multiple retries") - - if return_calc_time: - return drift_est, time.time() - start_time, itr_counter - else: - return drift_est - - -def segmentation_and_pair_indices_wrapper(dataset, segmentation_var, segmentation_mode, max_drift_nm, - max_locs_per_segment, pair_indices_safety_check=False, hard_limit_pairs=None): - if not segmentation_mode == -1: # -1 is for pre-segmented data - result = segmentation_wrapper(dataset[:, -1], segmentation_var, segmentation_mode, - max_locs_per_segment, return_param_dict=True) - else: - # pre segmented data, anyway we set these values in case auto downsampling is needed - segmentation_mode = 2 # dummy --> segment per frame ... - segmentation_var = 1 # dummy --> ... using 1 frame per segment - if pair_indices_safety_check: - n_pairs_est = estimate_pairs(dataset[result.loc_valid, :3], max_drift_nm) - print(f"Estimated number of pairs within {max_drift_nm} nm: {n_pairs_est:,}") - if hard_limit_pairs is not None and n_pairs_est > hard_limit_pairs: - raise RuntimeError(f"Estimated number of pairs {n_pairs_est} exceeds hard limit of {hard_limit_pairs}. " - f"Aborting to avoid crash.") - if n_pairs_est > 5e8: - print(f"Estimated number of pairs is very large {n_pairs_est}. " - f"Automatic down-sampling is usually required above 500 mil." - f"Billions of pairs can lead to crash.") - ans = input("Continue anyway? (y/n): ") - if ans.lower() != 'y': - raise RuntimeError("Aborted by user due to large estimated number of pairs.") - idx_i, idx_j, successful = pair_indices_kdtree(dataset[result.loc_valid, :3], max_drift_nm) - if not successful: - if max_locs_per_segment is None: - max_locs_per_segment = int(result.out_dict['locs_per_segment'].max()) - while not successful: - max_locs_per_segment = int(max_locs_per_segment * 0.9) - print(f"Segmentation and Pairing attempt failed, automatic down-sampling active...") - print(f"Retrying segmentation with max_locs_per_segment={max_locs_per_segment}...") - result = segmentation_wrapper(dataset[:, -1], segmentation_var, segmentation_mode, - max_locs_per_segment, return_param_dict=True) - sorted_dataset = dataset.copy() - sorted_dataset[:, -1] = result.loc_segments - sorted_dataset = sorted_dataset[result.loc_valid] - idx_i, idx_j, successful = pair_indices_kdtree(sorted_dataset[:, :3], max_drift_nm) - print(f"Segmentation and Pairing successful resulting in {result.n_segments:,} time windows with on average " - f"{int(np.median(result.out_dict['locs_per_segment']))} locs per time window. " - f"{len(idx_i):,} Pairs where found.") - sorted_dataset = dataset.copy() - sorted_dataset[:, -1] = result.loc_segments - sorted_dataset = sorted_dataset[result.loc_valid] - return result, sorted_dataset, idx_i, idx_j - - - -def estimate_pairs(coordinates, distance): - for i in range(len(coordinates[0])): - coordinates[:, i] -= np.min(coordinates[:, i]) - coordinates = np.array(np.floor(coordinates / distance), dtype=int) - coordinates = np.array(list(map(tuple, coordinates))) - sort_indices = np.lexsort(coordinates.T)# get the unique tuples and their counts - unique_tuples, counts = np.unique(coordinates[sort_indices], axis=0, return_counts=True) - # get the indices of the similar tuples - similar_indices = np.split(sort_indices, np.cumsum(counts[:-1])) - idx_i = [] - idx_j = [] - pair_idc_estimate = 0 - for i in range(len(similar_indices)): - n_elements = len(similar_indices[i]) - pair_idc_estimate += n_elements * (n_elements - 1) - rounded = round(pair_idc_estimate,-4) - return rounded - - -def interpolate_drift(center_frames, drift_est, frame_range, method='cubic'): - """ - Interpolates drift estimates to all frames using specified method. - Parameters: - - center_frames: np.ndarray of shape (M,), frames corresponding to drift estimates. - - drift_est: np.ndarray of shape (M, 3), drift estimates at center frames. - - frame_range: array-like, frames to interpolate drift estimates to. - - method: str, interpolation method ('cubic' or 'catmull-rom'). - Returns: - - drift_interp: np.ndarray of shape (len(frame_range), 3), interpolated drift estimates. - """ - if method == 'cubic': - return _interpolate_cubic(center_frames, drift_est, frame_range) - elif method == 'catmull-rom': - return _interpolate_catmull_rom(center_frames, drift_est, frame_range) - elif method == 'linear': - drift_x = np.interp(frame_range, center_frames, drift_est[:, 0]) - drift_y = np.interp(frame_range, center_frames, drift_est[:, 1]) - drift_z = np.interp(frame_range, center_frames, drift_est[:, 2]) - return np.vstack([drift_x, drift_y, drift_z]).T - else: - raise ValueError(f"Unknown interpolation method: {method}") - - -def _interpolate_cubic(center_frames, drift_est, frame_range): - from scipy.interpolate import CubicSpline - drift_x = CubicSpline(center_frames, drift_est[:, 0])(frame_range) - drift_y = CubicSpline(center_frames, drift_est[:, 1])(frame_range) - drift_z = CubicSpline(center_frames, drift_est[:, 2])(frame_range) - return np.vstack([drift_x, drift_y, drift_z]).T - - -def _interpolate_catmull_rom(center_frames, drift_est, frame_range): - def catmull_rom_1d(x, y, x_interp): - result = np.zeros_like(x_interp) - for i in range(1, len(x) - 2): - x0, x1, x2, x3 = x[i-1], x[i], x[i+1], x[i+2] - y0, y1, y2, y3 = y[i-1], y[i], y[i+1], y[i+2] - - mask = (x_interp >= x1) & (x_interp <= x2) - t = (x_interp[mask] - x1) / (x2 - x1) - - result[mask] = ( - 0.5 * ( - (2 * y1) + - (-y0 + y2) * t + - (2*y0 - 5*y1 + 4*y2 - y3) * t**2 + - (-y0 + 3*y1 - 3*y2 + y3) * t**3 - ) - ) - return result - - x_interp = np.asarray(frame_range) - x = np.asarray(center_frames) - - if len(x) < 4: - raise ValueError("Catmull-Rom interpolation requires at least 4 points.") - - drift_x = catmull_rom_1d(x, drift_est[:, 0], x_interp) - drift_y = catmull_rom_1d(x, drift_est[:, 1], x_interp) - drift_z = catmull_rom_1d(x, drift_est[:, 2], x_interp) - - return np.vstack([drift_x, drift_y, drift_z]).T - - -# Numba cuda code of cost-function --> quantized overlap of pairs of localizations, gaussian mixture model -@cuda.jit -def cost_function_full_3d_chunked(d_locs_time, start_idx, chunk_size, d_idx_i, d_idx_j, d_sigma, d_sigma_factor, d_val, - d_val_sum, d_deri, d_locs_coords, mu): - """Compute negative-overlap cost and gradient for 3D localizations (GPU/CPU backend; internal use).""" - tx = cuda.threadIdx.x - ty = cuda.blockIdx.x - bw = cuda.blockDim.x - pos = tx + ty * bw - - if pos < chunk_size: - i = d_idx_i[pos + start_idx] - j = d_idx_j[pos + start_idx] - - ti = d_locs_time[i] - tj = d_locs_time[j] - - dx = (d_locs_coords[i, 0] - mu[ti, 0]) - (d_locs_coords[j, 0] - mu[tj, 0]) - dy = (d_locs_coords[i, 1] - mu[ti, 1]) - (d_locs_coords[j, 1] - mu[tj, 1]) - dz = (d_locs_coords[i, 2] - mu[ti, 2]) - (d_locs_coords[j, 2] - mu[tj, 2]) - sigma_sq = (2 * d_sigma * d_sigma_factor) ** 2 - - diff_sq = dx * dx + dy * dy + dz * dz - val = 1 / (d_sigma * d_sigma_factor) * math.exp(-diff_sq / sigma_sq) - d_val[pos] = val - - # Update derivatives - cuda.atomic.add(d_deri, (tj, 0), 2 * val * (d_locs_coords[j, 0] - d_locs_coords[i, 0] + mu[ti, 0] - mu[tj, 0]) / sigma_sq) - cuda.atomic.add(d_deri, (tj, 1), 2 * val * (d_locs_coords[j, 1] - d_locs_coords[i, 1] + mu[ti, 1] - mu[tj, 1]) / sigma_sq) - cuda.atomic.add(d_deri, (tj, 2), 2 * val * (d_locs_coords[j, 2] - d_locs_coords[i, 2] + mu[ti, 2] - mu[tj, 2]) / sigma_sq) - - cuda.atomic.add(d_deri, (ti, 0), 2 * val * (d_locs_coords[i, 0] - d_locs_coords[j, 0] + mu[tj, 0] - mu[ti, 0]) / sigma_sq) - cuda.atomic.add(d_deri, (ti, 1), 2 * val * (d_locs_coords[i, 1] - d_locs_coords[j, 1] + mu[tj, 1] - mu[ti, 1]) / sigma_sq) - cuda.atomic.add(d_deri, (ti, 2), 2 * val * (d_locs_coords[i, 2] - d_locs_coords[j, 2] + mu[tj, 2] - mu[ti, 2]) / sigma_sq) - - cuda.atomic.add(d_val_sum, 0, val) - d_val[pos] = 0 - - -# Interface between the Python code and the CUDA kernel, mainly for chunking the data to avoid memory issues -def cuda_wrapper_chunked(mu, d_locs_coords, d_locs_time, d_idx_i, d_idx_j, d_sigma, d_sigma_factor, d_val, d_deri, - chunk_size): - val_total = 0 - d_val_sum = cuda.to_device(np.zeros(1, dtype=np.float64)) - mu_dev = cuda.to_device(np.asarray(mu.reshape(int(mu.size / 3), 3), dtype=np.float64)) - - n_chunks = int(np.ceil(d_idx_i.size / chunk_size)) - threadsperblock = 128 - - for i in range(n_chunks - 1): - idc_start = i*chunk_size - blockspergrid = (chunk_size + (threadsperblock - 1)) // threadsperblock - cost_function_full_3d_chunked[blockspergrid, threadsperblock]( - d_locs_time, idc_start, chunk_size, d_idx_i, d_idx_j, - d_sigma, d_sigma_factor, d_val, d_val_sum, d_deri, d_locs_coords, mu_dev - ) - val_total += d_val_sum.copy_to_host() - - # Final chunk - n_remaining = d_idx_i.size - (n_chunks - 1) * chunk_size - idc_start = (n_chunks - 1) * chunk_size - blockspergrid = (n_remaining + (threadsperblock - 1)) // threadsperblock - cost_function_full_3d_chunked[blockspergrid, threadsperblock]( - d_locs_time, idc_start, n_remaining, d_idx_i, d_idx_j, - d_sigma, d_sigma_factor, d_val, d_val_sum, d_deri, d_locs_coords, mu_dev - ) - val_total += d_val_sum.copy_to_host() - deri = d_deri.copy_to_host() - d_deri[:] = 0 - - return -np.nansum(val_total), -deri.flatten() - -from dataclasses import dataclass -from typing import Optional, Dict - - -@dataclass -class SegmentationResult: - """Container for segmentation results.""" - loc_segments: np.ndarray - loc_valid: np.ndarray - center_frames: np.ndarray - n_segments: int - out_dict: Optional[Dict] = None - - -def _group_by_frame(loc_frames: np.ndarray): - """Returns a dict {frame_number: indices_in_loc_frames} efficiently.""" - sort_idx = np.argsort(loc_frames) - sorted_frames = loc_frames[sort_idx] - unique_frames, start_idx, counts = np.unique(sorted_frames, return_index=True, return_counts=True) - frame_to_indices = { - frame: sort_idx[start:start + count] - for frame, start, count in zip(unique_frames, start_idx, counts) - } - return unique_frames, frame_to_indices - - -def segment_by_num_locs_per_window(loc_frames: np.ndarray, min_n_locs_per_window: int, - max_locs_per_segment = None, - return_param_dict: bool = False) -> SegmentationResult: - """ - Segments by collecting a minimum number of localizations per window. - Once the threshold is met and enough locs remain, a new segment is created. - This method ensures that each segment has at least `min_n_locs_per_window` localizations, - while also trying to avoid creating segments that are too small at the end of the dataset. - If `max_locs_per_segment` is set, a random subset of that size is chosen from each segment. - This method is particularly useful for datasets with varying localization densities over time. - Parameters: - loc_frames (np.ndarray): Array of frame numbers for each localization. - min_n_locs_per_window (int): Minimum number of localizations per segment. - max_locs_per_segment (Optional[int]): Maximum number of localizations per segment. If None, all locs are used. - return_param_dict (bool): Whether to return a dictionary of segmentation parameters. - Returns: - SegmentationResult: A dataclass containing segmentation results and parameters. - """ - loc_frames = np.asarray(loc_frames, dtype=int) - n_locs = len(loc_frames) - - if max_locs_per_segment is not None and max_locs_per_segment < 1: # downsampling in percentage - max_locs_per_segment = int(min_n_locs_per_window * max_locs_per_segment) - - unique_frames, frame_to_indices = _group_by_frame(loc_frames) - loc_segments = np.full(n_locs, -1, dtype=int) # Default to -1 for safety - segment_counter = 0 - n_locs_in_current_segment = 0 - current_segment_indices = [] - start_frames, end_frames, locs_per_segment = [], [], [] - - for i, frame in enumerate(unique_frames): - indices = frame_to_indices[frame] - n_locs_this_frame = len(indices) - remaining_locs = n_locs - (len(current_segment_indices) + n_locs_this_frame + np.sum(locs_per_segment)) - - # Add frame to current segment if: - # - It fills the current segment to threshold - # - AND there are enough locs left for another segment (or it's the last frame) - if (n_locs_in_current_segment + n_locs_this_frame >= min_n_locs_per_window) and \ - (remaining_locs >= min_n_locs_per_window or i == len(unique_frames) - 1): - current_segment_indices.extend(indices) - n_locs_in_current_segment += n_locs_this_frame - - loc_segments[current_segment_indices] = segment_counter - start_frames.append(loc_frames[current_segment_indices[0]]) - end_frames.append(loc_frames[current_segment_indices[-1]]) - locs_per_segment.append(len(current_segment_indices)) - - segment_counter += 1 - current_segment_indices = [] - n_locs_in_current_segment = 0 - else: - # Defer frame to current segment - current_segment_indices.extend(indices) - n_locs_in_current_segment += n_locs_this_frame - - n_segments = segment_counter - center_frames = np.zeros(n_segments) - loc_valid = np.zeros(n_locs, dtype=bool) - - for i in range(n_segments): - segment_indices = np.where(loc_segments == i)[0] - if max_locs_per_segment and len(segment_indices) > max_locs_per_segment: - selected = np.random.choice(segment_indices, max_locs_per_segment, replace=False) - else: - selected = segment_indices - loc_valid[selected] = True - locs_per_segment[i] = len(selected) - center_frames[i] = np.mean(loc_frames[selected]) - - out_dict = None - if return_param_dict: - n_locs_valid = loc_valid.sum() - out_dict = { - "n_segments": n_segments, - "min_n_locs_per_window": min_n_locs_per_window, - "frames_per_window": -1, - "start_frames": np.array(start_frames), - "end_frames": np.array(end_frames), - "locs_per_segment": np.array(locs_per_segment), - "n_locs": n_locs, - "n_locs_valid": n_locs_valid, - "n_locs_invalid": n_locs - n_locs_valid, - "center_frames": center_frames - } - - return SegmentationResult(loc_segments, loc_valid, center_frames, n_segments, out_dict) - - -def segment_by_frame_windows(loc_frames: np.ndarray, n_frames_per_window: int, - max_locs_per_segment = None, - return_param_dict: bool = False) -> SegmentationResult: - """ - Splits localization data into fixed-size windows of N frames. - All localizations in those frames are grouped into one segment. - """ - loc_frames = np.asarray(loc_frames, dtype=int) - frames, frame_to_indices = _group_by_frame(loc_frames) - n_locs = len(loc_frames) - n_segments = int(np.ceil(len(frames) / n_frames_per_window)) - - if max_locs_per_segment is not None and max_locs_per_segment < 1: # downsampling in percentage - max_locs_per_segment = n_locs / n_segments * max_locs_per_segment - - loc_segments = np.zeros(n_locs, dtype=int) - center_frames = np.zeros(n_segments) - loc_valid = np.ones(n_locs, dtype=bool) - start_frames, end_frames, locs_per_segment = [], [], [] - - for i in range(n_segments): - frame_window = frames[i * n_frames_per_window:(i + 1) * n_frames_per_window] - indices = np.concatenate([frame_to_indices[frame] for frame in frame_window if frame in frame_to_indices]) - if len(indices) == 0: - continue - loc_segments[indices] = i - start_frames.append(frame_window[0]) - end_frames.append(frame_window[-1]) - center_frames[i] = np.mean(loc_frames[indices]) - locs_per_segment.append(len(indices)) - if max_locs_per_segment and len(indices) > max_locs_per_segment: - mask = np.ones(len(indices), dtype=bool) - mask[np.random.choice(len(indices), len(indices) - max_locs_per_segment, replace=False)] = False - loc_valid[indices[~mask]] = False - - out_dict = None - if return_param_dict: - n_locs_valid = loc_valid.sum() - out_dict = { - "n_segments": n_segments, - "min_n_locs_per_window": -1, - "frames_per_window": n_frames_per_window, - "start_frames": np.array(start_frames), - "end_frames": np.array(end_frames), - "locs_per_segment": np.array(locs_per_segment), - "n_locs": n_locs, - "n_locs_valid": n_locs_valid, - "n_locs_invalid": n_locs - n_locs_valid, - "center_frames": center_frames - } - - return SegmentationResult(loc_segments, loc_valid, center_frames, n_segments, out_dict) - - -def segment_by_num_windows(loc_frames: np.ndarray, n_windows: int, max_locs_per_segment = None, - return_param_dict: bool = False) -> SegmentationResult: - """ - Converts number of windows into an equivalent minimum locs per window, - then calls `segment_by_num_locs_per_window`. - """ - n_locs = len(loc_frames) - n_locs_per_window = int(np.ceil(n_locs / n_windows)) - if max_locs_per_segment is not None and max_locs_per_segment < 1: # downsampling in percentage - max_locs_per_segment = int(n_locs_per_window * max_locs_per_segment) - return segment_by_num_locs_per_window(loc_frames, n_locs_per_window, max_locs_per_segment, return_param_dict) - - -def segmentation_wrapper(loc_frames: np.ndarray, segmentation_var: int, segmentation_mode: int = 2, - max_locs_per_segment = None, - return_param_dict: bool = False) -> SegmentationResult: - """ - Dispatch function that selects segmentation method: - 0 → fixed number of windows - 1 → fixed number of localizations per window - 2 → fixed number of frames per window (default) - """ - if segmentation_mode == 0: - return segment_by_num_windows(loc_frames, segmentation_var, max_locs_per_segment, return_param_dict) - elif segmentation_mode == 1: - return segment_by_num_locs_per_window(loc_frames, segmentation_var, max_locs_per_segment, return_param_dict) - else: - return segment_by_frame_windows(loc_frames, segmentation_var, max_locs_per_segment, return_param_dict) - - -def pair_indices_kdtree(coordinates, distance): - """ - Find all pairs of points within a certain distance using a KD-tree. - Parameters: - - coordinates: np.ndarray of shape (N, D) where N is the number of points and D is the dimensionality. - - distance: float, the maximum distance to consider points as a pair. - Returns: - - idx1: np.ndarray of shape (M,), indices of the first point in each pair. - - idx2: np.ndarray of shape (M,), indices of the second point in each pair. - """ - tree = cKDTree(coordinates) - try: - pairs = tree.query_pairs(r=distance, output_type='ndarray') - except MemoryError: - print("[pair_indices_kdtree] MemoryError encountered") - return [], [], False - return np.ascontiguousarray(pairs[:, 0], dtype=np.int32), np.ascontiguousarray(pairs[:, 1], dtype=np.int32), True \ No newline at end of file diff --git a/picasso/gui/render.py b/picasso/gui/render.py index 4952a21b..1f1bb79b 100644 --- a/picasso/gui/render.py +++ b/picasso/gui/render.py @@ -40,7 +40,6 @@ from PyQt5 import QtCore, QtGui, QtWidgets from .. import ( aim, - comet, clusterer, g5m, imageprocess, @@ -53,8 +52,9 @@ ) from .rotation import RotationWindow -# PyImarisWrite works on windows only +# Optional modules with external/hardware dependencies live in ext/ from ..ext.bitplane import IMSWRITER +from ..ext import comet if IMSWRITER: from ..ext.bitplane import numpy_to_imaris @@ -10685,11 +10685,23 @@ def undrift_comet(self) -> None: params, ok = COMETDialog.getParams(self.window) if ok: - locs, new_info, drift = comet.comet( - locs, - info, - **params, - ) + try: + locs, new_info, drift = comet.comet( + locs, + info, + **params, + ) + except RuntimeError as e: + QtWidgets.QMessageBox.warning( + self.window, + "COMET not available", + ( + f"{e}\n\n" + "Please use a different drift-correction method, " + "such as Undrift by AIM or Undrift by RCC." + ), + ) + return locs = lib.ensure_sanity(locs, info) self.all_locs[channel] = locs @@ -11855,10 +11867,10 @@ def initUI(self, plugins_loaded: bool) -> None: # menu bar - Postprocess postprocess_menu = self.menu_bar.addMenu("Postprocess") undrift_comet_action = postprocess_menu.addAction("Undrift by COMET") - undrift_comet_action.setShortcut("Ctrl+U") undrift_comet_action.triggered.connect(self.view.undrift_comet) undrift_aim_action = postprocess_menu.addAction("Undrift by AIM") + undrift_aim_action.setShortcut("Ctrl+U") undrift_aim_action.triggered.connect(self.view.undrift_aim) undrift_from_picked_action = postprocess_menu.addAction( "Undrift from picked" diff --git a/tests/test_comet.py b/tests/test_comet.py new file mode 100644 index 00000000..976b3399 --- /dev/null +++ b/tests/test_comet.py @@ -0,0 +1,322 @@ +"""Tests for picasso.comet — drift-correction with COMET. + +Tests are split into two groups: +- Pure-Python helpers (segmentation, pairing, interpolation): always run when + numba is installed, no GPU required. +- Full pipeline (comet.comet): requires CUDA hardware; skipped automatically + on CPU-only machines, but the no-GPU error path IS tested on any machine. + +:author: Lenny Reinkensmeier, 2026 +""" + +import numpy as np +import pandas as pd +import pytest + +# The whole picasso package requires numba (lib.py imports it). +# Skip this module cleanly when numba is absent (e.g. CI without GPU stack). +pytest.importorskip("numba") + +from picasso.ext import comet # noqa: E402 — import after guard + + +# --------------------------------------------------------------------------- +# Fixtures +# --------------------------------------------------------------------------- + +@pytest.fixture +def minimal_info(): + """Minimal Picasso info list with required Pixelsize entry.""" + return [{"Pixelsize": 130, "Frames": 200}] + + +@pytest.fixture +def minimal_locs_df(): + """Small synthetic localisation DataFrame spread over 200 frames.""" + rng = np.random.default_rng(0) + n_frames = 200 + n_per_frame = 5 + frames = np.repeat(np.arange(n_frames), n_per_frame) + x = rng.uniform(5, 25, size=len(frames)).astype(np.float32) + y = rng.uniform(5, 25, size=len(frames)).astype(np.float32) + return pd.DataFrame({"frame": frames, "x": x, "y": y}) + + +@pytest.fixture +def minimal_locs_structured(minimal_locs_df): + """Same data as a numpy structured array (Picasso's native format).""" + df = minimal_locs_df + dtype = np.dtype([("frame", np.uint32), ("x", np.float32), ("y", np.float32)]) + arr = np.empty(len(df), dtype=dtype) + arr["frame"] = df["frame"].to_numpy(np.uint32) + arr["x"] = df["x"].to_numpy(np.float32) + arr["y"] = df["y"].to_numpy(np.float32) + return arr + + +# --------------------------------------------------------------------------- +# Segmentation helpers +# --------------------------------------------------------------------------- + +class TestSegmentationByLocs: + def test_produces_correct_segment_count(self): + frames = np.repeat(np.arange(100), 5) # 500 locs, 5/frame + result = comet.segmentation_wrapper( + frames, segmentation_var=50, segmentation_mode=1, return_param_dict=True + ) + assert result.n_segments == 10 + + def test_all_locs_assigned(self): + frames = np.repeat(np.arange(50), 10) + result = comet.segmentation_wrapper( + frames, segmentation_var=100, segmentation_mode=1 + ) + assert (result.loc_segments >= 0).all() + + def test_loc_valid_subset_of_locs(self): + frames = np.repeat(np.arange(60), 8) + result = comet.segmentation_wrapper( + frames, segmentation_var=40, segmentation_mode=1 + ) + assert result.loc_valid.shape == frames.shape + assert result.loc_valid.dtype == bool + + def test_center_frames_length_matches_n_segments(self): + frames = np.repeat(np.arange(80), 5) + result = comet.segmentation_wrapper( + frames, segmentation_var=50, segmentation_mode=1 + ) + assert len(result.center_frames) == result.n_segments + + def test_max_locs_per_segment_respected(self): + frames = np.repeat(np.arange(20), 100) # 2000 locs, 100/frame + cap = 30 + result = comet.segmentation_wrapper( + frames, segmentation_var=200, segmentation_mode=1, + max_locs_per_segment=cap, return_param_dict=True + ) + assert result.loc_valid.sum() <= result.n_segments * cap + + +class TestSegmentationByFrameWindows: + def test_frame_window_mode(self): + frames = np.repeat(np.arange(100), 3) + result = comet.segmentation_wrapper( + frames, segmentation_var=10, segmentation_mode=2 + ) + assert result.n_segments == 10 # 100 frames / 10 per window + + def test_single_window(self): + frames = np.arange(50) + result = comet.segmentation_wrapper( + frames, segmentation_var=100, segmentation_mode=2 + ) + assert result.n_segments == 1 + + +class TestSegmentationByNumWindows: + def test_num_windows_mode_creates_requested_count(self): + frames = np.repeat(np.arange(100), 5) # 500 locs + result = comet.segmentation_wrapper( + frames, segmentation_var=5, segmentation_mode=0, return_param_dict=True + ) + # mode 0 derives min-locs-per-window from total/n_windows; the resulting + # segment count should be approximately n_windows (off-by-one tolerated + # because the last segment is folded into the previous if too small). + assert abs(result.n_segments - 5) <= 1 + + def test_all_locs_assigned(self): + frames = np.repeat(np.arange(40), 4) + result = comet.segmentation_wrapper( + frames, segmentation_var=4, segmentation_mode=0 + ) + assert (result.loc_segments >= 0).all() + + +# --------------------------------------------------------------------------- +# Pair-finding +# --------------------------------------------------------------------------- + +class TestPairIndicesKdtree: + def test_finds_close_pairs(self): + coords = np.array([[0, 0, 0], [1, 0, 0], [100, 0, 0]], dtype=np.float32) + idx_i, idx_j, ok = comet.pair_indices_kdtree(coords, distance=5) + assert ok + assert len(idx_i) == 1 # only (0,1) are within distance 5 + + def test_no_pairs_when_far_apart(self): + coords = np.eye(3, dtype=np.float32) * 1000 + idx_i, idx_j, ok = comet.pair_indices_kdtree(coords, distance=1) + assert ok + assert len(idx_i) == 0 + + def test_all_pairs_when_all_close(self): + coords = np.zeros((5, 3), dtype=np.float32) + idx_i, idx_j, ok = comet.pair_indices_kdtree(coords, distance=1) + assert ok + assert len(idx_i) == 5 * 4 // 2 # C(5,2) + + def test_returns_int32_arrays(self): + coords = np.random.rand(20, 3).astype(np.float32) + idx_i, idx_j, ok = comet.pair_indices_kdtree(coords, distance=0.5) + assert ok + assert idx_i.dtype == np.int32 + assert idx_j.dtype == np.int32 + + +# --------------------------------------------------------------------------- +# Drift interpolation +# --------------------------------------------------------------------------- + +class TestInterpolateDrift: + def _center_frames_and_drift(self): + center_frames = np.linspace(10, 90, 6) + drift_est = np.column_stack([ + np.sin(center_frames / 10), + np.cos(center_frames / 10), + np.zeros_like(center_frames), + ]) + return center_frames, drift_est + + def test_cubic_output_shape(self): + cf, de = self._center_frames_and_drift() + frame_range = np.arange(0, 100) + out = comet.interpolate_drift(cf, de, frame_range, method="cubic") + assert out.shape == (100, 3) + + def test_linear_output_shape(self): + cf, de = self._center_frames_and_drift() + frame_range = np.arange(0, 100) + out = comet.interpolate_drift(cf, de, frame_range, method="linear") + assert out.shape == (100, 3) + + def test_cubic_passes_through_control_points(self): + cf = np.array([0.0, 25.0, 50.0, 75.0, 100.0]) + de = np.column_stack([cf * 0.01, cf * 0.005, np.zeros_like(cf)]) + out = comet.interpolate_drift(cf, de, cf.astype(int), method="cubic") + np.testing.assert_allclose(out[:, 0], de[:, 0], atol=1e-6) + + def test_unknown_method_raises(self): + cf, de = self._center_frames_and_drift() + with pytest.raises(ValueError, match="Unknown interpolation method"): + comet.interpolate_drift(cf, de, np.arange(100), method="spagetti") + + def test_catmull_rom_output_shape(self): + cf, de = self._center_frames_and_drift() + frame_range = np.arange(10, 90) + out = comet.interpolate_drift(cf, de, frame_range, method="catmull-rom") + assert out.shape == (len(frame_range), 3) + + def test_catmull_rom_requires_min_points(self): + cf = np.array([0.0, 10.0, 20.0]) # only 3 points + de = np.zeros((3, 3)) + with pytest.raises(ValueError, match="at least 4 points"): + comet.interpolate_drift(cf, de, np.arange(20), method="catmull-rom") + + +# --------------------------------------------------------------------------- +# comet() public API — no-GPU error path +# --------------------------------------------------------------------------- + +class TestCometNoCuda: + """These tests run on any machine; they verify the failure mode when + CUDA is not present.""" + + def test_raises_runtime_error_without_cuda(self, minimal_locs_df, minimal_info): + if comet._CUDA_AVAILABLE: + pytest.skip("CUDA is present; testing no-GPU path only on CPU machines") + with pytest.raises(RuntimeError, match="CUDA"): + comet.comet(minimal_locs_df, minimal_info, locs_per_segment=100, max_drift_nm=200) + + def test_error_message_mentions_alternatives(self, minimal_locs_df, minimal_info): + if comet._CUDA_AVAILABLE: + pytest.skip("CUDA is present; testing no-GPU path only on CPU machines") + with pytest.raises(RuntimeError, match="AIM|RCC"): + comet.comet(minimal_locs_df, minimal_info, locs_per_segment=100, max_drift_nm=200) + + def test_structured_array_input_accepted_before_gpu_check(self, minimal_locs_structured, minimal_info): + """Numpy structured arrays should be converted to DataFrame before the + CUDA check, so the error is still a RuntimeError (not an AttributeError + from missing .columns).""" + if comet._CUDA_AVAILABLE: + pytest.skip("CUDA is present; testing no-GPU path only on CPU machines") + with pytest.raises(RuntimeError, match="CUDA"): + comet.comet(minimal_locs_structured, minimal_info, locs_per_segment=100, max_drift_nm=200) + + def test_missing_columns_raises_key_error(self, minimal_info): + """Locs missing required columns should raise KeyError, not crash on .columns.""" + if comet._CUDA_AVAILABLE: + pytest.skip("CUDA is present; GPU path reached before column check") + bad_locs = pd.DataFrame({"frame": np.arange(10), "x": np.zeros(10)}) + # No 'y' column — should raise KeyError from the column check + with pytest.raises((KeyError, RuntimeError)): + comet.comet(bad_locs, minimal_info, locs_per_segment=5, max_drift_nm=100) + + def test_z_column_accepted(self, minimal_locs_df, minimal_info): + """3D input (with z) should not crash before the GPU check.""" + if comet._CUDA_AVAILABLE: + pytest.skip("CUDA is present; testing no-GPU path only on CPU machines") + df = minimal_locs_df.copy() + df["z"] = np.zeros(len(df), dtype=np.float32) + with pytest.raises(RuntimeError, match="CUDA"): + comet.comet(df, minimal_info, locs_per_segment=100, max_drift_nm=200) + + def test_missing_pixelsize_raises_on_gpu(self, minimal_locs_df): + """info without 'Pixelsize' should fail with an informative error. + Only meaningful on GPU since the CUDA check fires first on CPU.""" + if not comet._CUDA_AVAILABLE: + pytest.skip("CUDA check fires before pixelsize check on CPU") + bad_info = [{"Frames": 200}] + with pytest.raises((KeyError, RuntimeError, ValueError)): + comet.comet(minimal_locs_df, bad_info, locs_per_segment=100, max_drift_nm=200) + + +# --------------------------------------------------------------------------- +# comet() public API — full pipeline (GPU required) +# --------------------------------------------------------------------------- + +@pytest.mark.skipif( + not comet._CUDA_AVAILABLE, + reason="CUDA GPU not available", +) +class TestCometFullPipeline: + """Integration tests that run the full optimisation loop on GPU.""" + + def test_output_shapes(self, minimal_locs_df, minimal_info): + undrifted, new_info, drift = comet.comet( + minimal_locs_df, minimal_info, + locs_per_segment=100, max_drift_nm=300, + target_sigma_nm=50, + ) + assert len(undrifted) == len(minimal_locs_df) + assert set(undrifted.columns) >= {"x", "y", "frame"} + assert "x" in drift.columns and "y" in drift.columns + + def test_drift_covers_frame_range(self, minimal_locs_df, minimal_info): + frame0 = minimal_locs_df["frame"] - minimal_locs_df["frame"].min() + _, _, drift = comet.comet( + minimal_locs_df, minimal_info, + locs_per_segment=100, max_drift_nm=300, + target_sigma_nm=50, + ) + assert len(drift) == int(frame0.max()) + 1 + + def test_new_info_appended(self, minimal_locs_df, minimal_info): + _, new_info, _ = comet.comet( + minimal_locs_df, minimal_info, + locs_per_segment=100, max_drift_nm=300, + target_sigma_nm=50, + ) + assert len(new_info) == len(minimal_info) + 1 + assert "COMET" in new_info[-1]["Generated by"] + + def test_structured_array_gives_same_result_as_dataframe(self, minimal_locs_df, minimal_locs_structured, minimal_info): + locs_df, _, drift_df = comet.comet( + minimal_locs_df, minimal_info, + locs_per_segment=100, max_drift_nm=300, target_sigma_nm=50, + ) + locs_arr, _, drift_arr = comet.comet( + minimal_locs_structured, minimal_info, + locs_per_segment=100, max_drift_nm=300, target_sigma_nm=50, + ) + np.testing.assert_allclose(drift_df["x"].to_numpy(), drift_arr["x"].to_numpy(), rtol=1e-5) From 431b22a38bc74c4b0ce09cce584df2fa5e745141 Mon Sep 17 00:00:00 2001 From: Lenny Reinkensmeier Date: Thu, 28 May 2026 17:17:43 +0200 Subject: [PATCH 03/19] Handle COMET availability in drift correction and update menu shortcuts. Added proper testing for COMET. --- picasso/comet.py | 852 ------------------------------------------ picasso/gui/render.py | 28 +- tests/test_comet.py | 322 ++++++++++++++++ 3 files changed, 342 insertions(+), 860 deletions(-) delete mode 100644 picasso/comet.py create mode 100644 tests/test_comet.py diff --git a/picasso/comet.py b/picasso/comet.py deleted file mode 100644 index 11321927..00000000 --- a/picasso/comet.py +++ /dev/null @@ -1,852 +0,0 @@ -import numpy as np -import pandas as pd -import matplotlib.pyplot as plt -from numba import cuda -from scipy.ndimage import convolve -from scipy.optimize import minimize -import time -import math -from scipy.spatial import cKDTree - -from . import lib, __version__ - - - - -def comet( - locs: pd.DataFrame, - info: list[dict], - locs_per_segment: int, - max_drift_nm: float, - max_locs_per_segment: int | None = None, - *, - initial_sigma_nm: float | None = None, - target_sigma_nm: float = 10.0, - boxcar_width: int = 3, - drift_max_bound_factor: float = 2.0, - interpolation_method: str = "cubic", - mode: str = "cuda", - display: bool = False, - progress=None, # currently unused -) -> tuple[pd.DataFrame, list[dict], pd.DataFrame]: - """Apply COMET undrifting to the localizations. - - Parameters - ---------- - locs : pd.DataFrame - Localization table. Must contain columns 'x', 'y', 'frame'. - Optional: 'z'. - x/y are expected in camera pixels, z in the native z-unit used by Picasso. - info : list of dict - Localization metadata. - locs_per_segment : int - Target number of localizations per segment for COMET segmentation. - max_drift_nm : float - Maximum expected drift in nm. - max_locs_per_segment : int or None, optional - Optional cap for localizations per segment. If negative, treated as None. - initial_sigma_nm : float or None, optional - Initial Gaussian sigma for COMET. If None, defaults to max_drift_nm / 3. - target_sigma_nm : float, optional - Final Gaussian sigma for refinement. - boxcar_width : int, optional - Smoothing width over segments. - drift_max_bound_factor : float, optional - Bound factor for optimizer. - interpolation_method : str, optional - Drift interpolation method. - mode : str, optional - Backend mode, only "cuda". - display : bool, optional - Whether to show diagnostic plots. - progress : optional - Placeholder for API compatibility with AIM. - - Returns - ------- - locs : pd.DataFrame - Undrift-corrected localizations. - new_info : list[dict] - Updated metadata. - drift : pd.DataFrame - Per-frame drift table with columns x, y, and optionally z. - Drift is returned in the same units as locs: - x/y in pixels, z in the original z-unit. - """ - locs = locs.copy() - - if max_locs_per_segment is not None and max_locs_per_segment < 0: - max_locs_per_segment = None - - pixelsize_nm = lib.get_from_metadata(info, "Pixelsize", raise_error=True) - - required_cols = {"x", "y", "frame"} - missing = required_cols - set(locs.columns) - if missing: - raise KeyError(f"Missing required columns for COMET: {sorted(missing)}") - - has_z = "z" in locs.columns - - # Normalize frames so they start at 0 for internal indexing. - # This is important because COMET later indexes drift by frame number. - frame = locs["frame"].to_numpy(dtype=np.int64) - frame0 = frame - frame.min() - - # Build COMET input array in nm - dataset = np.zeros((len(locs), 4), dtype=np.float64) - dataset[:, 0] = locs["x"].to_numpy(dtype=np.float64) * pixelsize_nm - dataset[:, 1] = locs["y"].to_numpy(dtype=np.float64) * pixelsize_nm - dataset[:, 2] = ( - locs["z"].to_numpy(dtype=np.float64) if has_z else 0.0 - ) - dataset[:, 3] = frame0 - - if initial_sigma_nm is None: - initial_sigma_nm = max_drift_nm / 3 - - drift_nm_with_frame = comet_run_kd( - dataset=dataset, - segmentation_mode=1, # fixed locs per segment - segmentation_var=locs_per_segment, - max_locs_per_segment=max_locs_per_segment, - initial_sigma_nm=initial_sigma_nm, - gt_drift=None, - display=display, - return_corrected_locs=False, - max_drift_nm=max_drift_nm, - target_sigma_nm=target_sigma_nm, - boxcar_width=boxcar_width, - drift_max_bound_factor=drift_max_bound_factor, - interpolation_method=interpolation_method, - mode=mode, - min_max_frames=(int(frame0.min()), int(frame0.max())), - ) - - # drift_nm_with_frame columns: dx_nm, dy_nm, dz_nm, frame0 - drift_nm = drift_nm_with_frame[:, :3] - - # Apply drift back to dataframe in original units - frame_idx = frame0.astype(np.int64) - locs["x"] = locs["x"].to_numpy(dtype=np.float64) - drift_nm[frame_idx, 0] / pixelsize_nm - locs["y"] = locs["y"].to_numpy(dtype=np.float64) - drift_nm[frame_idx, 1] / pixelsize_nm - if has_z: - locs["z"] = locs["z"].to_numpy(dtype=np.float64) - drift_nm[frame_idx, 2] - - # Build drift dataframe in Picasso-style units - drift_dict = { - "x": (drift_nm[:, 0] / pixelsize_nm).astype("float32"), - "y": (drift_nm[:, 1] / pixelsize_nm).astype("float32"), - } - if has_z: - drift_dict["z"] = drift_nm[:, 2].astype("float32") - - drift = pd.DataFrame(drift_dict) - - new_info_entry = { - "Generated by": f"Picasso v{__version__} COMET", - "Segmentation mode": "localizations per segment", - "Localizations per segment": locs_per_segment, - "Maximum drift (nm)": max_drift_nm, - "Initial sigma (nm)": float(initial_sigma_nm), - "Target sigma (nm)": float(target_sigma_nm), - "Boxcar width": int(boxcar_width), - "Max localizations per segment": ( - None if max_locs_per_segment is None else int(max_locs_per_segment) - ), - "Interpolation method": interpolation_method, - "Backend": mode, - } - new_info = info + [new_info_entry] - - return locs, new_info, drift - -def comet_run_kd(dataset, segmentation_mode, segmentation_var, max_locs_per_segment=None, - initial_sigma_nm=None, gt_drift=None, display=False, return_corrected_locs=False, - max_drift_nm=300, target_sigma_nm=1, boxcar_width=1, drift_max_bound_factor=2, - save_intermediate_results=False, - interpolation_method='cubic', mode="cuda", min_max_frames=None, - pair_indices_safety_check=False): - """ - Run COMET drift correction end-to-end. - - Pipeline: temporal segmentation -> KD-tree neighbor pairs -> cost optimization (L-BFGS-B) - -> optional temporal smoothing -> spline interpolation to per-frame drift -> (optional) subtract drift. - - Parameters - ---------- - dataset : ndarray of shape (N, 4) - Localization array with columns [x_nm, y_nm, z_nm, frame]. Units in nm; frame is int. - For 2D CSVs, insert a zero z column to get (N, 4). - segmentation_mode : {0, 1, 2} - Temporal segmentation mode: - 0 = number of windows (choose S directly), - 1 = localizations per window (accumulate frames until >= X locs), - 2 = fixed frame window size (default). - segmentation_var : int - Mode-dependent value (S, locs per window, or frames per window). - initial_sigma_nm : float, default=100 - Initial Gaussian length scale for the overlap kernel (coarse scale). - target_sigma_nm : float, default=1 - Target (final) Gaussian length scale for fine refinement. - max_drift_nm : float or None, default=None - Pair radius in nm used for neighbor search. If None, uses 3 * initial_sigma_nm. - drift_max_bound_factor : float, default=1.0 - Multiplicative factor for L-BFGS-B box bounds around +-max_drift. - boxcar_width : int, default=1 - Temporal smoothing width (segments) applied to the estimated drift between optimizer steps. - interpolation_method : {"cubic", "catmull-rom"}, default="cubic" - Spline used to convert per-segment drift to per-frame drift. - max_locs_per_segment : int or None, default=None - Optional downsampling cap per segment (to control memory/time). - mode : str, "cuda" - only Cuda in this version - return_corrected_locs : bool, default=False - If True, also return drift-corrected localizations. - - Returns - ------- - drift_interp_with_frames : ndarray of shape (F, 4) - Per-frame drift with columns [dx_nm, dy_nm, dz_nm, frame]. - corrected_locs : ndarray of shape (N, 4), optional - Only if return_corrected_locs=True. Columns are [x_nm, y_nm, z_nm, segment_id]. - """ - - loc_frames = dataset[:, -1] - if min_max_frames is None: - min_max_frames = (loc_frames.min(), loc_frames.max()) - - # Segment the dataset based on frame numbers into time windows - - result, sorted_dataset, idx_i, idx_j = segmentation_and_pair_indices_wrapper( - dataset, segmentation_var, segmentation_mode, max_drift_nm, max_locs_per_segment, - pair_indices_safety_check=pair_indices_safety_check) - - # Set default initial sigma if not provided - if initial_sigma_nm is None: - initial_sigma_nm = max_drift_nm // 3 - - # Run drift optimization - t0 = time.time() - drift_est = optimize_3d_chunked_better_moving_avg_kd( - result.n_segments, sorted_dataset, - idx_i, idx_j, - sigma_nm=initial_sigma_nm, - target_sigma_nm=target_sigma_nm, - drift_max_nm=max_drift_nm, - drift_max_bound_factor=drift_max_bound_factor, - display_steps=display, - boxcar_width=boxcar_width, - segmentation_result=result, - mode=mode - ) - elapsed = time.time() - t0 - - # Reshape and interpolate drift across all frames - drift_est = drift_est.reshape((result.n_segments, 3)) - vld_tp = np.where(~np.isnan(drift_est[:, 0])) - - frame_interp = np.arange(0, min_max_frames[1] + 1, dtype=int) - drift_interp = interpolate_drift(result.center_frames[vld_tp], drift_est[vld_tp], frame_interp, - method=interpolation_method) - drift_interp_with_frames = np.hstack((drift_interp, frame_interp[:, np.newaxis])) - - # Apply drift correction to original localizations - for i in range(3): - dataset[:, i] = dataset[:, i] - drift_interp[dataset[:, -1].astype(int), i] - - # Optionally show estimated drift curve - if display: - print(f"Drift estimation completed in {elapsed:.2f} seconds.") - plt.figure() - plt.plot(frame_interp, drift_interp) - plt.title("Estimated Drift") - plt.xlabel("Frames") - plt.ylabel("Drift (nm)") - plt.legend(['X', 'Y', 'Z']) - plt.show() - - - # Optional GT comparison plot - if display and gt_drift is not None: - fig, ax = plt.subplots(3, 1) - ax[0].plot(gt_drift[:, 3], gt_drift[:, 0], label='GT Drift X') - ax[1].plot(gt_drift[:, 3], gt_drift[:, 1], label='GT Drift Y') - ax[2].plot(gt_drift[:, 3], gt_drift[:, 2], label='GT Drift Z') - ax[0].plot(frame_interp, drift_interp[:, 0], label='Estimated Drift X', linestyle='--') - ax[1].plot(frame_interp, drift_interp[:, 1], label='Estimated Drift Y', linestyle='--') - ax[2].plot(frame_interp, drift_interp[:, 2], label='Estimated Drift Z', linestyle='--') - ax[1].set_title("Ground Truth vs Estimated Drift") - ax[1].set_xlabel("Frames") - ax[0].set_ylabel("Drift (nm)") - plt.legend() - plt.show() - - # Return corrected locs + drift - if return_corrected_locs: - return drift_interp_with_frames, dataset - else: - return drift_interp_with_frames - - -def optimize_3d_chunked_better_moving_avg_kd(n_segments, locs_nm, idx_i, idx_j, sigma_nm=30, drift_max_nm=300, - target_sigma_nm=30, display_steps=False, - boxcar_width=3, drift_max_bound_factor=2, - segmentation_result=None, - mode="cuda", return_calc_time=False): - """ - Estimate per-segment drift (mu) by minimizing the negative Gaussian-overlap cost - with an L-BFGS-B optimizer and a coarse-to-fine schedule on sigma. - - The routine operates on temporally segmented localizations and reuses a static - neighbor graph (pairs within drift_max_nm). Between optimizer steps, a moving - average (boxcar) can be applied to mu as temporal regularization. Sigma is - reduced iteratively from `sigma_nm` toward `target_sigma_nm` for robust - convergence. - - Parameters - ---------- - n_segments : int - Number of temporal segments S (0..S-1). One 3D drift vector is estimated per segment. - locs_nm : ndarray of shape (M, 3) - localizations in nanometers, columns [x, y, z]; - sigma_nm : float, default=30 - Initial Gaussian width (coarse scale) for the overlap kernel. - drift_max_nm : float, default=300 - Maximum expected drift (nm). Also used as the radius for neighbor pairs and - as the L-BFGS-B bound scale (see `drift_max_bound_factor`). - target_sigma_nm : float, default=30 - Target / final Gaussian width (fine scale). The optimizer reduces sigma toward - this value over iterations. - display_steps : bool, default=False - If True, print or log intermediate progress per iteration/scale. - - boxcar_width : int, default=3 - Temporal smoothing width (in segments) for a moving average applied to mu between steps. - Use 0 or 1 to disable smoothing. - drift_max_bound_factor : float, default=2 - Multiplier for L-BFGS-B bounds around +- drift_max_nm to keep updates physically reasonable. - segmentation_result : object or None, default=None - Segmentation info and metadata. Expected to provide, at minimum: - - segment IDs per localization - - center frames per segment - - any additional structures required by backend (e.g., pair indices) - If None, pairs/ids may be built internally depending on implementation. - mode : str, default=cuda - If True, use the CPU backend; otherwise try GPU (CUDA) and fall back to CPU if unavailable. - return_calc_time : bool, default=False - If True, also return the total computation time in seconds. - - Returns - ------- - mu : ndarray of shape (S, 3) - Estimated per-segment drift (dx, dy, dz) in nanometers. - calc_time_s : float, optional - Only when `return_calc_time=True`. Wall-clock time for the optimization. - """ - - if segmentation_result is None: - segmentation_result = {} - intermediate_results_filehandle = None - sigma_factor = 1.0 - - # Extract coordinate + time arrays, convert to device if CUDA - coords = locs_nm[:, :3].astype(np.float32).copy() - times = locs_nm[:, 3].astype(np.int32).copy() - - chunk_size = int(1E8) # 1E7 - - quality_control = mode == "torch_qc" or mode == "cuda_qc" - - d_coords = cuda.to_device(coords) - d_times = cuda.to_device(times) - if len(idx_i) * 4 > 2e9: - # Use mapped memory if index arrays are large - print("Large index arrays — using mapped memory.") - d_idx_i = cuda.mapped_array_like(idx_i.astype(np.int32), wc=True) - d_idx_j = cuda.mapped_array_like(idx_j.astype(np.int32), wc=True) - d_idx_i[:] = idx_i - d_idx_j[:] = idx_j - else: - d_idx_i = cuda.to_device(idx_i.astype(np.int32)) - d_idx_j = cuda.to_device(idx_j.astype(np.int32)) - # Preallocate device arrays - d_sigma = np.float64(sigma_nm) - d_val = cuda.to_device(np.zeros(chunk_size)) - d_deri = cuda.to_device(np.zeros((n_segments, 3), dtype=np.float64)) - wrapper = cuda_wrapper_chunked - - # Initial drift estimate + bounds - drift_est = np.zeros(n_segments * 3) - bounds = [(-drift_max_nm * drift_max_bound_factor, drift_max_nm * drift_max_bound_factor)] * (3 * n_segments) - - drift_est_gradient = np.inf - fails = 0 - done = False - itr_counter = 0 - start_time = time.time() - - while not done: - # Apply boxcar smoothing to current estimate - tmp = drift_est.reshape((-1, 3)) - for i in range(3): - tmp[:, i] = convolve(tmp[:, i], np.ones(boxcar_width) / boxcar_width) - drift_est = tmp.flatten() - - # Run L-BFGS-B optimization step - result = minimize(wrapper, drift_est, method='L-BFGS-B', - args=( - d_coords, d_times, d_idx_i, d_idx_j, d_sigma, sigma_factor, d_val, d_deri, chunk_size), - jac=True, bounds=bounds, - options={'disp': display_steps, 'gtol': 1E-5, 'ftol': 1E3 * np.finfo(float).eps, - 'maxls': 40}) - itr_counter += 1 - print(f"Iteration {itr_counter}: status = {result.status}, success = {result.success}") - print(f" current sigma: {np.round(sigma_nm * sigma_factor, 2)} nm") - - # Update if successful - if result.success: - delta = np.median((result.x - drift_est) ** 2) - print(f" drift estimate gradient: {delta}") - print(f" previous gradient: {drift_est_gradient}") - # Check convergence - if (delta > drift_est_gradient or sigma_nm * sigma_factor <= 1.0) and sigma_nm * sigma_factor <= target_sigma_nm: - done = True - calc_time = time.time() - start_time - print(f"Optimization completed in {calc_time:.2f} s") - else: - sigma_factor /= 1.5 - drift_est_gradient = delta - drift_est = result.x - else: - fails += 1 - if fails > 2: - sigma_factor *= 2 - print("Restarting with larger sigma_factor") - if fails > 5: - raise RuntimeError("L-BFGS-B Optimization failed after multiple retries") - - if return_calc_time: - return drift_est, time.time() - start_time, itr_counter - else: - return drift_est - - -def segmentation_and_pair_indices_wrapper(dataset, segmentation_var, segmentation_mode, max_drift_nm, - max_locs_per_segment, pair_indices_safety_check=False, hard_limit_pairs=None): - if not segmentation_mode == -1: # -1 is for pre-segmented data - result = segmentation_wrapper(dataset[:, -1], segmentation_var, segmentation_mode, - max_locs_per_segment, return_param_dict=True) - else: - # pre segmented data, anyway we set these values in case auto downsampling is needed - segmentation_mode = 2 # dummy --> segment per frame ... - segmentation_var = 1 # dummy --> ... using 1 frame per segment - if pair_indices_safety_check: - n_pairs_est = estimate_pairs(dataset[result.loc_valid, :3], max_drift_nm) - print(f"Estimated number of pairs within {max_drift_nm} nm: {n_pairs_est:,}") - if hard_limit_pairs is not None and n_pairs_est > hard_limit_pairs: - raise RuntimeError(f"Estimated number of pairs {n_pairs_est} exceeds hard limit of {hard_limit_pairs}. " - f"Aborting to avoid crash.") - if n_pairs_est > 5e8: - print(f"Estimated number of pairs is very large {n_pairs_est}. " - f"Automatic down-sampling is usually required above 500 mil." - f"Billions of pairs can lead to crash.") - ans = input("Continue anyway? (y/n): ") - if ans.lower() != 'y': - raise RuntimeError("Aborted by user due to large estimated number of pairs.") - idx_i, idx_j, successful = pair_indices_kdtree(dataset[result.loc_valid, :3], max_drift_nm) - if not successful: - if max_locs_per_segment is None: - max_locs_per_segment = int(result.out_dict['locs_per_segment'].max()) - while not successful: - max_locs_per_segment = int(max_locs_per_segment * 0.9) - print(f"Segmentation and Pairing attempt failed, automatic down-sampling active...") - print(f"Retrying segmentation with max_locs_per_segment={max_locs_per_segment}...") - result = segmentation_wrapper(dataset[:, -1], segmentation_var, segmentation_mode, - max_locs_per_segment, return_param_dict=True) - sorted_dataset = dataset.copy() - sorted_dataset[:, -1] = result.loc_segments - sorted_dataset = sorted_dataset[result.loc_valid] - idx_i, idx_j, successful = pair_indices_kdtree(sorted_dataset[:, :3], max_drift_nm) - print(f"Segmentation and Pairing successful resulting in {result.n_segments:,} time windows with on average " - f"{int(np.median(result.out_dict['locs_per_segment']))} locs per time window. " - f"{len(idx_i):,} Pairs where found.") - sorted_dataset = dataset.copy() - sorted_dataset[:, -1] = result.loc_segments - sorted_dataset = sorted_dataset[result.loc_valid] - return result, sorted_dataset, idx_i, idx_j - - - -def estimate_pairs(coordinates, distance): - for i in range(len(coordinates[0])): - coordinates[:, i] -= np.min(coordinates[:, i]) - coordinates = np.array(np.floor(coordinates / distance), dtype=int) - coordinates = np.array(list(map(tuple, coordinates))) - sort_indices = np.lexsort(coordinates.T)# get the unique tuples and their counts - unique_tuples, counts = np.unique(coordinates[sort_indices], axis=0, return_counts=True) - # get the indices of the similar tuples - similar_indices = np.split(sort_indices, np.cumsum(counts[:-1])) - idx_i = [] - idx_j = [] - pair_idc_estimate = 0 - for i in range(len(similar_indices)): - n_elements = len(similar_indices[i]) - pair_idc_estimate += n_elements * (n_elements - 1) - rounded = round(pair_idc_estimate,-4) - return rounded - - -def interpolate_drift(center_frames, drift_est, frame_range, method='cubic'): - """ - Interpolates drift estimates to all frames using specified method. - Parameters: - - center_frames: np.ndarray of shape (M,), frames corresponding to drift estimates. - - drift_est: np.ndarray of shape (M, 3), drift estimates at center frames. - - frame_range: array-like, frames to interpolate drift estimates to. - - method: str, interpolation method ('cubic' or 'catmull-rom'). - Returns: - - drift_interp: np.ndarray of shape (len(frame_range), 3), interpolated drift estimates. - """ - if method == 'cubic': - return _interpolate_cubic(center_frames, drift_est, frame_range) - elif method == 'catmull-rom': - return _interpolate_catmull_rom(center_frames, drift_est, frame_range) - elif method == 'linear': - drift_x = np.interp(frame_range, center_frames, drift_est[:, 0]) - drift_y = np.interp(frame_range, center_frames, drift_est[:, 1]) - drift_z = np.interp(frame_range, center_frames, drift_est[:, 2]) - return np.vstack([drift_x, drift_y, drift_z]).T - else: - raise ValueError(f"Unknown interpolation method: {method}") - - -def _interpolate_cubic(center_frames, drift_est, frame_range): - from scipy.interpolate import CubicSpline - drift_x = CubicSpline(center_frames, drift_est[:, 0])(frame_range) - drift_y = CubicSpline(center_frames, drift_est[:, 1])(frame_range) - drift_z = CubicSpline(center_frames, drift_est[:, 2])(frame_range) - return np.vstack([drift_x, drift_y, drift_z]).T - - -def _interpolate_catmull_rom(center_frames, drift_est, frame_range): - def catmull_rom_1d(x, y, x_interp): - result = np.zeros_like(x_interp) - for i in range(1, len(x) - 2): - x0, x1, x2, x3 = x[i-1], x[i], x[i+1], x[i+2] - y0, y1, y2, y3 = y[i-1], y[i], y[i+1], y[i+2] - - mask = (x_interp >= x1) & (x_interp <= x2) - t = (x_interp[mask] - x1) / (x2 - x1) - - result[mask] = ( - 0.5 * ( - (2 * y1) + - (-y0 + y2) * t + - (2*y0 - 5*y1 + 4*y2 - y3) * t**2 + - (-y0 + 3*y1 - 3*y2 + y3) * t**3 - ) - ) - return result - - x_interp = np.asarray(frame_range) - x = np.asarray(center_frames) - - if len(x) < 4: - raise ValueError("Catmull-Rom interpolation requires at least 4 points.") - - drift_x = catmull_rom_1d(x, drift_est[:, 0], x_interp) - drift_y = catmull_rom_1d(x, drift_est[:, 1], x_interp) - drift_z = catmull_rom_1d(x, drift_est[:, 2], x_interp) - - return np.vstack([drift_x, drift_y, drift_z]).T - - -# Numba cuda code of cost-function --> quantized overlap of pairs of localizations, gaussian mixture model -@cuda.jit -def cost_function_full_3d_chunked(d_locs_time, start_idx, chunk_size, d_idx_i, d_idx_j, d_sigma, d_sigma_factor, d_val, - d_val_sum, d_deri, d_locs_coords, mu): - """Compute negative-overlap cost and gradient for 3D localizations (GPU/CPU backend; internal use).""" - tx = cuda.threadIdx.x - ty = cuda.blockIdx.x - bw = cuda.blockDim.x - pos = tx + ty * bw - - if pos < chunk_size: - i = d_idx_i[pos + start_idx] - j = d_idx_j[pos + start_idx] - - ti = d_locs_time[i] - tj = d_locs_time[j] - - dx = (d_locs_coords[i, 0] - mu[ti, 0]) - (d_locs_coords[j, 0] - mu[tj, 0]) - dy = (d_locs_coords[i, 1] - mu[ti, 1]) - (d_locs_coords[j, 1] - mu[tj, 1]) - dz = (d_locs_coords[i, 2] - mu[ti, 2]) - (d_locs_coords[j, 2] - mu[tj, 2]) - sigma_sq = (2 * d_sigma * d_sigma_factor) ** 2 - - diff_sq = dx * dx + dy * dy + dz * dz - val = 1 / (d_sigma * d_sigma_factor) * math.exp(-diff_sq / sigma_sq) - d_val[pos] = val - - # Update derivatives - cuda.atomic.add(d_deri, (tj, 0), 2 * val * (d_locs_coords[j, 0] - d_locs_coords[i, 0] + mu[ti, 0] - mu[tj, 0]) / sigma_sq) - cuda.atomic.add(d_deri, (tj, 1), 2 * val * (d_locs_coords[j, 1] - d_locs_coords[i, 1] + mu[ti, 1] - mu[tj, 1]) / sigma_sq) - cuda.atomic.add(d_deri, (tj, 2), 2 * val * (d_locs_coords[j, 2] - d_locs_coords[i, 2] + mu[ti, 2] - mu[tj, 2]) / sigma_sq) - - cuda.atomic.add(d_deri, (ti, 0), 2 * val * (d_locs_coords[i, 0] - d_locs_coords[j, 0] + mu[tj, 0] - mu[ti, 0]) / sigma_sq) - cuda.atomic.add(d_deri, (ti, 1), 2 * val * (d_locs_coords[i, 1] - d_locs_coords[j, 1] + mu[tj, 1] - mu[ti, 1]) / sigma_sq) - cuda.atomic.add(d_deri, (ti, 2), 2 * val * (d_locs_coords[i, 2] - d_locs_coords[j, 2] + mu[tj, 2] - mu[ti, 2]) / sigma_sq) - - cuda.atomic.add(d_val_sum, 0, val) - d_val[pos] = 0 - - -# Interface between the Python code and the CUDA kernel, mainly for chunking the data to avoid memory issues -def cuda_wrapper_chunked(mu, d_locs_coords, d_locs_time, d_idx_i, d_idx_j, d_sigma, d_sigma_factor, d_val, d_deri, - chunk_size): - val_total = 0 - d_val_sum = cuda.to_device(np.zeros(1, dtype=np.float64)) - mu_dev = cuda.to_device(np.asarray(mu.reshape(int(mu.size / 3), 3), dtype=np.float64)) - - n_chunks = int(np.ceil(d_idx_i.size / chunk_size)) - threadsperblock = 128 - - for i in range(n_chunks - 1): - idc_start = i*chunk_size - blockspergrid = (chunk_size + (threadsperblock - 1)) // threadsperblock - cost_function_full_3d_chunked[blockspergrid, threadsperblock]( - d_locs_time, idc_start, chunk_size, d_idx_i, d_idx_j, - d_sigma, d_sigma_factor, d_val, d_val_sum, d_deri, d_locs_coords, mu_dev - ) - val_total += d_val_sum.copy_to_host() - - # Final chunk - n_remaining = d_idx_i.size - (n_chunks - 1) * chunk_size - idc_start = (n_chunks - 1) * chunk_size - blockspergrid = (n_remaining + (threadsperblock - 1)) // threadsperblock - cost_function_full_3d_chunked[blockspergrid, threadsperblock]( - d_locs_time, idc_start, n_remaining, d_idx_i, d_idx_j, - d_sigma, d_sigma_factor, d_val, d_val_sum, d_deri, d_locs_coords, mu_dev - ) - val_total += d_val_sum.copy_to_host() - deri = d_deri.copy_to_host() - d_deri[:] = 0 - - return -np.nansum(val_total), -deri.flatten() - -from dataclasses import dataclass -from typing import Optional, Dict - - -@dataclass -class SegmentationResult: - """Container for segmentation results.""" - loc_segments: np.ndarray - loc_valid: np.ndarray - center_frames: np.ndarray - n_segments: int - out_dict: Optional[Dict] = None - - -def _group_by_frame(loc_frames: np.ndarray): - """Returns a dict {frame_number: indices_in_loc_frames} efficiently.""" - sort_idx = np.argsort(loc_frames) - sorted_frames = loc_frames[sort_idx] - unique_frames, start_idx, counts = np.unique(sorted_frames, return_index=True, return_counts=True) - frame_to_indices = { - frame: sort_idx[start:start + count] - for frame, start, count in zip(unique_frames, start_idx, counts) - } - return unique_frames, frame_to_indices - - -def segment_by_num_locs_per_window(loc_frames: np.ndarray, min_n_locs_per_window: int, - max_locs_per_segment = None, - return_param_dict: bool = False) -> SegmentationResult: - """ - Segments by collecting a minimum number of localizations per window. - Once the threshold is met and enough locs remain, a new segment is created. - This method ensures that each segment has at least `min_n_locs_per_window` localizations, - while also trying to avoid creating segments that are too small at the end of the dataset. - If `max_locs_per_segment` is set, a random subset of that size is chosen from each segment. - This method is particularly useful for datasets with varying localization densities over time. - Parameters: - loc_frames (np.ndarray): Array of frame numbers for each localization. - min_n_locs_per_window (int): Minimum number of localizations per segment. - max_locs_per_segment (Optional[int]): Maximum number of localizations per segment. If None, all locs are used. - return_param_dict (bool): Whether to return a dictionary of segmentation parameters. - Returns: - SegmentationResult: A dataclass containing segmentation results and parameters. - """ - loc_frames = np.asarray(loc_frames, dtype=int) - n_locs = len(loc_frames) - - if max_locs_per_segment is not None and max_locs_per_segment < 1: # downsampling in percentage - max_locs_per_segment = int(min_n_locs_per_window * max_locs_per_segment) - - unique_frames, frame_to_indices = _group_by_frame(loc_frames) - loc_segments = np.full(n_locs, -1, dtype=int) # Default to -1 for safety - segment_counter = 0 - n_locs_in_current_segment = 0 - current_segment_indices = [] - start_frames, end_frames, locs_per_segment = [], [], [] - - for i, frame in enumerate(unique_frames): - indices = frame_to_indices[frame] - n_locs_this_frame = len(indices) - remaining_locs = n_locs - (len(current_segment_indices) + n_locs_this_frame + np.sum(locs_per_segment)) - - # Add frame to current segment if: - # - It fills the current segment to threshold - # - AND there are enough locs left for another segment (or it's the last frame) - if (n_locs_in_current_segment + n_locs_this_frame >= min_n_locs_per_window) and \ - (remaining_locs >= min_n_locs_per_window or i == len(unique_frames) - 1): - current_segment_indices.extend(indices) - n_locs_in_current_segment += n_locs_this_frame - - loc_segments[current_segment_indices] = segment_counter - start_frames.append(loc_frames[current_segment_indices[0]]) - end_frames.append(loc_frames[current_segment_indices[-1]]) - locs_per_segment.append(len(current_segment_indices)) - - segment_counter += 1 - current_segment_indices = [] - n_locs_in_current_segment = 0 - else: - # Defer frame to current segment - current_segment_indices.extend(indices) - n_locs_in_current_segment += n_locs_this_frame - - n_segments = segment_counter - center_frames = np.zeros(n_segments) - loc_valid = np.zeros(n_locs, dtype=bool) - - for i in range(n_segments): - segment_indices = np.where(loc_segments == i)[0] - if max_locs_per_segment and len(segment_indices) > max_locs_per_segment: - selected = np.random.choice(segment_indices, max_locs_per_segment, replace=False) - else: - selected = segment_indices - loc_valid[selected] = True - locs_per_segment[i] = len(selected) - center_frames[i] = np.mean(loc_frames[selected]) - - out_dict = None - if return_param_dict: - n_locs_valid = loc_valid.sum() - out_dict = { - "n_segments": n_segments, - "min_n_locs_per_window": min_n_locs_per_window, - "frames_per_window": -1, - "start_frames": np.array(start_frames), - "end_frames": np.array(end_frames), - "locs_per_segment": np.array(locs_per_segment), - "n_locs": n_locs, - "n_locs_valid": n_locs_valid, - "n_locs_invalid": n_locs - n_locs_valid, - "center_frames": center_frames - } - - return SegmentationResult(loc_segments, loc_valid, center_frames, n_segments, out_dict) - - -def segment_by_frame_windows(loc_frames: np.ndarray, n_frames_per_window: int, - max_locs_per_segment = None, - return_param_dict: bool = False) -> SegmentationResult: - """ - Splits localization data into fixed-size windows of N frames. - All localizations in those frames are grouped into one segment. - """ - loc_frames = np.asarray(loc_frames, dtype=int) - frames, frame_to_indices = _group_by_frame(loc_frames) - n_locs = len(loc_frames) - n_segments = int(np.ceil(len(frames) / n_frames_per_window)) - - if max_locs_per_segment is not None and max_locs_per_segment < 1: # downsampling in percentage - max_locs_per_segment = n_locs / n_segments * max_locs_per_segment - - loc_segments = np.zeros(n_locs, dtype=int) - center_frames = np.zeros(n_segments) - loc_valid = np.ones(n_locs, dtype=bool) - start_frames, end_frames, locs_per_segment = [], [], [] - - for i in range(n_segments): - frame_window = frames[i * n_frames_per_window:(i + 1) * n_frames_per_window] - indices = np.concatenate([frame_to_indices[frame] for frame in frame_window if frame in frame_to_indices]) - if len(indices) == 0: - continue - loc_segments[indices] = i - start_frames.append(frame_window[0]) - end_frames.append(frame_window[-1]) - center_frames[i] = np.mean(loc_frames[indices]) - locs_per_segment.append(len(indices)) - if max_locs_per_segment and len(indices) > max_locs_per_segment: - mask = np.ones(len(indices), dtype=bool) - mask[np.random.choice(len(indices), len(indices) - max_locs_per_segment, replace=False)] = False - loc_valid[indices[~mask]] = False - - out_dict = None - if return_param_dict: - n_locs_valid = loc_valid.sum() - out_dict = { - "n_segments": n_segments, - "min_n_locs_per_window": -1, - "frames_per_window": n_frames_per_window, - "start_frames": np.array(start_frames), - "end_frames": np.array(end_frames), - "locs_per_segment": np.array(locs_per_segment), - "n_locs": n_locs, - "n_locs_valid": n_locs_valid, - "n_locs_invalid": n_locs - n_locs_valid, - "center_frames": center_frames - } - - return SegmentationResult(loc_segments, loc_valid, center_frames, n_segments, out_dict) - - -def segment_by_num_windows(loc_frames: np.ndarray, n_windows: int, max_locs_per_segment = None, - return_param_dict: bool = False) -> SegmentationResult: - """ - Converts number of windows into an equivalent minimum locs per window, - then calls `segment_by_num_locs_per_window`. - """ - n_locs = len(loc_frames) - n_locs_per_window = int(np.ceil(n_locs / n_windows)) - if max_locs_per_segment is not None and max_locs_per_segment < 1: # downsampling in percentage - max_locs_per_segment = int(n_locs_per_window * max_locs_per_segment) - return segment_by_num_locs_per_window(loc_frames, n_locs_per_window, max_locs_per_segment, return_param_dict) - - -def segmentation_wrapper(loc_frames: np.ndarray, segmentation_var: int, segmentation_mode: int = 2, - max_locs_per_segment = None, - return_param_dict: bool = False) -> SegmentationResult: - """ - Dispatch function that selects segmentation method: - 0 → fixed number of windows - 1 → fixed number of localizations per window - 2 → fixed number of frames per window (default) - """ - if segmentation_mode == 0: - return segment_by_num_windows(loc_frames, segmentation_var, max_locs_per_segment, return_param_dict) - elif segmentation_mode == 1: - return segment_by_num_locs_per_window(loc_frames, segmentation_var, max_locs_per_segment, return_param_dict) - else: - return segment_by_frame_windows(loc_frames, segmentation_var, max_locs_per_segment, return_param_dict) - - -def pair_indices_kdtree(coordinates, distance): - """ - Find all pairs of points within a certain distance using a KD-tree. - Parameters: - - coordinates: np.ndarray of shape (N, D) where N is the number of points and D is the dimensionality. - - distance: float, the maximum distance to consider points as a pair. - Returns: - - idx1: np.ndarray of shape (M,), indices of the first point in each pair. - - idx2: np.ndarray of shape (M,), indices of the second point in each pair. - """ - tree = cKDTree(coordinates) - try: - pairs = tree.query_pairs(r=distance, output_type='ndarray') - except MemoryError: - print("[pair_indices_kdtree] MemoryError encountered") - return [], [], False - return np.ascontiguousarray(pairs[:, 0], dtype=np.int32), np.ascontiguousarray(pairs[:, 1], dtype=np.int32), True \ No newline at end of file diff --git a/picasso/gui/render.py b/picasso/gui/render.py index 4952a21b..1f1bb79b 100644 --- a/picasso/gui/render.py +++ b/picasso/gui/render.py @@ -40,7 +40,6 @@ from PyQt5 import QtCore, QtGui, QtWidgets from .. import ( aim, - comet, clusterer, g5m, imageprocess, @@ -53,8 +52,9 @@ ) from .rotation import RotationWindow -# PyImarisWrite works on windows only +# Optional modules with external/hardware dependencies live in ext/ from ..ext.bitplane import IMSWRITER +from ..ext import comet if IMSWRITER: from ..ext.bitplane import numpy_to_imaris @@ -10685,11 +10685,23 @@ def undrift_comet(self) -> None: params, ok = COMETDialog.getParams(self.window) if ok: - locs, new_info, drift = comet.comet( - locs, - info, - **params, - ) + try: + locs, new_info, drift = comet.comet( + locs, + info, + **params, + ) + except RuntimeError as e: + QtWidgets.QMessageBox.warning( + self.window, + "COMET not available", + ( + f"{e}\n\n" + "Please use a different drift-correction method, " + "such as Undrift by AIM or Undrift by RCC." + ), + ) + return locs = lib.ensure_sanity(locs, info) self.all_locs[channel] = locs @@ -11855,10 +11867,10 @@ def initUI(self, plugins_loaded: bool) -> None: # menu bar - Postprocess postprocess_menu = self.menu_bar.addMenu("Postprocess") undrift_comet_action = postprocess_menu.addAction("Undrift by COMET") - undrift_comet_action.setShortcut("Ctrl+U") undrift_comet_action.triggered.connect(self.view.undrift_comet) undrift_aim_action = postprocess_menu.addAction("Undrift by AIM") + undrift_aim_action.setShortcut("Ctrl+U") undrift_aim_action.triggered.connect(self.view.undrift_aim) undrift_from_picked_action = postprocess_menu.addAction( "Undrift from picked" diff --git a/tests/test_comet.py b/tests/test_comet.py new file mode 100644 index 00000000..976b3399 --- /dev/null +++ b/tests/test_comet.py @@ -0,0 +1,322 @@ +"""Tests for picasso.comet — drift-correction with COMET. + +Tests are split into two groups: +- Pure-Python helpers (segmentation, pairing, interpolation): always run when + numba is installed, no GPU required. +- Full pipeline (comet.comet): requires CUDA hardware; skipped automatically + on CPU-only machines, but the no-GPU error path IS tested on any machine. + +:author: Lenny Reinkensmeier, 2026 +""" + +import numpy as np +import pandas as pd +import pytest + +# The whole picasso package requires numba (lib.py imports it). +# Skip this module cleanly when numba is absent (e.g. CI without GPU stack). +pytest.importorskip("numba") + +from picasso.ext import comet # noqa: E402 — import after guard + + +# --------------------------------------------------------------------------- +# Fixtures +# --------------------------------------------------------------------------- + +@pytest.fixture +def minimal_info(): + """Minimal Picasso info list with required Pixelsize entry.""" + return [{"Pixelsize": 130, "Frames": 200}] + + +@pytest.fixture +def minimal_locs_df(): + """Small synthetic localisation DataFrame spread over 200 frames.""" + rng = np.random.default_rng(0) + n_frames = 200 + n_per_frame = 5 + frames = np.repeat(np.arange(n_frames), n_per_frame) + x = rng.uniform(5, 25, size=len(frames)).astype(np.float32) + y = rng.uniform(5, 25, size=len(frames)).astype(np.float32) + return pd.DataFrame({"frame": frames, "x": x, "y": y}) + + +@pytest.fixture +def minimal_locs_structured(minimal_locs_df): + """Same data as a numpy structured array (Picasso's native format).""" + df = minimal_locs_df + dtype = np.dtype([("frame", np.uint32), ("x", np.float32), ("y", np.float32)]) + arr = np.empty(len(df), dtype=dtype) + arr["frame"] = df["frame"].to_numpy(np.uint32) + arr["x"] = df["x"].to_numpy(np.float32) + arr["y"] = df["y"].to_numpy(np.float32) + return arr + + +# --------------------------------------------------------------------------- +# Segmentation helpers +# --------------------------------------------------------------------------- + +class TestSegmentationByLocs: + def test_produces_correct_segment_count(self): + frames = np.repeat(np.arange(100), 5) # 500 locs, 5/frame + result = comet.segmentation_wrapper( + frames, segmentation_var=50, segmentation_mode=1, return_param_dict=True + ) + assert result.n_segments == 10 + + def test_all_locs_assigned(self): + frames = np.repeat(np.arange(50), 10) + result = comet.segmentation_wrapper( + frames, segmentation_var=100, segmentation_mode=1 + ) + assert (result.loc_segments >= 0).all() + + def test_loc_valid_subset_of_locs(self): + frames = np.repeat(np.arange(60), 8) + result = comet.segmentation_wrapper( + frames, segmentation_var=40, segmentation_mode=1 + ) + assert result.loc_valid.shape == frames.shape + assert result.loc_valid.dtype == bool + + def test_center_frames_length_matches_n_segments(self): + frames = np.repeat(np.arange(80), 5) + result = comet.segmentation_wrapper( + frames, segmentation_var=50, segmentation_mode=1 + ) + assert len(result.center_frames) == result.n_segments + + def test_max_locs_per_segment_respected(self): + frames = np.repeat(np.arange(20), 100) # 2000 locs, 100/frame + cap = 30 + result = comet.segmentation_wrapper( + frames, segmentation_var=200, segmentation_mode=1, + max_locs_per_segment=cap, return_param_dict=True + ) + assert result.loc_valid.sum() <= result.n_segments * cap + + +class TestSegmentationByFrameWindows: + def test_frame_window_mode(self): + frames = np.repeat(np.arange(100), 3) + result = comet.segmentation_wrapper( + frames, segmentation_var=10, segmentation_mode=2 + ) + assert result.n_segments == 10 # 100 frames / 10 per window + + def test_single_window(self): + frames = np.arange(50) + result = comet.segmentation_wrapper( + frames, segmentation_var=100, segmentation_mode=2 + ) + assert result.n_segments == 1 + + +class TestSegmentationByNumWindows: + def test_num_windows_mode_creates_requested_count(self): + frames = np.repeat(np.arange(100), 5) # 500 locs + result = comet.segmentation_wrapper( + frames, segmentation_var=5, segmentation_mode=0, return_param_dict=True + ) + # mode 0 derives min-locs-per-window from total/n_windows; the resulting + # segment count should be approximately n_windows (off-by-one tolerated + # because the last segment is folded into the previous if too small). + assert abs(result.n_segments - 5) <= 1 + + def test_all_locs_assigned(self): + frames = np.repeat(np.arange(40), 4) + result = comet.segmentation_wrapper( + frames, segmentation_var=4, segmentation_mode=0 + ) + assert (result.loc_segments >= 0).all() + + +# --------------------------------------------------------------------------- +# Pair-finding +# --------------------------------------------------------------------------- + +class TestPairIndicesKdtree: + def test_finds_close_pairs(self): + coords = np.array([[0, 0, 0], [1, 0, 0], [100, 0, 0]], dtype=np.float32) + idx_i, idx_j, ok = comet.pair_indices_kdtree(coords, distance=5) + assert ok + assert len(idx_i) == 1 # only (0,1) are within distance 5 + + def test_no_pairs_when_far_apart(self): + coords = np.eye(3, dtype=np.float32) * 1000 + idx_i, idx_j, ok = comet.pair_indices_kdtree(coords, distance=1) + assert ok + assert len(idx_i) == 0 + + def test_all_pairs_when_all_close(self): + coords = np.zeros((5, 3), dtype=np.float32) + idx_i, idx_j, ok = comet.pair_indices_kdtree(coords, distance=1) + assert ok + assert len(idx_i) == 5 * 4 // 2 # C(5,2) + + def test_returns_int32_arrays(self): + coords = np.random.rand(20, 3).astype(np.float32) + idx_i, idx_j, ok = comet.pair_indices_kdtree(coords, distance=0.5) + assert ok + assert idx_i.dtype == np.int32 + assert idx_j.dtype == np.int32 + + +# --------------------------------------------------------------------------- +# Drift interpolation +# --------------------------------------------------------------------------- + +class TestInterpolateDrift: + def _center_frames_and_drift(self): + center_frames = np.linspace(10, 90, 6) + drift_est = np.column_stack([ + np.sin(center_frames / 10), + np.cos(center_frames / 10), + np.zeros_like(center_frames), + ]) + return center_frames, drift_est + + def test_cubic_output_shape(self): + cf, de = self._center_frames_and_drift() + frame_range = np.arange(0, 100) + out = comet.interpolate_drift(cf, de, frame_range, method="cubic") + assert out.shape == (100, 3) + + def test_linear_output_shape(self): + cf, de = self._center_frames_and_drift() + frame_range = np.arange(0, 100) + out = comet.interpolate_drift(cf, de, frame_range, method="linear") + assert out.shape == (100, 3) + + def test_cubic_passes_through_control_points(self): + cf = np.array([0.0, 25.0, 50.0, 75.0, 100.0]) + de = np.column_stack([cf * 0.01, cf * 0.005, np.zeros_like(cf)]) + out = comet.interpolate_drift(cf, de, cf.astype(int), method="cubic") + np.testing.assert_allclose(out[:, 0], de[:, 0], atol=1e-6) + + def test_unknown_method_raises(self): + cf, de = self._center_frames_and_drift() + with pytest.raises(ValueError, match="Unknown interpolation method"): + comet.interpolate_drift(cf, de, np.arange(100), method="spagetti") + + def test_catmull_rom_output_shape(self): + cf, de = self._center_frames_and_drift() + frame_range = np.arange(10, 90) + out = comet.interpolate_drift(cf, de, frame_range, method="catmull-rom") + assert out.shape == (len(frame_range), 3) + + def test_catmull_rom_requires_min_points(self): + cf = np.array([0.0, 10.0, 20.0]) # only 3 points + de = np.zeros((3, 3)) + with pytest.raises(ValueError, match="at least 4 points"): + comet.interpolate_drift(cf, de, np.arange(20), method="catmull-rom") + + +# --------------------------------------------------------------------------- +# comet() public API — no-GPU error path +# --------------------------------------------------------------------------- + +class TestCometNoCuda: + """These tests run on any machine; they verify the failure mode when + CUDA is not present.""" + + def test_raises_runtime_error_without_cuda(self, minimal_locs_df, minimal_info): + if comet._CUDA_AVAILABLE: + pytest.skip("CUDA is present; testing no-GPU path only on CPU machines") + with pytest.raises(RuntimeError, match="CUDA"): + comet.comet(minimal_locs_df, minimal_info, locs_per_segment=100, max_drift_nm=200) + + def test_error_message_mentions_alternatives(self, minimal_locs_df, minimal_info): + if comet._CUDA_AVAILABLE: + pytest.skip("CUDA is present; testing no-GPU path only on CPU machines") + with pytest.raises(RuntimeError, match="AIM|RCC"): + comet.comet(minimal_locs_df, minimal_info, locs_per_segment=100, max_drift_nm=200) + + def test_structured_array_input_accepted_before_gpu_check(self, minimal_locs_structured, minimal_info): + """Numpy structured arrays should be converted to DataFrame before the + CUDA check, so the error is still a RuntimeError (not an AttributeError + from missing .columns).""" + if comet._CUDA_AVAILABLE: + pytest.skip("CUDA is present; testing no-GPU path only on CPU machines") + with pytest.raises(RuntimeError, match="CUDA"): + comet.comet(minimal_locs_structured, minimal_info, locs_per_segment=100, max_drift_nm=200) + + def test_missing_columns_raises_key_error(self, minimal_info): + """Locs missing required columns should raise KeyError, not crash on .columns.""" + if comet._CUDA_AVAILABLE: + pytest.skip("CUDA is present; GPU path reached before column check") + bad_locs = pd.DataFrame({"frame": np.arange(10), "x": np.zeros(10)}) + # No 'y' column — should raise KeyError from the column check + with pytest.raises((KeyError, RuntimeError)): + comet.comet(bad_locs, minimal_info, locs_per_segment=5, max_drift_nm=100) + + def test_z_column_accepted(self, minimal_locs_df, minimal_info): + """3D input (with z) should not crash before the GPU check.""" + if comet._CUDA_AVAILABLE: + pytest.skip("CUDA is present; testing no-GPU path only on CPU machines") + df = minimal_locs_df.copy() + df["z"] = np.zeros(len(df), dtype=np.float32) + with pytest.raises(RuntimeError, match="CUDA"): + comet.comet(df, minimal_info, locs_per_segment=100, max_drift_nm=200) + + def test_missing_pixelsize_raises_on_gpu(self, minimal_locs_df): + """info without 'Pixelsize' should fail with an informative error. + Only meaningful on GPU since the CUDA check fires first on CPU.""" + if not comet._CUDA_AVAILABLE: + pytest.skip("CUDA check fires before pixelsize check on CPU") + bad_info = [{"Frames": 200}] + with pytest.raises((KeyError, RuntimeError, ValueError)): + comet.comet(minimal_locs_df, bad_info, locs_per_segment=100, max_drift_nm=200) + + +# --------------------------------------------------------------------------- +# comet() public API — full pipeline (GPU required) +# --------------------------------------------------------------------------- + +@pytest.mark.skipif( + not comet._CUDA_AVAILABLE, + reason="CUDA GPU not available", +) +class TestCometFullPipeline: + """Integration tests that run the full optimisation loop on GPU.""" + + def test_output_shapes(self, minimal_locs_df, minimal_info): + undrifted, new_info, drift = comet.comet( + minimal_locs_df, minimal_info, + locs_per_segment=100, max_drift_nm=300, + target_sigma_nm=50, + ) + assert len(undrifted) == len(minimal_locs_df) + assert set(undrifted.columns) >= {"x", "y", "frame"} + assert "x" in drift.columns and "y" in drift.columns + + def test_drift_covers_frame_range(self, minimal_locs_df, minimal_info): + frame0 = minimal_locs_df["frame"] - minimal_locs_df["frame"].min() + _, _, drift = comet.comet( + minimal_locs_df, minimal_info, + locs_per_segment=100, max_drift_nm=300, + target_sigma_nm=50, + ) + assert len(drift) == int(frame0.max()) + 1 + + def test_new_info_appended(self, minimal_locs_df, minimal_info): + _, new_info, _ = comet.comet( + minimal_locs_df, minimal_info, + locs_per_segment=100, max_drift_nm=300, + target_sigma_nm=50, + ) + assert len(new_info) == len(minimal_info) + 1 + assert "COMET" in new_info[-1]["Generated by"] + + def test_structured_array_gives_same_result_as_dataframe(self, minimal_locs_df, minimal_locs_structured, minimal_info): + locs_df, _, drift_df = comet.comet( + minimal_locs_df, minimal_info, + locs_per_segment=100, max_drift_nm=300, target_sigma_nm=50, + ) + locs_arr, _, drift_arr = comet.comet( + minimal_locs_structured, minimal_info, + locs_per_segment=100, max_drift_nm=300, target_sigma_nm=50, + ) + np.testing.assert_allclose(drift_df["x"].to_numpy(), drift_arr["x"].to_numpy(), rtol=1e-5) From a66d513e2cb4ca5c0100cee2259ebe9dcc45c4d3 Mon Sep 17 00:00:00 2001 From: Lenny Reinkensmeier Date: Thu, 28 May 2026 17:19:34 +0200 Subject: [PATCH 04/19] Handle COMET availability in drift correction and update menu shortcuts. Added proper testing for COMET. --- picasso/ext/comet.py | 870 +++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 870 insertions(+) create mode 100644 picasso/ext/comet.py diff --git a/picasso/ext/comet.py b/picasso/ext/comet.py new file mode 100644 index 00000000..a014354e --- /dev/null +++ b/picasso/ext/comet.py @@ -0,0 +1,870 @@ +import math +import time +from dataclasses import dataclass +from typing import Dict, Optional + +import matplotlib.pyplot as plt +import numpy as np +import pandas as pd +from scipy.ndimage import convolve +from scipy.optimize import minimize +from scipy.spatial import cKDTree + +from . import lib, __version__ + +try: + from numba import cuda as _cuda + _CUDA_AVAILABLE = _cuda.is_available() +except Exception: + _cuda = None + _CUDA_AVAILABLE = False + + + + +def comet( + locs: pd.DataFrame, + info: list[dict], + locs_per_segment: int, + max_drift_nm: float, + max_locs_per_segment: int | None = None, + *, + initial_sigma_nm: float | None = None, + target_sigma_nm: float = 10.0, + boxcar_width: int = 3, + drift_max_bound_factor: float = 2.0, + interpolation_method: str = "cubic", + mode: str = "cuda", + display: bool = False, + progress=None, # currently unused +) -> tuple[pd.DataFrame, list[dict], pd.DataFrame]: + """Apply COMET undrifting to the localizations. + + Parameters + ---------- + locs : pd.DataFrame + Localization table. Must contain columns 'x', 'y', 'frame'. + Optional: 'z'. + x/y are expected in camera pixels, z in the native z-unit used by Picasso. + info : list of dict + Localization metadata. + locs_per_segment : int + Target number of localizations per segment for COMET segmentation. + max_drift_nm : float + Maximum expected drift in nm. + max_locs_per_segment : int or None, optional + Optional cap for localizations per segment. If negative, treated as None. + initial_sigma_nm : float or None, optional + Initial Gaussian sigma for COMET. If None, defaults to max_drift_nm / 3. + target_sigma_nm : float, optional + Final Gaussian sigma for refinement. + boxcar_width : int, optional + Smoothing width over segments. + drift_max_bound_factor : float, optional + Bound factor for optimizer. + interpolation_method : str, optional + Drift interpolation method. + mode : str, optional + Backend mode. Currently only ``"cuda"`` is supported, which requires + a CUDA-capable NVIDIA GPU and numba installed with CUDA support. + Raises ``RuntimeError`` when CUDA is not available. + display : bool, optional + Whether to show diagnostic plots. + progress : optional + Placeholder for API compatibility with AIM. + + Returns + ------- + locs : pd.DataFrame + Undrift-corrected localizations. + new_info : list[dict] + Updated metadata. + drift : pd.DataFrame + Per-frame drift table with columns x, y, and optionally z. + Drift is returned in the same units as locs: + x/y in pixels, z in the original z-unit. + """ + # Accept numpy structured arrays as well as pandas DataFrames. + if isinstance(locs, np.ndarray) and locs.dtype.names is not None: + locs = pd.DataFrame({name: locs[name] for name in locs.dtype.names}) + + locs = locs.copy() + + if not _CUDA_AVAILABLE: + raise RuntimeError( + "COMET requires a CUDA-capable NVIDIA GPU and numba with CUDA support, " + "which are not available on this system. " + "Please use a different drift-correction method (e.g. AIM or RCC)." + ) + + if max_locs_per_segment is not None and max_locs_per_segment < 0: + max_locs_per_segment = None + + pixelsize_nm = lib.get_from_metadata(info, "Pixelsize", raise_error=True) + + required_cols = {"x", "y", "frame"} + missing = required_cols - set(locs.columns) + if missing: + raise KeyError(f"Missing required columns for COMET: {sorted(missing)}") + + has_z = "z" in locs.columns + + # Normalize frames so they start at 0 for internal indexing. + # This is important because COMET later indexes drift by frame number. + frame = locs["frame"].to_numpy(dtype=np.int64) + frame0 = frame - frame.min() + + # Build COMET input array in nm + dataset = np.zeros((len(locs), 4), dtype=np.float64) + dataset[:, 0] = locs["x"].to_numpy(dtype=np.float64) * pixelsize_nm + dataset[:, 1] = locs["y"].to_numpy(dtype=np.float64) * pixelsize_nm + dataset[:, 2] = ( + locs["z"].to_numpy(dtype=np.float64) if has_z else 0.0 + ) + dataset[:, 3] = frame0 + + if initial_sigma_nm is None: + initial_sigma_nm = max_drift_nm / 3 + + drift_nm_with_frame = comet_run_kd( + dataset=dataset, + segmentation_mode=1, # fixed locs per segment + segmentation_var=locs_per_segment, + max_locs_per_segment=max_locs_per_segment, + initial_sigma_nm=initial_sigma_nm, + gt_drift=None, + display=display, + return_corrected_locs=False, + max_drift_nm=max_drift_nm, + target_sigma_nm=target_sigma_nm, + boxcar_width=boxcar_width, + drift_max_bound_factor=drift_max_bound_factor, + interpolation_method=interpolation_method, + mode=mode, + min_max_frames=(int(frame0.min()), int(frame0.max())), + ) + + # drift_nm_with_frame columns: dx_nm, dy_nm, dz_nm, frame0 + drift_nm = drift_nm_with_frame[:, :3] + + # Apply drift back to dataframe in original units + frame_idx = frame0.astype(np.int64) + locs["x"] = locs["x"].to_numpy(dtype=np.float64) - drift_nm[frame_idx, 0] / pixelsize_nm + locs["y"] = locs["y"].to_numpy(dtype=np.float64) - drift_nm[frame_idx, 1] / pixelsize_nm + if has_z: + locs["z"] = locs["z"].to_numpy(dtype=np.float64) - drift_nm[frame_idx, 2] + + # Build drift dataframe in Picasso-style units + drift_dict = { + "x": (drift_nm[:, 0] / pixelsize_nm).astype("float32"), + "y": (drift_nm[:, 1] / pixelsize_nm).astype("float32"), + } + if has_z: + drift_dict["z"] = drift_nm[:, 2].astype("float32") + + drift = pd.DataFrame(drift_dict) + + new_info_entry = { + "Generated by": f"Picasso v{__version__} COMET", + "Segmentation mode": "localizations per segment", + "Localizations per segment": locs_per_segment, + "Maximum drift (nm)": max_drift_nm, + "Initial sigma (nm)": float(initial_sigma_nm), + "Target sigma (nm)": float(target_sigma_nm), + "Boxcar width": int(boxcar_width), + "Max localizations per segment": ( + None if max_locs_per_segment is None else int(max_locs_per_segment) + ), + "Interpolation method": interpolation_method, + "Backend": mode, + } + new_info = info + [new_info_entry] + + return locs, new_info, drift + +def comet_run_kd(dataset, segmentation_mode, segmentation_var, max_locs_per_segment=None, + initial_sigma_nm=None, gt_drift=None, display=False, return_corrected_locs=False, + max_drift_nm=300, target_sigma_nm=1, boxcar_width=1, drift_max_bound_factor=2, + save_intermediate_results=False, + interpolation_method='cubic', mode="cuda", min_max_frames=None, + pair_indices_safety_check=False): + """ + Run COMET drift correction end-to-end. + + Pipeline: temporal segmentation -> KD-tree neighbor pairs -> cost optimization (L-BFGS-B) + -> optional temporal smoothing -> spline interpolation to per-frame drift -> (optional) subtract drift. + + Parameters + ---------- + dataset : ndarray of shape (N, 4) + Localization array with columns [x_nm, y_nm, z_nm, frame]. Units in nm; frame is int. + For 2D CSVs, insert a zero z column to get (N, 4). + segmentation_mode : {0, 1, 2} + Temporal segmentation mode: + 0 = number of windows (choose S directly), + 1 = localizations per window (accumulate frames until >= X locs), + 2 = fixed frame window size (default). + segmentation_var : int + Mode-dependent value (S, locs per window, or frames per window). + initial_sigma_nm : float, default=100 + Initial Gaussian length scale for the overlap kernel (coarse scale). + target_sigma_nm : float, default=1 + Target (final) Gaussian length scale for fine refinement. + max_drift_nm : float or None, default=None + Pair radius in nm used for neighbor search. If None, uses 3 * initial_sigma_nm. + drift_max_bound_factor : float, default=1.0 + Multiplicative factor for L-BFGS-B box bounds around +-max_drift. + boxcar_width : int, default=1 + Temporal smoothing width (segments) applied to the estimated drift between optimizer steps. + interpolation_method : {"cubic", "catmull-rom"}, default="cubic" + Spline used to convert per-segment drift to per-frame drift. + max_locs_per_segment : int or None, default=None + Optional downsampling cap per segment (to control memory/time). + mode : str, "cuda" + only Cuda in this version + return_corrected_locs : bool, default=False + If True, also return drift-corrected localizations. + + Returns + ------- + drift_interp_with_frames : ndarray of shape (F, 4) + Per-frame drift with columns [dx_nm, dy_nm, dz_nm, frame]. + corrected_locs : ndarray of shape (N, 4), optional + Only if return_corrected_locs=True. Columns are [x_nm, y_nm, z_nm, segment_id]. + """ + + loc_frames = dataset[:, -1] + if min_max_frames is None: + min_max_frames = (loc_frames.min(), loc_frames.max()) + + # Segment the dataset based on frame numbers into time windows + + result, sorted_dataset, idx_i, idx_j = segmentation_and_pair_indices_wrapper( + dataset, segmentation_var, segmentation_mode, max_drift_nm, max_locs_per_segment, + pair_indices_safety_check=pair_indices_safety_check) + + # Set default initial sigma if not provided + if initial_sigma_nm is None: + initial_sigma_nm = max_drift_nm // 3 + + # Run drift optimization + t0 = time.time() + drift_est = optimize_3d_chunked_better_moving_avg_kd( + result.n_segments, sorted_dataset, + idx_i, idx_j, + sigma_nm=initial_sigma_nm, + target_sigma_nm=target_sigma_nm, + drift_max_nm=max_drift_nm, + drift_max_bound_factor=drift_max_bound_factor, + display_steps=display, + boxcar_width=boxcar_width, + segmentation_result=result, + mode=mode + ) + elapsed = time.time() - t0 + + # Reshape and interpolate drift across all frames + drift_est = drift_est.reshape((result.n_segments, 3)) + vld_tp = np.where(~np.isnan(drift_est[:, 0])) + + frame_interp = np.arange(0, min_max_frames[1] + 1, dtype=int) + drift_interp = interpolate_drift(result.center_frames[vld_tp], drift_est[vld_tp], frame_interp, + method=interpolation_method) + drift_interp_with_frames = np.hstack((drift_interp, frame_interp[:, np.newaxis])) + + # Apply drift correction to original localizations + for i in range(3): + dataset[:, i] = dataset[:, i] - drift_interp[dataset[:, -1].astype(int), i] + + # Optionally show estimated drift curve + if display: + print(f"Drift estimation completed in {elapsed:.2f} seconds.") + plt.figure() + plt.plot(frame_interp, drift_interp) + plt.title("Estimated Drift") + plt.xlabel("Frames") + plt.ylabel("Drift (nm)") + plt.legend(['X', 'Y', 'Z']) + plt.show() + + + # Optional GT comparison plot + if display and gt_drift is not None: + fig, ax = plt.subplots(3, 1) + ax[0].plot(gt_drift[:, 3], gt_drift[:, 0], label='GT Drift X') + ax[1].plot(gt_drift[:, 3], gt_drift[:, 1], label='GT Drift Y') + ax[2].plot(gt_drift[:, 3], gt_drift[:, 2], label='GT Drift Z') + ax[0].plot(frame_interp, drift_interp[:, 0], label='Estimated Drift X', linestyle='--') + ax[1].plot(frame_interp, drift_interp[:, 1], label='Estimated Drift Y', linestyle='--') + ax[2].plot(frame_interp, drift_interp[:, 2], label='Estimated Drift Z', linestyle='--') + ax[1].set_title("Ground Truth vs Estimated Drift") + ax[1].set_xlabel("Frames") + ax[0].set_ylabel("Drift (nm)") + plt.legend() + plt.show() + + # Return corrected locs + drift + if return_corrected_locs: + return drift_interp_with_frames, dataset + else: + return drift_interp_with_frames + + +def optimize_3d_chunked_better_moving_avg_kd(n_segments, locs_nm, idx_i, idx_j, sigma_nm=30, drift_max_nm=300, + target_sigma_nm=30, display_steps=False, + boxcar_width=3, drift_max_bound_factor=2, + segmentation_result=None, + mode="cuda", return_calc_time=False): + """ + Estimate per-segment drift (mu) by minimizing the negative Gaussian-overlap cost + with an L-BFGS-B optimizer and a coarse-to-fine schedule on sigma. + + The routine operates on temporally segmented localizations and reuses a static + neighbor graph (pairs within drift_max_nm). Between optimizer steps, a moving + average (boxcar) can be applied to mu as temporal regularization. Sigma is + reduced iteratively from `sigma_nm` toward `target_sigma_nm` for robust + convergence. + + Parameters + ---------- + n_segments : int + Number of temporal segments S (0..S-1). One 3D drift vector is estimated per segment. + locs_nm : ndarray of shape (M, 3) + localizations in nanometers, columns [x, y, z]; + sigma_nm : float, default=30 + Initial Gaussian width (coarse scale) for the overlap kernel. + drift_max_nm : float, default=300 + Maximum expected drift (nm). Also used as the radius for neighbor pairs and + as the L-BFGS-B bound scale (see `drift_max_bound_factor`). + target_sigma_nm : float, default=30 + Target / final Gaussian width (fine scale). The optimizer reduces sigma toward + this value over iterations. + display_steps : bool, default=False + If True, print or log intermediate progress per iteration/scale. + + boxcar_width : int, default=3 + Temporal smoothing width (in segments) for a moving average applied to mu between steps. + Use 0 or 1 to disable smoothing. + drift_max_bound_factor : float, default=2 + Multiplier for L-BFGS-B bounds around +- drift_max_nm to keep updates physically reasonable. + segmentation_result : object or None, default=None + Segmentation info and metadata. Expected to provide, at minimum: + - segment IDs per localization + - center frames per segment + - any additional structures required by backend (e.g., pair indices) + If None, pairs/ids may be built internally depending on implementation. + mode : str, default=cuda + If True, use the CPU backend; otherwise try GPU (CUDA) and fall back to CPU if unavailable. + return_calc_time : bool, default=False + If True, also return the total computation time in seconds. + + Returns + ------- + mu : ndarray of shape (S, 3) + Estimated per-segment drift (dx, dy, dz) in nanometers. + calc_time_s : float, optional + Only when `return_calc_time=True`. Wall-clock time for the optimization. + """ + + if segmentation_result is None: + segmentation_result = {} + intermediate_results_filehandle = None + sigma_factor = 1.0 + + # Extract coordinate + time arrays, convert to device if CUDA + coords = locs_nm[:, :3].astype(np.float32).copy() + times = locs_nm[:, 3].astype(np.int32).copy() + + chunk_size = int(1E8) # 1E7 + + quality_control = mode == "torch_qc" or mode == "cuda_qc" + + d_coords = _cuda.to_device(coords) + d_times = _cuda.to_device(times) + if len(idx_i) * 4 > 2e9: + # Use mapped memory if index arrays are large + print("Large index arrays — using mapped memory.") + d_idx_i = _cuda.mapped_array_like(idx_i.astype(np.int32), wc=True) + d_idx_j = _cuda.mapped_array_like(idx_j.astype(np.int32), wc=True) + d_idx_i[:] = idx_i + d_idx_j[:] = idx_j + else: + d_idx_i = _cuda.to_device(idx_i.astype(np.int32)) + d_idx_j = _cuda.to_device(idx_j.astype(np.int32)) + # Preallocate device arrays + d_sigma = np.float64(sigma_nm) + d_val = _cuda.to_device(np.zeros(chunk_size)) + d_deri = _cuda.to_device(np.zeros((n_segments, 3), dtype=np.float64)) + wrapper = cuda_wrapper_chunked + + # Initial drift estimate + bounds + drift_est = np.zeros(n_segments * 3) + bounds = [(-drift_max_nm * drift_max_bound_factor, drift_max_nm * drift_max_bound_factor)] * (3 * n_segments) + + drift_est_gradient = np.inf + fails = 0 + done = False + itr_counter = 0 + start_time = time.time() + + while not done: + # Apply boxcar smoothing to current estimate + tmp = drift_est.reshape((-1, 3)) + for i in range(3): + tmp[:, i] = convolve(tmp[:, i], np.ones(boxcar_width) / boxcar_width) + drift_est = tmp.flatten() + + # Run L-BFGS-B optimization step + result = minimize(wrapper, drift_est, method='L-BFGS-B', + args=( + d_coords, d_times, d_idx_i, d_idx_j, d_sigma, sigma_factor, d_val, d_deri, chunk_size), + jac=True, bounds=bounds, + options={'disp': display_steps, 'gtol': 1E-5, 'ftol': 1E3 * np.finfo(float).eps, + 'maxls': 40}) + itr_counter += 1 + print(f"Iteration {itr_counter}: status = {result.status}, success = {result.success}") + print(f" current sigma: {np.round(sigma_nm * sigma_factor, 2)} nm") + + # Update if successful + if result.success: + delta = np.median((result.x - drift_est) ** 2) + print(f" drift estimate gradient: {delta}") + print(f" previous gradient: {drift_est_gradient}") + # Check convergence + if (delta > drift_est_gradient or sigma_nm * sigma_factor <= 1.0) and sigma_nm * sigma_factor <= target_sigma_nm: + done = True + calc_time = time.time() - start_time + print(f"Optimization completed in {calc_time:.2f} s") + else: + sigma_factor /= 1.5 + drift_est_gradient = delta + drift_est = result.x + else: + fails += 1 + if fails > 2: + sigma_factor *= 2 + print("Restarting with larger sigma_factor") + if fails > 5: + raise RuntimeError("L-BFGS-B Optimization failed after multiple retries") + + if return_calc_time: + return drift_est, time.time() - start_time, itr_counter + else: + return drift_est + + +def segmentation_and_pair_indices_wrapper(dataset, segmentation_var, segmentation_mode, max_drift_nm, + max_locs_per_segment, pair_indices_safety_check=False, hard_limit_pairs=None): + if not segmentation_mode == -1: # -1 is for pre-segmented data + result = segmentation_wrapper(dataset[:, -1], segmentation_var, segmentation_mode, + max_locs_per_segment, return_param_dict=True) + else: + # pre segmented data, anyway we set these values in case auto downsampling is needed + segmentation_mode = 2 # dummy --> segment per frame ... + segmentation_var = 1 # dummy --> ... using 1 frame per segment + if pair_indices_safety_check: + n_pairs_est = estimate_pairs(dataset[result.loc_valid, :3], max_drift_nm) + print(f"Estimated number of pairs within {max_drift_nm} nm: {n_pairs_est:,}") + if hard_limit_pairs is not None and n_pairs_est > hard_limit_pairs: + raise RuntimeError(f"Estimated number of pairs {n_pairs_est} exceeds hard limit of {hard_limit_pairs}. " + f"Aborting to avoid crash.") + if n_pairs_est > 5e8: + print(f"Estimated number of pairs is very large {n_pairs_est}. " + f"Automatic down-sampling is usually required above 500 mil." + f"Billions of pairs can lead to crash.") + ans = input("Continue anyway? (y/n): ") + if ans.lower() != 'y': + raise RuntimeError("Aborted by user due to large estimated number of pairs.") + idx_i, idx_j, successful = pair_indices_kdtree(dataset[result.loc_valid, :3], max_drift_nm) + if not successful: + if max_locs_per_segment is None: + max_locs_per_segment = int(result.out_dict['locs_per_segment'].max()) + while not successful: + max_locs_per_segment = int(max_locs_per_segment * 0.9) + print(f"Segmentation and Pairing attempt failed, automatic down-sampling active...") + print(f"Retrying segmentation with max_locs_per_segment={max_locs_per_segment}...") + result = segmentation_wrapper(dataset[:, -1], segmentation_var, segmentation_mode, + max_locs_per_segment, return_param_dict=True) + sorted_dataset = dataset.copy() + sorted_dataset[:, -1] = result.loc_segments + sorted_dataset = sorted_dataset[result.loc_valid] + idx_i, idx_j, successful = pair_indices_kdtree(sorted_dataset[:, :3], max_drift_nm) + print(f"Segmentation and Pairing successful resulting in {result.n_segments:,} time windows with on average " + f"{int(np.median(result.out_dict['locs_per_segment']))} locs per time window. " + f"{len(idx_i):,} Pairs where found.") + sorted_dataset = dataset.copy() + sorted_dataset[:, -1] = result.loc_segments + sorted_dataset = sorted_dataset[result.loc_valid] + return result, sorted_dataset, idx_i, idx_j + + + +def estimate_pairs(coordinates, distance): + for i in range(len(coordinates[0])): + coordinates[:, i] -= np.min(coordinates[:, i]) + coordinates = np.array(np.floor(coordinates / distance), dtype=int) + coordinates = np.array(list(map(tuple, coordinates))) + sort_indices = np.lexsort(coordinates.T)# get the unique tuples and their counts + unique_tuples, counts = np.unique(coordinates[sort_indices], axis=0, return_counts=True) + # get the indices of the similar tuples + similar_indices = np.split(sort_indices, np.cumsum(counts[:-1])) + idx_i = [] + idx_j = [] + pair_idc_estimate = 0 + for i in range(len(similar_indices)): + n_elements = len(similar_indices[i]) + pair_idc_estimate += n_elements * (n_elements - 1) + rounded = round(pair_idc_estimate,-4) + return rounded + + +def interpolate_drift(center_frames, drift_est, frame_range, method='cubic'): + """ + Interpolates drift estimates to all frames using specified method. + Parameters: + - center_frames: np.ndarray of shape (M,), frames corresponding to drift estimates. + - drift_est: np.ndarray of shape (M, 3), drift estimates at center frames. + - frame_range: array-like, frames to interpolate drift estimates to. + - method: str, interpolation method ('cubic' or 'catmull-rom'). + Returns: + - drift_interp: np.ndarray of shape (len(frame_range), 3), interpolated drift estimates. + """ + if method == 'cubic': + return _interpolate_cubic(center_frames, drift_est, frame_range) + elif method == 'catmull-rom': + return _interpolate_catmull_rom(center_frames, drift_est, frame_range) + elif method == 'linear': + drift_x = np.interp(frame_range, center_frames, drift_est[:, 0]) + drift_y = np.interp(frame_range, center_frames, drift_est[:, 1]) + drift_z = np.interp(frame_range, center_frames, drift_est[:, 2]) + return np.vstack([drift_x, drift_y, drift_z]).T + else: + raise ValueError(f"Unknown interpolation method: {method}") + + +def _interpolate_cubic(center_frames, drift_est, frame_range): + from scipy.interpolate import CubicSpline + drift_x = CubicSpline(center_frames, drift_est[:, 0])(frame_range) + drift_y = CubicSpline(center_frames, drift_est[:, 1])(frame_range) + drift_z = CubicSpline(center_frames, drift_est[:, 2])(frame_range) + return np.vstack([drift_x, drift_y, drift_z]).T + + +def _interpolate_catmull_rom(center_frames, drift_est, frame_range): + def catmull_rom_1d(x, y, x_interp): + result = np.zeros_like(x_interp) + for i in range(1, len(x) - 2): + x0, x1, x2, x3 = x[i-1], x[i], x[i+1], x[i+2] + y0, y1, y2, y3 = y[i-1], y[i], y[i+1], y[i+2] + + mask = (x_interp >= x1) & (x_interp <= x2) + t = (x_interp[mask] - x1) / (x2 - x1) + + result[mask] = ( + 0.5 * ( + (2 * y1) + + (-y0 + y2) * t + + (2*y0 - 5*y1 + 4*y2 - y3) * t**2 + + (-y0 + 3*y1 - 3*y2 + y3) * t**3 + ) + ) + return result + + x_interp = np.asarray(frame_range) + x = np.asarray(center_frames) + + if len(x) < 4: + raise ValueError("Catmull-Rom interpolation requires at least 4 points.") + + drift_x = catmull_rom_1d(x, drift_est[:, 0], x_interp) + drift_y = catmull_rom_1d(x, drift_est[:, 1], x_interp) + drift_z = catmull_rom_1d(x, drift_est[:, 2], x_interp) + + return np.vstack([drift_x, drift_y, drift_z]).T + + +# CUDA kernel and wrapper — only defined when CUDA is available at import time. +if _CUDA_AVAILABLE: + @_cuda.jit + def cost_function_full_3d_chunked(d_locs_time, start_idx, chunk_size, d_idx_i, d_idx_j, d_sigma, + d_sigma_factor, d_val, d_val_sum, d_deri, d_locs_coords, mu): + """Compute negative-overlap cost and gradient for 3D localizations (CUDA kernel).""" + tx = _cuda.threadIdx.x + ty = _cuda.blockIdx.x + bw = _cuda.blockDim.x + pos = tx + ty * bw + + if pos < chunk_size: + i = d_idx_i[pos + start_idx] + j = d_idx_j[pos + start_idx] + + ti = d_locs_time[i] + tj = d_locs_time[j] + + dx = (d_locs_coords[i, 0] - mu[ti, 0]) - (d_locs_coords[j, 0] - mu[tj, 0]) + dy = (d_locs_coords[i, 1] - mu[ti, 1]) - (d_locs_coords[j, 1] - mu[tj, 1]) + dz = (d_locs_coords[i, 2] - mu[ti, 2]) - (d_locs_coords[j, 2] - mu[tj, 2]) + sigma_sq = (2 * d_sigma * d_sigma_factor) ** 2 + + diff_sq = dx * dx + dy * dy + dz * dz + val = 1 / (d_sigma * d_sigma_factor) * math.exp(-diff_sq / sigma_sq) + d_val[pos] = val + + _cuda.atomic.add(d_deri, (tj, 0), 2 * val * (d_locs_coords[j, 0] - d_locs_coords[i, 0] + mu[ti, 0] - mu[tj, 0]) / sigma_sq) + _cuda.atomic.add(d_deri, (tj, 1), 2 * val * (d_locs_coords[j, 1] - d_locs_coords[i, 1] + mu[ti, 1] - mu[tj, 1]) / sigma_sq) + _cuda.atomic.add(d_deri, (tj, 2), 2 * val * (d_locs_coords[j, 2] - d_locs_coords[i, 2] + mu[ti, 2] - mu[tj, 2]) / sigma_sq) + + _cuda.atomic.add(d_deri, (ti, 0), 2 * val * (d_locs_coords[i, 0] - d_locs_coords[j, 0] + mu[tj, 0] - mu[ti, 0]) / sigma_sq) + _cuda.atomic.add(d_deri, (ti, 1), 2 * val * (d_locs_coords[i, 1] - d_locs_coords[j, 1] + mu[tj, 1] - mu[ti, 1]) / sigma_sq) + _cuda.atomic.add(d_deri, (ti, 2), 2 * val * (d_locs_coords[i, 2] - d_locs_coords[j, 2] + mu[tj, 2] - mu[ti, 2]) / sigma_sq) + + _cuda.atomic.add(d_val_sum, 0, val) + d_val[pos] = 0 + + def cuda_wrapper_chunked(mu, d_locs_coords, d_locs_time, d_idx_i, d_idx_j, d_sigma, d_sigma_factor, d_val, + d_deri, chunk_size): + """Interface between Python optimizer and the CUDA kernel (chunked to manage memory).""" + val_total = 0 + d_val_sum = _cuda.to_device(np.zeros(1, dtype=np.float64)) + mu_dev = _cuda.to_device(np.asarray(mu.reshape(int(mu.size / 3), 3), dtype=np.float64)) + + n_chunks = int(np.ceil(d_idx_i.size / chunk_size)) + threadsperblock = 128 + + for i in range(n_chunks - 1): + idc_start = i * chunk_size + blockspergrid = (chunk_size + (threadsperblock - 1)) // threadsperblock + cost_function_full_3d_chunked[blockspergrid, threadsperblock]( + d_locs_time, idc_start, chunk_size, d_idx_i, d_idx_j, + d_sigma, d_sigma_factor, d_val, d_val_sum, d_deri, d_locs_coords, mu_dev + ) + val_total += d_val_sum.copy_to_host() + + # Final chunk + n_remaining = d_idx_i.size - (n_chunks - 1) * chunk_size + idc_start = (n_chunks - 1) * chunk_size + blockspergrid = (n_remaining + (threadsperblock - 1)) // threadsperblock + cost_function_full_3d_chunked[blockspergrid, threadsperblock]( + d_locs_time, idc_start, n_remaining, d_idx_i, d_idx_j, + d_sigma, d_sigma_factor, d_val, d_val_sum, d_deri, d_locs_coords, mu_dev + ) + val_total += d_val_sum.copy_to_host() + deri = d_deri.copy_to_host() + d_deri[:] = 0 + + return -np.nansum(val_total), -deri.flatten() + + +@dataclass +class SegmentationResult: + """Container for segmentation results.""" + loc_segments: np.ndarray + loc_valid: np.ndarray + center_frames: np.ndarray + n_segments: int + out_dict: Optional[Dict] = None + + +def _group_by_frame(loc_frames: np.ndarray): + """Returns a dict {frame_number: indices_in_loc_frames} efficiently.""" + sort_idx = np.argsort(loc_frames) + sorted_frames = loc_frames[sort_idx] + unique_frames, start_idx, counts = np.unique(sorted_frames, return_index=True, return_counts=True) + frame_to_indices = { + frame: sort_idx[start:start + count] + for frame, start, count in zip(unique_frames, start_idx, counts) + } + return unique_frames, frame_to_indices + + +def segment_by_num_locs_per_window(loc_frames: np.ndarray, min_n_locs_per_window: int, + max_locs_per_segment = None, + return_param_dict: bool = False) -> SegmentationResult: + """ + Segments by collecting a minimum number of localizations per window. + Once the threshold is met and enough locs remain, a new segment is created. + This method ensures that each segment has at least `min_n_locs_per_window` localizations, + while also trying to avoid creating segments that are too small at the end of the dataset. + If `max_locs_per_segment` is set, a random subset of that size is chosen from each segment. + This method is particularly useful for datasets with varying localization densities over time. + Parameters: + loc_frames (np.ndarray): Array of frame numbers for each localization. + min_n_locs_per_window (int): Minimum number of localizations per segment. + max_locs_per_segment (Optional[int]): Maximum number of localizations per segment. If None, all locs are used. + return_param_dict (bool): Whether to return a dictionary of segmentation parameters. + Returns: + SegmentationResult: A dataclass containing segmentation results and parameters. + """ + loc_frames = np.asarray(loc_frames, dtype=int) + n_locs = len(loc_frames) + + if max_locs_per_segment is not None and max_locs_per_segment < 1: # downsampling in percentage + max_locs_per_segment = int(min_n_locs_per_window * max_locs_per_segment) + + unique_frames, frame_to_indices = _group_by_frame(loc_frames) + loc_segments = np.full(n_locs, -1, dtype=int) # Default to -1 for safety + segment_counter = 0 + n_locs_in_current_segment = 0 + current_segment_indices = [] + start_frames, end_frames, locs_per_segment = [], [], [] + + for i, frame in enumerate(unique_frames): + indices = frame_to_indices[frame] + n_locs_this_frame = len(indices) + remaining_locs = n_locs - (len(current_segment_indices) + n_locs_this_frame + np.sum(locs_per_segment)) + + # Add frame to current segment if: + # - It fills the current segment to threshold + # - AND there are enough locs left for another segment (or it's the last frame) + if (n_locs_in_current_segment + n_locs_this_frame >= min_n_locs_per_window) and \ + (remaining_locs >= min_n_locs_per_window or i == len(unique_frames) - 1): + current_segment_indices.extend(indices) + n_locs_in_current_segment += n_locs_this_frame + + loc_segments[current_segment_indices] = segment_counter + start_frames.append(loc_frames[current_segment_indices[0]]) + end_frames.append(loc_frames[current_segment_indices[-1]]) + locs_per_segment.append(len(current_segment_indices)) + + segment_counter += 1 + current_segment_indices = [] + n_locs_in_current_segment = 0 + else: + # Defer frame to current segment + current_segment_indices.extend(indices) + n_locs_in_current_segment += n_locs_this_frame + + n_segments = segment_counter + center_frames = np.zeros(n_segments) + loc_valid = np.zeros(n_locs, dtype=bool) + + for i in range(n_segments): + segment_indices = np.where(loc_segments == i)[0] + if max_locs_per_segment and len(segment_indices) > max_locs_per_segment: + selected = np.random.choice(segment_indices, max_locs_per_segment, replace=False) + else: + selected = segment_indices + loc_valid[selected] = True + locs_per_segment[i] = len(selected) + center_frames[i] = np.mean(loc_frames[selected]) + + out_dict = None + if return_param_dict: + n_locs_valid = loc_valid.sum() + out_dict = { + "n_segments": n_segments, + "min_n_locs_per_window": min_n_locs_per_window, + "frames_per_window": -1, + "start_frames": np.array(start_frames), + "end_frames": np.array(end_frames), + "locs_per_segment": np.array(locs_per_segment), + "n_locs": n_locs, + "n_locs_valid": n_locs_valid, + "n_locs_invalid": n_locs - n_locs_valid, + "center_frames": center_frames + } + + return SegmentationResult(loc_segments, loc_valid, center_frames, n_segments, out_dict) + + +def segment_by_frame_windows(loc_frames: np.ndarray, n_frames_per_window: int, + max_locs_per_segment = None, + return_param_dict: bool = False) -> SegmentationResult: + """ + Splits localization data into fixed-size windows of N frames. + All localizations in those frames are grouped into one segment. + """ + loc_frames = np.asarray(loc_frames, dtype=int) + frames, frame_to_indices = _group_by_frame(loc_frames) + n_locs = len(loc_frames) + n_segments = int(np.ceil(len(frames) / n_frames_per_window)) + + if max_locs_per_segment is not None and max_locs_per_segment < 1: # downsampling in percentage + max_locs_per_segment = n_locs / n_segments * max_locs_per_segment + + loc_segments = np.zeros(n_locs, dtype=int) + center_frames = np.zeros(n_segments) + loc_valid = np.ones(n_locs, dtype=bool) + start_frames, end_frames, locs_per_segment = [], [], [] + + for i in range(n_segments): + frame_window = frames[i * n_frames_per_window:(i + 1) * n_frames_per_window] + indices = np.concatenate([frame_to_indices[frame] for frame in frame_window if frame in frame_to_indices]) + if len(indices) == 0: + continue + loc_segments[indices] = i + start_frames.append(frame_window[0]) + end_frames.append(frame_window[-1]) + center_frames[i] = np.mean(loc_frames[indices]) + locs_per_segment.append(len(indices)) + if max_locs_per_segment and len(indices) > max_locs_per_segment: + mask = np.ones(len(indices), dtype=bool) + mask[np.random.choice(len(indices), len(indices) - max_locs_per_segment, replace=False)] = False + loc_valid[indices[~mask]] = False + + out_dict = None + if return_param_dict: + n_locs_valid = loc_valid.sum() + out_dict = { + "n_segments": n_segments, + "min_n_locs_per_window": -1, + "frames_per_window": n_frames_per_window, + "start_frames": np.array(start_frames), + "end_frames": np.array(end_frames), + "locs_per_segment": np.array(locs_per_segment), + "n_locs": n_locs, + "n_locs_valid": n_locs_valid, + "n_locs_invalid": n_locs - n_locs_valid, + "center_frames": center_frames + } + + return SegmentationResult(loc_segments, loc_valid, center_frames, n_segments, out_dict) + + +def segment_by_num_windows(loc_frames: np.ndarray, n_windows: int, max_locs_per_segment = None, + return_param_dict: bool = False) -> SegmentationResult: + """ + Converts number of windows into an equivalent minimum locs per window, + then calls `segment_by_num_locs_per_window`. + """ + n_locs = len(loc_frames) + n_locs_per_window = int(np.ceil(n_locs / n_windows)) + if max_locs_per_segment is not None and max_locs_per_segment < 1: # downsampling in percentage + max_locs_per_segment = int(n_locs_per_window * max_locs_per_segment) + return segment_by_num_locs_per_window(loc_frames, n_locs_per_window, max_locs_per_segment, return_param_dict) + + +def segmentation_wrapper(loc_frames: np.ndarray, segmentation_var: int, segmentation_mode: int = 2, + max_locs_per_segment = None, + return_param_dict: bool = False) -> SegmentationResult: + """ + Dispatch function that selects segmentation method: + 0 → fixed number of windows + 1 → fixed number of localizations per window + 2 → fixed number of frames per window (default) + """ + if segmentation_mode == 0: + return segment_by_num_windows(loc_frames, segmentation_var, max_locs_per_segment, return_param_dict) + elif segmentation_mode == 1: + return segment_by_num_locs_per_window(loc_frames, segmentation_var, max_locs_per_segment, return_param_dict) + else: + return segment_by_frame_windows(loc_frames, segmentation_var, max_locs_per_segment, return_param_dict) + + +def pair_indices_kdtree(coordinates, distance): + """ + Find all pairs of points within a certain distance using a KD-tree. + Parameters: + - coordinates: np.ndarray of shape (N, D) where N is the number of points and D is the dimensionality. + - distance: float, the maximum distance to consider points as a pair. + Returns: + - idx1: np.ndarray of shape (M,), indices of the first point in each pair. + - idx2: np.ndarray of shape (M,), indices of the second point in each pair. + """ + tree = cKDTree(coordinates) + try: + pairs = tree.query_pairs(r=distance, output_type='ndarray') + except MemoryError: + print("[pair_indices_kdtree] MemoryError encountered") + return [], [], False + return np.ascontiguousarray(pairs[:, 0], dtype=np.int32), np.ascontiguousarray(pairs[:, 1], dtype=np.int32), True \ No newline at end of file From a921780e9cc62a0a3384c0eebf22cead419e9920 Mon Sep 17 00:00:00 2001 From: Lenny Reinkensmeier Date: Thu, 28 May 2026 17:21:01 +0200 Subject: [PATCH 05/19] Fix relative lib import after moving comet to ext folder. --- picasso/ext/comet.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/picasso/ext/comet.py b/picasso/ext/comet.py index a014354e..ff074e4e 100644 --- a/picasso/ext/comet.py +++ b/picasso/ext/comet.py @@ -10,7 +10,7 @@ from scipy.optimize import minimize from scipy.spatial import cKDTree -from . import lib, __version__ +from .. import lib, __version__ try: from numba import cuda as _cuda From 9e6ce80f1c755f8be4490437ae3a0446165cf6d5 Mon Sep 17 00:00:00 2001 From: rafalkowalewski1 Date: Fri, 29 May 2026 13:11:15 +0200 Subject: [PATCH 06/19] clean up the comet integration --- changelog.md | 4 +++- picasso/gui/render.py | 27 +++++++++++++----------- readme.rst | 1 + release/one_click_macos_gui/readme.txt | 1 + release/one_click_windows_gui/readme.txt | 1 + 5 files changed, 21 insertions(+), 13 deletions(-) diff --git a/changelog.md b/changelog.md index bb5c30b0..49cc37e4 100644 --- a/changelog.md +++ b/changelog.md @@ -1,6 +1,8 @@ # Changelog -Last change: 28-MAY-2026 CEST +Last change: 29-MAY-2026 CEST + +- COMET ([DOI: 10.64898/2026.03.27.714864](https://doi.org/10.64898/2026.03.27.714864)) integrated (#671), big thanks to @LREIN663 ## 0.10.1 diff --git a/picasso/gui/render.py b/picasso/gui/render.py index b142e4ba..b474b52b 100644 --- a/picasso/gui/render.py +++ b/picasso/gui/render.py @@ -55,8 +55,8 @@ ) from .rotation import RotationWindow -# Optional modules with external/hardware dependencies live in ext/ -from ..ext.bitplane import IMSWRITER +# Optional modules with external/hardware dependencies live in ext +from ..ext.bitplane import IMSWRITER # PyImarisWrite works on Windows only from ..ext import comet if IMSWRITER: @@ -2219,6 +2219,7 @@ def getParams( clustered_locs, ) + class COMETDialog(QtWidgets.QDialog): """Dialog to choose parameters for COMET undrifting. @@ -11078,10 +11079,9 @@ def show_drift(self) -> None: self.plot_window.show() def undrift_comet(self) -> None: - """Undrift with COMET. - - See https://comet.smlm.tools - """ + """Undrift with COMET. See https://comet.smlm.tools and + Reinkensmeier L., Aufmkolk S., Farabella I., Egner A., and + Bates M. biorxiv, 2026.""" channel = self.get_channel("Undrift by COMET") if channel is not None: locs = self.all_locs[channel] @@ -11108,12 +11108,12 @@ def undrift_comet(self) -> None: return locs = lib.ensure_sanity(locs, info) - self.all_locs[channel] = locs - self.locs[channel] = copy.copy(locs) + self.locs[channel] = locs self.infos[channel] = new_info self.index_blocks[channel] = None + self.render_index[channel] = None self.add_drift(channel, drift) - self.update_scene() + self.update_scene(resample_locs=True) self.show_drift() def undrift_aim(self) -> None: @@ -11364,7 +11364,7 @@ def _apply_drift(self, channel: int, drift: pd.DataFrame) -> None: self.index_blocks[channel] = None self.render_index[channel] = None self.update_scene(resample_locs=True) - + def sync_groups(self) -> None: """Remove localizations whose group field is not found in all channels.""" @@ -12116,8 +12116,6 @@ def initUI(self, plugins_loaded: bool) -> None: # menu bar - Postprocess postprocess_menu = self.menu_bar.addMenu("Postprocess") - undrift_comet_action = postprocess_menu.addAction("Undrift by COMET") - undrift_comet_action.triggered.connect(self.view.undrift_comet) undrift_aim_action = postprocess_menu.addAction("Undrift by AIM") undrift_aim_action.setShortcut("Ctrl+U") @@ -12135,6 +12133,11 @@ def initUI(self, plugins_loaded: bool) -> None: undrift_from_picked2d_action.triggered.connect( self.view.undrift_from_picked2d ) + undrift_comet_action = postprocess_menu.addAction("Undrift by COMET") + undrift_comet_action.triggered.connect(self.view.undrift_comet) + if not comet._CUDA_AVAILABLE: + undrift_comet_action.setVisible(False) + undrift_action = postprocess_menu.addAction("Undrift by RCC") undrift_action.triggered.connect(self.view.undrift_rcc) drift_action = postprocess_menu.addAction("Undo drift") diff --git a/readme.rst b/readme.rst index 8d51e25e..c65f093d 100644 --- a/readme.rst +++ b/readme.rst @@ -137,6 +137,7 @@ If you use Picasso in your research, please cite our Nature Protocols publicatio - 3D fitting via astigmatism. DOI: `10.1126/science.1153529 `__. - RCC undrifting: DOI: `10.1364/OE.22.015982 `__ - AIM undrifting. DOI: `10.1126/sciadv.adm776 `__ +- COMET undrifting. DOI: `10.64898/2026.03.27.714864 `__ - SMLM clusterer. DOIs: `10.1038/s41467-021-22606-1 `__ and `10.1038/s41586-023-05925-9 `__ - DBSCAN: Ester, et al. Inkdd, 1996. (Vol. 96, No. 34, pp. 226-231). - Anisotropic DBSCAN inspired by: `10.1021/acs.jpcb.4c02030 `__ diff --git a/release/one_click_macos_gui/readme.txt b/release/one_click_macos_gui/readme.txt index fc2cf68e..d0c6a105 100644 --- a/release/one_click_macos_gui/readme.txt +++ b/release/one_click_macos_gui/readme.txt @@ -62,6 +62,7 @@ If you use some of the functionalities provided by Picasso, please also cite the - 3D fitting via astigmatism. DOI: 10.1126/science.1153529 (https://www.science.org/doi/10.1126/science.1153529). - RCC undrifting: DOI: 10.1364/OE.22.015982 (https://doi.org/10.1364/OE.22.015982) - AIM undrifting. DOI: 10.1126/sciadv.adm776 (https://www.science.org/doi/10.1126/sciadv.adm7765) +- COMET undrifting. DOI: 10.64898/2026.03.27.714864 (https://doi.org/10.64898/2026.03.27.714864) - SMLM clusterer. DOIs: 10.1038/s41467-021-22606-1 (https://doi.org/10.1038/s41467-021-22606-1) and 10.1038/s41586-023-05925-9 (https://doi.org/10.1038/s41586-023-05925-9) - DBSCAN: Ester, et al. Inkdd, 1996. (Vol. 96, No. 34, pp. 226-231). - Anisotropic DBSCAN inspired by: 10.1021/acs.jpcb.4c02030 (https://doi.org/10.1021/acs.jpcb.4c02030) diff --git a/release/one_click_windows_gui/readme.txt b/release/one_click_windows_gui/readme.txt index 039f1303..78f30fd8 100644 --- a/release/one_click_windows_gui/readme.txt +++ b/release/one_click_windows_gui/readme.txt @@ -67,6 +67,7 @@ If you use some of the functionalities provided by Picasso, please also cite the - 3D fitting via astigmatism. DOI: 10.1126/science.1153529 (https://www.science.org/doi/10.1126/science.1153529). - RCC undrifting: DOI: 10.1364/OE.22.015982 (https://doi.org/10.1364/OE.22.015982) - AIM undrifting. DOI: 10.1126/sciadv.adm776 (https://www.science.org/doi/10.1126/sciadv.adm7765) +- COMET undrifting. DOI: 10.64898/2026.03.27.714864 (https://doi.org/10.64898/2026.03.27.714864) - SMLM clusterer. DOIs: 10.1038/s41467-021-22606-1 (https://doi.org/10.1038/s41467-021-22606-1) and 10.1038/s41586-023-05925-9 (https://doi.org/10.1038/s41586-023-05925-9) - DBSCAN: Ester, et al. Inkdd, 1996. (Vol. 96, No. 34, pp. 226-231). - Anisotropic DBSCAN inspired by: 10.1021/acs.jpcb.4c02030 (https://doi.org/10.1021/acs.jpcb.4c02030) From 1c4f1976d7926d46dea42c232fe69b6b7a9c0e80 Mon Sep 17 00:00:00 2001 From: rafalkowalewski1 Date: Fri, 29 May 2026 13:53:04 +0200 Subject: [PATCH 07/19] add more tests to comet --- tests/test_comet.py | 332 ++++++++++++++++++++++++++++++++++++++------ 1 file changed, 291 insertions(+), 41 deletions(-) diff --git a/tests/test_comet.py b/tests/test_comet.py index 976b3399..3050ef4a 100644 --- a/tests/test_comet.py +++ b/tests/test_comet.py @@ -24,6 +24,7 @@ # Fixtures # --------------------------------------------------------------------------- + @pytest.fixture def minimal_info(): """Minimal Picasso info list with required Pixelsize entry.""" @@ -46,7 +47,9 @@ def minimal_locs_df(): def minimal_locs_structured(minimal_locs_df): """Same data as a numpy structured array (Picasso's native format).""" df = minimal_locs_df - dtype = np.dtype([("frame", np.uint32), ("x", np.float32), ("y", np.float32)]) + dtype = np.dtype( + [("frame", np.uint32), ("x", np.float32), ("y", np.float32)] + ) arr = np.empty(len(df), dtype=dtype) arr["frame"] = df["frame"].to_numpy(np.uint32) arr["x"] = df["x"].to_numpy(np.float32) @@ -58,11 +61,15 @@ def minimal_locs_structured(minimal_locs_df): # Segmentation helpers # --------------------------------------------------------------------------- + class TestSegmentationByLocs: def test_produces_correct_segment_count(self): frames = np.repeat(np.arange(100), 5) # 500 locs, 5/frame result = comet.segmentation_wrapper( - frames, segmentation_var=50, segmentation_mode=1, return_param_dict=True + frames, + segmentation_var=50, + segmentation_mode=1, + return_param_dict=True, ) assert result.n_segments == 10 @@ -92,8 +99,11 @@ def test_max_locs_per_segment_respected(self): frames = np.repeat(np.arange(20), 100) # 2000 locs, 100/frame cap = 30 result = comet.segmentation_wrapper( - frames, segmentation_var=200, segmentation_mode=1, - max_locs_per_segment=cap, return_param_dict=True + frames, + segmentation_var=200, + segmentation_mode=1, + max_locs_per_segment=cap, + return_param_dict=True, ) assert result.loc_valid.sum() <= result.n_segments * cap @@ -118,7 +128,10 @@ class TestSegmentationByNumWindows: def test_num_windows_mode_creates_requested_count(self): frames = np.repeat(np.arange(100), 5) # 500 locs result = comet.segmentation_wrapper( - frames, segmentation_var=5, segmentation_mode=0, return_param_dict=True + frames, + segmentation_var=5, + segmentation_mode=0, + return_param_dict=True, ) # mode 0 derives min-locs-per-window from total/n_windows; the resulting # segment count should be approximately n_windows (off-by-one tolerated @@ -133,13 +146,115 @@ def test_all_locs_assigned(self): assert (result.loc_segments >= 0).all() +class TestSegmentationOutDict: + """The param dict reports per-segment bookkeeping; verify the values, not + just that segmentation runs.""" + + def test_valid_plus_invalid_equals_total(self): + frames = np.repeat(np.arange(60), 8) + result = comet.segmentation_wrapper( + frames, + segmentation_var=40, + segmentation_mode=1, + return_param_dict=True, + ) + d = result.out_dict + assert d["n_locs"] == len(frames) + assert d["n_locs_valid"] + d["n_locs_invalid"] == d["n_locs"] + + def test_per_segment_arrays_have_consistent_length(self): + frames = np.repeat(np.arange(80), 5) + result = comet.segmentation_wrapper( + frames, + segmentation_var=50, + segmentation_mode=1, + return_param_dict=True, + ) + d = result.out_dict + assert len(d["start_frames"]) == result.n_segments + assert len(d["end_frames"]) == result.n_segments + assert len(d["locs_per_segment"]) == result.n_segments + + def test_start_frames_not_after_end_frames(self): + frames = np.repeat(np.arange(80), 5) + result = comet.segmentation_wrapper( + frames, + segmentation_var=50, + segmentation_mode=1, + return_param_dict=True, + ) + d = result.out_dict + assert (d["start_frames"] <= d["end_frames"]).all() + + def test_center_frames_within_segment_bounds(self): + frames = np.repeat(np.arange(80), 5) + result = comet.segmentation_wrapper( + frames, + segmentation_var=50, + segmentation_mode=1, + return_param_dict=True, + ) + d = result.out_dict + assert (result.center_frames >= d["start_frames"]).all() + assert (result.center_frames <= d["end_frames"]).all() + + def test_locs_per_segment_sums_to_valid_count(self): + frames = np.repeat(np.arange(60), 6) + result = comet.segmentation_wrapper( + frames, + segmentation_var=40, + segmentation_mode=1, + return_param_dict=True, + ) + d = result.out_dict + # locs_per_segment is rewritten to the post-downsampling (valid) count. + assert d["locs_per_segment"].sum() == d["n_locs_valid"] + + def test_frame_window_center_frames_within_bounds(self): + frames = np.repeat(np.arange(100), 3) + result = comet.segmentation_wrapper( + frames, + segmentation_var=10, + segmentation_mode=2, + return_param_dict=True, + ) + d = result.out_dict + assert d["frames_per_window"] == 10 + assert (result.center_frames >= d["start_frames"]).all() + assert (result.center_frames <= d["end_frames"]).all() + + +class TestSegmentationDownsamplingByPercentage: + """max_locs_per_segment < 1 is interpreted as a fraction of the window size.""" + + def test_fractional_cap_reduces_valid_locs(self): + frames = np.repeat(np.arange(20), 50) # 1000 locs, 50/frame + full = comet.segmentation_wrapper( + frames, + segmentation_var=200, + segmentation_mode=1, + return_param_dict=True, + ) + half = comet.segmentation_wrapper( + frames, + segmentation_var=200, + segmentation_mode=1, + max_locs_per_segment=0.5, + return_param_dict=True, + ) + assert half.out_dict["n_locs_valid"] < full.out_dict["n_locs_valid"] + + # --------------------------------------------------------------------------- # Pair-finding # --------------------------------------------------------------------------- + class TestPairIndicesKdtree: def test_finds_close_pairs(self): - coords = np.array([[0, 0, 0], [1, 0, 0], [100, 0, 0]], dtype=np.float32) + coords = np.array( + [[0, 0, 0], [1, 0, 0], [100, 0, 0]], dtype=np.float32 + ) idx_i, idx_j, ok = comet.pair_indices_kdtree(coords, distance=5) assert ok assert len(idx_i) == 1 # only (0,1) are within distance 5 @@ -163,19 +278,68 @@ def test_returns_int32_arrays(self): assert idx_i.dtype == np.int32 assert idx_j.dtype == np.int32 + def test_2d_coordinates_supported(self): + coords = np.array([[0, 0], [1, 0], [100, 0]], dtype=np.float32) + idx_i, idx_j, ok = comet.pair_indices_kdtree(coords, distance=5) + assert ok + assert len(idx_i) == 1 # only (0,1) within distance 5 + + def test_pairs_are_lower_index_first(self): + coords = np.zeros((6, 3), dtype=np.float32) + idx_i, idx_j, ok = comet.pair_indices_kdtree(coords, distance=1) + assert ok + # query_pairs returns each pair once with i < j; relied on by the kernel. + assert (idx_i < idx_j).all() + + def test_distance_is_inclusive(self): + coords = np.array([[0, 0, 0], [10, 0, 0]], dtype=np.float32) + idx_i, idx_j, ok = comet.pair_indices_kdtree(coords, distance=10) + assert ok + assert ( + len(idx_i) == 1 + ) # points exactly `distance` apart count as a pair + + +# --------------------------------------------------------------------------- +# Pair-count estimation +# --------------------------------------------------------------------------- + + +class TestEstimatePairs: + def test_far_apart_points_estimate_zero(self): + # Each point lands in its own grid cell -> no within-cell pairs. + coords = ( + (np.arange(5)[:, None] * 1000.0).repeat(3, axis=1).astype(float) + ) + assert comet.estimate_pairs(coords, distance=10) == 0 + + def test_clustered_points_estimate_nonzero(self): + # 200 coincident points -> 200*199 = 39800 -> rounds to 40000. + coords = np.zeros((200, 3), dtype=float) + est = comet.estimate_pairs(coords, distance=10) + assert est == 40000 + + def test_result_rounded_to_nearest_ten_thousand(self): + coords = np.zeros((150, 3), dtype=float) + est = comet.estimate_pairs(coords, distance=10) + assert est % 10000 == 0 + # --------------------------------------------------------------------------- # Drift interpolation # --------------------------------------------------------------------------- + class TestInterpolateDrift: def _center_frames_and_drift(self): center_frames = np.linspace(10, 90, 6) - drift_est = np.column_stack([ - np.sin(center_frames / 10), - np.cos(center_frames / 10), - np.zeros_like(center_frames), - ]) + drift_est = np.column_stack( + [ + np.sin(center_frames / 10), + np.cos(center_frames / 10), + np.zeros_like(center_frames), + ] + ) return center_frames, drift_est def test_cubic_output_shape(self): @@ -204,62 +368,126 @@ def test_unknown_method_raises(self): def test_catmull_rom_output_shape(self): cf, de = self._center_frames_and_drift() frame_range = np.arange(10, 90) - out = comet.interpolate_drift(cf, de, frame_range, method="catmull-rom") + out = comet.interpolate_drift( + cf, de, frame_range, method="catmull-rom" + ) assert out.shape == (len(frame_range), 3) def test_catmull_rom_requires_min_points(self): cf = np.array([0.0, 10.0, 20.0]) # only 3 points de = np.zeros((3, 3)) with pytest.raises(ValueError, match="at least 4 points"): - comet.interpolate_drift(cf, de, np.arange(20), method="catmull-rom") + comet.interpolate_drift( + cf, de, np.arange(20), method="catmull-rom" + ) + + def test_linear_passes_through_control_points(self): + cf = np.array([0.0, 25.0, 50.0, 75.0, 100.0]) + de = np.column_stack([cf * 0.01, cf * 0.005, cf * 0.002]) + out = comet.interpolate_drift(cf, de, cf.astype(int), method="linear") + np.testing.assert_allclose(out, de, atol=1e-6) + + def test_linear_clamps_outside_range(self): + # np.interp clamps to edge values beyond the control-point range. + cf = np.array([10.0, 20.0, 30.0]) + de = np.column_stack([cf, cf, cf]) + out = comet.interpolate_drift( + cf, de, np.array([0, 40]), method="linear" + ) + np.testing.assert_allclose(out[0], [10, 10, 10]) + np.testing.assert_allclose(out[1], [30, 30, 30]) + + def test_z_channel_interpolated_independently(self): + cf = np.array([0.0, 25.0, 50.0, 75.0, 100.0]) + de = np.column_stack([np.zeros_like(cf), np.zeros_like(cf), cf * 0.03]) + out = comet.interpolate_drift(cf, de, cf.astype(int), method="cubic") + np.testing.assert_allclose(out[:, 0], 0.0, atol=1e-6) + np.testing.assert_allclose(out[:, 1], 0.0, atol=1e-6) + np.testing.assert_allclose(out[:, 2], de[:, 2], atol=1e-6) # --------------------------------------------------------------------------- # comet() public API — no-GPU error path # --------------------------------------------------------------------------- + class TestCometNoCuda: """These tests run on any machine; they verify the failure mode when CUDA is not present.""" - def test_raises_runtime_error_without_cuda(self, minimal_locs_df, minimal_info): + def test_raises_runtime_error_without_cuda( + self, minimal_locs_df, minimal_info + ): if comet._CUDA_AVAILABLE: - pytest.skip("CUDA is present; testing no-GPU path only on CPU machines") + pytest.skip( + "CUDA is present; testing no-GPU path only on CPU machines" + ) with pytest.raises(RuntimeError, match="CUDA"): - comet.comet(minimal_locs_df, minimal_info, locs_per_segment=100, max_drift_nm=200) - - def test_error_message_mentions_alternatives(self, minimal_locs_df, minimal_info): + comet.comet( + minimal_locs_df, + minimal_info, + locs_per_segment=100, + max_drift_nm=200, + ) + + def test_error_message_mentions_alternatives( + self, minimal_locs_df, minimal_info + ): if comet._CUDA_AVAILABLE: - pytest.skip("CUDA is present; testing no-GPU path only on CPU machines") + pytest.skip( + "CUDA is present; testing no-GPU path only on CPU machines" + ) with pytest.raises(RuntimeError, match="AIM|RCC"): - comet.comet(minimal_locs_df, minimal_info, locs_per_segment=100, max_drift_nm=200) - - def test_structured_array_input_accepted_before_gpu_check(self, minimal_locs_structured, minimal_info): + comet.comet( + minimal_locs_df, + minimal_info, + locs_per_segment=100, + max_drift_nm=200, + ) + + def test_structured_array_input_accepted_before_gpu_check( + self, minimal_locs_structured, minimal_info + ): """Numpy structured arrays should be converted to DataFrame before the CUDA check, so the error is still a RuntimeError (not an AttributeError from missing .columns).""" if comet._CUDA_AVAILABLE: - pytest.skip("CUDA is present; testing no-GPU path only on CPU machines") + pytest.skip( + "CUDA is present; testing no-GPU path only on CPU machines" + ) with pytest.raises(RuntimeError, match="CUDA"): - comet.comet(minimal_locs_structured, minimal_info, locs_per_segment=100, max_drift_nm=200) + comet.comet( + minimal_locs_structured, + minimal_info, + locs_per_segment=100, + max_drift_nm=200, + ) def test_missing_columns_raises_key_error(self, minimal_info): """Locs missing required columns should raise KeyError, not crash on .columns.""" if comet._CUDA_AVAILABLE: - pytest.skip("CUDA is present; GPU path reached before column check") + pytest.skip( + "CUDA is present; GPU path reached before column check" + ) bad_locs = pd.DataFrame({"frame": np.arange(10), "x": np.zeros(10)}) # No 'y' column — should raise KeyError from the column check with pytest.raises((KeyError, RuntimeError)): - comet.comet(bad_locs, minimal_info, locs_per_segment=5, max_drift_nm=100) + comet.comet( + bad_locs, minimal_info, locs_per_segment=5, max_drift_nm=100 + ) def test_z_column_accepted(self, minimal_locs_df, minimal_info): """3D input (with z) should not crash before the GPU check.""" if comet._CUDA_AVAILABLE: - pytest.skip("CUDA is present; testing no-GPU path only on CPU machines") + pytest.skip( + "CUDA is present; testing no-GPU path only on CPU machines" + ) df = minimal_locs_df.copy() df["z"] = np.zeros(len(df), dtype=np.float32) with pytest.raises(RuntimeError, match="CUDA"): - comet.comet(df, minimal_info, locs_per_segment=100, max_drift_nm=200) + comet.comet( + df, minimal_info, locs_per_segment=100, max_drift_nm=200 + ) def test_missing_pixelsize_raises_on_gpu(self, minimal_locs_df): """info without 'Pixelsize' should fail with an informative error. @@ -268,13 +496,19 @@ def test_missing_pixelsize_raises_on_gpu(self, minimal_locs_df): pytest.skip("CUDA check fires before pixelsize check on CPU") bad_info = [{"Frames": 200}] with pytest.raises((KeyError, RuntimeError, ValueError)): - comet.comet(minimal_locs_df, bad_info, locs_per_segment=100, max_drift_nm=200) + comet.comet( + minimal_locs_df, + bad_info, + locs_per_segment=100, + max_drift_nm=200, + ) # --------------------------------------------------------------------------- # comet() public API — full pipeline (GPU required) # --------------------------------------------------------------------------- + @pytest.mark.skipif( not comet._CUDA_AVAILABLE, reason="CUDA GPU not available", @@ -284,8 +518,10 @@ class TestCometFullPipeline: def test_output_shapes(self, minimal_locs_df, minimal_info): undrifted, new_info, drift = comet.comet( - minimal_locs_df, minimal_info, - locs_per_segment=100, max_drift_nm=300, + minimal_locs_df, + minimal_info, + locs_per_segment=100, + max_drift_nm=300, target_sigma_nm=50, ) assert len(undrifted) == len(minimal_locs_df) @@ -295,28 +531,42 @@ def test_output_shapes(self, minimal_locs_df, minimal_info): def test_drift_covers_frame_range(self, minimal_locs_df, minimal_info): frame0 = minimal_locs_df["frame"] - minimal_locs_df["frame"].min() _, _, drift = comet.comet( - minimal_locs_df, minimal_info, - locs_per_segment=100, max_drift_nm=300, + minimal_locs_df, + minimal_info, + locs_per_segment=100, + max_drift_nm=300, target_sigma_nm=50, ) assert len(drift) == int(frame0.max()) + 1 def test_new_info_appended(self, minimal_locs_df, minimal_info): _, new_info, _ = comet.comet( - minimal_locs_df, minimal_info, - locs_per_segment=100, max_drift_nm=300, + minimal_locs_df, + minimal_info, + locs_per_segment=100, + max_drift_nm=300, target_sigma_nm=50, ) assert len(new_info) == len(minimal_info) + 1 assert "COMET" in new_info[-1]["Generated by"] - def test_structured_array_gives_same_result_as_dataframe(self, minimal_locs_df, minimal_locs_structured, minimal_info): + def test_structured_array_gives_same_result_as_dataframe( + self, minimal_locs_df, minimal_locs_structured, minimal_info + ): locs_df, _, drift_df = comet.comet( - minimal_locs_df, minimal_info, - locs_per_segment=100, max_drift_nm=300, target_sigma_nm=50, + minimal_locs_df, + minimal_info, + locs_per_segment=100, + max_drift_nm=300, + target_sigma_nm=50, ) locs_arr, _, drift_arr = comet.comet( - minimal_locs_structured, minimal_info, - locs_per_segment=100, max_drift_nm=300, target_sigma_nm=50, + minimal_locs_structured, + minimal_info, + locs_per_segment=100, + max_drift_nm=300, + target_sigma_nm=50, + ) + np.testing.assert_allclose( + drift_df["x"].to_numpy(), drift_arr["x"].to_numpy(), rtol=1e-5 ) - np.testing.assert_allclose(drift_df["x"].to_numpy(), drift_arr["x"].to_numpy(), rtol=1e-5) From e438de60815ffada224e1b9cbbd5a6b08faf7dc6 Mon Sep 17 00:00:00 2001 From: rafalkowalewski1 Date: Fri, 29 May 2026 15:25:57 +0200 Subject: [PATCH 08/19] fix comet integration to gui (qt6 update and no View.all_locs) + allow higher GPU usage + add new deps to pyproject --- picasso/ext/comet.py | 824 +++++++++++++++++++++++++++++++----------- picasso/gui/render.py | 11 +- pyproject.toml | 36 ++ 3 files changed, 649 insertions(+), 222 deletions(-) diff --git a/picasso/ext/comet.py b/picasso/ext/comet.py index ff074e4e..c172b1a0 100644 --- a/picasso/ext/comet.py +++ b/picasso/ext/comet.py @@ -14,14 +14,13 @@ try: from numba import cuda as _cuda + _CUDA_AVAILABLE = _cuda.is_available() except Exception: _cuda = None _CUDA_AVAILABLE = False - - def comet( locs: pd.DataFrame, info: list[dict], @@ -105,7 +104,9 @@ def comet( required_cols = {"x", "y", "frame"} missing = required_cols - set(locs.columns) if missing: - raise KeyError(f"Missing required columns for COMET: {sorted(missing)}") + raise KeyError( + f"Missing required columns for COMET: {sorted(missing)}" + ) has_z = "z" in locs.columns @@ -118,9 +119,7 @@ def comet( dataset = np.zeros((len(locs), 4), dtype=np.float64) dataset[:, 0] = locs["x"].to_numpy(dtype=np.float64) * pixelsize_nm dataset[:, 1] = locs["y"].to_numpy(dtype=np.float64) * pixelsize_nm - dataset[:, 2] = ( - locs["z"].to_numpy(dtype=np.float64) if has_z else 0.0 - ) + dataset[:, 2] = locs["z"].to_numpy(dtype=np.float64) if has_z else 0.0 dataset[:, 3] = frame0 if initial_sigma_nm is None: @@ -149,10 +148,18 @@ def comet( # Apply drift back to dataframe in original units frame_idx = frame0.astype(np.int64) - locs["x"] = locs["x"].to_numpy(dtype=np.float64) - drift_nm[frame_idx, 0] / pixelsize_nm - locs["y"] = locs["y"].to_numpy(dtype=np.float64) - drift_nm[frame_idx, 1] / pixelsize_nm + locs["x"] = ( + locs["x"].to_numpy(dtype=np.float64) + - drift_nm[frame_idx, 0] / pixelsize_nm + ) + locs["y"] = ( + locs["y"].to_numpy(dtype=np.float64) + - drift_nm[frame_idx, 1] / pixelsize_nm + ) if has_z: - locs["z"] = locs["z"].to_numpy(dtype=np.float64) - drift_nm[frame_idx, 2] + locs["z"] = ( + locs["z"].to_numpy(dtype=np.float64) - drift_nm[frame_idx, 2] + ) # Build drift dataframe in Picasso-style units drift_dict = { @@ -182,56 +189,81 @@ def comet( return locs, new_info, drift -def comet_run_kd(dataset, segmentation_mode, segmentation_var, max_locs_per_segment=None, - initial_sigma_nm=None, gt_drift=None, display=False, return_corrected_locs=False, - max_drift_nm=300, target_sigma_nm=1, boxcar_width=1, drift_max_bound_factor=2, - save_intermediate_results=False, - interpolation_method='cubic', mode="cuda", min_max_frames=None, - pair_indices_safety_check=False): + +def comet_run_kd( + dataset, + segmentation_mode, + segmentation_var, + max_locs_per_segment=None, + initial_sigma_nm=None, + gt_drift=None, + display=False, + return_corrected_locs=False, + max_drift_nm=300, + target_sigma_nm=1, + boxcar_width=1, + drift_max_bound_factor=2, + save_intermediate_results=False, + interpolation_method="cubic", + mode="cuda", + min_max_frames=None, + pair_indices_safety_check=False, + mapped_memory_threshold_bytes=None, + device_mem_fraction=0.5, +): + """ + Run COMET drift correction end-to-end. + + Pipeline: temporal segmentation -> KD-tree neighbor pairs -> cost optimization (L-BFGS-B) + -> optional temporal smoothing -> spline interpolation to per-frame drift -> (optional) subtract drift. + + Parameters + ---------- + dataset : ndarray of shape (N, 4) + Localization array with columns [x_nm, y_nm, z_nm, frame]. Units in nm; frame is int. + For 2D CSVs, insert a zero z column to get (N, 4). + segmentation_mode : {0, 1, 2} + Temporal segmentation mode: + 0 = number of windows (choose S directly), + 1 = localizations per window (accumulate frames until >= X locs), + 2 = fixed frame window size (default). + segmentation_var : int + Mode-dependent value (S, locs per window, or frames per window). + initial_sigma_nm : float, default=100 + Initial Gaussian length scale for the overlap kernel (coarse scale). + target_sigma_nm : float, default=1 + Target (final) Gaussian length scale for fine refinement. + max_drift_nm : float or None, default=None + Pair radius in nm used for neighbor search. If None, uses 3 * initial_sigma_nm. + drift_max_bound_factor : float, default=1.0 + Multiplicative factor for L-BFGS-B box bounds around +-max_drift. + boxcar_width : int, default=1 + Temporal smoothing width (segments) applied to the estimated drift between optimizer steps. + interpolation_method : {"cubic", "catmull-rom"}, default="cubic" + Spline used to convert per-segment drift to per-frame drift. + max_locs_per_segment : int or None, default=None + Optional downsampling cap per segment (to control memory/time). + mode : str, "cuda" + only Cuda in this version + return_corrected_locs : bool, default=False + If True, also return drift-corrected localizations. + mapped_memory_threshold_bytes : int or None, default=None + Byte budget for the pair-index arrays on the GPU. If exceeded, the + arrays are staged in mapped host memory instead of VRAM. None + auto-detects the budget from free VRAM (see `device_mem_fraction`). + Set explicitly to force deterministic behavior across machines/GPUs. + device_mem_fraction : float, default=0.5 + Fraction of free VRAM used as the index-array budget when + `mapped_memory_threshold_bytes` is None. Lower on small GPUs to avoid + out-of-memory; raise toward 1.0 to keep more in faster device memory. + + Returns + ------- + drift_interp_with_frames : ndarray of shape (F, 4) + Per-frame drift with columns [dx_nm, dy_nm, dz_nm, frame]. + corrected_locs : ndarray of shape (N, 4), optional + Only if return_corrected_locs=True. Columns are [x_nm, y_nm, z_nm, segment_id]. """ - Run COMET drift correction end-to-end. - - Pipeline: temporal segmentation -> KD-tree neighbor pairs -> cost optimization (L-BFGS-B) - -> optional temporal smoothing -> spline interpolation to per-frame drift -> (optional) subtract drift. - - Parameters - ---------- - dataset : ndarray of shape (N, 4) - Localization array with columns [x_nm, y_nm, z_nm, frame]. Units in nm; frame is int. - For 2D CSVs, insert a zero z column to get (N, 4). - segmentation_mode : {0, 1, 2} - Temporal segmentation mode: - 0 = number of windows (choose S directly), - 1 = localizations per window (accumulate frames until >= X locs), - 2 = fixed frame window size (default). - segmentation_var : int - Mode-dependent value (S, locs per window, or frames per window). - initial_sigma_nm : float, default=100 - Initial Gaussian length scale for the overlap kernel (coarse scale). - target_sigma_nm : float, default=1 - Target (final) Gaussian length scale for fine refinement. - max_drift_nm : float or None, default=None - Pair radius in nm used for neighbor search. If None, uses 3 * initial_sigma_nm. - drift_max_bound_factor : float, default=1.0 - Multiplicative factor for L-BFGS-B box bounds around +-max_drift. - boxcar_width : int, default=1 - Temporal smoothing width (segments) applied to the estimated drift between optimizer steps. - interpolation_method : {"cubic", "catmull-rom"}, default="cubic" - Spline used to convert per-segment drift to per-frame drift. - max_locs_per_segment : int or None, default=None - Optional downsampling cap per segment (to control memory/time). - mode : str, "cuda" - only Cuda in this version - return_corrected_locs : bool, default=False - If True, also return drift-corrected localizations. - - Returns - ------- - drift_interp_with_frames : ndarray of shape (F, 4) - Per-frame drift with columns [dx_nm, dy_nm, dz_nm, frame]. - corrected_locs : ndarray of shape (N, 4), optional - Only if return_corrected_locs=True. Columns are [x_nm, y_nm, z_nm, segment_id]. - """ loc_frames = dataset[:, -1] if min_max_frames is None: @@ -239,9 +271,16 @@ def comet_run_kd(dataset, segmentation_mode, segmentation_var, max_locs_per_segm # Segment the dataset based on frame numbers into time windows - result, sorted_dataset, idx_i, idx_j = segmentation_and_pair_indices_wrapper( - dataset, segmentation_var, segmentation_mode, max_drift_nm, max_locs_per_segment, - pair_indices_safety_check=pair_indices_safety_check) + result, sorted_dataset, idx_i, idx_j = ( + segmentation_and_pair_indices_wrapper( + dataset, + segmentation_var, + segmentation_mode, + max_drift_nm, + max_locs_per_segment, + pair_indices_safety_check=pair_indices_safety_check, + ) + ) # Set default initial sigma if not provided if initial_sigma_nm is None: @@ -250,8 +289,10 @@ def comet_run_kd(dataset, segmentation_mode, segmentation_var, max_locs_per_segm # Run drift optimization t0 = time.time() drift_est = optimize_3d_chunked_better_moving_avg_kd( - result.n_segments, sorted_dataset, - idx_i, idx_j, + result.n_segments, + sorted_dataset, + idx_i, + idx_j, sigma_nm=initial_sigma_nm, target_sigma_nm=target_sigma_nm, drift_max_nm=max_drift_nm, @@ -259,7 +300,9 @@ def comet_run_kd(dataset, segmentation_mode, segmentation_var, max_locs_per_segm display_steps=display, boxcar_width=boxcar_width, segmentation_result=result, - mode=mode + mode=mode, + mapped_memory_threshold_bytes=mapped_memory_threshold_bytes, + device_mem_fraction=device_mem_fraction, ) elapsed = time.time() - t0 @@ -268,13 +311,21 @@ def comet_run_kd(dataset, segmentation_mode, segmentation_var, max_locs_per_segm vld_tp = np.where(~np.isnan(drift_est[:, 0])) frame_interp = np.arange(0, min_max_frames[1] + 1, dtype=int) - drift_interp = interpolate_drift(result.center_frames[vld_tp], drift_est[vld_tp], frame_interp, - method=interpolation_method) - drift_interp_with_frames = np.hstack((drift_interp, frame_interp[:, np.newaxis])) + drift_interp = interpolate_drift( + result.center_frames[vld_tp], + drift_est[vld_tp], + frame_interp, + method=interpolation_method, + ) + drift_interp_with_frames = np.hstack( + (drift_interp, frame_interp[:, np.newaxis]) + ) # Apply drift correction to original localizations for i in range(3): - dataset[:, i] = dataset[:, i] - drift_interp[dataset[:, -1].astype(int), i] + dataset[:, i] = ( + dataset[:, i] - drift_interp[dataset[:, -1].astype(int), i] + ) # Optionally show estimated drift curve if display: @@ -284,19 +335,33 @@ def comet_run_kd(dataset, segmentation_mode, segmentation_var, max_locs_per_segm plt.title("Estimated Drift") plt.xlabel("Frames") plt.ylabel("Drift (nm)") - plt.legend(['X', 'Y', 'Z']) + plt.legend(["X", "Y", "Z"]) plt.show() - # Optional GT comparison plot if display and gt_drift is not None: fig, ax = plt.subplots(3, 1) - ax[0].plot(gt_drift[:, 3], gt_drift[:, 0], label='GT Drift X') - ax[1].plot(gt_drift[:, 3], gt_drift[:, 1], label='GT Drift Y') - ax[2].plot(gt_drift[:, 3], gt_drift[:, 2], label='GT Drift Z') - ax[0].plot(frame_interp, drift_interp[:, 0], label='Estimated Drift X', linestyle='--') - ax[1].plot(frame_interp, drift_interp[:, 1], label='Estimated Drift Y', linestyle='--') - ax[2].plot(frame_interp, drift_interp[:, 2], label='Estimated Drift Z', linestyle='--') + ax[0].plot(gt_drift[:, 3], gt_drift[:, 0], label="GT Drift X") + ax[1].plot(gt_drift[:, 3], gt_drift[:, 1], label="GT Drift Y") + ax[2].plot(gt_drift[:, 3], gt_drift[:, 2], label="GT Drift Z") + ax[0].plot( + frame_interp, + drift_interp[:, 0], + label="Estimated Drift X", + linestyle="--", + ) + ax[1].plot( + frame_interp, + drift_interp[:, 1], + label="Estimated Drift Y", + linestyle="--", + ) + ax[2].plot( + frame_interp, + drift_interp[:, 2], + label="Estimated Drift Z", + linestyle="--", + ) ax[1].set_title("Ground Truth vs Estimated Drift") ax[1].set_xlabel("Frames") ax[0].set_ylabel("Drift (nm)") @@ -310,11 +375,23 @@ def comet_run_kd(dataset, segmentation_mode, segmentation_var, max_locs_per_segm return drift_interp_with_frames -def optimize_3d_chunked_better_moving_avg_kd(n_segments, locs_nm, idx_i, idx_j, sigma_nm=30, drift_max_nm=300, - target_sigma_nm=30, display_steps=False, - boxcar_width=3, drift_max_bound_factor=2, - segmentation_result=None, - mode="cuda", return_calc_time=False): +def optimize_3d_chunked_better_moving_avg_kd( + n_segments, + locs_nm, + idx_i, + idx_j, + sigma_nm=30, + drift_max_nm=300, + target_sigma_nm=30, + display_steps=False, + boxcar_width=3, + drift_max_bound_factor=2, + segmentation_result=None, + mode="cuda", + return_calc_time=False, + mapped_memory_threshold_bytes=None, + device_mem_fraction=0.5, +): """ Estimate per-segment drift (mu) by minimizing the negative Gaussian-overlap cost with an L-BFGS-B optimizer and a coarse-to-fine schedule on sigma. @@ -357,6 +434,18 @@ def optimize_3d_chunked_better_moving_avg_kd(n_segments, locs_nm, idx_i, idx_j, If True, use the CPU backend; otherwise try GPU (CUDA) and fall back to CPU if unavailable. return_calc_time : bool, default=False If True, also return the total computation time in seconds. + mapped_memory_threshold_bytes : int or None, default=None + Size budget (in bytes) for the pair-index arrays on the GPU. If the two + int32 index arrays together exceed this, they are staged in mapped host + memory instead of device VRAM. If None, the budget is auto-detected from + currently-free VRAM (see `device_mem_fraction`). Provide an explicit value + to force deterministic behavior regardless of GPU state. + device_mem_fraction : float, default=0.5 + Fraction of currently-free VRAM used as the index-array budget when + `mapped_memory_threshold_bytes` is None. Lower it if you hit out-of-memory + errors (more aggressive fallback to mapped memory); raise it toward 1.0 to + keep larger arrays in faster device memory. Ignored when an explicit + threshold is given. Returns ------- @@ -364,7 +453,7 @@ def optimize_3d_chunked_better_moving_avg_kd(n_segments, locs_nm, idx_i, idx_j, Estimated per-segment drift (dx, dy, dz) in nanometers. calc_time_s : float, optional Only when `return_calc_time=True`. Wall-clock time for the optimization. - """ + """ if segmentation_result is None: segmentation_result = {} @@ -375,15 +464,35 @@ def optimize_3d_chunked_better_moving_avg_kd(n_segments, locs_nm, idx_i, idx_j, coords = locs_nm[:, :3].astype(np.float32).copy() times = locs_nm[:, 3].astype(np.int32).copy() - chunk_size = int(1E8) # 1E7 + chunk_size = int(1e8) # 1E7 quality_control = mode == "torch_qc" or mode == "cuda_qc" d_coords = _cuda.to_device(coords) d_times = _cuda.to_device(times) - if len(idx_i) * 4 > 2e9: - # Use mapped memory if index arrays are large - print("Large index arrays — using mapped memory.") + # The pair-index arrays can dwarf available VRAM on large datasets. Decide + # whether to stage them in device memory (fast) or fall back to mapped host + # memory (handles arrays larger than VRAM, slower per access). The threshold + # auto-adapts to the current GPU unless an explicit override is given. + idx_bytes = len(idx_i) * 4 * 2 # two int32 arrays (idx_i + idx_j) + if mapped_memory_threshold_bytes is None: + # Budget a fraction of currently-free VRAM, leaving headroom for d_val, + # d_coords, d_times and the optimizer working set. + free_bytes, _ = _cuda.current_context().get_memory_info() + threshold = free_bytes * device_mem_fraction + threshold_src = ( + f"{free_bytes / 1e9:.1f} GB free x {device_mem_fraction:g} (auto)" + ) + else: + threshold = float(mapped_memory_threshold_bytes) + threshold_src = f"{threshold / 1e9:.1f} GB (override)" + + if idx_bytes > threshold: + # Use mapped memory if index arrays are large relative to the budget. + print( + f"Large index arrays ({idx_bytes / 1e9:.1f} GB) exceed device budget " + f"({threshold_src}) — using mapped memory." + ) d_idx_i = _cuda.mapped_array_like(idx_i.astype(np.int32), wc=True) d_idx_j = _cuda.mapped_array_like(idx_j.astype(np.int32), wc=True) d_idx_i[:] = idx_i @@ -399,7 +508,12 @@ def optimize_3d_chunked_better_moving_avg_kd(n_segments, locs_nm, idx_i, idx_j, # Initial drift estimate + bounds drift_est = np.zeros(n_segments * 3) - bounds = [(-drift_max_nm * drift_max_bound_factor, drift_max_nm * drift_max_bound_factor)] * (3 * n_segments) + bounds = [ + ( + -drift_max_nm * drift_max_bound_factor, + drift_max_nm * drift_max_bound_factor, + ) + ] * (3 * n_segments) drift_est_gradient = np.inf fails = 0 @@ -411,18 +525,40 @@ def optimize_3d_chunked_better_moving_avg_kd(n_segments, locs_nm, idx_i, idx_j, # Apply boxcar smoothing to current estimate tmp = drift_est.reshape((-1, 3)) for i in range(3): - tmp[:, i] = convolve(tmp[:, i], np.ones(boxcar_width) / boxcar_width) + tmp[:, i] = convolve( + tmp[:, i], np.ones(boxcar_width) / boxcar_width + ) drift_est = tmp.flatten() # Run L-BFGS-B optimization step - result = minimize(wrapper, drift_est, method='L-BFGS-B', - args=( - d_coords, d_times, d_idx_i, d_idx_j, d_sigma, sigma_factor, d_val, d_deri, chunk_size), - jac=True, bounds=bounds, - options={'disp': display_steps, 'gtol': 1E-5, 'ftol': 1E3 * np.finfo(float).eps, - 'maxls': 40}) + result = minimize( + wrapper, + drift_est, + method="L-BFGS-B", + args=( + d_coords, + d_times, + d_idx_i, + d_idx_j, + d_sigma, + sigma_factor, + d_val, + d_deri, + chunk_size, + ), + jac=True, + bounds=bounds, + options={ + "disp": display_steps, + "gtol": 1e-5, + "ftol": 1e3 * np.finfo(float).eps, + "maxls": 40, + }, + ) itr_counter += 1 - print(f"Iteration {itr_counter}: status = {result.status}, success = {result.success}") + print( + f"Iteration {itr_counter}: status = {result.status}, success = {result.success}" + ) print(f" current sigma: {np.round(sigma_nm * sigma_factor, 2)} nm") # Update if successful @@ -431,7 +567,9 @@ def optimize_3d_chunked_better_moving_avg_kd(n_segments, locs_nm, idx_i, idx_j, print(f" drift estimate gradient: {delta}") print(f" previous gradient: {drift_est_gradient}") # Check convergence - if (delta > drift_est_gradient or sigma_nm * sigma_factor <= 1.0) and sigma_nm * sigma_factor <= target_sigma_nm: + if ( + delta > drift_est_gradient or sigma_nm * sigma_factor <= 1.0 + ) and sigma_nm * sigma_factor <= target_sigma_nm: done = True calc_time = time.time() - start_time print(f"Optimization completed in {calc_time:.2f} s") @@ -445,7 +583,9 @@ def optimize_3d_chunked_better_moving_avg_kd(n_segments, locs_nm, idx_i, idx_j, sigma_factor *= 2 print("Restarting with larger sigma_factor") if fails > 5: - raise RuntimeError("L-BFGS-B Optimization failed after multiple retries") + raise RuntimeError( + "L-BFGS-B Optimization failed after multiple retries" + ) if return_calc_time: return drift_est, time.time() - start_time, itr_counter @@ -453,72 +593,114 @@ def optimize_3d_chunked_better_moving_avg_kd(n_segments, locs_nm, idx_i, idx_j, return drift_est -def segmentation_and_pair_indices_wrapper(dataset, segmentation_var, segmentation_mode, max_drift_nm, - max_locs_per_segment, pair_indices_safety_check=False, hard_limit_pairs=None): - if not segmentation_mode == -1: # -1 is for pre-segmented data - result = segmentation_wrapper(dataset[:, -1], segmentation_var, segmentation_mode, - max_locs_per_segment, return_param_dict=True) +def segmentation_and_pair_indices_wrapper( + dataset, + segmentation_var, + segmentation_mode, + max_drift_nm, + max_locs_per_segment, + pair_indices_safety_check=False, + hard_limit_pairs=None, +): + if not segmentation_mode == -1: # -1 is for pre-segmented data + result = segmentation_wrapper( + dataset[:, -1], + segmentation_var, + segmentation_mode, + max_locs_per_segment, + return_param_dict=True, + ) else: # pre segmented data, anyway we set these values in case auto downsampling is needed segmentation_mode = 2 # dummy --> segment per frame ... - segmentation_var = 1 # dummy --> ... using 1 frame per segment + segmentation_var = 1 # dummy --> ... using 1 frame per segment if pair_indices_safety_check: - n_pairs_est = estimate_pairs(dataset[result.loc_valid, :3], max_drift_nm) - print(f"Estimated number of pairs within {max_drift_nm} nm: {n_pairs_est:,}") + n_pairs_est = estimate_pairs( + dataset[result.loc_valid, :3], max_drift_nm + ) + print( + f"Estimated number of pairs within {max_drift_nm} nm: {n_pairs_est:,}" + ) if hard_limit_pairs is not None and n_pairs_est > hard_limit_pairs: - raise RuntimeError(f"Estimated number of pairs {n_pairs_est} exceeds hard limit of {hard_limit_pairs}. " - f"Aborting to avoid crash.") + raise RuntimeError( + f"Estimated number of pairs {n_pairs_est} exceeds hard limit of {hard_limit_pairs}. " + f"Aborting to avoid crash." + ) if n_pairs_est > 5e8: - print(f"Estimated number of pairs is very large {n_pairs_est}. " - f"Automatic down-sampling is usually required above 500 mil." - f"Billions of pairs can lead to crash.") + print( + f"Estimated number of pairs is very large {n_pairs_est}. " + f"Automatic down-sampling is usually required above 500 mil." + f"Billions of pairs can lead to crash." + ) ans = input("Continue anyway? (y/n): ") - if ans.lower() != 'y': - raise RuntimeError("Aborted by user due to large estimated number of pairs.") - idx_i, idx_j, successful = pair_indices_kdtree(dataset[result.loc_valid, :3], max_drift_nm) + if ans.lower() != "y": + raise RuntimeError( + "Aborted by user due to large estimated number of pairs." + ) + idx_i, idx_j, successful = pair_indices_kdtree( + dataset[result.loc_valid, :3], max_drift_nm + ) if not successful: if max_locs_per_segment is None: - max_locs_per_segment = int(result.out_dict['locs_per_segment'].max()) + max_locs_per_segment = int( + result.out_dict["locs_per_segment"].max() + ) while not successful: max_locs_per_segment = int(max_locs_per_segment * 0.9) - print(f"Segmentation and Pairing attempt failed, automatic down-sampling active...") - print(f"Retrying segmentation with max_locs_per_segment={max_locs_per_segment}...") - result = segmentation_wrapper(dataset[:, -1], segmentation_var, segmentation_mode, - max_locs_per_segment, return_param_dict=True) + print( + f"Segmentation and Pairing attempt failed, automatic down-sampling active..." + ) + print( + f"Retrying segmentation with max_locs_per_segment={max_locs_per_segment}..." + ) + result = segmentation_wrapper( + dataset[:, -1], + segmentation_var, + segmentation_mode, + max_locs_per_segment, + return_param_dict=True, + ) sorted_dataset = dataset.copy() sorted_dataset[:, -1] = result.loc_segments sorted_dataset = sorted_dataset[result.loc_valid] - idx_i, idx_j, successful = pair_indices_kdtree(sorted_dataset[:, :3], max_drift_nm) - print(f"Segmentation and Pairing successful resulting in {result.n_segments:,} time windows with on average " - f"{int(np.median(result.out_dict['locs_per_segment']))} locs per time window. " - f"{len(idx_i):,} Pairs where found.") + idx_i, idx_j, successful = pair_indices_kdtree( + sorted_dataset[:, :3], max_drift_nm + ) + print( + f"Segmentation and Pairing successful resulting in {result.n_segments:,} time windows with on average " + f"{int(np.median(result.out_dict['locs_per_segment']))} locs per time window. " + f"{len(idx_i):,} Pairs where found." + ) sorted_dataset = dataset.copy() sorted_dataset[:, -1] = result.loc_segments sorted_dataset = sorted_dataset[result.loc_valid] return result, sorted_dataset, idx_i, idx_j - def estimate_pairs(coordinates, distance): - for i in range(len(coordinates[0])): - coordinates[:, i] -= np.min(coordinates[:, i]) - coordinates = np.array(np.floor(coordinates / distance), dtype=int) - coordinates = np.array(list(map(tuple, coordinates))) - sort_indices = np.lexsort(coordinates.T)# get the unique tuples and their counts - unique_tuples, counts = np.unique(coordinates[sort_indices], axis=0, return_counts=True) - # get the indices of the similar tuples - similar_indices = np.split(sort_indices, np.cumsum(counts[:-1])) - idx_i = [] - idx_j = [] - pair_idc_estimate = 0 - for i in range(len(similar_indices)): - n_elements = len(similar_indices[i]) - pair_idc_estimate += n_elements * (n_elements - 1) - rounded = round(pair_idc_estimate,-4) - return rounded - - -def interpolate_drift(center_frames, drift_est, frame_range, method='cubic'): + for i in range(len(coordinates[0])): + coordinates[:, i] -= np.min(coordinates[:, i]) + coordinates = np.array(np.floor(coordinates / distance), dtype=int) + coordinates = np.array(list(map(tuple, coordinates))) + sort_indices = np.lexsort( + coordinates.T + ) # get the unique tuples and their counts + unique_tuples, counts = np.unique( + coordinates[sort_indices], axis=0, return_counts=True + ) + # get the indices of the similar tuples + similar_indices = np.split(sort_indices, np.cumsum(counts[:-1])) + idx_i = [] + idx_j = [] + pair_idc_estimate = 0 + for i in range(len(similar_indices)): + n_elements = len(similar_indices[i]) + pair_idc_estimate += n_elements * (n_elements - 1) + rounded = round(pair_idc_estimate, -4) + return rounded + + +def interpolate_drift(center_frames, drift_est, frame_range, method="cubic"): """ Interpolates drift estimates to all frames using specified method. Parameters: @@ -529,11 +711,11 @@ def interpolate_drift(center_frames, drift_est, frame_range, method='cubic'): Returns: - drift_interp: np.ndarray of shape (len(frame_range), 3), interpolated drift estimates. """ - if method == 'cubic': + if method == "cubic": return _interpolate_cubic(center_frames, drift_est, frame_range) - elif method == 'catmull-rom': + elif method == "catmull-rom": return _interpolate_catmull_rom(center_frames, drift_est, frame_range) - elif method == 'linear': + elif method == "linear": drift_x = np.interp(frame_range, center_frames, drift_est[:, 0]) drift_y = np.interp(frame_range, center_frames, drift_est[:, 1]) drift_z = np.interp(frame_range, center_frames, drift_est[:, 2]) @@ -544,6 +726,7 @@ def interpolate_drift(center_frames, drift_est, frame_range, method='cubic'): def _interpolate_cubic(center_frames, drift_est, frame_range): from scipy.interpolate import CubicSpline + drift_x = CubicSpline(center_frames, drift_est[:, 0])(frame_range) drift_y = CubicSpline(center_frames, drift_est[:, 1])(frame_range) drift_z = CubicSpline(center_frames, drift_est[:, 2])(frame_range) @@ -554,19 +737,17 @@ def _interpolate_catmull_rom(center_frames, drift_est, frame_range): def catmull_rom_1d(x, y, x_interp): result = np.zeros_like(x_interp) for i in range(1, len(x) - 2): - x0, x1, x2, x3 = x[i-1], x[i], x[i+1], x[i+2] - y0, y1, y2, y3 = y[i-1], y[i], y[i+1], y[i+2] + x0, x1, x2, x3 = x[i - 1], x[i], x[i + 1], x[i + 2] + y0, y1, y2, y3 = y[i - 1], y[i], y[i + 1], y[i + 2] mask = (x_interp >= x1) & (x_interp <= x2) t = (x_interp[mask] - x1) / (x2 - x1) - result[mask] = ( - 0.5 * ( - (2 * y1) + - (-y0 + y2) * t + - (2*y0 - 5*y1 + 4*y2 - y3) * t**2 + - (-y0 + 3*y1 - 3*y2 + y3) * t**3 - ) + result[mask] = 0.5 * ( + (2 * y1) + + (-y0 + y2) * t + + (2 * y0 - 5 * y1 + 4 * y2 - y3) * t**2 + + (-y0 + 3 * y1 - 3 * y2 + y3) * t**3 ) return result @@ -574,7 +755,9 @@ def catmull_rom_1d(x, y, x_interp): x = np.asarray(center_frames) if len(x) < 4: - raise ValueError("Catmull-Rom interpolation requires at least 4 points.") + raise ValueError( + "Catmull-Rom interpolation requires at least 4 points." + ) drift_x = catmull_rom_1d(x, drift_est[:, 0], x_interp) drift_y = catmull_rom_1d(x, drift_est[:, 1], x_interp) @@ -585,9 +768,22 @@ def catmull_rom_1d(x, y, x_interp): # CUDA kernel and wrapper — only defined when CUDA is available at import time. if _CUDA_AVAILABLE: + @_cuda.jit - def cost_function_full_3d_chunked(d_locs_time, start_idx, chunk_size, d_idx_i, d_idx_j, d_sigma, - d_sigma_factor, d_val, d_val_sum, d_deri, d_locs_coords, mu): + def cost_function_full_3d_chunked( + d_locs_time, + start_idx, + chunk_size, + d_idx_i, + d_idx_j, + d_sigma, + d_sigma_factor, + d_val, + d_val_sum, + d_deri, + d_locs_coords, + mu, + ): """Compute negative-overlap cost and gradient for 3D localizations (CUDA kernel).""" tx = _cuda.threadIdx.x ty = _cuda.blockIdx.x @@ -601,52 +797,168 @@ def cost_function_full_3d_chunked(d_locs_time, start_idx, chunk_size, d_idx_i, d ti = d_locs_time[i] tj = d_locs_time[j] - dx = (d_locs_coords[i, 0] - mu[ti, 0]) - (d_locs_coords[j, 0] - mu[tj, 0]) - dy = (d_locs_coords[i, 1] - mu[ti, 1]) - (d_locs_coords[j, 1] - mu[tj, 1]) - dz = (d_locs_coords[i, 2] - mu[ti, 2]) - (d_locs_coords[j, 2] - mu[tj, 2]) + dx = (d_locs_coords[i, 0] - mu[ti, 0]) - ( + d_locs_coords[j, 0] - mu[tj, 0] + ) + dy = (d_locs_coords[i, 1] - mu[ti, 1]) - ( + d_locs_coords[j, 1] - mu[tj, 1] + ) + dz = (d_locs_coords[i, 2] - mu[ti, 2]) - ( + d_locs_coords[j, 2] - mu[tj, 2] + ) sigma_sq = (2 * d_sigma * d_sigma_factor) ** 2 diff_sq = dx * dx + dy * dy + dz * dz - val = 1 / (d_sigma * d_sigma_factor) * math.exp(-diff_sq / sigma_sq) + val = ( + 1 / (d_sigma * d_sigma_factor) * math.exp(-diff_sq / sigma_sq) + ) d_val[pos] = val - _cuda.atomic.add(d_deri, (tj, 0), 2 * val * (d_locs_coords[j, 0] - d_locs_coords[i, 0] + mu[ti, 0] - mu[tj, 0]) / sigma_sq) - _cuda.atomic.add(d_deri, (tj, 1), 2 * val * (d_locs_coords[j, 1] - d_locs_coords[i, 1] + mu[ti, 1] - mu[tj, 1]) / sigma_sq) - _cuda.atomic.add(d_deri, (tj, 2), 2 * val * (d_locs_coords[j, 2] - d_locs_coords[i, 2] + mu[ti, 2] - mu[tj, 2]) / sigma_sq) + _cuda.atomic.add( + d_deri, + (tj, 0), + 2 + * val + * ( + d_locs_coords[j, 0] + - d_locs_coords[i, 0] + + mu[ti, 0] + - mu[tj, 0] + ) + / sigma_sq, + ) + _cuda.atomic.add( + d_deri, + (tj, 1), + 2 + * val + * ( + d_locs_coords[j, 1] + - d_locs_coords[i, 1] + + mu[ti, 1] + - mu[tj, 1] + ) + / sigma_sq, + ) + _cuda.atomic.add( + d_deri, + (tj, 2), + 2 + * val + * ( + d_locs_coords[j, 2] + - d_locs_coords[i, 2] + + mu[ti, 2] + - mu[tj, 2] + ) + / sigma_sq, + ) - _cuda.atomic.add(d_deri, (ti, 0), 2 * val * (d_locs_coords[i, 0] - d_locs_coords[j, 0] + mu[tj, 0] - mu[ti, 0]) / sigma_sq) - _cuda.atomic.add(d_deri, (ti, 1), 2 * val * (d_locs_coords[i, 1] - d_locs_coords[j, 1] + mu[tj, 1] - mu[ti, 1]) / sigma_sq) - _cuda.atomic.add(d_deri, (ti, 2), 2 * val * (d_locs_coords[i, 2] - d_locs_coords[j, 2] + mu[tj, 2] - mu[ti, 2]) / sigma_sq) + _cuda.atomic.add( + d_deri, + (ti, 0), + 2 + * val + * ( + d_locs_coords[i, 0] + - d_locs_coords[j, 0] + + mu[tj, 0] + - mu[ti, 0] + ) + / sigma_sq, + ) + _cuda.atomic.add( + d_deri, + (ti, 1), + 2 + * val + * ( + d_locs_coords[i, 1] + - d_locs_coords[j, 1] + + mu[tj, 1] + - mu[ti, 1] + ) + / sigma_sq, + ) + _cuda.atomic.add( + d_deri, + (ti, 2), + 2 + * val + * ( + d_locs_coords[i, 2] + - d_locs_coords[j, 2] + + mu[tj, 2] + - mu[ti, 2] + ) + / sigma_sq, + ) _cuda.atomic.add(d_val_sum, 0, val) d_val[pos] = 0 - def cuda_wrapper_chunked(mu, d_locs_coords, d_locs_time, d_idx_i, d_idx_j, d_sigma, d_sigma_factor, d_val, - d_deri, chunk_size): + def cuda_wrapper_chunked( + mu, + d_locs_coords, + d_locs_time, + d_idx_i, + d_idx_j, + d_sigma, + d_sigma_factor, + d_val, + d_deri, + chunk_size, + ): """Interface between Python optimizer and the CUDA kernel (chunked to manage memory).""" val_total = 0 d_val_sum = _cuda.to_device(np.zeros(1, dtype=np.float64)) - mu_dev = _cuda.to_device(np.asarray(mu.reshape(int(mu.size / 3), 3), dtype=np.float64)) + mu_dev = _cuda.to_device( + np.asarray(mu.reshape(int(mu.size / 3), 3), dtype=np.float64) + ) n_chunks = int(np.ceil(d_idx_i.size / chunk_size)) threadsperblock = 128 for i in range(n_chunks - 1): idc_start = i * chunk_size - blockspergrid = (chunk_size + (threadsperblock - 1)) // threadsperblock + blockspergrid = ( + chunk_size + (threadsperblock - 1) + ) // threadsperblock cost_function_full_3d_chunked[blockspergrid, threadsperblock]( - d_locs_time, idc_start, chunk_size, d_idx_i, d_idx_j, - d_sigma, d_sigma_factor, d_val, d_val_sum, d_deri, d_locs_coords, mu_dev + d_locs_time, + idc_start, + chunk_size, + d_idx_i, + d_idx_j, + d_sigma, + d_sigma_factor, + d_val, + d_val_sum, + d_deri, + d_locs_coords, + mu_dev, ) val_total += d_val_sum.copy_to_host() # Final chunk n_remaining = d_idx_i.size - (n_chunks - 1) * chunk_size idc_start = (n_chunks - 1) * chunk_size - blockspergrid = (n_remaining + (threadsperblock - 1)) // threadsperblock + blockspergrid = ( + n_remaining + (threadsperblock - 1) + ) // threadsperblock cost_function_full_3d_chunked[blockspergrid, threadsperblock]( - d_locs_time, idc_start, n_remaining, d_idx_i, d_idx_j, - d_sigma, d_sigma_factor, d_val, d_val_sum, d_deri, d_locs_coords, mu_dev + d_locs_time, + idc_start, + n_remaining, + d_idx_i, + d_idx_j, + d_sigma, + d_sigma_factor, + d_val, + d_val_sum, + d_deri, + d_locs_coords, + mu_dev, ) val_total += d_val_sum.copy_to_host() deri = d_deri.copy_to_host() @@ -658,6 +970,7 @@ def cuda_wrapper_chunked(mu, d_locs_coords, d_locs_time, d_idx_i, d_idx_j, d_sig @dataclass class SegmentationResult: """Container for segmentation results.""" + loc_segments: np.ndarray loc_valid: np.ndarray center_frames: np.ndarray @@ -669,17 +982,22 @@ def _group_by_frame(loc_frames: np.ndarray): """Returns a dict {frame_number: indices_in_loc_frames} efficiently.""" sort_idx = np.argsort(loc_frames) sorted_frames = loc_frames[sort_idx] - unique_frames, start_idx, counts = np.unique(sorted_frames, return_index=True, return_counts=True) + unique_frames, start_idx, counts = np.unique( + sorted_frames, return_index=True, return_counts=True + ) frame_to_indices = { - frame: sort_idx[start:start + count] + frame: sort_idx[start : start + count] for frame, start, count in zip(unique_frames, start_idx, counts) } return unique_frames, frame_to_indices -def segment_by_num_locs_per_window(loc_frames: np.ndarray, min_n_locs_per_window: int, - max_locs_per_segment = None, - return_param_dict: bool = False) -> SegmentationResult: +def segment_by_num_locs_per_window( + loc_frames: np.ndarray, + min_n_locs_per_window: int, + max_locs_per_segment=None, + return_param_dict: bool = False, +) -> SegmentationResult: """ Segments by collecting a minimum number of localizations per window. Once the threshold is met and enough locs remain, a new segment is created. @@ -698,8 +1016,12 @@ def segment_by_num_locs_per_window(loc_frames: np.ndarray, min_n_locs_per_window loc_frames = np.asarray(loc_frames, dtype=int) n_locs = len(loc_frames) - if max_locs_per_segment is not None and max_locs_per_segment < 1: # downsampling in percentage - max_locs_per_segment = int(min_n_locs_per_window * max_locs_per_segment) + if ( + max_locs_per_segment is not None and max_locs_per_segment < 1 + ): # downsampling in percentage + max_locs_per_segment = int( + min_n_locs_per_window * max_locs_per_segment + ) unique_frames, frame_to_indices = _group_by_frame(loc_frames) loc_segments = np.full(n_locs, -1, dtype=int) # Default to -1 for safety @@ -711,13 +1033,22 @@ def segment_by_num_locs_per_window(loc_frames: np.ndarray, min_n_locs_per_window for i, frame in enumerate(unique_frames): indices = frame_to_indices[frame] n_locs_this_frame = len(indices) - remaining_locs = n_locs - (len(current_segment_indices) + n_locs_this_frame + np.sum(locs_per_segment)) + remaining_locs = n_locs - ( + len(current_segment_indices) + + n_locs_this_frame + + np.sum(locs_per_segment) + ) # Add frame to current segment if: # - It fills the current segment to threshold # - AND there are enough locs left for another segment (or it's the last frame) - if (n_locs_in_current_segment + n_locs_this_frame >= min_n_locs_per_window) and \ - (remaining_locs >= min_n_locs_per_window or i == len(unique_frames) - 1): + if ( + n_locs_in_current_segment + n_locs_this_frame + >= min_n_locs_per_window + ) and ( + remaining_locs >= min_n_locs_per_window + or i == len(unique_frames) - 1 + ): current_segment_indices.extend(indices) n_locs_in_current_segment += n_locs_this_frame @@ -740,8 +1071,13 @@ def segment_by_num_locs_per_window(loc_frames: np.ndarray, min_n_locs_per_window for i in range(n_segments): segment_indices = np.where(loc_segments == i)[0] - if max_locs_per_segment and len(segment_indices) > max_locs_per_segment: - selected = np.random.choice(segment_indices, max_locs_per_segment, replace=False) + if ( + max_locs_per_segment + and len(segment_indices) > max_locs_per_segment + ): + selected = np.random.choice( + segment_indices, max_locs_per_segment, replace=False + ) else: selected = segment_indices loc_valid[selected] = True @@ -761,15 +1097,20 @@ def segment_by_num_locs_per_window(loc_frames: np.ndarray, min_n_locs_per_window "n_locs": n_locs, "n_locs_valid": n_locs_valid, "n_locs_invalid": n_locs - n_locs_valid, - "center_frames": center_frames + "center_frames": center_frames, } - return SegmentationResult(loc_segments, loc_valid, center_frames, n_segments, out_dict) + return SegmentationResult( + loc_segments, loc_valid, center_frames, n_segments, out_dict + ) -def segment_by_frame_windows(loc_frames: np.ndarray, n_frames_per_window: int, - max_locs_per_segment = None, - return_param_dict: bool = False) -> SegmentationResult: +def segment_by_frame_windows( + loc_frames: np.ndarray, + n_frames_per_window: int, + max_locs_per_segment=None, + return_param_dict: bool = False, +) -> SegmentationResult: """ Splits localization data into fixed-size windows of N frames. All localizations in those frames are grouped into one segment. @@ -779,7 +1120,9 @@ def segment_by_frame_windows(loc_frames: np.ndarray, n_frames_per_window: int, n_locs = len(loc_frames) n_segments = int(np.ceil(len(frames) / n_frames_per_window)) - if max_locs_per_segment is not None and max_locs_per_segment < 1: # downsampling in percentage + if ( + max_locs_per_segment is not None and max_locs_per_segment < 1 + ): # downsampling in percentage max_locs_per_segment = n_locs / n_segments * max_locs_per_segment loc_segments = np.zeros(n_locs, dtype=int) @@ -788,8 +1131,16 @@ def segment_by_frame_windows(loc_frames: np.ndarray, n_frames_per_window: int, start_frames, end_frames, locs_per_segment = [], [], [] for i in range(n_segments): - frame_window = frames[i * n_frames_per_window:(i + 1) * n_frames_per_window] - indices = np.concatenate([frame_to_indices[frame] for frame in frame_window if frame in frame_to_indices]) + frame_window = frames[ + i * n_frames_per_window : (i + 1) * n_frames_per_window + ] + indices = np.concatenate( + [ + frame_to_indices[frame] + for frame in frame_window + if frame in frame_to_indices + ] + ) if len(indices) == 0: continue loc_segments[indices] = i @@ -799,7 +1150,13 @@ def segment_by_frame_windows(loc_frames: np.ndarray, n_frames_per_window: int, locs_per_segment.append(len(indices)) if max_locs_per_segment and len(indices) > max_locs_per_segment: mask = np.ones(len(indices), dtype=bool) - mask[np.random.choice(len(indices), len(indices) - max_locs_per_segment, replace=False)] = False + mask[ + np.random.choice( + len(indices), + len(indices) - max_locs_per_segment, + replace=False, + ) + ] = False loc_valid[indices[~mask]] = False out_dict = None @@ -815,40 +1172,69 @@ def segment_by_frame_windows(loc_frames: np.ndarray, n_frames_per_window: int, "n_locs": n_locs, "n_locs_valid": n_locs_valid, "n_locs_invalid": n_locs - n_locs_valid, - "center_frames": center_frames + "center_frames": center_frames, } - return SegmentationResult(loc_segments, loc_valid, center_frames, n_segments, out_dict) + return SegmentationResult( + loc_segments, loc_valid, center_frames, n_segments, out_dict + ) -def segment_by_num_windows(loc_frames: np.ndarray, n_windows: int, max_locs_per_segment = None, - return_param_dict: bool = False) -> SegmentationResult: +def segment_by_num_windows( + loc_frames: np.ndarray, + n_windows: int, + max_locs_per_segment=None, + return_param_dict: bool = False, +) -> SegmentationResult: """ - Converts number of windows into an equivalent minimum locs per window, - then calls `segment_by_num_locs_per_window`. + Converts number of windows into an equivalent minimum locs per window, + then calls `segment_by_num_locs_per_window`. """ n_locs = len(loc_frames) n_locs_per_window = int(np.ceil(n_locs / n_windows)) - if max_locs_per_segment is not None and max_locs_per_segment < 1: # downsampling in percentage + if ( + max_locs_per_segment is not None and max_locs_per_segment < 1 + ): # downsampling in percentage max_locs_per_segment = int(n_locs_per_window * max_locs_per_segment) - return segment_by_num_locs_per_window(loc_frames, n_locs_per_window, max_locs_per_segment, return_param_dict) + return segment_by_num_locs_per_window( + loc_frames, n_locs_per_window, max_locs_per_segment, return_param_dict + ) -def segmentation_wrapper(loc_frames: np.ndarray, segmentation_var: int, segmentation_mode: int = 2, - max_locs_per_segment = None, - return_param_dict: bool = False) -> SegmentationResult: +def segmentation_wrapper( + loc_frames: np.ndarray, + segmentation_var: int, + segmentation_mode: int = 2, + max_locs_per_segment=None, + return_param_dict: bool = False, +) -> SegmentationResult: """ - Dispatch function that selects segmentation method: - 0 → fixed number of windows - 1 → fixed number of localizations per window - 2 → fixed number of frames per window (default) + Dispatch function that selects segmentation method: + 0 → fixed number of windows + 1 → fixed number of localizations per window + 2 → fixed number of frames per window (default) """ if segmentation_mode == 0: - return segment_by_num_windows(loc_frames, segmentation_var, max_locs_per_segment, return_param_dict) + return segment_by_num_windows( + loc_frames, + segmentation_var, + max_locs_per_segment, + return_param_dict, + ) elif segmentation_mode == 1: - return segment_by_num_locs_per_window(loc_frames, segmentation_var, max_locs_per_segment, return_param_dict) + return segment_by_num_locs_per_window( + loc_frames, + segmentation_var, + max_locs_per_segment, + return_param_dict, + ) else: - return segment_by_frame_windows(loc_frames, segmentation_var, max_locs_per_segment, return_param_dict) + return segment_by_frame_windows( + loc_frames, + segmentation_var, + max_locs_per_segment, + return_param_dict, + ) def pair_indices_kdtree(coordinates, distance): @@ -863,8 +1249,12 @@ def pair_indices_kdtree(coordinates, distance): """ tree = cKDTree(coordinates) try: - pairs = tree.query_pairs(r=distance, output_type='ndarray') + pairs = tree.query_pairs(r=distance, output_type="ndarray") except MemoryError: print("[pair_indices_kdtree] MemoryError encountered") return [], [], False - return np.ascontiguousarray(pairs[:, 0], dtype=np.int32), np.ascontiguousarray(pairs[:, 1], dtype=np.int32), True \ No newline at end of file + return ( + np.ascontiguousarray(pairs[:, 0], dtype=np.int32), + np.ascontiguousarray(pairs[:, 1], dtype=np.int32), + True, + ) diff --git a/picasso/gui/render.py b/picasso/gui/render.py index b474b52b..42ebb72e 100644 --- a/picasso/gui/render.py +++ b/picasso/gui/render.py @@ -2279,8 +2279,9 @@ def __init__(self, window: QtWidgets.QMainWindow) -> None: vbox.addLayout(grid) self.buttons = QtWidgets.QDialogButtonBox( - QtWidgets.QDialogButtonBox.Ok | QtWidgets.QDialogButtonBox.Cancel, - QtCore.Qt.Horizontal, + QtWidgets.QDialogButtonBox.StandardButton.Ok + | QtWidgets.QDialogButtonBox.StandardButton.Cancel, + QtCore.Qt.Orientation.Horizontal, self, ) vbox.addWidget(self.buttons) @@ -2293,13 +2294,13 @@ def getParams( ) -> tuple[dict, bool]: """Create the dialog and return the requested COMET parameters.""" dialog = COMETDialog(parent) - result = dialog.exec_() + result = dialog.exec() params = { "locs_per_segment": dialog.locs_per_segment.value(), "max_drift_nm": dialog.max_drift_nm.value(), "max_locs_per_segment": dialog.max_locs_per_segment.value(), } - return params, result == QtWidgets.QDialog.Accepted + return params, result == QtWidgets.QDialog.DialogCode.Accepted class AIMDialog(lib.Dialog): @@ -11084,7 +11085,7 @@ def undrift_comet(self) -> None: Bates M. biorxiv, 2026.""" channel = self.get_channel("Undrift by COMET") if channel is not None: - locs = self.all_locs[channel] + locs = self.locs[channel] info = self.infos[channel] params, ok = COMETDialog.getParams(self.window) diff --git a/pyproject.toml b/pyproject.toml index 1e033b1b..4be6e658 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -54,6 +54,15 @@ dependencies = [ ] [project.optional-dependencies] +cuda11 = [ + "numba-cuda[cu11]" +] +cuda12 = [ + "numba-cuda[cu12]" +] +cuda13 = [ + "numba-cuda[cu13]" +] dev = [ "build", "black", @@ -86,6 +95,33 @@ installer = [ "PyImarisWriter==0.7.0; sys_platform=='win32'", "hdf5plugin==6.0.0; sys_platform=='win32'" ] +installer_gpu = [ + "PyQt6==6.11.0", + "numpy==2.4.4", + "matplotlib==3.10.8", + "h5py==3.16.0", + "numba==0.65.0", + "scipy==1.17.1", + "pandas==2.3.3", + "tables==3.11.1", + "pyyaml==6.0.3", + "scikit-learn==1.7.2", + "tqdm==4.67.3", + "streamlit==1.56.0", + "nd2==0.11.3", + "sqlalchemy==2.0.49", + "plotly-express==0.4.1", + "watchdog==6.0.0", + "psutil==7.2.2", + "playsound3==3.3.1", + "imageio==2.37.3", + "imageio-ffmpeg==0.6.0", + "tifffile==2026.4.11", + "pyinstaller==6.19.0", + "PyImarisWriter==0.7.0; sys_platform=='win32'", + "hdf5plugin==6.0.0; sys_platform=='win32'", + "numba-cuda[cu12]==0.30.2" +] [project.urls] Homepage = "https://github.com/jungmannlab/picasso" From 6c68f1b132bd48e6185c78d6d1e63801d969e0da Mon Sep 17 00:00:00 2001 From: rafalkowalewski1 Date: Sun, 31 May 2026 18:55:27 +0200 Subject: [PATCH 09/19] remove comet temporarily --- changelog.md | 2 - picasso/ext/comet.py | 1260 ---------------------- picasso/gui/render.py | 126 --- readme.rst | 1 - release/one_click_macos_gui/readme.txt | 1 - release/one_click_windows_gui/readme.txt | 1 - tests/test_comet.py | 572 ---------- 7 files changed, 1963 deletions(-) delete mode 100644 picasso/ext/comet.py delete mode 100644 tests/test_comet.py diff --git a/changelog.md b/changelog.md index 49cc37e4..f104bd0d 100644 --- a/changelog.md +++ b/changelog.md @@ -2,8 +2,6 @@ Last change: 29-MAY-2026 CEST -- COMET ([DOI: 10.64898/2026.03.27.714864](https://doi.org/10.64898/2026.03.27.714864)) integrated (#671), big thanks to @LREIN663 - ## 0.10.1 #### Localize diff --git a/picasso/ext/comet.py b/picasso/ext/comet.py deleted file mode 100644 index c172b1a0..00000000 --- a/picasso/ext/comet.py +++ /dev/null @@ -1,1260 +0,0 @@ -import math -import time -from dataclasses import dataclass -from typing import Dict, Optional - -import matplotlib.pyplot as plt -import numpy as np -import pandas as pd -from scipy.ndimage import convolve -from scipy.optimize import minimize -from scipy.spatial import cKDTree - -from .. import lib, __version__ - -try: - from numba import cuda as _cuda - - _CUDA_AVAILABLE = _cuda.is_available() -except Exception: - _cuda = None - _CUDA_AVAILABLE = False - - -def comet( - locs: pd.DataFrame, - info: list[dict], - locs_per_segment: int, - max_drift_nm: float, - max_locs_per_segment: int | None = None, - *, - initial_sigma_nm: float | None = None, - target_sigma_nm: float = 10.0, - boxcar_width: int = 3, - drift_max_bound_factor: float = 2.0, - interpolation_method: str = "cubic", - mode: str = "cuda", - display: bool = False, - progress=None, # currently unused -) -> tuple[pd.DataFrame, list[dict], pd.DataFrame]: - """Apply COMET undrifting to the localizations. - - Parameters - ---------- - locs : pd.DataFrame - Localization table. Must contain columns 'x', 'y', 'frame'. - Optional: 'z'. - x/y are expected in camera pixels, z in the native z-unit used by Picasso. - info : list of dict - Localization metadata. - locs_per_segment : int - Target number of localizations per segment for COMET segmentation. - max_drift_nm : float - Maximum expected drift in nm. - max_locs_per_segment : int or None, optional - Optional cap for localizations per segment. If negative, treated as None. - initial_sigma_nm : float or None, optional - Initial Gaussian sigma for COMET. If None, defaults to max_drift_nm / 3. - target_sigma_nm : float, optional - Final Gaussian sigma for refinement. - boxcar_width : int, optional - Smoothing width over segments. - drift_max_bound_factor : float, optional - Bound factor for optimizer. - interpolation_method : str, optional - Drift interpolation method. - mode : str, optional - Backend mode. Currently only ``"cuda"`` is supported, which requires - a CUDA-capable NVIDIA GPU and numba installed with CUDA support. - Raises ``RuntimeError`` when CUDA is not available. - display : bool, optional - Whether to show diagnostic plots. - progress : optional - Placeholder for API compatibility with AIM. - - Returns - ------- - locs : pd.DataFrame - Undrift-corrected localizations. - new_info : list[dict] - Updated metadata. - drift : pd.DataFrame - Per-frame drift table with columns x, y, and optionally z. - Drift is returned in the same units as locs: - x/y in pixels, z in the original z-unit. - """ - # Accept numpy structured arrays as well as pandas DataFrames. - if isinstance(locs, np.ndarray) and locs.dtype.names is not None: - locs = pd.DataFrame({name: locs[name] for name in locs.dtype.names}) - - locs = locs.copy() - - if not _CUDA_AVAILABLE: - raise RuntimeError( - "COMET requires a CUDA-capable NVIDIA GPU and numba with CUDA support, " - "which are not available on this system. " - "Please use a different drift-correction method (e.g. AIM or RCC)." - ) - - if max_locs_per_segment is not None and max_locs_per_segment < 0: - max_locs_per_segment = None - - pixelsize_nm = lib.get_from_metadata(info, "Pixelsize", raise_error=True) - - required_cols = {"x", "y", "frame"} - missing = required_cols - set(locs.columns) - if missing: - raise KeyError( - f"Missing required columns for COMET: {sorted(missing)}" - ) - - has_z = "z" in locs.columns - - # Normalize frames so they start at 0 for internal indexing. - # This is important because COMET later indexes drift by frame number. - frame = locs["frame"].to_numpy(dtype=np.int64) - frame0 = frame - frame.min() - - # Build COMET input array in nm - dataset = np.zeros((len(locs), 4), dtype=np.float64) - dataset[:, 0] = locs["x"].to_numpy(dtype=np.float64) * pixelsize_nm - dataset[:, 1] = locs["y"].to_numpy(dtype=np.float64) * pixelsize_nm - dataset[:, 2] = locs["z"].to_numpy(dtype=np.float64) if has_z else 0.0 - dataset[:, 3] = frame0 - - if initial_sigma_nm is None: - initial_sigma_nm = max_drift_nm / 3 - - drift_nm_with_frame = comet_run_kd( - dataset=dataset, - segmentation_mode=1, # fixed locs per segment - segmentation_var=locs_per_segment, - max_locs_per_segment=max_locs_per_segment, - initial_sigma_nm=initial_sigma_nm, - gt_drift=None, - display=display, - return_corrected_locs=False, - max_drift_nm=max_drift_nm, - target_sigma_nm=target_sigma_nm, - boxcar_width=boxcar_width, - drift_max_bound_factor=drift_max_bound_factor, - interpolation_method=interpolation_method, - mode=mode, - min_max_frames=(int(frame0.min()), int(frame0.max())), - ) - - # drift_nm_with_frame columns: dx_nm, dy_nm, dz_nm, frame0 - drift_nm = drift_nm_with_frame[:, :3] - - # Apply drift back to dataframe in original units - frame_idx = frame0.astype(np.int64) - locs["x"] = ( - locs["x"].to_numpy(dtype=np.float64) - - drift_nm[frame_idx, 0] / pixelsize_nm - ) - locs["y"] = ( - locs["y"].to_numpy(dtype=np.float64) - - drift_nm[frame_idx, 1] / pixelsize_nm - ) - if has_z: - locs["z"] = ( - locs["z"].to_numpy(dtype=np.float64) - drift_nm[frame_idx, 2] - ) - - # Build drift dataframe in Picasso-style units - drift_dict = { - "x": (drift_nm[:, 0] / pixelsize_nm).astype("float32"), - "y": (drift_nm[:, 1] / pixelsize_nm).astype("float32"), - } - if has_z: - drift_dict["z"] = drift_nm[:, 2].astype("float32") - - drift = pd.DataFrame(drift_dict) - - new_info_entry = { - "Generated by": f"Picasso v{__version__} COMET", - "Segmentation mode": "localizations per segment", - "Localizations per segment": locs_per_segment, - "Maximum drift (nm)": max_drift_nm, - "Initial sigma (nm)": float(initial_sigma_nm), - "Target sigma (nm)": float(target_sigma_nm), - "Boxcar width": int(boxcar_width), - "Max localizations per segment": ( - None if max_locs_per_segment is None else int(max_locs_per_segment) - ), - "Interpolation method": interpolation_method, - "Backend": mode, - } - new_info = info + [new_info_entry] - - return locs, new_info, drift - - -def comet_run_kd( - dataset, - segmentation_mode, - segmentation_var, - max_locs_per_segment=None, - initial_sigma_nm=None, - gt_drift=None, - display=False, - return_corrected_locs=False, - max_drift_nm=300, - target_sigma_nm=1, - boxcar_width=1, - drift_max_bound_factor=2, - save_intermediate_results=False, - interpolation_method="cubic", - mode="cuda", - min_max_frames=None, - pair_indices_safety_check=False, - mapped_memory_threshold_bytes=None, - device_mem_fraction=0.5, -): - """ - Run COMET drift correction end-to-end. - - Pipeline: temporal segmentation -> KD-tree neighbor pairs -> cost optimization (L-BFGS-B) - -> optional temporal smoothing -> spline interpolation to per-frame drift -> (optional) subtract drift. - - Parameters - ---------- - dataset : ndarray of shape (N, 4) - Localization array with columns [x_nm, y_nm, z_nm, frame]. Units in nm; frame is int. - For 2D CSVs, insert a zero z column to get (N, 4). - segmentation_mode : {0, 1, 2} - Temporal segmentation mode: - 0 = number of windows (choose S directly), - 1 = localizations per window (accumulate frames until >= X locs), - 2 = fixed frame window size (default). - segmentation_var : int - Mode-dependent value (S, locs per window, or frames per window). - initial_sigma_nm : float, default=100 - Initial Gaussian length scale for the overlap kernel (coarse scale). - target_sigma_nm : float, default=1 - Target (final) Gaussian length scale for fine refinement. - max_drift_nm : float or None, default=None - Pair radius in nm used for neighbor search. If None, uses 3 * initial_sigma_nm. - drift_max_bound_factor : float, default=1.0 - Multiplicative factor for L-BFGS-B box bounds around +-max_drift. - boxcar_width : int, default=1 - Temporal smoothing width (segments) applied to the estimated drift between optimizer steps. - interpolation_method : {"cubic", "catmull-rom"}, default="cubic" - Spline used to convert per-segment drift to per-frame drift. - max_locs_per_segment : int or None, default=None - Optional downsampling cap per segment (to control memory/time). - mode : str, "cuda" - only Cuda in this version - return_corrected_locs : bool, default=False - If True, also return drift-corrected localizations. - mapped_memory_threshold_bytes : int or None, default=None - Byte budget for the pair-index arrays on the GPU. If exceeded, the - arrays are staged in mapped host memory instead of VRAM. None - auto-detects the budget from free VRAM (see `device_mem_fraction`). - Set explicitly to force deterministic behavior across machines/GPUs. - device_mem_fraction : float, default=0.5 - Fraction of free VRAM used as the index-array budget when - `mapped_memory_threshold_bytes` is None. Lower on small GPUs to avoid - out-of-memory; raise toward 1.0 to keep more in faster device memory. - - Returns - ------- - drift_interp_with_frames : ndarray of shape (F, 4) - Per-frame drift with columns [dx_nm, dy_nm, dz_nm, frame]. - corrected_locs : ndarray of shape (N, 4), optional - Only if return_corrected_locs=True. Columns are [x_nm, y_nm, z_nm, segment_id]. - """ - - loc_frames = dataset[:, -1] - if min_max_frames is None: - min_max_frames = (loc_frames.min(), loc_frames.max()) - - # Segment the dataset based on frame numbers into time windows - - result, sorted_dataset, idx_i, idx_j = ( - segmentation_and_pair_indices_wrapper( - dataset, - segmentation_var, - segmentation_mode, - max_drift_nm, - max_locs_per_segment, - pair_indices_safety_check=pair_indices_safety_check, - ) - ) - - # Set default initial sigma if not provided - if initial_sigma_nm is None: - initial_sigma_nm = max_drift_nm // 3 - - # Run drift optimization - t0 = time.time() - drift_est = optimize_3d_chunked_better_moving_avg_kd( - result.n_segments, - sorted_dataset, - idx_i, - idx_j, - sigma_nm=initial_sigma_nm, - target_sigma_nm=target_sigma_nm, - drift_max_nm=max_drift_nm, - drift_max_bound_factor=drift_max_bound_factor, - display_steps=display, - boxcar_width=boxcar_width, - segmentation_result=result, - mode=mode, - mapped_memory_threshold_bytes=mapped_memory_threshold_bytes, - device_mem_fraction=device_mem_fraction, - ) - elapsed = time.time() - t0 - - # Reshape and interpolate drift across all frames - drift_est = drift_est.reshape((result.n_segments, 3)) - vld_tp = np.where(~np.isnan(drift_est[:, 0])) - - frame_interp = np.arange(0, min_max_frames[1] + 1, dtype=int) - drift_interp = interpolate_drift( - result.center_frames[vld_tp], - drift_est[vld_tp], - frame_interp, - method=interpolation_method, - ) - drift_interp_with_frames = np.hstack( - (drift_interp, frame_interp[:, np.newaxis]) - ) - - # Apply drift correction to original localizations - for i in range(3): - dataset[:, i] = ( - dataset[:, i] - drift_interp[dataset[:, -1].astype(int), i] - ) - - # Optionally show estimated drift curve - if display: - print(f"Drift estimation completed in {elapsed:.2f} seconds.") - plt.figure() - plt.plot(frame_interp, drift_interp) - plt.title("Estimated Drift") - plt.xlabel("Frames") - plt.ylabel("Drift (nm)") - plt.legend(["X", "Y", "Z"]) - plt.show() - - # Optional GT comparison plot - if display and gt_drift is not None: - fig, ax = plt.subplots(3, 1) - ax[0].plot(gt_drift[:, 3], gt_drift[:, 0], label="GT Drift X") - ax[1].plot(gt_drift[:, 3], gt_drift[:, 1], label="GT Drift Y") - ax[2].plot(gt_drift[:, 3], gt_drift[:, 2], label="GT Drift Z") - ax[0].plot( - frame_interp, - drift_interp[:, 0], - label="Estimated Drift X", - linestyle="--", - ) - ax[1].plot( - frame_interp, - drift_interp[:, 1], - label="Estimated Drift Y", - linestyle="--", - ) - ax[2].plot( - frame_interp, - drift_interp[:, 2], - label="Estimated Drift Z", - linestyle="--", - ) - ax[1].set_title("Ground Truth vs Estimated Drift") - ax[1].set_xlabel("Frames") - ax[0].set_ylabel("Drift (nm)") - plt.legend() - plt.show() - - # Return corrected locs + drift - if return_corrected_locs: - return drift_interp_with_frames, dataset - else: - return drift_interp_with_frames - - -def optimize_3d_chunked_better_moving_avg_kd( - n_segments, - locs_nm, - idx_i, - idx_j, - sigma_nm=30, - drift_max_nm=300, - target_sigma_nm=30, - display_steps=False, - boxcar_width=3, - drift_max_bound_factor=2, - segmentation_result=None, - mode="cuda", - return_calc_time=False, - mapped_memory_threshold_bytes=None, - device_mem_fraction=0.5, -): - """ - Estimate per-segment drift (mu) by minimizing the negative Gaussian-overlap cost - with an L-BFGS-B optimizer and a coarse-to-fine schedule on sigma. - - The routine operates on temporally segmented localizations and reuses a static - neighbor graph (pairs within drift_max_nm). Between optimizer steps, a moving - average (boxcar) can be applied to mu as temporal regularization. Sigma is - reduced iteratively from `sigma_nm` toward `target_sigma_nm` for robust - convergence. - - Parameters - ---------- - n_segments : int - Number of temporal segments S (0..S-1). One 3D drift vector is estimated per segment. - locs_nm : ndarray of shape (M, 3) - localizations in nanometers, columns [x, y, z]; - sigma_nm : float, default=30 - Initial Gaussian width (coarse scale) for the overlap kernel. - drift_max_nm : float, default=300 - Maximum expected drift (nm). Also used as the radius for neighbor pairs and - as the L-BFGS-B bound scale (see `drift_max_bound_factor`). - target_sigma_nm : float, default=30 - Target / final Gaussian width (fine scale). The optimizer reduces sigma toward - this value over iterations. - display_steps : bool, default=False - If True, print or log intermediate progress per iteration/scale. - - boxcar_width : int, default=3 - Temporal smoothing width (in segments) for a moving average applied to mu between steps. - Use 0 or 1 to disable smoothing. - drift_max_bound_factor : float, default=2 - Multiplier for L-BFGS-B bounds around +- drift_max_nm to keep updates physically reasonable. - segmentation_result : object or None, default=None - Segmentation info and metadata. Expected to provide, at minimum: - - segment IDs per localization - - center frames per segment - - any additional structures required by backend (e.g., pair indices) - If None, pairs/ids may be built internally depending on implementation. - mode : str, default=cuda - If True, use the CPU backend; otherwise try GPU (CUDA) and fall back to CPU if unavailable. - return_calc_time : bool, default=False - If True, also return the total computation time in seconds. - mapped_memory_threshold_bytes : int or None, default=None - Size budget (in bytes) for the pair-index arrays on the GPU. If the two - int32 index arrays together exceed this, they are staged in mapped host - memory instead of device VRAM. If None, the budget is auto-detected from - currently-free VRAM (see `device_mem_fraction`). Provide an explicit value - to force deterministic behavior regardless of GPU state. - device_mem_fraction : float, default=0.5 - Fraction of currently-free VRAM used as the index-array budget when - `mapped_memory_threshold_bytes` is None. Lower it if you hit out-of-memory - errors (more aggressive fallback to mapped memory); raise it toward 1.0 to - keep larger arrays in faster device memory. Ignored when an explicit - threshold is given. - - Returns - ------- - mu : ndarray of shape (S, 3) - Estimated per-segment drift (dx, dy, dz) in nanometers. - calc_time_s : float, optional - Only when `return_calc_time=True`. Wall-clock time for the optimization. - """ - - if segmentation_result is None: - segmentation_result = {} - intermediate_results_filehandle = None - sigma_factor = 1.0 - - # Extract coordinate + time arrays, convert to device if CUDA - coords = locs_nm[:, :3].astype(np.float32).copy() - times = locs_nm[:, 3].astype(np.int32).copy() - - chunk_size = int(1e8) # 1E7 - - quality_control = mode == "torch_qc" or mode == "cuda_qc" - - d_coords = _cuda.to_device(coords) - d_times = _cuda.to_device(times) - # The pair-index arrays can dwarf available VRAM on large datasets. Decide - # whether to stage them in device memory (fast) or fall back to mapped host - # memory (handles arrays larger than VRAM, slower per access). The threshold - # auto-adapts to the current GPU unless an explicit override is given. - idx_bytes = len(idx_i) * 4 * 2 # two int32 arrays (idx_i + idx_j) - if mapped_memory_threshold_bytes is None: - # Budget a fraction of currently-free VRAM, leaving headroom for d_val, - # d_coords, d_times and the optimizer working set. - free_bytes, _ = _cuda.current_context().get_memory_info() - threshold = free_bytes * device_mem_fraction - threshold_src = ( - f"{free_bytes / 1e9:.1f} GB free x {device_mem_fraction:g} (auto)" - ) - else: - threshold = float(mapped_memory_threshold_bytes) - threshold_src = f"{threshold / 1e9:.1f} GB (override)" - - if idx_bytes > threshold: - # Use mapped memory if index arrays are large relative to the budget. - print( - f"Large index arrays ({idx_bytes / 1e9:.1f} GB) exceed device budget " - f"({threshold_src}) — using mapped memory." - ) - d_idx_i = _cuda.mapped_array_like(idx_i.astype(np.int32), wc=True) - d_idx_j = _cuda.mapped_array_like(idx_j.astype(np.int32), wc=True) - d_idx_i[:] = idx_i - d_idx_j[:] = idx_j - else: - d_idx_i = _cuda.to_device(idx_i.astype(np.int32)) - d_idx_j = _cuda.to_device(idx_j.astype(np.int32)) - # Preallocate device arrays - d_sigma = np.float64(sigma_nm) - d_val = _cuda.to_device(np.zeros(chunk_size)) - d_deri = _cuda.to_device(np.zeros((n_segments, 3), dtype=np.float64)) - wrapper = cuda_wrapper_chunked - - # Initial drift estimate + bounds - drift_est = np.zeros(n_segments * 3) - bounds = [ - ( - -drift_max_nm * drift_max_bound_factor, - drift_max_nm * drift_max_bound_factor, - ) - ] * (3 * n_segments) - - drift_est_gradient = np.inf - fails = 0 - done = False - itr_counter = 0 - start_time = time.time() - - while not done: - # Apply boxcar smoothing to current estimate - tmp = drift_est.reshape((-1, 3)) - for i in range(3): - tmp[:, i] = convolve( - tmp[:, i], np.ones(boxcar_width) / boxcar_width - ) - drift_est = tmp.flatten() - - # Run L-BFGS-B optimization step - result = minimize( - wrapper, - drift_est, - method="L-BFGS-B", - args=( - d_coords, - d_times, - d_idx_i, - d_idx_j, - d_sigma, - sigma_factor, - d_val, - d_deri, - chunk_size, - ), - jac=True, - bounds=bounds, - options={ - "disp": display_steps, - "gtol": 1e-5, - "ftol": 1e3 * np.finfo(float).eps, - "maxls": 40, - }, - ) - itr_counter += 1 - print( - f"Iteration {itr_counter}: status = {result.status}, success = {result.success}" - ) - print(f" current sigma: {np.round(sigma_nm * sigma_factor, 2)} nm") - - # Update if successful - if result.success: - delta = np.median((result.x - drift_est) ** 2) - print(f" drift estimate gradient: {delta}") - print(f" previous gradient: {drift_est_gradient}") - # Check convergence - if ( - delta > drift_est_gradient or sigma_nm * sigma_factor <= 1.0 - ) and sigma_nm * sigma_factor <= target_sigma_nm: - done = True - calc_time = time.time() - start_time - print(f"Optimization completed in {calc_time:.2f} s") - else: - sigma_factor /= 1.5 - drift_est_gradient = delta - drift_est = result.x - else: - fails += 1 - if fails > 2: - sigma_factor *= 2 - print("Restarting with larger sigma_factor") - if fails > 5: - raise RuntimeError( - "L-BFGS-B Optimization failed after multiple retries" - ) - - if return_calc_time: - return drift_est, time.time() - start_time, itr_counter - else: - return drift_est - - -def segmentation_and_pair_indices_wrapper( - dataset, - segmentation_var, - segmentation_mode, - max_drift_nm, - max_locs_per_segment, - pair_indices_safety_check=False, - hard_limit_pairs=None, -): - if not segmentation_mode == -1: # -1 is for pre-segmented data - result = segmentation_wrapper( - dataset[:, -1], - segmentation_var, - segmentation_mode, - max_locs_per_segment, - return_param_dict=True, - ) - else: - # pre segmented data, anyway we set these values in case auto downsampling is needed - segmentation_mode = 2 # dummy --> segment per frame ... - segmentation_var = 1 # dummy --> ... using 1 frame per segment - if pair_indices_safety_check: - n_pairs_est = estimate_pairs( - dataset[result.loc_valid, :3], max_drift_nm - ) - print( - f"Estimated number of pairs within {max_drift_nm} nm: {n_pairs_est:,}" - ) - if hard_limit_pairs is not None and n_pairs_est > hard_limit_pairs: - raise RuntimeError( - f"Estimated number of pairs {n_pairs_est} exceeds hard limit of {hard_limit_pairs}. " - f"Aborting to avoid crash." - ) - if n_pairs_est > 5e8: - print( - f"Estimated number of pairs is very large {n_pairs_est}. " - f"Automatic down-sampling is usually required above 500 mil." - f"Billions of pairs can lead to crash." - ) - ans = input("Continue anyway? (y/n): ") - if ans.lower() != "y": - raise RuntimeError( - "Aborted by user due to large estimated number of pairs." - ) - idx_i, idx_j, successful = pair_indices_kdtree( - dataset[result.loc_valid, :3], max_drift_nm - ) - if not successful: - if max_locs_per_segment is None: - max_locs_per_segment = int( - result.out_dict["locs_per_segment"].max() - ) - while not successful: - max_locs_per_segment = int(max_locs_per_segment * 0.9) - print( - f"Segmentation and Pairing attempt failed, automatic down-sampling active..." - ) - print( - f"Retrying segmentation with max_locs_per_segment={max_locs_per_segment}..." - ) - result = segmentation_wrapper( - dataset[:, -1], - segmentation_var, - segmentation_mode, - max_locs_per_segment, - return_param_dict=True, - ) - sorted_dataset = dataset.copy() - sorted_dataset[:, -1] = result.loc_segments - sorted_dataset = sorted_dataset[result.loc_valid] - idx_i, idx_j, successful = pair_indices_kdtree( - sorted_dataset[:, :3], max_drift_nm - ) - print( - f"Segmentation and Pairing successful resulting in {result.n_segments:,} time windows with on average " - f"{int(np.median(result.out_dict['locs_per_segment']))} locs per time window. " - f"{len(idx_i):,} Pairs where found." - ) - sorted_dataset = dataset.copy() - sorted_dataset[:, -1] = result.loc_segments - sorted_dataset = sorted_dataset[result.loc_valid] - return result, sorted_dataset, idx_i, idx_j - - -def estimate_pairs(coordinates, distance): - for i in range(len(coordinates[0])): - coordinates[:, i] -= np.min(coordinates[:, i]) - coordinates = np.array(np.floor(coordinates / distance), dtype=int) - coordinates = np.array(list(map(tuple, coordinates))) - sort_indices = np.lexsort( - coordinates.T - ) # get the unique tuples and their counts - unique_tuples, counts = np.unique( - coordinates[sort_indices], axis=0, return_counts=True - ) - # get the indices of the similar tuples - similar_indices = np.split(sort_indices, np.cumsum(counts[:-1])) - idx_i = [] - idx_j = [] - pair_idc_estimate = 0 - for i in range(len(similar_indices)): - n_elements = len(similar_indices[i]) - pair_idc_estimate += n_elements * (n_elements - 1) - rounded = round(pair_idc_estimate, -4) - return rounded - - -def interpolate_drift(center_frames, drift_est, frame_range, method="cubic"): - """ - Interpolates drift estimates to all frames using specified method. - Parameters: - - center_frames: np.ndarray of shape (M,), frames corresponding to drift estimates. - - drift_est: np.ndarray of shape (M, 3), drift estimates at center frames. - - frame_range: array-like, frames to interpolate drift estimates to. - - method: str, interpolation method ('cubic' or 'catmull-rom'). - Returns: - - drift_interp: np.ndarray of shape (len(frame_range), 3), interpolated drift estimates. - """ - if method == "cubic": - return _interpolate_cubic(center_frames, drift_est, frame_range) - elif method == "catmull-rom": - return _interpolate_catmull_rom(center_frames, drift_est, frame_range) - elif method == "linear": - drift_x = np.interp(frame_range, center_frames, drift_est[:, 0]) - drift_y = np.interp(frame_range, center_frames, drift_est[:, 1]) - drift_z = np.interp(frame_range, center_frames, drift_est[:, 2]) - return np.vstack([drift_x, drift_y, drift_z]).T - else: - raise ValueError(f"Unknown interpolation method: {method}") - - -def _interpolate_cubic(center_frames, drift_est, frame_range): - from scipy.interpolate import CubicSpline - - drift_x = CubicSpline(center_frames, drift_est[:, 0])(frame_range) - drift_y = CubicSpline(center_frames, drift_est[:, 1])(frame_range) - drift_z = CubicSpline(center_frames, drift_est[:, 2])(frame_range) - return np.vstack([drift_x, drift_y, drift_z]).T - - -def _interpolate_catmull_rom(center_frames, drift_est, frame_range): - def catmull_rom_1d(x, y, x_interp): - result = np.zeros_like(x_interp) - for i in range(1, len(x) - 2): - x0, x1, x2, x3 = x[i - 1], x[i], x[i + 1], x[i + 2] - y0, y1, y2, y3 = y[i - 1], y[i], y[i + 1], y[i + 2] - - mask = (x_interp >= x1) & (x_interp <= x2) - t = (x_interp[mask] - x1) / (x2 - x1) - - result[mask] = 0.5 * ( - (2 * y1) - + (-y0 + y2) * t - + (2 * y0 - 5 * y1 + 4 * y2 - y3) * t**2 - + (-y0 + 3 * y1 - 3 * y2 + y3) * t**3 - ) - return result - - x_interp = np.asarray(frame_range) - x = np.asarray(center_frames) - - if len(x) < 4: - raise ValueError( - "Catmull-Rom interpolation requires at least 4 points." - ) - - drift_x = catmull_rom_1d(x, drift_est[:, 0], x_interp) - drift_y = catmull_rom_1d(x, drift_est[:, 1], x_interp) - drift_z = catmull_rom_1d(x, drift_est[:, 2], x_interp) - - return np.vstack([drift_x, drift_y, drift_z]).T - - -# CUDA kernel and wrapper — only defined when CUDA is available at import time. -if _CUDA_AVAILABLE: - - @_cuda.jit - def cost_function_full_3d_chunked( - d_locs_time, - start_idx, - chunk_size, - d_idx_i, - d_idx_j, - d_sigma, - d_sigma_factor, - d_val, - d_val_sum, - d_deri, - d_locs_coords, - mu, - ): - """Compute negative-overlap cost and gradient for 3D localizations (CUDA kernel).""" - tx = _cuda.threadIdx.x - ty = _cuda.blockIdx.x - bw = _cuda.blockDim.x - pos = tx + ty * bw - - if pos < chunk_size: - i = d_idx_i[pos + start_idx] - j = d_idx_j[pos + start_idx] - - ti = d_locs_time[i] - tj = d_locs_time[j] - - dx = (d_locs_coords[i, 0] - mu[ti, 0]) - ( - d_locs_coords[j, 0] - mu[tj, 0] - ) - dy = (d_locs_coords[i, 1] - mu[ti, 1]) - ( - d_locs_coords[j, 1] - mu[tj, 1] - ) - dz = (d_locs_coords[i, 2] - mu[ti, 2]) - ( - d_locs_coords[j, 2] - mu[tj, 2] - ) - sigma_sq = (2 * d_sigma * d_sigma_factor) ** 2 - - diff_sq = dx * dx + dy * dy + dz * dz - val = ( - 1 / (d_sigma * d_sigma_factor) * math.exp(-diff_sq / sigma_sq) - ) - d_val[pos] = val - - _cuda.atomic.add( - d_deri, - (tj, 0), - 2 - * val - * ( - d_locs_coords[j, 0] - - d_locs_coords[i, 0] - + mu[ti, 0] - - mu[tj, 0] - ) - / sigma_sq, - ) - _cuda.atomic.add( - d_deri, - (tj, 1), - 2 - * val - * ( - d_locs_coords[j, 1] - - d_locs_coords[i, 1] - + mu[ti, 1] - - mu[tj, 1] - ) - / sigma_sq, - ) - _cuda.atomic.add( - d_deri, - (tj, 2), - 2 - * val - * ( - d_locs_coords[j, 2] - - d_locs_coords[i, 2] - + mu[ti, 2] - - mu[tj, 2] - ) - / sigma_sq, - ) - - _cuda.atomic.add( - d_deri, - (ti, 0), - 2 - * val - * ( - d_locs_coords[i, 0] - - d_locs_coords[j, 0] - + mu[tj, 0] - - mu[ti, 0] - ) - / sigma_sq, - ) - _cuda.atomic.add( - d_deri, - (ti, 1), - 2 - * val - * ( - d_locs_coords[i, 1] - - d_locs_coords[j, 1] - + mu[tj, 1] - - mu[ti, 1] - ) - / sigma_sq, - ) - _cuda.atomic.add( - d_deri, - (ti, 2), - 2 - * val - * ( - d_locs_coords[i, 2] - - d_locs_coords[j, 2] - + mu[tj, 2] - - mu[ti, 2] - ) - / sigma_sq, - ) - - _cuda.atomic.add(d_val_sum, 0, val) - d_val[pos] = 0 - - def cuda_wrapper_chunked( - mu, - d_locs_coords, - d_locs_time, - d_idx_i, - d_idx_j, - d_sigma, - d_sigma_factor, - d_val, - d_deri, - chunk_size, - ): - """Interface between Python optimizer and the CUDA kernel (chunked to manage memory).""" - val_total = 0 - d_val_sum = _cuda.to_device(np.zeros(1, dtype=np.float64)) - mu_dev = _cuda.to_device( - np.asarray(mu.reshape(int(mu.size / 3), 3), dtype=np.float64) - ) - - n_chunks = int(np.ceil(d_idx_i.size / chunk_size)) - threadsperblock = 128 - - for i in range(n_chunks - 1): - idc_start = i * chunk_size - blockspergrid = ( - chunk_size + (threadsperblock - 1) - ) // threadsperblock - cost_function_full_3d_chunked[blockspergrid, threadsperblock]( - d_locs_time, - idc_start, - chunk_size, - d_idx_i, - d_idx_j, - d_sigma, - d_sigma_factor, - d_val, - d_val_sum, - d_deri, - d_locs_coords, - mu_dev, - ) - val_total += d_val_sum.copy_to_host() - - # Final chunk - n_remaining = d_idx_i.size - (n_chunks - 1) * chunk_size - idc_start = (n_chunks - 1) * chunk_size - blockspergrid = ( - n_remaining + (threadsperblock - 1) - ) // threadsperblock - cost_function_full_3d_chunked[blockspergrid, threadsperblock]( - d_locs_time, - idc_start, - n_remaining, - d_idx_i, - d_idx_j, - d_sigma, - d_sigma_factor, - d_val, - d_val_sum, - d_deri, - d_locs_coords, - mu_dev, - ) - val_total += d_val_sum.copy_to_host() - deri = d_deri.copy_to_host() - d_deri[:] = 0 - - return -np.nansum(val_total), -deri.flatten() - - -@dataclass -class SegmentationResult: - """Container for segmentation results.""" - - loc_segments: np.ndarray - loc_valid: np.ndarray - center_frames: np.ndarray - n_segments: int - out_dict: Optional[Dict] = None - - -def _group_by_frame(loc_frames: np.ndarray): - """Returns a dict {frame_number: indices_in_loc_frames} efficiently.""" - sort_idx = np.argsort(loc_frames) - sorted_frames = loc_frames[sort_idx] - unique_frames, start_idx, counts = np.unique( - sorted_frames, return_index=True, return_counts=True - ) - frame_to_indices = { - frame: sort_idx[start : start + count] - for frame, start, count in zip(unique_frames, start_idx, counts) - } - return unique_frames, frame_to_indices - - -def segment_by_num_locs_per_window( - loc_frames: np.ndarray, - min_n_locs_per_window: int, - max_locs_per_segment=None, - return_param_dict: bool = False, -) -> SegmentationResult: - """ - Segments by collecting a minimum number of localizations per window. - Once the threshold is met and enough locs remain, a new segment is created. - This method ensures that each segment has at least `min_n_locs_per_window` localizations, - while also trying to avoid creating segments that are too small at the end of the dataset. - If `max_locs_per_segment` is set, a random subset of that size is chosen from each segment. - This method is particularly useful for datasets with varying localization densities over time. - Parameters: - loc_frames (np.ndarray): Array of frame numbers for each localization. - min_n_locs_per_window (int): Minimum number of localizations per segment. - max_locs_per_segment (Optional[int]): Maximum number of localizations per segment. If None, all locs are used. - return_param_dict (bool): Whether to return a dictionary of segmentation parameters. - Returns: - SegmentationResult: A dataclass containing segmentation results and parameters. - """ - loc_frames = np.asarray(loc_frames, dtype=int) - n_locs = len(loc_frames) - - if ( - max_locs_per_segment is not None and max_locs_per_segment < 1 - ): # downsampling in percentage - max_locs_per_segment = int( - min_n_locs_per_window * max_locs_per_segment - ) - - unique_frames, frame_to_indices = _group_by_frame(loc_frames) - loc_segments = np.full(n_locs, -1, dtype=int) # Default to -1 for safety - segment_counter = 0 - n_locs_in_current_segment = 0 - current_segment_indices = [] - start_frames, end_frames, locs_per_segment = [], [], [] - - for i, frame in enumerate(unique_frames): - indices = frame_to_indices[frame] - n_locs_this_frame = len(indices) - remaining_locs = n_locs - ( - len(current_segment_indices) - + n_locs_this_frame - + np.sum(locs_per_segment) - ) - - # Add frame to current segment if: - # - It fills the current segment to threshold - # - AND there are enough locs left for another segment (or it's the last frame) - if ( - n_locs_in_current_segment + n_locs_this_frame - >= min_n_locs_per_window - ) and ( - remaining_locs >= min_n_locs_per_window - or i == len(unique_frames) - 1 - ): - current_segment_indices.extend(indices) - n_locs_in_current_segment += n_locs_this_frame - - loc_segments[current_segment_indices] = segment_counter - start_frames.append(loc_frames[current_segment_indices[0]]) - end_frames.append(loc_frames[current_segment_indices[-1]]) - locs_per_segment.append(len(current_segment_indices)) - - segment_counter += 1 - current_segment_indices = [] - n_locs_in_current_segment = 0 - else: - # Defer frame to current segment - current_segment_indices.extend(indices) - n_locs_in_current_segment += n_locs_this_frame - - n_segments = segment_counter - center_frames = np.zeros(n_segments) - loc_valid = np.zeros(n_locs, dtype=bool) - - for i in range(n_segments): - segment_indices = np.where(loc_segments == i)[0] - if ( - max_locs_per_segment - and len(segment_indices) > max_locs_per_segment - ): - selected = np.random.choice( - segment_indices, max_locs_per_segment, replace=False - ) - else: - selected = segment_indices - loc_valid[selected] = True - locs_per_segment[i] = len(selected) - center_frames[i] = np.mean(loc_frames[selected]) - - out_dict = None - if return_param_dict: - n_locs_valid = loc_valid.sum() - out_dict = { - "n_segments": n_segments, - "min_n_locs_per_window": min_n_locs_per_window, - "frames_per_window": -1, - "start_frames": np.array(start_frames), - "end_frames": np.array(end_frames), - "locs_per_segment": np.array(locs_per_segment), - "n_locs": n_locs, - "n_locs_valid": n_locs_valid, - "n_locs_invalid": n_locs - n_locs_valid, - "center_frames": center_frames, - } - - return SegmentationResult( - loc_segments, loc_valid, center_frames, n_segments, out_dict - ) - - -def segment_by_frame_windows( - loc_frames: np.ndarray, - n_frames_per_window: int, - max_locs_per_segment=None, - return_param_dict: bool = False, -) -> SegmentationResult: - """ - Splits localization data into fixed-size windows of N frames. - All localizations in those frames are grouped into one segment. - """ - loc_frames = np.asarray(loc_frames, dtype=int) - frames, frame_to_indices = _group_by_frame(loc_frames) - n_locs = len(loc_frames) - n_segments = int(np.ceil(len(frames) / n_frames_per_window)) - - if ( - max_locs_per_segment is not None and max_locs_per_segment < 1 - ): # downsampling in percentage - max_locs_per_segment = n_locs / n_segments * max_locs_per_segment - - loc_segments = np.zeros(n_locs, dtype=int) - center_frames = np.zeros(n_segments) - loc_valid = np.ones(n_locs, dtype=bool) - start_frames, end_frames, locs_per_segment = [], [], [] - - for i in range(n_segments): - frame_window = frames[ - i * n_frames_per_window : (i + 1) * n_frames_per_window - ] - indices = np.concatenate( - [ - frame_to_indices[frame] - for frame in frame_window - if frame in frame_to_indices - ] - ) - if len(indices) == 0: - continue - loc_segments[indices] = i - start_frames.append(frame_window[0]) - end_frames.append(frame_window[-1]) - center_frames[i] = np.mean(loc_frames[indices]) - locs_per_segment.append(len(indices)) - if max_locs_per_segment and len(indices) > max_locs_per_segment: - mask = np.ones(len(indices), dtype=bool) - mask[ - np.random.choice( - len(indices), - len(indices) - max_locs_per_segment, - replace=False, - ) - ] = False - loc_valid[indices[~mask]] = False - - out_dict = None - if return_param_dict: - n_locs_valid = loc_valid.sum() - out_dict = { - "n_segments": n_segments, - "min_n_locs_per_window": -1, - "frames_per_window": n_frames_per_window, - "start_frames": np.array(start_frames), - "end_frames": np.array(end_frames), - "locs_per_segment": np.array(locs_per_segment), - "n_locs": n_locs, - "n_locs_valid": n_locs_valid, - "n_locs_invalid": n_locs - n_locs_valid, - "center_frames": center_frames, - } - - return SegmentationResult( - loc_segments, loc_valid, center_frames, n_segments, out_dict - ) - - -def segment_by_num_windows( - loc_frames: np.ndarray, - n_windows: int, - max_locs_per_segment=None, - return_param_dict: bool = False, -) -> SegmentationResult: - """ - Converts number of windows into an equivalent minimum locs per window, - then calls `segment_by_num_locs_per_window`. - """ - n_locs = len(loc_frames) - n_locs_per_window = int(np.ceil(n_locs / n_windows)) - if ( - max_locs_per_segment is not None and max_locs_per_segment < 1 - ): # downsampling in percentage - max_locs_per_segment = int(n_locs_per_window * max_locs_per_segment) - return segment_by_num_locs_per_window( - loc_frames, n_locs_per_window, max_locs_per_segment, return_param_dict - ) - - -def segmentation_wrapper( - loc_frames: np.ndarray, - segmentation_var: int, - segmentation_mode: int = 2, - max_locs_per_segment=None, - return_param_dict: bool = False, -) -> SegmentationResult: - """ - Dispatch function that selects segmentation method: - 0 → fixed number of windows - 1 → fixed number of localizations per window - 2 → fixed number of frames per window (default) - """ - if segmentation_mode == 0: - return segment_by_num_windows( - loc_frames, - segmentation_var, - max_locs_per_segment, - return_param_dict, - ) - elif segmentation_mode == 1: - return segment_by_num_locs_per_window( - loc_frames, - segmentation_var, - max_locs_per_segment, - return_param_dict, - ) - else: - return segment_by_frame_windows( - loc_frames, - segmentation_var, - max_locs_per_segment, - return_param_dict, - ) - - -def pair_indices_kdtree(coordinates, distance): - """ - Find all pairs of points within a certain distance using a KD-tree. - Parameters: - - coordinates: np.ndarray of shape (N, D) where N is the number of points and D is the dimensionality. - - distance: float, the maximum distance to consider points as a pair. - Returns: - - idx1: np.ndarray of shape (M,), indices of the first point in each pair. - - idx2: np.ndarray of shape (M,), indices of the second point in each pair. - """ - tree = cKDTree(coordinates) - try: - pairs = tree.query_pairs(r=distance, output_type="ndarray") - except MemoryError: - print("[pair_indices_kdtree] MemoryError encountered") - return [], [], False - return ( - np.ascontiguousarray(pairs[:, 0], dtype=np.int32), - np.ascontiguousarray(pairs[:, 1], dtype=np.int32), - True, - ) diff --git a/picasso/gui/render.py b/picasso/gui/render.py index 42ebb72e..d5f31ee5 100644 --- a/picasso/gui/render.py +++ b/picasso/gui/render.py @@ -57,7 +57,6 @@ # Optional modules with external/hardware dependencies live in ext from ..ext.bitplane import IMSWRITER # PyImarisWrite works on Windows only -from ..ext import comet if IMSWRITER: from ..ext.bitplane import numpy_to_imaris @@ -2220,89 +2219,6 @@ def getParams( ) -class COMETDialog(QtWidgets.QDialog): - """Dialog to choose parameters for COMET undrifting. - - Attributes - ---------- - locs_per_segment : QSpinBox - Target number of localizations per temporal segment. - max_drift_nm : QDoubleSpinBox - Maximum expected drift in nm. - max_locs_per_segment : QSpinBox - Optional cap for downsampling localizations per segment. - Use -1 to disable downsampling. - """ - - def __init__(self, window: QtWidgets.QMainWindow) -> None: - super().__init__(window) - self.window = window - self.setWindowTitle("COMET undrifting") - - vbox = QtWidgets.QVBoxLayout(self) - grid = QtWidgets.QGridLayout() - - locs_label = QtWidgets.QLabel("Localizations per segment:") - locs_label.setToolTip( - "Target number of localizations used to form each temporal segment." - ) - grid.addWidget(locs_label, 0, 0) - self.locs_per_segment = QtWidgets.QSpinBox() - self.locs_per_segment.setRange(1, int(1e6)) - self.locs_per_segment.setValue(500) - grid.addWidget(self.locs_per_segment, 0, 1) - - max_drift_label = QtWidgets.QLabel("Maximum drift (nm):") - max_drift_label.setToolTip( - "Maximum expected drift over the dataset. " - "Used for pairing and optimization bounds." - ) - grid.addWidget(max_drift_label, 1, 0) - self.max_drift_nm = QtWidgets.QDoubleSpinBox() - self.max_drift_nm.setRange(0.1, 1e6) - self.max_drift_nm.setValue(60.0) - self.max_drift_nm.setDecimals(1) - self.max_drift_nm.setSingleStep(1.0) - grid.addWidget(self.max_drift_nm, 1, 1) - - downsample_label = QtWidgets.QLabel("Max. localizations per segment:") - downsample_label.setToolTip( - "Optional downsampling cap per segment. " - "Set to -1 to keep all localizations." - ) - grid.addWidget(downsample_label, 2, 0) - self.max_locs_per_segment = QtWidgets.QSpinBox() - self.max_locs_per_segment.setRange(-1, int(1e6)) - self.max_locs_per_segment.setValue(-1) - grid.addWidget(self.max_locs_per_segment, 2, 1) - - vbox.addLayout(grid) - - self.buttons = QtWidgets.QDialogButtonBox( - QtWidgets.QDialogButtonBox.StandardButton.Ok - | QtWidgets.QDialogButtonBox.StandardButton.Cancel, - QtCore.Qt.Orientation.Horizontal, - self, - ) - vbox.addWidget(self.buttons) - self.buttons.accepted.connect(self.accept) - self.buttons.rejected.connect(self.reject) - - @staticmethod - def getParams( - parent: QtWidgets.QMainWindow | None = None, - ) -> tuple[dict, bool]: - """Create the dialog and return the requested COMET parameters.""" - dialog = COMETDialog(parent) - result = dialog.exec() - params = { - "locs_per_segment": dialog.locs_per_segment.value(), - "max_drift_nm": dialog.max_drift_nm.value(), - "max_locs_per_segment": dialog.max_locs_per_segment.value(), - } - return params, result == QtWidgets.QDialog.DialogCode.Accepted - - class AIMDialog(lib.Dialog): """Choose parameters for AIM undrifting. @@ -11079,44 +10995,6 @@ def show_drift(self) -> None: self.plot_window.plot(drift) self.plot_window.show() - def undrift_comet(self) -> None: - """Undrift with COMET. See https://comet.smlm.tools and - Reinkensmeier L., Aufmkolk S., Farabella I., Egner A., and - Bates M. biorxiv, 2026.""" - channel = self.get_channel("Undrift by COMET") - if channel is not None: - locs = self.locs[channel] - info = self.infos[channel] - - params, ok = COMETDialog.getParams(self.window) - if ok: - try: - locs, new_info, drift = comet.comet( - locs, - info, - **params, - ) - except RuntimeError as e: - QtWidgets.QMessageBox.warning( - self.window, - "COMET not available", - ( - f"{e}\n\n" - "Please use a different drift-correction method, " - "such as Undrift by AIM or Undrift by RCC." - ), - ) - return - - locs = lib.ensure_sanity(locs, info) - self.locs[channel] = locs - self.infos[channel] = new_info - self.index_blocks[channel] = None - self.render_index[channel] = None - self.add_drift(channel, drift) - self.update_scene(resample_locs=True) - self.show_drift() - def undrift_aim(self) -> None: """Undrift with Adaptive Intersection Maximization (AIM). @@ -12134,10 +12012,6 @@ def initUI(self, plugins_loaded: bool) -> None: undrift_from_picked2d_action.triggered.connect( self.view.undrift_from_picked2d ) - undrift_comet_action = postprocess_menu.addAction("Undrift by COMET") - undrift_comet_action.triggered.connect(self.view.undrift_comet) - if not comet._CUDA_AVAILABLE: - undrift_comet_action.setVisible(False) undrift_action = postprocess_menu.addAction("Undrift by RCC") undrift_action.triggered.connect(self.view.undrift_rcc) diff --git a/readme.rst b/readme.rst index c65f093d..8d51e25e 100644 --- a/readme.rst +++ b/readme.rst @@ -137,7 +137,6 @@ If you use Picasso in your research, please cite our Nature Protocols publicatio - 3D fitting via astigmatism. DOI: `10.1126/science.1153529 `__. - RCC undrifting: DOI: `10.1364/OE.22.015982 `__ - AIM undrifting. DOI: `10.1126/sciadv.adm776 `__ -- COMET undrifting. DOI: `10.64898/2026.03.27.714864 `__ - SMLM clusterer. DOIs: `10.1038/s41467-021-22606-1 `__ and `10.1038/s41586-023-05925-9 `__ - DBSCAN: Ester, et al. Inkdd, 1996. (Vol. 96, No. 34, pp. 226-231). - Anisotropic DBSCAN inspired by: `10.1021/acs.jpcb.4c02030 `__ diff --git a/release/one_click_macos_gui/readme.txt b/release/one_click_macos_gui/readme.txt index d0c6a105..fc2cf68e 100644 --- a/release/one_click_macos_gui/readme.txt +++ b/release/one_click_macos_gui/readme.txt @@ -62,7 +62,6 @@ If you use some of the functionalities provided by Picasso, please also cite the - 3D fitting via astigmatism. DOI: 10.1126/science.1153529 (https://www.science.org/doi/10.1126/science.1153529). - RCC undrifting: DOI: 10.1364/OE.22.015982 (https://doi.org/10.1364/OE.22.015982) - AIM undrifting. DOI: 10.1126/sciadv.adm776 (https://www.science.org/doi/10.1126/sciadv.adm7765) -- COMET undrifting. DOI: 10.64898/2026.03.27.714864 (https://doi.org/10.64898/2026.03.27.714864) - SMLM clusterer. DOIs: 10.1038/s41467-021-22606-1 (https://doi.org/10.1038/s41467-021-22606-1) and 10.1038/s41586-023-05925-9 (https://doi.org/10.1038/s41586-023-05925-9) - DBSCAN: Ester, et al. Inkdd, 1996. (Vol. 96, No. 34, pp. 226-231). - Anisotropic DBSCAN inspired by: 10.1021/acs.jpcb.4c02030 (https://doi.org/10.1021/acs.jpcb.4c02030) diff --git a/release/one_click_windows_gui/readme.txt b/release/one_click_windows_gui/readme.txt index 78f30fd8..039f1303 100644 --- a/release/one_click_windows_gui/readme.txt +++ b/release/one_click_windows_gui/readme.txt @@ -67,7 +67,6 @@ If you use some of the functionalities provided by Picasso, please also cite the - 3D fitting via astigmatism. DOI: 10.1126/science.1153529 (https://www.science.org/doi/10.1126/science.1153529). - RCC undrifting: DOI: 10.1364/OE.22.015982 (https://doi.org/10.1364/OE.22.015982) - AIM undrifting. DOI: 10.1126/sciadv.adm776 (https://www.science.org/doi/10.1126/sciadv.adm7765) -- COMET undrifting. DOI: 10.64898/2026.03.27.714864 (https://doi.org/10.64898/2026.03.27.714864) - SMLM clusterer. DOIs: 10.1038/s41467-021-22606-1 (https://doi.org/10.1038/s41467-021-22606-1) and 10.1038/s41586-023-05925-9 (https://doi.org/10.1038/s41586-023-05925-9) - DBSCAN: Ester, et al. Inkdd, 1996. (Vol. 96, No. 34, pp. 226-231). - Anisotropic DBSCAN inspired by: 10.1021/acs.jpcb.4c02030 (https://doi.org/10.1021/acs.jpcb.4c02030) diff --git a/tests/test_comet.py b/tests/test_comet.py deleted file mode 100644 index 3050ef4a..00000000 --- a/tests/test_comet.py +++ /dev/null @@ -1,572 +0,0 @@ -"""Tests for picasso.comet — drift-correction with COMET. - -Tests are split into two groups: -- Pure-Python helpers (segmentation, pairing, interpolation): always run when - numba is installed, no GPU required. -- Full pipeline (comet.comet): requires CUDA hardware; skipped automatically - on CPU-only machines, but the no-GPU error path IS tested on any machine. - -:author: Lenny Reinkensmeier, 2026 -""" - -import numpy as np -import pandas as pd -import pytest - -# The whole picasso package requires numba (lib.py imports it). -# Skip this module cleanly when numba is absent (e.g. CI without GPU stack). -pytest.importorskip("numba") - -from picasso.ext import comet # noqa: E402 — import after guard - - -# --------------------------------------------------------------------------- -# Fixtures -# --------------------------------------------------------------------------- - - -@pytest.fixture -def minimal_info(): - """Minimal Picasso info list with required Pixelsize entry.""" - return [{"Pixelsize": 130, "Frames": 200}] - - -@pytest.fixture -def minimal_locs_df(): - """Small synthetic localisation DataFrame spread over 200 frames.""" - rng = np.random.default_rng(0) - n_frames = 200 - n_per_frame = 5 - frames = np.repeat(np.arange(n_frames), n_per_frame) - x = rng.uniform(5, 25, size=len(frames)).astype(np.float32) - y = rng.uniform(5, 25, size=len(frames)).astype(np.float32) - return pd.DataFrame({"frame": frames, "x": x, "y": y}) - - -@pytest.fixture -def minimal_locs_structured(minimal_locs_df): - """Same data as a numpy structured array (Picasso's native format).""" - df = minimal_locs_df - dtype = np.dtype( - [("frame", np.uint32), ("x", np.float32), ("y", np.float32)] - ) - arr = np.empty(len(df), dtype=dtype) - arr["frame"] = df["frame"].to_numpy(np.uint32) - arr["x"] = df["x"].to_numpy(np.float32) - arr["y"] = df["y"].to_numpy(np.float32) - return arr - - -# --------------------------------------------------------------------------- -# Segmentation helpers -# --------------------------------------------------------------------------- - - -class TestSegmentationByLocs: - def test_produces_correct_segment_count(self): - frames = np.repeat(np.arange(100), 5) # 500 locs, 5/frame - result = comet.segmentation_wrapper( - frames, - segmentation_var=50, - segmentation_mode=1, - return_param_dict=True, - ) - assert result.n_segments == 10 - - def test_all_locs_assigned(self): - frames = np.repeat(np.arange(50), 10) - result = comet.segmentation_wrapper( - frames, segmentation_var=100, segmentation_mode=1 - ) - assert (result.loc_segments >= 0).all() - - def test_loc_valid_subset_of_locs(self): - frames = np.repeat(np.arange(60), 8) - result = comet.segmentation_wrapper( - frames, segmentation_var=40, segmentation_mode=1 - ) - assert result.loc_valid.shape == frames.shape - assert result.loc_valid.dtype == bool - - def test_center_frames_length_matches_n_segments(self): - frames = np.repeat(np.arange(80), 5) - result = comet.segmentation_wrapper( - frames, segmentation_var=50, segmentation_mode=1 - ) - assert len(result.center_frames) == result.n_segments - - def test_max_locs_per_segment_respected(self): - frames = np.repeat(np.arange(20), 100) # 2000 locs, 100/frame - cap = 30 - result = comet.segmentation_wrapper( - frames, - segmentation_var=200, - segmentation_mode=1, - max_locs_per_segment=cap, - return_param_dict=True, - ) - assert result.loc_valid.sum() <= result.n_segments * cap - - -class TestSegmentationByFrameWindows: - def test_frame_window_mode(self): - frames = np.repeat(np.arange(100), 3) - result = comet.segmentation_wrapper( - frames, segmentation_var=10, segmentation_mode=2 - ) - assert result.n_segments == 10 # 100 frames / 10 per window - - def test_single_window(self): - frames = np.arange(50) - result = comet.segmentation_wrapper( - frames, segmentation_var=100, segmentation_mode=2 - ) - assert result.n_segments == 1 - - -class TestSegmentationByNumWindows: - def test_num_windows_mode_creates_requested_count(self): - frames = np.repeat(np.arange(100), 5) # 500 locs - result = comet.segmentation_wrapper( - frames, - segmentation_var=5, - segmentation_mode=0, - return_param_dict=True, - ) - # mode 0 derives min-locs-per-window from total/n_windows; the resulting - # segment count should be approximately n_windows (off-by-one tolerated - # because the last segment is folded into the previous if too small). - assert abs(result.n_segments - 5) <= 1 - - def test_all_locs_assigned(self): - frames = np.repeat(np.arange(40), 4) - result = comet.segmentation_wrapper( - frames, segmentation_var=4, segmentation_mode=0 - ) - assert (result.loc_segments >= 0).all() - - -class TestSegmentationOutDict: - """The param dict reports per-segment bookkeeping; verify the values, not - just that segmentation runs.""" - - def test_valid_plus_invalid_equals_total(self): - frames = np.repeat(np.arange(60), 8) - result = comet.segmentation_wrapper( - frames, - segmentation_var=40, - segmentation_mode=1, - return_param_dict=True, - ) - d = result.out_dict - assert d["n_locs"] == len(frames) - assert d["n_locs_valid"] + d["n_locs_invalid"] == d["n_locs"] - - def test_per_segment_arrays_have_consistent_length(self): - frames = np.repeat(np.arange(80), 5) - result = comet.segmentation_wrapper( - frames, - segmentation_var=50, - segmentation_mode=1, - return_param_dict=True, - ) - d = result.out_dict - assert len(d["start_frames"]) == result.n_segments - assert len(d["end_frames"]) == result.n_segments - assert len(d["locs_per_segment"]) == result.n_segments - - def test_start_frames_not_after_end_frames(self): - frames = np.repeat(np.arange(80), 5) - result = comet.segmentation_wrapper( - frames, - segmentation_var=50, - segmentation_mode=1, - return_param_dict=True, - ) - d = result.out_dict - assert (d["start_frames"] <= d["end_frames"]).all() - - def test_center_frames_within_segment_bounds(self): - frames = np.repeat(np.arange(80), 5) - result = comet.segmentation_wrapper( - frames, - segmentation_var=50, - segmentation_mode=1, - return_param_dict=True, - ) - d = result.out_dict - assert (result.center_frames >= d["start_frames"]).all() - assert (result.center_frames <= d["end_frames"]).all() - - def test_locs_per_segment_sums_to_valid_count(self): - frames = np.repeat(np.arange(60), 6) - result = comet.segmentation_wrapper( - frames, - segmentation_var=40, - segmentation_mode=1, - return_param_dict=True, - ) - d = result.out_dict - # locs_per_segment is rewritten to the post-downsampling (valid) count. - assert d["locs_per_segment"].sum() == d["n_locs_valid"] - - def test_frame_window_center_frames_within_bounds(self): - frames = np.repeat(np.arange(100), 3) - result = comet.segmentation_wrapper( - frames, - segmentation_var=10, - segmentation_mode=2, - return_param_dict=True, - ) - d = result.out_dict - assert d["frames_per_window"] == 10 - assert (result.center_frames >= d["start_frames"]).all() - assert (result.center_frames <= d["end_frames"]).all() - - -class TestSegmentationDownsamplingByPercentage: - """max_locs_per_segment < 1 is interpreted as a fraction of the window size.""" - - def test_fractional_cap_reduces_valid_locs(self): - frames = np.repeat(np.arange(20), 50) # 1000 locs, 50/frame - full = comet.segmentation_wrapper( - frames, - segmentation_var=200, - segmentation_mode=1, - return_param_dict=True, - ) - half = comet.segmentation_wrapper( - frames, - segmentation_var=200, - segmentation_mode=1, - max_locs_per_segment=0.5, - return_param_dict=True, - ) - assert half.out_dict["n_locs_valid"] < full.out_dict["n_locs_valid"] - - -# --------------------------------------------------------------------------- -# Pair-finding -# --------------------------------------------------------------------------- - - -class TestPairIndicesKdtree: - def test_finds_close_pairs(self): - coords = np.array( - [[0, 0, 0], [1, 0, 0], [100, 0, 0]], dtype=np.float32 - ) - idx_i, idx_j, ok = comet.pair_indices_kdtree(coords, distance=5) - assert ok - assert len(idx_i) == 1 # only (0,1) are within distance 5 - - def test_no_pairs_when_far_apart(self): - coords = np.eye(3, dtype=np.float32) * 1000 - idx_i, idx_j, ok = comet.pair_indices_kdtree(coords, distance=1) - assert ok - assert len(idx_i) == 0 - - def test_all_pairs_when_all_close(self): - coords = np.zeros((5, 3), dtype=np.float32) - idx_i, idx_j, ok = comet.pair_indices_kdtree(coords, distance=1) - assert ok - assert len(idx_i) == 5 * 4 // 2 # C(5,2) - - def test_returns_int32_arrays(self): - coords = np.random.rand(20, 3).astype(np.float32) - idx_i, idx_j, ok = comet.pair_indices_kdtree(coords, distance=0.5) - assert ok - assert idx_i.dtype == np.int32 - assert idx_j.dtype == np.int32 - - def test_2d_coordinates_supported(self): - coords = np.array([[0, 0], [1, 0], [100, 0]], dtype=np.float32) - idx_i, idx_j, ok = comet.pair_indices_kdtree(coords, distance=5) - assert ok - assert len(idx_i) == 1 # only (0,1) within distance 5 - - def test_pairs_are_lower_index_first(self): - coords = np.zeros((6, 3), dtype=np.float32) - idx_i, idx_j, ok = comet.pair_indices_kdtree(coords, distance=1) - assert ok - # query_pairs returns each pair once with i < j; relied on by the kernel. - assert (idx_i < idx_j).all() - - def test_distance_is_inclusive(self): - coords = np.array([[0, 0, 0], [10, 0, 0]], dtype=np.float32) - idx_i, idx_j, ok = comet.pair_indices_kdtree(coords, distance=10) - assert ok - assert ( - len(idx_i) == 1 - ) # points exactly `distance` apart count as a pair - - -# --------------------------------------------------------------------------- -# Pair-count estimation -# --------------------------------------------------------------------------- - - -class TestEstimatePairs: - def test_far_apart_points_estimate_zero(self): - # Each point lands in its own grid cell -> no within-cell pairs. - coords = ( - (np.arange(5)[:, None] * 1000.0).repeat(3, axis=1).astype(float) - ) - assert comet.estimate_pairs(coords, distance=10) == 0 - - def test_clustered_points_estimate_nonzero(self): - # 200 coincident points -> 200*199 = 39800 -> rounds to 40000. - coords = np.zeros((200, 3), dtype=float) - est = comet.estimate_pairs(coords, distance=10) - assert est == 40000 - - def test_result_rounded_to_nearest_ten_thousand(self): - coords = np.zeros((150, 3), dtype=float) - est = comet.estimate_pairs(coords, distance=10) - assert est % 10000 == 0 - - -# --------------------------------------------------------------------------- -# Drift interpolation -# --------------------------------------------------------------------------- - - -class TestInterpolateDrift: - def _center_frames_and_drift(self): - center_frames = np.linspace(10, 90, 6) - drift_est = np.column_stack( - [ - np.sin(center_frames / 10), - np.cos(center_frames / 10), - np.zeros_like(center_frames), - ] - ) - return center_frames, drift_est - - def test_cubic_output_shape(self): - cf, de = self._center_frames_and_drift() - frame_range = np.arange(0, 100) - out = comet.interpolate_drift(cf, de, frame_range, method="cubic") - assert out.shape == (100, 3) - - def test_linear_output_shape(self): - cf, de = self._center_frames_and_drift() - frame_range = np.arange(0, 100) - out = comet.interpolate_drift(cf, de, frame_range, method="linear") - assert out.shape == (100, 3) - - def test_cubic_passes_through_control_points(self): - cf = np.array([0.0, 25.0, 50.0, 75.0, 100.0]) - de = np.column_stack([cf * 0.01, cf * 0.005, np.zeros_like(cf)]) - out = comet.interpolate_drift(cf, de, cf.astype(int), method="cubic") - np.testing.assert_allclose(out[:, 0], de[:, 0], atol=1e-6) - - def test_unknown_method_raises(self): - cf, de = self._center_frames_and_drift() - with pytest.raises(ValueError, match="Unknown interpolation method"): - comet.interpolate_drift(cf, de, np.arange(100), method="spagetti") - - def test_catmull_rom_output_shape(self): - cf, de = self._center_frames_and_drift() - frame_range = np.arange(10, 90) - out = comet.interpolate_drift( - cf, de, frame_range, method="catmull-rom" - ) - assert out.shape == (len(frame_range), 3) - - def test_catmull_rom_requires_min_points(self): - cf = np.array([0.0, 10.0, 20.0]) # only 3 points - de = np.zeros((3, 3)) - with pytest.raises(ValueError, match="at least 4 points"): - comet.interpolate_drift( - cf, de, np.arange(20), method="catmull-rom" - ) - - def test_linear_passes_through_control_points(self): - cf = np.array([0.0, 25.0, 50.0, 75.0, 100.0]) - de = np.column_stack([cf * 0.01, cf * 0.005, cf * 0.002]) - out = comet.interpolate_drift(cf, de, cf.astype(int), method="linear") - np.testing.assert_allclose(out, de, atol=1e-6) - - def test_linear_clamps_outside_range(self): - # np.interp clamps to edge values beyond the control-point range. - cf = np.array([10.0, 20.0, 30.0]) - de = np.column_stack([cf, cf, cf]) - out = comet.interpolate_drift( - cf, de, np.array([0, 40]), method="linear" - ) - np.testing.assert_allclose(out[0], [10, 10, 10]) - np.testing.assert_allclose(out[1], [30, 30, 30]) - - def test_z_channel_interpolated_independently(self): - cf = np.array([0.0, 25.0, 50.0, 75.0, 100.0]) - de = np.column_stack([np.zeros_like(cf), np.zeros_like(cf), cf * 0.03]) - out = comet.interpolate_drift(cf, de, cf.astype(int), method="cubic") - np.testing.assert_allclose(out[:, 0], 0.0, atol=1e-6) - np.testing.assert_allclose(out[:, 1], 0.0, atol=1e-6) - np.testing.assert_allclose(out[:, 2], de[:, 2], atol=1e-6) - - -# --------------------------------------------------------------------------- -# comet() public API — no-GPU error path -# --------------------------------------------------------------------------- - - -class TestCometNoCuda: - """These tests run on any machine; they verify the failure mode when - CUDA is not present.""" - - def test_raises_runtime_error_without_cuda( - self, minimal_locs_df, minimal_info - ): - if comet._CUDA_AVAILABLE: - pytest.skip( - "CUDA is present; testing no-GPU path only on CPU machines" - ) - with pytest.raises(RuntimeError, match="CUDA"): - comet.comet( - minimal_locs_df, - minimal_info, - locs_per_segment=100, - max_drift_nm=200, - ) - - def test_error_message_mentions_alternatives( - self, minimal_locs_df, minimal_info - ): - if comet._CUDA_AVAILABLE: - pytest.skip( - "CUDA is present; testing no-GPU path only on CPU machines" - ) - with pytest.raises(RuntimeError, match="AIM|RCC"): - comet.comet( - minimal_locs_df, - minimal_info, - locs_per_segment=100, - max_drift_nm=200, - ) - - def test_structured_array_input_accepted_before_gpu_check( - self, minimal_locs_structured, minimal_info - ): - """Numpy structured arrays should be converted to DataFrame before the - CUDA check, so the error is still a RuntimeError (not an AttributeError - from missing .columns).""" - if comet._CUDA_AVAILABLE: - pytest.skip( - "CUDA is present; testing no-GPU path only on CPU machines" - ) - with pytest.raises(RuntimeError, match="CUDA"): - comet.comet( - minimal_locs_structured, - minimal_info, - locs_per_segment=100, - max_drift_nm=200, - ) - - def test_missing_columns_raises_key_error(self, minimal_info): - """Locs missing required columns should raise KeyError, not crash on .columns.""" - if comet._CUDA_AVAILABLE: - pytest.skip( - "CUDA is present; GPU path reached before column check" - ) - bad_locs = pd.DataFrame({"frame": np.arange(10), "x": np.zeros(10)}) - # No 'y' column — should raise KeyError from the column check - with pytest.raises((KeyError, RuntimeError)): - comet.comet( - bad_locs, minimal_info, locs_per_segment=5, max_drift_nm=100 - ) - - def test_z_column_accepted(self, minimal_locs_df, minimal_info): - """3D input (with z) should not crash before the GPU check.""" - if comet._CUDA_AVAILABLE: - pytest.skip( - "CUDA is present; testing no-GPU path only on CPU machines" - ) - df = minimal_locs_df.copy() - df["z"] = np.zeros(len(df), dtype=np.float32) - with pytest.raises(RuntimeError, match="CUDA"): - comet.comet( - df, minimal_info, locs_per_segment=100, max_drift_nm=200 - ) - - def test_missing_pixelsize_raises_on_gpu(self, minimal_locs_df): - """info without 'Pixelsize' should fail with an informative error. - Only meaningful on GPU since the CUDA check fires first on CPU.""" - if not comet._CUDA_AVAILABLE: - pytest.skip("CUDA check fires before pixelsize check on CPU") - bad_info = [{"Frames": 200}] - with pytest.raises((KeyError, RuntimeError, ValueError)): - comet.comet( - minimal_locs_df, - bad_info, - locs_per_segment=100, - max_drift_nm=200, - ) - - -# --------------------------------------------------------------------------- -# comet() public API — full pipeline (GPU required) -# --------------------------------------------------------------------------- - - -@pytest.mark.skipif( - not comet._CUDA_AVAILABLE, - reason="CUDA GPU not available", -) -class TestCometFullPipeline: - """Integration tests that run the full optimisation loop on GPU.""" - - def test_output_shapes(self, minimal_locs_df, minimal_info): - undrifted, new_info, drift = comet.comet( - minimal_locs_df, - minimal_info, - locs_per_segment=100, - max_drift_nm=300, - target_sigma_nm=50, - ) - assert len(undrifted) == len(minimal_locs_df) - assert set(undrifted.columns) >= {"x", "y", "frame"} - assert "x" in drift.columns and "y" in drift.columns - - def test_drift_covers_frame_range(self, minimal_locs_df, minimal_info): - frame0 = minimal_locs_df["frame"] - minimal_locs_df["frame"].min() - _, _, drift = comet.comet( - minimal_locs_df, - minimal_info, - locs_per_segment=100, - max_drift_nm=300, - target_sigma_nm=50, - ) - assert len(drift) == int(frame0.max()) + 1 - - def test_new_info_appended(self, minimal_locs_df, minimal_info): - _, new_info, _ = comet.comet( - minimal_locs_df, - minimal_info, - locs_per_segment=100, - max_drift_nm=300, - target_sigma_nm=50, - ) - assert len(new_info) == len(minimal_info) + 1 - assert "COMET" in new_info[-1]["Generated by"] - - def test_structured_array_gives_same_result_as_dataframe( - self, minimal_locs_df, minimal_locs_structured, minimal_info - ): - locs_df, _, drift_df = comet.comet( - minimal_locs_df, - minimal_info, - locs_per_segment=100, - max_drift_nm=300, - target_sigma_nm=50, - ) - locs_arr, _, drift_arr = comet.comet( - minimal_locs_structured, - minimal_info, - locs_per_segment=100, - max_drift_nm=300, - target_sigma_nm=50, - ) - np.testing.assert_allclose( - drift_df["x"].to_numpy(), drift_arr["x"].to_numpy(), rtol=1e-5 - ) From 7a88437d7952e5f7b77c8d0e87d1a6550022be36 Mon Sep 17 00:00:00 2001 From: rafalkowalewski1 Date: Wed, 10 Jun 2026 09:14:36 +0200 Subject: [PATCH 10/19] Fixed "Best fitting combination" button in SPINNA (#676) --- changelog.md | 5 ++++- picasso/gui/spinna.py | 2 ++ 2 files changed, 6 insertions(+), 1 deletion(-) diff --git a/changelog.md b/changelog.md index f104bd0d..744109f0 100644 --- a/changelog.md +++ b/changelog.md @@ -1,6 +1,9 @@ # Changelog -Last change: 29-MAY-2026 CEST +Last change: 10-JUN-2026 CEST + +## 0.10.2 +- Fixed "Best fitting combination" button in SPINNA ([#676](https://github.com/jungmannlab/picasso/issues/676)) ## 0.10.1 diff --git a/picasso/gui/spinna.py b/picasso/gui/spinna.py index 5cbf1d75..92646d2f 100644 --- a/picasso/gui/spinna.py +++ b/picasso/gui/spinna.py @@ -3893,7 +3893,9 @@ def update_prop_str_input_spins( def retrieve_results(): for i, ps in enumerate(prop_str): spin = self.prop_str_input_spins[i] + spin.blockSignals(True) spin.setValue(ps) + spin.blockSignals(False) button.released.connect(retrieve_results) self.prop_str_input.add_widget( From 7e1d16facc3fe22b3fc858a7a1742c21bee3540c Mon Sep 17 00:00:00 2001 From: rafalkowalewski1 Date: Wed, 10 Jun 2026 16:30:35 +0200 Subject: [PATCH 11/19] Fixed 3D calibration when frame range is user-specified --- changelog.md | 1 + picasso/gui/localize.py | 1 + picasso/zfit.py | 30 ++++++++++++++++++++++++++---- 3 files changed, 28 insertions(+), 4 deletions(-) diff --git a/changelog.md b/changelog.md index 744109f0..bfafc8b6 100644 --- a/changelog.md +++ b/changelog.md @@ -4,6 +4,7 @@ Last change: 10-JUN-2026 CEST ## 0.10.2 - Fixed "Best fitting combination" button in SPINNA ([#676](https://github.com/jungmannlab/picasso/issues/676)) +- Fixed 3D calibration when frame range is user-specified ## 0.10.1 diff --git a/picasso/gui/localize.py b/picasso/gui/localize.py index 33717b15..84f7d2a5 100644 --- a/picasso/gui/localize.py +++ b/picasso/gui/localize.py @@ -2572,6 +2572,7 @@ def on_fit_finished( step, self.parameters_dialog.magnification_factor.value(), path=path, + frame_bounds=self.frame_range, ) dt = time.time() - t0 if dt > lib.SOUND_NOTIFICATION_DURATION: diff --git a/picasso/zfit.py b/picasso/zfit.py index 0dd6a8b8..27576db7 100644 --- a/picasso/zfit.py +++ b/picasso/zfit.py @@ -49,6 +49,7 @@ def calibrate_z( d: float, magnification_factor: float, path: str | None = None, + frame_bounds: tuple[int, int] | None = None, ) -> dict: """Given localizations of a calibration sample (e.g., gold beads at different z positions), calibrate the z-axis by fitting a polynomial @@ -72,6 +73,12 @@ def calibrate_z( path : str, optional Path to save the calibration data as a YAML file. If None, the calibration data will not be saved. Default is None. + frame_bounds : tuple, optional + Minimum and maximum frame numbers to consider for the + calibration. If None, all frames are used. If only min or max + is to be specified, the other is to be set to None, for example, + ``(5, None)`` sets minimum frame to 5 without maximum frame. + Default is None. Returns ------- @@ -85,6 +92,16 @@ def calibrate_z( frame_range = np.arange(n_frames) z_range = -(frame_range * d - range / 2) # negative so that the # first frames of a bottom-to-up scan are positive z coordinates. + if frame_bounds is not None: + frame_min, frame_max = frame_bounds + frame_min = frame_min or 0 + frame_max = frame_max or (n_frames - 1) + # frame bounds are inclusive, like in picasso.localize + frame_range = frame_range[frame_min : frame_max + 1] + z_range = z_range[frame_min : frame_max + 1] + locs = locs[ + (locs["frame"] >= frame_min) & (locs["frame"] <= frame_max) + ] mean_sx = np.array( [locs["sx"][locs["frame"] == _].mean() for _ in frame_range] @@ -99,8 +116,11 @@ def calibrate_z( [locs["sy"][locs["frame"] == _].var() for _ in frame_range] ) - keep_x = (locs["sx"] - mean_sx[locs["frame"]]) ** 2 < var_sx[locs["frame"]] - keep_y = (locs["sy"] - mean_sy[locs["frame"]]) ** 2 < var_sy[locs["frame"]] + # index of each localization's frame in the (possibly bounded) + # frame_range + frame_idx = locs["frame"] - frame_range[0] + keep_x = (locs["sx"] - mean_sx[frame_idx]) ** 2 < var_sx[frame_idx] + keep_y = (locs["sy"] - mean_sy[frame_idx]) ** 2 < var_sy[frame_idx] keep = keep_x & keep_y locs = locs[keep] @@ -134,6 +154,7 @@ def calibrate_z( "Step size in nm": float(d), "Magnification factor": float(magnification_factor), "Path": path if path is not None else "N/A", + "Frame bounds": frame_bounds, } if path is not None: with open(path, "w") as f: @@ -142,6 +163,7 @@ def calibrate_z( # pixelsize does not matter here anyway locs = _fit_z(locs, info, calibration, magnification_factor, pixelsize=130) locs["z"] /= magnification_factor + frame_idx = locs["frame"] - frame_range[0] plt.figure(figsize=(18, 10)) @@ -186,7 +208,7 @@ def calibrate_z( plt.legend(loc="best") ax = plt.subplot(234) - plt.plot(z_range[locs["frame"]], locs["z"], ".k", alpha=0.1) + plt.plot(z_range[frame_idx], locs["z"], ".k", alpha=0.1) plt.plot( [z_range.min(), z_range.max()], [z_range.min(), z_range.max()], @@ -201,7 +223,7 @@ def calibrate_z( plt.legend(loc="best") ax = plt.subplot(235) - deviation = locs["z"] - z_range[locs["frame"]] + deviation = locs["z"] - z_range[frame_idx] bins = lib.calculate_optimal_bins(deviation, max_n_bins=1000) plt.hist(deviation, bins) plt.xlabel("Deviation to true position") From e7aea5120586169ae9ea5e7730362986612686bc Mon Sep 17 00:00:00 2001 From: rafalkowalewski1 Date: Wed, 10 Jun 2026 16:44:47 +0200 Subject: [PATCH 12/19] add tests for frame-bounded calibration --- tests/test_zfit.py | 215 +++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 215 insertions(+) diff --git a/tests/test_zfit.py b/tests/test_zfit.py index cb0dd808..b6d0d588 100644 --- a/tests/test_zfit.py +++ b/tests/test_zfit.py @@ -497,6 +497,221 @@ def test_writes_yaml_and_png_when_path_given( assert loaded["Step size in nm"] == calib["Step size in nm"] +class TestCalibrateZFrameBounds: + """``frame_bounds`` restricts which frames enter the calibration. + Bounds are inclusive on both ends, matching ``picasso.localize``.""" + + N_FRAMES = 50 + D = 10.0 # nm step + + @pytest.fixture(autouse=True) + def _no_show(self, monkeypatch): + monkeypatch.setattr("matplotlib.pyplot.show", lambda *a, **k: None) + + @pytest.fixture + def bead_stack(self): + """Same synthetic bead stack as in ``TestCalibrateZ``: sx/sy + driven by known polynomials of stage z plus tiny noise.""" + rng = np.random.default_rng(0) + rows = [] + z_total = (self.N_FRAMES - 1) * self.D + for fi in range(self.N_FRAMES): + z = -(fi * self.D - z_total / 2) + sx_mean = 1.5 + 1e-3 * z + 1e-5 * z**2 + sy_mean = 1.5 - 1e-3 * z + 1e-5 * z**2 + for _ in range(40): + rows.append( + { + "frame": fi, + "x": 16.0, + "y": 16.0, + "sx": sx_mean + rng.normal(0, 0.02), + "sy": sy_mean + rng.normal(0, 0.02), + "photons": 5000.0, + "bg": 10.0, + "lpx": 0.01, + "lpy": 0.01, + } + ) + locs = pd.DataFrame(rows) + info = [ + { + "Frames": self.N_FRAMES, + "Pixelsize": 130, + "Width": 32, + "Height": 32, + } + ] + return locs, info + + def _polyfit_lengths(self, monkeypatch): + """Spy on ``np.polyfit`` to record how many frames enter each + calibration fit.""" + lengths = [] + orig = np.polyfit + + def spy(x, y, deg, **kwargs): + lengths.append(len(x)) + return orig(x, y, deg, **kwargs) + + monkeypatch.setattr(np, "polyfit", spy) + return lengths + + def test_none_and_none_none_equivalent(self, bead_stack): + """``frame_bounds=(None, None)`` must behave exactly like + ``frame_bounds=None`` (all frames used).""" + locs, info = bead_stack + calib_none = zfit.calibrate_z( + locs, info, self.D, magnification_factor=0.79, frame_bounds=None + ) + calib_nn = zfit.calibrate_z( + locs, + info, + self.D, + magnification_factor=0.79, + frame_bounds=(None, None), + ) + np.testing.assert_allclose( + calib_none["X Coefficients"], calib_nn["X Coefficients"] + ) + np.testing.assert_allclose( + calib_none["Y Coefficients"], calib_nn["Y Coefficients"] + ) + + def test_full_range_bounds_equivalent_to_none(self, bead_stack): + """``(0, n_frames - 1)`` covers all frames, so it must match the + unbounded calibration.""" + locs, info = bead_stack + calib_none = zfit.calibrate_z( + locs, info, self.D, magnification_factor=0.79 + ) + calib_full = zfit.calibrate_z( + locs, + info, + self.D, + magnification_factor=0.79, + frame_bounds=(0, self.N_FRAMES - 1), + ) + np.testing.assert_allclose( + calib_none["X Coefficients"], calib_full["X Coefficients"] + ) + np.testing.assert_allclose( + calib_none["Y Coefficients"], calib_full["Y Coefficients"] + ) + + def test_bounds_are_inclusive(self, bead_stack, monkeypatch): + """``(10, 39)`` keeps frames 10..39 inclusive -> 30 frames enter + every polynomial fit.""" + locs, info = bead_stack + lengths = self._polyfit_lengths(monkeypatch) + zfit.calibrate_z( + locs, + info, + self.D, + magnification_factor=0.79, + frame_bounds=(10, 39), + ) + assert lengths and all(n == 30 for n in lengths) + + def test_one_sided_bounds(self, bead_stack, monkeypatch): + """``(None, max)`` and ``(min, None)`` each leave the other side + unbounded.""" + locs, info = bead_stack + lengths = self._polyfit_lengths(monkeypatch) + zfit.calibrate_z( + locs, + info, + self.D, + magnification_factor=0.79, + frame_bounds=(None, 29), # frames 0..29 + ) + assert lengths and all(n == 30 for n in lengths) + lengths.clear() + zfit.calibrate_z( + locs, + info, + self.D, + magnification_factor=0.79, + frame_bounds=(20, None), # frames 20..49 + ) + assert lengths and all(n == 30 for n in lengths) + + def test_bounds_including_last_frame(self, bead_stack): + """Regression test: a nonzero minimum together with the last + frame as maximum used to raise IndexError (localizations' + frame numbers were used as indices into the sliced per-frame + arrays without offsetting).""" + locs, info = bead_stack + calib = zfit.calibrate_z( + locs, + info, + self.D, + magnification_factor=0.79, + frame_bounds=(10, self.N_FRAMES - 1), + ) + assert len(calib["X Coefficients"]) == 7 + assert len(calib["Y Coefficients"]) == 7 + + def test_locs_outside_bounds_are_ignored(self, bead_stack): + """Passing the full locs with bounds must equal passing locs + already restricted to those frames.""" + locs, info = bead_stack + bounds = (10, 39) + pre_filtered = locs[ + (locs["frame"] >= bounds[0]) & (locs["frame"] <= bounds[1]) + ] + calib_full = zfit.calibrate_z( + locs, + info, + self.D, + magnification_factor=0.79, + frame_bounds=bounds, + ) + calib_pre = zfit.calibrate_z( + pre_filtered, + info, + self.D, + magnification_factor=0.79, + frame_bounds=bounds, + ) + np.testing.assert_allclose( + calib_full["X Coefficients"], calib_pre["X Coefficients"] + ) + np.testing.assert_allclose( + calib_full["Y Coefficients"], calib_pre["Y Coefficients"] + ) + + def test_bounded_calibration_differs_from_unbounded(self, bead_stack): + """Restricting the frames must actually change the fit.""" + locs, info = bead_stack + calib_none = zfit.calibrate_z( + locs, info, self.D, magnification_factor=0.79 + ) + calib_bounded = zfit.calibrate_z( + locs, + info, + self.D, + magnification_factor=0.79, + frame_bounds=(10, 39), + ) + assert not np.allclose( + calib_none["X Coefficients"], calib_bounded["X Coefficients"] + ) + + def test_frame_bounds_stored_in_calibration(self, bead_stack): + locs, info = bead_stack + calib = zfit.calibrate_z( + locs, + info, + self.D, + magnification_factor=0.79, + frame_bounds=(10, 39), + ) + assert calib["Frame bounds"] == (10, 39) + calib = zfit.calibrate_z(locs, info, self.D, magnification_factor=0.79) + assert calib["Frame bounds"] is None + + # --------------------------------------------------------------------------- # locs_from_futures # --------------------------------------------------------------------------- From 77288320632a8a38ccbf05694d00c4f19ae2795f Mon Sep 17 00:00:00 2001 From: rafalkowalewski1 Date: Wed, 10 Jun 2026 21:05:59 +0200 Subject: [PATCH 13/19] fixed redoing 3D calibration when identification parameters change --- changelog.md | 1 + picasso/gui/localize.py | 28 +++++++++++++++++++++++++++- 2 files changed, 28 insertions(+), 1 deletion(-) diff --git a/changelog.md b/changelog.md index bfafc8b6..92eb567d 100644 --- a/changelog.md +++ b/changelog.md @@ -5,6 +5,7 @@ Last change: 10-JUN-2026 CEST ## 0.10.2 - Fixed "Best fitting combination" button in SPINNA ([#676](https://github.com/jungmannlab/picasso/issues/676)) - Fixed 3D calibration when frame range is user-specified +- Fixed redoing 3D calibration when identification parameters change ## 0.10.1 diff --git a/picasso/gui/localize.py b/picasso/gui/localize.py index 84f7d2a5..85d6848c 100644 --- a/picasso/gui/localize.py +++ b/picasso/gui/localize.py @@ -2103,6 +2103,8 @@ def _clean_up_external_ids(self) -> None: self.last_identification_info = { "Box Size": self.parameters_dialog.box_spinbox.value(), "Min. Net Gradient": self.parameters_dialog.mng_slider.value(), + "ROI": self.view.roi, + "Frame bounds": self.frame_range, } self.ready_for_fit = True self.draw_frame() @@ -2551,6 +2553,7 @@ def on_fit_finished( playsound(sound_path, block=False) base, ext = os.path.splitext(self.movie_path) if calibrate_z: + self.parameters_dialog.gpufit_checkbox.setDisabled(False) step, ok = QtWidgets.QInputDialog.getDouble( self, "3D Calibration", @@ -2903,11 +2906,34 @@ def localize(self, calibrate_z: bool = False) -> None: Default is False. """ self.parameters_dialog.gpufit_checkbox.setDisabled(True) - if self.identifications is not None and calibrate_z is True: + if ( + calibrate_z + and self.identifications is not None + and self.ready_for_fit + and not self.identifications_outdated() + ): + # reuse existing identifications (e.g., loaded picks/locs) self.fit(calibrate_z=calibrate_z) else: self.identify(fit_afterwards=True, calibrate_z=calibrate_z) + def identifications_outdated(self) -> bool: + """Check whether the identification settings (box size, min. net + gradient, ROI, frame range) have changed since + ``self.identifications`` was created.""" + last = self.last_identification_info + if last is None: + return True + if any( + last.get(key) != value for key, value in self.parameters.items() + ): + return True + if last.get("ROI") != self.view.roi: + return True + if last.get("Frame bounds") != self.frame_range: + return True + return False + def select_locs_columns(self) -> None: """Select only the columns that are checked in the corresponding dialog.""" From e4e3777a0212e34124e853b17bdcc38d96b05deb Mon Sep 17 00:00:00 2001 From: rafalkowalewski1 Date: Thu, 11 Jun 2026 13:14:13 +0200 Subject: [PATCH 14/19] change rotations to quaterions (#673, #674, #675) --- changelog.md | 3 +- picasso/gui/render.py | 6 +- picasso/gui/rotation.py | 352 ++++++++++++++++++++++++----------- picasso/render.py | 401 +++++++++++++++++++++++++++------------- tests/test_render.py | 215 ++++++++++++++++++++- 5 files changed, 740 insertions(+), 237 deletions(-) diff --git a/changelog.md b/changelog.md index 92eb567d..a3321b32 100644 --- a/changelog.md +++ b/changelog.md @@ -1,8 +1,9 @@ # Changelog -Last change: 10-JUN-2026 CEST +Last change: 11-JUN-2026 CEST ## 0.10.2 +- Rotations use quaterions for unambiguous workflow, fixing bugs [#673](https://github.com/jungmannlab/picasso/issues/673), [#674](https://github.com/jungmannlab/picasso/issues/674) and [#675](https://github.com/jungmannlab/picasso/issues/675) - Fixed "Best fitting combination" button in SPINNA ([#676](https://github.com/jungmannlab/picasso/issues/676)) - Fixed 3D calibration when frame range is user-specified - Fixed redoing 3D calibration when identification parameters change diff --git a/picasso/gui/render.py b/picasso/gui/render.py index d5f31ee5..b985aded 100644 --- a/picasso/gui/render.py +++ b/picasso/gui/render.py @@ -12831,9 +12831,9 @@ def open_rotated_locs(self) -> None: self.tools_settings_dialog.pick_side_length.setValue( self.view.infos[0][-1]["Pick size (nm)"] ) - self.window_rot.view_rot.angx = self.view.infos[0][-1]["angx"] - self.window_rot.view_rot.angy = self.view.infos[0][-1]["angy"] - self.window_rot.view_rot.angz = self.view.infos[0][-1]["angz"] + self.window_rot.view_rot.load_saved_rotation( + self.view.infos[0][-1] + ) self.rot_win() def resize_view_to_fov(self, w: float, h: float) -> None: diff --git a/picasso/gui/rotation.py b/picasso/gui/rotation.py index f7a59ce7..68c421d0 100644 --- a/picasso/gui/rotation.py +++ b/picasso/gui/rotation.py @@ -311,9 +311,14 @@ class AnimationDialog(lib.Dialog): Click to delete the last saved position. fps : QSpinBox Contains frames per second used in the animation. - positions : list - Contains all positions in the animation sequence; each includes - 3 rotations angles and viewport. + positions : list of dict + Contains all positions in the animation sequence. Each holds + the rotation (scipy Rotation, key ``"R"``), the accumulated + rotation vector (key ``"rotvec"``, radians, scipy convention, + keeps full turns), the rotation accumulated since the previous + position (key ``"segment_rotvec"``, radians, scipy convention; + may encode turns beyond 180 degrees) and the viewport (key + ``"viewport"``). rows : list of dict One entry per saved position, holding the row's widgets: ``p_label``, ``angle_label``, ``show_btn``, ``d_label``, @@ -352,7 +357,10 @@ def __init__(self, window: QtWidgets.QMainWindow) -> None: # Header: current position header = QtWidgets.QHBoxLayout() cp_label = QtWidgets.QLabel("Current position:") - cp_label.setToolTip("Current rotation angles in x, y, z (deg).") + cp_label.setToolTip( + "Current rotation in x, y, z (deg). The angles keep track " + "of full turns, e.g. 720 encodes two full rotations." + ) header.addWidget(cp_label) angx = np.round(self.window.view_rot.angx * 180 / np.pi, 1) angy = np.round(self.window.view_rot.angy * 180 / np.pi, 1) @@ -453,25 +461,35 @@ def add_position(self, freeze: bool = False) -> None: True when the new position is the same as the last one, i.e., when self.stay is clicked. """ - # check that the viewport or angle(s) have changed - cond1 = cond2 = cond3 = False + view = self.window.view_rot + # rotation accumulated since the last position (keeps full + # turns, e.g. 720 deg); zero when freezing in place + segment_rotvec = ( + np.zeros(3) if freeze else view.rotation_since_anchor() + ) + + # check that the viewport or the rotation have changed if not freeze and self.positions: - cond1 = self.window.view_rot.angx == self.positions[-1][0] - cond2 = self.window.view_rot.angy == self.positions[-1][1] - cond3 = self.window.view_rot.angz == self.positions[-1][2] - cond4 = self.window.view_rot.viewport == self.positions[-1][3] - if all([cond1, cond2, cond3, cond4]): + last = self.positions[-1] + same_rotation = ( + view.rotation * last["R"].inv() + ).magnitude() < 1e-9 + no_turns = np.linalg.norm(segment_rotvec) < 1e-9 + same_viewport = view.viewport == last["viewport"] + if same_rotation and no_turns and same_viewport: return # add a new position to the attribute self.positions.append( - [ - self.window.view_rot.angx, - self.window.view_rot.angy, - self.window.view_rot.angz, - self.window.view_rot.viewport, - ] + { + "R": view.rotation, + "rotvec": view._rotvec.copy(), + "segment_rotvec": segment_rotvec, + "viewport": view.viewport, + } ) + # track the next segment's rotation from this position onwards + view.reset_rotation_anchor() index = len(self.positions) - 1 grid_row = index @@ -483,9 +501,9 @@ def add_position(self, freeze: bool = False) -> None: ) self.rows_layout.addWidget(p_label, grid_row, 0) - angx = np.round(self.window.view_rot.angx * 180 / np.pi, 1) - angy = np.round(self.window.view_rot.angy * 180 / np.pi, 1) - angz = np.round(self.window.view_rot.angz * 180 / np.pi, 1) + angx = np.round(view.angx * 180 / np.pi, 1) + angy = np.round(view.angy * 180 / np.pi, 1) + angz = np.round(view.angz * 180 / np.pi, 1) angle_label = QtWidgets.QLabel("{}, {}, {}".format(angx, angy, angz)) self.rows_layout.addWidget(angle_label, grid_row, 1) @@ -507,14 +525,12 @@ def add_position(self, freeze: bool = False) -> None: duration.setDecimals(2) self.rows_layout.addWidget(duration, grid_row, 4) - # calculate recommended duration - if not freeze and not all([cond1, cond2, cond3]): - dx = self.positions[-1][0] - self.positions[-2][0] - dy = self.positions[-1][1] - self.positions[-2][1] - dz = self.positions[-1][2] - self.positions[-2][2] - dmax = np.max(np.abs([dx, dy, dz])) + # calculate recommended duration from the total rotation + # (including full turns) at the requested rotation speed + total_angle = float(np.linalg.norm(segment_rotvec)) + if not freeze and total_angle > 1e-9: rot_speed = self.rot_speed.value() * np.pi / 180 - duration.setValue(dmax / rot_speed) + duration.setValue(total_angle / rot_speed) self.rows.append( { @@ -553,10 +569,11 @@ def retrieve_position(self, i: int) -> None: Index of the position to be displayed. """ if i <= len(self.positions) - 1: - self.window.view_rot.angx = self.positions[i][0] - self.window.view_rot.angy = self.positions[i][1] - self.window.view_rot.angz = self.positions[i][2] - self.window.view_rot.update_scene(viewport=self.positions[i][3]) + position = self.positions[i] + self.window.view_rot.set_rotation( + position["R"], rotvec=position["rotvec"] + ) + self.window.view_rot.update_scene(viewport=position["viewport"]) def build_animation(self) -> None: """Create an animation as an .mp4 file using the positions from @@ -592,12 +609,17 @@ def build_animation(self) -> None: ) adjust_display_pixel = disp_dlg.dynamic_disp_px.isChecked() intensities = self.window.window.view.read_relative_intensities() + positions = [(p["R"], p["viewport"]) for p in self.positions] + segment_rotations = [ + p["segment_rotvec"] for p in self.positions[1:] + ] render.build_animation( path, locs, infos, - positions=self.positions, + positions=positions, durations=durations, + segment_rotations=segment_rotations, disp_px_size=disp_dlg.disp_px_size.value(), image_size=( self.window.view_rot.width(), @@ -629,8 +651,10 @@ class ViewRotation(QtWidgets.QLabel): Attributes ---------- angx, angy, angz : float - Current rotation angle around x, y, and z axes (read/write - properties; backed by ``self._R``). + Accumulated rotation around the x, y and z axes (radians) in + the codebase's angle convention (read-only properties; backed + by ``self._rotvec``). They keep track of full turns, e.g. read + 720 degrees after two full rotations. block_x, block_y, block_z : bool True if rotate only around x, y, or z axis respectively. group_color : lib.IntArray1D @@ -660,8 +684,21 @@ class ViewRotation(QtWidgets.QLabel): qimage : QImage Current image of rendered locs, picks and other drawings. _R : scipy.spatial.transform.Rotation - Source of truth for the current rotation. ``angx/angy/angz`` - are derived from this via ``_codebase_euler``. + Source of truth for the current rotation (quaternion-backed). + _rotvec : np.ndarray + Accumulated rotation vector (radians, scipy convention) - the + sum of all applied rotation vectors (a path integral), so full + turns are preserved, e.g. two full rotations around z read + (0, 0, 4 pi). For rotations around a single axis it represents + ``_R`` exactly; in general the exact orientation is ``_R``. + ``angx/angy/angz`` are derived from this. + _anchor_R : scipy.spatial.transform.Rotation + Reference orientation for per-segment rotation tracking (set + when an animation position is added or shown). + _anchor_rotvec : np.ndarray + Unwrapped rotation vector (radians, scipy convention) of the + rotation accumulated since ``_anchor_R``; used to encode + rotations beyond 180 degrees in animation segments. _last_mouse_x, _last_mouse_y : int Previous mouse position (Qt coords) during a trackball drag. viewport : tuple @@ -674,6 +711,9 @@ def __init__(self, window: QtWidgets.QMainWindow) -> None: super().__init__(window) self.window = window self._R = Rotation.identity() + self._rotvec = np.zeros(3) + self._anchor_R = Rotation.identity() + self._anchor_rotvec = np.zeros(3) self._pan_z = 0.0 self._last_mouse_x = 0 self._last_mouse_y = 0 @@ -692,37 +732,139 @@ def __init__(self, window: QtWidgets.QMainWindow) -> None: self.block_z = False self.setFocusPolicy(QtCore.Qt.FocusPolicy.ClickFocus) - # --- rotation angles --- # - def _codebase_euler(self) -> tuple[float, float, float]: - e = self._R.as_euler("XYZ") - return -float(e[0]), float(e[1]), float(e[2]) - + # --- rotation state --- # + # The codebase's angle convention (see ``render.rotation_matrix``) + # flips the sign of the x rotation relative to scipy's right-hand + # rule, i.e., rotation_matrix(a, 0, 0).as_rotvec() == (-a, 0, 0). + # The angx/angy/angz properties report the accumulated rotation + # vector in that convention. + # + # ``_rotvec`` accumulates the applied rotation vectors (a path + # integral) rather than re-deriving angles from ``_R``. An absolute + # orientation cannot count full turns - spinning a tilted structure + # by 360 degrees returns to the same orientation - so this is the + # only way the displayed angles can keep track of the number of + # turns (e.g. read 720 after two full rotations). For rotations + # around a single axis the accumulated values match the orientation + # exactly; for composed rotations around different axes they are a + # record of the applied rotations and the exact orientation is + # given by ``_R``. @property def angx(self) -> float: - return self._codebase_euler()[0] - - @angx.setter - def angx(self, value: float) -> None: - _, b, c = self._codebase_euler() - self._R = render.rotation_matrix(float(value), b, c) + return -float(self._rotvec[0]) @property def angy(self) -> float: - return self._codebase_euler()[1] - - @angy.setter - def angy(self, value: float) -> None: - a, _, c = self._codebase_euler() - self._R = render.rotation_matrix(a, float(value), c) + return float(self._rotvec[1]) @property def angz(self) -> float: - return self._codebase_euler()[2] + return float(self._rotvec[2]) + + @property + def rotation(self) -> Rotation: + """Current rotation of the localizations.""" + return self._R - @angz.setter - def angz(self, value: float) -> None: - a, b, _ = self._codebase_euler() - self._R = render.rotation_matrix(a, b, float(value)) + def set_rotation( + self, + R: Rotation, + rotvec: np.ndarray | None = None, + ) -> None: + """Set the absolute rotation and reset the per-segment anchor. + + Parameters + ---------- + R : scipy.spatial.transform.Rotation + New rotation. + rotvec : np.ndarray, optional + Accumulated rotation vector to restore as the displayed + angles (radians, scipy convention), e.g. saved earlier + from ``self._rotvec``; keeps full-turn information. If + None, the shortest rotation vector of ``R`` is used. + """ + self._R = R + if rotvec is None: + self._rotvec = R.as_rotvec() + else: + self._rotvec = np.asarray(rotvec, dtype=float).copy() + self.reset_rotation_anchor() + + def reset_rotation_anchor(self) -> None: + """Start tracking the rotation accumulated from the current + orientation (used for animation segments).""" + self._anchor_R = self._R + self._anchor_rotvec = np.zeros(3) + + def rotation_since_anchor(self) -> np.ndarray: + """Unwrapped rotation vector (radians, scipy convention) + accumulated since the last anchor; magnitude may exceed pi if + rotated beyond 180 degrees.""" + return self._anchor_rotvec.copy() + + def apply_rotation( + self, + rotvec: np.ndarray, + frame: str = "world", + ) -> None: + """Compose a rotation onto the current one and track turns. + + Parameters + ---------- + rotvec : np.ndarray + Rotation vector (radians, scipy convention). Its magnitude + may exceed pi; full turns are preserved in the displayed + angles and in the per-segment tracking. + frame : {"world", "object"}, optional + Frame in which ``rotvec`` is given. "world": the fixed + world/screen frame (rotation composed as ``delta * R``). + "object": the data's own (rotated) frame, i.e., rotation + around the data axes shown by the axes icon (composed as + ``R * delta``). Default is "world". + """ + rotvec = np.asarray(rotvec, dtype=float) + magnitude = float(np.linalg.norm(rotvec)) + if magnitude < 1e-12: + return + # accumulate the displayed angles (path integral; an absolute + # orientation cannot count full turns, so the applied rotation + # vectors are summed instead) + self._rotvec = self._rotvec + rotvec + # apply in sub-steps small enough for unambiguous unwrapping of + # the per-segment rotation + n_steps = max(1, int(np.ceil(magnitude / (np.pi / 2)))) + step = Rotation.from_rotvec(rotvec / n_steps) + for _ in range(n_steps): + if frame == "object": + self._R = self._R * step + else: + self._R = step * self._R + relative = self._R * self._anchor_R.inv() + self._anchor_rotvec = render.closest_rotvec( + relative, self._anchor_rotvec + ) + + def load_saved_rotation(self, info: dict) -> None: + """Restore the rotation saved by + ``RotationWindow.save_locs_rotated``. + + New files store a quaternion together with the accumulated + rotation angles (keeping full turns); legacy files store Euler + angles (angx, angy, angz, radians, codebase convention). + """ + angles = ( + info.get("angx", 0.0), + info.get("angy", 0.0), + info.get("angz", 0.0), + ) + if "Quaternion (x, y, z, w)" in info: + R = Rotation.from_quat(info["Quaternion (x, y, z, w)"]) + # angx is stored in the codebase convention (x negated + # relative to scipy's rotation vector) + rotvec = np.array([-angles[0], angles[1], angles[2]]) + self.set_rotation(R, rotvec=rotvec) + else: # legacy: Euler angles + self.set_rotation(render.rotation_matrix(*angles)) def sizeHint(self) -> QtCore.QSize: return QtCore.QSize(*self._size_hint) @@ -851,9 +993,6 @@ def render_scene( viewport : tuple, optional Viewport to be rendered ``((y_min, x_min), (y_max, x_max))``. If None, takes current viewport. - ang : tuple, optional - Rotation angles to be rendered. If None, takes the current - angles. autoscale : bool, optional If True, optimally adjust contrast. use_cache : bool, optional @@ -882,7 +1021,7 @@ def render_scene( locs=locs, info=infos, **kwargs, - ang=(self.angx, self.angy, self.angz), + ang=self._R, contrast=contrast, invert_colors=self.window.dataset_dialog.wbackground.isChecked(), single_channel_colormap=cmap, @@ -1053,9 +1192,7 @@ def draw_rotation(self, image: QtGui.QImage) -> QtGui.QImage: Image with the drawn rotation axes icon. """ if self.window.rotation_action.isChecked(): - image = render.draw_rotation( - image=image, ang=(self.angx, self.angy, self.angz) - ) + image = render.draw_rotation(image=image, ang=self._R) return image def draw_rotation_angles(self, image: QtGui.QImage) -> QtGui.QImage: @@ -1098,7 +1235,15 @@ def draw_points(self, image: QtGui.QImage) -> QtGui.QImage: ) def rotation_input(self) -> None: - """Ask the user to input 3 rotation angles manually.""" + """Ask the user to input 3 rotation angles manually. + + The rotations are applied sequentially around the data's own + (rotated) x, y and z axes - the axes shown by the axes icon - + so each displayed angle changes by exactly the entered amount + regardless of the current orientation. Angles beyond +/- 180 + degrees are preserved, e.g. 720 degrees encodes two full turns + (relevant for animations). + """ angx, ok = QtWidgets.QInputDialog.getDouble( self, "Rotation angle x", @@ -1123,15 +1268,24 @@ def rotation_input(self) -> None: decimals=2, ) if ok3: - self.angx += np.pi * angx / 180 - self.angy += np.pi * angy / 180 - self.angz += np.pi * angz / 180 + # codebase convention: x angle is the negative of + # scipy's right-handed x rotation; "object" frame + # rotates around the data's own axes + self.apply_rotation( + [-np.radians(angx), 0.0, 0.0], frame="object" + ) + self.apply_rotation( + [0.0, np.radians(angy), 0.0], frame="object" + ) + self.apply_rotation( + [0.0, 0.0, np.radians(angz)], frame="object" + ) self.update_scene() def delete_rotation(self) -> None: """Reset rotation and any accumulated pan offset.""" - self._R = Rotation.identity() + self.set_rotation(Rotation.identity()) self._pan_z = 0.0 self.update_scene() @@ -1182,24 +1336,18 @@ def fit_in_view_rotated(self, get_viewport: bool = False) -> None: self.update_scene() def xy_projection(self) -> None: - """Set angles to 0 to get XY projection.""" - self.angx = 0 - self.angy = 0 - self.angz = 0 + """Reset rotation to get XY projection.""" + self.set_rotation(Rotation.identity()) self.update_scene() def xz_projection(self) -> None: - """Set angles to get XZ projection.""" - self.angx = np.pi / 2 - self.angy = 0 - self.angz = 0 + """Set rotation to get XZ projection.""" + self.set_rotation(render.rotation_matrix(np.pi / 2, 0, 0)) self.update_scene() def yz_projection(self) -> None: - """Set angles to get YZ projection.""" - self.angx = 0 - self.angy = np.pi / 2 - self.angz = 0 + """Set rotation to get YZ projection.""" + self.set_rotation(render.rotation_matrix(0, np.pi / 2, 0)) self.update_scene() def _arrow_pan(self, sx: float, sy: float) -> None: @@ -1332,12 +1480,13 @@ def _rotate_drag(self, event: QtGui.QMouseEvent) -> None: ax = 0.0 # Axis locks: pressing X/Y/Z constrains rotation to the - # corresponding *world* axis. Project the screen-frame rotation - # vector into world frame, zero the components we don't want, and - # rotate it back. + # corresponding *data* axis (the axes shown by the axes icon). + # Project the screen-frame rotation vector into the data frame, + # keep only the locked component and apply it around the data + # axis; the displayed angle then changes only for that axis. if self.block_x or self.block_y or self.block_z: - v_screen = np.array([ax, ay, az], dtype=float) - v_world = self._R.inv().apply(v_screen) + v_screen = render.rotation_matrix(ax, ay, az).as_rotvec() + v_object = self._R.inv().apply(v_screen) keep = np.array( [ 1.0 if self.block_x else 0.0, @@ -1345,20 +1494,10 @@ def _rotate_drag(self, event: QtGui.QMouseEvent) -> None: 1.0 if self.block_z else 0.0, ] ) - v_world = v_world * keep - v_screen = self._R.apply(v_world) - ax, ay, az = ( - float(v_screen[0]), - float(v_screen[1]), - float(v_screen[2]), - ) - - # The codebase's render.rotation_matrix flips the sign of the X - # rotation relative to scipy's right-hand-rule convention; using - # it here keeps screen-aligned tilts feeling identical to the old - # Euler-update behaviour at R = I. - delta_R = render.rotation_matrix(ax, ay, az) - self._R = delta_R * self._R + self.apply_rotation(v_object * keep, frame="object") + else: + delta_R = render.rotation_matrix(ax, ay, az) + self.apply_rotation(delta_R.as_rotvec()) self.update_scene() def _pan_drag(self, event: QtGui.QMouseEvent) -> None: @@ -1949,9 +2088,14 @@ def save_locs_rotated(self) -> None: "Pick": pick, "Pick shape": self.view_rot.pick_shape, "Pick size (nm)": size * pixelsize, - "angx": self.view_rot.angx, - "angy": self.view_rot.angy, - "angz": self.view_rot.angz, + # accumulated rotation angles incl. full turns + # (radians, codebase convention); legacy keys + "angx": float(self.view_rot.angx), + "angy": float(self.view_rot.angy), + "angz": float(self.view_rot.angz), + "Quaternion (x, y, z, w)": [ + float(q) for q in self.view_rot.rotation.as_quat() + ], } ] diff --git a/picasso/render.py b/picasso/render.py index 22d1bc6e..a64014ce 100755 --- a/picasso/render.py +++ b/picasso/render.py @@ -43,7 +43,7 @@ def render( Literal["gaussian", "gaussian_iso", "smooth", "convolve"] | None ) = None, min_blur_width: float = 0.0, - ang: tuple | None = None, + ang: tuple | Rotation | None = None, disp_px_size: float | None = None, ) -> tuple[int, lib.FloatArray2D]: """Render localizations given FOV and blur method. @@ -74,9 +74,11 @@ def render( localizations which is the median localization precision. min_blur_width : float, optional Minimum size of blur (camera pixels). - ang : tuple, optional - Rotation angles of locs around x, y and z axes in radians. If - None, locs are not rotated. + ang : tuple or scipy.spatial.transform.Rotation, optional + Rotation of locs; either a scipy Rotation (e.g. built from a + quaternion) or a tuple of 3 rotation angles around the x, y + and z axes in radians (legacy Euler convention, see + ``rotation_matrix``). If None, locs are not rotated. disp_px_size : float, optional Display pixel size in nm. Will replace oversampling in v0.11.0. @@ -637,7 +639,7 @@ def _fill_gaussian_rot( sz: lib.FloatArray1D, n_pixel_x: int, n_pixel_y: int, - ang: tuple[float, float, float], + rot_matrix: lib.Array3x3, ) -> None: """Fill image with rotated gaussian-blurred localizations. @@ -654,38 +656,12 @@ def _fill_gaussian_rot( Localization precision in x, y and z for each localization. n_pixel_x, n_pixel_y : int Number of pixels in x and y. - ang : tuple - Rotation angles of locs around x, y and z axes (radians). + rot_matrix : lib.Array3x3 + Rotation matrix (float32) applied to the localizations. """ n_locs = len(x) if n_locs == 0: return - angx, angy, angz = ang - rot_mat_x = np.array( - [ - [1.0, 0.0, 0.0], - [0.0, np.cos(angx), np.sin(angx)], - [0.0, -np.sin(angx), np.cos(angx)], - ], - dtype=np.float32, - ) - rot_mat_y = np.array( - [ - [np.cos(angy), 0.0, np.sin(angy)], - [0.0, 1.0, 0.0], - [-np.sin(angy), 0.0, np.cos(angy)], - ], - dtype=np.float32, - ) - rot_mat_z = np.array( - [ - [np.cos(angz), -np.sin(angz), 0.0], - [np.sin(angz), np.cos(angz), 0.0], - [0.0, 0.0, 1.0], - ], - dtype=np.float32, - ) - rot_matrix = (rot_mat_x @ rot_mat_y @ rot_mat_z).astype(np.float32) rot_matrixT = np.ascontiguousarray(rot_matrix.T) for i in range(n_locs): @@ -804,7 +780,7 @@ def render_hist( x_min: float, y_max: float, x_max: float, - ang: tuple[float, float, float] | None = None, + ang: tuple[float, float, float] | Rotation | None = None, ) -> tuple[int, lib.FloatArray2D]: """Alias for _render_hist which will be a private function in v0.11.0. Kept for backward compatibility but will be removed in @@ -826,7 +802,7 @@ def _render_hist( x_min: float, y_max: float, x_max: float, - ang: tuple[float, float, float] | None = None, + ang: tuple[float, float, float] | Rotation | None = None, ) -> tuple[int, lib.FloatArray2D]: """Render localizations with no blur by assigning them to pixels. @@ -840,9 +816,11 @@ def _render_hist( Minimum y and x coordinates to be rendered (camera pixels) y_max, x_max : float Maximum y and x coordinates to be rendered (camera pixels) - ang : tuple (default=None) - Rotation angles of locs around x, y and z axes in radians. If - None, locs are not rotated. + ang : tuple or scipy.spatial.transform.Rotation, optional + Rotation of locs; either a scipy Rotation (e.g. built from a + quaternion) or a tuple of 3 rotation angles around the x, y + and z axes in radians (legacy Euler convention, see + ``rotation_matrix``). If None, locs are not rotated. Returns ------- @@ -1017,7 +995,7 @@ def render_gaussian( y_max: float, x_max: float, min_blur_width: float, - ang: tuple[float, float, float] | None = None, + ang: tuple[float, float, float] | Rotation | None = None, ) -> tuple[int, lib.FloatArray2D]: """Alias for _render_gaussian which will be a private function in v0.11.0. Kept for backward compatibility but will be removed in v0.11.0. Use @@ -1047,7 +1025,7 @@ def _render_gaussian( y_max: float, x_max: float, min_blur_width: float, - ang: tuple[float, float, float] | None = None, + ang: tuple[float, float, float] | Rotation | None = None, ) -> tuple[int, lib.FloatArray2D]: """Render localizations with with individual localization precision which differs in x and y. @@ -1064,9 +1042,11 @@ def _render_gaussian( Minimum and maximum x coordinates to be rendered (camera pixels). min_blur_width : float Minimum localization precision (camera pixels). - ang : tuple, optional - Rotation angles of localizations around x, y and z axes in - radians. If None, localizations are not rotated. + ang : tuple or scipy.spatial.transform.Rotation, optional + Rotation of localizations; either a scipy Rotation (e.g. built + from a quaternion) or a tuple of 3 rotation angles around the + x, y and z axes in radians (legacy Euler convention, see + ``rotation_matrix``). If None, localizations are not rotated. Returns ------- @@ -1124,7 +1104,12 @@ def _render_gaussian( sx = blur_width[in_view] sz = blur_depth[in_view] - _fill_gaussian_rot(image, x, y, sx, sy, sz, n_pixel_x, n_pixel_y, ang) + rot_matrix = np.ascontiguousarray( + to_rotation(ang).as_matrix(), dtype=np.float32 + ) + _fill_gaussian_rot( + image, x, y, sx, sy, sz, n_pixel_x, n_pixel_y, rot_matrix + ) n = len(x) return n, image @@ -1138,7 +1123,7 @@ def render_gaussian_iso( y_max: float, x_max: float, min_blur_width: float, - ang: tuple[float, float, float] | None = None, + ang: tuple[float, float, float] | Rotation | None = None, ) -> tuple[int, lib.FloatArray2D]: """Alias for _render_gaussian_iso which will be a private function in v0.11.0. Kept for backward compatibility but will be removed in v0.11.0. Use @@ -1168,7 +1153,7 @@ def _render_gaussian_iso( y_max: float, x_max: float, min_blur_width: float, - ang: tuple[float, float, float] | None = None, + ang: tuple[float, float, float] | Rotation | None = None, ) -> tuple[int, lib.FloatArray2D]: """Same as ``_render_gaussian``, but uses the same localization precision in x and y.""" @@ -1221,7 +1206,12 @@ def _render_gaussian_iso( sx = sy sz = blur_depth[in_view] - _fill_gaussian_rot(image, x, y, sx, sy, sz, n_pixel_x, n_pixel_y, ang) + rot_matrix = np.ascontiguousarray( + to_rotation(ang).as_matrix(), dtype=np.float32 + ) + _fill_gaussian_rot( + image, x, y, sx, sy, sz, n_pixel_x, n_pixel_y, rot_matrix + ) return len(x), image @@ -1234,7 +1224,7 @@ def render_convolve( y_max: float, x_max: float, min_blur_width: float, - ang: tuple[float, float, float] | None = None, + ang: tuple[float, float, float] | Rotation | None = None, ) -> tuple[int, lib.FloatArray2D]: """Alias for _render_convolve which will be a private function in v0.11.0. Kept for backward compatibility but will be removed in v0.11.0. Use @@ -1264,7 +1254,7 @@ def _render_convolve( y_max: float, x_max: float, min_blur_width: float, - ang: tuple[float, float, float] | None = None, + ang: tuple[float, float, float] | Rotation | None = None, ) -> tuple[int, lib.FloatArray2D]: """Render localizations with with global localization precision, i.e. each localization is blurred by the median localization @@ -1282,9 +1272,11 @@ def _render_convolve( Maximum y and x coordinates to be rendered (camera pixels). min_blur_width : float Minimum localization precision (camera pixels). - ang : tuple, optional - Rotation angles of localizations around x, y and z axes in - radians. If None, localizations are not rotated. + ang : tuple or scipy.spatial.transform.Rotation, optional + Rotation of localizations; either a scipy Rotation (e.g. built + from a quaternion) or a tuple of 3 rotation angles around the + x, y and z axes in radians (legacy Euler convention, see + ``rotation_matrix``). If None, localizations are not rotated. Returns ------- @@ -1334,7 +1326,7 @@ def render_smooth( x_min: float, y_max: float, x_max: float, - ang: tuple[float, float, float] | None = None, + ang: tuple[float, float, float] | Rotation | None = None, ) -> tuple[int, lib.FloatArray2D]: """Alias for _render_smooth which will be a private function in v0.11.0. Kept for backward compatibility but will be removed in v0.11.0. Use _render_smooth @@ -1361,7 +1353,7 @@ def _render_smooth( x_min: float, y_max: float, x_max: float, - ang: tuple[float, float, float] | None = None, + ang: tuple[float, float, float] | Rotation | None = None, ) -> tuple[int, lib.FloatArray2D]: """Render localizations with with blur of one display pixel (set by oversampling). @@ -1376,9 +1368,11 @@ def _render_smooth( Minimum y and x coordinates to be rendered (camera pixels). y_max, x_max : float Maximum y and x coordinates to be rendered (camera pixels). - ang : tuple, optional - Rotation angles of localizations around x, y and z axes in - radians. If None, localizations are not rotated. + ang : tuple or scipy.spatial.transform.Rotation, optional + Rotation of localizations; either a scipy Rotation (e.g. built + from a quaternion) or a tuple of 3 rotation angles around the + x, y and z axes in radians (legacy Euler convention, see + ``rotation_matrix``). If None, localizations are not rotated. Returns ------- @@ -1504,6 +1498,76 @@ def rotation_matrix(angx: float, angy: float, angz: float) -> Rotation: return Rotation.from_matrix(rot_mat) +def to_rotation( + ang: tuple[float, float, float] | Rotation | None, +) -> Rotation | None: + """Normalize a rotation input to a scipy Rotation. + + Parameters + ---------- + ang : tuple, Rotation or None + Rotation to be normalized. Either a scipy Rotation (e.g. built + from a quaternion; used as is), a tuple of 3 rotation angles + around the x, y and z axes in radians (legacy Euler convention, + see ``rotation_matrix``), or None. + + Returns + ------- + scipy.spatial.transform.Rotation or None + The rotation as a scipy Rotation, or None if ``ang`` is None. + """ + if ang is None: + return None + if isinstance(ang, Rotation): + return ang + return rotation_matrix(*ang) + + +def closest_rotvec( + rotation: Rotation, + reference: lib.FloatArray1D, +) -> lib.FloatArray1D: + """Find the rotation vector representing ``rotation`` that is + closest to ``reference``. + + A rotation vector (axis * angle, radians) is not unique: adding + full turns (2 pi) around the same axis yields the same rotation. + This function picks the representation closest to ``reference``, + which allows keeping track of rotations beyond +/- 180 degrees + (e.g. unwrapping a continuously updated rotation, or encoding + multiple full turns in an animation segment). + + Parameters + ---------- + rotation : scipy.spatial.transform.Rotation + The rotation to be represented. + reference : lib.FloatArray1D + Rotation vector (radians) to which the result should be + closest. + + Returns + ------- + rotvec : lib.FloatArray1D + Rotation vector (radians) such that + ``Rotation.from_rotvec(rotvec) == rotation``, with the number + of full turns chosen to match ``reference``. Its magnitude may + exceed pi. + """ + reference = np.asarray(reference, dtype=float) + base = rotation.as_rotvec() # magnitude <= pi + theta = np.linalg.norm(base) + if theta < 1e-9: + # identity rotation; keep the full turns along reference's axis + ref_magnitude = np.linalg.norm(reference) + if ref_magnitude < 1e-9: + return np.zeros(3) + n_turns = np.round(ref_magnitude / (2 * np.pi)) + return reference * (2 * np.pi * n_turns / ref_magnitude) + axis = base / theta + n_turns = np.round((axis @ reference - theta) / (2 * np.pi)) + return axis * (theta + 2 * np.pi * n_turns) + + def locs_rotation( locs: pd.DataFrame, oversampling: float, @@ -1511,7 +1575,7 @@ def locs_rotation( x_max: float, y_min: float, y_max: float, - ang: tuple[float, float, float], + ang: tuple[float, float, float] | Rotation, ) -> tuple[ lib.FloatArray1D, lib.FloatArray1D, lib.BoolArray1D, lib.FloatArray1D ]: @@ -1527,9 +1591,10 @@ def locs_rotation( Minimum y and x coordinate to be rendered (camera pixels). y_max, x_max : float Maximum y and x coordinate to be rendered (camera pixels). - ang : tuple - Rotation angles of localizations around x, y and z axes in - radians. + ang : tuple or scipy.spatial.transform.Rotation + Rotation of localizations; either a scipy Rotation or a tuple + of 3 rotation angles around the x, y and z axes in radians + (legacy Euler convention, see ``rotation_matrix``). Returns ------- @@ -1551,7 +1616,7 @@ def locs_rotation( locs_coord[:, 1] -= y_min + (y_max - y_min) / 2 # rotate locs - R = rotation_matrix(ang[0], ang[1], ang[2]) + R = to_rotation(ang) locs_coord = R.apply(locs_coord) # unshift locs @@ -2538,7 +2603,7 @@ def draw_minimap( def draw_rotation( image: QtGui.QImage, - ang: tuple[float, float, float], + ang: tuple[float, float, float] | Rotation, axis_length: int = 30, axis_center: tuple[int, int] = (50, -50), # bottom left ) -> QtGui.QImage: @@ -2548,8 +2613,10 @@ def draw_rotation( ---------- image : QImage Image containing rendered localizations. - ang : tuple of float - Rotation angles around x, y, and z axes in radians. + ang : tuple of float or scipy.spatial.transform.Rotation + Rotation of the localizations; either a scipy Rotation or a + tuple of 3 rotation angles around the x, y and z axes in + radians (legacy Euler convention, see ``rotation_matrix``). axis_length : int, optional Length of the rotation axes in display pixels. Default is 30. axis_center : tuple of int, optional @@ -2592,7 +2659,7 @@ def draw_rotation( # rotate these points coordinates = [[xx, xy, xz], [yx, yy, yz], [zx, zy, zz]] - R = rotation_matrix(*ang) + R = to_rotation(ang) coordinates = R.apply(coordinates).astype(int) (xx, xy, xz) = coordinates[0] (yx, yy, yz) = coordinates[1] @@ -2635,7 +2702,8 @@ def draw_rotation_angles( image : QImage Image containing rendered localizations. ang : tuple of float - Rotation angles around x, y, and z axes in radians. + Rotation angles (or rotation-vector components) around x, y, + and z axes in radians, used only for the displayed text. color : QColor, optional Color of the text. Default is white. @@ -2667,7 +2735,7 @@ def render_scene( Literal["gaussian", "gaussian_iso", "smooth", "convolve"] | None ) = None, min_blur_width: float = 0.0, - ang: tuple | None = None, + ang: tuple | Rotation | None = None, contrast: tuple[float, float] | None = None, invert_colors: bool = False, single_channel_colormap: str | lib.FloatArray2D = "magma", @@ -2738,9 +2806,11 @@ def render_scene( localizations which is the median localization precision. min_blur_width : float, optional Minimum size of blur (camera pixels). - ang : tuple, optional - Rotation angles of locs around x, y and z axes in radians. If - None, locs are not rotated. + ang : tuple or scipy.spatial.transform.Rotation, optional + Rotation of locs; either a scipy Rotation (e.g. built from a + quaternion) or a tuple of 3 rotation angles around the x, y + and z axes in radians (legacy Euler convention, see + ``rotation_matrix``). If None, locs are not rotated. contrast : tuple of float, optional Contrast limits for scaling. If None, contrast is automatically determined. @@ -2854,7 +2924,7 @@ def _render_multi_channel( Literal["gaussian", "gaussian_iso", "smooth", "convolve"] | None ) = None, min_blur_width: float = 0.0, - ang: tuple | None = None, + ang: tuple | Rotation | None = None, contrast: tuple[float, float] | None = None, relative_intensities: list[float] | None = None, invert_colors: bool = False, @@ -2933,7 +3003,7 @@ def _render_single_channel( Literal["gaussian", "gaussian_iso", "smooth", "convolve"] | None ) = None, min_blur_width: float = 0.0, - ang: tuple | None = None, + ang: tuple | Rotation | None = None, contrast: tuple[float, float] | None = None, invert_colors: bool = False, single_channel_colormap: str = "magma", @@ -3244,41 +3314,91 @@ def optimal_scalebar_length(pixelsize: int | float, width: int | float) -> int: return scalebar +def _normalize_animation_positions( + positions: list, +) -> list[tuple[Rotation, tuple]]: + """Normalize animation checkpoints to (Rotation, viewport) tuples. + + Accepts two formats: each position is (rotation, viewport) with + rotation being a scipy Rotation, as well as the legacy format + (angle_x, angle_y, angle_z, viewport) with Euler angles in radians + (see ``rotation_matrix``), which is deprecated. + """ + normalized = [] + legacy = False + for p in positions: + if len(p) == 2 and isinstance(p[0], Rotation): + normalized.append((p[0], p[1])) + elif len(p) == 4: + legacy = True + normalized.append((rotation_matrix(p[0], p[1], p[2]), p[3])) + else: + raise ValueError( + "Each position must be a tuple (rotation, viewport) with " + "rotation being a scipy Rotation, or the deprecated " + "4-element form (angle_x, angle_y, angle_z, viewport)." + ) + if legacy: + lib.deprecation_warning( + "Deprecation warning: passing animation positions as Euler " + "angles (angle_x, angle_y, angle_z, viewport) is deprecated " + "and will be removed in v0.12.0. Pass (rotation, viewport) " + "with rotation being a scipy.spatial.transform.Rotation " + "instead." + ) + return normalized + + def _animation_sequence( - positions: list[list[float, float, float, tuple]], + positions: list[tuple[Rotation, tuple]], durations: list[float], fps: int, -) -> tuple[list, list]: - """Calculate the sequence of angles and viewports for the animation. - See ``build_animation`` for more details.""" - n_frames = [0] - for i in range(len(positions) - 1): - n_frames.append(int(fps * durations[i])) - - # find rotation angles and viewport for each frame - angles = [] + segment_rotations: list | None = None, +) -> tuple[list[Rotation], list]: + """Calculate the sequence of rotations and viewports for the + animation. See ``build_animation`` for more details. + + Each segment is interpolated along the geodesic between the two + checkpoint rotations at constant angular velocity (slerp). If + ``segment_rotations`` is given, the corresponding rotation vector + defines the rotation path of each segment and may include full + turns (magnitude beyond pi), e.g. 4 pi for two full turns. The + vector is snapped to the true relative rotation between the two + checkpoints (see ``closest_rotvec``) so that each segment always + ends exactly at the next checkpoint. This is equivalent to + splitting a segment with a rotation larger than 180 degrees into + sub-180-degree pieces and applying slerp to each piece.""" + rotations = [] viewports = [] for i in range(len(positions) - 1): - # angles - x1, y1, z1 = positions[i][:3] - x2, y2, z2 = positions[i + 1][:3] - current_angles = np.linspace( - [x1, y1, z1], [x2, y2, z2], n_frames[i + 1] + n_frames = int(fps * durations[i]) + + # rotations + R1, vp1 = positions[i] + R2, vp2 = positions[i + 1] + relative = R2 * R1.inv() + if segment_rotations is not None: + rotvec = closest_rotvec( + relative, np.asarray(segment_rotations[i], dtype=float) + ) + else: + rotvec = relative.as_rotvec() + fractions = np.linspace(0, 1, n_frames) + rotations.extend( + Rotation.from_rotvec(fraction * rotvec) * R1 + for fraction in fractions ) - angles.extend(current_angles) # viewports - vp1 = positions[i][3] - vp2 = positions[i + 1][3] - ymin = np.linspace(vp1[0][0], vp2[0][0], n_frames[i + 1]) - xmin = np.linspace(vp1[0][1], vp2[0][1], n_frames[i + 1]) - ymax = np.linspace(vp1[1][0], vp2[1][0], n_frames[i + 1]) - xmax = np.linspace(vp1[1][1], vp2[1][1], n_frames[i + 1]) + ymin = np.linspace(vp1[0][0], vp2[0][0], n_frames) + xmin = np.linspace(vp1[0][1], vp2[0][1], n_frames) + ymax = np.linspace(vp1[1][0], vp2[1][0], n_frames) + xmax = np.linspace(vp1[1][1], vp2[1][1], n_frames) current_viewports = [ ((ymin[j], xmin[j]), (ymax[j], xmax[j])) for j in range(len(ymin)) ] viewports.extend(current_viewports) - return angles, viewports + return rotations, viewports def build_animation( @@ -3286,10 +3406,14 @@ def build_animation( locs: pd.DataFrame | list[pd.DataFrame], info: list[dict] | list[list[dict]], *, - positions: list[tuple[float, float, float, tuple]], + positions: ( + list[tuple[Rotation, tuple]] + | list[tuple[tuple[float, float, float], tuple]] + ), durations: list[float], disp_px_size: int | float, # nm image_size: tuple[int, int], + segment_rotations: list | None = None, blur_method: ( Literal["gaussian", "gaussian_iso", "smooth", "convolve"] | None ) = None, @@ -3306,7 +3430,7 @@ def build_animation( ) = None, ) -> None: """Build an animation of rendered localizations given the - checkpoints (angle, viewport, etc) and the time between them. + checkpoints (rotation, viewport, etc) and the time between them. Parameters ---------- @@ -3328,13 +3452,27 @@ def build_animation( image_size : tuple of int Size of the rendered image in pixels, given as (width, height). positions : list - Each element determines the checkpoint of the animation, which - is a tuple of 4 elements: (angle_x, angle_y, angle_z, viewport). - Angles are in radians. Viewport is given as ((y_min, x_min), - (y_max, x_max)) in camera pixels. + Each element determines a checkpoint of the animation, which + is a tuple of 2 elements: (rotation, viewport). Rotation is a + ``scipy.spatial.transform.Rotation`` defining the orientation + of the localizations at the checkpoint. Viewport is given as + ((y_min, x_min), (y_max, x_max)) in camera pixels. The + deprecated legacy format (angle_x, angle_y, angle_z, viewport) + with Euler angles in radians (see ``rotation_matrix``) is also + accepted and will be removed in v0.12.0. durations : list List of durations in seconds between the checkpoints. Must have the same length as positions - 1. + segment_rotations : list, optional + One rotation vector (3 floats, radians, scipy convention) per + segment, i.e., of length len(positions) - 1, describing the + full rotation path from one checkpoint to the next. The + magnitude may exceed pi to encode rotations larger than 180 + degrees, e.g. (0, 0, 4 * pi) for two full turns around the z + axis. Each vector is snapped to the true relative rotation + between its two checkpoints, so the segment always ends + exactly at the next checkpoint. If None, each segment follows + the shortest path (slerp). Default is None. blur_method : {"gaussian", "gaussian_iso", "smooth", "convolve"} or None, \ optional Defines localizations' blur. The string has to be one of @@ -3397,10 +3535,16 @@ def build_animation( assert ( isinstance(positions, list) and len(positions) >= 2 ), "positions must be a list with at least 2 elements." - assert all(len(p) == 4 for p in positions), ( - "Each position must be a tuple/list of 4 elements: " - "(angle_x, angle_y, angle_z, viewport)." - ) + positions = _normalize_animation_positions(positions) + if segment_rotations is not None: + assert ( + isinstance(segment_rotations, list) + and len(segment_rotations) == len(positions) - 1 + ), "segment_rotations must be a list of length len(positions) - 1." + assert all( + np.asarray(rotvec, dtype=float).shape == (3,) + for rotvec in segment_rotations + ), "Each segment rotation must be a rotation vector of 3 floats." assert ( isinstance(durations, list) and len(durations) == len(positions) - 1 ), "durations must be a list of length len(positions) - 1." @@ -3448,9 +3592,6 @@ def build_animation( assert ( len(colors) == n_channels ), "colors must have one entry per channel." - assert all( - len(c) == 3 and all(0.0 <= v <= 1.0 for v in c) for c in colors - ), "Each color must be an RGB tuple with values between 0 and 1." if relative_intensities is not None: n_channels = len(locs) if isinstance(locs, list) else 1 assert ( @@ -3475,6 +3616,7 @@ def build_animation( info=info, positions=positions, durations=durations, + segment_rotations=segment_rotations, disp_px_size=disp_px_size, image_size=image_size, blur_method=blur_method, @@ -3494,8 +3636,9 @@ def _build_animation( path: str, locs: pd.DataFrame | list[pd.DataFrame], info: list[dict] | list[list[dict]], - positions: list[tuple[float, float, float, tuple]], + positions: list[tuple[Rotation, tuple]], durations: list[float], + segment_rotations: list | None, disp_px_size: int | float, image_size: tuple[int, int], blur_method: ( @@ -3513,7 +3656,9 @@ def _build_animation( ) -> None: """Internal function to build an animation of rendered localizations given the checkpoints. See ``build_animation`` for more details.""" - angles, viewports = _animation_sequence(positions, durations, fps) + rotations, viewports = _animation_sequence( + positions, durations, fps, segment_rotations=segment_rotations + ) # width and height for building the animation; must be divisible by 16 # as ffmpeg codecs require this for proper encoding @@ -3526,10 +3671,10 @@ def _build_animation( use_tqdm = progress_callback == "console" if use_tqdm: iter_range = tqdm( - range(len(angles)), desc="Building animation", unit="frame" + range(len(rotations)), desc="Building animation", unit="frame" ) else: - iter_range = range(len(angles)) + iter_range = range(len(rotations)) for i in iter_range: if callable(progress_callback): @@ -3546,7 +3691,7 @@ def _build_animation( info=info, disp_px_size=disp_px_size_, viewport=viewports[i], - ang=angles[i], + ang=rotations[i], blur_method=blur_method, min_blur_width=min_blur_width, contrast=contrast_, @@ -3570,30 +3715,32 @@ def _build_animation( video_writer.append_data(frame) if callable(progress_callback): - progress_callback(len(angles)) + progress_callback(len(rotations)) video_writer.close() # save a yaml with animation settings, note that yaml does not support # numpy types and arrays - angles_yaml = [ - ( - float(np.degrees(p[0])), - float(np.degrees(p[1])), - float(np.degrees(p[2])), - ) - for p in positions - ] + quaternions_yaml = [[float(q) for q in R.as_quat()] for R, _ in positions] viewports_yaml = [ ( - (float(p[3][0][0]), float(p[3][0][1])), - (float(p[3][1][0]), float(p[3][1][1])), + (float(vp[0][0]), float(vp[0][1])), + (float(vp[1][0]), float(vp[1][1])), ) - for p in positions + for _, vp in positions + ] + if segment_rotations is None: + segment_rotations = [ + (positions[i + 1][0] * positions[i][0].inv()).as_rotvec() + for i in range(len(positions) - 1) + ] + segments_yaml = [ + [float(np.degrees(v)) for v in rotvec] for rotvec in segment_rotations ] anim_settings = { "Generated by": f"Picasso v{__version__} Render 3D Animation", "FPS": fps, - "Angles at checkpoints (x, y, z) (deg)": angles_yaml, + "Quaternions at checkpoints (x, y, z, w)": quaternions_yaml, + "Rotations between checkpoints (x, y, z) (deg)": segments_yaml, "Viewports at checkpoints (camera pixels)": viewports_yaml, "Durations (s)": durations, } diff --git a/tests/test_render.py b/tests/test_render.py index 6669c40f..b342df12 100644 --- a/tests/test_render.py +++ b/tests/test_render.py @@ -11,6 +11,7 @@ import pandas as pd import pytest from PyQt6 import QtCore, QtGui +from scipy.spatial.transform import Rotation from picasso import io, masking, render @@ -540,6 +541,99 @@ def test_locs_rotation_in_view_consistency(self, locs_3d): assert len(y_out) == n_in_view assert len(z_out) == n_in_view + def test_to_rotation_none(self): + assert render.to_rotation(None) is None + + def test_to_rotation_passes_rotation_through(self): + R = Rotation.from_rotvec([0.1, -0.2, 0.3]) + assert render.to_rotation(R) is R + + def test_to_rotation_legacy_euler_equivalence(self): + ang = (0.4, -0.7, 1.1) + R = render.to_rotation(ang) + expected = render.rotation_matrix(*ang) + assert np.allclose(R.as_matrix(), expected.as_matrix()) + + def test_locs_rotation_accepts_rotation_object(self, locs_3d): + ang = (0.1, 0.2, 0.3) + out_tuple = render.locs_rotation( + locs_3d, 5.0, 0.0, 32.0, 0.0, 32.0, ang + ) + out_rotation = render.locs_rotation( + locs_3d, 5.0, 0.0, 32.0, 0.0, 32.0, render.rotation_matrix(*ang) + ) + for a, b in zip(out_tuple, out_rotation): + assert np.allclose(a, b) + + def test_render_accepts_rotation_object(self, locs_3d, info): + """Rendering with a scipy Rotation matches the legacy Euler + tuple input for all rotation-aware code paths.""" + ang = (0.5, 0.3, 0.2) + for blur_method in (None, "gaussian"): + _, im_tuple = render.render( + locs_3d, + info, + oversampling=5, + viewport=FULL_VIEWPORT, + blur_method=blur_method, + ang=ang, + ) + _, im_rotation = render.render( + locs_3d, + info, + oversampling=5, + viewport=FULL_VIEWPORT, + blur_method=blur_method, + ang=render.rotation_matrix(*ang), + ) + assert np.allclose(im_tuple, im_rotation) + + +class TestClosestRotvec: + def test_zero_reference_returns_base(self): + R = Rotation.from_rotvec([0.3, 0.0, 0.0]) + out = render.closest_rotvec(R, np.zeros(3)) + assert np.allclose(out, [0.3, 0.0, 0.0]) + + def test_unwraps_full_turns(self): + """A 10-degree rotation with a reference near 370 degrees must + unwrap to 370 degrees.""" + axis = np.array([0.0, 0.0, 1.0]) + R = Rotation.from_rotvec(np.radians(10) * axis) + reference = np.radians(365) * axis + out = render.closest_rotvec(R, reference) + assert np.allclose(out, np.radians(370) * axis) + + def test_unwraps_across_pi(self): + """Crossing 180 degrees must continue counting up instead of + flipping the axis.""" + axis = np.array([1.0, 0.0, 0.0]) + R = Rotation.from_rotvec(np.radians(181) * axis) # wraps to -179 + reference = np.radians(180) * axis + out = render.closest_rotvec(R, reference) + assert np.allclose(out, np.radians(181) * axis) + + def test_identity_keeps_turns_of_reference(self): + """The identity rotation with a reference of ~2 turns must keep + the full turns along the reference axis.""" + axis = np.array([0.0, 1.0, 0.0]) + out = render.closest_rotvec( + Rotation.identity(), np.radians(719) * axis + ) + assert np.allclose(out, np.radians(720) * axis) + + def test_identity_zero_reference(self): + out = render.closest_rotvec(Rotation.identity(), np.zeros(3)) + assert np.allclose(out, np.zeros(3)) + + def test_result_represents_same_rotation(self): + R = Rotation.from_rotvec([0.2, -0.4, 0.6]) + reference = np.array([2.5, -4.0, 6.5]) + out = render.closest_rotvec(R, reference) + assert np.allclose( + (Rotation.from_rotvec(out) * R.inv()).magnitude(), 0.0, atol=1e-9 + ) + # --------------------------------------------------------------------------- # 3x3 matrix helpers @@ -1366,20 +1460,137 @@ def test_return_bgra(self): # --------------------------------------------------------------------------- +class TestAnimationSequence: + def test_slerp_endpoints_and_midpoint(self): + """Default interpolation is slerp: endpoints exact, midpoint at + half the rotation angle.""" + R1 = Rotation.identity() + R2 = Rotation.from_rotvec([0.0, 0.0, np.pi / 2]) + positions = [(R1, FULL_VIEWPORT), (R2, FULL_VIEWPORT)] + rotations, viewports = render._animation_sequence( + positions, [1.0], fps=3 + ) + assert len(rotations) == 3 + assert len(viewports) == 3 + assert (rotations[0] * R1.inv()).magnitude() == pytest.approx( + 0.0, abs=1e-9 + ) + assert (rotations[-1] * R2.inv()).magnitude() == pytest.approx( + 0.0, abs=1e-9 + ) + assert np.allclose( + rotations[1].as_rotvec(), [0.0, 0.0, np.pi / 4], atol=1e-9 + ) + + def test_full_turn_segment(self): + """A 360-degree segment rotation spins the full turn even though + both checkpoints have the same orientation.""" + R = Rotation.identity() + positions = [(R, FULL_VIEWPORT), (R, FULL_VIEWPORT)] + rotations, _ = render._animation_sequence( + positions, + [1.0], + fps=5, + segment_rotations=[np.array([0.0, 0.0, 2 * np.pi])], + ) + assert len(rotations) == 5 + # quarter of the way through: 90 degrees around z + assert np.allclose( + rotations[1].as_rotvec(), [0.0, 0.0, np.pi / 2], atol=1e-9 + ) + # halfway through: 180 degrees around z + assert rotations[2].magnitude() == pytest.approx(np.pi, abs=1e-9) + # ends exactly at the checkpoint again + assert rotations[-1].magnitude() == pytest.approx(0.0, abs=1e-9) + + def test_segment_rotation_snaps_to_checkpoints(self): + """The segment rotation vector is snapped so that the segment + ends exactly at the next checkpoint, keeping its turns.""" + R1 = Rotation.identity() + R2 = Rotation.from_rotvec([0.0, 0.0, np.pi / 2]) + positions = [(R1, FULL_VIEWPORT), (R2, FULL_VIEWPORT)] + # 450 degrees = 90 degrees + 1 full turn; pass a slightly off + # target to verify snapping + target = np.array([0.0, 0.0, np.radians(449)]) + rotations, _ = render._animation_sequence( + positions, [1.0], fps=11, segment_rotations=[target] + ) + assert (rotations[-1] * R2.inv()).magnitude() == pytest.approx( + 0.0, abs=1e-9 + ) + # total rotation path is 450 degrees: 1/5 of the way is 90 deg + assert np.allclose( + rotations[2].as_rotvec(), [0.0, 0.0, np.pi / 2], atol=1e-6 + ) + + def test_multi_segment_viewports(self): + """Viewports interpolate linearly per segment.""" + R = Rotation.identity() + vp1 = ((0.0, 0.0), (32.0, 32.0)) + vp2 = ((8.0, 8.0), (24.0, 24.0)) + positions = [(R, vp1), (R, vp2)] + _, viewports = render._animation_sequence(positions, [1.0], fps=3) + assert np.allclose(viewports[0], vp1) + assert np.allclose(viewports[1], ((4.0, 4.0), (28.0, 28.0))) + assert np.allclose(viewports[-1], vp2) + + def test_normalize_legacy_positions(self): + """Legacy Euler positions are converted with a deprecation + warning and match rotation_matrix.""" + legacy = [(0.1, 0.2, 0.3, FULL_VIEWPORT)] + with pytest.warns(DeprecationWarning): + normalized = render._normalize_animation_positions(legacy) + R, vp = normalized[0] + expected = render.rotation_matrix(0.1, 0.2, 0.3) + assert (R * expected.inv()).magnitude() == pytest.approx(0.0, abs=1e-9) + assert vp == FULL_VIEWPORT + + def test_normalize_invalid_position_raises(self): + with pytest.raises(ValueError): + render._normalize_animation_positions([(0.1, FULL_VIEWPORT)]) + + class TestBuildAnimation: - def test_smoke_two_frames(self, locs_3d, info, tmp_path): - """A 2-frame animation writes both an .mp4 and the sidecar .yaml.""" + def test_smoke_two_frames_legacy_euler(self, locs_3d, info, tmp_path): + """A 2-frame animation from legacy Euler positions writes both + an .mp4 and the sidecar .yaml (deprecated input format).""" out_path = tmp_path / "anim.mp4" positions = [ (0.0, 0.0, 0.0, FULL_VIEWPORT), (0.1, 0.0, 0.0, FULL_VIEWPORT), ] + with pytest.warns(DeprecationWarning): + render.build_animation( + str(out_path), + locs_3d, + info, + positions=positions, + durations=[1.0], + disp_px_size=PIXELSIZE, + image_size=(64, 64), + fps=2, + ) + assert out_path.exists() + assert out_path.stat().st_size > 0 + yaml_path = out_path.with_suffix(".yaml") + assert yaml_path.exists() + assert yaml_path.stat().st_size > 0 + + def test_smoke_quaternions_with_turns(self, locs_3d, info, tmp_path): + """An animation from scipy Rotations with a multi-turn segment + writes both an .mp4 and the sidecar .yaml.""" + out_path = tmp_path / "anim.mp4" + positions = [ + (Rotation.identity(), FULL_VIEWPORT), + (Rotation.from_rotvec([0.1, 0.0, 0.0]), FULL_VIEWPORT), + ] render.build_animation( str(out_path), locs_3d, info, positions=positions, durations=[1.0], + segment_rotations=[np.array([0.1 + 2 * np.pi, 0.0, 0.0])], disp_px_size=PIXELSIZE, image_size=(64, 64), fps=2, From 10e93df2b6b01dd2ce1f914f692cf31df77440ec Mon Sep 17 00:00:00 2001 From: rafalkowalewski1 Date: Thu, 11 Jun 2026 14:33:21 +0200 Subject: [PATCH 15/19] update 3d rotation docs --- docs/render.rst | 2 -- 1 file changed, 2 deletions(-) diff --git a/docs/render.rst b/docs/render.rst index 85edc7a2..edf10572 100644 --- a/docs/render.rst +++ b/docs/render.rst @@ -67,8 +67,6 @@ Note that to build animations, the user must have ``ffmpeg`` installed on their Rotation around z-axis is available by pressing Ctrl/Command. Rotation axis can be frozen by pressing x/y/z to freeze around the corresponding axes (to freeze around the z-axis, Ctrl/Command must be pressed as well). -There are several things to keep in mind when using the rotation window. Firstly, using individual localization precision is very slow and is not recommended as a default blur method. Also, the size of the rotation window can be altered, however, if it becomes too large, rendering may start to lag. - RESI ---- .. image:: ../docs/render_resi.png From ab7aba1c85a2aa61a5bc2e03048308ab72837ba3 Mon Sep 17 00:00:00 2001 From: rafalkowalewski1 Date: Thu, 11 Jun 2026 14:46:07 +0200 Subject: [PATCH 16/19] All 3 angles in "Rotate by angle" in the 3D rotation window are accumulated into a single widget --- changelog.md | 1 + picasso/gui/rotation.py | 101 +++++++++++++++++++++++++--------------- 2 files changed, 65 insertions(+), 37 deletions(-) diff --git a/changelog.md b/changelog.md index a3321b32..3a7b50b4 100644 --- a/changelog.md +++ b/changelog.md @@ -4,6 +4,7 @@ Last change: 11-JUN-2026 CEST ## 0.10.2 - Rotations use quaterions for unambiguous workflow, fixing bugs [#673](https://github.com/jungmannlab/picasso/issues/673), [#674](https://github.com/jungmannlab/picasso/issues/674) and [#675](https://github.com/jungmannlab/picasso/issues/675) +- All 3 angles in "Rotate by angle" in the 3D rotation window are accumulated into a single widget - Fixed "Best fitting combination" button in SPINNA ([#676](https://github.com/jungmannlab/picasso/issues/676)) - Fixed 3D calibration when frame range is user-specified - Fixed redoing 3D calibration when identification parameters change diff --git a/picasso/gui/rotation.py b/picasso/gui/rotation.py index 68c421d0..fc1bb310 100644 --- a/picasso/gui/rotation.py +++ b/picasso/gui/rotation.py @@ -641,6 +641,62 @@ def build_animation(self) -> None: progress.close() +class RotateByAngleDialog(lib.Dialog): + """Choose rotations angles. + + ... + + Attributes + ---------- + angx, angy, angz: QtWidgets.QDoubleSpinBoxes + Store the rotation angles input by the user. + """ + + def __init__(self, window: QtWidgets.QMainWindow) -> None: + super().__init__(window) + self.window = window + self.setWindowTitle(f"Enter rotation angles") + layout = QtWidgets.QFormLayout(self) + self.angx = QtWidgets.QDoubleSpinBox() + self.angx.setValue(0) + self.angx.setRange(-999999, 999999) + self.angx.setSingleStep(1) + layout.addRow("Angle x (deg)", self.angx) + self.angy = QtWidgets.QDoubleSpinBox() + self.angy.setValue(0) + self.angy.setRange(-999999, 999999) + self.angy.setSingleStep(1) + layout.addRow("Angle y (deg)", self.angy) + self.angz = QtWidgets.QDoubleSpinBox() + self.angz.setValue(0) + self.angz.setRange(-999999, 999999) + self.angz.setSingleStep(1) + layout.addRow("Angle z (deg)", self.angz) + + # OK and Cancel buttons + self.buttons = QtWidgets.QDialogButtonBox( + QtWidgets.QDialogButtonBox.StandardButton.Ok + | QtWidgets.QDialogButtonBox.StandardButton.Cancel, + QtCore.Qt.Orientation.Horizontal, + self, + ) + layout.addRow(self.buttons) + self.buttons.accepted.connect(self.accept) + self.buttons.rejected.connect(self.reject) + + @staticmethod + def getParams(parent: QtWidgets.QMainWindow | None = None) -> tuple: + """Create the dialog and return the requested rot. angles.""" + dialog = RotateByAngleDialog(parent) + result = dialog.exec() + return ( + dialog.angx.value(), + dialog.angy.value(), + dialog.angz.value(), + result == QtWidgets.QDialog.DialogCode.Accepted, + ) + + class ViewRotation(QtWidgets.QLabel): """Display rotated super-resolution datasets. @@ -1244,44 +1300,15 @@ def rotation_input(self) -> None: degrees are preserved, e.g. 720 degrees encodes two full turns (relevant for animations). """ - angx, ok = QtWidgets.QInputDialog.getDouble( - self, - "Rotation angle x", - "Angle x (degrees):", - 0, - decimals=2, - ) + angx, angy, angz, ok = RotateByAngleDialog.getParams(self) if ok: - angy, ok2 = QtWidgets.QInputDialog.getDouble( - self, - "Rotation angle y", - "Angle y (degrees):", - 0, - decimals=2, - ) - if ok2: - angz, ok3 = QtWidgets.QInputDialog.getDouble( - self, - "Rotation angle z", - "Angle z (degrees):", - 0, - decimals=2, - ) - if ok3: - # codebase convention: x angle is the negative of - # scipy's right-handed x rotation; "object" frame - # rotates around the data's own axes - self.apply_rotation( - [-np.radians(angx), 0.0, 0.0], frame="object" - ) - self.apply_rotation( - [0.0, np.radians(angy), 0.0], frame="object" - ) - self.apply_rotation( - [0.0, 0.0, np.radians(angz)], frame="object" - ) - - self.update_scene() + # codebase convention: x angle is the negative of + # scipy's right-handed x rotation; "object" frame + # rotates around the data's own axes + self.apply_rotation([-np.radians(angx), 0.0, 0.0], frame="object") + self.apply_rotation([0.0, np.radians(angy), 0.0], frame="object") + self.apply_rotation([0.0, 0.0, np.radians(angz)], frame="object") + self.update_scene() def delete_rotation(self) -> None: """Reset rotation and any accumulated pan offset.""" From 8cf83be1a4add02b5603841818147c1562f9b166 Mon Sep 17 00:00:00 2001 From: rafalkowalewski1 Date: Thu, 11 Jun 2026 15:18:21 +0200 Subject: [PATCH 17/19] Render by property fixed for large files (#677, #672 possibly too) --- changelog.md | 1 + picasso/lib.py | 1 + 2 files changed, 2 insertions(+) diff --git a/changelog.md b/changelog.md index 3a7b50b4..73cca0f5 100644 --- a/changelog.md +++ b/changelog.md @@ -5,6 +5,7 @@ Last change: 11-JUN-2026 CEST ## 0.10.2 - Rotations use quaterions for unambiguous workflow, fixing bugs [#673](https://github.com/jungmannlab/picasso/issues/673), [#674](https://github.com/jungmannlab/picasso/issues/674) and [#675](https://github.com/jungmannlab/picasso/issues/675) - All 3 angles in "Rotate by angle" in the 3D rotation window are accumulated into a single widget +- Render by property fixed for large files ([#677](https://github.com/jungmannlab/picasso/issues/677)), possibly related to [#672](https://github.com/jungmannlab/picasso/issues/672) - Fixed "Best fitting combination" button in SPINNA ([#676](https://github.com/jungmannlab/picasso/issues/676)) - Fixed 3D calibration when frame range is user-specified - Fixed redoing 3D calibration when identification parameters change diff --git a/picasso/lib.py b/picasso/lib.py index 4a1512cc..330ba9e6 100644 --- a/picasso/lib.py +++ b/picasso/lib.py @@ -1562,6 +1562,7 @@ def calculate_optimal_bins( bins : FloatArray1D Bins for display. """ + data = np.asarray(data) # positional indexing below; Series use labels n = len(data) if n == 0: return np.array([0.0, 1.0]) From 7c68b268ae963d92f0358fbc9a6dbf342083215e Mon Sep 17 00:00:00 2001 From: rafalkowalewski1 Date: Thu, 11 Jun 2026 15:22:05 +0200 Subject: [PATCH 18/19] update docstrings for multichannel rendering (new LUTs) --- picasso/render.py | 30 ++++++++++++++++++++++-------- 1 file changed, 22 insertions(+), 8 deletions(-) diff --git a/picasso/render.py b/picasso/render.py index a64014ce..11146913 100755 --- a/picasso/render.py +++ b/picasso/render.py @@ -2821,10 +2821,17 @@ def render_scene( corresponding pyplot colormap is selected. If a 2D array, a 256x4 array is expected with values between 0 and 1. Default is 'magma'. - colors : list of tuples, optional - List of RGB tuples corresponding to the colors of the channels. - Only needs to be specified for multi-channel data. Must range - between 0 and 1. Default is None. + colors : list of tuples or list of lib.FloatArray2D, optional + Colors of the channels, one entry per channel. Each entry is + either an ``(r, g, b)`` tuple with values between 0 and 1 + (the channel is rendered as ``intensity * rgb``) or a + ``(256, 3)`` lookup table with values between 0 and 1 (the + channel intensity is indexed into the LUT, allowing + per-channel colormaps; see ``solid_to_lut`` and + ``stops_to_lut``). The two forms cannot be mixed. Channels are + blended additively. Only needs to be specified for + multi-channel data. Default is None, in which case colors are + taken from ``lib.get_colors``. relative_intensities : list of float, optional List of relative intensities for each channel. Only needs to be specified for multi-channel data. Default is None, in which @@ -3497,10 +3504,17 @@ def build_animation( corresponding pyplot colormap is selected. If a 2D array, a 256x4 array is expected with values between 0 and 1. Default is 'magma'. - colors : list of tuples, optional - List of RGB tuples corresponding to the colors of the channels. - Only needs to be specified for multi-channel data. Must range - between 0 and 1. Default is None. + colors : list of tuples or list of lib.FloatArray2D, optional + Colors of the channels, one entry per channel. Each entry is + either an ``(r, g, b)`` tuple with values between 0 and 1 + (the channel is rendered as ``intensity * rgb``) or a + ``(256, 3)`` lookup table with values between 0 and 1 (the + channel intensity is indexed into the LUT, allowing + per-channel colormaps; see ``solid_to_lut`` and + ``stops_to_lut``). The two forms cannot be mixed. Channels are + blended additively. Only needs to be specified for + multi-channel data. Default is None, in which case colors are + taken from ``lib.get_colors``. relative_intensities : list of float, optional List of relative intensities for each channel. Only needs to be specified for multi-channel data. Default is None, in which From 64a2f531685e5ab141e701213a768f6c05e402ed Mon Sep 17 00:00:00 2001 From: rafalkowalewski1 Date: Thu, 11 Jun 2026 15:33:37 +0200 Subject: [PATCH 19/19] update version for release --- picasso/version.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/picasso/version.py b/picasso/version.py index 1f4c4d43..17c1a626 100644 --- a/picasso/version.py +++ b/picasso/version.py @@ -1 +1 @@ -__version__ = "0.10.1" +__version__ = "0.10.2"