osiris-utils 1.2.1__py3-none-any.whl → 1.2.3__py3-none-any.whl

osiris_utils/postprocessing/field_centering.py

@@ -1,7 +1,5 @@
 from __future__ import annotations
 
-from collections.abc import Generator
-
 import numpy as np
 
 from ..data.diagnostic import OSIRIS_FLD, Diagnostic
@@ -36,15 +34,8 @@ class FieldCentering_Simulation(PostProcess):
         simulation : Simulation
             The simulation object.
         """
-        super().__init__("FieldCentering Simulation")
-
-        # Accept Simulation-compatible objects (Simulation or other PostProcess subclasses)
-        if not isinstance(simulation, Simulation):
-            raise ValueError("simulation must be a Simulation-compatible object.")
-        self._simulation = simulation
-
+        super().__init__("FieldCentering Simulation", simulation)
         self._field_centered = {}
-        # no need to create a species handler for field centering since fields are not species related
 
     def __getitem__(self, key: str) -> FieldCentering_Diagnostic:
         if key not in OSIRIS_FLD:
@@ -62,10 +53,6 @@ class FieldCentering_Simulation(PostProcess):
         else:
             print(f"Field {key} not found in simulation")
 
-    def process(self, diagnostic: Diagnostic) -> FieldCentering_Diagnostic:
-        """Apply field centering to a diagnostic"""
-        return FieldCentering_Diagnostic(diagnostic)
-
 
 class FieldCentering_Diagnostic(Diagnostic):
     def __init__(self, diagnostic: Diagnostic):
@@ -109,242 +96,117 @@ class FieldCentering_Diagnostic(Diagnostic):
 
         self._original_name = diagnostic._name
         self._name = diagnostic._name + "_centered"
-
-        self._data: np.ndarray | None = None
         self._all_loaded = False
+        self._data = None
 
-    def load_all(self) -> np.ndarray:
+    def _needs_center_axes_osiris(self) -> tuple[int, ...]:
         """
-        Load all data and center the fields.
+        Return OSIRIS spatial axes (1..3) along which this field must be averaged
+        to move it to cell centers.
+        """
+        name = self._original_name.lower()
 
-        Returns
-        -------
-        data : np.ndarray
-            The centered field data.
+        if self._dim == 1:
+            # data shape: (t, x1) or per-frame (x1)
+            # E1 is staggered in x1; B2/B3 staggered in x1 in 1D Yee
+            if name in {"e1", "part_e1", "ext_e1", "b2", "part_b2", "ext_b2", "b3", "part_b3", "ext_b3"}:
+                return (1,)
+            return ()
+
+        if self._dim == 2:
+            # E1,B2 staggered in x1; E2,B1 staggered in x2; B3 staggered in both x1 and x2; E3 centered
+            if name in {"e1", "part_e1", "ext_e1", "b2", "part_b2", "ext_b2"}:
+                return (1,)
+            if name in {"e2", "part_e2", "ext_e2", "b1", "part_b1", "ext_b1"}:
+                return (2,)
+            if name in {"b3", "part_b3", "ext_b3"}:
+                return (1, 2)
+            return ()
+
+        if self._dim == 3:
+            # Typical Yee: E components are staggered along their own axis; B components staggered in the other two axes.
+            # Your original code implements that. We'll preserve it:
+            if name in {"e1", "part_e1", "ext_e1"}:
+                return (1,)
+            if name in {"e2", "part_e2", "ext_e2"}:
+                return (2,)
+            if name in {"e3", "part_e3", "ext_e3"}:
+                return (3,)
+
+            if name in {"b1", "part_b1", "ext_b1"}:
+                return (2, 3)
+            if name in {"b2", "part_b2", "ext_b2"}:
+                return (1, 3)
+            if name in {"b3", "part_b3", "ext_b3"}:
+                return (1, 2)
+
+            return ()
+
+        raise ValueError(f"Unknown dimension {self._dim}.")
+
+    @staticmethod
+    def _center_along_axis(data: np.ndarray, axis: int) -> np.ndarray:
+        """0.5 * (f + roll(f, +1)) along the given numpy axis."""
+        return 0.5 * (data + np.roll(data, shift=1, axis=axis))
+
+    def _check_slice_safe(self, data_slice: tuple[slice, ...] | None, axes_to_center_osiris: tuple[int, ...]) -> None:
         """
+        For correctness without halos: if user slices along a centering axis,
+        we require full slice(None) along that axis.
+        """
+        if data_slice is None:
+            return
+        # data_slice is spatial-only, ordered (x1,x2,x3) depending on dim
+        for ax_osiris in axes_to_center_osiris:
+            spatial_index = ax_osiris - 1  # x1->0, x2->1, x3->2 within data_slice
+            if spatial_index < len(data_slice):
+                s = data_slice[spatial_index]
+                if isinstance(s, slice) and (s.start is None and s.stop is None and s.step is None):
+                    continue
+                raise ValueError(
+                    "FieldCentering with spatial slicing requires full slices (:) along axes that need centering. "
+                    f"Field {self._original_name} needs centering along x{ax_osiris}, but got slice {s}. "
+                    "Either use full domain on that axis, or implement halo reads."
+                )
+
+        # ---- eager path ----
+
+    def load_all(self) -> np.ndarray:
         if self._data is not None:
             return self._data
 
-        if not hasattr(self._diag, "_data") or self._diag._data is None:
-            self._diag.load_all()
+        self._diag.load_all()
+        data = self._diag.data  # shape: (t, x1[,x2[,x3]])
 
-        if self._dim == 1:
-            if self._original_name.lower() in [
-                "b2",
-                "part_b2",
-                "ext_b2",
-                "b3",
-                "part_b3",
-                "ext_b3",
-                "e1",
-                "part_e1",
-                "ext_e1",
-            ]:
-                result = 0.5 * (np.roll(self._diag.data, shift=1, axis=1) + self._diag.data)
-            elif self._original_name.lower() in [
-                "b1",
-                "part_b1",
-                "ext_b1",
-                "e2",
-                "part_e2",
-                "ext_e2",
-                "e3",
-                "part_e3",
-                "ext_e3",
-            ]:
-                result = self._diag.data
-
-        elif self._dim == 2:
-            if self._original_name.lower() in [
-                "e1",
-                "part_e1",
-                "ext_e1",
-                "b2",
-                "part_b2",
-                "ext_b2",
-            ]:
-                result = 0.5 * (np.roll(self._diag.data, shift=1, axis=1) + self._diag.data)
-            elif self._original_name.lower() in [
-                "e2",
-                "part_e2",
-                "ext_e2",
-                "b1",
-                "part_b1",
-                "ext_b1",
-            ]:
-                result = 0.5 * (np.roll(self._diag.data, shift=1, axis=2) + self._diag.data)
-            elif self._original_name.lower() in ["b3", "part_b3", "ext_b3"]:
-                result = 0.5 * (
-                    np.roll(
-                        (0.5 * (np.roll(self._diag.data, shift=1, axis=1) + self._diag.data)),
-                        shift=1,
-                        axis=2,
-                    )
-                    + (0.5 * (np.roll(self._diag.data, shift=1, axis=1) + self._diag.data))
-                )
-            elif self._original_name.lower() in ["e3", "part_e3", "ext_e3"]:
-                result = self._diag.data
-
-        elif self._dim == 3:
-            # TODO test this
-            if self._original_name in ["b1", "part_b1", "ext_b1"]:
-                result = 0.5 * (
-                    0.5
-                    * np.roll(
-                        (np.roll(self._diag.data, shift=1, axis=2) + self._diag.data),
-                        shift=1,
-                        axis=3,
-                    )
-                    + 0.5 * (np.roll(self._diag.data, shift=1, axis=2) + self._diag.data)
-                )
-            elif self._original_name in ["b2", "part_b2", "ext_b2"]:
-                result = 0.5 * (
-                    0.5
-                    * np.roll(
-                        (np.roll(self._diag.data, shift=1, axis=1) + self._diag.data),
-                        shift=1,
-                        axis=3,
-                    )
-                    + 0.5 * (np.roll(self._diag.data, shift=1, axis=1) + self._diag.data)
-                )
-            elif self._original_name in ["b3", "part_b3", "ext_b3"]:
-                result = 0.5 * (
-                    0.5
-                    * np.roll(
-                        (np.roll(self._diag.data, shift=1, axis=1) + self._diag.data),
-                        shift=1,
-                        axis=2,
-                    )
-                    + 0.5 * (np.roll(self._diag.data, shift=1, axis=1) + self._diag.data)
-                )
-            elif self._original_name in ["e1", "part_e1", "ext_e1"]:
-                result = 0.5 * (np.roll(self._diag.data, shift=1, axis=1) + self._diag.data)
-            elif self._original_name in ["e2", "part_e2", "ext_e2"]:
-                result = 0.5 * (np.roll(self._diag.data, shift=1, axis=2) + self._diag.data)
-            elif self._original_name in ["e3", "part_e3", "ext_e3"]:
-                result = 0.5 * (np.roll(self._diag.data, shift=1, axis=3) + self._diag.data)
+        axes_to_center = self._needs_center_axes_osiris()
 
-        else:
-            raise ValueError(f"Unknown dimension {self._dim}.")
+        # Map OSIRIS spatial axes (1..3) to numpy axes in the loaded array:
+        # loaded includes time at axis 0, so x1->1, x2->2, x3->3
+        result = data
+        for ax_osiris in axes_to_center:
+            np_axis = ax_osiris  # because time is axis 0
+            result = self._center_along_axis(result, axis=np_axis)
 
         self._data = result
         self._all_loaded = True
         return self._data
 
-    def __getitem__(self, index: int | slice) -> np.ndarray:
-        """Get data at a specific index"""
-        if self._all_loaded and self._data is not None:
-            return self._data[index]
-
-        if isinstance(index, int):
-            return next(self._data_generator(index))
-        elif isinstance(index, slice):
-            start = 0 if index.start is None else index.start
-            step = 1 if index.step is None else index.step
-            stop = self._diag._maxiter if index.stop is None else index.stop
-            return np.array([next(self._data_generator(i)) for i in range(start, stop, step)])
-        else:
-            raise ValueError("Invalid index type. Use int or slice.")
+    # ---- lazy path ----
 
-    def _data_generator(self, index: int) -> Generator[np.ndarray, None, None]:
-        if self._dim == 1:
-            if self._original_name.lower() in [
-                "b2",
-                "part_b2",
-                "ext_b2",
-                "b3",
-                "part_b3",
-                "ext_b3",
-                "e1",
-                "part_e1",
-                "ext_e1",
-            ]:
-                yield 0.5 * (np.roll(self._diag[index], shift=1) + self._diag[index])
-            elif self._original_name.lower() in [
-                "b1",
-                "part_b1",
-                "ext_b1",
-                "e2",
-                "part_e2",
-                "ext_e2",
-                "e3",
-                "part_e3",
-                "ext_e3",
-            ]:  # it's already centered but self._data does not exist
-                yield self._diag[index]
-            else:
-                raise ValueError(f"Unknown field {self._original_name}.")
-
-        elif self._dim == 2:
-            if self._original_name in [
-                "e1",
-                "part_e1",
-                "ext_e1",
-                "b2",
-                "part_b2",
-                "ext_b2",
-            ]:
-                yield 0.5 * (np.roll(self._diag[index], shift=1, axis=0) + self._diag[index])
-            elif self._original_name in [
-                "e2",
-                "part_e2",
-                "ext_e2",
-                "b1",
-                "part_b1",
-                "ext_b1",
-            ]:
-                yield 0.5 * (np.roll(self._diag[index], shift=1, axis=1) + self._diag[index])
-            elif self._original_name in ["b3", "part_b3", "ext_b3"]:
-                yield 0.5 * (
-                    np.roll(
-                        (0.5 * (np.roll(self._diag[index], shift=1, axis=0) + self._diag[index])),
-                        shift=1,
-                        axis=1,
-                    )
-                    + (0.5 * (np.roll(self._diag[index], shift=1, axis=0) + self._diag[index]))
-                )
-            elif self._original_name in ["e3", "part_e3", "ext_e3"]:
-                yield self._diag[index]
-            else:
-                raise ValueError(f"Unknown field {self._original_name}.")
-
-        elif self._dim == 3:
-            if self._original_name in ["b1", "part_b1", "ext_b1"]:
-                yield 0.5 * (
-                    0.5
-                    * np.roll(
-                        (np.roll(self._diag[index], shift=1, axis=1) + self._diag[index]),
-                        shift=1,
-                        axis=2,
-                    )
-                    + 0.5 * (np.roll(self._diag[index], shift=1, axis=1) + self._diag[index])
-                )
-            elif self._original_name in ["b2", "part_b2", "ext_b2"]:
-                yield 0.5 * (
-                    0.5
-                    * np.roll(
-                        (np.roll(self._diag[index], shift=1, axis=0) + self._diag[index]),
-                        shift=1,
-                        axis=2,
-                    )
-                    + 0.5 * (np.roll(self._diag[index], shift=1, axis=0) + self._diag[index])
-                )
-            elif self._original_name in ["b3", "part_b3", "ext_b3"]:
-                yield 0.5 * (
-                    0.5
-                    * np.roll(
-                        (np.roll(self._diag[index], shift=1, axis=0) + self._diag[index]),
-                        shift=1,
-                        axis=1,
-                    )
-                    + 0.5 * (np.roll(self._diag[index], shift=1, axis=0) + self._diag[index])
-                )
-            elif self._original_name in ["e1", "part_e1", "ext_e1"]:
-                yield 0.5 * (np.roll(self._diag[index], shift=1, axis=0) + self._diag[index])
-            elif self._original_name in ["e2", "part_e2", "ext_e2"]:
-                yield 0.5 * (np.roll(self._diag[index], shift=1, axis=1) + self._diag[index])
-            elif self._original_name in ["e3", "part_e3", "ext_e3"]:
-                yield 0.5 * (np.roll(self._diag[index], shift=1, axis=2) + self._diag[index])
+    def _frame(self, index: int, data_slice: tuple | None = None) -> np.ndarray:
+        axes_to_center = self._needs_center_axes_osiris()
 
-        else:
-            raise ValueError(f"Unknown dimension {self._dim}.")
+        # Ensure slicing won't produce wrong values without halos
+        self._check_slice_safe(data_slice, axes_to_center)
+
+        # read one timestep lazily; returned array shape: (x1[,x2[,x3]])
+        f = self._diag._frame(index, data_slice=data_slice)
+
+        # Map OSIRIS spatial axes to numpy axes in a single-timestep array:
+        # per-frame has no time dim, so x1->0, x2->1, x3->2
+        result = f
+        for ax_osiris in axes_to_center:
+            np_axis = ax_osiris - 1
+            result = self._center_along_axis(result, axis=np_axis)
+
+        return result

osiris_utils/postprocessing/heatflux_correction.py

@@ -1,6 +1,5 @@
 from __future__ import annotations
 
-from collections.abc import Generator
 from typing import Any
 
 import numpy as np
@@ -21,7 +20,7 @@ __all__ = [
 
 class HeatfluxCorrection_Simulation(PostProcess):
     def __init__(self, simulation: Simulation):
-        super().__init__("HeatfluxCorrection Simulation")
+        super().__init__("HeatfluxCorrection Simulation", simulation)
         """
         Class to correct pressure tensor components by subtracting Reynolds stress.
 
@@ -32,9 +31,6 @@ class HeatfluxCorrection_Simulation(PostProcess):
         heatflux : str
             The heatflux component to center.
         """
-        if not isinstance(simulation, Simulation):
-            raise ValueError("simulation must be a Simulation-compatible object.")
-        self._simulation = simulation
         self._heatflux_corrected: dict[str, HeatfluxCorrection_Diagnostic] = {}
         self._species_handler: dict[str, HeatfluxCorrection_Species_Handler] = {}
 
@@ -61,11 +57,6 @@ class HeatfluxCorrection_Simulation(PostProcess):
         else:
             print(f"Heatflux {key} not found in simulation")
 
-    def process(self, diagnostic: Diagnostic) -> HeatfluxCorrection_Diagnostic:
-        """Apply heatflux correction to a diagnostic"""
-        # FIX: This is a bit of a hack, but it works for now
-        return HeatfluxCorrection_Diagnostic(diagnostic, self._simulation)
-
 
 class HeatfluxCorrection_Diagnostic(Diagnostic):
     def __init__(
@@ -126,6 +117,18 @@ class HeatfluxCorrection_Diagnostic(Diagnostic):
         self._data: np.ndarray | None = None
         self._all_loaded = False
 
+    def _compute(
+        self,
+        q: np.ndarray,
+        vfl_i: np.ndarray,
+        trace_P: np.ndarray,
+        vfl_dot_Pji: np.ndarray,
+    ) -> np.ndarray:
+        """
+        Centralize the correction formula so eager/lazy always match.
+        """
+        return 2.0 * q - vfl_i * trace_P - 2.0 * vfl_dot_Pji
+
     def load_all(self) -> np.ndarray:
         if self._data is not None:
             return self._data
@@ -147,37 +150,41 @@ class HeatfluxCorrection_Diagnostic(Diagnostic):
         q = self._diag.data
         vfl_i = self._vfl_i.data
 
-        trace_P = sum(Pjj.data for Pjj in self._Pjj_list)
+        trace_P = None
+        for Pjj in self._Pjj_list:
+            trace_P = Pjj.data if trace_P is None else trace_P + Pjj.data
 
-        # Sum over j: vfl_j * Pji
-        vfl_dot_Pji = sum(vfl_j.data * Pji.data for vfl_j, Pji in zip(self._vfl_j_list, self._Pji_list, strict=False))
+        vfl_dot_Pji = None
+        for vfl_j, Pji in zip(self._vfl_j_list, self._Pji_list, strict=False):
+            term = vfl_j.data * Pji.data
+            vfl_dot_Pji = term if vfl_dot_Pji is None else vfl_dot_Pji + term
 
-        self._data = 2 * q - vfl_i * trace_P - 2 * vfl_dot_Pji
+        self._data = self._compute(q, vfl_i, trace_P, vfl_dot_Pji)
         self._all_loaded = True
-
         return self._data
 
-    def __getitem__(self, index: int | slice) -> np.ndarray:
-        """Get data at a specific index"""
-        if self._all_loaded and self._data is not None:
-            return self._data[index]
-
-        if isinstance(index, int):
-            return next(self._data_generator(index))
-        elif isinstance(index, slice):
-            start = 0 if index.start is None else index.start
-            step = 1 if index.step is None else index.step
-            stop = self._diag._maxiter if index.stop is None else index.stop
-            return np.array([next(self._data_generator(i)) for i in range(start, stop, step)])
-        else:
-            raise ValueError("Invalid index type. Use int or slice.")
+    def _frame(self, index: int, data_slice: tuple | None = None) -> np.ndarray:
+        """
+        Lazy per-timestep correction. Reads only requested slices from disk,
+        provided the underlying diagnostics support data_slice.
+        """
+        # Read one timestep (and only the requested spatial slice)
+        q = self._diag._frame(index, data_slice=data_slice)
+        vfl_i = self._vfl_i._frame(index, data_slice=data_slice)
+
+        # trace_P = sum_j Pjj
+        trace_P = None
+        for Pjj in self._Pjj_list:
+            arr = Pjj._frame(index, data_slice=data_slice)
+            trace_P = arr if trace_P is None else trace_P + arr
+
+        # vfl_dot_Pji = sum_j vfl_j * Pji
+        vfl_dot_Pji = None
+        for vfl_j, Pji in zip(self._vfl_j_list, self._Pji_list, strict=False):
+            term = vfl_j._frame(index, data_slice=data_slice) * Pji._frame(index, data_slice=data_slice)
+            vfl_dot_Pji = term if vfl_dot_Pji is None else vfl_dot_Pji + term
 
-    def _data_generator(self, index: int) -> Generator[np.ndarray, None, None]:
-        q = self._diag[index]
-        vfl_i = self._vfl_i[index]
-        trace_P = sum(Pjj[index] for Pjj in self._Pjj_list)
-        vfl_dot_Pji = sum(vfl_j[index] * Pji[index] for vfl_j, Pji in zip(self._vfl_j_list, self._Pji_list, strict=False))
-        yield 2 * q - 0.5 * vfl_i * trace_P - vfl_dot_Pji
+        return self._compute(q, vfl_i, trace_P, vfl_dot_Pji)
 
 
 class HeatfluxCorrection_Species_Handler:
@@ -215,7 +222,8 @@ class HeatfluxCorrection_Species_Handler:
 
             # Compute quantities for vfl_j * P_{ji}
             vfl_j_list = [self._species_handler[f"vfl{j}"] for j in range(1, diag._dim + 1)]
-            Pji_list = [PressureCorrection_Simulation(self._simulation)[diag._species._name][f"P{j}{i}"] for j in range(1, diag._dim + 1)]
+            pc = PressureCorrection_Simulation(self._simulation)
+            Pji_list = [pc[diag._species._name][f"P{j}{i}"] for j in range(1, diag._dim + 1)]
 
             self._heatflux_corrected[key] = HeatfluxCorrection_Diagnostic(diag, vfl_i, Pjj_list, vfl_j_list, Pji_list)
         return self._heatflux_corrected[key]