osiris-utils 1.1.10a0__py3-none-any.whl → 1.2.0__py3-none-any.whl

osiris_utils/data/simulation.py

@@ -1,8 +1,13 @@
+from __future__ import annotations
+
 import os
+from typing import Any
 
 from ..data.diagnostic import Diagnostic
 from ..decks.decks import InputDeckIO
 
+__all__ = ["Simulation", "Species_Handler"]
+
 
 class Simulation:
     """
@@ -17,7 +22,7 @@ class Simulation:
     ----------
     simulation_folder : str
         The simulation folder.
-    species : Specie object
+    species : Species object
         The species to analyze.
     diagnostics : dict
         Dictionary to store diagnostics for each quantity when `load_all` method is used.
@@ -34,7 +39,7 @@ class Simulation:
 
     """
 
-    def __init__(self, input_deck_path):
+    def __init__(self, input_deck_path: str) -> None:
         folder_path = os.path.dirname(input_deck_path)
         self._input_deck_path = input_deck_path
         self._input_deck = InputDeckIO(self._input_deck_path, verbose=False)
@@ -60,7 +65,7 @@ class Simulation:
         else:
             print(f"Diagnostic {key} not found in simulation")
 
-    def __getitem__(self, key):
+    def __getitem__(self, key: str) -> Diagnostic | Species_Handler:
         # check if key is a species
         if key in self._species:
             # check if species handler already exists
@@ -94,7 +99,7 @@ class Simulation:
 
         return diag
 
-    def add_diagnostic(self, diagnostic, name=None):
+    def add_diagnostic(self, diagnostic: Diagnostic, name: str | None = None) -> str:
         """
         Add a custom diagnostic to the simulation.
 
@@ -130,27 +135,28 @@ class Simulation:
         # If already a Diagnostic, store directly
         if isinstance(diagnostic, Diagnostic):
             self._diagnostics[name] = diagnostic
+            return name
         else:
             raise ValueError("Only Diagnostic objects are supported for now")
 
     @property
-    def species(self):
+    def species(self) -> list[str]:
         return self._species
 
     @property
-    def loaded_diagnostics(self):
+    def loaded_diagnostics(self) -> dict[str, Diagnostic]:
         return self._diagnostics
 
 
 # This is to handle species related diagnostics
 class Species_Handler:
-    def __init__(self, simulation_folder, species_name, input_deck):
+    def __init__(self, simulation_folder: str, species_name: Any, input_deck: Any) -> None:
         self._simulation_folder = simulation_folder
         self._species_name = species_name
         self._input_deck = input_deck
         self._diagnostics = {}
 
-    def __getitem__(self, key):
+    def __getitem__(self, key: str) -> Diagnostic:
         if key in self._diagnostics:
             return self._diagnostics[key]
 
@@ -173,7 +179,7 @@ class Species_Handler:
 
         return diag
 
-    def add_diagnostic(self, diagnostic, name=None):
+    def add_diagnostic(self, diagnostic: Diagnostic, name: str | None = None) -> str:
         """
         Add a custom diagnostic to the simulation.
 
@@ -203,10 +209,11 @@ class Species_Handler:
         # If already a Diagnostic, store directly
         if isinstance(diagnostic, Diagnostic):
             self._diagnostics[name] = diagnostic
+            return name
         else:
             raise ValueError("Only Diagnostic objects are supported for now")
 
-    def delete_diagnostic(self, key):
+    def delete_diagnostic(self, key: str) -> None:
         """
         Delete a diagnostic.
         """
@@ -216,16 +223,16 @@ class Species_Handler:
             print(f"Diagnostic {key} not found in species {self._species_name}")
             return None
 
-    def delete_all_diagnostics(self):
+    def delete_all_diagnostics(self) -> None:
         """
         Delete all diagnostics.
         """
         self._diagnostics = {}
 
     @property
-    def species(self):
+    def species(self) -> Any:
         return self._species_name
 
     @property
-    def loaded_diagnostics(self):
+    def loaded_diagnostics(self) -> dict[str, Diagnostic]:
         return self._diagnostics

osiris_utils/decks/__init__.py

@@ -0,0 +1,4 @@
+from .decks import InputDeckIO, deval
+from .species import Species
+
+__all__ = ["InputDeckIO", "deval", "Species"]

osiris_utils/decks/decks.py

@@ -4,7 +4,9 @@ import re
 
 import numpy as np
 
-from .species import Specie
+from .species import Species
+
+__all__ = ["InputDeckIO", "deval"]
 
 
 def deval(x):
@@ -42,7 +44,7 @@ class InputDeckIO:
         Number of dimensions in the simulation (1, 2, or 3).
     """
 
-    def __init__(self, filename: str, verbose: bool = True):
+    def __init__(self, filename: str, verbose: bool = False):
         self._filename = str(filename)
         self._sections = self._parse_input_deck(verbose)
         self._dim = self._get_dim()
@@ -54,7 +56,7 @@ class InputDeckIO:
         if verbose:
             print(f"\nParsing input deck : {self._filename}")
 
-        with open(self._filename, "r", encoding="utf-8") as f:
+        with open(self._filename, encoding="utf-8") as f:
             lines = f.readlines()
 
         # remove comments
@@ -162,14 +164,14 @@ 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(f"Number of specie names does not match number of species: {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(f"Number of specie rqm does not match number of species: {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(f"Number of specie rqm does not match number of species: {len(s_qreal)} != {len(self.n_species)}.")
 
         return {
-            ast.literal_eval(s_names[i]): Specie(
+            ast.literal_eval(s_names[i]): Species(
                 name=ast.literal_eval(s_names[i]),
                 rqm=float(s_rqm[i]),
                 q=int(s_qreal[0]) * np.sign(float(s_rqm[i])),
@@ -178,6 +180,29 @@ class InputDeckIO:
         }
 
     def set_param(self, section, param, value, i_use=None, unexistent_ok=False):
+        """
+        Set a parameter value in the input deck.
+
+        Parameters
+        ----------
+        section : str
+            The name of the section where the parameter is located.
+        param : str
+            The name of the parameter to set.
+        value : str or list or int or float
+            The value to set. If a list, it is converted to a comma-separated string.
+        i_use : int or list of int, optional
+            The index(es) of the section(s) to modify if multiple sections with the same name exist.
+            If None, all sections with the matching name are modified.
+        unexistent_ok : bool, optional
+            If True, allows setting a parameter that does not currently exist in the section.
+            Default is False.
+
+        Raises
+        ------
+        KeyError
+            If the section is not found, or if the parameter is not found (and unexistent_ok is False).
+        """
         # get all sections with the same name
         # (e.g. there might be multiple 'species')
         i_sections = [i for i, m in enumerate(self._sections) if m[0] == section]
@@ -193,7 +218,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):
@@ -202,11 +227,43 @@ class InputDeckIO:
                 self._sections[i][1][param] = str(value)
 
     def set_tag(self, tag, value):
+        """
+        Replace a tag in all parameters within the input deck.
+        This is useful for template replacement.
+
+        Parameters
+        ----------
+        tag : str
+            The tag substring to search for (e.g., "<TAG>").
+        value : str or int or float
+            The value to replace the tag with.
+        """
         for im, (_, params) in enumerate(self._sections):
             for p, v in params.items():
                 self._sections[im][1][p] = v.replace(tag, str(value))
 
     def get_param(self, section, param):
+        """
+        Get the value(s) of a parameter from a specific section.
+
+        Parameters
+        ----------
+        section : str
+            The name of the section.
+        param : str
+            The name of the parameter.
+
+        Returns
+        -------
+        values : list
+            A list of values for the parameter. Returns a list because there can be multiple
+            sections with the same name.
+
+        Raises
+        ------
+        KeyError
+            If the parameter is not found in one of the sections.
+        """
         i_sections = [i for i, m in enumerate(self._sections) if m[0] == section]
 
         if len(i_sections) == 0:
@@ -222,6 +279,16 @@ class InputDeckIO:
         return values
 
     def delete_param(self, section, param):
+        """
+        Delete a parameter from a specific section.
+
+        Parameters
+        ----------
+        section : str
+            The name of the section.
+        param : str
+            The name of the parameter to delete.
+        """
         sections_new = []
         for m_name, m_dict in self._sections:
             if m_name == section and param in m_dict:
@@ -230,6 +297,14 @@ class InputDeckIO:
         self._sections = sections_new
 
     def print_to_file(self, filename):
+        """
+        Write the current state of the input deck to a file.
+
+        Parameters
+        ----------
+        filename : str
+            The path to the output file.
+        """
         with open(filename, "w", encoding="utf-8") as f:
             for section, section_dict in self._sections:
                 f.write(f"{section}\n{{\n")

osiris_utils/decks/species.py

@@ -1,32 +1,35 @@
-class Specie:
+__all__ = ["Species"]
+
+
+class Species:
     """
     Class to store OSIRIS species object.
 
     Parameters
     ----------
     name : str
-        Specie name.
+        Species name.
 
     rqm : float
-        Specie charge to mass ratio.
+        Species charge to mass ratio.
 
     q : int
-        Specie charge in units of the electron charge.
+        Species charge in units of the electron charge.
         Electrons would be represented by q=-1 and protons q=1.
 
     Attributes
     ----------
     name : str
-        Specie name.
+        Species name.
 
     rqm : float
-        Specie charge to mass ratio.
+        Species charge to mass ratio.
 
     q : int
-        Specie charge in units of the electron charge.
+        Species charge in units of the electron charge.
 
     m : float
-        Specie mass in units of the electron mass.
+        Species mass in units of the electron mass.
     """
 
     def __init__(self, name, rqm, q: int = 1):
@@ -36,7 +39,7 @@ class Specie:
         self._m = rqm * q
 
     def __repr__(self) -> str:
-        return f"Specie(name={self._name}, rqm={self._rqm}, q={self._q}, m={self._m})"
+        return f"Species(name={self._name}, rqm={self._rqm}, q={self._q}, m={self._m})"
 
     @property
     def name(self):

osiris_utils/postprocessing/__init__.py

@@ -0,0 +1,28 @@
+from .derivative import Derivative_Diagnostic, Derivative_Simulation, Derivative_Species_Handler
+from .fft import FFT_Diagnostic, FFT_Simulation, FFT_Species_Handler
+from .field_centering import FieldCentering_Diagnostic, FieldCentering_Simulation
+from .heatflux_correction import HeatfluxCorrection_Diagnostic, HeatfluxCorrection_Simulation, HeatfluxCorrection_Species_Handler
+from .mft import MFT_Diagnostic, MFT_Simulation, MFT_Species_Handler
+from .postprocess import PostProcess
+from .pressure_correction import PressureCorrection_Diagnostic, PressureCorrection_Simulation, PressureCorrection_Species_Handler
+
+__all__ = [
+    "Derivative_Simulation",
+    "Derivative_Diagnostic",
+    "Derivative_Species_Handler",
+    "FFT_Simulation",
+    "FFT_Diagnostic",
+    "FFT_Species_Handler",
+    "FieldCentering_Simulation",
+    "FieldCentering_Diagnostic",
+    "HeatfluxCorrection_Simulation",
+    "HeatfluxCorrection_Diagnostic",
+    "HeatfluxCorrection_Species_Handler",
+    "MFT_Simulation",
+    "MFT_Diagnostic",
+    "MFT_Species_Handler",
+    "PostProcess",
+    "PressureCorrection_Simulation",
+    "PressureCorrection_Diagnostic",
+    "PressureCorrection_Species_Handler",
+]