diff --git a/cmd/hepmc2root/main.go b/cmd/hepmc2root/main.go index fdc9729c5..c94e08dee 100644 --- a/cmd/hepmc2root/main.go +++ b/cmd/hepmc2root/main.go @@ -22,15 +22,13 @@ 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" - "go-hep.org/x/hep/sliceop" + "go-hep.org/x/hep/hepmc/rootcnv" ) func main() { @@ -82,47 +80,14 @@ func process(oname, tname, fname string) error { } defer o.Close() - var ( - revt hepmc.Event - wevt Event - wvars = rtree.WriteVarsFromStruct(&wevt) - ) - - tree, err := rtree.NewWriter(o, tname, wvars, rtree.WithTitle(tname)) + tree, err := rootcnv.NewFlatTreeWriter(o, tname, 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 = hepmc.Copy(tree, hepmc.NewASCIIReader(f)) + if err != nil { + return fmt.Errorf("could not write HepMC events to ROOT: %w", err) } err = tree.Close() @@ -137,294 +102,3 @@ func process(oname, tname, fname string) error { return nil } - -type Event struct { - SignalProcessID int32 `groot:"Event_processID"` // id of the signal process - Event_number int32 `groot:"Event_nbr"` // event number - Event_mpi int32 `groot:"Event_mpi"` // 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_barcodeBP1 int32 `groot:"Event_barcodeBP1"` - Event_barcodeBP2 int32 `groot:"Event_barcodeBP2"` - Event_nvtx int32 `groot:"Event_nvtx"` - Event_npart int32 `groot:"Event_npart"` - Event_inbcs []int32 `groot:"Event_inbcs"` // Event barcodes of (p-in) for each vertex - Event_outbcs []int32 `groot:"Event_outbcs"` // Event barcodes of (p-out) for each vertex - - WeightsSlice []float64 `groot:"Weights_slice"` - WeightsMapKeys []string `groot:"Weights_keys"` - WeightsMapNames []int32 `groot:"Weights_names"` - RandomStates []int64 `groot:"Random_states"` - - XsectValue float64 `groot:"Xsection_value"` - XsectError float64 `groot:"Xsection_error"` - - HI_ncollHard int32 `groot:"HI_ncoll_hard"` - HI_npartProj int32 `groot:"HI_npart_proj"` - HI_npartTarg int32 `groot:"HI_npart_targ"` - HI_ncoll int32 `groot:"HI_ncoll"` - HI_nnwColl int32 `groot:"HI_nnw_coll"` - HI_nwNColl int32 `groot:"HI_nwn_coll"` - HI_nwNwColl int32 `groot:"HI_nwnw_coll"` - HI_spectatorNeutrons int32 `groot:"HI_spect_neutrons"` - HI_spectatorProtons int32 `groot:"HI_spect_protons"` - HI_impactParameter float32 `groot:"HI_impact_param"` - HI_eventPlaneAngle float32 `groot:"HI_evt_plane_angle"` - HI_eccentricity float32 `groot:"HI_eccentricity"` - HI_sigmaInelNN float32 `groot:"HI_sigma_inel_nn"` - - 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"` - - MomentumUnit int8 `groot:"Momentum_unit"` - LengthUnit int8 `groot:"Length_unit"` - - Vertex_x []float64 `groot:"Vertex_x"` - Vertex_y []float64 `groot:"Vertex_y"` - Vertex_z []float64 `groot:"Vertex_z"` - Vertex_t []float64 `groot:"Vertex_t"` - Vertex_id []int32 `groot:"Vertex_id"` - Vertex_bc []int32 `groot:"Vertex_bc"` - Vertex_nin []int32 `groot:"Vertex_nin"` - Vertex_nout []int32 `groot:"Vertex_nout"` - - Particle_bc []int32 `groot:"Particle_bc"` - 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_ene"` - Particle_mass []float64 `groot:"Particle_mass"` - Particle_nflow []int32 `groot:"Particle_nflow"` - Particle_flow [][]int32 `groot:"Particle_flow"` - Particle_theta []float64 `groot:"Particle_theta"` - Particle_phi []float64 `groot:"Particle_phi"` - Particle_status []int32 `groot:"Particle_status"` - Particle_pvtx []int32 `groot:"Particle_pvtx"` - Particle_evtx []int32 `groot:"Particle_evtx"` - - barcodes []int // work buffer -} - -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)) - evt.Event_npart = int32(len(h.Particles)) - 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.WeightsSlice = sliceop.Resize(evt.WeightsSlice, len(h.Weights.Slice)) - copy(evt.WeightsSlice, h.Weights.Slice) - evt.WeightsMapKeys = sliceop.Resize(evt.WeightsMapKeys, len(h.Weights.Map))[:0] - evt.WeightsMapNames = sliceop.Resize(evt.WeightsMapNames, len(h.Weights.Map))[:0] - for k, v := range h.Weights.Map { - evt.WeightsMapKeys = append(evt.WeightsMapKeys, k) - evt.WeightsMapNames = append(evt.WeightsMapNames, int32(v)) - } - evt.RandomStates = sliceop.Resize(evt.RandomStates, len(h.RandomStates)) - copy(evt.RandomStates, h.RandomStates) - - switch xsect := h.CrossSection; xsect { - case nil: - evt.XsectValue = 0 - evt.XsectError = 0 - default: - evt.XsectValue = h.CrossSection.Value - evt.XsectError = h.CrossSection.Error - } - - switch hi := h.HeavyIon; hi { - case nil: - evt.HI_ncollHard = 0 - evt.HI_npartProj = 0 - evt.HI_npartTarg = 0 - evt.HI_ncoll = 0 - evt.HI_nnwColl = 0 - evt.HI_nwNColl = 0 - evt.HI_nwNwColl = 0 - evt.HI_spectatorNeutrons = 0 - evt.HI_spectatorProtons = 0 - evt.HI_impactParameter = 0 - evt.HI_eventPlaneAngle = 0 - evt.HI_eccentricity = 0 - evt.HI_sigmaInelNN = 0 - default: - evt.HI_ncollHard = int32(hi.NCollHard) - evt.HI_npartProj = int32(hi.NPartProj) - evt.HI_npartTarg = int32(hi.NPartTarg) - evt.HI_ncoll = int32(hi.NColl) - evt.HI_nnwColl = int32(hi.NNwColl) - evt.HI_nwNColl = int32(hi.NwNColl) - evt.HI_nwNwColl = int32(hi.NwNwColl) - evt.HI_spectatorNeutrons = int32(hi.SpectatorNeutrons) - evt.HI_spectatorProtons = int32(hi.SpectatorProtons) - evt.HI_impactParameter = hi.ImpactParameter - evt.HI_eventPlaneAngle = hi.EventPlaneAngle - evt.HI_eccentricity = hi.Eccentricity - evt.HI_sigmaInelNN = hi.SigmaInelNN - } - - 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) - - evt.barcodes = sliceop.Resize(evt.barcodes, len(h.Vertices))[:0] - for bc := range h.Vertices { - evt.barcodes = append(evt.barcodes, bc) - } - sort.Sort(sort.Reverse(sort.IntSlice(evt.barcodes))) - - n := len(h.Vertices) - evt.Vertex_x = sliceop.Resize(evt.Vertex_x, n)[:0] - evt.Vertex_y = sliceop.Resize(evt.Vertex_y, n)[:0] - evt.Vertex_z = sliceop.Resize(evt.Vertex_z, n)[:0] - evt.Vertex_t = sliceop.Resize(evt.Vertex_t, n)[:0] - evt.Vertex_id = sliceop.Resize(evt.Vertex_id, n)[:0] - evt.Vertex_bc = sliceop.Resize(evt.Vertex_bc, n)[:0] - evt.Vertex_nin = sliceop.Resize(evt.Vertex_nin, n)[:0] - evt.Vertex_nout = sliceop.Resize(evt.Vertex_nout, n)[:0] - - for _, bc := range evt.barcodes { - vtx := h.Vertices[bc] - evt.Vertex_x = append(evt.Vertex_x, vtx.Position.X()) - evt.Vertex_y = append(evt.Vertex_y, vtx.Position.Y()) - evt.Vertex_z = append(evt.Vertex_z, vtx.Position.Z()) - evt.Vertex_t = append(evt.Vertex_t, vtx.Position.T()) - evt.Vertex_id = append(evt.Vertex_id, int32(vtx.ID)) - evt.Vertex_bc = append(evt.Vertex_bc, int32(vtx.Barcode)) - evt.Vertex_nin = append(evt.Vertex_nin, int32(len(vtx.ParticlesIn))) - evt.Vertex_nout = append(evt.Vertex_nout, int32(len(vtx.ParticlesOut))) - for _, p := range vtx.ParticlesIn { - evt.Event_inbcs = append(evt.Event_inbcs, int32(p.Barcode)) - } - for _, p := range vtx.ParticlesOut { - evt.Event_outbcs = append(evt.Event_outbcs, int32(p.Barcode)) - } - } - - evt.barcodes = sliceop.Resize(evt.barcodes, len(h.Particles))[:0] - for bc := range h.Particles { - evt.barcodes = append(evt.barcodes, bc) - } - sort.Ints(evt.barcodes) - - n = len(h.Particles) - evt.Particle_bc = sliceop.Resize(evt.Particle_bc, n)[:0] - evt.Particle_pid = sliceop.Resize(evt.Particle_pid, n)[:0] - evt.Particle_px = sliceop.Resize(evt.Particle_px, n)[:0] - evt.Particle_py = sliceop.Resize(evt.Particle_py, n)[:0] - evt.Particle_pz = sliceop.Resize(evt.Particle_pz, n)[:0] - evt.Particle_ene = sliceop.Resize(evt.Particle_ene, n)[:0] - evt.Particle_mass = sliceop.Resize(evt.Particle_mass, n)[:0] - evt.Particle_nflow = sliceop.Resize(evt.Particle_nflow, n)[:0] - evt.Particle_flow = sliceop.Resize(evt.Particle_flow, n)[:0] - evt.Particle_theta = sliceop.Resize(evt.Particle_theta, n)[:0] - evt.Particle_phi = sliceop.Resize(evt.Particle_phi, n)[:0] - evt.Particle_status = sliceop.Resize(evt.Particle_status, n)[:0] - evt.Particle_pvtx = sliceop.Resize(evt.Particle_pvtx, n)[:0] - evt.Particle_evtx = sliceop.Resize(evt.Particle_evtx, n)[:0] - - for _, bc := range evt.barcodes { - p := h.Particles[bc] - switch vtx := p.ProdVertex; vtx { - case nil: - evt.Particle_pvtx = append(evt.Particle_pvtx, 0) - default: - evt.Particle_pvtx = append(evt.Particle_pvtx, int32(vtx.Barcode)) - } - - evt.Particle_bc = append(evt.Particle_bc, 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.GeneratedMass) - evt.Particle_nflow = append(evt.Particle_nflow, int32(len(p.Flow.Icode))) - flow := make([]int32, 0, 2*len(p.Flow.Icode)) - for k, v := range p.Flow.Icode { - flow = append(flow, int32(k), int32(v)) - } - evt.Particle_flow = append(evt.Particle_flow, flow) - evt.Particle_theta = append(evt.Particle_theta, p.Polarization.Theta) - evt.Particle_phi = append(evt.Particle_phi, p.Polarization.Phi) - evt.Particle_status = append(evt.Particle_status, int32(p.Status)) - switch vtx := p.EndVertex; vtx { - case nil: - evt.Particle_evtx = append(evt.Particle_evtx, 0) - default: - evt.Particle_evtx = append(evt.Particle_evtx, int32(vtx.Barcode)) - } - } - - return nil -} - -func (evt *Event) reset() { - evt.Event_inbcs = evt.Event_inbcs[:0] - evt.Event_outbcs = evt.Event_outbcs[:0] - evt.WeightsSlice = evt.WeightsSlice[:0] - evt.WeightsMapKeys = evt.WeightsMapKeys[:0] - evt.WeightsMapNames = evt.WeightsMapNames[:0] - evt.RandomStates = evt.RandomStates[:0] - - evt.Vertex_x = evt.Vertex_x[:0] - evt.Vertex_y = evt.Vertex_y[:0] - evt.Vertex_z = evt.Vertex_z[:0] - evt.Vertex_t = evt.Vertex_t[:0] - evt.Vertex_id = evt.Vertex_id[:0] - evt.Vertex_bc = evt.Vertex_bc[:0] - evt.Vertex_nin = evt.Vertex_nin[:0] - evt.Vertex_nout = evt.Vertex_nout[:0] - - evt.Particle_bc = evt.Particle_bc[: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_nflow = evt.Particle_nflow[:0] - evt.Particle_flow = evt.Particle_flow[:0] - evt.Particle_theta = evt.Particle_theta[:0] - evt.Particle_phi = evt.Particle_phi[:0] - evt.Particle_status = evt.Particle_status[:0] - evt.Particle_pvtx = evt.Particle_pvtx[:0] - evt.Particle_evtx = evt.Particle_evtx[:0] -} diff --git a/cmd/hepmc2root/testdata/small.hepmc.txt b/cmd/hepmc2root/testdata/small.hepmc.txt index 4526e6e52..f445a1a1d 100644 --- a/cmd/hepmc2root/testdata/small.hepmc.txt +++ b/cmd/hepmc2root/testdata/small.hepmc.txt @@ -40,13 +40,16 @@ key[000]: tree;1 "tree" (TTree) [000][PDF_x2f]: 0 [000][PDF_id1]: 0 [000][PDF_id2]: 0 -[000][Momentum_unit]: 0 +[000][Momentum_unit]: 1 [000][Length_unit]: 0 [000][Vertex_x]: [0 0 0 0.12] [000][Vertex_y]: [0 0 0 -0.3] [000][Vertex_z]: [0 0 0 0.05] [000][Vertex_t]: [0 0 0 0.004] [000][Vertex_id]: [0 0 0 0] +[000][Vertex_wsli]: [[] [] [] []] +[000][Vertex_wkey]: [[] [] [] []] +[000][Vertex_wval]: [[] [] [] []] [000][Vertex_bc]: [-1 -2 -3 -4] [000][Vertex_nin]: [1 1 2 1] [000][Vertex_nout]: [1 1 2 2] @@ -57,7 +60,6 @@ key[000]: tree;1 "tree" (TTree) [000][Particle_pz]: [7000 -7000 32.191 -54.629 -1.833 -20.605 6.082 -26.687] [000][Particle_ene]: [7000 7000 32.238 57.92 4.233 85.925 29.552 56.373] [000][Particle_mass]: [0 0 0 0 0 0 0 0] -[000][Particle_nflow]: [0 0 0 0 0 0 0 0] [000][Particle_flow]: [[] [] [] [] [] [] [] []] [000][Particle_theta]: [0 0 0 0 0 0 0 0] [000][Particle_phi]: [0 0 0 0 0 0 0 0] diff --git a/hepmc/io.go b/hepmc/io.go index adb901d36..bdf78c213 100644 --- a/hepmc/io.go +++ b/hepmc/io.go @@ -4,12 +4,93 @@ package hepmc -// // An Encoder writes and encodes hepmc.Event objects into an output stream. -// type Encoder interface { -// Encode(evt *Event) error -// } - -// // A Decoder reads and decodes hepmc.Event objects from an input stream. -// type Decoder interface { -// Decode(evt *Event) error -// } +import ( + "errors" + "io" +) + +// Reader is the interface that wraps the Read method. +// +// Read reads an HepMC event from the underlying storage into +// the provided event. +// Read returns io.EOF when no more events are available. +type Reader interface { + Read(evt *Event) error +} + +// Writer is the interface that wraps the Write method. +// +// Write writes the provided HepMC event to the underlying storage. +type Writer interface { + Write(evt Event) error +} + +// Copy copies events from r to w until either EOF is reached or an +// error occurs. It returns the number of events copied and the first error +// encountered while copying, if any. +// +// A successful Copy returns err == nil, not err == io.EOF. +func Copy(w Writer, r Reader) (n int64, err error) { + var evt Event + +loop: + for { + err = r.Read(&evt) + if err != nil { + if errors.Is(err, io.EOF) { + err = nil + } + break loop + } + + err = w.Write(evt) + if err != nil { + break loop + } + + n++ + } + + return n, err +} + +// ASCIIReader reads ASCII HepMC-v2 data. +type ASCIIReader struct { + dec *Decoder +} + +// NewASCIIReader creates a HepMC reader that reads data +// from r, in the HepMC-v2 ASCII format. +func NewASCIIReader(r io.Reader) *ASCIIReader { + return &ASCIIReader{dec: NewDecoder(r)} +} + +func (r *ASCIIReader) Read(evt *Event) error { + return r.dec.Decode(evt) +} + +// ASCIIWriter writes ASCII HepMC-v2 data. +type ASCIIWriter struct { + enc *Encoder +} + +// NewASCIIWriter creates a HepMC writer that writes data +// to w, in the HepMC-v2 ASCII format. +func NewASCIIWriter(w io.Writer) *ASCIIWriter { + return &ASCIIWriter{enc: NewEncoder(w)} +} + +func (w *ASCIIWriter) Write(evt Event) error { + return w.enc.Encode(&evt) +} + +func (w *ASCIIWriter) Close() error { + return w.enc.Close() +} + +var ( + _ Reader = (*ASCIIReader)(nil) + + _ Writer = (*ASCIIWriter)(nil) + _ io.Closer = (*ASCIIWriter)(nil) +) diff --git a/hepmc/rootcnv/reader.go b/hepmc/rootcnv/reader.go new file mode 100644 index 000000000..f6e8e3f71 --- /dev/null +++ b/hepmc/rootcnv/reader.go @@ -0,0 +1,79 @@ +// 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. + +package rootcnv + +import ( + "errors" + "fmt" + "io" + + "go-hep.org/x/hep/groot/rtree" + "go-hep.org/x/hep/hepmc" +) + +type rstream struct { + evt hepmc.Event + err error +} + +type FlatTreeReader struct { + r *rtree.Reader + + evt event + rvars []rtree.ReadVar + + evts chan rstream +} + +func NewFlatTreeReader(t rtree.Tree, opts ...rtree.ReadOption) (*FlatTreeReader, error) { + r := FlatTreeReader{ + evts: make(chan rstream), + } + + r.rvars = rtree.ReadVarsFromStruct(&r.evt) + + rr, err := rtree.NewReader(t, r.rvars, opts...) + if err != nil { + return nil, fmt.Errorf("hepmc: could not create ROOT Tree reader: %w", err) + } + r.r = rr + + go func() { + defer close(r.evts) + err := rr.Read(func(ctx rtree.RCtx) error { + defer r.evt.reset() + var o rstream + err := r.evt.write(&o.evt) + if err != nil { + return err + } + r.evts <- o + return nil + }) + if err != nil { + r.evts <- rstream{err: err} + return + } + r.evts <- rstream{err: io.EOF} + }() + + return &r, nil +} + +func (r *FlatTreeReader) Close() error { + return r.r.Close() +} + +func (r *FlatTreeReader) Read(evt *hepmc.Event) error { + o := <-r.evts + if o.err != nil { + if errors.Is(o.err, io.EOF) { + return io.EOF + } + return fmt.Errorf("hepmc: could not read event from ROOT: %w", o.err) + } + *evt = o.evt + return nil +} diff --git a/hepmc/rootcnv/rootcnv.go b/hepmc/rootcnv/rootcnv.go new file mode 100644 index 000000000..ad7b95f12 --- /dev/null +++ b/hepmc/rootcnv/rootcnv.go @@ -0,0 +1,509 @@ +// 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. + +// Package rootcnv provides tools to convert between HepMC2 data types +// and ROOT data structures. +package rootcnv // import "go-hep.org/x/hep/hepmc/rootcnv" + +import ( + "sort" + + "go-hep.org/x/hep/fmom" + "go-hep.org/x/hep/hepmc" + "go-hep.org/x/hep/sliceop" +) + +type event struct { + SignalProcessID int32 `groot:"Event_processID"` // id of the signal process + Event_number int32 `groot:"Event_nbr"` // event number + Event_mpi int32 `groot:"Event_mpi"` // 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_barcodeBP1 int32 `groot:"Event_barcodeBP1"` + Event_barcodeBP2 int32 `groot:"Event_barcodeBP2"` + Event_nvtx int32 `groot:"Event_nvtx"` + Event_npart int32 `groot:"Event_npart"` + Event_inbcs []int32 `groot:"Event_inbcs"` // Event barcodes of (p-in) for each vertex + Event_outbcs []int32 `groot:"Event_outbcs"` // Event barcodes of (p-out) for each vertex + + WeightsSlice []float64 `groot:"Weights_slice"` + WeightsMapKeys []string `groot:"Weights_keys"` + WeightsMapNames []int32 `groot:"Weights_names"` + RandomStates []int64 `groot:"Random_states"` + + XsectValue float64 `groot:"Xsection_value"` + XsectError float64 `groot:"Xsection_error"` + + HI_ncollHard int32 `groot:"HI_ncoll_hard"` + HI_npartProj int32 `groot:"HI_npart_proj"` + HI_npartTarg int32 `groot:"HI_npart_targ"` + HI_ncoll int32 `groot:"HI_ncoll"` + HI_nnwColl int32 `groot:"HI_nnw_coll"` + HI_nwNColl int32 `groot:"HI_nwn_coll"` + HI_nwNwColl int32 `groot:"HI_nwnw_coll"` + HI_spectatorNeutrons int32 `groot:"HI_spect_neutrons"` + HI_spectatorProtons int32 `groot:"HI_spect_protons"` + HI_impactParameter float32 `groot:"HI_impact_param"` + HI_eventPlaneAngle float32 `groot:"HI_evt_plane_angle"` + HI_eccentricity float32 `groot:"HI_eccentricity"` + HI_sigmaInelNN float32 `groot:"HI_sigma_inel_nn"` + + 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"` + + MomentumUnit int8 `groot:"Momentum_unit"` + LengthUnit int8 `groot:"Length_unit"` + + Vertex_x []float64 `groot:"Vertex_x"` + Vertex_y []float64 `groot:"Vertex_y"` + Vertex_z []float64 `groot:"Vertex_z"` + Vertex_t []float64 `groot:"Vertex_t"` + Vertex_id []int32 `groot:"Vertex_id"` + Vertex_wsli [][]float64 `groot:"Vertex_wsli"` + Vertex_wkey [][]string `groot:"Vertex_wkey"` + Vertex_wval [][]int32 `groot:"Vertex_wval"` + Vertex_bc []int32 `groot:"Vertex_bc"` + Vertex_nin []int32 `groot:"Vertex_nin"` + Vertex_nout []int32 `groot:"Vertex_nout"` + + Particle_bc []int32 `groot:"Particle_bc"` + 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_ene"` + Particle_mass []float64 `groot:"Particle_mass"` + // Particle_nflow []int32 `groot:"Particle_nflow"` + Particle_flow [][]int32 `groot:"Particle_flow"` + Particle_theta []float64 `groot:"Particle_theta"` + Particle_phi []float64 `groot:"Particle_phi"` + Particle_status []int32 `groot:"Particle_status"` + Particle_pvtx []int32 `groot:"Particle_pvtx"` + Particle_evtx []int32 `groot:"Particle_evtx"` + + barcodes []int // work buffer +} + +func (evt *event) write(h *hepmc.Event) error { + h.SignalProcessID = int(evt.SignalProcessID) + h.EventNumber = int(evt.Event_number) + h.Mpi = int(evt.Event_mpi) + h.Scale = evt.Event_scale + h.AlphaQCD = evt.Event_alphaQCD + h.AlphaQED = evt.Event_alphaQED + h.Vertices = make(map[int]*hepmc.Vertex, evt.Event_nvtx) + h.Particles = make(map[int]*hepmc.Particle, evt.Event_npart) + + h.Weights.Slice = sliceop.Resize(h.Weights.Slice, len(evt.WeightsSlice)) + copy(h.Weights.Slice, evt.WeightsSlice) + h.Weights.Map = make(map[string]int, len(evt.WeightsMapKeys)) + for i, k := range evt.WeightsMapKeys { + v := evt.WeightsMapNames[i] + h.Weights.Map[k] = int(v) + } + h.RandomStates = sliceop.Resize(h.RandomStates, len(evt.RandomStates)) + copy(h.RandomStates, evt.RandomStates) + + switch { + case !evt.hasXSect(): + h.CrossSection = nil + default: + h.CrossSection = &hepmc.CrossSection{ + Value: evt.XsectValue, + Error: evt.XsectError, + } + } + + switch { + case !evt.hasHI(): + h.HeavyIon = nil + default: + h.HeavyIon = &hepmc.HeavyIon{ + NCollHard: int(evt.HI_ncollHard), + NPartProj: int(evt.HI_npartProj), + NPartTarg: int(evt.HI_npartTarg), + NColl: int(evt.HI_ncoll), + NNwColl: int(evt.HI_nnwColl), + NwNColl: int(evt.HI_nwNColl), + NwNwColl: int(evt.HI_nwNwColl), + SpectatorNeutrons: int(evt.HI_spectatorNeutrons), + SpectatorProtons: int(evt.HI_spectatorProtons), + ImpactParameter: evt.HI_impactParameter, + EventPlaneAngle: evt.HI_eventPlaneAngle, + Eccentricity: evt.HI_eccentricity, + SigmaInelNN: evt.HI_sigmaInelNN, + } + } + + if h.PdfInfo == nil { + h.PdfInfo = new(hepmc.PdfInfo) + } + h.PdfInfo.ID1 = int(evt.PDF_Parton1) + h.PdfInfo.ID2 = int(evt.PDF_Parton2) + h.PdfInfo.X1 = evt.PDF_X1 + h.PdfInfo.X2 = evt.PDF_X2 + h.PdfInfo.ScalePDF = evt.PDF_Q2 + h.PdfInfo.Pdf1 = evt.PDF_X1f + h.PdfInfo.Pdf2 = evt.PDF_X2f + h.PdfInfo.LHAPdf1 = int(evt.PDF_ID1) + h.PdfInfo.LHAPdf2 = int(evt.PDF_ID2) + + h.MomentumUnit = hepmc.MomentumUnit(evt.MomentumUnit) + h.LengthUnit = hepmc.LengthUnit(evt.LengthUnit) + + // flow := evt.Particle_flow + for i := range evt.Particle_bc { + flow := evt.Particle_flow[i] + p := &hepmc.Particle{ + Momentum: fmom.NewPxPyPzE( + evt.Particle_px[i], + evt.Particle_py[i], + evt.Particle_pz[i], + evt.Particle_ene[i], + ), + PdgID: evt.Particle_pid[i], + Status: int(evt.Particle_status[i]), + Flow: hepmc.Flow{ + Icode: make(map[int]int, len(flow)), + }, + Polarization: hepmc.Polarization{ + Theta: evt.Particle_theta[i], + Phi: evt.Particle_phi[i], + }, + Barcode: int(evt.Particle_bc[i]), + GeneratedMass: evt.Particle_mass[i], + } + p.Flow.Particle = p + for ii := 0; ii < len(flow); ii += 2 { + p.Flow.Icode[int(flow[ii])] = int(flow[ii+1]) + } + h.Particles[p.Barcode] = p + } + + var ( + pin = evt.Event_inbcs + pout = evt.Event_outbcs + ) + for i := range evt.Vertex_bc { + wsli := evt.Vertex_wsli[i] + wkey := evt.Vertex_wkey[i] + wval := evt.Vertex_wval[i] + vtx := &hepmc.Vertex{ + Position: fmom.NewPxPyPzE( + evt.Vertex_x[i], + evt.Vertex_y[i], + evt.Vertex_z[i], + evt.Vertex_t[i], + ), + ParticlesIn: make([]*hepmc.Particle, evt.Vertex_nin[i]), + ParticlesOut: make([]*hepmc.Particle, evt.Vertex_nout[i]), + ID: int(evt.Vertex_id[i]), + Weights: hepmc.Weights{ + Slice: make([]float64, len(wsli)), + Map: make(map[string]int, len(wkey)), + }, + Event: h, + Barcode: int(evt.Vertex_bc[i]), + } + for i := range vtx.ParticlesIn { + p := h.Particles[int(pin[i])] + p.EndVertex = vtx + vtx.ParticlesIn[i] = p + } + pin = pin[len(vtx.ParticlesIn):] + + for i := range vtx.ParticlesOut { + p := h.Particles[int(pout[i])] + p.ProdVertex = vtx + vtx.ParticlesOut[i] = p + } + pout = pout[len(vtx.ParticlesOut):] + + copy(vtx.Weights.Slice, wsli) + for i, k := range wkey { + v := wval[i] + vtx.Weights.Map[k] = int(v) + } + + h.Vertices[vtx.Barcode] = vtx + } + + switch evt.Event_barcodeSPV { + case 0: + h.SignalVertex = nil + default: + h.SignalVertex = h.Vertices[int(evt.Event_barcodeSPV)] + } + + switch evt.Event_barcodeBP1 { + case 0: + h.Beams[0] = nil + default: + h.Beams[0] = h.Particles[int(evt.Event_barcodeBP1)] + } + + switch evt.Event_barcodeBP2 { + case 0: + h.Beams[1] = nil + default: + h.Beams[1] = h.Particles[int(evt.Event_barcodeBP2)] + } + + return nil +} + +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)) + evt.Event_npart = int32(len(h.Particles)) + 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.WeightsSlice = sliceop.Resize(evt.WeightsSlice, len(h.Weights.Slice)) + copy(evt.WeightsSlice, h.Weights.Slice) + evt.WeightsMapKeys = sliceop.Resize(evt.WeightsMapKeys, len(h.Weights.Map))[:0] + evt.WeightsMapNames = sliceop.Resize(evt.WeightsMapNames, len(h.Weights.Map))[:0] + for k, v := range h.Weights.Map { + evt.WeightsMapKeys = append(evt.WeightsMapKeys, k) + evt.WeightsMapNames = append(evt.WeightsMapNames, int32(v)) + } + evt.RandomStates = sliceop.Resize(evt.RandomStates, len(h.RandomStates)) + copy(evt.RandomStates, h.RandomStates) + + switch xsect := h.CrossSection; xsect { + case nil: + evt.XsectValue = 0 + evt.XsectError = 0 + default: + evt.XsectValue = h.CrossSection.Value + evt.XsectError = h.CrossSection.Error + } + + switch hi := h.HeavyIon; hi { + case nil: + evt.HI_ncollHard = 0 + evt.HI_npartProj = 0 + evt.HI_npartTarg = 0 + evt.HI_ncoll = 0 + evt.HI_nnwColl = 0 + evt.HI_nwNColl = 0 + evt.HI_nwNwColl = 0 + evt.HI_spectatorNeutrons = 0 + evt.HI_spectatorProtons = 0 + evt.HI_impactParameter = 0 + evt.HI_eventPlaneAngle = 0 + evt.HI_eccentricity = 0 + evt.HI_sigmaInelNN = 0 + default: + evt.HI_ncollHard = int32(hi.NCollHard) + evt.HI_npartProj = int32(hi.NPartProj) + evt.HI_npartTarg = int32(hi.NPartTarg) + evt.HI_ncoll = int32(hi.NColl) + evt.HI_nnwColl = int32(hi.NNwColl) + evt.HI_nwNColl = int32(hi.NwNColl) + evt.HI_nwNwColl = int32(hi.NwNwColl) + evt.HI_spectatorNeutrons = int32(hi.SpectatorNeutrons) + evt.HI_spectatorProtons = int32(hi.SpectatorProtons) + evt.HI_impactParameter = hi.ImpactParameter + evt.HI_eventPlaneAngle = hi.EventPlaneAngle + evt.HI_eccentricity = hi.Eccentricity + evt.HI_sigmaInelNN = hi.SigmaInelNN + } + + 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) + + evt.MomentumUnit = int8(h.MomentumUnit) + evt.LengthUnit = int8(h.LengthUnit) + + evt.barcodes = sliceop.Resize(evt.barcodes, len(h.Vertices))[:0] + for bc := range h.Vertices { + evt.barcodes = append(evt.barcodes, bc) + } + sort.Sort(sort.Reverse(sort.IntSlice(evt.barcodes))) + + n := len(h.Vertices) + evt.Vertex_x = sliceop.Resize(evt.Vertex_x, n)[:0] + evt.Vertex_y = sliceop.Resize(evt.Vertex_y, n)[:0] + evt.Vertex_z = sliceop.Resize(evt.Vertex_z, n)[:0] + evt.Vertex_t = sliceop.Resize(evt.Vertex_t, n)[:0] + evt.Vertex_id = sliceop.Resize(evt.Vertex_id, n)[:0] + evt.Vertex_wsli = sliceop.Resize(evt.Vertex_wsli, n)[:0] + evt.Vertex_wkey = sliceop.Resize(evt.Vertex_wkey, n)[:0] + evt.Vertex_wval = sliceop.Resize(evt.Vertex_wval, n)[:0] + evt.Vertex_bc = sliceop.Resize(evt.Vertex_bc, n)[:0] + evt.Vertex_nin = sliceop.Resize(evt.Vertex_nin, n)[:0] + evt.Vertex_nout = sliceop.Resize(evt.Vertex_nout, n)[:0] + + for _, bc := range evt.barcodes { + vtx := h.Vertices[bc] + evt.Vertex_x = append(evt.Vertex_x, vtx.Position.X()) + evt.Vertex_y = append(evt.Vertex_y, vtx.Position.Y()) + evt.Vertex_z = append(evt.Vertex_z, vtx.Position.Z()) + evt.Vertex_t = append(evt.Vertex_t, vtx.Position.T()) + evt.Vertex_id = append(evt.Vertex_id, int32(vtx.ID)) + evt.Vertex_bc = append(evt.Vertex_bc, int32(vtx.Barcode)) + evt.Vertex_nin = append(evt.Vertex_nin, int32(len(vtx.ParticlesIn))) + evt.Vertex_nout = append(evt.Vertex_nout, int32(len(vtx.ParticlesOut))) + for _, p := range vtx.ParticlesIn { + evt.Event_inbcs = append(evt.Event_inbcs, int32(p.Barcode)) + } + for _, p := range vtx.ParticlesOut { + evt.Event_outbcs = append(evt.Event_outbcs, int32(p.Barcode)) + } + + wsli := make([]float64, len(vtx.Weights.Slice)) + copy(wsli, vtx.Weights.Slice) + wkey := make([]string, 0, len(vtx.Weights.Map)) + wval := make([]int32, 0, len(vtx.Weights.Map)) + for k, v := range vtx.Weights.Map { + wkey = append(wkey, k) + wval = append(wval, int32(v)) + } + evt.Vertex_wsli = append(evt.Vertex_wsli, wsli) + evt.Vertex_wkey = append(evt.Vertex_wkey, wkey) + evt.Vertex_wval = append(evt.Vertex_wval, wval) + } + + evt.barcodes = sliceop.Resize(evt.barcodes, len(h.Particles))[:0] + for bc := range h.Particles { + evt.barcodes = append(evt.barcodes, bc) + } + sort.Ints(evt.barcodes) + + n = len(h.Particles) + evt.Particle_bc = sliceop.Resize(evt.Particle_bc, n)[:0] + evt.Particle_pid = sliceop.Resize(evt.Particle_pid, n)[:0] + evt.Particle_px = sliceop.Resize(evt.Particle_px, n)[:0] + evt.Particle_py = sliceop.Resize(evt.Particle_py, n)[:0] + evt.Particle_pz = sliceop.Resize(evt.Particle_pz, n)[:0] + evt.Particle_ene = sliceop.Resize(evt.Particle_ene, n)[:0] + evt.Particle_mass = sliceop.Resize(evt.Particle_mass, n)[:0] + // evt.Particle_nflow = sliceop.Resize(evt.Particle_nflow, n)[:0] + evt.Particle_flow = sliceop.Resize(evt.Particle_flow, n)[:0] + evt.Particle_theta = sliceop.Resize(evt.Particle_theta, n)[:0] + evt.Particle_phi = sliceop.Resize(evt.Particle_phi, n)[:0] + evt.Particle_status = sliceop.Resize(evt.Particle_status, n)[:0] + evt.Particle_pvtx = sliceop.Resize(evt.Particle_pvtx, n)[:0] + evt.Particle_evtx = sliceop.Resize(evt.Particle_evtx, n)[:0] + + for _, bc := range evt.barcodes { + p := h.Particles[bc] + switch vtx := p.ProdVertex; vtx { + case nil: + evt.Particle_pvtx = append(evt.Particle_pvtx, 0) + default: + evt.Particle_pvtx = append(evt.Particle_pvtx, int32(vtx.Barcode)) + } + + evt.Particle_bc = append(evt.Particle_bc, 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.GeneratedMass) + // evt.Particle_nflow = append(evt.Particle_nflow, int32(len(p.Flow.Icode))) + flow := make([]int32, 0, 2*len(p.Flow.Icode)) + for k, v := range p.Flow.Icode { + flow = append(flow, int32(k), int32(v)) + } + evt.Particle_flow = append(evt.Particle_flow, flow) + evt.Particle_theta = append(evt.Particle_theta, p.Polarization.Theta) + evt.Particle_phi = append(evt.Particle_phi, p.Polarization.Phi) + evt.Particle_status = append(evt.Particle_status, int32(p.Status)) + switch vtx := p.EndVertex; vtx { + case nil: + evt.Particle_evtx = append(evt.Particle_evtx, 0) + default: + evt.Particle_evtx = append(evt.Particle_evtx, int32(vtx.Barcode)) + } + } + + return nil +} + +func (evt *event) reset() { + evt.Event_inbcs = evt.Event_inbcs[:0] + evt.Event_outbcs = evt.Event_outbcs[:0] + evt.WeightsSlice = evt.WeightsSlice[:0] + evt.WeightsMapKeys = evt.WeightsMapKeys[:0] + evt.WeightsMapNames = evt.WeightsMapNames[:0] + evt.RandomStates = evt.RandomStates[:0] + + evt.Vertex_x = evt.Vertex_x[:0] + evt.Vertex_y = evt.Vertex_y[:0] + evt.Vertex_z = evt.Vertex_z[:0] + evt.Vertex_t = evt.Vertex_t[:0] + evt.Vertex_wsli = evt.Vertex_wsli[:0] + evt.Vertex_wkey = evt.Vertex_wkey[:0] + evt.Vertex_wval = evt.Vertex_wval[:0] + evt.Vertex_id = evt.Vertex_id[:0] + evt.Vertex_bc = evt.Vertex_bc[:0] + evt.Vertex_nin = evt.Vertex_nin[:0] + evt.Vertex_nout = evt.Vertex_nout[:0] + + evt.Particle_bc = evt.Particle_bc[: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_nflow = evt.Particle_nflow[:0] + evt.Particle_flow = evt.Particle_flow[:0] + evt.Particle_theta = evt.Particle_theta[:0] + evt.Particle_phi = evt.Particle_phi[:0] + evt.Particle_status = evt.Particle_status[:0] + evt.Particle_pvtx = evt.Particle_pvtx[:0] + evt.Particle_evtx = evt.Particle_evtx[:0] +} + +func (evt *event) hasXSect() bool { + return !(evt.XsectError == 0 && evt.XsectValue == 0) +} + +func (evt *event) hasHI() bool { + return !(evt.HI_ncollHard == 0 && evt.HI_npartProj == 0 && evt.HI_npartTarg == 0 && evt.HI_ncoll == 0 && + evt.HI_nnwColl == 0 && evt.HI_nwNColl == 0 && evt.HI_nwNwColl == 0 && evt.HI_spectatorNeutrons == 0 && + evt.HI_spectatorProtons == 0 && evt.HI_impactParameter == 0 && evt.HI_eventPlaneAngle == 0 && + evt.HI_eccentricity == 0 && evt.HI_sigmaInelNN == 0) +} diff --git a/hepmc/rootcnv/rootcnv_test.go b/hepmc/rootcnv/rootcnv_test.go new file mode 100644 index 000000000..9687f40c0 --- /dev/null +++ b/hepmc/rootcnv/rootcnv_test.go @@ -0,0 +1,101 @@ +// 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. + +package rootcnv + +import ( + "bytes" + "os" + "path/filepath" + "reflect" + "testing" + + "github.com/google/go-cmp/cmp" + "go-hep.org/x/hep/groot" + "go-hep.org/x/hep/groot/riofs" + "go-hep.org/x/hep/groot/rtree" + "go-hep.org/x/hep/hepmc" +) + +func TestRW(t *testing.T) { + dir, err := os.MkdirTemp("", "hepmc-rootcnv-") + if err != nil { + t.Fatal(err) + } + defer os.RemoveAll(dir) + + for _, tc := range []string{ + "../testdata/small.hepmc", + "../testdata/test.hepmc", + } { + t.Run(tc, func(t *testing.T) { + raw, err := os.ReadFile(tc) + if err != nil { + t.Fatal(err) + } + + fname := filepath.Join(dir, filepath.Base(tc)+".root") + o, err := groot.Create(fname) + if err != nil { + t.Fatalf("could not create output ROOT file: %+v", err) + } + defer o.Close() + + w, err := NewFlatTreeWriter(o, "tree", rtree.WithTitle("HepMC ROOT tree")) + if err != nil { + t.Fatalf("could not create ROOT tree writer: %+v", err) + } + + _, err = hepmc.Copy(w, hepmc.NewASCIIReader(bytes.NewReader(raw))) + if err != nil { + t.Fatalf("could not copy hepmc event to ROOT: %+v", err) + } + + err = w.Close() + if err != nil { + t.Fatalf("could not close ROOT tree writer: %+v", err) + } + + err = o.Close() + if err != nil { + t.Fatalf("could not close ROOT file: %+v", err) + } + + f, err := groot.Open(fname) + if err != nil { + t.Fatalf("could not open ROOT file: %+v", err) + } + defer f.Close() + + r, err := riofs.Get[rtree.Tree](f, "tree") + if err != nil { + t.Fatalf("could not retrieve ROOT tree: %+v", err) + } + + rr, err := NewFlatTreeReader(r) + if err != nil { + t.Fatalf("could not create ROOT tree reader: %+v", err) + } + defer rr.Close() + + buf := new(bytes.Buffer) + + ww := hepmc.NewASCIIWriter(buf) + _, err = hepmc.Copy(ww, rr) + if err != nil { + t.Fatalf("could not copy hepmc event from ROOT: %+v", err) + } + + err = ww.Close() + if err != nil { + t.Fatalf("could not close hepmc writer: %+v", err) + } + + if got, want := buf.Bytes(), raw; !reflect.DeepEqual(got, want) { + diff := cmp.Diff(string(want), string(got)) + t.Fatalf("invalid r/w round trip:\n%s", diff) + } + }) + } +} diff --git a/hepmc/rootcnv/writer.go b/hepmc/rootcnv/writer.go new file mode 100644 index 000000000..dda7525cf --- /dev/null +++ b/hepmc/rootcnv/writer.go @@ -0,0 +1,60 @@ +// 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. + +package rootcnv + +import ( + "fmt" + + "go-hep.org/x/hep/groot/riofs" + "go-hep.org/x/hep/groot/rtree" + "go-hep.org/x/hep/hepmc" +) + +// FlatTreeWriter writes HepMC events as a flat ROOT TTree. +type FlatTreeWriter struct { + w rtree.Writer + + evt event + wvars []rtree.WriteVar +} + +// NewFlatTreeWriter creates a new named tree under the dir directory. +func NewFlatTreeWriter(dir riofs.Directory, name string, opts ...rtree.WriteOption) (*FlatTreeWriter, error) { + var w FlatTreeWriter + w.wvars = rtree.WriteVarsFromStruct(&w.evt) + tree, err := rtree.NewWriter(dir, name, w.wvars, opts...) + if err != nil { + return nil, fmt.Errorf("hepmc: could not create flat-tree writer %q: %w", name, err) + } + + w.w = tree + + return &w, nil +} + +func (w *FlatTreeWriter) Close() error { + w.wvars = nil + return w.w.Close() +} + +func (w *FlatTreeWriter) Write(evt hepmc.Event) error { + w.evt.reset() + + err := w.evt.read(&evt) + if err != nil { + return fmt.Errorf("hepmc: could not encode event to ROOT: %w", err) + } + + _, err = w.w.Write() + if err != nil { + return fmt.Errorf("hepmc: could not write event to ROOT: %w", err) + } + + return nil +} + +var ( + _ hepmc.Writer = (*FlatTreeWriter)(nil) +)