coincidence calculation

[jizhang02]

Hello,

Since Gate10 doesn't implement the coincidence function yet. I write a simplified coincidence algorithm for the purpose of pet reconstruction. I don't know if it is reasonable or not. So welcome to discuss and make it robust.

• The main function:

import torch
from tqdm import tqdm


def calculate_coincidences_gpu(time_window, singles_data, no_repeat=True, keep='all'):
    """
    calculate coincidences on GPU.

    Args:
        time_window (float): used to determine whether the time difference between events is within the window.
        singles_data (dict): A dictionary containing various single event data
        no_repeat (bool, optional): Whether to avoid repeating events. 
            'True': Do not repeat singles. 
            'False': Repeat singles.
        keep (str, optional): The ways of keeping coincidences. #*In case no_repeat=True:
            'all': keep all coincidences that within time window. e.g. (1,2),(1,3),(1,4)|(5,6),(5,7),(5,8)
            'first': keep the first coincidences. e.g. (1,2)|(3,4)
            'energy': keep the coincidences that energy is high. e.g. (1,3)|(2,4)
        keep (str, optional): The ways of keeping coincidences. #*In case no_repeat=False:
            'all': keep all coincidences that within time window. e.g. (1,2),(1,3),(1,4)|(2,3),(2,4),(2,5)
            'first': keep the first coincidences. e.g. (1,2)|(2,3)
            'energy': keep the coincidences that energy is high. e.g. (1,3)|(2,3)

    Returns:
        the number of coincidences (int) and dictionary of coincidences (dict).
    """

    # initialize
    num_coincidences = 0
    single_time1 = []
    single_time2 = []
    array_eventid_1 = []
    array_eventid_2 = []
    array_event_x_1 = []
    array_event_x_2 = []
    array_event_y_1 = []
    array_event_y_2 = []
    array_event_z_1 = []
    array_event_z_2 = []
    array_x_1 = []
    array_x_2 = []
    array_y_1 = []
    array_y_2 = []
    array_z_1 = []
    array_z_2 = []
    array_post_z_1 = []
    array_post_z_2 = []
    array_post_x_1 = []
    array_post_x_2 = []
    array_post_y_1 = []
    array_post_y_2 = []
    array_energy_1 = []
    array_energy_2 = []
    array_prevolumeid_1 = []
    array_prevolumeid_2 = []

    # read from dictionary singles_data
    singles_time = singles_data['singles_time']
    energy = singles_data['energy']
    eventid = singles_data['eventid']
    event_X = singles_data['event_X']
    event_Y = singles_data['event_Y']
    event_Z = singles_data['event_Z']
    detector_X = singles_data['detector_X']
    detector_Y = singles_data['detector_Y']
    detector_Z = singles_data['detector_Z']
    detector_post_X = singles_data['detector_post_X']
    detector_post_Y = singles_data['detector_post_Y']
    detector_post_Z = singles_data['detector_post_Z']
    volumeid = singles_data['volumeid'] # crystal block id

    # move to GPU
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    singles_time = torch.tensor(singles_time, device=device)
    energy = torch.tensor(energy, device=device)
    #eventid = torch.tensor(eventid, device=device) # not correspond to castor
    event_X = torch.tensor(event_X, device=device)
    event_Y = torch.tensor(event_Y, device=device)
    event_Z = torch.tensor(event_Z, device=device)
    detector_X = torch.tensor(detector_X, device=device)
    detector_Y = torch.tensor(detector_Y, device=device)
    detector_Z = torch.tensor(detector_Z, device=device)
    detector_post_X = torch.tensor(detector_post_X, device=device)
    detector_post_Y = torch.tensor(detector_post_Y, device=device)
    detector_post_Z = torch.tensor(detector_post_Z, device=device)

    is_coincidence = torch.zeros_like(singles_time, dtype=torch.bool, device=device)
    total_length_time = len(singles_time)
    
    for i in tqdm(range(total_length_time - 1)):
        if is_coincidence[i]: 
            continue
        
        current_time = singles_time[i]
        ####################keep=='first'##############################
        if keep == 'first':

            for j in range(i + 1, total_length_time):
                if is_coincidence[j]: 
                    continue

                time_difference = singles_time[j] - current_time

                if 0 <= time_difference <= time_window:
                    # if two detectors are the same, continue to the next iteration
                    if volumeid[i] == volumeid[j]: #0_0_14_10_6
                        continue
                    else:
                        is_coincidence[j] = no_repeat # mark it as coincidence or not
                        num_coincidences += 1
                        single_time1.append(current_time)
                        single_time2.append(singles_time[j])
                        array_eventid_1.append(eventid[i])
                        array_eventid_2.append(eventid[i]) # same event
                        array_event_x_1.append(event_X[i])
                        array_event_x_2.append(event_X[j])
                        array_event_y_1.append(event_Y[i])
                        array_event_y_2.append(event_Y[j])
                        array_event_z_1.append(event_Z[i])
                        array_event_z_2.append(event_Z[j])
                        array_x_1.append(detector_X[i])
                        array_x_2.append(detector_X[j])
                        array_y_1.append(detector_Y[i])
                        array_y_2.append(detector_Y[j])
                        array_z_1.append(detector_Z[i])
                        array_z_2.append(detector_Z[j])
                        array_post_x_1.append(detector_post_X[i])
                        array_post_x_2.append(detector_post_X[j])
                        array_post_y_1.append(detector_post_Y[i])
                        array_post_y_2.append(detector_post_Y[j])
                        array_post_z_1.append(detector_post_Z[i])
                        array_post_z_2.append(detector_post_Z[j])
                        array_energy_1.append(energy[i])
                        array_energy_2.append(energy[j])
                        array_prevolumeid_1.append(volumeid[i])
                        array_prevolumeid_2.append(volumeid[j])
                        break # once found, do not need find again
                else: break  # if the time difference exceeds the time window, break the inner loop

        ####################keep=='energy'##############################
        elif keep == 'energy':

            max_energy = 0
            j_max = 0

            for j in range(i + 1, total_length_time):
                if is_coincidence[j]: 
                    continue

                time_difference = singles_time[j] - current_time

                if 0 <= time_difference <= time_window:
                    # if two detectors are the same, continue to the next iteration
                    if volumeid[i] != volumeid[j] and energy[j] > max_energy:
                        max_energy = energy[j]
                        j_max = j
                    
                elif j_max !=0: # means j_max is updated
                    is_coincidence[j_max] = no_repeat # mark it as coincidence or not
                    num_coincidences += 1
                    single_time1.append(current_time)
                    single_time2.append(singles_time[j_max])
                    array_eventid_1.append(eventid[i])
                    array_eventid_2.append(eventid[i]) # same event
                    array_event_x_1.append(event_X[i])
                    array_event_x_2.append(event_X[j_max])
                    array_event_y_1.append(event_Y[i])
                    array_event_y_2.append(event_Y[j_max])
                    array_event_z_1.append(event_Z[i])
                    array_event_z_2.append(event_Z[j_max])
                    array_x_1.append(detector_X[i])
                    array_x_2.append(detector_X[j_max])
                    array_y_1.append(detector_Y[i])
                    array_y_2.append(detector_Y[j_max])
                    array_z_1.append(detector_Z[i])
                    array_z_2.append(detector_Z[j_max])
                    array_post_x_1.append(detector_post_X[i])
                    array_post_x_2.append(detector_post_X[j_max])
                    array_post_y_1.append(detector_post_Y[i])
                    array_post_y_2.append(detector_post_Y[j_max])
                    array_post_z_1.append(detector_post_Z[i])
                    array_post_z_2.append(detector_post_Z[j_max])
                    array_energy_1.append(energy[i])
                    array_energy_2.append(energy[j_max])
                    array_prevolumeid_1.append(volumeid[i])
                    array_prevolumeid_2.append(volumeid[j_max])
                    break # if the time difference exceeds the time window, break the inner loop
                else:    
                    break  # if the time difference exceeds the time window, break the inner loop

        ####################keep=='all'##############################
        elif keep == 'all':

            for j in range(i + 1, total_length_time):
                if is_coincidence[j]: 
                    continue

                time_difference = singles_time[j] - current_time

                if 0 <= time_difference <= time_window:
                    # if two detectors are the same, continue to the next iteration
                    if volumeid[i] == volumeid[j]: #0_0_14_10_6
                        continue
                    else:
                        is_coincidence[j] = no_repeat # mark it as coincidence or not
                        num_coincidences += 1
                        single_time1.append(current_time)
                        single_time2.append(singles_time[j])
                        array_eventid_1.append(eventid[i])
                        array_eventid_2.append(eventid[i]) # same event
                        array_event_x_1.append(event_X[i])
                        array_event_x_2.append(event_X[j])
                        array_event_y_1.append(event_Y[i])
                        array_event_y_2.append(event_Y[j])
                        array_event_z_1.append(event_Z[i])
                        array_event_z_2.append(event_Z[j])
                        array_x_1.append(detector_X[i])
                        array_x_2.append(detector_X[j])
                        array_y_1.append(detector_Y[i])
                        array_y_2.append(detector_Y[j])
                        array_z_1.append(detector_Z[i])
                        array_z_2.append(detector_Z[j])
                        array_post_x_1.append(detector_post_X[i])
                        array_post_x_2.append(detector_post_X[j])
                        array_post_y_1.append(detector_post_Y[i])
                        array_post_y_2.append(detector_post_Y[j])
                        array_post_z_1.append(detector_post_Z[i])
                        array_post_z_2.append(detector_post_Z[j])
                        array_energy_1.append(energy[i])
                        array_energy_2.append(energy[j])
                        array_prevolumeid_1.append(volumeid[i])
                        array_prevolumeid_2.append(volumeid[j])
                else: 
                    break  # if the time difference exceeds the time window, break the inner loop
        else:
            print(f'Please give right word "all", "first", "energy".')
            return

    # make it a dictionary
    coincidence_data = {
        'single_time1': single_time1,
        'single_time2': single_time2,
        'array_eventid_1': array_eventid_1,
        'array_eventid_2': array_eventid_2,
        'array_event_x_1': array_event_x_1,
        'array_event_x_2': array_event_x_2,
        'array_event_y_1': array_event_y_1,
        'array_event_y_2': array_event_y_2,
        'array_event_z_1': array_event_z_1,
        'array_event_z_2': array_event_z_2,
        'array_x_1': array_x_1,
        'array_x_2': array_x_2,
        'array_y_1': array_y_1,
        'array_y_2': array_y_2,
        'array_z_1': array_z_1,
        'array_z_2': array_z_2,
        'array_post_x_1': array_post_x_1,
        'array_post_x_2': array_post_x_2,
        'array_post_y_1': array_post_y_1,
        'array_post_y_2': array_post_y_2,
        'array_post_z_1': array_post_z_1,
        'array_post_z_2': array_post_z_2,
        'array_energy_1': array_energy_1,
        'array_energy_2': array_energy_2,
        'array_prevolumeid_1': array_prevolumeid_1,
        'array_prevolumeid_2': array_prevolumeid_2,
    }
    torch.cuda.empty_cache()
    return num_coincidences, coincidence_data

• The batch processing function

import uproot
import numpy as np

def roots_coincidence(start_file, end_file, file_path, save_path, save_file, time_window, no_repeat=True, keep='all'):

    #initialize
    all_coincidences = 0
    all_array_eventid_1 = []
    all_array_eventid_2 = []
    all_array_event_x_1 = []
    all_array_event_x_2 = []
    all_array_event_y_1 = []
    all_array_event_y_2 = []
    all_array_event_z_1 = []
    all_array_event_z_2 = []
    all_singles_time1 = []
    all_singles_time2 = []
    all_array_x_1 = []
    all_array_x_2 = []
    all_array_y_1 = []
    all_array_y_2 = []
    all_array_z_1 = []
    all_array_z_2 = []
    all_array_post_x_1 = []
    all_array_post_x_2 = []
    all_array_post_y_1 = []
    all_array_post_y_2 = []
    all_array_post_z_1 = []
    all_array_post_z_2 = []
    all_array_energy_1 = []
    all_array_energy_2 = []
    all_array_volumeid_1 = []
    all_array_volumeid_2 = []

    # process 10 root files
    for i in range(start_file, end_file):
        #root_file = file_path + 'pet_'+str(i)+'.root'
        root_file = file_path + 'pet_20r5s.root'
        print(f'{root_file} is being processed...')

        branches = ["EventID", "EventPosition_X", "EventPosition_Y", "EventPosition_Z", 
                    "Position_X", "Position_Y", "Position_Z", "PostPosition_X", "PostPosition_Y", "PostPosition_Z", 
                    "TotalEnergyDeposit", "PreStepUniqueVolumeID", "GlobalTime"]

        Singles_Tree = uproot.open(root_file)['EnergyWindows'].arrays(branches, library="numpy")
        
        # make it a dictionary
        singles_data = {
            'singles_time': Singles_Tree['GlobalTime'],
            'eventid': Singles_Tree['EventID'],
            'event_X': Singles_Tree['EventPosition_X'],
            'event_Y': Singles_Tree['EventPosition_Y'],
            'event_Z': Singles_Tree['EventPosition_Z'],
            'detector_X': Singles_Tree['Position_X'],
            'detector_Y': Singles_Tree['Position_Y'],
            'detector_Z': Singles_Tree['Position_Z'],
            'detector_post_X': Singles_Tree['PostPosition_X'],
            'detector_post_Y': Singles_Tree['PostPosition_Y'],
            'detector_post_Z': Singles_Tree['PostPosition_Z'],
            'energy': Singles_Tree['TotalEnergyDeposit'],
            'volumeid': Singles_Tree['PreStepUniqueVolumeID']
        }

        coincidences, coincidence_data = calculate_coincidences_gpu(time_window, singles_data, no_repeat, keep)

        print(f'The number of coincidences {coincidences}')
        all_coincidences += coincidences
        all_singles_time1.extend(coincidence_data['single_time1'])
        all_singles_time2.extend(coincidence_data['single_time2'])
        all_array_eventid_1.extend(coincidence_data['array_eventid_1'])
        all_array_eventid_2.extend(coincidence_data['array_eventid_2'])
        all_array_event_x_1.extend(coincidence_data['array_event_x_1'])
        all_array_event_x_2.extend(coincidence_data['array_event_x_2'])
        all_array_event_y_1.extend(coincidence_data['array_event_y_1'])
        all_array_event_y_2.extend(coincidence_data['array_event_y_2'])
        all_array_event_z_1.extend(coincidence_data['array_event_z_1'])
        all_array_event_z_2.extend(coincidence_data['array_event_z_2'])
        all_array_x_1.extend(coincidence_data['array_x_1'])
        all_array_x_2.extend(coincidence_data['array_x_2'])
        all_array_y_1.extend(coincidence_data['array_y_1'])
        all_array_y_2.extend(coincidence_data['array_y_2'])
        all_array_z_1.extend(coincidence_data['array_z_1'])
        all_array_z_2.extend(coincidence_data['array_z_2'])
        all_array_post_x_1.extend(coincidence_data['array_post_x_1'])
        all_array_post_x_2.extend(coincidence_data['array_post_x_2'])
        all_array_post_y_1.extend(coincidence_data['array_post_y_1'])
        all_array_post_y_2.extend(coincidence_data['array_post_y_2'])
        all_array_post_z_1.extend(coincidence_data['array_post_z_1'])
        all_array_post_z_2.extend(coincidence_data['array_post_z_2'])
        all_array_energy_1.extend(coincidence_data['array_energy_1'])
        all_array_energy_2.extend(coincidence_data['array_energy_2'])
        all_array_volumeid_1.extend(coincidence_data['array_prevolumeid_1'])
        all_array_volumeid_2.extend(coincidence_data['array_prevolumeid_2'])


    tfile = uproot.recreate(save_path+save_file)
    tfile["Coincidence"] = {
    'EventID1': all_array_eventid_1,
    'EventID2': all_array_eventid_2,
    'EventPostPosition_X1': all_array_event_x_1,
    'EventPostPosition_X2': all_array_event_x_2,
    'EventPostPosition_Y1': all_array_event_y_1,
    'EventPostPosition_Y2': all_array_event_y_2,
    'EventPostPosition_Z1': all_array_event_z_1,
    'EventPostPosition_Z2': all_array_event_z_2,
    'Position_X1': all_array_x_1,
    'Position_X2': all_array_x_2,
    'Position_Y1': all_array_y_1,
    'Position_Y2': all_array_y_2,
    'Position_Z1': all_array_z_1,
    'Position_Z2': all_array_z_2,
    'PostPosition_X1': all_array_post_x_1,
    'PostPosition_X2': all_array_post_x_2,
    'PostPosition_Y1': all_array_post_y_1,
    'PostPosition_Y2': all_array_post_y_2,
    'PostPosition_Z1': all_array_post_z_1,
    'PostPosition_Z2': all_array_post_z_2,
    'TotalEnergyDeposit1': all_array_energy_1,
    'TotalEnergyDeposit2': all_array_energy_2,
    'PreStepUniqueVolumeID1': all_array_volumeid_1,
    'PreStepUniqueVolumeID2': all_array_volumeid_2,
    'GlobalTime1': all_singles_time1,
    'GlobalTime2': all_singles_time2,
    }
    tfile.close()
    print(f'Done, the number of total coincidences {all_coincidences}')

• call the function:

start_file = 0
end_file = 1
file_path = 'MC/roots/'
save_path = 'MC/coincidences/'
save_file = 'coincidence_test.root'
time_window = 100 # nanosecondes
roots_coincidence(start_file, end_file, file_path, save_path, save_file, time_window, no_repeat=True, keep='energy')

[dsarrut]

This is very welcome, thank you! There is a work in progress, started in the PR #355. We will contact you to participate to the group working on the coinc sorter. Thanks !