Allow combination of ESS and automatic scaling

desc/equilibrium/equilibrium.py

@@ -2243,7 +2243,7 @@ def solve(
             stopping tolerances. `None` will use defaults for given optimizer.
         maxiter : int
             Maximum number of solver steps.
-        x_scale : array, list[dict | ``'ess'``], ``'ess'`` or ``'auto'``, optional
+        x_scale : array, list[dict | ``'ess'``, ``'auto'``], ``'ess'`` or ``'auto'``
             Characteristic scale of each variable. Setting ``x_scale`` is equivalent
             to reformulating the problem in scaled variables ``xs = x / x_scale``.
             An alternative view is that the size of a trust region along jth
@@ -2253,17 +2253,20 @@ def solve(
             function. Default is ``'auto'``, which iteratively updates the scale using
             the inverse norms of the columns of the Jacobian or Hessian matrix.
             If set to ``'ess'``, the scale is set using Exponential Spectral Scaling,
-            this scaling is set with two parameters, ``ess_alpha`` and ``ess_order``
-            which are passed through ``options``. ``ess_alpha`` is the decay rate of
-            the scaling, and ``ess_order`` is the norm order for multi-index modes,
-            which can be ``1``, ``2``, or ``np.inf``. If not provided in ``options``,
-            the defaults are: ``ess_alpha=1.2``, ``ess_order=np.inf'`` and
-            ``ess_min_value=1e-7`` (minimum allowed scale value). If an array, should
-            be the same size as sum(thing.dim_x for thing in things). If a list, the
-            list should have 1 element for each thing, and each element should either
-            be ``'ess'`` to use exponential spectral scaling for that thing, or a dict
+            this scaling is set with parameters, ``ess_alpha``, ``ess_order``,
+            ``ess_min_value`` and ``ess_default`` which are passed through ``options``.
+            ``ess_alpha`` is the decay rate of the scaling, ``ess_order`` is the norm
+            order for multi-index modes, which can be ``1``, ``2``, or ``np.inf``.
+            ``ess_min_value`` is the minimum allowed scale value, and ``ess_default``
+            sets the default scale for variables without an ess rule defined.
+            If not provided in ``options``, the defaults are: ``ess_alpha=1.2``,
+            ``ess_order=np.inf'``, ``ess_min_value=1e-7``, ``ess_default=0.0``.
+            If an array, should be the same size as sum(thing.dim_x for thing in
+            things). If a list, the list should have 1 element for each thing, and
+            each element should either be ``'ess'``, ``'auto'`` to use exponential
+            spectral scaling or automatic jacobian scaling for that thing, or a dict
             with the same keys and dimensions as thing.params_dict to specify scales
-            manually.
+            manually. Anywhere ``x_scale==0``, automatic jacobian scaling will be used.
         options : dict
             Dictionary of additional options to pass to optimizer.
         verbose : int
@@ -2361,7 +2364,7 @@ def optimize(
             stopping tolerances. `None` will use defaults for given optimizer.
         maxiter : int
             Maximum number of solver steps.
-        x_scale : array, list[dict | ``'ess'``], ``'ess'`` or ``'auto'``, optional
+        x_scale : array, list[dict | ``'ess'``, ``'auto'``], ``'ess'`` or ``'auto'``
             Characteristic scale of each variable. Setting ``x_scale`` is equivalent
             to reformulating the problem in scaled variables ``xs = x / x_scale``.
             An alternative view is that the size of a trust region along jth
@@ -2371,17 +2374,20 @@ def optimize(
             function. Default is ``'auto'``, which iteratively updates the scale using
             the inverse norms of the columns of the Jacobian or Hessian matrix.
             If set to ``'ess'``, the scale is set using Exponential Spectral Scaling,
-            this scaling is set with two parameters, ``ess_alpha`` and ``ess_order``
-            which are passed through ``options``. ``ess_alpha`` is the decay rate of
-            the scaling, and ``ess_order`` is the norm order for multi-index modes,
-            which can be ``1``, ``2``, or ``np.inf``. If not provided in ``options``,
-            the defaults are: ``ess_alpha=1.2``, ``ess_order=np.inf'`` and
-            ``ess_min_value=1e-7`` (minimum allowed scale value). If an array, should
-            be the same size as sum(thing.dim_x for thing in things). If a list, the
-            list should have 1 element for each thing, and each element should either
-            be ``'ess'`` to use exponential spectral scaling for that thing, or a dict
+            this scaling is set with parameters, ``ess_alpha``, ``ess_order``,
+            ``ess_min_value`` and ``ess_default`` which are passed through ``options``.
+            ``ess_alpha`` is the decay rate of the scaling, ``ess_order`` is the norm
+            order for multi-index modes, which can be ``1``, ``2``, or ``np.inf``.
+            ``ess_min_value`` is the minimum allowed scale value, and ``ess_default``
+            sets the default scale for variables without an ess rule defined.
+            If not provided in ``options``, the defaults are: ``ess_alpha=1.2``,
+            ``ess_order=np.inf'``, ``ess_min_value=1e-7``, ``ess_default=0.0``.
+            If an array, should be the same size as sum(thing.dim_x for thing in
+            things). If a list, the list should have 1 element for each thing, and
+            each element should either be ``'ess'``, ``'auto'`` to use exponential
+            spectral scaling or automatic jacobian scaling for that thing, or a dict
             with the same keys and dimensions as thing.params_dict to specify scales
-            manually.
+            manually. Anywhere ``x_scale==0``, automatic jacobian scaling will be used.
         options : dict
             Dictionary of additional options to pass to optimizer.
         verbose : int
@@ -2515,7 +2521,7 @@ def perturb(
 
         return eq
 
-    def _get_ess_scale(self, alpha=1.2, order=np.inf, min_value=1e-7):
+    def _get_ess_scale(self, alpha=1.2, order=np.inf, min_value=1e-7, default=0.0):
         """Create x_scale using exponential spectral scaling.
 
         Parameters
@@ -2530,16 +2536,19 @@ def _get_ess_scale(self, alpha=1.2, order=np.inf, min_value=1e-7):
             Default is 'np.inf'
         min_value : float, optional
             Minimum allowed scale value. Default is 1e-7
+        default : float, optional
+            Default scale for variables that don't have an ess rule defined. 0 means
+            use automatic jacobian scaling.
 
         Returns
         -------
         dict of ndarray
             Array of scale values for each parameter
         """
         # this is the all ones scale:
-        scales = super()._get_ess_scale(alpha, order, min_value)
-        bdry_scale = self.surface._get_ess_scale(alpha, order, min_value)
-        axis_scale = self.axis._get_ess_scale(alpha, order, min_value)
+        scales = super()._get_ess_scale(alpha, order, min_value, default)
+        bdry_scale = self.surface._get_ess_scale(alpha, order, min_value, default)
+        axis_scale = self.axis._get_ess_scale(alpha, order, min_value, default)
         # we use ESS for the following:
         modes = {
             "R_lmn": self.R_basis.modes,

desc/geometry/curve.py

@@ -334,7 +334,7 @@ def from_values(cls, coords, N=10, NFP=1, sym=False, basis="rpz", name=""):
             name=name,
         )
 
-    def _get_ess_scale(self, alpha=1.2, order=np.inf, min_value=1e-7):
+    def _get_ess_scale(self, alpha=1.2, order=np.inf, min_value=1e-7, default=0.0):
         """Create x_scale using exponential spectral scaling.
 
         Parameters
@@ -349,14 +349,17 @@ def _get_ess_scale(self, alpha=1.2, order=np.inf, min_value=1e-7):
             Default is 'np.inf'
         min_value : float, optional
             Minimum allowed scale value. Default is 1e-7
+        default : float, optional
+            Default scale for variables that don't have an ess rule defined. 0 means
+            use automatic jacobian scaling.
 
         Returns
         -------
         dict of ndarray
             Array of scale values for each parameter
         """
         # this is the base class scale:
-        scales = super()._get_ess_scale(alpha, order, min_value)
+        scales = super()._get_ess_scale(alpha, order, min_value, default)
         # we use ESS for the following:
         modes = {"R_n": self.R_basis.modes, "Z_n": self.Z_basis.modes}
         scales.update(get_ess_scale(modes, alpha, order, min_value))
@@ -631,7 +634,7 @@ def from_values(cls, coords, N=10, s=None, basis="xyz", name=""):
             X_n=X_n, Y_n=Y_n, Z_n=Z_n, modes=basis.modes[:, 2], name=name
         )
 
-    def _get_ess_scale(self, alpha=1.2, order=np.inf, min_value=1e-7):
+    def _get_ess_scale(self, alpha=1.2, order=np.inf, min_value=1e-7, default=0.0):
         """Create x_scale using exponential spectral scaling.
 
         Parameters
@@ -646,14 +649,17 @@ def _get_ess_scale(self, alpha=1.2, order=np.inf, min_value=1e-7):
             Default is 'np.inf'
         min_value : float, optional
             Minimum allowed scale value. Default is 1e-7
+        default : float, optional
+            Default scale for variables that don't have an ess rule defined. 0 means
+            use automatic jacobian scaling.
 
         Returns
         -------
         dict of ndarray
             Array of scale values for each parameter
         """
         # this is the base class scale:
-        scales = super()._get_ess_scale(alpha, order, min_value)
+        scales = super()._get_ess_scale(alpha, order, min_value, default)
         # we use ESS for the following:
         modes = {
             "X_n": self.X_basis.modes,
@@ -979,7 +985,7 @@ def from_values(cls, coords, N=10, basis="xyz", name=""):
             name=name,
         )
 
-    def _get_ess_scale(self, alpha=1.2, order=np.inf, min_value=1e-7):
+    def _get_ess_scale(self, alpha=1.2, order=np.inf, min_value=1e-7, default=0.0):
         """Create x_scale using exponential spectral scaling.
 
         Parameters
@@ -994,14 +1000,17 @@ def _get_ess_scale(self, alpha=1.2, order=np.inf, min_value=1e-7):
             Default is 'np.inf'
         min_value : float, optional
             Minimum allowed scale value. Default is 1e-7
+        default : float, optional
+            Default scale for variables that don't have an ess rule defined. 0 means
+            use automatic jacobian scaling.
 
         Returns
         -------
         dict of ndarray
             Array of scale values for each parameter
         """
         # this is the base class scale:
-        scales = super()._get_ess_scale(alpha, order, min_value)
+        scales = super()._get_ess_scale(alpha, order, min_value, default)
         # we use ESS for the following:
         modes = {"r_n": self.r_basis.modes}
         scales.update(get_ess_scale(modes, alpha, order, min_value))
@@ -1384,7 +1393,7 @@ def from_values(cls, coords, N=10, s=None, basis="xyz", name=""):
             name=name,
         )
 
-    def _get_ess_scale(self, alpha=1.2, order=np.inf, min_value=1e-7):
+    def _get_ess_scale(self, alpha=1.2, order=np.inf, min_value=1e-7, default=0.0):
         """Create x_scale using exponential spectral scaling.
 
         Parameters
@@ -1399,14 +1408,17 @@ def _get_ess_scale(self, alpha=1.2, order=np.inf, min_value=1e-7):
             Default is 'np.inf'
         min_value : float, optional
             Minimum allowed scale value. Default is 1e-7
+        default : float, optional
+            Default scale for variables that don't have an ess rule defined. 0 means
+            use automatic jacobian scaling.
 
         Returns
         -------
         dict of ndarray
             Array of scale values for each parameter
         """
         # this is the base class scale:
-        scales = super()._get_ess_scale(alpha, order, min_value)
+        scales = super()._get_ess_scale(alpha, order, min_value, default)
         # we use ESS for the following:
         modes = {"X_n": self.X_basis.modes, "Y_n": self.Y_basis.modes}
         scales.update(get_ess_scale(modes, alpha, order, min_value))

desc/geometry/surface.py

@@ -837,7 +837,7 @@ def get_axis(self):
         )
         return axis
 
-    def _get_ess_scale(self, alpha=1.2, order=np.inf, min_value=1e-7):
+    def _get_ess_scale(self, alpha=1.2, order=np.inf, min_value=1e-7, default=0.0):
         """Create x_scale using exponential spectral scaling.
 
         Parameters
@@ -852,14 +852,17 @@ def _get_ess_scale(self, alpha=1.2, order=np.inf, min_value=1e-7):
             Default is 'np.inf'
         min_value : float, optional
             Minimum allowed scale value. Default is 1e-7
+        default : float, optional
+            Default scale for variables that don't have an ess rule defined. 0 means
+            use automatic jacobian scaling.
 
         Returns
         -------
         dict of ndarray
             Array of scale values for each parameter
         """
         # this is the base class scale:
-        scales = super()._get_ess_scale(alpha, order, min_value)
+        scales = super()._get_ess_scale(alpha, order, min_value, default)
         # we use ESS for the following:
         modes = {"R_lmn": self.R_basis.modes, "Z_lmn": self.Z_basis.modes}
         scales.update(get_ess_scale(modes, alpha, order, min_value))
@@ -1177,7 +1180,7 @@ def get_axis(self):
         axis = FourierRZCurve(R_n=data["R"][0], Z_n=data["Z"][0], sym=self.sym)
         return axis
 
-    def _get_ess_scale(self, alpha=1.2, order=np.inf, min_value=1e-7):
+    def _get_ess_scale(self, alpha=1.2, order=np.inf, min_value=1e-7, default=0.0):
         """Create x_scale using exponential spectral scaling.
 
         Parameters
@@ -1192,14 +1195,17 @@ def _get_ess_scale(self, alpha=1.2, order=np.inf, min_value=1e-7):
             Default is 'np.inf'
         min_value : float, optional
             Minimum allowed scale value. Default is 1e-7
+        default : float, optional
+            Default scale for variables that don't have an ess rule defined. 0 means
+            use automatic jacobian scaling.
 
         Returns
         -------
         dict of ndarray
             Array of scale values for each parameter
         """
         # this is the base class scale:
-        scales = super()._get_ess_scale(alpha, order, min_value)
+        scales = super()._get_ess_scale(alpha, order, min_value, default)
         # we use ESS for the following:
         modes = {"R_lmn": self.R_basis.modes, "Z_lmn": self.Z_basis.modes}
         scales.update(get_ess_scale(modes, alpha, order, min_value))

desc/magnetic_fields/_current_potential.py

@@ -1028,7 +1028,7 @@ def to_CoilSet(  # noqa: C901 - FIXME: simplify this
             final_coilset = CoilSet(*coils, check_intersection=check_intersection)
         return final_coilset
 
-    def _get_ess_scale(self, alpha=1.2, order=np.inf, min_value=1e-7):
+    def _get_ess_scale(self, alpha=1.2, order=np.inf, min_value=1e-7, default=0.0):
         """Create x_scale using exponential spectral scaling.
 
         Parameters
@@ -1043,14 +1043,17 @@ def _get_ess_scale(self, alpha=1.2, order=np.inf, min_value=1e-7):
             Default is 'np.inf'
         min_value : float, optional
             Minimum allowed scale value. Default is 1e-7
+        default : float, optional
+            Default scale for variables that don't have an ess rule defined. 0 means
+            use automatic jacobian scaling.
 
         Returns
         -------
         dict of ndarray
             Array of scale values for each parameter
         """
         # this is the base class scale:
-        scales = super()._get_ess_scale(alpha, order, min_value)
+        scales = super()._get_ess_scale(alpha, order, min_value, default)
         # we use ESS for the following, the R,Z scales are already in the base class:
         modes = {"Phi_mn": self.Phi_basis.modes}
         scales.update(get_ess_scale(modes, alpha, order, min_value))

desc/objectives/utils.py

@@ -121,17 +121,18 @@ def factorize_linear_constraints(objective, constraint, x_scale="auto"):  # noqa
     A, b, xp, unfixed_idx, fixed_idx = remove_fixed_parameters(A, b, xp)
 
     # compute x_scale if not provided
-    # Note: this x_scale is not the same as the x_scale as in solve_options["x_scale"]
-    # but the one given as solve_options["linear_constraint_options"]["x_scale"]
+    x0 = objective.x(*objective.things)
+    auto_x_scale = np.where(np.abs(x0) < 1e2, 1, np.abs(x0))
     if x_scale == "auto":
-        x_scale = objective.x(*objective.things)
+        x_scale = auto_x_scale
     errorif(
         x_scale.shape != xp.shape,
         ValueError,
         "x_scale must be the same size as the full state vector. "
         + f"Got size {x_scale.size} for state vector of size {xp.size}.",
     )
-    D = np.where(np.abs(x_scale) < 1e2, 1, np.abs(x_scale))
+    # x_scale==0 means use auto scale, otherwise use user scale
+    D = np.where(x_scale == 0, auto_x_scale, x_scale)
 
     # null space & particular solution
     A = A * D[None, unfixed_idx]

desc/optimizable.py

@@ -126,7 +126,7 @@
         """
         return sorted(set(list(args)))
 
-    def _get_ess_scale(self, alpha=1.2, order=np.inf, min_value=1e-7):
+    def _get_ess_scale(self, alpha=1.2, order=np.inf, min_value=1e-7, default=0.0):
         """Create x_scale using exponential spectral scaling.
 
         Parameters
@@ -141,6 +141,9 @@
             Default is 'np.inf'
         min_value : float, optional
             Minimum allowed scale value. Default is 1e-7
+        default : float, optional
+            Default scale for variables that don't have an ess rule defined. 0 means
+            use automatic jacobian scaling.
 
         Returns
         -------
@@ -149,7 +152,7 @@
         """
         # we don't know anything about the object so just assume scale is all 1s.
         # subclasses can implement their own logic.
-        return tree_map(jnp.ones_like, self.params_dict)
+        return tree_map(lambda x: default * jnp.ones_like(x), self.params_dict)
 
 
 class OptimizableCollection(Optimizable):
@@ -233,7 +236,7 @@
         params = [s.unpack_params(xi) for s, xi in zip(self, xs)]
         return params
 
-    def _get_ess_scale(self, alpha=1.2, order=np.inf, min_value=1e-7):
+    def _get_ess_scale(self, alpha=1.2, order=np.inf, min_value=1e-7, default=0.0):
         """Create x_scale using exponential spectral scaling.
 
         Parameters
@@ -248,13 +251,16 @@
             Default is 'np.inf'
         min_value : float, optional
             Minimum allowed scale value. Default is 1e-7
+        default : float, optional
+            Default scale for variables that don't have an ess rule defined. 0 means
+            use automatic jacobian scaling.
 
         Returns
         -------
         list of dict of ndarray
             Array of scale values for each parameter
         """
-        return [s._get_ess_scale(alpha, order, min_value) for s in self]
+        return [s._get_ess_scale(alpha, order, min_value, default) for s in self]
 
 
 def optimizable_parameter(f):