Repository: PlasmaControl/DESC

Issue:
Scaling with quadrature weights messes up profile constraints
I'm running an optimization, using `RotationalTransform` as a constraint:
`RotationalTransform(eq=eqfam[-1], bounds=(0.41, 0.43)),`

However, the resulting rotational transform I get is more like 1.7, the issue is seemingly because the profile objectives get multiplied by `sqrt(w)`, which in this case is `0.242`, so effectively its targeting iota of `0.4/0/242 ~ 1.7`

When used as a constraint, we use `Objective.compute_scaled` because we often want to deal with constraints in normalized units to keep everything a similar order of magnitude. However, this also applies quadrature weights which may mess up bounds/targets for non-scalar objectives with non-zero targets/bounds.

It's not an issue if we use it as an objective, because in that case we call `Objective.compute_scaled_error` which applies the target/bounds to the unscaled values, then applies the scaling

Possibly related to #484 ?

Discussion:

Patch:
diff --git a/desc/objectives/objective_funs.py b/desc/objectives/objective_funs.py
index bc030c557..ba9f5f642 100644
--- a/desc/objectives/objective_funs.py
+++ b/desc/objectives/objective_funs.py
@@ -707,8 +707,9 @@ class ObjectiveFunction(IOAble):
             else:
                 # need to return something, so use midpoint of bounds as approx target
                 target_i = jnp.ones(obj.dim_f) * (obj.bounds[0] + obj.bounds[1]) / 2
-            if obj._normalize_target:
-                target_i /= obj.normalization
+            target_i = obj._scale(target_i)
+            if not obj._normalize_target:
+                target_i *= obj.normalization
             target += [target_i]
         return jnp.concatenate(target)
 
@@ -723,9 +724,11 @@ class ObjectiveFunction(IOAble):
             else:
                 lb_i = jnp.ones(obj.dim_f) * obj.target
                 ub_i = jnp.ones(obj.dim_f) * obj.target
-            if obj._normalize_target:
-                lb_i /= obj.normalization
-                ub_i /= obj.normalization
+            lb_i = obj._scale(lb_i)
+            ub_i = obj._scale(ub_i)
+            if not obj._normalize_target:
+                lb_i *= obj.normalization
+                ub_i *= obj.normalization
             lb += [lb_i]
             ub += [ub_i]
         return (jnp.concatenate(lb), jnp.concatenate(ub))


Test patch:
diff --git a/tests/test_objective_funs.py b/tests/test_objective_funs.py
index 9e12fde36..d35c0629f 100644
--- a/tests/test_objective_funs.py
+++ b/tests/test_objective_funs.py
@@ -846,8 +846,8 @@ def test_objective_target_bounds():
     """Test that the target_scaled and bounds_scaled etc. return the right things."""
     eq = Equilibrium()
 
-    vol = Volume(target=3, normalize=True, eq=eq)
-    asp = AspectRatio(bounds=(2, 3), normalize=False, eq=eq)
+    vol = Volume(target=3, normalize=True, weight=2, eq=eq)
+    asp = AspectRatio(bounds=(2, 3), normalize=False, weight=3, eq=eq)
     fbl = ForceBalance(normalize=True, bounds=(-1, 2), weight=5, eq=eq)
 
     objective = ObjectiveFunction((vol, asp, fbl))
@@ -857,21 +857,30 @@ def test_objective_target_bounds():
     bounds = objective.bounds_scaled
     weight = objective.weights
 
-    assert bounds[0][0] == 3 / vol.normalization
-    assert bounds[1][0] == 3 / vol.normalization
-    assert bounds[0][1] == 2
-    assert bounds[1][1] == 3
-    assert np.all(bounds[0][2:] == -1 / fbl.normalization)
-    assert np.all(bounds[1][2:] == 2 / fbl.normalization)
+    assert bounds[0][0] == 3 / vol.normalization * vol.weight
+    assert bounds[1][0] == 3 / vol.normalization * vol.weight
+    assert bounds[0][1] == 2 * asp.weight
+    assert bounds[1][1] == 3 * asp.weight
+    assert np.all(bounds[0][2:] == -1 / fbl.normalization * fbl.weight)
+    assert np.all(bounds[1][2:] == 2 / fbl.normalization * fbl.weight)
 
-    assert target[0] == 3 / vol.normalization
-    assert target[1] == 2.5
-    assert np.all(target[2:] == 0.5 / fbl.normalization)
+    assert target[0] == 3 / vol.normalization * vol.weight
+    assert target[1] == 2.5 * asp.weight
+    assert np.all(target[2:] == 0.5 / fbl.normalization * fbl.weight)
 
-    assert weight[0] == 1
-    assert weight[1] == 1
+    assert weight[0] == 2
+    assert weight[1] == 3
     assert np.all(weight[2:] == 5)
 
+    eq = Equilibrium(L=8, M=2, N=2, iota=PowerSeriesProfile(0.42))
+
+    con = ObjectiveFunction(RotationalTransform(eq=eq, bounds=(0.41, 0.43)))
+    con.build()
+
+    np.testing.assert_allclose(con.compute_scaled_error(con.x(eq)), 0)
+    np.testing.assert_array_less(con.bounds_scaled[0], con.compute_scaled(con.x(eq)))
+    np.testing.assert_array_less(con.compute_scaled(con.x(eq)), con.bounds_scaled[1])
+
 
 @pytest.mark.unit
 def test_softmax_and_softmin():