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

osiris_utils/data/simulation.py

@@ -1,18 +1,18 @@
-from matplotlib.style import available
-from ..data.diagnostic import *
-from ..utils import *
+import os
+
+from ..data.diagnostic import Diagnostic
 from ..decks.decks import InputDeckIO
 
 
 class Simulation:
-    '''
+    """
     Class to handle the simulation data. It is a wrapper for the Diagnostic class.'
-    
+
     Parameters
     ----------
     input_deck_path : str
         Path to the input deck (It must be in the folder where the simulation was run)
-    
+
     Attributes
     ----------
     simulation_folder : str
@@ -31,16 +31,9 @@ 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)
         self._input_deck_path = input_deck_path
@@ -51,7 +44,7 @@ class Simulation:
         self._simulation_folder = folder_path
         self._diagnostics = {}  # Dictionary to store diagnostics for each quantity
         self._species_handler = {}
-    
+
     def delete_all_diagnostics(self):
         """
         Delete all diagnostics.
@@ -72,31 +65,39 @@ class Simulation:
         if key in self._species:
             # check if species handler already exists
             if key not in self._species_handler:
-                self._species_handler[key] = Species_Handler(self._simulation_folder, self._input_deck.species[key], self._input_deck)
+                self._species_handler[key] = Species_Handler(
+                    self._simulation_folder,
+                    self._input_deck.species[key],
+                    self._input_deck,
+                )
             return self._species_handler[key]
-        
+
         if key in self._diagnostics:
             return self._diagnostics[key]
-        
+
         # Create a temporary diagnostic for this quantity - this is for quantities that are not species related
-        diag = Diagnostic(simulation_folder=self._simulation_folder, species=None, input_deck=self._input_deck)
+        diag = Diagnostic(
+            simulation_folder=self._simulation_folder,
+            species=None,
+            input_deck=self._input_deck,
+        )
         diag.get_quantity(key)
-        
+
         original_load_all = diag.load_all
-        
+
         def patched_load_all(*args, **kwargs):
-            result = original_load_all(*args, **kwargs)
+            result = original_load_all(*args, **kwargs)  # noqa: F841
             self._diagnostics[key] = diag
             return diag
-        
+
         diag.load_all = patched_load_all
-        
+
         return diag
-    
+
     def add_diagnostic(self, diagnostic, name=None):
         """
         Add a custom diagnostic to the simulation.
-        
+
         Parameters
         ----------
         diagnostic : Diagnostic or array-like
@@ -105,12 +106,12 @@ class Simulation:
         name : str, optional
             The name to use as the key for accessing the diagnostic.
             If None, an auto-generated name will be used.
-            
+
         Returns
         -------
         str
             The name (key) used to store the diagnostic
-            
+
         Example
         -------
         >>> sim = Simulation('path/to/simulation', 'input_deck.txt')
@@ -125,17 +126,22 @@ class Simulation:
             while f"custom_diag_{i}" in self._diagnostics:
                 i += 1
             name = f"custom_diag_{i}"
-        
+
         # If already a Diagnostic, store directly
         if isinstance(diagnostic, Diagnostic):
             self._diagnostics[name] = diagnostic
         else:
             raise ValueError("Only Diagnostic objects are supported for now")
-    
+
     @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:
     def __init__(self, simulation_folder, species_name, input_deck):
@@ -143,30 +149,34 @@ class Species_Handler:
         self._species_name = species_name
         self._input_deck = input_deck
         self._diagnostics = {}
-    
+
     def __getitem__(self, key):
         if key in self._diagnostics:
             return self._diagnostics[key]
-        
+
         # Create a temporary diagnostic for this quantity
-        diag = Diagnostic(simulation_folder=self._simulation_folder, species=self._species_name, input_deck=self._input_deck)
+        diag = Diagnostic(
+            simulation_folder=self._simulation_folder,
+            species=self._species_name,
+            input_deck=self._input_deck,
+        )
         diag.get_quantity(key)
 
         original_load_all = diag.load_all
 
         def patched_load_all(*args, **kwargs):
-            result = original_load_all(*args, **kwargs)
+            result = original_load_all(*args, **kwargs)  # noqa: F841
             self._diagnostics[key] = diag
             return diag
-        
+
         diag.load_all = patched_load_all
 
         return diag
-    
+
     def add_diagnostic(self, diagnostic, name=None):
         """
         Add a custom diagnostic to the simulation.
-        
+
         Parameters
         ----------
         diagnostic : Diagnostic or array-like
@@ -175,18 +185,12 @@ class Species_Handler:
         name : str, optional
             The name to use as the key for accessing the diagnostic.
             If None, an auto-generated name will be used.
-            
+
         Returns
         -------
         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:
@@ -195,9 +199,33 @@ class Species_Handler:
             while f"custom_diag_{i}" in self._diagnostics:
                 i += 1
             name = f"custom_diag_{i}"
-        
+
         # If already a Diagnostic, store directly
         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/decks/decks.py

@@ -1,6 +1,7 @@
-import re
 import ast
 import copy
+import re
+
 import numpy as np
 
 from .species import Specie
@@ -48,7 +49,6 @@ class InputDeckIO:
         self._species = self._get_species()
 
     def _parse_input_deck(self, verbose):
-
         section_list = []
 
         if verbose:
@@ -58,15 +58,13 @@ class InputDeckIO:
             lines = f.readlines()
 
         # remove comments
-        lines = [l[: l.find("!")] if "!" in l else l for l in lines]
+        lines = [l[: l.find("!")] if "!" in l else l for l in lines]  # noqa: E741
 
         # join into single string (makes it easier to parse using regex)
         lines = "".join(lines)
 
         # remove tabs/spaces/paragraphs (except spaces inside "")
-        lines = re.sub(
-            r'"[^"]*"|(\s+)', lambda x: "" if x.group(1) else x.group(0), lines
-        )
+        lines = re.sub(r'"[^"]*"|(\s+)', lambda x: "" if x.group(1) else x.group(0), lines)
 
         # split sections
         # get name before brackets
@@ -84,7 +82,7 @@ class InputDeckIO:
             )
 
         # parse section information
-        for section, info in zip(section_names, section_infos):
+        for section, info in zip(section_names, section_infos, strict=False):
             if verbose:
                 print(f"Reading {section}")
 
@@ -149,9 +147,7 @@ class InputDeckIO:
                 dim = i
                 break
         if dim is None:
-            raise RuntimeError(
-                "Error parsing grid dimension. Grid dimension could not be estabilished."
-            )
+            raise RuntimeError("Error parsing grid dimension. Grid dimension could not be estabilished.")
         return dim
 
     def _get_species(self):
@@ -166,20 +162,11 @@ class InputDeckIO:
             s_qreal = np.ones(len(s_names))
         # check if we have information for all species
         if len(s_names) != self.n_species:
-            raise RuntimeError(
-                "Number of specie names does not match number of species: "
-                f"{len(s_names)} != {len(self.n_species)}."
-            )
+            raise RuntimeError("Number of specie names does not match number of species: " f"{len(s_names)} != {len(self.n_species)}.")
         if len(s_rqm) != self.n_species:
-            raise RuntimeError(
-                "Number of specie rqm does not match number of species: "
-                f"{len(s_rqm)} != {len(self.n_species)}."
-            )
+            raise RuntimeError("Number of specie rqm does not match number of species: " f"{len(s_rqm)} != {len(self.n_species)}.")
         if len(s_qreal) != self.n_species:
-            raise RuntimeError(
-                "Number of specie rqm does not match number of species: "
-                f"{len(s_qreal)} != {len(self.n_species)}."
-            )
+            raise RuntimeError("Number of specie rqm does not match number of species: " f"{len(s_qreal)} != {len(self.n_species)}.")
 
         return {
             ast.literal_eval(s_names[i]): Specie(
@@ -206,10 +193,7 @@ class InputDeckIO:
 
         for i in i_sections:
             if not unexistent_ok and param not in self._sections[i][1]:
-                raise KeyError(
-                    f'"{param}" not yet inside section "{section}" '
-                    "(set unexistent_ok=True to ignore)."
-                )
+                raise KeyError(f'"{param}" not yet inside section "{section}" ' "(set unexistent_ok=True to ignore).")
             if isinstance(value, str):
                 self._sections[i][1][param] = str(f'"{value}"')
             elif isinstance(value, list):
@@ -279,9 +263,7 @@ class InputDeckIO:
                 return int(self["particles"][0]["num_cathode"])
             except (KeyError, IndexError):
                 # If neither exists, raise an informative error
-                raise KeyError(
-                    "Could not find 'num_species' or 'num_cathode' in the particles section"
-                )
+                raise KeyError("Could not find 'num_species' or 'num_cathode' in the particles section") from None
 
     @property
     def species(self):

osiris_utils/postprocessing/derivative.py

@@ -1,7 +1,9 @@
-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 Derivative_Simulation(PostProcess):
     """
@@ -12,7 +14,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 +26,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,14 +41,17 @@ 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)
+            self._derivatives_computed[key] = Derivative_Diagnostic(
+                diagnostic=self._simulation[key],
+                deriv_type=self._deriv_type,
+                axis=self._axis,
+            )
         return self._derivatives_computed[key]
-    
+
     def delete_all(self):
         self._derivatives_computed = {}
 
@@ -60,12 +60,12 @@ class Derivative_Simulation(PostProcess):
             del self._derivatives_computed[key]
         else:
             print(f"Derivative {key} not found in simulation")
-    
+
     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):
     """
     Auxiliar class to compute the derivative of a diagnostic, for it to be similar in behavior to a Diagnostic object.
@@ -75,7 +75,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,89 +87,114 @@ 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)
+        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 = "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:
             print("Using cached derivative")
             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()
+            self._data = self._diag._data
 
-        if self._type == "t":
-            result = np.gradient(self._diag._data, self._diag._dt * self._diag._ndump, axis=0, edge_order=2)
+        if self._diag._all_loaded is True:
+            print("Using cached data from diagnostic")
+            self._data = self._diag._data
 
-        elif self._type == "x1":
+        if self._deriv_type == "t":
+            result = np.gradient(self._data, self._diag._dt * self._diag._ndump, axis=0, edge_order=2)
+
+        elif self._deriv_type == "x1":
             if self._dim == 1:
-                result = np.gradient(self._diag._data, self._diag._dx, axis=1, edge_order=2)
+                result = np.gradient(self._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)
+                result = np.gradient(self._data, self._diag._dx[0], axis=1, edge_order=2)
+
+        elif self._deriv_type == "x2":
+            result = np.gradient(self._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._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":
+            result = np.gradient(
+                np.gradient(
+                    self._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._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":
+            result = np.gradient(
+                np.gradient(self._data, self._diag._dt, axis=0, edge_order=2),
+                self._diag._dx[self._axis - 1],
+                axis=self._axis[0],
+                edge_order=2,
+            )
+
+        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), 
-                                self._diag._dt, axis=0, edge_order=2)
+            result = np.gradient(
+                np.gradient(
+                    self._data,
+                    self._diag._dx[self._axis - 1],
+                    axis=self._axis,
+                    edge_order=2,
+                ),
+                self._diag._dt,
+                axis=0,
+                edge_order=2,
+            )
         else:
             raise ValueError("Invalid derivative type.")
-        
+
         # Store the result in the cache
         self._all_loaded = True
         self._data = result
@@ -177,21 +202,23 @@ 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)
+                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:
@@ -203,7 +230,7 @@ class Derivative_Diagnostic(Diagnostic):
         """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):
@@ -214,6 +241,7 @@ class Derivative_Diagnostic(Diagnostic):
         else:
             raise ValueError("Invalid index type. Use int or slice.")
 
+
 class Derivative_Species_Handler:
     """
     Class to handle derivatives for a species.
@@ -230,14 +258,15 @@ 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)
-        return self._derivatives_computed[key]
+            self._derivatives_computed[key] = Derivative_Diagnostic(diag, self._deriv_type, self._axis)
+        return self._derivatives_computed[key]