From c884677a69a3cb673240de25179da137b40a6f4f Mon Sep 17 00:00:00 2001 From: Ryan Abernathey Date: Thu, 12 Nov 2020 11:49:17 -0500 Subject: [PATCH] resolve final bugs --- xmitgcm/llcreader/llcmodel.py | 23 ++++++++++++++--------- 1 file changed, 14 insertions(+), 9 deletions(-) diff --git a/xmitgcm/llcreader/llcmodel.py b/xmitgcm/llcreader/llcmodel.py index 5a5a3697..58edea96 100644 --- a/xmitgcm/llcreader/llcmodel.py +++ b/xmitgcm/llcreader/llcmodel.py @@ -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""" @@ -582,20 +583,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)