diff --git a/mlair/data_handler/data_handler_wrf_chem.py b/mlair/data_handler/data_handler_wrf_chem.py index 81a1458a400fd6eb3e7e586c7f2c0f5ade039af2..227a609b33c9ecd10c88cd89147e366fe8c34539 100644 --- a/mlair/data_handler/data_handler_wrf_chem.py +++ b/mlair/data_handler/data_handler_wrf_chem.py @@ -45,8 +45,10 @@ class BaseWrfChemDataLoader: DEFAULT_STAGED_ROTATION_opts = dict(cen_lon=12., cen_lat=52.5, truelat1=30., truelat2=60., stand_lon=12.) - DEFAULT_VARS_TO_ROTATE = ((('U', 'V'), ('Ull', 'Vll')), (('U10', 'V10'), ('U10ll', 'V10ll')), - (('U', 'V'), ('wspdll', 'wdirll')), (('U10', 'V10'), ('wspd10ll', 'wdir10ll')) + DEFAULT_VARS_TO_ROTATE = ((('U', 'V'), ('Ull', 'Vll')), + (('U10', 'V10'), ('U10ll', 'V10ll')), + (('U', 'V'), ('wspdll', 'wdirll')), + (('U10', 'V10'), ('wspd10ll', 'wdir10ll')) ) def __init__(self, data_path: str, common_file_starter: str = DEFAULT_FILE_STARTER, @@ -384,6 +386,8 @@ class DataHandlerSingleGridColumn(DataHandlerSingleStation): self.targetvar_logical_z_coord_selector = self._ret_z_coord_select_if_valid(targetvar_logical_z_coord_selector, as_input=False) self._logical_z_coord_name = None + self._joint_z_coord_selector = self._extract_largest_coord_extractor(self.var_logical_z_coord_selector, + self.targetvar_logical_z_coord_selector) super().__init__(*args, **kwargs) @staticmethod @@ -431,10 +435,13 @@ class DataHandlerSingleGridColumn(DataHandlerSingleStation): self._logical_z_coord_name = sgc_loader.logical_z_coord_name # # select defined variables at grid box or grid coloumn based on nearest icoords data = self.extract_data_from_loader(sgc_loader, station) + if self._joint_z_coord_selector is not None: + data = data.sel({self._logical_z_coord_name: self._joint_z_coord_selector}) # expand dimesion for iterdim data = data.expand_dims({self.iter_dim: station}).to_array(self.target_dim) # transpose dataarray: set first three fixed and keep remaining as is data = data.transpose(self.iter_dim, self.time_dim, self.target_dim, ...) + data = dask.compute(self._slice_prep(data, start=start, end=end))[0] sgc_loader.data.close() # data = self.check_for_negative_concentrations(data) @@ -445,6 +452,14 @@ class DataHandlerSingleGridColumn(DataHandlerSingleStation): meta = None return data, meta + @staticmethod + def _extract_largest_coord_extractor(var_extarctor, target_extractor) -> Union[List, None]: + if var_extarctor is not None and target_extractor is not None: + res = list(set(to_list(var_extarctor) + to_list(target_extractor))) + else: + res = None + return res + def get_X(self, upsampling=False, as_numpy=False): if as_numpy is True: # return None