def test_evaluate_dot(wrapped_tree):
wrapped_tree.new_variable("Muon.Px", wrapped_tree.array("Muon_Px"))
all_true = expressions.evaluate(wrapped_tree, "Muon.Px == Muon_Px")
assert all(all_true.all())
def test_evaulate_matches_array(wrapped_tree):
mu_px_array = wrapped_tree.array("Muon_Px") < 0.3
mu_px_evalu = expressions.evaluate(wrapped_tree, "Muon_Px < 0.3")
assert (mu_px_evalu == mu_px_array).all().all()
def test_evaluate_bool(wrapped_tree):
all_true = expressions.evaluate(wrapped_tree, "Muon_Px == Muon_Px")
assert all(all_true.all())
def test_evaluate(wrapped_tree):
Muon_py, Muon_pz = wrapped_tree.arrays(["Muon_Py", "Muon_Pz"],
outputtype=tuple)
mu_pt = expressions.evaluate(wrapped_tree, "sqrt(Muon_Px**2 + Muon_Py**2)")
assert len(mu_pt) == 100
assert all(mu_pt.counts == Muon_py.counts)
def test_broadcast(wrapped_tree):
expressions.evaluate(wrapped_tree, "NJet * Jet_Py + NElectron * Jet_Px")
with pytest.raises(ValueError):
expressions.evaluate(wrapped_tree, "Jet_Py + Muon_Px")