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

osiris_utils/postprocessing/fft.py

@@ -0,0 +1,234 @@
+from ..utils import *
+from ..data.simulation import Simulation
+from .postprocess import PostProcess
+from ..data.diagnostic import Diagnostic
+import numpy as np
+import tqdm as tqdm
+
+
+class FastFourierTransform_Simulation(PostProcess):
+    """
+    Class to handle the Fast Fourier Transform on data. Works as a wrapper for the FFT_Diagnostic class.
+    Inherits from PostProcess to ensure all operation overloads work properly.
+    
+    Parameters
+    ----------
+    
+    simulation : Simulation
+        The simulation object.
+    axis : int
+        The axis to compute the FFT.
+
+    """
+    def __init__(self, simulation, fft_axis):
+        super().__init__("FFT")
+        if not isinstance(simulation, Simulation):
+            raise ValueError("Simulation must be a Simulation object.")
+        self._simulation = simulation
+        self._fft_axis = fft_axis
+        self._fft_computed = {}
+        self._species_handler = {}
+
+    def __getitem__(self, key):
+        if key in self._simulation._species:
+            if key not in self._species_handler:
+                self._species_handler[key] = FFT_Species_Handler(self._simulation[key], self._fft_axis)
+            return self._species_handler[key]
+        
+        if key not in self._fft_computed:
+            self._fft_computed[key] = FFT_Diagnostic(self._simulation[key], self._fft_axis)
+        return self._fft_computed[key]
+    
+    def delete_all(self):
+        self._fft_computed = {}
+
+    def delete(self, key):
+        if key in self._fft_computed:
+            del self._fft_computed[key]
+        else:
+            print(f"FFT {key} not found in simulation")
+
+    def process(self, diagnostic):
+        """Apply FFT to a diagnostic"""
+        return FFT_Diagnostic(diagnostic, self._fft_axis)
+    
+
+class FFT_Diagnostic(Diagnostic):
+    """
+    Auxiliar class to compute the FFT of a diagnostic, for it to be similar in behavior to a Diagnostic object.
+    Inherits directly from Diagnostic to ensure all operation overloads work properly.
+
+    Parameters
+    ---------- 
+    diagnostic : Diagnostic
+        The diagnostic to compute the FFT.
+    axis : int
+        The axis to compute the FFT.
+
+    Methods
+    -------
+    load_all()
+        Load all the data and compute the FFT.
+    omega()
+        Get the angular frequency array for the FFT.
+    __getitem__(index)
+        Get data at a specific index.
+
+    """
+    def __init__(self, diagnostic, fft_axis):
+        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"FFT"
+
+        self._name = f"FFT[{diagnostic._name}, {fft_axis}]"
+        self._diag = diagnostic
+        self._fft_axis = fft_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']:
+            if hasattr(diagnostic, attr):
+                setattr(self, attr, getattr(diagnostic, attr))
+
+        if isinstance(self._dx, (int, float)):
+            self._kmax = np.pi / (self._dx) 
+        else:
+            self._kmax = np.pi / np.array([self._dx[ax-1] for ax in self._fft_axis if ax != 0])
+    
+    def load_all(self):
+        if self._data is not None:
+            print("Using cached data.")
+            return self._data
+        
+        if not hasattr(self._diag, '_data') or self._diag._data is None:
+            self._diag.load_all()
+            self._diag._data = np.nan_to_num(self._diag._data)
+
+       # Apply appropriate windows based on which axes we're transforming
+        if isinstance(self._fft_axis, (list, tuple)):
+            if self._diag._data is None:
+                raise ValueError(f"Unable to load data for diagnostic {self._diag._name}. The data is None even after loading.")
+
+            result = self._diag._data.copy()
+            
+            for axis in self._fft_axis:
+                if axis == 0:  # Time axis
+                    window = np.hanning(result.shape[0]).reshape(-1, *([1] * (result.ndim - 1)))
+                    result = result * window
+                else:  # Spatial axis
+                    window = self._get_window(result.shape[axis], axis)
+                    result = self._apply_window(result, window, axis)
+                    
+            with tqdm.tqdm(total=1, desc="FFT calculation") as pbar:
+                data_fft = np.fft.fftn(result, axes=self._fft_axis)
+                pbar.update(0.5)
+                result = np.fft.fftshift(data_fft, axes=self._fft_axis)
+                pbar.update(0.5)
+        
+        else:
+            if self._fft_axis == 0:
+                hanning_window = np.hanning(self._diag._data.shape[0]).reshape(-1, *([1] * (self._diag._data.ndim - 1)))
+                data_windowed = hanning_window * self._diag._data
+            else: 
+                window = self._get_window(self._diag._data.shape[self._fft_axis], self._fft_axis)
+                data_windowed = self._apply_window(self._diag._data, window, self._fft_axis)
+            
+            with tqdm.tqdm(total=1, desc="FFT calculation") as pbar:
+                data_fft = np.fft.fft(data_windowed, axis=self._fft_axis)
+                pbar.update(0.5)
+                result = np.fft.fftshift(data_fft, axes=self._fft_axis)
+                pbar.update(0.5)
+
+        self.omega_max = np.pi / self._dt / self._ndump
+
+        self._all_loaded = True
+        self._data = np.abs(result)**2
+        return self._data
+    
+    def _data_generator(self, index):
+        # Get the data for this index
+        original_data = self._diag[index]
+        
+        if self._fft_axis == 0:
+            raise ValueError("Cannot generate FFT along time axis for a single timestep. Use load_all() instead.")
+        
+        # For spatial FFT, we can apply a spatial window if desired
+        if isinstance(self._fft_axis, (list, tuple)):
+            result = original_data
+            for axis in self._fft_axis:
+                if axis != 0:  # Skip time axis
+                    # Apply window along this spatial dimension
+                    window = self._get_window(original_data.shape[axis-1], axis-1)
+                    result = self._apply_window(result, window, axis-1)
+                    
+            # Compute FFT
+            result_fft = np.fft.fftn(result, axes=[ax-1 for ax in self._fft_axis if ax != 0])
+            result_fft = np.fft.fftshift(result_fft, axes=[ax-1 for ax in self._fft_axis if ax != 0])
+            
+        else:
+            if self._fft_axis > 0:  # Spatial axis
+                window = self._get_window(original_data.shape[self._fft_axis-1], self._fft_axis-1)
+                windowed_data = self._apply_window(original_data, window, self._fft_axis-1)
+                
+                result_fft = np.fft.fft(windowed_data, axis=self._fft_axis-1)
+                result_fft = np.fft.fftshift(result_fft, axes=self._fft_axis-1)
+        
+        yield np.abs(result_fft)**2
+        
+    def _get_window(self, length, axis):
+        return np.hanning(length)
+
+    def _apply_window(self, data, window, axis):
+        ndim = data.ndim
+        window_shape = [1] * ndim
+        window_shape[axis] = len(window)
+        
+        reshaped_window = window.reshape(window_shape)
+        
+        return data * reshaped_window
+
+    def __getitem__(self, 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 omega(self):
+        """
+        Get the angular frequency array for the FFT.
+        """
+        if not self._all_loaded:
+            raise ValueError("Load the data first using load_all() method.")
+        
+        omega = np.fft.fftfreq(self._data.shape[self._fft_axis], d=self._dx[self._fft_axis-1])
+        omega = np.fft.fftshift(omega)
+        return omega
+    
+    @property
+    def kmax(self):
+        return self._kmax
+    
+class FFT_Species_Handler:
+    def __init__(self, species_handler, fft_axis):
+        self._species_handler = species_handler
+        self._fft_axis = fft_axis
+        self._fft_computed = {}
+
+    def __getitem__(self, key):
+        if key not in self._fft_computed:
+            diag = self._species_handler[key]
+            self._fft_computed[key] = FFT_Diagnostic(diag, self._fft_axis)
+        return self._fft_computed[key]

osiris_utils/postprocessing/field_centering.py

@@ -0,0 +1,168 @@
+from ..utils import *
+from ..data.simulation import Simulation
+from .postprocess import PostProcess
+from ..data.diagnostic import Diagnostic
+
+OSIRIS_FLD = ["e1", "e2", "e3", "b1", "b2", "b3"]
+
+
+class FieldCentering_Simulation(PostProcess):
+    """
+    Class to handle the field centering on data. Works as a wrapper for the FieldCentering_Diagnostic class.
+    Inherits from PostProcess to ensure all operation overloads work properly.
+    
+    Parameters
+    ----------
+    simulation : Simulation
+        The simulation object.
+    field : str
+        The field to center.
+    """
+    
+    def __init__(self, simulation: Simulation):
+        super().__init__(f"FieldCentering Simulation")
+        """
+        Class to center the field in the simulation.
+
+        Parameters
+        ----------
+        sim : Simulation
+            The simulation object.
+        field : str
+            The field to center.
+        """
+        if not isinstance(simulation, Simulation):
+            raise ValueError("Simulation must be a Simulation object.")
+        self._simulation = simulation
+
+        self._field_centered = {}
+        # no need to create a species handler for field centering since fields are not species related 
+
+    def __getitem__(self, key):
+        if key not in OSIRIS_FLD:
+            raise ValueError(f"Does it make sense to center {key} field? Only {OSIRIS_FLD} are supported.")
+        if key not in self._field_centered:
+            self._field_centered[key] = FieldCentering_Diagnostic(self._simulation[key])
+        return self._field_centered[key]
+    
+    def delete_all(self):
+        self._field_centered = {}
+
+    def delete(self, key):
+        if key in self._field_centered:
+            del self._field_centered[key]
+        else:
+            print(f"Field {key} not found in simulation")
+    
+    def process(self, diagnostic):
+        """Apply field centering to a diagnostic"""
+        return FieldCentering_Diagnostic(diagnostic)
+    
+class FieldCentering_Diagnostic(Diagnostic):
+    def __init__(self, diagnostic):
+
+        """
+        Class to center the field 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 = "FLD_CTR"
+
+        if diagnostic._name not in OSIRIS_FLD:
+            raise ValueError(f"Does it make sense to center {diagnostic._name} field? Only {OSIRIS_FLD} are supported.")
+        
+        self._diag = diagnostic
+
+        for attr in ['_dt', '_dx', '_ndump', '_axis', '_nx', '_x', '_grid', '_dim', '_maxiter']:
+            if hasattr(diagnostic, attr):
+                setattr(self, attr, getattr(diagnostic, attr))
+
+        self._original_name = diagnostic._name
+        self._name = diagnostic._name + "_centered"
+
+        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()
+
+        if self._dim == 1:
+            if self._original_name.lower() in ['b2', 'b3', 'e1']:
+                result = 0.5 * (np.roll(self._diag.data, shift=1, axis=1) + self._diag.data) 
+            elif self._original_name.lower() in ['b1', 'e2', 'e3']:
+                result = self._diag.data 
+
+        elif self._dim == 2:
+            if self._original_name.lower() in ['e1', 'b2']:
+                result = 0.5 * (np.roll(self._diag.data, shift=1, axis=1) + self._diag.data)
+            elif self._original_name.lower() in ['e2', 'b1']:
+                result = 0.5 * (np.roll(self._diag.data, shift=1, axis=2) + self._diag.data)
+            elif self._original_name.lower() in ['b3']:
+                result = 0.5 * (np.roll((0.5 * (np.roll(self._diag.data, shift=1, axis=1) + self._diag.data)), shift=1, axis=2) + (0.5 * (np.roll(self._diag.data, shift=1, axis=1) + self._diag.data)))
+            elif self._original_name.lower() in ['e3']:
+                result = self._diag.data
+            
+        elif self._dim == 3:
+            raise NotImplementedError("3D field centering is not implemented yet.")
+        
+        else:
+            raise ValueError(f"Unknown dimension {self._dim}.")
+        
+        self._data = result
+        self._all_loaded = True
+        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):
+        if self._dim == 1:
+            if self._original_name.lower() in ['b2', 'b3', 'e1']:
+                yield 0.5 * (np.roll(self._diag[index], shift=1) + self._diag[index])
+            elif self._original_name.lower() in ['b1', 'e2', 'e3']: # it's already centered but self._data does not exist 
+                yield self._diag[index]
+            else:
+                raise ValueError(f"Unknown field {self._original_name}.")
+            
+        elif self._dim == 2:
+            if self._original_name in ['e1', 'b2']:
+                yield 0.5 * (np.roll(self._diag[index], shift=1, axis=0) + self._diag[index])
+            elif self._original_name in ['e2', 'b1']:
+                yield 0.5 * (np.roll(self._diag[index], shift=1, axis=1) + self._diag[index])
+            elif self._original_name in ['b3']:
+                yield 0.5 * (np.roll((0.5 * (np.roll(self._diag[index], shift=1, axis=0) + self._diag[index])), shift=1, axis=1) + (0.5 * (np.roll(self._diag[index], shift=1, axis=0) + self._diag[index])))
+            elif self._original_name in ['e3']:
+                yield self._diag[index]
+            else:
+                raise ValueError(f"Unknown field {self._original_name}.")
+            
+        elif self._dim == 3:
+            raise NotImplementedError("3D field centering is not implemented yet.")
+        
+        else:
+            raise ValueError(f"Unknown dimension {self._dim}.")

osiris_utils/postprocessing/heatflux_correction.py

@@ -0,0 +1,193 @@
+from ..utils import *
+from ..data.simulation import Simulation
+from .postprocess import PostProcess
+from ..data.diagnostic import Diagnostic
+
+from .pressure_correction import *
+
+OSIRIS_H = ["q1", "q2", "q3"]
+
+class HeatfluxCorrection_Simulation(PostProcess):
+    def __init__(self, simulation):
+        super().__init__(f"HeatfluxCorrection Simulation")
+        """
+        Class to correct pressure tensor components by subtracting Reynolds stress.
+
+        Parameters
+        ----------
+        sim : Simulation
+            The simulation object.
+        heatflux : str
+            The heatflux component to center.
+        """
+        if not isinstance(simulation, Simulation):
+            raise ValueError("Simulation must be a Simulation object.")
+        self._simulation = simulation
+        self._heatflux_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] = HeatfluxCorrection_Species_Handler(self._simulation[key], self._simulation)
+            return self._species_handler[key]
+        if key not in OSIRIS_H:
+            raise ValueError(f"Invalid heatflux component {key}. Supported: {OSIRIS_H}.")
+        if key not in self._heatflux_corrected:
+            print("Weird that it got here - heatflux is always species dependent on OSIRIS")
+            self._heatflux_corrected[key] = HeatfluxCorrection_Diagnostic(self._simulation[key], self._simulation)
+        return self._heatflux_corrected[key]
+
+    
+    def delete_all(self):
+        self._heatflux_corrected = {}
+
+    def delete(self, key):
+        if key in self._heatflux_corrected:
+            del self._heatflux_corrected[key]
+        else:
+            print(f"Heatflux {key} not found in simulation")
+    
+    def process(self, diagnostic):
+        """Apply heatflux correction to a diagnostic"""
+        return HeatfluxCorrection_Diagnostic(diagnostic, self._simulation)
+    
+class HeatfluxCorrection_Diagnostic(Diagnostic):
+    def __init__(self, diagnostic, vfl_i, Pjj_list, vfl_j_list, Pji_list):
+
+        """
+        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 = "HFL_CORR"
+
+        if diagnostic._name not in OSIRIS_H:
+            raise ValueError(f"Invalid heatflux component {diagnostic._name}. Supported: {OSIRIS_H}")
+        
+        self._diag = diagnostic
+
+        # The density and velocities are now passed as arguments (so it can doesn't depend on the simulation)
+        self._vfl_i = vfl_i
+        self._Pjj_list = Pjj_list
+        self._vfl_j_list = vfl_j_list
+        self._Pji_list = Pji_list
+
+        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._vfl_i.load_all()
+        
+
+        for vfl_j in self._vfl_j_list:
+            vfl_j.load_all()
+        for Pji in self._Pji_list:
+            Pji.load_all()
+        for Pjj in self._Pjj_list:
+            Pjj.load_all()
+
+        q = self._diag.data
+        vfl_i = self._vfl_i.data
+        
+        trace_P = sum(Pjj.data for Pjj in self._Pjj_list)
+
+        # Sum over j: vfl_j * Pji
+        vfl_dot_Pji = sum(vfl_j.data * Pji.data for vfl_j, Pji in zip(self._vfl_j_list, self._Pji_list))
+
+        self._data = 2 * q - 0.5 * vfl_i * trace_P - vfl_dot_Pji
+        self._all_loaded = True
+
+
+        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):
+        q = self._diag[index]
+        vfl_i = self._vfl_i[index]
+        trace_P = sum(Pjj[index] for Pjj in self._Pjj_list)
+        vfl_dot_Pji = sum(vfl_j[index] * Pji[index] for vfl_j, Pji in zip(self._vfl_j_list, self._Pji_list))
+        yield 2 * q - 0.5 * vfl_i * trace_P - vfl_dot_Pji
+        
+class HeatfluxCorrection_Species_Handler:
+    """
+    Class to handle heatflux correction for a species.
+    Acts as a wrapper for the HeatfluxCorrection_Diagnostic class.
+
+    Not intended to be used directly, but through the HeatfluxCorrection_Simulation class.
+
+    Parameters
+    ----------
+    species_handler : Species_Handler
+        The species handler object.
+    simulation : Simulation
+        The simulation object.
+    """
+    def __init__(self, species_handler, simulation):
+        self._species_handler = species_handler
+        self._simulation = simulation
+        self._heatflux_corrected = {}
+
+    def __getitem__(self, key):
+        if key not in self._heatflux_corrected:
+            diag = self._species_handler[key]
+
+            # Velocities alwayes depend on the species so this can be done here
+
+            i = int(key[-1])  # Get i from 'q1', 'q2', etc.
+            
+            vfl_i = self._species_handler[f"vfl{i}"]
+
+            # Load trace(P): sum over Pjj
+            Pjj_list = [self._species_handler[f"P{j}{j}"] for j in range(1, diag._dim + 1)]
+
+            # Compute quantities for vfl_j * P_{ji}
+            vfl_j_list = [self._species_handler[f"vfl{j}"] for j in range(1, diag._dim + 1)]
+            Pji_list = [PressureCorrection_Simulation(self._simulation)[diag._species._name][f"P{j}{i}"] for j in range(1, diag._dim + 1)]
+
+            self._heatflux_corrected[key] = HeatfluxCorrection_Diagnostic(
+                diag,
+                vfl_i,
+                Pjj_list,
+                vfl_j_list,
+                Pji_list
+            )
+        return self._heatflux_corrected[key]