ak_stack.py

@ak._connect._numpy.implements("stack")
def stack(arrays, axis=0, merge=True, mergebool=True, highlevel=True, behavior=None):
    """
    Args:
        arrays: Arrays to concatenate along any dimension.
        axis (int): The dimension at which this operation is applied. The
            outermost dimension is `0`, followed by `1`, etc., and negative
            values count backward from the innermost: `-1` is the innermost
            dimension, `-2` is the next level up, etc.
        merge (bool): If True, combine data into the same buffers wherever
            possible, eliminating unnecessary #ak.layout.UnionArray8_64 types
            at the expense of materializing #ak.layout.VirtualArray nodes.
        mergebool (bool): If True, boolean and nummeric data can be combined
            into the same buffer, losing information about False vs `0` and
            True vs `1`; otherwise, they are kept in separate buffers with
            distinct types (using an #ak.layout.UnionArray8_64).
        highlevel (bool): If True, return an #ak.Array; otherwise, return
            a low-level #ak.layout.Content subclass.
        behavior (None or dict): Custom #ak.behavior for the output array, if
            high-level.

    Returns an array with `arrays` concatenated. For `axis=0`, this means that
    one whole array follows another. For `axis=1`, it means that the `arrays`
    must have the same lengths and nested lists are each concatenated,
    element for element, and similarly for deeper levels.
    """
    nplike = ak.nplike.of(*arrays)
    behavior = ak._util.behaviorof(*arrays, behavior=behavior)
    layouts = [
        ak.operations.convert.to_layout(
            x, allow_record=False if axis == 0 else True, allow_other=True
        )
        for x in arrays
    ]

    # We need at least one array
    contents = [x for x in layouts if isinstance(x, ak.layout.Content)]
    if not contents:
        raise ValueError(
            "need at least one array to concatenate"
            + ak._util.exception_suffix(__file__)
        )

    # `axis` should lie within the range of possible axes
    posaxis = contents[0].axis_wrap_if_negative(axis)
    maxdepth = max(x.minmax_depth[1] for x in contents)
    if not 0 <= posaxis <= maxdepth:
        raise ValueError(
            "axis={} is beyond the depth of this array or the depth of this array "
            "is ambiguous".format(axis) + ak._util.exception_suffix(__file__)
        )
    # All arrays should understand the same value for `axis`
    if any((x.axis_wrap_if_negative(axis) != posaxis) for x in contents):
        raise ValueError(
            "arrays to concatenate do not have the same depth for negative "
            "axis={}".format(axis) + ak._util.exception_suffix(__file__)
        )

    # New first axis
    if posaxis == 0:
        layouts = [
            x
            if isinstance(x, ak.layout.Content)
            else ak.operations.convert.to_layout([x])
            for x in layouts
        ]
        length = len(layouts[0])

        tags = nplike.repeat(nplike.arange(len(layouts)), length)
        index = nplike.broadcast_to(
            nplike.arange(length), (len(layouts), length)
        ).ravel()

        inner = ak.layout.UnionArray8_64(
            ak.layout.Index8(tags), ak.layout.Index64(index), layouts
        ).simplify(merge=merge, mergebool=mergebool)

        offset = nplike.arange(0, len(index) + 1, length)
        out = ak.layout.ListOffsetArray64(ak.layout.Index64(offset), inner)

    else:

        def stack_contents(contents, length):
            nplike = ak.nplike.of(*contents)
            tags = nplike.broadcast_to(
                nplike.arange(len(contents)), (length, len(contents))
            ).ravel()
            index = nplike.repeat(nplike.arange(length), len(contents))

            inner = ak.layout.UnionArray8_64(
                ak.layout.Index8(tags), ak.layout.Index64(index), contents
            ).simplify(merge=merge, mergebool=mergebool)

            offset = nplike.arange(0, len(index) + 1, len(contents))
            return ak.layout.ListOffsetArray64(ak.layout.Index64(offset), inner)

        def getfunction(inputs, depth):
            if depth == posaxis and any(
                isinstance(x, ak._util.optiontypes) for x in inputs
            ):
                nextinputs = []
                for x in inputs:
                    if isinstance(x, ak._util.optiontypes) and isinstance(
                        x.content, ak._util.listtypes
                    ):
                        nextinputs.append(
                            ak.operations.structure.fill_none(
                                x, [], axis=0, highlevel=False
                            )
                        )
                    else:
                        nextinputs.append(x)
                inputs = nextinputs

            # New axis above existing
            if depth == posaxis and all(
                isinstance(x, ak._util.listtypes)
                or not isinstance(x, ak.layout.Content)
                for x in inputs
            ):
                length = max(len(x) for x in inputs if isinstance(x, ak.layout.Content))
                nextinputs = []
                for x in inputs:
                    if isinstance(x, ak.layout.Content):
                        nextinputs.append(x)
                    else:
                        nextinputs.append(
                            ak.layout.ListOffsetArray64(
                                ak.layout.Index64(
                                    nplike.arange(length + 1, dtype=np.int64)
                                ),
                                ak.layout.NumpyArray(
                                    nplike.broadcast_to(nplike.array([x]), (length,))
                                ),
                            )
                        )

                return lambda: (stack_contents(nextinputs, length),)

            # New axis at end
            elif depth == posaxis and all(
                isinstance(x, ak.layout.NumpyArray)
                or not isinstance(x, ak.layout.Content)
                for x in inputs
            ):
                length = max(len(x) for x in inputs if isinstance(x, ak.layout.Content))
                nextinputs = []
                for x in inputs:
                    if isinstance(x, ak.layout.Content):
                        nextinputs.append(x)
                    else:
                        nextinputs.append(
                            ak.layout.NumpyArray(
                                nplike.broadcast_to(nplike.array([x]), (length,))
                            )
                        )

                return lambda: (stack_contents(nextinputs, length),)

            elif any(
                x.purelist_depth == 1
                for x in inputs
                if isinstance(x, ak.layout.Content)
            ):
                raise ValueError(
                    "at least one array is not deep enough to concatenate at "
                    "axis={}".format(axis) + ak._util.exception_suffix(__file__)
                )

            else:
                return None

        (out,) = ak._util.broadcast_and_apply(
            layouts,
            getfunction,
            behavior=behavior,
            right_broadcast=False,
            numpy_to_regular=True,
        )

    return ak._util.maybe_wrap(out, behavior, highlevel)