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

osiris_utils/gui/gui.py

@@ -0,0 +1,266 @@
+import sys
+import os
+from PySide6.QtWidgets import (QApplication, QMainWindow, QWidget, QPushButton, 
+                              QFileDialog, QMessageBox, QComboBox, QHBoxLayout,
+                              QVBoxLayout, QLabel, QLineEdit, QFrame, QDoubleSpinBox)
+from PySide6.QtCore import Qt
+from matplotlib.backends.backend_qt5agg import FigureCanvasQTAgg as FigureCanvas
+import matplotlib.pyplot as plt
+from matplotlib.colors import LogNorm
+import numpy as np
+from ..data.data import OsirisGridFile  # Update import as needed
+from ..utils import integrate, transverse_average  # Update import as needed
+
+class LAVA_Qt(QMainWindow):
+    def __init__(self):
+        super().__init__()
+        self.setWindowTitle('LAVA (LabAstro Visualization Assistant) - OSIRIS Data Grid Viewer')
+        self.setGeometry(100, 100, 1000, 600)
+        
+        # Initialize data
+        self.data_info = None
+        self.dims = 0
+        self.current_ax = None
+        self.current_folder = None
+        
+        # Main widget and layout
+        self.main_widget = QWidget()
+        self.setCentralWidget(self.main_widget)
+        self.main_layout = QVBoxLayout(self.main_widget)
+        
+        # Create UI elements
+        self.create_controls()
+        self.create_labels_section()
+        self.create_plot_area()
+
+    def create_controls(self):
+        # Control buttons frame
+        control_frame = QWidget()
+        control_layout = QHBoxLayout(control_frame)
+        
+        # Buttons
+        self.browse_btn = QPushButton('Browse Folder')
+        self.browse_btn.clicked.connect(self.load_folder)
+        self.save_btn = QPushButton('Save Plot')
+        self.save_btn.clicked.connect(self.save_plot)
+        
+        # File selector
+        self.file_selector = QComboBox()
+        self.file_selector.setPlaceholderText('Select file...')
+        self.file_selector.currentIndexChanged.connect(self.file_selection_changed)
+        self.file_selector.view().setVerticalScrollBarPolicy(Qt.ScrollBarAlwaysOn)
+        self.file_selector.setSizeAdjustPolicy(QComboBox.AdjustToContents)
+        
+        # Plot type combo box
+        self.plot_combo = QComboBox()
+        self.plot_combo.addItem('Select Plot Type')
+        self.plot_combo.currentTextChanged.connect(self.plot_data)
+        
+        control_layout.addWidget(self.browse_btn)
+        control_layout.addWidget(self.save_btn)
+        control_layout.addWidget(QLabel('Files:'))
+        control_layout.addWidget(self.file_selector)
+        control_layout.addWidget(QLabel('Plot Type:'))
+        control_layout.addWidget(self.plot_combo)
+        self.main_layout.addWidget(control_frame)
+
+    def create_labels_section(self):
+        # Labels frame
+        labels_frame = QWidget()
+        labels_layout = QHBoxLayout(labels_frame)
+        
+        # Title and labels
+        self.title_edit = QLineEdit()
+        self.xlabel_edit = QLineEdit()
+        self.ylabel_edit = QLineEdit()
+        
+        # Connect text changes
+        self.title_edit.textChanged.connect(self.update_plot_labels)
+        self.xlabel_edit.textChanged.connect(self.update_plot_labels)
+        self.ylabel_edit.textChanged.connect(self.update_plot_labels)
+        
+        
+        labels_layout.addWidget(QLabel('Title:'))
+        labels_layout.addWidget(self.title_edit)
+        labels_layout.addWidget(QLabel('X Label:'))
+        labels_layout.addWidget(self.xlabel_edit)
+        labels_layout.addWidget(QLabel('Y Label:'))
+        labels_layout.addWidget(self.ylabel_edit)
+        
+        # define the size of the labels frame
+        self.main_layout.addWidget(labels_frame)
+
+    
+    def load_folder(self):
+        folder_dialog = QFileDialog()
+        folderpath = folder_dialog.getExistingDirectory(
+            self, 'Select Folder with HDF5 Files'
+        )
+        
+        if not folderpath:
+            return
+            
+        try:
+            self.current_folder = folderpath
+            self.file_selector.clear()
+            
+            # Find all .h5 files
+            h5_files = [f for f in os.listdir(folderpath) if f.endswith('.h5')]
+            # all the files end with xxxxxx.h5 so we can use this to order them by the number
+            def sort_key(filename):
+                try:
+                    # Split filename into parts and get the numeric portion
+                    base = os.path.splitext(filename)[0]  # Remove .h5
+                    numeric_part = base.split('-')[-1]    # Get last part after -
+                    return int(numeric_part)
+                except (IndexError, ValueError):
+                    return 0  # Fallback for malformed filenames
+
+            h5_files.sort(key=sort_key)
+            
+            if not h5_files:
+                raise ValueError('No HDF5 files found in selected folder')
+            
+            self.file_selector.addItems(h5_files)
+            self.file_selector.setCurrentIndex(0)
+            
+        except Exception as e:
+            QMessageBox.critical(self, 'Error', str(e))
+
+    def file_selection_changed(self, index):
+        '''Handle file selection change in the combo box'''
+        if index >= 0 and self.current_folder:
+            filename = self.file_selector.itemText(index)
+            self.process_file(filename)
+
+    def process_file(self, filename):
+        try:
+            filepath = os.path.join(self.current_folder, filename)
+            gridfile = OsirisGridFile(filepath)
+            self.dims = len(gridfile.axis)
+            self.type = gridfile.type
+            
+            if self.type == 'grid':
+                if self.dims == 1:
+                    x = np.arange(gridfile.grid[0], gridfile.grid[1], gridfile.dx)
+                    self.xlabel_edit.setText(r'$%s$ [$%s$]' % (gridfile.axis[0]['long_name'], gridfile.axis[0]['units']))
+                    self.ylabel_edit.setText(r'$%s$ [$%s$]' % (gridfile.label, gridfile.units))
+                    self.data_info = (x, gridfile.data)
+                elif self.dims == 2:
+                    x = np.arange(gridfile.grid[0][0], gridfile.grid[0][1], gridfile.dx[0])
+                    y = np.arange(gridfile.grid[1][0], gridfile.grid[1][1], gridfile.dx[1])
+                    self.xlabel_edit.setText(r'$%s$ [$%s$]' % (gridfile.axis[0]['long_name'], gridfile.axis[0]['units']))
+                    self.ylabel_edit.setText(r'$%s$ [$%s$]' % (gridfile.axis[1]['long_name'], gridfile.axis[1]['units']))
+                    self.data_info = (x, y, gridfile.data)
+                elif self.dims == 3:
+                    raise ValueError('3D not supported yet')
+                else:
+                    raise ValueError('Unsupported dimensionality')
+                
+                self.title_edit.setText(r'$%s$ [$%s$]' %( gridfile.label, gridfile.units))
+                self.update_plot_menu()
+                self.plot_data()
+                
+            else:
+                QMessageBox.information(self, 'Info', f'{self.type} data not supported yet')
+                
+        except Exception as e:
+            QMessageBox.critical(self, 'Error', str(e))
+
+    def create_plot_area(self):
+        # Matplotlib figure and canvas
+        self.figure = plt.figure(figsize=(8, 6))
+        self.canvas = FigureCanvas(self.figure)
+        self.main_layout.addWidget(self.canvas)
+
+    def update_plot_labels(self):
+        if self.current_ax:
+            self.current_ax.set_xlabel(self.xlabel_edit.text())
+            self.current_ax.set_ylabel(self.ylabel_edit.text())
+            self.figure.suptitle(self.title_edit.text())
+            self.canvas.draw()
+
+    def plot_data(self):
+        self.figure.clear()
+        if self.dims == 1:
+            self.plot_1d()
+        elif self.dims == 2:
+            self.plot_2d()
+        self.update_plot_labels()
+        self.canvas.draw()
+
+    def plot_1d(self):
+        x, data = self.data_info
+        self.current_ax = self.figure.add_subplot(111)
+        plot_type = self.plot_combo.currentText()
+        
+        if 'Line' in plot_type:
+            self.current_ax.plot(x, data)
+        elif 'Scatter' in plot_type:
+            self.current_ax.scatter(x, data)
+            
+        self.current_ax.set_xlabel(self.xlabel_edit.text())
+        self.current_ax.set_ylabel(self.ylabel_edit.text())
+        self.figure.suptitle(self.title_edit.text())
+
+    def plot_2d(self):
+        x, y, data = self.data_info
+        self.current_ax = self.figure.add_subplot(111)
+        plot_type = self.plot_combo.currentText()
+        
+        if 'Quantity' in plot_type:
+            img = self.current_ax.imshow(data.T, extent=(x[0], x[-1], y[0], y[-1]), origin='lower', aspect='auto')
+            self.figure.colorbar(img)
+        elif 'Integral' in plot_type:
+            avg = integrate(transverse_average(data), x[-1]/len(x))
+            self.current_ax.plot(x, avg)
+        elif 'Transverse' in plot_type:
+            avg = transverse_average(data)
+            self.current_ax.plot(x, avg)
+        elif 'Phase' in plot_type:
+            img = self.current_ax.imshow(np.abs(-data.T), extent=(x[0], x[-1], y[0], y[-1]), origin='lower', aspect='auto', norm=LogNorm())
+            self.figure.colorbar(img)
+        
+        self.current_ax.set_xlabel(self.xlabel_edit.text())
+        self.current_ax.set_ylabel(self.ylabel_edit.text())
+        self.figure.suptitle(self.title_edit.text())
+
+    def update_plot_menu(self):
+        
+        # Save current plot type before clearing
+        current_plot_type = self.plot_combo.currentText()
+        self.plot_combo.clear()
+        
+        # Determine items based on dimensions
+        if self.dims == 1:
+            items = ['Line Plot', 'Scatter Plot']
+        elif self.dims == 2:
+            items = ['Quantity Plot', 'T. Average Integral', 'Transverse Average', 'Phase Space']
+        else:
+            items = []
+        
+        self.plot_combo.addItems(items)
+    
+        # Restore previous selection if possible
+        if current_plot_type in items:
+            self.plot_combo.setCurrentText(current_plot_type)
+        else:
+            self.plot_combo.setCurrentIndex(0 if items else -1)
+
+    def save_plot(self):
+        file_dialog = QFileDialog()
+        filepath, _ = file_dialog.getSaveFileName(
+            self, 'Save Plot', '', 'PNG Files (*.png);;PDF Files (*.pdf)'
+        )
+        
+        if filepath:
+            self.figure.savefig(filepath, dpi=800, bbox_inches='tight')
+
+def LAVA():
+    app = QApplication(sys.argv)
+    window = LAVA_Qt()
+    window.show()
+    sys.exit(app.exec())
+
+if __name__ == '__main__':
+    LAVA()

osiris_utils/postprocessing/derivative.py

@@ -0,0 +1,243 @@
+from ..utils import *
+from ..data.simulation import Simulation
+from .postprocess import PostProcess
+from ..data.diagnostic import Diagnostic
+
+class Derivative_Simulation(PostProcess):
+    """
+    Class to compute the derivative of a diagnostic. Works as a wrapper for the Derivative_Diagnostic class.
+    Inherits from PostProcess to ensure all operation overloads work properly.
+
+    Parameters
+    ----------
+    simulation : Simulation
+        The simulation object.
+    type : str
+        The type of derivative to compute. Options are:
+        - 't' for time derivative.
+        - 'x1' for first spatial derivative.
+        - 'x2' for second spatial derivative.
+        - 'x3' for third spatial derivative.
+        - 'xx' for second spatial derivative in two axis.
+        - 'xt' for mixed derivative in time and one spatial axis.
+        - 'tx' for mixed derivative in one spatial axis and time.
+    axis : int or tuple
+        The axis to compute the derivative. Only used for 'xx', 'xt' and 'tx' types.
+
+    Example
+    -------
+    >>> sim = Simulation('electrons', 'path/to/simulation')
+    >>> derivative = Derivative(sim, 'x1')
+    >>> deriv_e1_wrt_x1 = derivative['e1']
+    """
+
+    def __init__(self, simulation, type, axis=None):
+        super().__init__(f"Derivative({type})")
+        if not isinstance(simulation, Simulation):
+            raise ValueError("Simulation must be a Simulation object.")
+        self._simulation = simulation
+        self._type = type
+        self._axis = axis
+        self._derivatives_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] = Derivative_Species_Handler(self._simulation[key], self._type, self._axis)
+            return self._species_handler[key]
+                
+        if key not in self._derivatives_computed:
+            self._derivatives_computed[key] = Derivative_Diagnostic(diagnostic=self._simulation[key], 
+                                                                    type=self._type, axis=self._axis)
+        return self._derivatives_computed[key]
+    
+    def delete_all(self):
+        self._derivatives_computed = {}
+
+    def delete(self, key):
+        if key in self._derivatives_computed:
+            del self._derivatives_computed[key]
+        else:
+            print(f"Derivative {key} not found in simulation")
+    
+    def process(self, diagnostic):
+        """Apply derivative to a diagnostic"""
+        return Derivative_Diagnostic(diagnostic, self._type, self._axis)
+
+    
+class Derivative_Diagnostic(Diagnostic):
+    """
+    Auxiliar class to compute the derivative 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 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
+
+    Methods
+    -------
+    load_all()
+        Load all the data and compute the derivative.
+    __getitem__(index)
+        Get data at a specific index.
+
+    Example
+    -------
+    >>> sim = Simulation('electrons', 'path/to/simulation')
+    >>> diag = sim['e1']
+    >>> derivative = Derivative_Diagnostic(diag, 'x1')
+    """
+
+    def __init__(self, diagnostic, type, axis=None):
+        # Initialize using parent's __init__ with the same species
+        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._name = f"D[{diagnostic._name}, {type}]"
+        self._diag = diagnostic
+        self._type = type
+        self._axis = axis if axis is not None else diagnostic._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']:
+            if hasattr(diagnostic, attr):
+                setattr(self, attr, getattr(diagnostic, attr))
+
+    def load_metadata(self):
+        """Copy metadata from original diagnostic to ensure consistency"""
+        self._dt = self._diag._dt
+        self._dx = self._diag._dx
+        self._ndump = self._diag._ndump
+        self._axis = self._diag._axis
+        self._nx = self._diag._nx
+        self._x = self._diag._x
+        self._grid = self._diag._grid
+        self._dim = self._diag._dim
+        self._maxiter = self._diag._maxiter
+
+    def load_all(self):
+        """Load all data and compute the derivative"""
+        if self._data is not None:
+            print("Using cached derivative")
+            return self._data
+        
+        if not hasattr(self._diag, '_data') or self._diag._data is None:
+            self._diag.load_all()
+
+        if self._type == "t":
+            result = np.gradient(self._diag._data, self._diag._dt * self._diag._ndump, axis=0, edge_order=2)
+
+        elif self._type == "x1":
+            if self._dim == 1:
+                result = np.gradient(self._diag._data, self._diag._dx, axis=1, edge_order=2)
+            else:
+                result = np.gradient(self._diag._data, self._diag._dx[0], axis=1, edge_order=2)
+                    
+        elif self._type == "x2":
+            result = np.gradient(self._diag._data, self._diag._dx[0], axis=2, edge_order=2)
+
+        elif self._type == "x3":
+            result = np.gradient(self._diag._data, self._diag._dx[0], axis=3, edge_order=2)
+
+        elif self._type == "xx":
+            if len(self._axis) != 2:
+                raise ValueError("Axis must be a tuple with two elements.")
+            result = np.gradient(np.gradient(self._diag._data, self._diag._dx[self._axis[0]-1], axis=self._axis[0], edge_order=2), 
+                                self._diag._dx[self._axis[1]-1], axis=self._axis[1], edge_order=2)
+            
+        elif self._type == "xt":
+            if not isinstance(self._axis, int):
+                raise ValueError("Axis must be an integer.")
+            result = np.gradient(np.gradient(self._diag._data, self._diag._dt, axis=0, edge_order=2), 
+                                self._diag._dx[self._axis-1], axis=self._axis[0], edge_order=2)
+            
+        elif self._type == "tx":
+            if not isinstance(self._axis, int):
+                raise ValueError("Axis must be an integer.")
+            result = np.gradient(np.gradient(self._diag._data, self._diag._dx[self._axis-1], axis=self._axis, edge_order=2), 
+                                self._diag._dt, axis=0, edge_order=2)
+        else:
+            raise ValueError("Invalid derivative type.")
+        
+        # Store the result in the cache
+        self._all_loaded = True
+        self._data = result
+        return self._data
+
+    def _data_generator(self, index):
+        """Generate data for a specific index on-demand"""
+        if self._type == "x1":
+            if self._dim == 1:
+                yield np.gradient(self._diag[index], self._diag._dx, axis=0, edge_order=2)
+            else:
+                yield np.gradient(self._diag[index], self._diag._dx[0], axis=0, edge_order=2)
+        
+        elif self._type == "x2":
+            yield np.gradient(self._diag[index], self._diag._dx[1], axis=1, edge_order=2)
+        
+        elif self._type == "x3":
+            yield np.gradient(self._diag[index], self._diag._dx[2], axis=2, edge_order=2)
+            
+        elif self._type == "t":
+            if index == 0:
+                yield (-3 * self._diag[index] + 4 * self._diag[index + 1] - self._diag[index + 2]) / (2 * self._diag._dt * self._diag._ndump)
+            elif index == self._diag._maxiter - 1:
+                yield (3 * self._diag[index] - 4 * self._diag[index - 1] + self._diag[index - 2]) / (2 * self._diag._dt * self._diag._ndump)
+            else:
+                yield (self._diag[index + 1] - self._diag[index - 1]) / (2 * self._diag._dt * self._diag._ndump)
+        else:
+            raise ValueError("Invalid derivative type. Use 'x1', 'x2', 'x3' or 't'.")
+
+    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.")
+
+class Derivative_Species_Handler:
+    """
+    Class to handle derivatives for a species.
+    Acts as a wrapper for the Derivative_Diagnostic class.
+
+    Not intended to be used directly, but through the Derivative_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, type, axis=None):
+        self._species_handler = species_handler
+        self._type = type
+        self._axis = axis
+        self._derivatives_computed = {}
+
+    def __getitem__(self, key):
+        if key not in self._derivatives_computed:
+            diag = self._species_handler[key]
+            self._derivatives_computed[key] = Derivative_Diagnostic(diag, self._type, self._axis)
+        return self._derivatives_computed[key]