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.

    ):
        """
        Initialization.
        """
        super().__init__()
        
        self.particle_lists = particle_lists
        
        self.features = features
        
        self.mcparticle_list = mcparticle_list
        
        self.output_file = output_file
 
        
        self.max_particles = 500
 

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.

    def _reset_LCA(self):
        """
        Resets the value of the LCA to 0.
        """
        for p_index in [1, 2]:
            self.truth_dict[f"LCAS_{p_index}"][: self.max_particles**2] *= 0
            self.truth_dict[f"n_LCA_leaves_{p_index}"][0] *= 0
            self.truth_dict[f"LCA_leaves_{p_index}"][: self.max_particles] *= 0
            self.truth_dict[f"n_LCA_{p_index}"][0] *= 0
 

event()

def event

(

self

)

Called for each event.

Definition at line 330 of file LCASaverModule.py.

    def event(self):
        """
        Called for each event.
        """
        from ROOT.Belle2 import PyStoreObj
 
        # Get the particle list (note this is a regular Particle list, not MCParticle)
        p_list = PyStoreObj(self.mcparticle_list)
 
        self.event_num[0] = self.eventinfo.getEvent()
 
        # Is Upsilon flag
        if self.mcparticle_list == "Upsilon(4S):MC":
            self.isB[0] = 0
        elif self.mcparticle_list == "B0:MC":
            self.isB[0] = 1
        elif self.mcparticle_list == "B+:MC":
            self.isB[0] = 2
 
        # Create the LCA
        # IMPORTANT: The ArrayIndex is 0-based.
        # mcplist contains the root particles we are to create LCAs from
        # Reset the LCA list so if only one B is there it does not carry an older version over
        self._reset_LCA()
 
        if p_list.getListSize() > 0:
            for part in p_list.obj():
                # Get the corresponding MCParticle
                mcp = part.getMCParticle()
                # In this way the LCA variables in the ntuples will be labelled _1 and _2
                # If we train on B decays these will correspond to the two B's
                # while if we train on the Upsilon, _1 will correspond to it and _2 will remain empty
                # becaus getArrayIndex() gives 0 for the Upsilon and 1, 2 for the B's
                array_index = 1 if self.isB[0] == 0 else mcp.getArrayIndex()
 
                # Call function to write history of leaves in the tree.
                # It internally calls function update_levels to find and save the level of each particle in the tree.
                # Necessary step to build the LCA.
                (
                    lcas_leaf_hist,
                    lcas_levels,
                    _, _, _,
                ) = write_hist(
                    particle=mcp,
                    leaf_hist={},
                    levels={},
                    hist=[],
                    pdg={},
                    leaf_pdg={},
                    semilep_flag=False,
                )
 
                lcas = np.zeros([len(lcas_leaf_hist), len(lcas_leaf_hist)])
 
                for x, y in combinations(enumerate(lcas_leaf_hist), 2):
                    lcas_intersection = [
                        i for i in lcas_leaf_hist[x[1]] if i in lcas_leaf_hist[y[1]]
                    ]
                    lcas[x[0], y[0]] = lcas_levels[lcas_intersection[-1]]
                    lcas[y[0], x[0]] = lcas_levels[lcas_intersection[-1]]
 
                self.truth_dict[f"LCAS_{array_index}"][: lcas.size] = lcas.flatten()
 
                self.truth_dict[f"n_LCA_leaves_{array_index}"][0] = len(lcas_leaf_hist.keys())
                self.truth_dict[f"LCA_leaves_{array_index}"][
                    : len(lcas_leaf_hist.keys())
                ] = list(lcas_leaf_hist.keys())
 
                self.truth_dict[f"n_LCA_{array_index}"][0] = lcas.size
 
            # ### Create the features
            # Where we'll append features
            evt_feats = {f: [] for f in self.features}
 
            # Ideally this would be saved in evt_feat_dict but the -1 for unmatched
            # particles messes that up
            evt_leaf_dict = {
                "leaves": [],
                "primary": [],
                "b_index": [],
                "mc_pdg": [],
            }
 
            # IMPORTANT: The ArrayIndex is 0-based.
            # mcplist contains the root particles we are to create LCAs from
            for p_list_name in self.particle_lists:
                # Get the particle list (note this is a regular Particle list, not MCParticle)
                p_list = PyStoreObj(p_list_name)
 
                for particle in p_list.obj():
                    # Get the B parent index, set to -1 if particle has no MC match
                    b_index = int(vm.evaluate("ancestorBIndex", particle))
                    # Need this to reorder LCA later, returns -1 if no MC match
                    if b_index >= 0:
                        p_index = particle.getMCParticle().getArrayIndex()
                        p_primary = particle.getMCParticle().isPrimaryParticle()
                        # Save particle's PDG code
                        mc_pdg = particle.getMCParticle().getPDG()
                    else:
                        p_index = -1
                        p_primary = False
                        mc_pdg = 0
 
                    evt_leaf_dict["primary"].append(p_primary)
                    evt_leaf_dict["leaves"].append(p_index)
                    evt_leaf_dict["b_index"].append(b_index)
                    evt_leaf_dict["mc_pdg"].append(mc_pdg)
                    for feat in self.features:
                        evt_feats[feat].append(vm.evaluate(feat, particle))
 
            n_particles = len(evt_leaf_dict["primary"])
            self.n_particles[0] = n_particles
 
            self.primary[:n_particles] = evt_leaf_dict["primary"]
            self.leaves[:n_particles] = evt_leaf_dict["leaves"]
            self.b_index[:n_particles] = evt_leaf_dict["b_index"]
            self.mc_pdg[:n_particles] = evt_leaf_dict["mc_pdg"]
 
            for feat in self.features:
                self.feat_dict[feat][:n_particles] = evt_feats[feat]
            # Add number of final state particles
 
            # Write everything to the TTree
            self.tree.Fill()
 

initialize()

def initialize

(

self

)

Called once at the beginning.

Definition at line 241 of file LCASaverModule.py.

    def initialize(self):
        """
        Called once at the beginning.
        """
        from ROOT import Belle2, TFile, TTree
 
        
        self.eventinfo = Belle2.PyStoreObj("EventMetaData")
 
        
        self.root_outfile = TFile(self.output_file, "recreate")
        
        self.tree = TTree("Tree", "tree")
 
        
        self.event_num = np.zeros(1, dtype=np.int32)
        self.tree.Branch("event", self.event_num, "event/I")
        
        self.isB = np.zeros(1, dtype=bool)
        self.tree.Branch("isB", self.isB, "isB/b")
 
        
        self.truth_dict = {}
 
        # We assume at most two LCA matrices for event
        for i in [1, 2]:
            self.truth_dict[f"n_LCA_leaves_{i}"] = np.zeros(1, dtype=np.int32)
            self.truth_dict[f"LCA_leaves_{i}"] = np.zeros(
                self.max_particles, dtype=np.int32
            )
            self.tree.Branch(
                f"n_LCA_leaves_{i}",
                self.truth_dict[f"n_LCA_leaves_{i}"],
                f"n_LCA_leaves_{i}/I",
            )
            self.tree.Branch(
                f"LCA_leaves_{i}",
                self.truth_dict[f"LCA_leaves_{i}"],
                f"LCA_leaves_{i}[n_LCA_leaves_{i}]/I",
            )
 
            self.truth_dict[f"n_LCA_{i}"] = np.zeros(1, dtype=np.int32)
            self.truth_dict[f"LCAS_{i}"] = np.zeros(
                self.max_particles**2, dtype=np.uint8
            )
            self.tree.Branch(f"n_LCA_{i}", self.truth_dict[f"n_LCA_{i}"], f"n_LCA_{i}/I")
            self.tree.Branch(
                f"LCAS_{i}", self.truth_dict[f"LCAS_{i}"], f"LCAS_{i}[n_LCA_{i}]/b"
            )
 
        # Feature data
        
        self.n_particles = np.zeros(1, dtype=np.int32)
        self.tree.Branch("n_particles", self.n_particles, "n_particles/I")
 
        
        self.primary = np.zeros(self.max_particles, dtype=np.bool)
        self.tree.Branch("primary", self.primary, "primary[n_particles]/O")
 
        
        self.leaves = np.zeros(self.max_particles, dtype=np.int32)
        self.tree.Branch("leaves", self.leaves, "leaves[n_particles]/I")
 
        
        self.b_index = np.zeros(self.max_particles, dtype=np.int32)
        self.tree.Branch("b_index", self.b_index, "b_index[n_particles]/I")
 
        
        self.feat_dict = {}
        for feat in self.features:
            self.feat_dict[feat] = np.zeros(self.max_particles, np.float32)
            self.tree.Branch(
                f"feat_{feat}", self.feat_dict[feat], f"feat_{feat}[n_particles]/F"
            )
 
        
        self.mc_pdg = np.zeros(self.max_particles, np.float32)
        self.tree.Branch("mcPDG", self.mc_pdg, "mcPDG[n_particles]/F")
 

terminate()

def terminate

(

self

)

Called at the end.

Definition at line 455 of file LCASaverModule.py.

    def terminate(self):
        """
        Called at the end.
        """
        self.root_outfile.Write()
        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:

• analysis/scripts/grafei/modules/LCASaverModule.py