diff --git a/pymech/dataset.py b/pymech/dataset.py index 1d5f5ae..32ac6ee 100644 --- a/pymech/dataset.py +++ b/pymech/dataset.py @@ -10,6 +10,7 @@ import xarray as xr from .neksuite import readnek +from .log import logger def open_dataset(path, **kwargs): @@ -50,6 +51,7 @@ def open_mfdataset( parallel=False, join="outer", attrs_file=None, + box_mesh=True, **kwargs, ): """Helper function for opening multiple files as an xarray_ dataset. @@ -79,7 +81,7 @@ def open_mfdataset( if isinstance(concat_dim, (str, xr.DataArray)) or concat_dim is None: concat_dim = [concat_dim] - open_kwargs = dict() + open_kwargs = dict(box_mesh=box_mesh) if parallel: import dask @@ -140,44 +142,55 @@ def open_mfdataset( return combined -def _open_nek_dataset(path): +def _open_nek_dataset(path, box_mesh=True): """Interface for converting Nek field files into xarray_ datasets.""" field = readnek(path) if isinstance(field, int): raise IOError("Failed to load {}".format(path)) elements = field.elem - elem_stores = [_NekDataStore(elem) for elem in elements] - elem_dsets = [ - xr.Dataset.load_store(store).set_coords(store.axes) - for store in elem_stores - ] + if box_mesh: + DataStore = _NekDataStoreBox + coords = ('z', 'y', 'x') + else: + DataStore = _NekDataStoreGeneric + coords = ('zmesh', 'ymesh', 'xmesh') + + elem_stores = [DataStore(elem) for elem in elements] + try: + elem_dsets = [ + xr.Dataset.load_store(store).set_coords(coords) + for store in elem_stores + ] + except ValueError: + logger.critical( + "Cannot load as a box mesh. Try passing the keyword argument: " + "`box_mesh=False` while opening the dataset." + ) + 1/0 # See: https://github.com/MITgcm/xmitgcm/pull/200 - if xr.__version__ < '0.15.2': - ds = xr.combine_by_coords(elem_dsets) + if box_mesh: + if xr.__version__ < '0.15.2': + ds = xr.combine_by_coords(elem_dsets) + else: + ds = xr.combine_by_coords(elem_dsets, combine_attrs="drop") else: - ds = xr.combine_by_coords(elem_dsets, combine_attrs="drop") + # This would work if we pre-sort elem_dsets and pass as a nested list / + # iterator + ds = xr.combine_nested(elem_dsets, concat_dim=('z', 'y', 'x')) ds.coords.update({"time": field.time}) return ds -class _NekDataStore(xr.backends.common.AbstractDataStore): +class _NekDataStoreGeneric(xr.backends.common.AbstractDataStore): """Xarray store for a Nek field element.""" def __init__(self, elem): self.elem = elem self.axes = ("z", "y", "x") - def meshgrid_to_dim(self, mesh): - """Reverse of np.meshgrid. This method extracts one-dimensional - coordinates from a cubical array format for every direction - - """ - dim = np.unique(np.round(mesh, 8)) - return dim - def get_dimensions(self): return self.axes @@ -196,9 +209,6 @@ def get_variables(self): elem = self.elem data_vars = { - ax[2]: self.meshgrid_to_dim(elem.pos[0]), # x - ax[1]: self.meshgrid_to_dim(elem.pos[1]), # y - ax[0]: self.meshgrid_to_dim(elem.pos[2]), # z "xmesh": (ax, elem.pos[0]), "ymesh": (ax, elem.pos[1]), "zmesh": (ax, elem.pos[2]), @@ -221,3 +231,34 @@ def get_variables(self): ) return data_vars + + +class _NekDataStoreBox(_NekDataStoreGeneric): + """Useful for box meshes where coordinates can be simplified into 1D + arrays. + + """ + + def meshgrid_to_dim(self, mesh): + """Reverse of np.meshgrid. This method extracts one-dimensional + coordinates from a cubical array format for every direction + + """ + dim = np.unique(np.round(mesh, 8)) + return dim + + def get_variables(self): + """Generate an xarray dataset from a single element.""" + ax = self.axes + elem = self.elem + + data_vars = super().get_variables() + + data_vars.update( + { + ax[2]: self.meshgrid_to_dim(elem.pos[0]), # x + ax[1]: self.meshgrid_to_dim(elem.pos[1]), # y + ax[0]: self.meshgrid_to_dim(elem.pos[2]), # z + } + ) + return data_vars