diff --git a/run_selector.py b/run_selector.py index ea27e7d..2ee63b1 100644 --- a/run_selector.py +++ b/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): - if do_tests: + """ + 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 + return success - if isinstance(in_nums, int): - in_nums = [in_nums] + else: + 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: