Convergence restructure

tardis/io/configuration/config_reader.py

@@ -429,7 +429,7 @@ def parse_convergence_section(convergence_section_dict):
         convergence_section_dict : dict
             dictionary
         """
-        convergence_parameters = ["damping_constant", "threshold"]
+        convergence_parameters = ["damping_constant", "threshold", "type"]
 
         for convergence_variable in ["t_inner", "t_rad", "w"]:
             if convergence_variable not in convergence_section_dict:

tardis/simulation/base.py

@@ -20,6 +20,7 @@
 )
 from tardis.montecarlo.base import MonteCarloTransportSolver
 from tardis.plasma.standard_plasmas import assemble_plasma
+from tardis.simulation.convergence import ConvergenceSolver
 from tardis.util.base import is_notebook
 from tardis.visualization import ConvergencePlots
 
@@ -155,21 +156,21 @@ def __init__(
         self.show_progress_bars = show_progress_bars
         self.version = tardis.__version__
 
-        if convergence_strategy.type in ("damped"):
-            self.convergence_strategy = convergence_strategy
-            self.converged = False
-            self.consecutive_converges_count = 0
-        elif convergence_strategy.type in ("custom"):
-            raise NotImplementedError(
-                "Convergence strategy type is custom; "
-                "you need to implement your specific treatment!"
-            )
-        else:
-            raise ValueError(
-                f"Convergence strategy type is "
-                f"not damped or custom "
-                f"- input is {convergence_strategy.type}"
-            )
+        # Convergence
+        self.convergence_strategy = convergence_strategy
+        self.converged = False
+        self.consecutive_converges_count = 0
+
+        # Convergence solvers
+        self.t_rad_convergence_solver = ConvergenceSolver(
+            self.convergence_strategy.t_rad
+        )
+        self.w_convergence_solver = ConvergenceSolver(
+            self.convergence_strategy.w
+        )
+        self.t_inner_convergence_solver = ConvergenceSolver(
+            self.convergence_strategy.t_inner
+        )
 
         if show_convergence_plots:
             self.convergence_plots = ConvergencePlots(
@@ -211,48 +212,25 @@ def estimate_t_inner(
 
         return input_t_inner * luminosity_ratios**t_inner_update_exponent
 
-    @staticmethod
-    def damped_converge(value, estimated_value, damping_factor):
-        # FIXME: Should convergence strategy have its own class containing this
-        # as a method
-        return value + damping_factor * (estimated_value - value)
-
     def _get_convergence_status(
         self, t_rad, w, t_inner, estimated_t_rad, estimated_w, estimated_t_inner
     ):
-        # FIXME: Move the convergence checking in its own class.
-        no_of_shells = self.simulation_state.no_of_shells
-
-        convergence_t_rad = (
-            abs(t_rad - estimated_t_rad) / estimated_t_rad
-        ).value
-        convergence_w = abs(w - estimated_w) / estimated_w
-        convergence_t_inner = (
-            abs(t_inner - estimated_t_inner) / estimated_t_inner
-        ).value
-
-        fraction_t_rad_converged = (
-            np.count_nonzero(
-                convergence_t_rad < self.convergence_strategy.t_rad.threshold
-            )
-            / no_of_shells
-        )
-
-        t_rad_converged = (
-            fraction_t_rad_converged > self.convergence_strategy.fraction
+        t_rad_converged = self.t_rad_convergence_solver.get_convergence_status(
+            t_rad.value,
+            estimated_t_rad.value,
+            self.simulation_state.no_of_shells,
         )
 
-        fraction_w_converged = (
-            np.count_nonzero(
-                convergence_w < self.convergence_strategy.w.threshold
-            )
-            / no_of_shells
+        w_converged = self.w_convergence_solver.get_convergence_status(
+            w, estimated_w, self.simulation_state.no_of_shells
         )
 
-        w_converged = fraction_w_converged > self.convergence_strategy.fraction
-
         t_inner_converged = (
-            convergence_t_inner < self.convergence_strategy.t_inner.threshold
+            self.t_inner_convergence_solver.get_convergence_status(
+                t_inner.value,
+                estimated_t_inner.value,
+                self.simulation_state.no_of_shells,
+            )
         )
 
         if np.all([t_rad_converged, w_converged, t_inner_converged]):
@@ -300,24 +278,20 @@ def advance_state(self):
         )
 
         # calculate_next_plasma_state equivalent
-        # FIXME: Should convergence strategy have its own class?
-        next_t_radiative = self.damped_converge(
+        next_t_radiative = self.t_rad_convergence_solver.converge(
             self.simulation_state.t_radiative,
             estimated_t_rad,
-            self.convergence_strategy.t_rad.damping_constant,
         )
-        next_dilution_factor = self.damped_converge(
+        next_dilution_factor = self.w_convergence_solver.converge(
             self.simulation_state.dilution_factor,
             estimated_dilution_factor,
-            self.convergence_strategy.w.damping_constant,
         )
         if (
             self.iterations_executed + 1
         ) % self.convergence_strategy.lock_t_inner_cycles == 0:
-            next_t_inner = self.damped_converge(
+            next_t_inner = self.t_inner_convergence_solver.converge(
                 self.simulation_state.t_inner,
                 estimated_t_inner,
-                self.convergence_strategy.t_inner.damping_constant,
             )
         else:
             next_t_inner = self.simulation_state.t_inner

tardis/simulation/convergence.py

@@ -0,0 +1,77 @@
+import numpy as np
+
+
+class ConvergenceSolver:
+    def __init__(self, strategy):
+        """_summary_
+
+        Parameters
+        ----------
+        strategy : _type_
+            Convergence strategy for the physical property
+
+        Raises
+        ------
+        NotImplementedError
+            Custom convergence type specified
+        ValueError
+            Unknown convergence type specified
+        """
+        self.convergence_strategy = strategy
+        self.damping_factor = self.convergence_strategy.damping_constant
+        self.threshold = self.convergence_strategy.threshold
+
+        if self.convergence_strategy.type in ("damped"):
+            self.converge = self.damped_converge
+        elif self.convergence_strategy.type in ("custom"):
+            raise NotImplementedError(
+                "Convergence strategy type is custom; "
+                "you need to implement your specific treatment!"
+            )
+        else:
+            raise ValueError(
+                f"Convergence strategy type is "
+                f"not damped or custom "
+                f"- input is {self.convergence_strategy.type}"
+            )
+
+    def damped_converge(self, value, estimated_value):
+        """Damped convergence solver
+
+        Parameters
+        ----------
+        value : np.float64
+            The current value of the physical property
+        estimated_value : np.float64
+            The estimated value of the physical property
+
+        Returns
+        -------
+        np.float64
+            The converged value
+        """
+        return value + self.damping_factor * (estimated_value - value)
+
+    def get_convergence_status(self, value, estimated_value, no_of_cells):
+        """Get the status of convergence for the physical property
+
+        Parameters
+        ----------
+        value : np.float64, Quantity
+            The current value of the physical property
+        estimated_value : np.float64, Quantity
+            The estimated value of the physical property
+        no_of_cells : np.int64
+            The number of cells to measure convergence over
+
+        Returns
+        -------
+        bool
+            True if convergence is reached
+        """
+        convergence = abs(value - estimated_value) / estimated_value
+
+        fraction_converged = (
+            np.count_nonzero(convergence < self.threshold) / no_of_cells
+        )
+        return fraction_converged > self.threshold