Asif Rahman

Select a random window of maximum duration in NumPy

Given a sequence of indices \(x\) and times \(t\), we want to select a window
of data from \(x\) with a maximum duration and a maximum number of measurements.

def select_random_window(
    x: np.ndarray, t: np.ndarray, max_window_size: int, dur: float
):
    """Select a random window of consecutive elements with a maximum duration.
    This function has extra logic that ensures the selected windows have length
    max_window_size whenever possible.

    Args:
        x (np.ndarray): sequence of indices
        max_window_size (int): maximum number of observations
    
    Returns:
        1D array of samples in the observation window.
    """
    L = len(x)
    if L <= max_window_size:
        return x
    if (t[-1] - t[0]) <= dur:
        return x
    L_tmp = L
    if L < max_window_size - 6:
        L_tmp = L - int(max_window_size / 2)
    else:
        L_tmp = L - 1
    start = random.randint(0, L_tmp)
    t_start = t[start]
    t_end = t_start + dur
    if t_end > t[-1]:
        return x[start:]
    end = np.argwhere((t - t_end) >= 0)[0][0]
    start = start if end <= L else max(0, start - (end - L))
    return x[start:end]