Pydantic ColocationSetup

pyaerocom/aeroval/_processing_base.py

@@ -4,6 +4,7 @@
 from pyaerocom.aeroval import EvalSetup
 from pyaerocom.aeroval.experiment_output import ExperimentOutput
 from pyaerocom.colocation_auto import Colocator
+from pyaerocom.colocation_setup import ColocationSetup
 
 
 class HasConfig:
@@ -113,21 +114,37 @@ def get_colocator(self, model_name: str = None, obs_name: str = None) -> Colocat
         Colocator
 
         """
+        col_cfg = {**self.cfg.colocation_opts.model_dump()}
+        outdir = self.cfg.path_manager.get_coldata_dir()
+        col_cfg["basedir_coldata"] = outdir
+
+        if not model_name and not obs_name:
+            col_stp = ColocationSetup(**col_cfg)
+            return Colocator(col_stp)
+
+        if model_name:
+            mod_cfg = self.cfg.get_model_entry(model_name)
+            col_cfg["model_cfg"] = mod_cfg
+
+            # LB: Hack and at what lowlevel_helpers's import_from was doing
+            for key, val in mod_cfg.items():
+                if key in ColocationSetup.model_fields:
+                    col_cfg[key] = val
         if obs_name:
             obs_cfg = self.cfg.get_obs_entry(obs_name)
             pyaro_config = obs_cfg["obs_config"] if "obs_config" in obs_cfg else None
-            col_cfg = {**self.cfg.colocation_opts}
             col_cfg["obs_config"] = pyaro_config
-            col = Colocator(**col_cfg)
-            col.import_from(obs_cfg)
-            col.add_glob_meta(diurnal_only=self._get_diurnal_only(obs_name))
-        else:
-            col = Colocator(**self.cfg.colocation_opts)
-        if model_name:
-            mod_cfg = self.cfg.get_model_entry(model_name)
-            col.import_from(mod_cfg)
-        outdir = self.cfg.path_manager.get_coldata_dir()
-        col.basedir_coldata = outdir
+
+            # LB: Hack and at what lowlevel_helpers's import_from was doing
+            for key, val in obs_cfg.items():
+                if key in ColocationSetup.model_fields:
+                    col_cfg[key] = val
+
+            col_cfg["add_meta"].update(diurnal_only=self._get_diurnal_only(obs_name))
+
+        col_stp = ColocationSetup(**col_cfg)
+        col = Colocator(col_stp)
+
         return col
 
 

pyaerocom/aeroval/experiment_processor.py

@@ -56,7 +56,7 @@ def _run_single_entry(self, model_name, obs_name, var_list):
         elif ocfg["only_json"]:
             if not ocfg["coldata_dir"]:
                 raise Exception(
-                    f"No coldata_dir provided for an obs network for which only_json=True. The assumption of setting only_json=True is that colocated files already exist, and so a directory for these files must be provided."
+                    "No coldata_dir provided for an obs network for which only_json=True. The assumption of setting only_json=True is that colocated files already exist, and so a directory for these files must be provided."
                 )
             else:
                 preprocessed_coldata_dir = ocfg["coldata_dir"]
@@ -70,8 +70,7 @@ def _run_single_entry(self, model_name, obs_name, var_list):
             if self.cfg.processing_opts.only_json:
                 files_to_convert = col.get_available_coldata_files(var_list)
             else:
-                model_read_kwargs = self.cfg.model_cfg[model_name]["kwargs"]
-                col.run(var_list, model_read_kwargs=model_read_kwargs)
+                col.run(var_list)
                 files_to_convert = col.files_written
 
             if self.cfg.processing_opts.only_colocation:

pyaerocom/aeroval/helpers.py

@@ -13,8 +13,6 @@
     start_stop,
     to_pandas_timestamp,
 )
-from pyaerocom.io import ReadGridded
-from pyaerocom.tstype import TsType
 from pyaerocom.variable import Variable
 
 logger = logging.getLogger(__name__)

pyaerocom/aeroval/setupclasses.py

@@ -33,7 +33,7 @@
     check_if_year,
 )
 from pyaerocom.aeroval.json_utils import read_json, set_float_serialization_precision, write_json
-from pyaerocom.colocation_auto import ColocationSetup
+from pyaerocom.colocation_setup import ColocationSetup
 
 logger = logging.getLogger(__name__)
 
@@ -50,7 +50,7 @@ class OutputPaths(BaseModel):
         project ID
     exp_id : str
         experiment ID
-    json_basedir : str
+    json_basedir : str, Path
 
     """
 
@@ -422,13 +422,6 @@ def statistics_opts(self) -> StatisticsSetup:
         }
         return StatisticsSetup(**model_args)
 
-    ##################################
-    ## Non-BaseModel-based attributes
-    ##################################
-
-    # These attributes require special attention b/c they're not based on Pydantic's BaseModel class.
-
-    # TODO: Use Pydantic for ColocationSetup
     @computed_field
     @cached_property
     def colocation_opts(self) -> ColocationSetup:
@@ -438,7 +431,7 @@ def colocation_opts(self) -> ColocationSetup:
         model_args = {
             key: val
             for key, val in self.model_extra.items()
-            if key in ColocationSetup().__dict__.keys()
+            if key in ColocationSetup.model_fields
         }
         # need to pass some default values to the ColocationSetup if not provided in config
         default_dict = {"save_coldata": True, "keep_data": False, "resample_how": "mean"}
@@ -448,12 +441,11 @@ def colocation_opts(self) -> ColocationSetup:
 
         return ColocationSetup(**model_args)
 
-    @field_serializer("colocation_opts")
-    def serialize_colocation_opts(self, colocation_opts: ColocationSetup):
-        return colocation_opts.json_repr()
+    ##################################
+    ## Non-BaseModel-based attributes
+    ##################################
 
-    # ObsCollection and ModelCollection
-    # TODO Use Pydantic for ObsCollection and ModelCollection
+    # These attributes require special attention b/c they're not based on Pydantic's BaseModel class.
 
     obs_cfg: ObsCollection | dict = ObsCollection()
 
@@ -550,8 +542,8 @@ def _import_aux_funs(self) -> None:
 
     def _check_time_config(self) -> None:
         periods = self.time_cfg.periods
-        colstart = self.colocation_opts["start"]
-        colstop = self.colocation_opts["stop"]
+        colstart = self.colocation_opts.start
+        colstop = self.colocation_opts.stop
 
         if len(periods) == 0:
             if colstart is None:
@@ -576,15 +568,15 @@ def _check_time_config(self) -> None:
             if stop_yr == start_yr:
                 stop_yr += 1
             if colstart is None:
-                self.colocation_opts["start"] = start.strftime("%Y/%m/%d %H:%M:%S")
+                self.colocation_opts.start = start.strftime("%Y/%m/%d %H:%M:%S")
             if colstop is None:
-                self.colocation_opts["stop"] = stop.strftime(
+                self.colocation_opts.stop = stop.strftime(
                     "%Y/%m/%d %H:%M:%S"
                 )  # + 1  # add 1 year since we want to include stop year
         else:
             if colstart is None:
-                self.colocation_opts["start"] = start_yr
+                self.colocation_opts.start = start_yr
             if colstop is None:
-                self.colocation_opts["stop"] = (
+                self.colocation_opts.stop = (
                     stop_yr + 1
                 )  # add 1 year since we want to include stop year