diff --git a/mlair/data_handler/data_handler_wrf_chem.py b/mlair/data_handler/data_handler_wrf_chem.py index 82bf3fe17f65be5d2422f8bde7ca42356cf0f67b..c89442c0b4e19a0f523c89c31e3d7abd3c58065b 100644 --- a/mlair/data_handler/data_handler_wrf_chem.py +++ b/mlair/data_handler/data_handler_wrf_chem.py @@ -3,7 +3,9 @@ __data__ = "2021-02-03" import xarray as xr import numpy as np +import pandas as pd import itertools +import glob import matplotlib.pyplot as plt from dask.diagnostics import ProgressBar @@ -60,6 +62,7 @@ class BaseWrfChemDataLoader: physical_t_coord_name: str = DEFAULT_PHYSICAL_TIME_COORD_NAME, staged_vars: List[str] = DEFAULT_STAGED_VARS, variables=None, z_coord_selector=None, + start_time=None, end_time=None, staged_rotation_opts: Dict = DEFAULT_STAGED_ROTATION_opts, vars_to_rotate: Tuple[Tuple[Tuple[str, str], Tuple[str, str]]] = DEFAULT_VARS_TO_ROTATE, staged_dimension_mapping=None, stag_ending='_stag', @@ -84,6 +87,9 @@ class BaseWrfChemDataLoader: self.variables = variables self.z_coord_selector = z_coord_selector + self.start_time = start_time + self.end_time = end_time + if rechunk_values is None: self.rechunk_values = {self.time_dim_name: 1} else: @@ -118,7 +124,19 @@ class BaseWrfChemDataLoader: @property def dataset_search_str(self): - return os.path.join(self.data_path, self.common_file_starter + '*') + if (self.start_time is None) and (self.end_time is None): + return os.path.join(self.data_path, self.common_file_starter + '*') + elif (self.start_time is not None) and (self.end_time is not None): + path_list = [] + for day in pd.date_range(self.start_time, self.end_time): + path_list += sorted(glob.glob(os.path.join(self.data_path, + self.common_file_starter) + '*' + day.strftime( + "%Y-%m-%d") + '*' + )) + + return path_list + else: + raise ValueError(f"`start_time' and `end_time' must both be given or None.") @TimeTrackingWrapper def open_data(self): @@ -496,7 +514,9 @@ class DataHandlerSingleGridColumn(DataHandlerSingleStation): time_dim_name=self.time_dim, rechunk_values=self.rechunk_values, variables=self.variables, - z_coord_selector=self._joint_z_coord_selector + z_coord_selector=self._joint_z_coord_selector, + start_time=self.start, + end_time=self.end, # preprocess_open_mfdataset=preprocess_fkt_for_loader, ) @@ -937,7 +957,7 @@ if __name__ == '__main__': wrf_new = SingleGridColumnWrfChemDataLoader((lat_xr, lon_xr), data_path='/home/felix/Data/WRF-Chem/upload_aura_2021-02-24/2009/', - common_file_starter='wrfout_d0', + common_file_starter='wrfout_d01_', time_dim_name='XTIME', logical_x_coord_name='west_east', logical_y_coord_name='south_north',