aiaspikes-postprocess.py

import cupy as cp
 
def get_dist_mat_gpu(coords):
# Calculate distance matrix for all elements in coords. 
# ‘coords’ is a cupy array of [1D image coordinates, wavelengths] within a group of 7 files (7 wavelengths), in GPU memory. The 2D images coordinates are 1st converted from the 1D values “on the fly” and the Euclidian distance matrix is returned. 
 
    coords_x = coords % 4096
    coords_y = coords // 4096
    coords_xb = coords_x[:, cp.newaxis]
    coords_yb = coords_y[:, cp.newaxis]
    return (coords_x - coords_xb)**2 + (coords_y - coords_yb)**2
 
def get_rows_list_gpu(array, w1_idx, w2_idx):
# Get pairs of spikes whose coordinates are within a distance of 2, which defines an “overlap” with a tolerance of 2 pixels. That distance could be change if we want to change the tolerance but here it is a hard-coded number. 
# We start by selecting which pair of wavelength we are considering before calling the distance matrix calculation.  
# We also filter in the same-wavelength coordinates within the 2-pixel tolerance, as it is as important as cross-wavelength overlap, so in effect, this looks at same-wavelength and 2-wavelength overlap at a time. 
 
    mask = (array[:, w1_idx] == 1) & (array[:, w2_idx] == 1)
    df_idx = cp.nonzero(mask)[0]
    coords = array[df_idx, 0]
    dist_matrix = get_dist_mat_gpu(coords)
    select = dist_matrix < 2 
    select2 = cp.triu(select, k=1)
    r,c = cp.nonzero(select2)
    idx1, idx2 = df_idx[r], df_idx[c]
    return idx1, idx2