diff --git a/cmd/hepmc2root/main.go b/cmd/hepmc2root/main.go new file mode 100644 index 000000000..d5c0cd21d --- /dev/null +++ b/cmd/hepmc2root/main.go @@ -0,0 +1,264 @@ +// Copyright ©2022 The go-hep Authors. All rights reserved. +// Use of this source code is governed by a BSD-style +// license that can be found in the LICENSE file. + +// hepmc2root converts HepMC files into ROOT files. +// +// Example: +// +// $> hepmc2root -o hepmc.root hepmc.ascii +// $> hepmc2root -o hepmc.root -t mytree hepmc.ascii +package main // import "go-hep.org/x/hep/cmd/hepmc2root" + +import ( + "flag" + "fmt" + "io" + "log" + "os" + "sort" + + "go-hep.org/x/hep/groot" + "go-hep.org/x/hep/groot/rtree" + "go-hep.org/x/hep/hepmc" +) + +func main() { + log.SetPrefix("hepmc2root: ") + log.SetFlags(0) + + oname := flag.String("o", "out.root", "path to output ROOT file name") + tname := flag.String("t", "tree", "name of the output tree") + + flag.Parse() + + if flag.NArg() != 1 { + flag.Usage() + log.Fatalf("missing input HepMC filename argument") + } + fname := flag.Arg(0) + + err := process(*oname, *tname, fname) + if err != nil { + log.Fatalf("%+v", err) + } +} + +func process(oname, tname, fname string) error { + f, err := os.Open(fname) + if err != nil { + return fmt.Errorf("could not open HepMC file %q: %w", fname, err) + } + defer f.Close() + + o, err := groot.Create(oname) + if err != nil { + return fmt.Errorf("could not create output ROOT file %q: %w", oname, err) + } + defer o.Close() + + var ( + revt hepmc.Event + wevt Event + wvars = rtree.WriteVarsFromStruct(&wevt) + ) + + tree, err := rtree.NewWriter(o, tname, wvars, rtree.WithTitle(tname)) + if err != nil { + return fmt.Errorf("could not create output ROOT tree %q: %w", tname, err) + } + + var ( + ievt int + dec = hepmc.NewDecoder(f) + ) + for { + err := dec.Decode(&revt) + if err == io.EOF { + break + } + if err != nil { + return fmt.Errorf("could not decode event %d from %q: %w", ievt, fname, err) + } + + err = wevt.read(&revt) + if err != nil { + return fmt.Errorf("could not convert event %d to ROOT: %w", ievt, err) + } + + _, err = tree.Write() + if err != nil { + return fmt.Errorf("could not write event %d to ROOT: %w", ievt, err) + } + + wevt.reset() + + err = hepmc.Delete(&revt) + if err != nil { + return fmt.Errorf("could not gc event %d from %q: %w", ievt, fname, err) + } + ievt++ + } + + err = tree.Close() + if err != nil { + return fmt.Errorf("could not close ROOT tree writer: %w", err) + } + + err = o.Close() + if err != nil { + return fmt.Errorf("could not close output ROOT file %q: %w", oname, err) + } + + return nil +} + +type Event struct { + SignalProcessID int32 `groot:"Event_processID"` // id of the signal process + Event_number int32 `groot:"Event_number"` // event number + Event_mpi int32 `groot:"Event_numberMP"` // number of multi particle interactions + Event_scale float64 `groot:"Event_scale"` // energy scale, + Event_alphaQCD float64 `groot:"Event_alphaQCD"` // QCD coupling, see hep-ph/0109068 + Event_alphaQED float64 `groot:"Event_alphaQED"` // QED coupling, see hep-ph/0109068 + Event_barcodeSPV int32 `groot:"Event_barcodeSPV"` + Event_nvtx int32 `groot:"Event_numberV"` + Event_barcodeBP1 int32 `groot:"Event_barcodeBP1"` + Event_barcodeBP2 int32 `groot:"Event_barcodeBP2"` + Event_npart int32 `groot:"Event_numberP"` + + XsectValue float64 `groot:"Xsection_value"` + XsectError float64 `groot:"Xsection_error"` + + PDF_Parton1 int32 `groot:"PDF_parton1"` + PDF_Parton2 int32 `groot:"PDF_parton2"` + PDF_X1 float64 `groot:"PDF_x1"` + PDF_X2 float64 `groot:"PDF_x2"` + PDF_Q2 float64 `groot:"PDF_Q2"` + PDF_X1f float64 `groot:"PDF_x1f"` + PDF_X2f float64 `groot:"PDF_x2f"` + PDF_ID1 int32 `groot:"PDF_id1"` + PDF_ID2 int32 `groot:"PDF_id2"` + + Particle_x []float64 `groot:"Particle_x"` + Particle_y []float64 `groot:"Particle_y"` + Particle_z []float64 `groot:"Particle_z"` + Particle_ctau []float64 `groot:"Particle_ctau"` + + Particle_barcode []int32 `groot:"Particle_barcode"` + Particle_pid []int64 `groot:"Particle_pid"` + Particle_px []float64 `groot:"Particle_px"` + Particle_py []float64 `groot:"Particle_py"` + Particle_pz []float64 `groot:"Particle_pz"` + Particle_ene []float64 `groot:"Particle_energy"` + Particle_mass []float64 `groot:"Particle_mass"` + Particle_status []int32 `groot:"Particle_status"` + Particle_d1 []int32 `groot:"Particle_d1"` + Particle_d2 []int32 `groot:"Particle_d2"` +} + +func (evt *Event) read(h *hepmc.Event) error { + evt.SignalProcessID = int32(h.SignalProcessID) + evt.Event_number = int32(h.EventNumber) + evt.Event_mpi = int32(h.Mpi) + evt.Event_scale = h.Scale + evt.Event_alphaQCD = h.AlphaQCD + evt.Event_alphaQED = h.AlphaQED + switch { + case h.SignalVertex != nil: + evt.Event_barcodeSPV = int32(h.SignalVertex.Barcode) + default: + evt.Event_barcodeSPV = 0 + } + evt.Event_nvtx = int32(len(h.Vertices)) + switch { + case h.Beams[0] != nil: + evt.Event_barcodeBP1 = int32(h.Beams[0].Barcode) + default: + evt.Event_barcodeBP1 = 0 + } + switch { + case h.Beams[1] != nil: + evt.Event_barcodeBP2 = int32(h.Beams[1].Barcode) + default: + evt.Event_barcodeBP2 = 0 + } + + evt.Event_npart = int32(len(h.Particles)) + + switch xsect := h.CrossSection; xsect { + case nil: + evt.XsectValue = 0 + evt.XsectError = 0 + default: + evt.XsectValue = h.CrossSection.Value + evt.XsectError = h.CrossSection.Error + } + + evt.PDF_Parton1 = int32(h.PdfInfo.ID1) + evt.PDF_Parton2 = int32(h.PdfInfo.ID2) + evt.PDF_X1 = h.PdfInfo.X1 + evt.PDF_X2 = h.PdfInfo.X2 + evt.PDF_Q2 = h.PdfInfo.ScalePDF + evt.PDF_X1f = h.PdfInfo.Pdf1 + evt.PDF_X2f = h.PdfInfo.Pdf2 + evt.PDF_ID1 = int32(h.PdfInfo.LHAPdf1) + evt.PDF_ID2 = int32(h.PdfInfo.LHAPdf2) + + barcodes := make([]int, 0, len(h.Particles)) + for bc := range h.Particles { + barcodes = append(barcodes, bc) + } + sort.Ints(barcodes) + for _, bc := range barcodes { + p := h.Particles[bc] + switch vtx := p.ProdVertex; vtx { + case nil: + evt.Particle_x = append(evt.Particle_x, 0) + evt.Particle_y = append(evt.Particle_y, 0) + evt.Particle_z = append(evt.Particle_z, 0) + evt.Particle_ctau = append(evt.Particle_ctau, 0) + evt.Particle_d1 = append(evt.Particle_d1, 0) + default: + evt.Particle_x = append(evt.Particle_x, vtx.Position.X()) + evt.Particle_y = append(evt.Particle_y, vtx.Position.Y()) + evt.Particle_z = append(evt.Particle_z, vtx.Position.Z()) + evt.Particle_ctau = append(evt.Particle_ctau, vtx.Position.T()) + evt.Particle_d1 = append(evt.Particle_d1, int32(vtx.Barcode)) + } + + evt.Particle_barcode = append(evt.Particle_barcode, int32(p.Barcode)) + evt.Particle_pid = append(evt.Particle_pid, p.PdgID) + evt.Particle_px = append(evt.Particle_px, p.Momentum.Px()) + evt.Particle_py = append(evt.Particle_py, p.Momentum.Py()) + evt.Particle_pz = append(evt.Particle_pz, p.Momentum.Pz()) + evt.Particle_ene = append(evt.Particle_ene, p.Momentum.E()) + evt.Particle_mass = append(evt.Particle_mass, p.Momentum.M()) + evt.Particle_status = append(evt.Particle_status, int32(p.Status)) + switch vtx := p.EndVertex; vtx { + case nil: + evt.Particle_d2 = append(evt.Particle_d2, 0) + default: + evt.Particle_d2 = append(evt.Particle_d2, int32(vtx.Barcode)) + } + } + + return nil +} + +func (evt *Event) reset() { + evt.Particle_x = evt.Particle_x[:0] + evt.Particle_y = evt.Particle_y[:0] + evt.Particle_z = evt.Particle_z[:0] + evt.Particle_ctau = evt.Particle_ctau[:0] + + evt.Particle_barcode = evt.Particle_barcode[:0] + evt.Particle_pid = evt.Particle_pid[:0] + evt.Particle_px = evt.Particle_px[:0] + evt.Particle_py = evt.Particle_py[:0] + evt.Particle_pz = evt.Particle_pz[:0] + evt.Particle_ene = evt.Particle_ene[:0] + evt.Particle_mass = evt.Particle_mass[:0] + evt.Particle_status = evt.Particle_status[:0] + evt.Particle_d1 = evt.Particle_d1[:0] + evt.Particle_d2 = evt.Particle_d2[:0] +}