Skip to content
Snippets Groups Projects
DriftEstimate.py 1.46 KiB
Newer Older
  • Learn to ignore specific revisions
  • 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