def _string_equal(one, two):
nplike = awkward1.nplike.of(one, two)
behavior = awkward1._util.behaviorof(one, two)
one, two = one.layout, two.layout
# first condition: string lengths must be the same
counts1 = nplike.asarray(one.count(axis=-1))
counts2 = nplike.asarray(two.count(axis=-1))
out = counts1 == counts2
# only compare characters in strings that are possibly equal (same length)
possible = nplike.logical_and(out, counts1)
possible_counts = counts1[possible]
if len(possible_counts) > 0:
onepossible = awkward1.without_parameters(one[possible])
twopossible = awkward1.without_parameters(two[possible])
reduced = awkward1.all(onepossible == twopossible, axis=-1).layout
# update same-length strings with a verdict about their characters
out[possible] = reduced
return awkward1._util.wrap(awkward1.layout.NumpyArray(out), behavior)
def getfunction(inputs, depth):
signature = [ufunc]
for x in inputs:
if isinstance(x, awkward1.layout.Content):
record = x.parameters.get("__record__")
array = x.parameters.get("__array__")
if record is not None:
signature.append(record)
elif array is not None:
signature.append(array)
elif isinstance(x, awkward1.layout.NumpyArray):
signature.append(awkward1.nplike.of(x).asarray(x).dtype.type)
else:
signature.append(None)
else:
signature.append(type(x))
custom = awkward1._util.overload(behavior, signature)
if custom is not None:
return lambda: adjust(custom, inputs, kwargs)
if all(
isinstance(x, awkward1.layout.NumpyArray)
or not isinstance(
x, (awkward1.layout.Content, awkward1.partition.PartitionedArray)
)
for x in inputs
):
nplike = awkward1.nplike.of(*inputs)
result = getattr(ufunc, method)(
*[nplike.asarray(x) for x in inputs], **kwargs
)
return lambda: (awkward1.layout.NumpyArray(result),)
return None
def _string_broadcast(layout, offsets):
nplike = awkward1.nplike.of(offsets)
offsets = nplike.asarray(offsets)
counts = offsets[1:] - offsets[:-1]
if awkward1._util.win:
counts = counts.astype(np.int32)
parents = nplike.repeat(nplike.arange(len(counts), dtype=counts.dtype),
counts)
return awkward1.layout.IndexedArray64(awkward1.layout.Index64(parents),
layout).project()
def take(self, indices, *args, **kwargs):
# https://pandas.pydata.org/pandas-docs/version/1.0.0/reference/api/pandas.api.extensions.ExtensionArray.take.html
allow_fill, fill_value = awkward1._util.extra(args, kwargs,
[("allow_fill", False),
("fill_value", None)])
vote()
nplike = awkward1.nplike.of(self)
if allow_fill:
content1 = self.layout
if isinstance(content1, awkward1.partition.PartitionedArray):
content1 = content1.toContent()
indices = nplike.asarray(indices, dtype=np.int64)
if fill_value is None:
index = awkward1.layout.Index64(indices)
layout = awkward1.layout.IndexedOptionArray64(
index, content1, parameters=self.layout.parameters)
return awkward1._util.wrap(layout,
awkward1._util.behaviorof(self))
else:
tags = (indices >= 0).view(np.int8)
index = indices.copy()
index[~tags] = 0
content0 = awkward1.operations.convert.from_iter(
[fill_value], highlevel=False)
tags = awkward1.layout.Index8(tags)
index = awkward1.layout.Index64(index)
layout = awkward1.layout.UnionArray8_64(
tags, index, [content0, content1])
return awkward1._util.wrap(layout,
awkward1._util.behaviorof(self))
else:
return self[indices]
def apply(inputs, depth):
nplike = awkward1.nplike.of(*inputs)
# handle implicit right-broadcasting (i.e. NumPy-like)
if any(isinstance(x, listtypes) for x in inputs) and not any(
isinstance(x, (awkward1.layout.Content,
awkward1.layout.Record)) and x.has_virtual_form
for x in inputs):
maxdepth = max(x.purelist_depth for x in inputs
if isinstance(x, awkward1.layout.Content))
if maxdepth > 0 and all(x.purelist_isregular for x in inputs
if isinstance(x, awkward1.layout.Content)):
nextinputs = []
for x in inputs:
if isinstance(x, awkward1.layout.Content):
while x.purelist_depth < maxdepth:
x = awkward1.layout.RegularArray(x, 1)
nextinputs.append(x)
if any(x is not y for x, y in zip(inputs, nextinputs)):
return apply(nextinputs, depth)
# now all lengths must agree
checklength(
[x for x in inputs if isinstance(x, awkward1.layout.Content)])
function = getfunction(inputs, depth)
# the rest of this is one switch statement
if function is not None:
return function()
elif any(isinstance(x, virtualtypes) for x in inputs):
return apply(
[
x if not isinstance(x, virtualtypes) else x.array
for x in inputs
],
depth,
)
elif any(isinstance(x, unknowntypes) for x in inputs):
return apply(
[
x if not isinstance(x, unknowntypes)
else awkward1.layout.NumpyArray(
nplike.array([], dtype=np.bool_)) for x in inputs
],
depth,
)
elif any(
isinstance(x, awkward1.layout.NumpyArray) and x.ndim > 1
for x in inputs):
return apply(
[
x if not (isinstance(x, awkward1.layout.NumpyArray)
and x.ndim > 1) else x.toRegularArray()
for x in inputs
],
depth,
)
elif any(isinstance(x, indexedtypes) for x in inputs):
return apply(
[
x if not isinstance(x, indexedtypes) else x.project()
for x in inputs
],
depth,
)
elif any(isinstance(x, uniontypes) for x in inputs):
tagslist = []
length = None
for x in inputs:
if isinstance(x, uniontypes):
tagslist.append(nplike.asarray(x.tags))
if length is None:
length = len(tagslist[-1])
elif length != len(tagslist[-1]):
raise ValueError(
"cannot broadcast UnionArray of length {0} "
"with UnionArray of length {1}".format(
length, len(tagslist[-1])) +
exception_suffix(__file__))
combos = nplike.stack(tagslist, axis=-1)
combos = combos.view([
(str(i), combos.dtype) for i in range(len(tagslist))
]).reshape(length)
tags = nplike.empty(length, dtype=np.int8)
index = nplike.empty(length, dtype=np.int64)
outcontents = []
for tag, combo in enumerate(nplike.unique(combos)):
mask = combos == combo
tags[mask] = tag
index[mask] = nplike.arange(nplike.count_nonzero(mask))
nextinputs = []
numoutputs = None
for i, x in enumerate(inputs):
if isinstance(x, uniontypes):
nextinputs.append(x[mask].project(combo[str(i)]))
elif isinstance(x, awkward1.layout.Content):
nextinputs.append(x[mask])
else:
nextinputs.append(x)
outcontents.append(apply(nextinputs, depth))
assert isinstance(outcontents[-1], tuple)
if numoutputs is not None:
assert numoutputs == len(outcontents[-1])
numoutputs = len(outcontents[-1])
tags = awkward1.layout.Index8(tags)
index = awkward1.layout.Index64(index)
return tuple(
awkward1.layout.UnionArray8_64(
tags, index, [x[i] for x in outcontents]).simplify()
for i in range(numoutputs))
elif any(isinstance(x, optiontypes) for x in inputs):
mask = None
for x in inputs:
if isinstance(x, optiontypes):
m = nplike.asarray(x.bytemask()).view(np.bool_)
if mask is None:
mask = m
else:
nplike.bitwise_or(mask, m, out=mask)
nextmask = awkward1.layout.Index8(mask.view(np.int8))
index = nplike.full(len(mask), -1, dtype=np.int64)
index[~mask] = nplike.arange(len(mask) -
nplike.count_nonzero(mask),
dtype=np.int64)
index = awkward1.layout.Index64(index)
if any(not isinstance(x, optiontypes) for x in inputs):
nextindex = nplike.arange(len(mask), dtype=np.int64)
nextindex[mask] = -1
nextindex = awkward1.layout.Index64(nextindex)
nextinputs = []
for x in inputs:
if isinstance(x, optiontypes):
nextinputs.append(x.project(nextmask))
elif isinstance(x, awkward1.layout.Content):
nextinputs.append(
awkward1.layout.IndexedOptionArray64(
nextindex, x).project(nextmask))
else:
nextinputs.append(x)
outcontent = apply(nextinputs, depth)
assert isinstance(outcontent, tuple)
return tuple(
awkward1.layout.IndexedOptionArray64(index, x).simplify()
for x in outcontent)
elif any(isinstance(x, listtypes) for x in inputs):
if all(
isinstance(x, awkward1.layout.RegularArray)
or not isinstance(x, listtypes) for x in inputs):
maxsize = max([
x.size for x in inputs
if isinstance(x, awkward1.layout.RegularArray)
])
for x in inputs:
if isinstance(x, awkward1.layout.RegularArray):
if maxsize > 1 and x.size == 1:
tmpindex = awkward1.layout.Index64(
nplike.repeat(
nplike.arange(len(x), dtype=np.int64),
maxsize))
nextinputs = []
for x in inputs:
if isinstance(x, awkward1.layout.RegularArray):
if maxsize > 1 and x.size == 1:
nextinputs.append(
awkward1.layout.IndexedArray64(
tmpindex,
x.content[:len(x) * x.size]).project())
elif x.size == maxsize:
nextinputs.append(x.content[:len(x) * x.size])
else:
raise ValueError(
"cannot broadcast RegularArray of size "
"{0} with RegularArray of size {1}".format(
x.size, maxsize) +
exception_suffix(__file__))
else:
nextinputs.append(x)
outcontent = apply(nextinputs, depth + 1)
assert isinstance(outcontent, tuple)
return tuple(
awkward1.layout.RegularArray(x, maxsize)
for x in outcontent)
else:
fcns = [
custom_broadcast(x, behavior) if isinstance(
x, awkward1.layout.Content) else None for x in inputs
]
first, secondround = None, False
for x, fcn in zip(inputs, fcns):
if (isinstance(x, listtypes)
and not isinstance(x, awkward1.layout.RegularArray)
and fcn is None):
first = x
break
if first is None:
secondround = True
for x in inputs:
if isinstance(x, listtypes) and not isinstance(
x, awkward1.layout.RegularArray):
first = x
break
offsets = first.compact_offsets64(True)
nextinputs = []
for x, fcn in zip(inputs, fcns):
if callable(fcn) and not secondround:
nextinputs.append(fcn(x, offsets))
elif isinstance(x, listtypes):
nextinputs.append(
x.broadcast_tooffsets64(offsets).content)
# handle implicit left-broadcasting (unlike NumPy)
elif isinstance(x, awkward1.layout.Content):
nextinputs.append(
awkward1.layout.RegularArray(
x, 1).broadcast_tooffsets64(offsets).content)
else:
nextinputs.append(x)
outcontent = apply(nextinputs, depth + 1)
assert isinstance(outcontent, tuple)
return tuple(
awkward1.layout.ListOffsetArray64(offsets, x)
for x in outcontent)
elif any(isinstance(x, recordtypes) for x in inputs):
keys = None
length = None
istuple = True
for x in inputs:
if isinstance(x, recordtypes):
if keys is None:
keys = x.keys()
elif set(keys) != set(x.keys()):
raise ValueError(
"cannot broadcast records because keys don't "
"match:\n {0}\n {1}".format(
", ".join(sorted(keys)), ", ".join(
sorted(x.keys()))) +
exception_suffix(__file__))
if length is None:
length = len(x)
elif length != len(x):
raise ValueError(
"cannot broadcast RecordArray of length {0} "
"with RecordArray of length {1}".format(
length, len(x)) + exception_suffix(__file__))
if not x.istuple:
istuple = False
outcontents = []
numoutputs = None
for key in keys:
outcontents.append(
apply(
[
x if not isinstance(x, recordtypes) else x[key]
for x in inputs
],
depth,
))
assert isinstance(outcontents[-1], tuple)
if numoutputs is not None:
assert numoutputs == len(outcontents[-1])
numoutputs = len(outcontents[-1])
return tuple(
awkward1.layout.RecordArray([x[i] for x in outcontents],
None if istuple else keys, length)
for i in range(numoutputs))
else:
raise ValueError("cannot broadcast: {0}".format(", ".join(
repr(type(x)) for x in inputs)) + exception_suffix(__file__))
def getfunction(inputs, depth):
signature = [ufunc]
for x in inputs:
if isinstance(x, awkward1.layout.Content):
record = x.parameter("__record__")
array = x.parameter("__array__")
if record is not None:
signature.append(record)
elif array is not None:
signature.append(array)
elif isinstance(x, awkward1.layout.NumpyArray):
signature.append(awkward1.nplike.of(x).asarray(x).dtype.type)
else:
signature.append(None)
else:
signature.append(type(x))
custom = awkward1._util.overload(behavior, signature)
if custom is not None:
return lambda: adjust(custom, inputs, kwargs)
inputs = [deregulate(x) for x in inputs]
if ufunc is numpy.matmul:
custom_matmul = getfunction_matmul(inputs)
if custom_matmul is not None:
return custom_matmul
if all(
isinstance(x, awkward1.layout.NumpyArray)
or not isinstance(
x, (awkward1.layout.Content, awkward1.partition.PartitionedArray)
)
for x in inputs
):
nplike = awkward1.nplike.of(*inputs)
result = getattr(ufunc, method)(
*[nplike.asarray(x) for x in inputs], **kwargs
)
return lambda: (
awkward1.operations.convert.from_numpy(result, highlevel=False),
)
for x in inputs:
if isinstance(x, awkward1.layout.Content):
chained_behavior = awkward1._util.Behavior(awkward1.behavior, behavior)
apply_ufunc = chained_behavior[numpy.ufunc, x.parameter("__array__")]
if apply_ufunc is not None:
out = adjust_apply_ufunc(
apply_ufunc, ufunc, method, inputs, kwargs
)
if out is not None:
return out
apply_ufunc = chained_behavior[numpy.ufunc, x.parameter("__record__")]
if apply_ufunc is not None:
out = adjust_apply_ufunc(
apply_ufunc, ufunc, method, inputs, kwargs
)
if out is not None:
return out
if all(
x.parameter("__array__") is not None
or x.parameter("__record__") is not None
for x in inputs if isinstance(x, awkward1.layout.Content)
):
custom_types = []
for x in inputs:
if isinstance(x, awkward1.layout.Content):
if x.parameter("__array__") is not None:
custom_types.append(x.parameter("__array__"))
elif x.parameter("__record__") is not None:
custom_types.append(x.parameter("__record__"))
else:
custom_types.append(type(x).__name__)
else:
custom_types.append(type(x).__name__)
exception = ValueError(
"no overloads for custom types: {0}({1})".format(
ufunc.__name__,
", ".join(custom_types),
)
+ awkward1._util.exception_suffix(__file__)
)
awkward1._util.deprecate(exception, "1.0.0", date="2020-12-01")
return None