Using DESC Interactively (Using Spline Basis for profiles)
[2]:
%matplotlib inline
import numpy as np
from desc.equilibrium import Equilibrium
from desc.geometry import FourierRZToroidalSurface
from desc.profiles import SplineProfile
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 5.04 GB available memory
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 \(\phi\). We specify the mode numbers for R and Z as 2d arrays of [m,n] pairs, and the coefficients as a 1d array.
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, which gives the following surface:
[3]:
surface = FourierRZToroidalSurface(
R_lmn=[10, 1],
modes_R=[[0, 0], [1, 0]],
Z_lmn=[0, -1],
modes_Z=[[0, 0], [-1, 0]],
)
Next, we need to define the profiles. This time we will use a spline profile. We’ll take a vacuum case to start (pressure=0), and a simple quadratic \(\iota\) profile. The profiles are given as SplineProfile objects, which are specified by values at the knots (in \(\rho\)) locations. The default splines used are cubic \(C^2\) (natural) splines, though other options are available:
[4]:
knots = np.linspace(0, 1, 20)
pressure_values = np.zeros_like(knots)
iota_values = 1 + 1.5 * knots**2
pressure = SplineProfile(values=pressure_values, knots=knots)
iota = SplineProfile(values=iota_values, knots=knots)
print("Pressure rho knots:\n", pressure._knots)
# params here are the values at the knot locations
print("Pressure values at these knots:\n", pressure.params)
print("Rotational Transform rho knots:\n", iota._knots)
# params here are the values at the knot locations
print("Rotational Transform values at these knots:\n", iota.params)
Pressure rho knots:
[0. 0.05263158 0.10526316 0.15789474 0.21052632 0.26315789
0.31578947 0.36842105 0.42105263 0.47368421 0.52631579 0.57894737
0.63157895 0.68421053 0.73684211 0.78947368 0.84210526 0.89473684
0.94736842 1. ]
Pressure values at these knots:
[0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0.]
Rotational Transform rho knots:
[0. 0.05263158 0.10526316 0.15789474 0.21052632 0.26315789
0.31578947 0.36842105 0.42105263 0.47368421 0.52631579 0.57894737
0.63157895 0.68421053 0.73684211 0.78947368 0.84210526 0.89473684
0.94736842 1. ]
Rotational Transform values at these knots:
[1. 1.00415512 1.0166205 1.03739612 1.06648199 1.10387812
1.14958449 1.20360111 1.26592798 1.3365651 1.41551247 1.50277008
1.59833795 1.70221607 1.81440443 1.93490305 2.06371191 2.20083102
2.34626039 2.5 ]
Finally, we create an Equilibrium by giving it the surface and profiles, as well as specifying what resolution we want to use and a few other parameters:
[5]:
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=6, # radial spectral resolution
M=6, # poloidal spectral resolution
N=0, # toroidal spectral resolution (it's a tokamak, so we don't need any toroidal modes)
L_grid=12, # real space radial resolution, slightly oversampled
M_grid=9, # real space poloidal resolution, slightly oversampled
N_grid=0, # real space toroidal resolution
sym=True, # explicitly enforce stellarator symmetry
)
we can plot these profiles and see that they fit the original power-series expressions well
[6]:
fig, ax = plot_1d(eq=eq, name="iota", label="spline")
rho = np.linspace(0, 1, 100)
ax.plot(rho, 1 + 1.5 * rho**2, "--", c="r", label="power series")
ax.legend()
# make sure values are the same
np.testing.assert_array_almost_equal(iota.params, 1 + 1.5 * iota._knots**2)
fig, ax2 = plot_1d(eq=eq, name="p", label="spline")
ax2.plot(rho, np.zeros_like(rho), "--", c="r", label="power series")
ax2.legend()
# make sure values are the same
np.testing.assert_array_almost_equal(pressure.params, np.zeros_like(pressure.params))
This automatically creates the spectral bases and generates an initial guess for the flux surfaces by scaling the boundary surface, which we can plot below:
Plotting
[7]:
# plot_surfaces generates poincare plots of the flux surfaces
plot_surfaces(eq);
We can also look at the force balance error:
[8]:
# 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 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
[9]:
from desc.optimize import Optimizer
# create an Optimizer object using the lsq-exact option
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)
[10]:
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 final ObjectiveFunction object which we can pass to the eq.solve method
obj = ObjectiveFunction(objectives=objectives)
a utility function exists that will create the typical objective function for a fixed boundary equilibrium (the above cell’s contents)
[11]:
constraints2 = get_fixed_boundary_constraints()
print("Same constraints for each")
for c1, c2 in zip(constraints, constraints2):
print("type of ", c1, "=", "type of ", c2, ":", type(c1) == type(c2))
Same constraints for each
type of <desc.objectives.linear_objectives.FixBoundaryR object at 0x7fb0dffb6b20> = type of <desc.objectives.linear_objectives.FixBoundaryR object at 0x7fb0dffd34c0> : True
type of <desc.objectives.linear_objectives.FixBoundaryZ object at 0x7fb0dffb61f0> = type of <desc.objectives.linear_objectives.FixBoundaryZ object at 0x7fb0dffd32b0> : True
type of <desc.objectives.linear_objectives.FixPressure object at 0x7fb0dffb6340> = type of <desc.objectives.linear_objectives.FixPsi object at 0x7fb0dffd3fa0> : False
type of <desc.objectives.linear_objectives.FixIota object at 0x7fb0dffb6760> = type of <desc.objectives.linear_objectives.FixPressure object at 0x7fb11c266b80> : False
type of <desc.objectives.linear_objectives.FixPsi object at 0x7fb0dffb6070> = type of <desc.objectives.linear_objectives.FixIota object at 0x7fb11c266160> : False
Next, we simply call eq.solve() to minimize the force balance error. Here we can also pass in arguments to the optimizer such as maximum number of iterations or stopping tolerances (we only run a few iterations here to show the idea). 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.
[12]:
eq.solve(
verbose=2, ftol=1e-8, objective=obj, optimizer=optimizer, constraints=constraints
);
Building objective: force
Precomputing transforms
Timer: Precomputing transforms = 894 ms
Timer: Objective build = 5.90 sec
Timer: Linear constraint projection build = 13.9 sec
Compiling objective function and derivatives
Timer: Objective compilation time = 4.83 sec
Timer: Jacobian compilation time = 8.47 sec
Timer: Total compilation time = 13.3 sec
Number of parameters: 27
Number of objectives: 80
Starting optimization
Iteration Total nfev Cost Cost reduction Step norm Optimality
0 1 6.8906e-03 1.40e+00
1 2 2.4861e-03 4.40e-03 2.96e-01 2.73e-01
2 3 1.9680e-04 2.29e-03 1.38e-01 8.07e-02
3 4 8.0793e-06 1.89e-04 2.69e-02 1.17e-02
4 5 4.7022e-07 7.61e-06 1.33e-02 2.06e-03
5 6 6.1976e-08 4.08e-07 4.63e-03 2.67e-04
6 7 4.6435e-08 1.55e-08 4.96e-03 2.71e-04
7 8 4.3575e-08 2.86e-09 4.86e-03 7.32e-05
8 9 4.3225e-08 3.50e-10 1.01e-03 7.27e-06
9 10 4.3116e-08 1.09e-10 1.17e-03 7.90e-06
10 11 4.3072e-08 4.44e-11 5.31e-04 2.33e-06
11 12 4.3051e-08 2.05e-11 4.24e-04 1.39e-06
12 13 4.3041e-08 1.02e-11 2.64e-04 6.14e-07
13 14 4.3036e-08 5.39e-12 1.92e-04 3.44e-07
14 15 4.3033e-08 2.98e-12 1.35e-04 1.86e-07
15 16 4.3031e-08 1.70e-12 9.93e-05 1.10e-07
16 17 4.3030e-08 1.00e-12 7.36e-05 6.68e-08
17 18 4.3029e-08 6.00e-13 5.57e-05 4.46e-08
18 19 4.3029e-08 3.67e-13 4.26e-05 3.17e-08
19 20 4.3029e-08 2.27e-13 3.30e-05 2.30e-08
20 21 4.3029e-08 1.42e-13 2.57e-05 1.71e-08
21 22 4.3029e-08 8.95e-14 2.02e-05 1.29e-08
22 23 4.3028e-08 5.68e-14 1.60e-05 9.83e-09
`gtol` condition satisfied.
Current function value: 4.303e-08
Total delta_x: 1.631e-01
Iterations: 22
Function evaluations: 23
Jacobian evaluations: 23
Timer: Solution time = 3.03 sec
Timer: Avg time per step = 131 ms
Start of solver
Total (sum of squares): 6.891e-03,
Total force: 3.436e+05 (N)
Total force: 1.174e-01 (normalized)
End of solver
Total (sum of squares): 4.303e-08,
Total force: 8.586e+02 (N)
Total force: 2.934e-04 (normalized)
We can then look at the flux surfaces and force error again:
[13]:
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.
[14]:
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:
[15]:
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 spline in \(\rho\), specified by the values at the knot locations)
If we want a quadratic pressure profile of the form \(p(\rho) = 1000(1-\rho^2)\), we can evaluate this profile at the knot locations and then pass those in as the delta values
[16]:
delta_p = np.zeros_like(eq.p_l)
p_values = 1000 * (1 - eq.pressure._knots**2)
delta_p = p_values
[17]:
eq1 = eq.perturb(deltas={"p_l": delta_p}, order=2)
Building objective: force
Precomputing transforms
Perturbing p_l
Factorizing linear constraints
Computing df
Factoring df
Computing d^2f
||dx||/||x|| = 2.741e-03
Note that this gives us back a new equilibrium, so that the original is saved for future study.
[18]:
# make sure values are the same
np.testing.assert_array_almost_equal(
eq1.pressure.params, 1000 * (1 - eq.pressure._knots**2)
)
plot_1d(eq1, "p");
[19]:
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%
[20]:
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
[21]:
# need to remake the constraints and objective, as the equilibrium now has finite pressure and so
# the constraints and objective must be rebuilt to account for that
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)
# the default ObjectiveFunction is force balance, and the default optimizer is 'lsq-exact'
eq1.solve(
verbose=2, ftol=1e-4, optimizer=optimizer, constraints=constraints, objective=obj
);
Building objective: force
Precomputing transforms
Timer: Precomputing transforms = 395 ms
Timer: Objective build = 750 ms
Timer: Linear constraint projection build = 2.69 sec
Compiling objective function and derivatives
Timer: Objective compilation time = 4.43 sec
Timer: Jacobian compilation time = 10.1 sec
Timer: Total compilation time = 14.5 sec
Number of parameters: 27
Number of objectives: 80
Starting optimization
Iteration Total nfev Cost Cost reduction Step norm Optimality
0 1 2.5027e-08 3.08e-04
1 2 2.4047e-08 9.81e-10 2.04e-03 3.03e-05
2 3 2.3817e-08 2.29e-10 1.58e-03 1.67e-05
3 4 2.3635e-08 1.82e-10 9.50e-04 8.03e-06
4 5 2.3379e-08 2.56e-10 1.11e-03 1.21e-05
5 6 2.2796e-08 5.84e-10 1.66e-03 2.77e-05
6 7 2.0643e-08 2.15e-09 3.58e-03 1.16e-04
7 9 1.7416e-08 3.23e-09 3.04e-03 4.98e-05
8 10 1.1598e-08 5.82e-09 6.33e-03 1.61e-04
9 11 7.6607e-09 3.94e-09 1.43e-02 3.19e-04
10 12 4.2197e-09 3.44e-09 3.08e-03 1.79e-05
11 13 4.2105e-09 9.23e-12 4.42e-04 2.52e-07
12 14 4.2103e-09 1.85e-13 8.13e-05 2.12e-08
Optimization terminated successfully.
`ftol` condition satisfied.
Current function value: 4.210e-09
Total delta_x: 3.022e-02
Iterations: 12
Function evaluations: 14
Jacobian evaluations: 13
Timer: Solution time = 1.24 sec
Timer: Avg time per step = 95.6 ms
Start of solver
Total (sum of squares): 2.503e-08,
Total force: 6.548e+02 (N)
Total force: 2.237e-04 (normalized)
End of solver
Total (sum of squares): 4.210e-09,
Total force: 2.686e+02 (N)
Total force: 9.176e-05 (normalized)
[22]:
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
[23]:
plot_section(eq1, "|F|", norm_F=True, log=True);
