Solving Current-Constrained Equilibria

DESC is able to solve equilibria with a prescribed toroidal current profile. We will solve two different zero-current solutions: first an axisymmetric zero-pressure equilibrium, then perturb that solution to a zero-pressure, zero-current stellarator. In this example, the 2D solution is a simple circular tokamak and the 3D solution is from the HELIOTRON example.

If you have access to a GPU, uncomment the following two lines. 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")

%matplotlib inline
import numpy as np

from desc.equilibrium import Equilibrium
from desc.geometry import FourierRZToroidalSurface
from desc.objectives import (
    ObjectiveFunction,
    ForceBalance,
    get_fixed_boundary_constraints,
)
from desc.optimize import Optimizer
from desc.plotting import plot_1d, plot_section, plot_surfaces
from desc.profiles import PowerSeriesProfile

DESC version 0.7.2+334.g1fe6b36e,using JAX backend, jax version=0.4.1, jaxlib version=0.4.1, dtype=float64
Using device: CPU, with 10.58 GB available memory

2D Equilibrium

We start by creating the surface object that represents the axisymmetric boundary.

surface_2D = FourierRZToroidalSurface(
    R_lmn=np.array([10, -1]),  # boundary coefficients
    Z_lmn=np.array([1]),
    modes_R=np.array([[0, 0], [1, 0]]),  # [M, N] boundary Fourier modes
    modes_Z=np.array([[-1, 0]]),
    NFP=5,  # number of (toroidal) field periods
)

Now we can initialize an Equilibrium with this boundary surface. We will prescribe profiles that are zero, since we want zero pressure and zero current. We note that if we do not pass anything to an Equilibrium when we initialize it, it will default to having a zero pressure profile and a zero current profile.We also increase the resolution and use a collocation grid that oversamples by a factor of two.

current = PowerSeriesProfile(modes=np.array([0]), params=np.array([0]))
pressure = PowerSeriesProfile(modes=np.array([0]), params=np.array([0]))

# axisymmetric & stellarator symmetric equilibrium
eq = Equilibrium(current=current, pressure=pressure, surface=surface_2D, sym=True)
eq.change_resolution(L=6, M=6, L_grid=12, M_grid=12)

objective = ObjectiveFunction(ForceBalance())

Next we need to specify the optimization constraints, which indicate what parameters that will remain fixed during the optimization process. For this fixed-boundary problem we can call the utility function get_fixed_boundary_constraints that returns a list of the desired constraints. The profiles=True option tells it to return constraints that we have fixed profiles of pressure and either iota or current, and iota=False says to fix current instead of iota.

constraints = get_fixed_boundary_constraints(profiles=True, iota=False)

constraints  # we see FixPressure and FixCurrent in the constraints

(<desc.objectives.linear_objectives.FixBoundaryR at 0x7f25ec199940>,
 <desc.objectives.linear_objectives.FixBoundaryZ at 0x7f25ec199c40>,
 <desc.objectives.linear_objectives.FixLambdaGauge at 0x7f25ec199250>,
 <desc.objectives.linear_objectives.FixPsi at 0x7f25ec1990d0>,
 <desc.objectives.linear_objectives.FixPressure at 0x7f25ec199df0>,
 <desc.objectives.linear_objectives.FixCurrent at 0x7f25ec199580>)

Finally, we can solve the equilibrium with the objective and constraints specified above. We also chose an optimization algorithm by initializing an Optimizer object. The verbose=3 option will display output at each optimization step.

# this is a port of scipy's trust region least squares algorithm but using jax functions for better performance
optimizer = Optimizer("lsq-exact")
eq, solver_outputs = eq.solve(
    objective=objective, constraints=constraints, optimizer=optimizer, verbose=3
)

Building objective: force
Precomputing transforms
Timer: Precomputing transforms = 691 ms
Timer: Objective build = 3.46 sec
Timer: Linear constraint projection build = 5.20 sec
Compiling objective function and derivatives
Timer: Objective compilation time = 1.47 sec
Timer: Jacobian compilation time = 3.72 sec
Timer: Total compilation time = 5.20 sec
Number of parameters: 27
Number of objectives: 98
Starting optimization
   Iteration     Total nfev        Cost      Cost reduction    Step norm     Optimality
       0              1         1.1345e-02                                    1.26e+00
       1              2         2.0864e-03      9.26e-03       1.93e-01       8.91e-01
       2              3         2.1361e-04      1.87e-03       7.73e-02       2.71e-01
       3              4         4.5299e-06      2.09e-04       5.07e-02       2.86e-02
       4              5         3.1993e-07      4.21e-06       3.73e-02       7.61e-03
       5              6         2.1350e-08      2.99e-07       1.61e-02       3.60e-03
       6              7         1.8191e-10      2.12e-08       3.14e-03       2.17e-04
       7              8         1.0815e-10      7.38e-11       7.72e-03       4.99e-04
       8              9         4.0570e-13      1.08e-10       9.62e-04       1.51e-05
       9             11         8.0533e-14      3.25e-13       8.85e-04       4.44e-06
      10             13         1.7655e-14      6.29e-14       4.27e-04       1.08e-06
      11             15         2.0994e-15      1.56e-14       2.26e-04       8.46e-07
      12             17         3.3019e-16      1.77e-15       9.24e-05       3.85e-07
      13             19         8.8731e-18      3.21e-16       4.13e-05       3.43e-08
      14             21         6.6568e-19      8.21e-18       2.06e-05       1.15e-08
      15             23         3.4753e-19      3.18e-19       1.13e-05       1.69e-09
`gtol` condition satisfied.
         Current function value: 3.475e-19
         Total delta_x: 1.461e-01
         Iterations: 15
         Function evaluations: 23
         Jacobian evaluations: 16
Timer: Solution time = 1.32 sec
Timer: Avg time per step = 82.9 ms
Start of solver
Total (sum of squares):  1.135e-02,
Total force:  2.422e+05 (N)
Total force:  1.506e-01 (normalized)
End of solver
Total (sum of squares):  3.475e-19,
Total force:  1.340e-03 (N)
Total force:  8.337e-10 (normalized)

We can analyze the equilibrium solution using the available plotting routines. . Here we plot the net toroidal current profile (unsurprisingly is zero, since we constrained it) and the normalized force balance error.

plot_1d(eq, "current")
plot_section(eq, "|F|", norm_F=True, log=True);

Since this is an axisymmetric vacuum equilibrium, there should be no pressure (since we fixed to zero) or rotational transform. We can check the profiles of both quantities as follows:

plot_1d(eq, "p")
plot_1d(eq, "iota")

(<Figure size 384x384 with 1 Axes>,
 <AxesSubplot:xlabel='$\\rho$', ylabel='$ \\iota ~(~)$'>)

3D Equilibrium

Now we want to solve a stellarator vacuum equilibrium by perturbing the existing tokamak solution we already found. We start by creating a new surface to represent the 3D (non-axisymmetric) stellarator boundary.

surface_3D = FourierRZToroidalSurface(
    R_lmn=np.array([10, -1, -0.3, 0.3]),  # boundary coefficients
    Z_lmn=np.array([1, -0.3, -0.3]),
    modes_R=np.array(
        [[0, 0], [1, 0], [1, 1], [-1, -1]]
    ),  # [M, N] boundary Fourier modes
    modes_Z=np.array([[-1, 0], [-1, 1], [1, -1]]),
    NFP=5,  # number of (toroidal) field periods
)

In the previous solution we did not use any toroidal Fourier modes because they were unnecessary for axisymmetry. Now we need to increase the toroidal resolution, and we will also increase the radial and poloidal resolutions as well. Again we oversample the collocation grid by a factor of two.

We will also update the resolution of the 3D surface to match the new resolution of the Equilibrium.

eq.change_resolution(L=10, M=10, N=6, L_grid=20, M_grid=20, N_grid=12)
surface_3D.change_resolution(eq.L, eq.M, eq.N)

We need to initialize new instances of the objective and constraints. This is necessary because the original instances got built for a specific resolution during the previous 2D equilibrium solve, and are no longer compatible with the Equilibrium after increasing the resolution.

objective = ObjectiveFunction(ForceBalance())
constraints = get_fixed_boundary_constraints(profiles=True, iota=False)

Next is the boundary perturbation. In this step, we approximate the heliotron equilibrium solution from a 2nd-order Taylor expansion about the axisymmetric solution. This is possible thanks to the wealth of derivative information provided by automatic differentiation.

eq.perturb(
    deltas={
        "Rb_lmn": surface_3D.R_lmn - eq.Rb_lmn,  # change in the R boundary coefficients
        "Zb_lmn": surface_3D.Z_lmn - eq.Zb_lmn,
    },  # change in the Z boundary coefficients
    objective=objective,  # perturb the solution such that J=0 is maintained
    constraints=constraints,  # same constraints used in the equilibrium solve
    order=2,  # use a 2nd-order Taylor expansion
    verbose=2,  # display timing data
)

Building objective: force
Precomputing transforms
Timer: Precomputing transforms = 603 ms
Timer: Objective build = 3.82 sec
Perturbing Rb_lmn, Zb_lmn
Factorizing linear constraints
Timer: linear constraint factorize = 1.28 sec
Computing df
Timer: df computation = 18.9 sec
Factoring df
Timer: df/dx factorization = 1.25 sec
Computing d^2f
Timer: d^2f computation = 19.9 sec
||dx||/||x|| =  5.505e-02
Timer: Total perturbation = 42.8 sec

Equilibrium at 0x7f25ec225190 (L=10, M=10, N=6, NFP=5, sym=True, spectral_indexing=ansi)

We now have an approximation of the stellarator equilibrium from the tokamak solution! Let us look at the 3D surfaces and rotational transform profile to check that the perturbation actually updated the solution:

plot_surfaces(eq)
plot_1d(eq, "iota");

The surfaces match the heliotron boundary we want and there is non-zero rotational transform as expected. But the equilibrium error is now large because the perturbed solution is only an approximation to the true equilibrium:

plot_section(eq, "|F|", norm_F=True, log=True);

We can re-solve the equilibrium using the new 3D boundary constraint. This should converge in only a few Newton iterations because we are starting from a good initial guess.

eq, solver_outputs = eq.solve(
    objective=objective,  # solve F=0
    constraints=constraints,  # fixed-boundary constraints, except iota is a free parameter
    optimizer=optimizer,  # we can use the same optimizer as before
    ftol=1e-2,  # stopping tolerance on the function value
    xtol=1e-4,  # stopping tolerance on the step size
    gtol=1e-6,  # stopping tolerance on the gradient
    maxiter=20,  # maximum number of iterations
    verbose=3,  # display output at each iteration
)

Timer: Linear constraint projection build = 436 ms
Compiling objective function and derivatives
Timer: Objective compilation time = 2.78 ms
Timer: Jacobian compilation time = 2.87 sec
Timer: Total compilation time = 2.87 sec
Number of parameters: 1005
Number of objectives: 6050
Starting optimization
   Iteration     Total nfev        Cost      Cost reduction    Step norm     Optimality
       0              1         9.6555e-02                                    2.37e+01
       1              2         2.4304e-03      9.41e-02       1.39e-01       2.58e+00
       2              3         4.6211e-04      1.97e-03       1.29e-01       1.78e+00
       3              4         9.5056e-05      3.67e-04       7.85e-02       9.02e-01
       4              5         6.1316e-05      3.37e-05       1.24e-01       3.04e-01
       5              6         9.9169e-07      6.03e-05       4.22e-02       4.98e-02
       6              7         4.0397e-07      5.88e-07       3.26e-02       2.69e-02
       7              8         6.7418e-09      3.97e-07       6.23e-03       7.24e-03
       8              9         3.4765e-10      6.39e-09       2.57e-03       2.09e-03
       9             10         8.2181e-11      2.65e-10       1.95e-03       9.74e-04
      10             11         4.7683e-11      3.45e-11       1.62e-03       8.02e-04
      11             13         7.0083e-13      4.70e-11       5.83e-04       7.95e-05
Optimization terminated successfully.
`xtol` condition satisfied.
         Current function value: 7.008e-13
         Total delta_x: 3.089e-01
         Iterations: 11
         Function evaluations: 13
         Jacobian evaluations: 12
Timer: Solution time = 53.5 sec
Timer: Avg time per step = 4.46 sec
Start of solver
Total (sum of squares):  9.656e-02,
Total force:  8.992e+04 (N)
Total force:  4.394e-01 (normalized)
End of solver
Total (sum of squares):  7.008e-13,
Total force:  2.422e-01 (N)
Total force:  1.184e-06 (normalized)

We can analyze our final solution using the same plotting commands as before. Note that the flux surfaces and rotational transform profile only had small corrections compared to the perturbed solution, but the equilibrium error was significantly improved.

plot_surfaces(eq)
plot_1d(eq, "iota")
plot_1d(eq, "current")
plot_section(eq, "|F|", norm_F=True, log=True);