def apply_mask(Y, mask, axis=0):
    """Masks the data along a specified axis

    Parameters
    ----------
    Y : array-like
        Data to be masked.
    mask : array-like
        Boolean vector containing True for each element to keep.
    axis : int, optional
        Axis along which to operate, by default 0.

    Returns
    -------
    numpy.array
        Masked data.
    """
    Y = Y.swapaxes(0, axis)
    Y = Y[mask, ...]
    return Y.swapaxes(0, axis)