Feature: optional loading of an EPOCH input deck

CITATION.cff

@@ -24,5 +24,9 @@ authors:
     given-names: Liam
     orcid: 'https://orcid.org/0000-0001-8604-6904'
     affiliation: University of York
+  - family-names: Shekhanov
+    given-names: Sviatoslav
+    orcid: 'https://orcid.org/0000-0002-2125-8962'
+    affiliation: University of York
 doi: 10.5281/zenodo.15351323
 date-released: '2024-07-25'

docs/conf.py

@@ -50,6 +50,7 @@
     "myst_parser",
     "jupyter_sphinx",
     "sphinx_copybutton",
+    "sphinx_togglebutton",
 ]
 
 autosummary_generate = True

docs/key_functionality.rst

@@ -161,6 +161,47 @@ consumption.
 
    sdfxr.open_mfdataset("tutorial_dataset_1d/*.sdf", data_vars=["Electric_Field_Ex"])
 
+.. _loading-input-deck:
+
+Loading the input.deck
+~~~~~~~~~~~~~~~~~~~~~~
+
+When loading SDF files, `sdf_xarray` will attempt to automatically load
+the ``input.deck`` file used to initialise the simulation from the same
+directory as the SDF file. If the file is not found, it will silently fail
+and continue loading the SDF file as normal. This file contains the initial
+simulation setup information which is not present in SDF outputs. By loading
+this file, you can access these parameters as part of your dataset's metadata.
+To do this, use the ``deck_path`` parameter when loading an SDF file with
+`xarray.open_dataset`, `sdf_xarray.open_datatree`, `sdf_xarray.open_mfdataset`
+or `sdf_xarray.open_mfdatatree`.
+
+There are a few ways you can load an input deck:
+
+- **Default behaviour**: The input deck is loaded from the same directory
+  as the SDF file if it exists. If it does not exist, it will silently fail.
+- **Relative path**: (e.g. ``"template.deck"``) Searches for that specific filename
+  within the same directory as the SDF file.
+- **Absolute path**: (e.g. ``"/path/to/input.deck"``) Uses the full, specified path
+  to locate the file.
+
+An example of loading a deck can be seen below
+
+.. toggle::
+
+   .. jupyter-execute::
+
+      import json
+      from IPython.display import Code
+
+      ds = xr.open_dataset("tutorial_dataset_1d/0010.sdf")
+      # The results are accessible by calling 
+      deck = ds.attrs["deck"]
+
+      # Some prettification to make it looks nice in jupyter notebooks
+      json_str = json.dumps(deck, indent=4)
+      Code(json_str, language='json')
+
 Data interaction examples
 -------------------------
 

pyproject.toml

@@ -18,9 +18,15 @@ authors = [
   { name = "Shaun Doherty", email = "shaun.doherty@york.ac.uk" },
   { name = "Chris Herdman", email = "chris.herdman@york.ac.uk" },
   { name = "Liam Pattinson", email = "liam.pattinson@york.ac.uk" },
+  { name = "Sviatoslav Shekhanov", email = "sviatoslav.shekhanov@york.ac.uk" },
 ]
 requires-python = ">=3.11,<3.15"
-dependencies = ["numpy>=2.0.0", "xarray>=2024.1.0", "dask>=2024.7.1"]
+dependencies = [
+  "numpy>=2.0.0",
+  "xarray>=2024.1.0",
+  "dask>=2024.7.1",
+  "epydeck~=1.0"
+]
 description = "Provides a backend for xarray to read SDF files as created by the EPOCH plasma PIC code."
 classifiers = [
   "Development Status :: 5 - Production/Stable",
@@ -49,6 +55,7 @@ dev = [
 docs = [
   "sphinx>=5.3",
   "sphinx_autodoc_typehints>=1.19",
+  "sphinx-togglebutton",
   "sphinx-book-theme>=0.4.0rc1",
   "sphinx-argparse-cli>=1.10.0",
   "sphinx-inline-tabs",

src/sdf_xarray/__init__.py

@@ -9,6 +9,7 @@
 from pathlib import Path
 from typing import ClassVar
 
+import epydeck
 import numpy as np
 import xarray as xr
 from packaging.version import Version
@@ -46,6 +47,32 @@ def _rename_with_underscore(name: str) -> str:
     return name.replace("/", "_").replace(" ", "_").replace("-", "_")
 
 
+def _load_deck(
+    root_dir: PathLike,
+    filename: PathLike | None,
+) -> dict:
+    """Load and attach an EPOCH input deck to the dataset.
+
+    A provided filename is resolved relative to the SDF file directory and must
+    exist, otherwise a FileNotFoundError is raised. If no filename is given, a
+    default ``input.deck`` is searched for and silently ignored if missing.
+
+    When found, the parsed deck is stored in ``ds.attrs["deck"]``.
+    """
+
+    root_dir = Path(root_dir).parent
+    target = Path("input.deck") if filename is None else Path(filename)
+    deck_path = target if target.is_absolute() else root_dir / target
+
+    if not deck_path.exists():
+        if filename is not None:
+            raise FileNotFoundError(f"Deck file not found: {deck_path}")
+        return {}
+
+    with deck_path.open() as f:
+        return epydeck.load(f)
+
+
 def _process_latex_name(variable_name: str) -> str:
     """Converts variable names to LaTeX format where possible
     using the following rules:
@@ -165,15 +192,18 @@ def purge_unselected_data_vars(ds: xr.Dataset, data_vars: list[str]) -> xr.Datas
 
 
 def combine_datasets(
-    path_glob: Iterable | str, data_vars: list[str], **kwargs
+    path_glob: Iterable | str,
+    data_vars: list[str] | None = None,
+    deck_path: PathLike | None = None,
+    **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(
+        ds = xr.open_mfdataset(
             path_glob,
             join="outer",
             coords="different",
@@ -183,16 +213,20 @@ def combine_datasets(
             preprocess=SDFPreprocess(data_vars=data_vars),
             **kwargs,
         )
+    else:
+        ds = xr.open_mfdataset(
+            path_glob,
+            data_vars="all",
+            coords="different",
+            compat="no_conflicts",
+            join="outer",
+            preprocess=SDFPreprocess(),
+            **kwargs,
+        )
 
-    return xr.open_mfdataset(
-        path_glob,
-        data_vars="all",
-        coords="different",
-        compat="no_conflicts",
-        join="outer",
-        preprocess=SDFPreprocess(),
-        **kwargs,
-    )
+    ds.attrs["deck"] = _load_deck(ds.attrs["filename"], deck_path)
+
+    return ds
 
 
 def open_mfdataset(
@@ -203,6 +237,7 @@ def open_mfdataset(
     probe_names: list[str] | None = None,
     data_vars: list[str] | None = None,
     chunks: T_Chunks = "auto",
+    deck_path: PathLike | None = None,
 ) -> xr.Dataset:
     """Open a set of EPOCH SDF files as one `xarray.Dataset`
 
@@ -244,6 +279,10 @@ def open_mfdataset(
         <https://docs.xarray.dev/en/stable/user-guide/dask.html#chunking-and-performance>`_
         for details on why this is useful for large datasets. The default behaviour is
         to do this automatically and can be disabled by ``chunks=None``.
+    deck_path :
+        If ``None``, attempt to load the ``"input.deck"`` from the same directory as the SDF files
+        and silently fail if it does not exist. If a path is given, load the specified deck
+        from a relative or absolute file path. See :ref:`loading-input-deck` for details.
     """
 
     path_glob = _resolve_glob(path_glob)
@@ -255,14 +294,19 @@ def open_mfdataset(
             keep_particles=keep_particles,
             probe_names=probe_names,
             chunks=chunks,
+            deck_path=deck_path,
         )
 
     _, var_times_map = make_time_dims(path_glob)
 
     all_dfs = []
     for f in path_glob:
         ds = xr.open_dataset(
-            f, keep_particles=keep_particles, probe_names=probe_names, chunks=chunks
+            f,
+            keep_particles=keep_particles,
+            probe_names=probe_names,
+            chunks=chunks,
+            deck_path=deck_path,
         )
 
         # If the data_vars are specified then only load them in and disregard the rest.
@@ -302,6 +346,7 @@ def open_datatree(
     *,
     keep_particles: bool = False,
     probe_names: list[str] | None = None,
+    deck_path: PathLike | None = None,
 ) -> xr.DataTree:
     """
     An `xarray.DataTree` is constructed utilising the original names in the SDF
@@ -338,15 +383,21 @@ def open_datatree(
         If ``True``, also load particle data (this may use a lot of memory!)
     probe_names
         List of EPOCH probe names
-
+    deck_path
+        If ``None``, attempt to load the ``"input.deck"`` from the same directory as the SDF files
+        and silently fail if it does not exist. If a path is given, load the specified deck
+        from a relative or absolute file path. See :ref:`loading-input-deck` for details.
     Examples
     --------
     >>> dt = open_datatree("0000.sdf")
     >>> dt["Electric_Field"]["Ex"].values  # Access all Electric_Field_Ex data
     """
 
     return xr.open_datatree(
-        path, keep_particles=keep_particles, probe_names=probe_names
+        path,
+        keep_particles=keep_particles,
+        probe_names=probe_names,
+        deck_path=deck_path,
     )
 
 
@@ -357,6 +408,7 @@ def open_mfdatatree(
     keep_particles: bool = False,
     probe_names: list[str] | None = None,
     data_vars: list[str] | None = None,
+    deck_path: PathLike | None = None,
 ) -> xr.DataTree:
     """Open a set of EPOCH SDF files as one `xarray.DataTree`
 
@@ -419,6 +471,10 @@ def open_mfdatatree(
         List of EPOCH probe names
     data_vars
         List of data vars to load in (If not specified loads in all variables)
+    deck_path
+        If ``None``, attempt to load the ``"input.deck"`` from the same directory as the SDF files
+        and silently fail if it does not exist. If a path is given, load the specified deck
+        from a relative or absolute file path. See :ref:`loading-input-deck` for details.
 
     Examples
     --------
@@ -433,6 +489,7 @@ def open_mfdatatree(
         keep_particles=keep_particles,
         probe_names=probe_names,
         data_vars=data_vars,
+        deck_path=deck_path,
     )
 
     return _build_datatree_from_dataset(combined_ds)
@@ -512,6 +569,7 @@ class SDFDataStore(AbstractDataStore):
     __slots__ = (
         "_filename",
         "_manager",
+        "deck_path",
         "drop_variables",
         "keep_particles",
         "lock",
@@ -523,13 +581,15 @@ def __init__(
         manager,
         drop_variables=None,
         keep_particles=False,
+        deck_path=None,
         lock=None,
         probe_names=None,
     ):
         self._manager = manager
-        self._filename = self.ds.filename
+        self._filename = self.ds.header["filename"]
         self.drop_variables = drop_variables
         self.keep_particles = keep_particles
+        self.deck_path = deck_path
         self.lock = ensure_lock(lock)
         self.probe_names = probe_names
 
@@ -541,6 +601,7 @@ def open(
         drop_variables=None,
         keep_particles=False,
         probe_names=None,
+        deck_path=None,
     ):
         if isinstance(filename, os.PathLike):
             filename = os.fspath(filename)
@@ -552,6 +613,7 @@ def open(
             drop_variables=drop_variables,
             keep_particles=keep_particles,
             probe_names=probe_names,
+            deck_path=deck_path,
         )
 
     def _acquire(self, needs_lock=True):
@@ -594,7 +656,7 @@ def _norm_grid_name(grid_name: str) -> str:
             return grid_name.split("/", maxsplit=1)[-1]
 
         def _grid_species_name(grid_name: str) -> str:
-            return grid_name.split("/")[-1]
+            return grid_name.rsplit("/", maxsplit=1)[-1]
 
         def _process_grid_name(grid_name: str, transform_func) -> str:
             """Apply the given transformation function and then rename with underscores."""
@@ -756,6 +818,7 @@ def _process_grid_name(grid_name: str, transform_func) -> str:
         # )
 
         ds = xr.Dataset(data_vars, attrs=attrs, coords=coords)
+        ds.attrs["deck"] = _load_deck(ds.attrs["filename"], self.deck_path)
         ds.set_close(self.ds.close)
 
         return ds
@@ -771,6 +834,7 @@ class SDFEntrypoint(BackendEntrypoint):
         "drop_variables",
         "keep_particles",
         "probe_names",
+        "deck_path",
     ]
 
     def open_dataset(
@@ -780,6 +844,7 @@ def open_dataset(
         drop_variables=None,
         keep_particles=False,
         probe_names=None,
+        deck_path=None,
     ):
         if isinstance(filename_or_obj, Path):
             # sdf library takes a filename only
@@ -791,6 +856,7 @@ def open_dataset(
             drop_variables=drop_variables,
             keep_particles=keep_particles,
             probe_names=probe_names,
+            deck_path=deck_path,
         )
         with close_on_error(store):
             return store.load()
@@ -800,6 +866,7 @@ def open_dataset(
         "drop_variables",
         "keep_particles",
         "probe_names",
+        "deck_path",
     ]
 
     def open_datatree(
@@ -809,12 +876,14 @@ def open_datatree(
         drop_variables=None,
         keep_particles=False,
         probe_names=None,
+        deck_path=None,
     ):
         ds = self.open_dataset(
             filename_or_obj,
             drop_variables=drop_variables,
             keep_particles=keep_particles,
             probe_names=probe_names,
+            deck_path=deck_path,
         )
         return _build_datatree_from_dataset(ds)