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

osiris_utils/data/diagnostic.py

@@ -7,30 +7,17 @@ Took some inspiration from Diogo and Madox's work.
 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
-import matplotlib.pyplot as plt
+import os
 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})")
-            return dim
-        except:
-            continue
-
-    raise Exception("Error parsing grid dimension")
+import h5py
+import matplotlib.pyplot as plt
+import numpy as np
+import tqdm
 
+from .data import OsirisGridFile
 
 OSIRIS_DENSITY = ["n"]
 OSIRIS_SPECIE_REPORTS = ["charge", "q1", "q2", "q3", "j1", "j2", "j3"]
@@ -63,7 +50,7 @@ OSIRIS_FLD = [
     "b3",
     "part_e1",
     "part_e2",
-    "epart_3",
+    "part_e3",
     "part_b1",
     "part_b2",
     "part_b3",
@@ -88,13 +75,7 @@ OSIRIS_PHA = [
     "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_ALL = OSIRIS_DENSITY + OSIRIS_SPECIE_REPORTS + OSIRIS_SPECIE_REP_UDIST + OSIRIS_FLD + OSIRIS_PHA
 
 
 def which_quantities():
@@ -179,7 +160,7 @@ class Diagnostic:
         Plot a 3D scatter plot of the diagnostic data.
     time(index)
         Get the time for a given index.
-        
+
     """
 
     def __init__(self, simulation_folder=None, species=None, input_deck=None):
@@ -202,9 +183,7 @@ class Diagnostic:
         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
 
@@ -229,9 +208,7 @@ 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.")
@@ -257,9 +234,7 @@ class Diagnostic:
 
     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 = glob.glob(f"{self._path}/*.h5")[0][:-9]
         self._maxiter = len(glob.glob(f"{self._path}/*.h5"))
@@ -267,9 +242,7 @@ class Diagnostic:
 
     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 = glob.glob(f"{self._path}/*.h5")[0][:-9]
         self._maxiter = len(glob.glob(f"{self._path}/*.h5"))
@@ -277,9 +250,7 @@ class Diagnostic:
 
     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 = glob.glob(f"{self._path}/*.h5")[0][:-9]
         self._maxiter = len(glob.glob(f"{self._path}/*.h5"))
@@ -287,21 +258,16 @@ class Diagnostic:
 
     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 = 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._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)]
         #     xmax = [deval(input_deck["space"][0][f"xmax(1:{self._dim})"].split(',')[i]) for i in range(self._dim)]
@@ -310,7 +276,7 @@ class Diagnostic:
         #     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:
             # Try files 000001, 000002, etc. until one is found
             found_file = False
@@ -333,7 +299,7 @@ class Diagnostic:
                     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:
@@ -344,11 +310,7 @@ class Diagnostic:
             raise ValueError("Simulation folder not set.")
         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):
         """
@@ -368,9 +330,7 @@ class Diagnostic:
         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:
                     size = self._maxiter
@@ -381,14 +341,10 @@ class Diagnostic:
                     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
 
@@ -398,9 +354,7 @@ class Diagnostic:
         # Original implementation for file-based diagnostics
         print("Loading all data from files. This may take a while.")
         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._data = np.stack([self[i] for i in tqdm.tqdm(range(size), desc="Loading data")])
         self._all_loaded = True
         return self._data
 
@@ -409,7 +363,7 @@ class Diagnostic:
         Unload data from memory. This is useful to free memory when the data is not needed anymore.
         """
         print("Unloading data from memory.")
-        if self._all_loaded == False:
+        if self._all_loaded is False:
             print("Data is not loaded.")
             return
         self._data = None
@@ -421,47 +375,78 @@ class Diagnostic:
         """
         self._data = next(self._data_generator(index))
 
+    # def __getitem__(self, index):
+    #     # 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"):
+    #             return next(self._data_generator(index))
+
+    #         # For derived diagnostics with custom generators
+    #         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:
+    #                 stop = self._maxiter
+    #             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}."
+    #                 )
+    #         else:
+    #             stop = index.stop
+
+    #         indices = range(start, stop, step)
+    #         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):
+    #             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."
+    #     )
+
     def __getitem__(self, index):
-        # 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"):
-                return next(self._data_generator(index))
+        data_gen = getattr(self, "_data_generator", None)
+        has_gen = callable(data_gen)
 
-            # For derived diagnostics with custom generators
-            if hasattr(self, "_data_generator") and callable(self._data_generator):
-                return next(self._data_generator(index))
+        if isinstance(index, int):
+            if has_gen:
+                try:
+                    return next(data_gen(index))
+                except Exception as e:
+                    raise RuntimeError(f"Error loading data at index {index}: {e}")
 
         elif isinstance(index, slice):
-            start = 0 if index.start is None else index.start
-            step = 1 if index.step is None else index.step
+            start = index.start or 0
+            step = index.step or 1
+            stop = index.stop if index.stop is not None else self._maxiter
+            indices = range(start, stop, step)
 
-            if index.stop is 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(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}."
-                    )
-            else:
-                stop = index.stop
+            if has_gen:
+                data_list = []
+                for i in indices:
+                    try:
+                        data_list.append(next(data_gen(i)))
+                    except Exception as e:
+                        raise RuntimeError(f"Error loading slice at index {i}: {e}")
+                return np.stack(data_list)
 
-            indices = range(start, stop, step)
-            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):
-                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 index {index}. No data loaded and no generator available.")
 
     def __iter__(self):
         # If this is a file-based diagnostic
@@ -483,18 +468,14 @@ class Diagnostic:
                     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)):
@@ -534,12 +515,10 @@ class Diagnostic:
                         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):
@@ -580,9 +559,7 @@ class Diagnostic:
 
                 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"
 
@@ -625,16 +602,13 @@ class Diagnostic:
                         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 [
@@ -672,9 +646,7 @@ class Diagnostic:
 
                 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"
 
@@ -717,10 +689,8 @@ class Diagnostic:
                         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
@@ -763,10 +733,8 @@ class Diagnostic:
 
                 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
@@ -808,16 +776,13 @@ class Diagnostic:
                         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 [
@@ -855,10 +820,8 @@ class Diagnostic:
 
                 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
@@ -902,26 +865,20 @@ class Diagnostic:
                         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
+    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):
@@ -964,16 +921,13 @@ class Diagnostic:
                         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 [
@@ -1011,10 +965,8 @@ class Diagnostic:
 
                 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
@@ -1036,16 +988,16 @@ class Diagnostic:
         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}"
+        if path is None:
+            path = self._simulation_folder
+            self._save_path = path + f"/MS/MISC/{self._default_save}/{savename}"
         else:
-            self._save_path = folder
+            self._save_path = path
         # 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"):
@@ -1057,64 +1009,73 @@ class Diagnostic:
             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:
+            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(
+                    "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("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"")])
-                
+                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_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).")
+
+        print("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 not os.path.exists(path):
+            print(f"Creating folder {path}...")
+            os.makedirs(path)
+        if not os.path.isdir(path):
+            raise ValueError(f"{path} is not a directory.")
 
-        if all == False:
+        if all is False:
             if isinstance(index, int):
                 filename = self._save_path + f"/{savename}-{index:06d}.h5"
                 savefile(filename, index)
@@ -1122,13 +1083,13 @@ class Diagnostic:
                 for i in index:
                     filename = self._save_path + f"/{savename}-{i:06d}.h5"
                     savefile(filename, i)
-        elif all == True:
+        elif all is 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,
@@ -1174,11 +1135,9 @@ class Diagnostic:
         if not isinstance(boundaries, np.ndarray):
             try:
                 boundaries = np.array(boundaries)
-            except:
+            except Exception:
                 boundaries = self._grid
-                warnings.warn(
-                    "boundaries cannot be accessed as a numpy array with shape (3, 2), using default instead"
-                )
+                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")
@@ -1259,22 +1218,10 @@ 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
 
@@ -1282,9 +1229,7 @@ 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
@@ -1338,7 +1283,7 @@ class Diagnostic:
     @property
     def ndump(self):
         return self._ndump
-    
+
     # @property
     # def iter(self):
     #     return self._iter
@@ -1354,7 +1299,7 @@ class Diagnostic:
     @property
     def label(self):
         return self._label
-    
+
     @property
     def type(self):
         return self._type
@@ -1370,15 +1315,17 @@ class Diagnostic:
         """
         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}")
+        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):