How to use the awkward1.layout.ArrayGenerator function in awkward1

To help you get started, we’ve selected a few awkward1 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 CoffeaTeam / coffea / coffea / nanoevents / factory.py View on Github external
def _local2global(self, index, source_offsets, target_offsets, name):
        def globalindex():
            gidx = awkward1.Array(
                awkward1.layout.ListOffsetArray32(
                    awkward1.layout.Index32(source_offsets), index.generator(),
                )
            )
            gidx = gidx.mask[gidx >= 0] + target_offsets[:-1]
            return awkward1.fill_none(awkward1.flatten(gidx), -1)

        generator = awkward1.layout.ArrayGenerator(
            globalindex,
            (),
            {},
            form=awkward1.forms.Form.fromjson('"int64"'),
            length=index.generator.length,
        )
        return awkward1.layout.VirtualArray(
            generator,
            index.cache,
            "/".join([self._keyprefix, "local2global", name]),
            index.identities,
            index.parameters,
        )
github CoffeaTeam / coffea / coffea / nanoevents / factory.py View on Github external
def _nestedindex(self, indexers, name):
        def nestedindex():
            # idx = awkward1.concatenate([idx[:, None] for idx in indexers], axis=1)
            n = len(indexers)
            out = numpy.empty(n * len(indexers[0]), dtype="int64")
            for i, idx in enumerate(indexers):
                out[i::n] = idx
            offsets = numpy.arange(0, len(out) + 1, n, dtype=numpy.int64)
            return awkward1.layout.ListOffsetArray64(
                awkward1.layout.Index64(offsets), awkward1.layout.NumpyArray(out),
            )

        form = awkward1.forms.Form.fromjson(
            '{"class": "ListOffsetArray64", "offsets": "i64", "content": "int64"}'
        )
        generator = awkward1.layout.ArrayGenerator(
            nestedindex, (), {}, form=form, length=indexers[0].generator.length,
        )
        return awkward1.layout.VirtualArray(
            generator,
            self._cache,
            cache_key="/".join([self._keyprefix, "nestedindex", name]),
        )
github CoffeaTeam / coffea / coffea / nanoaod / nanoawkward1.py View on Github external
k[len(cname) :]: arrays[k]
                    for k in arrays
                    if k.startswith(name + "_")
                }
                recordparams = {
                    "__doc__": counts.parameters["__doc__"],
                    "__record__": mixin,
                    "events_key": self._keyprefix,
                }
                form = awkward1.forms.ListOffsetForm(
                    "i32",
                    awkward1.forms.RecordForm(
                        {k: v.form for k, v in content.items()}, parameters=recordparams
                    ),
                )
                generator = awkward1.layout.ArrayGenerator(
                    self._listarray,
                    (counts, content, recordparams),
                    {},
                    form=form,
                    length=len(self),
                )
                return awkward1.layout.VirtualArray(
                    generator,
                    self._cache,
                    cache_key="/".join([self._keyprefix, "file", name]),
                    parameters=recordparams,
                )
            elif name in arrays:
                # singleton
                return arrays[name]
            else:
github CoffeaTeam / coffea / coffea / nanoevents / factory.py View on Github external
)
            elif "o" + name in arrays:
                # list singleton
                offsets = arrays["o" + name]
                content = arrays[name]
                params = {
                    # This makes more sense as offsets doc but it seems that is empty
                    "__doc__": content.parameters["__doc__"],
                    "__array__": mixin,
                    "events_key": self._keyprefix,
                    "collection_name": name,
                }
                form = awkward1.forms.ListOffsetForm(
                    "i32", content.form, parameters=params
                )
                generator = awkward1.layout.ArrayGenerator(
                    self._listarray,
                    (offsets, content, params),
                    {},
                    form=form,
                    length=len(self),
                )
                source = "runtime"
                return awkward1.layout.VirtualArray(
                    generator,
                    self._cache,
                    cache_key="/".join([self._keyprefix, source, name]),
                    parameters=params,
                )
            elif name in arrays:
                # singleton
                return arrays[name]
github CoffeaTeam / coffea / coffea / nanoaod / nanoawkward1.py View on Github external
def _array(self, branch_name: bytes):
        interpretation = uproot.interpret(self._tree[branch_name])
        if isinstance(interpretation, uproot.asjagged):
            dtype = interpretation.content.type
            length = None
        else:
            dtype = interpretation.type
            length = len(self)
        form = awkward1.forms.Form.fromjson('"%s"' % dtype)
        generator = awkward1.layout.ArrayGenerator(
            self.reader, (branch_name,), {}, form=form, length=length,
        )
        return awkward1.layout.VirtualArray(
            generator,
            self._cache,
            cache_key="/".join([self._keyprefix, "file", branch_name.decode("ascii")]),
            parameters={"__doc__": self._tree[branch_name].title.decode("ascii"),},
        )