[run_selector] add an option to choose the time unit

run_selector.py

@@ -50,6 +50,7 @@ class RunSelector:
         time_min=None,
         time_max=None,
         time=None,
+        unit_time=None,
     ):
         """
         Select runs and outputs with several filter options.
@@ -83,6 +84,7 @@ class RunSelector:
         time_min    : float, select output where time >= time_min (in code units)
         time_max    : float, select output where time <= time_min (in code units)
         time        : float or list of float. For each value, select the output closer to it.
+        unit_time   : astrophysix.Unit, unit for the time above. None is code unit.
 
         3. Sort the runs
 
@@ -114,6 +116,7 @@ class RunSelector:
                 time_min,
                 time_max,
                 time,
+                unit_time,
             )
 
         i = 0
@@ -138,6 +141,7 @@ class RunSelector:
         time_min=None,
         time_max=None,
         time=None,
+        unit_time=None,
     ):
         """
         Sub-select runs and outputs from already selected runs and outputs
@@ -164,6 +168,8 @@ class RunSelector:
         time_min    : float, select output where time >= time_min (in code units)
         time_max    : float, select output where time <= time_min (in code units)
         time        : float or list of float. For each value, select the output closer to it.
+        unit_time   : astrophysix.Unit, unit for the time above. None is code unit.
+
 
         sort_run_by : str, a key from the namelist used to sort the runs (by ascending order)
 
@@ -193,7 +199,7 @@ class RunSelector:
 
         for i, run in enumerate(selected_runs):
             selected_nums[run] = self.get_nums(
-                run, nums[run], time_min, time_max, time, do_tests=False
+                run, nums[run], time_min, time_max, time, unit_time, do_tests=False
             )
 
         return selected_runs, selected_nums
@@ -278,22 +284,63 @@ class RunSelector:
         return info
 
     def get_nums(
-        self, run, in_nums=None, time_min=None, time_max=None, time=None, do_tests=True
+        self,
+        run,
+        in_nums=None,
+        time_min=None,
+        time_max=None,
+        time=None,
+        unit_time=None,
+        do_tests=True,
     ):
-        def try_load_info(num):
+        """
+        Select snapshots from the disk
+
+        Parameters
+        ---------
+
+        in_nums : int or list of int or str.
+             The output numbers to consider.
+             "last" select only the last output.
+             "all" preselect all outputs (default)
+
+        time_min    : float, select output where time >= time_min (in code units)
+        time_max    : float, select output where time <= time_min (in code units)
+        time        : float or list of float. For each value, select the output closer to it.
+        unit_time   : astrophysix.Unit, unit for the time above. None is code unit.
+
+        do_tests    : test if the snapshots are actually on disk. Not needed when subselecting snapshots.
+        """
+
+        # -- Initialize info loader --
         if do_tests:
+
+            def try_load_info(num):
                 try:
                     self.info[run][num] = self.load_info(run, num)
                     success = True
                 except (IOError, AttributeError):
                     success = False
-            else:
-                success = True
                 return success
 
-        if isinstance(in_nums, int):
+        else:
-            in_nums = [in_nums]
 
+            def try_load_info(num):
+                return True
+
+        # -- Time getter according to unit_times
+        if unit_time is None:
+
+            def get_time(num):
+                return self.info[run][num]["time"]
+
+        else:
+
+            def get_time(num):
+                time_code = self.info[run][num]["time"]
+                return time_code * self.info[run][num]["unit_time"].express(unit_time)
+
+        # -- Get the list of seemingly available snapshots on the disk or already selected --
         if do_tests:
             names = glob.glob(
                 self.path_in + "/" + run + "/output_[0-9][0-9][0-9][0-9][0-9]"
@@ -302,11 +349,17 @@ class RunSelector:
         else:
             nums = self.nums[run]
 
+        # -- Filter with the provided in_nums array
+
+        if isinstance(in_nums, int):
+            in_nums = [in_nums]
+
         if isinstance(in_nums, list):
             nums = list(filter(lambda n: n in nums, in_nums))
 
         nums = np.sort(nums)
 
+        # --  Select either the first or last output from the list, or all the valid ones --
         if in_nums == "first":
             i = 0
             while i < len(nums) and not try_load_info(nums[i]):
@@ -326,17 +379,18 @@ class RunSelector:
         else:
             nums = list(filter(try_load_info, nums))
 
+        # -- Select according to time --
         if time_min is not None:
-            nums = list(filter(lambda n: self.info[run][n]["time"] >= time_min, nums))
+            nums = list(filter(lambda n: get_time(n) >= time_min, nums))
         if time_max is not None:
-            nums = list(filter(lambda n: self.info[run][n]["time"] <= time_max, nums))
+            nums = list(filter(lambda n: get_time(n) <= time_max, nums))
 
         if time is not None and len(nums) > 0:
             filtered_nums = []
             if not isinstance(time, list):
                 time = [time]
             # Get time for all already selected nums
-            time_all = np.asarray([[self.info[run][n]["time"], n] for n in nums])
+            time_all = np.asarray([[get_time(n), n] for n in nums])
 
             # For all times provided by the user, select the output closer to it
             for t in time: