can run through w/o crashing but modify output

pyaerocom/colocateddata.py

@@ -1056,12 +1056,22 @@ def rename_variable(self, var_name, new_var_name, data_source, inplace=True):
 
     @staticmethod
     def _aerocom_savename(
-        obs_var, obs_id, mod_var, mod_id, start_str, stop_str, ts_type, filter_name
+        obs_var, obs_id, mod_var, mod_id, start_str, stop_str, ts_type, filter_name, vertical_layer
     ):
-        return (
-            f"{mod_var}_{obs_var}_MOD-{mod_id}_REF-{obs_id}_"
-            f"{start_str}_{stop_str}_{ts_type}_{filter_name}"
-        )
+        if (
+            not vertical_layer is None
+        ):  # LB: Note this is in beta and needs testing. Probably some positional issues
+            start = vertical_layer["start"] / 1000
+            end = vertical_layer["end"] / 1000
+            return (
+                f"{mod_var}_{obs_var}_MOD-{mod_id}_REF-{obs_id}_"
+                f"{start_str}_{stop_str}_{ts_type}_{filter_name}_{start}-{end}km"
+            )
+        else:
+            return (
+                f"{mod_var}_{obs_var}_MOD-{mod_id}_REF-{obs_id}_"
+                f"{start_str}_{stop_str}_{ts_type}_{filter_name}"
+            )
 
     @property
     def savename_aerocom(self):
@@ -1715,7 +1725,7 @@ def apply_region_mask(self, region_id, inplace=False):
 
             drop_idx = []
             nstats = len(arr.station_name)
-            for (lat, lon, stat) in data._iter_stats():
+            for lat, lon, stat in data._iter_stats():
                 if get_mask_value(lat, lon, mask) < 1:
                     drop_idx.append(stat)
 

pyaerocom/colocation_auto.py

@@ -813,7 +813,9 @@ def run(self, var_list: list = None, **opts):
         self._print_coloc_info(vars_to_process)
         for mod_var, obs_var in vars_to_process.items():
             try:
-                coldata = self._run_helper(mod_var, obs_var)
+                coldata = self._run_helper(
+                    mod_var, obs_var
+                )  # note this can be ColocatedData or ColocatedDataLists
                 if not mod_var in data_out:
                     data_out[mod_var] = {}
                 data_out[mod_var][obs_var] = coldata
@@ -1287,7 +1289,13 @@ def _save_coldata(self, coldata):
             coldata.rename_variable(mod_var, mvar, self.model_id)
         else:
             mvar = mod_var
-        savename = self._coldata_savename(obs_var, mvar, coldata.ts_type)
+        if coldata.vertical_layer:
+            savename = self._coldata_savename(
+                obs_var, mvar, coldata.ts_type, vertical_layer=coldata.vertical_layer
+            )
+
+        else:
+            savename = self._coldata_savename(obs_var, mvar, coldata.ts_type)
         fp = coldata.to_netcdf(self.output_dir, savename=savename)
         self.files_written.append(fp)
         msg = f"WRITE: {fp}\n"
@@ -1341,8 +1349,12 @@ def _check_set_start_stop(self):
                 )
         self.start, self.stop = start_stop(self.start, self.stop)
 
-    def _coldata_savename(self, obs_var, mod_var, ts_type):
+    def _coldata_savename(self, obs_var, mod_var, ts_type, **kwargs):
         """Get filename of colocated data file for saving"""
+        if "vertical_layer" in kwargs:
+            vertical_layer = kwargs["vertical_layer"]
+        else:
+            vertical_layer = None
         name = ColocatedData._aerocom_savename(
             obs_var=obs_var,
             obs_id=self.get_obs_name(),
@@ -1352,6 +1364,7 @@ def _coldata_savename(self, obs_var, mod_var, ts_type):
             stop_str=self.get_stop_str(),
             ts_type=ts_type,
             filter_name=self.filter_name,
+            vertical_layer=vertical_layer,
         )
         return f"{name}.nc"
 
@@ -1450,8 +1463,6 @@ def _run_helper(self, model_var: str, obs_var: str):
         args = self._check_dimensionality(args)
         coldata = self._colocation_func(**args)
 
-        breakpoint()
-
         if isinstance(coldata, ColocatedData):
             coldata.data.attrs["model_name"] = self.get_model_name()
             coldata.data.attrs["obs_name"] = self.get_obs_name()
@@ -1465,8 +1476,6 @@ def _run_helper(self, model_var: str, obs_var: str):
             if self.save_coldata:
                 self._save_coldata(coldata)
 
-            return coldata
-
         elif isinstance(coldata, ColocatedDataLists):
             breakpoint()
             for i_list in coldata:
@@ -1477,15 +1486,19 @@ def _run_helper(self, model_var: str, obs_var: str):
                     coldata_obj.data.attrs.update(**self.add_meta)
                     if self.zeros_to_nan:
                         coldata_obj = coldata_obj.set_zeros_nan()
-                    if self.model_to_stp:
+                    if self.model_to_stp:  # Lb: check is this needs modifying
                         coldata = correct_model_stp_coldata(coldata_obj)
                     if self.save_coldata:
                         self._save_coldata(coldata_obj)
+            breakpoint()
+
         else:
             raise Exception(
                 f"Invalid coldata type returned by colocation function {self._colocation_func}"
             )
 
+        return coldata
+
     def _print_coloc_info(self, var_matches):
         if not var_matches:
             logger.info("Nothing to colocate")