diff --git a/pyaerocom/colocateddata.py b/pyaerocom/colocateddata.py index 5f0cdacf4..aca715534 100644 --- a/pyaerocom/colocateddata.py +++ b/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) diff --git a/pyaerocom/colocation_auto.py b/pyaerocom/colocation_auto.py index bbcf13368..24dd1a9c2 100644 --- a/pyaerocom/colocation_auto.py +++ b/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")