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

osiris_utils/data/diagnostic.py

@@ -4,10 +4,14 @@ of just a single file. This is what this file is for - deal with ''folders'' of
 
 Took some inspiration from Diogo and Madox's work.
 
-This would be awsome to compute time derivatives. 
+This would be awsome to compute time derivatives.
 """
+
+from torch import isin
 import numpy as np
 import os
+import glob
+import h5py
 
 from .data import OsirisGridFile
 import tqdm
@@ -16,15 +20,17 @@ import warnings
 from typing import Literal
 from ..decks.decks import InputDeckIO, deval
 
+
 def get_dimension_from_deck(deck: InputDeckIO) -> int:
     for dim in range(1, 4):
         try:
-            deck.get_param(section='grid', param=f'nx_p(1:{dim})')
+            deck.get_param(section="grid", param=f"nx_p(1:{dim})")
             return dim
         except:
             continue
-    
-    raise Exception('Error parsing grid dimension')
+
+    raise Exception("Error parsing grid dimension")
+
 
 OSIRIS_DENSITY = ["n"]
 OSIRIS_SPECIE_REPORTS = ["charge", "q1", "q2", "q3", "j1", "j2", "j3"]
@@ -48,9 +54,48 @@ OSIRIS_SPECIE_REP_UDIST = [
     "T23",
     "T33",
 ]
-OSIRIS_FLD = ["e1", "e2", "e3", "b1", "b2", "b3", "part_e1", "part_e2", "epart_3", "part_b1", "part_b2", "part_b3", "ext_e1", "ext_e2", "ext_e3", "ext_b1", "ext_b2", "ext_b3"]
-OSIRIS_PHA = ["p1x1", "p1x2", "p1x3", "p2x1", "p2x2", "p2x3", "p3x1", "p3x2", "p3x3", "gammax1", "gammax2", "gammax3"] # there may be more that I don't know
-OSIRIS_ALL = OSIRIS_DENSITY + OSIRIS_SPECIE_REPORTS + OSIRIS_SPECIE_REP_UDIST + OSIRIS_FLD + OSIRIS_PHA
+OSIRIS_FLD = [
+    "e1",
+    "e2",
+    "e3",
+    "b1",
+    "b2",
+    "b3",
+    "part_e1",
+    "part_e2",
+    "epart_3",
+    "part_b1",
+    "part_b2",
+    "part_b3",
+    "ext_e1",
+    "ext_e2",
+    "ext_e3",
+    "ext_b1",
+    "ext_b2",
+    "ext_b3",
+]
+OSIRIS_PHA = [
+    "p1x1",
+    "p1x2",
+    "p1x3",
+    "p2x1",
+    "p2x2",
+    "p2x3",
+    "p3x1",
+    "p3x2",
+    "p3x3",
+    "gammax1",
+    "gammax2",
+    "gammax3",
+]  # there may be more that I don't know
+OSIRIS_ALL = (
+    OSIRIS_DENSITY
+    + OSIRIS_SPECIE_REPORTS
+    + OSIRIS_SPECIE_REP_UDIST
+    + OSIRIS_FLD
+    + OSIRIS_PHA
+)
+
 
 def which_quantities():
     print("Available quantities:")
@@ -134,20 +179,9 @@ class Diagnostic:
         Plot a 3D scatter plot of the diagnostic data.
     time(index)
         Get the time for a given index.
-
-    Examples
-    --------
-    >>> sim = Simulation("electrons", "path/to/simulation")
-    >>> sim.get_quantity("charge")
-    >>> sim.load_all()
-    >>> print(sim.data.shape)
-    (100, 100, 100)
-
-    >>> sim = Simulation("electrons", "path/to/simulation")
-    >>> sim.get_quantity("charge")
-    >>> sim[0]
-    array with the data for the first timestep
+        
     """
+
     def __init__(self, simulation_folder=None, species=None, input_deck=None):
         self._species = species if species else None
 
@@ -164,11 +198,13 @@ class Diagnostic:
         self._ndump = None
         self._maxiter = None
         self._tunits = None
-        
+
         if simulation_folder:
             self._simulation_folder = simulation_folder
             if not os.path.isdir(simulation_folder):
-                raise FileNotFoundError(f"Simulation folder {simulation_folder} not found.")
+                raise FileNotFoundError(
+                    f"Simulation folder {simulation_folder} not found."
+                )
         else:
             self._simulation_folder = None
 
@@ -180,7 +216,7 @@ class Diagnostic:
 
         self._all_loaded = False
         self._quantity = None
-    
+
     def get_quantity(self, quantity):
         """
         Get the data for a given quantity.
@@ -193,7 +229,9 @@ class Diagnostic:
         self._quantity = quantity
 
         if self._quantity not in OSIRIS_ALL:
-            raise ValueError(f"Invalid quantity {self._quantity}. Use which_quantities() to see the available quantities.")
+            raise ValueError(
+                f"Invalid quantity {self._quantity}. Use which_quantities() to see the available quantities."
+            )
         if self._quantity in OSIRIS_SPECIE_REP_UDIST:
             if self._species is None:
                 raise ValueError("Species not set.")
@@ -213,45 +251,56 @@ class Diagnostic:
                 raise ValueError("Species not set.")
             self._get_density(self._species.name, "charge")
         else:
-            raise ValueError(f"Invalid quantity {self._quantity}. Or it's not implemented yet (this may happen for phase space quantities).")
+            raise ValueError(
+                f"Invalid quantity {self._quantity}. Or it's not implemented yet (this may happen for phase space quantities)."
+            )
 
     def _get_moment(self, species, moment):
         if self._simulation_folder is None:
-            raise ValueError("Simulation folder not set. If you're using CustomDiagnostic, this method is not available.")
+            raise ValueError(
+                "Simulation folder not set. If you're using CustomDiagnostic, this method is not available."
+            )
         self._path = f"{self._simulation_folder}/MS/UDIST/{species}/{moment}/"
-        self._file_template = os.listdir(self._path)[0][:-9]
-        self._maxiter = len(os.listdir(self._path))
+        self._file_template = glob.glob(f"{self._path}/*.h5")[0][:-9]
+        self._maxiter = len(glob.glob(f"{self._path}/*.h5"))
         self._load_attributes(self._file_template, self._input_deck)
-    
+
     def _get_field(self, field):
         if self._simulation_folder is None:
-            raise ValueError("Simulation folder not set. If you're using CustomDiagnostic, this method is not available.")
+            raise ValueError(
+                "Simulation folder not set. If you're using CustomDiagnostic, this method is not available."
+            )
         self._path = f"{self._simulation_folder}/MS/FLD/{field}/"
-        self._file_template = os.listdir(self._path)[0][:-9]
-        self._maxiter = len(os.listdir(self._path))
+        self._file_template = glob.glob(f"{self._path}/*.h5")[0][:-9]
+        self._maxiter = len(glob.glob(f"{self._path}/*.h5"))
         self._load_attributes(self._file_template, self._input_deck)
-        
+
     def _get_density(self, species, quantity):
         if self._simulation_folder is None:
-            raise ValueError("Simulation folder not set. If you're using CustomDiagnostic, this method is not available.")
+            raise ValueError(
+                "Simulation folder not set. If you're using CustomDiagnostic, this method is not available."
+            )
         self._path = f"{self._simulation_folder}/MS/DENSITY/{species}/{quantity}/"
-        self._file_template = os.listdir(self._path)[0][:-9]
-        self._maxiter = len(os.listdir(self._path))
+        self._file_template = glob.glob(f"{self._path}/*.h5")[0][:-9]
+        self._maxiter = len(glob.glob(f"{self._path}/*.h5"))
         self._load_attributes(self._file_template, self._input_deck)
 
     def _get_phase_space(self, species, type):
         if self._simulation_folder is None:
-            raise ValueError("Simulation folder not set. If you're using CustomDiagnostic, this method is not available.")
+            raise ValueError(
+                "Simulation folder not set. If you're using CustomDiagnostic, this method is not available."
+            )
         self._path = f"{self._simulation_folder}/MS/PHA/{type}/{species}/"
-        self._file_template = os.listdir(self._path)[0][:-9]
-        self._maxiter = len(os.listdir(self._path))
+        self._file_template = glob.glob(f"{self._path}/*.h5")[0][:-9]
+        self._maxiter = len(glob.glob(f"{self._path}/*.h5"))
         self._load_attributes(self._file_template, self._input_deck)
 
-    def _load_attributes(self, file_template, input_deck): # this will be replaced by reading the input deck
+    def _load_attributes(
+        self, file_template, input_deck
+    ):  # this will be replaced by reading the input deck
         # This can go wrong! NDUMP
         # if input_deck is not None:
         #     self._dt = float(input_deck["time_step"][0]["dt"])
-        #     self._ndump = int(input_deck["time_step"][0]["ndump"])
         #     self._dim = get_dimension_from_deck(input_deck)
         #     self._nx = np.array(list(map(int, input_deck["grid"][0][f"nx_p(1:{self._dim})"].split(','))))
         #     xmin = [deval(input_deck["space"][0][f"xmin(1:{self._dim})"].split(',')[i]) for i in range(self._dim)]
@@ -260,30 +309,46 @@ class Diagnostic:
         #     self._dx = (self._grid[:,1] - self._grid[:,0])/self._nx
         #     self._x = [np.arange(self._grid[i,0], self._grid[i,1], self._dx[i]) for i in range(self._dim)]
 
+        self._ndump = int(input_deck["time_step"][0]["ndump"])
+        
         try:
-            path_file1 = os.path.join(self._path, file_template + "000001.h5")
-            dump1 = OsirisGridFile(path_file1)
-            self._dx = dump1.dx
-            self._nx = dump1.nx
-            self._x = dump1.x
-            self._dt = dump1.dt
-            self._grid = dump1.grid
-            self._axis = dump1.axis
-            self._units = dump1.units
-            self._name = dump1.name
-            self._label = dump1.label
-            self._dim = dump1.dim
-            self._ndump = dump1.iter
-            self._tunits = dump1.time[1]
-        except:
-            pass
-    
+            # Try files 000001, 000002, etc. until one is found
+            found_file = False
+            for file_num in range(1, self._maxiter + 1):
+                path_file = os.path.join(file_template + f"{file_num:06d}.h5")
+                if os.path.exists(path_file):
+                    dump = OsirisGridFile(path_file)
+                    self._dx = dump.dx
+                    self._nx = dump.nx
+                    self._x = dump.x
+                    self._dt = dump.dt
+                    self._grid = dump.grid
+                    self._axis = dump.axis
+                    self._units = dump.units
+                    self._name = dump.name
+                    self._label = dump.label
+                    self._dim = dump.dim
+                    # self._iter = dump.iter
+                    self._tunits = dump.time[1]
+                    self._type = dump.type
+                    found_file = True
+                    break
+            
+            if not found_file:
+                warnings.warn(f"No valid data files found in {self._path} to read metadata from.")
+        except Exception as e:
+            warnings.warn(f"Error loading diagnostic attributes: {str(e)}. Please verify it there's any file in the folder.")
+
     def _data_generator(self, index):
         if self._simulation_folder is None:
             raise ValueError("Simulation folder not set.")
-        file = os.path.join(self._path, self._file_template + f"{index:06d}.h5")
+        file = os.path.join(self._file_template + f"{index:06d}.h5")
         data_object = OsirisGridFile(file)
-        yield data_object.data if self._quantity not in OSIRIS_DENSITY else self._species.rqm * data_object.data
+        yield (
+            data_object.data
+            if self._quantity not in OSIRIS_DENSITY
+            else self._species.rqm * data_object.data
+        )
 
     def load_all(self):
         """
@@ -298,39 +363,47 @@ class Diagnostic:
         if self._all_loaded and self._data is not None:
             print("Data already loaded.")
             return self._data
-        
+
         # If this is a derived diagnostic without files
-        if self._simulation_folder is None:
+        if hasattr(self, "postprocess_name") or hasattr(self, "created_diagnostic_name"):
             # If it has a data generator but no direct files
             try:
-                print("This appears to be a derived diagnostic. Loading data from generators...")
+                print(
+                    "This appears to be a derived diagnostic. Loading data from generators..."
+                )
                 # Get the maximum size from the diagnostic attributes
-                if hasattr(self, '_maxiter') and self._maxiter is not None:
+                if hasattr(self, "_maxiter") and self._maxiter is not None:
                     size = self._maxiter
                 else:
                     # Try to infer from a related diagnostic
-                    if hasattr(self, '_diag') and hasattr(self._diag, '_maxiter'):
+                    if hasattr(self, "_diag") and hasattr(self._diag, "_maxiter"):
                         size = self._diag._maxiter
                     else:
                         # Default to a reasonable number if we can't determine
                         size = 100
-                        print(f"Warning: Could not determine timestep count, using {size}.")
-                
+                        print(
+                            f"Warning: Could not determine timestep count, using {size}."
+                        )
+
                 # Load data for all timesteps using the generator - this may take a while
-                self._data = np.stack([self[i] for i in tqdm.tqdm(range(size), desc="Loading data")])
+                self._data = np.stack(
+                    [self[i] for i in tqdm.tqdm(range(size), desc="Loading data")]
+                )
                 self._all_loaded = True
                 return self._data
-            
+
             except Exception as e:
                 raise ValueError(f"Could not load derived diagnostic data: {str(e)}")
-        
+
         # Original implementation for file-based diagnostics
         print("Loading all data from files. This may take a while.")
-        size = len(sorted(os.listdir(self._path)))
-        self._data = np.stack([self[i] for i in tqdm.tqdm(range(size), desc="Loading data")])
+        size = len(sorted(glob.glob(f"{self._path}/*.h5")))
+        self._data = np.stack(
+            [self[i] for i in tqdm.tqdm(range(size), desc="Loading data")]
+        )
         self._all_loaded = True
         return self._data
-    
+
     def unload(self):
         """
         Unload data from memory. This is useful to free memory when the data is not needed anymore.
@@ -352,108 +425,147 @@ class Diagnostic:
         # For derived diagnostics with cached data
         if self._all_loaded and self._data is not None:
             return self._data[index]
-        
+
         # For standard diagnostics with files
         if isinstance(index, int):
-            if self._simulation_folder is not None and hasattr(self, '_data_generator'):
+            if self._simulation_folder is not None and hasattr(self, "_data_generator"):
                 return next(self._data_generator(index))
-            
+
             # For derived diagnostics with custom generators
-            if hasattr(self, '_data_generator') and callable(self._data_generator):
+            if hasattr(self, "_data_generator") and callable(self._data_generator):
                 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
 
             if index.stop is None:
-                if hasattr(self, '_maxiter') and self._maxiter is not None:
+                if hasattr(self, "_maxiter") and self._maxiter is not None:
                     stop = self._maxiter
-                elif self._simulation_folder is not None and hasattr(self, '_path'):
-                    stop = len(sorted(os.listdir(self._path)))
+                elif self._simulation_folder is not None and hasattr(self, "_path"):
+                    stop = len(sorted(glob.glob(f"{self._path}/*.h5")))
                 else:
                     stop = 100  # Default if we can't determine
-                    print(f"Warning: Could not determine iteration count for iteration, using {stop}.")
+                    print(
+                        f"Warning: Could not determine iteration count for iteration, using {stop}."
+                    )
             else:
                 stop = index.stop
 
             indices = range(start, stop, step)
-            if self._simulation_folder is not None and hasattr(self, '_data_generator'):
+            if self._simulation_folder is not None and hasattr(self, "_data_generator"):
                 return np.stack([next(self._data_generator(i)) for i in indices])
-            elif hasattr(self, '_data_generator') and callable(self._data_generator):
+            elif hasattr(self, "_data_generator") and callable(self._data_generator):
                 return np.stack([next(self._data_generator(i)) for i in indices])
 
         # If we get here, we don't know how to get data for this index
-        raise ValueError(f"Cannot retrieve data for this diagnostic at index {index}. No data loaded and no generator available.")   
-    
+        raise ValueError(
+            f"Cannot retrieve data for this diagnostic at index {index}. No data loaded and no generator available."
+        )
+
     def __iter__(self):
         # If this is a file-based diagnostic
         if self._simulation_folder is not None:
-            for i in range(len(sorted(os.listdir(self._path)))):
+            for i in range(len(sorted(glob.glob(f"{self._path}/*.h5")))):
                 yield next(self._data_generator(i))
-        
+
         # If this is a derived diagnostic and data is already loaded
         elif self._all_loaded and self._data is not None:
             for i in range(self._data.shape[0]):
                 yield self._data[i]
-        
+
         # If this is a derived diagnostic with custom generator but no loaded data
-        elif hasattr(self, '_data_generator') and callable(self._data_generator):
+        elif hasattr(self, "_data_generator") and callable(self._data_generator):
             # Determine how many iterations to go through
             max_iter = self._maxiter
             if max_iter is None:
-                if hasattr(self, '_diag') and hasattr(self._diag, '_maxiter'):
+                if hasattr(self, "_diag") and hasattr(self._diag, "_maxiter"):
                     max_iter = self._diag._maxiter
                 else:
                     max_iter = 100  # Default if we can't determine
-                    print(f"Warning: Could not determine iteration count for iteration, using {max_iter}.")
-            
+                    print(
+                        f"Warning: Could not determine iteration count for iteration, using {max_iter}."
+                    )
+
             for i in range(max_iter):
                 yield next(self._data_generator(i))
-        
+
         # If we don't know how to handle this
         else:
-            raise ValueError("Cannot iterate over this diagnostic. No data loaded and no generator available.")
+            raise ValueError(
+                "Cannot iterate over this diagnostic. No data loaded and no generator available."
+            )
 
     def __add__(self, other):
         if isinstance(other, (int, float, np.ndarray)):
             result = Diagnostic(species=self._species)
-        
-            for attr in ['_dx', '_nx', '_x', '_dt', '_grid', '_axis', '_dim', '_ndump', '_maxiter']:
+
+            for attr in [
+                "_dx",
+                "_nx",
+                "_x",
+                "_dt",
+                "_grid",
+                "_axis",
+                "_dim",
+                "_ndump",
+                "_maxiter",
+                "_tunits",
+                "_type",
+                "_simulation_folder",
+            ]:
                 if hasattr(self, attr):
                     setattr(result, attr, getattr(self, attr))
-        
+
             # Make sure _maxiter is set even for derived diagnostics
-            if not hasattr(result, '_maxiter') or result._maxiter is None:
-                if hasattr(self, '_maxiter') and self._maxiter is not None:
+            if not hasattr(result, "_maxiter") or result._maxiter is None:
+                if hasattr(self, "_maxiter") and self._maxiter is not None:
                     result._maxiter = self._maxiter
 
             # result._name = self._name + " + " + str(other) if isinstance(other, (int, float)) else self._name + " + np.ndarray"
-            
+
             if self._all_loaded:
                 result._data = self._data + other
                 result._all_loaded = True
             else:
+
                 def gen_scalar_add(original_gen, scalar):
                     for val in original_gen:
                         yield val + scalar
-                
+
                 original_generator = self._data_generator
-                result._data_generator = lambda index: gen_scalar_add(original_generator(index), other)
-                
+                result._data_generator = lambda index: gen_scalar_add(
+                    original_generator(index), other
+                )
+
+            result.created_diagnostic_name = "MISC"
+    
             return result
 
         elif isinstance(other, Diagnostic):
             result = Diagnostic(species=self._species)
 
-            for attr in ['_dx', '_nx', '_x', '_dt', '_grid', '_axis', '_dim', '_ndump', '_maxiter']:
+            for attr in [
+                "_dx",
+                "_nx",
+                "_x",
+                "_dt",
+                "_grid",
+                "_axis",
+                "_dim",
+                "_ndump",
+                "_maxiter",
+                "_tunits",
+                "_type",
+                "_simulation_folder",
+            ]:
                 if hasattr(self, attr):
                     setattr(result, attr, getattr(self, attr))
-        
-            if not hasattr(result, '_maxiter') or result._maxiter is None:
-                if hasattr(self, '_maxiter') and self._maxiter is not None:
+
+            if not hasattr(result, "_maxiter") or result._maxiter is None:
+                if hasattr(self, "_maxiter") and self._maxiter is not None:
                     result._maxiter = self._maxiter
-            
+
             # result._name = self._name + " + " + str(other._name)
 
             if self._all_loaded:
@@ -461,26 +573,44 @@ class Diagnostic:
                 result._data = self._data + other._data
                 result._all_loaded = True
             else:
+
                 def gen_diag_add(original_gen1, original_gen2):
                     for val1, val2 in zip(original_gen1, original_gen2):
                         yield val1 + val2
-                
+
                 original_generator = self._data_generator
                 other_generator = other._data_generator
-                result._data_generator = lambda index: gen_diag_add(original_generator(index), other_generator(index))
+                result._data_generator = lambda index: gen_diag_add(
+                    original_generator(index), other_generator(index)
+                )
+
+            result.created_diagnostic_name = "MISC"
 
             return result
-    
+
     def __sub__(self, other):
         if isinstance(other, (int, float, np.ndarray)):
             result = Diagnostic(species=self._species)
-        
-            for attr in ['_dx', '_nx', '_x', '_dt', '_grid', '_axis', '_dim', '_ndump', '_maxiter']:
+
+            for attr in [
+                "_dx",
+                "_nx",
+                "_x",
+                "_dt",
+                "_grid",
+                "_axis",
+                "_dim",
+                "_ndump",
+                "_maxiter",
+                "_tunits",
+                "_type",
+                "_simulation_folder",
+            ]:
                 if hasattr(self, attr):
                     setattr(result, attr, getattr(self, attr))
-        
-            if not hasattr(result, '_maxiter') or result._maxiter is None:
-                if hasattr(self, '_maxiter') and self._maxiter is not None:
+
+            if not hasattr(result, "_maxiter") or result._maxiter is None:
+                if hasattr(self, "_maxiter") and self._maxiter is not None:
                     result._maxiter = self._maxiter
 
             # result._name = self._name + " - " + str(other) if isinstance(other, (int, float)) else self._name + " - np.ndarray"
@@ -489,28 +619,45 @@ class Diagnostic:
                 result._data = self._data - other
                 result._all_loaded = True
             else:
+
                 def gen_scalar_sub(original_gen, scalar):
                     for val in original_gen:
                         yield val - scalar
-                
+
                 original_generator = self._data_generator
-                result._data_generator = lambda index: gen_scalar_sub(original_generator(index), other)
-                
+                result._data_generator = lambda index: gen_scalar_sub(
+                    original_generator(index), other
+                )
+
+            result.created_diagnostic_name = "MISC"
+
             return result
 
         elif isinstance(other, Diagnostic):
-                
-            
+
             result = Diagnostic(species=self._species)
 
-            for attr in ['_dx', '_nx', '_x', '_dt', '_grid', '_axis', '_dim', '_ndump', '_maxiter']:
+            for attr in [
+                "_dx",
+                "_nx",
+                "_x",
+                "_dt",
+                "_grid",
+                "_axis",
+                "_dim",
+                "_ndump",
+                "_maxiter",
+                "_tunits",
+                "_type",
+                "_simulation_folder",
+            ]:
                 if hasattr(self, attr):
                     setattr(result, attr, getattr(self, attr))
-        
-            if not hasattr(result, '_maxiter') or result._maxiter is None:
-                if hasattr(self, '_maxiter') and self._maxiter is not None:
+
+            if not hasattr(result, "_maxiter") or result._maxiter is None:
+                if hasattr(self, "_maxiter") and self._maxiter is not None:
                     result._maxiter = self._maxiter
-            
+
             # result._name = self._name + " - " + str(other._name)
 
             if self._all_loaded:
@@ -518,54 +665,90 @@ class Diagnostic:
                 result._data = self._data - other._data
                 result._all_loaded = True
             else:
+
                 def gen_diag_sub(original_gen1, original_gen2):
                     for val1, val2 in zip(original_gen1, original_gen2):
                         yield val1 - val2
-                
+
                 original_generator = self._data_generator
                 other_generator = other._data_generator
-                result._data_generator = lambda index: gen_diag_sub(original_generator(index), other_generator(index))
+                result._data_generator = lambda index: gen_diag_sub(
+                    original_generator(index), other_generator(index)
+                )
+
+            result.created_diagnostic_name = "MISC"
 
             return result
-    
+
     def __mul__(self, other):
         if isinstance(other, (int, float, np.ndarray)):
             result = Diagnostic(species=self._species)
-        
-            for attr in ['_dx', '_nx', '_x', '_dt', '_grid', '_axis', '_dim', '_ndump', '_maxiter']:
+
+            for attr in [
+                "_dx",
+                "_nx",
+                "_x",
+                "_dt",
+                "_grid",
+                "_axis",
+                "_dim",
+                "_ndump",
+                "_maxiter",
+                "_tunits",
+                "_type",
+                "_simulation_folder",
+            ]:
                 if hasattr(self, attr):
                     setattr(result, attr, getattr(self, attr))
-        
-            if not hasattr(result, '_maxiter') or result._maxiter is None:
-                if hasattr(self, '_maxiter') and self._maxiter is not None:
+
+            if not hasattr(result, "_maxiter") or result._maxiter is None:
+                if hasattr(self, "_maxiter") and self._maxiter is not None:
                     result._maxiter = self._maxiter
 
             # result._name = self._name + " * " + str(other) if isinstance(other, (int, float)) else self._name + " * np.ndarray"
-            
+
             if self._all_loaded:
                 result._data = self._data * other
                 result._all_loaded = True
             else:
+
                 def gen_scalar_mul(original_gen, scalar):
                     for val in original_gen:
                         yield val * scalar
-                
+
                 original_generator = self._data_generator
-                result._data_generator = lambda index: gen_scalar_mul(original_generator(index), other)
+                result._data_generator = lambda index: gen_scalar_mul(
+                    original_generator(index), other
+                )
                 
+            result.created_diagnostic_name = "MISC"
+
             return result
 
         elif isinstance(other, Diagnostic):
             result = Diagnostic(species=self._species)
 
-            for attr in ['_dx', '_nx', '_x', '_dt', '_grid', '_axis', '_dim', '_ndump', '_maxiter']:
+            for attr in [
+                "_dx",
+                "_nx",
+                "_x",
+                "_dt",
+                "_grid",
+                "_axis",
+                "_dim",
+                "_ndump",
+                "_maxiter",
+                "_tunits",
+                "_type",
+                "_simulation_folder",
+            ]:
                 if hasattr(self, attr):
                     setattr(result, attr, getattr(self, attr))
-        
-            if not hasattr(result, '_maxiter') or result._maxiter is None:
-                if hasattr(self, '_maxiter') and self._maxiter is not None:
+
+            if not hasattr(result, "_maxiter") or result._maxiter is None:
+                if hasattr(self, "_maxiter") and self._maxiter is not None:
                     result._maxiter = self._maxiter
-            
+
             # result._name = self._name + " * " + str(other._name)
 
             if self._all_loaded:
@@ -573,55 +756,91 @@ class Diagnostic:
                 result._data = self._data * other._data
                 result._all_loaded = True
             else:
+
                 def gen_diag_mul(original_gen1, original_gen2):
                     for val1, val2 in zip(original_gen1, original_gen2):
                         yield val1 * val2
-                
+
                 original_generator = self._data_generator
                 other_generator = other._data_generator
-                result._data_generator = lambda index: gen_diag_mul(original_generator(index), other_generator(index))
+                result._data_generator = lambda index: gen_diag_mul(
+                    original_generator(index), other_generator(index)
+                )
+                
+            result.created_diagnostic_name = "MISC"
 
             return result
-    
+
     def __truediv__(self, other):
         if isinstance(other, (int, float, np.ndarray)):
             result = Diagnostic(species=self._species)
-        
-            for attr in ['_dx', '_nx', '_x', '_dt', '_grid', '_axis', '_dim', '_ndump', '_maxiter']:
+
+            for attr in [
+                "_dx",
+                "_nx",
+                "_x",
+                "_dt",
+                "_grid",
+                "_axis",
+                "_dim",
+                "_ndump",
+                "_maxiter",
+                "_tunits",
+                "_type",
+                "_simulation_folder",
+            ]:
                 if hasattr(self, attr):
                     setattr(result, attr, getattr(self, attr))
-        
-            if not hasattr(result, '_maxiter') or result._maxiter is None:
-                if hasattr(self, '_maxiter') and self._maxiter is not None:
+
+            if not hasattr(result, "_maxiter") or result._maxiter is None:
+                if hasattr(self, "_maxiter") and self._maxiter is not None:
                     result._maxiter = self._maxiter
 
             # result._name = self._name + " / " + str(other) if isinstance(other, (int, float)) else self._name + " / np.ndarray"
-            
+
             if self._all_loaded:
                 result._data = self._data / other
                 result._all_loaded = True
             else:
+
                 def gen_scalar_div(original_gen, scalar):
                     for val in original_gen:
                         yield val / scalar
-                
+
                 original_generator = self._data_generator
-                result._data_generator = lambda index: gen_scalar_div(original_generator(index), other)
+                result._data_generator = lambda index: gen_scalar_div(
+                    original_generator(index), other
+                )
                 
+            result.created_diagnostic_name = "MISC"
+
             return result
 
         elif isinstance(other, Diagnostic):
-                
+
             result = Diagnostic(species=self._species)
 
-            for attr in ['_dx', '_nx', '_x', '_dt', '_grid', '_axis', '_dim', '_ndump', '_maxiter']:
+            for attr in [
+                "_dx",
+                "_nx",
+                "_x",
+                "_dt",
+                "_grid",
+                "_axis",
+                "_dim",
+                "_ndump",
+                "_maxiter",
+                "_tunits",
+                "_type",
+                "_simulation_folder",
+            ]:
                 if hasattr(self, attr):
                     setattr(result, attr, getattr(self, attr))
-        
-            if not hasattr(result, '_maxiter') or result._maxiter is None:
-                if hasattr(self, '_maxiter') and self._maxiter is not None:
+
+            if not hasattr(result, "_maxiter") or result._maxiter is None:
+                if hasattr(self, "_maxiter") and self._maxiter is not None:
                     result._maxiter = self._maxiter
-            
+
             # result._name = self._name + " / " + str(other._name)
 
             if self._all_loaded:
@@ -629,115 +848,293 @@ class Diagnostic:
                 result._data = self._data / other._data
                 result._all_loaded = True
             else:
+
                 def gen_diag_div(original_gen1, original_gen2):
                     for val1, val2 in zip(original_gen1, original_gen2):
                         yield val1 / val2
-                
+
                 original_generator = self._data_generator
                 other_generator = other._data_generator
-                result._data_generator = lambda index: gen_diag_div(original_generator(index), other_generator(index))
+                result._data_generator = lambda index: gen_diag_div(
+                    original_generator(index), other_generator(index)
+                )
+                
+            result.created_diagnostic_name = "MISC"
 
             return result
-        
+
     def __pow__(self, other):
-       # power by scalar
+        # power by scalar
         if isinstance(other, (int, float)):
             result = Diagnostic(species=self._species)
 
-            for attr in ['_dx', '_nx', '_x', '_dt', '_grid', '_axis', '_dim', '_ndump', '_maxiter']:
+            for attr in [
+                "_dx",
+                "_nx",
+                "_x",
+                "_dt",
+                "_grid",
+                "_axis",
+                "_dim",
+                "_ndump",
+                "_maxiter",
+                "_tunits",
+                "_type",
+                "_simulation_folder",
+            ]:
                 if hasattr(self, attr):
                     setattr(result, attr, getattr(self, attr))
 
-            if not hasattr(result, '_maxiter') or result._maxiter is None:
-                if hasattr(self, '_maxiter') and self._maxiter is not None:
+            if not hasattr(result, "_maxiter") or result._maxiter is None:
+                if hasattr(self, "_maxiter") and self._maxiter is not None:
                     result._maxiter = self._maxiter
 
             # result._name = self._name + " ^(" + str(other) + ")"
             # result._label = self._label + rf"$ ^{other}$"
 
             if self._all_loaded:
-                result._data = self._data ** other
+                result._data = self._data**other
                 result._all_loaded = True
             else:
+
                 def gen_scalar_pow(original_gen, scalar):
                     for val in original_gen:
-                        yield val ** scalar
+                        yield val**scalar
 
                 original_generator = self._data_generator
-                result._data_generator = lambda index: gen_scalar_pow(original_generator(index), other)
+                result._data_generator = lambda index: gen_scalar_pow(
+                    original_generator(index), other
+                )
+                
+            result.created_diagnostic_name = "MISC"
 
             return result
-        
+
         # power by another diagnostic
         elif isinstance(other, Diagnostic):
-            raise ValueError("Power by another diagnostic is not supported. Why would you do that?")
+            raise ValueError(
+                "Power by another diagnostic is not supported. Why would you do that?"
+            )
 
     def __radd__(self, other):
         return self + other
-    
-    def __rsub__(self, other): # I don't know if this is correct because I'm not sure if the order of the subtraction is correct
-        return - self + other
-    
+
+    def __rsub__(
+        self, other
+    ):  # I don't know if this is correct because I'm not sure if the order of the subtraction is correct
+        return -self + other
+
     def __rmul__(self, other):
         return self * other
-    
-    def __rtruediv__(self, other): # division is not commutative
+
+    def __rtruediv__(self, other):  # division is not commutative
         if isinstance(other, (int, float, np.ndarray)):
             result = Diagnostic(species=self._species)
-        
-            for attr in ['_dx', '_nx', '_x', '_dt', '_grid', '_axis', '_dim', '_ndump', '_maxiter']:
+
+            for attr in [
+                "_dx",
+                "_nx",
+                "_x",
+                "_dt",
+                "_grid",
+                "_axis",
+                "_dim",
+                "_ndump",
+                "_maxiter",
+                "_tunits",
+                "_type",
+                "_simulation_folder",
+            ]:
                 if hasattr(self, attr):
                     setattr(result, attr, getattr(self, attr))
-        
-            if not hasattr(result, '_maxiter') or result._maxiter is None:
-                if hasattr(self, '_maxiter') and self._maxiter is not None:
+
+            if not hasattr(result, "_maxiter") or result._maxiter is None:
+                if hasattr(self, "_maxiter") and self._maxiter is not None:
                     result._maxiter = self._maxiter
 
             # result._name = str(other) + " / " + self._name if isinstance(other, (int, float)) else "np.ndarray / " + self._name
-            
+
             if self._all_loaded:
                 result._data = other / self._data
                 result._all_loaded = True
             else:
+
                 def gen_scalar_rdiv(scalar, original_gen):
                     for val in original_gen:
                         yield scalar / val
-                
+
                 original_generator = self._data_generator
-                result._data_generator = lambda index: gen_scalar_rdiv(other, original_generator(index))
+                result._data_generator = lambda index: gen_scalar_rdiv(
+                    other, original_generator(index)
+                )
                 
+            result.created_diagnostic_name = "MISC"
+
             return result
-        
+
         elif isinstance(other, Diagnostic):
-            
+
             result = Diagnostic(species=self._species)
 
-            for attr in ['_dx', '_nx', '_x', '_dt', '_grid', '_axis', '_dim', '_ndump', '_maxiter']:
+            for attr in [
+                "_dx",
+                "_nx",
+                "_x",
+                "_dt",
+                "_grid",
+                "_axis",
+                "_dim",
+                "_ndump",
+                "_maxiter",
+                "_tunits",
+                "_type",
+                "_simulation_folder",
+            ]:
                 if hasattr(self, attr):
                     setattr(result, attr, getattr(self, attr))
-        
-            if not hasattr(result, '_maxiter') or result._maxiter is None:
-                if hasattr(self, '_maxiter') and self._maxiter is not None:
+
+            if not hasattr(result, "_maxiter") or result._maxiter is None:
+                if hasattr(self, "_maxiter") and self._maxiter is not None:
                     result._maxiter = self._maxiter
-            
+
             # result._name =  str(other._name) + " / " + self._name
 
             if self._all_loaded:
                 other.load_all()
-                result._data =  other._data / self._data
+                result._data = other._data / self._data
                 result._all_loaded = True
             else:
+
                 def gen_diag_div(original_gen1, original_gen2):
                     for val1, val2 in zip(original_gen1, original_gen2):
-                        yield  val2 / val1
-                
+                        yield val2 / val1
+
                 original_generator = self._data_generator
                 other_generator = other._data_generator
-                result._data_generator = lambda index: gen_diag_div(original_generator(index), other_generator(index))
+                result._data_generator = lambda index: gen_diag_div(
+                    original_generator(index), other_generator(index)
+                )
+                
+            result.created_diagnostic_name = "MISC"
 
             return result
 
-    def plot_3d(self, idx, scale_type: Literal["zero_centered", "pos", "neg", "default"] = "default", boundaries: np.ndarray = None):
+    def to_h5(self, savename=None, index=None, all=False, verbose=False, path=None):
+        """
+        Save the diagnostic data to HDF5 files.
+
+        Parameters
+        ----------
+        savename : str, optional
+            The name of the HDF5 file. If None, uses the diagnostic name.
+        index : int, or list of ints, optional
+            The index or indices of the data to save.
+        all : bool, optional
+            If True, save all data. Default is False.
+        verbose : bool, optional
+            If True, print messages about the saving process.
+        path : str, optional
+            The path to save the HDF5 files. If None, uses the default save path (in simulation folder).
+        """
+        if folder is None:
+            folder = self._simulation_folder
+            self._save_path = folder + f"/MS/MISC/{self._default_save}/{savename}"
+        else:
+            self._save_path = folder
+        # Check if is has attribute created_diagnostic_name or postprocess_name
+        if savename is None:
+            print(f"No savename provided. Using {self._name}.")
+            savename = self._name
+        
+        if hasattr(self, "created_diagnostic_name"):
+            self._default_save = self.created_diagnostic_name
+        elif hasattr(self, "postprocess_name"):
+            self._default_save = self.postprocess_name
+        else:
+            self._default_save = "DIR_" + self._name
+
+        if not os.path.exists(self._save_path):
+            os.makedirs(self._save_path)
+            if verbose:
+                print(f"Created folder {self._save_path}")
+        
+        if verbose:
+            print(f"Save Path: {self._save_path}")
+
+        def savefile(filename, i):
+            with h5py.File(filename, 'w') as f:
+                # Create SIMULATION group with attributes
+                sim_group = f.create_group("SIMULATION")
+                sim_group.attrs.create("DT", [self._dt])
+                sim_group.attrs.create("NDIMS", [self._dim])
+                
+                # Set file attributes
+                f.attrs.create("TIME", [self.time(i)[0]])
+                f.attrs.create("TIME UNITS", [np.bytes_(self.time(i)[1].encode()) if self.time(i)[1] else np.bytes_(b"")])
+                f.attrs.create("ITER", [self._ndump * i])
+                f.attrs.create("NAME", [np.bytes_(self._name.encode())]) 
+                f.attrs.create("TYPE", [np.bytes_(self._type.encode())])
+                f.attrs.create("UNITS", [np.bytes_(self._units.encode()) if self._units else np.bytes_(b"")])
+                f.attrs.create("LABEL", [np.bytes_(self._label.encode()) if self._label else np.bytes_(b"")])
+                
+                # Create dataset with data (transposed to match convention)
+                f.create_dataset(savename, data=self[i].T)
+                
+                # Create AXIS group
+                axis_group = f.create_group("AXIS")
+                
+                # Create axis datasets
+                axis_names = ["AXIS1", "AXIS2", "AXIS3"][:self._dim]
+                axis_shortnames = [self._axis[i]["name"] for i in range(self._dim)]
+                axis_longnames = [self._axis[i]["long_name"] for i in range(self._dim)]
+                axis_units = [self._axis[i]["units"] for i in range(self._dim)]
+                
+                for i, axis_name in enumerate(axis_names):
+                    # Create axis dataset
+                    axis_dataset = axis_group.create_dataset(axis_name, data=np.array(self._grid[i]))
+                    
+                    # Set axis attributes
+                    axis_dataset.attrs.create("NAME", [np.bytes_(axis_shortnames[i].encode())])
+                    axis_dataset.attrs.create("UNITS", [np.bytes_(axis_units[i].encode())])
+                    axis_dataset.attrs.create("LONG_NAME", [np.bytes_(axis_longnames[i].encode())])
+                    axis_dataset.attrs.create("TYPE", [np.bytes_("linear".encode())])
+                
+                if verbose:
+                    print(f"File created: {filename}")
+
+        print(f"The savename of the diagnostic is {savename}. Files will be saves as {savename}-000001.h5, {savename}-000002.h5, etc.")
+        
+        print(f"If you desire a different name, please set it with the 'name' method (setter).")
+
+        if self._name is None:
+            raise ValueError("Diagnostic name is not set. Cannot save to HDF5.")
+        if not os.path.exists(folder):
+            print(f"Creating folder {folder}...")
+            os.makedirs(folder)
+        if not os.path.isdir(folder):
+            raise ValueError(f"{folder} is not a directory.")
+
+        if all == False:
+            if isinstance(index, int):
+                filename = self._save_path + f"/{savename}-{index:06d}.h5"
+                savefile(filename, index)
+            elif isinstance(index, list) or isinstance(index, tuple):
+                for i in index:
+                    filename = self._save_path + f"/{savename}-{i:06d}.h5"
+                    savefile(filename, i)
+        elif all == True:
+            for i in range(self._maxiter):
+                filename = self._save_path + f"/{savename}-{i:06d}.h5"
+                savefile(filename, i)
+        else:
+            raise ValueError("index should be an int, slice, or list of ints, or all should be True")
+                    
+    def plot_3d(
+        self,
+        idx,
+        scale_type: Literal["zero_centered", "pos", "neg", "default"] = "default",
+        boundaries: np.ndarray = None,
+    ):
         """
         Plots a 3D scatter plot of the diagnostic data (grid data).
 
@@ -768,23 +1165,24 @@ class Diagnostic:
         plt.show()
         """
 
-
         if self._dim != 3:
             raise ValueError("This method is only available for 3D diagnostics.")
-        
+
         if boundaries is None:
             boundaries = self._grid
 
         if not isinstance(boundaries, np.ndarray):
-            try :
+            try:
                 boundaries = np.array(boundaries)
             except:
-                boundaries = self._grid 
-                warnings.warn("boundaries cannot be accessed as a numpy array with shape (3, 2), using default instead")
+                boundaries = self._grid
+                warnings.warn(
+                    "boundaries cannot be accessed as a numpy array with shape (3, 2), using default instead"
+                )
 
         if boundaries.shape != (3, 2):
             warnings.warn("boundaries should have shape (3, 2), using default instead")
-            boundaries = self._grid 
+            boundaries = self._grid
 
         # Load data
         if self._all_loaded:
@@ -795,12 +1193,32 @@ class Diagnostic:
         X, Y, Z = np.meshgrid(self._x[0], self._x[1], self._x[2], indexing="ij")
 
         # Flatten arrays for scatter plot
-        X_flat, Y_flat, Z_flat, = X.ravel(), Y.ravel(), Z.ravel()
+        (
+            X_flat,
+            Y_flat,
+            Z_flat,
+        ) = (
+            X.ravel(),
+            Y.ravel(),
+            Z.ravel(),
+        )
         data_flat = data.ravel()
 
         # Apply filter: Keep only chosen points
-        mask = (X_flat > boundaries[0][0]) & (X_flat < boundaries[0][1]) & (Y_flat > boundaries[1][0]) & (Y_flat < boundaries[1][1]) & (Z_flat > boundaries[2][0]) & (Z_flat < boundaries[2][1])
-        X_cut, Y_cut, Z_cut, data_cut = X_flat[mask], Y_flat[mask], Z_flat[mask], data_flat[mask]
+        mask = (
+            (X_flat > boundaries[0][0])
+            & (X_flat < boundaries[0][1])
+            & (Y_flat > boundaries[1][0])
+            & (Y_flat < boundaries[1][1])
+            & (Z_flat > boundaries[2][0])
+            & (Z_flat < boundaries[2][1])
+        )
+        X_cut, Y_cut, Z_cut, data_cut = (
+            X_flat[mask],
+            Y_flat[mask],
+            Z_flat[mask],
+            data_flat[mask],
+        )
 
         if scale_type == "zero_centered":
             # Center colormap around zero
@@ -841,10 +1259,22 @@ class Diagnostic:
         # Labels
         # TODO try to use a latex label instaead of _name
         cbar.set_label(r"${}$".format(self._name) + r"$\  [{}]$".format(self._units))
-        ax.set_title(r"$t={:.2f}$".format(self.time(idx)[0]) + r"$\  [{}]$".format(self.time(idx)[1]))
-        ax.set_xlabel(r"${}$".format(self.axis[0]["long_name"]) + r"$\  [{}]$".format(self.axis[0]["units"]))
-        ax.set_ylabel(r"${}$".format(self.axis[1]["long_name"]) + r"$\  [{}]$".format(self.axis[1]["units"]))
-        ax.set_zlabel(r"${}$".format(self.axis[2]["long_name"]) + r"$\  [{}]$".format(self.axis[2]["units"]))
+        ax.set_title(
+            r"$t={:.2f}$".format(self.time(idx)[0])
+            + r"$\  [{}]$".format(self.time(idx)[1])
+        )
+        ax.set_xlabel(
+            r"${}$".format(self.axis[0]["long_name"])
+            + r"$\  [{}]$".format(self.axis[0]["units"])
+        )
+        ax.set_ylabel(
+            r"${}$".format(self.axis[1]["long_name"])
+            + r"$\  [{}]$".format(self.axis[1]["units"])
+        )
+        ax.set_zlabel(
+            r"${}$".format(self.axis[2]["long_name"])
+            + r"$\  [{}]$".format(self.axis[2]["units"])
+        )
 
         return fig, ax
 
@@ -852,84 +1282,108 @@ class Diagnostic:
     @property
     def data(self):
         if self._data is None:
-            raise ValueError("Data not loaded into memory. Use get_* method with load_all=True or access via generator/index.")
+            raise ValueError(
+                "Data not loaded into memory. Use get_* method with load_all=True or access via generator/index."
+            )
         return self._data
 
     @property
     def dx(self):
         return self._dx
-    
+
     @property
     def nx(self):
         return self._nx
-    
+
     @property
     def x(self):
         return self._x
-    
+
     @property
     def dt(self):
         return self._dt
-    
+
     @property
     def grid(self):
         return self._grid
-    
+
     @property
     def axis(self):
         return self._axis
-    
+
     @property
     def units(self):
         return self._units
-    
+
     @property
     def tunits(self):
         return self._tunits
-    
+
     @property
     def name(self):
         return self._name
-    
+
     @property
     def dim(self):
         return self._dim
-    
+
     @property
     def path(self):
         return self
-    
+
     @property
     def simulation_folder(self):
         return self._simulation_folder
-    
+
     @property
     def ndump(self):
         return self._ndump
     
+    # @property
+    # def iter(self):
+    #     return self._iter
+
     @property
     def all_loaded(self):
         return self._all_loaded
-    
+
     @property
     def maxiter(self):
         return self._maxiter
-    
+
     @property
     def label(self):
         return self._label
+    
+    @property
+    def type(self):
+        return self._type
 
     @property
     def quantity(self):
         return self._quantity
-    
+
     def time(self, index):
         return [index * self._dt * self._ndump, self._tunits]
-    
+
+    def attributes_to_save(self, index):
+        """
+        Prints the attributes of the diagnostic.
+        """
+        print(f"dt: {self._dt}\n"
+              f"dim: {self._dim}\n"
+              f"time: {self.time(index)[0]}\n"
+              f"tunits: {self.time(index)[1]}\n"
+              f"iter: {self._ndump * index}\n"
+              f"name: {self._name}\n"
+              f"type: {self._type}\n"
+              f"label: {self._label}\n"
+              f"units: {self._units}")
+
     @dx.setter
     def dx(self, value):
         self._dx = value
-    
+
     @nx.setter
     def nx(self, value):
         self._nx = value
@@ -957,7 +1411,7 @@ class Diagnostic:
     @tunits.setter
     def tunits(self, value):
         self._tunits = value
-    
+
     @name.setter
     def name(self, value):
         self._name = value
@@ -969,11 +1423,15 @@ class Diagnostic:
     @ndump.setter
     def ndump(self, value):
         self._ndump = value
- 
+
     @data.setter
     def data(self, value):
-        self._data = value       
+        self._data = value
 
     @quantity.setter
     def quantity(self, key):
-        self._quantity = key
+        self._quantity = key
+
+    @label.setter
+    def label(self, value):
+        self._label = value