Skip to content

Commit

Permalink
Merge branch 'master' into dd/external
Browse files Browse the repository at this point in the history
  • Loading branch information
@ddudt
ddudt authored Dec 19, 2024
2 parents 4db5c9f + 34dbac0 commit 96aec58
Show file tree
Hide file tree
Showing 3 changed files with 52 additions and 71 deletions.
2 changes: 2 additions & 0 deletions desc/examples/__init__.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@
import os

import desc.io
from desc.backend import execute_on_cpu


def listall():
Expand All @@ -13,6 +14,7 @@ def listall():
return names_stripped


@execute_on_cpu
def get(name, data=None):
"""Get example equilibria and data.
Expand Down
118 changes: 48 additions & 70 deletions desc/optimize/tr_subproblems.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand All @@ -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).
Expand Down Expand Up @@ -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)
Expand All @@ -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
Expand All @@ -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
Expand Down Expand Up @@ -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
Expand All @@ -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).
Expand Down
3 changes: 2 additions & 1 deletion tests/test_optimizer.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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(),
Expand Down

0 comments on commit 96aec58

Please sign in to comment.