diff --git a/grand/dataio/root_trees.py b/grand/dataio/root_trees.py index 74d46181..0398e82c 100644 --- a/grand/dataio/root_trees.py +++ b/grand/dataio/root_trees.py @@ -1398,6 +1398,11 @@ def get_list_of_all_used_dus(self): else: return None + def get_dus_indices_in_run(self, trun): + """Gets an array of the indices of DUs of the current event in the TRun tree""" + + return np.nonzero(np.isin(np.asarray(trun.du_id), np.asarray(self.du_id)))[0] + ## A class wrapping around a TTree holding values common for the whole run @dataclass diff --git a/grand/sim/efield2voltage.py b/grand/sim/efield2voltage.py index 27dfaa8c..6a136981 100644 --- a/grand/sim/efield2voltage.py +++ b/grand/sim/efield2voltage.py @@ -117,10 +117,12 @@ def get_event(self, event_idx=None, event_number=None, run_number=None): # stack efield traces trace_shape = np.asarray(self.events.trace).shape # (nb_du, 3, tbins of a trace) self.du_id = np.asarray(self.events.du_id) # used for printing info and saving in voltage tree. + self.event_dus_indices = self.events.get_dus_indices_in_run(self.run) self.nb_du = trace_shape[0] self.sig_size = trace_shape[-1] self.traces = np.asarray(self.events.trace, dtype=np.float32) # x,y,z components are stored in events.trace. shape (nb_du, 3, tbins) - self.du_pos = np.asarray(self.run.du_xyz) # (nb_du, 3) antenna position wrt local grand coordinate + # self.du_pos = np.asarray(self.run.du_xyz) # (nb_du, 3) antenna position wrt local grand coordinate + self.du_pos = np.asarray(self.run.du_xyz)[self.event_dus_indices] # (nb_du, 3) antenna position wrt local grand coordinate # shower information like theta, phi, xmax etc for one event. shower = ShowerEvent() @@ -129,7 +131,8 @@ def get_event(self, event_idx=None, event_number=None, run_number=None): self.evt_shower = shower # Note that 'shower' is an instance of 'self.shower' for one event. logger.info(f"shower origin in Geodetic: {self.run.origin_geoid}") - self.dt_ns = np.asarray(self.run.t_bin_size) # sampling time in ns, sampling freq = 1e9/dt_ns. #MATIAS: Why cast this to an array if it is a constant? + # self.dt_ns = np.asarray(self.run.t_bin_size) # sampling time in ns, sampling freq = 1e9/dt_ns. #MATIAS: Why cast this to an array if it is a constant? + self.dt_ns = np.asarray(self.run.t_bin_size)[self.event_dus_indices] # sampling time in ns, sampling freq = 1e9/dt_ns. #MATIAS: Why cast this to an array if it is a constant? self.f_samp_mhz = 1e3/self.dt_ns # MHz #MATIAS: this gets casted too! # comupte time samples in ns for all antennas in event with index event_idx. self.time_samples = self.get_time_samples() # t_samples.shape = (nb_du, self.sig_size) diff --git a/sim2root/Common/sim2root.py b/sim2root/Common/sim2root.py index 12f4f0b3..acdddccd 100755 --- a/sim2root/Common/sim2root.py +++ b/sim2root/Common/sim2root.py @@ -14,11 +14,14 @@ #ToDo:latitude,longitude and altitude are available in ZHAireS .sry file, and could be added to the RawRoot file. Site too. # Command line argument parsing clparser = argparse.ArgumentParser(description="Convert simulation data in rawroot format into GRANDROOT format") -clparser.add_argument("filename", nargs='+', help="ROOT file containing GRANDRaw data TTrees") +clparser.add_argument("file_dir_name", nargs='+', help="ROOT files containing GRANDRaw data TTrees or a directory with GRANDraw files") clparser.add_argument("-o", "--output_parent_directory", help="Output parent directory", default="") clparser.add_argument("-fo", "--forced_output_directory", help="Force this option as the output directory", default=None) clparser.add_argument("-s", "--site_name", help="The name of the site", default="nosite") -clparser.add_argument("-d", "--sim_date", help="The date of the simulation", default="19000101") +clparser.add_argument("-d", "--sim_date", help="The date of the simulation", default=None) +clparser.add_argument("-t", "--sim_time", help="The time of the simulation", default=None) +# clparser.add_argument("-d", "--sim_date", help="The date of the simulation", default="19000101") +# clparser.add_argument("-t", "--sim_time", help="The time of the simulation", default="000000") clparser.add_argument("-e", "--extra", help="Extra information to store in the directory name", default="") clparser.add_argument("-av", "--analysis_level", help="Analysis level of the data", default=0, type=int) # clparser.add_argument("-se", "--serial", help="Serial number of the simulation", default=0) @@ -50,20 +53,20 @@ def main(): if clargs.start_event is not None: ext_event_number = int(clargs.start_event) - # Create the appropriate output directory - if clargs.forced_output_directory is None: - out_dir_name = form_directory_name(clargs) - print("Storing files in directory ", out_dir_name) - out_dir_name.mkdir() - # If another directory was forced as the output directory, create it - else: - out_dir_name = Path(clargs.output_parent_directory, clargs.forced_output_directory) - out_dir_name.mkdir(exist_ok=True) - start_event_number = 0 + # Namespace for holding output trees + gt = SimpleNamespace() + + # Check if a directory was given as input + if Path(clargs.file_dir_name[0]).is_dir(): + file_list = glob.glob(clargs.file_dir_name[0]+"/*.RawRoot") + else: + file_list = clargs.file_dir_name + # Loop through the files specified on command line - for file_num, filename in enumerate(clargs.filename): + # for file_num, filename in enumerate(clargs.filename): + for file_num, filename in enumerate(file_list): # Output filename for GRAND Trees # if clargs.output_filename is None: @@ -76,15 +79,6 @@ def main(): trawefield = RawTrees.RawEfieldTree(filename) trawmeta = RawTrees.RawMetaTree(filename) - # Create appropriate GRANDROOT trees in temporary file names (event range not known until the end of the loop) - gt = SimpleNamespace() - gt.trun = TRun((out_dir_name/"trun.root").as_posix()) - gt.trunshowersim = TRunShowerSim((out_dir_name/"trunshowersim.root").as_posix()) - gt.trunefieldsim = TRunEfieldSim((out_dir_name/"trunefieldsim.root").as_posix()) - gt.tshower = TShower((out_dir_name/"tshower.root").as_posix()) - gt.tshowersim = TShowerSim((out_dir_name/"tshowersim.root").as_posix()) - gt.tefield = TEfield((out_dir_name/"tefield.root").as_posix()) - # Loop through entries - assuming same number of entries in each tree # ToDo: this should be a tree iterator through one tree and getting the other through friends. Need to make friends working... nentries = trawshower.get_entries() @@ -95,6 +89,13 @@ def main(): # If the first entry or (run number enforced and first file), fill the run trees if (i==0 and ext_run_number is None) or (ext_run_number is not None and file_num==0): + + # Overwrite the run number if specified on command line + run_number = ext_run_number if ext_run_number is not None else gt.trun.run_number + + # Init output trees in the proper directory + out_dir_name = init_trees(clargs, trawshower.unix_date, run_number, gt) + # Convert the RawShower entries rawshower2grandrootrun(trawshower, gt) # Convert the RawEfield entries @@ -105,17 +106,15 @@ def main(): # Set the origin geoid gt.trun.origin_geoid = get_origin_geoid(clargs) - # Overwrite the run number if specified on command line - if ext_run_number is not None: - gt.trun.run_number = ext_run_number - gt.trunshowersim.run_number = ext_run_number - gt.trunefieldsim.run_number = ext_run_number + gt.trun.run_number = run_number + gt.trunshowersim.run_number = run_number + gt.trunefieldsim.run_number = run_number # Fill the run trees and write - gt.trun.fill() + # gt.trun.fill() gt.trunshowersim.fill() gt.trunefieldsim.fill() - gt.trun.write() + # gt.trun.write() gt.trunshowersim.write() gt.trunefieldsim.write() @@ -150,6 +149,15 @@ def main(): gt.tshowersim.fill() gt.tefield.fill() + # For the first file, get all the file's events du ids and pos + if file_num==0: + du_ids, du_xyzs = get_tree_du_id_and_xyz(trawefield) + # For other files, append du ids and pos to the ones already retrieved + else: + tdu_ids, tdu_xyzs = get_tree_du_id_and_xyz(trawefield) + du_ids = np.append(du_ids, tdu_ids) + du_xyzs = np.vstack([du_xyzs, tdu_xyzs]) + # Write the event trees gt.tshower.write() gt.tshowersim.write() @@ -160,11 +168,66 @@ def main(): if ext_event_number is not None: ext_event_number += 1 + # Fill the trun with antenna positions and ids from ALL the events + # ToDo: this should be done with TChain in one loop over all the files... maybe (which would be faster?) + + # Get indices of the unique du_ids + unique_dus_idx = np.unique(du_ids, return_index=True)[1] + # Leave only the unique du_ids + du_ids = du_ids[unique_dus_idx] + # Sort the DUs + sorted_idx = np.argsort(du_ids) + du_ids = du_ids[sorted_idx] + # Stack x/y/z together and leave only the ones for unique du_ids, sort + du_xyzs = du_xyzs[unique_dus_idx][sorted_idx] + + # Assign the du ids and positions to the trun tree + gt.trun.du_id = du_ids + gt.trun.du_xyz = du_xyzs + + # ToDo: shouldn't this and above be created for every DU in sims? + gt.trun.t_bin_size = [trawefield.t_bin_size]*len(du_ids) #Matias Question: Why is this being mutiplied here? + + # Fill and write the TRun + gt.trun.fill() + gt.trun.write() + + # Rename the created files to appropriate names end_event_number = gt.tshower.event_number print("Renaming files to proper file names") rename_files(clargs, out_dir_name, start_event_number, end_event_number) +# Initialise output trees and their directory +def init_trees(clargs, unix_date, run_number, gt): + + # Use date/time from command line argument if specified, otherwise the unix time + date, time = datetime.datetime.utcfromtimestamp(unix_date).strftime('%Y%m%d_%H%M%S').split("_") + if clargs.sim_date is not None: + date = clargs.sim_date + if clargs.sim_time is not None: + time = clargs.sim_time + + # Create the appropriate output directory + if clargs.forced_output_directory is None: + out_dir_name = form_directory_name(clargs, date, time, run_number) + print("Storing files in directory ", out_dir_name) + out_dir_name.mkdir() + # If another directory was forced as the output directory, create it + else: + out_dir_name = Path(clargs.output_parent_directory, clargs.forced_output_directory) + out_dir_name.mkdir(exist_ok=True) + + # Create appropriate GRANDROOT trees in temporary file names (event range not known until the end of the loop) + gt.trun = TRun((out_dir_name / "run.root").as_posix()) + gt.trunshowersim = TRunShowerSim((out_dir_name / "runshowersim.root").as_posix()) + gt.trunefieldsim = TRunEfieldSim((out_dir_name / "runefieldsim.root").as_posix()) + gt.tshower = TShower((out_dir_name / "shower.root").as_posix()) + gt.tshowersim = TShowerSim((out_dir_name / "showersim.root").as_posix()) + gt.tefield = TEfield((out_dir_name / "efield.root").as_posix()) + + return out_dir_name + # Convert the RawShowerTree first entry to run values def rawshower2grandrootrun(trawshower, gt): @@ -223,6 +286,15 @@ def rawefield2grandrootrun(trawefield, gt): gt.trunefieldsim.refractivity_model = trawefield.refractivity_model gt.trunefieldsim.refractivity_model_parameters = trawefield.refractivity_model_parameters + # The TRun run number + gt.trun.run_number = trawefield.run_number + + ## The antenna time window is defined around a t0 that changes with the antenna, starts on t0+t_pre (thus t_pre is usually negative) and ends on t0+post + gt.trunefieldsim.t_pre = trawefield.t_pre + gt.trunefieldsim.t_post = trawefield.t_post + + +def get_tree_du_id_and_xyz(trawefield): # *** Store the DU's to run - they needed to be collected from all events *** # Get the ids and positions from all the events count = trawefield.draw("du_id:du_x:du_y:du_z", "", "goff") @@ -239,19 +311,7 @@ def rawefield2grandrootrun(trawefield, gt): # Stack x/y/z together and leave only the ones for unique du_ids du_xyzs = np.column_stack([du_xs, du_ys, du_zs])[unique_dus_idx] - # The TRun run number - gt.trun.run_number = trawefield.run_number - - # Assign the du ids and positions to the trun tree - gt.trun.du_id = du_ids - gt.trun.du_xyz = du_xyzs - - ## The antenna time window is defined around a t0 that changes with the antenna, starts on t0+t_pre (thus t_pre is usually negative) and ends on t0+post - gt.trunefieldsim.t_pre = trawefield.t_pre - gt.trunefieldsim.t_post = trawefield.t_post - # ToDo: shouldn't this and above be created for every DU in sims? - # gt.trun.t_bin_size = [trawefield.t_bin_size*1e9]*len(du_ids) - gt.trun.t_bin_size = [trawefield.t_bin_size]*len(du_ids) #Matias Question: Why is this being mutiplied here? + return np.asarray(du_ids), np.asarray(du_xyzs) # Convert the RawShowerTree entries @@ -462,13 +522,13 @@ def get_origin_geoid(clargs): return origin_geoid # Form the proper output directory name from command line arguments -def form_directory_name(clargs): +def form_directory_name(clargs, date, time, run_number): # Change possible underscores in extra into - extra = clargs.extra.replace("_", "-") # Go through serial numbers in directory names to find a one that does not exist for sn in range(1000): - dir_name = Path(clargs.output_parent_directory, f"sim_{clargs.site_name}_{clargs.sim_date}_{extra}_{sn:0>4}") + dir_name = Path(clargs.output_parent_directory, f"sim_{clargs.site_name}_{date}_{time}_RUN{run_number}_CD_{extra}_{sn:0>4}") if not dir_name.exists(): break # If directories with serial number up to 1000 already created @@ -481,7 +541,7 @@ def form_directory_name(clargs): # Rename the created files to appropriate names def rename_files(clargs, path, start_event_number, end_event_number): # Go through output files - for fn_start in ["trun", "trunshowersim", "trunefieldsim", "tshower", "tshowersim", "tefield"]: + for fn_start in ["runshowersim", "runefieldsim", "shower", "showersim", "efield"]: # Go through serial numbers in directory names to find a one that does not exist for sn in range(1000): fn_in = Path(path, f"{fn_start}.root")