interpolate_scandata_2d.py

def interpolate_scandata_2d(x,y,z, fast_axis=0, min_signal_size=500, points=100, filter_positive=False, center_to_peak=False, normalize=False):

    if fast_axis == 0:
        vector = x
    else:
        vector = y

    # consider only the positive direction of the x values
    if filter_positive:
        positive_x_index = np.where(np.sign(np.append(0,np.diff(vector)))>0) 
        x = x[positive_x_index]
        y = y[positive_x_index]
        z = z[positive_x_index]

    # center x and y to the maximum peak
    if center_to_peak:
        print("Centering xy to peak")
        peak_z_idx = np.argmax(z)
        peak_x = x[peak_z_idx]
        peak_y = y[peak_z_idx]
        x = x-peak_x
        y = y-peak_y

    if normalize:
        z = z/z.max()

    # Find uniques on the slow axis
    if fast_axis == 0:
        vector = y
    else:
        vector = x

    _,indexes_inverse = np.unique(vector, return_inverse=True)
    _,indexes_start = np.unique(vector, return_index=True)
    _,counts = np.unique(vector, return_counts=True)


    indexes_start = indexes_start[np.where(counts>min_signal_size)]


    split_x = np.split(x, indexes_start)
    split_y = np.split(y, indexes_start)
    split_z = np.split(z, indexes_start)

    split_x = split_x[1:]
    split_y = split_y[1:]
    split_z = split_z[1:]


    if fast_axis == 0:
        split = split_x
    else:
        split = split_y

    bound_min, bound_max = find_bounds(split)

    linear = np.linspace(bound_min, bound_max, points)

    new_x = []
    new_y = []
    new_z = []
    for i, (aux_x, aux_y, aux_z) in enumerate(zip(split_x, split_y, split_z)):
        
        
        if fast_axis == 0:
            vector = aux_x
            z_func = interpolate.interp1d(vector, aux_z)
            aux_y = np.ones(points)*aux_y[0]
            new_x.append(np.array(linear))
            new_y.append(np.array(aux_y))
        else:
            vector = aux_y
            z_func = interpolate.interp1d(vector, aux_z)
            aux_x = np.ones(points)*aux_x[0]
                    
            new_x.append(np.array(aux_x))
            new_y.append(np.array(linear))
        interpz = z_func(linear)
        new_z.append(interpz)

    new_x = np.array(new_x)
    new_y = np.array(new_y)
    new_z = np.array(new_z)

    return new_x, new_y, new_z