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

osiris_utils/postprocessing/fft.py

@@ -1,25 +1,26 @@
-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
 
+from ..data.diagnostic import Diagnostic
+from ..data.simulation import Simulation
+from .postprocess import PostProcess
+
 
-class FastFourierTransform_Simulation(PostProcess):
+class FFT_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):
@@ -34,11 +35,11 @@ class FastFourierTransform_Simulation(PostProcess):
             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 = {}
 
@@ -51,7 +52,7 @@ class FastFourierTransform_Simulation(PostProcess):
     def process(self, diagnostic):
         """Apply FFT to a diagnostic"""
         return FFT_Diagnostic(diagnostic, self._fft_axis)
-    
+
 
 class FFT_Diagnostic(Diagnostic):
     """
@@ -59,7 +60,7 @@ class FFT_Diagnostic(Diagnostic):
     Inherits directly from Diagnostic to ensure all operation overloads work properly.
 
     Parameters
-    ---------- 
+    ----------
     diagnostic : Diagnostic
         The diagnostic to compute the FFT.
     axis : int
@@ -75,47 +76,61 @@ class FFT_Diagnostic(Diagnostic):
         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)
+        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.postprocess_name = "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']:
+        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) 
+            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])
-    
+            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:
+
+        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
+        # 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)))
@@ -123,21 +138,21 @@ class FFT_Diagnostic(Diagnostic):
                 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: 
+            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)
@@ -147,39 +162,39 @@ class FFT_Diagnostic(Diagnostic):
         self.omega_max = np.pi / self._dt / self._ndump
 
         self._all_loaded = True
-        self._data = np.abs(result)**2
+        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)
-                    
+                    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])
-            
+            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
-        
+                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)
 
@@ -187,15 +202,15 @@ class FFT_Diagnostic(Diagnostic):
         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):
@@ -205,22 +220,23 @@ class FFT_Diagnostic(Diagnostic):
             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.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
@@ -231,4 +247,4 @@ class FFT_Species_Handler:
         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]
+        return self._fft_computed[key]

osiris_utils/postprocessing/field_centering.py

@@ -1,16 +1,18 @@
-from ..utils import *
+import numpy as np
+
+from ..data.diagnostic import OSIRIS_FLD, Diagnostic
 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.
+    Class to handle the field centering on data.It converts fields from the Osiris yee mesh to the center of the cells.
+    Works as a wrapper for the FieldCentering_Diagnostic class.
     Inherits from PostProcess to ensure all operation overloads work properly.
-    
+
+    It only works for periodic boundaries.
+
     Parameters
     ----------
     simulation : Simulation
@@ -18,9 +20,9 @@ class FieldCentering_Simulation(PostProcess):
     field : str
         The field to center.
     """
-    
+
     def __init__(self, simulation: Simulation):
-        super().__init__(f"FieldCentering Simulation")
+        super().__init__("FieldCentering Simulation")
         """
         Class to center the field in the simulation.
 
@@ -36,7 +38,7 @@ class FieldCentering_Simulation(PostProcess):
         self._simulation = simulation
 
         self._field_centered = {}
-        # no need to create a species handler for field centering since fields are not species related 
+        # 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:
@@ -44,7 +46,7 @@ class FieldCentering_Simulation(PostProcess):
         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 = {}
 
@@ -53,36 +55,49 @@ class FieldCentering_Simulation(PostProcess):
             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.
+        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.
 
         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)
+        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']:
+        for attr in [
+            "_dt",
+            "_dx",
+            "_ndump",
+            "_axis",
+            "_nx",
+            "_x",
+            "_grid",
+            "_dim",
+            "_maxiter",
+        ]:
             if hasattr(diagnostic, attr):
                 setattr(self, attr, getattr(diagnostic, attr))
 
@@ -95,41 +110,118 @@ class FieldCentering_Diagnostic(Diagnostic):
     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:
+
+        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 
+            if self._original_name.lower() in [
+                "b2",
+                "part_b2",
+                "ext_b2",
+                "b3",
+                "part_b3",
+                "ext_b3",
+                "e1",
+                "part_e1",
+                "ext_e1",
+            ]:
+                result = 0.5 * (np.roll(self._diag.data, shift=1, axis=1) + self._diag.data)
+            elif self._original_name.lower() in [
+                "b1",
+                "part_b1",
+                "ext_b1",
+                "e2",
+                "part_e2",
+                "ext_e2",
+                "e3",
+                "part_e3",
+                "ext_e3",
+            ]:
+                result = self._diag.data
 
         elif self._dim == 2:
-            if self._original_name.lower() in ['e1', 'b2']:
+            if self._original_name.lower() in [
+                "e1",
+                "part_e1",
+                "ext_e1",
+                "b2",
+                "part_b2",
+                "ext_b2",
+            ]:
                 result = 0.5 * (np.roll(self._diag.data, shift=1, axis=1) + self._diag.data)
-            elif self._original_name.lower() in ['e2', 'b1']:
+            elif self._original_name.lower() in [
+                "e2",
+                "part_e2",
+                "ext_e2",
+                "b1",
+                "part_b1",
+                "ext_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']:
+            elif self._original_name.lower() in ["b3", "part_b3", "ext_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", "part_e3", "ext_e3"]:
                 result = self._diag.data
-            
+
         elif self._dim == 3:
-            raise NotImplementedError("3D field centering is not implemented yet.")
-        
+            # TODO test this
+            if self._original_name in ["b1", "part_b1", "ext_b1"]:
+                result = 0.5 * (
+                    0.5
+                    * np.roll(
+                        (np.roll(self._diag.data, shift=1, axis=2) + self._diag.data),
+                        shift=1,
+                        axis=3,
+                    )
+                    + 0.5 * (np.roll(self._diag.data, shift=1, axis=2) + self._diag.data)
+                )
+            elif self._original_name in ["b2", "part_b2", "ext_b2"]:
+                result = 0.5 * (
+                    0.5
+                    * np.roll(
+                        (np.roll(self._diag.data, shift=1, axis=1) + self._diag.data),
+                        shift=1,
+                        axis=3,
+                    )
+                    + 0.5 * (np.roll(self._diag.data, shift=1, axis=1) + self._diag.data)
+                )
+            elif self._original_name in ["b3", "part_b3", "ext_b3"]:
+                result = 0.5 * (
+                    0.5
+                    * np.roll(
+                        (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 in ["e1", "part_e1", "ext_e1"]:
+                result = 0.5 * (np.roll(self._diag.data, shift=1, axis=1) + self._diag.data)
+            elif self._original_name in ["e2", "part_e2", "ext_e2"]:
+                result = 0.5 * (np.roll(self._diag.data, shift=1, axis=2) + self._diag.data)
+            elif self._original_name in ["e3", "part_e3", "ext_e3"]:
+                result = 0.5 * (np.roll(self._diag.data, shift=1, axis=3) + self._diag.data)
+
         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):
@@ -142,27 +234,103 @@ class FieldCentering_Diagnostic(Diagnostic):
 
     def _data_generator(self, index):
         if self._dim == 1:
-            if self._original_name.lower() in ['b2', 'b3', 'e1']:
+            if self._original_name.lower() in [
+                "b2",
+                "part_b2",
+                "ext_b2",
+                "b3",
+                "part_b3",
+                "ext_b3",
+                "e1",
+                "part_e1",
+                "ext_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 
+            elif self._original_name.lower() in [
+                "b1",
+                "part_b1",
+                "ext_b1",
+                "e2",
+                "part_e2",
+                "ext_e2",
+                "e3",
+                "part_e3",
+                "ext_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']:
+            if self._original_name in [
+                "e1",
+                "part_e1",
+                "ext_e1",
+                "b2",
+                "part_b2",
+                "ext_b2",
+            ]:
                 yield 0.5 * (np.roll(self._diag[index], shift=1, axis=0) + self._diag[index])
-            elif self._original_name in ['e2', 'b1']:
+            elif self._original_name in [
+                "e2",
+                "part_e2",
+                "ext_e2",
+                "b1",
+                "part_b1",
+                "ext_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']:
+            elif self._original_name in ["b3", "part_b3", "ext_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", "part_e3", "ext_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.")
-        
+            if self._original_name in ["b1", "part_b1", "ext_b1"]:
+                yield 0.5 * (
+                    0.5
+                    * np.roll(
+                        (np.roll(self._diag[index], shift=1, axis=1) + self._diag[index]),
+                        shift=1,
+                        axis=2,
+                    )
+                    + 0.5 * (np.roll(self._diag[index], shift=1, axis=1) + self._diag[index])
+                )
+            elif self._original_name in ["b2", "part_b2", "ext_b2"]:
+                yield 0.5 * (
+                    0.5
+                    * np.roll(
+                        (np.roll(self._diag[index], shift=1, axis=0) + self._diag[index]),
+                        shift=1,
+                        axis=2,
+                    )
+                    + 0.5 * (np.roll(self._diag[index], shift=1, axis=0) + self._diag[index])
+                )
+            elif self._original_name in ["b3", "part_b3", "ext_b3"]:
+                yield 0.5 * (
+                    0.5
+                    * np.roll(
+                        (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 ["e1", "part_e1", "ext_e1"]:
+                yield 0.5 * (np.roll(self._diag[index], shift=1, axis=0) + self._diag[index])
+            elif self._original_name in ["e2", "part_e2", "ext_e2"]:
+                yield 0.5 * (np.roll(self._diag[index], shift=1, axis=1) + self._diag[index])
+            elif self._original_name in ["e3", "part_e3", "ext_e3"]:
+                yield 0.5 * (np.roll(self._diag[index], shift=1, axis=2) + self._diag[index])
+
         else:
             raise ValueError(f"Unknown dimension {self._dim}.")