diff --git a/mlair/data_handler/data_handler_wrf_chem.py b/mlair/data_handler/data_handler_wrf_chem.py index ae57c93f9d303fec4df3ba7c5dd051e064374d85..03bf50f1bd24f15252867ac52bd633d76865155b 100644 --- a/mlair/data_handler/data_handler_wrf_chem.py +++ b/mlair/data_handler/data_handler_wrf_chem.py @@ -9,6 +9,7 @@ import dask import dask.array as da import os from mlair.helpers.geofunctions import haversine_dist +from mlair.helpers.helpers import convert2xrda from typing import Tuple, Union, List, Dict from mlair.data_handler.abstract_data_handler import AbstractDataHandler @@ -131,6 +132,8 @@ class DataHandlerSingleGridCoulumn(WrfChemDataHandler): self.target_var = target_var self.window_history_size = window_history_size self.window_lead_time = window_lead_time + self._nearest_icoords = None + self._nearest_coords = None self.open_data() self.rechunk_data(self.rechunk_values) @@ -160,29 +163,31 @@ class DataHandlerSingleGridCoulumn(WrfChemDataHandler): 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] + self._nearest_icoords = dask.compute(self.compute_nearest_icoordinates(lat, lon, dim))[0] def get_nearest_icoords(self, as_arrays=False): if as_arrays: - return (self.__nearest_icoords[self.logical_y_coord_name].values, - self.__nearest_icoords[self.logical_x_coord_name].values, + return (self._nearest_icoords[self.logical_y_coord_name].values, + self._nearest_icoords[self.logical_x_coord_name].values, ) else: - return {k: list(v.values) for k, v in self.__nearest_icoords.items()} + return {k: list(v.values) for k, v in self._nearest_icoords.items()} def _set_nearest_coords(self): icoords = self.get_nearest_icoords() - lat = self._data[self.physical_y_coord_name].sel(icoords) - lon = self._data[self.physical_x_coord_name].sel(icoords) - self.__nearest_coords = dask.compute(dict(lat=lat, lon=lon))[0] + ilat = convert2xrda(np.array(icoords[self.logical_y_coord_name]), use_1d_default=True) + ilon = convert2xrda(np.array(icoords[self.logical_x_coord_name]), use_1d_default=True) + lat = self._data[self.physical_y_coord_name].isel({'x': ilon, 'y': ilat}) + lon = self._data[self.physical_x_coord_name].isel({'x': ilon, 'y': ilat}) + self._nearest_coords = dict(lat=lat, lon=lon) def get_nearest_coords(self, as_arrays=False): if as_arrays: - return (self.__nearest_coords['lat'].values, - self.__nearest_coords['lon'].values, + return (self._nearest_coords['lat'].values, + self._nearest_coords['lon'].values, ) else: - return {k: v.values for k, v in self.__nearest_coords.items()} + return {k: list(v.values) for k, v in self._nearest_coords.items()} @@ -214,24 +219,26 @@ if __name__ == '__main__': plt.close('all') - lat_np = np.array([50.73333, 45.0]) - lon_np = np.array([7.1, 0.0]) + lat_np = np.array([50.73333]) + lon_np = np.array([7.1]) lat_xr = xr.DataArray(lat_np, dims=["points"], coords={'points': range(len(lat_np))}) lon_xr = xr.DataArray(lon_np, dims=["points"], coords={'points': range(len(lon_np))}) use_first_dummy_dataset = True if use_first_dummy_dataset: - wrf_dh = WrfChemDataHandler(data_path='/home/felix/Data/WRF-Chem/', - common_file_starter='wrfout_d01_2010-', - time_dim_name='Time', - ) + wrf_gridcol = DataHandlerSingleGridCoulumn((lat_xr, lon_xr), data_path='/home/felix/Data/WRF-Chem/', common_file_starter='wrfout_d01_2010-', time_dim_name='Time', ) + wrf_gridcol.get_nearest_coords() + wrf_dh = WrfChemDataHandler(data_path='/home/felix/Data/WRF-Chem/', + common_file_starter='wrfout_d01_2010-', + time_dim_name='Time', + ) wrf_dh.open_data() wrf_dh.rechunk_data({"XTIME": 1, "y": 36, "x": 40}) T2 = wrf_dh._data.T2