osiris-utils 1.2.2__py3-none-any.whl → 1.2.4__py3-none-any.whl

osiris_utils/postprocessing/fft.py

@@ -1,6 +1,5 @@
 from __future__ import annotations
 
-from collections.abc import Generator
 from typing import Any
 
 import numpy as np
@@ -20,16 +19,14 @@ class FFT_Simulation(PostProcess):
 
     Parameters
     ----------
-
     simulation : Simulation
         The simulation object.
-    axis : int
-        The axis to compute the FFT.
-
+    fft_axis : int or list of int
+        The axis or axes to compute the FFT.
     """
 
     def __init__(self, simulation: Simulation, fft_axis: int | list[int]) -> None:
-        super().__init__("FFT")
+        super().__init__("FFT", simulation)
         if not isinstance(simulation, Simulation):
             raise ValueError("simulation must be a Simulation-compatible object.")
         self._simulation = simulation
@@ -72,19 +69,40 @@ class FFT_Diagnostic(Diagnostic):
         The diagnostic to compute the FFT.
     axis : int
         The axis to compute the FFT.
+    window : str or None, optional
+        The window to apply before computing the FFT. Can be "hann" or None (default "hann").
+    detrend : str or None, optional
+        The detrending method to apply before computing the FFT. Can be "mean" or None (default "mean").
+    normalize : str, optional
+        The normalization method for the FFT. Can be "none", "ortho", or "density" (default "ortho").
+    assume_periodic : bool, optional
+        Whether to assume the data is periodic when choosing default windowing for spatial FFTs (default True).
+        If True, no spatial window is applied by default; if False, the same window is
 
     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.
-
+    load_all() -> np.ndarray
+        Load all data and compute FFT (possibly including time axis). Returns power spectrum \|FFT\|^2 stored in self._data.
+    _frame(index: int, data_slice: tuple \| None = None) -> np.ndarray
+        Per-timestep (lazy) FFT. Only allowed for spatial FFT. If fft_axis includes time (0), user must call load_all().
+        Returns power spectrum \|FFT\|^2.
+    omega() -> np.ndarray
+        Get the angular frequency array for the FFT along the time dimension.
+    k(axis: int \| None = None) -> np.ndarray \| dict[int, np.ndarray]
+        Get the wavenumber array for the FFT along spatial dimension(s).
     """
 
-    def __init__(self, diagnostic: Diagnostic, fft_axis: int | list[int]) -> None:
+    def __init__(
+        self,
+        diagnostic: Diagnostic,
+        fft_axis: int | list[int],
+        *,
+        window: str | None = "hann",  # "hann" or None
+        window_for_spatial: str | None = None,  # default None if periodic; override if desired
+        detrend: str | None = "mean",  # "mean" or None
+        normalize: str = "ortho",  # "none" or "ortho" or "density"
+        assume_periodic: bool = False,  # affects default windowing for spatial FFT)
+    ) -> None:
         if hasattr(diagnostic, "_species"):
             super().__init__(
                 simulation_folder=(diagnostic._simulation_folder if hasattr(diagnostic, "_simulation_folder") else None),
@@ -101,6 +119,20 @@ class FFT_Diagnostic(Diagnostic):
         self._data = None
         self._all_loaded = False
 
+        self._assume_periodic = assume_periodic
+        self._detrend = detrend
+        self._normalize = normalize
+        self._window_time = window
+
+        if window_for_spatial is not None:
+            self._window_spatial = window_for_spatial
+        else:
+            # default: no spatial window for periodic domains
+            self._window_spatial = None if assume_periodic else window
+
+        # numpy supports norm=None or "ortho"
+        self._np_norm = "ortho" if self._normalize == "ortho" else None
+
         # Copy all relevant attributes from diagnostic
         for attr in [
             "_dt",
@@ -125,142 +157,222 @@ class FFT_Diagnostic(Diagnostic):
             self._kmax = np.pi / np.array([self._dx[ax - 1] for ax in axes if ax != 0])
 
     def load_all(self) -> np.ndarray:
+        """
+        Load all data and compute FFT (possibly including time axis).
+        Returns power spectrum \|FFT\|^2 stored in self._data.
+        """
         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:
+        # Ensure base diagnostic is loaded
+        if self._diag._data is None or not getattr(self._diag, "_all_loaded", False):
             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.")
+        data = np.nan_to_num(self._diag._data, copy=True)
 
-            result = self._diag._data.copy()
+        axes_osiris = [self._fft_axis] if isinstance(self._fft_axis, int) else list(self._fft_axis)
 
-            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)
+        # Detrend across transform axes in the FULL array
+        # (here axes_osiris match numpy axes because loaded data includes time at axis 0)
+        data = self._detrend_data(data, axes_osiris)
 
-            with tqdm.tqdm(total=1, desc="FFT calculation") as pbar:
-                data_fft = np.fft.fftn(result, axes=self._fft_axis)
+        # Apply windows along the transform axes
+        result = data
+        for ax in axes_osiris:
+            if ax == 0:
+                w = self._get_window(result.shape[0], self._window_time)
+                result = self._apply_window(result, w, axis=0)
+            else:
+                w = self._get_window(result.shape[ax], self._window_spatial)
+                result = self._apply_window(result, w, axis=ax)
+
+        # FFT + shift
+        with tqdm.tqdm(total=1, desc="FFT calculation") as pbar:
+            if len(axes_osiris) == 1:
+                ax = axes_osiris[0]
+                fft = np.fft.fft(result, axis=ax, norm=self._np_norm)
                 pbar.update(0.5)
-                result = np.fft.fftshift(data_fft, axes=self._fft_axis)
+                fft = np.fft.fftshift(fft, axes=ax)
                 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)
+                fft = np.fft.fftn(result, axes=axes_osiris, norm=self._np_norm)
                 pbar.update(0.5)
-                result = np.fft.fftshift(data_fft, axes=self._fft_axis)
+                fft = np.fft.fftshift(fft, axes=axes_osiris)
                 pbar.update(0.5)
 
-        self.omega_max = np.pi / self._dt / self._ndump
+        # Useful for time FFT
+        self.omega_max = np.pi / (self._dt * self._ndump)
+
+        # Optional "density" scaling (includes dt/dx factors)
+        if self._normalize == "density":
+            fft = fft * self._spacing_product(axes_osiris)
 
+        self._data = np.abs(fft) ** 2
         self._all_loaded = True
-        self._data = np.abs(result) ** 2
         return self._data
 
-    def _data_generator(self, index: int) -> Generator[np.ndarray, None, None]:
-        # Get the data for this index
-        original_data = self._diag[index]
+    def _frame(self, index: int, data_slice: tuple | None = None) -> np.ndarray:
+        """
+        Per-timestep (lazy) FFT. Only allowed for spatial FFT.
+        If fft_axis includes time (0), user must call load_all().
+        Returns power spectrum \|FFT\|^2.
+        """
+        axes_osiris = [self._fft_axis] if isinstance(self._fft_axis, int) else list(self._fft_axis)
+
+        # Time FFT cannot be computed from a single timestep
+        if 0 in axes_osiris:
+            raise ValueError("Cannot compute FFT along time axis for a single timestep. Use load_all() instead.")
+
+        # Read ONLY requested spatial slice from the underlying diagnostic
+        # data_slice refers to spatial dimensions only (no time dimension here)
+        f = self._diag._frame(index, data_slice=data_slice)
+
+        # Map OSIRIS axes (1,2,3) -> per-timestep numpy axes (0,1,2)
+        spatial_axes_osiris = [ax for ax in axes_osiris if ax != 0]
+        data_axes = [ax - 1 for ax in spatial_axes_osiris]  # 1->0, 2->1, 3->2
+
+        # Detrend over transform axes (e.g., remove mean)
+        f = self._detrend_data(f, data_axes)
+
+        # Choose spatial window; if user passed a slice and no window, you may want to force one
+        # (uncomment if you want forced windowing on cropped FFTs)
+        # spatial_window_kind = self._window_spatial
+        # if data_slice is not None and spatial_window_kind is None:
+        #     spatial_window_kind = "hann"
+        spatial_window_kind = self._window_spatial
+        if data_slice is not None and spatial_window_kind is None:
+            spatial_window_kind = "hann"
+
+        # Apply window(s) along transformed axes
+        result = f
+        for dax in data_axes:
+            w = self._get_window(result.shape[dax], spatial_window_kind)
+            result = self._apply_window(result, w, dax)
+
+        # FFT + shift
+        if len(data_axes) == 1:
+            fft = np.fft.fft(result, axis=data_axes[0], norm=self._np_norm)
+            fft = np.fft.fftshift(fft, axes=data_axes[0])
+        else:
+            fft = np.fft.fftn(result, axes=data_axes, norm=self._np_norm)
+            fft = np.fft.fftshift(fft, axes=data_axes)
 
-        if self._fft_axis == 0:
-            raise ValueError("Cannot generate FFT along time axis for a single timestep. Use load_all() instead.")
+        # Optional "density" scaling (includes dt/dx factors)
+        if self._normalize == "density":
+            fft = fft * self._spacing_product(spatial_axes_osiris)
 
-        # 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)
+        return np.abs(fft) ** 2
 
-            # 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])
+    def _get_window(self, length: int, kind: str | None) -> np.ndarray | None:
+        """
+        Get the window array of given kind and length.
 
-        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)
+        Parameters
+        ----------
+        length : int
+            Length of the window.
+        kind : str or None
+            Type of window ("hann", "hanning", or None).
 
-                result_fft = np.fft.fft(windowed_data, axis=self._fft_axis - 1)
-                result_fft = np.fft.fftshift(result_fft, axes=self._fft_axis - 1)
+        Returns
+        -------
+        np.ndarray or None
+            The window array, or None if no window is applied.
+        """
+        if kind is None or kind == "none":
+            return None
+        if kind in ("hann", "hanning"):
+            return np.hanning(length)
+        raise ValueError(f"Unknown window: {kind}")
 
-        yield np.abs(result_fft) ** 2
+    def _apply_window(self, data: np.ndarray, window: np.ndarray | None, axis: int) -> np.ndarray:
+        """
+        Apply the given window along the specified axis of the data.
 
-    def _get_window(self, length: int, axis: int) -> np.ndarray:
-        return np.hanning(length)
+        Parameters
+        ----------
+        data : np.ndarray
+            The input data array.
+        window : np.ndarray or None
+            The window array to apply, or None for no window.
+        axis : int
+            The axis along which to apply the window.
 
-    def _apply_window(self, data: np.ndarray, window: np.ndarray, axis: int) -> np.ndarray:
-        ndim = data.ndim
-        window_shape = [1] * ndim
+        Returns
+        -------
+        np.ndarray
+            The windowed data array.
+        """
+        if window is None:
+            return data
+        window_shape = [1] * data.ndim
         window_shape[axis] = len(window)
+        return data * window.reshape(window_shape)
 
-        reshaped_window = window.reshape(window_shape)
-
-        return data * reshaped_window
-
-    def __getitem__(self, index: int | slice) -> np.ndarray:
-        if self._all_loaded and self._data is not None:
-            return self._data[index]
+    def _detrend_data(self, data: np.ndarray, axes: list[int]) -> np.ndarray:
+        """
+        Detrend the data along the specified axes.
 
-        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.")
+        Parameters
+        ----------
+        data : np.ndarray
+            The input data array.
+        axes : list of int
+            The axes along which to detrend.
 
-    def omega(self) -> np.ndarray:
+        Returns
+        -------
+        np.ndarray
+            The detrended data array.
+        """
+        if self._detrend is None or self._detrend == "none":
+            return data
+        if self._detrend == "mean":
+            # remove mean over the transform axes (keeps other dims intact)
+            mean = data.mean(axis=tuple(axes), keepdims=True)
+            return data - mean
+        raise ValueError(f"Unknown detrend: {self._detrend}")
+
+    def _spacing_product(self, axes_osiris: list[int]) -> float:
         """
-        Get the angular frequency array for the FFT along the time dimension (axis 0).
+        Compute the product of spacing factors (dt, dx, dy, dz) for the given OSIRIS axes.
+
+        Parameters
+        ----------
+        axes_osiris : list of int
+            The OSIRIS axes along which the FFT was computed.
+            (0=time, 1=x, 2=y, 3=z)
 
         Returns
         -------
-        np.ndarray
-            Angular frequency array for the time axis.
+        float
+        The product of spacing factors.
+        """
+        # axes_osiris uses: time=0, space=1,2,3
+        prod = 1.0
+        for ax in axes_osiris:
+            if ax == 0:
+                prod *= float(self._dt * self._ndump)
+            else:
+                dx = self._dx if isinstance(self._dx, (int, float)) else self._dx[ax - 1]
+                prod *= float(dx)
+        return prod
+
+    def omega(self) -> np.ndarray:
+        """
+        Get the angular frequency array for the FFT along the time dimension.
         """
         if not self._all_loaded:
             raise ValueError("Load the data first using load_all() method.")
 
-        # 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)
+        axes_osiris = [self._fft_axis] if isinstance(self._fft_axis, int) else list(self._fft_axis)
+        if 0 not in axes_osiris:
+            raise ValueError(f"FFT was not computed along time axis (0). fft_axis={self._fft_axis}. Use k() instead.")
+
+        dt = self._dt * self._ndump
+        omega = np.fft.fftfreq(self._data.shape[0], d=dt) * 2 * np.pi
+        return np.fft.fftshift(omega)
 
     def k(self, axis: int | None = None) -> np.ndarray | dict[int, np.ndarray]:
         """
@@ -286,6 +398,14 @@ class FFT_Diagnostic(Diagnostic):
         if self._data is None:
             raise ValueError("Load the data first using load_all() or access via indexing.")
 
+        # Which OSIRIS axes did we transform?
+        axes_osiris = [self._fft_axis] if isinstance(self._fft_axis, int) else list(self._fft_axis)
+        spatial_fft_axes = {ax for ax in axes_osiris if ax != 0}  # {1,2,3}
+
+        # If user requests a specific k-axis, it must be one we transformed
+        if axis is not None and axis not in spatial_fft_axes:
+            raise ValueError(f"Requested k for axis {axis}, but FFT was not computed along that axis. fft_axis={self._fft_axis}")
+
         # 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
@@ -345,8 +465,14 @@ class FFT_Diagnostic(Diagnostic):
 
     @property
     def kmax(self) -> float | np.ndarray:
+        """Nyquist wavenumber (rad/m). Upper limit of representable k."""
         return self._kmax
 
+    @property
+    def omega_max(self):
+        """Nyquist angular frequency (rad/s). Upper limit of representable ω."""
+        return np.pi / (self._dt * self._ndump)
+
 
 class FFT_Species_Handler:
     def __init__(self, species_handler: Any, fft_axis: int | list[int]) -> None: