Belle II Software light-2406-ragdoll
LCASaverModule Class Reference
Inheritance diagram for LCASaverModule:
Collaboration diagram for LCASaverModule:

Public Member Functions

def __init__ (self, particle_lists, features, mcparticle_list, output_file)
 
def initialize (self)
 
def event (self)
 
def terminate (self)
 

Public Attributes

 particle_lists
 Input particle lists.
 
 features
 Features to save.
 
 mcparticle_list
 MC particle list (Upsilon, B0, B+)
 
 output_file
 Output file name.
 
 max_particles
 Max number of particles It doesn't actually matter what this is as long as it's bigger than the number of particles in any events we'll encounter.
 
 eventinfo
 Event info.
 
 root_outfile
 ROOT output file.
 
 tree
 ROOT tree.
 
 event_num
 Event number.
 
 isB
 bool containing information whether we reconstruct B or Upsilon
 
 truth_dict
 Truth information.
 
 n_particles
 Number of particles.
 
 primary
 Particle-is-primary flag.
 
 leaves
 Leaves in event.
 
 b_index
 B index.
 
 feat_dict
 Features dictionary.
 
 mc_pdg
 True MC PDG.
 

Protected Member Functions

def _reset_LCA (self)
 

Detailed Description

Save Lowest Common Ancestor matrix of each MC Particle in the given list.

Args:
    particle_lists (list): Name of particle lists to save features of.
    features (list): List of features to save for each particle.
    mcparticle_list (str): Name of particle list to build LCAs from (will use as root).
    output_file (str): Path to output file to save.

Definition at line 205 of file LCASaverModule.py.

Constructor & Destructor Documentation

◆ __init__()

def __init__ (   self,
  particle_lists,
  features,
  mcparticle_list,
  output_file 
)
Initialization.

Definition at line 216 of file LCASaverModule.py.

222 ):
223 """
224 Initialization.
225 """
226 super().__init__()
227
228 self.particle_lists = particle_lists
229
230 self.features = features
231
232 self.mcparticle_list = mcparticle_list
233
234 self.output_file = output_file
235
236
239 self.max_particles = 500
240

Member Function Documentation

◆ _reset_LCA()

def _reset_LCA (   self)
protected
Resets the value of the LCA to 0.

Definition at line 320 of file LCASaverModule.py.

320 def _reset_LCA(self):
321 """
322 Resets the value of the LCA to 0.
323 """
324 for p_index in [1, 2]:
325 self.truth_dict[f"LCAS_{p_index}"][: self.max_particles**2] *= 0
326 self.truth_dict[f"n_LCA_leaves_{p_index}"][0] *= 0
327 self.truth_dict[f"LCA_leaves_{p_index}"][: self.max_particles] *= 0
328 self.truth_dict[f"n_LCA_{p_index}"][0] *= 0
329

◆ event()

def event (   self)
Called for each event.

Definition at line 330 of file LCASaverModule.py.

330 def event(self):
331 """
332 Called for each event.
333 """
334 from ROOT.Belle2 import PyStoreObj
335
336 # Get the particle list (note this is a regular Particle list, not MCParticle)
337 p_list = PyStoreObj(self.mcparticle_list)
338
339 self.event_num[0] = self.eventinfo.getEvent()
340
341 # Is Upsilon flag
342 if self.mcparticle_list == "Upsilon(4S):MC":
343 self.isB[0] = 0
344 elif self.mcparticle_list == "B0:MC":
345 self.isB[0] = 1
346 elif self.mcparticle_list == "B+:MC":
347 self.isB[0] = 2
348
349 # Create the LCA
350 # IMPORTANT: The ArrayIndex is 0-based.
351 # mcplist contains the root particles we are to create LCAs from
352 # Reset the LCA list so if only one B is there it does not carry an older version over
353 self._reset_LCA()
354
355 if p_list.getListSize() > 0:
356 for part in p_list.obj():
357 # Get the corresponding MCParticle
358 mcp = part.getMCParticle()
359 # In this way the LCA variables in the ntuples will be labelled _1 and _2
360 # If we train on B decays these will correspond to the two B's
361 # while if we train on the Upsilon, _1 will correspond to it and _2 will remain empty
362 # becaus getArrayIndex() gives 0 for the Upsilon and 1, 2 for the B's
363 array_index = 1 if self.isB[0] == 0 else mcp.getArrayIndex()
364
365 # Call function to write history of leaves in the tree.
366 # It internally calls function update_levels to find and save the level of each particle in the tree.
367 # Necessary step to build the LCA.
368 (
369 lcas_leaf_hist,
370 lcas_levels,
371 _, _, _,
372 ) = write_hist(
373 particle=mcp,
374 leaf_hist={},
375 levels={},
376 hist=[],
377 pdg={},
378 leaf_pdg={},
379 semilep_flag=False,
380 )
381
382 lcas = np.zeros([len(lcas_leaf_hist), len(lcas_leaf_hist)])
383
384 for x, y in combinations(enumerate(lcas_leaf_hist), 2):
385 lcas_intersection = [
386 i for i in lcas_leaf_hist[x[1]] if i in lcas_leaf_hist[y[1]]
387 ]
388 lcas[x[0], y[0]] = lcas_levels[lcas_intersection[-1]]
389 lcas[y[0], x[0]] = lcas_levels[lcas_intersection[-1]]
390
391 self.truth_dict[f"LCAS_{array_index}"][: lcas.size] = lcas.flatten()
392
393 self.truth_dict[f"n_LCA_leaves_{array_index}"][0] = len(lcas_leaf_hist.keys())
394 self.truth_dict[f"LCA_leaves_{array_index}"][
395 : len(lcas_leaf_hist.keys())
396 ] = list(lcas_leaf_hist.keys())
397
398 self.truth_dict[f"n_LCA_{array_index}"][0] = lcas.size
399
400 # ### Create the features
401 # Where we'll append features
402 evt_feats = {f: [] for f in self.features}
403
404 # Ideally this would be saved in evt_feat_dict but the -1 for unmatched
405 # particles messes that up
406 evt_leaf_dict = {
407 "leaves": [],
408 "primary": [],
409 "b_index": [],
410 "mc_pdg": [],
411 }
412
413 # IMPORTANT: The ArrayIndex is 0-based.
414 # mcplist contains the root particles we are to create LCAs from
415 for p_list_name in self.particle_lists:
416 # Get the particle list (note this is a regular Particle list, not MCParticle)
417 p_list = PyStoreObj(p_list_name)
418
419 for particle in p_list.obj():
420 # Get the B parent index, set to -1 if particle has no MC match
421 b_index = int(vm.evaluate("ancestorBIndex", particle))
422 # Need this to reorder LCA later, returns -1 if no MC match
423 if b_index >= 0:
424 p_index = particle.getMCParticle().getArrayIndex()
425 p_primary = particle.getMCParticle().isPrimaryParticle()
426 # Save particle's PDG code
427 mc_pdg = particle.getMCParticle().getPDG()
428 else:
429 p_index = -1
430 p_primary = False
431 mc_pdg = 0
432
433 evt_leaf_dict["primary"].append(p_primary)
434 evt_leaf_dict["leaves"].append(p_index)
435 evt_leaf_dict["b_index"].append(b_index)
436 evt_leaf_dict["mc_pdg"].append(mc_pdg)
437 for feat in self.features:
438 evt_feats[feat].append(vm.evaluate(feat, particle))
439
440 n_particles = len(evt_leaf_dict["primary"])
441 self.n_particles[0] = n_particles
442
443 self.primary[:n_particles] = evt_leaf_dict["primary"]
444 self.leaves[:n_particles] = evt_leaf_dict["leaves"]
445 self.b_index[:n_particles] = evt_leaf_dict["b_index"]
446 self.mc_pdg[:n_particles] = evt_leaf_dict["mc_pdg"]
447
448 for feat in self.features:
449 self.feat_dict[feat][:n_particles] = evt_feats[feat]
450 # Add number of final state particles
451
452 # Write everything to the TTree
453 self.tree.Fill()
454

◆ initialize()

def initialize (   self)
Called once at the beginning.

Definition at line 241 of file LCASaverModule.py.

241 def initialize(self):
242 """
243 Called once at the beginning.
244 """
245 from ROOT import Belle2, TFile, TTree
246
247
248 self.eventinfo = Belle2.PyStoreObj("EventMetaData")
249
250
251 self.root_outfile = TFile(self.output_file, "recreate")
252
253 self.tree = TTree("Tree", "tree")
254
255
256 self.event_num = np.zeros(1, dtype=np.int32)
257 self.tree.Branch("event", self.event_num, "event/I")
258
259 self.isB = np.zeros(1, dtype=bool)
260 self.tree.Branch("isB", self.isB, "isB/b")
261
262
263 self.truth_dict = {}
264
265 # We assume at most two LCA matrices for event
266 for i in [1, 2]:
267 self.truth_dict[f"n_LCA_leaves_{i}"] = np.zeros(1, dtype=np.int32)
268 self.truth_dict[f"LCA_leaves_{i}"] = np.zeros(
269 self.max_particles, dtype=np.int32
270 )
271 self.tree.Branch(
272 f"n_LCA_leaves_{i}",
273 self.truth_dict[f"n_LCA_leaves_{i}"],
274 f"n_LCA_leaves_{i}/I",
275 )
276 self.tree.Branch(
277 f"LCA_leaves_{i}",
278 self.truth_dict[f"LCA_leaves_{i}"],
279 f"LCA_leaves_{i}[n_LCA_leaves_{i}]/I",
280 )
281
282 self.truth_dict[f"n_LCA_{i}"] = np.zeros(1, dtype=np.int32)
283 self.truth_dict[f"LCAS_{i}"] = np.zeros(
284 self.max_particles**2, dtype=np.uint8
285 )
286 self.tree.Branch(f"n_LCA_{i}", self.truth_dict[f"n_LCA_{i}"], f"n_LCA_{i}/I")
287 self.tree.Branch(
288 f"LCAS_{i}", self.truth_dict[f"LCAS_{i}"], f"LCAS_{i}[n_LCA_{i}]/b"
289 )
290
291 # Feature data
292
293 self.n_particles = np.zeros(1, dtype=np.int32)
294 self.tree.Branch("n_particles", self.n_particles, "n_particles/I")
295
296
297 self.primary = np.zeros(self.max_particles, dtype=np.bool)
298 self.tree.Branch("primary", self.primary, "primary[n_particles]/O")
299
300
301 self.leaves = np.zeros(self.max_particles, dtype=np.int32)
302 self.tree.Branch("leaves", self.leaves, "leaves[n_particles]/I")
303
304
305 self.b_index = np.zeros(self.max_particles, dtype=np.int32)
306 self.tree.Branch("b_index", self.b_index, "b_index[n_particles]/I")
307
308
309 self.feat_dict = {}
310 for feat in self.features:
311 self.feat_dict[feat] = np.zeros(self.max_particles, np.float32)
312 self.tree.Branch(
313 f"feat_{feat}", self.feat_dict[feat], f"feat_{feat}[n_particles]/F"
314 )
315
316
317 self.mc_pdg = np.zeros(self.max_particles, np.float32)
318 self.tree.Branch("mcPDG", self.mc_pdg, "mcPDG[n_particles]/F")
319
a (simplified) python wrapper for StoreObjPtr.
Definition: PyStoreObj.h:67

◆ terminate()

def terminate (   self)
Called at the end.

Definition at line 455 of file LCASaverModule.py.

455 def terminate(self):
456 """
457 Called at the end.
458 """
459 self.root_outfile.Write()
460 self.root_outfile.Close()

Member Data Documentation

◆ b_index

b_index

B index.

Definition at line 305 of file LCASaverModule.py.

◆ event_num

event_num

Event number.

Definition at line 256 of file LCASaverModule.py.

◆ eventinfo

eventinfo

Event info.

Definition at line 248 of file LCASaverModule.py.

◆ feat_dict

feat_dict

Features dictionary.

Definition at line 309 of file LCASaverModule.py.

◆ features

features

Features to save.

Definition at line 230 of file LCASaverModule.py.

◆ isB

isB

bool containing information whether we reconstruct B or Upsilon

Definition at line 259 of file LCASaverModule.py.

◆ leaves

leaves

Leaves in event.

Definition at line 301 of file LCASaverModule.py.

◆ max_particles

max_particles

Max number of particles It doesn't actually matter what this is as long as it's bigger than the number of particles in any events we'll encounter.

ROOT won't save all entries.

Definition at line 239 of file LCASaverModule.py.

◆ mc_pdg

mc_pdg

True MC PDG.

Definition at line 317 of file LCASaverModule.py.

◆ mcparticle_list

mcparticle_list

MC particle list (Upsilon, B0, B+)

Definition at line 232 of file LCASaverModule.py.

◆ n_particles

n_particles

Number of particles.

Definition at line 293 of file LCASaverModule.py.

◆ output_file

output_file

Output file name.

Definition at line 234 of file LCASaverModule.py.

◆ particle_lists

particle_lists

Input particle lists.

Definition at line 228 of file LCASaverModule.py.

◆ primary

primary

Particle-is-primary flag.

Definition at line 297 of file LCASaverModule.py.

◆ root_outfile

root_outfile

ROOT output file.

Definition at line 251 of file LCASaverModule.py.

◆ tree

tree

ROOT tree.

Definition at line 253 of file LCASaverModule.py.

◆ truth_dict

truth_dict

Truth information.

Definition at line 263 of file LCASaverModule.py.


The documentation for this class was generated from the following file: