sdf-xarray 0.3.2__cp312-cp312-win_amd64.whl → 0.5.0__cp312-cp312-win_amd64.whl

sdf_xarray/__init__.py

@@ -16,12 +16,14 @@ from xarray.backends import AbstractDataStore, BackendArray, BackendEntrypoint
 from xarray.backends.file_manager import CachingFileManager
 from xarray.backends.locks import ensure_lock
 from xarray.core import indexing
+from xarray.core.types import T_Chunks
 from xarray.core.utils import close_on_error, try_read_magic_number_from_path
 from xarray.core.variable import Variable
 
 # NOTE: Do not delete these lines, otherwise the "epoch" dataset and dataarray
 # accessors will not be imported when the user imports sdf_xarray
 import sdf_xarray.dataset_accessor
+import sdf_xarray.download
 import sdf_xarray.plotting  # noqa: F401
 
 # NOTE: This attempts to initialise with the "pint" accessor if the user
@@ -84,6 +86,64 @@ def _resolve_glob(path_glob: PathLike | Iterable[PathLike]):
     return paths
 
 
+def _build_datatree_from_dataset(
+    ds: xr.Dataset,
+) -> xr.DataTree:
+    """
+    An `xarray.DataTree` is constructed utilising the original names in the SDF
+    file. This is due to the fact that these names include slashes which `xarray`
+    can use to automatically build up a datatree. We do additionally replace
+    spaces with underscores to be more pythonic. You can find the
+    `xarray.Dataset` name under the ``attrs["flat_structure_name"]`` for referencing.
+
+    In some cases the user may output the ``always + species`` dumpmask which
+    means that SDF variable will have species data plus a general one. When
+    defining a `xarray.DataTree` you cannot have a node of that tree contain both
+    variable information and have leaves with variables so we move the node
+    information to a leaf named ``node/All`` (see example of
+    ``Dervied/Number_Density/All`` in below table)
+
+    Below are some examples of how variable names are translated from the
+    regular `xarray.open_dataset` result into their more traditional names.
+
+    =================================== ===================================
+    Dataset variable name               DataTree variable name
+    =================================== ===================================
+    ``Derived_Number_Density``          ``Derived/Number_Density/All``
+    ``Derived_Number_Density_Electron`` ``Derived/Number_Density/Electron``
+    ``Derived_Number_Density_Ion``      ``Derived/Number_Density/Ion``
+    ``Derived_Number_Density_Photon``   ``Derived/Number_Density/Photon``
+    ``Derived_Average_Particle_Energy`` ``Derived/Average_Particle_Energy``
+    =================================== ===================================
+
+    Parameters
+    ----------
+    ds
+        Incoming `xarray.Dataset` to convert to a `xarray.DataTree`
+    """
+    renames = {}
+    for name, var in ds.data_vars.items():
+        # Append the current variable name to the attributes
+        var.attrs["flat_structure_name"] = name
+        renames.update({name: var.attrs["full_name"].replace(" ", "_")})
+
+    new_names = renames.values()
+
+    final_renames = {
+        key: (
+            f"{path}/All"
+            if any(other.startswith(f"{path}/") for other in new_names)
+            else path
+        )
+        for key, path in renames.items()
+    }
+
+    ds = ds.rename_vars(final_renames)
+    dt = xr.DataTree.from_dict(ds)
+    dt.attrs = ds.attrs
+    return dt
+
+
 def purge_unselected_data_vars(ds: xr.Dataset, data_vars: list[str]) -> xr.Dataset:
     """
     If the user has exclusively requested only certain variables be
@@ -142,6 +202,7 @@ def open_mfdataset(
     keep_particles: bool = False,
     probe_names: list[str] | None = None,
     data_vars: list[str] | None = None,
+    chunks: T_Chunks = "auto",
 ) -> xr.Dataset:
     """Open a set of EPOCH SDF files as one `xarray.Dataset`
 
@@ -175,6 +236,14 @@ def open_mfdataset(
         List of EPOCH probe names
     data_vars :
         List of data vars to load in (If not specified loads in all variables)
+    chunks :
+        Dictionary with keys given by dimension names and values given by chunk sizes.
+        In general, these should divide the dimensions of each dataset. By default
+        chunks are automatically set so that they are the same size as the dimensions
+        stored in each of the SDF files. See `Xarray chunking-and-performance
+        <https://docs.xarray.dev/en/stable/user-guide/dask.html#chunking-and-performance>`_
+        for details on why this is useful for large datasets. The default behaviour is
+        to do this automatically and can be disabled by ``chunks=None``.
     """
 
     path_glob = _resolve_glob(path_glob)
@@ -185,13 +254,16 @@ def open_mfdataset(
             data_vars=data_vars,
             keep_particles=keep_particles,
             probe_names=probe_names,
+            chunks=chunks,
         )
 
     _, var_times_map = make_time_dims(path_glob)
 
     all_dfs = []
     for f in path_glob:
-        ds = xr.open_dataset(f, keep_particles=keep_particles, probe_names=probe_names)
+        ds = xr.open_dataset(
+            f, keep_particles=keep_particles, probe_names=probe_names, chunks=chunks
+        )
 
         # If the data_vars are specified then only load them in and disregard the rest.
         # If there are no remaining data variables then skip adding the dataset to list
@@ -225,6 +297,147 @@ def open_mfdataset(
     )
 
 
+def open_datatree(
+    path: PathLike,
+    *,
+    keep_particles: bool = False,
+    probe_names: list[str] | None = None,
+) -> xr.DataTree:
+    """
+    An `xarray.DataTree` is constructed utilising the original names in the SDF
+    file. This is due to the fact that these names include slashes which `xarray`
+    can use to automatically build up a datatree. We do additionally replace
+    spaces with underscores to be more pythonic. You can find the
+    `xarray.Dataset` name under the ``attrs["flat_structure_name"]`` for referencing.
+
+    In some cases the user may output the ``always + species`` dumpmask which
+    means that SDF variable will have species data plus a general one. When
+    defining a `xarray.DataTree` you cannot have a node of that tree contain both
+    variable information and have leaves with variables so we move the node
+    information to a leaf named ``node/All`` (see example of
+    ``Dervied/Number_Density/All`` in below table)
+
+    Below are some examples of how variable names are translated from the
+    regular `xarray.open_dataset` result into their more traditional names.
+
+    =================================== ===================================
+    Dataset variable name               DataTree variable name
+    =================================== ===================================
+    ``Derived_Number_Density``          ``Derived/Number_Density/All``
+    ``Derived_Number_Density_Electron`` ``Derived/Number_Density/Electron``
+    ``Derived_Number_Density_Ion``      ``Derived/Number_Density/Ion``
+    ``Derived_Number_Density_Photon``   ``Derived/Number_Density/Photon``
+    ``Derived_Average_Particle_Energy`` ``Derived/Average_Particle_Energy``
+    =================================== ===================================
+
+    Parameters
+    ----------
+    path
+        The path to the SDF file
+    keep_particles
+        If ``True``, also load particle data (this may use a lot of memory!)
+    probe_names
+        List of EPOCH probe names
+
+    Examples
+    --------
+    >>> dt = open_datatree("0000.sdf")
+    >>> dt["Electric_Field"]["Ex"].values  # Access all Electric_Field_Ex data
+    """
+
+    return xr.open_datatree(
+        path, keep_particles=keep_particles, probe_names=probe_names
+    )
+
+
+def open_mfdatatree(
+    path_glob: Iterable | str | Path | Callable[..., Iterable[Path]],
+    *,
+    separate_times: bool = False,
+    keep_particles: bool = False,
+    probe_names: list[str] | None = None,
+    data_vars: list[str] | None = None,
+) -> xr.DataTree:
+    """Open a set of EPOCH SDF files as one `xarray.DataTree`
+
+    EPOCH can output variables at different periods, so each individal
+    SDF file from one EPOCH run may have different variables in it. In
+    order to combine all files into one `xarray.Dataset`, we need to
+    concatenate variables across their time dimension.
+
+    We have two choices:
+
+    1. One time dimension where some variables may not be defined at all time
+       points, and so will be filled with NaNs at missing points; or
+    2. Multiple time dimensions, one for each output frequency
+
+    The second option is better for memory consumption, as the missing data with
+    the first option still takes up space. However, proper lazy-loading may
+    mitigate this.
+
+    The ``separate_times`` argument can be used to switch between these choices.
+
+    An `xarray.DataTree` is constructed utilising the original names in the SDF
+    file. This is due to the fact that these names include slashes which `xarray`
+    can use to automatically build up a datatree. We do additionally replace
+    spaces with underscores to be more pythonic. You can find the
+    `xarray.Dataset` name under the ``attrs["flat_structure_name"]`` for referencing.
+
+    This function combines multiple SDF files into a single `xarray.DataTree` with a
+    unified time dimension and hierarchical organization of variables.
+
+    In some cases the user may output the ``always + species`` dumpmask which
+    means that SDF variable will have species data plus a general one. When
+    defining a `xarray.DataTree` you cannot have a node of that tree contain both
+    variable information and have leaves with variables so we move the node
+    information to a leaf named ``node/All`` (see example of
+    ``Dervied/Number_Density/All`` in below table)
+
+    Below are some examples of how variable names are translated from the
+    regular `xarray.open_dataset` result into their more traditional names.
+
+    =================================== ===================================
+    Dataset variable name               DataTree variable name
+    =================================== ===================================
+    ``Derived_Number_Density``          ``Derived/Number_Density/All``
+    ``Derived_Number_Density_Electron`` ``Derived/Number_Density/Electron``
+    ``Derived_Number_Density_Ion``      ``Derived/Number_Density/Ion``
+    ``Derived_Number_Density_Photon``   ``Derived/Number_Density/Photon``
+    ``Derived_Average_Particle_Energy`` ``Derived/Average_Particle_Energy``
+    =================================== ===================================
+
+    Parameters
+    ----------
+    path_glob
+        List of filenames or string glob pattern
+    separate_times
+        If ``True``, create separate time dimensions for variables defined at
+        different output frequencies
+    keep_particles
+        If ``True``, also load particle data (this may use a lot of memory!)
+    probe_names
+        List of EPOCH probe names
+    data_vars
+        List of data vars to load in (If not specified loads in all variables)
+
+    Examples
+    --------
+    >>> dt = open_mfdatatree("*.sdf")
+    >>> dt["Electric_Field"]["Ex"].values  # Access all Electric_Field_Ex data
+    >>> dt.coords["time"].values  # Access combined time dimension
+    """
+    # First, combine the datasets as usual
+    combined_ds = open_mfdataset(
+        path_glob,
+        separate_times=separate_times,
+        keep_particles=keep_particles,
+        probe_names=probe_names,
+        data_vars=data_vars,
+    )
+
+    return _build_datatree_from_dataset(combined_ds)
+
+
 def make_time_dims(path_glob):
     """Extract the distinct set of time arrays from a collection of
     SDF files, along with a mapping from variable names to their time
@@ -268,13 +481,11 @@ class SDFBackendArray(BackendArray):
 
     __slots__ = ("datastore", "dtype", "shape", "variable_name")
 
-    def __init__(self, variable_name, datastore):
+    def __init__(self, variable_name, datastore, shape, dtype):
         self.datastore = datastore
         self.variable_name = variable_name
-
-        array = self.get_array()
-        self.shape = array.shape
-        self.dtype = array.dtype
+        self.shape = shape
+        self.dtype = dtype
 
     def get_array(self, needs_lock=True):
         with self.datastore.acquire_context(needs_lock) as ds:
@@ -446,7 +657,12 @@ class SDFDataStore(AbstractDataStore):
                 if value.units is not None:
                     data_attrs["units"] = value.units
 
-                data_vars[base_name] = Variable(dims, value.data, attrs=data_attrs)
+                var = Variable(dims, value.data, attrs=data_attrs)
+
+                # Provide preferred_chunks for constants so dask aligns to natural shapes
+                var.encoding["preferred_chunks"] = dict(zip(dims, shape))
+
+                data_vars[base_name] = var
                 continue
 
             if value.is_point_data:
@@ -495,9 +711,9 @@ class SDFDataStore(AbstractDataStore):
                 grid_mid = self.ds.grids[value.grid_mid]
                 grid_mid_base_name = _process_grid_name(grid_mid.name, _norm_grid_name)
                 for dim_size, dim_name in zip(grid_mid.shape, grid_mid.labels):
-                    dim_size_lookup[dim_name][
-                        dim_size
-                    ] = f"{dim_name}_{grid_mid_base_name}"
+                    dim_size_lookup[dim_name][dim_size] = (
+                        f"{dim_name}_{grid_mid_base_name}"
+                    )
 
                 var_coords = [
                     dim_size_lookup[dim_name][dim_size]
@@ -513,8 +729,24 @@ class SDFDataStore(AbstractDataStore):
                 "full_name": key,
                 "long_name": long_name,
             }
-            lazy_data = indexing.LazilyIndexedArray(SDFBackendArray(key, self))
-            data_vars[base_name] = Variable(var_coords, lazy_data, data_attrs)
+            lazy_data = indexing.LazilyIndexedArray(
+                SDFBackendArray(key, self, shape=value.shape, dtype=value.data.dtype)
+            )
+            var = Variable(var_coords, lazy_data, data_attrs)
+            # Set preferred chunks to match on-disk layout
+            # For point data (1D): full dimension
+            # For grid data (N-D): individual grid chunk sizes
+            if value.is_point_data:
+                var.encoding["preferred_chunks"] = {var_coords[0]: len(value.data)}
+            else:
+                # Align with on-disk grid structure
+                chunk_dict = {}
+                for dim_name, size in zip(var_coords, value.shape):
+                    # Use natural on-disk boundaries
+                    chunk_dict[dim_name] = size
+                var.encoding["preferred_chunks"] = chunk_dict
+
+            data_vars[base_name] = var
 
         # TODO: might need to decode if mult is set?
 
@@ -533,6 +765,14 @@ class SDFDataStore(AbstractDataStore):
 
 
 class SDFEntrypoint(BackendEntrypoint):
+    supports_groups = True
+    open_dataset_parameters: ClassVar[list[str]] = [
+        "filename_or_obj",
+        "drop_variables",
+        "keep_particles",
+        "probe_names",
+    ]
+
     def open_dataset(
         self,
         filename_or_obj,
@@ -555,13 +795,29 @@ class SDFEntrypoint(BackendEntrypoint):
         with close_on_error(store):
             return store.load()
 
-    open_dataset_parameters: ClassVar[list[str]] = [
+    open_datatree_parameters: ClassVar[list[str]] = [
         "filename_or_obj",
         "drop_variables",
         "keep_particles",
         "probe_names",
     ]
 
+    def open_datatree(
+        self,
+        filename_or_obj,
+        *,
+        drop_variables=None,
+        keep_particles=False,
+        probe_names=None,
+    ):
+        ds = self.open_dataset(
+            filename_or_obj,
+            drop_variables=drop_variables,
+            keep_particles=keep_particles,
+            probe_names=probe_names,
+        )
+        return _build_datatree_from_dataset(ds)
+
     def guess_can_open(self, filename_or_obj):
         magic_number = try_read_magic_number_from_path(filename_or_obj)
         if magic_number is not None:
@@ -574,6 +830,11 @@ class SDFEntrypoint(BackendEntrypoint):
     url = "https://epochpic.github.io/documentation/visualising_output/python_beam.html"
 
 
+class XrTUIEntrpoint:
+    def open_mfdatatree(self, paths: list[Path]) -> xr.DataTree:
+        return open_mfdatatree(paths)
+
+
 class SDFPreprocess:
     """Preprocess SDF files for xarray ensuring matching job ids and sets
     time dimension.

sdf_xarray/_version.py

@@ -28,7 +28,7 @@ version_tuple: VERSION_TUPLE
 commit_id: COMMIT_ID
 __commit_id__: COMMIT_ID
 
-__version__ = version = '0.3.2'
-__version_tuple__ = version_tuple = (0, 3, 2)
+__version__ = version = '0.5.0'
+__version_tuple__ = version_tuple = (0, 5, 0)
 
-__commit_id__ = commit_id = 'g331520e50'
+__commit_id__ = commit_id = 'g8a1775409'

sdf_xarray/dataset_accessor.py

@@ -1,7 +1,15 @@
-from typing import Union
+from __future__ import annotations
+
+from types import MethodType
+from typing import TYPE_CHECKING
 
 import xarray as xr
 
+from .plotting import animate_multiple, show
+
+if TYPE_CHECKING:
+    from matplotlib.animation import FuncAnimation
+
 
 @xr.register_dataset_accessor("epoch")
 class EpochAccessor:
@@ -13,7 +21,7 @@ class EpochAccessor:
         self,
         multiplier: float,
         unit_label: str,
-        coord_names: Union[str, list[str]],
+        coord_names: str | list[str],
     ) -> xr.Dataset:
         """
         Rescales specified X and Y coordinates in the Dataset by a given multiplier
@@ -71,3 +79,46 @@ class EpochAccessor:
             new_coords[coord_name] = coord_rescaled
 
         return ds.assign_coords(new_coords)
+
+    def animate_multiple(
+        self,
+        *variables: str | xr.DataArray,
+        datasets_kwargs: list[dict] | None = None,
+        **kwargs,
+    ) -> FuncAnimation:
+        """
+        Animate multiple Dataset variables on the same axes.
+
+        Parameters
+        ----------
+        variables
+            The variables to animate.
+        datasets_kwargs
+            Per-dataset keyword arguments passed to plotting.
+        kwargs
+            Common keyword arguments forwarded to animation.
+
+        Examples
+        --------
+        >>> anim = ds.epoch.animate_multiple(
+                ds["Derived_Number_Density_Electron"],
+                ds["Derived_Number_Density_Ion"],
+                datasets_kwargs=[{"label": "Electron"}, {"label": "Ion"}],
+                ylabel="Derived Number Density [1/m$^3$]"
+            )
+        >>> anim.save("animation.gif")
+        >>> # Or in a jupyter notebook:
+        >>> anim.show()
+        """
+
+        dataarrays = [
+            self._obj[var] if isinstance(var, str) else var for var in variables
+        ]
+        anim = animate_multiple(
+            *dataarrays,
+            datasets_kwargs=datasets_kwargs,
+            **kwargs,
+        )
+        anim.show = MethodType(show, anim)
+
+        return anim

sdf_xarray/download.py

@@ -0,0 +1,88 @@
+from pathlib import Path
+from shutil import move
+from typing import TYPE_CHECKING, Literal, TypeAlias
+
+if TYPE_CHECKING:
+    import pooch  # noqa: F401
+
+DatasetName: TypeAlias = Literal[
+    "test_array_no_grids",
+    "test_dist_fn",
+    "test_files_1D",
+    "test_files_2D_moving_window",
+    "test_files_3D",
+    "test_mismatched_files",
+    "test_two_probes_2D",
+    "tutorial_dataset_1d",
+    "tutorial_dataset_2d",
+    "tutorial_dataset_2d_moving_window",
+    "tutorial_dataset_3d",
+]
+
+
+def fetch_dataset(
+    dataset_name: DatasetName, save_path: Path | str | None = None
+) -> Path:
+    """
+    Downloads the specified dataset from its Zenodo URL. If it is already
+    downloaded, then the path to the cached, unzipped directory is returned.
+
+    Parameters
+    ---------
+    dataset_name
+        The name of the dataset to download
+    save_path
+        The directory to save the dataset to (defaults to the cache folder ``"sdf_datasets"``.
+        See `pooch.os_cache` for details on how the cache works)
+
+    Returns
+    -------
+    Path
+        The path to the directory containing the unzipped dataset files
+
+    Examples
+    --------
+    >>> # Assuming the dataset has not been downloaded yet
+    >>> path = fetch_dataset("tutorial_dataset_1d")
+    Downloading file 'tutorial_dataset_1d.zip' ...
+    Unzipping contents of '.../sdf_datasets/tutorial_dataset_1d.zip' to '.../sdf_datasets/tutorial_dataset_1d'
+    >>> path
+    '.../sdf_datasets/tutorial_dataset_1d'
+    """
+    import pooch  # noqa: PLC0415
+
+    logger = pooch.get_logger()
+    datasets = pooch.create(
+        path=pooch.os_cache("sdf_datasets"),
+        base_url="https://zenodo.org/records/17991042/files",
+        registry={
+            "test_array_no_grids.zip": "md5:583c85ed8c31d0e34e7766b6d9f2d6da",
+            "test_dist_fn.zip": "md5:a582ff5e8c59bad62fe4897f65fc7a11",
+            "test_files_1D.zip": "md5:42e53b229556c174c538c5481c4d596a",
+            "test_files_2D_moving_window.zip": "md5:3744483bbf416936ad6df8847c54dad1",
+            "test_files_3D.zip": "md5:a679e71281bab1d373dc4980e6da1a7c",
+            "test_mismatched_files.zip": "md5:710fdc94666edf7777523e8fc9dd1bd4",
+            "test_two_probes_2D.zip": "md5:0f2a4fefe84a15292d066b3320d4d533",
+            "tutorial_dataset_1d.zip": "md5:7fad744d8b8b2b84bba5c0e705fdef7b",
+            "tutorial_dataset_2d.zip": "md5:b7f35c05703a48eb5128049cdd106ffa",
+            "tutorial_dataset_2d_moving_window.zip": "md5:a795f40d18df69263842055de4559501",
+            "tutorial_dataset_3d.zip": "md5:d9254648867016292440fdb028f717f7",
+        },
+        retry_if_failed=10,
+    )
+
+    datasets.fetch(
+        f"{dataset_name}.zip", processor=pooch.Unzip(extract_dir="."), progressbar=True
+    )
+    cache_path = Path(datasets.path) / dataset_name
+
+    if save_path is not None:
+        save_path = Path(save_path)
+        logger.info(
+            "Moving contents of '%s' to '%s'",
+            cache_path,
+            save_path / dataset_name,
+        )
+        return move(cache_path, save_path / dataset_name)
+
+    return cache_path