How to use the xhistogram.duck_array_ops.broadcast_arrays 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 / core.py View on Github external
ax_positive = ndim + ax
            assert ax_positive < ndim, 'axis must be less than ndim'
            axis_normed.append(ax_positive)
        axis =  np.atleast_1d(axis_normed)

    do_full_array = (axis is None) or (set(axis) == set(range(a0.ndim)))
    if do_full_array:
        kept_axes_shape = None
    else:
        kept_axes_shape = tuple([a0.shape[i]
                                 for i in range(a0.ndim) if i not in axis])

    all_args = list(args)
    if weights is not None:
         all_args += [weights]
    all_args_broadcast = broadcast_arrays(*all_args)

    def reshape_input(a):
        if do_full_array:
            d = a.ravel()[None, :]
        else:
            # reshape the array to 2D
            # axis 0: preserved axis after histogram
            # axis 1: calculate histogram along this axis
            new_pos = tuple(range(-len(axis), 0))
            c = np.moveaxis(a, axis, new_pos)
            split_idx = c.ndim - len(axis)
            dims_0 = c.shape[:split_idx]
            assert dims_0 == kept_axes_shape
            dims_1 = c.shape[split_idx:]
            new_dim_0 = np.prod(dims_0)
            new_dim_1 = np.prod(dims_1)