def test_abs_of_top_leveldata(servicex_ds):
f = ServiceXDatasetSource(servicex_ds)
df = xaod_table(f)
seq = abs(df.met)
a = make_local(seq)
assert a is not None
selection = extract_selection(servicex_ds)
txt = translate_linq(
f.Select("lambda e1: e1.met()").Select(
"lambda e2: abs(e2)").AsROOTTTree("file.root", "treeme", ['col1']))
assert clean_linq(selection) == txt
def test_numpy_abs(servicex_ds):
f = ServiceXDatasetSource(servicex_ds)
df = xaod_table(f)
import numpy as np
seq = np.abs(df.met)
make_local(seq)
selection = extract_selection(servicex_ds)
txt = translate_linq(
f.Select("lambda e1: e1.met()").Select(
"lambda e2: abs(e2)").AsROOTTTree("file.root", "treeme", ['col1']))
assert clean_linq(selection) == txt
def test_first_at_object_level(servicex_ds):
f = ServiceXDatasetSource(servicex_ds)
df = xaod_table(f)
seq = df.jets.First().pt
make_local(seq)
selection = extract_selection(servicex_ds)
txt = translate_linq(
f.Select("lambda e5: e5.jets()").Select(
"lambda e7: e7.First()").Select("lambda e8: e8.pt()").AsROOTTTree(
"file.root", "treeme", ['col1']))
assert clean_linq(selection) == txt
def test_count_of_objects(servicex_ds):
f = ServiceXDatasetSource(servicex_ds)
df = xaod_table(f)
seq = df.jets.Count()
make_local(seq)
selection = extract_selection(servicex_ds)
txt = translate_linq(
f.Select("lambda e1: e1.jets()").Select(
"lambda e2: e2.Count()").AsROOTTTree("file.root", "treeme",
['col1']))
assert clean_linq(selection) == txt
def test_binop_in_filter(servicex_ds):
f = ServiceXDatasetSource(servicex_ds)
df = xaod_table(f)
seq = df.jets[(df.jets.pt / 1000.0) > 30].pt
make_local(seq)
selection = extract_selection(servicex_ds)
txt = translate_linq(
f.Select("lambda e1: e1.jets()").Select(
"lambda e7: e7.Where(lambda e5: e5.pt()/1000.0 > 30)").Select(
"lambda e8: e8.Select(lambda e6: e6.pt())").AsROOTTTree(
"file.root", "treeme", ['col1']))
assert clean_linq(selection) == txt
def test_filter_and_abs(servicex_ds):
f = ServiceXDatasetSource(servicex_ds)
df = xaod_table(f)
seq = df.jets[(df.jets.pt > 30.0) & (abs(df.jets.eta) < 2.5)].pt
make_local(seq)
selection = extract_selection(servicex_ds)
txt = translate_linq(
f.Select("lambda e1: e1.jets()").Select(
"lambda e10: e10.Where(lambda e8: (e8.pt() > 30.0) and (abs(e8.eta()) < 2.5))"
).Select("lambda e11: e11.Select(lambda e9: e9.pt())").AsROOTTTree(
"file.root", "treeme", ['col1']))
assert clean_linq(selection) == txt
def test_filter_not(servicex_ds):
f = ServiceXDatasetSource(servicex_ds)
df = xaod_table(f)
seq = df.jets[~(df.jets.pt > 30.0)].pt
make_local(seq)
selection = extract_selection(servicex_ds)
txt = translate_linq(
f.Select("lambda e1: e1.jets()").Select(
"lambda e9: e9.Where(lambda e7: not (e7.pt() > 30.0))").Select(
"lambda e10: e10.Select(lambda e8: e8.pt())").AsROOTTTree(
"file.root", "treeme", ['col1']))
assert clean_linq(selection) == txt
def test_filter_jet_by_attributes(servicex_ds):
f = ServiceXDatasetSource(servicex_ds)
df = xaod_table(f)
seq = df.jets[df.jets.hasProdVtx & df.jets.hasDecayVtx].pt
make_local(seq)
selection = extract_selection(servicex_ds)
txt = translate_linq(
f.Select("lambda e1: e1.jets()").Select(
"lambda e7: e7.Where(lambda e2: e2.hasProdVtx() and e2.hasDecayVtx())"
).Select("lambda e8: e8.Select(lambda e6: e6.pt())").AsROOTTTree(
"file.root", "treeme", ['col1']))
assert clean_linq(selection) == txt
def test_jet_pt_filter_pts_gt(servicex_ds):
f = ServiceXDatasetSource(servicex_ds)
df = xaod_table(f)
seq = df.jets.pt[df.jets.pt > 30.0]
make_local(seq)
selection = extract_selection(servicex_ds)
txt = translate_linq(
f.Select("lambda e1: e1.jets()").Select(
"lambda e5: e5.Select(lambda e2: e2.pt())").Select(
"lambda e6: e6.Where(lambda e3: e3 > 30.0)").AsROOTTTree(
"file.root", "treeme", ['col1']))
assert clean_linq(selection) == txt
def test_pt_sub(servicex_ds):
f = ServiceXDatasetSource(servicex_ds)
df = xaod_table(f)
seq = df.jets.pt - 1000.0
make_local(seq)
selection = extract_selection(servicex_ds)
txt = translate_linq(
f.Select("lambda e1: e1.jets()").Select(
"lambda e4: e4.Select(lambda e2: e2.pt())").Select(
"lambda e5: e5.Select(lambda e3: e3 - 1000.0)").AsROOTTTree(
"file.root", "treeme", ['col1']))
assert clean_linq(selection) == txt
def test_collect_xaod_call_with_number(servicex_ds):
'Do this with the actual call we need in ATLAS'
f = ServiceXDatasetSource(servicex_ds)
df = xaod_table(f)
seq = df.Jets(22.0).pt
make_local(seq)
selection = extract_selection(servicex_ds)
txt = translate_linq(
f.Select("lambda e1: e1.Jets(22.0)").Select(
"lambda e3: e3.Select(lambda e2: e2.pt())").AsROOTTTree(
"file.root", "treeme", ['col1']))
assert clean_linq(selection) == txt
def test_filter_and_divide_with_call(servicex_ds):
f = ServiceXDatasetSource(servicex_ds)
df = xaod_table(f)
seq = df.jets().pt[df.jets().pt > 30.0] / 1000.0
make_local(seq)
selection = extract_selection(servicex_ds)
txt = translate_linq(
f.Select("lambda e1: e1.jets()").
Select("lambda e6: e6.Select(lambda e2: e2.pt())").Select(
"lambda e7: e7.Where(lambda e3: e3 > 30.0)").Select(
"lambda e8: e8.Select(lambda e5: e5 / 1000.0)").AsROOTTTree(
"file.root", "treeme", ['col1']))
assert clean_linq(selection) == txt
def test_jet_pt_filter_pts_ne(servicex_ds):
'Do this with the actual call we need in ATLAS'
f = ServiceXDatasetSource(servicex_ds)
df = xaod_table(f)
seq = df.jets.pt[df.jets.pt != 30.0]
make_local(seq)
selection = extract_selection(servicex_ds)
txt = translate_linq(
f.Select("lambda e1: e1.jets()").Select(
"lambda e5: e5.Select(lambda e2: e2.pt())").Select(
"lambda e6: e6.Where(lambda e4: e4 != 30.0)").AsROOTTTree(
"file.root", "treeme", ['col1']))
assert clean_linq(selection) == txt
def test_abs_of_data_with_calls(servicex_ds):
f = ServiceXDatasetSource(servicex_ds)
df = xaod_table(f)
seq = abs(df.jets().pt())
a = make_local(seq)
assert a is not None
selection = extract_selection(servicex_ds)
txt = translate_linq(
f.Select("lambda e1: e1.jets()").Select(
"lambda e4: e4.Select(lambda e2: e2.pt())").Select(
"lambda e5: e5.Select(lambda e3: abs(e3))").AsROOTTTree(
"file.root", "treeme", ['col1']))
assert clean_linq(selection) == txt
def test_collect_pts_as_call(servicex_ds):
f = ServiceXDatasetSource(servicex_ds)
df = xaod_table(f)
seq = df.jets().pt()
a = make_local(seq)
assert a is not None
assert len(a) == 283458
selection = extract_selection(servicex_ds)
txt = translate_linq(
f.Select("lambda e1: e1.jets()").Select(
"lambda e3: e3.Select(lambda e2: e2.pt())").AsROOTTTree(
"file.root", "treeme", ['col1']))
assert clean_linq(selection) == txt
def test_filter_chain(servicex_ds):
f = ServiceXDatasetSource(servicex_ds)
df = xaod_table(f)
seq1 = df.jets[df.jets.pt > 30.0]
seq = seq1[seq1.eta < 2.4].pt
make_local(seq)
selection = extract_selection(servicex_ds)
txt = translate_linq(
f.Select("lambda e1: e1.jets()").Select(
"lambda e6: e6.Where(lambda e3: e3.pt() > 30.0)").Select(
"lambda e7: e7.Where(lambda e4: e4.eta() < 2.4)").Select(
"lambda e5: e5.Select(lambda e2: e2.pt())").AsROOTTTree(
"file.root", "treeme", ['col1']))
assert clean_linq(selection) == txt
def test_numpy_histogram(servicex_ds):
f = ServiceXDatasetSource(servicex_ds)
df = xaod_table(f)
seq = histogram(df.met)
h = make_local(seq)
selection = extract_selection(servicex_ds)
txt = translate_linq(
f.Select("lambda e1: e1.met()").AsROOTTTree("file.root", "treeme",
['col1']))
assert clean_linq(selection) == txt
assert h is not None
assert isinstance(h, Tuple)
assert len(h) == 2
contents = h[0]
assert len(contents) == 10