Add filaments postproc, improve units detection, add automatic map rules, add selection

run_selector.py

@@ -106,10 +106,70 @@ class RunSelector:
             for run in self.runs:
                 in_nums[run] = nums_temp
 
-            for i, run in enumerate(self.runs):
-                self.nums[run] = self.get_nums(
-                    run, in_nums[run], time_min, time_max, time
-                )
+        for i, run in enumerate(self.runs):
+            self.nums[run] = self.get_nums(run, in_nums[run], time_min, time_max, time)
+
+    def select(
+        self,
+        runs=None,
+        nums="all",
+        filter_nml={},
+        sort_run_by=None,
+        time_min=None,
+        time_max=None,
+        time=None,
+    ):
+        """
+        Sub-select runs and outputs from already selected runs and outputs
+
+        Parameters
+        ---------
+        runs : str or list of str. The name runs to consider. Default: all.
+        nums : int or list of int or str.
+             The output numbers to consider.
+             "last" select only the last output.
+             "all" preselect all outputs (default)
+
+        filter_nml  : tuple or list of tupple.
+                     Filter runs by namelist.
+                     tuples are in the following form:
+                     (nml_key, operator, nml_value)
+                     with nml_key a key from the namelist (eg. "cloud_params/dens0")
+                     operator within ("=", "!=", "<", ">", "in")
+                     and nml_value a string, float or int
+        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.
+
+        sort_run_by : str, a key from the namelist used to sort the runs (by ascending order)
+
+        Returns
+        -------
+
+        (selected_runs, selected_nums)
+        """
+
+        selected_runs = self.get_runs(
+            runs, "*", filter_nml, sort_run_by, do_tests=False
+        )
+
+        if len(selected_runs) == 0:
+            raise ValueError("No runs found")
+
+        if not type(nums) == dict:
+            nums_temp = nums
+            nums = {}
+            for run in selected_runs:
+                nums[run] = nums_temp
+
+        selected_nums = {}
+
+        for i, run in enumerate(selected_runs):
+            selected_nums[run] = self.get_nums(
+                run, nums[run], time_min, time_max, time, do_tests=False
+            )
+
+        return selected_runs, selected_nums
 
     def load_namelist(self, run):
         path_run = self.path_in + "/" + run
@@ -139,7 +199,14 @@ class RunSelector:
                 runs = list(filter(lambda r: value[r] in operand, runs))
         return runs
 
-    def get_runs(self, in_runs=None, filter_name="*", filter_nml={}, sort_run_by=None):
+    def get_runs(
+        self,
+        in_runs=None,
+        filter_name="*",
+        filter_nml={},
+        sort_run_by=None,
+        do_tests=True,
+    ):
         def try_load_nml(run):
             try:
                 self.namelist[run] = self.load_namelist(run)
@@ -148,17 +215,25 @@ class RunSelector:
                 success = False
             return success
 
-        runs = list(
-            map(
-                os.path.basename,
-                list(
-                    filter(os.path.isdir, glob.glob(self.path_in + "/" + filter_name))
-                ),
+        if do_tests:
+            runs = list(
+                map(
+                    os.path.basename,
+                    list(
+                        filter(
+                            os.path.isdir, glob.glob(self.path_in + "/" + filter_name)
+                        )
+                    ),
+                )
             )
-        )
+        else:
+            runs = self.runs
+
         if in_runs is not None:
             runs = list(filter(lambda n: n in runs, in_runs))
-        runs = list(filter(try_load_nml, runs))
+
+        if do_tests:
+            runs = list(filter(try_load_nml, runs))
 
         # Select runs that match namelist conditions
         runs = self.nml_select(runs, filter_nml)
@@ -194,23 +269,32 @@ class RunSelector:
             info_file.close()
         return info
 
-    def get_nums(self, run, in_nums=None, time_min=None, time_max=None, time=None):
+    def get_nums(
+        self, run, in_nums=None, time_min=None, time_max=None, time=None, do_tests=True
+    ):
         def try_load_info(num):
-            try:
-                self.info[run][num] = self.load_info(run, num)
+            if do_tests:
+                try:
+                    self.info[run][num] = self.load_info(run, num)
+                    success = True
+                except IOError:
+                    success = False
+            else:
                 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 = list(map(lambda n: int(n.split("/")[-1].split("_")[1]), names))
-
-        if type(in_nums) == int:
+        if isinstance(in_nums, int):
             in_nums = [in_nums]
-        if type(in_nums) == list:
+
+        if do_tests:
+            names = glob.glob(
+                self.path_in + "/" + run + "/output_[0-9][0-9][0-9][0-9][0-9]"
+            )
+            nums = list(map(lambda n: int(n.split("/")[-1].split("_")[1]), names))
+        else:
+            nums = self.nums[run]
+
+        if isinstance(in_nums, list):
             nums = list(filter(lambda n: n in nums, in_nums))
 
         nums = np.sort(nums)