osiris-utils 1.2.2__py3-none-any.whl → 1.2.4__py3-none-any.whl

osiris_utils/postprocessing/mft.py

@@ -28,10 +28,7 @@ class MFT_Simulation(PostProcess):
     """
 
     def __init__(self, simulation: Simulation, mft_axis: int | None = None):
-        super().__init__(f"MeanFieldTheory({mft_axis})")
-        if not isinstance(simulation, Simulation):
-            raise ValueError("simulation must be a Simulation-compatible object.")
-        self._simulation = simulation
+        super().__init__(f"MeanFieldTheory({mft_axis})", simulation)
         self._mft_axis = mft_axis
         self._mft_computed = {}
         self._species_handler = {}
@@ -54,106 +51,44 @@ class MFT_Simulation(PostProcess):
         else:
             print(f"MeanFieldTheory {key} not found in simulation")
 
-    def process(self, diagnostic):
-        """Apply mean field theory to a diagnostic"""
-        return MFT_Diagnostic(diagnostic, self._mft_axis)
-
 
 class MFT_Diagnostic(Diagnostic):
     """
-    Class to compute mean field theory of a diagnostic.
-    Acts as a container for the average and fluctuation components.
-
-    Parameters
-    ----------
-    diagnostic : Diagnostic
-        The diagnostic object.
-    mft_axis : int
-        The axis to compute mean field theory along.
-
-
+    Container: gives access to "avg" and "delta" diagnostics.
+    Not itself a time-indexed diagnostic.
     """
 
-    def __init__(self, diagnostic, mft_axis):
-        # Initialize using parent's __init__ with the same species
+    def __init__(self, diagnostic: Diagnostic, mft_axis: int):
         if hasattr(diagnostic, "_species"):
-            super().__init__(
-                simulation_folder=(diagnostic._simulation_folder if hasattr(diagnostic, "_simulation_folder") else None),
-                species=diagnostic._species,
-            )
+            super().__init__(simulation_folder=getattr(diagnostic, "_simulation_folder", None), species=diagnostic._species)
         else:
             super().__init__(None)
 
+        if mft_axis is None:
+            raise ValueError("Mean field theory axis must be specified (1..dim).")
+        if not (1 <= mft_axis <= diagnostic._dim):
+            raise ValueError(f"mft_axis must be in 1..{diagnostic._dim}, got {mft_axis}")
+
         self._name = f"MFT[{diagnostic._name}]"
         self._diag = diagnostic
         self._mft_axis = mft_axis
-        self._data = None
-        self._all_loaded = False
+        self._components: dict[str, Diagnostic] = {}
 
-        # Components that will be lazily created
-        self._components = {}
-
-        # Copy all relevant attributes from diagnostic
-        for attr in [
-            "_dt",
-            "_dx",
-            "_ndump",
-            "_axis",
-            "_nx",
-            "_x",
-            "_grid",
-            "_dim",
-            "_maxiter",
-            "_tunits",
-            "_type",
-        ]:
+        # copy metadata if needed
+        for attr in ["_dt", "_dx", "_ndump", "_axis", "_nx", "_x", "_grid", "_dim", "_maxiter", "_tunits", "_type"]:
             if hasattr(diagnostic, attr):
                 setattr(self, attr, getattr(diagnostic, attr))
 
-    def __getitem__(self, key):
-        """
-        Get a component of the mean field theory.
-
-        Parameters
-        ----------
-        key : str
-            Either "avg" for average or "delta" for fluctuations.
-
-        Returns
-        -------
-        Diagnostic
-            The requested component.
-        """
+    def __getitem__(self, key: str) -> Diagnostic:
         if key == "avg":
             if "avg" not in self._components:
                 self._components["avg"] = MFT_Diagnostic_Average(self._diag, self._mft_axis)
             return self._components["avg"]
-
-        elif key == "delta":
+        if key == "delta":
             if "delta" not in self._components:
                 self._components["delta"] = MFT_Diagnostic_Fluctuations(self._diag, self._mft_axis)
             return self._components["delta"]
-
-        else:
-            raise ValueError("Invalid MFT component. Use 'avg' or 'delta'.")
-
-    def load_all(self):
-        """Load both average and fluctuation components"""
-        # This will compute both components at once for efficiency
-        if "avg" not in self._components:
-            self._components["avg"] = MFT_Diagnostic_Average(self._diag, self._mft_axis)
-
-        if "delta" not in self._components:
-            self._components["delta"] = MFT_Diagnostic_Fluctuations(self._diag, self._mft_axis)
-
-        # Load both components
-        self._components["avg"].load_all()
-        self._components["delta"].load_all()
-
-        # Mark this container as loaded
-        self._all_loaded = True
-
-        return self._components
+        raise ValueError("Invalid MFT component. Use 'avg' or 'delta'.")
 
 
 class MFT_Diagnostic_Average(Diagnostic):
@@ -170,201 +105,90 @@ class MFT_Diagnostic_Average(Diagnostic):
 
     """
 
-    def __init__(self, diagnostic, mft_axis):
-        # Initialize with the same species as the diagnostic
+    def __init__(self, diagnostic: Diagnostic, mft_axis: int):
         if hasattr(diagnostic, "_species"):
-            super().__init__(
-                simulation_folder=(diagnostic._simulation_folder if hasattr(diagnostic, "_simulation_folder") else None),
-                species=diagnostic._species,
-            )
+            super().__init__(simulation_folder=getattr(diagnostic, "_simulation_folder", None), species=diagnostic._species)
         else:
             super().__init__(None)
 
         if mft_axis is None:
-            raise ValueError("Mean field theory axis must be specified.")
+            raise ValueError("Mean field theory axis must be specified (1..dim).")
+        if not (1 <= mft_axis <= diagnostic._dim):
+            raise ValueError(f"mft_axis must be in 1..{diagnostic._dim}, got {mft_axis}")
 
         self.postprocess_name = "MFT_AVG"
-
-        self._name = f"MFT_avg[{diagnostic._name}, {mft_axis}]"
+        self._name = f"MFT_avg[{diagnostic._name}, x{mft_axis}]"
         self._diag = diagnostic
         self._mft_axis = mft_axis
         self._data = None
         self._all_loaded = False
 
-        # Copy all relevant attributes from diagnostic
-        for attr in [
-            "_dt",
-            "_dx",
-            "_ndump",
-            "_axis",
-            "_nx",
-            "_x",
-            "_grid",
-            "_dim",
-            "_maxiter",
-            "_type",
-        ]:
+        for attr in ["_dt", "_dx", "_ndump", "_axis", "_nx", "_x", "_grid", "_dim", "_maxiter", "_type"]:
             if hasattr(diagnostic, attr):
                 setattr(self, attr, getattr(diagnostic, attr))
 
-    def load_all(self):
-        """Load all data and compute the average"""
-        if self._diag._all_loaded is True:
-            print("Diagnostic data already loaded ... applyting MFT")
-            self._data = self._diag._data
+    def load_all(self) -> np.ndarray:
         if self._data is not None:
-            print("Data already loaded")
             return self._data
 
-        if not hasattr(self._diag, "_data") or self._diag._data is None:
-            self._diag.load_all()
-
-        if self._mft_axis is None:
-            raise ValueError("Mean field theory axis must be specified.")
-        else:
-            self._data = np.expand_dims(self._diag._data.mean(axis=self._mft_axis), axis=-1)
-
+        self._diag.load_all()
+        # loaded data includes time at axis 0; spatial x1..xd are axes 1..d
+        ax = self._mft_axis  # OSIRIS axis -> numpy axis in loaded array
+        self._data = self._diag.data.mean(axis=ax, keepdims=True)
         self._all_loaded = True
         return self._data
 
-    def _data_generator(self, index):
-        """Generate average data for a specific index"""
-        if self._mft_axis is not None:
-            # Get the data for this index
-            data = self._diag[index]
-            # Compute the average (mean) along the specified axis
-            # Note: When accessing a slice, axis numbering is 0-based
-            avg = np.expand_dims(data.mean(axis=self._mft_axis - 1), axis=-1)
-            yield avg
-        else:
-            raise ValueError("Invalid axis for mean field theory.")
-
-    def __getitem__(self, index):
-        """Get average at a specific index"""
-        if self._all_loaded and self._data is not None:
-            return self._data[index]
-
-        # Otherwise compute on-demand
-        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:
+        # per-frame array has only spatial dims; x1..xd -> numpy 0..d-1
+        ax = self._mft_axis - 1
+
+        f = self._diag._frame(index, data_slice=data_slice)
+        return f.mean(axis=ax, keepdims=True)
 
 
 class MFT_Diagnostic_Fluctuations(Diagnostic):
     """
-    Class to compute the fluctuation component of mean field theory.
-    Inherits from Diagnostic to ensure all operation overloads work properly.
-
-    Parameters
-    ----------
-    diagnostic : Diagnostic
-        The diagnostic object.
-    mft_axis : int
-        The axis to compute the mean field theory.
-
+    delta = f - avg, where avg uses keepdims to broadcast.
     """
 
-    def __init__(self, diagnostic, mft_axis):
-        # Initialize with the same species as the diagnostic
+    def __init__(self, diagnostic: Diagnostic, mft_axis: int):
         if hasattr(diagnostic, "_species"):
-            super().__init__(
-                simulation_folder=(diagnostic._simulation_folder if hasattr(diagnostic, "_simulation_folder") else None),
-                species=diagnostic._species,
-            )
+            super().__init__(simulation_folder=getattr(diagnostic, "_simulation_folder", None), species=diagnostic._species)
         else:
             super().__init__(None)
 
         if mft_axis is None:
-            raise ValueError("Mean field theory axis must be specified.")
+            raise ValueError("Mean field theory axis must be specified (1..dim).")
+        if not (1 <= mft_axis <= diagnostic._dim):
+            raise ValueError(f"mft_axis must be in 1..{diagnostic._dim}, got {mft_axis}")
 
         self.postprocess_name = "MFT_FLT"
-
-        self._name = f"MFT_delta[{diagnostic._name}, {mft_axis}]"
+        self._name = f"MFT_delta[{diagnostic._name}, x{mft_axis}]"
         self._diag = diagnostic
         self._mft_axis = mft_axis
         self._data = None
         self._all_loaded = False
 
-        # Copy all relevant attributes from diagnostic
-        for attr in [
-            "_dt",
-            "_dx",
-            "_ndump",
-            "_axis",
-            "_nx",
-            "_x",
-            "_grid",
-            "_dim",
-            "_maxiter",
-            "_type",
-        ]:
+        for attr in ["_dt", "_dx", "_ndump", "_axis", "_nx", "_x", "_grid", "_dim", "_maxiter", "_type"]:
             if hasattr(diagnostic, attr):
                 setattr(self, attr, getattr(diagnostic, attr))
 
-    def load_all(self):
-        """Load all data and compute the fluctuations"""
-        if self._diag._all_loaded is True:
-            print("Diagnostic data already loaded ... applyting MFT")
-            self._data = self._diag._data
+    def load_all(self) -> np.ndarray:
         if self._data is not None:
-            print("Data already loaded")
             return self._data
 
-        if not hasattr(self._diag, "_data") or self._diag._data is None:
-            self._diag.load_all()
-
-        if self._mft_axis is None:
-            raise ValueError("Mean field theory axis must be specified.")
-        else:
-            # Compute the average
-            avg = self._diag._data.mean(axis=self._mft_axis)
-            # Reshape avg for broadcasting
-            broadcast_shape = list(self._diag._data.shape)
-            broadcast_shape[self._mft_axis] = 1
-            avg_reshaped = avg.reshape(broadcast_shape)
-            # Compute the fluctuations
-            self._data = self._diag._data - avg_reshaped
-
+        self._diag.load_all()
+        ax = self._mft_axis  # spatial axis in loaded array (time is 0)
+        avg = self._diag.data.mean(axis=ax, keepdims=True)
+        self._data = self._diag.data - avg
         self._all_loaded = True
         return self._data
 
-    def _data_generator(self, index):
-        """Generate fluctuation data for a specific index"""
-        if self._mft_axis is not None:
-            # Get the data for this index
-            data = self._diag[index]
-            # Compute the average (mean) along the specified axis
-            # Note: When accessing a slice, axis numbering is 0-based
-            avg = data.mean(axis=self._mft_axis - 1)
-            # Expand dimensions to enable broadcasting
-            avg_reshaped = np.expand_dims(avg, axis=self._mft_axis - 1)
-            # Compute fluctuations
-            delta = data - avg_reshaped
-            yield delta
-        else:
-            raise ValueError("Invalid axis for mean field theory.")
-
-    def __getitem__(self, index):
-        """Get fluctuations at a specific index"""
-        if self._all_loaded and self._data is not None:
-            return self._data[index]
-
-        # Otherwise compute on-demand
-        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:
+        ax = self._mft_axis - 1
+        f = self._diag._frame(index, data_slice=data_slice)
+        avg = f.mean(axis=ax, keepdims=True)
+        return f - avg
 
 
 class MFT_Species_Handler:

osiris_utils/postprocessing/postprocess.py

@@ -1,46 +1,62 @@
 from __future__ import annotations
 
-from ..data.simulation import Simulation
+from typing import Any
+
+from ..data.diagnostic import Diagnostic
 
 __all__ = ["PostProcess"]
 
 
-class PostProcess(Simulation):
+class PostProcess:
     """
-    Base class for post-processing operations.
-    Inherits from Diagnostic to ensure all operation overloads work.
-
-    Parameters
-    ----------
-    name : str
-        Name of the post-processing operation.
-    species : str
-        The species to analyze.
+    Base class for post-processing wrappers (FFT_Simulation, Derivative_Simulation, etc).
+
+    It wraps a Simulation-like object (Simulation or another PostProcess) and exposes:
+      - .species
+      - .loaded_diagnostics
+      - .add_diagnostic()
+
+    Subclasses implement __getitem__ to return processed diagnostics.
     """
 
-    def __init__(self, name: str, species: str = None):
-        # Initialize with the same interface as Diagnostic
+    def __init__(self, name: str, simulation: Any):
         self._name = name
-        self._all_loaded = False
-        self._data = None
+        self._simulation = simulation
+
+        # Simulation-like capability checks (supports chaining)
+        if not hasattr(simulation, "__getitem__"):
+            raise TypeError("Wrapped object must support __getitem__ like Simulation.")
+        if not (hasattr(simulation, "_species") or hasattr(simulation, "species")):
+            raise TypeError("Wrapped object must have _species or species attribute.")
+
+        # normalize species
+        self._species = getattr(simulation, "_species", simulation.species)
+
+    @property
+    def species(self) -> list:
+        return self._species
+
+    @property
+    def loaded_diagnostics(self) -> dict:
+        """
+        Delegate to wrapped simulation if it has loaded_diagnostics.
+        Otherwise return empty dict.
+        """
+        return getattr(self._simulation, "loaded_diagnostics", {})
 
-    def process(self, simulation: Simulation) -> Simulation:
+    def add_diagnostic(self, diagnostic: Diagnostic, name: str | None = None) -> str:
         """
-        Apply the post-processing to a diagnostic.
-        Must be implemented by subclasses.
-
-        Parameters
-        ----------
-        diagnostic : Diagnostic
-            The diagnostic to process.
-
-        Returns
-        -------
-        Diagnostic or PostProcess
-            The processed diagnostic.
+        Delegate add_diagnostic downwards.
+        This lets you do: sim_fft.add_diagnostic(...)
         """
-        raise NotImplementedError("Subclasses must implement process method")
+        if not hasattr(self._simulation, "add_diagnostic"):
+            raise AttributeError("Wrapped simulation does not support add_diagnostic().")
+        return self._simulation.add_diagnostic(diagnostic, name=name)
 
+    def delete_all_diagnostics(self) -> None:
+        if hasattr(self._simulation, "delete_all_diagnostics"):
+            self._simulation.delete_all_diagnostics()
 
-# PostProcessing_Simulation
-# PostProcessing_Diagnostic
+    def delete_diagnostic(self, key: str) -> None:
+        if hasattr(self._simulation, "delete_diagnostic"):
+            self._simulation.delete_diagnostic(key)

osiris_utils/postprocessing/pressure_correction.py

@@ -1,6 +1,5 @@
 from __future__ import annotations
 
-from collections.abc import Generator
 from typing import Any
 
 import numpy as np
@@ -9,7 +8,7 @@ from ..data.diagnostic import Diagnostic
 from ..data.simulation import Simulation
 from .postprocess import PostProcess
 
-OSIRIS_P = ["P11", "P12", "P13", "P21", "P22", "P23", "P31", "P32", "P33"]
+OSIRIS_P = ["P11", "P12", "P13", "P22", "P23", "P33"]
 
 __all__ = [
     "PressureCorrection_Simulation",
@@ -20,7 +19,7 @@ __all__ = [
 
 class PressureCorrection_Simulation(PostProcess):
     def __init__(self, simulation: Simulation):
-        super().__init__("PressureCorrection Simulation")
+        super().__init__("PressureCorrection Simulation", simulation)
         """
         Class to correct pressure tensor components by subtracting Reynolds stress.
 
@@ -31,9 +30,6 @@ class PressureCorrection_Simulation(PostProcess):
         pressure : str
             The pressure component to center.
         """
-        if not isinstance(simulation, Simulation):
-            raise ValueError("simulation must be a Simulation-compatible object.")
-        self._simulation = simulation
         self._pressure_corrected: dict[str, PressureCorrection_Diagnostic] = {}
         self._species_handler: dict[str, PressureCorrection_Species_Handler] = {}
 
@@ -45,8 +41,7 @@ class PressureCorrection_Simulation(PostProcess):
         if key not in OSIRIS_P:
             raise ValueError(f"Invalid pressure component {key}. Supported: {OSIRIS_P}.")
         if key not in self._pressure_corrected:
-            print("Weird that it got here - pressure is always species dependent on OSIRIS")
-            # self._pressure_corrected[key] = PressureCorrection_Diagnostic(self._simulation[key], self._simulation)
+            raise KeyError("Pressure diagnostics are species-dependent. Use sim['species']['P12'] not sim['P12'].")
         return self._pressure_corrected[key]
 
     def delete_all(self) -> None:
@@ -58,26 +53,13 @@ class PressureCorrection_Simulation(PostProcess):
         else:
             print(f"Pressure {key} not found in simulation")
 
-    def process(self, diagnostic: Diagnostic) -> PressureCorrection_Diagnostic:
-        """Apply pressure correction to a diagnostic"""
-        # FIX: This is a bit of a hack, but it works for now
-        return PressureCorrection_Diagnostic(diagnostic, self._simulation)
-
 
 class PressureCorrection_Diagnostic(Diagnostic):
     def __init__(self, diagnostic: Diagnostic, n: Diagnostic, ufl_j: Diagnostic, vfl_k: Diagnostic):
-        """
-        Class to correct the pressure in the simulation.
-
-        Parameters
-        ----------
-        diagnostic : Diagnostic
-            The diagnostic object.
-        """
         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)
@@ -88,8 +70,6 @@ class PressureCorrection_Diagnostic(Diagnostic):
             raise ValueError(f"Invalid pressure component {diagnostic._name}. Supported: {OSIRIS_P}")
 
         self._diag = diagnostic
-
-        # The density and velocities are now passed as arguments (so it can doesn't depend on the simulation)
         self._n = n
         self._ufl_j = ufl_j
         self._vfl_k = vfl_k
@@ -119,46 +99,31 @@ class PressureCorrection_Diagnostic(Diagnostic):
         if self._data is not None:
             return self._data
 
-        if not hasattr(self._diag, "_data") or self._diag._data is None:
+        # Ensure all operands are loaded
+        if not getattr(self._diag, "_all_loaded", False) or self._diag._data is None:
             self._diag.load_all()
-
-        print(f"Loading {self._species._name} {self._original_name} diagnostic")
-        self._n.load_all()
-        self._ufl_j.load_all()
-        self._vfl_k.load_all()
-
-        # Then access the data
-        n = self._n.data
-        u = self._ufl_j.data
-        v = self._vfl_k.data
-
-        self._data = self._diag.data - n * v * u
+        if not getattr(self._n, "_all_loaded", False) or self._n._data is None:
+            self._n.load_all()
+        if not getattr(self._ufl_j, "_all_loaded", False) or self._ufl_j._data is None:
+            self._ufl_j.load_all()
+        if not getattr(self._vfl_k, "_all_loaded", False) or self._vfl_k._data is None:
+            self._vfl_k.load_all()
+
+        # Core formula: Pjk - n * u_j * v_k
+        self._data = self._diag._data - self._n._data * self._ufl_j._data * self._vfl_k._data
         self._all_loaded = True
-
-        # Unload the data to save memory
-        # self._n.unload()
-        # self._ufl_j.unload()
-        # self._vfl_k.unload()
-
         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 _data_generator(self, index: int) -> Generator[np.ndarray, None, None]:
-        yield (self._diag[index] - self._n[index] * self._vfl_k[index] * self._ufl_j[index])
+    def _frame(self, index: int, data_slice: tuple | None = None) -> np.ndarray:
+        """
+        Lazy per-timestep evaluation, supports spatial slicing.
+        data_slice applies to spatial axes only (same convention as Diagnostic).
+        """
+        P = self._diag._frame(index, data_slice=data_slice)
+        n = self._n._frame(index, data_slice=data_slice)
+        u = self._ufl_j._frame(index, data_slice=data_slice)
+        v = self._vfl_k._frame(index, data_slice=data_slice)
+        return P - n * u * v
 
 
 class PressureCorrection_Species_Handler:

osiris_utils/vis/plot3d.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import warnings
 from typing import Literal
 

{osiris_utils-1.2.2.dist-info → osiris_utils-1.2.4.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: osiris_utils
-Version: 1.2.2
+Version: 1.2.4
 Summary: Utilities to manipulate and visualise OSIRIS plasma PIC output data
 Author-email: João Pedro Ferreira Biu <joaopedrofbiu@tecnico.ulisboa.pt>
 License: MIT