From 0f35ace506b96d81f9cad8c712f489c8e77284ba Mon Sep 17 00:00:00 2001 From: Jonas Rembser Date: Fri, 12 Apr 2019 00:41:12 +0200 Subject: [PATCH 1/2] Receive wrapjagged functionality from uproot-methods --- awkward/util.py | 65 +++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 65 insertions(+) diff --git a/awkward/util.py b/awkward/util.py index 649e34da..cb6d66a0 100644 --- a/awkward/util.py +++ b/awkward/util.py @@ -6,6 +6,12 @@ import importlib import sys from collections import OrderedDict +try: + from collections.abc import Iterable +except ImportError: + from collections import Iterable + +from functools import wraps import numpy @@ -137,3 +143,62 @@ class NDArrayOperatorsMixin(object): __pos__ = _unary_method(um.positive, 'pos') __abs__ = _unary_method(um.absolute, 'abs') __invert__ = _unary_method(um.invert, 'invert') + + +def _normalize_arrays(cls, arrays): + length = None + for i in range(len(arrays)): + if isinstance(arrays[i], Iterable): + if length is None: + length = len(arrays[i]) + break + if length is None: + raise TypeError("cannot construct an array if all arguments are scalar") + + arrays = list(arrays) + jaggedtype = [cls.awkward.JaggedArray] * len(arrays) + starts, stops = None, None + for i in range(len(arrays)): + if starts is None and isinstance(arrays[i], cls.awkward.JaggedArray): + starts, stops = arrays[i].starts, arrays[i].stops + + if isinstance(arrays[i], cls.awkward.JaggedArray): + jaggedtype[i] = type(arrays[i]) + + if not isinstance(arrays[i], Iterable): + arrays[i] = cls.awkward.numpy.full(length, arrays[i]) + + arrays[i] = cls.awkward.util.toarray(arrays[i], cls.awkward.numpy.float64) + + if starts is None: + return arrays + + for i in range(len(arrays)): + if not isinstance(arrays[i], cls.awkward.JaggedArray) or not (cls.awkward.numpy.array_equal(starts, arrays[i].starts) and cls.awkward.numpy.array_equal(stops, arrays[i].stops)): + content = cls.awkward.numpy.zeros(stops.max(), dtype=cls.awkward.numpy.float64) + arrays[i] = jaggedtype[i](starts, stops, content) + arrays[i] # invoke jagged broadcasting to align arrays + + return arrays + + +def unwrap_jagged(cls, awkcls, arrays): + if not isinstance(arrays[0], cls.awkward.JaggedArray): + return lambda x: x, arrays + + counts = arrays[0].counts.reshape(-1) + offsets = awkcls.counts2offsets(counts) + starts, stops = offsets[:-1], offsets[1:] + starts = starts.reshape(arrays[0].starts.shape[:-1] + (-1,)) + stops = stops.reshape(arrays[0].stops.shape[:-1] + (-1,)) + wrap, arrays = unwrap_jagged(cls, awkcls, [x.flatten() for x in arrays]) + return lambda x: awkcls(starts, stops, wrap(x)), arrays + + +def wrapjaggedmethod(awkcls): + def wrapjagged_decorator(func): + @wraps(func) + def func_wrapper(cls, *arrays): + wrap, arrays = unwrap_jagged(cls, awkcls, _normalize_arrays(cls, arrays)) + return wrap(func(cls, *arrays)) + return func_wrapper + return wrapjagged_decorator From 69f41ba11719d10c2cdfbc5024f316cb9253a3e8 Mon Sep 17 00:00:00 2001 From: Jim Pivarski Date: Thu, 11 Apr 2019 20:26:44 -0500 Subject: [PATCH 2/2] move wrapping-jagged methods higher in the file and bump version number --- awkward/util.py | 117 ++++++++++++++++++++++----------------------- awkward/version.py | 2 +- 2 files changed, 59 insertions(+), 60 deletions(-) diff --git a/awkward/util.py b/awkward/util.py index cb6d66a0..d311bb3b 100644 --- a/awkward/util.py +++ b/awkward/util.py @@ -50,6 +50,64 @@ def __get__(self, ins, typ): else: return lambda *args, **kwargs: self.fcn(False, ins, *args, **kwargs) +################################################################ wrappers (used to be in uproot-methods) + +def _normalize_arrays(cls, arrays): + length = None + for i in range(len(arrays)): + if isinstance(arrays[i], Iterable): + if length is None: + length = len(arrays[i]) + break + if length is None: + raise TypeError("cannot construct an array if all arguments are scalar") + + arrays = list(arrays) + jaggedtype = [cls.awkward.JaggedArray] * len(arrays) + starts, stops = None, None + for i in range(len(arrays)): + if starts is None and isinstance(arrays[i], cls.awkward.JaggedArray): + starts, stops = arrays[i].starts, arrays[i].stops + + if isinstance(arrays[i], cls.awkward.JaggedArray): + jaggedtype[i] = type(arrays[i]) + + if not isinstance(arrays[i], Iterable): + arrays[i] = cls.awkward.numpy.full(length, arrays[i]) + + arrays[i] = cls.awkward.util.toarray(arrays[i], cls.awkward.numpy.float64) + + if starts is None: + return arrays + + for i in range(len(arrays)): + if not isinstance(arrays[i], cls.awkward.JaggedArray) or not (cls.awkward.numpy.array_equal(starts, arrays[i].starts) and cls.awkward.numpy.array_equal(stops, arrays[i].stops)): + content = cls.awkward.numpy.zeros(stops.max(), dtype=cls.awkward.numpy.float64) + arrays[i] = jaggedtype[i](starts, stops, content) + arrays[i] # invoke jagged broadcasting to align arrays + + return arrays + +def unwrap_jagged(cls, awkcls, arrays): + if not isinstance(arrays[0], cls.awkward.JaggedArray): + return lambda x: x, arrays + + counts = arrays[0].counts.reshape(-1) + offsets = awkcls.counts2offsets(counts) + starts, stops = offsets[:-1], offsets[1:] + starts = starts.reshape(arrays[0].starts.shape[:-1] + (-1,)) + stops = stops.reshape(arrays[0].stops.shape[:-1] + (-1,)) + wrap, arrays = unwrap_jagged(cls, awkcls, [x.flatten() for x in arrays]) + return lambda x: awkcls(starts, stops, wrap(x)), arrays + +def wrapjaggedmethod(awkcls): + def wrapjagged_decorator(func): + @wraps(func) + def func_wrapper(cls, *arrays): + wrap, arrays = unwrap_jagged(cls, awkcls, _normalize_arrays(cls, arrays)) + return wrap(func(cls, *arrays)) + return func_wrapper + return wrapjagged_decorator + ################################################################ array helpers try: @@ -143,62 +201,3 @@ class NDArrayOperatorsMixin(object): __pos__ = _unary_method(um.positive, 'pos') __abs__ = _unary_method(um.absolute, 'abs') __invert__ = _unary_method(um.invert, 'invert') - - -def _normalize_arrays(cls, arrays): - length = None - for i in range(len(arrays)): - if isinstance(arrays[i], Iterable): - if length is None: - length = len(arrays[i]) - break - if length is None: - raise TypeError("cannot construct an array if all arguments are scalar") - - arrays = list(arrays) - jaggedtype = [cls.awkward.JaggedArray] * len(arrays) - starts, stops = None, None - for i in range(len(arrays)): - if starts is None and isinstance(arrays[i], cls.awkward.JaggedArray): - starts, stops = arrays[i].starts, arrays[i].stops - - if isinstance(arrays[i], cls.awkward.JaggedArray): - jaggedtype[i] = type(arrays[i]) - - if not isinstance(arrays[i], Iterable): - arrays[i] = cls.awkward.numpy.full(length, arrays[i]) - - arrays[i] = cls.awkward.util.toarray(arrays[i], cls.awkward.numpy.float64) - - if starts is None: - return arrays - - for i in range(len(arrays)): - if not isinstance(arrays[i], cls.awkward.JaggedArray) or not (cls.awkward.numpy.array_equal(starts, arrays[i].starts) and cls.awkward.numpy.array_equal(stops, arrays[i].stops)): - content = cls.awkward.numpy.zeros(stops.max(), dtype=cls.awkward.numpy.float64) - arrays[i] = jaggedtype[i](starts, stops, content) + arrays[i] # invoke jagged broadcasting to align arrays - - return arrays - - -def unwrap_jagged(cls, awkcls, arrays): - if not isinstance(arrays[0], cls.awkward.JaggedArray): - return lambda x: x, arrays - - counts = arrays[0].counts.reshape(-1) - offsets = awkcls.counts2offsets(counts) - starts, stops = offsets[:-1], offsets[1:] - starts = starts.reshape(arrays[0].starts.shape[:-1] + (-1,)) - stops = stops.reshape(arrays[0].stops.shape[:-1] + (-1,)) - wrap, arrays = unwrap_jagged(cls, awkcls, [x.flatten() for x in arrays]) - return lambda x: awkcls(starts, stops, wrap(x)), arrays - - -def wrapjaggedmethod(awkcls): - def wrapjagged_decorator(func): - @wraps(func) - def func_wrapper(cls, *arrays): - wrap, arrays = unwrap_jagged(cls, awkcls, _normalize_arrays(cls, arrays)) - return wrap(func(cls, *arrays)) - return func_wrapper - return wrapjagged_decorator diff --git a/awkward/version.py b/awkward/version.py index aa0b2876..15c35926 100644 --- a/awkward/version.py +++ b/awkward/version.py @@ -4,7 +4,7 @@ import re -__version__ = "0.9.0" +__version__ = "0.9.0rc2" version = __version__ version_info = tuple(re.split(r"[-\.]", __version__))