diff --git a/mlair/data_handler/data_handler_wrf_chem.py b/mlair/data_handler/data_handler_wrf_chem.py index d1f9508d97037dcaab5faac956d66e916268623c..a1b7f96966cb15960096941b408fbd886df45749 100644 --- a/mlair/data_handler/data_handler_wrf_chem.py +++ b/mlair/data_handler/data_handler_wrf_chem.py @@ -148,10 +148,12 @@ class BaseWrfChemDataLoader: status = 'FAIL' print(f'{i}: {filenamestr} {status}') + @TimeTrackingWrapper def get_distances(self, lat, lon): dist = haversine_dist(lat1=self._data.XLAT, lon1=self._data.XLONG, lat2=lat, lon2=lon) return dist + @TimeTrackingWrapper def get_bearing(self, lat, lon, points_last=True): bearing = bearing_angle(lat1=lat, lon1=lon, lat2=self._data.XLAT, lon2=self._data.XLONG) if points_last: @@ -177,6 +179,7 @@ class BaseWrfChemDataLoader: self._data = data # logging.info('set dimensions as coordinates') + @TimeTrackingWrapper def apply_staged_transormation(self, mapping_of_stag2unstag=None): if mapping_of_stag2unstag is None: mapping_of_stag2unstag = {'U': 'U10', 'V': 'V10', 'U10': 'U10', 'V10': 'V10'} @@ -217,6 +220,7 @@ class BaseWrfChemDataLoader: self._data[u_ll_name] = ull self._data[v_ll_name] = vll + @TimeTrackingWrapper def set_interpolated_field(self, staged_field: xr.DataArray, target_field: xr.DataArray, dropped_staged_attrs: List[str] =None, **kwargs): stagger_attr_name = kwargs.pop('stagger_attr_name', 'stagger') @@ -277,6 +281,7 @@ class SingleGridColumnWrfChemDataLoader(BaseWrfChemDataLoader): # self._set_geoinfos() logging.debug("SingleGridColumnWrfChemDataLoader Initialised") + @TimeTrackingWrapper def __enter__(self): self.open_data() @@ -296,6 +301,7 @@ class SingleGridColumnWrfChemDataLoader(BaseWrfChemDataLoader): self.data.close() gc.collect() + @TimeTrackingWrapper def _set_geoinfos(self): # identify nearest coords self._set_nearest_icoords(dim=[self.logical_x_coord_name, self.logical_y_coord_name]) @@ -330,6 +336,7 @@ class SingleGridColumnWrfChemDataLoader(BaseWrfChemDataLoader): def geo_infos(self): return self._geo_infos + @TimeTrackingWrapper def _apply_external_coordinates(self): ds_coords = xr.open_dataset(self.external_coords_file, chunks={'south_north': 36, 'west_east': 40}) data = self._data @@ -337,7 +344,7 @@ class SingleGridColumnWrfChemDataLoader(BaseWrfChemDataLoader): data = data.assign_coords({k: (remove_items(list(v.dims), 'Time'), v.values.squeeze())}) self._data = data ds_coords.close() - logging.debug('setup external coords') + # logging.debug('setup external coords') def _set_coords(self, coords): __set_coords = dict(lat=None, lon=None) @@ -360,6 +367,7 @@ class SingleGridColumnWrfChemDataLoader(BaseWrfChemDataLoader): else: return self.__coords + @TimeTrackingWrapper def _set_nearest_icoords(self, dim=None): lat, lon = self.get_coordinates(as_arrays=True) self._nearest_icoords = dask.compute(self.compute_nearest_icoordinates(lat, lon, dim))[0] @@ -436,6 +444,7 @@ class DataHandlerSingleGridColumn(DataHandlerSingleStation): def setup_data_path(self, data_path: str, sampling: str): return data_path + @TimeTrackingWrapper def extract_data_from_loader(self, loader): data = loader.data.isel(loader.get_nearest_icoords()).squeeze()[self.variables] return data @@ -623,6 +632,7 @@ class DataHandlerSectorGrid(DataHandlerSingleGridColumn): self.wind_dir_name = None super().__init__(*args, **kwargs) + @TimeTrackingWrapper def extract_data_from_loader(self, loader): wind_dir_name = self._get_wind_dir_var_name(loader) full_data = loader.data.isel(loader.get_nearest_icoords()).squeeze() @@ -672,6 +682,7 @@ class DataHandlerSectorGrid(DataHandlerSingleGridColumn): opts=trafo_opts, transformation_dim=self.target_dim) return trafo + @TimeTrackingWrapper def modify_history(self): if self.transformation_is_applied: ws_edges = self.get_applied_transdormation_on_wind_sector_edges() @@ -720,6 +731,7 @@ class DataHandlerSectorGrid(DataHandlerSingleGridColumn): else: return self.history + @TimeTrackingWrapper def get_section_data_from_circle(self, grid_data, loader, sect): sec_data = grid_data.where( self.windsector.is_in_sector(sect, loader.geo_infos.bearing.drop('points').squeeze())) @@ -750,14 +762,17 @@ class DataHandlerSectorGrid(DataHandlerSingleGridColumn): return grid_data + @TimeTrackingWrapper def apply_aggregation_method_on_sector_data(self, data, loader): data = data.mean(dim=(loader.logical_x_coord_name, loader.logical_y_coord_name)) return data + @TimeTrackingWrapper def compute_wind_dir_of_interest(self): wind_dir_of_intrest = self.history.sel({self.target_dim: self.wind_dir_name, self.window_dim: 0}) return wind_dir_of_intrest + @TimeTrackingWrapper def get_applied_transdormation_on_wind_sector_edges(self): ws_edges = self._get_left_and_right_wind_sector_edges(return_as='xr.da', dim=self.wind_sector_edge_dim_name) ws_edges = self.apply_transformation_on_data(ws_edges)