open_datatree() from zarr creates issues with `kwargs`

[aladinor]

open_datatree() from zarr creates issues with `kwargs`

[TomNicholas]

Hi @aladinor - have you checked that this behaviour is actually different from what xr.open_dataset will do for one dataset? Because the aim here should be compliance with that behaviour, and clarity as to what options should be passed if you want different behaviour (e.g. passing chunks={}).

[aladinor]

Hi @TomNicholas. This is a really good question. So the main thing that I found out is that xarray.open_dataset has its default parameters. Thus, when passing the **kwargs dictionary from datatree.open_datatree to xarray.open_dataset, it will overwrite all the default parameters and set them None. At least in the test I performed, if none of those default arguments were passed, xr.open_dataset would behave as I mentioned in #328. I think that If we want to keep xr.open_dataset, all the parameters should be passed within the **kwargs dictionary or try to find a way to not overwrite the default args in xr.open_dataset. I could not find a way to do that without retyping/passing all of them from datatree.open_datatree to xr.open_dataset . That's why I felt using xr.open_zarr was more convenient. However, I will recheck it using one dataset

[aladinor]

Hi @TomNicholas. This could be a reproducible example for you to check.

import xarray as xr
import datatree
import numpy as np
import pandas as pd


def main():
    data = 283 + 5 * np.random.randn(5, 3, 4)
    times = pd.date_range("2018-01-01", periods=5)
    lons = np.linspace(-120, -90, 4)
    lats = np.linspace(25, 55, 3)
    temp = xr.DataArray(data, coords=[times, lats, lons], dims=["time", "lat", "lon"]).chunk({'time': 1,
                                                                                              'lat': -1,
                                                                                              'lon': -1})
    ds = xr.Dataset(data_vars={'Temp': temp})
    ds2 = xr.Dataset(data_vars={'Temp_C': temp - 273.15}).chunk({'time': 1, 'lat': 1, 'lon': 1})
    dtree = datatree.DataTree(data=ds)
    datatree.DataTree(ds2, parent=dtree, name='set2')
    print(dtree.Temp.chunksizes)
    print(dtree.set2.Temp_C.chunksizes)
    path_zarr = "test.zarr"
    dtree.to_zarr(path_zarr)

    dtree_back = datatree.open_datatree(path_zarr, engine='zarr')
    print(dtree_back.Temp.chunksizes)
    print(dtree_back.set2.Temp_C.chunksizes)

    ...


if __name__ == "__main__":
    main()

And the output

Frozen({'time': (1, 1, 1, 1, 1), 'lat': (3,), 'lon': (4,)})
Frozen({'time': (1, 1, 1, 1, 1), 'lat': (1, 1, 1), 'lon': (1, 1, 1, 1)})
Frozen({})
Frozen({})

The expected output without passing any parameters to open_datatree should be to keep the same chunk sizes. However, it didn't.
Let me know your thoughts.

[aladinor]

Hi @TomNicholas. I looked again at xr.open_dataset documentation and I found this

In order to reproduce the default behavior of xr.open_zarr(...) use xr.open_dataset(..., engine='zarr', chunks={}). chunks='auto' will use dask auto chunking taking into account the engine preferred chunks. See dask chunking for more details.

I think based on I will close this issue and the PR. Thanks for your time!