Fix np bug in LinearConstraintProjection

CHANGELOG.md

@@ -27,6 +27,7 @@ Bug Fixes
 
 - No longer uses the full Hessian to compute the scale when ``x_scale="auto"`` and using a scipy optimizer that approximates the hessian (e.g. if using ``"scipy-bfgs"``, no longer attempts the Hessian computation to get the x_scale).
 - ``SplineMagneticField.from_field()`` correctly uses the ``NFP`` input when given. Also adds this as a similar input option to ``MagneticField.save_mgrid()``.
+- Fixes bug that was overriding some components of the user supplied ``linear_constraint_options["x_scale"]``. Now, the given value is used without any alterations.
 
 Performance Improvements
 

desc/objectives/utils.py

@@ -123,15 +123,17 @@
     # 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"]
-    if x_scale == "auto":
+    if isinstance(x_scale, str) and x_scale == "auto":
         x_scale = objective.x(*objective.things)
+        D = np.where(np.abs(x_scale) < 1e2, 1, np.abs(x_scale))
+    else:
+        D = 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))
 
     # null space & particular solution
     A = A * D[None, unfixed_idx]

tests/test_linear_objectives.py

@@ -1233,3 +1233,33 @@ def test_NAE_asym_with_sym_axis():
     conZ.build()
     assert conR._A.shape[0] == conR.dim_f
     assert conZ._A.shape[0] == conZ.dim_f
+
+
+@pytest.mark.unit
+def test_linearconstraintprojection_xscale():
+    """Test that giving x_scale to LCP works."""
+    eq = desc.examples.get("DSHAPE")
+    with pytest.warns(UserWarning, match="Reducing radial"):
+        eq.change_resolution(L=2, M=2, L_grid=4, M_grid=4)
+    # use standard fixed boundary constraints as test
+    cons = get_fixed_boundary_constraints(eq)
+    cons = maybe_add_self_consistency(eq, cons)
+    con = ObjectiveFunction(cons)
+    obj = ObjectiveFunction(ForceBalance(eq))
+
+    # np and jnp arrays behaves differently for
+    # arr == "auto" kind of conditions. Numpy requires
+    # a type check, Jax gives False automatically
+    lcp_scale_np = np.ones(eq.dim_x)
+    lcp_scale_np[eq.x_idx["L_lmn"]] *= 5
+    lcp_scale_jnp = jnp.array(lcp_scale_np)
+
+    lcp = LinearConstraintProjection(obj, con, x_scale=lcp_scale_np)
+    lcp.build()
+
+    lcp2 = LinearConstraintProjection(obj, con, x_scale=lcp_scale_jnp)
+    lcp2.build()
+
+    # LCP shouldn't change the given x_scale
+    assert (lcp._D == lcp_scale_np).all()
+    assert (lcp2._D == lcp_scale_np).all()