From 3ac70269b1a419a1b28e75aadf141bd5255fc15a Mon Sep 17 00:00:00 2001 From: Brandon Talamini Date: Tue, 27 Feb 2024 10:56:57 -0800 Subject: [PATCH 1/3] Fix GMRES solver errors in tests triggered in new scipy version Scipy GMRES now modifies the preconditioned residual in place. This is perfectly reasonable, but we were passing in jax numpy arrays, which do not allow this. This was unsafe, but we got away with it. This commit makes deep copies of jax arrays to plain numpy arrays in the operators and initial data that the scipy gmres sees. --- optimism/NewtonSolver.py | 15 ++++++++++----- 1 file changed, 10 insertions(+), 5 deletions(-) diff --git a/optimism/NewtonSolver.py b/optimism/NewtonSolver.py index 22ad8fd3..e432f90c 100644 --- a/optimism/NewtonSolver.py +++ b/optimism/NewtonSolver.py @@ -1,3 +1,4 @@ +import numpy as onp from scipy.sparse.linalg import LinearOperator, gmres from optimism.JaxConfig import * @@ -27,9 +28,12 @@ def compute_min_p(ps, bounds): def newton_step(residual, linear_op, x, settings=Settings(1e-2,100), precond=None): - sz = x.size + sz = x.size + # The call to onp.array copies the jax array output into a plain numpy + # array. The copy is necessary for safety, since as far as scipy knows, + # it is allowed to modify the output in place. A = LinearOperator((sz,sz), - lambda v: linear_op(v)) + lambda v: onp.array(linear_op(v))) r = residual(x) rNorm = np.linalg.norm(r) @@ -42,11 +46,12 @@ def callback(xk): maxIters = settings.max_gmres_iters if precond==None: - dx, exitcode = gmres(A, r, tol=relTol*rNorm, atol=0, callback_type='legacy', callback=callback, maxiter=maxIters) + dx, exitcode = gmres(A, onp.array(r), tol=relTol*rNorm, atol=0, callback_type='legacy', callback=callback, maxiter=maxIters) else: + # Another copy to a plain numpy array, see comment for A above. M = LinearOperator((sz,sz), - lambda v: precond(v)) - dx, exitcode = gmres(A, r, tol=relTol*rNorm, atol=0, M=M, callback_type='legacy', callback=callback, maxiter=maxIters) + lambda v: onp.array(precond(v))) + dx, exitcode = gmres(A, onp.array(r), tol=relTol*rNorm, atol=0, M=M, callback_type='legacy', callback=callback, maxiter=maxIters) print('Number of GMRES iters = ', numIters) From 21b5f72ee2c625b85fb15c8ac0ff9fdd36b3c3ab Mon Sep 17 00:00:00 2001 From: Brandon Talamini Date: Tue, 27 Feb 2024 11:58:44 -0800 Subject: [PATCH 2/3] Make logic for preconditioning choice simpler in Newton solver --- optimism/NewtonSolver.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/optimism/NewtonSolver.py b/optimism/NewtonSolver.py index e432f90c..162e727a 100644 --- a/optimism/NewtonSolver.py +++ b/optimism/NewtonSolver.py @@ -45,14 +45,14 @@ def callback(xk): relTol = settings.relative_gmres_tol maxIters = settings.max_gmres_iters - if precond==None: - dx, exitcode = gmres(A, onp.array(r), tol=relTol*rNorm, atol=0, callback_type='legacy', callback=callback, maxiter=maxIters) - else: + if precond is not None: # Another copy to a plain numpy array, see comment for A above. M = LinearOperator((sz,sz), lambda v: onp.array(precond(v))) - dx, exitcode = gmres(A, onp.array(r), tol=relTol*rNorm, atol=0, M=M, callback_type='legacy', callback=callback, maxiter=maxIters) + else: + M = None + dx, exitcode = gmres(A, onp.array(r), tol=relTol*rNorm, atol=0, M=M, callback_type='legacy', callback=callback, maxiter=maxIters) print('Number of GMRES iters = ', numIters) return -dx, exitcode From 6025ebdeaecacbf148a1d96fd981be0ec4d3b20e Mon Sep 17 00:00:00 2001 From: Brandon Talamini Date: Tue, 27 Feb 2024 14:43:25 -0800 Subject: [PATCH 3/3] Fix conceptual error in GMRES tolerance specification The relative tolerance for convergence was being set to a fraction of the rhs norm, but the `tol` option is already a relative tolerance. This commit also fies a depracation warning, since the `tol` parameter is deprecated in facor of `rtol`. --- optimism/NewtonSolver.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/optimism/NewtonSolver.py b/optimism/NewtonSolver.py index 162e727a..f7b66949 100644 --- a/optimism/NewtonSolver.py +++ b/optimism/NewtonSolver.py @@ -34,8 +34,7 @@ def newton_step(residual, linear_op, x, settings=Settings(1e-2,100), precond=Non # it is allowed to modify the output in place. A = LinearOperator((sz,sz), lambda v: onp.array(linear_op(v))) - r = residual(x) - rNorm = np.linalg.norm(r) + r = onp.array(residual(x)) numIters = 0 def callback(xk): @@ -52,7 +51,7 @@ def callback(xk): else: M = None - dx, exitcode = gmres(A, onp.array(r), tol=relTol*rNorm, atol=0, M=M, callback_type='legacy', callback=callback, maxiter=maxIters) + dx, exitcode = gmres(A, r, rtol=relTol, atol=0, M=M, callback_type='legacy', callback=callback, maxiter=maxIters) print('Number of GMRES iters = ', numIters) return -dx, exitcode