def F(r, zold):
""" Cost function for implicit step in r """
[H, pth, vpar, dHdx, dHdpph, dpthdx, dpthdpph, dvpardx,
dvpardpph] = get_der(array([r[0], zold[1], zold[2], pph0]))
return (r - zold[0]) + dtau * dHdx[1]
def zdot(t, z):
global Ath, dAth, Aph, dAph, Bmod, dBmod, hth, dhth, hph, dhph,\
H, dH, pth, pph, dpth, vpar, dvpar
[H, pth, vpar, dHdx, dHdpph, dpthdx, dpthdpph, dvpardx,
dvpardpph] = get_der(array([z[0], z[1], ph0, pph0]))
ret = zeros(2)
ret[0] = -dHdx[1] / dpthdx[0]
ret[1] = dHdx[0] / dpthdx[0]
return ret
def F(x, thold, pthold):
""" Cost function in x=(r,th) ignoring toroidal angle ph in axisymmetry"""
[H, pth, vpar, dHdx, dHdpph, dpthdx, dpthdpph, dvpardx,
dvpardpph] = get_der(array([x[0], x[1], 0.0, pph0]))
ret = zeros(2)
ret[0] = pth - pthold
ret[1] = dpthdx[0] * (x[1] - thold) - dt * dHdx[0]
return ret
def Fr(r, q, pthold, pdotold):
""" Cost function in r for axisymmetric field with pph=const. """
# Tokamak, no change in p_phi
[H, pth, vpar, dHdx, dHdpph, dpthdx, dpthdpph, dvpardx,
dvpardpph] = get_der(array([r[0], q[0], 0.0, pph0]))
ret = pthold - pth + dt / 2.0 * pdotold - dt / 2.0 * (
dHdx[1] - dHdx[0] / dpthdx[0] * dpthdx[1])
return ret
def Fx(x, qold, pthold):
""" Cost function in x=(r,th) ignoring toroidal angle ph in axisymmetry"""
[H, pth, vpar, dHdx, dHdpph, dpthdx, dpthdpph, dvpardx,
dvpardpph] = get_der(array([x[0], x[1], 0.0, pph0]))
ret = zeros(2)
ret[0] = pthold - pth
ret[1] = qold[0] - x[1] + dt / 2.0 * dHdx[0] / dpthdx[0]
return ret
def F(x, xold, pthold):
""" Cost function for implicit midpoint rule in axisymmetric field """
global pth
ret = zeros(3)
# evaluate at midpoint
[H, pth, vpar, dHdx, dHdpph, dpthdx, dpthdpph, dvpardx,
dvpardpph] = get_der(array([x[2], 0.5 * (x[1] + xold[1]), 0.0, pph0]))
ret[0] = dpthdx[0] * (x[1] - xold[1]) - dt * dHdx[0]
dpthdrmid = dpthdx[0]
pthdotbar = dpthdrmid * dHdx[1] - dpthdx[1] * dHdx[0]
ret[1] = dpthdrmid * (pth - pthold) + dt / 2.0 * pthdotbar
# evaluate at final position
[H, pth, vpar, dHdx, dHdpph, dpthdx, dpthdpph, dvpardx,
dvpardpph] = get_der(array([x[0], x[1], x[2], pph0]))
ret[2] = dpthdrmid * (pth - pthold) + dt * pthdotbar
return ret
#%%
from time import time
tic = time()
[H, pth, vpar] = get_val(array([r0, th0, ph0, pph0]))
for kt in range(nt):
pthold = pth
# Initialize via Lagrange extrapolation
if (kt >= nlag):
x0 = zeros(2)
extrapr0 = lagrange(arange(-nlag, 1), z[0, kt - nlag:kt + 1])
extrapr1 = lagrange(arange(-nlag, 1), z[1, kt - nlag:kt + 1])
x0[0] = extrapr0(1.0)
x0[1] = extrapr1(1.0)
else:
x0 = array([z[0, kt], z[1, kt]])
sol = root(F, x0, method='hybr', tol=1e-12, args=(z[1, kt], pthold))
z[:2, kt + 1] = sol.x
# Tokamak, no change in p_phi
[H, pth, vpar, dHdx, dHdpph, dpthdx, dpthdpph, dvpardx,
dvpardpph] = get_der(array([sol.x[0], sol.x[1], 0.0, pph0]))
pth = pth - dt * (dHdx[1] - dHdx[0] * dpthdx[1] / dpthdx[0])
print('Field evaluations: {}'.format(common.neval))
print('Time taken: {}'.format(time() - tic))
plot_orbit(z)
from numpy import array, zeros, arange
from scipy.optimize import root
from scipy.interpolate import lagrange
import common
from common import r0, th0, ph0, pph0, timesteps, get_der
from plotting import plot_orbit
dt, nt = timesteps(steps_per_bounce=11, nbounce=100)
nlag = 2
z = zeros([3, nt + 1])
z[:, 0] = [r0, th0, ph0]
[H, pth, vpar, dHdx, dHdpph, dpthdx, dpthdpph, dvpardx,
dvpardpph] = get_der(array([r0, th0, ph0, pph0]))
dtau = dt / dpthdx[0] # step in orbit parameter tau
t = zeros(nt + 1)
def F(r, zold):
""" Cost function for implicit step in r """
[H, pth, vpar, dHdx, dHdpph, dpthdx, dpthdpph, dvpardx,
dvpardpph] = get_der(array([r[0], zold[1], zold[2], pph0]))
return (r - zold[0]) + dtau * dHdx[1]
#%%
def F(r, q, pthold):
""" Cost function in r for axisymmetric field with pph=const. """
[H, pth, vpar, dHdx, dHdpph, dpthdx,
dpthdpph, dvpardx, dvpardpph] = get_der(array([r[0], q[0], q[1], pph0]))
return dpthdx[0]*(pth - pthold) + dt*(dHdx[1]*dpthdx[0]-dHdx[0]*dpthdx[1])
#%%
from time import time
tic = time()
[H, pth, vpar] = get_val(array([r0,th0,ph0,pph0]))
for kt in range(nt):
pthold = pth
# Initialize via Lagrange extrapolation
if(kt>=nlag):
extrapr = lagrange(arange(-nlag, 1), z[0, kt-nlag:kt+1])
r0 = extrapr(1)
else:
r0 = z[0,kt]
# Implicit substep in r
sol = root(F, r0, method='hybr',tol=1e-12,args=(z[1:,kt], pthold))
z[0,kt+1] = sol.x
# Explicit substep in q = (th, ph)
[H, pth, vpar, dHdx, dHdpph, dpthdx, dpthdpph,
dvpardx, dvpardpph] = get_der(array([sol.x[0], z[1,kt], z[2,kt], pph0]))
z[1,kt+1] = z[1,kt] + dt*dHdx[0]/dpthdx[0]
z[2,kt+1] = z[2,kt] + dt*vpar/f.hph
print('Field evaluations: {}'.format(common.neval))
print('Time taken: {}'.format(time()-tic))
plot_orbit(z)
plot_cost_function(F, z[:,-2], z[:,-1], pthold)
extrapri = lagrange(arange(-nlag, 1), z[2, kt - nlag:kt + 1])
extrapthi = lagrange(arange(-nlag, 1), z[3, kt - nlag:kt + 1])
x0i = array([extrapri(1.0), extrapthi(1.0)])
extrapr = lagrange(arange(-nlag, 1), z[0, kt - nlag:kt + 1])
r0 = extrapr(1.0)
else:
x0i = z[:2, kt]
r0 = z[0, kt]
# Implicit step
sol = root(Fx, x0i, method='hybr', tol=1e-12, args=(zold[1:2], pthold))
xi = sol.x
z[2:, kt + 1] = xi
[H, pth, vpar, dHdx, dHdpph, dpthdx, dpthdpph, dvpardx,
dvpardpph] = get_der(array([xi[0], xi[1], 0.0, pph0]))
pdotold = -(dHdx[1] - dHdx[0] / dpthdx[0] * dpthdx[1])
pthold = pth
sol = root(Fr,
r0,
method='hybr',
tol=1e-12,
args=(xi[1:], pthold, pdotold))
z[0, kt + 1] = sol.x
[H, pth, vpar, dHdx, dHdpph, dpthdx, dpthdpph, dvpardx,
dvpardpph] = get_der(array([sol.x[0], xi[1], 0.0, pph0]))
z[1, kt + 1] = xi[1] + dt / 2.0 * dHdx[0] / dpthdx[0]
zold = z[:, kt + 1]