pipeline/analysis/target.py

import numpy as np


def build_tgt_matrix(tl, aeps_events, tgt_dur):
    # compute timeline / AEP indices of entrances / exist to the target
    tl_tgt_start_idxs = []  # timeline indices of entrance in target
    tl_tgt_end_idxs   = []  # timeline indices of exit from target

    for i in range(len(tl) - 1):
        if tl[i][6] < 2 and tl[i+1][6] == 2:
            tl_tgt_start_idxs.append(i + 1)
        if tl[i][6] == 2 and tl[i+1][6] < 2:
            tl_tgt_end_idxs.append(i)

    # ignore first/last target if not ended
    if tl_tgt_start_idxs[-1] > tl_tgt_end_idxs[-1]:
        tl_tgt_start_idxs = tl_tgt_start_idxs[:-1]
    if tl_tgt_end_idxs[0] < tl_tgt_start_idxs[0]:
        tl_tgt_end_idxs = tl_tgt_end_idxs[1:]
    tl_tgt_start_idxs = np.array(tl_tgt_start_idxs)
    tl_tgt_end_idxs   = np.array(tl_tgt_end_idxs)

    aeps_tgt_start_idxs = []  # indices of first AEPs in target
    aeps_tgt_end_idxs   = []  # indices of last AEPs in target

    for idx in tl_tgt_start_idxs:
        aeps_tgt_start_idxs.append(np.abs(aeps_events[:, 0] - tl[idx][0]).argmin())
    for idx in tl_tgt_end_idxs:
        aeps_tgt_end_idxs.append(np.abs(aeps_events[:, 0] - tl[idx][0]).argmin())
    aeps_tgt_start_idxs = np.array(aeps_tgt_start_idxs)
    aeps_tgt_end_idxs = np.array(aeps_tgt_end_idxs)

    # successful / missed
    tgt_results = np.zeros(len(tl_tgt_start_idxs))
    succ_idxs = np.where((tl_tgt_end_idxs - tl_tgt_start_idxs > tgt_dur * 100 - 10) & \
                         (tl_tgt_end_idxs - tl_tgt_start_idxs < tgt_dur * 100 + 10))[0]
    tgt_results[succ_idxs] = 1

    # tl_idx_start, tl_idx_end, aep_idx_start, aer_idx_end, success / miss
    return np.column_stack([
        tl_tgt_start_idxs,
        tl_tgt_end_idxs,
        aeps_tgt_start_idxs,
        aeps_tgt_end_idxs,
        tgt_results
    ]).astype(np.int32)