def apply_mask(data, mask):
    """ Mask all variables in a dataset with the provided mask. Routine uses
    mask(z=0) for variables without a vertical dimension.

    Parameters
    ----------
    data : xarray.Dataset
            Dataset containing 4D ('time_counter','x','y','z') and
            3D ('time_counter','x','y','z') variables.
    mask : xarray.Dataset
            Mask should not contain any dimensions that are not present in data
            to avoid broadcasting

    Returns
    -------
    xarray.Dataset
    """
    for v in data.data_vars:
        if 'z' in data[v].dims:
            data[v] = data[v].where(mask == 1)
        else:
            data[v] = data[v].where(mask.isel(z=0) == 1)
    return data