Skip to content

Commit

Permalink
Added muons and primary vertices
Browse files Browse the repository at this point in the history
  • Loading branch information
@gordonwatts
gordonwatts committed May 3, 2024
1 parent 977785b commit 41d775b
Showing 1 changed file with 61 additions and 151 deletions.
212 changes: 61 additions & 151 deletions servicex/query_library.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -46,13 +46,17 @@ def build_xaod_query(q_type: str) -> FuncADLQuery:
# Build the query
# TODO: The EventInfo argument should default correctly
# (that may just be a matter of using func_adl xaod r22)
# Same for Vertices
# TODO: dataclass should be supported so as not to lose type-following!

# Build the event model.
event_model = ds.Select(
lambda e: {
"evt": e.EventInfo("EventInfo"),
"jet": e.Jets(),
"ele": e.Electrons(),
"muons": e.Muons(),
"pv": e.Vertices("PrimaryVertices"),
}
)

Expand All @@ -69,6 +73,9 @@ def skim_jets(events, jet_cut: float):
"jet": e.jet.Where(
lambda j: (j.pt() / 1000 > jet_cut) and (abs(j.eta()) < 2.5)
),
"ele": e.ele,
"muons": e.muons,
"pv": e.pv,
}
)

Expand All @@ -85,101 +92,6 @@ def count_jets(events, n_jets: int):
event_model = count_jets(event_model, 4)

# Get out the data we need
# query = event_model.Select(
# lambda ei: {
# "event_number": ei.evt.eventNumber(), # type: ignore
# "run_number": ei.evt.runNumber(), # type: ignore
# "jet_pt": ei.jet.Select(lambda j: j.pt() / 1000), # type: ignore
# "jet_eta": ei.jet.Select(lambda j: j.eta()), # type: ignore
# "jet_phi": ei.jet.Select(lambda j: j.phi()), # type: ignore
# "jet_m": ei.jet.Select(lambda j: j.m()), # type: ignore
# "jet_EnergyPerSampling": ei.jet.Select( # type: ignore
# lambda j: j.getAttribute[cpp_vfloat]("EnergyPerSampling")
# ),
# "jet_SumPtTrkPt500": ei.jet.Select( # type: ignore
# lambda j: j.getAttribute[cpp_vfloat]("SumPtTrkPt500")
# ),
# "jet_TrackWidthPt1000": ei.jet.Select( # type: ignore
# lambda j: j.getAttribute[cpp_vfloat]("TrackWidthPt1000")
# ),
# "jet_NumTrkPt500": ei.jet.Select( # type: ignore
# lambda j: j.getAttribute[cpp_vint]("NumTrkPt500")
# ),
# "jet_NumTrkPt1000": ei.jet.Select( # type: ignore
# lambda j: j.getAttribute[cpp_vint]("NumTrkPt1000")
# ),
# "jet_SumPtChargedPFOPt500": ei.jet.Select( # type: ignore
# lambda j: j.getAttribute[cpp_vfloat]("SumPtChargedPFOPt500")
# ),
# "jet_Timing": ei.jet.Select( # type: ignore
# lambda j: j.getAttribute[cpp_float]("Timing")
# ),
# "jet_JetConstitScaleMomentum_eta": ei.jet.Select( # type: ignore
# lambda j: j.getAttribute[cpp_float]("JetConstitScaleMomentum_eta")
# ),
# "jet_ActiveArea4vec_eta": ei.jet.Select( # type: ignore
# lambda j: j.getAttribute[cpp_float]("ActiveArea4vec_eta")
# ),
# "jet_DetectorEta": ei.jet.Select( # type: ignore
# lambda j: j.getAttribute[cpp_float]("DetectorEta")
# ),
# "jet_JetConstitScaleMomentum_phi": ei.jet.Select( # type: ignore
# lambda j: j.getAttribute[cpp_float]("JetConstitScaleMomentum_phi")
# ),
# "jet_ActiveArea4vec_phi": ei.jet.Select( # type: ignore
# lambda j: j.getAttribute[cpp_float]("ActiveArea4vec_phi")
# ),
# "jet_JetConstitScaleMomentum_m": ei.jet.Select( # type: ignore
# lambda j: j.getAttribute[cpp_float]("JetConstitScaleMomentum_m")
# ),
# "jet_JetConstitScaleMomentum_pt": ei.jet.Select( # type: ignore
# lambda j: j.getAttribute[cpp_float]("JetConstitScaleMomentum_pt")
# ),
# "jet_Width": ei.jet.Select( # type: ignore
# lambda j: j.getAttribute[cpp_float]("Width")
# ),
# "jet_EMFrac": ei.jet.Select( # type: ignore
# lambda j: j.getAttribute[cpp_float]("EMFrac")
# ),
# "jet_ActiveArea4vec_m": ei.jet.Select( # type: ignore
# lambda j: j.getAttribute[cpp_float]("ActiveArea4vec_m")
# ),
# "jet_DFCommonJets_QGTagger_TracksWidth": ei.jet.Select( # type: ignore
# lambda j: j.getAttribute[cpp_float]("DFCommonJets_QGTagger_TracksWidth")
# ),
# "jet_JVFCorr": ei.jet.Select( # type: ignore
# lambda j: j.getAttribute[cpp_float]("JVFCorr")
# ),
# "jet_DFCommonJets_QGTagger_TracksC1": ei.jet.Select( # type: ignore
# lambda j: j.getAttribute[cpp_float]("DFCommonJets_QGTagger_TracksC1")
# ),
# "jet_PSFrac": ei.jet.Select( # type: ignore
# lambda j: j.getAttribute[cpp_float]("PSFrac")
# ),
# "jet_DFCommonJets_QGTagger_NTracks": ei.jet.Select( # type: ignore
# lambda j: j.getAttribute[cpp_int]("DFCommonJets_QGTagger_NTracks")
# ),
# "jet_DFCommonJets_fJvt": ei.jet.Select( # type: ignore
# lambda j: j.getAttribute[cpp_float]("DFCommonJets_fJvt")
# ),
# "jet_PartonTruthLabelID": ei.jet.Select( # type: ignore
# lambda j: j.getAttribute[cpp_int]("PartonTruthLabelID")
# ),
# "jet_HadronConeExclExtendedTruthLabelID": ei.jet.Select( # type: ignore
# lambda j: j.getAttribute[cpp_int]("HadronConeExclExtendedTruthLabelID")
# ),
# "jet_ConeTruthLabelID": ei.jet.Select( # type: ignore
# lambda j: j.getAttribute[cpp_int]("ConeTruthLabelID")
# ),
# "jet_HadronConeExclTruthLabelID": ei.jet.Select( # type: ignore
# lambda j: j.getAttribute[cpp_int]("HadronConeExclTruthLabelID")
# ),
# "jet_ActiveArea4vec_pt": ei.jet.Select( # type: ignore
# lambda j: j.getAttribute[cpp_float]("ActiveArea4vec_pt")
# ),
# }
# )

query = event_model.Select(
lambda ei: {
"event_number": ei.evt.eventNumber(), # type: ignore
Expand Down Expand Up @@ -267,63 +179,61 @@ def count_jets(events, n_jets: int):
"jet_GhostMuonSegmentCount": ei.jet.Select( # type: ignore
lambda j: j.getAttribute[cpp_int]("GhostMuonSegmentCount")
),
"ele_pt": ei.ele.Select(lambda e: e.pt() / 1000), # type: ignore
"ele_eta": ei.ele.Select(lambda e: e.eta()), # type: ignore
"ele_phi": ei.ele.Select(lambda e: e.phi()), # type: ignore
"ele_m": ei.ele.Select(lambda e: e.m() / 1000), # type: ignore
# TODO: Figure out what these represent in actual C++ calls
# "AnalysisElectronsAuxDyn.topoetcone20_CloseByCorr",
# "AnalysisElectronsAuxDyn.topoetcone20ptCorrection",
# "AnalysisElectronsAuxDyn.topoetcone20",
# "AnalysisElectronsAuxDyn.ambiguityLink",
# "AnalysisElectronsAuxDyn.DFCommonElectronsECIDSResult",
# "AnalysisElectronsAuxDyn.neflowisol20",
# "AnalysisElectronsAuxDyn.f1",
# "AnalysisElectronsAuxDyn.ptvarcone30_Nonprompt_All_MaxWeightTTVALooseCone_pt500",
"muons_pt": ei.muons.Select(lambda m: m.pt() / 1000), # type: ignore
"muons_eta": ei.muons.Select(lambda m: m.pt()), # type: ignore
"muons_phi": ei.muons.Select(lambda m: m.pt()), # type: ignore
"muons_m": ei.muons.Select(lambda m: m.m() / 1000), # type: ignore
# "AnalysisMuonsAuxDyn.muonSegmentLinks",
# "AnalysisMuonsAuxDyn.msOnlyExtrapolatedMuonSpectrometerTrackParticleLink",
# "AnalysisMuonsAuxDyn.extrapolatedMuonSpectrometerTrackParticleLink",
# "AnalysisMuonsAuxDyn.inDetTrackParticleLink",
# "AnalysisMuonsAuxDyn.muonSpectrometerTrackParticleLink",
# "AnalysisMuonsAuxDyn.momentumBalanceSignificance",
# "AnalysisMuonsAuxDyn.topoetcone20_CloseByCorr",
# "AnalysisMuonsAuxDyn.scatteringCurvatureSignificance",
# "AnalysisMuonsAuxDyn.scatteringNeighbourSignificance",
# "AnalysisMuonsAuxDyn.neflowisol20_CloseByCorr",
# "AnalysisMuonsAuxDyn.topoetcone20",
# "AnalysisMuonsAuxDyn.topoetcone30",
# "AnalysisMuonsAuxDyn.topoetcone40",
# "AnalysisMuonsAuxDyn.neflowisol20",
# "AnalysisMuonsAuxDyn.segmentDeltaEta",
# "AnalysisMuonsAuxDyn.DFCommonJetDr",
# "AnalysisMuonsAuxDyn.combinedTrackParticleLink",
# "AnalysisMuonsAuxDyn.InnerDetectorPt",
# "AnalysisMuonsAuxDyn.MuonSpectrometerPt",
# "AnalysisMuonsAuxDyn.clusterLink",
# "AnalysisMuonsAuxDyn.spectrometerFieldIntegral",
# "AnalysisMuonsAuxDyn.EnergyLoss",
# "AnalysisMuonsAuxDyn.ptvarcone30_Nonprompt_All_MaxWeightTTVA_pt500_CloseByCorr",
# "AnalysisMuonsAuxDyn.ptvarcone30_Nonprompt_All_MaxWeightTTVA_pt500",
# "AnalysisMuonsAuxDyn.ptcone40",
# "AnalysisMuonsAuxDyn.ptvarcone30_Nonprompt_All_MaxWeightTTVA_pt1000_CloseByCorr",
# "AnalysisMuonsAuxDyn.ptvarcone30_Nonprompt_All_MaxWeightTTVA_pt1000",
# "AnalysisMuonsAuxDyn.ptvarcone40",
# "AnalysisMuonsAuxDyn.ptcone20_Nonprompt_All_MaxWeightTTVA_pt500",
# "AnalysisMuonsAuxDyn.ptvarcone30",
# "AnalysisMuonsAuxDyn.ptcone30",
# "AnalysisMuonsAuxDyn.ptcone20_Nonprompt_All_MaxWeightTTVA_pt1000",
# "AnalysisMuonsAuxDyn.CaloLRLikelihood",
"pv_x": ei.pv.Select(lambda p: p.x() / 1000), # type: ignore
"pv_y": ei.pv.Select(lambda p: p.y() / 1000), # type: ignore
"pv_z": ei.pv.Select(lambda p: p.z() / 1000), # type: ignore
# "PrimaryVerticesAuxDyn.vertexType",
}
)

return query


BRANCH_LIST = [
"AnalysisElectronsAuxDyn.ambiguityLink",
"AnalysisElectronsAuxDyn.pt",
"AnalysisElectronsAuxDyn.eta",
"AnalysisElectronsAuxDyn.phi",
"AnalysisElectronsAuxDyn.m",
"AnalysisElectronsAuxDyn.topoetcone20_CloseByCorr",
"AnalysisElectronsAuxDyn.topoetcone20ptCorrection",
"AnalysisElectronsAuxDyn.topoetcone20",
"AnalysisElectronsAuxDyn.DFCommonElectronsECIDSResult",
"AnalysisElectronsAuxDyn.neflowisol20",
"AnalysisElectronsAuxDyn.f1",
"AnalysisElectronsAuxDyn.ptvarcone30_Nonprompt_All_MaxWeightTTVALooseCone_pt500",
"AnalysisMuonsAuxDyn.pt",
"AnalysisMuonsAuxDyn.eta",
"AnalysisMuonsAuxDyn.phi",
"AnalysisMuonsAuxDyn.muonSegmentLinks",
"AnalysisMuonsAuxDyn.msOnlyExtrapolatedMuonSpectrometerTrackParticleLink",
"AnalysisMuonsAuxDyn.extrapolatedMuonSpectrometerTrackParticleLink",
"AnalysisMuonsAuxDyn.inDetTrackParticleLink",
"AnalysisMuonsAuxDyn.muonSpectrometerTrackParticleLink",
"AnalysisMuonsAuxDyn.momentumBalanceSignificance",
"AnalysisMuonsAuxDyn.topoetcone20_CloseByCorr",
"AnalysisMuonsAuxDyn.scatteringCurvatureSignificance",
"AnalysisMuonsAuxDyn.scatteringNeighbourSignificance",
"AnalysisMuonsAuxDyn.neflowisol20_CloseByCorr",
"AnalysisMuonsAuxDyn.topoetcone20",
"AnalysisMuonsAuxDyn.topoetcone30",
"AnalysisMuonsAuxDyn.topoetcone40",
"AnalysisMuonsAuxDyn.neflowisol20",
"AnalysisMuonsAuxDyn.segmentDeltaEta",
"AnalysisMuonsAuxDyn.DFCommonJetDr",
"AnalysisMuonsAuxDyn.combinedTrackParticleLink",
"AnalysisMuonsAuxDyn.InnerDetectorPt",
"AnalysisMuonsAuxDyn.MuonSpectrometerPt",
"AnalysisMuonsAuxDyn.clusterLink",
"AnalysisMuonsAuxDyn.spectrometerFieldIntegral",
"AnalysisMuonsAuxDyn.EnergyLoss",
"AnalysisMuonsAuxDyn.ptvarcone30_Nonprompt_All_MaxWeightTTVA_pt500_CloseByCorr",
"AnalysisMuonsAuxDyn.ptvarcone30_Nonprompt_All_MaxWeightTTVA_pt500",
"AnalysisMuonsAuxDyn.ptcone40",
"AnalysisMuonsAuxDyn.ptvarcone30_Nonprompt_All_MaxWeightTTVA_pt1000_CloseByCorr",
"AnalysisMuonsAuxDyn.ptvarcone30_Nonprompt_All_MaxWeightTTVA_pt1000",
"AnalysisMuonsAuxDyn.ptvarcone40",
"AnalysisMuonsAuxDyn.ptcone20_Nonprompt_All_MaxWeightTTVA_pt500",
"AnalysisMuonsAuxDyn.ptvarcone30",
"AnalysisMuonsAuxDyn.ptcone30",
"AnalysisMuonsAuxDyn.ptcone20_Nonprompt_All_MaxWeightTTVA_pt1000",
"AnalysisMuonsAuxDyn.CaloLRLikelihood",
"PrimaryVerticesAuxDyn.z",
"PrimaryVerticesAuxDyn.x",
"PrimaryVerticesAuxDyn.y",
"PrimaryVerticesAuxDyn.vertexType",
]

0 comments on commit 41d775b

Please sign in to comment.