Merge pull request #742 from PlasmaControl/dp/optimizable_general

Allow Objectives to target Objects other than Equilibrium

CHANGELOG.md

@@ -1,6 +1,11 @@
 Changelog
 =========
 
+- Allows certain objectives to target ``FourierRZToroidalSurface`` objects as well as ``Equilibrium`` objects,
+such as ``MeanCurvature``, ``MeanCurvature``, and ``Volume``.
+- Allow optimizations where the only object being optimized is not an ``Equilibrium`` object e.g.
+optimizing only a ``FourierRZToroidalSurface`` object to have a certain ``Volume``.
+
 v0.10.3
 -------
 

desc/objectives/_geometry.py

@@ -16,6 +16,8 @@
 from .utils import softmin
 
 
+# TODO: add Surface parametrization to compute R0/a
+# so can use this objective with FourierRZToroidalSurface
 class AspectRatio(_Objective):
     """Aspect ratio = major radius / minor radius.
 
@@ -302,8 +304,9 @@ class Volume(_Objective):
 
     Parameters
     ----------
-    eq : Equilibrium
-        Equilibrium that will be optimized to satisfy the Objective.
+    eq : Equilibrium or FourierRZToroidalSurface
+        Equilibrium or FourierRZToroidalSurface that
+        will be optimized to satisfy the Objective.
     target : {float, ndarray}, optional
         Target value(s) of the objective. Only used if bounds is None.
         Must be broadcastable to Objective.dim_f.
@@ -381,7 +384,23 @@ def build(self, use_jit=True, verbose=1):
         """
         eq = self.things[0]
         if self._grid is None:
-            grid = QuadratureGrid(L=eq.L_grid, M=eq.M_grid, N=eq.N_grid, NFP=eq.NFP)
+            if hasattr(eq, "L_grid"):
+                grid = QuadratureGrid(
+                    L=eq.L_grid,
+                    M=eq.M_grid,
+                    N=eq.N_grid,
+                    NFP=eq.NFP,
+                )
+            else:
+                # if not an Equilibrium, is a Surface,
+                # has no radial resolution so just need
+                # the surface points
+                grid = LinearGrid(
+                    rho=1.0,
+                    M=eq.M * 2,
+                    N=eq.N * 2,
+                    NFP=eq.NFP,
+                )
         else:
             grid = self._grid
 
@@ -416,7 +435,8 @@ def compute(self, params, constants=None):
         Parameters
         ----------
         params : dict
-            Dictionary of equilibrium degrees of freedom, eg Equilibrium.params_dict
+            Dictionary of equilibrium or surface degrees of freedom,
+            eg Equilibrium.params_dict
         constants : dict
             Dictionary of constant data, eg transforms, profiles etc. Defaults to
             self.constants
@@ -430,7 +450,7 @@ def compute(self, params, constants=None):
         if constants is None:
             constants = self.constants
         data = compute_fun(
-            "desc.equilibrium.equilibrium.Equilibrium",
+            self.things[0],
             self._data_keys,
             params=params,
             transforms=constants["transforms"],
@@ -716,8 +736,9 @@ class MeanCurvature(_Objective):
 
     Parameters
     ----------
-    eq : Equilibrium
-        Equilibrium that will be optimized to satisfy the Objective.
+    eq : Equilibrium or FourierRZToroidalSurface
+        Equilibrium or FourierRZToroidalSurface that
+        will be optimized to satisfy the Objective.
     target : {float, ndarray}, optional
         Target value(s) of the objective. Only used if bounds is None.
         Must be broadcastable to Objective.dim_f.
@@ -794,7 +815,12 @@ def build(self, use_jit=True, verbose=1):
         """
         eq = self.things[0]
         if self._grid is None:
-            grid = LinearGrid(M=eq.M_grid, N=eq.N_grid, NFP=eq.NFP, sym=eq.sym)
+            grid = LinearGrid(  # getattr statements in case a surface is passed in
+                M=getattr(eq, "M_grid", eq.M * 2),
+                N=getattr(eq, "N_grid", eq.N * 2),
+                NFP=eq.NFP,
+                sym=eq.sym,
+            )
         else:
             grid = self._grid
 
@@ -829,7 +855,8 @@ def compute(self, params, constants=None):
         Parameters
         ----------
         params : dict
-            Dictionary of equilibrium degrees of freedom, eg Equilibrium.params_dict
+            Dictionary of equilibrium or surface degrees of freedom,
+            eg Equilibrium.params_dict
         constants : dict
             Dictionary of constant data, eg transforms, profiles etc. Defaults to
             self.constants
@@ -843,7 +870,7 @@ def compute(self, params, constants=None):
         if constants is None:
             constants = self.constants
         data = compute_fun(
-            "desc.equilibrium.equilibrium.Equilibrium",
+            self.things[0],
             self._data_keys,
             params=params,
             transforms=constants["transforms"],
@@ -866,8 +893,9 @@ class PrincipalCurvature(_Objective):
 
     Parameters
     ----------
-    eq : Equilibrium
-        Equilibrium that will be optimized to satisfy the Objective.
+    eq : Equilibrium or FourierRZToroidalSurface
+        Equilibrium or FourierRZToroidalSurface that
+        will be optimized to satisfy the Objective.
     target : {float, ndarray}, optional
         Target value(s) of the objective. Only used if bounds is None.
         Must be broadcastable to Objective.dim_f.
@@ -944,7 +972,12 @@ def build(self, use_jit=True, verbose=1):
         """
         eq = self.things[0]
         if self._grid is None:
-            grid = LinearGrid(M=eq.M_grid, N=eq.N_grid, NFP=eq.NFP, sym=eq.sym)
+            grid = LinearGrid(  # getattr statements in case a surface is passed in
+                M=getattr(eq, "M_grid", eq.M * 2),
+                N=getattr(eq, "N_grid", eq.N * 2),
+                NFP=eq.NFP,
+                sym=eq.sym,
+            )
         else:
             grid = self._grid
 
@@ -979,7 +1012,8 @@ def compute(self, params, constants=None):
         Parameters
         ----------
         params : dict
-            Dictionary of equilibrium degrees of freedom, eg Equilibrium.params_dict
+            Dictionary of equilibrium or surface degrees of freedom,
+            eg Equilibrium.params_dict
         constants : dict
             Dictionary of constant data, eg transforms, profiles etc. Defaults to
             self.constants
@@ -993,7 +1027,7 @@ def compute(self, params, constants=None):
         if constants is None:
             constants = self.constants
         data = compute_fun(
-            "desc.equilibrium.equilibrium.Equilibrium",
+            self.things[0],
             self._data_keys,
             params=params,
             transforms=constants["transforms"],

desc/objectives/normalization.py

@@ -4,33 +4,51 @@
 from scipy.constants import elementary_charge, mu_0
 
 
-def compute_scaling_factors(eq):
+# TODO: generalize this so that it will make different normalizations depending on the
+# given thing (not always an equilibrium, if just a surface)
+def compute_scaling_factors(thing):
     """Compute dimensional quantities for normalizations."""
+    # local import to avoid circular import
+    from desc.equilibrium import Equilibrium
+    from desc.geometry import FourierRZToroidalSurface
+
     scales = {}
-    R10 = eq.Rb_lmn[eq.surface.R_basis.get_idx(M=1, N=0)]
-    Z10 = eq.Zb_lmn[eq.surface.Z_basis.get_idx(M=-1, N=0)]
-    R00 = eq.Rb_lmn[eq.surface.R_basis.get_idx(M=0, N=0)]
-
-    scales["R0"] = R00
-    scales["a"] = np.sqrt(np.abs(R10 * Z10))
-    scales["A"] = np.pi * scales["a"] ** 2
-    scales["V"] = 2 * np.pi * scales["R0"] * scales["A"]
-    scales["B_T"] = abs(eq.Psi) / scales["A"]
-    iota_avg = np.mean(np.abs(eq.get_profile("iota")(np.linspace(0, 1, 20))))
-    if np.isclose(iota_avg, 0):
-        scales["B_P"] = scales["B_T"]
-    else:
-        scales["B_P"] = scales["B_T"] * iota_avg
-    scales["B"] = np.sqrt(scales["B_T"] ** 2 + scales["B_P"] ** 2)
-    scales["I"] = scales["B_P"] * 2 * np.pi / mu_0
-    scales["p"] = scales["B"] ** 2 / (2 * mu_0)
-    scales["W"] = scales["p"] * scales["V"]
-    scales["J"] = scales["B"] / scales["a"] / mu_0
-    scales["F"] = scales["p"] / scales["a"]
-    scales["f"] = scales["F"] * scales["V"]
-    scales["Psi"] = abs(eq.Psi)
-    scales["n"] = 1e19
-    scales["T"] = scales["p"] / (scales["n"] * elementary_charge)
+
+    if isinstance(thing, Equilibrium):
+        R10 = thing.Rb_lmn[thing.surface.R_basis.get_idx(M=1, N=0)]
+        Z10 = thing.Zb_lmn[thing.surface.Z_basis.get_idx(M=-1, N=0)]
+        R00 = thing.Rb_lmn[thing.surface.R_basis.get_idx(M=0, N=0)]
+
+        scales["R0"] = R00
+        scales["a"] = np.sqrt(np.abs(R10 * Z10))
+        scales["A"] = np.pi * scales["a"] ** 2
+        scales["V"] = 2 * np.pi * scales["R0"] * scales["A"]
+        scales["B_T"] = abs(thing.Psi) / scales["A"]
+        iota_avg = np.mean(np.abs(thing.get_profile("iota")(np.linspace(0, 1, 20))))
+        if np.isclose(iota_avg, 0):
+            scales["B_P"] = scales["B_T"]
+        else:
+            scales["B_P"] = scales["B_T"] * iota_avg
+        scales["B"] = np.sqrt(scales["B_T"] ** 2 + scales["B_P"] ** 2)
+        scales["I"] = scales["B_P"] * 2 * np.pi / mu_0
+        scales["p"] = scales["B"] ** 2 / (2 * mu_0)
+        scales["W"] = scales["p"] * scales["V"]
+        scales["J"] = scales["B"] / scales["a"] / mu_0
+        scales["F"] = scales["p"] / scales["a"]
+        scales["f"] = scales["F"] * scales["V"]
+        scales["Psi"] = abs(thing.Psi)
+        scales["n"] = 1e19
+        scales["T"] = scales["p"] / (scales["n"] * elementary_charge)
+    elif isinstance(thing, FourierRZToroidalSurface):
+        R10 = thing.R_lmn[thing.R_basis.get_idx(M=1, N=0)]
+        Z10 = thing.Z_lmn[thing.Z_basis.get_idx(M=-1, N=0)]
+        R00 = thing.R_lmn[thing.R_basis.get_idx(M=0, N=0)]
+
+        scales["R0"] = R00
+        scales["a"] = np.sqrt(np.abs(R10 * Z10))
+        scales["A"] = np.pi * scales["a"] ** 2
+        scales["V"] = 2 * np.pi * scales["R0"] * scales["A"]
+
     # replace 0 scales to avoid normalizing by zero
     for scale in scales.keys():
         if np.isclose(scales[scale], 0):

desc/optimize/optimizer.py

@@ -136,6 +136,8 @@ def optimize(  # noqa: C901 - FIXME: simplify this
 
         Returns
         -------
+        things : list,
+            list of optimized things
         res : OptimizeResult
             The optimization result represented as a ``OptimizeResult`` object.
             Important attributes are: ``x`` the solution array, ``success`` a
@@ -167,7 +169,7 @@ def optimize(  # noqa: C901 - FIXME: simplify this
         )
 
         # make sure everything is built
-        if not objective.built:
+        if objective is not None and not objective.built:
             objective.build(verbose=verbose)
         if nonlinear_constraint is not None and not nonlinear_constraint.built:
             nonlinear_constraint.build(verbose=verbose)
@@ -392,7 +394,7 @@ def _maybe_wrap_nonlinear_constraints(
 ):
     """Use ProximalProjection to handle nonlinear constraints."""
     if eq is None:  # not deal with an equilibrium problem -> no ProximalProjection
-        return eq, nonlinear_constraint
+        return objective, nonlinear_constraint
     wrapper, method = _parse_method(method)
     if nonlinear_constraint is None:
         if wrapper is not None:

docs/notebooks/tutorials/advanced_optimization.ipynb

@@ -1222,7 +1222,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.10.11"
+   "version": "3.8.10"
   }
  },
  "nbformat": 4,

tests/test_examples.py

@@ -29,8 +29,10 @@
     ForceBalance,
     ForceBalanceAnisotropic,
     HelicalForceBalance,
+    MeanCurvature,
     ObjectiveFunction,
     PlasmaVesselDistance,
+    PrincipalCurvature,
     QuasisymmetryBoozer,
     QuasisymmetryTwoTerm,
     RadialForceBalance,
@@ -807,6 +809,72 @@ def test_multiobject_optimization_prox():
     np.testing.assert_allclose(eq.i_l, [2, 0, 0])
 
 
+@pytest.mark.unit
+def test_non_eq_optimization():
+    """Test for optimizing a non-eq object by fixing all eq parameters."""
+    eq = desc.examples.get("DSHAPE")
+    Rmax = 4
+    Rmin = 2
+
+    a = 2
+    R0 = (Rmax + Rmin) / 2
+    surf = FourierRZToroidalSurface(
+        R_lmn=[R0, a],
+        Z_lmn=[0.0, -a],
+        modes_R=np.array([[0, 0], [1, 0]]),
+        modes_Z=np.array([[0, 0], [-1, 0]]),
+        sym=True,
+        NFP=eq.NFP,
+    )
+
+    surf.change_resolution(M=eq.M, N=eq.N)
+    constraints = (
+        FixParameter(eq),
+        MeanCurvature(surf, bounds=(-8, 8)),
+        PrincipalCurvature(surf, bounds=(0, 15)),
+    )
+
+    grid = LinearGrid(M=18, N=0, NFP=eq.NFP)
+    obj = PlasmaVesselDistance(
+        surface=surf,
+        eq=eq,
+        target=0.5,
+        use_softmin=True,
+        surface_grid=grid,
+        plasma_grid=grid,
+        alpha=5000,
+    )
+    objective = ObjectiveFunction((obj,))
+    optimizer = Optimizer("lsq-auglag")
+    (eq, surf), result = optimizer.optimize(
+        (eq, surf), objective, constraints, verbose=3, maxiter=100
+    )
+
+    np.testing.assert_allclose(obj.compute(*obj.xs(eq, surf)), 0.5, atol=1e-5)
+
+
+@pytest.mark.unit
+def test_only_non_eq_optimization():
+    """Test for optimizing only a non-eq object."""
+    eq = desc.examples.get("DSHAPE")
+    surf = eq.surface
+
+    surf.change_resolution(M=eq.M, N=eq.N)
+    constraints = (
+        FixParameter(surf, params="R_lmn", indices=surf.R_basis.get_idx(0, 0, 0)),
+    )
+
+    obj = PrincipalCurvature(surf, target=1)
+
+    objective = ObjectiveFunction((obj,))
+    optimizer = Optimizer("lsq-exact")
+    (surf), result = optimizer.optimize(
+        (surf), objective, constraints, verbose=3, maxiter=100
+    )
+    surf = surf[0]
+    np.testing.assert_allclose(obj.compute(*obj.xs(surf)), 1, atol=1e-5)
+
+
 class TestGetExample:
     """Tests for desc.examples.get."""