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

sdf_xarray/_version.py

@@ -1,16 +1,34 @@
-# file generated by setuptools_scm
+# file generated by setuptools-scm
 # don't change, don't track in version control
+
+__all__ = [
+    "__version__",
+    "__version_tuple__",
+    "version",
+    "version_tuple",
+    "__commit_id__",
+    "commit_id",
+]
+
 TYPE_CHECKING = False
 if TYPE_CHECKING:
-    from typing import Tuple, Union
+    from typing import Tuple
+    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.5.0'
+__version_tuple__ = version_tuple = (0, 5, 0)
 
-__version__ = version = '0.2.0'
-__version_tuple__ = version_tuple = (0, 2, 0)
+__commit_id__ = commit_id = 'g8a1775409'

sdf_xarray/dataset_accessor.py

@@ -0,0 +1,124 @@
+from __future__ import annotations
+
+from types import MethodType
+from typing import TYPE_CHECKING
+
+import xarray as xr
+
+from .plotting import animate_multiple, show
+
+if TYPE_CHECKING:
+    from matplotlib.animation import FuncAnimation
+
+
+@xr.register_dataset_accessor("epoch")
+class EpochAccessor:
+    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)
+
+    def animate_multiple(
+        self,
+        *variables: str | xr.DataArray,
+        datasets_kwargs: list[dict] | None = None,
+        **kwargs,
+    ) -> FuncAnimation:
+        """
+        Animate multiple Dataset variables on the same axes.
+
+        Parameters
+        ----------
+        variables
+            The variables to animate.
+        datasets_kwargs
+            Per-dataset keyword arguments passed to plotting.
+        kwargs
+            Common keyword arguments forwarded to animation.
+
+        Examples
+        --------
+        >>> anim = ds.epoch.animate_multiple(
+                ds["Derived_Number_Density_Electron"],
+                ds["Derived_Number_Density_Ion"],
+                datasets_kwargs=[{"label": "Electron"}, {"label": "Ion"}],
+                ylabel="Derived Number Density [1/m$^3$]"
+            )
+        >>> anim.save("animation.gif")
+        >>> # Or in a jupyter notebook:
+        >>> anim.show()
+        """
+
+        dataarrays = [
+            self._obj[var] if isinstance(var, str) else var for var in variables
+        ]
+        anim = animate_multiple(
+            *dataarrays,
+            datasets_kwargs=datasets_kwargs,
+            **kwargs,
+        )
+        anim.show = MethodType(show, anim)
+
+        return anim

sdf_xarray/download.py

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