Add run selector + unit handling + multiprocessing + examples

run_selector.py

@@ -0,0 +1,166 @@
+# coding: utf-8
+
+import os
+import glob
+from functools import partial
+from pp_params import *
+import f90nml
+
+
+class RunSelector:
+    def __init__(
+        self,
+        path_in,
+        in_runs=None,
+        in_nums="all",
+        pp_params=default_params(),
+        number_run="[0-9]*_",
+        name_run="*",
+        namelist_cond={},
+        sort_run_by=None,
+        time_min=None,
+        time_max=None,
+    ):
+        self.path_in = path_in
+        self.pp_params = pp_params
+
+        self.namelist = {}
+        self.runs = self.get_runs(
+            in_runs, number_run, name_run, namelist_cond, sort_run_by
+        )
+
+        self.info = {}
+        for run in self.runs:
+            self.info[run] = {}
+        self.nums = {}
+
+        if not type(in_nums) == dict:
+            nums_temp = in_nums
+            in_nums = {}
+            for run in self.runs:
+                in_nums[run] = nums_temp
+
+        for run in self.runs:
+            self.nums[run] = self.get_nums(run, in_nums[run], time_min, time_max)
+
+    def load_namelist(self, run):
+        path_run = self.path_in + "/" + run
+        path_nml = path_run + "/" + self.pp_params.input.nml_filename
+        return f90nml.read(path_nml)
+
+    def get_nml_value(self, nml_key, run):
+        res = self.namelist[run]
+        for key in nml_key.split("/"):
+            res = res[key]
+        return res
+
+    def get_runs(
+        self,
+        in_runs=None,
+        number_run="[0-9]*_",
+        name_run="*",
+        namelist_cond={},
+        sort_run_by=None,
+    ):
+        def try_load_nml(run):
+            try:
+                self.namelist[run] = self.load_namelist(run)
+                success = True
+            except IOError:
+                success = False
+            return success
+
+        runs = map(
+            os.path.basename, glob.glob(self.path_in + "/" + number_run + name_run)
+        )
+        if not in_runs is None:
+            runs = filter(lambda n: n in runs, in_runs)
+        runs = filter(try_load_nml, runs)
+
+        if type(namelist_cond) == tuple:
+            namelist_cond = [namelist_cond]
+
+        for (nml_key, operator, operand) in namelist_cond:
+            value = {}
+            for run in runs:
+                value[run] = self.get_nml_value(nml_key, run)
+            if operator == "=":
+                runs = filter(lambda r: value[r] == operand, runs)
+            if operator == "!=":
+                runs = filter(lambda r: not value[r] == operand, runs)
+            elif operator == ">":
+                runs = filter(lambda r: value[r] > operand, runs)
+            elif operator == "<":
+                runs = filter(lambda r: value[r] < operand, runs)
+            elif operator == "in":
+                runs = filter(lambda r: value[r] in operand, runs)
+
+        # Sort by the value in the namelist of sort_run_by
+        if not sort_run_by is None:
+            if type(sort_run_by) == str:
+                sort_run_by = [sort_run_by]
+            for nml_key in reversed(sort_run_by):
+                runs.sort(key=partial(self.get_nml_value, nml_key))
+
+        return runs
+
+    def load_info(self, run, num):
+        info_file = open(
+            self.path_in
+            + "/"
+            + run
+            + "/"
+            + "output_"
+            + str(num).zfill(5)
+            + "/"
+            + "info_"
+            + str(num).zfill(5)
+            + ".txt",
+            "r",
+        )
+        info = {}
+        for line in info_file.readlines():
+            parsed = yaml.safe_load(line.replace("=", ":"))
+            if type(parsed) == dict:
+                info.update(parsed)
+        info_file.close()
+        return info
+
+    def get_nums(self, run, in_nums=None, time_min=None, time_max=None):
+        def try_load_info(num):
+            try:
+                self.info[run][num] = self.load_info(run, num)
+                success = True
+            except IOError:
+                success = False
+            return success
+
+        names = glob.glob(
+            self.path_in + "/" + run + "/output_[0-9][0-9][0-9][0-9][0-9]"
+        )
+        nums = map(lambda n: int(n.split("/")[-1].split("_")[1]), names)
+
+        if type(in_nums) == list:
+            nums = filter(lambda n: n in nums, in_nums)
+
+        nums = np.sort(nums)
+
+        if in_nums == "first":
+            i = 0
+            while i < len(nums) and not try_load_info(nums[i]):
+                i = i + 1
+            nums = [nums[i]]
+        elif in_nums == "last":
+            i = len(nums) - 1
+            while i >= 0 and not try_load_info(nums[i]):
+                i = i - 1
+            nums = [nums[i]]
+        else:
+            nums = filter(try_load_info, nums)
+
+        if not time_min is None:
+            nums = filter(lambda n: self.info[run][n]["time"] >= time_min, nums)
+        if not time_max is None:
+            nums = filter(lambda n: self.info[run][n]["time"] <= time_max, nums)
+
+        return nums