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

{sdf_xarray-0.3.1 → 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.1 → sdf_xarray-0.3.2}/PKG-INFO

@@ -1,8 +1,8 @@
 Metadata-Version: 2.4
 Name: sdf-xarray
-Version: 0.3.1
+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.1 → 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.1 → sdf_xarray-0.3.2}/src/sdf_xarray/__init__.py

@@ -84,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,
@@ -104,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`
 
@@ -135,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:
@@ -165,7 +218,6 @@ def open_mfdataset(
 
     return xr.combine_by_coords(
         all_dfs,
-        data_vars="all",
         coords="different",
         combine_attrs="drop_conflicts",
         join="outer",
@@ -523,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:
@@ -537,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.1 → 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.1'
-__version_tuple__ = version_tuple = (0, 3, 1)
+__version__ = version = '0.3.2'
+__version_tuple__ = version_tuple = (0, 3, 2)
 
-__commit_id__ = commit_id = 'gce5426d4a'
+__commit_id__ = commit_id = 'g331520e50'

{sdf_xarray-0.3.1 → sdf_xarray-0.3.2}/tests/test_basic.py

@@ -5,7 +5,12 @@ import numpy.testing as npt
 import pytest
 import xarray as xr
 
-from sdf_xarray import SDFPreprocess, _process_latex_name, _resolve_glob, open_mfdataset
+from sdf_xarray import (
+    SDFPreprocess,
+    _process_latex_name,
+    _resolve_glob,
+    open_mfdataset,
+)
 
 EXAMPLE_FILES_DIR = pathlib.Path(__file__).parent / "example_files_1D"
 EXAMPLE_MISMATCHED_FILES_DIR = (
@@ -454,3 +459,229 @@ def test_xr_oading_one_probe_drop_second_probe():
         assert "X_Probe_Electron_Front_Probe" in df.coords
         assert "ID_Electron_Front_Probe_Px" in df.dims
         assert "ID_Electron_Back_Probe_Px" not in df.dims
+
+
+def test_open_mfdataset_data_vars_single():
+    with open_mfdataset(
+        EXAMPLE_FILES_DIR.glob("*.sdf"),
+        data_vars=["Electric_Field_Ex"],
+    ) as df:
+        ex_field = "Electric_Field_Ex"
+        x_coord = "X_Grid_mid"
+        assert ex_field in df
+        assert x_coord in df[ex_field].coords
+        assert "time" in df[ex_field].coords
+        assert df[x_coord].attrs["long_name"] == "X"
+
+        assert "Electric_Field_Ey" not in df
+
+
+def test_open_mfdataset_data_vars_multiple():
+    with open_mfdataset(
+        EXAMPLE_FILES_DIR.glob("*.sdf"),
+        data_vars=["Electric_Field_Ex", "Electric_Field_Ey"],
+    ) as df:
+        ex_field = "Electric_Field_Ex"
+        x_coord = "X_Grid_mid"
+        assert ex_field in df
+        assert x_coord in df[ex_field].coords
+        assert "time" in df[ex_field].coords
+        assert df[x_coord].attrs["long_name"] == "X"
+
+        ey_field = "Electric_Field_Ey"
+        x_coord = "X_Grid_mid"
+        assert ey_field in df
+        assert x_coord in df[ey_field].coords
+        assert "time" in df[ey_field].coords
+        assert df[x_coord].attrs["long_name"] == "X"
+
+
+def test_open_mfdataset_data_vars_sparse_multiple():
+    with open_mfdataset(
+        EXAMPLE_FILES_DIR.glob("*.sdf"),
+        keep_particles=True,
+        data_vars=[
+            "Particles_Particles_Per_Cell_proton",
+            "Electric_Field_Ez",
+            "dist_fn_x_px_proton",
+        ],
+    ) as df:
+        ppc_proton = "Particles_Particles_Per_Cell_proton"
+        assert ppc_proton in df
+        assert "time" in df[ppc_proton].coords
+        assert (
+            df[ppc_proton].attrs["long_name"] == "Particles Particles Per Cell proton"
+        )
+
+        ez_field = "Electric_Field_Ez"
+        assert ez_field in df
+        assert len(df[ez_field].coords) == 2
+        assert "time" in df[ez_field].coords
+        assert "X_Grid_mid" in df[ez_field].coords
+        assert df[ez_field].attrs["long_name"] == "Electric Field $E_z$"
+
+        dist_fn = "dist_fn_x_px_proton"
+        assert dist_fn in df
+        assert len(df[dist_fn].coords) == 3
+        assert "time" in df[dist_fn].coords
+        assert "X_x_px_proton" in df[dist_fn].coords
+        assert "Px_x_px_proton" in df[dist_fn].coords
+
+        assert df["time"].size == 11
+
+
+def test_open_mfdataset_data_vars_invalid_var():
+    with open_mfdataset(
+        EXAMPLE_FILES_DIR.glob("*.sdf"),
+        data_vars=["Electric_Field"],
+    ) as df:
+        assert len(df.variables.keys()) == 1
+        assert df["time"].size == 11
+
+
+def test_open_mfdataset_data_vars_time():
+    with open_mfdataset(
+        EXAMPLE_FILES_DIR.glob("*.sdf"),
+        data_vars=["Electric_Field_Ex"],
+    ) as df:
+        time = df["time"]
+        assert time.units == "s"
+        assert time.long_name == "Time"
+        assert time.full_name == "time"
+
+        time_values = np.array(
+            [
+                5.466993e-14,
+                2.417504e-10,
+                4.833915e-10,
+                7.251419e-10,
+                9.667830e-10,
+                1.208533e-09,
+                1.450175e-09,
+                1.691925e-09,
+                1.933566e-09,
+                2.175316e-09,
+                2.416958e-09,
+            ]
+        )
+
+        npt.assert_allclose(time_values, time.values, rtol=1e-6)
+
+
+def test_open_mfdataset_data_vars_sparse_time():
+    with open_mfdataset(
+        EXAMPLE_FILES_DIR.glob("*.sdf"),
+        data_vars=["Particles_Particles_Per_Cell_proton"],
+    ) as df:
+        time = df["time"]
+        assert time.units == "s"
+        assert time.long_name == "Time"
+        assert time.full_name == "time"
+
+        time_values = np.array(
+            [
+                5.466993e-14,
+                2.417504e-10,
+                4.833915e-10,
+                7.251419e-10,
+                9.667830e-10,
+                1.208533e-09,
+                1.450175e-09,
+                1.691925e-09,
+                1.933566e-09,
+                2.175316e-09,
+                2.416958e-09,
+            ]
+        )
+
+        npt.assert_allclose(time_values, time.values, rtol=1e-6)
+
+
+def test_open_mfdataset_data_vars_separate_times_single():
+    with open_mfdataset(
+        EXAMPLE_FILES_DIR.glob("*.sdf"),
+        data_vars=["Electric_Field_Ex"],
+        separate_times=True,
+    ) as df:
+        coords = df.coords.sizes
+        assert len(coords) == 2
+        assert coords["time0"] == 11
+        assert coords["X_Grid_mid"] == 16
+
+        elec_x = "Electric_Field_Ex"
+        elec_x_coords = df[elec_x].coords.sizes
+        assert elec_x in df
+        assert len(elec_x_coords) == 2
+        assert "time0" in elec_x_coords
+        assert "X_Grid_mid" in elec_x_coords
+
+        assert elec_x_coords["time0"] == 11
+        assert elec_x_coords["X_Grid_mid"] == 16
+
+
+def test_open_mfdataset_data_vars_separate_times_multiple():
+    with open_mfdataset(
+        EXAMPLE_FILES_DIR.glob("*.sdf"),
+        data_vars=["Electric_Field_Ex", "Electric_Field_Ey"],
+        separate_times=True,
+    ) as df:
+        coords = df.coords.sizes
+        assert len(coords) == 2
+        assert coords["time0"] == 11
+        assert coords["X_Grid_mid"] == 16
+
+        elec_x = "Electric_Field_Ex"
+        elec_x_coords = df[elec_x].coords.sizes
+        assert elec_x in df
+        assert len(elec_x_coords) == 2
+        assert "time0" in elec_x_coords
+        assert "X_Grid_mid" in elec_x_coords
+
+        assert elec_x_coords["time0"] == 11
+        assert elec_x_coords["X_Grid_mid"] == 16
+
+        elec_y = "Electric_Field_Ey"
+        elec_y_coords = df[elec_y].coords.sizes
+        assert elec_y in df
+        assert len(elec_y_coords) == 2
+        assert "time0" in elec_y_coords
+        assert "X_Grid_mid" in elec_y_coords
+
+        assert elec_y_coords["time0"] == 11
+        assert elec_y_coords["X_Grid_mid"] == 16
+
+
+def test_open_mfdataset_data_vars_separate_times_multiple_times_keep_particles():
+    with open_mfdataset(
+        EXAMPLE_FILES_DIR.glob("*.sdf"),
+        data_vars=["Electric_Field_Ex", "Particles_Px_electron_beam"],
+        separate_times=True,
+        keep_particles=True,
+    ) as df:
+        coords = df.coords.sizes
+        assert len(coords) == 5
+        assert coords["time0"] == 11
+        assert coords["time1"] == 1
+        assert coords["time2"] == 1
+        assert coords["X_Grid_mid"] == 16
+        assert coords["ID_electron_beam"] == 1440
+
+        elec_x = "Electric_Field_Ex"
+        elec_x_coords = df[elec_x].coords.sizes
+        assert elec_x in df
+        assert len(elec_x_coords) == 2
+        assert "time0" in elec_x_coords
+        assert "X_Grid_mid" in elec_x_coords
+
+        assert elec_x_coords["time0"] == 11
+        assert elec_x_coords["X_Grid_mid"] == 16
+
+        particle_px = "Particles_Px_electron_beam"
+        particle_px_coords = df[particle_px].coords.sizes
+        assert particle_px in df
+        assert len(particle_px_coords) == 2
+        assert "time2" in particle_px_coords
+        assert "ID_electron_beam" in particle_px_coords
+
+        assert particle_px_coords["time2"] == 1
+        assert particle_px_coords["ID_electron_beam"] == 1440