Merge pull request #10 from PlasmaFAIR/preprocessor

Add jobid checking to preprocess

README.md

@@ -28,6 +28,7 @@ xarray:
 ### Single file loading
 ```python
 import xarray as xr
+from sdf_xarray import SDFPreprocess
 
 df = xr.open_dataset("0010.sdf")
 
@@ -50,7 +51,7 @@ ds = xr.open_mfdataset(
     data_vars='minimal', 
     coords='minimal', 
     compat='override', 
-    preprocess=lambda ds: ds.expand_dims(time=[ds.attrs["time"]])
+    preprocess=SDFPreprocess()
 )
 
 print(ds)

src/sdf_xarray/__init__.py

@@ -228,3 +228,21 @@ def guess_can_open(self, filename_or_obj):
     description = "Use .sdf files in Xarray"
 
     url = "https://epochpic.github.io/documentation/visualising_output/python.html"
+
+
+class SDFPreprocess:
+    """Preprocess SDF files for xarray ensuring matching job ids and sets time dimension"""
+
+    def __init__(self):
+        self.job_id: int | None = None
+
+    def __call__(self, ds: xr.Dataset) -> xr.Dataset:
+        if self.job_id is None:
+            self.job_id = ds.attrs["jobid1"]
+
+        if self.job_id != ds.attrs["jobid1"]:
+            raise ValueError(
+                f"Mismatching job ids (got {ds.attrs['jobid1']}, expected {self.job_id})"
+            )
+
+        return ds.expand_dims(time=[ds.attrs["time"]])

tests/test_basic.py

@@ -1,10 +1,13 @@
 import pathlib
 
 import xarray as xr
-from sdf_xarray import open_mfdataset
-
+from sdf_xarray import open_mfdataset, SDFPreprocess
+import pytest
 
 EXAMPLE_FILES_DIR = pathlib.Path(__file__).parent / "example_files"
+EXAMPLE_MISMATCHED_FILES_DIR = (
+    pathlib.Path(__file__).parent / "example_mismatched_files"
+)
 
 
 def test_basic():
@@ -42,3 +45,16 @@ def test_multiple_files_multiple_time_dims():
     assert list(df["Electric Field/Ex"].coords) != list(df["Electric Field/Ez"].coords)
     assert df["Electric Field/Ex"].shape == (11, 16)
     assert df["Electric Field/Ez"].shape == (1, 16)
+
+
+def test_erroring_on_mismatched_jobid_files():
+    with pytest.raises(ValueError):
+        xr.open_mfdataset(
+            EXAMPLE_MISMATCHED_FILES_DIR.glob("*.sdf"),
+            concat_dim="time",
+            combine="nested",
+            data_vars="minimal",
+            coords="minimal",
+            compat="override",
+            preprocess=SDFPreprocess(),
+        )