osiris-utils 1.1.10__py3-none-any.whl → 1.2.0__py3-none-any.whl

osiris_utils/postprocessing/fft.py

@@ -1,3 +1,8 @@
+from __future__ import annotations
+
+from collections.abc import Generator
+from typing import Any
+
 import numpy as np
 import tqdm as tqdm
 
@@ -5,6 +10,8 @@ from ..data.diagnostic import Diagnostic
 from ..data.simulation import Simulation
 from .postprocess import PostProcess
 
+__all__ = ["FFT_Simulation", "FFT_Diagnostic", "FFT_Species_Handler"]
+
 
 class FFT_Simulation(PostProcess):
     """
@@ -21,16 +28,16 @@ class FFT_Simulation(PostProcess):
 
     """
 
-    def __init__(self, simulation, fft_axis):
+    def __init__(self, simulation: Simulation, fft_axis: int | list[int]) -> None:
         super().__init__("FFT")
         if not isinstance(simulation, Simulation):
-            raise ValueError("Simulation must be a Simulation object.")
+            raise ValueError("simulation must be a Simulation-compatible object.")
         self._simulation = simulation
         self._fft_axis = fft_axis
-        self._fft_computed = {}
-        self._species_handler = {}
+        self._fft_computed: dict[Any, FFT_Diagnostic] = {}
+        self._species_handler: dict[Any, FFT_Species_Handler] = {}
 
-    def __getitem__(self, key):
+    def __getitem__(self, key: Any) -> FFT_Species_Handler | FFT_Diagnostic:
         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)
@@ -40,16 +47,16 @@ class FFT_Simulation(PostProcess):
             self._fft_computed[key] = FFT_Diagnostic(self._simulation[key], self._fft_axis)
         return self._fft_computed[key]
 
-    def delete_all(self):
+    def delete_all(self) -> None:
         self._fft_computed = {}
 
-    def delete(self, key):
+    def delete(self, key: Any) -> None:
         if key in self._fft_computed:
             del self._fft_computed[key]
         else:
             print(f"FFT {key} not found in simulation")
 
-    def process(self, diagnostic):
+    def process(self, diagnostic: Diagnostic) -> FFT_Diagnostic:
         """Apply FFT to a diagnostic"""
         return FFT_Diagnostic(diagnostic, self._fft_axis)
 
@@ -77,7 +84,7 @@ class FFT_Diagnostic(Diagnostic):
 
     """
 
-    def __init__(self, diagnostic, fft_axis):
+    def __init__(self, diagnostic: Diagnostic, fft_axis: int | list[int]) -> None:
         if hasattr(diagnostic, "_species"):
             super().__init__(
                 simulation_folder=(diagnostic._simulation_folder if hasattr(diagnostic, "_simulation_folder") else None),
@@ -113,9 +120,11 @@ class FFT_Diagnostic(Diagnostic):
         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])
+            # Handle if fft_axis is int
+            axes = [self._fft_axis] if isinstance(self._fft_axis, int) else self._fft_axis
+            self._kmax = np.pi / np.array([self._dx[ax - 1] for ax in axes if ax != 0])
 
-    def load_all(self):
+    def load_all(self) -> np.ndarray:
         if self._data is not None:
             print("Using cached data.")
             return self._data
@@ -165,7 +174,7 @@ class FFT_Diagnostic(Diagnostic):
         self._data = np.abs(result) ** 2
         return self._data
 
-    def _data_generator(self, index):
+    def _data_generator(self, index: int) -> Generator[np.ndarray, None, None]:
         # Get the data for this index
         original_data = self._diag[index]
 
@@ -195,10 +204,10 @@ class FFT_Diagnostic(Diagnostic):
 
         yield np.abs(result_fft) ** 2
 
-    def _get_window(self, length, axis):
+    def _get_window(self, length: int, axis: int) -> np.ndarray:
         return np.hanning(length)
 
-    def _apply_window(self, data, window, axis):
+    def _apply_window(self, data: np.ndarray, window: np.ndarray, axis: int) -> np.ndarray:
         ndim = data.ndim
         window_shape = [1] * ndim
         window_shape[axis] = len(window)
@@ -207,7 +216,7 @@ class FFT_Diagnostic(Diagnostic):
 
         return data * reshaped_window
 
-    def __getitem__(self, index):
+    def __getitem__(self, index: int | slice) -> np.ndarray:
         if self._all_loaded and self._data is not None:
             return self._data[index]
 
@@ -221,29 +230,131 @@ class FFT_Diagnostic(Diagnostic):
         else:
             raise ValueError("Invalid index type. Use int or slice.")
 
-    def omega(self):
+    def omega(self) -> np.ndarray:
         """
-        Get the angular frequency array for the FFT.
+        Get the angular frequency array for the FFT along the time dimension (axis 0).
+
+        Returns
+        -------
+        np.ndarray
+            Angular frequency array for the time axis.
         """
         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
+        # If the FFT was computed along the time axis (0) return temporal frequencies
+        if isinstance(self._fft_axis, (list, tuple)):
+            if 0 in self._fft_axis:
+                dt = self._dt * self._ndump
+                omega = np.fft.fftfreq(self._data.shape[0], d=dt) * 2 * np.pi
+                return np.fft.fftshift(omega)
+            # If FFT was computed along spatial axes only and a single spatial axis
+            spatial_axes = [ax for ax in self._fft_axis if ax != 0]
+            if len(spatial_axes) == 1:
+                return self.k(spatial_axes[0])
+            # Multi-dimensional spatial FFT: return concatenated or dict of k arrays
+            return self.k()
+        else:
+            if self._fft_axis == 0:
+                dt = self._dt * self._ndump
+                omega = np.fft.fftfreq(self._data.shape[0], d=dt) * 2 * np.pi
+                return np.fft.fftshift(omega)
+            # Single spatial axis: return wavenumber array for that axis
+            return self.k(self._fft_axis)
+
+    def k(self, axis: int | None = None) -> np.ndarray | dict[int, np.ndarray]:
+        """
+        Get the wavenumber array for the FFT along spatial dimension(s).
+
+        Parameters
+        ----------
+        axis : int or None, optional
+            The spatial axis to compute wavenumber for (1, 2, or 3).
+            If None, returns wavenumbers for all spatial axes in fft_axis.
+
+        Returns
+        -------
+        np.ndarray or dict
+            If axis is specified: wavenumber array for that axis.
+            If axis is None: dictionary mapping axis -> wavenumber array.
+
+        Notes
+        -----
+        When load_all() is used, time axis is 0 and spatial axes are 1,2,3.
+        When accessing single timesteps, spatial axes are 0,1,2.
+        """
+        if self._data is None:
+            raise ValueError("Load the data first using load_all() or access via indexing.")
+
+        # Determine if we have the time dimension in the data
+        # If all_loaded is True, then axis 0 is time, spatial axes are 1,2,3
+        # If all_loaded is False, we're looking at a single timestep, spatial axes are 0,1,2
+        has_time_axis = self._all_loaded
+
+        # Determine which axes to compute k for
+        if axis is not None:
+            # Single axis specified
+            if axis == 0:
+                raise ValueError("axis must be a spatial dimension (1, 2, or 3), not 0 (time).")
+            if axis < 1 or axis > 3:
+                raise ValueError(f"axis must be 1, 2, or 3, got {axis}")
+
+            # Get dx for this axis
+            if isinstance(self._dx, (int, float)):
+                dx = self._dx
+            else:
+                dx = self._dx[axis - 1]
+
+            # Compute the actual data axis index
+            # If we have time axis, spatial axis N is at index N
+            # If no time axis (single timestep), spatial axis N is at index N-1
+            data_axis = axis if has_time_axis else axis - 1
+
+            # Compute wavenumber
+            k_array = np.fft.fftfreq(self._data.shape[data_axis], d=dx) * 2 * np.pi
+            k_array = np.fft.fftshift(k_array)
+            return k_array
+        else:
+            # axis is None: return k for all spatial axes in fft_axis
+            result = {}
+
+            if isinstance(self._fft_axis, (list, tuple)):
+                # Multi-axis FFT: return k for all spatial axes
+                spatial_axes = [ax for ax in self._fft_axis if ax != 0]
+            elif self._fft_axis == 0:
+                # Only time FFT, no spatial axes
+                raise ValueError("No spatial FFT axes to compute wavenumber for. fft_axis is 0 (time only).")
+            else:
+                # Single spatial axis
+                spatial_axes = [self._fft_axis]
+
+            for ax in spatial_axes:
+                if isinstance(self._dx, (int, float)):
+                    dx = self._dx
+                else:
+                    dx = self._dx[ax - 1]
+
+                # Compute the actual data axis index
+                data_axis = ax if has_time_axis else ax - 1
+
+                k_array = np.fft.fftfreq(self._data.shape[data_axis], d=dx) * 2 * np.pi
+                k_array = np.fft.fftshift(k_array)
+                result[ax] = k_array
+
+            return result
 
     @property
-    def kmax(self):
+    def kmax(self) -> float | np.ndarray:
         return self._kmax
 
 
 class FFT_Species_Handler:
-    def __init__(self, species_handler, fft_axis):
+    def __init__(self, species_handler: Any, fft_axis: int | list[int]) -> None:
         self._species_handler = species_handler
         self._fft_axis = fft_axis
-        self._fft_computed = {}
+        self._fft_computed: dict[Any, FFT_Diagnostic] = {}
 
-    def __getitem__(self, key):
+    def __getitem__(self, key: Any) -> FFT_Diagnostic:
         if key not in self._fft_computed:
             diag = self._species_handler[key]
             self._fft_computed[key] = FFT_Diagnostic(diag, self._fft_axis)

osiris_utils/postprocessing/field_centering.py

@@ -1,9 +1,15 @@
+from __future__ import annotations
+
+from collections.abc import Generator
+
 import numpy as np
 
 from ..data.diagnostic import OSIRIS_FLD, Diagnostic
 from ..data.simulation import Simulation
 from .postprocess import PostProcess
 
+__all__ = ["FieldCentering_Simulation", "FieldCentering_Diagnostic"]
+
 
 class FieldCentering_Simulation(PostProcess):
     """
@@ -22,47 +28,47 @@ class FieldCentering_Simulation(PostProcess):
     """
 
     def __init__(self, simulation: Simulation):
-        super().__init__("FieldCentering Simulation")
         """
         Class to center the field in the simulation.
 
         Parameters
         ----------
-        sim : Simulation
+        simulation : Simulation
             The simulation object.
-        field : str
-            The field to center.
         """
+        super().__init__("FieldCentering Simulation")
+
+        # Accept Simulation-compatible objects (Simulation or other PostProcess subclasses)
         if not isinstance(simulation, Simulation):
-            raise ValueError("Simulation must be a Simulation object.")
+            raise ValueError("simulation must be a Simulation-compatible 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):
+    def __getitem__(self, key: str) -> FieldCentering_Diagnostic:
         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):
+    def delete_all(self) -> None:
         self._field_centered = {}
 
-    def delete(self, key):
+    def delete(self, key: str) -> None:
         if key in self._field_centered:
             del self._field_centered[key]
         else:
             print(f"Field {key} not found in simulation")
 
-    def process(self, diagnostic):
+    def process(self, diagnostic: Diagnostic) -> FieldCentering_Diagnostic:
         """Apply field centering to a diagnostic"""
         return FieldCentering_Diagnostic(diagnostic)
 
 
 class FieldCentering_Diagnostic(Diagnostic):
-    def __init__(self, diagnostic):
+    def __init__(self, diagnostic: Diagnostic):
         """
         Class to center the field in the simulation. It converts fields from the Osiris yee mesh to the center of the cells.
         It only works for periodic boundaries.
@@ -104,10 +110,18 @@ class FieldCentering_Diagnostic(Diagnostic):
         self._original_name = diagnostic._name
         self._name = diagnostic._name + "_centered"
 
-        self._data = None
+        self._data: np.ndarray | None = None
         self._all_loaded = False
 
-    def load_all(self):
+    def load_all(self) -> np.ndarray:
+        """
+        Load all data and center the fields.
+
+        Returns
+        -------
+        data : np.ndarray
+            The centered field data.
+        """
         if self._data is not None:
             return self._data
 
@@ -217,7 +231,7 @@ class FieldCentering_Diagnostic(Diagnostic):
         self._all_loaded = True
         return self._data
 
-    def __getitem__(self, index):
+    def __getitem__(self, index: int | slice) -> np.ndarray:
         """Get data at a specific index"""
         if self._all_loaded and self._data is not None:
             return self._data[index]
@@ -232,7 +246,7 @@ class FieldCentering_Diagnostic(Diagnostic):
         else:
             raise ValueError("Invalid index type. Use int or slice.")
 
-    def _data_generator(self, index):
+    def _data_generator(self, index: int) -> Generator[np.ndarray, None, None]:
         if self._dim == 1:
             if self._original_name.lower() in [
                 "b2",

osiris_utils/postprocessing/heatflux_correction.py

@@ -1,3 +1,8 @@
+from __future__ import annotations
+
+from collections.abc import Generator
+from typing import Any
+
 import numpy as np
 
 from ..data.diagnostic import Diagnostic
@@ -7,9 +12,15 @@ from .pressure_correction import PressureCorrection_Simulation
 
 OSIRIS_H = ["q1", "q2", "q3"]
 
+__all__ = [
+    "HeatfluxCorrection_Simulation",
+    "HeatfluxCorrection_Diagnostic",
+    "HeatfluxCorrection_Species_Handler",
+]
+
 
 class HeatfluxCorrection_Simulation(PostProcess):
-    def __init__(self, simulation):
+    def __init__(self, simulation: Simulation):
         super().__init__("HeatfluxCorrection Simulation")
         """
         Class to correct pressure tensor components by subtracting Reynolds stress.
@@ -22,12 +33,12 @@ class HeatfluxCorrection_Simulation(PostProcess):
             The heatflux component to center.
         """
         if not isinstance(simulation, Simulation):
-            raise ValueError("Simulation must be a Simulation object.")
+            raise ValueError("simulation must be a Simulation-compatible object.")
         self._simulation = simulation
-        self._heatflux_corrected = {}
-        self._species_handler = {}
+        self._heatflux_corrected: dict[str, HeatfluxCorrection_Diagnostic] = {}
+        self._species_handler: dict[str, HeatfluxCorrection_Species_Handler] = {}
 
-    def __getitem__(self, key):
+    def __getitem__(self, key: str) -> HeatfluxCorrection_Species_Handler | HeatfluxCorrection_Diagnostic:
         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)
@@ -36,25 +47,35 @@ class HeatfluxCorrection_Simulation(PostProcess):
             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)
+            # This part seems to lack some arguments for HeatfluxCorrection_Diagnostic,
+            # but keeping as is for structural consistency if reached
+            # self._heatflux_corrected[key] = HeatfluxCorrection_Diagnostic(self._simulation[key], self._simulation)
         return self._heatflux_corrected[key]
 
-    def delete_all(self):
+    def delete_all(self) -> None:
         self._heatflux_corrected = {}
 
-    def delete(self, key):
+    def delete(self, key: str) -> None:
         if key in self._heatflux_corrected:
             del self._heatflux_corrected[key]
         else:
             print(f"Heatflux {key} not found in simulation")
 
-    def process(self, diagnostic):
+    def process(self, diagnostic: Diagnostic) -> HeatfluxCorrection_Diagnostic:
         """Apply heatflux correction to a diagnostic"""
+        # FIX: This is a bit of a hack, but it works for now
         return HeatfluxCorrection_Diagnostic(diagnostic, self._simulation)
 
 
 class HeatfluxCorrection_Diagnostic(Diagnostic):
-    def __init__(self, diagnostic, vfl_i, Pjj_list, vfl_j_list, Pji_list):
+    def __init__(
+        self,
+        diagnostic: Diagnostic,
+        vfl_i: Diagnostic,
+        Pjj_list: list[Diagnostic],
+        vfl_j_list: list[Diagnostic],
+        Pji_list: list[Diagnostic],
+    ):
         """
         Class to correct the pressure in the simulation.
 
@@ -102,10 +123,10 @@ class HeatfluxCorrection_Diagnostic(Diagnostic):
         self._original_name = diagnostic._name
         self._name = diagnostic._name + "_corrected"
 
-        self._data = None
+        self._data: np.ndarray | None = None
         self._all_loaded = False
 
-    def load_all(self):
+    def load_all(self) -> np.ndarray:
         if self._data is not None:
             return self._data
 
@@ -131,12 +152,12 @@ class HeatfluxCorrection_Diagnostic(Diagnostic):
         # 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, strict=False))
 
-        self._data = 2 * q - 0.5 * vfl_i * trace_P - vfl_dot_Pji
+        self._data = 2 * q - vfl_i * trace_P - 2 * vfl_dot_Pji
         self._all_loaded = True
 
         return self._data
 
-    def __getitem__(self, index):
+    def __getitem__(self, index: int | slice) -> np.ndarray:
         """Get data at a specific index"""
         if self._all_loaded and self._data is not None:
             return self._data[index]
@@ -151,7 +172,7 @@ class HeatfluxCorrection_Diagnostic(Diagnostic):
         else:
             raise ValueError("Invalid index type. Use int or slice.")
 
-    def _data_generator(self, index):
+    def _data_generator(self, index: int) -> Generator[np.ndarray, None, None]:
         q = self._diag[index]
         vfl_i = self._vfl_i[index]
         trace_P = sum(Pjj[index] for Pjj in self._Pjj_list)
@@ -174,12 +195,12 @@ class HeatfluxCorrection_Species_Handler:
         The simulation object.
     """
 
-    def __init__(self, species_handler, simulation):
+    def __init__(self, species_handler: Any, simulation: Simulation):
         self._species_handler = species_handler
         self._simulation = simulation
-        self._heatflux_corrected = {}
+        self._heatflux_corrected: dict[str, HeatfluxCorrection_Diagnostic] = {}
 
-    def __getitem__(self, key):
+    def __getitem__(self, key: str) -> HeatfluxCorrection_Diagnostic:
         if key not in self._heatflux_corrected:
             diag = self._species_handler[key]
 

osiris_utils/postprocessing/mft.py

@@ -1,9 +1,17 @@
+from __future__ import annotations
+
 import numpy as np
 
 from ..data.diagnostic import Diagnostic
 from ..data.simulation import Simulation
 from .postprocess import PostProcess
 
+__all__ = [
+    "MFT_Simulation",
+    "MFT_Diagnostic",
+    "MFT_Species_Handler",
+]
+
 
 class MFT_Simulation(PostProcess):
     """
@@ -19,10 +27,10 @@ class MFT_Simulation(PostProcess):
 
     """
 
-    def __init__(self, simulation, mft_axis=None):
+    def __init__(self, simulation: Simulation, mft_axis: int | None = None):
         super().__init__(f"MeanFieldTheory({mft_axis})")
         if not isinstance(simulation, Simulation):
-            raise ValueError("Simulation must be a Simulation object.")
+            raise ValueError("simulation must be a Simulation-compatible object.")
         self._simulation = simulation
         self._mft_axis = mft_axis
         self._mft_computed = {}

osiris_utils/postprocessing/postprocess.py

@@ -1,7 +1,11 @@
-from ..data.diagnostic import Diagnostic
+from __future__ import annotations
 
+from ..data.simulation import Simulation
 
-class PostProcess(Diagnostic):
+__all__ = ["PostProcess"]
+
+
+class PostProcess(Simulation):
     """
     Base class for post-processing operations.
     Inherits from Diagnostic to ensure all operation overloads work.
@@ -14,14 +18,13 @@ class PostProcess(Diagnostic):
         The species to analyze.
     """
 
-    def __init__(self, name, species=None):
+    def __init__(self, name: str, species: str = None):
         # Initialize with the same interface as Diagnostic
-        super().__init__(species)
         self._name = name
         self._all_loaded = False
         self._data = None
 
-    def process(self, diagnostic):
+    def process(self, simulation: Simulation) -> Simulation:
         """
         Apply the post-processing to a diagnostic.
         Must be implemented by subclasses.

osiris_utils/postprocessing/pressure_correction.py

@@ -1,3 +1,8 @@
+from __future__ import annotations
+
+from collections.abc import Generator
+from typing import Any
+
 import numpy as np
 
 from ..data.diagnostic import Diagnostic
@@ -6,9 +11,15 @@ from .postprocess import PostProcess
 
 OSIRIS_P = ["P11", "P12", "P13", "P21", "P22", "P23", "P31", "P32", "P33"]
 
+__all__ = [
+    "PressureCorrection_Simulation",
+    "PressureCorrection_Diagnostic",
+    "PressureCorrection_Species_Handler",
+]
+
 
 class PressureCorrection_Simulation(PostProcess):
-    def __init__(self, simulation):
+    def __init__(self, simulation: Simulation):
         super().__init__("PressureCorrection Simulation")
         """
         Class to correct pressure tensor components by subtracting Reynolds stress.
@@ -21,12 +32,12 @@ class PressureCorrection_Simulation(PostProcess):
             The pressure component to center.
         """
         if not isinstance(simulation, Simulation):
-            raise ValueError("Simulation must be a Simulation object.")
+            raise ValueError("simulation must be a Simulation-compatible object.")
         self._simulation = simulation
-        self._pressure_corrected = {}
-        self._species_handler = {}
+        self._pressure_corrected: dict[str, PressureCorrection_Diagnostic] = {}
+        self._species_handler: dict[str, PressureCorrection_Species_Handler] = {}
 
-    def __getitem__(self, key):
+    def __getitem__(self, key: str) -> PressureCorrection_Species_Handler | PressureCorrection_Diagnostic:
         if key in self._simulation._species:
             if key not in self._species_handler:
                 self._species_handler[key] = PressureCorrection_Species_Handler(self._simulation[key])
@@ -35,25 +46,26 @@ class PressureCorrection_Simulation(PostProcess):
             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)
+            # self._pressure_corrected[key] = PressureCorrection_Diagnostic(self._simulation[key], self._simulation)
         return self._pressure_corrected[key]
 
-    def delete_all(self):
+    def delete_all(self) -> None:
         self._pressure_corrected = {}
 
-    def delete(self, key):
+    def delete(self, key: str) -> None:
         if key in self._pressure_corrected:
             del self._pressure_corrected[key]
         else:
             print(f"Pressure {key} not found in simulation")
 
-    def process(self, diagnostic):
+    def process(self, diagnostic: Diagnostic) -> PressureCorrection_Diagnostic:
         """Apply pressure correction to a diagnostic"""
+        # FIX: This is a bit of a hack, but it works for now
         return PressureCorrection_Diagnostic(diagnostic, self._simulation)
 
 
 class PressureCorrection_Diagnostic(Diagnostic):
-    def __init__(self, diagnostic, n, ufl_j, vfl_k):
+    def __init__(self, diagnostic: Diagnostic, n: Diagnostic, ufl_j: Diagnostic, vfl_k: Diagnostic):
         """
         Class to correct the pressure in the simulation.
 
@@ -100,10 +112,10 @@ class PressureCorrection_Diagnostic(Diagnostic):
         self._original_name = diagnostic._name
         self._name = diagnostic._name + "_corrected"
 
-        self._data = None
+        self._data: np.ndarray | None = None
         self._all_loaded = False
 
-    def load_all(self):
+    def load_all(self) -> np.ndarray:
         if self._data is not None:
             return self._data
 
@@ -130,7 +142,7 @@ class PressureCorrection_Diagnostic(Diagnostic):
 
         return self._data
 
-    def __getitem__(self, index):
+    def __getitem__(self, index: int | slice) -> np.ndarray:
         """Get data at a specific index"""
         if self._all_loaded and self._data is not None:
             return self._data[index]
@@ -145,7 +157,7 @@ class PressureCorrection_Diagnostic(Diagnostic):
         else:
             raise ValueError("Invalid index type. Use int or slice.")
 
-    def _data_generator(self, index):
+    def _data_generator(self, index: int) -> Generator[np.ndarray, None, None]:
         yield (self._diag[index] - self._n[index] * self._vfl_k[index] * self._ufl_j[index])
 
 
@@ -166,11 +178,11 @@ class PressureCorrection_Species_Handler:
         The axis to compute the derivative. Only used for 'xx', 'xt' and 'tx' types.
     """
 
-    def __init__(self, species_handler):
+    def __init__(self, species_handler: Any):
         self._species_handler = species_handler
-        self._pressure_corrected = {}
+        self._pressure_corrected: dict[str, PressureCorrection_Diagnostic] = {}
 
-    def __getitem__(self, key):
+    def __getitem__(self, key: str) -> PressureCorrection_Diagnostic:
         if key not in self._pressure_corrected:
             diag = self._species_handler[key]