osiris-utils 1.1.4__py3-none-any.whl → 1.1.7__py3-none-any.whl

osiris_utils/postprocessing/heatflux_correction.py

@@ -0,0 +1,200 @@
+import numpy as np
+
+from ..data.diagnostic import Diagnostic
+from ..data.simulation import Simulation
+from .postprocess import PostProcess
+from .pressure_correction import PressureCorrection_Simulation
+
+OSIRIS_H = ["q1", "q2", "q3"]
+
+
+class HeatfluxCorrection_Simulation(PostProcess):
+    def __init__(self, simulation):
+        super().__init__("HeatfluxCorrection Simulation")
+        """
+        Class to correct pressure tensor components by subtracting Reynolds stress.
+
+        Parameters
+        ----------
+        sim : Simulation
+            The simulation object.
+        heatflux : str
+            The heatflux component to center.
+        """
+        if not isinstance(simulation, Simulation):
+            raise ValueError("Simulation must be a Simulation object.")
+        self._simulation = simulation
+        self._heatflux_corrected = {}
+        self._species_handler = {}
+
+    def __getitem__(self, key):
+        if key in self._simulation._species:
+            if key not in self._species_handler:
+                self._species_handler[key] = HeatfluxCorrection_Species_Handler(self._simulation[key], self._simulation)
+            return self._species_handler[key]
+        if key not in OSIRIS_H:
+            raise ValueError(f"Invalid heatflux component {key}. Supported: {OSIRIS_H}.")
+        if key not in self._heatflux_corrected:
+            print("Weird that it got here - heatflux is always species dependent on OSIRIS")
+            self._heatflux_corrected[key] = HeatfluxCorrection_Diagnostic(self._simulation[key], self._simulation)
+        return self._heatflux_corrected[key]
+
+    def delete_all(self):
+        self._heatflux_corrected = {}
+
+    def delete(self, key):
+        if key in self._heatflux_corrected:
+            del self._heatflux_corrected[key]
+        else:
+            print(f"Heatflux {key} not found in simulation")
+
+    def process(self, diagnostic):
+        """Apply heatflux correction to a diagnostic"""
+        return HeatfluxCorrection_Diagnostic(diagnostic, self._simulation)
+
+
+class HeatfluxCorrection_Diagnostic(Diagnostic):
+    def __init__(self, diagnostic, vfl_i, Pjj_list, vfl_j_list, Pji_list):
+        """
+        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,
+            )
+        else:
+            super().__init__(None)
+
+        self.postprocess_name = "HFL_CORR"
+
+        if diagnostic._name not in OSIRIS_H:
+            raise ValueError(f"Invalid heatflux component {diagnostic._name}. Supported: {OSIRIS_H}")
+
+        self._diag = diagnostic
+
+        # The density and velocities are now passed as arguments (so it can doesn't depend on the simulation)
+        self._vfl_i = vfl_i
+        self._Pjj_list = Pjj_list
+        self._vfl_j_list = vfl_j_list
+        self._Pji_list = Pji_list
+
+        for attr in [
+            "_dt",
+            "_dx",
+            "_ndump",
+            "_axis",
+            "_nx",
+            "_x",
+            "_grid",
+            "_dim",
+            "_maxiter",
+            "_type",
+        ]:
+            if hasattr(diagnostic, attr):
+                setattr(self, attr, getattr(diagnostic, attr))
+
+        self._original_name = diagnostic._name
+        self._name = diagnostic._name + "_corrected"
+
+        self._data = None
+        self._all_loaded = False
+
+    def load_all(self):
+        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()
+
+        print(f"Loading {self._species._name} {self._original_name} diagnostic")
+
+        self._vfl_i.load_all()
+
+        for vfl_j in self._vfl_j_list:
+            vfl_j.load_all()
+        for Pji in self._Pji_list:
+            Pji.load_all()
+        for Pjj in self._Pjj_list:
+            Pjj.load_all()
+
+        q = self._diag.data
+        vfl_i = self._vfl_i.data
+
+        trace_P = sum(Pjj.data for Pjj in self._Pjj_list)
+
+        # 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))
+
+        self._data = 2 * q - 0.5 * vfl_i * trace_P - vfl_dot_Pji
+        self._all_loaded = True
+
+        return self._data
+
+    def __getitem__(self, index):
+        """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):
+        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
+
+
+class HeatfluxCorrection_Species_Handler:
+    """
+    Class to handle heatflux correction for a species.
+    Acts as a wrapper for the HeatfluxCorrection_Diagnostic class.
+
+    Not intended to be used directly, but through the HeatfluxCorrection_Simulation class.
+
+    Parameters
+    ----------
+    species_handler : Species_Handler
+        The species handler object.
+    simulation : Simulation
+        The simulation object.
+    """
+
+    def __init__(self, species_handler, simulation):
+        self._species_handler = species_handler
+        self._simulation = simulation
+        self._heatflux_corrected = {}
+
+    def __getitem__(self, key):
+        if key not in self._heatflux_corrected:
+            diag = self._species_handler[key]
+
+            # Velocities alwayes depend on the species so this can be done here
+
+            i = int(key[-1])  # Get i from 'q1', 'q2', etc.
+
+            vfl_i = self._species_handler[f"vfl{i}"]
+
+            # Load trace(P): sum over Pjj
+            Pjj_list = [self._species_handler[f"P{j}{j}"] for j in range(1, diag._dim + 1)]
+
+            # 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)]
+
+            self._heatflux_corrected[key] = HeatfluxCorrection_Diagnostic(diag, vfl_i, Pjj_list, vfl_j_list, Pji_list)
+        return self._heatflux_corrected[key]

osiris_utils/postprocessing/mft.py

@@ -1,11 +1,13 @@
-from ..utils import *
+import numpy as np
+
+from ..data.diagnostic import Diagnostic
 from ..data.simulation import Simulation
 from .postprocess import PostProcess
-from ..data.diagnostic import Diagnostic
 
-class MeanFieldTheory_Simulation(PostProcess):
+
+class MFT_Simulation(PostProcess):
     """
-    Class to compute the mean field theory of a diagnostic. Works as a wrapper for the MFT_Diagnostic class.
+    Class to compute the Mean Field Theory approximation of a diagnostic. Works as a wrapper for the MFT_Diagnostic class.
     Inherits from PostProcess to ensure all operation overloads work properly.
 
     Parameters
@@ -15,11 +17,6 @@ class MeanFieldTheory_Simulation(PostProcess):
     mft_axis : int
         The axis to compute the mean field theory.
 
-    Example
-    -------
-    >>> sim = Simulation('electrons', 'path/to/simulation')
-    >>> mft = MeanFieldTheory(sim, 1)
-    >>> mft_e1 = mft['e1']
     """
 
     def __init__(self, simulation, mft_axis=None):
@@ -39,7 +36,7 @@ class MeanFieldTheory_Simulation(PostProcess):
         if key not in self._mft_computed:
             self._mft_computed[key] = MFT_Diagnostic(self._simulation[key], self._mft_axis)
         return self._mft_computed[key]
-    
+
     def delete_all(self):
         self._mft_computed = {}
 
@@ -53,6 +50,7 @@ class MeanFieldTheory_Simulation(PostProcess):
         """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.
@@ -65,45 +63,54 @@ class MFT_Diagnostic(Diagnostic):
     mft_axis : int
         The axis to compute mean field theory along.
 
-    Example
-    -------
-    >>> sim = Simulation('electrons', 'path/to/simulation')
-    >>> diag = sim['e1']
-    >>> mft = MFT_Diagnostic(diag, 1)
-    >>> avg = mft['avg']
-    >>> delta = mft['delta']
+
     """
 
     def __init__(self, diagnostic, mft_axis):
         # Initialize using parent's __init__ with the same species
-        if hasattr(diagnostic, '_species'):
-            super().__init__(species=diagnostic._species)
+        if hasattr(diagnostic, "_species"):
+            super().__init__(
+                simulation_folder=(diagnostic._simulation_folder if hasattr(diagnostic, "_simulation_folder") else None),
+                species=diagnostic._species,
+            )
         else:
             super().__init__(None)
-            
+
         self._name = f"MFT[{diagnostic._name}]"
         self._diag = diagnostic
         self._mft_axis = mft_axis
         self._data = None
         self._all_loaded = False
-        
+
         # 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']:
+        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
@@ -113,33 +120,34 @@ class MFT_Diagnostic(Diagnostic):
             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 "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
 
+
 class MFT_Diagnostic_Average(Diagnostic):
     """
     Class to compute the average component of mean field theory.
@@ -152,23 +160,23 @@ class MFT_Diagnostic_Average(Diagnostic):
     mft_axis : int
         The axis to compute the mean field theory.
 
-    Example
-    -------
-    >>> sim = Simulation('electrons', 'path/to/simulation')
-    >>> diag = sim['e1']
-    >>> avg = MFT_Diagnostic_Average(diag, 1)
     """
 
     def __init__(self, diagnostic, mft_axis):
         # Initialize with the same species as the diagnostic
-        if hasattr(diagnostic, '_species'):
-            super().__init__(species=diagnostic._species)
+        if hasattr(diagnostic, "_species"):
+            super().__init__(
+                simulation_folder=(diagnostic._simulation_folder if hasattr(diagnostic, "_simulation_folder") else None),
+                species=diagnostic._species,
+            )
         else:
             super().__init__(None)
-            
+
         if mft_axis is None:
             raise ValueError("Mean field theory axis must be specified.")
-        
+
+        self.postprocess_name = "MFT_AVG"
+
         self._name = f"MFT_avg[{diagnostic._name}, {mft_axis}]"
         self._diag = diagnostic
         self._mft_axis = mft_axis
@@ -176,17 +184,31 @@ class MFT_Diagnostic_Average(Diagnostic):
         self._all_loaded = False
 
         # Copy all relevant attributes from diagnostic
-        for attr in ['_dt', '_dx', '_ndump', '_axis', '_nx', '_x', '_grid', '_dim', '_maxiter']:
+        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
         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:
+
+        if not hasattr(self._diag, "_data") or self._diag._data is None:
             self._diag.load_all()
 
         if self._mft_axis is None:
@@ -196,7 +218,7 @@ class MFT_Diagnostic_Average(Diagnostic):
 
         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:
@@ -204,16 +226,16 @@ class MFT_Diagnostic_Average(Diagnostic):
             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)
+            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))
@@ -224,7 +246,8 @@ class MFT_Diagnostic_Average(Diagnostic):
             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.")
-    
+
+
 class MFT_Diagnostic_Fluctuations(Diagnostic):
     """
     Class to compute the fluctuation component of mean field theory.
@@ -237,23 +260,23 @@ class MFT_Diagnostic_Fluctuations(Diagnostic):
     mft_axis : int
         The axis to compute the mean field theory.
 
-    Example
-    -------
-    >>> sim = Simulation('electrons', 'path/to/simulation')
-    >>> diag = sim['e1']
-    >>> delta = MFT_Diagnostic_Fluctuations(diag, 1)
     """
-    
+
     def __init__(self, diagnostic, mft_axis):
         # Initialize with the same species as the diagnostic
-        if hasattr(diagnostic, '_species'):
-            super().__init__(species=diagnostic._species)
+        if hasattr(diagnostic, "_species"):
+            super().__init__(
+                simulation_folder=(diagnostic._simulation_folder if hasattr(diagnostic, "_simulation_folder") else None),
+                species=diagnostic._species,
+            )
         else:
             super().__init__(None)
-            
+
         if mft_axis is None:
             raise ValueError("Mean field theory axis must be specified.")
-        
+
+        self.postprocess_name = "MFT_FLT"
+
         self._name = f"MFT_delta[{diagnostic._name}, {mft_axis}]"
         self._diag = diagnostic
         self._mft_axis = mft_axis
@@ -261,17 +284,31 @@ class MFT_Diagnostic_Fluctuations(Diagnostic):
         self._all_loaded = False
 
         # Copy all relevant attributes from diagnostic
-        for attr in ['_dt', '_dx', '_ndump', '_axis', '_nx', '_x', '_grid', '_dim', '_maxiter']:
+        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
         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:
+
+        if not hasattr(self._diag, "_data") or self._diag._data is None:
             self._diag.load_all()
 
         if self._mft_axis is None:
@@ -288,7 +325,7 @@ class MFT_Diagnostic_Fluctuations(Diagnostic):
 
         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:
@@ -296,20 +333,20 @@ class MFT_Diagnostic_Fluctuations(Diagnostic):
             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)
+            avg = data.mean(axis=self._mft_axis - 1)
             # Expand dimensions to enable broadcasting
-            avg_reshaped = np.expand_dims(avg, axis=self._mft_axis-1)
+            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))
@@ -320,7 +357,8 @@ class MFT_Diagnostic_Fluctuations(Diagnostic):
             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.")
-    
+
+
 class MFT_Species_Handler:
     """
     Class to handle mean field theory for a species.
@@ -335,7 +373,7 @@ class MFT_Species_Handler:
     mft_axis : int
         The axis to compute the mean field theory.
     """
-    
+
     def __init__(self, species_handler, mft_axis):
         self._species_handler = species_handler
         self._mft_axis = mft_axis
@@ -345,4 +383,4 @@ class MFT_Species_Handler:
         if key not in self._mft_computed:
             diag = self._species_handler[key]
             self._mft_computed[key] = MFT_Diagnostic(diag, self._mft_axis)
-        return self._mft_computed[key]
+        return self._mft_computed[key]

osiris_utils/postprocessing/mft_for_gridfile.py

@@ -1,18 +1,21 @@
 import numpy as np
+
 from ..data.data import OsirisGridFile
 
+
 # Deprecated
 class MFT_Single(OsirisGridFile):
-    '''
+    """
     Class to handle the mean field theory on data. Inherits from OsirisGridFile.
-    
+
     Parameters
     ----------
     source : str or OsirisGridFile
         The filename or an OsirisGridFile object.
     axis : int
         The axis to average over.
-    '''
+    """
+
     def __init__(self, source, axis=1):
         if isinstance(source, OsirisGridFile):
             self.__dict__.update(source.__dict__)
@@ -30,23 +33,23 @@ class MFT_Single(OsirisGridFile):
     @property
     def average(self):
         return self._average
-    
+
     @property
     def delta(self):
         return self._fluctuations
-    
+
     def __str__(self):
-        return super().__str__() + f'\nAverage: {self.average.shape}\nDelta: {self.delta.shape}'
-    
+        return super().__str__() + f"\nAverage: {self.average.shape}\nDelta: {self.delta.shape}"
+
     def derivative(self, field, axis=0):
-        '''
+        """
         Compute the derivative of the average or the fluctuations.
-        
+
         Parameters
         ----------
         field : MeanFieldTheory.average or MeanFieldTheory.delta
             The field to compute the derivative.
         axis : int
             The axis to compute the derivative.
-        '''
-        return np.gradient(field, self.dx[axis], axis=0)
+        """
+        return np.gradient(field, self.dx[axis], axis=0)

osiris_utils/postprocessing/postprocess.py

@@ -1,10 +1,11 @@
 from ..data.diagnostic import Diagnostic
 
+
 class PostProcess(Diagnostic):
     """
     Base class for post-processing operations.
     Inherits from Diagnostic to ensure all operation overloads work.
-    
+
     Parameters
     ----------
     name : str
@@ -12,31 +13,31 @@ class PostProcess(Diagnostic):
     species : str
         The species to analyze.
     """
-    
+
     def __init__(self, name, species=None):
         # Initialize with the same interface as Diagnostic
         super().__init__(species)
         self._name = name
         self._all_loaded = False
         self._data = None
-        
+
     def process(self, diagnostic):
         """
         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.
         """
         raise NotImplementedError("Subclasses must implement process method")
-    
+
 
 # PostProcessing_Simulation
-# PostProcessing_Diagnostic
+# PostProcessing_Diagnostic