Merge pull request geodynamics#676 from bobmyhill/equilibration_tolerances

Improve equilibration tolerances

Author: bobmyhill

burnman/optimize/nonlinear_solvers.py

@@ -119,9 +119,11 @@ def __init__(
         ), "The starting guess is outside the supplied constraints."
 
         if not isinstance(self.tol, float):
-            assert len(self.tol) < len(
+            self.tol = np.asarray(self.tol, dtype=np.float64)
+            assert self.tol.ndim == 1, "tol must be a float or a 1D array."
+            assert len(self.tol) == len(
                 self.guess
-            ), "tol must either be a float or an array like guess."
+            ), f"tol must either be a float or an array like guess. Got len(tol)={len(self.tol)} and len(guess)={len(self.guess)}."
 
     def _constraints(self, x: npt.NDArray[np.float64]) -> npt.NDArray[np.float64]:
         return np.dot(self.linear_constraints[0], x) + self.linear_constraints[1]

burnman/tools/equilibration.py

@@ -607,29 +607,48 @@ def get_equilibration_parameters(assemblage, composition, free_compositional_vec
         degrees of freedom for the equilibrium problem.
     :type free_compositional_vectors: list of dictionaries
 
-    :returns: A tuple with attributes n_parameters
-        (the number of parameters for the current equilibrium problem)
-        and phase_amount_indices (the indices of the parameters that
-        correspond to phase amounts).
+    :returns: A tuple with attributes:
+        - parameter_names: list of strings, the names of the parameters
+          for the current equilibrium problem
+        - default_tolerances: numpy.array, the default tolerances for each parameter
+        - n_parameters: int, the number of parameters for the current equilibrium problem
+        - phase_amount_indices: list of int, the indices of the parameters that
+          correspond to phase amounts
+        - bulk_composition_vector: numpy.array, the bulk composition vector
+        - free_compositional_vectors: 2D numpy.array, the free compositional vectors
+        - reduced_composition_vector: numpy.array, the bulk composition vector
+          reduced to the independent elements
+        - reduced_free_composition_vectors: 2D numpy.array, the free compositional vectors
+          reduced to the independent elements
+        - constraint_vector: numpy.array, vector for the linear constraints bounding the
+          valid parameter space (b in Ax + b < eps)
+        - constraint_matrix: 2D numpy.array, matrix for the linear constraints bounding the
+          valid parameter space (A in Ax + b < eps)
     :rtype: namedtuple
     """
     # Initialize a named tuple for the equilibration parameters
     prm = namedtuple("assemblage_parameters", [])
 
     # Process parameter names
     prm.parameter_names = ["Pressure (Pa)", "Temperature (K)"]
+    prm.default_tolerances = [1.0, 1.0e-3]
     for i, n in enumerate(assemblage.endmembers_per_phase):
         prm.parameter_names.append("x({0})".format(assemblage.phases[i].name))
+        prm.default_tolerances.append(1.0e-6)
         if n > 1:
             p_names = [
                 "p({0} in {1})".format(n, assemblage.phases[i].name)
                 for n in assemblage.phases[i].endmember_names[1:]
             ]
             prm.parameter_names.extend(p_names)
+            prm.default_tolerances.extend([1.0e-6] * len(p_names))
 
     n_free_compositional_vectors = len(free_compositional_vectors)
     for i in range(n_free_compositional_vectors):
         prm.parameter_names.append(f"v_{i}")
+        prm.default_tolerances.append(1.0e-6)
+
+    prm.default_tolerances = np.array(prm.default_tolerances)
 
     prm.n_parameters = len(prm.parameter_names)
     prm.phase_amount_indices = [
@@ -797,7 +816,7 @@ def equilibrate(
     assemblage,
     equality_constraints,
     free_compositional_vectors=[],
-    tol=1.0e-3,
+    tol=None,
     store_iterates=False,
     store_assemblage=True,
     max_iterations=100.0,
@@ -866,8 +885,16 @@ def equilibrate(
         Vector given in atomic (molar) units of elements.
     :type free_compositional_vectors: list of dict
 
-    :param tol: The tolerance for the nonlinear solver.
-    :type tol: float
+    :param tol: The tolerance for the nonlinear solver. This can be a single float,
+        which is then applied to all parameters, or a numpy array with the
+        same length as the number of parameters in the solve (2 + number of
+        phases + number of free compositional vectors). The default is None,
+        which sets the tolerances to 1 Pa for pressure, 1e-3 K for temperature,
+        1e-6 for phase amounts and 1e-6 for free compositional vectors.
+        One of the termination criteria for the nonlinear solver is that
+        the absolute change in each parameter is less than the corresponding
+        tolerance.
+    :type tol: float or numpy.array
 
     :param store_iterates: Whether to store the parameter values for
         each iteration in each solution object.
@@ -932,6 +959,10 @@ def equilibrate(
         assemblage, composition, free_compositional_vectors
     )
 
+    # Set default tolerances if not provided
+    if tol is None:
+        tol = prm.default_tolerances
+
     # Check equality constraints have the correct structure
     # Convert into the format readable by the function and jacobian functions
     eq_constraint_lists = process_eq_constraints(equality_constraints, assemblage, prm)

misc/benchmarks/slb_2024_benchmarks.py

@@ -196,13 +196,13 @@ def check_fig_3_fcc_ferric_fper():
 
     assemblage = Composite([bcc, fcc], [0.5, 0.5])
     equality_constraints = [["P", 1.0e5], ["phase_fraction", (bcc, 0.0)]]
-    sol = equilibrate(composition, assemblage, equality_constraints, tol=1.0e-5)
+    sol = equilibrate(composition, assemblage, equality_constraints)
     T_a_g = sol[0].assemblage.temperature
 
     assemblage = Composite([bcc, fper], [0.5, 0.5])
     temperatures = np.linspace(T_a_wu_mag, T_a_g, 31)
     equality_constraints = [["P", 1.0e5], ["T", temperatures]]
-    sol = equilibrate(composition, assemblage, equality_constraints, tol=1.0e-5)
+    sol = equilibrate(composition, assemblage, equality_constraints)
     xs = np.empty_like(temperatures)
     for i, s in enumerate(sol[0]):
         xs[i] = -s.assemblage.phases[1].formula["Fe"] / 4.0
@@ -211,16 +211,17 @@ def check_fig_3_fcc_ferric_fper():
     assemblage = Composite([fcc, fper], [0.5, 0.5])
     temperatures = np.linspace(T_a_g, 1800.0, 31)
     equality_constraints = [["P", 1.0e5], ["T", temperatures]]
-    sol = equilibrate(composition, assemblage, equality_constraints, tol=1.0e-5)
+    sol = equilibrate(composition, assemblage, equality_constraints)
     xs = np.empty_like(temperatures)
     for i, s in enumerate(sol[0]):
         xs[i] = -s.assemblage.phases[1].formula["Fe"] / 4.0
     ax[0].plot(xs, temperatures, linestyle=":", linewidth=3.0)
 
+    composition = {"Fe": 1.0, "O": 1.3}
     assemblage = Composite([mag, fper], [0.5, 0.5])
     temperatures = np.linspace(T_a_wu_mag, 1800.0, 31)
     equality_constraints = [["P", 1.0e5], ["T", temperatures]]
-    sol = equilibrate(composition, assemblage, equality_constraints, tol=1.0e-5)
+    sol = equilibrate(composition, assemblage, equality_constraints)
     xs = np.empty_like(temperatures)
     for i, s in enumerate(sol[0]):
         xs[i] = -s.assemblage.phases[1].formula["Fe"] / 4.0
@@ -246,9 +247,7 @@ def check_fig_3_fcc_ferric_fper():
         for i, x_MgO in enumerate(x_MgOs):
             composition = {"Mg": x_MgO, "Fe": 2.0, "O": 1.0}
             try:
-                sol = equilibrate(
-                    composition, assemblage, equality_constraints, tol=1.0e-5
-                )
+                sol = equilibrate(composition, assemblage, equality_constraints)
                 if sol[0].success:
                     f = assemblage.phases[0].formula
                     x_Mgs[i] = f["Mg"] / 4.0
@@ -408,7 +407,7 @@ def check_fig_7_fO2():
                 ["phase_fraction", [a.phases[0], 0.0]],
                 ["T", 1500.0],
             ]
-            equilibrate(a.phases[0].formula, a, equality_constraints, tol=1.0e-5)
+            equilibrate(a.phases[0].formula, a, equality_constraints)
             P_bounds[i][1] = a.pressure
             P_bounds[i + 1][0] = a.pressure
 
@@ -437,9 +436,9 @@ def check_fig_7_fO2():
         for i, P in enumerate(pressures):
             assemblage.set_state(P, T)
             try:
-                sol = equilibrate(
-                    assemblage.formula, assemblage, [["P", P], ["T", T]], tol=1.0e-7
-                )[0]
+                sol = equilibrate(assemblage.formula, assemblage, [["P", P], ["T", T]])[
+                    0
+                ]
                 if sol.success:
                     mu_O2 = assemblage.chemical_potential([{"O": 2.0}])[0]
                     O2_gas.set_state(1.0e5, T)

tests/test_equilibration.py

@@ -144,6 +144,44 @@ def test_binary_solution_convergence(self):
         sol, prm = equilibrate(composition, assemblage, equality_constraints)
         self.assertFalse(any(not s.success for s in sol))
 
+    def test_incorrect_tol(self):
+        ol = SLB_2011.mg_fe_olivine()
+        wad = SLB_2011.mg_fe_wadsleyite()
+
+        assemblage = burnman.Composite([ol, wad], [0.7, 0.3])
+        ol.set_composition([0.5, 0.5])
+        wad.set_composition([0.6, 0.4])
+
+        assemblage.set_state(10.0e9, 1200.0)
+        equality_constraints = [
+            ("P", 10.0e9),
+            (
+                "phase_composition",
+                (ol, (["Mg_A", "Fe_A"], [0.0, 1.0], [1.0, 1.0], 0.45)),
+            ),
+        ]
+        composition = {"Mg": 1.0, "Fe": 1.0, "Si": 1.0, "O": 4.0}
+
+        # These should raise errors
+        with self.assertRaises(AssertionError):
+            _ = equilibrate(
+                composition, assemblage, equality_constraints, tol=[1.0e-3, 1.0e-3]
+            )
+        with self.assertRaises(AssertionError):
+            _ = equilibrate(
+                composition, assemblage, equality_constraints, tol=[[1.0e-3] * 6] * 2
+            )
+
+        # These should work
+        _ = equilibrate(composition, assemblage, equality_constraints, tol=1.0e-3)
+        _ = equilibrate(
+            composition,
+            assemblage,
+            equality_constraints,
+            tol=[1.0e-3] * 6,
+        )
+        _ = equilibrate(composition, assemblage, equality_constraints)
+
 
 if __name__ == "__main__":
     unittest.main()