Skip to content

Commit

Permalink
Fix llcreader klevels bug (MITgcm#225)
Browse files Browse the repository at this point in the history 
* empty commit to trigger CI

* add test for k_levels=[1]

* add explicit test for MITgcm#233

* test for bug in llc90

* fixes MITgcm#224

* resolve final bugs
  • Loading branch information
@rabernat
rabernat authored Nov 13, 2020
1 parent 1a2990f commit e68b3e5
Show file tree
Hide file tree
Showing 2 changed files with 43 additions and 21 deletions.
32 changes: 18 additions & 14 deletions xmitgcm/llcreader/llcmodel.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -436,7 +436,8 @@ def _get_facet_chunk(store, varname, iternum, nfacet, klevels, nx, nz, dtype,
data.shape = facet_shape
level_data.append(data)

return np.concatenate(level_data, axis=1)
out = np.concatenate(level_data, axis=-4)
return out

def _get_1d_chunk(store, varname, klevels, nz, dtype):
"""for 1D vertical grid variables"""
Expand All @@ -450,7 +451,6 @@ def _get_1d_chunk(store, varname, klevels, nz, dtype):
buffer = file.read(read_length)
data = np.frombuffer(buffer,dtype=dtype)

# now subset
return data[klevels]

class BaseLLCModel:
Expand Down Expand Up @@ -533,12 +533,12 @@ 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]])
ku = 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 += [x] if x not in kp1 else []
kp1 += [y] if y-x==1 else [x+1]


kp1=np.array(kp1)
Expand Down Expand Up @@ -590,20 +590,24 @@ def _dask_array(self, nfacet, varname, iters, klevels, k_chunksize):
name = '-'.join([varname, token])

# iters == None for grid variables
def _key_and_task(n_k, these_klevels, n_iter=None, iternum=None):
if n_iter is None:
key = name, n_k, 0, 0, 0
else:
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)
return key, task

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)
key, task = _key_and_task(n_k, these_klevels, n_iter, iternum)
dsk[key] = task
else:
for n_k, these_klevels in enumerate(_chunks(klevels, k_chunksize)):
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)
key, task = _key_and_task(n_k, these_klevels)
dsk[key] = task

return dsa.Array(dsk, name, chunks, self.dtype)
Expand Down Expand Up @@ -711,7 +715,7 @@ def _if_not_none(a, b):
if type=='latlon':
ds = _faces_coords_to_latlon(ds)

k_levels = k_levels or np.arange(self.nz)
k_levels = k_levels or list(range(self.nz))
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)
Expand Down
32 changes: 25 additions & 7 deletions xmitgcm/test/test_llcreader.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,10 @@
EXPECTED_VARS = ['Eta', 'KPPhbl', 'oceFWflx', 'oceQnet', 'oceQsw', 'oceSflux',
'oceTAUX', 'oceTAUY', 'PhiBot', 'Salt', 'SIarea', 'SIheff',
'SIhsalt', 'SIhsnow', 'SIuice', 'SIvice', 'Theta', 'U', 'V', 'W']
GRID_VARNAMES = ['AngleCS', 'AngleSN', 'DRC', 'DRF', 'DXC', 'DXG', 'DYC', 'DYG',
'Depth', 'PHrefC', 'PHrefF', 'RAC', 'RAS', 'RAW', 'RAZ', 'RC', 'RF',
'RhoRef', 'XC', 'XG', 'YC', 'YG', 'hFacC', 'hFacS', 'hFacW']


EXPECTED_COORDS = {2160: ['CS','SN','Depth',
'drC','drF','dxC','dxF','dxG','dyC','dyF','dyG',
Expand All @@ -28,7 +32,8 @@ def local_llc90_store(llc_mds_datadirs):
from fsspec.implementations.local import LocalFileSystem
dirname, expected = llc_mds_datadirs
fs = LocalFileSystem()
return llcreader.BaseStore(fs, base_path=dirname)
store = llcreader.BaseStore(fs, base_path=dirname, grid_path=dirname)
return store

@pytest.fixture(scope='module')
def llc90_kwargs():
Expand Down Expand Up @@ -58,18 +63,30 @@ def test_llc90_local_latlon(local_llc90_store, llc90_kwargs):
'i_g': 360, 'k_u': 50, 'k': 50, 'k_l': 50,
'j_g': 270, 'j': 270}


# includes regression test for https://github.com/MITgcm/xmitgcm/issues/233
@pytest.mark.parametrize('rettype', ['faces', 'latlon'])
@pytest.mark.parametrize('k_levels, kp1_levels',
[(None,None),
@pytest.mark.parametrize('k_levels, kp1_levels, k_chunksize',
[(None, None, 1),
([1], [1, 2], 1),
([0, 2, 7, 9, 10, 20],
[0,1,2,3,7,8,9,10,11,20,21])])
@pytest.mark.parametrize('k_chunksize', [1, 2])
[0,1,2,3,7,8,9,10,11,20,21], 1),
([0, 2, 7, 9, 10, 20],
[0,1,2,3,7,8,9,10,11,20,21], 2)
])
@pytest.mark.parametrize('read_grid', [False, True]
)
def test_llc90_local_faces_load(local_llc90_store, llc90_kwargs, rettype, k_levels,
kp1_levels, k_chunksize):
kp1_levels, k_chunksize, read_grid):
store = local_llc90_store
model = llcreader.LLC90Model(store)
model.grid_varnames = GRID_VARNAMES
ds = model.get_dataset(k_levels=k_levels, k_chunksize=k_chunksize,
type=rettype, **llc90_kwargs)
type=rettype, read_grid=read_grid, **llc90_kwargs)
if read_grid:
# doesn't work because the variables change name
# assert set(GRID_VARNAMES).issubset(set(ds.coords))
pass
if k_levels is None:
assert list(ds.k.values) == list(range(50))
assert list(ds.k_p1.values) == list(range(51))
Expand All @@ -80,6 +97,7 @@ def test_llc90_local_faces_load(local_llc90_store, llc90_kwargs, rettype, k_leve

ds.load()


########### ECCO Portal Tests ##################################################

@pytest.fixture(scope='module', params=[2160, 4320])
Expand Down

0 comments on commit e68b3e5

Please sign in to comment.