Merge pull request #94 from rubisco-sfa/nersc

FIX: Bugs encountered while running E3SM studies at NERSC

src/ILAMB/Confrontation.py

@@ -345,7 +345,8 @@ def stageData(self, m):
             t0=None if len(self.study_limits) != 2 else self.study_limits[0],
             tf=None if len(self.study_limits) != 2 else self.study_limits[1],
         )
-        obs.data *= self.scale_factor
+        with np.errstate(all="ignore"):
+            obs.data *= self.scale_factor
         if obs.time is None:
             raise il.NotTemporalVariable()
         self.pruneRegions(obs)

src/ILAMB/Regions.py

@@ -218,6 +218,9 @@ def getMask(self, label, var):
                 ).argmax(axis=1)
             else:
                 # if more globally defined, nearest neighbor is fine
+                if lat is None or var.lat is None:
+                    msg = f"No latitudes are given either in the variable:\n{var}\n{var.lat}\nor in the region: {label} {lat}"
+                    logger.debug(msg)
                 rows = (np.abs(lat[:, np.newaxis] - var.lat)).argmin(axis=0)
                 cols = (np.abs(lon[:, np.newaxis] - var.lon)).argmin(axis=0)
             if var.ndata:

src/ILAMB/data/ilamb.cfg

@@ -103,7 +103,7 @@ force_emulation = True
 [HIPPOAToM]
 ctype      = "ConfGSNF"
 variable   = "gsnf"
-source     = "DATA/co2/HIPPO_AToM/HIPPO_AToM.nc"
+source     = "DATA/co2/HIPPOATom/HIPPO_AToM.nc"
 model_flux = "ra+rh-gpp"
 
 #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

src/ILAMB/ilamblib.py

@@ -1004,9 +1004,9 @@ def FromNetCDF4(
         lat_name = []
         lon_name = []
         for key in grp.variables.keys():
-            if "lat" in key:
+            if key.lower().startswith("lat"):
                 lat_name.append(key)
-            if "lon" in key:
+            if key.lower().startswith("lon"):
                 lon_name.append(key)
             if "altitude" in key:
                 depth_name = key

src/ILAMB/run.py

@@ -2,14 +2,17 @@
 import os
 import pickle
 from traceback import format_exc
+import yaml
+import inspect
 
 from mpi4py import MPI
 
 from ILAMB import ilamblib as il
 from ILAMB.ModelResult import ModelResult
+from matplotlib import colors as mplclrs
 
 logger = logging.getLogger("%i" % MPI.COMM_WORLD.rank)
-
+rank = 0
 
 def InitializeModels(
     model_root,
@@ -107,14 +110,89 @@ def InitializeModels(
     return M
 
 
+def _parse_model_yaml(filename: str, cache_path: str = "./", only_models: list = []):
+    """Setup models using a yaml file."""
+    model_classes = {
+        "ModelResult": ModelResult,
+    }
+    models = []
+    with open(filename, encoding="utf-8") as fin:
+        yml = yaml.safe_load(fin)
+    for name, opts in yml.items():
+        # optionally filter models
+        if len(only_models) > 0 and name not in only_models:
+            continue
+
+        if "name" not in opts:
+            opts["name"] = name
+
+        # if the model_year option is given, convert to lits of floats
+        if "model_year" in opts:
+            opts["model_year"] = [
+                float(y.strip()) for y in opts["model_year"].split(",")
+            ]
+
+        # select the class type
+        cls = model_classes[opts["type"]] if "type" in opts else ModelResult
+        if cls is None:
+            typ = opts["type"]
+            raise ValueError(f"The model type '{typ}' is not available")
+        fcns = dir(cls)
+
+        # if the pickle file exists, just load it
+        cache = os.path.join(cache_path, f"{name}.pkl")
+        if os.path.exists(cache):
+            if "read_pickle" in fcns:
+                model = cls().read_pickle(cache)
+            else:
+                with open(cache, mode="rb") as fin:
+                    model = pickle.load(fin)
+            models.append(model)
+            continue
+
+        # call the constructor using keywords defined in the YAML file
+        cls = model_classes[opts["type"]] if "type" in opts else ModelResult
+        model = cls(
+            **{
+                key: opts[key]
+                for key in inspect.getfullargspec(cls).args
+                if key in opts
+            }
+        )
+
+        # some model types have a find_files() method, call if present loading
+        # proper keywords from the YAML file
+        if "find_files" in fcns:
+            model.find_files(
+                **{
+                    key: opts[key]
+                    for key in inspect.getfullargspec(model.find_files).args
+                    if key in opts
+                }
+            )
+
+        # some model types allow you to specify snynonms
+        if "add_synonym" in fcns and "synonyms" in opts:
+            for mvar, syn in opts["synonyms"].items():
+                model.add_synonym(mvar, syn)
+
+        # cache the model result
+        if rank == 0:
+            if "read_pickle" in fcns:
+                model.to_pickle(cache)
+            else:
+                with open(cache, mode="wb") as fin:
+                    pickle.dump(model, fin)
+
+        models.append(model)
+
+    for model in models:
+        if isinstance(model.color, str) and model.color.startswith("#"):
+            model.color = mplclrs.hex2color(model.color)
+    return models
+
 def ParseModelSetup(
-    model_setup,
-    models=[],
-    verbose=False,
-    filter="",
-    regex="",
-    models_path="./",
-    comm=MPI.COMM_WORLD,
+        model_setup, models=[], verbose=False, filter="", regex="", models_path="./", comm=MPI.COMM_WORLD
 ):
     """Initializes a list of models
 
@@ -137,12 +215,24 @@ def ParseModelSetup(
        a list of the model results, sorted alphabetically by name
 
     """
-    rank = comm.rank
+    if rank == 0 and verbose:
+        print("\nSetting up model results from %s\n" % model_setup)
+
+    # intercept if this is a yaml file
+    if model_setup.endswith(".yaml"):
+        M = _parse_model_yaml(model_setup, cache_path=models_path, only_models=models)
+        if rank == 0 and verbose:
+            for m in M:
+                print(("    {0:>45}").format(m.name))
+            if len(M) == 0:
+                print("No model results found")
+                comm.Barrier()
+                comm.Abort(0)
+        return M
+
     # initialize the models
     M = []
     max_model_name_len = 0
-    if rank == 0 and verbose:
-        print("\nSetting up model results from %s\n" % model_setup)
     with open(model_setup) as f:
         for line in f.readlines():
             if line.strip().startswith("#"):
@@ -151,12 +241,15 @@ def ParseModelSetup(
             mname = None
             mdir = None
             model_year = None
+            mgrp = ""
             if len(line) >= 2:
                 mname = line[0].strip()
                 mdir = line[1].strip()
                 # if mdir not a directory, then maybe path is relative to ILAMB_ROOT
                 if not os.path.isdir(mdir):
                     mdir = os.path.join(os.environ["ILAMB_ROOT"], mdir).strip()
+                if len(line) == 3:
+                    mgrp = line[2].strip()
             if len(line) == 4:
                 model_year = [float(line[2].strip()), float(line[3].strip())]
             max_model_name_len = max(max_model_name_len, len(mname))
@@ -174,6 +267,7 @@ def ParseModelSetup(
                         filter=filter,
                         regex=regex,
                         model_year=model_year,
+                        group=mgrp,
                     )
                 except Exception as ex:
                     logger.debug("[%s]" % mname, format_exc())
@@ -183,8 +277,7 @@ def ParseModelSetup(
     # assign unique colors
     clrs = il.GenerateDistinctColors(len(M))
     for m in M:
-        clr = clrs.pop(0)
-        m.color = clr
+        m.color = clrs.pop(0)
 
     # save model objects as pickle files
     comm.Barrier()