Continuation method step by step
In this example we step through the continuation method step by step - starting from an axisymmetric vacuum configuration, adding 3D shaping, and then adding pressure.
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]:
%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
2D Equilibrium
We start by creating the surface object that represents the axisymmetric boundary.
[5]:
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 (does not matter for 2D, but will for 3D solution)
)
Now we can initialize an Equilibrium with this boundary surface. By default, Equilibrium objects have pressure and net toroidal current profiles of 0 assigned. We also increase the resolution and use a collocation grid that oversamples by a factor of two.
[6]:
# axisymmetric & stellarator symmetric equilibrium
eq = Equilibrium(surface=surface_2D, sym=True)
eq.change_resolution(L=6, M=6, L_grid=12, M_grid=12)
Next we create our objective function, ForceBalance which will seek to make \(\mathbf{F} \equiv \mathbf{J} \times \mathbf{B} - \nabla p = 0\)
[7]:
objective = ObjectiveFunction(ForceBalance(eq=eq))
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.
[8]:
constraints = get_fixed_boundary_constraints(eq=eq)
for c in constraints:
print(c)
<desc.objectives.linear_objectives.FixBoundaryR object at 0x7b11942601d0>
<desc.objectives.linear_objectives.FixBoundaryZ object at 0x7b11703bef00>
<desc.objectives.linear_objectives.FixPsi object at 0x7b1170112db0>
<desc.objectives.linear_objectives.FixPressure object at 0x7b1110394140>
<desc.objectives.linear_objectives.FixCurrent object at 0x7b111020ba10>
<desc.objectives.linear_objectives.FixSheetCurrent object at 0x7b119419a870>
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.
[9]:
# 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 = 1.16 sec
Timer: Objective build = 1.60 sec
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
Timer: Objective build = 463 ms
Timer: LinearConstraintProjection build = 4.18 sec
Number of parameters: 27
Number of objectives: 98
Timer: Initializing the optimization = 6.31 sec
Starting optimization
Using method: lsq-exact
Solver options:
------------------------------------------------------------
Maximum Function Evaluations : 501
Maximum Allowed Total Δx Norm : inf
Scaled Termination : True
Trust Region Method : qr
Initial Trust Radius : 3.289e+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 9.294e-03 1.350e-01
1 2 1.709e-03 7.585e-03 2.108e-01 3.122e-02
2 3 1.750e-04 1.534e-03 8.901e-02 1.071e-02
3 4 3.711e-06 1.713e-04 5.145e-02 1.255e-03
4 5 2.621e-07 3.449e-06 3.861e-02 3.428e-04
5 6 1.749e-08 2.446e-07 1.658e-02 1.173e-04
6 7 1.490e-10 1.734e-08 3.866e-03 8.023e-06
7 8 8.859e-11 6.042e-11 8.400e-03 1.001e-05
8 9 3.323e-13 8.826e-11 9.622e-04 4.061e-07
9 11 6.363e-14 2.687e-13 9.000e-04 1.436e-07
10 14 1.721e-15 6.191e-14 1.368e-04 1.778e-08
11 16 1.543e-16 1.567e-15 6.479e-05 1.022e-08
12 18 1.170e-17 1.426e-16 3.527e-05 2.347e-09
`gtol` condition satisfied. (gtol=1.00e-08)
Current function value: 1.170e-17
Total delta_x: 1.839e-01
Iterations: 12
Function evaluations: 18
Jacobian evaluations: 13
Timer: Solution time = 6.85 sec
Timer: Avg time per step = 527 ms
==============================================================================================================
Start --> End
Total (sum of squares): 9.294e-03 --> 1.170e-17,
Maximum absolute Force error: 7.882e+04 --> 1.352e-02 (N)
Minimum absolute Force error: 1.111e-12 --> 1.809e-05 (N)
Average absolute Force error: 2.555e+04 --> 7.139e-04 (N)
Maximum absolute Force error: 6.339e-03 --> 1.088e-09 (normalized)
Minimum absolute Force error: 8.936e-20 --> 1.455e-12 (normalized)
Average absolute Force error: 2.055e-03 --> 5.741e-11 (normalized)
R boundary error: 0.000e+00 --> 0.000e+00 (m)
Z boundary error: 0.000e+00 --> 0.000e+00 (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 (~)
==============================================================================================================
We can analyze the equilibrium solution using the available plotting routines. . Here we plot the magnitude of the total current density \(|\mathbf{J}|\) and the normalized force balance error. We expect both quantities to be low for this vacuum solution.
[10]:
plot_section(eq, "|J|")
plot_section(eq, "|F|_normalized", log=True);
Since this is an axisymmetric vacuum equilibrium, there should be no pressure or rotational transform. We can plot both quantities as follows:
[11]:
plot_1d(eq, "p")
plot_1d(eq, "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.
[12]:
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.
[13]:
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.
[14]:
objective = ObjectiveFunction(ForceBalance(eq=eq))
constraints = get_fixed_boundary_constraints(eq=eq)
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.
[15]:
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 = 871 ms
Timer: Objective build = 1.21 sec
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
Timer: Objective build = 511 ms
Perturbing Rb_lmn, Zb_lmn
Factorizing linear constraints
Timer: linear constraint factorize = 4.45 sec
Computing df
Timer: df computation = 8.34 sec
Factoring df
Timer: df/dx factorization = 756 ms
Computing d^2f
Timer: d^2f computation = 5.14 sec
Timer: Objective build = 1.27 ms
||dx||/||x|| = 3.834e-02
Timer: Total perturbation = 24.3 sec
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:
[16]:
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:
[17]:
plot_section(eq, "|F|_normalized", 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.
[18]:
eq, solver_outputs = eq.solve(
objective=objective, # solve JxB-grad(p)=0
constraints=constraints, # fixed-boundary constraints
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
)
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 = 21.6 ms
Timer: LinearConstraintProjection build = 420 ms
Number of parameters: 1011
Number of objectives: 6050
Timer: Initializing the optimization = 485 ms
Starting optimization
Using method: 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.263e+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 7.630e-02 1.326e-01
1 2 1.319e-03 7.498e-02 8.552e-02 1.773e-02
2 3 8.872e-05 1.231e-03 6.482e-02 3.950e-03
3 4 2.443e-05 6.428e-05 9.706e-02 2.051e-03
4 5 1.290e-05 1.154e-05 1.070e-01 1.552e-03
5 6 2.198e-06 1.070e-05 9.033e-02 6.943e-04
6 7 4.612e-08 2.152e-06 1.941e-02 7.357e-05
7 8 4.054e-09 4.207e-08 5.577e-03 2.658e-05
8 9 1.638e-09 2.416e-09 3.241e-03 1.748e-05
9 10 6.422e-10 9.957e-10 2.654e-03 1.206e-05
10 11 1.043e-10 5.379e-10 1.736e-03 5.176e-06
Optimization terminated successfully.
`xtol` condition satisfied. (xtol=1.00e-04)
Current function value: 1.043e-10
Total delta_x: 3.088e-01
Iterations: 10
Function evaluations: 11
Jacobian evaluations: 11
Timer: Solution time = 9.38 sec
Timer: Avg time per step = 853 ms
==============================================================================================================
Start --> End
Total (sum of squares): 7.630e-02 --> 1.043e-10,
Maximum absolute Force error: 2.886e+06 --> 1.513e+02 (N)
Minimum absolute Force error: 6.297e-01 --> 2.008e-05 (N)
Average absolute Force error: 5.082e+04 --> 1.264e+00 (N)
Maximum absolute Force error: 2.321e-01 --> 1.217e-05 (normalized)
Minimum absolute Force error: 5.064e-08 --> 1.615e-12 (normalized)
Average absolute Force error: 4.087e-03 --> 1.016e-07 (normalized)
R boundary error: 0.000e+00 --> 0.000e+00 (m)
Z boundary error: 0.000e+00 --> 0.000e+00 (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 (~)
==============================================================================================================
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.
[19]:
plot_surfaces(eq)
plot_1d(eq, "iota")
plot_section(eq, "|J|", log=True)
plot_section(eq, "|F|_normalized", log=True);
Finite beta stellarator
We’ve solved for a vacuum stellarator, but what if we now want to look at what happens at finite beta? We can simply apply a pressure perturbation.
[20]:
from desc.objectives import ForceBalance
objective = ObjectiveFunction(ForceBalance(eq=eq))
constraints = get_fixed_boundary_constraints(eq=eq, profiles=True)
Next we’ll make our desired pressure profile, corresponding to a profile of the form \(p(\rho) = 2000(1 - 2 \rho^2 + \rho^4) ~\text{Pa}\)
[21]:
from desc.profiles import PowerSeriesProfile
pressure = PowerSeriesProfile([2000, 0, -4000, 0, 2000])
pressure.change_resolution(eq.L)
Now, we will perturb our equilibrium’s pressure coefficients \(p_l\) to match those of the target pressure
[22]:
eq.perturb(
deltas={"p_l": pressure.params - eq.p_l}, # change in the pressure coefficients
objective=objective, # perturb the solution such that JxB-grad(p)=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 = 65.2 ms
Timer: Objective build = 84.3 ms
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
Timer: Objective build = 221 ms
Perturbing p_l
Factorizing linear constraints
Timer: linear constraint factorize = 312 ms
Computing df
Timer: df computation = 18.6 ms
Factoring df
Timer: df/dx factorization = 781 ms
Computing d^2f
Timer: d^2f computation = 7.07 ms
Timer: Objective build = 1.09 ms
||dx||/||x|| = 6.321e-02
Timer: Total perturbation = 2.90 sec
We can see that the axis has moved due to the Shafranov shift, and the pressure profile is now nonzero. The force balance error is significantly larger, however.
[23]:
plot_surfaces(eq)
plot_1d(eq, "p")
plot_section(eq, "|F|_normalized", log=True);
[24]:
eq, solver_outputs = eq.solve(
objective=objective, # solve JxB-grad(p)=0
constraints=constraints, # fixed-boundary and profile constraints
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=75, # maximum number of iterations
verbose=3, # display output at each iteration
)
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 = 18.4 ms
Timer: LinearConstraintProjection build = 226 ms
Number of parameters: 1011
Number of objectives: 6050
Timer: Initializing the optimization = 294 ms
Starting optimization
Using method: lsq-exact
Solver options:
------------------------------------------------------------
Maximum Function Evaluations : 376
Maximum Allowed Total Δx Norm : inf
Scaled Termination : True
Trust Region Method : qr
Initial Trust Radius : 1.928e+02
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 6.443e-02 1.963e-01
1 2 5.381e-04 6.389e-02 1.059e-01 9.898e-03
2 3 3.343e-04 2.038e-04 1.181e-01 6.760e-03
3 5 1.593e-04 1.749e-04 6.789e-02 8.414e-03
4 6 1.874e-05 1.406e-04 5.976e-02 1.847e-03
5 8 5.495e-06 1.325e-05 2.706e-02 1.457e-03
6 9 2.554e-06 2.941e-06 3.454e-02 9.769e-04
7 10 2.053e-06 5.017e-07 2.960e-02 8.769e-04
8 11 1.141e-07 1.938e-06 7.853e-03 7.425e-05
9 13 9.766e-08 1.646e-08 4.010e-03 1.907e-05
10 15 9.624e-08 1.422e-09 2.021e-03 4.599e-06
Optimization terminated successfully.
`xtol` condition satisfied. (xtol=1.00e-04)
Current function value: 9.624e-08
Total delta_x: 2.304e-01
Iterations: 10
Function evaluations: 15
Jacobian evaluations: 11
Timer: Solution time = 5.23 sec
Timer: Avg time per step = 475 ms
==============================================================================================================
Start --> End
Total (sum of squares): 6.443e-02 --> 9.624e-08,
Maximum absolute Force error: 1.023e+06 --> 4.454e+02 (N)
Minimum absolute Force error: 6.742e+00 --> 1.196e-03 (N)
Average absolute Force error: 5.609e+04 --> 6.862e+01 (N)
Maximum absolute Force error: 8.227e-02 --> 3.582e-05 (normalized)
Minimum absolute Force error: 5.422e-07 --> 9.616e-11 (normalized)
Average absolute Force error: 4.511e-03 --> 5.519e-06 (normalized)
R boundary error: 0.000e+00 --> 0.000e+00 (m)
Z boundary error: 0.000e+00 --> 0.000e+00 (m)
Fixed Psi error: 0.000e+00 --> 0.000e+00 (Wb)
Fixed pressure profile error: 0.000e+00 --> 5.569e-13 (Pa)
Fixed current profile error: 0.000e+00 --> 0.000e+00 (A)
Fixed sheet current error: 0.000e+00 --> 0.000e+00 (~)
==============================================================================================================
After re-solving, we find the force blance residuals are much lower, less than 1% normalized errors throughout the volume:
[25]:
plot_section(eq, "|F|_normalized", log=True);
