@@ -88,72 +88,72 @@ def open(filename, mode="r", precision=None):
8888
8989 elif mode == "w" :
9090 return WrappedAsciiWriter (filename , precision )
91-
92-
93- def fill_genevent_from_hepevent_ptr ( evt , ptr_as_int , max_size ,
94- int_type = ctypes . c_int ,
95- float_type = ctypes . c_double ,
96- hep_status_decoder = None ,
97- momentum_unit = 1 ,
98- length_unit = 1 ):
99- # struct HEPEVT // Fortran common block HEPEVT
100- # {
101- # int nevhep; // Event number
102- # int nhep; // Number of entries in the event
103- # int isthep[NMXHEP]; // Status code
104- # int idhep [NMXHEP]; // PDG ID
105- # int jmohep[NMXHEP][2]; // Idx of first and last mother
106- # int jdahep[NMXHEP][2]; // Idx of first and last daughter
107- # momentum_t phep [NMXHEP][5]; // Momentum: px, py, pz, e, m
108- # momentum_t vhep [NMXHEP][4]; // Vertex: x, y, z, t
109- # };
110-
111- IntArray = int_type * max_size
112- Int2 = int_type * 2
113- Int2Array = Int2 * max_size
114- Float4 = float_type * 4
115- Float5 = float_type * 5
116- Float4Array = Float4 * max_size
117- Float5Array = Float5 * max_size
118-
119- class HEPEVT ( ctypes . Structure ):
120- _fields_ = (
121- ( "event_number" , int_type ),
122- ( "nentries" , int_type ),
123- ( "status" , IntArray ),
124- ( "pid" , IntArray ),
125- ( "parents" , Int2Array ),
126- ( "children" , Int2Array ),
127- ( "pm" , Float5Array ),
128- ( "v" , Float4Array )
129- )
130-
131- h = ctypes . cast ( ptr_as_int , ctypes . POINTER ( HEPEVT ))[ 0 ]
132-
133- event_number = h . event_number
134- n = h . nentries
135-
136- import numpy as np
137-
138- status = np . asarray ( h . status )[: n ]
139- pm = np . asarray ( h . pm )
140-
141- if hep_status_decoder is None :
142- particle_status = status
143- vertex_status = np . zeros_like ( status )
144- else :
145- particle_status , vertex_status = hep_status_decoder ( status )
146-
147- fill_genevent_from_hepevt ( evt ,
148- event_number ,
149- pm [: n ,: 4 ],
150- pm [:n ,4 ],
151- h . v [:n ],
152- h . pid [:n ],
153- h . parents [:n ],
154- h . children [:n ],
155- particle_status ,
156- vertex_status ,
157- momentum_unit ,
158- length_unit )
159- return evt
91+ raise ValueError ( "mode must be r or w" )
92+
93+
94+ _hepevt_buffer = HEPEVT ()
95+
96+ def fill_genevent_from_hepevt ( evt , ** kwargs ):
97+ """
98+ Fills GenEvent from HEPEVT data.
99+
100+ Parameters
101+ ----------
102+ event_number : int
103+ Current Event Number. Starts with 1.
104+ p : array(float) N x 4
105+ Array of 4-vectors (px, py, py, e).
106+ m : array(float) N
107+ Array of generated masses (virtual masses).
108+ v : array(float) N x 4
109+ Array of 4-vectors (x, y, z, t).
110+ pid : array(int) N
111+ Array of PDG IDs.
112+ parents : array(int) N x 2
113+ Array of (parent first index, parent last index). Indices start at 1.
114+ children : array(int) N x 2 (optional)
115+ Array of (child first index, child last index). Indices start at 1. May be none.
116+ status : array(int) N
117+ Array of particle status. See HEPMC3 docs for definition.
118+ momentum_scaling : float (optional)
119+ Momentum coordinates are stored in units of this number.
120+ vertex_scaling : float (optional)
121+ Vertex coordinates are stored in units of this number.
122+
123+ Notes
124+ -----
125+ The current implementation copies the input into an internal buffer. This is not as efficient as possible, but currently the only way to use the HepMC3 C++ code with input data which does not have exactly the memory layout of the HEPEVT struct defined in HepMC3/HEPEVT_Wrapper.h.
126+
127+ Calling this function is not thread-safe.
128+ """
129+ event_number = kwargs [ "event_number" ]
130+ p = kwargs [ "p" ]
131+ m = kwargs [ "m" ]
132+ v = kwargs [ "v" ]
133+ pid = kwargs [ "pid" ]
134+ parents = kwargs [ "parents" ]
135+ children = kwargs . get ( "children" , 0 )
136+ status = kwargs [ "status" ]
137+ momentum_scaling = kwargs . get ( "momentum_scaling" , 1.0 )
138+ vertex_scaling = kwargs . get ( "vertex_scaling" , 1.0 )
139+
140+ global _hepevt_buffer
141+ n = pid . shape [ 0 ]
142+ if n > _hepevt_buffer . max_size :
143+ raise ValueError (
144+ ( 'Number of particles in event (%i) exceeds HepMC3 buffer size (%i). \n '
145+ 'Change the line `define_macros={"HEPMC3_HEPEVT_NMXHEP": 50000}` in setup.py \n '
146+ 'to a larger value and (re)compile pyhepmc_ng from scratch.' ) %
147+ ( n , _hepevt_buffer . max_size ))
148+ _hepevt_buffer . event_number = event_number
149+ _hepevt_buffer . nentries = n
150+ _hepevt_buffer . pm () [:n , : 4 ] = p / momentum_scaling
151+ _hepevt_buffer . pm () [:n , 4 ] = m / momentum_scaling
152+ _hepevt_buffer . v () [:n ] = v / vertex_scaling
153+ _hepevt_buffer . pid () [:n ] = pid
154+ _hepevt_buffer . parents () [:n ] = parents
155+ _hepevt_buffer . children ()[: n ] = children
156+ _hepevt_buffer . status ()[: n ] = status
157+
158+ evt . clear ( )
159+ _hepevt_buffer . to_genevent ( evt )
0 commit comments