From 39c60825f014611a94ecac86e3fa40d7bd40901c Mon Sep 17 00:00:00 2001 From: Timothy Smith Date: Wed, 1 Apr 2020 18:42:09 -0600 Subject: [PATCH 01/33] access to llc2160 grid is a go, with a couple caveats for now --- xmitgcm/llcreader/known_models.py | 11 +++++- xmitgcm/llcreader/llcmodel.py | 57 ++++++++++++++++++++++++++----- xmitgcm/llcreader/stores.py | 12 +++++-- xmitgcm/variables.py | 2 +- 4 files changed, 69 insertions(+), 13 deletions(-) diff --git a/xmitgcm/llcreader/known_models.py b/xmitgcm/llcreader/known_models.py index 3b2ea59b..b1f9c632 100644 --- a/xmitgcm/llcreader/known_models.py +++ b/xmitgcm/llcreader/known_models.py @@ -38,6 +38,14 @@ class LLC2160Model(BaseLLCModel): varnames = ['Eta', 'KPPhbl', 'oceFWflx', 'oceQnet', 'oceQsw', 'oceSflux', 'oceTAUX', 'oceTAUY', 'PhiBot', 'Salt', 'SIarea', 'SIheff', 'SIhsalt', 'SIhsnow', 'SIuice', 'SIvice', 'Theta', 'U', 'V', 'W'] + grid_varnames = ['AngleCS','AngleSN','Depth','DXC','DXG','DYC','DYG', + 'hFacC','hFacS','hFacW','RAC','RAS','RAW', + 'rLowC','rLowS','rLowW','rSurfC', + 'rSurfS','rSurfW','XC','YC'] + # unrecognized name: 'RhoRef' + # corner point problems: 'RAZ','XG','YG' + # k_p1 problems: 'DRC', 'PHrefF', 'RF' + # k problems: 'DRF', 'PHrefC', 'RC' mask_override = {'oceTAUX': 'c', 'oceTAUY': 'c'} @@ -61,9 +69,10 @@ class ECCOPortalLLC2160Model(LLC2160Model): def __init__(self): fs = _make_http_filesystem() base_path = 'https://data.nas.nasa.gov/ecco/download_data.php?file=/eccodata/llc_2160/compressed' + grid_path = 'https://data.nas.nasa.gov/ecco/download_data.php?file=/eccodata/llc_2160/grid' mask_path = 'https://storage.googleapis.com/pangeo-ecco/llc/masks/llc_2160_masks.zarr/' store = stores.NestedStore(fs, base_path=base_path, mask_path=mask_path, - shrunk=True, join_char='/') + grid_path=grid_path,shrunk=True, join_char='/') super(ECCOPortalLLC2160Model, self).__init__(store) diff --git a/xmitgcm/llcreader/llcmodel.py b/xmitgcm/llcreader/llcmodel.py index 7b302744..b12e370e 100644 --- a/xmitgcm/llcreader/llcmodel.py +++ b/xmitgcm/llcreader/llcmodel.py @@ -11,17 +11,38 @@ def _get_var_metadata(): # The LLC run data comes with zero metadata. So we import metadata from # the xmitgcm package. - from ..variables import state_variables, package_state_variables + from ..variables import (state_variables, package_state_variables, + horizontal_coordinates_llc, vertical_coordinates, + horizontal_grid_variables, vertical_grid_variables, + volume_grid_variables, mask_variables, + extra_grid_variables) from ..utils import parse_available_diagnostics from ..default_diagnostics import diagnostics + from ..mds_store import _get_all_grid_variables from io import StringIO + # get grid info + grid_metadata = horizontal_coordinates_llc + grid_metadata.update(vertical_coordinates) + grid_metadata.update(horizontal_grid_variables) + grid_metadata.update(vertical_grid_variables) + grid_metadata.update(volume_grid_variables) + grid_metadata.update(mask_variables) + grid_metadata.update(extra_grid_variables) + diag_file = StringIO(diagnostics) available_diags = parse_available_diagnostics(diag_file) var_metadata = state_variables var_metadata.update(package_state_variables) var_metadata.update(available_diags) + # add grid vars + for key,val in grid_metadata.items(): + if 'filename' in val: + var_metadata[val['filename']] = val + else: + var_metadata[key] = val + # even the file names from the LLC data differ from standard MITgcm output aliases = {'Eta': 'ETAN', 'PhiBot': 'PHIBOT', 'Salt': 'SALT', 'Theta': 'THETA'} @@ -445,6 +466,7 @@ class BaseLLCModel: iter_stop = None iter_step = None varnames = [] + grid_varnames = [] mask_override = {} def __init__(self, store): @@ -525,7 +547,7 @@ def _get_mask_and_index_for_variable(self, vname): def _dask_array(self, nfacet, varname, iters, klevels, k_chunksize): # return a dask array for a single facet facet_shape = _facet_shape(nfacet, self.nx) - time_chunks = (len(iters) * (1,),) + time_chunks = (len(iters) * (1,),) if iters is not None else () k_chunks = (tuple([len(c) for c in _chunks(klevels, k_chunksize)]),) chunks = time_chunks + k_chunks + tuple([(s,) for s in facet_shape]) @@ -534,10 +556,18 @@ def _dask_array(self, nfacet, varname, iters, klevels, k_chunksize): dsk = {} token = tokenize(varname, self.store, nfacet) name = '-'.join([varname, token]) - for n_iter, iternum in enumerate(iters): + if iters is not None: + for n_iter, iternum in enumerate(iters): + for n_k, these_klevels in enumerate(_chunks(klevels, k_chunksize)): + key = name, n_iter, n_k, 0, 0, 0 + task = (_get_facet_chunk, self.store, varname, iternum, + nfacet, these_klevels, self.nx, self.nz, self.dtype, + self.mask_override) + dsk[key] = task + else: for n_k, these_klevels in enumerate(_chunks(klevels, k_chunksize)): - key = name, n_iter, n_k, 0, 0, 0 - task = (_get_facet_chunk, self.store, varname, iternum, + key = name, n_k, 0, 0, 0 + task = (_get_facet_chunk, self.store, varname, None, nfacet, these_klevels, self.nx, self.nz, self.dtype, self.mask_override) dsk[key] = task @@ -565,7 +595,7 @@ def _get_facet_data(self, varname, iters, klevels, k_chunksize): def get_dataset(self, varnames=None, iter_start=None, iter_stop=None, iter_step=None, k_levels=None, k_chunksize=1, - type='faces'): + type='faces',grid_vars_to_coords=True): """ Create an xarray Dataset object for this model. @@ -625,22 +655,33 @@ def _if_not_none(a, b): self._get_facet_data(vname, iters, k_levels, k_chunksize) for vname in varnames} + # get the grid in facet form + grid_facets = {vname: + self._get_facet_data(vname, None, k_levels, k_chunksize) + for vname in self.grid_varnames} + # transform it into faces or latlon data_transformers = {'faces': _all_facets_to_faces, 'latlon': _all_facets_to_latlon} transformer = data_transformers[type] data = transformer(data_facets, _VAR_METADATA) + data.update(transformer(grid_facets, _VAR_METADATA)) variables = {} - for vname in varnames: + for vname in varnames+self.grid_varnames: meta = _VAR_METADATA[vname] dims = meta['dims'] if type=='faces': dims = _add_face_to_dims(dims) - dims = ['time',] + dims + dims = ['time',] + dims if vname not in self.grid_varnames else dims attrs = meta['attrs'] + print('vname: ',vname) variables[vname] = xr.Variable(dims, data[vname], attrs) ds = ds.update(variables) + if grid_vars_to_coords: + for gname in self.grid_varnames: + ds = ds.set_coords(gname) + return ds diff --git a/xmitgcm/llcreader/stores.py b/xmitgcm/llcreader/stores.py index 6748cce2..eda9292d 100644 --- a/xmitgcm/llcreader/stores.py +++ b/xmitgcm/llcreader/stores.py @@ -22,22 +22,28 @@ class BaseStore: """ def __init__(self, fs, base_path='/', shrunk=False, - mask_fs=None, mask_path=None, join_char=None): + mask_fs=None, mask_path=None, + grid_fs=None, grid_path=None, join_char=None): self.base_path = base_path self.fs = fs self.shrunk = shrunk self.mask_fs = mask_fs or self.fs self.mask_path = mask_path + self.grid_fs = grid_fs or self.fs + self.grid_path = grid_path self.join_char = join_char if shrunk and (mask_path is None): raise ValueError("`mask_path` can't be None if `shrunk` is True") def _directory(self, varname, iternum): - return self.base_path + if iternum is not None: + return self.base_path + else: + return self.grid_path def _fname(self, varname, iternum): - fname = varname + '.%010d.data' % iternum + fname = varname + '.%010d.data' % iternum if iternum is not None else varname if self.shrunk: fname += '.shrunk' return fname diff --git a/xmitgcm/variables.py b/xmitgcm/variables.py index 24f16916..9d9defac 100644 --- a/xmitgcm/variables.py +++ b/xmitgcm/variables.py @@ -187,7 +187,7 @@ rAs=dict(dims=["j_g", "i"], attrs=dict( standard_name="cell_area_at_v_location", long_name="cell area", units="m2", coordinates="YG XC"), - filename='RAZ'), + filename='RAS'), ) vertical_grid_variables = OrderedDict( From f8609b8df6cf490a9fb07a5494b8a8079282e22a Mon Sep 17 00:00:00 2001 From: Timothy Smith Date: Thu, 2 Apr 2020 11:10:15 -0600 Subject: [PATCH 02/33] use grid var names not filenames, get vertical grid vars working --- xmitgcm/llcreader/known_models.py | 9 ++-- xmitgcm/llcreader/llcmodel.py | 80 ++++++++++++++++++++++++++++--- 2 files changed, 78 insertions(+), 11 deletions(-) diff --git a/xmitgcm/llcreader/known_models.py b/xmitgcm/llcreader/known_models.py index b1f9c632..3dff1198 100644 --- a/xmitgcm/llcreader/known_models.py +++ b/xmitgcm/llcreader/known_models.py @@ -38,14 +38,13 @@ class LLC2160Model(BaseLLCModel): varnames = ['Eta', 'KPPhbl', 'oceFWflx', 'oceQnet', 'oceQsw', 'oceSflux', 'oceTAUX', 'oceTAUY', 'PhiBot', 'Salt', 'SIarea', 'SIheff', 'SIhsalt', 'SIhsnow', 'SIuice', 'SIvice', 'Theta', 'U', 'V', 'W'] - grid_varnames = ['AngleCS','AngleSN','Depth','DXC','DXG','DYC','DYG', - 'hFacC','hFacS','hFacW','RAC','RAS','RAW', - 'rLowC','rLowS','rLowW','rSurfC', + grid_varnames = ['AngleCS','AngleSN','Depth', + 'DRC','DRF','DXC','DXG','DYC','DYG', + 'hFacC','hFacS','hFacW','PHrefC','PHrefF','RAC','RAS','RAW', + 'RC','RF','rLowC','rLowS','rLowW','rSurfC', 'rSurfS','rSurfW','XC','YC'] # unrecognized name: 'RhoRef' # corner point problems: 'RAZ','XG','YG' - # k_p1 problems: 'DRC', 'PHrefF', 'RF' - # k problems: 'DRF', 'PHrefC', 'RC' mask_override = {'oceTAUX': 'c', 'oceTAUY': 'c'} diff --git a/xmitgcm/llcreader/llcmodel.py b/xmitgcm/llcreader/llcmodel.py index b12e370e..3077a29d 100644 --- a/xmitgcm/llcreader/llcmodel.py +++ b/xmitgcm/llcreader/llcmodel.py @@ -39,6 +39,7 @@ def _get_var_metadata(): # add grid vars for key,val in grid_metadata.items(): if 'filename' in val: + val.update({'real_name':key}) var_metadata[val['filename']] = val else: var_metadata[key] = val @@ -108,6 +109,7 @@ def _decompress(data, mask, dtype): _facet_reshape = (False, False, False, True, True) _nfaces = 13 _nfacets = 5 +_vgridvars = ['DRF','PHrefC','RC','DRF','PHrefF','RF'] def _uncompressed_facet_index(nfacet, nside): face_size = nside**2 @@ -420,6 +422,33 @@ def _get_facet_chunk(store, varname, iternum, nfacet, klevels, nx, nz, dtype, return np.concatenate(level_data, axis=1) +def _get_1d_chunk(store, varname, klevels, nz, dtype): + """for 1D vertical grid variables""" + + fs, path = store.get_fs_and_full_path(varname, None) + + file = fs.open(path) + + # insert singleton axis for k level + facet_shape = (1,) + level_data = [] + + for k in klevels: + assert (k >= 0) & (k < nz) + + # with a 1D file, k level gives "where to read" in file + read_offset = k * dtype.itemsize # in bytes + read_length = dtype.itemsize # in bytes + file.seek(read_offset) + buffer = file.read(read_length) + data = np.frombuffer(buffer, dtype=dtype) + assert len(data) == 1 + + # this is the shape this facet is supposed to have + data.shape = facet_shape + level_data.append(data) + + return np.concatenate(level_data, axis=0) class BaseLLCModel: """Class representing an LLC Model Dataset. @@ -556,6 +585,8 @@ def _dask_array(self, nfacet, varname, iters, klevels, k_chunksize): dsk = {} token = tokenize(varname, self.store, nfacet) name = '-'.join([varname, token]) + + # iters == None for grid variables if iters is not None: for n_iter, iternum in enumerate(iters): for n_k, these_klevels in enumerate(_chunks(klevels, k_chunksize)): @@ -574,6 +605,23 @@ def _dask_array(self, nfacet, varname, iters, klevels, k_chunksize): return dsa.Array(dsk, name, chunks, self.dtype) + def _dask_array_vgrid(self, varname, klevels, k_chunksize): + # return a dask array for a 1D vertical grid var + chunks = (tuple([len(c) + for c in _chunks(klevels, k_chunksize)]),) + + # manually build dask graph + dsk = {} + token = tokenize(varname, self.store) + name = '-'.join([varname, token]) + + for n_k, these_klevels in enumerate(_chunks(klevels, k_chunksize)): + key = name, n_k + task = (_get_1d_chunk, self.store, varname, + these_klevels, self.nz, self.dtype) + dsk[key] = task + + return dsa.Array(dsk, name, chunks, self.dtype) def _get_facet_data(self, varname, iters, klevels, k_chunksize): mask, index = self._get_mask_and_index_for_variable(varname) @@ -583,7 +631,10 @@ def _get_facet_data(self, varname, iters, klevels, k_chunksize): if len(dims)==2: klevels = [0,] - data_facets = [self._dask_array(nfacet, varname, iters, klevels, k_chunksize) + if varname in _vgridvars: + data_facets = self._dask_array_vgrid(varname,klevels,k_chunksize) + else: + data_facets = [self._dask_array(nfacet, varname, iters, klevels, k_chunksize) for nfacet in range(5)] if len(dims)==2: @@ -666,9 +717,20 @@ def _if_not_none(a, b): transformer = data_transformers[type] data = transformer(data_facets, _VAR_METADATA) - data.update(transformer(grid_facets, _VAR_METADATA)) + + # separate horizontal and vertical grid variables + hgrid_names = [x for x in self.grid_varnames if x not in _vgridvars] + vgrid_names = [x for x in self.grid_varnames if x in _vgridvars] + + hgrid_facets = {key: grid_facets[key] for key in hgrid_names} + vgrid_facets = {key: grid_facets[key] for key in vgrid_names} + + # do not transform vertical grid variables + data.update(transformer(hgrid_facets, _VAR_METADATA)) + data.update(vgrid_facets) variables = {} + gridvar_names = [] for vname in varnames+self.grid_varnames: meta = _VAR_METADATA[vname] dims = meta['dims'] @@ -676,12 +738,18 @@ def _if_not_none(a, b): dims = _add_face_to_dims(dims) dims = ['time',] + dims if vname not in self.grid_varnames else dims attrs = meta['attrs'] - print('vname: ',vname) - variables[vname] = xr.Variable(dims, data[vname], attrs) + + # Handle grid names different from filenames + fname = vname + vname = meta['real_name'] if 'real_name' in meta else vname + if fname in self.grid_varnames: + gridvar_names.append(vname) + + variables[vname] = xr.Variable(dims, data[fname], attrs) ds = ds.update(variables) + if grid_vars_to_coords: - for gname in self.grid_varnames: - ds = ds.set_coords(gname) + ds = ds.set_coords(gridvar_names) return ds From 83e39eb7eee0ad63dc3e4b6d323ab24a14c70a6e Mon Sep 17 00:00:00 2001 From: Timothy Smith Date: Thu, 2 Apr 2020 11:16:01 -0600 Subject: [PATCH 03/33] lil typo --- xmitgcm/llcreader/llcmodel.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/xmitgcm/llcreader/llcmodel.py b/xmitgcm/llcreader/llcmodel.py index 3077a29d..440a7aeb 100644 --- a/xmitgcm/llcreader/llcmodel.py +++ b/xmitgcm/llcreader/llcmodel.py @@ -109,7 +109,7 @@ def _decompress(data, mask, dtype): _facet_reshape = (False, False, False, True, True) _nfaces = 13 _nfacets = 5 -_vgridvars = ['DRF','PHrefC','RC','DRF','PHrefF','RF'] +_vgridvars = ['DRC','DRF','PHrefC','PHrefF','RC','RF'] def _uncompressed_facet_index(nfacet, nside): face_size = nside**2 From 0c761447d79b1cb2c7b836c9d5121f344e8044c8 Mon Sep 17 00:00:00 2001 From: Timothy Smith Date: Thu, 2 Apr 2020 11:36:44 -0600 Subject: [PATCH 04/33] fixes loading grid dir issues --- xmitgcm/llcreader/stores.py | 16 ++++++++++++---- 1 file changed, 12 insertions(+), 4 deletions(-) diff --git a/xmitgcm/llcreader/stores.py b/xmitgcm/llcreader/stores.py index eda9292d..09a69f8a 100644 --- a/xmitgcm/llcreader/stores.py +++ b/xmitgcm/llcreader/stores.py @@ -43,9 +43,14 @@ def _directory(self, varname, iternum): return self.grid_path def _fname(self, varname, iternum): - fname = varname + '.%010d.data' % iternum if iternum is not None else varname - if self.shrunk: - fname += '.shrunk' + + if iternum is not None: + fname = varname + '.%010d.data' % iternum + if self.shrunk: + fname += '.shrunk' + else: + fname = varname + '.data' + return fname def _join(self, *args): @@ -109,4 +114,7 @@ class NestedStore(BaseStore): iteration number.""" def _directory(self, varname, iternum): - return self._join(self.base_path, '%010d' % iternum) + if iternum is not None: + return self._join(self.base_path, '%010d' % iternum) + else: + return self.grid_path From cf0b7e4ef3a8bf8f09ad7590fd12d66c02b86cf4 Mon Sep 17 00:00:00 2001 From: Timothy Smith Date: Thu, 2 Apr 2020 12:12:30 -0600 Subject: [PATCH 05/33] it works for llc2160 --- xmitgcm/llcreader/llcmodel.py | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/xmitgcm/llcreader/llcmodel.py b/xmitgcm/llcreader/llcmodel.py index 440a7aeb..b5aca83c 100644 --- a/xmitgcm/llcreader/llcmodel.py +++ b/xmitgcm/llcreader/llcmodel.py @@ -373,14 +373,18 @@ def _get_facet_chunk(store, varname, iternum, nfacet, klevels, nx, nz, dtype, file = fs.open(path) - # insert singleton axis for time and k level - facet_shape = (1, 1,) + _facet_shape(nfacet, nx) + # insert singleton axis for time (if not grid var) and k level + facet_shape = (1,) + _facet_shape(nfacet, nx) + facet_shape = (1,) + facet_shape if iternum is not None else facet_shape level_data = [] # the store tells us whether we need a mask or not point = _get_variable_point(varname, mask_override) - if store.shrunk: + + # it seems grid variables are not shrunk even though + # data variables are... + if store.shrunk and iternum is not None: index = all_index_data[nx][point] zgroup = store.open_mask_group() mask = zgroup['mask_' + point].astype('bool') From 70766dfcb00f9f69746f6f7118481567745d29c8 Mon Sep 17 00:00:00 2001 From: Timothy Smith Date: Thu, 2 Apr 2020 14:21:02 -0600 Subject: [PATCH 06/33] add directory to some LLC4320 grid vars --- xmitgcm/llcreader/known_models.py | 12 +++++++++--- 1 file changed, 9 insertions(+), 3 deletions(-) diff --git a/xmitgcm/llcreader/known_models.py b/xmitgcm/llcreader/known_models.py index 3dff1198..9167adb3 100644 --- a/xmitgcm/llcreader/known_models.py +++ b/xmitgcm/llcreader/known_models.py @@ -43,7 +43,7 @@ class LLC2160Model(BaseLLCModel): 'hFacC','hFacS','hFacW','PHrefC','PHrefF','RAC','RAS','RAW', 'RC','RF','rLowC','rLowS','rLowW','rSurfC', 'rSurfS','rSurfW','XC','YC'] - # unrecognized name: 'RhoRef' + # unrecognized name: 'RhoRef' # corner point problems: 'RAZ','XG','YG' mask_override = {'oceTAUX': 'c', 'oceTAUY': 'c'} @@ -60,6 +60,11 @@ class LLC4320Model(BaseLLCModel): varnames = ['Eta', 'KPPhbl', 'oceFWflx', 'oceQnet', 'oceQsw', 'oceSflux', 'oceTAUX', 'oceTAUY', 'PhiBot', 'Salt', 'SIarea', 'SIheff', 'SIhsalt', 'SIhsnow', 'SIuice', 'SIvice', 'Theta', 'U', 'V', 'W'] + grid_varnames = ['DRC','DRF','DXC','DXG','DYC','DYG', + 'hFacC','hFacS','hFacW','PHrefC','PHrefF', + 'RAC','RAS','RAW','RC','RF','XC','YC'] + # unrecognized name: 'RhoRef','DXV','DYU','DXF','DYF' + # corner point problems: 'RAZ','XG','YG' mask_override = {'oceTAUX': 'c', 'oceTAUY': 'c'} @@ -71,7 +76,7 @@ def __init__(self): grid_path = 'https://data.nas.nasa.gov/ecco/download_data.php?file=/eccodata/llc_2160/grid' mask_path = 'https://storage.googleapis.com/pangeo-ecco/llc/masks/llc_2160_masks.zarr/' store = stores.NestedStore(fs, base_path=base_path, mask_path=mask_path, - grid_path=grid_path,shrunk=True, join_char='/') + grid_path=grid_path, shrunk=True, join_char='/') super(ECCOPortalLLC2160Model, self).__init__(store) @@ -80,9 +85,10 @@ class ECCOPortalLLC4320Model(LLC4320Model): def __init__(self): fs = _make_http_filesystem() base_path = 'https://data.nas.nasa.gov/ecco/download_data.php?file=/eccodata/llc_4320/compressed' + grid_path = 'https://data.nas.nasa.gov/ecco/download_data.php?file=/eccodata/llc_4320/grid' mask_path = 'https://storage.googleapis.com/pangeo-ecco/llc/masks/llc_4320_masks.zarr/' store = stores.NestedStore(fs, base_path=base_path, mask_path=mask_path, - shrunk=True, join_char='/') + grid_path=grid_path, shrunk=True, join_char='/') super(ECCOPortalLLC4320Model, self).__init__(store) From 06a0e35580cc3ad0abf8f67ebbbd45db1c84c41c Mon Sep 17 00:00:00 2001 From: Timothy Smith Date: Thu, 2 Apr 2020 14:21:23 -0600 Subject: [PATCH 07/33] separate grid metadata function, add shrunk_grid flag --- xmitgcm/llcreader/llcmodel.py | 70 ++++++++++++++++++++--------------- xmitgcm/llcreader/stores.py | 17 ++++++--- 2 files changed, 52 insertions(+), 35 deletions(-) diff --git a/xmitgcm/llcreader/llcmodel.py b/xmitgcm/llcreader/llcmodel.py index b5aca83c..32fab4a8 100644 --- a/xmitgcm/llcreader/llcmodel.py +++ b/xmitgcm/llcreader/llcmodel.py @@ -8,48 +8,59 @@ from .duck_array_ops import concatenate from .shrunk_index import all_index_data +def _get_grid_metadata(): + # keep this separate from get_var_metadata + # because grid stuff is weird + from ..variables import (horizontal_coordinates_llc, + vertical_coordinates, vertical_grid_variables, + horizontal_grid_variables, volume_grid_variables, + mask_variables, extra_grid_variables) + + # get grid info + # keys are variable names + grid_vars = horizontal_coordinates_llc + grid_vars.update(vertical_coordinates) + grid_vars.update(horizontal_grid_variables) + grid_vars.update(vertical_grid_variables) + grid_vars.update(volume_grid_variables) + grid_vars.update(mask_variables) + grid_vars.update(extra_grid_variables) + + # make dictionary with keys as filenames + grid_metadata = {} + for key,val in grid_vars.items(): + # masks use hFac filename to be computed in mds_store + if 'filename' in val and key[:4]!='mask': + val.update({'real_name':key}) + grid_metadata[val['filename']] = val + else: + grid_metadata[key] = val + + return grid_metadata + def _get_var_metadata(): # The LLC run data comes with zero metadata. So we import metadata from # the xmitgcm package. - from ..variables import (state_variables, package_state_variables, - horizontal_coordinates_llc, vertical_coordinates, - horizontal_grid_variables, vertical_grid_variables, - volume_grid_variables, mask_variables, - extra_grid_variables) + from ..variables import state_variables, package_state_variables from ..utils import parse_available_diagnostics from ..default_diagnostics import diagnostics - from ..mds_store import _get_all_grid_variables from io import StringIO - # get grid info - grid_metadata = horizontal_coordinates_llc - grid_metadata.update(vertical_coordinates) - grid_metadata.update(horizontal_grid_variables) - grid_metadata.update(vertical_grid_variables) - grid_metadata.update(volume_grid_variables) - grid_metadata.update(mask_variables) - grid_metadata.update(extra_grid_variables) - diag_file = StringIO(diagnostics) available_diags = parse_available_diagnostics(diag_file) var_metadata = state_variables var_metadata.update(package_state_variables) var_metadata.update(available_diags) - # add grid vars - for key,val in grid_metadata.items(): - if 'filename' in val: - val.update({'real_name':key}) - var_metadata[val['filename']] = val - else: - var_metadata[key] = val - # even the file names from the LLC data differ from standard MITgcm output aliases = {'Eta': 'ETAN', 'PhiBot': 'PHIBOT', 'Salt': 'SALT', 'Theta': 'THETA'} for a, b in aliases.items(): var_metadata[a] = var_metadata[b] + # add grid metadata + var_metadata.update(_get_grid_metadata()) + return var_metadata _VAR_METADATA = _get_var_metadata() @@ -109,7 +120,7 @@ def _decompress(data, mask, dtype): _facet_reshape = (False, False, False, True, True) _nfaces = 13 _nfacets = 5 -_vgridvars = ['DRC','DRF','PHrefC','PHrefF','RC','RF'] +_vgrid_prefixes = ['DRC','DRF','PHrefC','PHrefF','RC','RF'] def _uncompressed_facet_index(nfacet, nside): face_size = nside**2 @@ -382,9 +393,8 @@ def _get_facet_chunk(store, varname, iternum, nfacet, klevels, nx, nz, dtype, # the store tells us whether we need a mask or not point = _get_variable_point(varname, mask_override) - # it seems grid variables are not shrunk even though - # data variables are... - if store.shrunk and iternum is not None: + if (store.shrunk and iternum is not None) or \ + (store.shrunk_grid and iternum is None): index = all_index_data[nx][point] zgroup = store.open_mask_group() mask = zgroup['mask_' + point].astype('bool') @@ -635,7 +645,7 @@ def _get_facet_data(self, varname, iters, klevels, k_chunksize): if len(dims)==2: klevels = [0,] - if varname in _vgridvars: + if varname in _vgrid_prefixes: data_facets = self._dask_array_vgrid(varname,klevels,k_chunksize) else: data_facets = [self._dask_array(nfacet, varname, iters, klevels, k_chunksize) @@ -723,8 +733,8 @@ def _if_not_none(a, b): data = transformer(data_facets, _VAR_METADATA) # separate horizontal and vertical grid variables - hgrid_names = [x for x in self.grid_varnames if x not in _vgridvars] - vgrid_names = [x for x in self.grid_varnames if x in _vgridvars] + hgrid_names = [x for x in self.grid_varnames if x not in _vgrid_prefixes] + vgrid_names = [x for x in self.grid_varnames if x in _vgrid_prefixes] hgrid_facets = {key: grid_facets[key] for key in hgrid_names} vgrid_facets = {key: grid_facets[key] for key in vgrid_names} diff --git a/xmitgcm/llcreader/stores.py b/xmitgcm/llcreader/stores.py index 09a69f8a..a8bff129 100644 --- a/xmitgcm/llcreader/stores.py +++ b/xmitgcm/llcreader/stores.py @@ -13,17 +13,21 @@ class BaseStore: Where to find the data within the filesystem shrunk : bool, optional Whether the data files have been tagged with `.shrunk` - mask_fs : fsspec.AbstractFileSystem, optional - Where to find the mask datasets to decode the compression - mask_path : str, optional - Path the the mask datasets on the ``mask_fs`` filesystem + mask_fs, grid_fs : fsspec.AbstractFileSystem, optional + Where to find the mask or grid datasets to decode the compression + mask_path, grid_path : str, optional + Path to the mask or grid datasets on the ``mask_fs`` or ``grid_fs`` filesystem + shrunk_grid : bool, optional + Whether the grid files have been tagged with `.shrunk` + not always the same as for data variables join_char : str or None Character to use to join paths. Falls back on os.path.join if None. """ def __init__(self, fs, base_path='/', shrunk=False, mask_fs=None, mask_path=None, - grid_fs=None, grid_path=None, join_char=None): + grid_fs=None, grid_path=None, + shrunk_grid=False, join_char=None): self.base_path = base_path self.fs = fs self.shrunk = shrunk @@ -31,6 +35,7 @@ def __init__(self, fs, base_path='/', shrunk=False, self.mask_path = mask_path self.grid_fs = grid_fs or self.fs self.grid_path = grid_path + self.shrunk_grid = shrunk_grid self.join_char = join_char if shrunk and (mask_path is None): raise ValueError("`mask_path` can't be None if `shrunk` is True") @@ -50,6 +55,8 @@ def _fname(self, varname, iternum): fname += '.shrunk' else: fname = varname + '.data' + if self.shrunk_grid: + fname = varname + '.shrunk' return fname From 0910fce18570c849df3ee0bab7b017034b177238 Mon Sep 17 00:00:00 2001 From: Timothy Smith Date: Thu, 2 Apr 2020 14:36:52 -0600 Subject: [PATCH 08/33] add grid_varnames to paramter list --- xmitgcm/llcreader/llcmodel.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/xmitgcm/llcreader/llcmodel.py b/xmitgcm/llcreader/llcmodel.py index 32fab4a8..e6771189 100644 --- a/xmitgcm/llcreader/llcmodel.py +++ b/xmitgcm/llcreader/llcmodel.py @@ -492,8 +492,8 @@ class BaseLLCModel: Final model iteration number (exclusive; follows python range conventions) iter_step : int Spacing between iterations - varnames : list - List of variable names contained in the dataset + varnames, grid_varnames : list + List of data variable and grid variable names contained in the dataset mask_override : dict Override inference of masking variable, e.g. ``{'oceTAUX': 'c'}`` """ From fc8b9b328318644e96a8099474f1bcb1cec699ac Mon Sep 17 00:00:00 2001 From: Timothy Smith Date: Thu, 2 Apr 2020 16:15:28 -0600 Subject: [PATCH 09/33] test that grid variables show up in portal datasets --- xmitgcm/test/test_llcreader.py | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/xmitgcm/test/test_llcreader.py b/xmitgcm/test/test_llcreader.py index f2c1d160..390835e6 100644 --- a/xmitgcm/test/test_llcreader.py +++ b/xmitgcm/test/test_llcreader.py @@ -8,6 +8,15 @@ 'oceTAUX', 'oceTAUY', 'PhiBot', 'Salt', 'SIarea', 'SIheff', 'SIhsalt', 'SIhsnow', 'SIuice', 'SIvice', 'Theta', 'U', 'V', 'W'] +EXPECTED_COORDS = {2160: ['AngleCS','AngleSN','Depth', + 'DRC','DRF','DXC','DXG','DYC','DYG', + 'hFacC','hFacS','hFacW','PHrefC','PHrefF','RAC','RAS','RAW', + 'RC','RF','rLowC','rLowS','rLowW','rSurfC', + 'rSurfS','rSurfW','XC','YC'], + 4320: ['DRC','DRF','DXC','DXG','DYC','DYG', + 'hFacC','hFacS','hFacW','PHrefC','PHrefF', + 'RAC','RAS','RAW','RC','RF','XC','YC']} + ########### Generic llcreader tests on local data ############################## @pytest.fixture(scope='module') @@ -74,6 +83,7 @@ def test_ecco_portal_faces(ecco_portal_model): 'j_g': nx, 'k': 90, 'k_u': 90, 'k_l': 90, 'k_p1': 90, 'time': 3} assert set(EXPECTED_VARS) == set(ds_faces.data_vars) + assert set(EXPECTED_COORDS[nx]).issubset(set(ds_faces.coords)) def test_ecco_portal_load(ecco_portal_model): # an expensive test because it actually loads data @@ -91,3 +101,4 @@ def test_ecco_portal_latlon(ecco_portal_model): 'k': 90, 'j_g': 3*nx, 'i_g': 4*nx, 'k_p1': 90, 'j': 3*nx, 'face': 13} assert set(EXPECTED_VARS) == set(ds_ll.data_vars) + assert set(EXPECTED_COORDS[nx]).issubset(set(ds_ll.coords)) From 824bf7be18692e8d712c7ce812015040d0575ba4 Mon Sep 17 00:00:00 2001 From: Timothy Smith Date: Thu, 2 Apr 2020 20:42:35 -0600 Subject: [PATCH 10/33] prefixes->variables --- xmitgcm/test/test_llcreader.py | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/xmitgcm/test/test_llcreader.py b/xmitgcm/test/test_llcreader.py index 390835e6..5a9d5c1a 100644 --- a/xmitgcm/test/test_llcreader.py +++ b/xmitgcm/test/test_llcreader.py @@ -8,14 +8,14 @@ 'oceTAUX', 'oceTAUY', 'PhiBot', 'Salt', 'SIarea', 'SIheff', 'SIhsalt', 'SIhsnow', 'SIuice', 'SIvice', 'Theta', 'U', 'V', 'W'] -EXPECTED_COORDS = {2160: ['AngleCS','AngleSN','Depth', - 'DRC','DRF','DXC','DXG','DYC','DYG', - 'hFacC','hFacS','hFacW','PHrefC','PHrefF','RAC','RAS','RAW', - 'RC','RF','rLowC','rLowS','rLowW','rSurfC', +EXPECTED_COORDS = {2160: ['CS','SN','Depth', + 'drC','drF','dxC','dxG','dyC','dyG', + 'hFacC','hFacS','hFacW','PHrefC','PHrefF','rA','rAs','rAw', + 'Z','Zp1','rLowC','rLowS','rLowW','rSurfC', 'rSurfS','rSurfW','XC','YC'], - 4320: ['DRC','DRF','DXC','DXG','DYC','DYG', + 4320: ['drC','drF','dxC','dxG','dyC','dyG', 'hFacC','hFacS','hFacW','PHrefC','PHrefF', - 'RAC','RAS','RAW','RC','RF','XC','YC']} + 'rA','rAs','rAw','Z','Zp1','XC','YC']} ########### Generic llcreader tests on local data ############################## From 8229c6a96365956b5668a53f4a2680d44cf0dfcb Mon Sep 17 00:00:00 2001 From: Timothy Smith Date: Thu, 2 Apr 2020 21:26:27 -0600 Subject: [PATCH 11/33] try Zl... --- xmitgcm/test/test_llcreader.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/xmitgcm/test/test_llcreader.py b/xmitgcm/test/test_llcreader.py index 5a9d5c1a..9f8a662d 100644 --- a/xmitgcm/test/test_llcreader.py +++ b/xmitgcm/test/test_llcreader.py @@ -11,11 +11,11 @@ EXPECTED_COORDS = {2160: ['CS','SN','Depth', 'drC','drF','dxC','dxG','dyC','dyG', 'hFacC','hFacS','hFacW','PHrefC','PHrefF','rA','rAs','rAw', - 'Z','Zp1','rLowC','rLowS','rLowW','rSurfC', + 'Z','Zl','rLowC','rLowS','rLowW','rSurfC', 'rSurfS','rSurfW','XC','YC'], 4320: ['drC','drF','dxC','dxG','dyC','dyG', 'hFacC','hFacS','hFacW','PHrefC','PHrefF', - 'rA','rAs','rAw','Z','Zp1','XC','YC']} + 'rA','rAs','rAw','Z','Zl','XC','YC']} ########### Generic llcreader tests on local data ############################## From 5c61b63b80bdf5534c4626dfd7150cdcdf36fa52 Mon Sep 17 00:00:00 2001 From: Timothy Smith Date: Fri, 3 Apr 2020 11:44:20 -0600 Subject: [PATCH 12/33] extra logic to create Zl, Zu, and Zp1 --- xmitgcm/llcreader/llcmodel.py | 48 +++++++++++++++++++++++++++++----- xmitgcm/test/test_llcreader.py | 4 +-- 2 files changed, 43 insertions(+), 9 deletions(-) diff --git a/xmitgcm/llcreader/llcmodel.py b/xmitgcm/llcreader/llcmodel.py index e6771189..e5cc262f 100644 --- a/xmitgcm/llcreader/llcmodel.py +++ b/xmitgcm/llcreader/llcmodel.py @@ -36,6 +36,12 @@ def _get_grid_metadata(): else: grid_metadata[key] = val + # force RF to point to Zp1, deal with this manually.. + grid_metadata['RF']=vertical_coordinates['Zp1'] + grid_metadata['RF']['real_name'] = 'Zp1' + for zv in ['Zu','Zl']: + grid_metadata[zv] = vertical_coordinates[zv] + return grid_metadata def _get_var_metadata(): @@ -121,6 +127,7 @@ def _decompress(data, mask, dtype): _nfaces = 13 _nfacets = 5 _vgrid_prefixes = ['DRC','DRF','PHrefC','PHrefF','RC','RF'] +_vgrid_p1_prefixes = ['DRC','PHrefF','RF'] def _uncompressed_facet_index(nfacet, nside): face_size = nside**2 @@ -519,7 +526,7 @@ def __init__(self, store): ---------- store : llcreader.BaseStore mask_ds : zarr.Group - Must contain variables `mask_c`, `masc_w`, `mask_s` + Must contain variables `mask_c`, `mask_w`, `mask_s` """ self.store = store self.shape = (self.nz, self.nface, self.nx, self.nx) @@ -540,6 +547,21 @@ def _get_masks(self): masks[point] = _faces_to_facets(mask_faces) return masks + def _get_kp1_levels(self,k_levels): + # determine kp1 levels + # get borders to all k (center) levels + # ki used to get Zu, Zl later + ku = np.concatenate([k_levels[1:],[k_levels[-1]+1]]) + kp1 = [] + ki=[] + for i,(x,y) in enumerate(zip(k_levels,ku)): + kp1+= [x] if x not in kp1 else [] + kp1+= [y] if y-x==1 else [x+1] + + + kp1=np.array(kp1) + + return kp1 def _make_coords_faces(self, all_iters): time = self.delta_t * all_iters @@ -631,8 +653,9 @@ def _dask_array_vgrid(self, varname, klevels, k_chunksize): for n_k, these_klevels in enumerate(_chunks(klevels, k_chunksize)): key = name, n_k + nz=self.nz if varname not in _vgrid_p1_prefixes else self.nz+1 task = (_get_1d_chunk, self.store, varname, - these_klevels, self.nz, self.dtype) + these_klevels, nz, self.dtype) dsk[key] = task return dsa.Array(dsk, name, chunks, self.dtype) @@ -713,7 +736,9 @@ def _if_not_none(a, b): ds = _faces_coords_to_latlon(ds) k_levels = k_levels or np.arange(self.nz) - ds = ds.sel(k=k_levels, k_l=k_levels, k_u=k_levels, k_p1=k_levels) + kp1_levels = self._get_kp1_levels(k_levels) + + ds = ds.sel(k=k_levels, k_l=k_levels, k_u=k_levels, k_p1=kp1_levels) # get the data in facet form data_facets = {vname: @@ -721,9 +746,11 @@ def _if_not_none(a, b): for vname in varnames} # get the grid in facet form - grid_facets = {vname: - self._get_facet_data(vname, None, k_levels, k_chunksize) - for vname in self.grid_varnames} + # do separately for vertical coords on kp1_levels + grid_facets = {} + for vname in self.grid_varnames: + my_k_levels = k_levels if vname not in _vgrid_p1_prefixes else kp1_levels + grid_facets[vname] = self._get_facet_data(vname, None, my_k_levels, k_chunksize) # transform it into faces or latlon data_transformers = {'faces': _all_facets_to_faces, @@ -744,7 +771,7 @@ def _if_not_none(a, b): data.update(vgrid_facets) variables = {} - gridvar_names = [] + gridvar_names = ['Zl','Zu'] for vname in varnames+self.grid_varnames: meta = _VAR_METADATA[vname] dims = meta['dims'] @@ -761,6 +788,13 @@ def _if_not_none(a, b): variables[vname] = xr.Variable(dims, data[fname], attrs) + # handle vertical coordinate after the fact.. + ki = np.array([list(kp1_levels).index(x) for x in k_levels]) + for zv,sl in zip(['Zl','Zu'],[ki,ki+1]): + variables[zv] = xr.Variable(_VAR_METADATA[zv]['dims'], + data['RF'][sl], + _VAR_METADATA[zv]['attrs']) + ds = ds.update(variables) if grid_vars_to_coords: diff --git a/xmitgcm/test/test_llcreader.py b/xmitgcm/test/test_llcreader.py index 9f8a662d..5df44eca 100644 --- a/xmitgcm/test/test_llcreader.py +++ b/xmitgcm/test/test_llcreader.py @@ -11,11 +11,11 @@ EXPECTED_COORDS = {2160: ['CS','SN','Depth', 'drC','drF','dxC','dxG','dyC','dyG', 'hFacC','hFacS','hFacW','PHrefC','PHrefF','rA','rAs','rAw', - 'Z','Zl','rLowC','rLowS','rLowW','rSurfC', + 'Z','Zp1','Zl','Zu','rLowC','rLowS','rLowW','rSurfC', 'rSurfS','rSurfW','XC','YC'], 4320: ['drC','drF','dxC','dxG','dyC','dyG', 'hFacC','hFacS','hFacW','PHrefC','PHrefF', - 'rA','rAs','rAw','Z','Zl','XC','YC']} + 'rA','rAs','rAw','Z','Zp1','Zl','Zu','XC','YC']} ########### Generic llcreader tests on local data ############################## From 1ac31a20357278c56c962fc398a2e97aac2189ea Mon Sep 17 00:00:00 2001 From: Timothy Smith Date: Fri, 3 Apr 2020 12:03:25 -0600 Subject: [PATCH 13/33] make reading grid optional --- xmitgcm/llcreader/llcmodel.py | 41 +++++++++++++++++++++-------------- 1 file changed, 25 insertions(+), 16 deletions(-) diff --git a/xmitgcm/llcreader/llcmodel.py b/xmitgcm/llcreader/llcmodel.py index e5cc262f..ff9e7af7 100644 --- a/xmitgcm/llcreader/llcmodel.py +++ b/xmitgcm/llcreader/llcmodel.py @@ -683,7 +683,7 @@ def _get_facet_data(self, varname, iters, klevels, k_chunksize): def get_dataset(self, varnames=None, iter_start=None, iter_stop=None, iter_step=None, k_levels=None, k_chunksize=1, - type='faces',grid_vars_to_coords=True): + type='faces', read_grid=True, grid_vars_to_coords=True): """ Create an xarray Dataset object for this model. @@ -706,6 +706,10 @@ def get_dataset(self, varnames=None, iter_start=None, iter_stop=None, How many vertical levels per Dask chunk. type : {'faces', 'latlon'}, optional What type of dataset to create + read_grid : bool, optional + Whether to read the grid info + grid_vars_to_coords : bool, optional + Whether to promote grid variables to coordinate status Returns ------- @@ -731,6 +735,10 @@ def _if_not_none(a, b): varnames = varnames or self.varnames + # grid stuff + grid_vars_to_coords = read_grid and grid_vars_to_coords + grid_varnames = self.grid_varnames if read_grid else [] + ds = self._make_coords_faces(iters) if type=='latlon': ds = _faces_coords_to_latlon(ds) @@ -748,7 +756,7 @@ def _if_not_none(a, b): # get the grid in facet form # do separately for vertical coords on kp1_levels grid_facets = {} - for vname in self.grid_varnames: + for vname in grid_varnames: my_k_levels = k_levels if vname not in _vgrid_p1_prefixes else kp1_levels grid_facets[vname] = self._get_facet_data(vname, None, my_k_levels, k_chunksize) @@ -760,8 +768,8 @@ def _if_not_none(a, b): data = transformer(data_facets, _VAR_METADATA) # separate horizontal and vertical grid variables - hgrid_names = [x for x in self.grid_varnames if x not in _vgrid_prefixes] - vgrid_names = [x for x in self.grid_varnames if x in _vgrid_prefixes] + hgrid_names = [x for x in grid_varnames if x not in _vgrid_prefixes] + vgrid_names = [x for x in grid_varnames if x in _vgrid_prefixes] hgrid_facets = {key: grid_facets[key] for key in hgrid_names} vgrid_facets = {key: grid_facets[key] for key in vgrid_names} @@ -771,33 +779,34 @@ def _if_not_none(a, b): data.update(vgrid_facets) variables = {} - gridvar_names = ['Zl','Zu'] - for vname in varnames+self.grid_varnames: + gridlist = ['Zl','Zu'] if read_grid else [] + for vname in varnames+grid_varnames: meta = _VAR_METADATA[vname] dims = meta['dims'] if type=='faces': dims = _add_face_to_dims(dims) - dims = ['time',] + dims if vname not in self.grid_varnames else dims + dims = ['time',] + dims if vname not in grid_varnames else dims attrs = meta['attrs'] # Handle grid names different from filenames fname = vname vname = meta['real_name'] if 'real_name' in meta else vname - if fname in self.grid_varnames: - gridvar_names.append(vname) + if fname in grid_varnames: + gridlist.append(vname) variables[vname] = xr.Variable(dims, data[fname], attrs) - # handle vertical coordinate after the fact.. - ki = np.array([list(kp1_levels).index(x) for x in k_levels]) - for zv,sl in zip(['Zl','Zu'],[ki,ki+1]): - variables[zv] = xr.Variable(_VAR_METADATA[zv]['dims'], - data['RF'][sl], - _VAR_METADATA[zv]['attrs']) + # handle vertical coordinate after the fact... + if read_grid: + ki = np.array([list(kp1_levels).index(x) for x in k_levels]) + for zv,sl in zip(['Zl','Zu'],[ki,ki+1]): + variables[zv] = xr.Variable(_VAR_METADATA[zv]['dims'], + data['RF'][sl], + _VAR_METADATA[zv]['attrs']) ds = ds.update(variables) if grid_vars_to_coords: - ds = ds.set_coords(gridvar_names) + ds = ds.set_coords(gridlist) return ds From 3ad9e0016d18e86c8fd15d7cffed1c5b16023609 Mon Sep 17 00:00:00 2001 From: Timothy Smith Date: Fri, 3 Apr 2020 12:35:10 -0600 Subject: [PATCH 14/33] wow, not copying dicts really messes things up --- xmitgcm/llcreader/llcmodel.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/xmitgcm/llcreader/llcmodel.py b/xmitgcm/llcreader/llcmodel.py index ff9e7af7..b71e1576 100644 --- a/xmitgcm/llcreader/llcmodel.py +++ b/xmitgcm/llcreader/llcmodel.py @@ -18,9 +18,9 @@ def _get_grid_metadata(): # get grid info # keys are variable names - grid_vars = horizontal_coordinates_llc - grid_vars.update(vertical_coordinates) + grid_vars = horizontal_coordinates_llc.copy() grid_vars.update(horizontal_grid_variables) + grid_vars.update(vertical_coordinates) grid_vars.update(vertical_grid_variables) grid_vars.update(volume_grid_variables) grid_vars.update(mask_variables) @@ -54,7 +54,7 @@ def _get_var_metadata(): diag_file = StringIO(diagnostics) available_diags = parse_available_diagnostics(diag_file) - var_metadata = state_variables + var_metadata = state_variables.copy() var_metadata.update(package_state_variables) var_metadata.update(available_diags) From d1b5ec10325aac8fadfa15e20d0dc3fe0ed81a10 Mon Sep 17 00:00:00 2001 From: Timothy Smith Date: Fri, 3 Apr 2020 16:05:11 -0600 Subject: [PATCH 15/33] only build Zu, Zl if RF is there --- xmitgcm/llcreader/llcmodel.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/xmitgcm/llcreader/llcmodel.py b/xmitgcm/llcreader/llcmodel.py index b71e1576..a500ed43 100644 --- a/xmitgcm/llcreader/llcmodel.py +++ b/xmitgcm/llcreader/llcmodel.py @@ -796,8 +796,8 @@ def _if_not_none(a, b): variables[vname] = xr.Variable(dims, data[fname], attrs) - # handle vertical coordinate after the fact... - if read_grid: + # handle vertical coordinate after the fact + if read_grid and 'RF' in grid_varnames: ki = np.array([list(kp1_levels).index(x) for x in k_levels]) for zv,sl in zip(['Zl','Zu'],[ki,ki+1]): variables[zv] = xr.Variable(_VAR_METADATA[zv]['dims'], From bb0df28e54a96ae3c7b6cb0845b26deaedaa3068 Mon Sep 17 00:00:00 2001 From: Timothy Smith Date: Fri, 3 Apr 2020 16:15:59 -0600 Subject: [PATCH 16/33] shouldn't kp1 have nz+1 levels --- xmitgcm/test/test_llcreader.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/xmitgcm/test/test_llcreader.py b/xmitgcm/test/test_llcreader.py index 5df44eca..71518226 100644 --- a/xmitgcm/test/test_llcreader.py +++ b/xmitgcm/test/test_llcreader.py @@ -37,14 +37,14 @@ def test_llc90_local_faces(local_llc90_store, llc90_kwargs): ds_faces = model.get_dataset(**llc90_kwargs) assert set(llc90_kwargs['varnames']) == set(ds_faces.data_vars) assert ds_faces.dims == {'face': 13, 'i': 90, 'i_g': 90, 'j': 90, 'j_g': 90, - 'k': 50, 'k_u': 50, 'k_l': 50, 'k_p1': 50, 'time': 2} + 'k': 50, 'k_u': 50, 'k_l': 50, 'k_p1': 51, 'time': 2} def test_llc90_local_latlon(local_llc90_store, llc90_kwargs): store = local_llc90_store model = llcreader.LLC90Model(store) ds_latlon = model.get_dataset(type='latlon', **llc90_kwargs) assert set(llc90_kwargs['varnames']) == set(ds_latlon.data_vars) - assert ds_latlon.dims == {'i': 360, 'time': 2, 'k_p1': 50, 'face': 13, + assert ds_latlon.dims == {'i': 360, 'time': 2, 'k_p1': 51, 'face': 13, 'i_g': 360, 'k_u': 50, 'k': 50, 'k_l': 50, 'j_g': 270, 'j': 270} @@ -81,7 +81,7 @@ def test_ecco_portal_faces(ecco_portal_model): nx = ecco_portal_model.nx assert ds_faces.dims == {'face': 13, 'i': nx, 'i_g': nx, 'j': nx, 'j_g': nx, 'k': 90, 'k_u': 90, 'k_l': 90, - 'k_p1': 90, 'time': 3} + 'k_p1': 91, 'time': 3} assert set(EXPECTED_VARS) == set(ds_faces.data_vars) assert set(EXPECTED_COORDS[nx]).issubset(set(ds_faces.coords)) @@ -98,7 +98,7 @@ def test_ecco_portal_latlon(ecco_portal_model): ds_ll = ecco_portal_model.get_dataset(iter_stop=iter_stop, type='latlon') nx = ecco_portal_model.nx assert ds_ll.dims == {'i': 4*nx, 'k_u': 90, 'k_l': 90, 'time': 3, - 'k': 90, 'j_g': 3*nx, 'i_g': 4*nx, 'k_p1': 90, + 'k': 90, 'j_g': 3*nx, 'i_g': 4*nx, 'k_p1': 91, 'j': 3*nx, 'face': 13} assert set(EXPECTED_VARS) == set(ds_ll.data_vars) assert set(EXPECTED_COORDS[nx]).issubset(set(ds_ll.coords)) From 4b8879878e88ceb42f70c9a63cad1a6e8cf89f8f Mon Sep 17 00:00:00 2001 From: Timothy Smith Date: Fri, 3 Apr 2020 16:25:52 -0600 Subject: [PATCH 17/33] only read grid if grid_varnames has entries --- xmitgcm/llcreader/llcmodel.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/xmitgcm/llcreader/llcmodel.py b/xmitgcm/llcreader/llcmodel.py index a500ed43..8a15cd54 100644 --- a/xmitgcm/llcreader/llcmodel.py +++ b/xmitgcm/llcreader/llcmodel.py @@ -736,7 +736,8 @@ def _if_not_none(a, b): varnames = varnames or self.varnames # grid stuff - grid_vars_to_coords = read_grid and grid_vars_to_coords + read_grid = read_grid and len(self.grid_varnames)!=0 + grid_vars_to_coords = grid_vars_to_coords and read_grid grid_varnames = self.grid_varnames if read_grid else [] ds = self._make_coords_faces(iters) From 8f2c616ba5da07713a07649314a88e95275de48c Mon Sep 17 00:00:00 2001 From: Timothy Smith Date: Fri, 3 Apr 2020 16:36:30 -0600 Subject: [PATCH 18/33] also test selected k_p1 levels --- xmitgcm/test/test_llcreader.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/xmitgcm/test/test_llcreader.py b/xmitgcm/test/test_llcreader.py index 71518226..fca0e125 100644 --- a/xmitgcm/test/test_llcreader.py +++ b/xmitgcm/test/test_llcreader.py @@ -49,18 +49,22 @@ def test_llc90_local_latlon(local_llc90_store, llc90_kwargs): 'j_g': 270, 'j': 270} @pytest.mark.parametrize('rettype', ['faces', 'latlon']) -@pytest.mark.parametrize('k_levels', [None, [0, 2, 7, 9, 10, 20]]) +@pytest.mark.parametrize('k_levels,kp1_levels', + [None, [0, 2, 7, 9, 10, 20], + [None, [0,1,2,3,7,8,9,10,11,20,21]) @pytest.mark.parametrize('k_chunksize', [1, 2]) def test_llc90_local_faces_load(local_llc90_store, llc90_kwargs, rettype, k_levels, - k_chunksize): + kp1_levels, k_chunksize): store = local_llc90_store model = llcreader.LLC90Model(store) ds = model.get_dataset(k_levels=k_levels, k_chunksize=k_chunksize, type=rettype, **llc90_kwargs) if k_levels is None: assert list(ds.k.values) == list(range(50)) + assert list(ds.k_p1.values) == list(range(50)) else: assert list(ds.k.values) == k_levels + assert list(ds.k_p1.values) == k_levels assert all([cs==k_chunksize for cs in ds['T'].data.chunks[1]]) ds.load() From 7b0b89fc9fe799560b85d363395fe36f3f111fe9 Mon Sep 17 00:00:00 2001 From: Timothy Smith Date: Fri, 3 Apr 2020 16:43:05 -0600 Subject: [PATCH 19/33] typo... --- xmitgcm/test/test_llcreader.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/xmitgcm/test/test_llcreader.py b/xmitgcm/test/test_llcreader.py index fca0e125..28d35443 100644 --- a/xmitgcm/test/test_llcreader.py +++ b/xmitgcm/test/test_llcreader.py @@ -49,9 +49,10 @@ def test_llc90_local_latlon(local_llc90_store, llc90_kwargs): 'j_g': 270, 'j': 270} @pytest.mark.parametrize('rettype', ['faces', 'latlon']) -@pytest.mark.parametrize('k_levels,kp1_levels', - [None, [0, 2, 7, 9, 10, 20], - [None, [0,1,2,3,7,8,9,10,11,20,21]) +@pytest.mark.parametrize('k_levels, kp1_levels', + [(None,None), + ([0, 2, 7, 9, 10, 20], + [0,1,2,3,7,8,9,10,11,20,21])]) @pytest.mark.parametrize('k_chunksize', [1, 2]) def test_llc90_local_faces_load(local_llc90_store, llc90_kwargs, rettype, k_levels, kp1_levels, k_chunksize): From e51de9f6f56e52f5fb93f3d0945de63913c0cf92 Mon Sep 17 00:00:00 2001 From: Timothy Smith Date: Fri, 3 Apr 2020 16:56:03 -0600 Subject: [PATCH 20/33] ... more typos... --- xmitgcm/test/test_llcreader.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/xmitgcm/test/test_llcreader.py b/xmitgcm/test/test_llcreader.py index 28d35443..99c9f750 100644 --- a/xmitgcm/test/test_llcreader.py +++ b/xmitgcm/test/test_llcreader.py @@ -62,10 +62,10 @@ def test_llc90_local_faces_load(local_llc90_store, llc90_kwargs, rettype, k_leve type=rettype, **llc90_kwargs) if k_levels is None: assert list(ds.k.values) == list(range(50)) - assert list(ds.k_p1.values) == list(range(50)) + assert list(ds.k_p1.values) == list(range(51)) else: assert list(ds.k.values) == k_levels - assert list(ds.k_p1.values) == k_levels + assert list(ds.k_p1.values) == kp1_levels assert all([cs==k_chunksize for cs in ds['T'].data.chunks[1]]) ds.load() From 237b82db5436efc98a89411fa587b8e812163b89 Mon Sep 17 00:00:00 2001 From: Timothy Smith Date: Sat, 4 Apr 2020 11:50:45 -0600 Subject: [PATCH 21/33] rhoref, extra dx/dy metrics --- xmitgcm/llcreader/known_models.py | 11 +++++------ xmitgcm/variables.py | 24 ++++++++++++++++++++++++ 2 files changed, 29 insertions(+), 6 deletions(-) diff --git a/xmitgcm/llcreader/known_models.py b/xmitgcm/llcreader/known_models.py index 9167adb3..ec53183e 100644 --- a/xmitgcm/llcreader/known_models.py +++ b/xmitgcm/llcreader/known_models.py @@ -42,8 +42,7 @@ class LLC2160Model(BaseLLCModel): 'DRC','DRF','DXC','DXG','DYC','DYG', 'hFacC','hFacS','hFacW','PHrefC','PHrefF','RAC','RAS','RAW', 'RC','RF','rLowC','rLowS','rLowW','rSurfC', - 'rSurfS','rSurfW','XC','YC'] - # unrecognized name: 'RhoRef' + 'RhoRef','rSurfS','rSurfW','XC','YC'] # corner point problems: 'RAZ','XG','YG' mask_override = {'oceTAUX': 'c', 'oceTAUY': 'c'} @@ -60,11 +59,11 @@ class LLC4320Model(BaseLLCModel): varnames = ['Eta', 'KPPhbl', 'oceFWflx', 'oceQnet', 'oceQsw', 'oceSflux', 'oceTAUX', 'oceTAUY', 'PhiBot', 'Salt', 'SIarea', 'SIheff', 'SIhsalt', 'SIhsnow', 'SIuice', 'SIvice', 'Theta', 'U', 'V', 'W'] - grid_varnames = ['DRC','DRF','DXC','DXG','DYC','DYG', + grid_varnames = ['DRC','DRF','DXC','DXF','DXG','DYC','DYF','DYG', 'hFacC','hFacS','hFacW','PHrefC','PHrefF', - 'RAC','RAS','RAW','RC','RF','XC','YC'] - # unrecognized name: 'RhoRef','DXV','DYU','DXF','DYF' - # corner point problems: 'RAZ','XG','YG' + 'RAC','RAS','RAW','RC','RF', + 'RhoRef','XC','YC'] + # corner point problems: 'RAZ','XG','YG','DXV','DYU' mask_override = {'oceTAUX': 'c', 'oceTAUY': 'c'} diff --git a/xmitgcm/variables.py b/xmitgcm/variables.py index 9d9defac..bd35c5ba 100644 --- a/xmitgcm/variables.py +++ b/xmitgcm/variables.py @@ -652,6 +652,30 @@ standard_name="ctrl_vector_3d_mask_at_v_location", long_name='CTRL 3D mask where ctrl vector is active at v location', units='')) + # Reference density profile + rhoRef=dict(dims=['k'],attrs=dict( + standard_name="reference_density_profile", + long_name="1D, vertical reference density profile", + coordinate="Z", + units='kg m-3'), + filename='RhoRef') + # Additional grid metrics + dxF=dict(dims=['j','i'],attrs=dict( + standard_name="cell_x_size_at_t_location", + long_name="cell x size", units="m", coordinate="YC XC"), + filename='DXF'), + dyF=dict(dims=['j','i'],attrs=dict( + standard_name="cell_y_size_at_t_location", + long_name="cell y size", units="m", coordinate="YC XC"), + filename='DYF'), + dxV=dict(dims=['j_g','i_g'],attrs=dict( + standard_name="cell_x_size_at_f_location", + long_name="cell x size", units="m", coordinate="YG XG"), + filename='DXV'), + dyU=dict(dims=['j_g','i_g'],attrs=dict( + standard_name="cell_y_size_at_f_location", + long_name="cell y size", units="m", coordinate="YG XG"), + filename='DYU') # Printed from write_grid when sigma coordinates are used # AHybSigmF, BHybSigF, ... # Unclear where on the grid these variables exist, From 19ba68c89c931ac3762c4e642ecaf3bcc39787d4 Mon Sep 17 00:00:00 2001 From: Timothy Smith Date: Sat, 4 Apr 2020 11:56:04 -0600 Subject: [PATCH 22/33] add new metrics to test --- xmitgcm/llcreader/known_models.py | 4 ++-- xmitgcm/test/test_llcreader.py | 8 ++++---- 2 files changed, 6 insertions(+), 6 deletions(-) diff --git a/xmitgcm/llcreader/known_models.py b/xmitgcm/llcreader/known_models.py index ec53183e..60ec407c 100644 --- a/xmitgcm/llcreader/known_models.py +++ b/xmitgcm/llcreader/known_models.py @@ -41,8 +41,8 @@ class LLC2160Model(BaseLLCModel): grid_varnames = ['AngleCS','AngleSN','Depth', 'DRC','DRF','DXC','DXG','DYC','DYG', 'hFacC','hFacS','hFacW','PHrefC','PHrefF','RAC','RAS','RAW', - 'RC','RF','rLowC','rLowS','rLowW','rSurfC', - 'RhoRef','rSurfS','rSurfW','XC','YC'] + 'RC','RF','rhoRef','rLowC','rLowS','rLowW', + 'rSurfC','rSurfS','rSurfW','XC','YC'] # corner point problems: 'RAZ','XG','YG' mask_override = {'oceTAUX': 'c', 'oceTAUY': 'c'} diff --git a/xmitgcm/test/test_llcreader.py b/xmitgcm/test/test_llcreader.py index 99c9f750..d2eca3bb 100644 --- a/xmitgcm/test/test_llcreader.py +++ b/xmitgcm/test/test_llcreader.py @@ -11,11 +11,11 @@ EXPECTED_COORDS = {2160: ['CS','SN','Depth', 'drC','drF','dxC','dxG','dyC','dyG', 'hFacC','hFacS','hFacW','PHrefC','PHrefF','rA','rAs','rAw', - 'Z','Zp1','Zl','Zu','rLowC','rLowS','rLowW','rSurfC', - 'rSurfS','rSurfW','XC','YC'], - 4320: ['drC','drF','dxC','dxG','dyC','dyG', + 'Z','Zp1','Zl','Zu','rhoRef','rLowC','rLowS','rLowW', + 'rSurfC','rSurfS','rSurfW','XC','YC'], + 4320: ['drC','drF','dxC','dxF','dxG','dyC','dyF','dyG', 'hFacC','hFacS','hFacW','PHrefC','PHrefF', - 'rA','rAs','rAw','Z','Zp1','Zl','Zu','XC','YC']} + 'rA','rAs','rAw','rhoRef','Z','Zp1','Zl','Zu','XC','YC']} ########### Generic llcreader tests on local data ############################## From e2da50e05dc8006d18ad6c0ce5dc3c5da1760461 Mon Sep 17 00:00:00 2001 From: Timothy Smith Date: Sat, 4 Apr 2020 12:15:00 -0600 Subject: [PATCH 23/33] typos, and add rhoRef to vgrid prefixes --- xmitgcm/llcreader/known_models.py | 2 +- xmitgcm/llcreader/llcmodel.py | 2 +- xmitgcm/variables.py | 4 ++-- 3 files changed, 4 insertions(+), 4 deletions(-) diff --git a/xmitgcm/llcreader/known_models.py b/xmitgcm/llcreader/known_models.py index 60ec407c..645eac70 100644 --- a/xmitgcm/llcreader/known_models.py +++ b/xmitgcm/llcreader/known_models.py @@ -41,7 +41,7 @@ class LLC2160Model(BaseLLCModel): grid_varnames = ['AngleCS','AngleSN','Depth', 'DRC','DRF','DXC','DXG','DYC','DYG', 'hFacC','hFacS','hFacW','PHrefC','PHrefF','RAC','RAS','RAW', - 'RC','RF','rhoRef','rLowC','rLowS','rLowW', + 'RC','RF','RhoRef','rLowC','rLowS','rLowW', 'rSurfC','rSurfS','rSurfW','XC','YC'] # corner point problems: 'RAZ','XG','YG' mask_override = {'oceTAUX': 'c', 'oceTAUY': 'c'} diff --git a/xmitgcm/llcreader/llcmodel.py b/xmitgcm/llcreader/llcmodel.py index 8a15cd54..90c7968e 100644 --- a/xmitgcm/llcreader/llcmodel.py +++ b/xmitgcm/llcreader/llcmodel.py @@ -126,7 +126,7 @@ def _decompress(data, mask, dtype): _facet_reshape = (False, False, False, True, True) _nfaces = 13 _nfacets = 5 -_vgrid_prefixes = ['DRC','DRF','PHrefC','PHrefF','RC','RF'] +_vgrid_prefixes = ['DRC','DRF','PHrefC','PHrefF','RC','RF','RhoRef'] _vgrid_p1_prefixes = ['DRC','PHrefF','RF'] def _uncompressed_facet_index(nfacet, nside): diff --git a/xmitgcm/variables.py b/xmitgcm/variables.py index bd35c5ba..56aa0b6c 100644 --- a/xmitgcm/variables.py +++ b/xmitgcm/variables.py @@ -651,14 +651,14 @@ maskCtrlS=dict(dims=['k', 'j_g', 'i'], attrs=dict( standard_name="ctrl_vector_3d_mask_at_v_location", long_name='CTRL 3D mask where ctrl vector is active at v location', - units='')) + units='')), # Reference density profile rhoRef=dict(dims=['k'],attrs=dict( standard_name="reference_density_profile", long_name="1D, vertical reference density profile", coordinate="Z", units='kg m-3'), - filename='RhoRef') + filename='RhoRef'), # Additional grid metrics dxF=dict(dims=['j','i'],attrs=dict( standard_name="cell_x_size_at_t_location", From 17e34e136cd560b5662ba509ef67c8f40d0eca16 Mon Sep 17 00:00:00 2001 From: Timothy Smith Date: Sun, 12 Apr 2020 18:33:54 -0500 Subject: [PATCH 24/33] load vgrid as one big chunk, all at once. test it too! --- xmitgcm/llcreader/llcmodel.py | 42 ++++++++++++---------------------- xmitgcm/test/test_llcreader.py | 7 ++++++ 2 files changed, 22 insertions(+), 27 deletions(-) diff --git a/xmitgcm/llcreader/llcmodel.py b/xmitgcm/llcreader/llcmodel.py index 90c7968e..e25c313a 100644 --- a/xmitgcm/llcreader/llcmodel.py +++ b/xmitgcm/llcreader/llcmodel.py @@ -450,26 +450,13 @@ def _get_1d_chunk(store, varname, klevels, nz, dtype): file = fs.open(path) - # insert singleton axis for k level - facet_shape = (1,) - level_data = [] - - for k in klevels: - assert (k >= 0) & (k < nz) + # read all levels for 1D variables + read_length = nz*dtype.itemsize # all levels in bytes + buffer = file.read(read_length) + data = np.frombuffer(buffer,dtype=dtype) - # with a 1D file, k level gives "where to read" in file - read_offset = k * dtype.itemsize # in bytes - read_length = dtype.itemsize # in bytes - file.seek(read_offset) - buffer = file.read(read_length) - data = np.frombuffer(buffer, dtype=dtype) - assert len(data) == 1 - - # this is the shape this facet is supposed to have - data.shape = facet_shape - level_data.append(data) - - return np.concatenate(level_data, axis=0) + # now subset + return data[klevels] class BaseLLCModel: """Class representing an LLC Model Dataset. @@ -643,20 +630,21 @@ def _dask_array(self, nfacet, varname, iters, klevels, k_chunksize): def _dask_array_vgrid(self, varname, klevels, k_chunksize): # return a dask array for a 1D vertical grid var - chunks = (tuple([len(c) - for c in _chunks(klevels, k_chunksize)]),) + + # single chunk for 1D variables + chunks = ((len(klevels),),) # manually build dask graph dsk = {} token = tokenize(varname, self.store) name = '-'.join([varname, token]) - for n_k, these_klevels in enumerate(_chunks(klevels, k_chunksize)): - key = name, n_k - nz=self.nz if varname not in _vgrid_p1_prefixes else self.nz+1 - task = (_get_1d_chunk, self.store, varname, - these_klevels, nz, self.dtype) - dsk[key] = task + nz = self.nz if varname not in _vgrid_p1_prefixes else self.nz+1 + task = (_get_1d_chunk, self.store, varname, + list(klevels), nz, self.dtype) + + key = name, 0 + dsk[key] = task return dsa.Array(dsk, name, chunks, self.dtype) diff --git a/xmitgcm/test/test_llcreader.py b/xmitgcm/test/test_llcreader.py index d2eca3bb..7f598a35 100644 --- a/xmitgcm/test/test_llcreader.py +++ b/xmitgcm/test/test_llcreader.py @@ -1,4 +1,5 @@ import pytest +from dask.array.core import Array as dsa llcreader = pytest.importorskip("xmitgcm.llcreader") @@ -68,6 +69,12 @@ def test_llc90_local_faces_load(local_llc90_store, llc90_kwargs, rettype, k_leve assert list(ds.k_p1.values) == kp1_levels assert all([cs==k_chunksize for cs in ds['T'].data.chunks[1]]) + # make sure vertical coordinates are in one single chunk + for fld in ds['Z','Zl','Zu','Zp1'].coords: + if isinstance(ds[fld].data,dsa): + assert len(ds[fld].data.chunks)==1 + assert (len(ds[fld]),)==ds[fld].data.chunks[0] + ds.load() ########### ECCO Portal Tests ################################################## From f8a1d713c32ef0ad35c455e7ddb4e2a56a29018d Mon Sep 17 00:00:00 2001 From: Timothy Smith Date: Sun, 12 Apr 2020 18:47:09 -0500 Subject: [PATCH 25/33] oops, fixed test --- xmitgcm/test/test_llcreader.py | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/xmitgcm/test/test_llcreader.py b/xmitgcm/test/test_llcreader.py index 7f598a35..eaa772db 100644 --- a/xmitgcm/test/test_llcreader.py +++ b/xmitgcm/test/test_llcreader.py @@ -69,12 +69,6 @@ def test_llc90_local_faces_load(local_llc90_store, llc90_kwargs, rettype, k_leve assert list(ds.k_p1.values) == kp1_levels assert all([cs==k_chunksize for cs in ds['T'].data.chunks[1]]) - # make sure vertical coordinates are in one single chunk - for fld in ds['Z','Zl','Zu','Zp1'].coords: - if isinstance(ds[fld].data,dsa): - assert len(ds[fld].data.chunks)==1 - assert (len(ds[fld]),)==ds[fld].data.chunks[0] - ds.load() ########### ECCO Portal Tests ################################################## @@ -97,6 +91,12 @@ def test_ecco_portal_faces(ecco_portal_model): assert set(EXPECTED_VARS) == set(ds_faces.data_vars) assert set(EXPECTED_COORDS[nx]).issubset(set(ds_faces.coords)) + # make sure vertical coordinates are in one single chunk + for fld in ds_faces[['Z','Zl','Zu','Zp1']].coords: + if isinstance(ds_faces[fld].data,dsa): + assert len(ds_faces[fld].data.chunks)==1 + assert (len(ds_faces[fld]),)==ds_faces[fld].data.chunks[0] + def test_ecco_portal_load(ecco_portal_model): # an expensive test because it actually loads data iter_stop = ecco_portal_model.iter_start + 2 * ecco_portal_model.iter_step + 1 From 800f902d48eeb80e0010539ecae1887aa69e7aea Mon Sep 17 00:00:00 2001 From: Timothy Smith Date: Sun, 12 Apr 2020 19:24:19 -0500 Subject: [PATCH 26/33] fix for global vgrid vars... what about layers.. tbd --- xmitgcm/llcreader/llcmodel.py | 12 +++++------- 1 file changed, 5 insertions(+), 7 deletions(-) diff --git a/xmitgcm/llcreader/llcmodel.py b/xmitgcm/llcreader/llcmodel.py index e25c313a..2ad0fbfe 100644 --- a/xmitgcm/llcreader/llcmodel.py +++ b/xmitgcm/llcreader/llcmodel.py @@ -126,8 +126,6 @@ def _decompress(data, mask, dtype): _facet_reshape = (False, False, False, True, True) _nfaces = 13 _nfacets = 5 -_vgrid_prefixes = ['DRC','DRF','PHrefC','PHrefF','RC','RF','RhoRef'] -_vgrid_p1_prefixes = ['DRC','PHrefF','RF'] def _uncompressed_facet_index(nfacet, nside): face_size = nside**2 @@ -639,7 +637,7 @@ def _dask_array_vgrid(self, varname, klevels, k_chunksize): token = tokenize(varname, self.store) name = '-'.join([varname, token]) - nz = self.nz if varname not in _vgrid_p1_prefixes else self.nz+1 + nz = self.nz if _VAR_METADATA[varname]['dims'] != ['k_p1'] else self.nz+1 task = (_get_1d_chunk, self.store, varname, list(klevels), nz, self.dtype) @@ -656,7 +654,7 @@ def _get_facet_data(self, varname, iters, klevels, k_chunksize): if len(dims)==2: klevels = [0,] - if varname in _vgrid_prefixes: + if len(dims)==1: data_facets = self._dask_array_vgrid(varname,klevels,k_chunksize) else: data_facets = [self._dask_array(nfacet, varname, iters, klevels, k_chunksize) @@ -746,7 +744,7 @@ def _if_not_none(a, b): # do separately for vertical coords on kp1_levels grid_facets = {} for vname in grid_varnames: - my_k_levels = k_levels if vname not in _vgrid_p1_prefixes else kp1_levels + my_k_levels = k_levels if _VAR_METADATA[vname]['dims'] !=['k_p1'] else kp1_levels grid_facets[vname] = self._get_facet_data(vname, None, my_k_levels, k_chunksize) # transform it into faces or latlon @@ -757,8 +755,8 @@ def _if_not_none(a, b): data = transformer(data_facets, _VAR_METADATA) # separate horizontal and vertical grid variables - hgrid_names = [x for x in grid_varnames if x not in _vgrid_prefixes] - vgrid_names = [x for x in grid_varnames if x in _vgrid_prefixes] + hgrid_names = [x for x in grid_varnames if len(_VAR_METADATA[x]['dims'])!=1] + vgrid_names = [x for x in grid_varnames if len(_VAR_METADATA[x]['dims'])==1] hgrid_facets = {key: grid_facets[key] for key in hgrid_names} vgrid_facets = {key: grid_facets[key] for key in vgrid_names} From fb1c5fd83d3d2f7f05c38960ae6b01cc2dabd0fd Mon Sep 17 00:00:00 2001 From: Timothy Smith Date: Sun, 12 Apr 2020 23:08:25 -0500 Subject: [PATCH 27/33] cleanup vgrid check --- xmitgcm/llcreader/llcmodel.py | 16 ++++++++++------ 1 file changed, 10 insertions(+), 6 deletions(-) diff --git a/xmitgcm/llcreader/llcmodel.py b/xmitgcm/llcreader/llcmodel.py index 2ad0fbfe..32039c6f 100644 --- a/xmitgcm/llcreader/llcmodel.py +++ b/xmitgcm/llcreader/llcmodel.py @@ -71,6 +71,11 @@ def _get_var_metadata(): _VAR_METADATA = _get_var_metadata() +def _is_vgrid(vname): + # check for 1d, vertical grid variables + dims = _VAR_METADATA[vname]['dims'] + return len(dims)==1 and dims[0][0]=='k' + def _get_variable_point(vname, mask_override): # fix for https://github.com/MITgcm/xmitgcm/issues/191 if vname in mask_override: @@ -654,7 +659,7 @@ def _get_facet_data(self, varname, iters, klevels, k_chunksize): if len(dims)==2: klevels = [0,] - if len(dims)==1: + if _is_vgrid(varname): data_facets = self._dask_array_vgrid(varname,klevels,k_chunksize) else: data_facets = [self._dask_array(nfacet, varname, iters, klevels, k_chunksize) @@ -755,11 +760,10 @@ def _if_not_none(a, b): data = transformer(data_facets, _VAR_METADATA) # separate horizontal and vertical grid variables - hgrid_names = [x for x in grid_varnames if len(_VAR_METADATA[x]['dims'])!=1] - vgrid_names = [x for x in grid_varnames if len(_VAR_METADATA[x]['dims'])==1] - - hgrid_facets = {key: grid_facets[key] for key in hgrid_names} - vgrid_facets = {key: grid_facets[key] for key in vgrid_names} + hgrid_facets = {key: grid_facets[key] + for key in grid_varnames if not _is_vgrid(key)} + vgrid_facets = {key: grid_facets[key] + for key in grid_varnames if _is_vgrid(key)} # do not transform vertical grid variables data.update(transformer(hgrid_facets, _VAR_METADATA)) From 61405fb70f88f70710b18be27add3a494becce42 Mon Sep 17 00:00:00 2001 From: Timothy Smith Date: Sun, 12 Apr 2020 23:17:28 -0500 Subject: [PATCH 28/33] cleanup grid var dict situation --- xmitgcm/llcreader/llcmodel.py | 14 +++----------- xmitgcm/mds_store.py | 4 ++-- 2 files changed, 5 insertions(+), 13 deletions(-) diff --git a/xmitgcm/llcreader/llcmodel.py b/xmitgcm/llcreader/llcmodel.py index 32039c6f..0f4a0932 100644 --- a/xmitgcm/llcreader/llcmodel.py +++ b/xmitgcm/llcreader/llcmodel.py @@ -11,19 +11,11 @@ def _get_grid_metadata(): # keep this separate from get_var_metadata # because grid stuff is weird - from ..variables import (horizontal_coordinates_llc, - vertical_coordinates, vertical_grid_variables, - horizontal_grid_variables, volume_grid_variables, - mask_variables, extra_grid_variables) + from ..mds_store import _get_all_grid_variables + from ..variables import extra_grid_variables, vertical_coordinates # get grid info - # keys are variable names - grid_vars = horizontal_coordinates_llc.copy() - grid_vars.update(horizontal_grid_variables) - grid_vars.update(vertical_coordinates) - grid_vars.update(vertical_grid_variables) - grid_vars.update(volume_grid_variables) - grid_vars.update(mask_variables) + grid_vars = _get_all_grid_variables('llc') grid_vars.update(extra_grid_variables) # make dictionary with keys as filenames diff --git a/xmitgcm/mds_store.py b/xmitgcm/mds_store.py index 8b6fa1be..bf4bb3e3 100644 --- a/xmitgcm/mds_store.py +++ b/xmitgcm/mds_store.py @@ -763,7 +763,7 @@ def _guess_layers(data_dir): return all_layers -def _get_all_grid_variables(geometry, grid_dir, layers={}): +def _get_all_grid_variables(geometry, grid_dir=None, layers={}): """"Put all the relevant grid metadata into one big dictionary.""" possible_hcoords = {'cartesian': horizontal_coordinates_cartesian, 'llc': horizontal_coordinates_llc, @@ -772,7 +772,7 @@ def _get_all_grid_variables(geometry, grid_dir, layers={}): hcoords = possible_hcoords[geometry] # look for extra variables, if they exist in grid_dir - extravars = _get_extra_grid_variables(grid_dir) + extravars = _get_extra_grid_variables(grid_dir) if grid_dir is not None else {} allvars = [hcoords, vertical_coordinates, horizontal_grid_variables, vertical_grid_variables, volume_grid_variables, mask_variables, From fce198c622b064dd8cb159142f5299b971a96b23 Mon Sep 17 00:00:00 2001 From: Timothy Smith Date: Sun, 12 Apr 2020 23:31:03 -0500 Subject: [PATCH 29/33] also test latlon vgrid coord chunking --- xmitgcm/test/test_llcreader.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/xmitgcm/test/test_llcreader.py b/xmitgcm/test/test_llcreader.py index eaa772db..a15088ee 100644 --- a/xmitgcm/test/test_llcreader.py +++ b/xmitgcm/test/test_llcreader.py @@ -114,3 +114,9 @@ def test_ecco_portal_latlon(ecco_portal_model): 'j': 3*nx, 'face': 13} assert set(EXPECTED_VARS) == set(ds_ll.data_vars) assert set(EXPECTED_COORDS[nx]).issubset(set(ds_ll.coords)) + + # make sure vertical coordinates are in one single chunk + for fld in ds_faces[['Z','Zl','Zu','Zp1']].coords: + if isinstance(ds_faces[fld].data,dsa): + assert len(ds_faces[fld].data.chunks)==1 + assert (len(ds_faces[fld]),)==ds_faces[fld].data.chunks[0] From 154754ff68b9b9bb03ee79c99596e5e68e7e8a97 Mon Sep 17 00:00:00 2001 From: Timothy Smith Date: Sun, 12 Apr 2020 23:36:43 -0500 Subject: [PATCH 30/33] faces->ll --- xmitgcm/test/test_llcreader.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/xmitgcm/test/test_llcreader.py b/xmitgcm/test/test_llcreader.py index a15088ee..9ac3b04e 100644 --- a/xmitgcm/test/test_llcreader.py +++ b/xmitgcm/test/test_llcreader.py @@ -116,7 +116,7 @@ def test_ecco_portal_latlon(ecco_portal_model): assert set(EXPECTED_COORDS[nx]).issubset(set(ds_ll.coords)) # make sure vertical coordinates are in one single chunk - for fld in ds_faces[['Z','Zl','Zu','Zp1']].coords: - if isinstance(ds_faces[fld].data,dsa): - assert len(ds_faces[fld].data.chunks)==1 - assert (len(ds_faces[fld]),)==ds_faces[fld].data.chunks[0] + for fld in ds_ll[['Z','Zl','Zu','Zp1']].coords: + if isinstance(ds_ll[fld].data,dsa): + assert len(ds_ll[fld].data.chunks)==1 + assert (len(ds_ll[fld]),)==ds_ll[fld].data.chunks[0] From a1d713d712e5b338758c925a790d0f480cff5693 Mon Sep 17 00:00:00 2001 From: Timothy Smith Date: Mon, 13 Apr 2020 00:08:52 -0500 Subject: [PATCH 31/33] update docs --- doc/llcreader.rst | 133 ++++++++++++++++++++++++++++++++++++---------- 1 file changed, 104 insertions(+), 29 deletions(-) diff --git a/doc/llcreader.rst b/doc/llcreader.rst index b7e02d04..12d27ce8 100644 --- a/doc/llcreader.rst +++ b/doc/llcreader.rst @@ -98,7 +98,7 @@ By default, all variables and all timesteps are loaded:: >>> ds = model.get_dataset(k_chunksize=90) >>> print(ds) - Dimensions: (face: 13, i: 4320, i_g: 4320, j: 4320, j_g: 4320, k: 90, k_l: 90, k_p1: 90, k_u: 90, time: 9030) + Dimensions: (face: 13, i: 4320, i_g: 4320, j: 4320, j_g: 4320, k: 90, k_l: 90, k_p1: 91, k_u: 90, time: 9030) Coordinates: * face (face) int64 0 1 2 3 4 5 6 7 8 9 10 11 12 * i (i) int64 0 1 2 3 4 5 6 7 ... 4313 4314 4315 4316 4317 4318 4319 @@ -108,30 +108,54 @@ By default, all variables and all timesteps are loaded:: * k (k) int64 0 1 2 3 4 5 6 7 8 9 10 ... 80 81 82 83 84 85 86 87 88 89 * k_u (k_u) int64 0 1 2 3 4 5 6 7 8 9 ... 80 81 82 83 84 85 86 87 88 89 * k_l (k_l) int64 0 1 2 3 4 5 6 7 8 9 ... 80 81 82 83 84 85 86 87 88 89 - * k_p1 (k_p1) int64 0 1 2 3 4 5 6 7 8 9 ... 80 81 82 83 84 85 86 87 88 89 + * k_p1 (k_p1) int64 0 1 2 3 4 5 6 7 8 9 ... 81 82 83 84 85 86 87 88 89 90 niter (time) int64 ... * time (time) datetime64[ns] 2011-09-13 ... 2012-09-23T05:00:00 + drC (k_p1) >f4 dask.array + drF (k) >f4 dask.array + dxC (face, j, i_g) float32 dask.array + dxF (face, j, i) float32 dask.array + dxG (face, j_g, i) float32 dask.array + dyC (face, j_g, i) float32 dask.array + dyF (face, j, i) float32 dask.array + dyG (face, j, i_g) float32 dask.array + hFacC (k, face, j, i) float32 dask.array + hFacS (k, face, j_g, i) float32 dask.array + hFacW (k, face, j, i_g) float32 dask.array + PHrefC (k) >f4 dask.array + PHrefF (k_p1) >f4 dask.array + rA (face, j, i) float32 dask.array + rAs (face, j_g, i) float32 dask.array + rAw (face, j, i_g) float32 dask.array + Z (k) >f4 dask.array + Zp1 (k_p1) >f4 dask.array + rhoRef (k) >f4 dask.array + XC (face, j, i) float32 dask.array + YC (face, j, i) float32 dask.array + Zl (k_l) >f4 dask.array + Zu (k_u) >f4 dask.array Data variables: - Eta (time, face, j, i) >f4 dask.array - KPPhbl (time, face, j, i) >f4 dask.array - oceFWflx (time, face, j, i) >f4 dask.array - oceQnet (time, face, j, i) >f4 dask.array - oceQsw (time, face, j, i) >f4 dask.array - oceSflux (time, face, j, i) >f4 dask.array - oceTAUX (time, face, j, i_g) >f4 dask.array - oceTAUY (time, face, j_g, i) >f4 dask.array - PhiBot (time, face, j, i) >f4 dask.array - Salt (time, k, face, j, i) >f4 dask.array - SIarea (time, face, j, i) >f4 dask.array - SIheff (time, face, j, i) >f4 dask.array - SIhsalt (time, face, j, i) >f4 dask.array - SIhsnow (time, face, j, i) >f4 dask.array - SIuice (time, face, j, i_g) >f4 dask.array - SIvice (time, face, j_g, i) >f4 dask.array - Theta (time, k, face, j, i) >f4 dask.array - U (time, k, face, j, i_g) >f4 dask.array - V (time, k, face, j_g, i) >f4 dask.array - W (time, k_l, face, j, i) >f4 dask.array + Eta (time, face, j, i) float32 dask.array + KPPhbl (time, face, j, i) float32 dask.array + oceFWflx (time, face, j, i) float32 dask.array + oceQnet (time, face, j, i) float32 dask.array + oceQsw (time, face, j, i) float32 dask.array + oceSflux (time, face, j, i) float32 dask.array + oceTAUX (time, face, j, i_g) float32 dask.array + oceTAUY (time, face, j_g, i) float32 dask.array + PhiBot (time, face, j, i) float32 dask.array + Salt (time, k, face, j, i) float32 dask.array + SIarea (time, face, j, i) float32 dask.array + SIheff (time, face, j, i) float32 dask.array + SIhsalt (time, face, j, i) float32 dask.array + SIhsnow (time, face, j, i) float32 dask.array + SIuice (time, face, j, i_g) float32 dask.array + SIvice (time, face, j_g, i) float32 dask.array + Theta (time, k, face, j, i) float32 dask.array + U (time, k, face, j, i_g) float32 dask.array + V (time, k, face, j_g, i) float32 dask.array + W (time, k_l, face, j, i) float32 dask.array + This dataset is useless for computations on a laptop, because the individual chunks require nearly 20 GB of memory. @@ -141,8 +165,9 @@ Get a single 2D variable:: >>> ds = model.get_dataset(varnames=['Eta']) >>> print(ds) + - Dimensions: (face: 13, i: 4320, i_g: 4320, j: 4320, j_g: 4320, k: 90, k_l: 90, k_p1: 90, k_u: 90, time: 9030) + Dimensions: (face: 13, i: 4320, i_g: 4320, j: 4320, j_g: 4320, k: 90, k_l: 90, k_p1: 91, k_u: 90, time: 9030) Coordinates: * face (face) int64 0 1 2 3 4 5 6 7 8 9 10 11 12 * i (i) int64 0 1 2 3 4 5 6 7 ... 4313 4314 4315 4316 4317 4318 4319 @@ -152,18 +177,42 @@ Get a single 2D variable:: * k (k) int64 0 1 2 3 4 5 6 7 8 9 10 ... 80 81 82 83 84 85 86 87 88 89 * k_u (k_u) int64 0 1 2 3 4 5 6 7 8 9 ... 80 81 82 83 84 85 86 87 88 89 * k_l (k_l) int64 0 1 2 3 4 5 6 7 8 9 ... 80 81 82 83 84 85 86 87 88 89 - * k_p1 (k_p1) int64 0 1 2 3 4 5 6 7 8 9 ... 80 81 82 83 84 85 86 87 88 89 + * k_p1 (k_p1) int64 0 1 2 3 4 5 6 7 8 9 ... 81 82 83 84 85 86 87 88 89 90 niter (time) int64 ... * time (time) datetime64[ns] 2011-09-13 ... 2012-09-23T05:00:00 + drC (k_p1) >f4 dask.array + drF (k) >f4 dask.array + dxC (face, j, i_g) float32 dask.array + dxF (face, j, i) float32 dask.array + dxG (face, j_g, i) float32 dask.array + dyC (face, j_g, i) float32 dask.array + dyF (face, j, i) float32 dask.array + dyG (face, j, i_g) float32 dask.array + hFacC (k, face, j, i) float32 dask.array + hFacS (k, face, j_g, i) float32 dask.array + hFacW (k, face, j, i_g) float32 dask.array + PHrefC (k) >f4 dask.array + PHrefF (k_p1) >f4 dask.array + rA (face, j, i) float32 dask.array + rAs (face, j_g, i) float32 dask.array + rAw (face, j, i_g) float32 dask.array + Z (k) >f4 dask.array + Zp1 (k_p1) >f4 dask.array + rhoRef (k) >f4 dask.array + XC (face, j, i) float32 dask.array + YC (face, j, i) float32 dask.array + Zl (k_l) >f4 dask.array + Zu (k_u) >f4 dask.array Data variables: - Eta (time, face, j, i) >f4 dask.array + Eta (time, face, j, i) float32 dask.array + Get a few vertical levels from some 3D variables:: >>> ds = model.get_dataset(varnames=['Salt', 'Theta'], k_levels=[1, 10, 40]) >>> print(ds) - Dimensions: (face: 13, i: 4320, i_g: 4320, j: 4320, j_g: 4320, k: 3, k_l: 3, k_p1: 3, k_u: 3, time: 9030) + Dimensions: (face: 13, i: 4320, i_g: 4320, j: 4320, j_g: 4320, k: 3, k_l: 3, k_p1: 6, k_u: 3, time: 9030) Coordinates: * face (face) int64 0 1 2 3 4 5 6 7 8 9 10 11 12 * i (i) int64 0 1 2 3 4 5 6 7 ... 4313 4314 4315 4316 4317 4318 4319 @@ -173,13 +222,39 @@ Get a few vertical levels from some 3D variables:: * k (k) int64 1 10 40 * k_u (k_u) int64 1 10 40 * k_l (k_l) int64 1 10 40 - * k_p1 (k_p1) int64 1 10 40 + * k_p1 (k_p1) int64 1 2 10 11 40 41 niter (time) int64 ... * time (time) datetime64[ns] 2011-09-13 ... 2012-09-23T05:00:00 + drC (k_p1) >f4 dask.array + drF (k) >f4 dask.array + dxC (face, j, i_g) float32 dask.array + dxF (face, j, i) float32 dask.array + dxG (face, j_g, i) float32 dask.array + dyC (face, j_g, i) float32 dask.array + dyF (face, j, i) float32 dask.array + dyG (face, j, i_g) float32 dask.array + hFacC (k, face, j, i) float32 dask.array + hFacS (k, face, j_g, i) float32 dask.array + hFacW (k, face, j, i_g) float32 dask.array + PHrefC (k) >f4 dask.array + PHrefF (k_p1) >f4 dask.array + rA (face, j, i) float32 dask.array + rAs (face, j_g, i) float32 dask.array + rAw (face, j, i_g) float32 dask.array + Z (k) >f4 dask.array + Zp1 (k_p1) >f4 dask.array + rhoRef (k) >f4 dask.array + XC (face, j, i) float32 dask.array + YC (face, j, i) float32 dask.array + Zl (k_l) >f4 dask.array + Zu (k_u) >f4 dask.array Data variables: - Salt (time, k, face, j, i) >f4 dask.array - Theta (time, k, face, j, i) >f4 dask.array + Salt (time, k, face, j, i) float32 dask.array + Theta (time, k, face, j, i) float32 dask.array +Note that when vertical levels are subset like this, any vertical coordinate +associated with dimension `k_p1` will have levels above and below the selected +`k_levls`, which are at cell center. A list of all available variables can be seen as follows:: From ca80a05af7634625057109e052e21b8512505297 Mon Sep 17 00:00:00 2001 From: Timothy Smith Date: Sun, 12 Apr 2020 23:10:38 -0600 Subject: [PATCH 32/33] tiny typo --- doc/llcreader.rst | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/doc/llcreader.rst b/doc/llcreader.rst index 12d27ce8..372a8f0b 100644 --- a/doc/llcreader.rst +++ b/doc/llcreader.rst @@ -254,7 +254,7 @@ Get a few vertical levels from some 3D variables:: Note that when vertical levels are subset like this, any vertical coordinate associated with dimension `k_p1` will have levels above and below the selected -`k_levls`, which are at cell center. +`k_levels`, which are at cell center. A list of all available variables can be seen as follows:: From 24009a7ffc4c8d911cb23a8ba1dfe62f92e25daa Mon Sep 17 00:00:00 2001 From: Timothy Smith Date: Thu, 4 Jun 2020 09:11:49 -0500 Subject: [PATCH 33/33] added more 4320 grid files --- xmitgcm/llcreader/known_models.py | 14 ++++++++++---- xmitgcm/test/test_llcreader.py | 5 +++-- 2 files changed, 13 insertions(+), 6 deletions(-) diff --git a/xmitgcm/llcreader/known_models.py b/xmitgcm/llcreader/known_models.py index 645eac70..7d2bd13b 100644 --- a/xmitgcm/llcreader/known_models.py +++ b/xmitgcm/llcreader/known_models.py @@ -39,11 +39,14 @@ class LLC2160Model(BaseLLCModel): 'oceTAUX', 'oceTAUY', 'PhiBot', 'Salt', 'SIarea', 'SIheff', 'SIhsalt', 'SIhsnow', 'SIuice', 'SIvice', 'Theta', 'U', 'V', 'W'] grid_varnames = ['AngleCS','AngleSN','Depth', - 'DRC','DRF','DXC','DXG','DYC','DYG', - 'hFacC','hFacS','hFacW','PHrefC','PHrefF','RAC','RAS','RAW', + 'DRC','DRF', + 'DXC','DXF','DXG', + 'DYC','DYF','DYG', + 'hFacC','hFacS','hFacW','PHrefC','PHrefF', + 'RAC','RAS','RAW', 'RC','RF','RhoRef','rLowC','rLowS','rLowW', 'rSurfC','rSurfS','rSurfW','XC','YC'] - # corner point problems: 'RAZ','XG','YG' + # corner point problems: 'RAZ','XG','YG','DXV','DYU' mask_override = {'oceTAUX': 'c', 'oceTAUY': 'c'} @@ -59,7 +62,10 @@ class LLC4320Model(BaseLLCModel): varnames = ['Eta', 'KPPhbl', 'oceFWflx', 'oceQnet', 'oceQsw', 'oceSflux', 'oceTAUX', 'oceTAUY', 'PhiBot', 'Salt', 'SIarea', 'SIheff', 'SIhsalt', 'SIhsnow', 'SIuice', 'SIvice', 'Theta', 'U', 'V', 'W'] - grid_varnames = ['DRC','DRF','DXC','DXF','DXG','DYC','DYF','DYG', + grid_varnames = ['AngleCS','AngleSN','Depth', + 'DRC','DRF', + 'DXC','DXF','DXG', + 'DYC','DYF','DYG', 'hFacC','hFacS','hFacW','PHrefC','PHrefF', 'RAC','RAS','RAW','RC','RF', 'RhoRef','XC','YC'] diff --git a/xmitgcm/test/test_llcreader.py b/xmitgcm/test/test_llcreader.py index 9ac3b04e..c41e16be 100644 --- a/xmitgcm/test/test_llcreader.py +++ b/xmitgcm/test/test_llcreader.py @@ -10,11 +10,12 @@ 'SIhsalt', 'SIhsnow', 'SIuice', 'SIvice', 'Theta', 'U', 'V', 'W'] EXPECTED_COORDS = {2160: ['CS','SN','Depth', - 'drC','drF','dxC','dxG','dyC','dyG', + 'drC','drF','dxC','dxF','dxG','dyC','dyF','dyG', 'hFacC','hFacS','hFacW','PHrefC','PHrefF','rA','rAs','rAw', 'Z','Zp1','Zl','Zu','rhoRef','rLowC','rLowS','rLowW', 'rSurfC','rSurfS','rSurfW','XC','YC'], - 4320: ['drC','drF','dxC','dxF','dxG','dyC','dyF','dyG', + 4320: ['CS','SN','Depth', + 'drC','drF','dxC','dxF','dxG','dyC','dyF','dyG', 'hFacC','hFacS','hFacW','PHrefC','PHrefF', 'rA','rAs','rAw','rhoRef','Z','Zp1','Zl','Zu','XC','YC']}