Implement imas.util.to_xarray

Reuses most of the tensorization and metadata logic from the netCDF export.

Author: maarten-ic

imas/_to_xarray.py

@@ -0,0 +1,73 @@
+# xarray is an optional dependency, but this module won't be imported when xarray is not
+# available
+import numpy
+import xarray
+
+from imas.ids_toplevel import IDSToplevel
+from imas.backends.netcdf.ids_tensorizer import IDSTensorizer
+from imas.ids_data_type import IDSDataType
+
+fillvals = {
+    IDSDataType.INT: -(2**31) + 1,
+    IDSDataType.STR: "",
+    IDSDataType.FLT: numpy.nan,
+    IDSDataType.CPX: numpy.nan * (1 + 1j),
+}
+
+
+def to_xarray(ids: IDSToplevel, *paths: str) -> xarray.Dataset:
+    """See :func:`imas.util.to_xarray`"""
+    # We really need an IDS toplevel element
+    if not isinstance(ids, IDSToplevel):
+        raise TypeError(
+            f"to_xarray needs a toplevel IDS element as first argument, but got {ids!r}"
+        )
+
+    # Valid path can use / or . as separator, but IDSTensorizer expects /. The following
+    # block checks if the paths are valid, and by using "metadata.path_string" we ensure
+    # that / are used as separator.
+    try:
+        paths = [ids.metadata[path].path_string for path in paths]
+    except KeyError as exc:
+        raise ValueError(str(exc)) from None
+
+    # Converting lazy-loaded IDSs requires users to specify at least one path
+    if ids._lazy and not paths:
+        raise RuntimeError(
+            "This IDS is lazy loaded. Please provide at least one path to convert to"
+            " xarray."
+        )
+
+    # Use netcdf IDS Tensorizer to tensorize the data and determine metadata
+    tensorizer = IDSTensorizer(ids, paths)
+    tensorizer.include_coordinate_paths()
+    tensorizer.collect_filled_data()
+    tensorizer.determine_data_shapes()
+
+    data_vars = {}
+    coordinate_names = set()
+    for path in tensorizer.filled_data:
+        var_name = path.replace("/", ".")
+        metadata = ids.metadata[path]
+        if metadata.data_type in (IDSDataType.STRUCTURE, IDSDataType.STRUCT_ARRAY):
+            continue  # We don't store these in xarray
+
+        dimensions = tensorizer.ncmeta.get_dimensions(path, tensorizer.homogeneous_time)
+        data = tensorizer.tensorize(path, fillvals[metadata.data_type])
+
+        attrs = dict(documentation=metadata.documentation)
+        if metadata.units:
+            attrs["units"] = metadata.units
+        coordinates = tensorizer.filter_coordinates(path)
+        if coordinates:
+            coordinate_names.update(coordinates.split(" "))
+            attrs["coordinates"] = coordinates
+
+        data_vars[var_name] = (dimensions, data, attrs)
+
+    # Remove coordinates from data_vars and put in coordinates mapping:
+    coordinates = {}
+    for coordinate_name in coordinate_names:
+        coordinates[coordinate_name] = data_vars.pop(coordinate_name)
+
+    return xarray.Dataset(data_vars, coordinates)

imas/backends/netcdf/ids_tensorizer.py

@@ -2,6 +2,7 @@
 # You should have received the IMAS-Python LICENSE file with this project.
 """Tensorization logic to convert IDSs to netCDF files and/or xarray Datasets."""
 
+from collections import deque
 from typing import List
 
 import numpy
@@ -53,6 +54,26 @@ def __init__(self, ids: IDSToplevel, paths_to_tensorize: List[str]) -> None:
         self.shapes = {}
         """Map of IDS paths to data shape arrays."""
 
+    def include_coordinate_paths(self) -> None:
+        """Append all paths that are coordinates of self.paths_to_tensorize"""
+        # Use a queue so we can also take coordinates of coordinates into account
+        queue = deque(self.paths_to_tensorize)
+        # Include all parent AoS as well:
+        for path in self.paths_to_tensorize:
+            while path:
+                path, _, _ = path.rpartition("/")
+                if self.ncmeta.get_dimensions(path, self.homogeneous_time):
+                    queue.append(path)
+
+        self.paths_to_tensorize = []
+        while queue:
+            path = queue.popleft()
+            if path in self.paths_to_tensorize:
+                continue  # already processed
+            self.paths_to_tensorize.append(path)
+            for coordinate in self.ncmeta.get_coordinates(path, self.homogeneous_time):
+                queue.append(coordinate.replace(".", "/"))
+
     def collect_filled_data(self) -> None:
         """Collect all filled data in the IDS and determine dimension sizes.
 

imas/test/test_to_xarray.py

@@ -0,0 +1,94 @@
+import numpy as np
+import pytest
+
+import imas
+import imas.training
+from imas.util import to_xarray
+
+pytest.importorskip("xarray")
+
+
+@pytest.fixture
+def entry(requires_imas, monkeypatch):
+    monkeypatch.setenv("IMAS_VERSION", "3.39.0")  # Use fixed DD version
+    return imas.training.get_training_db_entry()
+
+
+def test_to_xarray_invalid_argtype():
+    ids = imas.IDSFactory("3.39.0").core_profiles()
+
+    with pytest.raises(TypeError):
+        to_xarray("test")
+    with pytest.raises(TypeError):
+        to_xarray(ids.time)
+    with pytest.raises(TypeError):
+        to_xarray(ids.ids_properties)
+
+
+def test_to_xarray_invalid_paths():
+    ids = imas.IDSFactory("3.39.0").core_profiles()
+
+    with pytest.raises(ValueError, match="xyz"):
+        to_xarray(ids, "xyz")
+    with pytest.raises(ValueError, match="ids_properties/xyz"):
+        to_xarray(ids, "ids_properties/xyz")
+    with pytest.raises(ValueError, match="Xtime"):
+        to_xarray(ids, "time", "Xtime")
+
+
+def validate_trainingdb_electron_temperature_dataset(ds):
+    assert ds.sizes == {"time": 3, "profiles_1d.grid.rho_tor_norm:i": 101}
+    assert ds.data_vars.keys() == {"profiles_1d.electrons.temperature"}
+    assert ds.coords.keys() == {"time", "profiles_1d.grid.rho_tor_norm"}
+
+    # Check that values are loaded as expected
+    assert np.allclose(ds["time"], [3.987222, 432.937598, 792.0])
+    assert np.allclose(
+        ds.isel(time=1)["profiles_1d.electrons.temperature"][10:13],
+        [17728.81703089, 17440.78020568, 17139.35431082],
+    )
+
+
+def test_to_xarray_lazy_loaded(entry):
+    ids = entry.get("core_profiles", lazy=True)
+
+    with pytest.raises(RuntimeError):
+        to_xarray(ids)
+
+    ds = to_xarray(ids, "profiles_1d.electrons.temperature")
+    validate_trainingdb_electron_temperature_dataset(ds)
+
+
+def test_to_xarray_from_trainingdb(entry):
+    ids = entry.get("core_profiles")
+
+    ds = to_xarray(ids)
+    validate_trainingdb_electron_temperature_dataset(
+        ds["profiles_1d.electrons.temperature"].to_dataset()
+    )
+    ds = to_xarray(ids, "profiles_1d.electrons.temperature")
+    validate_trainingdb_electron_temperature_dataset(ds)
+
+    ds = to_xarray(
+        ids, "profiles_1d.electrons.temperature", "profiles_1d/electrons/density"
+    )
+    assert ds.data_vars.keys() == {
+        "profiles_1d.electrons.temperature",
+        "profiles_1d.electrons.density",
+    }
+
+
+def test_to_xarray():
+    ids = imas.IDSFactory("3.39.0").core_profiles()
+
+    ids.profiles_1d.resize(2)
+    ids.profiles_1d[0].electrons.temperature = [1.0, 2.0]
+    ids.profiles_1d[0].grid.rho_tor_norm = [0.0, 1.0]
+    ids.profiles_1d[0].time = 0.0
+
+    # These should all be identical:
+    ds1 = to_xarray(ids)
+    ds2 = to_xarray(ids, "profiles_1d.electrons.temperature")
+    ds3 = to_xarray(ids, "profiles_1d/electrons/temperature")
+    assert ds1.equals(ds2)
+    assert ds2.equals(ds3)

imas/util.py

@@ -1,8 +1,6 @@
 # This file is part of IMAS-Python.
 # You should have received the IMAS-Python LICENSE file with this project.
-"""Collection of useful helper methods when working with IMAS-Python.
-"""
-
+"""Collection of useful helper methods when working with IMAS-Python."""
 
 import logging
 import re
@@ -524,3 +522,55 @@ def get_data_dictionary_version(obj: Union[IDSBase, DBEntry, IDSFactory]) -> str
     if isinstance(obj, IDSBase):
         return obj._version
     raise TypeError(f"Cannot get data dictionary version of '{type(obj)}'")
+
+
+def to_xarray(ids: IDSToplevel, *paths: str) -> Any:
+    """Convert an IDS to an xarray Dataset.
+
+    Args:
+        ids: An IDS toplevel element
+        paths: Optional list of element paths to convert to xarray. The full IDS will be
+            converted to an xarray Dataset if no paths are provided.
+
+            Paths must not contain indices, and may use a ``/`` or a ``.`` as separator.
+            For example, ``"profiles_1d(itime)/electrons/density"`` is not allowed as
+            path, use ``"profiles_1d/electrons/density"`` or
+            ``profiles_1d.electrons.density"`` instead.
+
+            Coordinates to the quantities in the requested paths will also be included
+            in the xarray Dataset.
+
+    Returns:
+        An ``xarray.Dataset`` object.
+
+    Examples:
+        .. code-block:: python
+
+            # Convert the whole IDS to an xarray Dataset
+            ds = imas.util.to_xarray(ids)
+
+            # Convert only some elements in the IDS (including their coordinates)
+            ds = imas.util.to_xarray(
+                ids,
+                "profiles_1d/electrons/density",
+                "profiles_1d/electrons/temperature",
+            )
+
+            # Paths can be provided with "/" or "." as separator
+            ds = imas.util.to_xarray(
+                ids,
+                "profiles_1d.electrons.density",
+                "profiles_1d.electrons.temperature",
+            )
+
+    See Also:
+        https://docs.xarray.dev/en/stable/generated/xarray.Dataset.html
+    """
+    try:
+        import xarray  # noqa: F401
+    except ImportError:
+        raise RuntimeError("xarray is not available, cannot convert the IDS to xarray.")
+
+    from imas._to_xarray import to_xarray
+
+    return to_xarray(ids, *paths)