Solving Equilibria with Fixed Axis and Fixed NAE O(rho) Behavior in DESC
This tutorial shows how to find equilibrium solutions in DESC which are constrained to have the same axis and near-axis behavior as a NAE solution from the pyQSC code
Creating a DESC Equilibrium from a pyQsc Near-Axis Equilibrium
Note that you must have pyQsc installed in order to make use of the Equilibrium.from_near_axis
method, do so with pip install qsc
[1]:
# must have installed pyQsc with `pip install qsc` in order to use this!
from qsc import Qsc
import numpy as np
import matplotlib.pyplot as plt
from desc.equilibrium import Equilibrium
from desc.objectives import get_fixed_boundary_constraints
from desc.plotting import (
plot_comparison,
plot_fsa,
plot_section,
plot_surfaces,
plot_qs_error,
)
DESC version 0.8.1+11.g3876303b.dirty,using JAX backend, jax version=0.2.25, jaxlib version=0.1.76, dtype=float64
Using device: CPU, with 22.92 GB available memory
DESC is able to create an equilibrium based off of a pyQsc
NAE equilibrium object. First, we’ll make the NAE equilibrium using pyQsc
[2]:
qsc_eq = Qsc.from_paper("precise QA")
Then, to make the DESC equilibrium, the Equilibrium
class has a method Equilibrium.from_near_axis
. This method creates a DESC Equilibrium
based off of the pyQsc
equilibrium. It requires as input the desired DESC FourierZernike resolution, as well as the radius at which you want to evaluate the qsc equilibrium at to make the DESC equilibrium’s boundary. The equilibrium’s initial R_lmn
, Z_lmn
Fourier-Zernike coefficients are fit to the R,Z
evaluated from the pyQsc
equilibrium, and the initial L_lmn
are 0 (because the pyQsc
equilibrium uses Boozer angles, so there is no poloidal stream function)
[3]:
ntheta = 75
r = 0.35
desc_eq = Equilibrium.from_near_axis(
qsc_eq, # the Qsc equilibrium object
r=r, # the finite radius (m) at which to evaluate the Qsc surface to use as the DESC boundary
L=8, # DESC radial resolution
M=8, # DESC poloidal resolution
N=8, # DESC toroidal resolution
ntheta=ntheta,
)
eq_fit = desc_eq.copy() # copy so we can see the original Qsc surfaces later
Now we solve the equilibrium as normal in DESC
[4]:
# get the fixed-boundary constraints, which include also fixing the pressure and fixing the current profile (iota=False flag means fix current)
constraints = get_fixed_boundary_constraints(iota=False)
print(constraints)
# solve the equilibrium
desc_eq.solve(
verbose=3,
ftol=1e-2,
objective="force",
maxiter=100,
xtol=1e-6,
constraints=constraints,
)
# Save equilibrium as .h5 file
desc_eq.save("DESC_from_NAE_precise_QA_output.h5")
(<desc.objectives.linear_objectives.FixBoundaryR object at 0x7f9dc0069e20>, <desc.objectives.linear_objectives.FixBoundaryZ object at 0x7f9dc0069eb0>, <desc.objectives.linear_objectives.FixLambdaGauge object at 0x7f9dc0069f40>, <desc.objectives.linear_objectives.FixPsi object at 0x7f9dc0069fa0>, <desc.objectives.linear_objectives.FixPressure object at 0x7f9dc0069fd0>, <desc.objectives.linear_objectives.FixCurrent object at 0x7f9dc0069ca0>)
Building objective: force
Precomputing transforms
Timer: Precomputing transforms = 307 ms
Timer: Objective build = 2.51 sec
Timer: Linear constraint projection build = 5.52 sec
Compiling objective function and derivatives
Timer: Objective compilation time = 3.43 sec
Timer: Jacobian compilation time = 9.10 sec
Timer: Total compilation time = 12.5 sec
Number of parameters: 856
Number of objectives: 5346
Starting optimization
Iteration Total nfev Cost Cost reduction Step norm Optimality
0 1 4.1581e+01 2.60e+04
1 2 4.7097e+00 3.69e+01 3.13e-01 2.71e+03
2 3 1.1553e-01 4.59e+00 1.51e-01 8.43e+02
3 4 1.8096e-02 9.74e-02 1.86e-01 1.04e+02
4 6 2.7943e-03 1.53e-02 4.63e-02 5.28e+01
5 8 2.1187e-03 6.76e-04 2.43e-02 1.03e+01
6 10 2.0643e-03 5.43e-05 1.09e-02 1.31e+00
7 11 2.0204e-03 4.40e-05 1.71e-02 6.25e+00
8 13 1.9786e-03 4.18e-05 8.38e-03 1.27e+00
9 14 1.9432e-03 3.54e-05 1.53e-02 5.92e+00
10 16 1.9076e-03 3.56e-05 8.29e-03 1.23e+00
11 17 1.8775e-03 3.01e-05 1.46e-02 4.24e+00
12 19 1.8522e-03 2.53e-05 7.45e-03 6.78e-01
13 20 1.8286e-03 2.36e-05 1.38e-02 2.13e+00
14 22 1.8113e-03 1.73e-05 7.04e-03 2.87e-01
Optimization terminated successfully.
`ftol` condition satisfied.
Current function value: 1.811e-03
Total delta_x: 3.010e-01
Iterations: 14
Function evaluations: 22
Jacobian evaluations: 15
Timer: Solution time = 49.9 sec
Timer: Avg time per step = 3.33 sec
Start of solver
Total (sum of squares): 4.158e+01,
Total force: 2.015e+05 (N)
Total force: 9.119e+00 (normalized)
End of solver
Total (sum of squares): 1.811e-03,
Total force: 1.330e+03 (N)
Total force: 6.019e-02 (normalized)
Now we have a DESC equilibrium solved with the boundary from pyQsc
. It has zero toroidal current as its profile constraint along with zero pressure since the original Qsc equilibrium had 0 pressure and current.
However, if we plot the surfaces, we see that while the boundary matches, the interior deviates slightly from the NAE, especially near the core. This means we may have lost some of the optimized properties from the QSC equilibrium.
[5]:
plot_comparison(
eqs=[desc_eq, eq_fit],
labels=["DESC-solved equilibrium", "NAE surfaces"],
figsize=(12, 12),
theta=0,
colors=["k", "r"],
linestyles=["-", ":"],
lws=[1, 2],
);
Instead of evaluating the NAE at a finite radius, fixing that boundary, and hoping the axis behavior stays the same, we can instead directly fix the axis and the \(O(\rho)\) asymptotic behavior.
Solving Equilibria with Fixed Axis and Fixed NAE \(O(\rho)\) Behavior in DESC
[6]:
# utility functions for getting the NAE constraints
from desc.objectives.utils import get_equilibrium_objective, get_NAE_constraints
eq_NAE = eq_fit.copy()
# this has all the constraints we need, iota=False specifies we want to fix current instead of iota
constraints = get_NAE_constraints(eq_NAE, qsc_eq, iota=False, order=1)
eq_NAE.solve(
verbose=3,
ftol=1e-2,
objective="force",
maxiter=50,
xtol=1e-6,
constraints=constraints,
);
Building objective: force
Precomputing transforms
Timer: Precomputing transforms = 209 ms
Timer: Objective build = 503 ms
Timer: Linear constraint projection build = 4.08 sec
Compiling objective function and derivatives
Timer: Objective compilation time = 3.52 sec
Timer: Jacobian compilation time = 9.43 sec
Timer: Total compilation time = 12.9 sec
Number of parameters: 1094
Number of objectives: 5346
Starting optimization
Iteration Total nfev Cost Cost reduction Step norm Optimality
0 1 4.1717e+01 3.62e+04
1 3 1.6806e+01 2.49e+01 2.92e-01 3.59e+04
2 4 8.5057e-01 1.60e+01 1.78e-01 2.92e+03
3 5 3.3639e-01 5.14e-01 2.18e-01 5.29e+02
4 6 8.0149e-02 2.56e-01 2.20e-01 2.95e+02
5 7 7.1672e-02 8.48e-03 1.15e-01 8.22e+02
6 8 7.8575e-04 7.09e-02 3.89e-02 4.86e+01
7 10 3.5139e-04 4.34e-04 1.72e-02 2.24e+01
8 11 1.6654e-04 1.85e-04 2.00e-02 1.28e+01
9 12 6.8569e-05 9.80e-05 1.92e-02 1.03e+01
10 14 1.3285e-06 6.72e-05 4.84e-03 6.41e-01
11 16 9.0805e-07 4.20e-07 2.31e-03 1.69e-01
12 18 8.5728e-07 5.08e-08 1.14e-03 3.76e-02
13 19 8.0738e-07 4.99e-08 2.23e-03 1.88e-01
14 21 7.6094e-07 4.64e-08 1.10e-03 4.52e-02
15 22 7.1965e-07 4.13e-08 2.16e-03 2.25e-01
16 23 6.6672e-07 5.29e-08 2.12e-03 2.45e-01
17 24 6.1808e-07 4.86e-08 2.09e-03 2.69e-01
18 25 5.7333e-07 4.47e-08 2.06e-03 2.93e-01
19 26 5.3226e-07 4.11e-08 2.03e-03 3.16e-01
20 27 4.9465e-07 3.76e-08 2.01e-03 3.38e-01
21 28 4.6021e-07 3.44e-08 1.98e-03 3.58e-01
22 29 4.2864e-07 3.16e-08 1.97e-03 3.75e-01
23 30 3.9957e-07 2.91e-08 1.95e-03 3.89e-01
24 31 3.7269e-07 2.69e-08 1.93e-03 3.99e-01
25 32 3.4768e-07 2.50e-08 1.91e-03 4.05e-01
26 33 3.2427e-07 2.34e-08 1.89e-03 4.07e-01
27 34 3.0228e-07 2.20e-08 1.87e-03 4.05e-01
28 35 2.8152e-07 2.08e-08 1.85e-03 3.99e-01
29 36 2.6189e-07 1.96e-08 1.83e-03 3.89e-01
30 37 2.4330e-07 1.86e-08 1.81e-03 3.77e-01
31 38 2.2572e-07 1.76e-08 1.79e-03 3.61e-01
32 39 2.0912e-07 1.66e-08 1.76e-03 3.44e-01
33 40 1.9352e-07 1.56e-08 1.74e-03 3.26e-01
34 41 1.7891e-07 1.46e-08 1.71e-03 3.07e-01
35 42 1.6525e-07 1.37e-08 1.68e-03 2.88e-01
36 43 1.5245e-07 1.28e-08 1.65e-03 2.70e-01
37 44 1.4038e-07 1.21e-08 1.62e-03 2.52e-01
38 45 1.2894e-07 1.14e-08 1.59e-03 2.34e-01
39 46 1.1815e-07 1.08e-08 1.57e-03 2.18e-01
40 47 1.0826e-07 9.89e-09 1.56e-03 2.07e-01
41 48 9.9560e-08 8.70e-09 1.56e-03 2.04e-01
42 49 9.2176e-08 7.38e-09 1.59e-03 2.08e-01
43 50 8.5988e-08 6.19e-09 1.62e-03 2.16e-01
Warning: Maximum number of function evaluations has been exceeded.
Current function value: 8.599e-08
Total delta_x: 3.795e-01
Iterations: 43
Function evaluations: 50
Jacobian evaluations: 44
Timer: Solution time = 8.73 min
Timer: Avg time per step = 11.9 sec
Start of solver
Total (sum of squares): 4.158e+01,
Total force: 2.015e+05 (N)
Total force: 9.119e+00 (normalized)
End of solver
Total (sum of squares): 8.599e-08,
Total force: 9.161e+00 (N)
Total force: 4.147e-04 (normalized)
Again we can plot the surfaces, and we see that in this case, the near axis behavior is preserved, while the outer surfaces differ. This is expected, as the NAE from QSC is only valid asymptotically near the axis. By constraining this behavior we are able to keep all the desireable properties of the NAE where they are valid without overly constraining the problem.
[7]:
fig, ax = plot_comparison(
eqs=[eq_fit, eq_NAE],
labels=[f"NAE Surfaces r={r}", f"DESC NAE constrained"],
colors=["k", "g"],
linestyles=["-", "--"],
figsize=(12, 12),
theta=0,
lws=[1, 2],
)
fig, ax = plot_comparison(
eqs=[desc_eq, eq_NAE],
labels=[f"DESC Fixed Surface", f"DESC NAE constrained"],
colors=["r", "g"],
linestyles=[":", "--"],
figsize=(12, 12),
theta=0,
lws=[1, 2],
);
As a final comparison, we can plot the QS error for the fixed boundary and fixed NAE solves. We see that by fixing the NAE behavior, we are able to preserve the QS from the original QSC equilibrium.
[8]:
fix, ax = plt.subplots()
plot_qs_error(
desc_eq,
fC=False,
fT=False,
log=True,
ax=ax,
colors=["r"],
legend=False,
labels=["Fixed boundary"],
)
plot_qs_error(
eq_NAE,
fC=False,
fT=False,
log=True,
ax=ax,
colors=["g"],
legend=False,
labels=["Fixed NAE"],
)
ax.legend()
ax.set_ylabel("$\sum |B_{mn}|$")
[8]:
Text(0, 0.5, '$\\sum |B_{mn}|$')
The correct behaviour near the axis may be also checked by assessing directly other features of the equilibrium; most notably, the behaviour of |B| on axis, the rotational transform on axis, as well as the form of \(\lambda\). Starting from the rotational transform, we may compare the value obtained to that of the near-axis expansion,
[9]:
import matplotlib.pyplot as plt
from desc.plotting import plot_fsa
rho = np.linspace(1e-3,1e-1)
fig, ax, iota_nae = plot_fsa(eq_NAE,"iota",rho=rho, return_data=True)
fig, ax, iota_surf = plot_fsa(desc_eq,"iota",rho=rho, ax=ax, return_data=True, linecolor='orange')
plt.plot(rho, np.ones(np.size(rho))*qsc_eq.iota,linestyle='--')
plt.legend(['Fixed NAE', 'Fixed boundary', 'NAE'])
print('Relative error in the rotational transform (fixed NAE): ', (iota_nae['iota'][0]-qsc_eq.iota)/qsc_eq.iota)
print('Relative error in the rotational transform (fixed surface): ', (iota_surf['iota'][0]-qsc_eq.iota)/qsc_eq.iota)
# print(iota[1].yaxis)
Relative error in the rotational transform (fixed NAE): -3.252459192718128e-05
Relative error in the rotational transform (fixed surface): 0.10621309580868808
This clearly shows that the constraint is working as it is intended. A similar expected behaviour can be found for other features of the equilibrium such as the magnetic field on axis.
[10]:
from desc.compute import data_index
from desc.grid import LinearGrid
grid = LinearGrid(M=desc_eq.M_grid, N=desc_eq.N_grid, NFP=desc_eq.NFP, rho=np.array(1e-6))
# Evaluate B modes near the axis
data_surf = desc_eq.compute(["|B|_mn", "B modes"], grid=grid)
data_nae = eq_NAE.compute(["|B|_mn", "B modes"], grid=grid)
modes = data_surf["B modes"]
B_mn_surf = data_surf["|B|_mn"]
B_mn_nae = data_nae["|B|_mn"]
# Evaluate B on an angular grid
theta = np.linspace(0,2*np.pi,150)
phi = np.linspace(0,2*np.pi,100)
th, ph = np.meshgrid(theta,phi)
B_surf = np.zeros((100,150))
B_nae = np.zeros((100,150))
idx = np.where(modes[:,2] !=0)[0]
# Print the deviation of B from QS (that is, the variation of B respect to its average along the axis)
print('Deviation from QS (fixed surface): ', np.sqrt(np.sum(B_mn_surf[idx]*B_mn_surf[idx]))/np.sqrt(np.sum(B_mn_surf*B_mn_surf)))
print('Deviation from QS (fixed nae): ', np.sqrt(np.sum(B_mn_nae[idx]*B_mn_nae[idx]))/np.sqrt(np.sum(B_mn_nae*B_mn_nae)))
for i, (l,m,n) in enumerate(modes):
if m>=0 and n>=0:
B_surf += B_mn_surf[i]*np.cos(m*th)*np.cos(n*ph)
B_nae += B_mn_nae[i]*np.cos(m*th)*np.cos(n*ph)
elif m>=0 and n<0:
B_surf += -B_mn_surf[i]*np.cos(m*th)*np.sin(n*ph)
B_nae += -B_mn_nae[i]*np.cos(m*th)*np.sin(n*ph)
elif m<0 and n>=0:
B_surf += -B_mn_surf[i]*np.sin(m*th)*np.cos(n*ph)
B_nae += -B_mn_nae[i]*np.sin(m*th)*np.cos(n*ph)
elif m<0 and n<0:
B_surf += B_mn_surf[i]*np.sin(m*th)*np.sin(n*ph)
B_nae += B_mn_nae[i]*np.sin(m*th)*np.sin(n*ph)
# Eliminate the poloidal angle to focus on the toroidal behaviour
B_av_surf = np.mean(B_surf,axis=1)
B_av_nae = np.mean(B_nae,axis=1)
plt.plot(phi,B_av_surf)
plt.plot(phi,B_av_nae)
plt.plot(phi,np.ones(np.size(phi))*qsc_eq.B0,linestyle='--')
plt.xlabel('$\phi$')
plt.ylabel('$|B|$')
plt.legend(['Fixed surface', 'Fixed NAE','NAE'])
plt.show()
Deviation from QS (fixed surface): 0.006061270206450826
Deviation from QS (fixed nae): 9.28402496619238e-06
Finally, we may check \(\lambda\) on axis and compare it to \(\nu\) from within the near-axis expansion.
[11]:
grid_2d_05 = LinearGrid(rho=np.array(1e-6), M=50, N=50, NFP=desc_eq.NFP, endpoint=True)
# Evaluate lambda near the axis
data_surf = desc_eq.compute("lambda", grid=grid_2d_05)
data_nae = eq_NAE.compute("lambda", grid=grid_2d_05)
lam_surf = data_surf["lambda"]
lam_nae = data_nae["lambda"]
# Reshape to form grids on theta and phi
zeta = (
grid_2d_05.nodes[:, 2]
.reshape((grid_2d_05.num_theta, grid_2d_05.num_rho, grid_2d_05.num_zeta), order="F")
.squeeze()
)
lam_surf = lam_surf.reshape(
(grid_2d_05.num_theta, grid_2d_05.num_rho, grid_2d_05.num_zeta), order="F"
)
lam_nae = lam_nae.reshape(
(grid_2d_05.num_theta, grid_2d_05.num_rho, grid_2d_05.num_zeta), order="F"
)
phi = np.squeeze(zeta[0, :])
lam_surf = np.squeeze(lam_surf[:, 0, :])
lam_nae = np.squeeze(lam_nae[:, 0, :])
# Eliminate the poloidal angle to focus on the toroidal behaviour
lam_av_surf = np.mean(lam_surf,axis=0)
lam_av_nae = np.mean(lam_nae,axis=0)
print('Deviation of theta from Boozer angle (fixed surface)', np.mean(np.abs(lam_av_surf+qsc_eq.iota*qsc_eq.nu_spline(phi))))
print('Deviation of theta from Boozer angle (fixed nae)', np.mean(np.abs(lam_av_nae+qsc_eq.iota*qsc_eq.nu_spline(phi))))
plt.plot(phi,lam_av_surf)
plt.plot(phi,lam_av_nae)
plt.plot(phi,-qsc_eq.iota*qsc_eq.nu_spline(phi),linestyle='--')
plt.xlabel('$\phi$')
plt.ylabel('$\lambda$')
plt.legend(['Fixed surface', 'Fixed NAE','NAE'])
plt.show()
Deviation of theta from Boozer angle (fixed surface) 0.08181754718373944
Deviation of theta from Boozer angle (fixed nae) 4.670576901714567e-05
The above comparison shows that the DESC equilibrium constructed with the near-axis constraint is consistent with the near-axis behaviour.
[ ]: