sdf-xarray 0.5.0__cp314-cp314t-win_amd64.whl

sdf_xarray/_version.py

@@ -0,0 +1,34 @@
+# 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
+    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)
+
+__commit_id__ = commit_id = 'g8a1775409'

sdf_xarray/csdf.pxd

@@ -0,0 +1,127 @@
+cdef extern from "<stdint.h>" nogil:
+    ctypedef signed int int32_t
+    ctypedef signed long int64_t
+
+cdef extern from "sdf.h":
+    cdef enum:
+        SDF_VERSION
+        SDF_REVISION
+        SDF_LIB_VERSION
+        SDF_LIB_REVISION
+        SDF_MAGIC
+        SDF_MAXDIMS
+        SDF_READ
+        SDF_BLOCKTYPE_SCRUBBED
+        SDF_BLOCKTYPE_NULL
+        SDF_BLOCKTYPE_PLAIN_MESH
+        SDF_BLOCKTYPE_POINT_MESH
+        SDF_BLOCKTYPE_PLAIN_VARIABLE
+        SDF_BLOCKTYPE_POINT_VARIABLE
+        SDF_BLOCKTYPE_CONSTANT
+        SDF_BLOCKTYPE_ARRAY
+        SDF_BLOCKTYPE_RUN_INFO
+        SDF_BLOCKTYPE_SOURCE
+        SDF_BLOCKTYPE_STITCHED_TENSOR
+        SDF_BLOCKTYPE_STITCHED_MATERIAL
+        SDF_BLOCKTYPE_STITCHED_MATVAR
+        SDF_BLOCKTYPE_STITCHED_SPECIES
+        SDF_BLOCKTYPE_SPECIES
+        SDF_BLOCKTYPE_PLAIN_DERIVED
+        SDF_BLOCKTYPE_POINT_DERIVED
+        SDF_BLOCKTYPE_CONTIGUOUS_TENSOR
+        SDF_BLOCKTYPE_CONTIGUOUS_MATERIAL
+        SDF_BLOCKTYPE_CONTIGUOUS_MATVAR
+        SDF_BLOCKTYPE_CONTIGUOUS_SPECIES
+        SDF_BLOCKTYPE_CPU_SPLIT
+        SDF_BLOCKTYPE_STITCHED_OBSTACLE_GROUP
+        SDF_BLOCKTYPE_UNSTRUCTURED_MESH
+        SDF_BLOCKTYPE_STITCHED
+        SDF_BLOCKTYPE_CONTIGUOUS
+        SDF_BLOCKTYPE_LAGRANGIAN_MESH
+        SDF_BLOCKTYPE_STATION
+        SDF_BLOCKTYPE_STATION_DERIVED
+        SDF_BLOCKTYPE_DATABLOCK
+        SDF_BLOCKTYPE_NAMEVALUE
+        SDF_DATATYPE_NULL
+        SDF_DATATYPE_INTEGER4
+        SDF_DATATYPE_INTEGER8
+        SDF_DATATYPE_REAL4
+        SDF_DATATYPE_REAL8
+        SDF_DATATYPE_REAL16
+        SDF_DATATYPE_CHARACTER
+        SDF_DATATYPE_LOGICAL
+        SDF_DATATYPE_OTHER
+
+    ctypedef int comm_t
+
+    ctypedef struct sdf_block_t:
+        double* extents
+        double* dim_mults
+        double mult, time, time_increment
+        int64_t dims[SDF_MAXDIMS]
+        int64_t local_dims[SDF_MAXDIMS]
+        int64_t data_length
+        int32_t ndims, geometry, datatype, blocktype
+        int32_t stagger, datatype_out
+        char const_value[16]
+        char* id
+        char* units
+        char* mesh_id
+        char* material_id
+        char* name
+        char* material_name
+        char** dim_labels
+        char** dim_units
+        char** variable_ids
+        void** grids
+        void* data
+        sdf_block_t* next
+
+    ctypedef struct sdf_file_t:
+        int32_t sdf_lib_version, sdf_lib_revision
+        int32_t sdf_extension_version, sdf_extension_revision
+        int32_t file_version, file_revision
+        double time
+        int64_t current_location
+        int32_t jobid1, jobid2, endianness, summary_size
+        int32_t block_header_length, string_length, id_length
+        int32_t code_io_version, step
+        int32_t nblocks
+        char* buffer
+        char* filename
+        bint restart_flag, other_domains
+        bint station_file
+        bint restart_flag
+        char* code_name
+        sdf_block_t* blocklist
+        sdf_block_t* current_block
+
+    cdef struct run_info:
+        int64_t defines
+        int32_t version, revision, compile_date, run_date, io_date, minor_rev
+        char* commit_id
+        char* sha1sum
+        char* compile_machine
+        char* compile_flags
+
+    sdf_file_t *sdf_open(const char *filename, comm_t comm, int mode, int use_mmap)
+    bint sdf_close(sdf_file_t *h)
+    sdf_block_t *sdf_find_block_by_name(sdf_file_t *h, const char *name)
+    bint sdf_read_header(sdf_file_t *h)
+    bint sdf_read_blocklist_all(sdf_file_t *h)
+    bint sdf_read_block_info(sdf_file_t *h)
+    bint sdf_read_data(sdf_file_t *h)
+    bint sdf_get_domain_bounds(sdf_file_t *h, int rank,
+                                 int *starts, int *local_dims)
+    int sdf_block_set_array_section(sdf_block_t *b, const int ndims,
+                                       const int64_t *starts, const int64_t *ends,
+                                       const int64_t *strides)
+
+
+cdef extern from "sdf_helper.h":
+    bint sdf_helper_read_data(sdf_file_t *h, sdf_block_t *b)
+
+
+cdef extern from "stack_allocator.h":
+    void sdf_stack_destroy(sdf_file_t *h)
+    void sdf_stack_init(sdf_file_t *h)

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