<reponame>OceanNuclear/event_classification_example<filename>data/temporal_peaks_cluster.py
from numpy import array as ary; from numpy import log as ln
from numpy import cos, sin, pi, sqrt, exp, arccos;
tau = 2*pi
import numpy as np;
from matplotlib import pyplot as plt
import pandas as pd
from mpl_toolkits.mplot3d import Axes3D
from sklearn.cluster import DBSCAN, OPTICS, SpectralClustering, AgglomerativeClustering
from sklearn.preprocessing import MinMaxScaler
weak_source = 50
strong_source = 120

def cross_overs(series, constant):
    """
    Find the exact coordinates where the series crosses over the horizontal line y=constant.
    """
    above_bool = constant<series
    indices = np.where(np.diff(above_bool))[0]
    # we assume that the series starts and ends with non-zero values
    if constant<series[0]: indices = indices[1:]
    if constant<series[-1]: indices = indices[:-1]

    in_front = series[indices] - constant # difference between y_in_front and constant
    # behind = series[indices+1]
    x_offset = -in_front/np.diff(series)[indices] # (difference between y_in_front and constant) / (difference between y_in_front and y_behind)
    slope_sign = np.diff(above_bool.astype(int))[indices]
    return ary([x_offset+indices, np.full_like(indices, constant), slope_sign]).T

def find_intersection(trace, lower_thres, upper_thres):
    """
    Project the peak location of the trace using only the cross-over information of the series (crossing-over an upper and a lower threshold voltage).
    """
    lower = cross_overs(trace, lower_thres)
    upper = cross_overs(trace, upper_thres)
    xo_point_list = np.append(
                    np.append(lower, np.zeros(len(lower)).reshape([-1,1]), axis=1),
                    np.append(upper, np.ones(len(upper)).reshape([-1,1]), axis=1),
                    axis=0)
    x_ord = np.argsort(xo_point_list[:,0])
    xo_point_list = xo_point_list[x_ord] # re-order the list so that we can see which side came up first

    # organize the list into [points that rises] and [points that falls]
    rising = np.where(np.logical_and(np.diff(xo_point_list[:,2])==0, np.diff(xo_point_list[:,3])==1))[0]
    falling= np.where(np.logical_and(np.diff(xo_point_list[:,2])==0, np.diff(xo_point_list[:,3])==-1))[0]
    assert len(rising)==len(falling), "The program should detect the same number of rising and falling edges" # Otherwise I've written the program wrong
    cross_pt_list = []
    for ind1, ind3 in zip(rising, falling):
        (x1, y1), (x2, y2) = xo_point_list[ind1, :2], xo_point_list[ind1+1, :2]
        (x3, y3), (x4, y4) = xo_point_list[ind3, :2], xo_point_list[ind3+1, :2]
        m12 = (y2-y1)/(x2-x1)
        m34 = (y4-y3)/(x4-x3)
        cross_pt = np.linalg.solve([[-m12, 1], [-m34, 1]], [-m12*x1+y1, -m34*x3+y3])
        cross_pt_list.append(cross_pt)
    return cross_pt_list

def scroll(trace, offset, y_width, y_interval=1):
    """
    find the intersection points of the extrapolated lines, 
    where the extrapolated line is generated by examining the cross-over points with a reference horizontal line.
    As we scroll the reference horizontal line upwards, the location of these extrapolated intersection points will change.

    This function outputs a list of points where the first two coordinates = the location of the extrapolated intersection points,
    the last coordinate = the height of the reference horizontal line.
    """
    points = []
    for band_height in np.arange(offset+y_width, trace.max(), y_interval):
        intersection_points = find_intersection(trace, band_height-y_width/2, band_height+y_width/2)
        for point in intersection_points:
            points.append(point.tolist()+[band_height])
    return ary(points)

if __name__=='__main__':
    df = pd.read_csv('training_pm_nosat_150k.dat', sep=' ', header=None) # lazy way of reading the data
    num_events = df[0]
    # df[1] is just an integer 200. I don't know why.

    # Information which won't be available in the experimental data (i.e. labels)
    amp1, amp2 = df[2], df[7]
    rise1, rise2 = df[3], df[8]                 # default decay2=0
    decay1, decay2 = df[4], df[9]               # default pos2=0
    offset, pos1, pos2 = df[6], df[5], df[10]   # default pos2=0

    # Information which will actually be avialble in the experiment (i.e. features)
    wave_forms = df[df.columns[11:]]
    print('Data Read. Extracting useful information out of the data...')
    wave_forms.columns = range(wave_forms.shape[1])
    print("Renamed columns")

    # Determine the standard deviation of the background
    stdev_bg = np.diff(wave_forms.T[:30], axis=0).std(axis=0)
    # when dealing with experimental data a more sophisticated method of determining the BG variation would be required,
    # since not all traces may have the first 30 values = background


    # classification step
    features = []
    for num, (ind, trace) in enumerate(wave_forms.iterrows()):
        if num>30: # don't have time to classify all len(wave_forms) samples in this demo.
            break
        print(f'{ind=}, {amp1[ind]=}, {amp2[ind]=}',)
        intersection_paths = scroll(trace.values, offset[ind], y_width=stdev_bg[ind])
        point_cloud = intersection_paths[:,:]
        minmax = MinMaxScaler()
        # minmax.fit([[0,trace.min()],[len(trace),trace.max()]])
        minmax.fit([[0,trace.min(), trace.min()],[len(trace),trace.max(), trace.max()]])
        features.append(minmax.transform(point_cloud))
        print(f'{len(intersection_paths)}')
    features, labels = ary(features), num_events.values


    # displaying the classification results
    for ind, feat in enumerate(features):
        # !!!! DBSCAN requres tuning of the eps !!!!
        dbs = DBSCAN(eps=0.1)
        pred = dbs.fit_predict(feat)
        fig = plt.figure()
        ax = fig.add_subplot(111, projection='3d')
        ax.set_title(f'{ind}: {num_events[ind]} events')
        for lab in range(pred.max()+1):
            ax.scatter(*(feat[pred==lab]).T, label=f'label={lab} has {(pred==lab).sum()} points')
        line = wave_forms.loc[ind].values
        scaled_line = (line-line.min())/(line.max()-line.min())
        plt.plot(np.linspace(0,1, len(scaled_line)), scaled_line, np.zeros(len(scaled_line)))
        plt.legend()
        plt.show()

    """
    Verdict: 
    Looks like this scroll appraoch is not very easy to use: 
    1. the DBSCAN eps value needs tuning. 
    2. the outputted labels require some filtering
        (e.g. remove labels which occurs less than N times, where N=100? or N=len(feat)/30?
        Or some other filtering condition depending on the location of the point cluster relative to the line?)
    3. Scolling algorithm parameter that we can tune:
        3.1 use a different 'scrolling' algorithm: e.g. connect the furthest two cross-over points to extrapolate, instead of the nearest
        3.2 width of the band used when 'scrolling'
    4. background determination method

    Log of failed methods:
    1. Other sklearn.cluster algorithms considered
        OPTICS: Clearly doesn't work, divides into too many different clusters.
        And I remember testing one of (SpectralClustering, AgglomerativeClustering) and it also doesn't work. 
    Note: only these four algorithms were considered because the rest doesn't seem to recognize line clusters.
    2. 
    """