import numpy as np def driftEstimate(ellipsoids, z_bounds, drift_spread=7, minimum_N=2): drift_parameters = [] for E in ellipsoids: v, radii, center = E A, B, C, D, E, F = v sx = (D*F - B*E) / (A*B - D**2) sy = (D*E - A*F) / (A*B - D**2) drift_parameters.append([sx, sy, center[0], center[1], center[2]]) z_lower, z_upper = z_bounds drift_bins = [[] for i in range(z_lower, z_upper)] drift_components = [] for param in drift_parameters: drift_component = param[:2] center_z = param[4] center_k = int(round(center_z)) drift_components.append(drift_component) for k in range(center_k - drift_spread, center_k + drift_spread + 1): if k >= z_lower and k < z_upper: k_array = k - z_lower drift_bins[k_array].append(drift_component) drift_bins_means = [np.mean(np.array(bin_group), axis=0) if len(bin_group) >= minimum_N else np.array([0, 0]) for bin_group in drift_bins] drift_bins_std = [np.std(np.array(bin_group), axis=0) if len(bin_group) >= minimum_N else np.array([-1, -1]) for bin_group in drift_bins] drift_bins_N = [len(bin_group) for bin_group in drift_bins] # Interpolation, filling missing values etc., here drift_estimate = np.array(drift_bins_means) drift_std = np.array(drift_bins_std) drift_bins_N = np.array(drift_bins_N) return drift_estimate, drift_std, drift_bins_N, drift_components