Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

ENH: 2D chunks for all #285

Open
wants to merge 2 commits into
base: master
Choose a base branch
from
Open

ENH: 2D chunks for all #285

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
10d3d72
ENH: 2D chunks for all
jklymak Sep 27, 2021
3adb40c
Fix base case
jklymak Sep 27, 2021
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
12 changes: 7 additions & 5 deletions xmitgcm/mds_store.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -283,12 +283,12 @@ def open_mdsdataset(data_dir, grid_dir=None,
if grid_vars_to_coords:
ds = _set_coords(ds)
return ds

store = _MDSDataStore(data_dir, grid_dir, iternum, delta_t, read_grid,
prefix, ref_date, calendar,
geometry, endian,
ignore_unknown_vars=ignore_unknown_vars,
default_dtype=default_dtype,
chunks=chunks,
nx=nx, ny=ny, nz=nz, llc_method=llc_method,
levels=levels, extra_metadata=extra_metadata,
extra_variables=extra_variables)
Expand All @@ -303,7 +303,7 @@ def open_mdsdataset(data_dir, grid_dir=None,
ds['time'] = xr.decode_cf(ds[['time']])['time']

# do we need more fancy logic (like open_dataset), or is this enough
if chunks is not None:
if chunks is not None and chunks != "2D":
ds = ds.chunk(chunks)

# set attributes for CF conventions
Expand Down Expand Up @@ -374,6 +374,7 @@ def __init__(self, data_dir, grid_dir=None,
geometry='sphericalpolar',
endian='>', ignore_unknown_vars=False,
default_dtype=np.dtype('f4'),
chunks=None,
nx=None, ny=None, nz=None, llc_method="smallchunks",
levels=None, extra_metadata=None,
extra_variables=None):
Expand Down Expand Up @@ -597,7 +598,7 @@ def __init__(self, data_dir, grid_dir=None,
for p in prefixes:
# use a generator to loop through the variables in each file
for (vname, dims, data, attrs) in \
self.load_from_prefix(p, iternum, extra_metadata):
self.load_from_prefix(p, iternum=iternum, extra_metadata=extra_metadata, chunks=chunks):
# print(vname, dims, data.shape)
# Sizes of grid variables can vary between mitgcm versions.
# Check for such inconsistency and correct if so
Expand Down Expand Up @@ -631,7 +632,7 @@ def calc_masks(self, vname, data):

return data

def load_from_prefix(self, prefix, iternum=None, extra_metadata=None):
def load_from_prefix(self, prefix, chunks=None, iternum=None, extra_metadata=None):
"""Read data and look up metadata for grid variable `name`.

Parameters
Expand Down Expand Up @@ -668,7 +669,8 @@ def load_from_prefix(self, prefix, iternum=None, extra_metadata=None):
ddir = self.data_dir

basename = os.path.join(ddir, fname_base)
chunks = "CS" if self.cs else "3D"
if chunks is None or isinstance(chunks, dict):
chunks = "CS" if self.cs else "3D"

try:
vardata = read_mds(basename, iternum, endian=self.endian,
Expand Down