Using DESC Interactively

%matplotlib inline
import numpy as np
from desc.equilibrium import Equilibrium
from desc.geometry import FourierRZToroidalSurface
from desc.profiles import PowerSeriesProfile
from desc.plotting import plot_1d, plot_section, plot_surfaces

DESC version 0.8.0+20.gd07a08b0.dirty,using JAX backend, jax version=0.4.1, jaxlib version=0.4.1, dtype=float64
Using device: CPU, with 6.73 GB available memory

Jupyter Documentation Tip: pressing ``shift+tab`` when the cursor is inside a function will show the documentation! Take advantage of DESC’s documentation!

See the DESC documentation page for more information on the functionalities shown here

Initializing an Equilibrium

Let’s start with a simple example of a circular tokamak. We’ll start by defining the boundary, which is represented by a double Fourier series for R and Z in terms of a poloidal angle \(\theta\) and the geometric toroidal angle \(\zeta\). We specify the mode numbers for R and Z as 2d arrays of [m,n] pairs, and the coefficients as a 1d array.

In DESC the double Fourier series for R and Z are defined in a slightly different manner than VMEC:

\begin{align} R^b(\theta,\zeta) &= \sum^M \sum^N R^b_{mn} \mathcal{G}^{m}_{n}(\theta,\zeta)\\ Z^b(\theta,\zeta) &= \sum^M \sum^N Z^b_{mn} \mathcal{G}^{m}_{n}(\theta,\zeta) \end{align}

where

\[\begin{split}\mathcal{G}^{m}_{n}(\theta,\zeta) = \begin{cases} \cos(|m|\theta)\cos(|n|N_{FP}\zeta) &\text{for }m\ge0, n\ge0 \\ \cos(|m|\theta)\sin(|n|N_{FP}\zeta) &\text{for }m\ge0, n<0 \\ \sin(|m|\theta)\cos(|n|N_{FP}\zeta) &\text{for }m<0, n\ge0 \\ \sin(|m|\theta)\sin(|n|N_{FP}\zeta) &\text{for }m<0, n<0. \end{cases}\end{split}\]

Note: in DESC, radial modes are indexed by l, poloidal modes by m, and toroidal modes by n

In this case we’ll use a major radius of 10 and a minor radius of 1:

\begin{align} R^b(\theta,\zeta) &= 10 + cos(\theta)\\ Z^b(\theta,\zeta) &= -sin(\theta) \end{align}

which we define in DESC with the following surface:

surface = FourierRZToroidalSurface(
    R_lmn=[10, 1],
    modes_R=[[0, 0], [1, 0]],  # modes given as [m,n] for each coefficient
    Z_lmn=[0, -1],
    modes_Z=[[0, 0], [-1, 0]],
)

Next, we need to define the profiles. We’ll take a vacuum case to start (pressure=0), and a simple quadratic \(\iota\) profile. The profiles are given in terms a simple monic power series in powers of \(\rho^l\):

\begin{align} p(\rho)&=0\\ \iota(\rho) &= 1 + 1.5\rho^2 \end{align}

We create these in DESC as follows. We will give the pressure profile both a constant and a quadratic term which we set to zero for now, as later we will then add pressure to this equilbrium.

# this is a constant pressure of 0
pressure = PowerSeriesProfile(params=[0, 0], modes=[0, 2])
iota = PowerSeriesProfile(params=[1, 1.5], modes=[0, 2])  # iota = 1 + 1.5 r^2

Finally, we create an Equilibrium object by giving it the surface and profiles, as well as specifying what resolution we want to use and a few other parameters:

eq = Equilibrium(
    surface=surface,
    pressure=pressure,
    iota=iota,
    Psi=1.0,  # flux (in Webers) within the last closed flux surface
    NFP=1,  # number of field periods
    L=7,  # radial spectral resolution
    M=7,  # poloidal spectral resolution
    N=0,  # toroidal spectral resolution (axisymmetric case, so we don't need any toroidal modes)
    L_grid=12,  # real space radial resolution, slightly oversampled
    M_grid=12,  # real space poloidal resolution, slightly oversampled
    N_grid=0,  # real space toroidal resolution (axisymmetric, so we don't need any grid points toroidally)
    sym=True,  # explicitly enforce stellarator symmetry
)

This will automatically create the needed spectral bases and generates an initial guess for the flux surfaces by scaling the boundary surface, which we can plot below:

Plotting

# plot_surfaces generates poincare plots of the flux surfaces
plot_surfaces(eq);

We can also look at the force balance error:

# plot_section plots various quantities on a toroidal cross-section
# the second argument is a string telling it what to plot, in this case the force error density
# we could also look at B_z, (the toroidal magnetic field), or sqrt(g) (the coordinate jacobian) etc.,
# here we also tell it to normalize the force error (relative to the magnetic pressure gradient or thermal pressure if present)
plot_section(eq, "|F|", norm_F=True, log=True);

We see that this is very far from an equilibrium. Let’s try to fix that.

Solving the Equilibrium

First, we need to give the equilibrium an objective that should be minimized, in this case we’ll choose to minimize the force balance error.

We also need to select an optimizer. Many options from scipy.optimize are available, as well as a few custom solvers that may be more efficient in some cases

In DESC this is done by creating an Optimizer object, which handles the logic of passing the derivative information from the equilibrium to the solving routines.

from desc.optimize import Optimizer

# create an Optimizer object using the lsq-exact optimizer
optimizer = Optimizer("lsq-exact")

To select an objective, we must create an ObjectiveFunction object with the desired objective (such as force balance error) and we must also define a tuple of the desired constraints (such as keeping the LCFS, pressure, iota and psi fixed, as in the case of a fixed boundary equilibrium solve)

from desc.objectives import (
    get_fixed_boundary_constraints,
    ObjectiveFunction,
    FixBoundaryR,
    FixBoundaryZ,
    FixPressure,
    FixIota,
    FixPsi,
    ForceBalance,
)

constraints = (
    FixBoundaryR(),  # enforce fixed  LCFS for R
    FixBoundaryZ(),  # enforce fixed  LCFS for R
    FixPressure(),  # enforce that the pressure profile stay fixed
    FixIota(),  # enforce that the rotational transform profile stay fixed
    FixPsi(),  # enforce that the enclosed toroidal stay fixed
)
# choose the objectives to be ForceBalance(), which is a wrapper function for RadialForceBalance() and HelicalForceBalance()
objectives = ForceBalance()
# the ObjectiveFunction object which we can pass to the eq.solve method
obj = ObjectiveFunction(objectives=objectives)

a utility function, get_fixed_boundary_constraints() exists that will create the typical constraints for a fixed boundary equilibrium

constraints2 = get_fixed_boundary_constraints()
print("Same constraints for each")
for c1, c2 in zip(constraints, constraints2):
    print("type of ", c1, "=", "type of ", c2, ":\n", type(c1) == type(c2))

Same constraints for each
type of  <desc.objectives.linear_objectives.FixBoundaryR object at 0x7f7890024280> = type of  <desc.objectives.linear_objectives.FixBoundaryR object at 0x7f7890024a60> :
 True
type of  <desc.objectives.linear_objectives.FixBoundaryZ object at 0x7f78900adb50> = type of  <desc.objectives.linear_objectives.FixBoundaryZ object at 0x7f78900248b0> :
 True
type of  <desc.objectives.linear_objectives.FixPressure object at 0x7f78900adf10> = type of  <desc.objectives.linear_objectives.FixPsi object at 0x7f7890024400> :
 False
type of  <desc.objectives.linear_objectives.FixIota object at 0x7f78900add30> = type of  <desc.objectives.linear_objectives.FixPressure object at 0x7f7890024bb0> :
 False
type of  <desc.objectives.linear_objectives.FixPsi object at 0x7f78900ad4f0> = type of  <desc.objectives.linear_objectives.FixIota object at 0x7f78900aaf10> :
 False

Next, we simply call eq.solve() to minimize the force balance error.

Here we can also pass in arguments such as the maximum number of iterations or stopping tolerances. We also must pass in our constraints for the equilibrium solve.

Under the hood, the objective function and its derivative are JIT compiled for the specific parameters we defined above before being passed to the optimizer.

eq.solve(
    verbose=3,
    ftol=1e-8,
    maxiter=50,
    constraints=constraints,
    optimizer=optimizer,
    objective=obj,
);

Building objective: force
Precomputing transforms
Timer: Precomputing transforms = 1.20 sec
Timer: Objective build = 5.65 sec
Timer: Linear constraint projection build = 13.3 sec
Compiling objective function and derivatives
Timer: Objective compilation time = 2.01 sec
Timer: Jacobian compilation time = 5.43 sec
Timer: Total compilation time = 7.45 sec
Number of parameters: 37
Number of objectives: 98
Starting optimization
   Iteration     Total nfev        Cost      Cost reduction    Step norm     Optimality
       0              1         6.8245e-03                                    2.08e+00
       1              3         5.7678e-05      6.77e-03       1.69e-01       2.57e-01
       2              4         1.3872e-06      5.63e-05       4.48e-02       2.91e-02
       3              5         5.3211e-08      1.33e-06       7.60e-03       5.31e-03
       4              6         2.9865e-09      5.02e-08       3.86e-03       1.30e-03
       5              7         1.1962e-10      2.87e-09       1.55e-03       2.26e-04
       6              8         4.9299e-12      1.15e-10       4.29e-04       1.15e-05
       7              9         3.6124e-12      1.32e-12       5.61e-04       1.14e-05
       8             10         3.3322e-12      2.80e-13       2.32e-04       2.35e-06
       9             11         3.3197e-12      1.24e-14       4.93e-05       1.09e-07
      10             12         3.3192e-12      5.10e-16       1.36e-05       8.70e-09
`gtol` condition satisfied.
         Current function value: 3.319e-12
         Total delta_x: 1.212e-01
         Iterations: 10
         Function evaluations: 12
         Jacobian evaluations: 11
Timer: Solution time = 2.08 sec
Timer: Avg time per step = 189 ms
Start of solver
Total (sum of squares):  6.825e-03,
Total force:  3.089e+05 (N)
Total force:  1.168e-01 (normalized)
End of solver
Total (sum of squares):  3.319e-12,
Total force:  6.813e+00 (N)
Total force:  2.577e-06 (normalized)

We can then look at the flux surfaces and force error again:

plot_surfaces(eq);

Perturbing the equilibrium

Now we have a solved zero beta equilibrium. If we want to see what this equilibrium solution looks like with finite pressure, we could redo the above steps, but with a different pressure profile. But, in DESC one can perturb an existing equilibrium in order to find nearby solution. Let’s do this to find a finite pressure solution of this equilibrium.

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

Perturbations

Next, we can ask “what would this equilibrium look like if we add pressure to it?”

We can answer this by applying a pressure perturbation to the current equilibrium.

First, let’s plot the pressure to make sure it’s really zero:

plot_1d(eq, "p");  # original eq with zero pressure

Next, let’s decide on how much we want to increase the pressure.

We’ll give it a quadratic profile, peaked at 1000 Pascals in the core and dropping to 0 at the edge.

The perturbation is given in the same form the original profiles were (a power series in \(\rho\))

delta_p = np.zeros_like(eq.p_l)
delta_p[0] = 1000.0
delta_p[1] = -1000.0

eq1 = eq.perturb(
    deltas={"p_l": delta_p}, order=2, objective=obj, constraints=constraints
)

Perturbing p_l
Factorizing linear constraints
Computing df
Factoring df
Computing d^2f
||dx||/||x|| =  3.944e-03

Note that this gives us back a new equilibrium, so that the original is saved for future study.

plot_1d(eq1, "p");

plot_surfaces(eq1);

Note the changes in the flux surface, and the slight outward movement of the axis

During the perturbation, the force error increased slightly, but still remains less than 10%

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

We can do a few Newton iterations to converge the solution again, this should go much faster than the initial solution

# the default ObjectiveFunction is force balance, and the default optimizer is 'lsq-exact'
eq1.solve(
    verbose=2, ftol=1e-4, optimizer=optimizer, objective=obj, constraints=constraints
);

Timer: Linear constraint projection build = 713 ms
Compiling objective function and derivatives
Timer: Objective compilation time = 802 us
Timer: Jacobian compilation time = 2.07 ms
Timer: Total compilation time = 6.97 ms
Number of parameters: 37
Number of objectives: 98
Starting optimization
   Iteration     Total nfev        Cost      Cost reduction    Step norm     Optimality
       0              1         2.2632e-08                                    4.16e-03
       1              2         3.0267e-10      2.23e-08       2.17e-03       3.86e-04
       2              3         1.8486e-10      1.18e-10       9.30e-04       7.73e-05
       3              4         1.8113e-10      3.72e-12       8.06e-05       4.37e-07
       4              5         1.8112e-10      1.83e-14       7.90e-05       5.35e-07
       5              6         1.8111e-10      2.00e-15       1.40e-05       1.84e-08
Optimization terminated successfully.
`ftol` condition satisfied.
         Current function value: 1.811e-10
         Total delta_x: 3.045e-03
         Iterations: 5
         Function evaluations: 6
         Jacobian evaluations: 6
Timer: Solution time = 290 ms
Timer: Avg time per step = 48.4 ms
Start of solver
Total (sum of squares):  2.263e-08,
Total force:  5.626e+02 (N)
Total force:  2.128e-04 (normalized)
End of solver
Total (sum of squares):  1.811e-10,
Total force:  5.033e+01 (N)
Total force:  1.903e-05 (normalized)

plot_surfaces(eq1);

Note that the flux surfaces and axis location only change by a tiny amount during the Newton iterations - the perturbation captured them accurately

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

Next, we might want to try twisting our tokamak into a stellarator. We can accomplish this in a similar way, by applying a 3d boundary perturbation.

However, first we need to give our tokamak some non-axisymmetric modes in its spectral basis:

eq1.change_resolution(N=2, N_grid=4)

# remake the constraints and objective, since we have changed the resolution of our equilibrium
constraints = (
    FixBoundaryR(),  # enforce fixed  LCFS for R
    FixBoundaryZ(),  # enforce fixed  LCFS for R
    FixPressure(),  # enforce that the pressure profile stay fixed
    FixIota(),  # enforce that the rotational transform profile stay fixed
    FixPsi(),  # enforce that the enclosed toroidal stay fixed
)
# choose the objectives to be ForceBalance(), which is a wrapper function for RadialForceBalance() and HelicalForceBalance()
objectives = ForceBalance()
# the final ObjectiveFunction object which we can pass to the eq.solve method
obj = ObjectiveFunction(objectives=objectives)

Now we can apply a 3d perturbation by perturbing the coefficients of the double Fourier series that defines the plasma boundary:

delta_R = np.zeros_like(eq1.Rb_lmn)
delta_Z = np.zeros_like(eq1.Zb_lmn)
delta_R[eq1.surface.R_basis.get_idx(M=1, N=1)] = -0.4
delta_Z[eq1.surface.Z_basis.get_idx(M=1, N=-1)] = -0.4

eq2 = eq1.perturb(deltas={"Rb_lmn": delta_R, "Zb_lmn": delta_Z}, order=2)

Building objective: force
Precomputing transforms
Perturbing Rb_lmn, Zb_lmn
Factorizing linear constraints
Computing df
Factoring df
Computing d^2f
||dx||/||x|| =  4.097e-02

plot_surfaces(eq2);

Again, note that the force error increases after the perturbation. How much it increases depends on the “size” of the perturbation (note that due to the highly nonlinear equations, the “size” of the perturbation required to maintain force balance depends on the starting equilibrium as well as the actual delta).

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

We can again run a few additional Newton iterations to improve the force balance error.

eq2.solve(
    verbose=2, ftol=1e-2, optimizer=optimizer, objective=obj, constraints=constraints
);

Building objective: force
Precomputing transforms
Timer: Precomputing transforms = 285 ms
Timer: Objective build = 969 ms
Timer: Linear constraint projection build = 3.93 sec
Compiling objective function and derivatives
Timer: Objective compilation time = 4.72 sec
Timer: Jacobian compilation time = 11.7 sec
Timer: Total compilation time = 16.5 sec
Number of parameters: 193
Number of objectives: 882
Starting optimization
   Iteration     Total nfev        Cost      Cost reduction    Step norm     Optimality
       0              1         3.3146e-02                                    5.34e+00
       1              2         2.0832e-04      3.29e-02       5.42e-01       2.60e-01
       2              3         5.8801e-05      1.50e-04       3.34e-01       1.15e-01
       3              4         3.8911e-06      5.49e-05       1.77e-01       2.64e-02
       4              5         2.7322e-06      1.16e-06       7.34e-02       3.69e-02
       5              6         1.1782e-06      1.55e-06       6.74e-02       2.71e-02
       6              7         1.2083e-07      1.06e-06       2.72e-02       4.74e-03
       7              9         7.7004e-08      4.38e-08       8.17e-03       5.42e-04
       8             11         7.6094e-08      9.10e-10       3.49e-03       1.20e-04
       9             12         7.5113e-08      9.81e-10       5.77e-03       4.18e-04
      10             14         7.4266e-08      8.47e-10       4.25e-03       1.23e-04
      11             15         7.3959e-08      3.07e-10       4.52e-03       4.12e-04
Optimization terminated successfully.
`ftol` condition satisfied.
         Current function value: 7.396e-08
         Total delta_x: 2.206e-01
         Iterations: 11
         Function evaluations: 15
         Jacobian evaluations: 12
Timer: Solution time = 6.05 sec
Timer: Avg time per step = 504 ms
Start of solver
Total (sum of squares):  3.315e-02,
Total force:  2.270e+05 (N)
Total force:  2.575e-01 (normalized)
End of solver
Total (sum of squares):  7.396e-08,
Total force:  3.390e+02 (N)
Total force:  3.846e-04 (normalized)

plot_surfaces(eq2);

Note again how the flux surfaces remain very similar, as the perturbation captures them accurately. Additional Newton iterations primarily serve to reduce the force error

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

Note that even after additional Newton iterations, the force error is still rather high (between 1% and 10% throughout most of the volume). This is because the resolution we used was rather low, meant to be a quick demonstration rather than a realistic high quality solution. As DESC uses a pseudo-spectral discretization, increasing the resolution should decrease the error roughly exponentially.