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

osiris_utils/data/simulation.py

@@ -31,15 +31,7 @@ class Simulation:
     __getitem__(key)
         Get a diagnostic.
 
-    Example
-    -------
-    >>> sim = Simulation('electrons', 'path/to/simulation')
-    >>> diag = sim['e1']
-    >>> diag.load_all()
-    
-    >>> sim = Simulation('electrons', 'path/to/simulation')
-    >>> diag = sim['e1']
-    >>> diag[<index>]
+        
     '''
     def __init__(self, input_deck_path):
         folder_path = os.path.dirname(input_deck_path)
@@ -135,6 +127,10 @@ class Simulation:
     @property
     def species(self):
         return self._species
+    
+    @property
+    def loaded_diagnostics(self):
+        return self._diagnostics
 
 # This is to handle species related diagnostics
 class Species_Handler:
@@ -180,13 +176,7 @@ class Species_Handler:
         -------
         str
             The name (key) used to store the diagnostic
-            
-        Example
-        -------
-        >>> sim = Simulation('path/to/simulation', 'input_deck.txt')
-        >>> nT = sim['electrons']['n'] * sim['electrons']['T11']
-        >>> sim.add_diagnostic(nT, 'nT')
-        >>> sim['nT']  # Access the custom diagnostic
+        
         """
         # Generate a name if none provided
         if name is None:
@@ -200,4 +190,27 @@ class Species_Handler:
         if isinstance(diagnostic, Diagnostic):
             self._diagnostics[name] = diagnostic
         else:
-            raise ValueError("Only Diagnostic objects are supported for now")
+            raise ValueError("Only Diagnostic objects are supported for now")
+        
+    def delete_diagnostic(self, key):
+        """
+        Delete a diagnostic.
+        """
+        if key in self._diagnostics:
+            del self._diagnostics[key]
+        else:
+            print(f"Diagnostic {key} not found in species {self._species_name}")
+            return None
+
+    def delete_all_diagnostics(self):
+        """
+        Delete all diagnostics.
+        """
+        self._diagnostics = {}
+
+    @property
+    def species(self):
+        return self._species_name
+    @property
+    def loaded_diagnostics(self):
+        return self._diagnostics

osiris_utils/postprocessing/derivative.py

@@ -12,7 +12,7 @@ class Derivative_Simulation(PostProcess):
     ----------
     simulation : Simulation
         The simulation object.
-    type : str
+    deriv_type : str
         The type of derivative to compute. Options are:
         - 't' for time derivative.
         - 'x1' for first spatial derivative.
@@ -24,19 +24,14 @@ class Derivative_Simulation(PostProcess):
     axis : int or tuple
         The axis to compute the derivative. Only used for 'xx', 'xt' and 'tx' types.
 
-    Example
-    -------
-    >>> sim = Simulation('electrons', 'path/to/simulation')
-    >>> derivative = Derivative(sim, 'x1')
-    >>> deriv_e1_wrt_x1 = derivative['e1']
     """
 
-    def __init__(self, simulation, type, axis=None):
-        super().__init__(f"Derivative({type})")
+    def __init__(self, simulation, deriv_type, axis=None):
+        super().__init__(f"Derivative({deriv_type})")
         if not isinstance(simulation, Simulation):
             raise ValueError("Simulation must be a Simulation object.")
         self._simulation = simulation
-        self._type = type
+        self._deriv_type = deriv_type
         self._axis = axis
         self._derivatives_computed = {}
         self._species_handler = {}
@@ -44,12 +39,12 @@ class Derivative_Simulation(PostProcess):
     def __getitem__(self, key):
         if key in self._simulation._species:
             if key not in self._species_handler:
-                self._species_handler[key] = Derivative_Species_Handler(self._simulation[key], self._type, self._axis)
+                self._species_handler[key] = Derivative_Species_Handler(self._simulation[key], self._deriv_type, self._axis)
             return self._species_handler[key]
                 
         if key not in self._derivatives_computed:
             self._derivatives_computed[key] = Derivative_Diagnostic(diagnostic=self._simulation[key], 
-                                                                    type=self._type, axis=self._axis)
+                                                                    deriv_type=self._deriv_type, axis=self._axis)
         return self._derivatives_computed[key]
     
     def delete_all(self):
@@ -63,7 +58,7 @@ class Derivative_Simulation(PostProcess):
     
     def process(self, diagnostic):
         """Apply derivative to a diagnostic"""
-        return Derivative_Diagnostic(diagnostic, self._type, self._axis)
+        return Derivative_Diagnostic(diagnostic, self._deriv_type, self._axis)
 
     
 class Derivative_Diagnostic(Diagnostic):
@@ -75,7 +70,7 @@ class Derivative_Diagnostic(Diagnostic):
     ----------
     diagnostic : Diagnostic
         The diagnostic object.
-    type : str
+    deriv_type : str
         The type of derivative to compute. Options are: 't', 'x1', 'x2', 'x3', 'xx', 'xt' and 'tx'.
     axis : int or tuple
         The axis to compute the derivative. Only used for 'xx', 'xt' and 'tx' types
@@ -87,45 +82,30 @@ class Derivative_Diagnostic(Diagnostic):
     __getitem__(index)
         Get data at a specific index.
 
-    Example
-    -------
-    >>> sim = Simulation('electrons', 'path/to/simulation')
-    >>> diag = sim['e1']
-    >>> derivative = Derivative_Diagnostic(diag, 'x1')
     """
 
-    def __init__(self, diagnostic, type, axis=None):
+    def __init__(self, diagnostic, deriv_type, axis=None):
         # 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)
         else:
             super().__init__(None)
+
+        self.postprocess_name = f"DERIV"
             
         # self._name = f"D[{diagnostic._name}, {type}]"
         self._diag = diagnostic
-        self._type = type
+        self._deriv_type = deriv_type
         self._axis = axis if axis is not None else diagnostic._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']:
+        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_metadata(self):
-        """Copy metadata from original diagnostic to ensure consistency"""
-        self._dt = self._diag._dt
-        self._dx = self._diag._dx
-        self._ndump = self._diag._ndump
-        self._axis = self._diag._axis
-        self._nx = self._diag._nx
-        self._x = self._diag._x
-        self._grid = self._diag._grid
-        self._dim = self._diag._dim
-        self._maxiter = self._diag._maxiter
-
     def load_all(self):
         """Load all data and compute the derivative"""
         if self._data is not None:
@@ -135,34 +115,34 @@ class Derivative_Diagnostic(Diagnostic):
         if not hasattr(self._diag, '_data') or self._diag._data is None:
             self._diag.load_all()
 
-        if self._type == "t":
+        if self._deriv_type == "t":
             result = np.gradient(self._diag._data, self._diag._dt * self._diag._ndump, axis=0, edge_order=2)
 
-        elif self._type == "x1":
+        elif self._deriv_type == "x1":
             if self._dim == 1:
                 result = np.gradient(self._diag._data, self._diag._dx, axis=1, edge_order=2)
             else:
                 result = np.gradient(self._diag._data, self._diag._dx[0], axis=1, edge_order=2)
                     
-        elif self._type == "x2":
-            result = np.gradient(self._diag._data, self._diag._dx[0], axis=2, edge_order=2)
+        elif self._deriv_type == "x2":
+            result = np.gradient(self._diag._data, self._diag._dx[1], axis=2, edge_order=2)
 
-        elif self._type == "x3":
-            result = np.gradient(self._diag._data, self._diag._dx[0], axis=3, edge_order=2)
+        elif self._deriv_type == "x3":
+            result = np.gradient(self._diag._data, self._diag._dx[2], axis=3, edge_order=2)
 
-        elif self._type == "xx":
+        elif self._deriv_type == "xx":
             if len(self._axis) != 2:
                 raise ValueError("Axis must be a tuple with two elements.")
             result = np.gradient(np.gradient(self._diag._data, self._diag._dx[self._axis[0]-1], axis=self._axis[0], edge_order=2), 
                                 self._diag._dx[self._axis[1]-1], axis=self._axis[1], edge_order=2)
             
-        elif self._type == "xt":
+        elif self._deriv_type == "xt":
             if not isinstance(self._axis, int):
                 raise ValueError("Axis must be an integer.")
             result = np.gradient(np.gradient(self._diag._data, self._diag._dt, axis=0, edge_order=2), 
                                 self._diag._dx[self._axis-1], axis=self._axis[0], edge_order=2)
             
-        elif self._type == "tx":
+        elif self._deriv_type == "tx":
             if not isinstance(self._axis, int):
                 raise ValueError("Axis must be an integer.")
             result = np.gradient(np.gradient(self._diag._data, self._diag._dx[self._axis-1], axis=self._axis, edge_order=2), 
@@ -177,19 +157,19 @@ class Derivative_Diagnostic(Diagnostic):
 
     def _data_generator(self, index):
         """Generate data for a specific index on-demand"""
-        if self._type == "x1":
+        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._type == "x2":
+        elif self._deriv_type == "x2":
             yield np.gradient(self._diag[index], self._diag._dx[1], axis=1, edge_order=2)
         
-        elif self._type == "x3":
+        elif self._deriv_type == "x3":
             yield np.gradient(self._diag[index], self._diag._dx[2], axis=2, edge_order=2)
             
-        elif self._type == "t":
+        elif self._deriv_type == "t":
             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:
@@ -230,14 +210,14 @@ class Derivative_Species_Handler:
     axis : int or tuple
         The axis to compute the derivative. Only used for 'xx', 'xt' and 'tx' types.
     """
-    def __init__(self, species_handler, type, axis=None):
+    def __init__(self, species_handler, deriv_type, axis=None):
         self._species_handler = species_handler
-        self._type = type
+        self._deriv_type = deriv_type
         self._axis = axis
         self._derivatives_computed = {}
 
     def __getitem__(self, key):
         if key not in self._derivatives_computed:
             diag = self._species_handler[key]
-            self._derivatives_computed[key] = Derivative_Diagnostic(diag, self._type, self._axis)
+            self._derivatives_computed[key] = Derivative_Diagnostic(diag, self._deriv_type, self._axis)
         return self._derivatives_computed[key]

osiris_utils/postprocessing/fft.py

@@ -19,11 +19,6 @@ class FastFourierTransform_Simulation(PostProcess):
     axis : int
         The axis to compute the FFT.
 
-    Example
-    -------
-    >>> sim = Simulation('electrons', 'path/to/simulation')
-    >>> fft = FastFourierTransform(sim, 1)
-    >>> fft_e1 = fft['e1']
     """
     def __init__(self, simulation, fft_axis):
         super().__init__("FFT")
@@ -78,12 +73,7 @@ class FFT_Diagnostic(Diagnostic):
         Get the angular frequency array for the FFT.
     __getitem__(index)
         Get data at a specific index.
-    
-    Example
-    -------
-    >>> sim = Simulation('electrons', 'path/to/simulation')
-    >>> diag = sim['e1']
-    >>> fft = FFT_Diagnostic(diag, 1)
+
     """
     def __init__(self, diagnostic, fft_axis):
         if hasattr(diagnostic, '_species'):
@@ -92,6 +82,8 @@ class FFT_Diagnostic(Diagnostic):
         else:
             super().__init__(None)
 
+        self.postprocess_name = f"FFT"
+
         self._name = f"FFT[{diagnostic._name}, {fft_axis}]"
         self._diag = diagnostic
         self._fft_axis = fft_axis
@@ -99,7 +91,7 @@ class FFT_Diagnostic(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))
 
@@ -110,7 +102,7 @@ class FFT_Diagnostic(Diagnostic):
     
     def load_all(self):
         if self._data is not None:
-            print("Using cached derivative")
+            print("Using cached data.")
             return self._data
         
         if not hasattr(self._diag, '_data') or self._diag._data is None:
@@ -119,7 +111,9 @@ class FFT_Diagnostic(Diagnostic):
 
        # Apply appropriate windows based on which axes we're transforming
         if isinstance(self._fft_axis, (list, tuple)):
-            # Multiple axes FFT
+            if self._diag._data is None:
+                raise ValueError(f"Unable to load data for diagnostic {self._diag._name}. The data is None even after loading.")
+
             result = self._diag._data.copy()
             
             for axis in self._fft_axis:

osiris_utils/postprocessing/field_centering.py

@@ -0,0 +1,168 @@
+from ..utils import *
+from ..data.simulation import Simulation
+from .postprocess import PostProcess
+from ..data.diagnostic import Diagnostic
+
+OSIRIS_FLD = ["e1", "e2", "e3", "b1", "b2", "b3"]
+
+
+class FieldCentering_Simulation(PostProcess):
+    """
+    Class to handle the field centering on data. Works as a wrapper for the FieldCentering_Diagnostic class.
+    Inherits from PostProcess to ensure all operation overloads work properly.
+    
+    Parameters
+    ----------
+    simulation : Simulation
+        The simulation object.
+    field : str
+        The field to center.
+    """
+    
+    def __init__(self, simulation: Simulation):
+        super().__init__(f"FieldCentering Simulation")
+        """
+        Class to center the field in the simulation.
+
+        Parameters
+        ----------
+        sim : Simulation
+            The simulation object.
+        field : str
+            The field to center.
+        """
+        if not isinstance(simulation, Simulation):
+            raise ValueError("Simulation must be a Simulation object.")
+        self._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):
+        if key not in OSIRIS_FLD:
+            raise ValueError(f"Does it make sense to center {key} field? Only {OSIRIS_FLD} are supported.")
+        if key not in self._field_centered:
+            self._field_centered[key] = FieldCentering_Diagnostic(self._simulation[key])
+        return self._field_centered[key]
+    
+    def delete_all(self):
+        self._field_centered = {}
+
+    def delete(self, key):
+        if key in self._field_centered:
+            del self._field_centered[key]
+        else:
+            print(f"Field {key} not found in simulation")
+    
+    def process(self, diagnostic):
+        """Apply field centering to a diagnostic"""
+        return FieldCentering_Diagnostic(diagnostic)
+    
+class FieldCentering_Diagnostic(Diagnostic):
+    def __init__(self, diagnostic):
+
+        """
+        Class to center the field 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 = "FLD_CTR"
+
+        if diagnostic._name not in OSIRIS_FLD:
+            raise ValueError(f"Does it make sense to center {diagnostic._name} field? Only {OSIRIS_FLD} are supported.")
+        
+        self._diag = diagnostic
+
+        for attr in ['_dt', '_dx', '_ndump', '_axis', '_nx', '_x', '_grid', '_dim', '_maxiter']:
+            if hasattr(diagnostic, attr):
+                setattr(self, attr, getattr(diagnostic, attr))
+
+        self._original_name = diagnostic._name
+        self._name = diagnostic._name + "_centered"
+
+        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()
+
+        if self._dim == 1:
+            if self._original_name.lower() in ['b2', 'b3', 'e1']:
+                result = 0.5 * (np.roll(self._diag.data, shift=1, axis=1) + self._diag.data) 
+            elif self._original_name.lower() in ['b1', 'e2', 'e3']:
+                result = self._diag.data 
+
+        elif self._dim == 2:
+            if self._original_name.lower() in ['e1', 'b2']:
+                result = 0.5 * (np.roll(self._diag.data, shift=1, axis=1) + self._diag.data)
+            elif self._original_name.lower() in ['e2', 'b1']:
+                result = 0.5 * (np.roll(self._diag.data, shift=1, axis=2) + self._diag.data)
+            elif self._original_name.lower() in ['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']:
+                result = self._diag.data
+            
+        elif self._dim == 3:
+            raise NotImplementedError("3D field centering is not implemented yet.")
+        
+        else:
+            raise ValueError(f"Unknown dimension {self._dim}.")
+        
+        self._data = result
+        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):
+        if self._dim == 1:
+            if self._original_name.lower() in ['b2', 'b3', 'e1']:
+                yield 0.5 * (np.roll(self._diag[index], shift=1) + self._diag[index])
+            elif self._original_name.lower() in ['b1', 'e2', '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', 'b2']:
+                yield 0.5 * (np.roll(self._diag[index], shift=1, axis=0) + self._diag[index])
+            elif self._original_name in ['e2', 'b1']:
+                yield 0.5 * (np.roll(self._diag[index], shift=1, axis=1) + self._diag[index])
+            elif self._original_name in ['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']:
+                yield self._diag[index]
+            else:
+                raise ValueError(f"Unknown field {self._original_name}.")
+            
+        elif self._dim == 3:
+            raise NotImplementedError("3D field centering is not implemented yet.")
+        
+        else:
+            raise ValueError(f"Unknown dimension {self._dim}.")