coincidence count

coincidence.py

import numpy as np


def coinc_2(Ch1_TIME, Ch2_TIME, Ch1_AMPL, Ch2_AMPL, t_window):
    Ch1_TIME = np.asarray(Ch1_TIME)
    Ch2_TIME = np.asarray(Ch2_TIME)
    Ch1_AMPL = np.asarray(Ch1_AMPL)
    Ch2_AMPL = np.asarray(Ch2_AMPL)

    # For each Ch1 time, find window in Ch2 where match is possible
    idx_start = np.searchsorted(Ch2_TIME, Ch1_TIME - t_window, side="left")
    idx_end = np.searchsorted(Ch2_TIME, Ch1_TIME + t_window, side="right")

    # Keep only those with at least one match
    has_match = idx_start < idx_end

    matched_Ch1_idx = np.flatnonzero(has_match)
    matched_Ch2_idx = idx_start[has_match]  # First match only

    return (
        Ch1_TIME[matched_Ch1_idx],
        Ch2_TIME[matched_Ch2_idx],
        Ch1_AMPL[matched_Ch1_idx],
        Ch2_AMPL[matched_Ch2_idx],
    )


def coinc_3(Ch1_TIME, Ch2_TIME, Ch3_TIME, Ch1_AMPL, Ch2_AMPL, Ch3_AMPL, t_window):
    Ch1_TIME = np.asarray(Ch1_TIME)
    Ch2_TIME = np.asarray(Ch2_TIME)
    Ch3_TIME = np.asarray(Ch3_TIME)
    Ch1_AMPL = np.asarray(Ch1_AMPL)
    Ch2_AMPL = np.asarray(Ch2_AMPL)
    Ch3_AMPL = np.asarray(Ch3_AMPL)

    # For each Ch1 time, find window in Ch2 and Ch3
    idx_start_2 = np.searchsorted(Ch2_TIME, Ch1_TIME - t_window, side="left")
    idx_end_2 = np.searchsorted(Ch2_TIME, Ch1_TIME + t_window, side="right")

    idx_start_3 = np.searchsorted(Ch3_TIME, Ch1_TIME - t_window, side="left")
    idx_end_3 = np.searchsorted(Ch3_TIME, Ch1_TIME + t_window, side="right")

    # Valid coincidences: Ch1 has at least one match in both Ch2 and Ch3
    has_match = (idx_start_2 < idx_end_2) & (idx_start_3 < idx_end_3)
    matched_Ch1_idx = np.flatnonzero(has_match)
    matched_Ch2_idx = idx_start_2[has_match]
    matched_Ch3_idx = idx_start_3[has_match]

    return (
        Ch1_TIME[matched_Ch1_idx],
        Ch2_TIME[matched_Ch2_idx],
        Ch3_TIME[matched_Ch3_idx],
        Ch1_AMPL[matched_Ch1_idx],
        Ch2_AMPL[matched_Ch2_idx],
        Ch3_AMPL[matched_Ch3_idx],
    )


def COINC_4_first_match(
    Ch1_TIME, Ch2_TIME, Ch3_TIME, Ch4_TIME,
    Ch1_AMPL, Ch2_AMPL, Ch3_AMPL, Ch4_AMPL,
    t_window
):
    # Convert to NumPy arrays
    Ch1_TIME = np.asarray(Ch1_TIME)
    Ch2_TIME = np.asarray(Ch2_TIME)
    Ch3_TIME = np.asarray(Ch3_TIME)
    Ch4_TIME = np.asarray(Ch4_TIME)
    Ch1_AMPL = np.asarray(Ch1_AMPL)
    Ch2_AMPL = np.asarray(Ch2_AMPL)
    Ch3_AMPL = np.asarray(Ch3_AMPL)
    Ch4_AMPL = np.asarray(Ch4_AMPL)

    # For each Ch1 event, find index range in Ch2/Ch3/Ch4 within time window
    idx_start_2 = np.searchsorted(Ch2_TIME, Ch1_TIME - t_window, side='left')
    idx_end_2   = np.searchsorted(Ch2_TIME, Ch1_TIME + t_window, side='right')

    idx_start_3 = np.searchsorted(Ch3_TIME, Ch1_TIME - t_window, side='left')
    idx_end_3   = np.searchsorted(Ch3_TIME, Ch1_TIME + t_window, side='right')

    idx_start_4 = np.searchsorted(Ch4_TIME, Ch1_TIME - t_window, side='left')
    idx_end_4   = np.searchsorted(Ch4_TIME, Ch1_TIME + t_window, side='right')

    # Valid coincidences must have at least one match in Ch2, Ch3, and Ch4
    has_match = (
        (idx_start_2 < idx_end_2) &
        (idx_start_3 < idx_end_3) &
        (idx_start_4 < idx_end_4)
    )

    matched_Ch1_idx = np.flatnonzero(has_match)
    matched_Ch2_idx = idx_start_2[has_match]
    matched_Ch3_idx = idx_start_3[has_match]
    matched_Ch4_idx = idx_start_4[has_match]

    return (
        Ch1_TIME[matched_Ch1_idx],
        Ch2_TIME[matched_Ch2_idx],
        Ch3_TIME[matched_Ch3_idx],
        Ch4_TIME[matched_Ch4_idx],
        Ch1_AMPL[matched_Ch1_idx],
        Ch2_AMPL[matched_Ch2_idx],
        Ch3_AMPL[matched_Ch3_idx],
        Ch4_AMPL[matched_Ch4_idx],
    )


if __name__ == "__main__":
    N_ch1 = 100000
    N_ch2 = 10000
    N_ch3 = 1000
    # Generate random time and amplitude data
    Ch1_TIME = np.random.random(N_ch1)
    Ch2_TIME = np.random.random(N_ch2)
    Ch3_TIME = np.random.random(N_ch3)

    # sort times
    Ch1_TIME.sort()
    Ch2_TIME.sort()
    Ch3_TIME.sort()

    Ch1_AMPL = np.random.random(N_ch1) * 1000
    Ch2_AMPL = np.random.random(N_ch2) * 1000
    Ch3_AMPL = np.random.random(N_ch3) * 1000

    print(
        coinc_2(
            Ch1_TIME,
            Ch2_TIME,
            Ch1_AMPL,
            Ch2_AMPL,
            t_window=0.001,
        )
    )

    print(
        coinc_3(
            Ch1_TIME,
            Ch2_TIME,
            Ch3_TIME,
            Ch1_AMPL,
            Ch2_AMPL,
            Ch3_AMPL,
            t_window=0.001,
        )
    )