sdf-xarray 0.4.0__cp314-cp314t-macosx_11_0_arm64.whl

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,71 @@
+import xarray as xr
+
+
+@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)

sdf_xarray/download.py

@@ -0,0 +1,87 @@
+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="doi:10.5281/zenodo.17618510",
+        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:1945ecdbc1ac1798164f83ea2b3d1b31",
+            "tutorial_dataset_2d_moving_window.zip": "md5:a795f40d18df69263842055de4559501",
+            "tutorial_dataset_3d.zip": "md5:d9254648867016292440fdb028f717f7",
+        },
+    )
+
+    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

sdf_xarray/plotting.py

@@ -0,0 +1,293 @@
+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
+
+from types import MethodType
+
+
+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
+
+    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 axis and associated units to the title
+    t_axis_value = data[t][frame].values
+
+    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_t_axis}{title_sdf}"
+
+
+def calculate_window_boundaries(
+    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.
+
+    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_axis_name].values
+    x_half_cell = (x_grid[1] - x_grid[0]) / 2
+    N_frames = data[t].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 is not None:
+        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.
+
+    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
+    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,
+    t: str | None = None,
+    ax: plt.Axes | None = None,
+    **kwargs,
+) -> FuncAnimation:
+    """Generate an animation using an xarray.DataArray
+
+    Parameters
+    ---------
+    data
+        DataArray containing the target data
+    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
+    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
+    --------
+    >>> ds["Derived_Number_Density_Electron"].epoch.animate()
+    """
+    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:
+        fig, ax = plt.subplots()
+        # Prevents figure from prematurely displaying in Jupyter notebook
+        plt.close(fig)
+
+    # 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
+
+    if data.ndim == 2:
+        kwargs.setdefault("x", coord_names[0])
+        plot = data.isel({t: 0}).plot(ax=ax, **kwargs)
+        ax.set_title(get_frame_title(data, 0, display_sdf_name, title, t))
+        global_min, global_max = compute_global_limits(
+            data, min_percentile, max_percentile
+        )
+        ax.set_ylim(global_min, global_max)
+
+    if data.ndim == 3:
+        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
+        # Set default x and y coordinates for 3D data if not provided
+        kwargs.setdefault("x", coord_names[0])
+        kwargs.setdefault("y", coord_names[1])
+
+        # Finds the time step with the minimum data value
+        # This is needed so that the animation can use the correct colour bar
+        argmin_time = np.unravel_index(data.argmin(), data.shape)[0]
+
+        # Initialize the plot, the final output will still start at the first time step
+        plot = data.isel({t: argmin_time}).plot(ax=ax, **kwargs)
+        ax.set_title(get_frame_title(data, 0, display_sdf_name, title, t))
+        kwargs["cmap"] = plot.cmap
+
+        # Add colorbar
+        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}]")
+
+    # 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"), kwargs["x"]
+        )
+
+    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()
+
+        plot = data.isel({t: frame}).plot(ax=ax, **kwargs)
+        ax.set_title(get_frame_title(data, frame, display_sdf_name, title, t))
+
+        if data.ndim == 2:
+            ax.set_ylim(global_min, global_max)
+        return plot
+
+    return FuncAnimation(
+        ax.get_figure(),
+        update,
+        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):
+        self._obj = xarray_obj
+
+    def animate(self, *args, **kwargs) -> FuncAnimation:
+        """Generate animations of Epoch data.
+
+        Parameters
+        ----------
+        args
+            Positional arguments passed to :func:`animation`.
+        kwargs
+            Keyword arguments passed to :func:`animation`.
+
+        Examples
+        --------
+        >>> anim = ds["Electric_Field_Ey"].epoch.animate()
+        >>> anim.save("myfile.mp4")
+        >>> # Or in a jupyter notebook:
+        >>> anim.show()
+        """
+
+        # 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