Update notebooks

.gitignore

@@ -136,3 +136,10 @@ dmypy.json
 
 # Pyre type checker
 .pyre/
+
+# data
+*.h5
+*.hdf5
+notebooks/1
+notebooks/2
+notebooks/3

.pre-commit-config.yaml

@@ -25,14 +25,14 @@ repos:
     hooks:
       - id: docformatter
   - repo: https://github.com/charliermarsh/ruff-pre-commit
-    rev: v0.1.6
+    rev: v0.1.13
     hooks:
       - id: ruff
         args: [--fix]
       - id: ruff-format
         types_or: [python, pyi, jupyter]
   - repo: https://github.com/pre-commit/mirrors-mypy
-    rev: v1.7.1
+    rev: v1.8.0
     hooks:
       - id: mypy
         additional_dependencies:

notebooks/prepare.py

@@ -0,0 +1,37 @@
+# Artificially modify data for the notebooks
+from __future__ import annotations
+
+from imas2xarray import Variable, to_imas, to_xarray
+
+variables = (
+    Variable(
+        name='ion_temperature',
+        ids='core_profiles',
+        path='profiles_1d/*/ion/*/temperature',
+        dims=['time', 'ion', '$rho_tor_norm'],
+    ),
+    't_i_ave',
+)
+
+ids = 'core_profiles'
+path = '.'
+
+for subdir, k in (
+    ('2', 1.2),
+    ('3', 1.4),
+):
+    dataset = to_xarray(
+        f'{path}/1/data',
+        ids=ids,
+        variables=variables,
+    )
+    print(subdir, k)
+    dataset['t_i_ave'] *= k
+    dataset['ion_temperature'] *= k
+
+    to_imas(
+        f'{path}/{subdir}/data',
+        dataset=dataset,
+        ids=ids,
+        variables=variables,
+    )

notebooks/prepare.sh

@@ -0,0 +1,8 @@
+# Small script to initialize a few datasets for the notebooks using the `hdf5_testdata`:
+# https://github.com/duqtools/hdf5_testdata/tree/main
+
+git clone ~/python/hdf5_testdata 1
+git clone ~/python/hdf5_testdata 2
+git clone ~/python/hdf5_testdata 3
+
+python prepare.py

src/imas2xarray/_io.py

@@ -6,7 +6,7 @@
 
 from contextlib import contextmanager
 from pathlib import Path
-from typing import TYPE_CHECKING, Collection
+from typing import TYPE_CHECKING, Iterable
 
 import h5py
 import numpy as np
@@ -64,7 +64,7 @@ def _var_path_to_hdf5_key_and_slices(path: str) -> tuple[str, tuple[slice | int,
 
 def _mapping_to_xarray(
     data_file: h5py.File,
-    variables: Collection[str | IDSVariableModel],
+    variables: Iterable[str | IDSVariableModel],
     missing_ok: bool = False,
     empty_ok: bool = False,
 ) -> xr.Dataset:
@@ -74,7 +74,7 @@ def _mapping_to_xarray(
     ----------
     data_file : h5py.File
         Open hdf5 file
-    variables : Collection[str | IDSVariableModel]]
+    variables : Collection[(str | IDSVariableModel)]]
         List of data variables
     missing_ok : bool
         Ignore missing variables from dataset
@@ -116,7 +116,7 @@ def _mapping_to_xarray(
 
 
 def to_xarray(
-    path: str | Path, *, ids: str, variables: None | Collection[str] = None
+    path: str | Path, *, ids: str, variables: None | Iterable[str | IDSVariableModel] = None
 ) -> xr.Dataset:
     """Load IDS from given path to IMAS data into an xarray dataset.
 
@@ -128,7 +128,7 @@ def to_xarray(
         Path to the data
     ids : str
         The IDS to load (i.e. 'core_profiles')
-    variables : None | list[str], optional
+    variables : None | Iterable[str | Variable], optional
         List of variables to load. If None, attempt to load
         all variables known to `imas2xarray`
 
@@ -146,7 +146,11 @@ def to_xarray(
 
 
 def to_imas(
-    path: str | Path, dataset: xr.Dataset, *, ids: str, variables: None | Collection[str] = None
+    path: str | Path,
+    dataset: xr.Dataset,
+    *,
+    ids: str,
+    variables: None | Iterable[str | IDSVariableModel] = None,
 ):
     """Write variables in xarray dataset back to IMAS data at given path.
 
@@ -160,7 +164,7 @@ def to_imas(
         Input dataset
     ids : str
         The IDS to write to (i.e. 'core_profiles')
-    variables : Collection[str]
+    variables : Iterable[str | Variable]
         List of variables to write back. If None, attempt to write back
         all variables known to `imas2xarray`
     """
@@ -196,7 +200,7 @@ def get_all_variables(
         self,
         *,
         ids: str,
-        extra_variables: None | Collection[IDSVariableModel] = None,
+        extra_variables: None | Iterable[IDSVariableModel] = None,
         squash: bool = True,
         **kwargs,
     ) -> xr.Dataset:
@@ -209,7 +213,7 @@ def get_all_variables(
         ----------
         ids : str
             The IDS to write to (i.e. 'core_profiles')
-        extra_variables : Collection[IDSVariableModel]
+        extra_variables : Iterable[Variable]
             Extra variables to load in addition to the ones known through the config
         squash : bool
             Squash placeholder variables
@@ -229,7 +233,7 @@ def get_all_variables(
 
     def get_variables(
         self,
-        variables: Collection[str | IDSVariableModel],
+        variables: Iterable[str | IDSVariableModel],
         *,
         ids: str,
         squash: bool = True,
@@ -243,7 +247,7 @@ def get_variables(
 
         Parameters
         ----------
-        variables : Collection[Union[str, IDSVariableModel]]
+        variables : Iterable[str | Variable]
             Variable names of the data to load.
         ids : str
             The IDS to write to (i.e. 'core_profiles')
@@ -281,7 +285,11 @@ def get_variables(
         return ds
 
     def set_variables(
-        self, dataset: xr.Dataset, *, ids: str, variables: None | Collection[str] = None
+        self,
+        dataset: xr.Dataset,
+        *,
+        ids: str,
+        variables: None | Iterable[str | IDSVariableModel] = None,
     ):
         """Update variables in corresponding ids datafile.
 
@@ -292,7 +300,7 @@ def set_variables(
             target dataset.
         ids : str
             IDS to write to.
-        variables : Collection[str], optional
+        variables : Iterable[str | Variable], optional
             List of data variables to write.
         """
         if not variables:

src/imas2xarray/_lookup.py

@@ -6,7 +6,7 @@
 import sys
 from collections import UserDict
 from pathlib import Path, PosixPath
-from typing import Any, Collection, Hashable
+from typing import Any, Hashable, Iterable
 
 from pydantic_yaml import parse_yaml_raw_as
 
@@ -80,13 +80,13 @@ def groupby_ids(self) -> dict[Hashable, list[IDSVariableModel]]:
         return grouped_ids_vars
 
     def lookup(
-        self, variables: Collection[(str | IDSVariableModel)], skip_missing: bool = False
+        self, variables: Iterable[(str | IDSVariableModel)], skip_missing: bool = False
     ) -> set[IDSVariableModel]:
         """Helper function to look up a bunch of variables.
 
         Parameters
         ----------
-        variables : Collection[(str | IDSVariableModel)]
+        variables : Iterable[(str | IDSVariableModel)]
             List of variables to load. If str, look up the variable from the `var_lookup`.
             Else, ensure the variable is an `IDSVariableModel`.
         skip_missing : bool

src/imas2xarray/_rebase.py

@@ -10,7 +10,7 @@
 logger = logging.getLogger(__name__)
 
 
-def rezero_time(ds: xr.Dataset, *, start: int = 0) -> None:
+def rezero_time(ds: xr.Dataset, *, start: int = 0, key: str = 'time') -> None:
     """Standardize the time within a dataset by setting the first timestep to
     0.
 
@@ -21,10 +21,12 @@ def rezero_time(ds: xr.Dataset, *, start: int = 0) -> None:
     ----------
     ds : xr.Dataset
         Source dataset
+    key : str
+        Name of the time dimension
     start : int, optional
         Where to start the returned time series
     """
-    ds['time'] = ds['time'] - ds['time'][0] + start
+    ds[key] = ds[key] - ds[key][0] + start
 
 
 def squash_placeholders(ds: xr.Dataset) -> xr.Dataset: