sdf-xarray 0.3.0__tar.gz → 0.3.2__tar.gz

{sdf_xarray-0.3.0 → sdf_xarray-0.3.2}/CITATION.cff

@@ -16,5 +16,9 @@ authors:
     given-names: Shaun
     orcid: 'https://orcid.org/0009-0005-0693-030X'
     affiliation: University of York
+  - family-names: Herdman
+    given-names: Chris
+    orcid: 'https://orcid.org/0000-0002-5159-0130'
+    affiliation: University of York
 doi: 10.5281/zenodo.15351323
 date-released: '2024-07-25'

{sdf_xarray-0.3.0 → sdf_xarray-0.3.2}/PKG-INFO

@@ -1,8 +1,8 @@
 Metadata-Version: 2.4
 Name: sdf-xarray
-Version: 0.3.0
+Version: 0.3.2
 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>
+Author-Email: Peter Hill <peter.hill@york.ac.uk>, Joel Adams <joel.adams@york.ac.uk>, Shaun Doherty <shaun.doherty@york.ac.uk>, Chris Herdman <chris.herdman@york.ac.uk>
 License-Expression: BSD-3-Clause
 Classifier: Development Status :: 5 - Production/Stable
 Classifier: Intended Audience :: Science/Research

sdf_xarray-0.3.2/docs/tutorial_dataset_2d/input.deck

@@ -0,0 +1,65 @@
+begin:control
+  nx = 64
+  ny = 64
+
+  # Final time of simulation
+  t_end = 5 * femto
+
+  # Size of domain
+  x_min = 0
+  x_max = 6 * micron
+
+  y_min = 0
+  y_max = 6 * micron
+
+  stdout_frequency = 1
+  nparticles = nx * ny * 50
+end:control
+
+begin:constant
+    n_elec = 1000
+    L_target_x = 2 * micron
+    L_target_y = 4 * micron
+
+    x_center = (x_min + x_max) / 2
+    y_center = (y_min + y_max) / 2
+
+    density_profile = if( (abs(x - x_center) lt L_target_x/2), if( (abs(y - y_center) lt L_target_y/2), 1, 0), 0)
+end:constant
+
+begin:boundaries
+  bc_x_min = periodic
+  bc_x_max = periodic
+  bc_y_min = periodic
+  bc_y_max = periodic
+end:boundaries
+
+
+begin:species
+  name = Electron
+  frac = 0.5
+  number_density = n_elec * density_profile
+  identify:electron
+end:species
+
+begin:species
+  name = Ion
+  frac = 0.5
+  number_density = n_elec * density_profile
+  identify:proton
+end:species
+
+begin:output_global
+    force_last_to_be_restartable = F
+end:output_global
+
+begin:output
+  name = normal
+
+  dt_snapshot = 1 * femto
+
+  grid = always
+  number_density = always + species
+
+end:output
+

sdf_xarray-0.3.2/docs/unit_conversion.rst

@@ -0,0 +1,228 @@
+.. _sec-unit-conversion:
+
+===============
+Unit Conversion
+===============
+
+The ``sdf-xarray`` package automatically extracts the units for each
+coordinate/variable/constant from an SDF file and stores them as an :class:`xarray.Dataset`
+attribute called ``"units"``. Sometimes we want to convert our data from one format to
+another, e.g. converting the grid coordinates from meters to microns, time from seconds 
+to femto-seconds or particle energy from Joules to electron-volts.
+
+.. ipython:: python
+
+    from sdf_xarray import open_mfdataset
+    import matplotlib.pyplot as plt
+    plt.rcParams.update({
+        "axes.labelsize": 16,
+        "xtick.labelsize": 14,
+        "ytick.labelsize": 14,
+        "axes.titlesize": 16
+    })
+
+
+=====================
+Rescaling Coordinates
+=====================
+
+For simple scaling and unit relabeling of coordinates (e.g., converting meters to microns),
+the most straightforward approach is to use the ``rescale_coords()`` method  
+via the custom ``xarray.Dataset.epoch`` dataset accessor.  
+
+This method scales the coordinate values by a given multiplier and updates the ``"units"``
+attribute in one step.
+
+Rescaling Grid Coordinates
+~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+We can use the ``xarray.Dataset.epoch.rescale_coords()`` method to convert X, Y, and Z
+coordinates from meters (m) to microns (µm) by applying a multiplier of ``1e6``.
+
+.. ipython:: python
+
+    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(16, 6))
+
+    with open_mfdataset("tutorial_dataset_2d/*.sdf") as ds:
+        ds_in_microns = ds.epoch.rescale_coords(
+            multiplier=1e6, 
+            unit_label="µm", 
+            coord_names=["X_Grid_mid", "Y_Grid_mid"]
+        )
+        derived_number_density = ds["Derived_Number_Density_Electron"].isel(time=0).compute()
+        derived_number_density_microns = ds_in_microns["Derived_Number_Density_Electron"].isel(time=0).compute()
+
+    derived_number_density.plot(ax=ax1, x="X_Grid_mid", y="Y_Grid_mid")
+    ax1.set_title("Original X Coordinate (m)")
+
+    derived_number_density_microns.plot(ax=ax2, x="X_Grid_mid", y="Y_Grid_mid")
+    ax2.set_title("Rescaled X Coordinate (µm)")
+
+    @savefig coordinate_conversion.png width=9in
+    fig.tight_layout()
+
+
+Rescaling Time Coordinate
+~~~~~~~~~~~~~~~~~~~~~~~~~
+
+We can also use the ``xarray.Dataset.epoch.rescale_coords()`` method to convert the time
+coordinate from seconds (s) to femto-seconds (fs) by applying a multiplier of ``1e15``.
+
+.. ipython:: python
+    
+    with open_mfdataset("tutorial_dataset_2d/*.sdf") as ds:
+        ds_time_in_femto = ds.epoch.rescale_coords(
+            multiplier=1e15, 
+            unit_label="fs", 
+            coord_names="time"
+        )
+
+        print(f"[Original] units: {ds['time'].attrs['units']}, values: {ds['time'].values}")
+        print(f"[Rescaled] units: {ds_time_in_femto['time'].attrs['units']}, values: {ds_time_in_femto['time'].values}")
+
+
+================================
+Unit Conversion with pint-xarray
+================================
+
+While this is sufficient for most use cases, we can enhance this functionality
+using the `pint <https://pint.readthedocs.io/en/stable/getting/index.html>`_ library.
+Pint allows us to specify the units of a given array and convert them
+to another, which is incredibly handy. We can take this a step further,
+however, and utilize the `pint-xarray
+<https://pint-xarray.readthedocs.io/en/latest/>`_ library. This library
+allows us to infer units directly from an `xarray.Dataset.attrs` while
+retaining all the information about the `xarray.Dataset`. This works
+very similarly to taking a NumPy array and multiplying it by a constant or
+another array, which returns a new array; however, this library will also
+retain the unit logic (specifically the ``"units"`` information).
+
+.. note::
+    Unit conversion is not supported on coordinates in ``pint-xarray`` which is due to an
+    underlying issue with how ``xarray`` implements indexes.
+
+Installation
+~~~~~~~~~~~~
+
+To install the pint libraries you can simply run the following optional
+dependency pip command which will install both the ``pint`` and ``pint-xarray``
+libraries. You can install these optional dependencies via pip:
+
+.. code:: console
+
+    $ pip install "sdf_xarray[pint]"
+
+.. note::
+    Once you install ``pint-xarray`` it is automatically picked up and loaded
+    by the code so you should have access to the ``xarray.Dataset.pint`` accessor.
+
+Quantifying Arrays
+~~~~~~~~~~~~~~~~~~
+
+When using ``pint-xarray``, the library attempts to infer units from the
+``"units"`` attribute on each `xarray.DataArray`. Alternatively, you can
+also specify the units yourself by passing a string into the 
+``xarray.Dataset.DataArray.pint.quantify()`` function call. Once the type is inferred
+the original `xarray.DataArray` will be converted to a `pint.Quantity`
+and the ``"units"`` attribute will
+be removed.
+
+In the following example we will extract the time-resolved total particle
+energy of electrons which is measured in Joules and convert it to electron
+volts.
+
+.. ipython:: python
+
+    with open_mfdataset("tutorial_dataset_1d/*.sdf") as ds:
+        total_particle_energy = ds["Total_Particle_Energy_Electron"]
+
+    total_particle_energy
+
+    total_particle_energy = ds["Total_Particle_Energy_Electron"].pint.quantify()
+
+    total_particle_energy
+
+
+Now that this dataset has been converted a `pint.Quantity`, we can check
+it's units and dimensionality
+
+.. ipython:: python
+
+    total_particle_energy.pint.units
+    total_particle_energy.pint.dimensionality
+
+
+Converting Units
+~~~~~~~~~~~~~~~~
+
+We can now convert it to electron volts utilising the `pint.Quantity.to`
+function
+
+.. ipython:: python
+
+    total_particle_energy_ev = total_particle_energy.pint.to("eV")
+
+Unit Propagation
+~~~~~~~~~~~~~~~~
+
+Suppose instead of converting to ``"eV"``, we want to convert to ``"W"``
+(watts). To do this, we divide the total particle energy by time. However,
+since coordinates in `xarray.Dataset` cannot be directly converted to
+`pint.Quantity`, we must first extract the coordinate values manually
+and create a new Pint quantity for time.
+
+Once both arrays are quantified, Pint will automatically handle the unit
+propagation when we perform arithmetic operations like division.
+
+.. note::
+    Pint does not automatically simplify ``"J/s"`` to ``"W"``, so we use
+    `pint.Quantity.to` to convert the unit string. Since these units are
+    the same it will not change the underlying data, only the units. This is
+    only a small formatting choice and is not required.
+
+.. ipython:: python
+
+    import pint
+    time_values = total_particle_energy.coords["time"].data
+    time = pint.Quantity(time_values, "s")
+    total_particle_energy_w = total_particle_energy / time
+    total_particle_energy_w.pint.units
+    total_particle_energy_w = total_particle_energy_w.pint.to("W")
+    total_particle_energy_w.pint.units
+
+Dequantifying and Restoring Units
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. note::
+    If this function is not called prior to plotting then the ``units`` will be
+    inferred from the `pint.Quantity` array which will return the long
+    name of the units. i.e. instead of returning ``"eV"`` it will return
+    ``"electron_volt"``.
+
+The ``xarray.Dataset.DataArray.pint.dequantify`` function converts the data from
+`pint.Quantity` back to the original `xarray.DataArray` and adds
+the ``"units"`` attribute back in. It also has an optional ``format`` parameter
+that allows you to specify the formatting type of ``"units"`` attribute. We
+have used the ``format="~P"`` option as it shortens the unit to its
+"short pretty" format (``"eV"``). For more options, see the `Pint formatting
+documentation <https://pint.readthedocs.io/en/stable/user/formatting.html>`_.
+
+.. ipython:: python
+
+    total_particle_energy_ev = total_particle_energy_ev.pint.dequantify(format="~P")
+    total_particle_energy_w = total_particle_energy_w.pint.dequantify(format="~P")
+    total_particle_energy_ev
+
+To confirm the conversion has worked correctly, we can plot the original and
+converted `xarray.Dataset` side by side:
+
+.. ipython:: python
+    
+    fig, ((ax1, ax2), (ax3, ax4)) = plt.subplots(2, 2, figsize=(16,8))
+    ds["Total_Particle_Energy_Electron"].plot(ax=ax1)
+    total_particle_energy_ev.plot(ax=ax2)
+    total_particle_energy_w.plot(ax=ax3)
+    ax4.set_visible(False)
+    fig.suptitle("Comparison of conversion from Joules to electron volts and watts", fontsize="18")
+    @savefig unit_conversion.png width=9in
+    fig.tight_layout()

{sdf_xarray-0.3.0 → sdf_xarray-0.3.2}/pyproject.toml

@@ -16,6 +16,7 @@ authors = [
   { name = "Peter Hill", email = "peter.hill@york.ac.uk" },
   { name = "Joel Adams", email = "joel.adams@york.ac.uk" },
   { name = "Shaun Doherty", email = "shaun.doherty@york.ac.uk" },
+  { name = "Chris Herdman", email = "chris.herdman@york.ac.uk" },
 ]
 requires-python = ">=3.10,<3.14"
 dependencies = ["numpy>=2.0.0", "xarray>=2024.1.0", "dask>=2024.7.1"]

{sdf_xarray-0.3.0 → sdf_xarray-0.3.2}/src/sdf_xarray/__init__.py

@@ -1,3 +1,4 @@
+import contextlib
 import os
 import re
 from collections import Counter, defaultdict
@@ -18,10 +19,16 @@ from xarray.core import indexing
 from xarray.core.utils import close_on_error, try_read_magic_number_from_path
 from xarray.core.variable import Variable
 
-# NOTE: Do not delete this line, otherwise the "epoch" accessor will not be
-# imported when the user imports sdf_xarray
+# NOTE: Do not delete these lines, otherwise the "epoch" dataset and dataarray
+# accessors will not be imported when the user imports sdf_xarray
+import sdf_xarray.dataset_accessor
 import sdf_xarray.plotting  # noqa: F401
 
+# NOTE: This attempts to initialise with the "pint" accessor if the user
+# has installed the package
+with contextlib.suppress(ImportError):
+    import pint_xarray  # noqa: F401
+
 from .sdf_interface import Constant, SDFFile  # type: ignore  # noqa: PGH003
 
 # TODO Remove this once the new kwarg options are fully implemented
@@ -77,8 +84,45 @@ def _resolve_glob(path_glob: PathLike | Iterable[PathLike]):
     return paths
 
 
-def combine_datasets(path_glob: Iterable | str, **kwargs) -> xr.Dataset:
-    """Combine all datasets using a single time dimension"""
+def purge_unselected_data_vars(ds: xr.Dataset, data_vars: list[str]) -> xr.Dataset:
+    """
+    If the user has exclusively requested only certain variables be
+    loaded in then we purge all other variables and dimensions
+    """
+    existing_data_vars = set(ds.data_vars.keys())
+    vars_to_keep = set(data_vars) & existing_data_vars
+    vars_to_drop = existing_data_vars - vars_to_keep
+    ds = ds.drop_vars(vars_to_drop)
+
+    existing_dims = set(ds.sizes)
+    dims_to_keep = set()
+    for var in vars_to_keep:
+        dims_to_keep.update(ds[var].coords._names)
+        dims_to_keep.update(ds[var].dims)
+
+    coords_to_drop = existing_dims - dims_to_keep
+    return ds.drop_dims(coords_to_drop)
+
+
+def combine_datasets(
+    path_glob: Iterable | str, data_vars: list[str], **kwargs
+) -> xr.Dataset:
+    """
+    Combine all datasets using a single time dimension, optionally extract
+    data from only the listed data_vars
+    """
+
+    if data_vars is not None:
+        return xr.open_mfdataset(
+            path_glob,
+            join="outer",
+            coords="different",
+            compat="no_conflicts",
+            combine="nested",
+            concat_dim="time",
+            preprocess=SDFPreprocess(data_vars=data_vars),
+            **kwargs,
+        )
 
     return xr.open_mfdataset(
         path_glob,
@@ -97,6 +141,7 @@ def open_mfdataset(
     separate_times: bool = False,
     keep_particles: bool = False,
     probe_names: list[str] | None = None,
+    data_vars: list[str] | None = None,
 ) -> xr.Dataset:
     """Open a set of EPOCH SDF files as one `xarray.Dataset`
 
@@ -128,19 +173,34 @@ def open_mfdataset(
         If ``True``, also load particle data (this may use a lot of memory!)
     probe_names :
         List of EPOCH probe names
+    data_vars :
+        List of data vars to load in (If not specified loads in all variables)
     """
 
     path_glob = _resolve_glob(path_glob)
+
     if not separate_times:
         return combine_datasets(
-            path_glob, keep_particles=keep_particles, probe_names=probe_names
+            path_glob,
+            data_vars=data_vars,
+            keep_particles=keep_particles,
+            probe_names=probe_names,
         )
 
     _, var_times_map = make_time_dims(path_glob)
-    all_dfs = [
-        xr.open_dataset(f, keep_particles=keep_particles, probe_names=probe_names)
-        for f in path_glob
-    ]
+
+    all_dfs = []
+    for f in path_glob:
+        ds = xr.open_dataset(f, keep_particles=keep_particles, probe_names=probe_names)
+
+        # If the data_vars are specified then only load them in and disregard the rest.
+        # If there are no remaining data variables then skip adding the dataset to list
+        if data_vars is not None:
+            ds = purge_unselected_data_vars(ds, data_vars)
+            if not ds.data_vars:
+                continue
+
+        all_dfs.append(ds)
 
     for df in all_dfs:
         for da in df:
@@ -158,7 +218,6 @@ def open_mfdataset(
 
     return xr.combine_by_coords(
         all_dfs,
-        data_vars="all",
         coords="different",
         combine_attrs="drop_conflicts",
         join="outer",
@@ -516,10 +575,43 @@ class SDFEntrypoint(BackendEntrypoint):
 
 
 class SDFPreprocess:
-    """Preprocess SDF files for xarray ensuring matching job ids and sets time dimension"""
+    """Preprocess SDF files for xarray ensuring matching job ids and sets
+    time dimension.
+
+    This class is used as a 'preprocess' function within ``xr.open_mfdataset``. It
+    performs three main duties on each individual file's Dataset:
+
+    1. Checks for a **matching job ID** across all files to ensure dataset consistency.
+    2. **Filters** the Dataset to keep only the variables specified in `data_vars`
+       and their required coordinates.
+    3. **Expands dimensions** to include a single 'time' coordinate, preparing the
+       Dataset for concatenation.
+
+    EPOCH can output variables at different intervals, so some SDF files
+    may not contain the requested variable. We combine this data into one
+    dataset by concatenating across the time dimension.
+
+    The combination is performed using ``join="outer"`` (in the calling ``open_mfdataset`` function),
+    meaning that the final combined dataset will contain the variable across the
+    entire time span, with NaNs filling the time steps where the variable was absent in
+    the individual file.
+
+    With large SDF files, this filtering method will save on memory consumption when
+    compared to loading all variables from all files before concatenation.
 
-    def __init__(self):
+    Parameters
+    ----------
+    data_vars :
+        A list of data variables to load in (If not specified loads
+        in all variables)
+    """
+
+    def __init__(
+        self,
+        data_vars: list[str] | None = None,
+    ):
         self.job_id: int | None = None
+        self.data_vars = data_vars
 
     def __call__(self, ds: xr.Dataset) -> xr.Dataset:
         if self.job_id is None:
@@ -530,17 +622,23 @@ class SDFPreprocess:
                 f"Mismatching job ids (got {ds.attrs['jobid1']}, expected {self.job_id})"
             )
 
-        ds = ds.expand_dims(time=[ds.attrs["time"]])
+        # If the user has exclusively requested only certain variables be
+        # loaded in then we purge all other variables and coordinates
+        if self.data_vars:
+            ds = purge_unselected_data_vars(ds, self.data_vars)
+
+        time_val = ds.attrs.get("time", np.nan)
+        ds = ds.expand_dims(time=[time_val])
         ds = ds.assign_coords(
             time=(
                 "time",
-                [ds.attrs["time"]],
+                [time_val],
                 {"units": "s", "long_name": "Time", "full_name": "time"},
             )
         )
         # Particles' spartial coordinates also evolve in time
         for coord, value in ds.coords.items():
             if value.attrs.get("point_data", False):
-                ds.coords[coord] = value.expand_dims(time=[ds.attrs["time"]])
+                ds.coords[coord] = value.expand_dims(time=[time_val])
 
         return ds

{sdf_xarray-0.3.0 → sdf_xarray-0.3.2}/src/sdf_xarray/_version.py

@@ -28,7 +28,7 @@ version_tuple: VERSION_TUPLE
 commit_id: COMMIT_ID
 __commit_id__: COMMIT_ID
 
-__version__ = version = '0.3.0'
-__version_tuple__ = version_tuple = (0, 3, 0)
+__version__ = version = '0.3.2'
+__version_tuple__ = version_tuple = (0, 3, 2)
 
-__commit_id__ = commit_id = 'gcca942b3f'
+__commit_id__ = commit_id = 'g331520e50'

sdf_xarray-0.3.2/src/sdf_xarray/dataset_accessor.py

@@ -0,0 +1,73 @@
+from typing import Union
+
+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: Union[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-0.3.2/tests/example_files_3D/input.deck

@@ -0,0 +1,54 @@
+begin:control
+  nx = 64
+  ny = 64
+  nz = 64
+
+  # Final time of simulation
+  t_end = 5 * femto
+
+  # Size of domain
+  x_min = 0
+  x_max = 5.0e5
+
+  y_min = 0
+  y_max = 5.0e5
+
+  z_min = 0
+  z_max = 5.0e5
+
+  stdout_frequency = 1
+end:control
+
+
+begin:boundaries
+  bc_x_min = periodic
+  bc_x_max = periodic
+  bc_y_min = periodic
+  bc_y_max = periodic
+  bc_z_min = periodic
+  bc_z_max = periodic
+end:boundaries
+
+
+begin:species
+  name = electron
+  charge = -1
+  mass = 1.0
+  temperature_x = 273
+  number_density = 10
+  nparticles = nx * ny * nz * 2
+end:species
+
+begin:output_global
+    force_last_to_be_restartable = F
+end:output_global
+
+begin:output
+  name = normal
+
+  dt_snapshot = 1 * femto
+
+  grid = always
+
+end:output
+