osiris-utils 1.1.4__tar.gz → 1.1.6__tar.gz

{osiris_utils-1.1.4/osiris_utils.egg-info → osiris_utils-1.1.6}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: osiris_utils
-Version: 1.1.4
+Version: 1.1.6
 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']

{osiris_utils-1.1.4 → osiris_utils-1.1.6}/osiris_utils/__init__.py

@@ -8,7 +8,7 @@ from .utils import (
     courant2D,
 )
 from .gui.gui import LAVA_Qt, LAVA
-from .data.data import OsirisGridFile, OsirisRawFile, OsirisData, OsirisHIST
+from .data.data import OsirisGridFile, OsirisRawFile, OsirisData, OsirisHIST, OsirisTrackFile
 from .data.simulation import Simulation, Species_Handler
 from .data.diagnostic import Diagnostic
 
@@ -25,4 +25,10 @@ from .postprocessing.mft import (
     MFT_Diagnostic,
     MFT_Diagnostic_Average,
     MFT_Diagnostic_Fluctuations,
-)
+)
+
+from .postprocessing.field_centering import FieldCentering_Simulation, FieldCentering_Diagnostic
+
+from .postprocessing.pressure_correction import PressureCorrection_Simulation, PressureCorrection_Diagnostic
+
+from .postprocessing.heatflux_correction import HeatfluxCorrection_Simulation, HeatfluxCorrection_Diagnostic

{osiris_utils-1.1.4 → osiris_utils-1.1.6}/osiris_utils/data/data.py

@@ -1,6 +1,8 @@
 import numpy as np
 import pandas as pd
 import h5py
+from osiris_utils.utils import create_file_tags
+from typing import Literal
 
 class OsirisData():
     """
@@ -343,42 +345,65 @@ class OsirisRawFile(OsirisData):
     '''
     Class to read the raw data from an OSIRIS HDF5 file.
     
-    Input:
-        - filename: the path to the HDF5 file
+    Parameters
+    ----------
+    filename : str
+        Path to OSIRIS HDF5 track file (.h5 extension)
     
     Attributes:
-        - axis - a dictionary where each key is a dataset name, and each value is another dictionary containing
-            name (str): The name of the quantity (e.g., r'x1', r'ene').
-            units (str): The units associated with that dataset in LaTeX (e.g., r'c/\\omega_p', r'm_e c^2').
-            long_name (str): The name of the quantity in LaTeX (e.g., r'x_1', r'En2').
-            dictionary of dictionaries
-        - data - a dictionary where each key is a dataset name, and each value is the data
-            dictionary of np.arrays
-        - dim - the number of dimensions
-            int
-        - dt - the time step
-            float
-        - grid - maximum and minimum coordinates of the box, for each axis 
-            numpy.ndarray(dim,2)
-        - iter - the iteration number
-            int
-        - name - the name of the species
-            str
-        - time - the time and its units
-            list [time, units]
-            list [float, str]
-        - type - type of data (particles in the case of raw files)
-            str
-
-    '''
+    -----------
+    axis : dict[str, dict[str, str]]
+        Dictionary where each key is a dataset name, and each value is another dictionary containing:
+            - 'name' (str): Short name of the quantity (e.g., 'x1', 'ene')
+            - 'units' (str): Units (LaTeX formatted, e.g., 'c/\\omega_p', 'm_e c^2')
+            - 'long_name' (str): Descriptive name (LaTeX formatted, e.g., 'x_1', 'En2')
+    data : dict[str, np.ndarray]
+        Dataset values indexed by dataset name (quants).
+    dim : int
+        Number of spatial dimensions.
+    dt : float
+        Time step between iterations.
+    grid : np.ndarray
+        Grid boundaries as ((x1_min, x1_max), (x2_min, x2_max), ...)
+    iter : int
+        Iteration number corresponding to the data.
+    name : str
+        Name of the species.
+    time : list[float, str]
+        Simulation time and its units (e.g., [12.5, '1/\\omega_p']).
+    type : str
+        Type of data (e.g., 'particles' for raw files).
+    labels : list[str]
+        Field labels/names (LaTeX formatted, e.g., 'x_1')
+    quants : list[str]
+        field names of the data
+    units : list[str]
+        Units of each field of the data (LaTeX formatted, e.g., 'c/\\omega_p')        
     
+    Example
+    -------
+        >>> import osiris_utils as ou  
+        >>> raw = ou.raw = ou.OsirisRawFile("path/to/raw/file.h5")
+        >>> print(raw.data.keys())
+        >>> # Access x1 position of first 10 particles
+        >>> print(raw.data[\"x1\"][0:10])
+        >>> # Write beautiful labels and units
+        >>> print("${} = $".format(raw.labels[\"x1\"]) + "$[{}]$".format(track.units[\"x1\"]))
+    '''
+
     def __init__(self, filename):
         super().__init__(filename)
 
-        self.grid = np.array([self._file['SIMULATION'].attrs['XMIN'], self._file['SIMULATION'].attrs['XMAX']]).T
+        self._grid = np.array([self._file['SIMULATION'].attrs['XMIN'], self._file['SIMULATION'].attrs['XMAX']]).T
 
-        self.data = {}
-        self.axis = {}
+        self._quants = [byte.decode('utf-8') for byte in self._file.attrs['QUANTS'][:]]
+        units_list = [byte.decode('utf-8') for byte in self._file.attrs['UNITS'][:]]
+        labels_list = [byte.decode('utf-8') for byte in self._file.attrs['LABELS'][:]]
+        self._units = dict(zip(self._quants, units_list))
+        self._labels = dict(zip(self._quants, labels_list))
+
+        self._data = {}
+        self._axis = {}
         for key in self._file.keys():
             if key == 'SIMULATION': continue
 
@@ -390,29 +415,278 @@ class OsirisRawFile(OsirisData):
                 'units': self._file.attrs['UNITS'][idx][0].decode('utf-8'),
                 'long_name': self._file.attrs['LABELS'][idx][0].decode('utf-8'),
             }
-            self.axis[key] = axis_data
+            self._axis[key] = axis_data
+
+    def raw_to_file_tags(self, filename, type: Literal["all", "random"] = "all", n_tags=10, mask=None):
+        """
+        Function to write a file_tags file from raw data.
+        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.
+        type : {'all', 'random'}, optional
+            Selection mode for tags:
+            - 'all': Includes all available tags.
+            - 'random': Randomly selects `n_tags` tags.
+        n_tags : int, optional
+            Number of tags to randomly select when `type` is 'random'. Default is 10.
+        mask : np.ndarray, optional
+            Boolean mask array applied to filter valid tags before selection.
+
+        Returns
+        ------
+        A 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
+
+        """
+            
+        if mask is not None:
+            # Apply mask to select certain tags
+            if not isinstance(mask, np.ndarray) or mask.dtype != bool or mask.shape[0] != self.data["tag"].shape[0]:
+                raise ValueError("Mask must be a boolean NumPy array of the same length as 'tag'.")
+            filtered_indices = np.where(mask)[0]
+            filtered_tags = self.data["tag"][filtered_indices]
+        else:
+            filtered_tags = self.data["tag"]
+        
+        if type == "all":
+            tags = filtered_tags
+        elif type == "random":
+            if len(filtered_tags) < n_tags:
+                raise ValueError("Not enough tags to sample from.")
+            random_indices = np.random.choice(len(filtered_tags), size=n_tags, replace=False)
+            tags = filtered_tags[random_indices]
+        else:
+            raise TypeError("Invalid type", type)
+
+        create_file_tags(filename, tags)
+        print("Tag_file created: ", filename)
+
+    # Getters
+    @property
+    def grid(self):
+        return self._grid
+    @property
+    def data(self):
+        return self._data
+    @property
+    def units(self):
+        return self._units
+    @property
+    def labels(self):
+        return self._labels
+    @property
+    def quants(self):
+        return self._quants
+    @property
+    def axis(self):
+        return self._axis    
 
 class OsirisHIST(OsirisData):
     ''''
     Class to read the data from an OSIRIS HIST file.'
 
-    Input:
-        - filename: the path to the HIST file
+    Input
+    -----
+    filename: the path to the HIST file
 
-    Attributes:
-        - filename - the path to the file
-            str
-        - verbose - if True, the class will print messages
-            bool
-        - df - the data in a pandas DataFrame
-            pandas.DataFrame
+    Attributes
+    ----------
+    filename: the path to the file
+        str
+    df: the data in a pandas DataFrame
+        pandas.DataFrame
     '''
     def __init__(self, filename):
         super().__init__(filename)
 
     @property
     def df(self):
-        """
-        Returns the data in a pandas DataFrame
-        """
         return self._df
+
+class OsirisTrackFile(OsirisData):
+    """
+    Handles structured track data from OSIRIS HDF5 simulations.
+
+    Parameters
+    ----------
+    filename : str
+        Path to OSIRIS HDF5 track file (.h5 extension)
+
+    Attributes
+    ----------
+    data: numpy.ndarray of shape (num_particles, num_time_iter), 
+        dtype = [(field_name, float) for field_name in field_names]
+        A structured numpy array with the track data
+        Accessed as data[particles, time_iters][quant]
+    grid : np.ndarray
+        Grid boundaries as ((x1_min, x1_max), (x2_min, x2_max), ...)
+    labels : list[str]
+        Field labels/names (LaTeX formatted, e.g., 'x_1')
+    num_particles : int
+        Number of particlest tracked, they are accessed from 0 to num_particles-1
+    num_time_iters : int
+        Number of time iteratis, they are accessed from 0 to num_time_iters-1
+    quants : list[str]
+        field names of the data
+    units : list[str]
+        Units of each field of the data (LaTeX formatted, e.g., 'c/\\omega_p')
+    
+    Example
+    -------
+        >>> import osiris_utils as ou
+        >>> track = ou.OsirisTrackFile(path/to/track_file.h5)
+        >>> print(track.data[0:10, :]["x1"]) # Access x1 position of first 10 particles over all time steps
+    """
+
+    def __init__(self, filename):
+        super().__init__(filename)
+        
+        self._grid = np.array([self._file['SIMULATION'].attrs['XMIN'], self._file['SIMULATION'].attrs['XMAX']]).T
+
+        self._quants = [byte.decode('utf-8') for byte in self._file.attrs['QUANTS'][1:]]
+        units_list = [byte.decode('utf-8') for byte in self._file.attrs['UNITS'][1:]]
+        labels_list = [byte.decode('utf-8') for byte in self._file.attrs['LABELS'][1:]]
+        self._units = dict(zip(self._quants, units_list))
+        self._labels = dict(zip(self._quants, labels_list))
+        
+        self._num_particles = self._file.attrs['NTRACKS'][0]
+
+        unordered_data = self._file['data'][:]
+        itermap = self._file['itermap'][:]
+        
+        idxs = get_track_indexes(itermap, self._num_particles)
+        self._data = reorder_track_data(unordered_data, idxs, self._quants)
+        self._time = self._data[0][:]["t"]
+        self._num_time_iters = np.shape(self._time.shape)
+        self._close_file()
+
+    def _load_basic_attributes(self, f: h5py.File) -> None:
+        '''Load common attributes from HDF5 file'''
+        self._dt = float(f['SIMULATION'].attrs['DT'][0])
+        self._dim = int(f['SIMULATION'].attrs['NDIMS'][0])
+        self._time = None
+        self._iter = None
+        self._name = f.attrs['NAME'][0].decode('utf-8')
+        self._type = f.attrs['TYPE'][0].decode('utf-8')
+
+
+    # Getters
+    @property
+    def grid(self):
+        return self._grid
+    @property
+    def data(self):
+        return self._data
+    @property
+    def units(self):
+        return self._units
+    @property
+    def labels(self):
+        return self._labels
+    @property
+    def quants(self):
+        return self._quants
+    @property
+    def num_particles(self):
+        return self._num_particles
+    @property
+    def num_time_iters(self):
+        return self._num_time_iters   
+
+
+    # Setters
+    @data.setter
+    def data(self, data):
+        self._data = data
+
+    def __str__(self):
+        # write me a template to print with the name, label, units, iter, grid, nx, dx, axis, dt, dim in a logical way
+        return rf'{self.name}' + f'\n' + f'Iteration: {self.iter}' + f'\n' + f'Grid: {self.grid}' + f'\n' + f'dx: {self.dx}' + f'\n' + f'Dimensions: {self.dim}D'
+
+    def __array__(self):
+        return np.asarray(self.data)
+
+
+
+def reorder_track_data(unordered_data, indexes, field_names):
+    '''
+    Reorder data from HDF5 track file such data it can be accessed more intuitively
+    
+    Parameters
+    ----------
+    unordered_data: np.array
+        The data from a HDF5 osiris track file
+    
+    indexes : list[list[int]]
+        Output of get_track_indexes(), list with the indexes associated with each particle
+
+    field_names: list[str]
+        Names for the quantities on the output file. 
+        Recommended: field_names = [byte.decode('utf-8') for byte in file.attrs['QUANTS'][1:]]
+    
+    Returns
+    -------
+    data_sorted: numpy.ndarray of shape (num_particles, num_time_iter), 
+                    dtype = [(field_name, float) for field_name in field_names]
+        A structured numpy array where data is reordered according to indexes.
+    
+    '''
+    # Initialize the sorted data structure
+    num_particles = len(indexes)
+    num_time_iter = len(indexes[0])
+    data_sorted = np.empty((num_particles, num_time_iter), dtype=[(name, float) for name in field_names])
+
+    # Fill the sorted data based on the indexes
+    for particle in range(num_particles):
+        for time_iter in range(num_time_iter):
+            index = indexes[particle][time_iter]
+            if len(unordered_data[index]) != len(field_names):
+                raise ValueError(f"Data at index {index} has {len(unordered_data[index])} elements, "
+                                f"but {len(field_names)} are expected.")
+            data_sorted[particle, time_iter] = tuple(unordered_data[index])
+
+    return (data_sorted)
+
+
+def get_track_indexes(itermap, num_particles):
+    '''
+    Returns the indexes for each particle to read track data directly from the hd5 file
+    (before it is ordered)
+
+    Parameters
+    ----------
+    itermap: np.array
+        Itermap from a HDF5 osiris track file
+    num_particles: int
+        num of particles tracked, recomended file.attrs['NTRACKS'][0]
+    
+    Returns
+    -------     
+    indexes : list[list[int]]
+        Returns a list with the indexes associated with each particle
+        shape(num_particles, num_time_iters)
+    '''
+
+    itermapshape = itermap.shape
+    for i in range(itermapshape[0]):
+        part_number,npoints,nstart = itermap[i,:]
+    track_indices = np.zeros(num_particles)
+
+    data_index = 0
+    indexes = [[] for _ in range(num_particles)]
+    for i in range(itermapshape[0]):
+        part_number,npoints,nstart = itermap[i,:]
+
+        indexes[part_number-1].extend(list(range(data_index, data_index + npoints)))
+
+        data_index += npoints
+        track_indices[part_number-1] += npoints
+
+    return indexes