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

imas/backends/netcdf/ids_tensorizer.py

@@ -0,0 +1,205 @@
+# This file is part of IMAS-Python.
+# 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
+
+from imas.backends.netcdf.iterators import indexed_tree_iter
+from imas.backends.netcdf.nc_metadata import NCMetadata
+from imas.ids_data_type import IDSDataType
+from imas.ids_defs import IDS_TIME_MODE_HOMOGENEOUS
+from imas.ids_toplevel import IDSToplevel
+
+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
+
+
+class IDSTensorizer:
+    """Common functionality for tensorizing IDSs. Used in IDS2NC and util.to_xarray."""
+
+    def __init__(self, ids: IDSToplevel, paths_to_tensorize: List[str]) -> None:
+        """Initialize IDSTensorizer.
+
+        Args:
+            ids: IDSToplevel to store in the netCDF group
+            paths_to_tensorize: Restrict tensorization to the provided paths. If an
+                empty list is provided, all filled quantities in the IDS will be
+                tensorized.
+        """
+        self.ids = ids
+        """IDS to tensorize."""
+        self.paths_to_tensorize = paths_to_tensorize
+        """List of paths to tensorize"""
+
+        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 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.
+
+        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
+
+        if self.paths_to_tensorize:
+            # Restrict tensorization to provided paths
+            iterator = (
+                item
+                for path in self.paths_to_tensorize
+                for item in indexed_tree_iter(self.ids, self.ids.metadata[path])
+                if item[1].has_value  # Skip nodes without value set
+            )
+        else:
+            # Tensorize all non-empty nodes
+            iterator = indexed_tree_iter(self.ids)
+
+        for aos_index, node in iterator:
+            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 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 tensorize(self, path, fillvalue):
+        """
+        Tensorizes the data at the given path with the specified fill value.
+
+        Args:
+            path: The path to the data in the IDS.
+            fillvalue: The value to fill the tensor with. Can be of any type,
+                             including strings.
+
+        Returns:
+            A tensor filled with the data from the specified path.
+        """
+        dimensions = self.ncmeta.get_dimensions(path, self.homogeneous_time)
+        shape = tuple(self.dimension_size[dim] for dim in dimensions)
+
+        # 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(shape, fillvalue)
+        if isinstance(fillvalue, str):
+            tmp_var = numpy.asarray(tmp_var, dtype=object)
+
+        shapes = self.shapes.get(path)
+        nodes_dict = self.filled_data[path]
+
+        # 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
+
+        return tmp_var

imas/backends/netcdf/iterators.py

@@ -0,0 +1,98 @@
+from typing import Iterator, List, Optional, Tuple
+
+from imas.ids_base import IDSBase
+from imas.ids_data_type import IDSDataType
+from imas.ids_metadata import IDSMetadata
+from imas.ids_struct_array import IDSStructArray
+from imas.ids_structure import IDSStructure
+from imas.ids_toplevel import IDSToplevel
+
+
+def _split_on_aos(metadata: IDSMetadata):
+    """Split paths per IDS."""
+    paths = []
+    curpath = metadata.name
+
+    item = metadata
+    while item._parent.data_type is not None:
+        item = item._parent
+        if item.data_type is IDSDataType.STRUCT_ARRAY:
+            paths.append(curpath)
+            curpath = item.name
+        else:
+            curpath = f"{item.name}/{curpath}"
+    paths.append(curpath)
+    return paths[::-1]
+
+
+IndexedNode = Tuple[Tuple[int, ...], IDSBase]
+
+
+def indexed_tree_iter(
+    ids: IDSToplevel, metadata: Optional[IDSMetadata] = None
+) -> Iterator[IndexedNode]:
+    """Tree iterator that tracks indices of all ancestor array of structures.
+
+    Args:
+        ids: IDS top level element to iterate over
+        metadata: Iterate over all nodes inside the IDS at the metadata object.
+            If ``None``, all filled items in the IDS are iterated over.
+
+    Yields:
+        (aos_indices, node) for all filled nodes.
+
+    Example:
+        >>> ids = imas.IDSFactory().new("core_profiles")
+        >>> ids.profiles_1d.resize(2)
+        >>> ids.profiles_1d[0].time = 1.0
+        >>> ids.profiles_1d[1].t_i_average = [1.0]
+        >>> list(indexed_tree_iter(ids))
+        [
+            ((), <IDSStructArray (IDS:core_profiles, profiles_1d with 2 items)>),
+            ((0,), <IDSFloat0D (IDS:core_profiles, profiles_1d[0]/time, FLT_0D)>),
+            ((1,), <IDSNumericArray (IDS:core_profiles, profiles_1d[1]/t_i_average, FLT_1D)>)
+        ]
+        >>> list(indexed_tree_iter(ids, ids.metadata["profiles_1d/time"]))
+        [
+            ((0,), <IDSFloat0D (IDS:core_profiles, profiles_1d[0]/time, FLT_0D)>),
+            ((1,), <IDSFloat0D (IDS:core_profiles, profiles_1d[1]/time, empty FLT_0D)>)
+        ]
+    """  # noqa: E501
+    if metadata is None:
+        # Iterate over all filled nodes in the IDS
+        yield from _full_tree_iter(ids, ())
+
+    else:
+        paths = _split_on_aos(metadata)
+        if len(paths) == 1:
+            yield (), ids[paths[0]]
+        else:
+            yield from _tree_iter(ids, paths, ())
+
+
+def _tree_iter(
+    structure: IDSStructure, paths: List[str], curindex: Tuple[int, ...]
+) -> Iterator[IndexedNode]:
+    aos_path, *paths = paths
+    aos = structure[aos_path]
+
+    if len(paths) == 1:
+        path = paths[0]
+        for i, node in enumerate(aos):
+            yield curindex + (i,), node[path]
+
+    else:
+        for i, node in enumerate(aos):
+            yield from _tree_iter(node, paths, curindex + (i,))
+
+
+def _full_tree_iter(
+    node: IDSStructure, cur_index: Tuple[int, ...]
+) -> Iterator[IndexedNode]:
+    for child in node.iter_nonempty_():
+        yield (cur_index, child)
+        if isinstance(child, IDSStructArray):
+            for i in range(len(child)):
+                yield from _full_tree_iter(child[i], cur_index + (i,))
+        elif isinstance(child, IDSStructure):
+            yield from _full_tree_iter(child, cur_index)

imas/backends/netcdf/nc2ids.py

@@ -1,17 +1,18 @@
 import logging
 import os
-from typing import Iterator, List, Optional, Tuple
+from typing import Optional
 
 import netCDF4
+import numpy as np
 
 from imas.backends.netcdf import ids2nc
 from imas.backends.netcdf.nc_metadata import NCMetadata
+from imas.backends.netcdf.iterators import indexed_tree_iter
 from imas.exception import InvalidNetCDFEntry
-from imas.ids_base import IDSBase
+from imas.ids_convert import NBCPathMap
 from imas.ids_data_type import IDSDataType
 from imas.ids_defs import IDS_TIME_MODE_HOMOGENEOUS
 from imas.ids_metadata import IDSMetadata
-from imas.ids_structure import IDSStructure
 from imas.ids_toplevel import IDSToplevel
 
 logger = logging.getLogger(__name__)
@@ -24,68 +25,40 @@ def variable_error(var, issue, value, expected=None) -> InvalidNetCDFEntry:
     )
 
 
-def split_on_aos(metadata: IDSMetadata):
-    paths = []
-    curpath = metadata.name
-
-    item = metadata
-    while item._parent.data_type is not None:
-        item = item._parent
-        if item.data_type is IDSDataType.STRUCT_ARRAY:
-            paths.append(curpath)
-            curpath = item.name
-        else:
-            curpath = f"{item.name}/{curpath}"
-    paths.append(curpath)
-    return paths[::-1]
-
-
-IndexedNode = Tuple[Tuple[int, ...], IDSBase]
-
-
-def tree_iter(structure: IDSStructure, metadata: IDSMetadata) -> Iterator[IndexedNode]:
-    paths = split_on_aos(metadata)
-    if len(paths) == 1:
-        yield (), structure[paths[0]]
-    else:
-        yield from _tree_iter(structure, paths, ())
-
-
-def _tree_iter(
-    structure: IDSStructure, paths: List[str], curindex: Tuple[int, ...]
-) -> Iterator[IndexedNode]:
-    aos_path, *paths = paths
-    aos = structure[aos_path]
-
-    if len(paths) == 1:
-        path = paths[0]
-        for i, node in enumerate(aos):
-            yield curindex + (i,), node[path]
-
-    else:
-        for i, node in enumerate(aos):
-            yield from _tree_iter(node, paths, curindex + (i,))
-
-
 class NC2IDS:
     """Class responsible for reading an IDS from a NetCDF group."""
 
-    def __init__(self, group: netCDF4.Group, ids: IDSToplevel) -> None:
+    def __init__(
+        self,
+        group: netCDF4.Group,
+        ids: IDSToplevel,
+        ids_metadata: IDSMetadata,
+        nbc_map: Optional[NBCPathMap],
+    ) -> None:
         """Initialize NC2IDS converter.
 
         Args:
             group: NetCDF group that stores the IDS data.
             ids: Corresponding IDS toplevel to store the data in.
+            ids_metadata: Metadata corresponding to the DD version that the data is
+                stored in.
+            nbc_map: Path map for implicit DD conversions.
         """
         self.group = group
         """NetCDF Group that the IDS is stored in."""
         self.ids = ids
         """IDS to store the data in."""
+        self.ids_metadata = ids_metadata
+        """Metadata of the IDS in the DD version that the data is stored in"""
+        self.nbc_map = nbc_map
+        """Path map for implicit DD conversions."""
 
-        self.ncmeta = NCMetadata(ids.metadata)
+        self.ncmeta = NCMetadata(ids_metadata)
         """NetCDF related metadata."""
         self.variables = list(group.variables)
         """List of variable names stored in the netCDF group."""
+
+        self._lazy_map = {}
         # Don't use masked arrays: they're slow and we'll handle most of the unset
         # values through the `:shape` arrays
         self.group.set_auto_mask(False)
@@ -99,7 +72,7 @@ class NC2IDS:
                 "Mandatory variable `ids_properties.homogeneous_time` does not exist."
             )
         var = group["ids_properties.homogeneous_time"]
-        self._validate_variable(var, ids.ids_properties.homogeneous_time.metadata)
+        self._validate_variable(var, ids.metadata["ids_properties/homogeneous_time"])
         if var[()] not in [0, 1, 2]:
             raise InvalidNetCDFEntry(
                 f"Invalid value for ids_properties.homogeneous_time: {var[()]}. "
@@ -107,23 +80,52 @@ class NC2IDS:
             )
         self.homogeneous_time = var[()] == IDS_TIME_MODE_HOMOGENEOUS
 
-    def run(self) -> None:
+    def run(self, lazy: bool) -> None:
         """Load the data from the netCDF group into the IDS."""
         self.variables.sort()
         self.validate_variables()
+        if lazy:
+            self.ids._set_lazy_context(LazyContext(self))
         for var_name in self.variables:
             if var_name.endswith(":shape"):
                 continue
-            metadata = self.ids.metadata[var_name]
+            metadata = self.ids_metadata[var_name]
 
             if metadata.data_type is IDSDataType.STRUCTURE:
                 continue  # This only contains DD metadata we already know
 
+            # Handle implicit DD version conversion
+            if self.nbc_map is None:
+                target_metadata = metadata  # no conversion
+            elif metadata.path_string in self.nbc_map:
+                new_path = self.nbc_map.path[metadata.path_string]
+                if new_path is None:
+                    logging.info(
+                        "Not loading data for %s: no equivalent data structure exists "
+                        "in the target Data Dictionary version.",
+                        metadata.path_string,
+                    )
+                    continue
+                target_metadata = self.ids.metadata[new_path]
+            elif metadata.path_string in self.nbc_map.type_change:
+                logging.info(
+                    "Not loading data for %s: cannot hanlde type changes when "
+                    "implicitly converting data to the target Data Dictionary version.",
+                    metadata.path_string,
+                )
+                continue
+            else:
+                target_metadata = metadata  # no conversion required
+
             var = self.group[var_name]
+            if lazy:
+                self._lazy_map[target_metadata.path_string] = var
+                continue
+
             if metadata.data_type is IDSDataType.STRUCT_ARRAY:
                 if "sparse" in var.ncattrs():
                     shapes = self.group[var_name + ":shape"][()]
-                    for index, node in tree_iter(self.ids, metadata):
+                    for index, node in indexed_tree_iter(self.ids, target_metadata):
                         node.resize(shapes[index][0])
 
                 else:
@@ -132,7 +134,7 @@ class NC2IDS:
                         metadata.path_string, self.homogeneous_time
                     )[-1]
                     size = self.group.dimensions[dim].size
-                    for _, node in tree_iter(self.ids, metadata):
+                    for _, node in indexed_tree_iter(self.ids, target_metadata):
                         node.resize(size)
 
                 continue
@@ -144,23 +146,30 @@ class NC2IDS:
             if "sparse" in var.ncattrs():
                 if metadata.ndim:
                     shapes = self.group[var_name + ":shape"][()]
-                    for index, node in tree_iter(self.ids, metadata):
+                    for index, node in indexed_tree_iter(self.ids, target_metadata):
                         shape = shapes[index]
                         if shape.all():
-                            node.value = data[index + tuple(map(slice, shapes[index]))]
+                            # NOTE: bypassing IDSPrimitive.value.setter logic
+                            node._IDSPrimitive__value = data[
+                                index + tuple(map(slice, shape))
+                            ]
                 else:
-                    for index, node in tree_iter(self.ids, metadata):
+                    for index, node in indexed_tree_iter(self.ids, target_metadata):
                         value = data[index]
                         if value != getattr(var, "_FillValue", None):
-                            node.value = data[index]
+                            # NOTE: bypassing IDSPrimitive.value.setter logic
+                            node._IDSPrimitive__value = value
 
             elif metadata.path_string not in self.ncmeta.aos:
                 # Shortcut for assigning untensorized data
-                self.ids[metadata.path] = data
+                # Note: var[()] can return 0D numpy arrays. Instead of handling this
+                # here, we'll let IDSPrimitive.value.setter take care of it:
+                self.ids[target_metadata.path].value = data
 
             else:
-                for index, node in tree_iter(self.ids, metadata):
-                    node.value = data[index]
+                for index, node in indexed_tree_iter(self.ids, target_metadata):
+                    # NOTE: bypassing IDSPrimitive.value.setter logic
+                    node._IDSPrimitive__value = data[index]
 
     def validate_variables(self) -> None:
         """Validate that all variables in the netCDF Group exist and match the DD."""
@@ -194,7 +203,7 @@ class NC2IDS:
             # Check that the DD defines this variable, and validate its metadata
             var = self.group[var_name]
             try:
-                metadata = self.ids.metadata[var_name]
+                metadata = self.ids_metadata[var_name]
             except KeyError:
                 raise InvalidNetCDFEntry(
                     f"Invalid variable {var_name}: no such variable exists in the "
@@ -300,3 +309,109 @@ class NC2IDS:
             raise variable_error(
                 shape_var, "dtype", shape_var.dtype, "any integer type"
             )
+
+
+class LazyContext:
+    def __init__(self, nc2ids, index=()):
+        self.nc2ids = nc2ids
+        self.index = index
+
+    def get_child(self, child):
+        """
+        Retrieves and sets the appropriate context or value for a given
+        child node based on its metadata.
+
+        Args:
+            child: The child IDS node which should be lazy loaded.
+
+        """
+        metadata = child.metadata
+        path = metadata.path_string
+        data_type = metadata.data_type
+        nc2ids = self.nc2ids
+        var = nc2ids._lazy_map.get(path)
+
+        if data_type is IDSDataType.STRUCT_ARRAY:
+            # Determine size of the aos
+            if var is None:
+                size = 0
+            elif "sparse" in var.ncattrs():
+                size = nc2ids.group[var.name + ":shape"][self.index][0]
+            else:
+                # FIXME: extract dimension name from nc file?
+                dim = nc2ids.ncmeta.get_dimensions(
+                    metadata.path_string, nc2ids.homogeneous_time
+                )[-1]
+                size = nc2ids.group.dimensions[dim].size
+
+            child._set_lazy_context(LazyArrayStructContext(nc2ids, self.index, size))
+
+        elif data_type is IDSDataType.STRUCTURE:
+            child._set_lazy_context(self)
+
+        elif var is not None:  # Data elements
+            value = None
+            if "sparse" in var.ncattrs():
+                if metadata.ndim:
+                    shape_var = nc2ids.group[var.name + ":shape"]
+                    shape = shape_var[self.index]
+                    if shape.all():
+                        value = var[self.index + tuple(map(slice, shape))]
+                else:
+                    value = var[self.index]
+                    if value == getattr(var, "_FillValue", None):
+                        value = None  # Skip setting
+            else:
+                value = var[self.index]
+
+            if value is not None:
+                if isinstance(value, np.ndarray):
+                    # Convert the numpy array to a read-only view
+                    value = value.view()
+                    value.flags.writeable = False
+                # NOTE: bypassing IDSPrimitive.value.setter logic
+                child._IDSPrimitive__value = value
+
+
+class LazyArrayStructContext(LazyContext):
+    """
+    LazyArrayStructContext is a subclass of LazyContext that provides a context for
+    handling structured arrays in a lazy manner. It is initialized with a NetCDF to
+    IDS mapping object, an index, and a size.
+    """
+
+    def __init__(self, nc2ids, index, size):
+        """
+        Initialize the instance with nc2ids, index, and size.
+
+        Args:
+            nc2ids: The NetCDF to IDS mapping object.
+            index: The index within the NetCDF file.
+            size: The size of the data to be processed.
+        """
+        super().__init__(nc2ids, index)
+        self.size = size
+
+    def get_context(self):
+        """
+        Returns the current context.
+
+        This method returns the current instance of the class, which is expected
+        to have a 'size' attribute as required by IDSStructArray.
+
+        Returns:
+            The current instance of the class.
+        """
+        return self  # IDSStructArray expects to get something with a size attribute
+
+    def iterate_to_index(self, index: int) -> LazyContext:
+        """
+        Iterates to a specified index and returns a LazyContext object.
+
+        Args:
+            index (int): The index to iterate to.
+
+        Returns:
+            LazyContext: A LazyContext object initialized with the updated index.
+        """
+        return LazyContext(self.nc2ids, self.index + (index,))

imas/backends/netcdf/nc_validate.py

@@ -47,8 +47,9 @@ def validate_netcdf_file(filename: str) -> None:
         for ids_name in ids_names:
             for occurrence in entry.list_all_occurrences(ids_name):
                 group = dataset[f"{ids_name}/{occurrence}"]
+                ids = factory.new(ids_name)
                 try:
-                    NC2IDS(group, factory.new(ids_name)).validate_variables()
+                    NC2IDS(group, ids, ids.metadata, None).validate_variables()
                 except InvalidNetCDFEntry as exc:
                     occ = f":{occurrence}" if occurrence else ""
                     raise InvalidNetCDFEntry(f"Invalid IDS {ids_name}{occ}: {exc}")

imas/ids_convert.py

@@ -533,7 +533,8 @@ def _add_provenance_entry(
         # DD version after IMAS-5304
         node.reference.resize(len(node.reference) + 1, keep=True)
         node.reference[-1].name = source_txt
-        timestamp = datetime.datetime.now(datetime.UTC).isoformat(timespec="seconds")
+        utc = getattr(datetime, "UTC", datetime.timezone.utc)
+        timestamp = datetime.datetime.now(utc).isoformat(timespec="seconds")
         node.reference[-1].timestamp = timestamp.replace("+00:00", "Z")
     else:
         # DD before IMAS-5304 (between 3.34.0 and 3.41.0)