osiris-utils 1.1.6__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,8 +31,9 @@ class Simulation:
     __getitem__(key)
         Get a diagnostic.
 
-        
-    '''
+
+    """
+
     def __init__(self, input_deck_path):
         folder_path = os.path.dirname(input_deck_path)
         self._input_deck_path = input_deck_path
@@ -43,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.
@@ -64,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
@@ -97,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')
@@ -117,21 +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):
@@ -139,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
@@ -171,12 +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
-        
+
         """
         # Generate a name if none provided
         if name is None:
@@ -185,13 +199,13 @@ 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")
-        
+
     def delete_diagnostic(self, key):
         """
         Delete a diagnostic.
@@ -211,6 +225,7 @@ class Species_Handler:
     @property
     def species(self):
         return self._species_name
+
     @property
     def loaded_diagnostics(self):
-        return self._diagnostics
+        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):
     """
@@ -41,12 +43,15 @@ class Derivative_Simulation(PostProcess):
             if key not in self._species_handler:
                 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], 
-                                                                    deriv_type=self._deriv_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 = {}
 
@@ -55,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._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.
@@ -86,23 +91,36 @@ class Derivative_Diagnostic(Diagnostic):
 
     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 = f"DERIV"
-            
+        self.postprocess_name = "DERIV"
+
         # self._name = f"D[{diagnostic._name}, {type}]"
         self._diag = diagnostic
         self._deriv_type = deriv_type
         self._axis = axis if axis is not None else diagnostic._axis
         self._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))
 
@@ -111,45 +129,72 @@ class Derivative_Diagnostic(Diagnostic):
         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._diag._all_loaded is True:
+            print("Using cached data from diagnostic")
+            self._data = self._diag._data
 
         if self._deriv_type == "t":
-            result = np.gradient(self._diag._data, self._diag._dt * self._diag._ndump, axis=0, edge_order=2)
+            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)
-                    
+                result = np.gradient(self._data, self._diag._dx[0], axis=1, edge_order=2)
+
         elif self._deriv_type == "x2":
-            result = np.gradient(self._diag._data, self._diag._dx[1], axis=2, edge_order=2)
+            result = np.gradient(self._data, self._diag._dx[1], axis=2, edge_order=2)
 
         elif self._deriv_type == "x3":
-            result = np.gradient(self._diag._data, self._diag._dx[2], axis=3, edge_order=2)
+            result = np.gradient(self._data, self._diag._dx[2], axis=3, edge_order=2)
 
         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)
-            
+            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)
-            
+            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
@@ -162,16 +207,18 @@ class Derivative_Diagnostic(Diagnostic):
                 yield np.gradient(self._diag[index], self._diag._dx, axis=0, edge_order=2)
             else:
                 yield np.gradient(self._diag[index], self._diag._dx[0], axis=0, edge_order=2)
-        
+
         elif self._deriv_type == "x2":
             yield np.gradient(self._diag[index], self._diag._dx[1], axis=1, edge_order=2)
-        
+
         elif self._deriv_type == "x3":
             yield np.gradient(self._diag[index], self._diag._dx[2], axis=2, edge_order=2)
-            
+
         elif self._deriv_type == "t":
             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:
@@ -183,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):
@@ -194,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.
@@ -210,6 +258,7 @@ 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, deriv_type, axis=None):
         self._species_handler = species_handler
         self._deriv_type = deriv_type
@@ -220,4 +269,4 @@ class Derivative_Species_Handler:
         if key not in self._derivatives_computed:
             diag = self._species_handler[key]
             self._derivatives_computed[key] = Derivative_Diagnostic(diag, self._deriv_type, self._axis)
-        return self._derivatives_computed[key]
+        return self._derivatives_computed[key]