def threshold_img(data, threshold, mask=None, mask_out='below'):
    """ Threshold data, setting all values in the array above/below threshold to zero.
    Optionally, can provide a mask (a 1D array with the same length as data), in which
    case the threshold is first applied to the mask, and the resulting indices are used
    to threshold the data. This is primarily useful when, e.g., applying a statistical
    threshold to a z-value image based on a p-value threshold. The mask_out argument
    indicates whether to zero out values 'below' the threshold (default) or 'above' the
    threshold. Note: use 'above' when masking based on p values. """
    if mask is not None:
        mask = threshold_img(mask, threshold, mask_out=mask_out)
        return data * mask.astype(bool)
    if mask_out.startswith('b'):
        data[data < threshold] = 0
    elif mask_out.startswith('a'):
        data[data > threshold] = 0
    return data