Add Xarray accessor for Epoch animation funciton

src/sdf_xarray/__init__.py

@@ -12,6 +12,8 @@
 from xarray.core.utils import close_on_error, try_read_magic_number_from_path
 from xarray.core.variable import Variable
 
+import sdf_xarray.plotting  # noqa: F401
+
 from .sdf_interface import Constant, SDFFile
 
 
@@ -331,9 +333,9 @@ def _process_grid_name(grid_name: str, transform_func) -> str:
                 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]

src/sdf_xarray/plotting.py

@@ -46,31 +46,29 @@ def generate_animation(
     ax: plt.Axes | None = None,
     **kwargs,
 ) -> FuncAnimation:
-    """Generate an animation for the given target attribute
+    """Generate an animation for the given target attribute.
 
-    Arguments
+    Parameters
     ---------
-        dataset:
-            The dataset containing the simulation data
-        target_attribute:
-            The attribute to plot for each timestep
-        folder_path:
-            The path to save the generated animation (default: None)
-        display:
-            Whether to display the animation in the notebook (default: False)
-        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:
-            Dictionary of variables from matplotlib
+    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
     --------
-    >>> generateAnimation(dataset, "Derived_Number_Density_Electron")
+    >>> generate_animation(dataset, "Derived_Number_Density_Electron")
     """
     import matplotlib.pyplot as plt
     from matplotlib.animation import FuncAnimation
@@ -146,3 +144,31 @@ def update(frame):
         interval=1000 / fps,
         repeat=True,
     )
+
+
+@xr.register_dataset_accessor("epoch")
+class EpochAccessor:
+    def __init__(self, xarray_obj):
+        self._obj = xarray_obj
+
+    def animate(self, target: str, *args, **kwargs) -> FuncAnimation:
+        """Generate animation of 2D Epoch data.
+
+        Parameters
+        ----------
+        target
+            Attribute to plot for each time step.
+        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.epoch.animate("Electric_Field_Ey")
+        >>> ani.save("myfile.mp4")
+        """
+        return generate_animation(self._obj, target, *args, **kwargs)