diff --git a/apps/ta_dump.py b/apps/ta_dump.py index 7bdaf94..8388dba 100644 --- a/apps/ta_dump.py +++ b/apps/ta_dump.py @@ -258,7 +258,7 @@ def plot_summary_stats(ta_data, no_anomaly=False, quiet=False): write_summary_stats(ta_data[ta_key], anomaly_filename, title) -def all_event_displays(tp_data, run_id, sub_run_id, seconds=False): +def all_event_displays(tp_data, run_id, file_index, seconds=False): """ Plot all event_displays as pages in a PDF. @@ -266,7 +266,7 @@ def all_event_displays(tp_data, run_id, sub_run_id, seconds=False): """ time_unit = 's' if seconds else 'Ticks' - with PdfPages(f"event_displays_{run_id}.{sub_run_id:04}.pdf") as pdf: + with PdfPages(f"event_displays_{run_id}.{file_index:04}.pdf") as pdf: for tadx, ta in enumerate(tp_data): if seconds: ta = ta * TICK_TO_SEC_SCALE @@ -280,7 +280,7 @@ def all_event_displays(tp_data, run_id, sub_run_id, seconds=False): time_diff = max_time - min_time plt.xlim((min_time - 0.1*time_diff, max_time + 0.1*time_diff)) - plt.title(f'Run {run_id}.{sub_run_id:04} Event Display: {tadx:03}') + plt.title(f'Run {run_id}.{file_index:04} Event Display: {tadx:03}') plt.xlabel(f"Peak Time ({time_unit})") plt.ylabel("Channel") @@ -358,6 +358,7 @@ def parse(): parser.add_argument("--end-frag", type=int, help="Fragment index to stop processing (i.e. not inclusive). Takes negative indexing. Default: 0.", default=0) parser.add_argument("--no-anomaly", action="store_true", help="Pass to not write 'ta_anomaly_summary.txt'. Default: False.") parser.add_argument("--seconds", action="store_true", help="Pass to use seconds instead of time ticks. Default: False.") + parser.add_argument("--overwrite", action="store_true", help="Overwrite old outputs. Default: False.") return parser.parse_args() @@ -374,6 +375,7 @@ def main(): end_frag = args.end_frag no_anomaly = args.no_anomaly seconds = args.seconds + overwrite = args.overwrite data = trgtools.TAData(filename, quiet) @@ -390,8 +392,18 @@ def main(): for path in frag_paths: data.load_frag(path) - run_id = data.run_id - sub_run_id = data.sub_run_id + # Try to find an empty plotting directory + plot_iter = 0 + plot_dir = f"{data.run_id}-{data.file_index}_figures_{plot_iter:04}" + while not overwrite and os.path.isdir(plot_dir): + plot_iter += 1 + plot_dir = f"{data.run_id}-{data.file_index}_figures_{plot_iter:04}" + print(f"Saving outputs to ./{plot_dir}/") + # If overwriting and it does exist, don't need to make it. + # So take the inverse to mkdir. + if not (overwrite and os.path.isdir(plot_dir)): + os.mkdir(plot_dir) + os.chdir(plot_dir) print(f"Number of TAs: {data.ta_data.shape[0]}") # Enforcing output for useful metric @@ -407,7 +419,7 @@ def main(): plot_summary_stats(data.ta_data, no_anomaly, quiet) if (not no_displays): - all_event_displays(data.tp_data, run_id, sub_run_id, seconds) + all_event_displays(data.tp_data, data.run_id, data.file_index, seconds) if __name__ == "__main__": main() diff --git a/apps/tp_dump.py b/apps/tp_dump.py index 5f38672..ed6aa92 100644 --- a/apps/tp_dump.py +++ b/apps/tp_dump.py @@ -346,6 +346,7 @@ def parse(): parser.add_argument("--start-frag", type=int, help="Fragment to start loading from (inclusive); can take negative integers. Default: -10", default=-10) parser.add_argument("--end-frag", type=int, help="Fragment to stop loading at (exclusive); can take negative integers. Default: 0", default=0) parser.add_argument("--no-anomaly", action="store_true", help="Pass to not write 'tp_anomaly_summary.txt'. Default: False.") + parser.add_argument("--overwrite", action="store_true", help="Overwrite old outputs. Default: False.") return parser.parse_args() @@ -360,6 +361,7 @@ def main(): start_frag = args.start_frag end_frag = args.end_frag no_anomaly = args.no_anomaly + overwrite = args.overwrite data = trgtools.TPData(filename, quiet) if end_frag == 0: # Ex: [-10:0] is bad. @@ -372,6 +374,19 @@ def main(): print("Fragment Path:", path) data.load_frag(path) + # Try to find an empty plotting directory + plot_iter = 0 + plot_dir = f"{data.run_id}-{data.file_index}_figures_{plot_iter:04}" + while not overwrite and os.path.isdir(plot_dir): + plot_iter += 1 + plot_dir = f"{data.run_id}-{data.file_index}_figures_{plot_iter:04}" + print(f"Saving outputs to ./{plot_dir}/") + # If overwriting and it does exist, don't need to make it. + # So take the inverse to mkdir. + if not (overwrite and os.path.isdir(plot_dir)): + os.mkdir(plot_dir) + os.chdir(plot_dir) + if (not quiet): print("Size of tp_data:", data.tp_data.shape) diff --git a/python/trgtools/TAData.py b/python/trgtools/TAData.py index 7d32cab..93648ad 100644 --- a/python/trgtools/TAData.py +++ b/python/trgtools/TAData.py @@ -75,26 +75,9 @@ def __init__(self, filename, quiet=False): self._nonempty_frags_mask = np.ones((len(self._frag_paths),), dtype=bool) self._num_empty = 0 - if "run" in filename: # Waiting on hdf5libs PR to use get_int_attribute - tmp_name = filename.split("run")[1] - try: - self.run_id = int(tmp_name.split("_")[0]) - except: - if not self._quiet: - print(self._WARNING_TEXT_COLOR + "WARNING: Couldn't find Run ID in file name. Using run id 0." + self._END_TEXT_COLOR) - self.run_id = 0 - try: - self.sub_run_id = int(tmp_name.split("_")[1]) - except: - if not self._quiet: - print(self._WARNING_TEXT_COLOR + "WARNING: Couldn't find SubRun ID in file name. Using SubRun ID 1000." + self._END_TEXT_COLOR) - self.sub_run_id = 1000 - else: - if not self._quiet: - print(self._WARNING_TEXT_COLOR + "WARNING: Couldn't find Run ID in file name. Using run id 0." + self._END_TEXT_COLOR) - print(self._WARNING_TEXT_COLOR + "WARNING: Couldn't find SubRun ID in file name. Using SubRun ID 1000." + self._END_TEXT_COLOR) - self.run_id = 0 - self.sub_run_id = 1000 + # File identification attributes + self.run_id = self._h5_file.get_int_attribute('run_number') + self.file_index = self._h5_file.get_int_attribute('file_index') def _set_ta_frag_paths(self, frag_paths) -> None: """ diff --git a/python/trgtools/TPData.py b/python/trgtools/TPData.py index 119fec9..8f71cd3 100644 --- a/python/trgtools/TPData.py +++ b/python/trgtools/TPData.py @@ -48,6 +48,10 @@ def __init__(self, filename, quiet=False): self.tp_data = np.array([], dtype=self.tp_dt) # Will concatenate TPs self._quiet = quiet + # File identification attributes + self.run_id = self._h5_file.get_int_attribute('run_number') + self.file_index = self._h5_file.get_int_attribute('file_index') + def _set_tp_frag_paths(self, frag_paths) -> None: """ Only collect the fragment paths that are for TAs.