From 272bec0f6ecf4485d1f346fde753699ade66dbb4 Mon Sep 17 00:00:00 2001 From: Jim Pivarski Date: Fri, 12 Jul 2019 10:28:00 -0500 Subject: [PATCH 1/2] add an axis parameter to the pad method --- awkward/array/base.py | 2 +- awkward/array/chunked.py | 4 ++-- awkward/array/indexed.py | 8 ++++---- awkward/array/jagged.py | 11 ++++++++++- awkward/array/masked.py | 4 ++-- awkward/array/objects.py | 12 +++++++++--- awkward/array/table.py | 4 ++-- awkward/array/union.py | 4 ++-- awkward/array/virtual.py | 4 ++-- awkward/version.py | 2 +- 10 files changed, 35 insertions(+), 20 deletions(-) diff --git a/awkward/array/base.py b/awkward/array/base.py index 8a9864a4..3d9b2fa9 100644 --- a/awkward/array/base.py +++ b/awkward/array/base.py @@ -473,7 +473,7 @@ def _util_flatten(cls, array, axis): return array.reshape(array.shape[:axis] + (-1,) + array.shape[axis + 2:]) @classmethod - def _util_pad(cls, array, length, maskedwhen, clip): + def _util_pad(cls, array, length, maskedwhen, clip, axis): if isinstance(array, AwkwardArray): return array.pad(length, maskedwhen=maskedwhen, clip=clip) diff --git a/awkward/array/chunked.py b/awkward/array/chunked.py index 833bc4f0..6b542822 100644 --- a/awkward/array/chunked.py +++ b/awkward/array/chunked.py @@ -630,8 +630,8 @@ def flatten(self, axis=0): out.knowchunksizes() return out - def pad(self, length, maskedwhen=True, clip=False): - return self.copy(chunks=[self._util_pad(x, length, maskedwhen, clip) for x in self._chunks], chunksizes=self.chunksizes) + def pad(self, length, maskedwhen=True, clip=False, axis=0): + return self.copy(chunks=[self._util_pad(x, length, maskedwhen, clip, axis) for x in self._chunks], chunksizes=self.chunksizes) def regular(self): self._valid() diff --git a/awkward/array/indexed.py b/awkward/array/indexed.py index ba6ab73a..85bc9056 100644 --- a/awkward/array/indexed.py +++ b/awkward/array/indexed.py @@ -247,8 +247,8 @@ def flatten(self, axis=0): self._valid() return self._util_flatten(self._content[self._index], axis) - def pad(self, length, maskedwhen=True, clip=False): - return self._util_pad(self._content[self._index], length, maskedwhen, clip) + def pad(self, length, maskedwhen=True, clip=False, axis=0): + return self._util_pad(self._content[self._index], length, maskedwhen, clip, axis) def regular(self): self._valid() @@ -725,8 +725,8 @@ def flatten(self, axis=0): out[mask] = content[self._inverse[mask]] return out - def pad(self, length, maskedwhen=True, clip=False): - return self._util_pad(self._content.dense, length, maskedwhen, clip) + def pad(self, length, maskedwhen=True, clip=False, axis=0): + return self._util_pad(self._content.dense, length, maskedwhen, clip, axis) def regular(self): self._valid() diff --git a/awkward/array/jagged.py b/awkward/array/jagged.py index 37a7da08..a5c7917d 100644 --- a/awkward/array/jagged.py +++ b/awkward/array/jagged.py @@ -1781,7 +1781,16 @@ def ready(x): table = first.Table.named("tuple", columns1, *columns2) return first.JaggedArray(first._starts, first._stops, table) - def pad(self, length, maskedwhen=True, clip=False): + def pad(self, length, maskedwhen=True, clip=False, axis=0): + if not self._util_isinteger(axis) or axis < 0: + raise TypeError("axis must be a non-negative integer (can't count from the end)") + + if axis > 0: + if not self._util_hasjagged(self._content): + raise ValueError("pad axis is too deep for this degree of jaggedness") + else: + return type(self).fromcounts(self.counts, self._content.pad(length, maskedwhen=maskedwhen, clip=clip, axis=(axis - 1))) + flatstarts = self._starts.reshape(-1) if clip: diff --git a/awkward/array/masked.py b/awkward/array/masked.py index 4520d299..1ad41d63 100644 --- a/awkward/array/masked.py +++ b/awkward/array/masked.py @@ -291,8 +291,8 @@ def flattentuple(self): def flatten(self, axis=0): return self._util_flatten(self._content[self.boolmask(maskedwhen=False)], axis) - def pad(self, length, maskedwhen=True, clip=False): - return self.copy(content=self._util_pad(self._content, length, maskedwhen, clip)) + def pad(self, length, maskedwhen=True, clip=False, axis=0): + return self.copy(content=self._util_pad(self._content, length, maskedwhen, clip, axis)) def regular(self): self._valid() diff --git a/awkward/array/objects.py b/awkward/array/objects.py index 453b1005..f811867d 100644 --- a/awkward/array/objects.py +++ b/awkward/array/objects.py @@ -262,8 +262,8 @@ def flattentuple(self): def flatten(self, axis=0): return self.copy(content=self._util_flatten(self._content, axis)) - def pad(self, length, maskedwhen=True, clip=False): - return self.copy(content=self._util_pad(self._content, length, maskedwhen, clip)) + def pad(self, length, maskedwhen=True, clip=False, axis=0): + return self.copy(content=self._util_pad(self._content, length, maskedwhen, clip, axis)) def regular(self): return self.numpy.array(self) @@ -652,7 +652,13 @@ def flatten(self, axis=0): else: return self.fromjagged(self.JaggedArray.fromcounts([len(content)], content)) - def pad(self, length, maskedwhen=None, clip=False): + def pad(self, length, maskedwhen=None, clip=False, axis=0): + if not self._util_isinteger(axis) or axis < 0: + raise TypeError("axis must be a non-negative integer (can't count from the end)") + + if axis > 0: + raise ValueError("axis too deep for StringArray") + if maskedwhen is None: maskedwhen = ord(b" ") elif not isinstance(maskedwhen, bytes) or not len(maskedwhen) == 1: diff --git a/awkward/array/table.py b/awkward/array/table.py index 2fc6b77a..5816a5be 100644 --- a/awkward/array/table.py +++ b/awkward/array/table.py @@ -767,10 +767,10 @@ def argcross(self, other, nested=False): def flatten(self, axis=0): raise ValueError("cannot flatten through a Table") - def pad(self, length, maskedwhen=True, clip=False): + def pad(self, length, maskedwhen=True, clip=False, axis=0): out = self.copy(contents={}) for n, x in self._contents.items(): - out[n] = self._util_pad(x, length, maskedwhen, clip) + out[n] = self._util_pad(x, length, maskedwhen, clip, axis) return out def regular(self): diff --git a/awkward/array/union.py b/awkward/array/union.py index 05593154..acab2dff 100644 --- a/awkward/array/union.py +++ b/awkward/array/union.py @@ -468,9 +468,9 @@ def flattentuple(self): def flatten(self, axis=0): raise NotImplementedError("flatten not yet implemented for UnionArray") - def pad(self, length, maskedwhen=True, clip=False): + def pad(self, length, maskedwhen=True, clip=False, axis=0): self._valid() - return self.copy(contents=[self._util_pad(x, length, maskedwhen, clip) for x in self._contents]) + return self.copy(contents=[self._util_pad(x, length, maskedwhen, clip, axis) for x in self._contents]) def regular(self): self._valid() diff --git a/awkward/array/virtual.py b/awkward/array/virtual.py index f5acf2a6..96c66d54 100644 --- a/awkward/array/virtual.py +++ b/awkward/array/virtual.py @@ -439,8 +439,8 @@ def flattentuple(self): def flatten(self, axis=0): return self._util_flatten(self.array, axis) - def pad(self, length, maskedwhen=True, clip=False): - return self._util_pad(self.array, length, maskedwhen, clip) + def pad(self, length, maskedwhen=True, clip=False, axis=0): + return self._util_pad(self.array, length, maskedwhen, clip, axis) def regular(self): return self._util_regular(self.array) diff --git a/awkward/version.py b/awkward/version.py index 7a464001..bbbe23da 100644 --- a/awkward/version.py +++ b/awkward/version.py @@ -4,7 +4,7 @@ import re -__version__ = "0.12.1" +__version__ = "0.12.2" version = __version__ version_info = tuple(re.split(r"[-\.]", __version__)) From 4c5ed76f2871b4bd26133820ec466cfbde0e0f7c Mon Sep 17 00:00:00 2001 From: Jim Pivarski Date: Fri, 12 Jul 2019 16:02:43 -0400 Subject: [PATCH 2/2] fixed problem with pad axis not propagating (as well as the hack that 'fixed' it before) --- awkward/array/base.py | 2 +- awkward/array/jagged.py | 5 +---- 2 files changed, 2 insertions(+), 5 deletions(-) diff --git a/awkward/array/base.py b/awkward/array/base.py index 3d9b2fa9..02ec22c0 100644 --- a/awkward/array/base.py +++ b/awkward/array/base.py @@ -475,7 +475,7 @@ def _util_flatten(cls, array, axis): @classmethod def _util_pad(cls, array, length, maskedwhen, clip, axis): if isinstance(array, AwkwardArray): - return array.pad(length, maskedwhen=maskedwhen, clip=clip) + return array.pad(length, maskedwhen=maskedwhen, clip=clip, axis=axis) elif len(array.shape) == 1: raise ValueError("pad cannot be applied to scalars") diff --git a/awkward/array/jagged.py b/awkward/array/jagged.py index a5c7917d..1b73255d 100644 --- a/awkward/array/jagged.py +++ b/awkward/array/jagged.py @@ -1786,10 +1786,7 @@ def pad(self, length, maskedwhen=True, clip=False, axis=0): raise TypeError("axis must be a non-negative integer (can't count from the end)") if axis > 0: - if not self._util_hasjagged(self._content): - raise ValueError("pad axis is too deep for this degree of jaggedness") - else: - return type(self).fromcounts(self.counts, self._content.pad(length, maskedwhen=maskedwhen, clip=clip, axis=(axis - 1))) + return type(self).fromcounts(self.counts, self._util_pad(self._content, length, maskedwhen, clip, axis - 1)) flatstarts = self._starts.reshape(-1)