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

osiris_utils/data/simulation.py

@@ -0,0 +1,216 @@
+from matplotlib.style import available
+from ..data.diagnostic import *
+from ..utils import *
+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
+        The simulation folder.
+    species : Specie object
+        The species to analyze.
+    diagnostics : dict
+        Dictionary to store diagnostics for each quantity when `load_all` method is used.
+
+    Methods
+    -------
+    delete_all_diagnostics()
+        Delete all diagnostics.
+    delete_diagnostic(key)
+        Delete a diagnostic.
+    __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
+        self._input_deck = InputDeckIO(self._input_deck_path, verbose=False)
+
+        self._species = list(self._input_deck.species.keys())
+
+        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.
+        """
+        self._diagnostics = {}
+
+    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 simulation")
+
+    def __getitem__(self, key):
+        # check if key is a species
+        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)
+            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.get_quantity(key)
+        
+        original_load_all = diag.load_all
+        
+        def patched_load_all(*args, **kwargs):
+            result = original_load_all(*args, **kwargs)
+            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
+            The diagnostic to add. If not a Diagnostic object, it will be wrapped
+            in a Diagnostic object.
+        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:
+            # Find an unused name
+            i = 1
+            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):
+        self._simulation_folder = simulation_folder
+        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.get_quantity(key)
+
+        original_load_all = diag.load_all
+
+        def patched_load_all(*args, **kwargs):
+            result = original_load_all(*args, **kwargs)
+            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
+            The diagnostic to add. If not a Diagnostic object, it will be wrapped
+            in a Diagnostic object.
+        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:
+            # Find an unused name
+            i = 1
+            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.
+        """
+        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

@@ -0,0 +1,288 @@
+import re
+import ast
+import copy
+import numpy as np
+
+from .species import Specie
+
+
+def deval(x):
+    """
+    Auxiliar to handle eval of Fortran formatted numbers (e.g. 1.4d-5)
+    """
+    if "d" in x:
+        x = x.replace("d", "e")
+    return float(x)
+
+
+class InputDeckIO:
+    """
+    Class to handle parsing/re-writing of OSIRIS input decks
+
+    Parameters
+    ----------
+    filename : str
+        Path to OSIRIS input deck file.
+
+    verbose : bool
+        If True, prints additional information when parsing file.
+        Helpful for debugging issues if input deck parsing fails.
+
+    Attributes
+    ----------
+    filename : str
+        Path to original input file used to create the InputDeckIO object.
+
+    sections : list[dict]
+        List of pairs (section_name: str, section_dict: dict) which contain
+        current state of InputDeckIO object.
+
+    dim : int
+        Number of dimensions in the simulation (1, 2, or 3).
+    """
+
+    def __init__(self, filename: str, verbose: bool = True):
+        self._filename = str(filename)
+        self._sections = self._parse_input_deck(verbose)
+        self._dim = self._get_dim()
+        self._species = self._get_species()
+
+    def _parse_input_deck(self, verbose):
+
+        section_list = []
+
+        if verbose:
+            print(f"\nParsing input deck : {self._filename}")
+
+        with open(self._filename, "r", encoding="utf-8") as f:
+            lines = f.readlines()
+
+        # remove comments
+        lines = [l[: l.find("!")] if "!" in l else l for l in lines]
+
+        # 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
+        )
+
+        # split sections
+        # get name before brackets
+        section_names = re.findall(r"(?:^|\})(.*?)(?:\{)", lines)
+        # get content inside brackets
+        section_infos = re.findall(r"(?:\{)(.*?)(?:\})", lines)
+
+        if len(section_names) != len(section_infos):
+            raise RuntimeError(
+                "Unexpected problems parsing the document!\n"
+                f"Number of section names detected ({len(section_names)}) "
+                f"is different from number of sections ({len(section_infos)}).\n"
+                "Might be a bug in the code, or problem with input deck format!"
+                "Check if you could run the deck with OSIRIS."
+            )
+
+        # parse section information
+        for section, info in zip(section_names, section_infos):
+            if verbose:
+                print(f"Reading {section}")
+
+            # split section contents at commas (unless comma inside brackets e.g. pmax(1:2, 1))
+            info = re.split(r",(?![^()]*\))\s*", info)
+            info = list(filter(None, info))
+
+            # save pairs of (param, values) to dict
+            section_dict = {}
+            param = ""
+            value = ""
+
+            for i in info:
+                aux = i.split("=")
+                # solution to deal with parameters given by multiple values
+                # (e.g. ps_np(1:3) = 512,128,128 which was split into:
+                #   ['ps_np(1:3)=512', '128', '128'])
+                # need to be able to regroup '128's with the previous value
+                if len(aux) == 1 and param != "":
+                    value = ",".join([value, aux[0]])
+                    section_dict[param] = value
+                # simplest case where we simply have "param=value"
+                elif len(aux) == 2:
+                    param, value = aux
+                    section_dict[param] = value
+                # case where we have multipel '=' inside strings
+                # happens for e.g. with mathfuncs in density profiles
+                else:
+                    param = aux[0]
+                    value = "".join(aux[1:])
+                    # check that value is actually wrapped inside a string
+                    if value[0] in ['"', "'"] and value[-1] in ['"', "'"]:
+                        section_dict[param] = value
+                    # because if not, then there is an error in the parser
+                    # or a problem with the input deck
+                    else:
+                        raise RuntimeError(
+                            f'Error parsing section: "{section}".\n'
+                            "Might be a bug in the code, or problem with input deck format!\n"
+                            "Check if you could run the deck with OSIRIS."
+                        )
+
+            section_list.append([section, section_dict])
+
+            if verbose:
+                for k, v in section_dict.items():
+                    print(f"  {k} = {v}")
+
+        if verbose:
+            print("Input deck successfully parsed\n")
+
+        return section_list
+
+    def _get_dim(self):
+        dim = None
+        for i in range(1, 4):
+            try:
+                self.get_param(section="grid", param=f"nx_p(1:{i})")
+            except KeyError:
+                pass
+            else:
+                dim = i
+                break
+        if dim is None:
+            raise RuntimeError(
+                "Error parsing grid dimension. Grid dimension could not be estabilished."
+            )
+        return dim
+
+    def _get_species(self):
+        s_names = self.get_param("species", "name")
+        s_rqm = self.get_param("species", "rqm")
+        # real charge is optional in OSIRIS
+        # if real charge not provided assume electron charge
+        try:
+            s_qreal = self.get_param("species", "q_real")
+            s_qreal = np.array([float(q) for q in s_qreal])
+        except KeyError:
+            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)}."
+            )
+        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)}."
+            )
+        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)}."
+            )
+
+        return {
+            ast.literal_eval(s_names[i]): Specie(
+                name=ast.literal_eval(s_names[i]),
+                rqm=float(s_rqm[i]),
+                q=int(s_qreal[0]) * np.sign(float(s_rqm[i])),
+            )
+            for i in range(self.n_species)
+        }
+
+    def set_param(self, section, param, value, i_use=None, unexistent_ok=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]
+
+        if len(i_sections) == 0:
+            raise KeyError(f'section "{section}" not found')
+
+        if i_use is not None:
+            try:
+                i_sections = [im for i, im in enumerate(i_sections) if i in i_use]
+            except TypeError:
+                i_sections = [i_sections[i_use]]
+
+        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)."
+                )
+            if isinstance(value, str):
+                self._sections[i][1][param] = str(f'"{value}"')
+            elif isinstance(value, list):
+                self._sections[i][1][param] = ",".join(map(str, value))
+            else:
+                self._sections[i][1][param] = str(value)
+
+    def set_tag(self, tag, value):
+        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):
+        i_sections = [i for i, m in enumerate(self._sections) if m[0] == section]
+
+        if len(i_sections) == 0:
+            print("section not found")
+            return []
+
+        values = []
+        for i in i_sections:
+            if param not in self._sections[i][1]:
+                raise KeyError(f'"{param}" not found inside section "{section}"')
+            values.append(copy.deepcopy(self._sections[i][1][param]))
+
+        return values
+
+    def delete_param(self, section, param):
+        sections_new = []
+        for m_name, m_dict in self._sections:
+            if m_name == section and param in m_dict:
+                m_dict.pop(param)
+            sections_new.append([m_name, m_dict])
+        self._sections = sections_new
+
+    def print_to_file(self, filename):
+        with open(filename, "w", encoding="utf-8") as f:
+            for section, section_dict in self._sections:
+                f.write(f"{section}\n{{\n")
+                for k, v in section_dict.items():
+                    f.write(f'\t{k} = {v.replace(",", ", ")},\n')
+                f.write("}\n\n")
+
+    def __getitem__(self, section):
+        return copy.deepcopy([m[1] for m in self._sections if m[0] == section])
+
+    # Getters
+    @property
+    def filename(self):
+        return self._filename
+
+    @property
+    def sections(self):
+        return self._sections
+
+    @property
+    def dim(self):
+        return self._dim
+
+    @property
+    def n_species(self):
+        try:
+            return int(self["particles"][0]["num_species"])
+        except (KeyError, IndexError):
+            # If num_species doesn't exist, try num_cathode
+            try:
+                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"
+                )
+
+    @property
+    def species(self):
+        return self._species

osiris_utils/decks/species.py

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