import enum import attrs from typing import List, Optional, Union import numpy as np import uproot from pathlib import Path from tqdm import tqdm import os import joblib import click class ParticleType(enum.Enum): NeutralHadron = 0 Photon = 1 Electron = 2 Muon = 3 Pion = 4 ChargedKaon = 5 Proton = 6 Others = 7 @attrs.define class ParticleBase: particle_id: int = attrs.field() # unique id for each particle part_charge: int = attrs.field() # charge of the particle part_energy: float = attrs.field() # energy of the particle part_px: float = attrs.field() # x-component of the momentum vector part_py: float = attrs.field() # y-component of the momentum vector part_pz: float = attrs.field() # z-component of the momentum vector log_energy: float = attrs.field() # log10(part_energy) log_pt: float = attrs.field() # log10(part_pt) part_deta: float = attrs.field() # pseudorapidity part_dphi: float = attrs.field() # azimuthal angle part_logptrel: float = attrs.field() # log10(pt(particle)/pt(jet)) part_logerel: float = attrs.field() # log10(energy(particle)/energy(jet)) part_deltaR: float = attrs.field() # distance between the particle and the jet part_d0: float = attrs.field() # tanh(d0) part_dz: float = attrs.field() # tanh(z0) particle_type: ParticleType = attrs.field() # type of the particle as an enum particle_pid: int = attrs.field() # pid of the particle jet_type: str = attrs.field() # type of the jet (e.g. b, bbbar) def properties_concatenated(self, selected_attrs: Optional[List[str]] = None, attr_sep: str = ',') -> str: if selected_attrs is None: selected_attrs = [field.name for field in attrs.fields(self) if field.name != "jet_type"] values = [getattr(self, attr) for attr in selected_attrs] values = [v.value if isinstance(v, ParticleType) else v for v in values] # Convert ParticleType to its numeric value return attr_sep.join(map(str, values)) @attrs.define class Jet: jet_energy: float = attrs.field() # energy of the jet jet_pt: float = attrs.field() # transverse momentum of the jet jet_eta: float = attrs.field() # pseudorapidity of the jet label: str = attrs.field() # type of bb or bbbar particles: List[ParticleBase] = attrs.field(factory=list) # list of particles in the jet def __len__(self): return len(self.particles) def particles_concatenated(self, selected_attrs: Optional[List[str]] = None, attr_sep: str = ',', part_sep: str = '|') -> str: return part_sep.join(map(lambda p: p.properties_concatenated(selected_attrs, attr_sep), self.particles)) @attrs.define class JetSet: jets_type: str = attrs.field() # type of the jets (e.g. bb, bbbar) jets: List[Jet] = attrs.field(factory=list) def __len__(self): return len(self.jets) @staticmethod def jud_type(jtmp): particle_dict = { 0: ParticleType.NeutralHadron, 1: ParticleType.Photon, 2: ParticleType.Electron, 3: ParticleType.Muon, 4: ParticleType.Pion, 5: ParticleType.ChargedKaon, 6: ParticleType.Proton, 7: ParticleType.Others } max_element = max(jtmp) idx = jtmp.index(max_element) return particle_dict.get(idx, ParticleType.Others), idx @classmethod def build_jetset_fast(cls, root_file: Union[str, Path]) -> "JetSet": print(f"Building JetSet from {root_file}") if not 'bbbar' in Path(root_file).stem and not 'bb' in Path(root_file).stem: raise ValueError("Invalid file name, should contain 'bb' or 'bbbar'") jets_type = "bbbar" if "bbbar" in Path(root_file).stem else "bb" print(f"jet type: {jets_type}") with uproot.open(root_file) as f: tree = f["tree"] a = tree.arrays(library="pd") # print(f"DataFrame structure:\n{a.head()}") if a.empty: raise ValueError("DataFrame is empty") label = Path(root_file).stem.split('_')[0] jet_type = "bbbar jet" if "bbbar" in root_file.stem else "b jet" jet_list = [] for j in a.itertuples(): part_pt = np.array(j.part_pt) jet_pt = np.array(j.jet_pt) part_logptrel = np.log(np.divide(part_pt, jet_pt)) part_energy = np.array(j.part_energy) jet_energy = np.array(j.jet_energy) part_logerel = np.log(np.divide(part_energy, jet_energy)) part_deta = np.array(j.part_deta) part_dphi = np.array(j.part_dphi) part_deltaR = np.hypot(part_deta, part_dphi) assert len(j.part_pt) == len(j.part_energy) == len(j.part_deta) particles = [] # particle_list = ['part_isNeutralHadron', 'part_isPhoton', 'part_isElectron', 'part_isMuon', 'part_isPion', 'part_isChargedKaon', 'part_isProton'] part_type = [] part_pid = [] for pn in range(len(j.part_pt)): jtmp = [j.part_isNeutralHadron[pn], j.part_isPhoton[pn], j.part_isElectron[pn], j.part_isMuon[pn], j.part_isPion[pn], j.part_isChargedKaon[pn], j.part_isProton[pn]] tmp_type, tmp_pid = cls.jud_type(jtmp) part_type.append(tmp_type) part_pid.append(tmp_pid) bag = zip(j.part_charge, j.part_energy, j.part_px, j.part_py, j.part_pz, np.log(j.part_energy), np.log(j.part_pt), j.part_deta, j.part_dphi, part_logptrel, part_logerel, part_deltaR, np.tanh(j.part_d0val), np.tanh(j.part_dzval), part_type, part_pid) #下边的代码是要对第 j 个喷注中的所有粒子做循环，将每个粒子都 存成 ParticleBase，然后 append 到 particles里， #所以 partices 存储了 第 j 个喷注中所有粒子的信息 for num, (c, en, px, py, pz, lEn, lPt, eta, phi, ii, jj, kk, d0, dz, ptype, pid) in enumerate(bag): particles.append(ParticleBase( particle_id=num, part_charge=c, part_energy=en, part_px=px, part_py=py, part_pz=pz, log_energy=lEn, log_pt=lPt, part_deta=eta, part_dphi=phi, part_logptrel=ii, part_logerel=jj, part_deltaR=kk, part_d0=d0, part_dz=dz, particle_type=ptype, particle_pid=pid, jet_type=jet_type )) # add jets jet = 喷注， jet = Jet( jet_energy=j.jet_energy, jet_pt=j.jet_pt, jet_eta=j.jet_eta, particles=particles, label=label ) jet_list.append(jet) return cls(jets=jet_list, jets_type=jets_type) @classmethod def build_jetset_full(cls, root_file: Union[str, Path]) -> "JetSet": print(f"Building JetSet from {root_file}") if not 'bbbar' in Path(root_file).stem and not 'bb' in Path(root_file).stem: raise ValueError("Invalid file name, should contain 'bb' or 'bbbar'") jets_type = "bbbar" if "bbbar" in Path(root_file).stem else "bb" print(f"jet type: {jets_type}") with uproot.open(root_file) as f: tree = f["tree"] a = tree.arrays(library="pd") label = Path(root_file).stem.split('_')[0] jet_type = "bbbar jet" if "bbbar" in root_file.stem else "b jet" jet_list = [] for i, j in a.iterrows(): part_pt = np.array(j['part_pt']) jet_pt = np.array(j['jet_pt']) part_logptrel = np.log(np.divide(part_pt, jet_pt)) part_energy = np.array(j['part_energy']) jet_energy = np.array(j['jet_energy']) part_logerel = np.log(np.divide(part_energy, jet_energy)) part_deta = np.array(j['part_deta']) part_dphi = np.array(j['part_dphi']) part_deltaR = np.hypot(part_deta, part_dphi) assert len(j['part_pt']) == len(j['part_energy']) == len(j['part_deta']) particles = [] particle_list = ['part_isNeutralHadron', 'part_isPhoton', 'part_isElectron', 'part_isMuon', 'part_isPion', 'part_isChargedKaon', 'part_isProton'] part_type = [] part_pid = [] for pn in range(len(j['part_pt'])): jtmp = [] for t in particle_list: jtmp.append(j[t][pn]) tmp_type, tmp_pid = cls.jud_type(jtmp) part_type.append(tmp_type) part_pid.append(tmp_pid) for pn in range(len(part_type)): particles.append(ParticleBase( particle_id=i, part_charge=j['part_charge'][pn], part_energy=j['part_energy'][pn], part_px=j['part_px'][pn], part_py=j['part_py'][pn], part_pz=j['part_pz'][pn], log_energy=np.log(j['part_energy'][pn]), log_pt=np.log(j['part_pt'][pn]), part_deta=j['part_deta'][pn], part_dphi=j['part_dphi'][pn], part_logptrel=part_logptrel[pn], part_logerel=part_logerel[pn], part_deltaR=part_deltaR[pn], part_d0=np.tanh(j['part_d0val'][pn]), part_dz=np.tanh(j['part_dzval'][pn]), particle_type=part_type[pn], particle_pid=part_pid[pn], jet_type=jet_type )) jet = Jet( jet_energy=j['jet_energy'], jet_pt=j['jet_pt'], jet_eta=j['jet_eta'], particles=particles, label=label ) jet_list.append(jet) return cls(jets=jet_list, jets_type=jets_type) def save_to_txt(self, save_dir, selected_attrs: Optional[List[str]] = None, attr_sep: str = ',', part_sep: str = '|', prefix: str = ""): for file_counter, jet in enumerate(tqdm(self.jets)): filename = f"{self.jets_type}_{file_counter}_{prefix}.txt" filepath = os.path.join(save_dir, filename) jet_data = jet.particles_concatenated(selected_attrs, attr_sep, part_sep) with open(filepath, 'w') as file: file.write(jet_data) @click.command() @click.argument('root_dir', type=click.Path(exists=True)) @click.argument('save_dir', type=click.Path()) @click.option('--data-type', default='fast', type=click.Choice(['fast', 'full']), help='Type of ROOT data: fast or full.') @click.option('--attr-sep', default=',', help='Separator for attributes.') @click.option('--part-sep', default='|', help='Separator for particles.') @click.option('--selected-attrs', default='particle_id,part_charge,part_energy,part_px,part_py,part_pz,log_energy,log_pt,part_deta,part_dphi,part_logptrel,part_logerel,part_deltaR,part_d0,part_dz,particle_type,particle_pid', help='Comma-separated list of selected attributes.') def main(root_dir, save_dir, data_type, attr_sep, part_sep, selected_attrs): selected_attrs = selected_attrs.split(',') preprocess(root_dir, save_dir, data_type, selected_attrs=selected_attrs, attr_sep=attr_sep, part_sep=part_sep) def preprocess(root_dir, save_dir, data_type, selected_attrs: Optional[List[str]] = None, attr_sep: str = ',', part_sep: str = '|'): root_files = list(Path(root_dir).glob('*.root')) os.makedirs(save_dir, exist_ok=True) # process_file(root_files[0], save_dir, data_type, selected_attrs, attr_sep, part_sep, 0) # one process for test joblib.Parallel(n_jobs=-1)( joblib.delayed(process_file)(root_file, save_dir, data_type, selected_attrs, attr_sep, part_sep, i) for i, root_file in enumerate(root_files) ) def process_file(root_file, save_dir, data_type, selected_attrs, attr_sep, part_sep, process_id): prefix = f"process_{process_id}" if data_type == 'fast': jet_set = JetSet.build_jetset_fast(root_file) else: jet_set = JetSet.build_jetset_full(root_file) jet_set.save_to_txt(save_dir, selected_attrs, attr_sep, part_sep, prefix) if __name__ == "__main__": main()