Merge pull request geodynamics#692 from bobmyhill/adaptive_logish_solve

Add adaptive logish epsilon fallback for difficult equilibrate solves

Author: bobmyhill

burnman/classes/solutionmodel.py

@@ -6,14 +6,13 @@
 
 import importlib
 import numpy as np
+from ..constants import logish_eps
 from ..utils.chemistry import process_solution_chemistry
 from .. import constants
 import warnings
 from .material import material_property, cached_property
 from dataclasses import make_dataclass
 
-np_eps = np.finfo(float).eps
-
 Interaction = make_dataclass(
     "Interaction", ["inds", "expts", "f_r", "m_jr", "f_rs", "m_jrs"]
 )
@@ -106,26 +105,26 @@ def _non_ideal_interactions_subreg(p, n_endmembers, Wijk):
     return Wint
 
 
-def logish(x, eps=np_eps):
+def logish(x, eps=logish_eps):
     """
     2nd order series expansion of log(x) about eps:
     log(eps) - sum_k=1^infty (f_eps)^k / k
     Prevents infinities at x=0
     """
-    f_eps = 1.0 - x / eps
-    mask = x > eps
-    ln = np.where(x <= eps, np.log(eps) - f_eps - f_eps * f_eps / 2.0, 0.0)
+    f_eps = 1.0 - x / eps[0]
+    mask = x > eps[0]
+    ln = np.where(x <= eps[0], np.log(eps[0]) - f_eps - f_eps * f_eps / 2.0, 0.0)
     ln[mask] = np.log(x[mask])
     return ln
 
 
-def inverseish(x, eps=np_eps):
+def inverseish(x, eps=logish_eps):
     """
     1st order series expansion of 1/x about eps: 2/eps - x/eps/eps
     Prevents infinities at x=0
     """
-    mask = x > eps
-    oneoverx = np.where(x <= eps, 2.0 / eps - x / eps / eps, 0.0)
+    mask = x > eps[0]
+    oneoverx = np.where(x <= eps[0], 2.0 / eps[0] - x / eps[0] / eps[0], 0.0)
     oneoverx[mask] = 1.0 / x[mask]
     return oneoverx
 

burnman/constants.py

@@ -3,6 +3,7 @@
 # GPL v2 or later.
 
 import scipy.constants
+import numpy as np
 
 """
 molar gas constant (R) in J mol^-1 K^-1
@@ -42,3 +43,10 @@
 1 cm^-1 in J/mol
 """
 invcm = 11.9627
+
+
+"""
+Small number for numerical stability
+Kept in list form for mutability
+"""
+logish_eps = [np.finfo(float).eps]

burnman/tools/equilibration.py

@@ -10,7 +10,8 @@
 from collections import namedtuple
 import logging
 
-from ..optimize.nonlinear_solvers import damped_newton_solve
+from ..constants import logish_eps
+from ..optimize.nonlinear_solvers import damped_newton_solve, TerminationCode
 from ..classes.solution import Solution
 
 P_scaling = 1.0e9  # Pa
@@ -1073,32 +1074,94 @@ def equilibrate(
 
         F_tol = default_F_tolerances(assemblage, equality_constraints, n_atoms)
 
-        # Set the initial fractions and compositions
-        # of the phases in the assemblage:
-        sol = damped_newton_solve(
-            F=lambda x: F(
+        def F_fn(x):
+            return F(
                 x,
                 assemblage,
                 equality_constraints,
                 prm.reduced_composition_vector,
                 prm.reduced_free_composition_vectors,
-            ),
-            J=lambda x: jacobian(
+            )
+
+        def J_fn(x):
+            return jacobian(
                 x,
                 assemblage,
                 equality_constraints,
                 prm.reduced_free_composition_vectors,
-            ),
-            lambda_bounds=lambda dx, x: lambda_bounds(
-                dx, x, assemblage.endmembers_per_phase
-            ),
+            )
+
+        def lambda_bounds_fn(dx, x):
+            return lambda_bounds(dx, x, assemblage.endmembers_per_phase)
+
+        linear_constraints = (prm.constraint_matrix, prm.constraint_vector)
+
+        sol = damped_newton_solve(
+            F=F_fn,
+            J=J_fn,
+            lambda_bounds=lambda_bounds_fn,
             guess=parameters,
-            linear_constraints=(prm.constraint_matrix, prm.constraint_vector),
+            linear_constraints=linear_constraints,
             tol=tol,
             F_tol=F_tol,
             store_iterates=store_iterates,
             max_iterations=max_iterations,
         )
+        # If we hit max iterations, we try relaxing the logish_eps
+        # and solving again. This helps convergence in hard problems because
+        # it allows larger steps when phase compositions are close to the
+        # boundaries of the solution (where the gradient in excess entropy is large).
+        if sol.code == TerminationCode.MAX_ITERATIONS:
+            old_logish_eps = logish_eps[0]
+            logish_eps[0] = 1.0e-5
+            c_shifts = sol.x[-n_free_compositional_vectors:]
+            new_parameters = get_parameters(assemblage, n_free_compositional_vectors)
+            new_parameters[-n_free_compositional_vectors:] = c_shifts
+            set_compositions_and_state_from_parameters(assemblage, parameters)  # reset
+
+            sol2 = damped_newton_solve(
+                F=F_fn,
+                J=J_fn,
+                lambda_bounds=lambda_bounds_fn,
+                guess=new_parameters,
+                linear_constraints=linear_constraints,
+                tol=tol,
+                F_tol=F_tol,
+                store_iterates=store_iterates,
+                max_iterations=max_iterations,
+            )
+
+            logish_eps[0] = old_logish_eps
+            c_shifts = sol2.x[-n_free_compositional_vectors:]
+            new_parameters = get_parameters(assemblage, n_free_compositional_vectors)
+            new_parameters[-n_free_compositional_vectors:] = c_shifts
+            set_compositions_and_state_from_parameters(assemblage, parameters)  # reset
+            sol3 = damped_newton_solve(
+                F=F_fn,
+                J=J_fn,
+                lambda_bounds=lambda_bounds_fn,
+                guess=new_parameters,
+                linear_constraints=linear_constraints,
+                tol=tol,
+                F_tol=F_tol,
+                store_iterates=store_iterates,
+                max_iterations=max_iterations,
+            )
+
+            # combine iterates from all three solves
+            sol3.n_it = sol.n_it + sol2.n_it + sol3.n_it
+            if store_iterates:
+                sol3.iterates.x = np.concatenate(
+                    (sol.iterates.x, sol2.iterates.x, sol3.iterates.x), axis=0
+                )
+                sol3.iterates.F = np.concatenate(
+                    (sol.iterates.F, sol2.iterates.F, sol3.iterates.F), axis=0
+                )
+                sol3.iterates.lmda = np.concatenate(
+                    (sol.iterates.lmda, sol2.iterates.lmda, sol3.iterates.lmda),
+                    axis=0,
+                )
+            sol = sol3
 
         if sol.success and len(assemblage.reaction_affinities) > 0.0:
             maxres = np.max(np.abs(assemblage.reaction_affinities)) + 1.0e-5

tests/test_equilibration.py

@@ -284,6 +284,45 @@ def test_mg_rich_ol_wad_eqm_with_free_compositional_vector(self):
         )
         self.assertEqual(sol.code, TerminationCode.SUCCESS)
 
+    def test_very_mg_rich_ol_wad_eqm_with_free_compositional_vector(self):
+        # Without both the scaling of parameters in equilibrate,
+        # and adaptive relaxation of logish_eps, this problem fails to converge.
+        assemblage = make_ol_wad_assemblage()
+        ol = assemblage.phases[0]
+        wad = assemblage.phases[1]
+
+        assemblage = burnman.Composite([ol, wad], [1.0, 0.0])
+
+        # Set the pressure and temperature
+        P = 12.0e9  # Pa
+        T = 1673.0  # K
+        assemblage.set_state(P, T)
+
+        # Define the starting compositions of the phases
+        x_fe_ol = 1.0e-10
+        ol.set_composition([0.8, 0.2])
+        wad.set_composition([0.7, 0.3])
+
+        # Set up the constraints and compositional degree of freedom
+        composition = {"Fe": 0.1, "Mg": 1.9, "Si": 1.0, "O": 4.0}
+        free_compositional_vectors = [{"Mg": 1.0, "Fe": -1.0}]
+
+        equality_constraints = [
+            ("T", T),
+            (
+                "phase_composition",
+                (ol, [["Mg_A", "Fe_A"], [0.0, 1.0], [1.0, 1.0], x_fe_ol]),
+            ),
+            ("phase_fraction", (wad, 0.5)),
+        ]
+        sol, _ = equilibrate(
+            composition,
+            assemblage,
+            equality_constraints,
+            free_compositional_vectors,
+        )
+        self.assertEqual(sol.code, TerminationCode.SUCCESS)
+
     def test_convergence_with_singular_system(self):
 
         composition = {"Mg": 1.0, "Fe": 1.0, "Si": 1.0, "O": 4.0}