def test_with_name_field():
array = awkward1.Array([{"x": 0.0, "y": []}, {"x": 1.1, "y": [1]}, {"x": 2.2, "y": [2, 2]}, {"x": 3.3, "y": [3, 3, 3]}])
array2 = awkward1.repartition(array, 2)
z = awkward1.Array([100, 200, 300, 400])
z2 = awkward1.repartition(z, 3)
one = awkward1.with_name(array, "Wilbur")
assert isinstance(one.layout, awkward1.layout.Content)
assert one.layout.parameters["__record__"] == "Wilbur"
one = awkward1.with_name(array2, "Wilbur")
assert isinstance(one.layout, awkward1.partition.PartitionedArray)
assert one.layout.partition(0).parameters["__record__"] == "Wilbur"
assert one.layout.partition(1).parameters["__record__"] == "Wilbur"
two = awkward1.with_field(array, z, "z")
assert isinstance(two.layout, awkward1.layout.Content)
assert awkward1.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 = awkward1.with_field(array2, z, "z")
assert isinstance(two.layout, awkward1.partition.PartitionedArray)
assert awkward1.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 = awkward1.with_field(array2, z2, "z")
assert isinstance(two.layout, awkward1.partition.PartitionedArray)
assert awkward1.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 = awkward1.with_field(array2, z2, "z")
assert isinstance(two.layout, awkward1.partition.PartitionedArray)
assert awkward1.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_broadcast_single_bool():
base = awkward1.Array([[{
"x": 0.1,
"y": 0.2,
"z": 0.3
}, {
"x": 0.4,
"y": 0.5,
"z": 0.6
}]])
base_new1 = awkward1.with_field(base, True, "always_true")
assert awkward1.to_list(base_new1.always_true) == [[True, True]]
base_new2 = awkward1.with_field(base_new1, base.x > 0.3, "sometimes_true")
assert awkward1.to_list(base_new2.always_true) == [[True, True]]
def test_listarray():
one = awkward1.Array([[{
"x": 1
}, {
"x": 2
}, {
"x": 3
}], [], [{
"x": 4
}, {
"x": 5
}]],
check_valid=True)
two = awkward1.Array([[1.1, 2.2, 3.3], [], [4.4, 5.5]], check_valid=True)
assert awkward1.to_list(awkward1.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 = awkward1.Array([100, 200, 300], check_valid=True)
assert awkward1.to_list(awkward1.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 awkward1.to_list(awkward1.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 = awkward1.layout.NumpyArray(
numpy.array([0.0, 1.1, 2.2, 3.3, 4.4, 5.5, 6.6, 7.7, 8.8, 9.9]))
recordarray = awkward1.layout.RecordArray({"x": content})
regulararray = awkward1.Array(awkward1.layout.RegularArray(recordarray, 3),
check_valid=True)
content2 = awkward1.layout.NumpyArray(numpy.array([100, 200, 300]))
regulararray2 = awkward1.Array(awkward1.layout.RegularArray(content2, 1),
check_valid=True)
assert awkward1.to_list(
awkward1.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 = awkward1.layout.NumpyArray(
numpy.array([100, 200, 300, 400, 500, 600, 700, 800, 900]))
regulararray2 = awkward1.Array(awkward1.layout.RegularArray(content2, 3),
check_valid=True)
assert awkward1.to_list(
awkward1.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 = awkward1.Array(awkward1.layout.NumpyArray(
numpy.array([[100], [200], [300]])),
check_valid=True)
assert awkward1.to_list(awkward1.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 = awkward1.Array(awkward1.layout.NumpyArray(
numpy.array([[100, 200, 300], [400, 500, 600], [700, 800, 900]])),
check_valid=True)
assert awkward1.to_list(awkward1.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 = awkward1.Array([{"x": 1}, {"x": 2}, {"x": 3}], check_valid=True)
what = awkward1.Array([1.1, 2.2, 3.3], check_valid=True)
assert awkward1.to_list(awkward1.with_field(base, what)) == [{
"x": 1,
"1": 1.1
}, {
"x": 2,
"1": 2.2
}, {
"x": 3,
"1": 3.3
}]
assert awkward1.to_list(awkward1.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 awkward1.to_list(base) == [{
"x": 1,
"z": 1.1
}, {
"x": 2,
"z": 2.2
}, {
"x": 3,
"z": 3.3
}]
base["q"] = 123
assert awkward1.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 = awkward1.Array([{"x": 1}, {"x": 2}, {"x": 3}], check_valid=True)[2]
assert awkward1.to_list(awkward1.with_field(base, 100, "y")) == {
"x": 3,
"y": 100
}
def test_make_mixins():
@awkward1.mixin_class(awkward1.behavior)
class Point(object):
def distance(self, other):
return numpy.sqrt((self.x - other.x)**2 + (self.y - other.y)**2)
@awkward1.mixin_class_method(numpy.equal, {"Point"})
def point_equal(self, other):
return numpy.logical_and(self.x == other.x, self.y == other.y)
@awkward1.mixin_class_method(numpy.abs)
def point_abs(self):
return numpy.sqrt(self.x**2 + self.y**2)
@awkward1.mixin_class_method(numpy.add, {"Point"})
def point_add(self, other):
print("hi")
return awkward1.zip(
{
"x": self.x + other.x,
"y": self.y + other.y
},
with_name="Point",
)
@awkward1.mixin_class(awkward1.behavior)
class WeightedPoint(Point):
@awkward1.mixin_class_method(numpy.equal, {"WeightedPoint"})
def weighted_equal(self, other):
return numpy.logical_and(self.point_equal(other),
self.weight == other.weight)
@awkward1.mixin_class_method(numpy.add, {"WeightedPoint"})
def weighted_add(self, other):
sumw = self.weight + other.weight
return awkward1.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 = awkward1.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 = awkward1.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 = awkward1.Array(awkward1.with_field(one, abs(one), "weight"),
with_name="WeightedPoint")
wtwo = awkward1.Array(awkward1.with_field(two, abs(two), "weight"),
with_name="WeightedPoint")
assert awkward1.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 awkward1.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 awkward1.to_list(abs(one)) == [
[1.4866068747318506, 2.973213749463701, 4.459820624195552],
[],
[5.946427498927402, 7.433034373659253],
]
assert awkward1.to_list(one.distance(wtwo)) == [
[0.14142135623730953, 0.0, 0.31622776601683783],
[],
[0.4123105625617664, 0.0],
]