diff --git a/run_selector.py b/run_selector.py index 2ee63b1..708848d 100644 --- a/run_selector.py +++ b/run_selector.py @@ -317,7 +317,8 @@ class RunSelector: def try_load_info(num): try: - self.info[run][num] = self.load_info(run, num) + if num not in self.info[run]: + self.info[run][num] = self.load_info(run, num) success = True except (IOError, AttributeError): success = False @@ -340,6 +341,51 @@ class RunSelector: time_code = self.info[run][num]["time"] return time_code * self.info[run][num]["unit_time"].express(unit_time) + # -- A function to search a given time using dichotomy + def search(nums, time, position="closest"): + + while len(nums) > 0 and not try_load_info(nums[0]): + del nums[0] + while len(nums) > 0 and not try_load_info(nums[-1]): + del nums[-1] + + if len(nums) == 0: + return None + ileft, iright = 0, len(nums) - 1 + if get_time(nums[ileft]) >= time: + if position in ["closest", "right"]: + return ileft + else: + return None + if get_time(nums[iright]) < time: + if position in ["closest", "left"]: + return iright + else: + return None + + while iright - ileft > 1: + imid = (ileft + iright) // 2 + while not try_load_info(nums[imid]): + del nums[imid] + iright -= 1 + imid = (ileft + iright) // 2 + if get_time(nums[imid]) < time: + ileft = imid + else: + iright = imid + + if position == "left": + return ileft + elif position == "right": + return nums[iright] + else: + dleft = np.abs(get_time(nums[ileft]) - time) + dright = np.abs(get_time(nums[iright]) - time) + if dleft <= dright: + return ileft + else: + return iright + # -- Get the list of seemingly available snapshots on the disk or already selected -- if do_tests: names = glob.glob( @@ -357,7 +403,7 @@ class RunSelector: if isinstance(in_nums, list): nums = list(filter(lambda n: n in nums, in_nums)) - nums = np.sort(nums) + nums.sort() # -- Select either the first or last output from the list, or all the valid ones -- if in_nums == "first": @@ -376,30 +422,33 @@ class RunSelector: nums = [nums[i]] else: nums = [] - else: - nums = list(filter(try_load_info, nums)) # -- Select according to time -- if time_min is not None: - nums = list(filter(lambda n: get_time(n) >= time_min, nums)) + imin = search(nums, time_min, "right") + if imin is not None: + nums = nums[imin:] + else: + nums = [] if time_max is not None: - nums = list(filter(lambda n: get_time(n) <= time_max, nums)) + imax = search(nums, time_min, "left") + if imax is not None: + nums = nums[: imax + 1] 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([[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: - # Index of this output in the time_all array - idx = (np.abs(time_all[:, 0] - t)).argmin() - num = int(time_all[idx, 1]) + iclose = search(nums, t) + num = nums[iclose] # Only add each selected output once if num not in filtered_nums: filtered_nums.append(num) nums = filtered_nums + else: + nums = list(filter(try_load_info, nums)) return nums