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

sdf_xarray/plotting.py

@@ -1,6 +1,10 @@
 from __future__ import annotations
 
-from typing import TYPE_CHECKING
+import warnings
+from collections.abc import Callable
+from dataclasses import dataclass
+from types import MethodType
+from typing import TYPE_CHECKING, Any
 
 import numpy as np
 import xarray as xr
@@ -10,41 +14,77 @@ if TYPE_CHECKING:
     from matplotlib.animation import FuncAnimation
 
 
+@dataclass
+class AnimationUnit:
+    update: Callable[[int], object]
+    n_frames: int
+
+
 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.
+    """Calculate the boundaries a moving window frame. If the user specifies xlim, this will
+    be used as the initial boundaries 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
+    # Find the window boundaries by finding the first and last non-NaN values in the 0th lineout
     # along the x-axis.
-    window_boundaries = np.zeros((N_frames, 2))
-    for i in range(N_frames):
+    window_boundaries = np.zeros((n_frames, 2))
+    for i in range(n_frames):
         # Check if data is 1D
         if data.ndim == 2:
             target_lineout = data[i].values
@@ -55,8 +95,8 @@ def calculate_window_boundaries(
         window_boundaries[i, 0] = x_grid_non_nan[0] - x_half_cell
         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:
+    # User's choice for initial window edge supercedes the one calculated
+    if xlim is not None:
         window_boundaries = window_boundaries + xlim - window_boundaries[0]
     return window_boundaries
 
@@ -68,6 +108,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
@@ -79,6 +128,199 @@ def compute_global_limits(
     return global_min, global_max
 
 
+def _set_axes_labels(ax: plt.Axes, axis_kwargs: dict) -> None:
+    """Set the labels for the x and y axes"""
+    if "xlabel" in axis_kwargs:
+        ax.set_xlabel(axis_kwargs["xlabel"])
+    if "ylabel" in axis_kwargs:
+        ax.set_ylabel(axis_kwargs["ylabel"])
+
+
+def _setup_2d_plot(
+    data: xr.DataArray,
+    ax: plt.Axes,
+    coord_names: list[str],
+    kwargs: dict,
+    axis_kwargs: dict,
+    min_percentile: float,
+    max_percentile: float,
+    t: str,
+) -> tuple[float, float]:
+    """Setup 2D plot initialization."""
+
+    kwargs.setdefault("x", coord_names[0])
+
+    data.isel({t: 0}).plot(ax=ax, **kwargs)
+
+    global_min, global_max = compute_global_limits(data, min_percentile, max_percentile)
+
+    _set_axes_labels(ax, axis_kwargs)
+
+    if "ylim" not in kwargs:
+        ax.set_ylim(global_min, global_max)
+
+    return global_min, global_max
+
+
+def _setup_3d_plot(
+    data: xr.DataArray,
+    ax: plt.Axes,
+    coord_names: list[str],
+    kwargs: dict,
+    kwargs_original: dict,
+    axis_kwargs: dict,
+    min_percentile: float,
+    max_percentile: float,
+    t: str,
+) -> None:
+    """Setup 3D plot initialization."""
+    import matplotlib.pyplot as plt  # noqa: PLC0415
+
+    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
+    kwargs.setdefault("x", coord_names[0])
+    kwargs.setdefault("y", coord_names[1])
+
+    argmin_time = np.unravel_index(np.argmin(data.values), data.shape)[0]
+    plot = data.isel({t: argmin_time}).plot(ax=ax, **kwargs)
+    kwargs["cmap"] = plot.cmap
+
+    _set_axes_labels(ax, axis_kwargs)
+
+    if kwargs_original.get("add_colorbar", True):
+        long_name = data.attrs.get("long_name")
+        units = data.attrs.get("units")
+        fig = plot.get_figure()
+        fig.colorbar(plot, ax=ax, label=f"{long_name} [{units}]")
+
+
+def _generate_animation(
+    data: xr.DataArray,
+    clear_axes: bool = False,
+    min_percentile: float = 0,
+    max_percentile: float = 100,
+    title: str | None = None,
+    display_sdf_name: bool = False,
+    move_window: bool = False,
+    t: str | None = None,
+    ax: plt.Axes | None = None,
+    kwargs: dict | None = None,
+) -> AnimationUnit:
+    """
+    Internal function for generating the plotting logic required for animations.
+
+    Parameters
+    ---------
+    data
+        DataArray containing the target data
+    clear_axes
+        Decide whether to run ``ax.clear()`` in every update
+    min_percentile
+        Minimum percentile of the data
+    max_percentile
+        Maximum percentile of the data
+    title
+        Custom title to add to the plot
+    display_sdf_name
+        Display the sdf file name in the animation title
+    move_window
+        Update the ``xlim`` to be only values that are not NaNs at each time interval
+    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
+    kwargs
+        Keyword arguments to be passed to matplotlib
+
+    Examples
+    --------
+    >>> anim = animate(ds["Derived_Number_Density_Electron"])
+    >>> anim.save("animation.gif")
+    """
+
+    if kwargs is None:
+        kwargs = {}
+    kwargs_original = kwargs.copy()
+
+    axis_kwargs = {}
+    for key in ("xlabel", "ylabel"):
+        if key in kwargs:
+            axis_kwargs[key] = kwargs.pop(key)
+
+    # 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
+
+    global_min = global_max = None
+    if data.ndim == 2:
+        global_min, global_max = _setup_2d_plot(
+            data=data,
+            ax=ax,
+            coord_names=coord_names,
+            kwargs=kwargs,
+            axis_kwargs=axis_kwargs,
+            min_percentile=min_percentile,
+            max_percentile=max_percentile,
+            t=t,
+        )
+    elif data.ndim == 3:
+        _setup_3d_plot(
+            data=data,
+            ax=ax,
+            coord_names=coord_names,
+            kwargs=kwargs,
+            kwargs_original=kwargs_original,
+            axis_kwargs=axis_kwargs,
+            min_percentile=min_percentile,
+            max_percentile=max_percentile,
+            t=t,
+        )
+
+    ax.set_title(get_frame_title(data, 0, display_sdf_name, title, t))
+
+    window_boundaries = None
+    if move_window:
+        window_boundaries = calculate_window_boundaries(
+            data, kwargs.get("xlim"), kwargs["x"]
+        )
+
+    def update(frame):
+        if clear_axes:
+            ax.clear()
+        # Set the xlim for each frame in the case of a moving window
+        if move_window:
+            kwargs["xlim"] = window_boundaries[frame]
+
+        plot = data.isel({t: frame}).plot(ax=ax, **kwargs)
+        ax.set_title(get_frame_title(data, frame, display_sdf_name, title, t))
+        _set_axes_labels(ax, axis_kwargs)
+
+        if data.ndim == 2 and "ylim" not in kwargs and global_min is not None:
+            ax.set_ylim(global_min, global_max)
+
+        return plot
+
+    return AnimationUnit(
+        update=update,
+        n_frames=N_frames,
+    )
+
+
 def animate(
     data: xr.DataArray,
     fps: float = 10,
@@ -86,99 +328,214 @@ def animate(
     max_percentile: float = 100,
     title: str | None = None,
     display_sdf_name: bool = False,
+    move_window: bool = False,
+    t: str | None = None,
     ax: plt.Axes | None = None,
     **kwargs,
 ) -> FuncAnimation:
-    """Generate an animation
+    """
+    Generate an animation using an `xarray.DataArray`. The intended use
+    of this function is via `sdf_xarray.plotting.EpochAccessor.animate`.
 
     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
+    move_window
+        Update the ``xlim`` to be only values that are not NaNs at each time interval
+    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()
+    >>> anim = animate(ds["Derived_Number_Density_Electron"])
+    >>> anim.save("animation.gif")
     """
     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)
+
+    animation = _generate_animation(
+        data,
+        clear_axes=True,
+        min_percentile=min_percentile,
+        max_percentile=max_percentile,
+        title=title,
+        display_sdf_name=display_sdf_name,
+        move_window=move_window,
+        t=t,
+        ax=ax,
+        kwargs=kwargs,
+    )
 
-    N_frames = data["time"].size
-    global_min, global_max = compute_global_limits(data, min_percentile, max_percentile)
+    return FuncAnimation(
+        ax.get_figure(),
+        animation.update,
+        frames=range(animation.n_frames),
+        interval=1000 / fps,
+        repeat=True,
+    )
 
-    # 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))
-        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)
-        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")
-
-        # 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))
-
-        # 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}$]")
-
-    # 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))
+def animate_multiple(
+    *datasets: xr.DataArray,
+    datasets_kwargs: list[dict[str, Any]] | None = None,
+    fps: float = 10,
+    min_percentile: float = 0,
+    max_percentile: float = 100,
+    title: str | None = None,
+    display_sdf_name: bool = False,
+    move_window: bool = False,
+    t: str | None = None,
+    ax: plt.Axes | None = None,
+    **common_kwargs,
+) -> FuncAnimation:
+    """
+    Generate an animation using multiple `xarray.DataArray`. The intended use
+    of this function is via `sdf_xarray.dataset_accessor.EpochAccessor.animate_multiple`.
 
-    def update(frame):
-        # Set the xlim for each frame in the case of a moving window
-        if move_window:
-            kwargs["xlim"] = window_boundaries[frame]
+    Parameters
+    ---------
+    datasets
+        `xarray.DataArray` objects containing the data to be animated
+    datasets_kwargs
+       A list of dictionaries, following the same order as ``datasets``, containing
+        per-dataset matplotlib keyword arguments. The list does not need to be the same
+        length as ``datasets``; missing entries are initialised as empty dictionaries
+    fps
+        Frames per second for the animation
+    min_percentile
+        Minimum percentile of the data
+    max_percentile
+        Maximum percentile of the data
+    title
+        Custom title to add to the plot
+    display_sdf_name
+        Display the sdf file name in the animation title
+    move_window
+        Update the ``xlim`` to be only values that are not NaNs at each time interval
+    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
+    common_kwargs
+        Matplotlib keyword arguments applied to all datasets. These are overridden by
+        per-dataset entries in ``datasets_kwargs``
 
-        # Update plot for the new frame
-        ax.clear()
+    Examples
+    --------
+    >>> anim = animate_multiple(
+        ds["Derived_Number_Density_Electron"],
+        ds["Derived_Number_Density_Ion"],
+        datasets_kwargs=[{"label": "Electron"}, {"label": "Ion"}],
+        ylim=(0e27,4e27),
+        display_sdf_name=True,
+        ylabel="Derived Number Density [1/m$^3$]"
+    )
+    >>> anim.save("animation.gif")
+    """
+    import matplotlib.pyplot as plt  # noqa: PLC0415
+    from matplotlib.animation import FuncAnimation  # noqa: PLC0415
 
-        data.isel(time=frame).plot(ax=ax, **kwargs)
-        ax.set_title(get_frame_title(data, frame, display_sdf_name, title))
+    if not datasets:
+        raise ValueError("At least one dataset must be provided")
 
-        # Update y-limits for 1D data
-        if data.ndim == 2:
-            ax.set_ylim(global_min, global_max)
+    # Create plot if no ax is provided
+    if ax is None:
+        fig, ax = plt.subplots()
+        # Prevents figure from prematurely displaying in Jupyter notebook
+        plt.close(fig)
+
+    n_datasets = len(datasets)
+    if datasets_kwargs is None:
+        # Initialise an empty series of dicts the same size as the number of datasets
+        datasets_kwargs = [{} for _ in range(n_datasets)]
+    else:
+        # The user might only want to use kwargs on some of the datasets so we make sure
+        # to initialise additional empty dicts and append them to the list
+        datasets_kwargs.extend({} for _ in range(n_datasets - len(datasets_kwargs)))
+
+    animations: list[AnimationUnit] = []
+    for da, kw in zip(datasets, datasets_kwargs):
+        animations.append(
+            _generate_animation(
+                da,
+                ax=ax,
+                min_percentile=min_percentile,
+                max_percentile=max_percentile,
+                title=title,
+                display_sdf_name=display_sdf_name,
+                move_window=move_window,
+                t=t,
+                # Per-dataset kwargs override common matplotlib kwargs
+                kwargs={**common_kwargs, **kw},
+            )
+        )
+
+    lengths = [anim.n_frames for anim in animations]
+    n_frames = min(lengths)
+
+    if len(set(lengths)) > 1:
+        warnings.warn(
+            "Datasets have different frame counts; truncating to the shortest",
+            stacklevel=2,
+        )
+
+    # Render the legend if a label exists for any 2D dataset
+    show_legend = any(
+        "label" in kw and da.ndim == 2 for da, kw in zip(datasets, datasets_kwargs)
+    )
+
+    def update(frame):
+        ax.clear()
+        for anim in animations:
+            anim.update(frame)
+        if show_legend:
+            ax.legend(loc="upper right")
 
     return FuncAnimation(
         ax.get_figure(),
         update,
-        frames=range(N_frames),
+        frames=range(n_frames),
         interval=1000 / fps,
         repeat=True,
     )
 
 
+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 +547,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("animation.gif")
+        >>> # 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,