chunks?

[jklymak]

If I do

with xm.open_mdsdataset(data_dir, prefix=['spinup'], endian='=',
                        geometry='cartesian') as ds:
    print(ds)

I get

<xarray.Dataset>
Dimensions:  (XC: 480, XG: 480, YC: 400, YG: 400, Z: 400, Zl: 400, Zp1: 401, Zu: 400, time: 3)
Coordinates:
  * XC       (XC) float64 1.5e+03 4.5e+03 7.5e+03 ... 1.436e+06 1.438e+06
  * YC       (YC) float64 1.5e+03 4.5e+03 7.5e+03 ... 1.196e+06 1.198e+06
  * XG       (XG) float64 0.0 3e+03 6e+03 ... 1.431e+06 1.434e+06 1.437e+06
  * YG       (YG) float64 0.0 3e+03 6e+03 ... 1.191e+06 1.194e+06 1.197e+06
  ...
Data variables:
    UVEL     (time, Z, YC, XG) float64 dask.array<chunksize=(1, 400, 400, 480), meta=np.ndarray>
    VVEL     (time, Z, YG, XC) float64 dask.array<chunksize=(1, 400, 400, 480), meta=np.ndarray>
...

This is great. However, to save memory I thought I could chunk:

with xm.open_mdsdataset(data_dir, prefix=['spinup'], endian='=',
                        geometry='cartesian', chunks={'YC':10, 'YG':10}) as ds:
   print(ds)

fails with

Traceback (most recent call last):
  File "/home/jklymak/venv/AbHillInter3/lib/python3.7/site-packages/xmitgcm/mds_store.py", line 184, in open_mdsdataset
    iternum = int(iters)
TypeError: int() argument must be a string, a bytes-like object or a number, not 'list'

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "getSlice.py", line 20, in <module>
    geometry='cartesian', chunks={'YC':10, 'YG':10}) as ds:
  File "/home/jklymak/venv/AbHillInter3/lib/python3.7/site-packages/xmitgcm/mds_store.py", line 222, in open_mdsdataset
    for iternum in iters]
  File "/home/jklymak/venv/AbHillInter3/lib/python3.7/site-packages/xmitgcm/mds_store.py", line 222, in <listcomp>
    for iternum in iters]
  File "/home/jklymak/venv/AbHillInter3/lib/python3.7/site-packages/xmitgcm/mds_store.py", line 270, in open_mdsdataset
    ds = ds.chunk(chunks)
  File "/home/jklymak/venv/AbHillInter3/lib/python3.7/site-packages/xarray/core/dataset.py", line 1867, in chunk
    "object: %s" % bad_dims
ValueError: some chunks keys are not dimensions on this object: {'YG', 'YC'}

Does chunking work or should I just not bother? I just spent a couple hours chasing down a job abruptly terminating, when I finally figured out that it was simply running out of memory.

[rabernat]

Hi Jody! Thanks for uncovering a bug. The chunks option doesn't currently work and should not be there. We should remove it.

IMO, there is no point in having chunks that do not align with how the data are laid out on disk. For standard singlecpuio MDS files, I think the best we could hope for is to have an option for chunking the vertical dimension (as currently done in the llcreader module), because we know how to seek within a 3D file to extract a particular vertical slice without loading the whole file into memory. This is feasible and a PR in this direction would be welcome.

Another option would be to turn off singlecpuio and work on #28.

Note you can always rechunk your dataset, e.g.

ds = ds.chunk({'YC':10, 'YG':10})

However, this creates a new layer in your task graph, rather than changing the size of the initial read chunks.

[jklymak]

It seems #28 would be the best approach because then I assume reading and chunking could be multi-processor. I'm not really up on how dask works, but the end goal here was a simple slice at constant Y:

    ds = ds.isel(YC=200, YG=200, time=-1)
    ds.to_netcdf('./SliceMid.nc', engine='netcdf4')

Can dask/xarray realize that it would only need to read the tiles at YC/YG=200?

Thanks!

[rabernat]

Dask lazily maps a function call to a chunk. The function is expected to return a numpy array with the expected shape and dtype. In this case Xmitgcm provides the function, which is ultimately

xmitgcm/xmitgcm/utils.py

Lines 310 to 311 in 4a9c748

	def read_raw_data(datafile, dtype, shape, use_mmap=False, offset=0,
	order='C', partial_read=False):

The actual reading happens here:

xmitgcm/xmitgcm/utils.py

Lines 361 to 368 in 4a9c748

	with open(datafile, 'rb') as f:
	if use_mmap:
	data = np.memmap(f, dtype=dtype, mode='r', offset=offset,
	shape=tuple(shape), order=order)
	else:
	f.seek(offset)
	data = np.fromfile(f, dtype=dtype, count=number_of_values)
	data = data.reshape(shape, order=order)

The problem is that there is no easy way to write a function to return a "simple slice at constant Y", because the data are not contiguous on disk. You can't make a call to np.fromfile to get the data, which is ultimately what needs to happen.

Changing the on-disk chunk structure of your data to support different access patterns is what we created rechunker for.