sdf-xarray 0.2.3__tar.gz → 0.2.4__tar.gz

{sdf_xarray-0.2.3 → sdf_xarray-0.2.4}/CITATION.cff

@@ -16,4 +16,5 @@ authors:
     given-names: Shaun
     orcid: 'https://orcid.org/0009-0005-0693-030X'
     affiliation: University of York
+doi: 10.5281/zenodo.15351323
 date-released: '2024-07-25'

{sdf_xarray-0.2.3 → sdf_xarray-0.2.4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: sdf-xarray
-Version: 0.2.3
+Version: 0.2.4
 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
@@ -45,7 +45,6 @@ Requires-Python: >=3.10
 Requires-Dist: numpy>=2.0.0
 Requires-Dist: xarray>=2024.1.0
 Requires-Dist: dask>=2024.7.1
-Requires-Dist: cython>=3.0
 Provides-Extra: docs
 Requires-Dist: sphinx>=5.3; extra == "docs"
 Requires-Dist: sphinx_autodoc_typehints>=1.19; extra == "docs"
@@ -78,6 +77,7 @@ Description-Content-Type: text/markdown
 
 ![Dynamic TOML Badge](https://img.shields.io/badge/dynamic/toml?url=https%3A%2F%2Fraw.githubusercontent.com%2Fepochpic%2Fsdf-xarray%2Frefs%2Fheads%2Fmain%2Fpyproject.toml&query=%24.project.requires-python&label=python&logo=python)
 [![Available on PyPI](https://img.shields.io/pypi/v/sdf-xarray?color=blue&logo=pypi)](https://pypi.org/project/sdf-xarray/)
+[![DOI](https://zenodo.org/badge/DOI/10.5281/zenodo.15351323.svg)](https://doi.org/10.5281/zenodo.15351323)
 ![Build/Publish](https://github.com/epochpic/sdf-xarray/actions/workflows/build_publish.yml/badge.svg)
 ![Tests](https://github.com/epochpic/sdf-xarray/actions/workflows/tests.yml/badge.svg)
 [![Read the Docs](https://img.shields.io/readthedocs/sdf-xarray?logo=readthedocs&link=https%3A%2F%2Fsdf-xarray.readthedocs.io%2F)](https://sdf-xarray.readthedocs.io)

{sdf_xarray-0.2.3 → sdf_xarray-0.2.4}/README.md

@@ -2,6 +2,7 @@
 
 ![Dynamic TOML Badge](https://img.shields.io/badge/dynamic/toml?url=https%3A%2F%2Fraw.githubusercontent.com%2Fepochpic%2Fsdf-xarray%2Frefs%2Fheads%2Fmain%2Fpyproject.toml&query=%24.project.requires-python&label=python&logo=python)
 [![Available on PyPI](https://img.shields.io/pypi/v/sdf-xarray?color=blue&logo=pypi)](https://pypi.org/project/sdf-xarray/)
+[![DOI](https://zenodo.org/badge/DOI/10.5281/zenodo.15351323.svg)](https://doi.org/10.5281/zenodo.15351323)
 ![Build/Publish](https://github.com/epochpic/sdf-xarray/actions/workflows/build_publish.yml/badge.svg)
 ![Tests](https://github.com/epochpic/sdf-xarray/actions/workflows/tests.yml/badge.svg)
 [![Read the Docs](https://img.shields.io/readthedocs/sdf-xarray?logo=readthedocs&link=https%3A%2F%2Fsdf-xarray.readthedocs.io%2F)](https://sdf-xarray.readthedocs.io)

{sdf_xarray-0.2.3 → sdf_xarray-0.2.4}/pyproject.toml

@@ -3,7 +3,7 @@ requires = [
   "scikit-build-core>=0.3",
   "setuptools_scm",
   "numpy>=2.0.0",
-  "cython>=3.0",
+  "cython~=3.0",
 ]
 build-backend = "scikit_build_core.build"
 
@@ -22,7 +22,6 @@ dependencies = [
   "numpy>=2.0.0",
   "xarray>=2024.1.0",
   "dask>=2024.7.1",
-  "cython>=3.0",
 ]
 description = "Provides a backend for xarray to read SDF files as created by the EPOCH plasma PIC code."
 classifiers = [
@@ -105,4 +104,5 @@ extend-select = [
 ignore = [
   "PLR2004", # magic-comparison
   "B9",      # flake8-bugbear opinionated warnings
+  "PLR0913", # remove maximum number of arguments in a function
 ]

{sdf_xarray-0.2.3 → sdf_xarray-0.2.4}/src/sdf_xarray/_version.py

@@ -17,5 +17,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '0.2.3'
-__version_tuple__ = version_tuple = (0, 2, 3)
+__version__ = version = '0.2.4'
+__version_tuple__ = version_tuple = (0, 2, 4)

sdf_xarray-0.2.4/src/sdf_xarray/plotting.py

@@ -0,0 +1,205 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import numpy as np
+import xarray as xr
+
+if TYPE_CHECKING:
+    import matplotlib.pyplot as plt
+    from matplotlib.animation import FuncAnimation
+
+
+def get_frame_title(
+    data: xr.DataArray,
+    frame: int,
+    display_sdf_name: bool = False,
+    title_custom: str | None = None,
+) -> str:
+    """Generate the title for a frame"""
+    # 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()
+
+    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}"
+
+    # 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}"
+
+
+def calculate_window_boundaries(
+    data: xr.DataArray, xlim: tuple[float, float] | False = False
+) -> 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.
+    """
+    x_grid = data["X_Grid_mid"].values
+    x_half_cell = (x_grid[1] - x_grid[0]) / 2
+    N_frames = data["time"].size
+
+    # Find the window bounderies 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):
+        # Check if data is 1D
+        if data.ndim == 2:
+            target_lineout = data[i].values
+        # Check if data is 2D
+        if data.ndim == 3:
+            target_lineout = data[i, :, 0].values
+        x_grid_non_nan = x_grid[~np.isnan(target_lineout)]
+        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:
+        window_boundaries = window_boundaries + xlim - window_boundaries[0]
+    return window_boundaries
+
+
+def compute_global_limits(
+    data: xr.DataArray,
+    min_percentile: float = 0,
+    max_percentile: float = 100,
+) -> 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.
+    """
+
+    # Removes NaN values, needed for moving windows
+    values_no_nan = data.values[~np.isnan(data.values)]
+
+    # Finds the global minimum and maximum of the plot, based on the percentile of the data
+    global_min = np.percentile(values_no_nan, min_percentile)
+    global_max = np.percentile(values_no_nan, max_percentile)
+    return global_min, global_max
+
+
+def animate(
+    data: xr.DataArray,
+    fps: float = 10,
+    min_percentile: float = 0,
+    max_percentile: float = 100,
+    title: str | None = None,
+    display_sdf_name: bool = False,
+    ax: plt.Axes | None = None,
+    **kwargs,
+) -> FuncAnimation:
+    """Generate an animation
+
+    Parameters
+    ---------
+    data
+        The dataarray containing the target data
+    fps
+        Frames per second for the animation (default: 10)
+    min_percentile
+        Minimum percentile of the data (default: 0)
+    max_percentile
+        Maximum percentile of the data (default: 100)
+    title
+        Custom title to add to the plot.
+    display_sdf_name
+        Display the sdf file name in the animation title
+    ax
+        Matplotlib axes on which to plot.
+    kwargs
+        Keyword arguments to be passed to matplotlib.
+
+    Examples
+    --------
+    >>> dataset["Derived_Number_Density_Electron"].epoch.animate()
+    """
+    import matplotlib.pyplot as plt
+    from matplotlib.animation import FuncAnimation
+
+    kwargs_original = kwargs.copy()
+
+    if ax is None:
+        _, ax = plt.subplots()
+
+    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))
+        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 update(frame):
+        # Set the xlim for each frame in the case of a moving window
+        if move_window:
+            kwargs["xlim"] = window_boundaries[frame]
+
+        # 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))
+
+        # Update y-limits for 1D data
+        if data.ndim == 2:
+            ax.set_ylim(global_min, global_max)
+
+    return FuncAnimation(
+        ax.get_figure(),
+        update,
+        frames=range(N_frames),
+        interval=1000 / fps,
+        repeat=True,
+    )
+
+
+@xr.register_dataarray_accessor("epoch")
+class EpochAccessor:
+    def __init__(self, xarray_obj):
+        self._obj = xarray_obj
+
+    def animate(self, *args, **kwargs) -> FuncAnimation:
+        """Generate animations of Epoch data.
+
+        Parameters
+        ----------
+        args
+            Positional arguments passed to :func:`generate_animation`.
+        kwargs
+            Keyword arguments passed to :func:`generate_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")
+        """
+        return animate(self._obj, *args, **kwargs)

{sdf_xarray-0.2.3 → sdf_xarray-0.2.4}/src/sdf_xarray/sdf_interface.pyx

@@ -39,7 +39,7 @@ cdef class Block:
 
 @dataclasses.dataclass
 cdef class Variable(Block):
-    units: tuple[str] | None
+    units: str | None
     mult: float | None
     grid: str | None
     grid_mid: str | None

sdf_xarray-0.2.4/tests/example_files_2D_moving_window/input.deck

@@ -0,0 +1,63 @@
+begin:control
+  nx = 100
+  ny = 100
+
+  # final time of simulation
+  t_end = 10e-9
+
+  # size of domain
+  x_min = 0
+  x_end = 1
+
+  y_min = 0
+  y_max = 1
+
+  stdout_frequency = 10
+end:control
+
+
+begin:boundaries
+  bc_x_min = simple_outflow
+  bc_x_max = simple_outflow
+  bc_y_min = periodic
+  bc_y_max = periodic
+end:boundaries
+
+
+begin:window
+  move_window = T
+  window_start_time = 0
+  bc_x_min_after_move = simple_outflow
+  bc_x_max_after_move = simple_outflow
+  window_v_x = 2e8
+end:window
+
+begin:constant
+  x0 = 1.31
+  y0 = 0.5
+  r0 = 0.4^2
+  r2 = (x - x0)^2 + (y - y0)^2
+end:constant
+
+begin:species
+  # electron
+  name = electron
+  charge = -1.0
+  mass = 1.0
+  nparticles_per_cell = 5
+
+  number_density = if(abs(x-x0) lt 0.3, 2, 1)
+  number_density = if(abs(y-y0) lt 0.2, number_density(electron), 1)
+
+  temperature_ev = 0
+end:species
+
+
+begin:output
+  name = normal
+  dt_snapshot = 10e-9/4
+
+  # Properties on grid
+  grid = always
+  number_density = always + species + no_sum
+end:output

{sdf_xarray-0.2.3 → sdf_xarray-0.2.4}/tests/test_basic.py

@@ -5,7 +5,7 @@ import xarray as xr
 
 from sdf_xarray import SDFPreprocess, _process_latex_name, open_mfdataset
 
-EXAMPLE_FILES_DIR = pathlib.Path(__file__).parent / "example_files"
+EXAMPLE_FILES_DIR = pathlib.Path(__file__).parent / "example_files_1D"
 EXAMPLE_MISMATCHED_FILES_DIR = (
     pathlib.Path(__file__).parent / "example_mismatched_files"
 )

{sdf_xarray-0.2.3 → sdf_xarray-0.2.4}/tests/test_cython.py

@@ -6,7 +6,7 @@ import pytest
 
 from sdf_xarray import SDFFile
 
-EXAMPLE_FILES_DIR = pathlib.Path(__file__).parent / "example_files"
+EXAMPLE_FILES_DIR = pathlib.Path(__file__).parent / "example_files_1D"
 
 
 def test_sdffile():

sdf_xarray-0.2.4/tests/test_epoch_accessor.py

@@ -0,0 +1,183 @@
+import pathlib
+import tempfile
+
+import matplotlib as mpl
+import numpy as np
+import pytest
+import xarray as xr
+from matplotlib.animation import PillowWriter
+
+import sdf_xarray.plotting as sxp
+from sdf_xarray import SDFPreprocess
+
+mpl.use("Agg")
+
+EXAMPLE_FILES_DIR_1D = pathlib.Path(__file__).parent / "example_files_1D"
+EXAMPLE_FILES_DIR_2D_MW = (
+    pathlib.Path(__file__).parent / "example_files_2D_moving_window"
+)
+
+
+def test_animation_accessor():
+    array = xr.DataArray(
+        [1, 2, 3],
+        dims=["x"],
+        coords={"x": [0, 1, 2]},
+        attrs={"long_name": "Test Array", "units": "m"},
+    )
+    assert hasattr(array, "epoch")
+    assert hasattr(array.epoch, "animate")
+
+
+def test_animate_headless():
+    with xr.open_mfdataset(
+        EXAMPLE_FILES_DIR_1D.glob("*.sdf"), preprocess=SDFPreprocess()
+    ) as ds:
+        anim = ds["Derived_Number_Density_electron"].epoch.animate()
+
+        # Specify a custom writable temporary directory
+        with tempfile.TemporaryDirectory() as temp_dir:
+            temp_file_path = f"{temp_dir}/output.gif"
+            try:
+                anim.save(temp_file_path, writer=PillowWriter(fps=2))
+            except Exception as e:
+                pytest.fail(f"animate().save() failed in headless mode: {e}")
+
+
+def test_get_frame_title_no_optional_params():
+    with xr.open_mfdataset(
+        EXAMPLE_FILES_DIR_1D.glob("*.sdf"), preprocess=SDFPreprocess()
+    ) as ds:
+        data = ds["Derived_Number_Density_electron"]
+        expected_result = "time = 5.47e-14 [s]"
+        result = sxp.get_frame_title(data, 0)
+        assert expected_result == result
+
+
+def test_get_frame_title_sdf_name():
+    with xr.open_mfdataset(
+        EXAMPLE_FILES_DIR_1D.glob("*.sdf"), preprocess=SDFPreprocess()
+    ) as ds:
+        data = ds["Derived_Number_Density_electron"]
+        expected_result = "time = 5.47e-14 [s], 0000.sdf"
+        result = sxp.get_frame_title(data, 0, display_sdf_name=True)
+        assert expected_result == result
+
+
+def test_get_frame_title_custom_title():
+    with xr.open_mfdataset(
+        EXAMPLE_FILES_DIR_1D.glob("*.sdf"), preprocess=SDFPreprocess()
+    ) as ds:
+        data = ds["Derived_Number_Density_electron"]
+        expected_result = "Test Title, time = 5.47e-14 [s]"
+        result = sxp.get_frame_title(data, 0, title_custom="Test Title")
+        assert expected_result == result
+
+
+def test_get_frame_title_custom_title_and_sdf_name():
+    with xr.open_mfdataset(
+        EXAMPLE_FILES_DIR_1D.glob("*.sdf"), preprocess=SDFPreprocess()
+    ) as ds:
+        data = ds["Derived_Number_Density_electron"]
+        expected_result = "Test Title, time = 5.47e-14 [s], 0000.sdf"
+        result = sxp.get_frame_title(
+            data, 0, display_sdf_name=True, title_custom="Test Title"
+        )
+        assert expected_result == result
+
+
+def test_calculate_window_boundaries_1D():
+    with xr.open_mfdataset(
+        EXAMPLE_FILES_DIR_2D_MW.glob("*.sdf"),
+        preprocess=SDFPreprocess(),
+        combine="nested",
+    ) as ds:
+        data = ds["Derived_Number_Density_electron"][:, :, 50]
+        expected_result = np.array(
+            [[0, 1], [0.49, 1.49], [0.99, 1.99], [1.49, 2.49], [1.99, 2.99]]
+        )
+        result = sxp.calculate_window_boundaries(data)
+        assert result == pytest.approx(expected_result, abs=0.1)
+
+
+def test_calculate_window_boundaries_2D():
+    with xr.open_mfdataset(
+        EXAMPLE_FILES_DIR_2D_MW.glob("*.sdf"),
+        preprocess=SDFPreprocess(),
+        combine="nested",
+    ) as ds:
+        data = ds["Derived_Number_Density_electron"]
+        expected_result = np.array(
+            [[0, 1], [0.49, 1.49], [0.99, 1.99], [1.49, 2.49], [1.99, 2.99]]
+        )
+        result = sxp.calculate_window_boundaries(data)
+        assert result == pytest.approx(expected_result, abs=0.1)
+
+
+def test_calculate_window_boundaries_1D_xlim():
+    with xr.open_mfdataset(
+        EXAMPLE_FILES_DIR_2D_MW.glob("*.sdf"),
+        preprocess=SDFPreprocess(),
+        combine="nested",
+    ) as ds:
+        data = ds["Derived_Number_Density_electron"][:, :, 50]
+        expected_result = np.array(
+            [[0.1, 0.9], [0.59, 1.39], [1.09, 1.89], [1.59, 2.39], [2.09, 2.89]]
+        )
+        result = sxp.calculate_window_boundaries(data, xlim=(0.1, 0.9))
+        assert result == pytest.approx(expected_result, abs=0.1)
+
+
+def test_calculate_window_boundaries_2D_xlim():
+    with xr.open_mfdataset(
+        EXAMPLE_FILES_DIR_2D_MW.glob("*.sdf"),
+        preprocess=SDFPreprocess(),
+        combine="nested",
+    ) as ds:
+        data = ds["Derived_Number_Density_electron"]
+        expected_result = np.array(
+            [[0.1, 0.9], [0.59, 1.39], [1.09, 1.89], [1.59, 2.39], [2.09, 2.89]]
+        )
+        result = sxp.calculate_window_boundaries(data, xlim=(0.1, 0.9))
+        assert result == pytest.approx(expected_result, abs=0.1)
+
+
+def test_compute_global_limits():
+    with xr.open_mfdataset(
+        EXAMPLE_FILES_DIR_1D.glob("*.sdf"), preprocess=SDFPreprocess()
+    ) as ds:
+        result_min, result_max = sxp.compute_global_limits(
+            ds["Derived_Number_Density_electron"]
+        )
+        expected_result_min = 8.07e19
+        expected_result_max = 1.17e20
+        assert result_min == pytest.approx(expected_result_min, abs=1e18)
+        assert result_max == pytest.approx(expected_result_max, abs=1e19)
+
+
+def test_compute_global_limits_percentile():
+    with xr.open_mfdataset(
+        EXAMPLE_FILES_DIR_1D.glob("*.sdf"), preprocess=SDFPreprocess()
+    ) as ds:
+        result_min, result_max = sxp.compute_global_limits(
+            ds["Derived_Number_Density_electron"], 40, 45
+        )
+        expected_result_min = 9.84e19
+        expected_result_max = 9.94e19
+        assert result_min == pytest.approx(expected_result_min, abs=1e18)
+        assert result_max == pytest.approx(expected_result_max, abs=1e18)
+
+
+def test_compute_global_limits_NaNs():
+    with xr.open_mfdataset(
+        EXAMPLE_FILES_DIR_2D_MW.glob("*.sdf"),
+        preprocess=SDFPreprocess(),
+        combine="nested",
+    ) as ds:
+        result_min, result_max = sxp.compute_global_limits(
+            ds["Derived_Number_Density_electron"]
+        )
+        expected_result_min = 5.51e-1
+        expected_result_max = 2.70
+        assert result_min == pytest.approx(expected_result_min, abs=1e-2)
+        assert result_max == pytest.approx(expected_result_max, abs=1e-1)

sdf_xarray-0.2.3/src/sdf_xarray/plotting.py

@@ -1,184 +0,0 @@
-from __future__ import annotations
-
-from typing import TYPE_CHECKING
-
-import numpy as np
-import xarray as xr
-
-if TYPE_CHECKING:
-    import matplotlib.pyplot as plt
-    from matplotlib.animation import FuncAnimation
-
-
-def get_frame_title(data: xr.DataArray, frame: int, display_sdf_name: bool) -> str:
-    """Generate the title for a frame"""
-    sdf_name = f", {frame:04d}.sdf" if display_sdf_name else ""
-    time = data["time"][frame].to_numpy()
-    return f"t = {time:.2e}s{sdf_name}"
-
-
-def calculate_window_velocity_and_edges(
-    data: xr.DataArray, x_axis_coord: str
-) -> tuple[float, tuple[float, float], np.ndarray]:
-    """Calculate the moving window's velocity and initial edges.
-
-    1. Finds a lineout of the target atribute in the x coordinate of the first frame
-    2. Removes the NaN values to isolate the simulation window
-    3. Produces the index size of the window, indexed at zero
-    4. Uses distance moved and final time of the simulation to calculate velocity and initial xlims
-    """
-    time_since_start = data["time"].values - data["time"].values[0]
-    initial_window_edge = (0, 0)
-    target_lineout = data.values[0, :, 0]
-    target_lineout_window = target_lineout[~np.isnan(target_lineout)]
-    x_grid = data[x_axis_coord].values
-    window_size_index = target_lineout_window.size - 1
-
-    velocity_window = (x_grid[-1] - x_grid[window_size_index]) / time_since_start[-1]
-    initial_window_edge = (x_grid[0], x_grid[window_size_index])
-    return velocity_window, initial_window_edge, time_since_start
-
-
-def compute_global_limits(data: xr.DataArray) -> tuple[float, float]:
-    """Remove all NaN values from the target data to calculate the 1st and 99th percentiles,
-    excluding extreme outliers.
-    """
-    values_no_nan = data.values[~np.isnan(data.values)]
-    global_min = np.percentile(values_no_nan, 1)
-    global_max = np.percentile(values_no_nan, 99)
-    return global_min, global_max
-
-
-def is_1d(data: xr.DataArray) -> bool:
-    """Check if the data is 1D."""
-    return len(data.shape) == 2
-
-
-def is_2d(data: xr.DataArray) -> bool:
-    """Check if the data is 2D or 3D."""
-    return len(data.shape) == 3
-
-
-def generate_animation(
-    data: xr.DataArray,
-    display_sdf_name: bool = False,
-    fps: int = 10,
-    move_window: bool = False,
-    ax: plt.Axes | None = None,
-    **kwargs,
-) -> FuncAnimation:
-    """Generate an animation
-
-    Parameters
-    ---------
-    dataset
-        The dataset containing the simulation data
-    target_attribute
-        The attribute to plot for each timestep
-    display_sdf_name
-        Display the sdf file name in the animation title
-    fps
-        Frames per second for the animation (default: 10)
-    move_window
-        If the simulation has a moving window, the animation will move along
-        with it (default: False)
-    ax
-        Matplotlib axes on which to plot.
-    kwargs
-        Keyword arguments to be passed to matplotlib.
-
-    Examples
-    --------
-    >>> generate_animation(dataset["Derived_Number_Density_Electron"])
-    """
-    import matplotlib.pyplot as plt
-    from matplotlib.animation import FuncAnimation
-
-    if ax is None:
-        _, ax = plt.subplots()
-
-    N_frames = data["time"].size
-    global_min, global_max = compute_global_limits(data)
-
-    if is_2d(data):
-        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))
-
-        # Add colorbar
-        long_name = data.attrs.get("long_name")
-        units = data.attrs.get("units")
-        plt.colorbar(plot, ax=ax, label=f"{long_name} [${units}$]")
-
-    # Initialise plo and set y-limits for 1D data
-    if is_1d(data):
-        plot = data.isel(time=0).plot(ax=ax, **kwargs)
-        ax.set_title(get_frame_title(data, 0, display_sdf_name))
-        ax.set_ylim(global_min, global_max)
-
-    if move_window:
-        window_velocity, window_initial_edge, time_since_start = (
-            calculate_window_velocity_and_edges(data, kwargs["x"])
-        )
-
-    # User's choice for initial window edge supercides the one calculated
-    if "xlim" in kwargs:
-        window_initial_edge = kwargs["xlim"]
-
-    def update(frame):
-        # Set the xlim for each frame in the case of a moving window
-        if move_window:
-            kwargs["xlim"] = (
-                window_initial_edge[0] + window_velocity * time_since_start[frame],
-                window_initial_edge[1] * 0.99
-                + window_velocity * time_since_start[frame],
-            )
-
-        # 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))
-
-        # # Update y-limits for 1D data
-        if is_1d(data):
-            ax.set_ylim(global_min, global_max)
-
-    return FuncAnimation(
-        ax.get_figure(),
-        update,
-        frames=range(N_frames),
-        interval=1000 / fps,
-        repeat=True,
-    )
-
-
-@xr.register_dataarray_accessor("epoch")
-class EpochAccessor:
-    def __init__(self, xarray_obj):
-        self._obj = xarray_obj
-
-    def animate(self, *args, **kwargs) -> FuncAnimation:
-        """Generate animations of Epoch data.
-
-        Parameters
-        ----------
-        args
-            Positional arguments passed to :func:`generate_animation`.
-        kwargs
-            Keyword arguments passed to :func:`generate_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")
-        """
-        return generate_animation(self._obj, *args, **kwargs)

sdf_xarray-0.2.3/tests/test_epoch_accessor.py

@@ -1,39 +0,0 @@
-import pathlib
-import tempfile
-
-import matplotlib as mpl
-import pytest
-import xarray as xr
-from matplotlib.animation import PillowWriter
-
-from sdf_xarray import SDFPreprocess
-
-mpl.use("Agg")
-
-EXAMPLE_FILES_DIR = pathlib.Path(__file__).parent / "example_files"
-
-
-def test_animation_accessor():
-    array = xr.DataArray(
-        [1, 2, 3],
-        dims=["x"],
-        coords={"x": [0, 1, 2]},
-        attrs={"long_name": "Test Array", "units": "m"},
-    )
-    assert hasattr(array, "epoch")
-    assert hasattr(array.epoch, "animate")
-
-
-def test_animate_headless():
-    with xr.open_mfdataset(
-        EXAMPLE_FILES_DIR.glob("*.sdf"), preprocess=SDFPreprocess()
-    ) as ds:
-        anim = ds["Derived_Number_Density_electron"].epoch.animate()
-
-        # Specify a custom writable temporary directory
-        with tempfile.TemporaryDirectory() as temp_dir:
-            temp_file_path = f"{temp_dir}/output.gif"
-            try:
-                anim.save(temp_file_path, writer=PillowWriter(fps=2))
-            except Exception as e:
-                pytest.fail(f"animate().save() failed in headless mode: {e}")