def flatten_spatial_dims(data):
    """Flatten all spatial dimensions of an input array.

    Assumes input array has shape (time, vert, lat, lon).
    """
    n_lon, n_lat, n_lev = data.shape[-1], data.shape[-2], data.shape[-3]
    n_pt = n_lon*n_lat*n_lev
    return data.reshape((-1, n_pt))