pad-irregular.py

def pad(array, pad_width, mode="edge", highlevel=True, behavior=None):
    if isinstance(pad_width, int):
        pad_width = (pad_width,)

    if len(pad_width) == 1:
        before = after = pad_width[0]
    else:
        before, after = pad_width

    def apply(array):
        if mode != "edge":
            raise NotImplementedError

        layout = ak.to_layout(array).toListOffsetArray64(True)
        content = np.asarray(layout.content)

        offsets = np.asarray(layout.offsets)
        starts = offsets[:-1]
        stops = offsets[1:]
        widths = np.diff(offsets)

        left = np.full_like(starts, before)
        right = np.full_like(starts, after)
        padding = left + right

        repeats = np.ones(len(content), dtype=np.int_)
        np.add.at(repeats, starts, left)
        np.add.at(repeats, stops - 1, right)
        new_content = np.repeat(content, repeats)

        new_widths = widths + padding
        new_offsets = np.zeros_like(offsets)
        np.cumsum(new_widths, out=new_offsets[1:])

        inner = ak.layout.ByteMaskedArray(
            ak.layout.Index8(np.repeat(widths > 0, new_widths)),
            ak.layout.NumpyArray(new_content),
            True,
        )

        return ak.layout.ListOffsetArray64(ak.layout.Index64(new_offsets), inner)

    return apply_at_axis(array, apply, axis=-1, highlevel=highlevel, behavior=behavior)


padded = pad(n.measured, 20)