forked from go-hep/hep
-
Notifications
You must be signed in to change notification settings - Fork 0
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Fixes go-hep#959. Signed-off-by: Sebastien Binet <binet@cern.ch>
- Loading branch information
Showing
1 changed file
with
264 additions
and
0 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -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] | ||
| } |