From 173bda1c1bb44f03232e63bcb7814365ca6e61fa Mon Sep 17 00:00:00 2001
From: Kai Germaschewski <kai.germaschewski@gmail.com>
Date: Mon, 8 Jun 2026 01:41:57 +0000
Subject: [PATCH] compat: psc_decode() handle new by-component output

psc now writes its output in a by-component format, which presumably
means xarray-adios2 will already handle things like the time dimension,
so in theory psc_decode() won't be needed anymore, but for compatibility,
this makes it so that psc_decode() can handle the new format, so that
old code that calls psc_decode() doesn't break.
---
 src/pscpy/psc.py | 23 +++++++++++++++--------
 1 file changed, 15 insertions(+), 8 deletions(-)

diff --git a/src/pscpy/psc.py b/src/pscpy/psc.py
index ad0cf2f..5dab3b2 100644
--- a/src/pscpy/psc.py
+++ b/src/pscpy/psc.py
@@ -34,9 +34,14 @@ def __init__(
             message = "Dataset is missing corner. A value must be manually provided."
             raise ValueError(message)
 
-        self.x = self._get_coord(0)
-        self.y = self._get_coord(1)
-        self.z = self._get_coord(2)
+        if "x" in ds.coords:
+            self.x = ds.coords["x"].values
+            self.y = ds.coords["y"].values
+            self.z = ds.coords["z"].values
+        else:
+            self.x = self._get_coord(0)
+            self.y = self._get_coord(1)
+            self.z = self._get_coord(2)
 
     def _get_coord(self, coord_idx: int) -> NDArray[Any]:
         return np.linspace(
@@ -87,14 +92,16 @@ def decode_psc(
                 dims[4]: "x",
             }
         )
-    ds = ds.squeeze("step")
+    if "step" in ds.dims:
+        ds = ds.squeeze("step")
 
     for var_name in ds:
         components = list(iter_components(var_name, species_names))
-        for component_idx, component in enumerate(components):
-            ds = ds.assign({component: ds[var_name].isel(component=component_idx)})
-        if var_name not in components:
-            ds = ds.drop_vars([var_name])
+        if components:
+            for component_idx, component in enumerate(components):
+                ds = ds.assign({component: ds[var_name].isel(component=component_idx)})
+            if var_name not in components:
+                ds = ds.drop_vars([var_name])
 
     run_info = RunInfo(ds, length=length, corner=corner)
     coords = {