How to use the xhistogram.xarray.align function in xhistogram

To help you get started, we’ve selected a few xhistogram examples, based on popular ways it is used in public projects.

Secure your code as it's written. Use Snyk Code to scan source code in minutes - no build needed - and fix issues immediately.

github xgcm / xhistogram / xhistogram / xarray.py View on Github external
# TODO: replace this with a more robust function
    assert len(bins)==N_args
    for bin in bins:
        assert isinstance(bin, np.ndarray), 'all bins must be numpy arrays'

    for a in args:
        # TODO: make this a more robust check
        assert a.name is not None, 'all arrays must have a name'

    # we drop coords to simplify alignment
    args = [da.reset_coords(drop=True) for da in args]
    if N_weights:
        args += [weights.reset_coords(drop=True)]
    # explicitly broadcast so we understand what is going into apply_ufunc
    # (apply_ufunc might be doing this by itself again)
    args = list(xr.align(*args, join='exact'))



    # what happens if we skip this?
    #args = list(xr.broadcast(*args))
    a0 = args[0]
    a_dims = a0.dims

    # roll our own broadcasting
    # now manually expand the arrays
    all_dims = [d for a in args for d in a.dims]
    all_dims_ordered = list(OrderedDict.fromkeys(all_dims))
    args_expanded = []
    for a in args:
        expand_keys = [d for d in all_dims_ordered if d not in a.dims]
        a_expanded = a.expand_dims({k: 1 for k in expand_keys})