Add shape information in to_xarray

- Centralize attribute metadata logic
- Include :shape arrays in imas.util.to_xarray
- Test that netCDF files written with xarray can be read as DBEntry

Author: maarten-ic

imas/_to_xarray.py

@@ -8,7 +8,7 @@
 from imas.ids_data_type import IDSDataType
 
 fillvals = {
-    IDSDataType.INT: -(2**31) + 1,
+    IDSDataType.INT: numpy.int32(-(2**31) + 1),
     IDSDataType.STR: "",
     IDSDataType.FLT: numpy.nan,
     IDSDataType.CPX: numpy.nan * (1 + 1j),
@@ -54,22 +54,24 @@ def to_xarray(ids: IDSToplevel, *paths: str) -> xarray.Dataset:
             if paths and path not in paths:
                 continue
             dimensions = ()
-            data = ""
+            data = b""
         else:
             dimensions = tensorizer.get_dimensions(path)
             data = tensorizer.tensorize(path, fillvals[metadata.data_type])
 
-        attrs = dict(documentation=metadata.documentation)
-        if metadata.units:
-            attrs["units"] = metadata.units
-        if dimensions:
-            coordinates = tensorizer.filter_coordinates(path)
-            if coordinates:
-                coordinate_names.update(coordinates.split(" "))
-                attrs["coordinates"] = coordinates
-
+        attrs = tensorizer.get_attributes(path, fillvals)
+        if "coordinates" in attrs:
+            coordinate_names.update(attrs["coordinates"].split(" "))
         data_vars[var_name] = (dimensions, data, attrs)
 
+        # :shape array for sparse data
+        if path in tensorizer.shapes and metadata.ndim:
+            shape_name = f"{var_name}:shape"
+            dimensions = tensorizer.get_shape_dimensions(path)
+            data = tensorizer.shapes[path]
+            attrs = tensorizer.get_shape_attributes(var_name)
+            data_vars[shape_name] = (dimensions, data, attrs)
+
     # Remove coordinates from data_vars and put in coordinates mapping:
     coordinates = {}
     for coordinate_name in coordinate_names:

imas/backends/netcdf/ids2nc.py

@@ -78,50 +78,17 @@ def create_variables(self) -> None:
                 var = self.group.createVariable(var_name, dtype, dimensions, **kwargs)
 
             # Fill metadata attributes
-            var.documentation = metadata.documentation
-            if metadata.units:
-                var.units = metadata.units
-
-            ancillary_variables = " ".join(
-                error_var
-                for error_var in [f"{var_name}_error_upper", f"{var_name}_error_lower"]
-                if error_var in self.filled_variables
-            )
-            if ancillary_variables:
-                var.ancillary_variables = ancillary_variables
-
-            if metadata.data_type is not IDSDataType.STRUCT_ARRAY:
-                coordinates = self.filter_coordinates(path)
-                if coordinates:
-                    var.coordinates = coordinates
-
-            # Sparsity and :shape array
-            if path in self.shapes:
-                if not metadata.ndim:
-                    # Doesn't need a :shape array:
-                    var.sparse = "Sparse data, missing data is filled with _FillValue"
-                    var.sparse += f" ({default_fillvals[metadata.data_type]})"
-
-                else:
-                    shape_name = f"{var_name}:shape"
-                    var.sparse = f"Sparse data, data shapes are stored in {shape_name}"
-
-                    # Create variable to store data shape
-                    dimensions = self.get_dimensions(self.ncmeta.aos.get(path)) + (
-                        f"{metadata.ndim}D",
-                    )
-                    shape_var = self.group.createVariable(
-                        shape_name,
-                        SHAPE_DTYPE,
-                        dimensions,
-                    )
-                    doc_indices = ",".join(chr(ord("i") + i) for i in range(3))
-                    shape_var.documentation = (
-                        f"Shape information for {var_name}.\n"
-                        f"{shape_name}[{doc_indices},:] describes the shape of filled "
-                        f"data of {var_name}[{doc_indices},...]. Data outside this "
-                        "shape is unset (i.e. filled with _Fillvalue)."
-                    )
+            var.setncatts(self.get_attributes(path, default_fillvals))
+
+            # :shape array for sparse data
+            if path in self.shapes and metadata.ndim:
+                shape_name = f"{var_name}:shape"
+                # Create variable to store data shape
+                dimensions = self.get_shape_dimensions(path)
+                shape_var = self.group.createVariable(
+                    shape_name, SHAPE_DTYPE, dimensions
+                )
+                shape_var.setncatts(self.get_shape_attributes(var_name))
 
     def store_data(self) -> None:
         """Store data in the netCDF variables"""

imas/backends/netcdf/ids_tensorizer.py

@@ -3,7 +3,7 @@
 """Tensorization logic to convert IDSs to netCDF files and/or xarray Datasets."""
 
 from collections import deque
-from typing import List, Tuple
+from typing import List, Tuple, Dict
 
 import numpy
 
@@ -62,6 +62,11 @@ def get_dimensions(self, path: str) -> Tuple[str, ...]:
         """
         return self.ncmeta.get_dimensions(path, self.homogeneous_time)
 
+    def get_shape_dimensions(self, path: str) -> Tuple[str, ...]:
+        """Get dimensions names for shape array of the tensorized variable"""
+        ndim = self.ids.metadata[path].ndim
+        return self.get_dimensions(self.ncmeta.aos.get(path, "")) + (f"{ndim}D",)
+
     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
@@ -173,6 +178,54 @@ def filter_coordinates(self, path: str) -> str:
             if coordinate in self.filled_variables
         )
 
+    def get_attributes(self, path: str, fillvals: dict) -> Dict[str, str]:
+        """Get metadata attributes of the tensorized variable"""
+        metadata = self.ids.metadata[path]
+        var_name = path.replace("/", ".")
+
+        assert metadata.documentation is not None
+        attrs = {"documentation": metadata.documentation}
+        if metadata.units:
+            attrs["units"] = metadata.units
+
+        ancillary_variables = " ".join(
+            error_var
+            for error_var in [f"{var_name}_error_upper", f"{var_name}_error_lower"]
+            if error_var in self.filled_variables
+        )
+        if ancillary_variables:
+            attrs["ancillary_variables"] = ancillary_variables
+
+        if metadata.data_type is not IDSDataType.STRUCT_ARRAY:
+            coordinates = self.filter_coordinates(path)
+            if coordinates:
+                attrs["coordinates"] = coordinates
+
+        # Sparsity
+        if path in self.shapes:
+            if not metadata.ndim:
+                # Doesn't need a :shape array
+                attrs["sparse"] = (
+                    "Sparse data, missing data is filled with _FillValue"
+                    f" ({fillvals[metadata.data_type]})"
+                )
+            else:
+                attrs["sparse"] = (
+                    f"Sparse data, data shapes are stored in {var_name}:shape"
+                )
+
+        return attrs
+
+    def get_shape_attributes(self, var_name: str) -> Dict[str, str]:
+        doc_indices = ",".join(chr(ord("i") + i) for i in range(3))
+        documentation = (
+            f"Shape information for {var_name}.\n"
+            f"{var_name}:shape[{doc_indices},:] describes the shape of filled "
+            f"data of {var_name}[{doc_indices},...]. Data outside this "
+            "shape is unset (i.e. filled with _Fillvalue)."
+        )
+        return {"documentation": documentation}
+
     def tensorize(self, path, fillvalue):
         """
         Tensorizes the data at the given path with the specified fill value.

imas/test/test_to_xarray.py

@@ -1,11 +1,13 @@
 import numpy as np
+import netCDF4
 import pytest
 
 import imas
 import imas.training
+from imas.test.test_helpers import compare_children
 from imas.util import to_xarray
 
-pytest.importorskip("xarray")
+xarray = pytest.importorskip("xarray")
 
 
 @pytest.fixture
@@ -94,3 +96,22 @@ def test_to_xarray():
     ds3 = to_xarray(ids, "profiles_1d/electrons/temperature")
     assert ds1.equals(ds2)
     assert ds2.equals(ds3)
+
+
+@pytest.mark.parametrize("idsname", ["core_profiles", "equilibrium"])
+def test_roundtrip_xarray_netcdf(tmp_path, entry, idsname):
+    ids = entry.get(idsname)
+    xrds = to_xarray(ids)
+    fname = f"{tmp_path}/test-{idsname}-xarray.nc"
+    # First write mandatory file-level metadata
+    with netCDF4.Dataset(fname, "x") as ds:
+        ds.data_dictionary_version = imas.util.get_data_dictionary_version(ids)
+    # Then use xarray to write the IDS
+    xrds.to_netcdf(fname, "a", format="NETCDF4", group=f"{idsname}/0")
+    # And read it back with a DBEntry
+    with imas.DBEntry(fname, "r") as entry:
+        ids2 = entry.get(idsname)
+    compare_children(ids, ids2)
+    # Reading the netCDF file with xarray should produce an identical dataset
+    ncxrds = xarray.load_dataset(fname, group=f"{idsname}/0")
+    assert xrds.equals(ncxrds)