sdf-xarray 0.2.5__cp310-cp310-win_amd64.whl → 0.4.0__cp310-cp310-win_amd64.whl

sdf_xarray/__init__.py

@@ -1,13 +1,17 @@
+import contextlib
 import os
 import re
 from collections import Counter, defaultdict
 from collections.abc import Callable, Iterable
+from importlib.metadata import version
 from itertools import product
+from os import PathLike as os_PathLike
 from pathlib import Path
 from typing import ClassVar
 
 import numpy as np
 import xarray as xr
+from packaging.version import Version
 from xarray.backends import AbstractDataStore, BackendArray, BackendEntrypoint
 from xarray.backends.file_manager import CachingFileManager
 from xarray.backends.locks import ensure_lock
@@ -15,12 +19,25 @@ from xarray.core import indexing
 from xarray.core.utils import close_on_error, try_read_magic_number_from_path
 from xarray.core.variable import Variable
 
-# NOTE: Do not delete this line, otherwise the "epoch" accessor will not be
-# imported when the user imports sdf_xarray
+# 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
+# has installed the package
+with contextlib.suppress(ImportError):
+    import pint_xarray  # noqa: F401
+
 from .sdf_interface import Constant, SDFFile  # type: ignore  # noqa: PGH003
 
+# TODO Remove this once the new kwarg options are fully implemented
+if Version(version("xarray")) >= Version("2025.8.0"):
+    xr.set_options(use_new_combine_kwarg_defaults=True)
+
+PathLike = str | os_PathLike
+
 
 def _rename_with_underscore(name: str) -> str:
     """A lot of the variable names have spaces, forward slashes and dashes in them, which
@@ -51,14 +68,69 @@ def _process_latex_name(variable_name: str) -> str:
     return variable_name
 
 
-def combine_datasets(path_glob: Iterable | str, **kwargs) -> xr.Dataset:
-    """Combine all datasets using a single time dimension"""
+def _resolve_glob(path_glob: PathLike | Iterable[PathLike]):
+    """
+    Normalise input path_glob into a sorted list of absolute, resolved Path objects.
+    """
+
+    try:
+        p = Path(path_glob)
+        paths = list(p.parent.glob(p.name)) if p.name == "*.sdf" else list(p)
+    except TypeError:
+        paths = list({Path(p) for p in path_glob})
+
+    paths = sorted(p.resolve() for p in paths)
+    if not paths:
+        raise FileNotFoundError(f"No files matched pattern or input: {path_glob!r}")
+    return paths
+
+
+def purge_unselected_data_vars(ds: xr.Dataset, data_vars: list[str]) -> xr.Dataset:
+    """
+    If the user has exclusively requested only certain variables be
+    loaded in then we purge all other variables and dimensions
+    """
+    existing_data_vars = set(ds.data_vars.keys())
+    vars_to_keep = set(data_vars) & existing_data_vars
+    vars_to_drop = existing_data_vars - vars_to_keep
+    ds = ds.drop_vars(vars_to_drop)
+
+    existing_dims = set(ds.sizes)
+    dims_to_keep = set()
+    for var in vars_to_keep:
+        dims_to_keep.update(ds[var].coords._names)
+        dims_to_keep.update(ds[var].dims)
+
+    coords_to_drop = existing_dims - dims_to_keep
+    return ds.drop_dims(coords_to_drop)
+
+
+def combine_datasets(
+    path_glob: Iterable | str, data_vars: list[str], **kwargs
+) -> xr.Dataset:
+    """
+    Combine all datasets using a single time dimension, optionally extract
+    data from only the listed data_vars
+    """
+
+    if data_vars is not None:
+        return xr.open_mfdataset(
+            path_glob,
+            join="outer",
+            coords="different",
+            compat="no_conflicts",
+            combine="nested",
+            concat_dim="time",
+            preprocess=SDFPreprocess(data_vars=data_vars),
+            **kwargs,
+        )
 
     return xr.open_mfdataset(
         path_glob,
-        data_vars="minimal",
-        coords="minimal",
-        compat="override",
+        data_vars="all",
+        coords="different",
+        compat="no_conflicts",
+        join="outer",
         preprocess=SDFPreprocess(),
         **kwargs,
     )
@@ -69,6 +141,8 @@ def open_mfdataset(
     *,
     separate_times: bool = False,
     keep_particles: bool = False,
+    probe_names: list[str] | None = None,
+    data_vars: list[str] | None = None,
 ) -> xr.Dataset:
     """Open a set of EPOCH SDF files as one `xarray.Dataset`
 
@@ -98,20 +172,36 @@ def open_mfdataset(
         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)
     """
 
-    # TODO: This is not very robust, look at how xarray.open_mfdataset does it
-    if isinstance(path_glob, str):
-        path_glob = Path().glob(path_glob)
-
-    # Coerce to list because we might need to use the sequence multiple times
-    path_glob = sorted(list(path_glob))  # noqa: C414
+    path_glob = _resolve_glob(path_glob)
 
     if not separate_times:
-        return combine_datasets(path_glob, keep_particles=keep_particles)
+        return combine_datasets(
+            path_glob,
+            data_vars=data_vars,
+            keep_particles=keep_particles,
+            probe_names=probe_names,
+        )
+
+    _, var_times_map = make_time_dims(path_glob)
 
-    time_dims, var_times_map = make_time_dims(path_glob)
-    all_dfs = [xr.open_dataset(f, keep_particles=keep_particles) for f in path_glob]
+    all_dfs = []
+    for f in path_glob:
+        ds = xr.open_dataset(f, keep_particles=keep_particles, probe_names=probe_names)
+
+        # 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
+        if data_vars is not None:
+            ds = purge_unselected_data_vars(ds, data_vars)
+            if not ds.data_vars:
+                continue
+
+        all_dfs.append(ds)
 
     for df in all_dfs:
         for da in df:
@@ -128,7 +218,11 @@ def open_mfdataset(
                 )
 
     return xr.combine_by_coords(
-        all_dfs, data_vars="minimal", combine_attrs="drop_conflicts"
+        all_dfs,
+        coords="different",
+        combine_attrs="drop_conflicts",
+        join="outer",
+        compat="no_conflicts",
     )
 
 
@@ -211,14 +305,23 @@ class SDFDataStore(AbstractDataStore):
         "drop_variables",
         "keep_particles",
         "lock",
+        "probe_names",
     )
 
-    def __init__(self, manager, drop_variables=None, keep_particles=False, lock=None):
+    def __init__(
+        self,
+        manager,
+        drop_variables=None,
+        keep_particles=False,
+        lock=None,
+        probe_names=None,
+    ):
         self._manager = manager
         self._filename = self.ds.filename
         self.drop_variables = drop_variables
         self.keep_particles = keep_particles
         self.lock = ensure_lock(lock)
+        self.probe_names = probe_names
 
     @classmethod
     def open(
@@ -227,6 +330,7 @@ class SDFDataStore(AbstractDataStore):
         lock=None,
         drop_variables=None,
         keep_particles=False,
+        probe_names=None,
     ):
         if isinstance(filename, os.PathLike):
             filename = os.fspath(filename)
@@ -237,6 +341,7 @@ class SDFDataStore(AbstractDataStore):
             lock=lock,
             drop_variables=drop_variables,
             keep_particles=keep_particles,
+            probe_names=probe_names,
         )
 
     def _acquire(self, needs_lock=True):
@@ -347,7 +452,28 @@ class SDFDataStore(AbstractDataStore):
 
             if value.is_point_data:
                 # Point (particle) variables are 1D
-                var_coords = (f"ID_{_process_grid_name(key, _grid_species_name)}",)
+
+                # Particle data does not maintain a fixed dimension size
+                # throughout the simulation. An example of a particle name comes
+                # in the form of `Particles/Px/Ion_H` which is then modified
+                # using `_process_grid_name()` into `Ion_H`. This is fine as the
+                # other components of the momentum (`Py`, `Pz`) will have the same
+                # size as they represent the same bunch of particles.
+
+                # Probes however have names in the form of `Electron_Front_Probe/Px`
+                # which are changed to just `Px`; this is fine when there is only one
+                # probe in the system but when there are multiple they will have
+                # conflicting sizes so we can't keep the names as simply `Px` so we
+                # instead set their dimension as the full name `Electron_Front_Probe_Px`.
+                is_probe_name_match = self.probe_names is not None and any(
+                    name in key for name in self.probe_names
+                )
+                name_processor = (
+                    _rename_with_underscore
+                    if is_probe_name_match
+                    else _grid_species_name
+                )
+                var_coords = (f"ID_{_process_grid_name(key, name_processor)}",)
             else:
                 # These are DataArrays
 
@@ -414,6 +540,7 @@ class SDFEntrypoint(BackendEntrypoint):
         *,
         drop_variables=None,
         keep_particles=False,
+        probe_names=None,
     ):
         if isinstance(filename_or_obj, Path):
             # sdf library takes a filename only
@@ -424,6 +551,7 @@ class SDFEntrypoint(BackendEntrypoint):
             filename_or_obj,
             drop_variables=drop_variables,
             keep_particles=keep_particles,
+            probe_names=probe_names,
         )
         with close_on_error(store):
             return store.load()
@@ -432,6 +560,7 @@ class SDFEntrypoint(BackendEntrypoint):
         "filename_or_obj",
         "drop_variables",
         "keep_particles",
+        "probe_names",
     ]
 
     def guess_can_open(self, filename_or_obj):
@@ -447,10 +576,43 @@ class SDFEntrypoint(BackendEntrypoint):
 
 
 class SDFPreprocess:
-    """Preprocess SDF files for xarray ensuring matching job ids and sets time dimension"""
+    """Preprocess SDF files for xarray ensuring matching job ids and sets
+    time dimension.
 
-    def __init__(self):
+    This class is used as a 'preprocess' function within ``xr.open_mfdataset``. It
+    performs three main duties on each individual file's Dataset:
+
+    1. Checks for a **matching job ID** across all files to ensure dataset consistency.
+    2. **Filters** the Dataset to keep only the variables specified in `data_vars`
+       and their required coordinates.
+    3. **Expands dimensions** to include a single 'time' coordinate, preparing the
+       Dataset for concatenation.
+
+    EPOCH can output variables at different intervals, so some SDF files
+    may not contain the requested variable. We combine this data into one
+    dataset by concatenating across the time dimension.
+
+    The combination is performed using ``join="outer"`` (in the calling ``open_mfdataset`` function),
+    meaning that the final combined dataset will contain the variable across the
+    entire time span, with NaNs filling the time steps where the variable was absent in
+    the individual file.
+
+    With large SDF files, this filtering method will save on memory consumption when
+    compared to loading all variables from all files before concatenation.
+
+    Parameters
+    ----------
+    data_vars :
+        A list of data variables to load in (If not specified loads
+        in all variables)
+    """
+
+    def __init__(
+        self,
+        data_vars: list[str] | None = None,
+    ):
         self.job_id: int | None = None
+        self.data_vars = data_vars
 
     def __call__(self, ds: xr.Dataset) -> xr.Dataset:
         if self.job_id is None:
@@ -461,17 +623,23 @@ class SDFPreprocess:
                 f"Mismatching job ids (got {ds.attrs['jobid1']}, expected {self.job_id})"
             )
 
-        ds = ds.expand_dims(time=[ds.attrs["time"]])
+        # If the user has exclusively requested only certain variables be
+        # loaded in then we purge all other variables and coordinates
+        if self.data_vars:
+            ds = purge_unselected_data_vars(ds, self.data_vars)
+
+        time_val = ds.attrs.get("time", np.nan)
+        ds = ds.expand_dims(time=[time_val])
         ds = ds.assign_coords(
             time=(
                 "time",
-                [ds.attrs["time"]],
+                [time_val],
                 {"units": "s", "long_name": "Time", "full_name": "time"},
             )
         )
         # Particles' spartial coordinates also evolve in time
         for coord, value in ds.coords.items():
             if value.attrs.get("point_data", False):
-                ds.coords[coord] = value.expand_dims(time=[ds.attrs["time"]])
+                ds.coords[coord] = value.expand_dims(time=[time_val])
 
         return ds

sdf_xarray/_version.py

@@ -1,7 +1,14 @@
 # file generated by setuptools-scm
 # don't change, don't track in version control
 
-__all__ = ["__version__", "__version_tuple__", "version", "version_tuple"]
+__all__ = [
+    "__version__",
+    "__version_tuple__",
+    "version",
+    "version_tuple",
+    "__commit_id__",
+    "commit_id",
+]
 
 TYPE_CHECKING = False
 if TYPE_CHECKING:
@@ -9,13 +16,19 @@ if TYPE_CHECKING:
     from typing import Union
 
     VERSION_TUPLE = Tuple[Union[int, str], ...]
+    COMMIT_ID = Union[str, None]
 else:
     VERSION_TUPLE = object
+    COMMIT_ID = object
 
 version: str
 __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
+commit_id: COMMIT_ID
+__commit_id__: COMMIT_ID
 
-__version__ = version = '0.2.5'
-__version_tuple__ = version_tuple = (0, 2, 5)
+__version__ = version = '0.4.0'
+__version_tuple__ = version_tuple = (0, 4, 0)
+
+__commit_id__ = commit_id = 'gc5cdb3bf9'

sdf_xarray/dataset_accessor.py

@@ -0,0 +1,71 @@
+import xarray as xr
+
+
+@xr.register_dataset_accessor("epoch")
+class EpochAccessor:
+    def __init__(self, xarray_obj: xr.Dataset):
+        # The xarray object is the Dataset, which we store as self._ds
+        self._ds = xarray_obj
+
+    def rescale_coords(
+        self,
+        multiplier: float,
+        unit_label: str,
+        coord_names: str | list[str],
+    ) -> xr.Dataset:
+        """
+        Rescales specified X and Y coordinates in the Dataset by a given multiplier
+        and updates the unit label attribute.
+
+        Parameters
+        ----------
+        multiplier : float
+            The factor by which to multiply the coordinate values (e.g., 1e6 for meters to microns).
+        unit_label : str
+            The new unit label for the coordinates (e.g., "µm").
+        coord_names : str or list of str
+            The name(s) of the coordinate variable(s) to rescale.
+            If a string, only that coordinate is rescaled.
+            If a list, all listed coordinates are rescaled.
+
+        Returns
+        -------
+        xr.Dataset
+            A new Dataset with the updated and rescaled coordinates.
+
+        Examples
+        --------
+        # Convert X, Y, and Z from meters to microns
+        >>> ds_in_microns = ds.epoch.rescale_coords(1e6, "µm", coord_names=["X_Grid", "Y_Grid", "Z_Grid"])
+
+        # Convert only X to millimeters
+        >>> ds_in_mm = ds.epoch.rescale_coords(1000, "mm", coord_names="X_Grid")
+        """
+
+        ds = self._ds
+        new_coords = {}
+
+        if isinstance(coord_names, str):
+            # Convert single string to a list
+            coords_to_process = [coord_names]
+        elif isinstance(coord_names, list):
+            # Use the provided list
+            coords_to_process = coord_names
+        else:
+            coords_to_process = list(coord_names)
+
+        for coord_name in coords_to_process:
+            if coord_name not in ds.coords:
+                raise ValueError(
+                    f"Coordinate '{coord_name}' not found in the Dataset. Cannot rescale."
+                )
+
+            coord_original = ds[coord_name]
+
+            coord_rescaled = coord_original * multiplier
+            coord_rescaled.attrs = coord_original.attrs.copy()
+            coord_rescaled.attrs["units"] = unit_label
+
+            new_coords[coord_name] = coord_rescaled
+
+        return ds.assign_coords(new_coords)

sdf_xarray/download.py

@@ -0,0 +1,87 @@
+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="doi:10.5281/zenodo.17618510",
+        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:1945ecdbc1ac1798164f83ea2b3d1b31",
+            "tutorial_dataset_2d_moving_window.zip": "md5:a795f40d18df69263842055de4559501",
+            "tutorial_dataset_3d.zip": "md5:d9254648867016292440fdb028f717f7",
+        },
+    )
+
+    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

sdf_xarray/plotting.py

@@ -9,37 +9,69 @@ if TYPE_CHECKING:
     import matplotlib.pyplot as plt
     from matplotlib.animation import FuncAnimation
 
+from types import MethodType
+
 
 def get_frame_title(
     data: xr.DataArray,
     frame: int,
     display_sdf_name: bool = False,
     title_custom: str | None = None,
+    t: str = "time",
 ) -> str:
-    """Generate the title for a frame"""
+    """Generate the title for a frame
+
+    Parameters
+    ----------
+    data
+        DataArray containing the target data
+    frame
+        Frame number
+    display_sdf_name
+        Display the sdf file name in the animation title
+    title_custom
+        Custom title to add to the plot
+    t
+        Time coordinate
+    """
+
     # Adds custom text to the start of the title, if specified
     title_custom = "" if title_custom is None else f"{title_custom}, "
-    # Adds the time and associated units to the title
-    time = data["time"][frame].to_numpy()
+    # Adds the time axis and associated units to the title
+    t_axis_value = data[t][frame].values
 
-    time_units = data["time"].attrs.get("units", False)
-    time_units_formatted = f" [{time_units}]" if time_units else ""
-    title_time = f"time = {time:.2e}{time_units_formatted}"
+    t_axis_units = data[t].attrs.get("units", False)
+    t_axis_units_formatted = f" [{t_axis_units}]" if t_axis_units else ""
+    title_t_axis = f"{data[t].long_name} = {t_axis_value:.2e}{t_axis_units_formatted}"
 
     # Adds sdf name to the title, if specifed
     title_sdf = f", {frame:04d}.sdf" if display_sdf_name else ""
-    return f"{title_custom}{title_time}{title_sdf}"
+    return f"{title_custom}{title_t_axis}{title_sdf}"
 
 
 def calculate_window_boundaries(
-    data: xr.DataArray, xlim: tuple[float, float] | False = False
+    data: xr.DataArray,
+    xlim: tuple[float, float] | None = None,
+    x_axis_name: str = "X_Grid_mid",
+    t: str = "time",
 ) -> np.ndarray:
     """Calculate the bounderies a moving window frame. If the user specifies xlim, this will
     be used as the initial bounderies and the window will move along acordingly.
+
+    Parameters
+    ----------
+    data
+        DataArray containing the target data
+    xlim
+        x limits
+    x_axis_name
+        Name of coordinate to assign to the x-axis
+    t
+        Time coordinate
     """
-    x_grid = data["X_Grid_mid"].values
+    x_grid = data[x_axis_name].values
     x_half_cell = (x_grid[1] - x_grid[0]) / 2
-    N_frames = data["time"].size
+    N_frames = data[t].size
 
     # Find the window bounderies by finding the first and last non-NaN values in the 0th lineout
     # along the x-axis.
@@ -56,7 +88,7 @@ def calculate_window_boundaries(
         window_boundaries[i, 1] = x_grid_non_nan[-1] + x_half_cell
 
     # User's choice for initial window edge supercides the one calculated
-    if xlim:
+    if xlim is not None:
         window_boundaries = window_boundaries + xlim - window_boundaries[0]
     return window_boundaries
 
@@ -68,6 +100,15 @@ def compute_global_limits(
 ) -> tuple[float, float]:
     """Remove all NaN values from the target data to calculate the global minimum and maximum of the data.
     User defined percentiles can remove extreme outliers.
+
+    Parameters
+    ----------
+    data
+        DataArray containing the target data
+    min_percentile
+        Minimum percentile of the data
+    max_percentile
+        Maximum percentile of the data
     """
 
     # Removes NaN values, needed for moving windows
@@ -86,74 +127,103 @@ def animate(
     max_percentile: float = 100,
     title: str | None = None,
     display_sdf_name: bool = False,
+    t: str | None = None,
     ax: plt.Axes | None = None,
     **kwargs,
 ) -> FuncAnimation:
-    """Generate an animation
+    """Generate an animation using an xarray.DataArray
 
     Parameters
     ---------
     data
-        The dataarray containing the target data
+        DataArray containing the target data
     fps
-        Frames per second for the animation (default: 10)
+        Frames per second for the animation
     min_percentile
-        Minimum percentile of the data (default: 0)
+        Minimum percentile of the data
     max_percentile
-        Maximum percentile of the data (default: 100)
+        Maximum percentile of the data
     title
-        Custom title to add to the plot.
+        Custom title to add to the plot
     display_sdf_name
         Display the sdf file name in the animation title
+    t
+        Coordinate for t axis (the coordinate which will be animated over). If `None`, use data.dims[0]
     ax
-        Matplotlib axes on which to plot.
+        Matplotlib axes on which to plot
     kwargs
-        Keyword arguments to be passed to matplotlib.
+        Keyword arguments to be passed to matplotlib
 
     Examples
     --------
-    >>> dataset["Derived_Number_Density_Electron"].epoch.animate()
+    >>> ds["Derived_Number_Density_Electron"].epoch.animate()
     """
-    import matplotlib.pyplot as plt
-    from matplotlib.animation import FuncAnimation
+    import matplotlib.pyplot as plt  # noqa: PLC0415
+    from matplotlib.animation import FuncAnimation  # noqa: PLC0415
 
     kwargs_original = kwargs.copy()
 
+    # Create plot if no ax is provided
     if ax is None:
-        _, ax = plt.subplots()
+        fig, ax = plt.subplots()
+        # Prevents figure from prematurely displaying in Jupyter notebook
+        plt.close(fig)
+
+    # Sets the animation coordinate (t) for iteration. If time is in the coords
+    # then it will set time to be t. If it is not it will fallback to the last
+    # coordinate passed in. By default coordinates are passed in from xarray in
+    # the form x, y, z so in order to preserve the x and y being on their
+    # respective axes we animate over the final coordinate that is passed in
+    # which in this example is z
+    coord_names = list(data.dims)
+    if t is None:
+        t = "time" if "time" in coord_names else coord_names[-1]
+    coord_names.remove(t)
+
+    N_frames = data[t].size
 
-    N_frames = data["time"].size
-    global_min, global_max = compute_global_limits(data, min_percentile, max_percentile)
-
-    # Initialise plot and set y-limits for 1D data
     if data.ndim == 2:
-        kwargs.setdefault("x", "X_Grid_mid")
-        plot = data.isel(time=0).plot(ax=ax, **kwargs)
-        ax.set_title(get_frame_title(data, 0, display_sdf_name, title))
+        kwargs.setdefault("x", coord_names[0])
+        plot = data.isel({t: 0}).plot(ax=ax, **kwargs)
+        ax.set_title(get_frame_title(data, 0, display_sdf_name, title, t))
+        global_min, global_max = compute_global_limits(
+            data, min_percentile, max_percentile
+        )
         ax.set_ylim(global_min, global_max)
 
-    # Initilise plot and set colour bar for 2D data
     if data.ndim == 3:
-        kwargs["norm"] = plt.Normalize(vmin=global_min, vmax=global_max)
+        if "norm" not in kwargs:
+            global_min, global_max = compute_global_limits(
+                data, min_percentile, max_percentile
+            )
+            kwargs["norm"] = plt.Normalize(vmin=global_min, vmax=global_max)
         kwargs["add_colorbar"] = False
-        # Set default x and y coordinates for 2D data if not provided
-        kwargs.setdefault("x", "X_Grid_mid")
-        kwargs.setdefault("y", "Y_Grid_mid")
+        # Set default x and y coordinates for 3D data if not provided
+        kwargs.setdefault("x", coord_names[0])
+        kwargs.setdefault("y", coord_names[1])
+
+        # Finds the time step with the minimum data value
+        # This is needed so that the animation can use the correct colour bar
+        argmin_time = np.unravel_index(data.argmin(), data.shape)[0]
 
-        # Initialize the plot with the first timestep
-        plot = data.isel(time=0).plot(ax=ax, **kwargs)
-        ax.set_title(get_frame_title(data, 0, display_sdf_name, title))
+        # Initialize the plot, the final output will still start at the first time step
+        plot = data.isel({t: argmin_time}).plot(ax=ax, **kwargs)
+        ax.set_title(get_frame_title(data, 0, display_sdf_name, title, t))
+        kwargs["cmap"] = plot.cmap
 
         # Add colorbar
         if kwargs_original.get("add_colorbar", True):
             long_name = data.attrs.get("long_name")
             units = data.attrs.get("units")
-            plt.colorbar(plot, ax=ax, label=f"{long_name} [${units}$]")
+            fig = plot.get_figure()
+            fig.colorbar(plot, ax=ax, label=f"{long_name} [{units}]")
 
     # check if there is a moving window by finding NaNs in the data
     move_window = np.isnan(np.sum(data.values))
     if move_window:
-        window_boundaries = calculate_window_boundaries(data, kwargs.get("xlim", False))
+        window_boundaries = calculate_window_boundaries(
+            data, kwargs.get("xlim"), kwargs["x"]
+        )
 
     def update(frame):
         # Set the xlim for each frame in the case of a moving window
@@ -163,12 +233,12 @@ def animate(
         # Update plot for the new frame
         ax.clear()
 
-        data.isel(time=frame).plot(ax=ax, **kwargs)
-        ax.set_title(get_frame_title(data, frame, display_sdf_name, title))
+        plot = data.isel({t: frame}).plot(ax=ax, **kwargs)
+        ax.set_title(get_frame_title(data, frame, display_sdf_name, title, t))
 
-        # Update y-limits for 1D data
         if data.ndim == 2:
             ax.set_ylim(global_min, global_max)
+        return plot
 
     return FuncAnimation(
         ax.get_figure(),
@@ -179,6 +249,19 @@ def animate(
     )
 
 
+def show(anim):
+    """Shows the FuncAnimation in a Jupyter notebook.
+
+    Parameters
+    ----------
+    anim
+        `matplotlib.animation.FuncAnimation`
+    """
+    from IPython.display import HTML  # noqa: PLC0415
+
+    return HTML(anim.to_jshtml())
+
+
 @xr.register_dataarray_accessor("epoch")
 class EpochAccessor:
     def __init__(self, xarray_obj):
@@ -190,16 +273,21 @@ class EpochAccessor:
         Parameters
         ----------
         args
-            Positional arguments passed to :func:`generate_animation`.
+            Positional arguments passed to :func:`animation`.
         kwargs
-            Keyword arguments passed to :func:`generate_animation`.
+            Keyword arguments passed to :func:`animation`.
 
         Examples
         --------
-        >>> import xarray as xr
-        >>> from sdf_xarray import SDFPreprocess
-        >>> ds = xr.open_mfdataset("*.sdf", preprocess=SDFPreprocess())
-        >>> ani = ds["Electric_Field_Ey"].epoch.animate()
-        >>> ani.save("myfile.mp4")
+        >>> anim = ds["Electric_Field_Ey"].epoch.animate()
+        >>> anim.save("myfile.mp4")
+        >>> # Or in a jupyter notebook:
+        >>> anim.show()
         """
-        return animate(self._obj, *args, **kwargs)
+
+        # Add anim.show() functionality
+        # anim.show() will display the animation in a jupyter notebook
+        anim = animate(self._obj, *args, **kwargs)
+        anim.show = MethodType(show, anim)
+
+        return anim

sdf_xarray/sdf_interface.pyx

@@ -110,12 +110,12 @@ _CONSTANT_UNITS_RE = re.compile(r"(?P<name>.*) \((?P<units>.*)\)$")
 cdef class Constant:
     _id: str
     name: str
-    data: int | str | float
+    data: int | str | float | bool
     units: str | None
 
     @staticmethod
     cdef Constant from_block(str name, csdf.sdf_block_t* block):
-        data: int | str | float | double
+        data: int | str | float | double | bool
 
         if block.datatype == csdf.SDF_DATATYPE_REAL4:
             data = (<float*>block.const_value)[0]
@@ -125,6 +125,8 @@ cdef class Constant:
             data = (<csdf.int32_t*>block.const_value)[0]
         if block.datatype == csdf.SDF_DATATYPE_INTEGER8:
             data = (<csdf.int64_t*>block.const_value)[0]
+        if block.datatype == csdf.SDF_DATATYPE_LOGICAL:
+            data = (<bint*>block.const_value)[0]
 
         # There's no metadata with e.g. units, but there's a
         # convention to put one in brackets at the end of the name,

{sdf_xarray-0.2.5.dist-info → sdf_xarray-0.4.0.dist-info}/METADATA

@@ -1,70 +1,24 @@
-Metadata-Version: 2.2
+Metadata-Version: 2.4
 Name: sdf-xarray
-Version: 0.2.5
+Version: 0.4.0
 Summary: Provides a backend for xarray to read SDF files as created by the EPOCH plasma PIC code.
-Author-Email: Peter Hill <peter.hill@york.ac.uk>, Joel Adams <joel.adams@york.ac.uk>, Shaun Doherty <shaun.doherty@york.ac.uk>
-License: Copyright 2024, Peter Hill, Joel Adams, epochpic team
-         
-         Redistribution and use in source and binary forms, with or without
-         modification, are permitted provided that the following conditions are
-         met:
-         
-         1. Redistributions of source code must retain the above copyright
-         notice, this list of conditions and the following disclaimer.
-         
-         2. Redistributions in binary form must reproduce the above copyright
-         notice, this list of conditions and the following disclaimer in the
-         documentation and/or other materials provided with the distribution.
-         
-         3. Neither the name of the copyright holder nor the names of its
-         contributors may be used to endorse or promote products derived from
-         this software without specific prior written permission.
-         
-         THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
-         “AS IS” AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
-         LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
-         A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
-         HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
-         SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
-         LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
-         DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
-         THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
-         (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
-         OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
-         
+Author-Email: Peter Hill <peter.hill@york.ac.uk>, Joel Adams <joel.adams@york.ac.uk>, Shaun Doherty <shaun.doherty@york.ac.uk>, Chris Herdman <chris.herdman@york.ac.uk>, Liam Pattinson <liam.pattinson@york.ac.uk>
+License-Expression: BSD-3-Clause
+Classifier: Development Status :: 5 - Production/Stable
+Classifier: Intended Audience :: Science/Research
+Classifier: Topic :: Scientific/Engineering
+Classifier: Operating System :: OS Independent
 Classifier: Programming Language :: Python
 Classifier: Programming Language :: Python :: 3
 Classifier: Programming Language :: Python :: 3.10
 Classifier: Programming Language :: Python :: 3.11
 Classifier: Programming Language :: Python :: 3.12
 Classifier: Programming Language :: Python :: 3.13
-Classifier: Intended Audience :: Science/Research
-Classifier: Topic :: Scientific/Engineering
-Classifier: Operating System :: OS Independent
-Requires-Python: >=3.10
+Classifier: Programming Language :: Python :: 3.14
+Requires-Python: <3.15,>=3.10
 Requires-Dist: numpy>=2.0.0
 Requires-Dist: xarray>=2024.1.0
 Requires-Dist: dask>=2024.7.1
-Provides-Extra: docs
-Requires-Dist: sphinx>=5.3; extra == "docs"
-Requires-Dist: sphinx_autodoc_typehints>=1.19; extra == "docs"
-Requires-Dist: sphinx-book-theme>=0.4.0rc1; extra == "docs"
-Requires-Dist: sphinx-argparse-cli>=1.10.0; extra == "docs"
-Requires-Dist: sphinx-inline-tabs; extra == "docs"
-Requires-Dist: pickleshare; extra == "docs"
-Requires-Dist: ipython; extra == "docs"
-Requires-Dist: matplotlib; extra == "docs"
-Requires-Dist: pint; extra == "docs"
-Requires-Dist: pint-xarray; extra == "docs"
-Requires-Dist: myst-parser; extra == "docs"
-Provides-Extra: test
-Requires-Dist: pytest>=3.3.0; extra == "test"
-Requires-Dist: dask[complete]; extra == "test"
-Requires-Dist: matplotlib; extra == "test"
-Provides-Extra: lint
-Requires-Dist: ruff; extra == "lint"
-Provides-Extra: build
-Requires-Dist: cibuildwheel[uv]; extra == "build"
 Provides-Extra: jupyter
 Requires-Dist: dask[diagnostics]; extra == "jupyter"
 Requires-Dist: ipykernel>=6.29.5; extra == "jupyter"
@@ -90,23 +44,19 @@ Part of [BEAM](#broad-epoch-analysis-modules-beam) (Broad EPOCH Analysis Modules
 
 ## Installation
 
+> [!IMPORTANT]
+> To install this package make sure you are using one of the Python versions listed above.
+
 Install from PyPI with:
 
 ```bash
 pip install sdf-xarray
 ```
 
-> [!NOTE]
-> For use within jupyter notebooks, run this additional command after installation:
->
-> ```bash
-> pip install "sdf-xarray[jupyter]"
-> ```
-
-or from a local checkout:
+or download this code locally:
 
 ```bash
-git clone https://github.com/epochpic/sdf-xarray.git
+git clone --recursive https://github.com/epochpic/sdf-xarray.git
 cd sdf-xarray
 pip install .
 ```
@@ -115,6 +65,9 @@ We recommend switching to [uv](https://docs.astral.sh/uv/) to manage packages.
 
 ## Usage
 
+Below are some simple examples to get you started. Please read the full
+documentation here <https://sdf-xarray.readthedocs.io>.
+
 ### Single file loading
 
 ```python
@@ -135,15 +88,22 @@ print(df["Electric_Field_Ex"])
 
 ### Multi-file loading
 
-To open a whole simulation at once, pass `preprocess=sdf_xarray.SDFPreprocess()`
-to `xarray.open_mfdataset`:
+You can open all the SDF files for a given simulation by calling the `open_mfdataset`
+function from `sdf_xarray`. This will additionally add a time dimension using the `"time"`
+value stored in each files attributes.
+
+> [!IMPORTANT]
+> If your simulation has multiple `output` blocks so that not all variables are
+> output at every time step, then at the timesteps where those variables are not
+> present they will have have a value of nan. To clean your dataset by removing
+> these nan values we suggest using the `xarray.DataArray.dropna` function or
+> loading sparse data along separate time dimensions using `separate_times=True`.
 
 ```python
-import xarray as xr
-from sdf_xarray import SDFPreprocess
+from sdf_xarray import open_mfdataset
 
-with xr.open_mfdataset("*.sdf", preprocess=SDFPreprocess()) as ds:
-    print(ds)
+ds = open_mfdataset("*.sdf")
+print(ds)
 
 # Dimensions:
 # time: 301, X_Grid_mid: 128, ...
@@ -153,15 +113,6 @@ with xr.open_mfdataset("*.sdf", preprocess=SDFPreprocess()) as ds:
 # Attributes: (22) ...
 ```
 
-`SDFPreprocess` checks that all the files are from the same simulation, as
-ensures there's a `time` dimension so the files are correctly concatenated.
-
-If your simulation has multiple `output` blocks so that not all variables are
-output at every time step, then those variables will have `NaN` values at the
-corresponding time points.
-
-For more in depth documentation please visit: <https://sdf-xarray.readthedocs.io/>
-
 ## Citing
 
 If sdf-xarray contributes to a project that leads to publication, please acknowledge this by citing sdf-xarray. This can be done by clicking the "cite this repository" button located near the top right of this page.

{sdf_xarray-0.2.5.dist-info → sdf_xarray-0.4.0.dist-info}/RECORD

@@ -6,19 +6,21 @@ include/SDFC_14.4.7/sdf_list_type.h,sha256=Quu8v0-SEsQuJpGtEZnm09tAyXqWNitx0sXl5
 include/SDFC_14.4.7/sdf_vector_type.h,sha256=dbKjhzRRsvhzrnTwVjtVlvnuisEnRMKY-vvdm94ok_Q,1595
 include/SDFC_14.4.7/stack_allocator.h,sha256=L7U9vmGiVSw3VQLIv9EzTaVq7JbFxs9aNonKStTkUSg,1335
 include/SDFC_14.4.7/uthash.h,sha256=rIyy_-ylY6S_7WaZCCC3VtvXaC9q37rFyA0f1U9xc4w,63030
-lib/SDFC_14.4.7/sdfc.lib,sha256=No_BVr7T1qROlbEiwg-jQL4JARLE-tW-C0uPc9wX-_A,350410
+lib/SDFC_14.4.7/sdfc.lib,sha256=FzSUfDKmu1bg4FtlzubnyO4SyrKmDTwNBQJVss7Gu5Y,350158
 lib/SDFC_14.4.7/SDFCConfig.cmake,sha256=IOA1eusC-KvUK4LNTEiOAmEdaPH1ZvNvbYPgiG1oZio,802
 lib/SDFC_14.4.7/SDFCConfigVersion.cmake,sha256=pN7Qqyf04s3izw7PYQ0XK6imvmhaVegSdR_nEl3Ok_o,2830
 lib/SDFC_14.4.7/SDFCTargets-release.cmake,sha256=G4zdx5PyjePigeD_a6rmZAxbk7L8Nf0klUnV78Lm2fI,828
 lib/SDFC_14.4.7/SDFCTargets.cmake,sha256=OVt1Gm8n7Ew4fiTmA9yHoef3vIIGwsXUZfqeG9p9Bys,4152
-sdf_xarray/__init__.py,sha256=MgATg9E6Jyo1fW3_li3lIEm-yxTHgJbNfZf5qeAvuyY,18073
-sdf_xarray/_version.py,sha256=XjF4m8_cMkXrbRU7Rim3aBSkgXL9DCvwDPeFFTVyWd8,532
+sdf_xarray/__init__.py,sha256=OAihj1CQVqoVvf-eEmwIttedjyplcoiTOfXpDafNqT8,24229
+sdf_xarray/_version.py,sha256=q0PPvfDga48CLepNF0cz12bk-1nMuF8tgn-8a0yG_w4,746
 sdf_xarray/csdf.pxd,sha256=ADPjAuHsodAvdOz96Z_XlFF7VL3KmVaXcTifWDP3rK0,4205
-sdf_xarray/plotting.py,sha256=ze1paC1Uw42WOWspdqkyNsUviDt-Z7AwQlPyO7JB90o,7007
-sdf_xarray/sdf_interface.cp310-win_amd64.pyd,sha256=XnLyq22rQxnDiE8Qrw1FeCEHxS6gOoL7U_rpwScnwjU,351232
-sdf_xarray/sdf_interface.pyx,sha256=PFC6upg14OZBqiGInLgBoxztIIKBk-HOh3WC9Ro4YUw,11975
-sdf_xarray-0.2.5.dist-info/METADATA,sha256=ybk80dZjGpoDSonOcmjyIO5HdCP4zUJRbqWjuxBn68c,9129
-sdf_xarray-0.2.5.dist-info/WHEEL,sha256=lOk88j15txMfymgGB6f9CZllTiUmb3GEt--Zg_z5WJg,106
-sdf_xarray-0.2.5.dist-info/entry_points.txt,sha256=gP7BIQpXNg6vIf7S7p-Rw_EJZTC1X50BsVTkK7dA7g0,57
-sdf_xarray-0.2.5.dist-info/licenses/LICENCE,sha256=aHWuyELjtzIL1jTXFHTbI3tr9vyVyhnw3I9_QYPdEX8,1515
-sdf_xarray-0.2.5.dist-info/RECORD,,
+sdf_xarray/dataset_accessor.py,sha256=IhROgmqgdD5SvuMbpZz-G6WoTre06-SOYgXVQSb7VAY,2450
+sdf_xarray/download.py,sha256=yT_z5q8KuGKe3yha_t7JW39IZjzdN2wczqRl8FIhgRA,3123
+sdf_xarray/plotting.py,sha256=FNaptcnHzvwY462JyfXSy9tMKNtuerKNn9A0sD8vbe4,9550
+sdf_xarray/sdf_interface.cp310-win_amd64.pyd,sha256=f1W3SDaA5jGniDeOY_-QkThgc_M64npCw16VNcQMN7s,340480
+sdf_xarray/sdf_interface.pyx,sha256=j0BpaltExUI_T-DMQvWaavKSrq9vqHkrmkJMJfCwDsA,12096
+sdf_xarray-0.4.0.dist-info/METADATA,sha256=9EoUJ22i_txDaZAX6ZIgnaQdogVv88M-L2VFIZDHofQ,6921
+sdf_xarray-0.4.0.dist-info/WHEEL,sha256=hrGeChGtn46HBGmzasO9QQDSLelRN-tUarBSv4gFcsI,106
+sdf_xarray-0.4.0.dist-info/entry_points.txt,sha256=gP7BIQpXNg6vIf7S7p-Rw_EJZTC1X50BsVTkK7dA7g0,57
+sdf_xarray-0.4.0.dist-info/licenses/LICENCE,sha256=aHWuyELjtzIL1jTXFHTbI3tr9vyVyhnw3I9_QYPdEX8,1515
+sdf_xarray-0.4.0.dist-info/RECORD,,

{sdf_xarray-0.2.5.dist-info → sdf_xarray-0.4.0.dist-info}/WHEEL

@@ -1,5 +1,5 @@
 Wheel-Version: 1.0
-Generator: scikit-build-core 0.11.2
+Generator: scikit-build-core 0.11.6
 Root-Is-Purelib: false
 Tag: cp310-cp310-win_amd64