Ensemble forecast fun

tools.py

import numpy as np
import functools
import operator
import dask.array as da


def estimate_cube_size(cube):
    num_points = functools.reduce(operator.mul, cube.shape, 1)
    if cube[(0,) * len(cube.shape)].data.dtype != 'float32':
        return Fals
    return human_bytes((num_points * 32) / 8)


def human_bytes(num, suffix='B'):
    for unit in ['','Ki','Mi','Gi','Ti','Pi','Ei','Zi']:
        if abs(num) < 1024.0:
            return "%3.1f%s%s" % (num, unit, suffix)
        num /= 1024.0
    return "%.1f%s%s" % (num, 'Yi', suffix)


def metadata_broadcast(wbpt, wbpt_truth):
    x = wbpt_truth.coord("time").points.flatten()
    y = wbpt.coord("time").points

    index = np.argsort(x)
    sorted_x = x[index]
    sorted_index = np.searchsorted(sorted_x, y)

    yindex = np.take(index, sorted_index, mode="clip")
    mask = x[yindex] != y

    idx = np.ma.array(yindex, mask=mask)

    wbpt_truth_data = wbpt_truth.lazy_data()
    data = wbpt_truth_data[..., idx.flatten()].reshape(wbpt.shape)
    mask = np.broadcast_to(idx.mask.flatten().reshape(idx.shape), wbpt.shape)

    # da.from_array(mask, -1)
    wbpt_truth_data_bc = data #da.ma.masked_array(data, mask)
    wbpt_truth_bc = wbpt.copy() # where bc indicates broadcasted
    wbpt_truth_bc.data = wbpt_truth_data_bc

    return wbpt_truth_bc

WBPT_skill.ipynb