Merge branch 'master' into select_gpu

PlasmaControl · Dec 24, 2024 · 23e2ece · 23e2ece
2 parents f90973b + e357a9b
commit 23e2ece
Show file tree

Hide file tree

Showing 4 changed files with 24 additions and 28 deletions.
diff --git a/desc/optimize/aug_lagrangian.py b/desc/optimize/aug_lagrangian.py
@@ -3,7 +3,7 @@
 from scipy.optimize import BFGS, NonlinearConstraint, OptimizeResult
 
 from desc.backend import jnp
-from desc.utils import errorif, setdefault
+from desc.utils import errorif, safediv, setdefault
 
 from .bound_utils import (
     cl_scaling_vector,
@@ -350,12 +350,12 @@ def laghess(z, y, mu, *args):
     # conngould : norm of the cauchy point, as recommended in ch17 of Conn & Gould
     # scipy : norm of the scaled x, as used in scipy
     # mix : geometric mean of conngould and scipy
+    tr_scipy = jnp.linalg.norm(z * scale_inv / v**0.5)
+    conngould = safediv(g_h @ g_h, abs(g_h @ H_h @ g_h))
     init_tr = {
-        "scipy": jnp.linalg.norm(z * scale_inv / v**0.5),
-        "conngould": (g_h @ g_h) / abs(g_h @ H_h @ g_h),
-        "mix": jnp.sqrt(
-            (g_h @ g_h) / abs(g_h @ H_h @ g_h) * jnp.linalg.norm(z * scale_inv / v**0.5)
-        ),
+        "scipy": tr_scipy,
+        "conngould": conngould,
+        "mix": jnp.sqrt(conngould * tr_scipy),
     }
     trust_radius = options.pop("initial_trust_radius", "conngould")
     tr_ratio = options.pop("initial_trust_ratio", 1.0)

diff --git a/desc/optimize/aug_lagrangian_ls.py b/desc/optimize/aug_lagrangian_ls.py
@@ -3,7 +3,7 @@
 from scipy.optimize import NonlinearConstraint, OptimizeResult
 
 from desc.backend import jnp, qr
-from desc.utils import errorif, setdefault
+from desc.utils import errorif, safediv, setdefault
 
 from .bound_utils import (
     cl_scaling_vector,
@@ -289,14 +289,12 @@ def lagjac(z, y, mu, *args):
     # conngould : norm of the cauchy point, as recommended in ch17 of Conn & Gould
     # scipy : norm of the scaled x, as used in scipy
     # mix : geometric mean of conngould and scipy
+    tr_scipy = jnp.linalg.norm(z * scale_inv / v**0.5)
+    conngould = safediv(jnp.sum(g_h**2), jnp.sum((J_h @ g_h) ** 2))
     init_tr = {
-        "scipy": jnp.linalg.norm(z * scale_inv / v**0.5),
-        "conngould": jnp.sum(g_h**2) / jnp.sum((J_h @ g_h) ** 2),
-        "mix": jnp.sqrt(
-            jnp.sum(g_h**2)
-            / jnp.sum((J_h @ g_h) ** 2)
-            * jnp.linalg.norm(z * scale_inv / v**0.5)
-        ),
+        "scipy": tr_scipy,
+        "conngould": conngould,
+        "mix": jnp.sqrt(conngould * tr_scipy),
     }
     trust_radius = options.pop("initial_trust_radius", "conngould")
     tr_ratio = options.pop("initial_trust_ratio", 1.0)

diff --git a/desc/optimize/fmin_scalar.py b/desc/optimize/fmin_scalar.py
@@ -3,7 +3,7 @@
 from scipy.optimize import BFGS, OptimizeResult
 
 from desc.backend import jnp
-from desc.utils import errorif, setdefault
+from desc.utils import errorif, safediv, setdefault
 
 from .bound_utils import (
     cl_scaling_vector,
@@ -247,12 +247,12 @@ def fmintr(  # noqa: C901
     # conngould : norm of the cauchy point, as recommended in ch17 of Conn & Gould
     # scipy : norm of the scaled x, as used in scipy
     # mix : geometric mean of conngould and scipy
+    tr_scipy = jnp.linalg.norm(x * scale_inv / v**0.5)
+    conngould = safediv(g_h @ g_h, abs(g_h @ H_h @ g_h))
     init_tr = {
-        "scipy": jnp.linalg.norm(x * scale_inv / v**0.5),
-        "conngould": (g_h @ g_h) / abs(g_h @ H_h @ g_h),
-        "mix": jnp.sqrt(
-            (g_h @ g_h) / abs(g_h @ H_h @ g_h) * jnp.linalg.norm(x * scale_inv / v**0.5)
-        ),
+        "scipy": tr_scipy,
+        "conngould": conngould,
+        "mix": jnp.sqrt(conngould * tr_scipy),
     }
     trust_radius = options.pop("initial_trust_radius", "scipy")
     tr_ratio = options.pop("initial_trust_ratio", 1.0)

diff --git a/desc/optimize/least_squares.py b/desc/optimize/least_squares.py
@@ -3,7 +3,7 @@
 from scipy.optimize import OptimizeResult
 
 from desc.backend import jnp, qr
-from desc.utils import errorif, setdefault
+from desc.utils import errorif, safediv, setdefault
 
 from .bound_utils import (
     cl_scaling_vector,
@@ -208,14 +208,12 @@ def lsqtr(  # noqa: C901
     # conngould : norm of the cauchy point, as recommended in ch17 of Conn & Gould
     # scipy : norm of the scaled x, as used in scipy
     # mix : geometric mean of conngould and scipy
+    tr_scipy = jnp.linalg.norm(x * scale_inv / v**0.5)
+    conngould = safediv(jnp.sum(g_h**2), jnp.sum((J_h @ g_h) ** 2))
     init_tr = {
-        "scipy": jnp.linalg.norm(x * scale_inv / v**0.5),
-        "conngould": jnp.sum(g_h**2) / jnp.sum((J_h @ g_h) ** 2),
-        "mix": jnp.sqrt(
-            jnp.sum(g_h**2)
-            / jnp.sum((J_h @ g_h) ** 2)
-            * jnp.linalg.norm(x * scale_inv / v**0.5)
-        ),
+        "scipy": tr_scipy,
+        "conngould": conngould,
+        "mix": jnp.sqrt(conngould * tr_scipy),
     }
     trust_radius = options.pop("initial_trust_radius", "scipy")
     tr_ratio = options.pop("initial_trust_ratio", 1.0)