diff --git a/desc/examples/__init__.py b/desc/examples/__init__.py index e8eeb5966..8c0f253ad 100644 --- a/desc/examples/__init__.py +++ b/desc/examples/__init__.py @@ -3,6 +3,7 @@ import os import desc.io +from desc.backend import execute_on_cpu def listall(): @@ -13,6 +14,7 @@ def listall(): return names_stripped +@execute_on_cpu def get(name, data=None): """Get example equilibria and data. diff --git a/desc/optimize/tr_subproblems.py b/desc/optimize/tr_subproblems.py index 8c39e8229..4cc8e58fb 100644 --- a/desc/optimize/tr_subproblems.py +++ b/desc/optimize/tr_subproblems.py @@ -205,10 +205,11 @@ def phi_and_derivative(alpha, suf, s, trust_radius): """ denom = s**2 + alpha denom = jnp.where(denom == 0, 1, denom) - p_norm = jnp.linalg.norm(suf / denom) + p = -v.dot(suf / denom) + p_norm = jnp.linalg.norm(p) phi = p_norm - trust_radius phi_prime = -jnp.sum(suf**2 / denom**3) / p_norm - return phi, phi_prime + return p, phi, phi_prime # Check if J has full rank and try Gauss-Newton step. threshold = setdefault(threshold, jnp.finfo(s.dtype).eps * f.size) @@ -222,43 +223,34 @@ def truefun(*_): return p_newton, False, 0.0 def falsefun(*_): - alpha_upper = jnp.linalg.norm(suf) / trust_radius alpha_lower = 0.0 - alpha = setdefault( - initial_alpha, - jnp.maximum(0.001 * alpha_upper, (alpha_lower * alpha_upper) ** 0.5), - ) + alpha = setdefault(initial_alpha, 0.0) + alpha = jnp.clip(alpha, alpha_lower, alpha_upper) - phi, phi_prime = phi_and_derivative(alpha, suf, s, trust_radius) + _, phi, phi_prime = phi_and_derivative(alpha, suf, s, trust_radius) k = 0 def loop_cond(state): - alpha, alpha_lower, alpha_upper, phi, k = state + p, alpha, alpha_lower, alpha_upper, phi, k = state return (jnp.abs(phi) > rtol * trust_radius) & (k < max_iter) def loop_body(state): - alpha, alpha_lower, alpha_upper, phi, k = state - alpha = jnp.where( - (alpha < alpha_lower) | (alpha > alpha_upper), - jnp.maximum(0.001 * alpha_upper, (alpha_lower * alpha_upper) ** 0.5), - alpha, - ) - - phi, phi_prime = phi_and_derivative(alpha, suf, s, trust_radius) + p, alpha, alpha_lower, alpha_upper, phi, k = state + + p, phi, phi_prime = phi_and_derivative(alpha, suf, s, trust_radius) alpha_upper = jnp.where(phi < 0, alpha, alpha_upper) ratio = phi / phi_prime alpha_lower = jnp.maximum(alpha_lower, alpha - ratio) alpha -= (phi + trust_radius) * ratio / trust_radius + alpha = jnp.clip(alpha, alpha_lower, alpha_upper) k += 1 - return alpha, alpha_lower, alpha_upper, phi, k + return p, alpha, alpha_lower, alpha_upper, phi, k - alpha, *_ = while_loop( - loop_cond, loop_body, (alpha, alpha_lower, alpha_upper, phi, k) + p, alpha, *_ = while_loop( + loop_cond, loop_body, (p_newton, alpha, alpha_lower, alpha_upper, phi, k) ) - p = -v.dot(suf / (s**2 + alpha)) - # Make the norm of p equal to trust_radius; p is changed only slightly. # This is done to prevent p from lying outside the trust region # (which can cause problems later). @@ -318,25 +310,19 @@ def truefun(*_): def falsefun(*_): alpha_upper = jnp.linalg.norm(g) / trust_radius alpha_lower = 0.0 - alpha = setdefault( - initial_alpha, - jnp.maximum(0.001 * alpha_upper, (alpha_lower * alpha_upper) ** 0.5), - ) + # the final alpha value is very small. So, starting from 0 + # is faster for root finding + alpha = setdefault(initial_alpha, alpha_lower) + alpha = jnp.clip(alpha, alpha_lower, alpha_upper) k = 0 # algorithm 4.3 from Nocedal & Wright def loop_cond(state): - alpha, alpha_lower, alpha_upper, phi, k = state + p, alpha, alpha_lower, alpha_upper, phi, k = state return (jnp.abs(phi) > rtol * trust_radius) & (k < max_iter) def loop_body(state): - alpha, alpha_lower, alpha_upper, phi, k = state - - alpha = jnp.where( - (alpha < alpha_lower) | (alpha > alpha_upper), - jnp.maximum(0.001 * alpha_upper, (alpha_lower * alpha_upper) ** 0.5), - alpha, - ) + p, alpha, alpha_lower, alpha_upper, phi, k = state Bi = B + alpha * jnp.eye(B.shape[0]) R = chol(Bi) @@ -350,21 +336,16 @@ def loop_body(state): q_norm = jnp.linalg.norm(q) alpha += (p_norm / q_norm) ** 2 * phi / trust_radius - alpha = jnp.where( - (alpha < alpha_lower) | (alpha > alpha_upper), - jnp.maximum(0.001 * alpha_upper, (alpha_lower * alpha_upper) ** 0.5), - alpha, - ) + alpha = jnp.clip(alpha, alpha_lower, alpha_upper) k += 1 - return alpha, alpha_lower, alpha_upper, phi, k + return p, alpha, alpha_lower, alpha_upper, phi, k - alpha, *_ = while_loop( - loop_cond, loop_body, (alpha, alpha_lower, alpha_upper, jnp.inf, k) + p, alpha, *_ = while_loop( + loop_cond, + loop_body, + (p_newton, alpha, alpha_lower, alpha_upper, jnp.inf, k), ) - Bi = B + alpha * jnp.eye(B.shape[0]) - R = chol(Bi) - p = cho_solve((R, True), -g) # Make the norm of p equal to trust_radius; p is changed only slightly. # This is done to prevent p from lying outside the trust region @@ -385,6 +366,15 @@ def trust_region_step_exact_qr( Solves problems of the form min_p ||J*p + f||^2, ||p|| < trust_radius + Introduces a Levenberg-Marquardt parameter alpha to make the problem + well-conditioned. + min_p ||J*p + f||^2 + alpha*||p||^2, ||p|| < trust_radius + + The objective function can be written as + p.T(J.T@J + alpha*I)p + 2f.TJp + f.Tf + which is equivalent to + || [J; sqrt(alpha)*I].Tp - [f; 0].T ||^2 + Parameters ---------- f : ndarray @@ -421,26 +411,20 @@ def truefun(*_): def falsefun(*_): alpha_upper = jnp.linalg.norm(J.T @ f) / trust_radius alpha_lower = 0.0 - alpha = setdefault( - initial_alpha, - jnp.maximum(0.001 * alpha_upper, (alpha_lower * alpha_upper) ** 0.5), - ) + # the final alpha value is very small. So, starting from 0 + # is faster for root finding + alpha = setdefault(initial_alpha, alpha_lower) + alpha = jnp.clip(alpha, alpha_lower, alpha_upper) k = 0 - # algorithm 4.3 from Nocedal & Wright + fp = jnp.pad(f, (0, J.shape[1])) def loop_cond(state): - alpha, alpha_lower, alpha_upper, phi, k = state + p, alpha, alpha_lower, alpha_upper, phi, k = state return (jnp.abs(phi) > rtol * trust_radius) & (k < max_iter) def loop_body(state): - alpha, alpha_lower, alpha_upper, phi, k = state - - alpha = jnp.where( - (alpha < alpha_lower) | (alpha > alpha_upper), - jnp.maximum(0.001 * alpha_upper, (alpha_lower * alpha_upper) ** 0.5), - alpha, - ) + p, alpha, alpha_lower, alpha_upper, phi, k = state Ji = jnp.vstack([J, jnp.sqrt(alpha) * jnp.eye(J.shape[1])]) # Ji is always tall since its padded by alpha*I @@ -456,22 +440,16 @@ def loop_body(state): q_norm = jnp.linalg.norm(q) alpha += (p_norm / q_norm) ** 2 * phi / trust_radius - alpha = jnp.where( - (alpha < alpha_lower) | (alpha > alpha_upper), - jnp.maximum(0.001 * alpha_upper, (alpha_lower * alpha_upper) ** 0.5), - alpha, - ) + alpha = jnp.clip(alpha, alpha_lower, alpha_upper) k += 1 - return alpha, alpha_lower, alpha_upper, phi, k + return p, alpha, alpha_lower, alpha_upper, phi, k - alpha, *_ = while_loop( - loop_cond, loop_body, (alpha, alpha_lower, alpha_upper, jnp.inf, k) + p, alpha, *_ = while_loop( + loop_cond, + loop_body, + (p_newton, alpha, alpha_lower, alpha_upper, jnp.inf, k), ) - Ji = jnp.vstack([J, jnp.sqrt(alpha) * jnp.eye(J.shape[1])]) - Q, R = qr(Ji, mode="economic") - p = solve_triangular(R, -Q.T @ fp) - # Make the norm of p equal to trust_radius; p is changed only slightly. # This is done to prevent p from lying outside the trust region # (which can cause problems later). diff --git a/tests/test_optimizer.py b/tests/test_optimizer.py index 637d247a8..82e9130ee 100644 --- a/tests/test_optimizer.py +++ b/tests/test_optimizer.py @@ -1416,7 +1416,8 @@ def test_optimize_coil_currents(DummyCoilSet): verbose=2, copy=True, ) - # check that optimized coil currents changed by more than 15% from initial values + # check that on average optimized coil currents changed by more than + # 15% from initial values np.testing.assert_array_less( np.asarray(coils.current).mean() * 0.15, np.abs(np.asarray(coils_opt.current) - np.asarray(coils.current)).mean(),