Merge remote-tracking branch 'origin/multievent_73_dev_efield2voltage…

…' into 73_dev_for_merge

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

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)

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")