imas-python 1.2.1.dev13__py3-none-any.whl → 2.0.0.post1__py3-none-any.whl

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/_version.py

@@ -1,8 +1,13 @@
-# file generated by setuptools_scm
+# file generated by setuptools-scm
 # don't change, don't track in version control
+
+__all__ = ["__version__", "__version_tuple__", "version", "version_tuple"]
+
 TYPE_CHECKING = False
 if TYPE_CHECKING:
-    from typing import Tuple, Union
+    from typing import Tuple
+    from typing import Union
+
     VERSION_TUPLE = Tuple[Union[int, str], ...]
 else:
     VERSION_TUPLE = object
@@ -12,5 +17,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '1.2.1.dev13'
-__version_tuple__ = version_tuple = (1, 2, 1, 'dev13')
+__version__ = version = '2.0.0.post1'
+__version_tuple__ = version_tuple = (2, 0, 0, 'post1')

imas/backends/imas_core/al_context.py

@@ -10,6 +10,7 @@ from typing import TYPE_CHECKING, Any, Callable, Iterator, List, Optional, Tuple
 
 import numpy
 
+import imas
 from imas.backends.imas_core.imas_interface import ll_interface
 from imas.exception import LowlevelError
 from imas.ids_defs import (
@@ -280,6 +281,18 @@ class LazyALContext:
         self.context = None
         """Potential weak reference to opened context."""
 
+    def get_child(self, child):
+        """
+        Retrieve a child entry from the field.
+
+        Args:
+            child (str): The name or identifier of the child entry to retrieve.
+
+        Returns:
+            The child entry retrieved from the database.
+        """
+        imas.backends.imas_core.db_entry_helpers._get_child(child, self)
+
     def get_context(self) -> ALContext:
         """Create and yield the actual ALContext."""
         if self.dbentry._db_ctx is not self.dbentry_ctx:

imas/backends/imas_core/db_entry_helpers.py

@@ -22,7 +22,7 @@ def get_children(
     structure: IDSStructure,
     ctx: ALContext,
     time_mode: int,
-    nbc_map: Optional[NBCPathMap],
+    nbc_map: Optional["NBCPathMap"],
 ) -> None:
     """Recursively get all children of an IDSStructure."""
     # NOTE: changes in this method must be propagated to _get_child and vice versa
@@ -77,15 +77,11 @@ def get_children(
                 getattr(structure, name)._IDSPrimitive__value = data
 
 
-def _get_child(child: IDSBase, ctx: Optional[LazyALContext]):
+def _get_child(child: IDSBase, ctx: LazyALContext):
     """Get a single child when required (lazy loading)."""
     # NOTE: changes in this method must be propagated to _get_children and vice versa
     #   Performance: this method is specialized for the lazy get
 
-    # ctx can be None when the parent structure does not exist in the on-disk DD version
-    if ctx is None:
-        return  # There is no data to be loaded
-
     time_mode = ctx.time_mode
     if time_mode == IDS_TIME_MODE_INDEPENDENT and child.metadata.type.is_dynamic:
         return  # skip dynamic (time-dependent) nodes
@@ -148,7 +144,7 @@ def put_children(
     ctx: ALContext,
     time_mode: int,
     is_slice: bool,
-    nbc_map: Optional[NBCPathMap],
+    nbc_map: Optional["NBCPathMap"],
     verify_maxoccur: bool,
 ) -> None:
     """Recursively put all children of an IDSStructure"""

imas/backends/imas_core/imas_interface.py

@@ -32,7 +32,7 @@ except ImportError as exc:
     imasdef = None
     lowlevel = None
     logger.critical(
-        "Could not import 'al_core': %s. Some functionality is not available.",
+        "Could not import 'imas_core': %s. Some functionality is not available.",
         exc,
     )
 

imas/backends/imas_core/mdsplus_model.py

@@ -1,7 +1,6 @@
 # Helper functions to create MDSPlus reference models
 # and store them in a cache directory (.cache/imas/MDSPlus/name-HASH/)
-"""Module for generating and working with MDSplus models.
-"""
+"""Module for generating and working with MDSplus models."""
 
 import errno
 import getpass
@@ -235,6 +234,14 @@ def model_exists(path: Path) -> bool:
     )
 
 
+def transform_with_xslt(xslt_processor, source, xslfile, output_file):
+    return xslt_processor.transform_to_file(
+        source_file=str(source),
+        stylesheet_file=str(xslfile),
+        output_file=str(output_file),
+    )
+
+
 def create_model_ids_xml(cache_dir_path, fname, version):
     """Use Saxon/C to compile an ids.xml suitable for creating an MDSplus model."""
     try:
@@ -243,40 +250,32 @@ def create_model_ids_xml(cache_dir_path, fname, version):
 
             with PySaxonProcessor(license=False) as proc:
                 xslt_processor = proc.new_xslt30_processor()
-
-                xslt_processor.compile_stylesheet(stylesheet_file=str(xslfile))
-
-                input_xml = get_dd_xml(version) if version else None
-                if fname:
-                    source_file = str(fname)
-                elif input_xml:
-                    source_file = input_xml  # Use standard input for the XML string
-                else:
-                    raise ValueError(
-                        "Either 'fname' or 'version' must be provided to generate XML."
-                    )
-
-                # xdm_ddgit = proc.make_string_value(str(version or fname))
-                # xsltproc.set_parameter("DD_GIT_DESCRIBE", xdm_ddgit)
-                # xdm_algit = proc.make_string_value(os.environ.get
-                # ("AL_VERSION", "0.0.0"))
-                # xsltproc.set_parameter("AL_GIT_DESCRIBE", xdm_algit)
-                # Transform XML
-                result = xslt_processor.transform_to_file(
-                    source_file=source_file,
-                    output_file=str(output_file),
-                    initial_template_params={
-                        "DD_GIT_DESCRIBE": str(version or fname),
-                        "AL_GIT_DESCRIBE": os.environ.get("AL_VERSION", "0.0.0"),
-                    },
+                xdm_ddgit = proc.make_string_value(str(version) or fname)
+                xslt_processor.set_parameter("DD_GIT_DESCRIBE", xdm_ddgit)
+                xdm_algit = proc.make_string_value(
+                    os.environ.get("AL_VERSION", "0.0.0")
                 )
-
-                if result is False:
-                    logger.error(
-                        "Transformation failed: Check Saxon/C logs for details."
-                    )
-                    raise RuntimeError("Saxon/C XSLT transformation failed.")
-
+                xslt_processor.set_parameter("AL_GIT_DESCRIBE", xdm_algit)
+                if (
+                    fname is not None
+                    and fname != "-"
+                    and fname != ""
+                    and os.path.exists(fname)
+                ):
+                    transform_with_xslt(xslt_processor, fname, xslfile, output_file)
+                elif version is not None and version != "":
+                    xml_string = get_dd_xml(version)
+
+                    with tempfile.NamedTemporaryFile(
+                        delete=True, mode="w+b"
+                    ) as temp_file:
+                        temp_file.write(xml_string)
+                        temp_file.seek(0)
+                        transform_with_xslt(
+                            xslt_processor, temp_file.name, xslfile, output_file
+                        )
+                else:
+                    raise MDSPlusModelError("Either fname or version must be provided")
     except Exception as e:
         if fname:
             logger.error("Error making MDSplus model IDS.xml for %s", fname)

imas/backends/netcdf/db_entry_nc.py

@@ -11,7 +11,7 @@ from imas.backends.db_entry_impl import (
 from imas.backends.netcdf.ids2nc import IDS2NC
 from imas.backends.netcdf.nc2ids import NC2IDS
 from imas.exception import DataEntryException, InvalidNetCDFEntry
-from imas.ids_convert import NBCPathMap, convert_ids
+from imas.ids_convert import NBCPathMap, dd_version_map_from_factories
 from imas.ids_factory import IDSFactory
 from imas.ids_toplevel import IDSToplevel
 
@@ -108,10 +108,6 @@ class NCDBEntryImpl(DBEntryImpl):
             else:
                 func = "get_sample"
             raise NotImplementedError(f"`{func}` is not available for netCDF files.")
-        if lazy:
-            raise NotImplementedError(
-                "Lazy loading is not implemented for netCDF files."
-            )
 
         # Check if the IDS/occurrence exists, and obtain the group it is stored in
         try:
@@ -123,14 +119,19 @@ class NCDBEntryImpl(DBEntryImpl):
 
         # Load data into the destination IDS
         if self._ds_factory.dd_version == destination._dd_version:
-            NC2IDS(group, destination).run()
+            NC2IDS(group, destination, destination.metadata, None).run(lazy)
         else:
-            # FIXME: implement automatic conversion using nbc_map
-            #   As a work-around: do an explicit conversion, but automatic conversion
-            #   will also be needed to implement lazy loading.
-            ids = self._ds_factory.new(ids_name)
-            NC2IDS(group, ids).run()
-            convert_ids(ids, None, target=destination)
+            # Construct relevant NBCPathMap, the one we get from DBEntry has the reverse
+            # mapping from what we need. The imas_core logic does the mapping from
+            # in-memory to on-disk, while we take what is on-disk and map it to
+            # in-memory.
+            ddmap, source_is_older = dd_version_map_from_factories(
+                ids_name, self._ds_factory, self._factory
+            )
+            nbc_map = ddmap.old_to_new if source_is_older else ddmap.new_to_old
+            NC2IDS(
+                group, destination, self._ds_factory.new(ids_name).metadata, nbc_map
+            ).run(lazy)
 
         return destination
 

imas/backends/netcdf/ids2nc.py

@@ -1,21 +1,13 @@
 # This file is part of IMAS-Python.
 # You should have received the IMAS-Python LICENSE file with this project.
-"""NetCDF IO support for IMAS-Python. Requires [netcdf] extra dependencies.
-"""
-
-from typing import Iterator, Tuple
+"""NetCDF IO support for IMAS-Python. Requires [netcdf] extra dependencies."""
 
 import netCDF4
-import numpy
 from packaging import version
 
-from imas.backends.netcdf.nc_metadata import NCMetadata
+from imas.backends.netcdf.ids_tensorizer import SHAPE_DTYPE, IDSTensorizer, dtypes
 from imas.exception import InvalidNetCDFEntry
-from imas.ids_base import IDSBase
 from imas.ids_data_type import IDSDataType
-from imas.ids_defs import IDS_TIME_MODE_HOMOGENEOUS
-from imas.ids_struct_array import IDSStructArray
-from imas.ids_structure import IDSStructure
 from imas.ids_toplevel import IDSToplevel
 
 default_fillvals = {
@@ -24,36 +16,9 @@ default_fillvals = {
     IDSDataType.FLT: netCDF4.default_fillvals["f8"],
     IDSDataType.CPX: netCDF4.default_fillvals["f8"] * (1 + 1j),
 }
-dtypes = {
-    IDSDataType.INT: numpy.dtype(numpy.int32),
-    IDSDataType.STR: str,
-    IDSDataType.FLT: numpy.dtype(numpy.float64),
-    IDSDataType.CPX: numpy.dtype(numpy.complex128),
-}
-SHAPE_DTYPE = numpy.int32
-
-
-def nc_tree_iter(
-    node: IDSStructure, aos_index: Tuple[int, ...] = ()
-) -> Iterator[Tuple[Tuple[int, ...], IDSBase]]:
-    """Tree iterator that tracks indices of all ancestor array of structures.
-
-    Args:
-        node: IDS node to iterate over
 
-    Yields:
-        (aos_index, node) for all filled nodes.
-    """
-    for child in node.iter_nonempty_():
-        yield (aos_index, child)
-        if isinstance(child, IDSStructArray):
-            for i in range(len(child)):
-                yield from nc_tree_iter(child[i], aos_index + (i,))
-        elif isinstance(child, IDSStructure):
-            yield from nc_tree_iter(child, aos_index)
 
-
-class IDS2NC:
+class IDS2NC(IDSTensorizer):
     """Class responsible for storing an IDS to a NetCDF file."""
 
     def __init__(self, ids: IDSToplevel, group: netCDF4.Group) -> None:
@@ -63,112 +28,18 @@ class IDS2NC:
             ids: IDSToplevel to store in the netCDF group
             group: Empty netCDF group to store the IDS in.
         """
-        self.ids = ids
-        """IDS to store."""
+        super().__init__(ids, [])  # pass empty list: tensorize full IDS
         self.group = group
         """NetCDF Group to store the IDS in."""
 
-        self.ncmeta = NCMetadata(ids.metadata)
-        """NetCDF related metadata."""
-        self.dimension_size = {}
-        """Map dimension name to its size."""
-        self.filled_data = {}
-        """Map of IDS paths to filled data nodes."""
-        self.filled_variables = set()
-        """Set of filled IDS variables"""
-        self.homogeneous_time = (
-            ids.ids_properties.homogeneous_time == IDS_TIME_MODE_HOMOGENEOUS
-        )
-        """True iff the IDS time mode is homogeneous."""
-        self.shapes = {}
-        """Map of IDS paths to data shape arrays."""
-
     def run(self) -> None:
         """Store the IDS in the NetCDF group."""
         self.collect_filled_data()
         self.determine_data_shapes()
         self.create_dimensions()
         self.create_variables()
-        # Synchronize variables to disk
-        # This is not strictly required (automatically done by netCDF4 when needed), but
-        # by separating it we get more meaningful profiling statistics
-        self.group.sync()
         self.store_data()
 
-    def collect_filled_data(self) -> None:
-        """Collect all filled data in the IDS and determine dimension sizes.
-
-        Results are stored in :attr:`filled_data` and :attr:`dimension_size`.
-        """
-        # Initialize dictionary with all paths that could exist in this IDS
-        filled_data = {path: {} for path in self.ncmeta.paths}
-        dimension_size = {}
-        get_dimensions = self.ncmeta.get_dimensions
-
-        for aos_index, node in nc_tree_iter(self.ids):
-            path = node.metadata.path_string
-            filled_data[path][aos_index] = node
-            ndim = node.metadata.ndim
-            if not ndim:
-                continue
-            dimensions = get_dimensions(path, self.homogeneous_time)
-            # We're only interested in the non-tensorized dimensions: [-ndim:]
-            for dim_name, size in zip(dimensions[-ndim:], node.shape):
-                dimension_size[dim_name] = max(dimension_size.get(dim_name, 0), size)
-
-        # Remove paths without data
-        self.filled_data = {path: data for path, data in filled_data.items() if data}
-        self.filled_variables = {path.replace("/", ".") for path in self.filled_data}
-        # Store dimension sizes
-        self.dimension_size = dimension_size
-
-    def determine_data_shapes(self) -> None:
-        """Determine tensorized data shapes and sparsity, save in :attr:`shapes`."""
-        get_dimensions = self.ncmeta.get_dimensions
-
-        for path, nodes_dict in self.filled_data.items():
-            metadata = self.ids.metadata[path]
-            # Structures don't have a size
-            if metadata.data_type is IDSDataType.STRUCTURE:
-                continue
-            ndim = metadata.ndim
-            dimensions = get_dimensions(path, self.homogeneous_time)
-
-            # node shape if it is completely filled
-            full_shape = tuple(self.dimension_size[dim] for dim in dimensions[-ndim:])
-
-            if len(dimensions) == ndim:
-                # Data at this path is not tensorized
-                node = nodes_dict[()]
-                sparse = node.shape != full_shape
-                if sparse:
-                    shapes = numpy.array(node.shape, dtype=SHAPE_DTYPE)
-
-            else:
-                # Data is tensorized, determine if it is homogeneously shaped
-                aos_dims = get_dimensions(self.ncmeta.aos[path], self.homogeneous_time)
-                shapes_shape = [self.dimension_size[dim] for dim in aos_dims]
-                if ndim:
-                    shapes_shape.append(ndim)
-                shapes = numpy.zeros(shapes_shape, dtype=SHAPE_DTYPE)
-
-                if ndim:  # ND types have a shape
-                    for aos_coords, node in nodes_dict.items():
-                        shapes[aos_coords] = node.shape
-                    sparse = not numpy.array_equiv(shapes, full_shape)
-
-                else:  # 0D types don't have a shape
-                    for aos_coords in nodes_dict.keys():
-                        shapes[aos_coords] = 1
-                    sparse = not shapes.all()
-                    shapes = None
-
-            if sparse:
-                self.shapes[path] = shapes
-                if ndim:
-                    # Ensure there is a pseudo-dimension f"{ndim}D" for shapes variable
-                    self.dimension_size[f"{ndim}D"] = ndim
-
     def create_dimensions(self) -> None:
         """Create netCDF dimensions."""
         for dimension, size in self.dimension_size.items():
@@ -253,14 +124,6 @@ class IDS2NC:
                         "shape is unset (i.e. filled with _Fillvalue)."
                     )
 
-    def filter_coordinates(self, path: str) -> str:
-        """Filter the coordinates list from NCMetadata to filled variables only."""
-        return " ".join(
-            coordinate
-            for coordinate in self.ncmeta.get_coordinates(path, self.homogeneous_time)
-            if coordinate in self.filled_variables
-        )
-
     def store_data(self) -> None:
         """Store data in the netCDF variables"""
         for path, nodes_dict in self.filled_data.items():
@@ -298,22 +161,4 @@ class IDS2NC:
 
             else:
                 # Data is tensorized: tensorize in-memory
-                # TODO: depending on the data, tmp_var may be HUGE, we may need a more
-                # efficient assignment algorithm for large and/or irregular data
-                tmp_var = numpy.full(var.shape, default_fillvals[metadata.data_type])
-                if metadata.data_type is IDSDataType.STR:
-                    tmp_var = numpy.asarray(tmp_var, dtype=object)
-
-                # Fill tmp_var
-                if shapes is None:
-                    # Data is not sparse, so we can assign to the aos_coords
-                    for aos_coords, node in nodes_dict.items():
-                        tmp_var[aos_coords] = node.value
-                else:
-                    # Data is sparse, so we must select a slice
-                    for aos_coords, node in nodes_dict.items():
-                        tmp_var[aos_coords + tuple(map(slice, node.shape))] = node.value
-
-                # Assign data to variable
-                var[()] = tmp_var
-                del tmp_var
+                var[()] = self.tensorize(path, default_fillvals[metadata.data_type])