DataTree: Align `from_dict` and `to_dict` behaviours to their Dataset equivalents

Is your feature request related to a problem?

This feature request arises from a "real-life" use case: I rely on Dataset.from_dict and Dataset.to_dict to convert Datasets to a dict before serializing them to JSON, and then loading back the JSON back to a Dataset with xarray.

JSON can be useful for small datasets, containing configuration with small values, that should be easily openable/modifiable by a human directly in a text editor, without using any library or script. Using xarray provide benefits as it solves questions like "how do I represent an array with coordinates in JSON": no need to reinvent super-languages above JSON, when the xarray serialization already does the job.

However, these capabilities do not exist (yet) for DataTree. It means that this "magic" method of using xarray as a way to dump to JSON is limited to flat structures.

Describe the solution you'd like

I would like the DataTree.from_dict and DataTree.to_dict to have a similar behaviour as their Dataset counterparts.

Currently the DataTree.from_dict method (https://xarray-datatree.readthedocs.io/en/stable/generated/datatree.DataTree.from_dict.html) expects A mapping from path names to xarray.Dataset, xarray.DataArray, or DataTree objects. It means a JSON cannot be reloaded back.

Currently the DataTree.to_dict method does not attempt to "serialize": the keys are paths and values are instance of Datasets. I would expect the Datasets to be replaced by their dictified version.

In [49]: xdt.to_dict()
Out[49]: 
{'/': <xarray.Dataset> 0B
 Dimensions:  ()
 Data variables:
     *empty*
 Attributes:
     top_level_attr:  Ho,
 '/parent': <xarray.Dataset> 96B
 Dimensions:  (dim_one: 3, dim_two: 2)
 Coordinates:
   * dim_one  (dim_one) int64 24B 10 20 30
 Dimensions without coordinates: dim_two
 Data variables:
     child_1  (dim_one) int64 24B 1 2 3
     child_2  (dim_two, dim_one) int64 48B 5 6 9 7 8 0}

The solution I would like resembled more this:

In [55]: datatree_dict = {path: xds.to_dict() for path, xds in xdt.to_dict().items()}

In [56]: datatree_dict
Out[56]: 
{'/': {'coords': {},
  'attrs': {'top_level_attr': 'Ho'},
  'dims': {},
  'data_vars': {}},
 '/parent': {'coords': {'dim_one': {'dims': ('dim_one',),
    'attrs': {},
    'data': [10, 20, 30]}},
  'attrs': {},
  'dims': {'dim_one': 3, 'dim_two': 2},
  'data_vars': {'child_1': {'dims': ('dim_one',),
    'attrs': {},
    'data': [1, 2, 3]},
   'child_2': {'dims': ('dim_two', 'dim_one'),
    'attrs': {'units': 'm', 'long_name': 'Hey'},
    'data': [[5, 6, 9], [7, 8, 0]]}}}}

In [58]: print(json.dumps(datatree_dict, indent=4))
{
    "/": {
        "coords": {},
        "attrs": {
            "top_level_attr": "Ho"
        },
        "dims": {},
        "data_vars": {}
    },
    "/parent": {
        "coords": {
            "dim_one": {
                "dims": [
                    "dim_one"
                ],
                "attrs": {},
                "data": [
                    10,
                    20,
                    30
                ]
            }
        },
        "attrs": {},
        "dims": {
            "dim_one": 3,
            "dim_two": 2
        },
        "data_vars": {
            "child_1": {
                "dims": [
                    "dim_one"
                ],
                "attrs": {},
                "data": [
                    1,
                    2,
                    3
                ]
            },
            "child_2": {
                "dims": [
                    "dim_two",
                    "dim_one"
                ],
                "attrs": {
                    "units": "m",
                    "long_name": "Hey"
                },
                "data": [
                    [
                        5,
                        6,
                        9
                    ],
                    [
                        7,
                        8,
                        0
                    ]
                ]
            }
        }
    }
}

Describe alternatives you've considered

Until now, I have been storing PurePosixPath-like variable names in Datasets. This helps organizing the configuration data, however, this loses the benefits of having scoped dimension names that DataTree provide.

Note: I did not want to add any custom parsing logic written by myself, not-standard and potentially breakable. The whole point of the from_dict and to_dict, to me, as I use them, is to be "universal-one-liners", a guarantee that an other xarray user can easily read the JSON I produced without writing themselves new parsing logic on their own.

Example:

• Create a Dataset with all xarray features (top-level attrs, variable-level attrs, 1D and 2D array with a shared dimension ; a dimension with coordinates and a dimension without coordinates), with PurePosixPath-like variable names
• Convert it to dict then dump to a JSON string
• Load back the JSON string to a Dataset
• Convert it to a DataTree to benefit from the tree hierarchy permitted by the path-like variable names

In [31]: xds = xr.Dataset({'parent/child_1': xr.DataArray([1,2,3], coords={"dim_one": [10,20,30]}), "parent/child_2": xr.DataArray([[5,6,9],[7,8,0]], di
    ...: ms=("dim_two", "dim_one"), attrs={"units": "m", "long_name": "Hey"})}, attrs={"top_level_attr": "Ho"})

In [32]: xds.to_dict()
Out[32]: 
{'coords': {'dim_one': {'dims': ('dim_one',),
   'attrs': {},
   'data': [10, 20, 30]}},
 'attrs': {'top_level_attr': 'Ho'},
 'dims': {'dim_one': 3, 'dim_two': 2},
 'data_vars': {'parent/child_1': {'dims': ('dim_one',),
   'attrs': {},
   'data': [1, 2, 3]},
  'parent/child_2': {'dims': ('dim_two', 'dim_one'),
   'attrs': {'units': 'm', 'long_name': 'Hey'},
   'data': [[5, 6, 9], [7, 8, 0]]}}}

In [33]: print(json.dumps(xds.to_dict(), indent=4))
{
    "coords": {
        "dim_one": {
            "dims": [
                "dim_one"
            ],
            "attrs": {},
            "data": [
                10,
                20,
                30
            ]
        }
    },
    "attrs": {
        "top_level_attr": "Ho"
    },
    "dims": {
        "dim_one": 3,
        "dim_two": 2
    },
    "data_vars": {
        "parent/child_1": {
            "dims": [
                "dim_one"
            ],
            "attrs": {},
            "data": [
                1,
                2,
                3
            ]
        },
        "parent/child_2": {
            "dims": [
                "dim_two",
                "dim_one"
            ],
            "attrs": {
                "units": "m",
                "long_name": "Hey"
            },
            "data": [
                [
                    5,
                    6,
                    9
                ],
                [
                    7,
                    8,
                    0
                ]
            ]
        }
    }
}

In [41]: reloaded = xr.Dataset.from_dict(json.loads(json.dumps(xds.to_dict(), indent=4)))

In [42]: reloaded
Out[42]: 
<xarray.Dataset> 96B
Dimensions:         (dim_one: 3, dim_two: 2)
Coordinates:
  * dim_one         (dim_one) int64 24B 10 20 30
Dimensions without coordinates: dim_two
Data variables:
    parent/child_1  (dim_one) int64 24B 1 2 3
    parent/child_2  (dim_two, dim_one) int64 48B 5 6 9 7 8 0
Attributes:
    top_level_attr:  Ho

In [43]: import xarray.core.datatree as dt

In [44]: xdt = dt.DataTree()

In [45]: for varname in reloaded: xdt[varname] = reloaded[varname]

In [46]: xdt
Out[46]: 
DataTree('None', parent=None)
└── DataTree('parent')
        Dimensions:  (dim_one: 3, dim_two: 2)
        Coordinates:
          * dim_one  (dim_one) int64 24B 10 20 30
        Dimensions without coordinates: dim_two
        Data variables:
            child_1  (dim_one) int64 24B 1 2 3
            child_2  (dim_two, dim_one) int64 48B 5 6 9 7 8 0

Root-level attrs are lost but can be added again.

In [47]: xdt.attrs.update(xds.attrs)

In [48]: xdt
Out[48]: 
DataTree('None', parent=None)
│   Dimensions:  ()
│   Data variables:
│       *empty*
│   Attributes:
│       top_level_attr:  Ho
└── DataTree('parent')
        Dimensions:  (dim_one: 3, dim_two: 2)
        Coordinates:
          * dim_one  (dim_one) int64 24B 10 20 30
        Dimensions without coordinates: dim_two
        Data variables:
            child_1  (dim_one) int64 24B 1 2 3
            child_2  (dim_two, dim_one) int64 48B 5 6 9 7 8 0

Additional context

No response

Start with the DataTree.from_dict and DataTree.to_dict entry points and compare their documented behavior with Dataset.from_dict and Dataset.to_dict. Trace how nested paths, datasets, dimensions, coordinates, variables, and attributes are represented. Done means a DataTree can round-trip through JSON-compatible dictionaries while preserving the demonstrated tree and dataset content, with tests covering the conversion.