def test_with_name_field():
array = ak.Array(
[
{"x": 0.0, "y": []},
{"x": 1.1, "y": [1]},
{"x": 2.2, "y": [2, 2]},
{"x": 3.3, "y": [3, 3, 3]},
]
)
array2 = ak.repartition(array, 2)
z = ak.Array([100, 200, 300, 400])
z2 = ak.repartition(z, 3)
one = ak.with_name(array, "Wilbur")
assert isinstance(one.layout, ak.layout.Content)
assert one.layout.parameters["__record__"] == "Wilbur"
one = ak.with_name(array2, "Wilbur")
assert isinstance(one.layout, ak.partition.PartitionedArray)
assert one.layout.partition(0).parameters["__record__"] == "Wilbur"
assert one.layout.partition(1).parameters["__record__"] == "Wilbur"
two = ak.with_field(array, z, "z")
assert isinstance(two.layout, ak.layout.Content)
assert ak.to_list(two) == [
{"x": 0.0, "y": [], "z": 100},
{"x": 1.1, "y": [1], "z": 200},
{"x": 2.2, "y": [2, 2], "z": 300},
{"x": 3.3, "y": [3, 3, 3], "z": 400},
]
two = ak.with_field(array2, z, "z")
assert isinstance(two.layout, ak.partition.PartitionedArray)
assert ak.to_list(two) == [
{"x": 0.0, "y": [], "z": 100},
{"x": 1.1, "y": [1], "z": 200},
{"x": 2.2, "y": [2, 2], "z": 300},
{"x": 3.3, "y": [3, 3, 3], "z": 400},
]
two = ak.with_field(array2, z2, "z")
assert isinstance(two.layout, ak.partition.PartitionedArray)
assert ak.to_list(two) == [
{"x": 0.0, "y": [], "z": 100},
{"x": 1.1, "y": [1], "z": 200},
{"x": 2.2, "y": [2, 2], "z": 300},
{"x": 3.3, "y": [3, 3, 3], "z": 400},
]
two = ak.with_field(array2, z2, "z")
assert isinstance(two.layout, ak.partition.PartitionedArray)
assert ak.to_list(two) == [
{"x": 0.0, "y": [], "z": 100},
{"x": 1.1, "y": [1], "z": 200},
{"x": 2.2, "y": [2, 2], "z": 300},
{"x": 3.3, "y": [3, 3, 3], "z": 400},
]
def test_unknown_type():
array = ak.Array({"x": np.arange(10)})
array = ak.with_field(base=array, what=None, where="unknown field1")
array = ak.with_field(base=array, what=[None], where="unknown field2")
# Try to access the type of a single element
# This raises a ValueError in #879
tpe1 = array["unknown field1"].type
tpe2 = array["unknown field2"].type
assert str(tpe1) == "10 * ?unknown"
assert str(tpe2) == "10 * ?unknown"
def test_broadcast_single_bool():
base = ak.Array([[{
"x": 0.1,
"y": 0.2,
"z": 0.3
}, {
"x": 0.4,
"y": 0.5,
"z": 0.6
}]])
base_new1 = ak.with_field(base, True, "always_true")
assert ak.to_list(base_new1.always_true) == [[True, True]]
base_new2 = ak.with_field(base_new1, base.x > 0.3, "sometimes_true")
assert ak.to_list(base_new2.always_true) == [[True, True]]
def test_one_level():
base = ak.zip({
"a": ak.zip({"x": [1, 2, 3]}),
"b": [1, 2, 3]
},
depth_limit=1)
what = ak.Array([1.1, 2.2, 3.3], check_valid=True)
assert ak.to_list(ak.with_field(base, what, where=["a", "y"])) == [
{
"b": 1,
"a": {
"x": 1,
"y": 1.1
}
},
{
"b": 2,
"a": {
"x": 2,
"y": 2.2
}
},
{
"b": 3,
"a": {
"x": 3,
"y": 3.3
}
},
]
base["a", "y"] = what
assert ak.to_list(base) == [
{
"b": 1,
"a": {
"x": 1,
"y": 1.1
}
},
{
"b": 2,
"a": {
"x": 2,
"y": 2.2
}
},
{
"b": 3,
"a": {
"x": 3,
"y": 3.3
}
},
]
def update(events, collections):
"""Return a shallow copy of events array with some collections swapped out"""
out = events
for name, value in collections.items():
out = ak.with_field(out, value, name)
return out
def test_listarray():
one = ak.Array([[{
"x": 1
}, {
"x": 2
}, {
"x": 3
}], [], [{
"x": 4
}, {
"x": 5
}]],
check_valid=True)
two = ak.Array([[1.1, 2.2, 3.3], [], [4.4, 5.5]], check_valid=True)
assert ak.to_list(ak.with_field(one, two, "y")) == [
[{
"x": 1,
"y": 1.1
}, {
"x": 2,
"y": 2.2
}, {
"x": 3,
"y": 3.3
}],
[],
[{
"x": 4,
"y": 4.4
}, {
"x": 5,
"y": 5.5
}],
]
three = ak.Array([100, 200, 300], check_valid=True)
assert ak.to_list(ak.with_field(one, three, "y")) == [
[{
"x": 1,
"y": 100
}, {
"x": 2,
"y": 100
}, {
"x": 3,
"y": 100
}],
[],
[{
"x": 4,
"y": 300
}, {
"x": 5,
"y": 300
}],
]
assert ak.to_list(ak.with_field(one, [100, 200, 300], "y")) == [
[{
"x": 1,
"y": 100
}, {
"x": 2,
"y": 100
}, {
"x": 3,
"y": 100
}],
[],
[{
"x": 4,
"y": 300
}, {
"x": 5,
"y": 300
}],
]
def test_regulararray():
content = ak.layout.NumpyArray(
np.array([0.0, 1.1, 2.2, 3.3, 4.4, 5.5, 6.6, 7.7, 8.8, 9.9]))
recordarray = ak.layout.RecordArray({"x": content})
regulararray = ak.Array(ak.layout.RegularArray(recordarray,
3,
zeros_length=0),
check_valid=True)
content2 = ak.layout.NumpyArray(np.array([100, 200, 300]))
regulararray2 = ak.Array(ak.layout.RegularArray(content2,
1,
zeros_length=0),
check_valid=True)
assert ak.to_list(ak.with_field(regulararray, regulararray2, "y")) == [
[{
"x": 0.0,
"y": 100
}, {
"x": 1.1,
"y": 100
}, {
"x": 2.2,
"y": 100
}],
[{
"x": 3.3,
"y": 200
}, {
"x": 4.4,
"y": 200
}, {
"x": 5.5,
"y": 200
}],
[{
"x": 6.6,
"y": 300
}, {
"x": 7.7,
"y": 300
}, {
"x": 8.8,
"y": 300
}],
]
content2 = ak.layout.NumpyArray(
np.array([100, 200, 300, 400, 500, 600, 700, 800, 900]))
regulararray2 = ak.Array(ak.layout.RegularArray(content2,
3,
zeros_length=0),
check_valid=True)
assert ak.to_list(ak.with_field(regulararray, regulararray2, "y")) == [
[{
"x": 0.0,
"y": 100
}, {
"x": 1.1,
"y": 200
}, {
"x": 2.2,
"y": 300
}],
[{
"x": 3.3,
"y": 400
}, {
"x": 4.4,
"y": 500
}, {
"x": 5.5,
"y": 600
}],
[{
"x": 6.6,
"y": 700
}, {
"x": 7.7,
"y": 800
}, {
"x": 8.8,
"y": 900
}],
]
content2 = ak.Array(ak.layout.NumpyArray(np.array([[100], [200], [300]])),
check_valid=True)
assert ak.to_list(ak.with_field(regulararray, content2, "y")) == [
[{
"x": 0.0,
"y": 100
}, {
"x": 1.1,
"y": 100
}, {
"x": 2.2,
"y": 100
}],
[{
"x": 3.3,
"y": 200
}, {
"x": 4.4,
"y": 200
}, {
"x": 5.5,
"y": 200
}],
[{
"x": 6.6,
"y": 300
}, {
"x": 7.7,
"y": 300
}, {
"x": 8.8,
"y": 300
}],
]
content2 = ak.Array(
ak.layout.NumpyArray(
np.array([[100, 200, 300], [400, 500, 600], [700, 800, 900]])),
check_valid=True,
)
assert ak.to_list(ak.with_field(regulararray, content2, "y")) == [
[{
"x": 0.0,
"y": 100
}, {
"x": 1.1,
"y": 200
}, {
"x": 2.2,
"y": 300
}],
[{
"x": 3.3,
"y": 400
}, {
"x": 4.4,
"y": 500
}, {
"x": 5.5,
"y": 600
}],
[{
"x": 6.6,
"y": 700
}, {
"x": 7.7,
"y": 800
}, {
"x": 8.8,
"y": 900
}],
]
def test_withfield():
base = ak.Array([{"x": 1}, {"x": 2}, {"x": 3}], check_valid=True)
what = ak.Array([1.1, 2.2, 3.3], check_valid=True)
assert ak.to_list(ak.with_field(base, what)) == [
{
"x": 1,
"1": 1.1
},
{
"x": 2,
"1": 2.2
},
{
"x": 3,
"1": 3.3
},
]
assert ak.to_list(ak.with_field(base, what, where="y")) == [
{
"x": 1,
"y": 1.1
},
{
"x": 2,
"y": 2.2
},
{
"x": 3,
"y": 3.3
},
]
base["z"] = what
assert ak.to_list(base) == [
{
"x": 1,
"z": 1.1
},
{
"x": 2,
"z": 2.2
},
{
"x": 3,
"z": 3.3
},
]
base["q"] = 123
assert ak.to_list(base) == [
{
"x": 1,
"z": 1.1,
"q": 123
},
{
"x": 2,
"z": 2.2,
"q": 123
},
{
"x": 3,
"z": 3.3,
"q": 123
},
]
base = ak.Array([{"x": 1}, {"x": 2}, {"x": 3}], check_valid=True)[2]
assert ak.to_list(ak.with_field(base, 100, "y")) == {"x": 3, "y": 100}
def test_make_mixins():
@ak.mixin_class(ak.behavior)
class Point(object):
def distance(self, other):
return np.sqrt((self.x - other.x)**2 + (self.y - other.y)**2)
@ak.mixin_class_method(np.equal, {"Point"})
def point_equal(self, other):
return np.logical_and(self.x == other.x, self.y == other.y)
@ak.mixin_class_method(np.abs)
def point_abs(self):
return np.sqrt(self.x**2 + self.y**2)
@ak.mixin_class_method(np.add, {"Point"})
def point_add(self, other):
return ak.zip(
{
"x": self.x + other.x,
"y": self.y + other.y
},
with_name="Point",
)
@ak.mixin_class(ak.behavior)
class WeightedPoint(Point):
@ak.mixin_class_method(np.equal, {"WeightedPoint"})
def weighted_equal(self, other):
return np.logical_and(self.point_equal(other),
self.weight == other.weight)
@ak.mixin_class_method(np.add, {"WeightedPoint"})
def weighted_add(self, other):
sumw = self.weight + other.weight
return ak.zip(
{
"x":
(self.x * self.weight + other.x * other.weight) / sumw,
"y":
(self.y * self.weight + other.y * other.weight) / sumw,
"weight": sumw,
},
with_name="WeightedPoint",
)
one = ak.Array(
[
[{
"x": 1,
"y": 1.1
}, {
"x": 2,
"y": 2.2
}, {
"x": 3,
"y": 3.3
}],
[],
[{
"x": 4,
"y": 4.4
}, {
"x": 5,
"y": 5.5
}],
],
with_name="Point",
)
two = ak.Array(
[
[{
"x": 0.9,
"y": 1
}, {
"x": 2,
"y": 2.2
}, {
"x": 2.9,
"y": 3
}],
[],
[{
"x": 3.9,
"y": 4
}, {
"x": 5,
"y": 5.5
}],
],
with_name="Point",
)
wone = ak.Array(ak.with_field(one, abs(one), "weight"),
with_name="WeightedPoint")
wtwo = ak.Array(ak.with_field(two, abs(two), "weight"),
with_name="WeightedPoint")
assert ak.to_list(one + wone) == [
[{
"x": 2,
"y": 2.2
}, {
"x": 4,
"y": 4.4
}, {
"x": 6,
"y": 6.6
}],
[],
[{
"x": 8,
"y": 8.8
}, {
"x": 10,
"y": 11.0
}],
]
assert ak.to_list(wone + wtwo) == [
[
{
"x": 0.9524937500390619,
"y": 1.052493750039062,
"weight": 2.831969279439222,
},
{
"x": 2.0,
"y": 2.2,
"weight": 5.946427498927402
},
{
"x": 2.9516640394605282,
"y": 3.1549921183815837,
"weight": 8.632349833200564,
},
],
[],
[
{
"x": 3.9515600270076154,
"y": 4.206240108030463,
"weight": 11.533018588312771,
},
{
"x": 5.0,
"y": 5.5,
"weight": 14.866068747318506
},
],
]
assert ak.to_list(abs(one)) == [
[1.4866068747318506, 2.973213749463701, 4.459820624195552],
[],
[5.946427498927402, 7.433034373659253],
]
assert ak.to_list(one.distance(wtwo)) == [
[0.14142135623730953, 0.0, 0.31622776601683783],
[],
[0.4123105625617664, 0.0],
]