import fastjet import os import copy import json import numpy as np import awkward as ak import torch.utils.data import time import pickle from collections import OrderedDict from functools import partial from concurrent.futures.thread import ThreadPoolExecutor from src.logger.logger import _logger, warn_once from src.data.tools import _pad, _repeat_pad, _clip, _pad_vector from src.data.fileio import _read_files from src.data.config import DataConfig, _md5 from src.data.preprocess import ( _apply_selection, _build_new_variables, _build_weights, AutoStandardizer, WeightMaker, ) from src.dataset.functions_data import to_tensor from src.layers.object_cond import calc_eta_phi from torch_scatter import scatter_sum from src.dataset.functions_graph import (create_graph, create_jets_outputs, create_jets_outputs_new, create_jets_outputs_Delphes, create_jets_outputs_Delphes2) from src.dataset.functions_data import Event, EventCollection, EventJets from src.utils.utils import CPU_Unpickler from src.dataset.functions_data import EventPFCands, concat_event_collection def get_pseudojets_fastjet(pfcands): pseudojets = [] for i in range(len(pfcands)): pseudojets.append(fastjet.PseudoJet(pfcands.pxyz[i, 0].item(), pfcands.pxyz[i, 1].item(), pfcands.pxyz[i, 2].item(), pfcands.E[i].item())) return pseudojets def _finalize_inputs(table, data_config): # transformation output = {} # transformation for k, params in data_config.preprocess_params.items(): if data_config._auto_standardization and params["center"] == "auto": raise ValueError("No valid standardization params for %s" % k) # if params["center"] is not None: # table[k] = (table[k] - params["center"]) * params["scale"] if params["length"] is not None: # if k == "hit_genlink": # pad_fn = partial(_pad_vector, value=-1) # table[k] = pad_fn(table[k]) # else: pad_fn = partial(_pad, value=0) table[k] = pad_fn(table[k], params["length"]) # stack variables for each input group for k, names in data_config.input_dicts.items(): if ( len(names) == 1 and data_config.preprocess_params[names[0]]["length"] is None ): output["_" + k] = ak.to_numpy(ak.values_astype(table[names[0]], "float32")) else: output["_" + k] = ak.to_numpy( np.stack( [ak.to_numpy(table[n]).astype("float32") for n in names], axis=1 ) ) # copy monitor variables for k in data_config.z_variables: if k not in output: output[k] = ak.to_numpy(table[k]) return output def _padlabel(table, _data_config): for k in _data_config.label_value: pad_fn = partial(_pad, value=0) table[k] = pad_fn(table[k], 400) return table def _preprocess(table, data_config, options): # apply selection table = _apply_selection( table, data_config.selection if options["training"] else data_config.test_time_selection, ) if len(table) == 0: return [] # table = _padlabel(table,data_config) # define new variables table = _build_new_variables(table, data_config.var_funcs) # else: indices = np.arange( len(table[table.fields[0]]) ) # np.arange(len(table[data_config.label_names[0]])) # shuffle if options["shuffle"]: np.random.shuffle(indices) # perform input variable standardization, clipping, padding and stacking table = _finalize_inputs(table, data_config) return table, indices def _load_next(data_config, filelist, load_range, options): table = _read_files( filelist, data_config.load_branches, load_range, treename=data_config.treename ) table, indices = _preprocess(table, data_config, options) return table, indices class _SimpleIter(object): r"""_SimpleIter Iterator object for ``SimpleIterDataset''. """ def __init__(self, **kwargs): # inherit all properties from SimpleIterDataset self.__dict__.update(**kwargs) self.iter_count = 0 # to raise StopIteration when dataset_cap is reached if "dataset_cap" in kwargs and kwargs["dataset_cap"] is not None: self.dataset_cap = kwargs["dataset_cap"] self._sampler_options["shuffle"] = False print("!!! Dataset_cap flag set, disabling shuffling") else: self.dataset_cap = None # executor to read files and run preprocessing asynchronously self.executor = ThreadPoolExecutor(max_workers=1) if self._async_load else None # init: prefetch holds table and indices for the next fetch self.prefetch = None self.table = None self.indices = [] self.cursor = 0 self._seed = None worker_info = torch.utils.data.get_worker_info() file_dict = self._init_file_dict.copy() if worker_info is not None: # in a worker process self._name += "_worker%d" % worker_info.id self._seed = worker_info.seed & 0xFFFFFFFF np.random.seed(self._seed) # split workload by files new_file_dict = {} for name, files in file_dict.items(): new_files = files[worker_info.id :: worker_info.num_workers] assert len(new_files) > 0 new_file_dict[name] = new_files file_dict = new_file_dict self.worker_file_dict = file_dict self.worker_filelist = sum(file_dict.values(), []) self.worker_info = worker_info self.restart() def restart(self): print("=== Restarting DataIter %s, seed=%s ===" % (self._name, self._seed)) # re-shuffle filelist and load range if for training filelist = self.worker_filelist.copy() if self._sampler_options["shuffle"]: np.random.shuffle(filelist) if self._file_fraction < 1: num_files = int(len(filelist) * self._file_fraction) filelist = filelist[:num_files] self.filelist = filelist if self._init_load_range_and_fraction is None: self.load_range = (0, 1) else: (start_pos, end_pos), load_frac = self._init_load_range_and_fraction interval = (end_pos - start_pos) * load_frac if self._sampler_options["shuffle"]: offset = np.random.uniform(start_pos, end_pos - interval) self.load_range = (offset, offset + interval) else: self.load_range = (start_pos, start_pos + interval) _logger.debug( "Init iter [%d], will load %d (out of %d*%s=%d) files with load_range=%s:\n%s", 0 if self.worker_info is None else self.worker_info.id, len(self.filelist), len(sum(self._init_file_dict.values(), [])), self._file_fraction, int(len(sum(self._init_file_dict.values(), [])) * self._file_fraction), str(self.load_range), ) # '\n'.join(self.filelist[: 3]) + '\n ... ' + self.filelist[-1],) _logger.info( "Restarted DataIter %s, load_range=%s, file_list:\n%s" % ( self._name, str(self.load_range), json.dumps(self.worker_file_dict, indent=2), ) ) # reset file fetching cursor self.ipos = 0 if self._fetch_by_files else self.load_range[0] # prefetch the first entry asynchronously self._try_get_next(init=True) def __next__(self): # print(self.ipos, self.cursor) graph_empty = True self.iter_count += 1 if self.dataset_cap is not None and self.iter_count > self.dataset_cap: raise StopIteration while graph_empty: if len(self.filelist) == 0: raise StopIteration try: i = self.indices[self.cursor] except IndexError: # case 1: first entry, `self.indices` is still empty # case 2: running out of entries, `self.indices` is not empty while True: if self._in_memory and len(self.indices) > 0: # only need to re-shuffle the indices, if this is not the first entry if self._sampler_options["shuffle"]: np.random.shuffle(self.indices) break if self.prefetch is None: # reaching the end as prefetch got nothing self.table = None if self._async_load: self.executor.shutdown(wait=False) raise StopIteration # get result from prefetch if self._async_load: self.table, self.indices = self.prefetch.result() else: self.table, self.indices = self.prefetch # try to load the next ones asynchronously self._try_get_next() # check if any entries are fetched (i.e., passing selection) -- if not, do another fetch if len(self.indices) > 0: break # reset cursor self.cursor = 0 i = self.indices[self.cursor] self.cursor += 1 data, graph_empty = self.get_data(i) return data def _try_get_next(self, init=False): end_of_list = ( self.ipos >= len(self.filelist) if self._fetch_by_files else self.ipos >= self.load_range[1] ) if end_of_list: if init: raise RuntimeError( "Nothing to load for worker %d" % 0 if self.worker_info is None else self.worker_info.id ) if self._infinity_mode and not self._in_memory: # infinity mode: re-start self.restart() return else: # finite mode: set prefetch to None, exit self.prefetch = None return if self._fetch_by_files: filelist = self.filelist[int(self.ipos) : int(self.ipos + self._fetch_step)] load_range = self.load_range else: filelist = self.filelist load_range = ( self.ipos, min(self.ipos + self._fetch_step, self.load_range[1]), ) # _logger.info('Start fetching next batch, len(filelist)=%d, load_range=%s'%(len(filelist), load_range)) if self._async_load: self.prefetch = self.executor.submit( _load_next, self._data_config, filelist, load_range, self._sampler_options, ) else: self.prefetch = _load_next( self._data_config, filelist, load_range, self._sampler_options ) self.ipos += self._fetch_step def get_data(self, i): # inputs X = {k: self.table["_" + k][i].copy() for k in self._data_config.input_names} if "EFlowPhoton" in X: return create_jets_outputs_Delphes(X), False elif "PFCands" in X: # v2 config return create_jets_outputs_Delphes2(X), False return create_jets_outputs_new(X), False class EventDatasetCollection(torch.utils.data.Dataset): def __init__(self, dir_list, args, aug_soft=False, aug_collinear=False, shuffle_seed=10): self.event_collections_dict = OrderedDict() if args: aug_soft = args.augment_soft_particles else: aug_soft=False for dir in dir_list: self.event_collections_dict[dir] = EventDataset.from_directory(dir, mmap=True, aug_soft=aug_soft or aug_soft, seed=0, aug_collinear=aug_collinear) self.n_events = sum([x.n_events for x in self.event_collections_dict.values()]) evt_idx = np.arange(0, self.n_events) # now shuffle this using the shuffle_seed and a separate random generator rng = np.random.default_rng(shuffle_seed) rng.shuffle(evt_idx) self.old_to_new_idx = evt_idx self.event_thresholds = [x.n_events for x in self.event_collections_dict.values()] self.event_thresholds = np.cumsum([0] + self.event_thresholds) self.dir_list = dir_list def __len__(self): return self.n_events def get_idx(self, i): assert i < self.n_events, "Index out of bounds: %d >= %d" % (i, self.n_events) for j in range(len(self.event_thresholds)-1): threshold = self.event_thresholds[j] if i >= threshold and i < self.event_thresholds[j+1]: #print("-------------", i, threshold, self.event_thresholds, j, self.dir_list[j]) return self.event_collections_dict[self.dir_list[j]][i - threshold] def getitem(self, i): return self.get_idx(i) def __iter__(self): for i in range(self.n_events): yield self.get_idx(self.old_to_new_idx[i]) def __getitem__(self, i): assert i < self.n_events, "Index out of bounds: %d >= %d" % (i, self.n_events) return self.get_idx(self.old_to_new_idx[i]) # A collection of EventDatasets. # You should use a sampler together with this, as by default it just concatenates the EventDatasets together! def get_batch_bounds(batch_idx): # batch_idx: tensor of format [0,0,0,0,1,1,1...] # returns tensor of format [0, 4, ...] print("Batch idx", batch_idx.shape, batch_idx[(batch_idx>3130) & (batch_idx < 3140)]) batches = sorted(batch_idx.unique().tolist()) skipped = [] for i in range(batch_idx.max().int().item()): if i not in batches: skipped.append(i) # reverse sort skipped skipped = sorted(skipped, reverse=True) result = torch.zeros(batch_idx.max().int().item() + 2 + len(skipped)) #for i, b in enumerate(batches): # assert i == b # result[i] = torch.where(batch_idx==b)[0].min() # result[i+1] = torch.where(batch_idx==b)[0].max() b_list = batch_idx.int().tolist() prev = -1 for i, b in enumerate(b_list): if b != prev: result[b] = i prev = b result[-1] = len(b_list) print("skipped", skipped) for s in skipped: if s == 0: result[s] = 0 else: result[s] = result[s+1] print("result", result.shape, result[3130:3140].tolist()) return result def filter_pfcands(pfcands): # filter the GenParticles so that dark matter particles are not present # dark matter particles are defined as those with abs(pdgId) > 10000 or pdgId between 50-60 # TODO: filter out high eta - temporarily this is done here, but it should be done in the ntuplizer in order to avoid big files mask = (torch.abs(pfcands.pid) < 10000) & ((torch.abs(pfcands.pid) < 50) | (torch.abs(pfcands.pid) > 60)) & (torch.abs(pfcands.eta) < 2.4) & (pfcands.pt > 0.5) #& (pfcands.pt > 0.5) pfcands.mask(mask) return pfcands class EventDataset(torch.utils.data.Dataset): @staticmethod def from_directory(dir, mmap=True, model_clusters_file=None, model_output_file=None, include_model_jets_unfiltered=False, fastjet_R=None, parton_level=False, gen_level=False, aug_soft=False, seed=0, aug_collinear=False, pt_jet_cutoff=100): result = {} for file in os.listdir(dir): if file == "metadata.pkl": metadata = pickle.load(open(os.path.join(dir, file), "rb")) else: print("File:", file) result[file.split(".")[0]] = np.load( os.path.join(dir, file), mmap_mode="r" if mmap else None ) dataset = EventDataset(result, metadata, model_clusters_file=model_clusters_file, model_output_file=model_output_file, include_model_jets_unfiltered=include_model_jets_unfiltered, fastjet_R=fastjet_R, parton_level=parton_level, gen_level=gen_level, aug_soft=aug_soft, seed=seed, aug_collinear=aug_collinear, pt_jet_cutoff=pt_jet_cutoff) return dataset def get_pfcands_key(self): pfcands_key = "pfcands" print("get_pfcands_key") if self.gen_level: return "final_gen_particles" if self.parton_level: return "final_parton_level_particles" if self.model_output is None: if self.gen_level: return "final_gen_particles" if self.parton_level: return "final_parton_level_particles" return pfcands_key # ignore for i in [0, 1, 2]: # try the first three if it fits start = {key: self.metadata[key + "_batch_idx"][i] for key in self.attrs} end = {key: self.metadata[key + "_batch_idx"][i + 1] for key in self.attrs} result = {key: self.events[key][start[key]:end[key]] for key in self.attrs} result = {key: EventCollection.deserialize(result[key], batch_number=None, cls=Event.evt_collections[key]) for key in self.attrs} if "final_parton_level_particles" in result: result["final_parton_level_particles"] = filter_pfcands(result["final_parton_level_particles"]) if "final_gen_particles" in result: result["final_gen_particles"] = filter_pfcands(result["final_gen_particles"]) event_filter_s, event_filter_e = self.model_output["event_idx_bounds"][i].int().item(), \ self.model_output["event_idx_bounds"][i + 1].int().item() diff = event_filter_e - event_filter_s if diff != len(result["pfcands"]): if diff == len(result["final_parton_level_particles"]): pfcands_key = "final_parton_level_particles" break if diff == len(result["final_gen_particles"]): pfcands_key = "final_gen_particles" break print("Found pfcands_key=%s" % pfcands_key) return pfcands_key def __init__(self, events, metadata, model_clusters_file=None, model_output_file=None, include_model_jets_unfiltered=False, fastjet_R=None, parton_level=False, gen_level=False, aug_soft=False, seed=0, aug_collinear=False, pt_jet_cutoff=100): # events: serialized events dict # metadata: dict with metadata self.events = events self.n_events = metadata["n_events"] self.attrs = metadata["attrs"] self.metadata = metadata self.include_model_jets_unfiltered = include_model_jets_unfiltered self.model_i = 0 self.parton_level = parton_level self.gen_level = gen_level self.augment_soft_particles = aug_soft self.aug_collinear = aug_collinear self.seed = seed self.pt_jet_cutoff = pt_jet_cutoff #self.pfcands_key = "pfcands" # set to final_parton_level_particles or final_gen_particles in case needed #for key in self.attrs: # self.evt_idx_to_batch_idx[key] = {} if model_output_file is not None: if type(model_output_file) == str: self.model_output = CPU_Unpickler(open(model_output_file, "rb")).load() else: self.model_output = model_output_file self.model_output["event_idx_bounds"] = get_batch_bounds(self.model_output["event_idx"]) self.n_events = self.model_output["event_idx"].max().int().item() # sometimes the last batch gets cut off, which causes problems if model_clusters_file is not None: self.model_clusters = to_tensor(pickle.load((open(model_clusters_file, "rb")))) else: self.model_clusters = self.model_output["model_cluster"] # model_output["batch_idx"] contains the batch index for each event. model_clusters is an array of the model labels for each event. else: self.model_output = None self.model_clusters = None if fastjet_R is not None: self.fastjet_jetdef = {r: fastjet.JetDefinition(fastjet.antikt_algorithm, r) for r in fastjet_R} ## fastjet_R is an array of radiuses for which to compute that self.pfcands_key = self.get_pfcands_key() def __len__(self): return self.n_events # def __next__(self): def add_model_output(self, model_output): if model_output is not None: if type(model_output) == str: self.model_output = CPU_Unpickler(open(model_output, "rb")).load() else: self.model_output = model_output self.model_output["event_idx_bounds"] = get_batch_bounds(self.model_output["event_idx"]) self.n_events = self.model_output["event_idx"].max().int().item() # sometimes the last batch gets cut off, which causes problems self.model_clusters = self.model_output["model_cluster"] # model_output["batch_idx"] contains the batch index for each event. model_clusters is an array of the model labels for each event. else: self.model_output = None self.model_clusters = None @staticmethod def pfcands_add_soft_particles(pfcands, n_soft, random_generator, add_original_particle_mapping=False): # augment the dataset with soft particles eta_bounds = [-2.4, 2.4] phi_bounds = [-3.14, 3.14] #pt_bounds = [0.02, 0.5] # choose random eta and phi # use the random generator for eta, phi eta = random_generator.uniform(eta_bounds[0], eta_bounds[1], n_soft).astype(np.double) phi = random_generator.uniform(phi_bounds[0], phi_bounds[1], n_soft).astype(np.double) #pt = random_generator.uniform(pt_bounds[0], pt_bounds[1], n_soft).astype(np.double) pt = np.ones(n_soft).astype(np.double) * 1e-2 charge = np.zeros(n_soft).astype(np.double) pid = np.zeros(n_soft).astype(np.double) mass = np.zeros(n_soft).astype(np.double) if hasattr(pfcands, "status"): status = np.zeros(n_soft) soft_pfcands = EventPFCands(pt, eta, phi, mass, charge, pid, pf_cand_jet_idx=-1 * torch.ones(n_soft), status=status) else: soft_pfcands = EventPFCands(pt, eta, phi, mass, charge, pid, pf_cand_jet_idx=-1*torch.ones(n_soft)) soft_pfcands.original_particle_mapping = torch.tensor([-1] * len(soft_pfcands)) pfcandsc = copy.deepcopy(pfcands) pfcandsc.original_particle_mapping = torch.arange(len(pfcands)) pfcandsc = concat_event_collection([pfcandsc, soft_pfcands], nobatch=1) if not add_original_particle_mapping: pfcandsc.original_particle_mapping = torch.arange(len(pfcandsc)) # For now, ignore the soft particles #print("Original PM:", pfcandsc.original_particle_mapping.max()) return pfcandsc @staticmethod def pfcands_split_particles(pfcands, random_generator): # Augment the dataset by spliting the harder particles # 5 highest pt particles k = min(5, len(pfcands)) highest_pt_idx = torch.topk(pfcands.pt, k)[1] weights = pfcands.pt[highest_pt_idx] # Pick a random particle to split according to weights n_to_split = random_generator.randint(0, k) #idx = random_generator.choice(highest_pt_idx, p=weights / weights.sum()) indices = highest_pt_idx[:n_to_split] pfcandsc = copy.deepcopy(pfcands) pfcandsc.original_particle_mapping = torch.arange(len(pfcands)) # assert that indices are all lower than len(pfcands) if not torch.all(indices < len(pfcands)): print("Indices:", indices) print("PFCands:", pfcands.pt) print("PFCands len:", len(pfcands.pt)) raise ValueError("Indices are out of bounds") for idx in indices: split_into = random_generator.randint(2, 5) # split the particle into eta = pfcands.eta[idx] phi = pfcands.phi[idx] pt = pfcands.pt[idx] / split_into charge = pfcands.charge[idx] mass = 0 pid = pfcands.pid[idx] colinear_pfcands = EventPFCands(pt=[pt], eta=[eta], phi=[phi], mass=[mass], charge=[charge], pid=[pid], pf_cand_jet_idx=[pfcands.pf_cand_jet_idx[idx]], original_particle_mapping=[idx]) #pfcandsc.original_particle_mapping[idx] = idx pfcandsc.pt[idx] = pt for _ in range(split_into-1): pfcandsc = concat_event_collection([pfcandsc, colinear_pfcands], nobatch=1) if pfcandsc.original_particle_mapping.max() >= len(pfcands): #print("Original PM:", pfcandsc.original_particle_mapping.max(), "len pfcands", len(pfcands)) raise ValueError("Original particle mapping is out of bounds") return pfcandsc def get_idx(self, i): #print("Getting idx", i) start = {key: self.metadata[key + "_batch_idx"][i] for key in self.attrs} end = {key: self.metadata[key + "_batch_idx"][i + 1] for key in self.attrs} result = {key: self.events[key][start[key]:end[key]] for key in self.attrs} result = {key: EventCollection.deserialize(result[key], batch_number=None, cls=Event.evt_collections[key]) for key in self.attrs} result["pfcands"] = filter_pfcands(result["pfcands"]) if "final_parton_level_particles" in result: #print("i=", i) #print("BEFORE:", len(result["final_parton_level_particles"])) result["final_parton_level_particles"] = filter_pfcands(result["final_parton_level_particles"]) #print("AFTER:", len(result["final_parton_level_particles"])) #print("------") if "final_gen_particles" in result: result["final_gen_particles"] = filter_pfcands(result["final_gen_particles"]) ## augment pfcands here if self.augment_soft_particles: random_generator = np.random.RandomState(seed=i + self.seed) #n_soft = int(random_generator.uniform(10, 1000)) n_soft = 500 #n_soft = 1000 result["pfcands"] = EventDataset.pfcands_add_soft_particles(result["pfcands"], n_soft, random_generator) if "final_parton_level_particles" in result: result["final_parton_level_particles"] = EventDataset.pfcands_add_soft_particles(result["final_parton_level_particles"], n_soft, random_generator) # Also augment parton-level event for testing if "final_gen_particles" in result: result["final_gen_particles"] = EventDataset.pfcands_add_soft_particles(result["final_gen_particles"], n_soft, random_generator) else: result["pfcands"].original_particle_mapping = torch.arange(len(result["pfcands"].pt)) if self.aug_collinear: random_generator = np.random.RandomState(seed=i + self.seed) if i % 2: # Every second one: result["pfcands"] = EventDataset.pfcands_split_particles(result["pfcands"], random_generator) if "final_parton_level_particles" in result: result["final_parton_level_particles"] = EventDataset.pfcands_split_particles( result["final_parton_level_particles"], random_generator ) # Also augment parton-level event for testing if "final_gen_particles" in result: result["final_gen_particles"] = EventDataset.pfcands_split_particles(result["final_gen_particles"], random_generator) else: n_soft = 500 result["pfcands"] = EventDataset.pfcands_add_soft_particles(result["pfcands"], n_soft, random_generator, add_original_particle_mapping=True) if "final_parton_level_particles" in result: result["final_parton_level_particles"] = EventDataset.pfcands_add_soft_particles( result["final_parton_level_particles"], n_soft, random_generator, add_original_particle_mapping=True ) # Also augment parton-level event for testing if "final_gen_particles" in result: result["final_gen_particles"] = EventDataset.pfcands_add_soft_particles( result["final_gen_particles"], n_soft, random_generator, add_original_particle_mapping=True ) if self.model_output is not None: #if "final_parton_level_particles" in result and len(result["final_parton_level_particles"]) == 0: # print("!!") # return None result["model_jets"], bc_scores_pfcands, bc_labels_pfcands = self.get_model_jets(i, pfcands=result[self.pfcands_key], include_target=1, dq=result["matrix_element_gen_particles"]) result[self.pfcands_key].bc_scores_pfcands = bc_scores_pfcands result[self.pfcands_key].bc_labels_pfcands = bc_labels_pfcands if self.include_model_jets_unfiltered: result["model_jets_unfiltered"], _, _ = self.get_model_jets(i, pfcands=result[self.pfcands_key], filter=False) if hasattr(self, "fastjet_jetdef") and self.fastjet_jetdef is not None: if self.gen_level: result["fastjet_jets"] = {key: EventDataset.get_fastjet_jets(result, self.fastjet_jetdef[key], key="final_gen_particles", pt_cutoff=self.pt_jet_cutoff) for key in self.fastjet_jetdef} elif self.parton_level: result["fastjet_jets"] = {key: EventDataset.get_fastjet_jets(result, self.fastjet_jetdef[key], key="final_parton_level_particles", pt_cutoff=self.pt_jet_cutoff) for key in self.fastjet_jetdef} else: result["fastjet_jets"] = {key: EventDataset.get_fastjet_jets(result, self.fastjet_jetdef[key], key="pfcands", pt_cutoff=self.pt_jet_cutoff) for key in self.fastjet_jetdef} if "genjets" in result: result["genjets"] = EventDataset.mask_jets(result["genjets"]) evt = Event(**result) assert evt.pfcands.original_particle_mapping.max() < len(evt.pfcands.pt), "Original particle mapping is out of bounds: " + str(evt.original_particle_mapping.max()) + " >= " + str(len(evt.pfcands.pt)) return evt @staticmethod def get_target_obj_score(clusters_eta, clusters_phi, clusters_pt, event_idx_clusters, dq_eta, dq_phi, dq_event_idx): # return the target scores for each cluster (reteurns list of 1's and 0's) # dq_coords: list of [eta, phi] for each dark quark # dq_event_idx: list of event_idx for each dark quarks target = [] for event in event_idx_clusters.unique(): filt = event_idx_clusters == event clusters = torch.stack([clusters_eta[filt], clusters_phi[filt], clusters_pt[filt]], dim=1) dq_coords_event = torch.stack([dq_eta[dq_event_idx == event], dq_phi[dq_event_idx == event]], dim=1) dist_matrix = torch.cdist( dq_coords_event, clusters[:, :2].to(dq_coords_event.device), p=2 ).T if len(dist_matrix) == 0: target.append(torch.zeros(len(clusters)).int().to(dist_matrix.device)) continue closest_quark_dist, closest_quark_idx = dist_matrix.min(dim=1) closest_quark_idx[closest_quark_dist > 0.8] = -1 target.append((closest_quark_idx != -1).float()) if len(target): return torch.cat(target).flatten() return torch.tensor([]) @staticmethod def mask_jets(jets, cutoff=100): mask = jets.pt >= cutoff return EventJets(jets.pt[mask], jets.eta[mask], jets.phi[mask], jets.mass[mask]) @staticmethod def get_model_jets_static(i, pfcands, model_output, model_clusters): event_filter_s, event_filter_e = model_output["event_idx_bounds"][i].int().item(), model_output["event_idx_bounds"][i + 1].int().item() pfcands_pt = pfcands.pt pfcands_pxyz = pfcands.pxyz pfcands_E = pfcands.E #assert len(pfcands_pt) == event_filter_e - event_filter_s, "Error!, len(pfcands_pt)==%d, event_filter_e-event_filter_s=%d" % (len(pfcands_pt), event_filter_e - event_filter_s) if not len(pfcands_pt) == event_filter_e - event_filter_s: return None # jets_pt = scatter_sum(to_tensor(pfcands_pt), self.model_clusters[event_filter] + 1, dim=0)[1:] jets_pxyz = scatter_sum(to_tensor(pfcands_pxyz), model_clusters[event_filter_s:event_filter_e] + 1, dim=0)[1:] jets_pt = torch.norm(jets_pxyz[:, :2], p=2, dim=-1) jets_eta, jets_phi = calc_eta_phi(jets_pxyz, False) # jets_mass = torch.zeros_like(jets_eta) jets_E = scatter_sum(to_tensor(pfcands_E), model_clusters[event_filter_s:event_filter_e] + 1, dim=0)[1:] jets_mass = torch.sqrt(jets_E ** 2 - jets_pxyz.norm(dim=-1) ** 2) cluster_labels = model_clusters[event_filter_s:event_filter_e] bc_scores = model_output["pred"][event_filter_s:event_filter_e, -1] cutoff = 100 mask = jets_pt >= cutoff return EventJets(jets_pt[mask], jets_eta[mask], jets_phi[mask], jets_mass[mask]) @staticmethod def get_jets_fastjets_raw_with_assignment(pfcands, jetdef, pt_cutoff=100): pt = [] eta = [] phis = [] mass = [] particle_to_jet = {} # this will map particle_idx -> jet_idx array = get_pseudojets_fastjet(pfcands) for idx, pseudojet in enumerate(array): pseudojet.set_user_index(idx) cluster = fastjet.ClusterSequence(array, jetdef) inc_jets = cluster.inclusive_jets() jet_idx = 0 for elem in inc_jets: if elem.pt() < pt_cutoff: continue # print("pt:", elem.pt(), "eta:", elem.rap(), "phi:", elem.phi())ž pt.append(elem.pt()) eta.append(elem.rap()) phi = elem.phi() if phi > np.pi: phi -= 2 * np.pi phis.append(phi) mass.append(elem.m()) # Get constituents of this jet constituents = cluster.constituents(elem) for constituent in constituents: particle_idx = constituent.user_index() particle_to_jet[particle_idx] = jet_idx jet_idx += 1 return pt, eta, phis, mass, particle_to_jet @staticmethod def get_jets_fastjets_raw(pfcands, jetdef, pt_cutoff=100): pt = [] eta = [] phis = [] mass = [] array = get_pseudojets_fastjet(pfcands) cluster = fastjet.ClusterSequence(array, jetdef) inc_jets = cluster.inclusive_jets() for elem in inc_jets: if elem.pt() < pt_cutoff: continue # print("pt:", elem.pt(), "eta:", elem.rap(), "phi:", elem.phi())ž pt.append(elem.pt()) eta.append(elem.rap()) phi = elem.phi() if phi > np.pi: phi -= 2 * np.pi phis.append(phi) mass.append(elem.m()) return pt, eta, phis, mass @staticmethod def get_fastjet_jets_with_assignment(event, jetdef, key="pfcands", pt_cutoff=100): if type(event) == dict: k = event[key] else: k = getattr(event, key) pt, eta, phi, m, assignment = EventDataset.get_jets_fastjets_raw_with_assignment(k, jetdef, pt_cutoff=pt_cutoff) return EventJets(torch.tensor(pt), torch.tensor(eta), torch.tensor(phi), torch.tensor(m)), assignment @staticmethod def get_fastjet_jets(event, jetdef, key="pfcands", pt_cutoff=100): if type(event) == dict: k = event[key] else: k = getattr(event, key) pt, eta, phi, m = EventDataset.get_jets_fastjets_raw(k, jetdef, pt_cutoff=pt_cutoff) return EventJets(torch.tensor(pt), torch.tensor(eta), torch.tensor(phi), torch.tensor(m)) def get_model_jets(self, i, pfcands, filter=True, dq=None, include_target=False): event_filter_s, event_filter_e = self.model_output["event_idx_bounds"][i].int().item(), self.model_output["event_idx_bounds"][i+1].int().item() pfcands_pt = pfcands.pt pfcands_pxyz = pfcands.pxyz pfcands_E = pfcands.E obj_score = None #print("Len pfcands_pt", len(pfcands_pt), "event_filter_e", event_filter_e, "event_filter_s", event_filter_s) if len(pfcands_pt) == 0: return EventJets(torch.tensor([]), torch.tensor([]), torch.tensor([]) ,torch.tensor([])), None, None assert len(pfcands_pt) == event_filter_e - event_filter_s, "Error! filter={} len(pfcands_pt)={} event_filter_e={} event_filter_s={}".format(filter, len(pfcands_pt), event_filter_e, event_filter_s) #jets_pt = scatter_sum(to_tensor(pfcands_pt), self.model_clusters[event_filter] + 1, dim=0)[1:] jets_pxyz = scatter_sum(to_tensor(pfcands_pxyz), self.model_clusters[event_filter_s:event_filter_e] + 1, dim=0)[1:] jets_pt = torch.norm(jets_pxyz[:, :2], p=2, dim=-1) jets_eta, jets_phi = calc_eta_phi(jets_pxyz, False) #jets_mass = torch.zeros_like(jets_eta) jets_E = scatter_sum(to_tensor(pfcands_E), self.model_clusters[event_filter_s:event_filter_e] + 1, dim=0)[1:] jets_mass = torch.sqrt(jets_E**2 - jets_pxyz.norm(dim=-1)**2) cluster_labels = self.model_clusters[event_filter_s:event_filter_e] bc_scores = self.model_output["pred"][event_filter_s:event_filter_e, -1] if "obj_score_pred" in self.model_output and not torch.is_tensor(self.model_output["obj_score_pred"]): self.model_output["obj_score_pred"] = torch.cat(self.model_output["obj_score_pred"]) print("Concatenated obj_score_pred") target_obj_score = None if filter: cutoff = self.pt_jet_cutoff mask = jets_pt >= cutoff if "obj_score_pred" in self.model_output: obj_score = self.model_output["obj_score_pred"][(self.model_output["event_clusters_idx"] == i)] #print("Jets pt", jets_pt, "obj score", obj_score) assert len(obj_score) == len(jets_pt), "Error! len(obj_score)=%d, len(jets_pt)=%d" % ( len(obj_score), len(jets_pt)) if include_target: target_obj_score = EventDataset.get_target_obj_score(jets_eta, jets_phi, jets_pt, torch.zeros(jets_pt.size(0)), dq.eta, dq.phi, torch.zeros(dq.eta.size(0))) else: mask = torch.ones_like(jets_pt, dtype=torch.bool) if obj_score is not None: obj_score = obj_score[mask] assert len(jets_pt[mask]) == len(obj_score), "Error! len(jets_pt[mask])=%d, len(obj_score)=%d" % (len(jets_pt[mask]), len(obj_score)) if target_obj_score is not None: target_obj_score = target_obj_score[mask] assert len(jets_pt[mask]) == len(target_obj_score), "Error! len(jets_pt[mask])=%d, len(obj_score)=%d" % (len(jets_pt[mask]), len(obj_score)) return EventJets(jets_pt[mask], jets_eta[mask], jets_phi[mask], jets_mass[mask], obj_score=obj_score, target_obj_score=target_obj_score), bc_scores, cluster_labels def get_iter(self): self.i = 0 while self.i < self.n_events: yield self.get_idx(self.i) self.i += 1 def __iter__(self): return self.get_iter() def __getitem__(self, i): assert i < self.n_events, "Index out of bounds: %d >= %d" % (i, self.n_events) return self.get_idx(i) class SimpleIterDataset(torch.utils.data.IterableDataset): r"""Base IterableDataset. Handles dataloading. Arguments: file_dict (dict): dictionary of lists of files to be loaded. data_config_file (str): YAML file containing data format information. for_training (bool): flag indicating whether the dataset is used for training or testing. When set to ``True``, will enable shuffling and sampling-based reweighting. When set to ``False``, will disable shuffling and reweighting, but will load the observer variables. load_range_and_fraction (tuple of tuples, ``((start_pos, end_pos), load_frac)``): fractional range of events to load from each file. E.g., setting load_range_and_fraction=((0, 0.8), 0.5) will randomly load 50% out of the first 80% events from each file (so load 50%*80% = 40% of the file). fetch_by_files (bool): flag to control how events are retrieved each time we fetch data from disk. When set to ``True``, will read only a small number (set by ``fetch_step``) of files each time, but load all the events in these files. When set to ``False``, will read from all input files, but load only a small fraction (set by ``fetch_step``) of events each time. Default is ``False``, which results in a more uniform sample distribution but reduces the data loading speed. fetch_step (float or int): fraction of events (when ``fetch_by_files=False``) or number of files (when ``fetch_by_files=True``) to load each time we fetch data from disk. Event shuffling and reweighting (sampling) is performed each time after we fetch data. So set this to a large enough value to avoid getting an imbalanced minibatch (due to reweighting/sampling), especially when ``fetch_by_files`` set to ``True``. Will load all events (files) at once if set to non-positive value. file_fraction (float): fraction of files to load. """ def __init__( self, file_dict, data_config_file, for_training=True, load_range_and_fraction=None, extra_selection=None, fetch_by_files=False, fetch_step=0.01, file_fraction=1, remake_weights=False, up_sample=True, weight_scale=1, max_resample=10, async_load=True, infinity_mode=False, in_memory=False, name="", laplace=False, edges=False, diffs=False, dataset_cap=None, n_noise=0, synthetic=False, synthetic_npart_min=2, synthetic_npart_max=5, jets=False, ): self._iters = {} if infinity_mode or in_memory else None _init_args = set(self.__dict__.keys()) self._init_file_dict = file_dict self._init_load_range_and_fraction = load_range_and_fraction self._fetch_by_files = fetch_by_files self._fetch_step = fetch_step self._file_fraction = file_fraction self._async_load = async_load self._infinity_mode = infinity_mode self._in_memory = in_memory self._name = name self.laplace = laplace self.edges = edges self.diffs = diffs self.synthetic = synthetic self.synthetic_npart_min = synthetic_npart_min self.synthetic_npart_max = synthetic_npart_max self.dataset_cap = dataset_cap # used to cap the dataset to some fixed number of events - used for debugging purposes self.n_noise = n_noise self.jets = jets # ==== sampling parameters ==== self._sampler_options = { "up_sample": up_sample, "weight_scale": weight_scale, "max_resample": max_resample, } if for_training: self._sampler_options.update(training=True, shuffle=False, reweight=True) else: self._sampler_options.update(training=False, shuffle=False, reweight=False) # discover auto-generated reweight file if ".auto.yaml" in data_config_file: data_config_autogen_file = data_config_file else: data_config_md5 = _md5(data_config_file) data_config_autogen_file = data_config_file.replace( ".yaml", ".%s.auto.yaml" % data_config_md5 ) if os.path.exists(data_config_autogen_file): data_config_file = data_config_autogen_file _logger.info( "Found file %s w/ auto-generated preprocessing information, will use that instead!" % data_config_file ) # load data config (w/ observers now -- so they will be included in the auto-generated yaml) self._data_config = DataConfig.load(data_config_file) if for_training: # produce variable standardization info if needed if self._data_config._missing_standardization_info: s = AutoStandardizer(file_dict, self._data_config) self._data_config = s.produce(data_config_autogen_file) # produce reweight info if needed # if self._sampler_options['reweight'] and self._data_config.weight_name and not self._data_config.use_precomputed_weights: # if remake_weights or self._data_config.reweight_hists is None: # w = WeightMaker(file_dict, self._data_config) # self._data_config = w.produce(data_config_autogen_file) # reload data_config w/o observers for training if ( os.path.exists(data_config_autogen_file) and data_config_file != data_config_autogen_file ): data_config_file = data_config_autogen_file _logger.info( "Found file %s w/ auto-generated preprocessing information, will use that instead!" % data_config_file ) self._data_config = DataConfig.load( data_config_file, load_observers=False, extra_selection=extra_selection ) else: self._data_config = DataConfig.load( data_config_file, load_reweight_info=False, extra_test_selection=extra_selection, ) # Derive all variables added to self.__dict__ self._init_args = set(self.__dict__.keys()) - _init_args @property def config(self): return self._data_config def __iter__(self): if self._iters is None: kwargs = {k: copy.deepcopy(self.__dict__[k]) for k in self._init_args} return _SimpleIter(**kwargs) else: worker_info = torch.utils.data.get_worker_info() worker_id = worker_info.id if worker_info is not None else 0 try: return self._iters[worker_id] except KeyError: kwargs = {k: copy.deepcopy(self.__dict__[k]) for k in self._init_args} self._iters[worker_id] = _SimpleIter(**kwargs) return self._iters[worker_id]