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

osiris_utils/postprocessing/derivative.py

@@ -1,12 +1,11 @@
 from __future__ import annotations
 
-from collections.abc import Generator
+from collections import OrderedDict
 from typing import Any
 
 import numpy as np
 
 from ..data.diagnostic import Diagnostic
-from ..data.simulation import Simulation
 from .postprocess import PostProcess
 
 __all__ = ["Derivative_Diagnostic", "Derivative_Simulation", "Derivative_Species_Handler"]
@@ -44,55 +43,60 @@ class Derivative_Simulation(PostProcess):
 
     def __init__(
         self,
-        simulation: Simulation | Derivative_Simulation,
+        simulation: Any,  # Simulation-like (Simulation or PostProcess wrapper)
         deriv_type: str,
         axis: int | tuple[int, int] | None = None,
         order: int = 2,
     ):
-        super().__init__(f"Derivative({deriv_type})")
-        # Accept both Simulation and Derivative_Simulation objects
-        if not isinstance(simulation, (Simulation, Derivative_Simulation)):
-            raise ValueError("simulation must be a Simulation or Derivative_Simulation object.")
-        self._simulation = simulation
+        super().__init__(f"Derivative({deriv_type})", simulation)
+
+        # Capability checks (supports chaining)
+        if not hasattr(simulation, "__getitem__"):
+            raise TypeError("simulation must be Simulation-like: supports __getitem__.")
+        if not (hasattr(simulation, "_species") or hasattr(simulation, "species")):
+            raise TypeError("simulation must be Simulation-like: has _species or species.")
+
         self._deriv_type = deriv_type
-        self._axis = axis
-        self._derivatives_computed = {}
-        self._species_handler = {}
+        self._op_axis = axis
         self._order = order
 
-        # Copy species list to make this class behave like a Simulation
-        if hasattr(simulation, "_species"):
-            self._species = simulation._species
-        else:
-            self._species = []
+        self._derivatives_computed: dict[Any, Derivative_Diagnostic] = {}
+        self._species_handler: dict[Any, Derivative_Species_Handler] = {}
 
-    def __getitem__(self, key: Any) -> Derivative_Species_Handler | Derivative_Diagnostic:
-        if key in self._simulation._species:
+        # species already set by PostProcess, but keep it explicit if you want:
+        self._species = getattr(simulation, "_species", getattr(simulation, "species", []))
+
+    def __getitem__(self, key: Any):
+        if key in self._species:
             if key not in self._species_handler:
-                self._species_handler[key] = Derivative_Species_Handler(self._simulation[key], self._deriv_type, self._axis, self._order)
+                self._species_handler[key] = Derivative_Species_Handler(
+                    self._simulation[key],  # species handler from wrapped sim
+                    self._deriv_type,
+                    self._op_axis,
+                    self._order,
+                )
             return self._species_handler[key]
 
         if key not in self._derivatives_computed:
             self._derivatives_computed[key] = Derivative_Diagnostic(
                 diagnostic=self._simulation[key],
                 deriv_type=self._deriv_type,
-                axis=self._axis,
+                axis=self._op_axis,
                 order=self._order,
             )
         return self._derivatives_computed[key]
 
     def delete_all(self) -> None:
         self._derivatives_computed = {}
+        self._species_handler = {}
 
     def delete(self, key: Any) -> None:
         if key in self._derivatives_computed:
             del self._derivatives_computed[key]
+        elif key in self._species_handler:
+            del self._species_handler[key]
         else:
-            print(f"Derivative {key} not found in simulation")
-
-    def process(self, diagnostic: Diagnostic) -> Derivative_Diagnostic:
-        """Apply derivative to a diagnostic"""
-        return Derivative_Diagnostic(diagnostic, self._deriv_type, self._axis)
+            print(f"Derivative {key} not found")
 
     @property
     def species(self) -> list:
@@ -159,8 +163,8 @@ class Derivative_Diagnostic(Diagnostic):
         # Initialize using parent's __init__ with the same species
         if hasattr(diagnostic, "_species"):
             super().__init__(
-                simulation_folder=(diagnostic._simulation_folder if hasattr(diagnostic, "_simulation_folder") else None),
-                species=diagnostic._species,
+                simulation_folder=(getattr(diagnostic, "_simulation_folder", None)),
+                species=getattr(diagnostic, "_species", None),
             )
         else:
             super().__init__(None)
@@ -170,11 +174,14 @@ class Derivative_Diagnostic(Diagnostic):
         # self._name = f"D[{diagnostic._name}, {type}]"
         self._diag = diagnostic
         self._deriv_type = deriv_type
-        self._axis = axis if axis is not None else diagnostic._axis
+        self._op_axis = axis
         self._data = None
         self._all_loaded = False
         self._order = order
 
+        self._cache = OrderedDict()
+        self._cache_max = 6  # Maximum number of items to keep in cache (enough for 4th-order time stencil)
+
         # Copy all relevant attributes from diagnostic
         for attr in [
             "_dt",
@@ -302,38 +309,43 @@ class Derivative_Diagnostic(Diagnostic):
                 result = np.gradient(self._data, self._diag._dx[2], axis=3, edge_order=2)
 
             elif self._deriv_type == "xx":
-                if len(self._axis) != 2:
+                if len(self._op_axis) != 2:
                     raise ValueError("Axis must be a tuple with two elements.")
                 result = np.gradient(
                     np.gradient(
                         self._data,
-                        self._diag._dx[self._axis[0] - 1],
-                        axis=self._axis[0],
+                        self._diag._dx[self._op_axis[0] - 1],
+                        axis=self._op_axis[0],
                         edge_order=2,
                     ),
-                    self._diag._dx[self._axis[1] - 1],
-                    axis=self._axis[1],
+                    self._diag._dx[self._op_axis[1] - 1],
+                    axis=self._op_axis[1],
                     edge_order=2,
                 )
 
-            elif self._deriv_type == "xt":
-                if not isinstance(self._axis, int):
-                    raise ValueError("Axis must be an integer.")
+            elif self._deriv_type == "xx":
+                if not isinstance(self._op_axis, (tuple, list)) or len(self._op_axis) != 2:
+                    raise ValueError("Axis must be a tuple with two elements.")
                 result = np.gradient(
-                    np.gradient(self._data, self._diag._dt, axis=0, edge_order=2),
-                    self._diag._dx[self._axis - 1],
-                    axis=self._axis,
+                    np.gradient(
+                        self._data,
+                        self._diag._dx[self._op_axis[0] - 1],
+                        axis=self._op_axis[0],
+                        edge_order=2,
+                    ),
+                    self._diag._dx[self._op_axis[1] - 1],
+                    axis=self._op_axis[1],
                     edge_order=2,
                 )
 
             elif self._deriv_type == "tx":
-                if not isinstance(self._axis, int):
+                if not isinstance(self._op_axis, int):
                     raise ValueError("Axis must be an integer.")
                 result = np.gradient(
                     np.gradient(
                         self._data,
-                        self._diag._dx[self._axis - 1],
-                        axis=self._axis,
+                        self._diag._dx[self._op_axis - 1],
+                        axis=self._op_axis,
                         edge_order=2,
                     ),
                     self._diag._dt,
@@ -365,91 +377,153 @@ class Derivative_Diagnostic(Diagnostic):
         self._data = result
         return self._data
 
-    def _data_generator(self, index: int) -> Generator[np.ndarray, None, None]:
-        """Generate data for a specific index on-demand"""
-        if self._order == 2:
-            if self._deriv_type == "x1":
-                if self._dim == 1:
-                    yield np.gradient(self._diag[index], self._diag._dx, axis=0, edge_order=2)
-                else:
-                    yield np.gradient(self._diag[index], self._diag._dx[0], axis=0, edge_order=2)
-
-            elif self._deriv_type == "x2":
-                yield np.gradient(self._diag[index], self._diag._dx[1], axis=1, edge_order=2)
-
-            elif self._deriv_type == "x3":
-                yield np.gradient(self._diag[index], self._diag._dx[2], axis=2, edge_order=2)
-
-            elif self._deriv_type == "t":
+    def _frame(self, index: int, data_slice: tuple | None = None) -> np.ndarray:
+        # dimension guards
+        if self._deriv_type in ("x2", "xt", "tx") and self._diag._dim < 2:
+            raise ValueError(f"{self._deriv_type} requested but diagnostic dim={self._diag._dim}")
+        if self._deriv_type in ("x3",) and self._diag._dim < 3:
+            raise ValueError(f"{self._deriv_type} requested but diagnostic dim={self._diag._dim}")
+        if self._deriv_type == "xx":
+            if not isinstance(self._op_axis, (tuple, list)) or len(self._op_axis) != 2:
+                raise ValueError("For 'xx', axis must be a tuple/list of two spatial axes (1..3).")
+            if max(self._op_axis) > self._diag._dim:
+                raise ValueError(f"xx requested for axes {self._op_axis} but dim={self._diag._dim}")
+
+        n = self._diag._maxiter
+        dt = float(self._diag._dt * self._diag._ndump)
+
+        # ---------- helpers ----------
+        def spatial_axis_np_from_osiris(ax_osiris: int) -> int:
+            # OSIRIS spatial axes are 1..3; per-timestep array axes are 0..2
+            if not isinstance(ax_osiris, int):
+                raise ValueError("Axis must be int for spatial derivatives.")
+            if ax_osiris < 1 or ax_osiris > 3:
+                raise ValueError(f"Spatial axis must be 1..3, got {ax_osiris}")
+            return ax_osiris - 1
+
+        def dx_for_np_axis(ax_np: int) -> float:
+            dx = self._diag._dx
+            if self._diag._dim == 1:
+                return float(dx[0] if isinstance(dx, (list, tuple, np.ndarray)) else dx)
+            return float(dx[ax_np])
+
+        def d_dx(f: np.ndarray, ax_np: int, order: int) -> np.ndarray:
+            dx = dx_for_np_axis(ax_np)
+            if order == 2:
+                return np.gradient(f, dx, axis=ax_np, edge_order=2)
+            if order == 4:
+                return self._compute_fourth_order_spatial(f, dx, ax_np)
+            raise ValueError("Only order 2 and 4 supported.")
+
+        def d_dt_at(i: int, order: int) -> np.ndarray:
+            # Uses upstream frames, all with identical data_slice
+            if order == 2:
+                if i == 0:
+                    f0 = self._base(0, data_slice)
+                    f1 = self._base(1, data_slice)
+                    f2 = self._base(2, data_slice)
+                    return (-3 * f0 + 4 * f1 - f2) / (2 * dt)
+                if i == n - 1:
+                    f0 = self._base(n - 1, data_slice)
+                    f1 = self._base(n - 2, data_slice)
+                    f2 = self._base(n - 3, data_slice)
+                    return (3 * f0 - 4 * f1 + f2) / (2 * dt)
+                fp = self._base(i + 1, data_slice)
+                fm = self._base(i - 1, data_slice)
+                return (fp - fm) / (2 * dt)
+
+            if order == 4:
+                # fall back near edges
+                if i < 2 or i > n - 3:
+                    return d_dt_at(i, order=2)
+                f_p2 = self._base(i + 2, data_slice)
+                f_p1 = self._base(i + 1, data_slice)
+                f_m1 = self._base(i - 1, data_slice)
+                f_m2 = self._base(i - 2, data_slice)
+                return (-f_p2 + 8 * f_p1 - 8 * f_m1 + f_m2) / (12 * dt)
+
+            raise ValueError("Only order 2 and 4 supported.")
+
+        # ---------- dispatch ----------
+        if self._deriv_type in ("x1", "x2", "x3"):
+            f = self._base(index, data_slice)
+            ax_np = {"x1": 0, "x2": 1, "x3": 2}[self._deriv_type]
+            return d_dx(f, ax_np, self._order)
+
+        if self._deriv_type == "t":
+            return d_dt_at(index, self._order)
+
+        if self._deriv_type == "xx":
+            # second derivative along two spatial axes (can be same or different)
+            if not isinstance(self._op_axis, (tuple, list)) or len(self._op_axis) != 2:
+                raise ValueError("For 'xx', axis must be a tuple/list of two OSIRIS spatial axes, e.g. (1,2).")
+            ax1_np = spatial_axis_np_from_osiris(self._op_axis[0])
+            ax2_np = spatial_axis_np_from_osiris(self._op_axis[1])
+
+            f = self._base(index, data_slice)
+            # order=4 here is not implemented as a true 4th-order second derivative; keep order=2 behavior
+            g = d_dx(f, ax1_np, order=2)
+            return d_dx(g, ax2_np, order=2)
+
+        if self._deriv_type == "xt":
+            # d/dx ( d/dt f )
+            if not isinstance(self._op_axis, int):
+                raise ValueError("For 'xt', axis must be an OSIRIS spatial axis int (1..3).")
+            ax_np = spatial_axis_np_from_osiris(self._op_axis)
+
+            ft = d_dt_at(index, order=2 if self._order == 4 else self._order)  # keep stable
+            return d_dx(ft, ax_np, order=2)
+
+        if self._deriv_type == "tx":
+            # d/dt ( d/dx f )
+            if not isinstance(self._op_axis, int):
+                raise ValueError("For 'tx', axis must be an OSIRIS spatial axis int (1..3).")
+            ax_np = spatial_axis_np_from_osiris(self._op_axis)
+
+            # do spatial derivative on each needed time frame, then time-stencil those
+            def spatial_frame(i: int) -> np.ndarray:
+                fi = self._base(i, data_slice)
+                return d_dx(fi, ax_np, order=2)
+
+            if self._order == 2:
                 if index == 0:
-                    yield (-3 * self._diag[index] + 4 * self._diag[index + 1] - self._diag[index + 2]) / (
-                        2 * self._diag._dt * self._diag._ndump
-                    )
-                elif index == self._diag._maxiter - 1:
-                    yield (3 * self._diag[index] - 4 * self._diag[index - 1] + self._diag[index - 2]) / (
-                        2 * self._diag._dt * self._diag._ndump
-                    )
-                else:
-                    yield (self._diag[index + 1] - self._diag[index - 1]) / (2 * self._diag._dt * self._diag._ndump)
-            else:
-                raise ValueError("Invalid derivative type. Use 'x1', 'x2', 'x3' or 't'.")
-
-        elif self._order == 4:
-            if self._deriv_type in ["x1", "x2", "x3"]:
-                # Use vectorized helper function
-                data = self._diag[index]
-                axis_map = {"x1": 0, "x2": 1, "x3": 2}
-                axis = axis_map[self._deriv_type]
-
-                if self._deriv_type == "x1":
-                    dx = self._diag._dx if self._dim == 1 else self._diag._dx[0]
-                elif self._deriv_type == "x2":
-                    dx = self._diag._dx[1]
-                else:  # x3
-                    dx = self._diag._dx[2]
-
-                yield self._compute_fourth_order_spatial(data, dx, axis)
-
-            elif self._deriv_type == "t":
-                idx = index
-                # Fourth-order time derivative
-                if idx < 2:
-                    # Forward difference for first two points
-                    yield (-3 * self._diag[idx] + 4 * self._diag[idx + 1] - self._diag[idx + 2]) / (2 * self._diag._dt * self._diag._ndump)
-                elif idx >= self._diag._maxiter - 2:
-                    # Backward difference for last two points
-                    yield (3 * self._diag[idx] - 4 * self._diag[idx - 1] + self._diag[idx - 2]) / (2 * self._diag._dt * self._diag._ndump)
-                else:
-                    # Fourth-order central: (-f[i+2] + 8*f[i+1] - 8*f[i-1] + f[i-2]) / (12*h)
-                    yield (-self._diag[idx + 2] + 8 * self._diag[idx + 1] - 8 * self._diag[idx - 1] + self._diag[idx - 2]) / (
-                        12 * self._diag._dt * self._diag._ndump
-                    )
-            else:
-                raise ValueError("Invalid derivative type. Use 'x1', 'x2', 'x3' or 't'.")
-
-    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))
-        if 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
-
-            # Pre-allocate array for better performance
-            indices = range(start, stop, step)
-            if len(indices) > 0:
-                first_result = next(self._data_generator(indices[0]))
-                result = np.empty((len(indices),) + first_result.shape, dtype=first_result.dtype)
-                result[0] = first_result
-                for i, idx in enumerate(indices[1:], start=1):
-                    result[i] = next(self._data_generator(idx))
-                return result
-            return np.array([])
-        raise ValueError("Invalid index type. Use int or slice.")
+                    return (-3 * spatial_frame(0) + 4 * spatial_frame(1) - spatial_frame(2)) / (2 * dt)
+                if index == n - 1:
+                    return (3 * spatial_frame(n - 1) - 4 * spatial_frame(n - 2) + spatial_frame(n - 3)) / (2 * dt)
+                return (spatial_frame(index + 1) - spatial_frame(index - 1)) / (2 * dt)
+
+            if self._order == 4:
+                if index < 2 or index > n - 3:
+                    # edge fallback
+                    return (spatial_frame(min(index + 1, n - 1)) - spatial_frame(max(index - 1, 0))) / (2 * dt)
+                return (
+                    -spatial_frame(index + 2) + 8 * spatial_frame(index + 1) - 8 * spatial_frame(index - 1) + spatial_frame(index - 2)
+                ) / (12 * dt)
+
+            raise ValueError("Only order 2 and 4 supported.")
+
+        raise ValueError("Invalid derivative type.")
+
+    def _cache_get(self, key):
+        if key in self._cache:
+            self._cache.move_to_end(key)
+            return self._cache[key]
+        return None
+
+    def _cache_put(self, key, val):
+        self._cache[key] = val
+        self._cache.move_to_end(key)
+        if len(self._cache) > self._cache_max:
+            self._cache.popitem(last=False)
+
+    def _base(self, idx: int, data_slice: tuple | None):
+        # base diagnostic frame WITH slice (this is the key!)
+        key = (idx, data_slice)
+        v = self._cache_get(key)
+        if v is None:
+            v = self._diag._frame(idx, data_slice=data_slice)
+            self._cache_put(key, v)
+        return v
 
 
 class Derivative_Species_Handler: