import pandas as pd
import servicex as sx
import awkward as ak
from func_adl_servicex import ServiceXSourceUpROOT
import matplotlib.pyplot as plt
dataset_name = "cms:DYJetsToLL_M-50_TuneCP5_13TeV-amcatnloFXFX-pythia8/RunIIAutumn18NanoAODv7-Nano02Apr2020_102X_upgrade2018_realistic_v21_ext2-v1/NANOAODSIM"
sx_dataset = sx.ServiceXDataset(dataset_name, "uproot")
ds = ServiceXSourceUpROOT(sx_dataset, "Events")
dimu_mass = ds.Where(lambda event: event.nMuon == 2)\
.Select(lambda event: Zip({'p4': Zip({'pt': event.Muon_pt,
'eta': event.Muon_eta,
'phi': event.Muon_phi,
'mass': event.Muon_mass}).ToFourMomenta(),
'charge': event.Muon_charge})
.Choose(2)
.Where(lambda dimu: dimu[0].charge * dimu[1].charge < 0)
.Select(lambda dimu: (dimu[0].p4 + dimu[1].p4).mass)).AsAwkwardArray().value()
cms:DYJetsToLL_M-50_...: 0%| | 0/9000000000.0 [00:00]
cms:DYJetsToLL_M-50_... Downloaded: 0%| | 0/9000000000.0 [00:00]
dimu_mass
<Array [[90.6], [71.5], ... [88.4], [87.3]] type='32974713 * var * float32'>
plt.hist(ak.flatten(dimu_mass), bins=100, range=(0, 100))
plt.xlabel(r'$\mu\mu$ Mass [GeV]')
plt.ylabel('Events')
plt.show()