Advanced QS Optimization

In this tutorial we will show an example of “precise” QS optimization using a multigrid approach, and using constrained optimization

[1]:
import sys
import os

sys.path.insert(0, os.path.abspath("."))
sys.path.append(os.path.abspath("../../../"))

If you have access to a GPU, uncomment the following two lines before any DESC or JAX related imports. You should see about an order of magnitude speed improvement with only these two lines of code!

[ ]:
# from desc import set_device
# set_device("gpu")

As mentioned in DESC Documentation on performance tips, one can use compilation cache directory to reduce the compilation overhead time. Note: One needs to create jax-caches folder manually.

[3]:
# import jax

# jax.config.update("jax_compilation_cache_dir", "../jax-caches")
# jax.config.update("jax_persistent_cache_min_entry_size_bytes", -1)
# jax.config.update("jax_persistent_cache_min_compile_time_secs", 0)
[4]:
import numpy as np

from desc.continuation import solve_continuation_automatic
from desc.equilibrium import EquilibriaFamily, Equilibrium
from desc.geometry import FourierRZToroidalSurface
from desc.grid import LinearGrid
from desc.objectives import (
    AspectRatio,
    FixBoundaryR,
    FixBoundaryZ,
    FixCurrent,
    FixPressure,
    FixPsi,
    ForceBalance,
    ObjectiveFunction,
    QuasisymmetryTwoTerm,
    GenericObjective,
    ObjectiveFromUser,
)
from desc.optimize import Optimizer

Initial Guess

We start by creating an initial equilibrium and solving a standard fixed boundary problem:

[9]:
# create initial surface. Aspect ratio ~ 8, circular cross section with slight
# axis torsion to make it nonplanar
surf = FourierRZToroidalSurface(
    R_lmn=[1, 0.125, 0.1],
    Z_lmn=[-0.125, -0.1],
    modes_R=[[0, 0], [1, 0], [0, 1]],
    modes_Z=[[-1, 0], [0, -1]],
    NFP=4,
)
# create initial equilibrium. Psi chosen to give B ~ 1 T. Could also give profiles here,
# default is zero pressure and zero current
eq = Equilibrium(M=4, N=4, Psi=0.04, surface=surf)
# this is usually all you need to solve a fixed boundary equilibrium
eq0 = solve_continuation_automatic(eq, verbose=0)[-1]

# it will be helpful to store intermediate results
eqfam = EquilibriaFamily(eq0)

Multigrid method with proximal optimizer

By “multigrid” method we mean we will start by optimizing over boundary modes with \(|m|, |n| \leq 1\), then \(|m|, |n| \leq 2\) and so on. To do this we’ll define a helper function that will create the necessary constraints and objectives for a given maximum mode number \(k\).

By a “proximal” optimizer we mean one that handles the equilibrium constraint by re-solving a fixed boundary equilibrium problem at each step, given the current position of the boundary. This is made more efficient by using a perturbed estimate based on the previous step as a warm start to the equilibrium sub-problem.

[10]:
def run_qh_step(k, eq):
    """Run a step of the precise QH optimization example from Landreman & Paul."""
    # this step will only optimize boundary modes with |m|,|n| <= k

    # create grid where we want to minimize QS error. Here we do it on 3 surfaces
    grid = LinearGrid(
        M=eq.M_grid, N=eq.N_grid, NFP=eq.NFP, rho=np.array([0.6, 0.8, 1.0]), sym=True
    )

    # we create an ObjectiveFunction, in this case made up of multiple objectives
    # which will be combined in a least squares sense
    objective = ObjectiveFunction(
        (
            # pass in the grid we defined, and don't forget the target helicity!
            QuasisymmetryTwoTerm(eq=eq, helicity=(1, eq.NFP), grid=grid),
            # try to keep the aspect ratio about the same
            AspectRatio(eq=eq, target=8, weight=100),
        ),
    )
    # as opposed to SIMSOPT and STELLOPT where variables are assumed fixed, in DESC
    # we assume variables are free. Here we decide which ones to fix, starting with
    # the major radius (R mode = [0,0,0]) and all modes with m,n > k
    R_modes = np.vstack(
        (
            [0, 0, 0],
            eq.surface.R_basis.modes[
                np.max(np.abs(eq.surface.R_basis.modes), 1) > k, :
            ],
        )
    )
    Z_modes = eq.surface.Z_basis.modes[
        np.max(np.abs(eq.surface.Z_basis.modes), 1) > k, :
    ]
    # next we create the constraints, using the mode number arrays just created
    # if we didn't pass those in, it would fix all the modes (like for the profiles)
    constraints = (
        ForceBalance(eq=eq),
        FixBoundaryR(eq=eq, modes=R_modes),
        FixBoundaryZ(eq=eq, modes=Z_modes),
        FixPressure(eq=eq),
        FixCurrent(eq=eq),
        FixPsi(eq=eq),
    )
    # this is the default optimizer, which re-solves the equilibrium at each step
    optimizer = Optimizer("proximal-lsq-exact")

    eq_new, history = eq.optimize(
        objective=objective,
        constraints=constraints,
        optimizer=optimizer,
        maxiter=20,  # we don't need to solve to optimality at each multigrid step
        verbose=3,
        copy=True,  # don't modify original, return a new optimized copy
        options={
            # Sometimes the default initial trust radius is too big, allowing the
            # optimizer to take too large a step in a bad direction. If this happens,
            # we can manually specify a smaller starting radius. Each optimizer has a
            # number of different options that can be used to tune the performance.
            # See the documentation for more info.
            "initial_trust_ratio": 0.1,
        },
    )

    return eq_new

Lets look at the initial field:

[11]:
from desc.plotting import plot_boozer_surface

plot_boozer_surface(eq0);
../../_images/notebooks_tutorials_advanced_optimization_15_0.png

We see that it is vaguely QH like, which is why we’re targeting QH symmetry. Now let’s run the optimization in steps and look at the intermediate result after each step:

[12]:
eq1 = run_qh_step(1, eq0)
eqfam.append(eq1)
plot_boozer_surface(eq1);
Building objective: QS two-term
Precomputing transforms
Timer: Precomputing transforms = 56.3 ms
Building objective: aspect ratio
Precomputing transforms
Timer: Precomputing transforms = 25.8 ms
Timer: Objective build = 102 ms
Building objective: force
Precomputing transforms
Timer: Precomputing transforms = 61.0 ms
Timer: Objective build = 78.1 ms
Timer: Objective build = 992 us
Timer: Eq Update LinearConstraintProjection build = 84.3 ms
Timer: Proximal projection build = 1.30 sec
Building objective: lcfs R
Building objective: lcfs Z
Building objective: fixed pressure
Building objective: fixed current
Building objective: fixed Psi
Timer: Objective build = 174 ms
Timer: LinearConstraintProjection build = 151 ms
Number of parameters: 8
Number of objectives: 460
Timer: Initializing the optimization = 1.67 sec

Starting optimization
Using method: proximal-lsq-exact
Solver options:
------------------------------------------------------------
Maximum Function Evaluations       : 101
Maximum Allowed Total Δx Norm      : inf
Scaled Termination                 : True
Trust Region Method                : qr
Initial Trust Radius               : 5.668e+01
Maximum Trust Radius               : inf
Minimum Trust Radius               : 2.220e-16
Trust Radius Increase Ratio        : 2.000e+00
Trust Radius Decrease Ratio        : 2.500e-01
Trust Radius Increase Threshold    : 7.500e-01
Trust Radius Decrease Threshold    : 2.500e-01
------------------------------------------------------------

   Iteration     Total nfev        Cost      Cost reduction    Step norm     Optimality
       0              1          2.949e+01                                    2.852e+00
       1              5          2.723e+01      2.259e+00      1.231e-02      1.226e+00
       2              6          2.582e+01      1.419e+00      4.602e-03      8.559e-01
       3              7          2.368e+01      2.133e+00      1.095e-02      6.697e-01
       4              8          2.071e+01      2.973e+00      2.264e-02      3.671e-01
       5              9          1.831e+01      2.397e+00      4.321e-02      4.348e-01
       6             10          1.739e+01      9.241e-01      2.409e-02      1.122e+00
       7             11          1.647e+01      9.203e-01      3.711e-02      1.547e-01
       8             13          1.587e+01      6.004e-01      1.712e-02      2.024e-01
       9             14          1.540e+01      4.630e-01      1.068e-02      2.417e-01
      10             15          1.502e+01      3.895e-01      1.666e-02      1.288e+00
      11             16          1.338e+01      1.636e+00      2.821e-02      7.097e-01
      12             18          1.251e+01      8.681e-01      8.719e-03      2.368e-01
      13             19          1.148e+01      1.036e+00      1.253e-02      7.064e-01
      14             21          1.068e+01      8.000e-01      1.095e-02      1.158e-01
      15             22          9.909e+00      7.674e-01      1.152e-02      5.918e-01
      16             23          9.685e+00      2.233e-01      2.157e-02      1.677e+00
      17             24          7.983e+00      1.702e+00      1.411e-02      1.112e-01
      18             25          7.403e+00      5.802e-01      1.141e-02      1.259e-01
      19             26          7.289e+00      1.134e-01      2.362e-02      7.641e-01
      20             27          6.777e+00      5.121e-01      8.160e-03      1.169e-01
Warning: Maximum number of iterations has been exceeded.
         Current function value: 6.777e+00
         Total delta_x: 1.590e-01
         Iterations: 20
         Function evaluations: 27
         Jacobian evaluations: 21
Timer: Solution time = 37.8 sec
Timer: Avg time per step = 1.80 sec
==============================================================================================================
                                                                 Start  -->   End
Total (sum of squares):                                      2.946e+01  -->   6.777e+00,
Maximum absolute Quasi-symmetry (1,4) two-term error:        6.229e-01  -->   5.480e-01 (T^3)
Minimum absolute Quasi-symmetry (1,4) two-term error:        1.964e-04  -->   4.448e-04 (T^3)
Average absolute Quasi-symmetry (1,4) two-term error:        1.501e-01  -->   7.104e-02 (T^3)
Maximum absolute Quasi-symmetry (1,4) two-term error:        5.894e-01  -->   5.186e-01 (normalized)
Minimum absolute Quasi-symmetry (1,4) two-term error:        1.859e-04  -->   4.209e-04 (normalized)
Average absolute Quasi-symmetry (1,4) two-term error:        1.420e-01  -->   6.722e-02 (normalized)
Aspect ratio:                                                8.000e+00  -->   7.998e+00 (dimensionless)
Maximum absolute Force error:                                4.342e+01  -->   1.616e+03 (N)
Minimum absolute Force error:                                4.390e-03  -->   2.552e-02 (N)
Average absolute Force error:                                3.601e+00  -->   8.598e+01 (N)
Maximum absolute Force error:                                4.262e-05  -->   1.587e-03 (normalized)
Minimum absolute Force error:                                4.310e-09  -->   2.506e-08 (normalized)
Average absolute Force error:                                3.535e-06  -->   8.441e-05 (normalized)
R boundary error:                                            0.000e+00  -->   0.000e+00 (m)
Z boundary error:                                            0.000e+00  -->   0.000e+00 (m)
Fixed pressure profile error:                                0.000e+00  -->   0.000e+00 (Pa)
Fixed current profile error:                                 0.000e+00  -->   0.000e+00 (A)
Fixed Psi error:                                             0.000e+00  -->   0.000e+00 (Wb)
==============================================================================================================

../../_images/notebooks_tutorials_advanced_optimization_17_1.png
[13]:
eq2 = run_qh_step(2, eq1)
eqfam.append(eq2)
plot_boozer_surface(eq2);
Building objective: QS two-term
Precomputing transforms
Timer: Precomputing transforms = 57.7 ms
Building objective: aspect ratio
Precomputing transforms
Timer: Precomputing transforms = 26.3 ms
Timer: Objective build = 103 ms
Building objective: force
Precomputing transforms
Timer: Precomputing transforms = 65.0 ms
Timer: Objective build = 81.9 ms
Timer: Objective build = 1.03 ms
Timer: Eq Update LinearConstraintProjection build = 102 ms
Timer: Proximal projection build = 1.46 sec
Building objective: lcfs R
Building objective: lcfs Z
Building objective: fixed pressure
Building objective: fixed current
Building objective: fixed Psi
Timer: Objective build = 436 ms
Timer: LinearConstraintProjection build = 1.54 sec
Number of parameters: 24
Number of objectives: 460
Timer: Initializing the optimization = 3.54 sec

Starting optimization
Using method: proximal-lsq-exact
Solver options:
------------------------------------------------------------
Maximum Function Evaluations       : 101
Maximum Allowed Total Δx Norm      : inf
Scaled Termination                 : True
Trust Region Method                : qr
Initial Trust Radius               : 8.710e+01
Maximum Trust Radius               : inf
Minimum Trust Radius               : 2.220e-16
Trust Radius Increase Ratio        : 2.000e+00
Trust Radius Decrease Ratio        : 2.500e-01
Trust Radius Increase Threshold    : 7.500e-01
Trust Radius Decrease Threshold    : 2.500e-01
------------------------------------------------------------

   Iteration     Total nfev        Cost      Cost reduction    Step norm     Optimality
       0              1          4.034e+01                                    1.093e+00
       1              4          3.562e+01      4.723e+00      6.933e-03      8.447e-01
       2              5          3.307e+01      2.554e+00      7.166e-03      6.279e-01
       3              6          2.983e+01      3.237e+00      6.908e-03      6.533e-01
       4              7          2.727e+01      2.560e+00      6.849e-03      6.366e-01
       5              8          2.529e+01      1.984e+00      6.607e-03      6.035e-01
       6              9          2.349e+01      1.793e+00      6.135e-03      5.429e-01
       7             10          2.052e+01      2.975e+00      5.482e-03      3.802e-01
       8             11          1.464e+01      5.883e+00      9.564e-03      4.578e-01
       9             12          1.021e+01      4.424e+00      1.814e-02      2.212e+00
      10             13          3.050e+00      7.161e+00      2.698e-02      6.911e-01
      11             15          2.100e+00      9.504e-01      1.283e-02      2.600e-01
      12             16          1.511e+00      5.888e-01      1.635e-02      5.015e-01
      13             18          1.184e+00      3.270e-01      8.574e-03      1.844e-01
      14             20          1.079e+00      1.048e-01      4.419e-03      4.029e-02
      15             21          9.752e-01      1.042e-01      9.340e-03      1.672e-01
      16             23          9.096e-01      6.564e-02      5.453e-03      5.591e-02
      17             25          8.836e-01      2.594e-02      3.010e-03      2.369e-02
      18             26          8.593e-01      2.435e-02      6.420e-03      1.504e-01
      19             28          8.286e-01      3.070e-02      3.110e-03      4.059e-02
      20             29          8.166e-01      1.199e-02      6.792e-03      1.979e-01
Warning: Maximum number of iterations has been exceeded.
         Current function value: 8.166e-01
         Total delta_x: 1.065e-01
         Iterations: 20
         Function evaluations: 29
         Jacobian evaluations: 21
Timer: Solution time = 50.4 sec
Timer: Avg time per step = 2.40 sec
==============================================================================================================
                                                                 Start  -->   End
Total (sum of squares):                                      4.166e+01  -->   8.166e-01,
Maximum absolute Quasi-symmetry (1,4) two-term error:        5.480e-01  -->   1.114e-01 (T^3)
Minimum absolute Quasi-symmetry (1,4) two-term error:        4.448e-04  -->   2.858e-06 (T^3)
Average absolute Quasi-symmetry (1,4) two-term error:        7.104e-02  -->   1.119e-02 (T^3)
Maximum absolute Quasi-symmetry (1,4) two-term error:        1.287e+00  -->   2.616e-01 (normalized)
Minimum absolute Quasi-symmetry (1,4) two-term error:        1.044e-03  -->   6.711e-06 (normalized)
Average absolute Quasi-symmetry (1,4) two-term error:        1.668e-01  -->   2.627e-02 (normalized)
Aspect ratio:                                                7.998e+00  -->   7.999e+00 (dimensionless)
Maximum absolute Force error:                                1.616e+03  -->   1.274e+03 (N)
Minimum absolute Force error:                                2.552e-02  -->   2.167e-01 (N)
Average absolute Force error:                                8.598e+01  -->   9.647e+01 (N)
Maximum absolute Force error:                                2.500e-03  -->   1.970e-03 (normalized)
Minimum absolute Force error:                                3.947e-08  -->   3.352e-07 (normalized)
Average absolute Force error:                                1.330e-04  -->   1.492e-04 (normalized)
R boundary error:                                            0.000e+00  -->   0.000e+00 (m)
Z boundary error:                                            0.000e+00  -->   0.000e+00 (m)
Fixed pressure profile error:                                0.000e+00  -->   0.000e+00 (Pa)
Fixed current profile error:                                 0.000e+00  -->   0.000e+00 (A)
Fixed Psi error:                                             0.000e+00  -->   0.000e+00 (Wb)
==============================================================================================================

../../_images/notebooks_tutorials_advanced_optimization_18_1.png
[14]:
eq3 = run_qh_step(3, eq2)
eqfam.append(eq3)
plot_boozer_surface(eq3);
Building objective: QS two-term
Precomputing transforms
Timer: Precomputing transforms = 58.5 ms
Building objective: aspect ratio
Precomputing transforms
Timer: Precomputing transforms = 26.9 ms
Timer: Objective build = 104 ms
Building objective: force
Precomputing transforms
Timer: Precomputing transforms = 62.3 ms
Timer: Objective build = 80.4 ms
Timer: Objective build = 971 us
Timer: Eq Update LinearConstraintProjection build = 76.5 ms
Timer: Proximal projection build = 1.15 sec
Building objective: lcfs R
Building objective: lcfs Z
Building objective: fixed pressure
Building objective: fixed current
Building objective: fixed Psi
Timer: Objective build = 364 ms
Timer: LinearConstraintProjection build = 1.48 sec
Number of parameters: 48
Number of objectives: 460
Timer: Initializing the optimization = 3.12 sec

Starting optimization
Using method: proximal-lsq-exact
Solver options:
------------------------------------------------------------
Maximum Function Evaluations       : 101
Maximum Allowed Total Δx Norm      : inf
Scaled Termination                 : True
Trust Region Method                : qr
Initial Trust Radius               : 9.336e+01
Maximum Trust Radius               : inf
Minimum Trust Radius               : 2.220e-16
Trust Radius Increase Ratio        : 2.000e+00
Trust Radius Decrease Ratio        : 2.500e-01
Trust Radius Increase Threshold    : 7.500e-01
Trust Radius Decrease Threshold    : 2.500e-01
------------------------------------------------------------

   Iteration     Total nfev        Cost      Cost reduction    Step norm     Optimality
       0              1          1.134e+00                                    2.279e-01
       1              3          1.095e+00      3.888e-02      1.263e-02      8.014e-01
       2              4          6.282e-01      4.668e-01      2.683e-03      2.024e-01
       3              5          4.388e-01      1.894e-01      2.806e-03      9.021e-02
       4              7          4.153e-01      2.345e-02      1.310e-03      3.544e-02
       5              8          4.096e-01      5.755e-03      2.519e-03      7.476e-02
       6              9          3.715e-01      3.810e-02      8.986e-04      2.562e-02
       7             10          3.671e-01      4.388e-03      1.384e-03      3.763e-02
       8             11          3.538e-01      1.329e-02      4.425e-04      1.416e-02
       9             12          3.493e-01      4.543e-03      7.308e-04      2.057e-02
      10             13          3.427e-01      6.537e-03      7.455e-04      1.431e-02
      11             14          3.372e-01      5.553e-03      7.441e-04      1.633e-02
      12             15          3.318e-01      5.315e-03      7.461e-04      1.676e-02
      13             16          3.267e-01      5.186e-03      7.478e-04      1.710e-02
      14             17          3.233e-01      3.400e-03      7.484e-04      1.538e-02
      15             18          3.184e-01      4.906e-03      7.457e-04      1.600e-02
      16             19          3.139e-01      4.448e-03      7.348e-04      1.694e-02
      17             20          3.093e-01      4.603e-03      7.309e-04      1.580e-02
      18             21          3.048e-01      4.486e-03      7.273e-04      1.542e-02
      19             22          3.006e-01      4.177e-03      7.231e-04      1.527e-02
      20             23          2.968e-01      3.822e-03      7.201e-04      1.512e-02
Warning: Maximum number of iterations has been exceeded.
         Current function value: 2.968e-01
         Total delta_x: 1.973e-02
         Iterations: 20
         Function evaluations: 23
         Jacobian evaluations: 21
Timer: Solution time = 42.3 sec
Timer: Avg time per step = 2.01 sec
==============================================================================================================
                                                                 Start  -->   End
Total (sum of squares):                                      1.133e+00  -->   2.968e-01,
Maximum absolute Quasi-symmetry (1,4) two-term error:        1.114e-01  -->   8.355e-02 (T^3)
Minimum absolute Quasi-symmetry (1,4) two-term error:        2.858e-06  -->   2.441e-06 (T^3)
Average absolute Quasi-symmetry (1,4) two-term error:        1.119e-02  -->   5.567e-03 (T^3)
Maximum absolute Quasi-symmetry (1,4) two-term error:        3.085e-01  -->   2.313e-01 (normalized)
Minimum absolute Quasi-symmetry (1,4) two-term error:        7.913e-06  -->   6.758e-06 (normalized)
Average absolute Quasi-symmetry (1,4) two-term error:        3.097e-02  -->   1.541e-02 (normalized)
Aspect ratio:                                                7.999e+00  -->   8.000e+00 (dimensionless)
Maximum absolute Force error:                                1.274e+03  -->   6.911e+02 (N)
Minimum absolute Force error:                                2.167e-01  -->   2.130e-02 (N)
Average absolute Force error:                                9.647e+01  -->   6.973e+01 (N)
Maximum absolute Force error:                                2.139e-03  -->   1.161e-03 (normalized)
Minimum absolute Force error:                                3.640e-07  -->   3.577e-08 (normalized)
Average absolute Force error:                                1.620e-04  -->   1.171e-04 (normalized)
R boundary error:                                            0.000e+00  -->   0.000e+00 (m)
Z boundary error:                                            0.000e+00  -->   0.000e+00 (m)
Fixed pressure profile error:                                0.000e+00  -->   0.000e+00 (Pa)
Fixed current profile error:                                 0.000e+00  -->   0.000e+00 (A)
Fixed Psi error:                                             0.000e+00  -->   0.000e+00 (Wb)
==============================================================================================================

../../_images/notebooks_tutorials_advanced_optimization_19_1.png

We see that after only 3 multigrid steps we have achieved very straight contours of magnetic field strength. These could be further refined by running for more iterations, using higher resolution, tighter tolerances, etc.

As a final comparison, we’ll look at the maximum symmetry breaking boozer harmonic for each step of the equilibrium

[15]:
import matplotlib.pyplot as plt
from desc.plotting import plot_boozer_modes, plot_boundaries

fig, ax = plt.subplots()
colors = ["r", "g", "c", "m"]

for i, (eq, color) in enumerate(zip(eqfam, colors)):
    plot_boozer_modes(
        eq, color=color, helicity=(1, eq.NFP), max_only=True, label=f"Step {i}", ax=ax
    );
../../_images/notebooks_tutorials_advanced_optimization_21_0.png

Constrained Optimization

Next, we’ll do a similar optimization but this time treating it as a constrained optimization problem, where we attempt to minimize QS error subject to more complicated constraints. We’ll start with the same QS objective:

[16]:
# create grid where we want to minimize QS error. Here we do it on 3 surfaces
grid = LinearGrid(
    M=eq0.M_grid, N=eq0.N_grid, NFP=eq0.NFP, rho=np.array([0.6, 0.8, 1.0]), sym=True
)

objective = ObjectiveFunction(
    (
        # we use a GenericObjective here for demonstration purposes, we could alternatively
        # have used `QuasisymmetryTwoTerm(eq=eq0, helicity=(1, eq.NFP), grid=grid)`,the already-existing objective for two-term QS
        GenericObjective(
            f="f_C",
            thing=eq0,
            # pass in the grid we defined
            grid=grid,
            # and don't forget the target helicity!
            # two-term QS is known in the data index as "f_C" and requires "helicity" as a kwarg,
            # we specify this for GenericObjective through compute_kwargs
            compute_kwargs={"helicity": (1, eq.NFP)},
            name="QS Two-Term",
        ),
    ),
)

For constraints, we’ll include the standard force balance to start. In the previous example, fixing certain boundary modes served as a form of regularization to prevent the solution from going into a bad local minimum. In this case however, instead of fixing a range of boundary modes we will only fix the \(R_{00}\) mode, and include constraints on aspect ratio, volume, and elongation to keep the solution from going off in a bad direction.

Finally, we also include a constraint on the average rotational transform, and on the mirror ratio at the surface (which we will manually compute to demonstrate using ObjectiveFromUser):

[17]:
from desc.objectives import Elongation, RotationalTransform, Volume, ObjectiveFromUser
from desc.integrals import surface_min, surface_max


# Mirror ratio, manually computed from "|B|"
def fun_mirror_ratio(grid, data):
    # compute max and min of |B| on each flux surface in the grid
    ## we have utility functions which, given a quantity computed on a grid,
    ## return the max or min of that quantity on each coordinate surface
    max_tz_B = surface_max(grid=grid, x=data["|B|"], surface_label="rho")
    min_tz_B = surface_min(grid=grid, x=data["|B|"], surface_label="rho")
    # to avoid issues with array shapes, these two arrays (max_tz_B and min_tz_B)
    # are still the same shape as data["|B|"] i.e. still are 1-D arrays of length grid.num_nodes,
    # (i.e. there are data corresponding to nodes (rho, theta, zeta) = (1.0, 0, 0) and (1.0, pi, 0),
    # which have the same max_tz_B). This is useful if we, for instance, wanted to  multiply
    # a flux-surface quantity like max_tz_B with a non-flux-surface quantity like sqrt(g)
    # without needing to worry about shape mismatches.

    # However, since we only need these quantities on each flux surface from
    # here on out, we can use the grid.compress function to reduce these quantities
    # down to just the values at each unique rho surface
    max_tz_B = grid.compress(max_tz_B, surface_label="rho")
    min_tz_B = grid.compress(min_tz_B, surface_label="rho")

    # now max_tz_B and min_tz_B are just 1-D arrays of size grid.num_rho

    # Finally, compute the mirror ratio using the above max/min on each flux surface
    mirror_ratio = (max_tz_B - min_tz_B) / (min_tz_B + max_tz_B)
    return mirror_ratio
    # alternatively, "mirror ratio" is something that can be computed in the data index
    # directly (see List of Variables docs), so can replace entirety of the above function code with this return statement
    # return grid.compress(data["mirror ratio"])
    # or can just use GenericObjective(f="mirror ratio", thing=eq) or the already-existing objective, MirrorRatio(eq)


obj_mirror_ratio = ObjectiveFromUser(
    fun=fun_mirror_ratio,
    thing=eq0,
    grid=LinearGrid(rho=1.0, M=eq.M_grid, N=eq.N_grid, NFP=eq.NFP),
    bounds=(0.18, 0.22),
    name="my mirror ratio",
)

constraints = (
    ForceBalance(eq=eq0),
    # try to keep the aspect ratio between 7 and 9
    AspectRatio(eq=eq0, bounds=(7, 9)),
    # similarly, try to keep it from getting too elongated
    Elongation(eq=eq0, bounds=(0, 3)),
    # Keep volume the same as the initial volume
    Volume(eq=eq0, target=eq0.compute("V")["V"]),
    # target for average iota
    RotationalTransform(eq=eq0, target=1.1, loss_function="mean"),
    # bounds for mirror ratio
    obj_mirror_ratio,
    # fix major radius
    FixBoundaryR(eq=eq0, modes=[0, 0, 0]),
    # fix vacuum profiles
    FixPressure(eq=eq0),
    FixCurrent(eq=eq0),
    FixPsi(eq=eq0),
)

Finally, we’ll use an optimizer that can handle general nonlinear constraints (the proximal-lsq-exact optimizer can only handle equilibrium constraints such as ForceBalance and regular linear constraints like Fix*). In this case we use a least-squares augmented Lagrangian method.

[18]:
optimizer = Optimizer("lsq-auglag")

eqa, history = eq0.optimize(
    objective=objective,
    constraints=constraints,
    optimizer=optimizer,
    # each iteration of the augmented Lagrangian optimizer is cheaper than a step of a
    # proximal optimizer, but it generally requires more iterations to converge
    maxiter=200,
    copy=True,
    verbose=3,
    options={},
)
Building objective: QS Two-Term
Timer: Objective build = 58.3 ms
Building objective: lcfs R
Building objective: fixed pressure
Building objective: fixed current
Building objective: fixed Psi
Building objective: self_consistency R
Building objective: self_consistency Z
Building objective: lambda gauge
Building objective: axis R self consistency
Building objective: axis Z self consistency
Timer: Objective build = 162 ms
Building objective: force
Precomputing transforms
Timer: Precomputing transforms = 67.0 ms
Building objective: aspect ratio
Precomputing transforms
Timer: Precomputing transforms = 25.9 ms
Building objective: elongation
Precomputing transforms
Timer: Precomputing transforms = 26.1 ms
Building objective: volume
Precomputing transforms
Timer: Precomputing transforms = 25.3 ms
Building objective: rotational transform
Precomputing transforms
Timer: Precomputing transforms = 716 ms
Building objective: my mirror ratio
Timer: Objective build = 1.66 sec
Timer: LinearConstraintProjection build = 2.11 sec
Timer: LinearConstraintProjection build = 99.5 ms
Number of parameters: 200
Number of objectives: 459
Number of equality constraints: 852
Number of inequality constraints: 3
Timer: Initializing the optimization = 4.62 sec

Starting optimization
Using method: lsq-auglag
Solver options:
------------------------------------------------------------
Maximum Function Evaluations       : 1001
Maximum Allowed Total Δx Norm      : inf
Scaled Termination                 : True
Trust Region Method                : qr
Initial Trust Radius               : 2.063e-01
Maximum Trust Radius               : inf
Minimum Trust Radius               : 2.220e-16
Trust Radius Increase Ratio        : 4.000e+00
Trust Radius Decrease Ratio        : 2.500e-01
Trust Radius Increase Threshold    : 7.500e-01
Trust Radius Decrease Threshold    : 5.000e-01
Alpha Omega                        : 1.000e+00
Beta Omega                         : 1.000e+00
Alpha Eta                          : 1.000e-01
Beta Eta                           : 9.000e-01
Omega                              : 1.000e-02
Eta                                : 1.000e-02
Tau                                : 9.000e-01
------------------------------------------------------------

   Iteration     Total nfev        Cost      Cost reduction    Step norm     Optimality    Constr viol.   Penalty param  max(|mltplr|)
       0              1          3.290e+01                                    7.512e+00      8.523e-01      1.000e+01      0.000e+00
       1              2          3.002e+01      2.884e+00      3.994e-03      7.097e+00      8.531e-01      1.000e+01      0.000e+00
       2              3          2.042e+01      9.597e+00      1.789e-02      5.589e+00      8.561e-01      1.000e+01      0.000e+00
       3              4          2.578e+00      1.784e+01      7.964e-02      1.547e+00      8.594e-01      1.000e+01      0.000e+00
       4              5          1.666e+00      9.120e-01      3.648e-01      9.408e-01      5.652e-01      1.000e+01      0.000e+00
       5              6          4.820e-01      1.184e+00      2.318e-01      2.535e-01      6.174e-01      1.000e+01      0.000e+00
       6              7          5.697e-01     -8.775e-02      3.411e-01      2.653e-01      4.944e-01      1.000e+01      0.000e+00
       7              9          4.685e-01      1.012e-01      2.445e-01      2.154e-01      4.668e-01      1.000e+01      0.000e+00
       8             10          4.190e-01      4.950e-02      2.197e-01      1.492e-01      4.222e-01      1.000e+01      0.000e+00
       9             11          2.683e-01      1.506e-01      8.007e-02      2.853e-02      4.355e-01      1.000e+01      0.000e+00
      10             12          2.252e-01      4.314e-02      5.635e-02      1.659e-01      4.026e-01      1.000e+01      0.000e+00
      11             13          2.366e-01     -1.139e-02      2.101e-01      2.607e+00      2.496e-01      1.000e+01      0.000e+00
      12             16          8.823e-02      1.484e-01      1.744e-01      1.995e-01      2.033e-01      1.000e+01      0.000e+00
      13             18          6.253e-02      2.569e-02      8.369e-02      2.145e-02      1.698e-01      1.000e+01      0.000e+00
      14             20          4.944e-02      1.309e-02      7.209e-02      2.199e-02      1.347e-01      1.000e+01      0.000e+00
      15             22          4.363e-02      5.807e-03      7.704e-02      2.400e-02      1.059e-01      1.000e+01      0.000e+00
      16             24          4.015e-02      3.485e-03      6.892e-02      2.300e-02      8.266e-02      1.000e+01      0.000e+00
      17             26          3.755e-02      2.602e-03      6.304e-02      2.035e-02      6.626e-02      1.000e+01      0.000e+00
      18             28          3.498e-02      2.571e-03      5.208e-02      3.000e-02      5.485e-02      1.000e+01      0.000e+00
      19             30          3.193e-02      3.048e-03      4.268e-02      2.856e-02      4.595e-02      1.000e+01      0.000e+00
      20             32          2.923e-02      2.697e-03      3.282e-02      2.294e-02      3.902e-02      1.000e+01      0.000e+00
      21             34          2.663e-02      2.601e-03      2.609e-02      1.836e-02      3.410e-02      1.000e+01      0.000e+00
      22             36          2.447e-02      2.163e-03      2.206e-02      1.807e-02      3.029e-02      1.000e+01      0.000e+00
      23             38          2.259e-02      1.873e-03      2.013e-02      1.700e-02      2.733e-02      1.000e+01      0.000e+00
      24             40          2.086e-02      1.734e-03      1.929e-02      1.642e-02      2.480e-02      1.000e+01      0.000e+00
      25             42          1.941e-02      1.449e-03      1.935e-02      1.554e-02      2.265e-02      1.000e+01      0.000e+00
      26             44          1.803e-02      1.383e-03      1.997e-02      1.534e-02      2.066e-02      1.000e+01      0.000e+00
      27             46          1.688e-02      1.147e-03      2.118e-02      2.575e-02      1.892e-02      1.000e+01      0.000e+00
      28             48          1.573e-02      1.148e-03      2.350e-02      1.501e-02      1.726e-02      1.000e+01      0.000e+00
      29             50          1.477e-02      9.644e-04      2.344e-02      1.607e-02      1.628e-02      1.000e+01      0.000e+00
      30             52          1.383e-02      9.359e-04      2.404e-02      2.249e-02      1.576e-02      1.000e+01      0.000e+00
      31             54          1.299e-02      8.423e-04      2.717e-02      1.902e-02      1.520e-02      1.000e+01      0.000e+00
      32             56          1.222e-02      7.730e-04      2.720e-02      2.091e-02      1.462e-02      1.000e+01      0.000e+00
      33             58          1.150e-02      7.130e-04      2.761e-02      2.240e-02      1.404e-02      1.000e+01      0.000e+00
      34             60          1.083e-02      6.743e-04      2.773e-02      2.327e-02      1.357e-02      1.000e+01      0.000e+00
      35             62          1.080e-02      3.351e-05      2.712e-02      2.536e-02      1.776e-02      1.000e+01      0.000e+00
      36             63          9.390e-03      1.408e-03      2.722e-02      2.119e-02      1.354e-02      1.000e+01      0.000e+00
      37             65          9.399e-03     -9.048e-06      2.599e-02      2.380e-02      1.658e-02      1.000e+01      0.000e+00
      38             66          8.067e-03      1.332e-03      2.556e-02      1.733e-02      1.355e-02      1.000e+01      0.000e+00
      39             68          7.956e-03      1.109e-04      2.391e-02      1.985e-02      1.563e-02      1.000e+01      0.000e+00
      40             69          6.775e-03      1.181e-03      2.363e-02      1.293e-02      1.348e-02      1.000e+01      0.000e+00
      41             71          6.537e-03      2.375e-04      2.175e-02      1.577e-02      1.494e-02      1.000e+01      0.000e+00
      42             73          5.566e-03      9.710e-04      2.144e-02      9.391e-03      1.367e-02      1.000e+01      0.000e+00
      43             75          5.251e-03      3.155e-04      2.032e-02      1.233e-02      1.430e-02      1.000e+01      0.000e+00
      44             77          4.532e-03      7.190e-04      1.878e-02      5.660e-03      1.407e-02      1.000e+01      0.000e+00
      45             79          4.271e-03      2.606e-04      1.865e-02      8.909e-03      1.428e-02      1.000e+01      0.000e+00
      46             81          3.879e-03      3.925e-04      1.557e-02      8.191e-03      1.502e-02      1.000e+01      0.000e+00
      47             83          3.616e-03      2.628e-04      1.624e-02      7.937e-03      1.417e-02      1.000e+01      0.000e+00
      48             85          3.307e-03      3.089e-04      1.400e-02      7.147e-03      1.496e-02      1.000e+01      0.000e+00
      49             87          3.211e-03      9.612e-05      1.286e-02      1.053e-02      1.499e-02      1.000e+01      0.000e+00
      50             89          2.979e-03      2.320e-04      1.141e-02      8.149e-03      1.505e-02      1.000e+01      0.000e+00
      51             91          2.879e-03      9.981e-05      1.091e-02      1.095e-02      1.542e-02      1.000e+01      0.000e+00
      52             93          2.715e-03      1.638e-04      9.967e-03      7.512e-03      1.542e-02      1.000e+01      0.000e+00
      53             95          2.590e-03      1.253e-04      9.469e-03      9.941e-03      1.549e-02      1.000e+01      0.000e+00
      54             97          2.485e-03      1.049e-04      8.950e-03      6.841e-03      1.559e-02      1.000e+01      0.000e+00
      55             99          2.366e-03      1.190e-04      8.592e-03      8.605e-03      1.551e-02      1.000e+01      0.000e+00
      56             101         2.295e-03      7.114e-05      9.146e-03      6.396e-03      1.522e-02      1.000e+01      0.000e+00
      57             103         2.188e-03      1.069e-04      8.589e-03      8.903e-03      1.536e-02      1.000e+01      0.000e+00
      58             104         2.147e-03      4.069e-05      8.792e-03      7.278e-03      1.519e-02      1.000e+01      0.000e+00
      59             106         2.057e-03      9.014e-05      9.037e-03      8.932e-03      1.541e-02      1.000e+01      0.000e+00
      60             107         2.051e-03      6.651e-06      9.072e-03      8.449e-03      1.508e-02      1.000e+01      0.000e+00
      61             108         1.930e-03      1.208e-04      9.044e-03      8.807e-03      1.516e-02      1.000e+01      0.000e+00
      62             109         1.991e-03     -6.069e-05      9.399e-03      1.056e-02      1.502e-02      1.000e+01      0.000e+00
      63             110         1.868e-03      1.229e-04      9.907e-03      9.287e-03      1.494e-02      1.000e+01      0.000e+00
      64             111         1.919e-03     -5.103e-05      9.639e-03      1.156e-02      1.483e-02      1.000e+01      0.000e+00
      65             112         1.776e-03      1.425e-04      1.017e-02      1.060e-02      1.454e-02      1.000e+01      0.000e+00
      66             113         1.810e-03     -3.416e-05      9.531e-03      1.147e-02      1.448e-02      1.000e+01      0.000e+00
      67             114         1.692e-03      1.188e-04      1.050e-02      1.122e-02      1.402e-02      1.000e+01      0.000e+00
      68             115         1.724e-03     -3.204e-05      9.713e-03      1.178e-02      1.408e-02      1.000e+01      0.000e+00
      69             116         1.598e-03      1.251e-04      1.076e-02      1.142e-02      1.335e-02      1.000e+01      0.000e+00
      70             117         1.612e-03     -1.319e-05      9.931e-03      1.201e-02      1.356e-02      1.000e+01      0.000e+00
      71             118         1.502e-03      1.098e-04      1.092e-02      1.159e-02      1.257e-02      1.000e+01      0.000e+00
      72             119         1.487e-03      1.510e-05      9.884e-03      1.323e-02      1.297e-02      1.000e+01      0.000e+00
      73             120         1.407e-03      8.022e-05      1.057e-02      1.136e-02      1.169e-02      1.000e+01      0.000e+00
      74             121         1.358e-03      4.881e-05      1.016e-02      1.259e-02      1.226e-02      1.000e+01      0.000e+00
      75             122         1.290e-03      6.759e-05      1.055e-02      1.084e-02      1.073e-02      1.000e+01      0.000e+00
      76             124         1.236e-03      5.370e-05      1.021e-02      1.236e-02      1.143e-02      1.000e+01      0.000e+00
      77             125         1.180e-03      5.682e-05      1.034e-02      1.036e-02      9.769e-03      1.000e+01      0.000e+00
      78             127         1.099e-03      8.032e-05      9.565e-03      1.135e-02      1.002e-02      1.000e+01      0.000e+00
      79             128         1.061e-03      3.878e-05      9.942e-03      1.004e-02      8.854e-03      1.000e+01      0.000e+00
      80             130         9.844e-04      7.618e-05      8.775e-03      1.051e-02      9.044e-03      1.000e+01      0.000e+00
      81             131         9.483e-04      3.612e-05      9.135e-03      8.726e-03      7.958e-03      1.000e+01      0.000e+00
      82             133         8.646e-04      8.362e-05      8.764e-03      7.923e-03      8.103e-03      1.000e+01      0.000e+00
      83             135         8.313e-04      3.334e-05      1.027e-02      6.531e-03      7.061e-03      1.000e+01      0.000e+00
      84             137         7.559e-04      7.539e-05      1.016e-02      5.440e-03      7.140e-03      1.000e+01      0.000e+00
      85             139         7.253e-04      3.065e-05      9.639e-03      4.505e-03      6.155e-03      1.000e+01      0.000e+00
      86             141         6.622e-04      6.306e-05      9.532e-03      3.458e-03      6.191e-03      1.000e+01      0.000e+00
      87             143         6.344e-04      2.775e-05      8.880e-03      3.854e-03      5.260e-03      1.000e+01      0.000e+00
      88             145         5.860e-04      4.843e-05      8.712e-03      4.525e-03      5.270e-03      1.000e+01      0.000e+00
      89             147         5.699e-04      1.613e-05      7.942e-03      5.521e-03      4.463e-03      1.000e+01      0.000e+00
      90             149         5.424e-04      2.748e-05      7.562e-03      5.707e-03      4.618e-03      1.000e+01      0.000e+00
      91             150         5.549e-04     -1.249e-05      6.699e-03      7.536e-03      4.170e-03      1.000e+01      0.000e+00
      92             152         5.429e-04      1.199e-05      6.659e-03      8.280e-03      4.074e-03      1.000e+01      0.000e+00
      93             153         5.714e-04     -2.845e-05      7.289e-03      8.983e-03      3.925e-03      1.000e+01      0.000e+00
      94             154         6.059e-04     -3.454e-05      7.340e-03      1.153e-02      3.821e-03      1.000e+01      0.000e+00
      95             155         6.467e-04     -4.083e-05      9.185e-03      1.391e-02      3.698e-03      1.000e+01      0.000e+00
      96             156         3.802e-04      2.665e-04      2.789e-03      5.100e-03      3.669e-03      1.000e+01      3.669e-02
      97             157         1.129e-03     -7.492e-04      1.091e-02      2.824e-02      3.411e-03      1.000e+01      3.669e-02
      98             158         9.289e-04      2.005e-04      6.960e-03      9.830e-04      3.314e-03      1.000e+01      3.669e-02
      99             159         1.056e-03     -1.275e-04      1.603e-02      1.388e-02      3.267e-03      1.000e+01      3.669e-02
      100            160         9.250e-04      1.314e-04      4.536e-03      9.379e-04      3.156e-03      1.000e+01      3.669e-02
      101            161         1.121e-03     -1.958e-04      1.631e-02      1.400e-02      3.131e-03      1.000e+01      3.669e-02
      102            162         9.298e-04      1.910e-04      4.667e-03      9.119e-04      3.008e-03      1.000e+01      3.669e-02
      103            163         1.176e-03     -2.458e-04      1.690e-02      1.491e-02      3.008e-03      1.000e+01      3.669e-02
      104            164         9.480e-04      2.276e-04      5.003e-03      8.811e-04      2.933e-03      1.000e+01      3.669e-02
      105            165         1.277e-03     -3.287e-04      1.800e-02      1.587e-02      2.917e-03      1.000e+01      3.669e-02
      106            166         9.758e-04      3.008e-04      5.688e-03      9.685e-04      2.888e-03      1.000e+01      3.669e-02
      107            168         9.826e-04     -6.794e-06      5.474e-03      9.985e-04      2.893e-03      1.000e+01      3.669e-02
      108            170         9.885e-04     -5.864e-06      5.471e-03      1.016e-03      2.900e-03      1.000e+01      3.669e-02
      109            172         9.938e-04     -5.325e-06      5.474e-03      1.026e-03      2.907e-03      1.000e+01      3.669e-02
      110            174         9.996e-04     -5.740e-06      5.465e-03      1.081e-03      2.911e-03      1.000e+01      3.669e-02
      111            176         1.006e-03     -6.619e-06      5.446e-03      1.115e-03      2.912e-03      1.000e+01      3.669e-02
      112            177         1.364e-03     -3.582e-04      2.025e-02      1.687e-02      2.840e-03      1.000e+01      3.669e-02
      113            178         1.029e-03      3.350e-04      5.829e-03      1.130e-03      2.825e-03      1.000e+01      3.669e-02
      114            179         1.232e-03     -2.027e-04      2.130e-02      1.419e-02      2.726e-03      1.000e+01      3.669e-02
      115            180         1.043e-03      1.894e-04      5.777e-03      9.490e-04      2.695e-03      1.000e+01      3.669e-02
      116            181         1.154e-03     -1.114e-04      2.151e-02      3.154e-02      4.587e-03      1.000e+01      3.669e-02
      117            182         1.046e-03      1.083e-04      8.989e-03      9.471e-04      2.549e-03      1.000e+01      3.669e-02
      118            183         1.104e-03     -5.775e-05      2.513e-02      1.136e-02      2.441e-03      1.000e+01      3.669e-02
      119            184         1.090e-03      1.321e-05      2.432e-02      1.144e-02      2.397e-03      1.000e+01      3.669e-02
      120            185         1.071e-03      1.956e-05      2.407e-02      1.125e-02      2.355e-03      1.000e+01      3.669e-02
      121            186         1.053e-03      1.786e-05      2.374e-02      1.117e-02      2.296e-03      1.000e+01      3.669e-02
      122            187         1.035e-03      1.826e-05      2.341e-02      1.261e-02      2.281e-03      1.000e+01      3.669e-02
      123            188         1.029e-03      5.248e-06      2.439e-02      1.257e-02      2.299e-03      1.000e+01      3.669e-02
      124            189         9.849e-04      4.452e-05      2.881e-02      1.099e-02      2.275e-03      1.000e+01      3.669e-02
      125            191         9.575e-04      2.745e-05      2.757e-02      1.135e-02      2.284e-03      1.000e+01      3.669e-02
      126            192         9.289e-04      2.863e-05      2.790e-02      1.139e-02      2.279e-03      1.000e+01      3.669e-02
      127            193         9.025e-04      2.640e-05      2.636e-02      1.168e-02      2.281e-03      1.000e+01      3.669e-02
      128            194         8.843e-04      1.814e-05      2.540e-02      1.202e-02      2.282e-03      1.000e+01      3.669e-02
      129            195         8.779e-04      6.402e-06      2.497e-02      1.238e-02      2.386e-03      1.000e+01      3.669e-02
      130            196         8.804e-04     -2.437e-06      2.512e-02      1.268e-02      2.494e-03      1.000e+01      3.669e-02
      131            197         8.853e-04     -4.976e-06      2.563e-02      1.279e-02      2.603e-03      1.000e+01      3.669e-02
      132            198         8.878e-04     -2.497e-06      2.620e-02      1.273e-02      2.710e-03      1.000e+01      3.669e-02
      133            199         8.841e-04      3.733e-06      2.660e-02      1.233e-02      2.812e-03      1.000e+01      3.669e-02
      134            200         8.723e-04      1.180e-05      2.674e-02      1.149e-02      2.904e-03      1.000e+01      3.669e-02
      135            201         8.539e-04      1.845e-05      2.665e-02      1.017e-02      2.983e-03      1.000e+01      3.669e-02
      136            202         8.324e-04      2.145e-05      2.637e-02      8.368e-03      3.045e-03      1.000e+01      3.669e-02
      137            203         8.111e-04      2.130e-05      2.597e-02      6.590e-03      3.086e-03      1.000e+01      3.669e-02
      138            204         7.912e-04      1.990e-05      2.549e-02      6.312e-03      3.106e-03      1.000e+01      3.669e-02
      139            205         7.727e-04      1.851e-05      2.500e-02      5.659e-03      3.108e-03      1.000e+01      3.669e-02
      140            206         7.555e-04      1.721e-05      2.450e-02      4.701e-03      3.094e-03      1.000e+01      3.669e-02
      141            207         7.403e-04      1.516e-05      2.396e-02      4.068e-03      3.066e-03      1.000e+01      3.669e-02
      142            208         7.316e-04      8.678e-06      2.353e-02      4.260e-03      3.026e-03      1.000e+01      3.669e-02
      143            209         7.096e-04      2.207e-05      2.454e-02      4.207e-03      3.010e-03      1.000e+01      3.669e-02
      144            210         6.837e-04      2.586e-05      2.471e-02      3.679e-03      3.030e-03      1.000e+01      3.669e-02
      145            212         6.734e-04      1.031e-05      2.315e-02      3.734e-03      3.087e-03      1.000e+01      3.669e-02
      146            213         6.653e-04      8.096e-06      2.093e-02      3.984e-03      3.140e-03      1.000e+01      3.669e-02
      147            214         6.564e-04      8.913e-06      1.802e-02      4.370e-03      3.189e-03      1.000e+01      3.669e-02
      148            215         6.260e-04      3.040e-05      9.257e-03      6.177e-04      3.229e-03      1.000e+01      3.669e-02
      149            216         6.288e-04     -2.764e-06      2.900e-02      3.408e-03      3.258e-03      1.000e+01      3.669e-02
      150            217         6.140e-04      1.477e-05      4.156e-03      9.100e-04      3.283e-03      1.000e+01      3.669e-02
      151            219         6.123e-04      1.628e-06      5.499e-03      5.262e-04      3.302e-03      1.000e+01      3.669e-02
      152            221         6.108e-04      1.564e-06      3.710e-03      6.327e-04      3.315e-03      1.000e+01      3.669e-02
      153            223         6.092e-04      1.602e-06      3.708e-03      5.823e-04      3.329e-03      1.000e+01      3.669e-02
      154            225         6.072e-04      1.936e-06      3.703e-03      5.498e-04      3.341e-03      1.000e+01      3.669e-02
      155            227         6.054e-04      1.803e-06      3.737e-03      5.209e-04      3.352e-03      1.000e+01      3.669e-02
      156            229         6.037e-04      1.772e-06      3.755e-03      5.115e-04      3.362e-03      1.000e+01      3.669e-02
      157            231         6.021e-04      1.611e-06      3.785e-03      5.054e-04      3.371e-03      1.000e+01      3.669e-02
      158            233         6.006e-04      1.421e-06      3.788e-03      5.201e-04      3.379e-03      1.000e+01      3.669e-02
      159            234         5.994e-04      1.195e-06      3.765e-03      5.865e-04      3.387e-03      1.000e+01      3.669e-02
      160            235         5.984e-04      1.047e-06      3.678e-03      6.426e-04      3.393e-03      1.000e+01      3.669e-02
      161            236         5.974e-04      1.037e-06      3.533e-03      6.977e-04      3.399e-03      1.000e+01      3.669e-02
      162            237         5.962e-04      1.176e-06      3.370e-03      7.532e-04      3.405e-03      1.000e+01      3.669e-02
      163            238         5.949e-04      1.323e-06      3.236e-03      7.926e-04      3.410e-03      1.000e+01      3.669e-02
      164            239         5.934e-04      1.457e-06      3.145e-03      8.250e-04      3.414e-03      1.000e+01      3.669e-02
      165            240         5.918e-04      1.581e-06      3.090e-03      8.573e-04      3.418e-03      1.000e+01      3.669e-02
      166            241         5.901e-04      1.715e-06      3.060e-03      8.912e-04      3.422e-03      1.000e+01      3.669e-02
      167            242         5.883e-04      1.857e-06      3.048e-03      9.274e-04      3.426e-03      1.000e+01      3.669e-02
      168            243         5.862e-04      2.009e-06      3.047e-03      9.652e-04      3.429e-03      1.000e+01      3.669e-02
      169            244         5.841e-04      2.166e-06      3.056e-03      1.004e-03      3.432e-03      1.000e+01      3.669e-02
      170            245         5.818e-04      2.324e-06      3.070e-03      1.042e-03      3.434e-03      1.000e+01      3.669e-02
      171            246         5.793e-04      2.478e-06      3.088e-03      1.079e-03      3.436e-03      1.000e+01      3.669e-02
      172            247         5.767e-04      2.620e-06      3.110e-03      1.113e-03      3.437e-03      1.000e+01      3.669e-02
      173            248         5.739e-04      2.745e-06      3.133e-03      1.144e-03      3.438e-03      1.000e+01      3.669e-02
      174            249         5.711e-04      2.845e-06      3.156e-03      1.170e-03      3.439e-03      1.000e+01      3.669e-02
      175            250         5.682e-04      2.915e-06      3.179e-03      1.191e-03      3.439e-03      1.000e+01      3.669e-02
      176            251         5.652e-04      2.953e-06      3.201e-03      1.207e-03      3.439e-03      1.000e+01      3.669e-02
      177            252         5.622e-04      2.958e-06      3.222e-03      1.218e-03      3.438e-03      1.000e+01      3.669e-02
      178            253         5.593e-04      2.933e-06      3.242e-03      1.224e-03      3.437e-03      1.000e+01      3.669e-02
      179            254         5.564e-04      2.884e-06      3.262e-03      1.226e-03      3.440e-03      1.000e+01      3.669e-02
      180            255         5.536e-04      2.816e-06      3.280e-03      1.225e-03      3.457e-03      1.000e+01      3.669e-02
      181            256         5.509e-04      2.736e-06      3.299e-03      1.221e-03      3.475e-03      1.000e+01      3.669e-02
      182            257         5.482e-04      2.651e-06      3.319e-03      1.214e-03      3.492e-03      1.000e+01      3.669e-02
      183            258         5.457e-04      2.566e-06      3.340e-03      1.206e-03      3.509e-03      1.000e+01      3.669e-02
      184            259         5.432e-04      2.435e-06      3.361e-03      1.205e-03      3.526e-03      1.000e+01      3.669e-02
      185            260         5.408e-04      2.426e-06      3.388e-03      1.192e-03      3.543e-03      1.000e+01      3.669e-02
      186            261         5.384e-04      2.355e-06      3.417e-03      1.179e-03      3.560e-03      1.000e+01      3.669e-02
      187            262         5.361e-04      2.299e-06      3.451e-03      1.165e-03      3.577e-03      1.000e+01      3.669e-02
      188            263         5.339e-04      2.249e-06      3.488e-03      1.152e-03      3.593e-03      1.000e+01      3.669e-02
      189            264         5.317e-04      2.208e-06      3.529e-03      1.139e-03      3.610e-03      1.000e+01      3.669e-02
      190            265         5.295e-04      2.175e-06      3.574e-03      1.139e-03      3.626e-03      1.000e+01      3.669e-02
      191            266         5.274e-04      2.150e-06      3.625e-03      1.149e-03      3.643e-03      1.000e+01      3.669e-02
      192            267         5.252e-04      2.133e-06      3.679e-03      1.158e-03      3.659e-03      1.000e+01      3.669e-02
      193            268         5.231e-04      2.124e-06      3.738e-03      1.165e-03      3.676e-03      1.000e+01      3.669e-02
      194            269         5.210e-04      2.121e-06      3.801e-03      1.170e-03      3.692e-03      1.000e+01      3.669e-02
      195            270         5.189e-04      2.124e-06      3.866e-03      1.172e-03      3.709e-03      1.000e+01      3.669e-02
      196            271         5.167e-04      2.132e-06      3.932e-03      1.170e-03      3.725e-03      1.000e+01      3.669e-02
      197            272         5.146e-04      2.144e-06      3.999e-03      1.163e-03      3.742e-03      1.000e+01      3.669e-02
      198            273         5.124e-04      2.219e-06      4.067e-03      1.145e-03      3.758e-03      1.000e+01      3.669e-02
      199            274         5.102e-04      2.193e-06      4.130e-03      1.129e-03      3.775e-03      1.000e+01      3.669e-02
      200            275         5.080e-04      2.207e-06      4.188e-03      1.109e-03      3.791e-03      1.000e+01      3.669e-02
Warning: Maximum number of iterations has been exceeded.
         Current function value: 5.080e-04
         Constraint violation: 3.791e-03
         Total delta_x: 3.471e-01
         Iterations: 200
         Function evaluations: 275
         Jacobian evaluations: 201
Timer: Solution time = 1.14 min
Timer: Avg time per step = 341 ms
==============================================================================================================
                                                                 Start  -->   End
Total (sum of squares):                                      3.290e+01  -->   5.080e-04,
Maximum absolute QS Two-Term objective value:                6.229e-01  -->   1.063e-02 (T^{3})
Minimum absolute QS Two-Term objective value:                1.964e-04  -->   3.044e-06 (T^{3})
Average absolute QS Two-Term objective value:                1.501e-01  -->   6.181e-04 (T^{3})
Maximum absolute QS Two-Term objective value:                6.229e-01  -->   1.063e-02 (normalized)
Minimum absolute QS Two-Term objective value:                1.964e-04  -->   3.044e-06 (normalized)
Average absolute QS Two-Term objective value:                1.501e-01  -->   6.181e-04 (normalized)
Maximum absolute Force error:                                4.342e+01  -->   3.688e+03 (N)
Minimum absolute Force error:                                4.390e-03  -->   2.197e-01 (N)
Average absolute Force error:                                3.601e+00  -->   3.633e+02 (N)
Maximum absolute Force error:                                4.262e-05  -->   3.620e-03 (normalized)
Minimum absolute Force error:                                4.310e-09  -->   2.157e-07 (normalized)
Average absolute Force error:                                3.535e-06  -->   3.566e-04 (normalized)
Aspect ratio:                                                8.000e+00  -->   8.018e+00 (dimensionless)
Elongation:                                                  1.096e+00  -->   3.000e+00 (dimensionless)
Plasma volume:                                               3.084e-01  -->   3.085e-01 (m^3)
Plasma volume:                                               0.000e+00  -->   3.079e-04 (normalized error)
Maximum Rotational transform:                                2.477e-01  -->   1.102e+00 (dimensionless)
Minimum Rotational transform:                                2.477e-01  -->   1.102e+00 (dimensionless)
Average Rotational transform:                                2.477e-01  -->   1.102e+00 (dimensionless)
my mirror ratio objective value:                             2.398e-01  -->   1.882e-01 (Unknown)
my mirror ratio objective value:                             1.982e-02  -->   0.000e+00 (normalized error)
R boundary error:                                            0.000e+00  -->   0.000e+00 (m)
Fixed pressure profile error:                                0.000e+00  -->   0.000e+00 (Pa)
Fixed current profile error:                                 0.000e+00  -->   0.000e+00 (A)
Fixed Psi error:                                             0.000e+00  -->   0.000e+00 (Wb)
==============================================================================================================

As before, results can be improved by running for more iterations. Note the constraint violation may be larger than desired, so it can be helpful to call eq.solve() at the end to decrease the force error without changing the boundary.

[19]:
eqa.solve();
Building objective: force
Precomputing transforms
Building objective: lcfs R
Building objective: lcfs Z
Building objective: fixed Psi
Building objective: fixed pressure
Building objective: fixed current
Building objective: fixed sheet current
Building objective: self_consistency R
Building objective: self_consistency Z
Building objective: lambda gauge
Building objective: axis R self consistency
Building objective: axis Z self consistency
Number of parameters: 120
Number of objectives: 850

Starting optimization
Using method: lsq-exact
Optimization terminated successfully.
`ftol` condition satisfied. (ftol=1.00e-02)
         Current function value: 1.151e-04
         Total delta_x: 1.653e-01
         Iterations: 4
         Function evaluations: 5
         Jacobian evaluations: 5
==============================================================================================================
                                                                 Start  -->   End
Total (sum of squares):                                      3.384e-03  -->   1.151e-04,
Maximum absolute Force error:                                3.688e+03  -->   7.428e+02 (N)
Minimum absolute Force error:                                2.197e-01  -->   8.577e-02 (N)
Average absolute Force error:                                3.633e+02  -->   6.959e+01 (N)
Maximum absolute Force error:                                1.158e-02  -->   2.332e-03 (normalized)
Minimum absolute Force error:                                6.898e-07  -->   2.692e-07 (normalized)
Average absolute Force error:                                1.140e-03  -->   2.184e-04 (normalized)
R boundary error:                                            0.000e+00  -->   1.841e-18 (m)
Z boundary error:                                            0.000e+00  -->   5.225e-18 (m)
Fixed Psi error:                                             0.000e+00  -->   0.000e+00 (Wb)
Fixed pressure profile error:                                0.000e+00  -->   0.000e+00 (Pa)
Fixed current profile error:                                 0.000e+00  -->   0.000e+00 (A)
Fixed sheet current error:                                   0.000e+00  -->   0.000e+00 (~)
==============================================================================================================

From the Boozer plot below we see that we are already doing fairly good for QS:

[20]:
plot_boozer_surface(eqa);
../../_images/notebooks_tutorials_advanced_optimization_32_0.png

As a final comparison, we can look at the boundary shapes obtained by the different methods. We see that the final shapes are fairly similar, with the proximal method giving slightly more elongation and tighter curvature. We could include additional objectives or constraints to try to reduce this if desired.

[21]:
plot_boundaries(
    [eq0, eqa, eqfam[-1]], labels=["Initial", "Augmented Lagrangian", "Proximal"]
);
../../_images/notebooks_tutorials_advanced_optimization_34_0.png

Further example scripts for precise QS optimization can be found in the desc/examples folder.