Allow for optional saving of background state

docs/general/parameter_file.md

@@ -130,6 +130,7 @@ The level of output is controlled through the integer `logging_level`:
 | write_matrices | logical | whether to write the matrices to the datfile | `.false.` |
 | write_eigenvectors | logical | whether to write the eigenvectors to the datfile | `.false.` |
 | write_residuals | logial | whether to write the residuals to the datfile | `.false.` |
+| write_background | logical | whether to write the background to the datfile | `.true.` |
 | write_eigenfunctions | logical | whether to write the eigenfunctions to the datfile | `.true.` |
 | write_derived_eigenfunctions | logical | whether to write the derived eigenfunctions to the datfile | `.false.` |
 | write_eigenfunction_subset | logical | if `.true.` only writes eigenfunctions for eigenvalues inside a given radius around a point in the complex plane | `.false.` |

post_processing/pylbo/automation/defaults.py

@@ -18,6 +18,7 @@
         ("write_matrices", bool),
         ("write_eigenvectors", bool),
         ("write_residuals", bool),
+        ("write_background", bool),
         ("write_eigenfunctions", bool),
         ("write_derived_eigenfunctions", bool),
         ("show_results", bool),

post_processing/pylbo/data_containers.py

@@ -7,6 +7,7 @@
 import numpy as np
 from pylbo._version import VersionHandler
 from pylbo.exceptions import (
+    BackgroundNotPresent,
     EigenfunctionsNotPresent,
     EigenvectorsNotPresent,
     MatricesNotPresent,
@@ -55,6 +56,10 @@ def ef_grid(self):
     def ef_names(self):
         pass
 
+    @abstractmethod
+    def has_background(self):
+        pass
+
     @abstractmethod
     def has_derived_efs(self):
         pass
@@ -229,6 +234,11 @@ def parameters(self) -> dict:
         """Returns the parameters in a dict with the parameter names as keys"""
         return self._parameters
 
+    @property
+    def has_background(self) -> bool:
+        """Returns `True` if background is present."""
+        return self.header["has_background"]
+
     @property
     def has_efs(self) -> bool:
         """Returns `True` if eigenfunctions are present."""
@@ -294,6 +304,8 @@ def get_sound_speed(self, which_values=None) -> Union[float, np.ndarray]:
             Array with the sound speed at every grid point, or a float corresponding
             to the value of `which_values` if provided.
         """
+        if not self.has_background:
+            raise BackgroundNotPresent(self.datfile, "get sound speed")
         pressure = self.equilibria["T0"] * self.equilibria["rho0"]
         cs = np.sqrt(self.gamma * pressure / self.equilibria["rho0"])
         return get_values(cs, which_values)
@@ -314,6 +326,8 @@ def get_alfven_speed(self, which_values=None) -> Union[float, np.ndarray]:
             Array with the Alfvén speed at every grid point,
             or a float corresponding to the value of `which_values` if provided.
         """
+        if not self.has_background:
+            raise BackgroundNotPresent(self.datfile, "get Alfvén speed")
         cA = self.equilibria["B0"] / np.sqrt(self.equilibria["rho0"])
         return get_values(cA, which_values)
 
@@ -334,16 +348,17 @@ def get_tube_speed(self, which_values=None) -> Union[float, np.ndarray]:
             or a float corresponding to the value of `which_values` if provided.
             Returns `None` if the geometry is not cylindrical.
         """
+        if not self.has_background:
+            raise BackgroundNotPresent(self.datfile, "get tube speed")
         if not self.geometry == "cylindrical":
             pylboLogger.warning(
                 "geometry is not cylindrical, unable to calculate tube speed"
             )
             return None
-        else:
-            cA = self.get_alfven_speed()
-            cs = self.get_sound_speed()
-            ct = cs * cA / np.sqrt(cs**2 + cA**2)
-            return get_values(ct, which_values)
+        cA = self.get_alfven_speed()
+        cs = self.get_sound_speed()
+        ct = cs * cA / np.sqrt(cs**2 + cA**2)
+        return get_values(ct, which_values)
 
     def get_reynolds_nb(self, which_values=None) -> Union[float, np.ndarray]:
         """
@@ -363,6 +378,8 @@ def get_reynolds_nb(self, which_values=None) -> Union[float, np.ndarray]:
             or a float corresponding to the value of `which_values` if provided.
             Returns `None` if the resistivity is zero somewhere on the domain.
         """
+        if not self.has_background:
+            raise BackgroundNotPresent(self.datfile, "get Reynolds number")
         cs = self.get_sound_speed()
         a = self.x_end - self.x_start
         eta = self.equilibria["eta"]
@@ -372,9 +389,8 @@ def get_reynolds_nb(self, which_values=None) -> Union[float, np.ndarray]:
                 "calculate the Reynolds number"
             )
             return None
-        else:
-            Re = a * cs / eta
-            return get_values(Re, which_values)
+        Re = a * cs / eta
+        return get_values(Re, which_values)
 
     def get_magnetic_reynolds_nb(self, which_values=None) -> Union[float, np.ndarray]:
         """
@@ -394,6 +410,8 @@ def get_magnetic_reynolds_nb(self, which_values=None) -> Union[float, np.ndarray
             or a float corresponding to the value of `which_values` if provided.
             Returns `None` if the resistivity is zero somewhere on the domain.
         """
+        if not self.has_background:
+            raise BackgroundNotPresent(self.datfile, "get magnetic Reynolds number")
         cA = self.get_alfven_speed()
         a = self.x_end - self.x_start
         eta = self.equilibria["eta"]
@@ -403,9 +421,8 @@ def get_magnetic_reynolds_nb(self, which_values=None) -> Union[float, np.ndarray
                 "calculate the magnetic Reynolds number"
             )
             return None
-        else:
-            Rm = a * cA / eta
-            return get_values(Rm, which_values)
+        Rm = a * cA / eta
+        return get_values(Rm, which_values)
 
     def get_k0_squared(self) -> float:
         """
@@ -715,7 +732,10 @@ def continua(self) -> dict:
         Returns the continua. Each key corresponds to a multiple Numpy arrays,
         one for each dataset.
         """
-        keys = self[0].continua.keys()
+        continua = self[0].continua
+        if continua is None:
+            return np.array([None] * len(self), dtype=object)
+        keys = continua.keys()
         _continua = {key: [] for key in keys}
         for ds in self:
             for key in keys:
@@ -735,6 +755,11 @@ def parameters(self) -> dict:
                 _params[key].append(ds.parameters[key])
         return {key: np.array(values) for key, values in _params.items()}
 
+    @property
+    def has_background(self) -> np.ndarray:
+        """Returns `True` if background is present."""
+        return np.array([ds.has_background for ds in self.datasets], dtype=bool)
+
     @property
     def has_efs(self) -> np.ndarray:
         """Returns `True` if eigenfunctions are present."""

post_processing/pylbo/exceptions.py

@@ -29,6 +29,29 @@ def __str__(self):
         return f"{self.file} is not a valid Legolas file"
 
 
+class BackgroundNotPresent(LegolasException):
+    """
+    Handles trying to query for the background when this is not present in the datfile.
+
+    Parameters
+    ----------
+    file : str, ~os.PathLike
+        The path to the file.
+    unable_to_do : str
+        The thing that was unable to be done.
+    """
+
+    def __init__(self, file, unable_to_get=None):
+        self.file = file
+        self.unable_to_do = unable_to_get
+
+    def __str__(self):
+        msg = f"No background present in {self.file}"
+        if self.unable_to_do is not None:
+            msg = f"{msg}, unable to {self.unable_to_do}"
+        return msg
+
+
 class EigenfunctionsNotPresent(LegolasException):
     """
     Handles trying to query for eigenfunctions when these

post_processing/pylbo/utilities/datfiles/file_reader.py

@@ -1,4 +1,5 @@
 from __future__ import annotations
+
 from os import PathLike
 from typing import BinaryIO
 

post_processing/pylbo/utilities/datfiles/header.py

@@ -217,7 +217,12 @@ def _read_parameters(self, istream: BinaryIO) -> dict:
 
     def _read_equilibrium_names(self, istream: BinaryIO) -> dict:
         nb_names, len_name = read_int_from_istream(istream, amount=2)
-        names = read_string_from_istream(istream, length=len_name, amount=nb_names)
+        self.data["has_background"] = nb_names > 0
+        names = (
+            read_string_from_istream(istream, length=len_name, amount=nb_names)
+            if self.data["has_background"]
+            else []
+        )
         return {"equilibrium_names": names}
 
     def _get_eigenfunction_offsets(self, istream: BinaryIO) -> dict:

post_processing/pylbo/utilities/datfiles/header_legacy.py

@@ -205,3 +205,4 @@ def _handle_entries_not_in_datfile_for_compatibility(self) -> None:
             "dim_quadblock": 2 * 8 * 2,
             "dim_matrix": self.data["gridpoints"] * 8 * 2,
         }
+        self.data["has_background"] = True

post_processing/pylbo/visualisation/api.py

@@ -196,6 +196,12 @@ def plot_equilibrium(data, figsize=None, interactive=True, **kwargs):
         The profile instance containing the equilibrium plots.
     """
     ensure_dataset(data)
+    if not data.has_background:
+        pylboLogger.warning(
+            "The dataset does not contain a background, equilibrium profiles "
+            "can not be plotted."
+        )
+        return
     p = EquilibriumProfile(data, figsize, interactive, **kwargs)
     return p
 
@@ -218,6 +224,13 @@ def plot_equilibrium_balance(data, figsize=None, **kwargs):
     p : ~pylbo.visualisation.profiles.EquilibriumBalance
         The profile instance containing the equilibrium balance plots.
     """
+    ensure_dataset(data)
+    if not data.has_background:
+        pylboLogger.warning(
+            "The dataset does not contain a background, equilibrium balance "
+            "can not be plotted."
+        )
+        return
     p = EquilibriumBalance(data, figsize, **kwargs)
     return p
 

post_processing/pylbo/visualisation/continua.py

@@ -35,9 +35,13 @@ def calculate_continua(ds):
 
     Returns
     -------
-    continua : dict
-        Dictonary containing the various continua as numpy arrays.
+    continua : dict, None
+        Dictonary containing the various continua as numpy arrays. If the
+        dataset does not contain a background, `None` is returned.
     """
+    if not ds.has_background:
+        return None
+
     rho = ds.equilibria["rho0"]
     B02 = ds.equilibria["B02"]
     B03 = ds.equilibria["B03"]

post_processing/pylbo/visualisation/modes/mode_data.py

@@ -5,6 +5,7 @@
 
 import numpy as np
 from pylbo.data_containers import LegolasDataSet
+from pylbo.exceptions import BackgroundNotPresent
 from pylbo.utilities.logger import pylboLogger
 from pylbo.visualisation.utils import ef_name_to_latex, validate_ef_name
 
@@ -56,6 +57,8 @@ def __init__(
         self.ds = ds
         self.use_real_part = use_real_part
         self.complex_factor = self._validate_complex_factor(complex_factor)
+        if add_background and not ds.has_background:
+            raise BackgroundNotPresent(ds.datfile, "add background to solution")
         self.add_background = add_background
         self._print_bg_info = True
 

post_processing/pylbo/visualisation/spectra/spectrum_figure.py

@@ -2,6 +2,7 @@
 
 from copy import copy
 
+import numpy as np
 from matplotlib.axes import Axes as mpl_axes
 from matplotlib.figure import Figure as mpl_fig
 from pylbo.visualisation.figure_window import InteractiveFigureWindow
@@ -125,6 +126,12 @@ def add_eigenfunctions(self):
     def add_derived_eigenfunctions(self):
         raise NotImplementedError()
 
+    def has_valid_continua(self, data):
+        continua = getattr(data, "continua", None)
+        if isinstance(continua, (list, np.ndarray)):
+            return all([c is not None for c in continua])
+        return continua is not None
+
     @property
     def c_handler(self):
         """Property, returns the continua handler."""

post_processing/pylbo/visualisation/spectra/spectrum_multi.py

@@ -176,6 +176,8 @@ def add_continua(self, interactive=True):
         interactive : bool
             If `True`, makes the legend interactive.
         """
+        if not self.has_valid_continua(self.dataseries):
+            return
         if self._c_handler is None:
             self._c_handler = ContinuaHandler(interactive=interactive)
 

post_processing/pylbo/visualisation/spectra/spectrum_single.py

@@ -89,6 +89,8 @@ def add_continua(self, interactive=True):
         c_handler : ~pylbo.continua.ContinuaHandler
             The legendhandler used to plot the continua.
         """
+        if not self.has_valid_continua(self.dataset):
+            return
         if self._c_handler is None:
             self._c_handler = ContinuaHandler(interactive=interactive)
 

src/dataIO/mod_console.f08

@@ -271,14 +271,19 @@ subroutine log_io_info(settings)
     call logger%info("            << DataIO settings >>")
     call logger%info("datfile name         : " // settings%io%get_basename_datfile())
     call logger%info("output folder        : " // settings%io%get_output_folder())
+    call logger%info("write background     : " // str(settings%io%write_background))
     call logger%info("write matrices       : " // str(settings%io%write_matrices))
-    call logger%info("write eigenvectors   : " // str(settings%io%write_eigenvectors))
-    call logger%info("write residuals      : " // str(settings%io%write_residuals))
     call logger%info("write eigenfunctions : " // str(settings%io%write_eigenfunctions))
     call logger%info( &
       "write derived eigenfunctions : " &
       // str(settings%io%write_derived_eigenfunctions) &
     )
+    if (settings%io%write_eigenvectors) then
+      call logger%info("write eigenvectors   : " // str(settings%io%write_eigenvectors))
+    end if
+    if (settings%io%write_residuals) then
+      call logger%info("write residuals      : " // str(settings%io%write_residuals))
+    end if
     if (.not. settings%io%write_ef_subset) return
     call logger%info( &
       "write eigenfunction subset : " // str(settings%io%write_ef_subset) &

src/dataIO/mod_input.f08

@@ -90,7 +90,7 @@ subroutine read_savelist(unit, settings)
     integer :: iostat
     character(str_len) :: iomsg
 
-    logical :: write_matrices, write_eigenvectors, write_residuals
+    logical :: write_matrices, write_eigenvectors, write_residuals, write_background
     logical :: write_eigenfunctions, write_derived_eigenfunctions
     logical :: write_eigenfunction_subset
     logical :: show_results
@@ -100,15 +100,16 @@ subroutine read_savelist(unit, settings)
     character(len=str_len) :: basename_datfile, output_folder
 
     namelist /savelist/ &
-      write_matrices, write_eigenvectors, write_residuals, write_eigenfunctions, &
-      write_derived_eigenfunctions, write_eigenfunction_subset, &
+      write_matrices, write_eigenvectors, write_residuals, write_background, &
+      write_eigenfunctions, write_derived_eigenfunctions, write_eigenfunction_subset, &
       show_results, basename_datfile, output_folder, logging_level, &
       eigenfunction_subset_radius, eigenfunction_subset_center
 
     ! get defaults
     write_matrices = settings%io%write_matrices
     write_eigenvectors = settings%io%write_eigenvectors
     write_residuals = settings%io%write_residuals
+    write_background = settings%io%write_background
     write_eigenfunctions = settings%io%write_eigenfunctions
     write_derived_eigenfunctions = settings%io%write_derived_eigenfunctions
     write_eigenfunction_subset = settings%io%write_ef_subset
@@ -127,6 +128,7 @@ subroutine read_savelist(unit, settings)
     settings%io%write_matrices = write_matrices
     settings%io%write_eigenvectors = write_eigenvectors
     settings%io%write_residuals = write_residuals
+    settings%io%write_background = write_background
     settings%io%write_eigenfunctions = write_eigenfunctions
     settings%io%write_derived_eigenfunctions = write_derived_eigenfunctions
     settings%io%write_ef_subset = write_eigenfunction_subset