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

osiris_utils/postprocessing/pressure_correction.py

@@ -0,0 +1,187 @@
+import numpy as np
+
+from ..data.diagnostic import Diagnostic
+from ..data.simulation import Simulation
+from .postprocess import PostProcess
+
+OSIRIS_P = ["P11", "P12", "P13", "P21", "P22", "P23", "P31", "P32", "P33"]
+
+
+class PressureCorrection_Simulation(PostProcess):
+    def __init__(self, simulation):
+        super().__init__("PressureCorrection Simulation")
+        """
+        Class to correct pressure tensor components by subtracting Reynolds stress.
+
+        Parameters
+        ----------
+        sim : Simulation
+            The simulation object.
+        pressure : str
+            The pressure component to center.
+        """
+        if not isinstance(simulation, Simulation):
+            raise ValueError("Simulation must be a Simulation object.")
+        self._simulation = simulation
+        self._pressure_corrected = {}
+        self._species_handler = {}
+
+    def __getitem__(self, key):
+        if key in self._simulation._species:
+            if key not in self._species_handler:
+                self._species_handler[key] = PressureCorrection_Species_Handler(self._simulation[key])
+            return self._species_handler[key]
+        if key not in OSIRIS_P:
+            raise ValueError(f"Invalid pressure component {key}. Supported: {OSIRIS_P}.")
+        if key not in self._pressure_corrected:
+            print("Weird that it got here - pressure is always species dependent on OSIRIS")
+            self._pressure_corrected[key] = PressureCorrection_Diagnostic(self._simulation[key], self._simulation)
+        return self._pressure_corrected[key]
+
+    def delete_all(self):
+        self._pressure_corrected = {}
+
+    def delete(self, key):
+        if key in self._pressure_corrected:
+            del self._pressure_corrected[key]
+        else:
+            print(f"Pressure {key} not found in simulation")
+
+    def process(self, diagnostic):
+        """Apply pressure correction to a diagnostic"""
+        return PressureCorrection_Diagnostic(diagnostic, self._simulation)
+
+
+class PressureCorrection_Diagnostic(Diagnostic):
+    def __init__(self, diagnostic, n, ufl_j, vfl_k):
+        """
+        Class to correct the pressure 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 = "P_CORR"
+
+        if diagnostic._name not in OSIRIS_P:
+            raise ValueError(f"Invalid pressure component {diagnostic._name}. Supported: {OSIRIS_P}")
+
+        self._diag = diagnostic
+
+        # The density and velocities are now passed as arguments (so it can doesn't depend on the simulation)
+        self._n = n
+        self._ufl_j = ufl_j
+        self._vfl_k = vfl_k
+
+        for attr in [
+            "_dt",
+            "_dx",
+            "_ndump",
+            "_axis",
+            "_nx",
+            "_x",
+            "_grid",
+            "_dim",
+            "_maxiter",
+            "_type",
+        ]:
+            if hasattr(diagnostic, attr):
+                setattr(self, attr, getattr(diagnostic, attr))
+
+        self._original_name = diagnostic._name
+        self._name = diagnostic._name + "_corrected"
+
+        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()
+
+        print(f"Loading {self._species._name} {self._original_name} diagnostic")
+        self._n.load_all()
+        self._ufl_j.load_all()
+        self._vfl_k.load_all()
+
+        # Then access the data
+        n = self._n.data
+        u = self._ufl_j.data
+        v = self._vfl_k.data
+
+        self._data = self._diag.data - n * v * u
+        self._all_loaded = True
+
+        # Unload the data to save memory
+        # self._n.unload()
+        # self._ufl_j.unload()
+        # self._vfl_k.unload()
+
+        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):
+        yield (self._diag[index] - self._n[index] * self._vfl_k[index] * self._ufl_j[index])
+
+
+class PressureCorrection_Species_Handler:
+    """
+    Class to handle pressure correction for a species.
+    Acts as a wrapper for the PressureCorrection_Diagnostic class.
+
+    Not intended to be used directly, but through the PressureCorrection_Simulation class.
+
+    Parameters
+    ----------
+    species_handler : Species_Handler
+        The species handler object.
+    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.
+    """
+
+    def __init__(self, species_handler):
+        self._species_handler = species_handler
+        self._pressure_corrected = {}
+
+    def __getitem__(self, key):
+        if key not in self._pressure_corrected:
+            diag = self._species_handler[key]
+
+            # Density and velocities alwayes depend on the species so this can be done here
+
+            n = self._species_handler["n"]
+            self._j, self._k = key[-2], key[-1]
+            try:
+                ufl = self._species_handler[f"ufl{self._j}"]
+            except Exception:
+                ufl = self._species_handler[f"vfl{self._j}"]
+            vfl = self._species_handler[f"vfl{self._k}"]
+            self._pressure_corrected[key] = PressureCorrection_Diagnostic(diag, n, ufl, vfl)
+        return self._pressure_corrected[key]

osiris_utils/utils.py

@@ -1,12 +1,13 @@
+from datetime import datetime
+
+import h5py
 import numpy as np
-import h5py  
-import matplotlib.pyplot as plt
-import matplotlib.animation as animation
-import scipy 
 import pandas as pd
+import scipy
+
 
 def courant2D(dx, dy):
-    '''
+    """
     Compute the Courant number for a 2D simulation.
 
     Parameters
@@ -20,12 +21,13 @@ def courant2D(dx, dy):
     -------
     float
         The limit for dt.
-    '''
-    dt = 1 / (np.sqrt(1/dx**2 + 1/dy**2))
+    """
+    dt = 1 / (np.sqrt(1 / dx**2 + 1 / dy**2))
     return dt
 
-def time_estimation(n_cells, ppc, t_steps, n_cpu, push_time = 1e-7, hours = False):
-    '''
+
+def time_estimation(n_cells, ppc, t_steps, n_cpu, push_time=1e-7, hours=False):
+    """
     Estimate the simulation time.
 
     Parameters
@@ -47,41 +49,44 @@ def time_estimation(n_cells, ppc, t_steps, n_cpu, push_time = 1e-7, hours = Fals
     -------
     float
         The estimated time in seconds or hours.
-    '''
-    time = (n_cells*ppc*push_time*t_steps)/n_cpu
+    """
+    time = (n_cells * ppc * push_time * t_steps) / n_cpu
     if hours:
-        return time/3600
+        return time / 3600
     else:
         return time
-    
+
+
 def filesize_estimation(n_gridpoints):
-    return n_gridpoints*4/(1024**2)
+    return n_gridpoints * 4 / (1024**2)
+
 
 def transverse_average(data):
-    '''
+    """
     Computes the transverse average of a 2D array.
-    
+
     Parameters
     ----------
     data : numpy.ndarray
         Dim: 2D.
         The input data.
-        
+
     Returns
     -------
     numpy.ndarray
         Dim: 1D.
         The transverse average.
 
-    '''
+    """
 
     if len(data.shape) != 2:
-        raise ValueError('The input data must be a 2D array.')
-    return np.mean(data, axis = 1)
+        raise ValueError("The input data must be a 2D array.")
+    return np.mean(data, axis=1)
+
 
 def integrate(array, dx):
-    '''
-    Integrate a 1D from the left to the right. This may be changed in the future to allow 
+    """
+    Integrate a 1D from the left to the right. This may be changed in the future to allow
     for integration in both directions or for other more general cases.
 
     Parameters
@@ -97,37 +102,39 @@ def integrate(array, dx):
     numpy.ndarray
         Dim: 1D.
         The integrated array.
-    '''
+    """
 
     if len(array.shape) != 1:
-        raise ValueError(f'Array must be 1D\n Array shape: {array.shape}')
+        raise ValueError(f"Array must be 1D\n Array shape: {array.shape}")
     flip_array = np.flip(array)
     # int = -scipy.integrate.cumulative_trapezoid(flip_array, dx = dx, initial = flip_array[0])
-    int = -scipy.integrate.cumulative_simpson(flip_array, dx = dx, initial = 0)
+    int = -scipy.integrate.cumulative_simpson(flip_array, dx=dx, initial=0)
     return np.flip(int)
 
-def save_data(data, savename, option='numpy'):
+
+def save_data(data, savename, option="numpy"):
     """
     Save the data to a .txt (with Numpy) or .csv (with Pandas) file.
 
     Parameters
     ----------
-    data : list 
+    data : list
         The data to be saved.
     savename : str
         The path to the file.
     option : str, optional
         The option for saving the data. The default is 'numpy'. Can be 'numpy' or 'pandas'.
     """
-    if option == 'numpy':
+    if option == "numpy":
         np.savetxt(savename, data)
-    elif option == 'pandas':
+    elif option == "pandas":
         pd.DataFrame(data).to_csv(savename, index=False)
     else:
         raise ValueError("Option must be 'numpy' or 'pandas'.")
 
-def read_data(filename, option='numpy'):
-    '''
+
+def read_data(filename, option="numpy"):
+    """
     Read the data from a .txt file.
 
     Parameters
@@ -140,5 +147,139 @@ def read_data(filename, option='numpy'):
     numpy.ndarray
         Dim: 2D.
         The data.
-    '''
-    return np.loadtxt(filename) if option == 'numpy' else pd.read_csv(filename).values
+    """
+    return np.loadtxt(filename) if option == "numpy" else pd.read_csv(filename).values
+
+
+def convert_tracks(filename_in):
+    """
+    Converts a new OSIRIS track file aka IDL-formatted aka tracks-2 to an older format that is more human-readable.
+    In the old format, each particle is stored in a separate folder, with datasets for each quantity.
+    The function reads data from the input file, processes it, and writes it to a new file with the suffix "-v2".
+
+     code from https://github.com/GoLP-IST/RaDi-x/blob/main/tools/convert_idl_tracks_py3.py
+
+     Parameters
+     ----------
+     filename_in : str
+         The path to the trackfile.
+
+     Returns
+     -------
+     The output file will be in the same folder as the input file with the same name with \"v2\" added
+
+    """
+
+    try:
+        file_in = h5py.File(filename_in, "r")
+    except IOError:
+        print("cannot open " + filename_in)
+        exit()
+
+    # read data from file
+    data = file_in["data"][:]
+    itermap = file_in["itermap"][:]
+    ntracks = file_in.attrs["NTRACKS"][0]
+    niter = file_in.attrs["NITER"][0]
+    quants = file_in.attrs["QUANTS"][:]
+    file_in_attr_keys = file_in.attrs.keys()
+    sim_attr_keys = file_in["SIMULATION"].attrs.keys()
+    nquants = len(quants)
+
+    # construct file out for new format
+    filename_out = filename_in[:-3] + "-v2" + filename_in[-3:]
+    file_out = h5py.File(filename_out, "w")
+
+    # copy attrs from file_in
+    for item in file_in_attr_keys:
+        file_out.attrs[item] = file_in.attrs[item]
+    for item in sim_attr_keys:
+        file_out.attrs[item] = file_in["SIMULATION"].attrs[item]
+
+    # first pass -- find total size of each track
+    # ----------------------------------------#
+    sizes = np.zeros(ntracks)
+
+    itermapshape = itermap.shape
+    for i in range(itermapshape[0]):
+        part_number, npoints, nstart = itermap[i, :]
+        sizes[part_number - 1] += npoints
+
+    # initialize ordered data buffer
+    # ----------------------------------------#
+    ordered_data = []
+    for i in range(ntracks):
+        ordered_data.append(np.zeros((int(sizes[i]), nquants)))
+    # ----------------------------------------#
+
+    # assign tracks to ordered data from file_in data
+    # ----------------------------------------#
+    track_indices = np.zeros(ntracks)
+    data_index = 0
+
+    for i in range(itermapshape[0]):
+        part_number, npoints, nstart = itermap[i, :]
+        track_index = int(track_indices[part_number - 1])
+
+        ordered_data[part_number - 1][track_index : track_index + npoints, 0] = nstart + np.arange(npoints) * niter
+
+        ordered_data[part_number - 1][track_index : track_index + npoints, 1:] = data[data_index : data_index + npoints, :]
+
+        data_index += npoints
+        track_indices[part_number - 1] += npoints
+
+    # ----------------------------------------#
+
+    # write to file out
+    for i in range(ntracks):
+        group = file_out.create_group(str(i + 1))
+        for j in range(nquants):
+            if j == 0:
+                group.create_dataset(quants[j], data=np.array(ordered_data[i][:, j], dtype=int))
+            else:
+                group.create_dataset(quants[j], data=ordered_data[i][:, j])
+
+    file_out.close()
+    file_in.close()
+    print("Track file converted to the old, more readable format: ", filename_out)
+    return filename_out
+
+
+def create_file_tags(filename, tags_array):
+    """
+    Function to write a file_tags file from a (number_of_tags, 2) NumPy array of tags.
+    this file is used to choose particles for the OSIRIS track diagnostic.
+
+    Parameters
+    ----------
+    filename : str
+        Path to the output file where tags will be stored.
+    tags_array: np.ndarray
+        shape (number_of_tags, 2), containing particle tags
+
+    Returns
+    -------
+    file_tags file with path \"filename\" to be used for the OSIRIS track diagnostic.
+
+    Notes
+    ------
+    The first element of the tag of a particle that is already being tracked is negative,
+        so we apply the absolute function when generating the file
+
+    """
+
+    # In case the particles chosen were already being tracked
+    tags_array[:, 0] = np.abs(tags_array[:, 0])
+    tags_array = tags_array[np.lexsort((tags_array[:, 1], tags_array[:, 0]))]
+    num_tags = tags_array.shape[0]
+
+    with open(filename, "w") as file:
+        file.write("! particle tag list\n")
+        file.write(f"! generated on {datetime.now().strftime('%a %b %d %H:%M:%S %Y')}\n")
+        file.write("! number of tags\n")
+        file.write(f"       {num_tags}\n")
+        file.write("! particle tag list\n")
+
+        for i in range(num_tags):
+            file.write(f"         {tags_array[i, 0]:<6}{tags_array[i, 1]:>6}\n")
+    return filename

{osiris_utils-1.1.4.dist-info → osiris_utils-1.1.7.dist-info}/METADATA

@@ -1,13 +1,13 @@
 Metadata-Version: 2.4
 Name: osiris_utils
-Version: 1.1.4
-Summary: Utilities to manipulate and visualize OSIRIS framework output data
-Author: ['João Pedro Ferreira Biu', 'João Cândido', 'Diogo Carvalho']
-Author-email: ['joaopedrofbiu@tecnico.ulisboa.pt']
+Version: 1.1.7
+Summary: Utilities to manipulate and visualise OSIRIS plasma PIC output data
+Author: João Cândido, Diogo Carvalho
+Author-email: João Pedro Ferreira Biu <joaopedrofbiu@tecnico.ulisboa.pt>
 License: MIT
-Project-URL: Issues Tracker, https://github.com/joaopedrobiu6/osiris_utils/issues
 Project-URL: Source Code, https://github.com/joaopedrobiu6/osiris_utils
-Keywords: OSIRIS,particle-in-cell,plasma,physics
+Project-URL: Issues Tracker, https://github.com/joaopedrobiu6/osiris_utils/issues
+Project-URL: Documentation, https://osiris-utils.readthedocs.io
 Classifier: Development Status :: 3 - Alpha
 Classifier: Intended Audience :: Science/Research
 Classifier: License :: OSI Approved :: MIT License
@@ -22,31 +22,45 @@ Classifier: Topic :: Scientific/Engineering
 Requires-Python: >=3.10
 Description-Content-Type: text/x-rst
 License-File: LICENSE.txt
-Requires-Dist: numpy
-Requires-Dist: matplotlib
-Requires-Dist: pandas
-Requires-Dist: scipy
-Requires-Dist: h5py
-Requires-Dist: pyside6
-Requires-Dist: pyarrow
-Requires-Dist: tqdm
-Dynamic: author
-Dynamic: author-email
-Dynamic: classifier
-Dynamic: description
-Dynamic: description-content-type
-Dynamic: keywords
-Dynamic: license
+Requires-Dist: numpy>=1.25
+Requires-Dist: matplotlib>=3.8
+Requires-Dist: pandas>=2.0
+Requires-Dist: scipy>=1.11
+Requires-Dist: h5py>=3.10
+Requires-Dist: tqdm>=4.66
+Provides-Extra: dev
+Requires-Dist: pytest>=8.1; extra == "dev"
+Requires-Dist: pytest-cov; extra == "dev"
+Requires-Dist: pre-commit>=3.7; extra == "dev"
+Requires-Dist: ruff>=0.4; extra == "dev"
+Requires-Dist: isort>=5.13; extra == "dev"
+Requires-Dist: mypy>=1.10; extra == "dev"
+Requires-Dist: types-tqdm; extra == "dev"
+Requires-Dist: types-Pillow; extra == "dev"
+Provides-Extra: docs
+Requires-Dist: sphinx>=8; extra == "docs"
+Requires-Dist: ipykernel; extra == "docs"
+Requires-Dist: myst-parser; extra == "docs"
+Requires-Dist: sphinx-copybutton<=0.5.2; extra == "docs"
+Requires-Dist: sphinx-rtd-theme<3.0.3; extra == "docs"
+Requires-Dist: sphinx-github-style<=1.2.2; extra == "docs"
+Provides-Extra: plot
+Requires-Dist: seaborn>=0.13; extra == "plot"
 Dynamic: license-file
-Dynamic: project-url
-Dynamic: requires-dist
-Dynamic: requires-python
-Dynamic: summary
 
 OSIRIS_UTILS
 ============
 |Pypi|
 
+.. image:: https://github.com/joaopedrobiu6/osiris_utils/actions/workflows/ci.yml/badge.svg
+   :target: https://github.com/joaopedrobiu6/osiris_utils/actions
+   :alt: CI status
+.. image:: https://codecov.io/gh/joaopedrobiu6/osiris_utils/branch/main/graph/badge.svg
+   :target: https://codecov.io/gh/joaopedrobiu6/osiris_utils
+   :alt: Coverage
+
+
+
 This package contains a set of utilities to open and analyze OSIRIS output files, using Python. All the methods implemented are fully integrated with `NumPy`, and use `np.ndarray` as the main data structure.
 High-level functions are provided to manipulate data from OSIRIS, from reading the data of the diagnostics, to making post-processing calculations.
 
@@ -70,6 +84,25 @@ Finally, you can install it in editor mode if you want to contribute to the code
     git clone https://github.com/joaopedrobiu6/osiris_utils.git
     pip install -e .
 
+Quick-start
+-----------
+
+.. image:: https://mybinder.org/badge_logo.svg
+   :target: https://mybinder.org/v2/gh/joaopedrobiu6/osiris_utils/main?filepath=examples%2Fquick_start.ipynb
+   :alt: Launch quick-start on Binder
+
+.. code-block:: bash
+
+   pip install osiris_utils              # from PyPI
+   python -m pip install matplotlib      # plotting backend (optional)
+   git clone https://github.com/joaopedrobiu6/osiris_utils  # sample data
+   cd osiris_utils
+   python examples/quick_start.py examples/example_data/thermal.1d
+
+.. image:: docs/source/_static/quick_start_ez.png
+   :alt: Example Ez plot produced by the quick-start
+   :width: 400px
+
 Documentation
 -------------