diff --git a/python/DELETE.py b/python/DELETE.py new file mode 100644 index 00000000..cafe8e5c --- /dev/null +++ b/python/DELETE.py @@ -0,0 +1,44 @@ +#%% +import matplotlib.pyplot as plt +import numpy as np +import common +from cpp_sym_midpoint import cpp_sym_mid +from common import (r0,th0,ph0,vpar0,timesteps,get_val,get_der,mu,qe,m,c,) +from field_correct_test import field +from plotale import plot_orbit, plot_mani, plot_cost_function +#from manifolds import surf_R + +f = field() + +dt = 600 +nt = 1000 + +# Initial Conditions +z = np.zeros([3, nt + 1]) +z[:, 0] = [r0, th0, ph0] + +f.evaluate(r0, th0, ph0) +v = np.zeros(3) +p = np.zeros(3) + +v[0] = vpar0*f.co_hr +v[1] = vpar0*f.co_hth +v[2] = vpar0*f.co_hph +p[0] = v[0] +p[1] = v[1] + qe/c*f.co_Ath +p[2] = v[2] + qe/c*f.co_Aph + +P = np.zeros([3,nt+1]) +P[:,0] = p + +print(P) + +# Symplectic midpoint +for kt in range(nt): + + z[:,kt+1] = cpp_sym_mid(z[:,kt]) + + +plot_orbit(z) +plt.plot(z[0,:5]*np.cos(z[1,:5]), z[0,:5]*np.sin(z[1,:5]), "o") +plt.show() \ No newline at end of file diff --git a/python/cp_runge_kutta.py b/python/cp_runge_kutta.py new file mode 100644 index 00000000..582b6820 --- /dev/null +++ b/python/cp_runge_kutta.py @@ -0,0 +1,121 @@ +# %% +from scipy.integrate import solve_ivp +import matplotlib.pyplot as plt +import numpy as np +from scipy.optimize import root +from scipy.interpolate import lagrange +import common +from common import (f,r0,th0,ph0,pph0,timesteps,get_val,get_der,mu,qe,m,c,) +from plotting import plot_orbit, plot_cost_function + +B0 = 1.0 # magnetic field modulus normalization +iota0 = 1.0 # constant part of rotational transform +a = 0.5 # (equivalent) minor radius +R0 = 1.0 # (equivalent) major radius + +z0_cpp = np.array([r0, th0, ph0]) + +metric = lambda z: { + "_11": 1, + "_22": z[0]**2, + "_33": (1 + z[0]*np.cos(z[1]))**2, + "^11": 1, + "^22": 1/z[0]**2, + "^33": 1/(1 + z[0]*np.cos(z[1]))**2, + "d_11": [0,0,0], + "d_22": [2*z[0], 0,0], + "d_33": [2*(1+z[0]*np.cos(z[1]))*np.cos(z[1]), -2*(1+z[0]*np.cos(z[1]))*np.sin(z[1]), 0], +} + +def field_Ath(r, th, ph): + Ath = B0*(r**2/2 - r**3/(3*R0)*np.cos(th)) + dAth = np.array((B0*(r - r**2/R0*np.cos(th)), B0*r**3*np.sin(th)/(3*R0), 0)) + return Ath, dAth + +def field_Aph(r, th, ph): + Aph = -B0*iota0*(r**2/2-r**4/(4*a**2)) + dAph = np.array((-B0*iota0*(r-r**3/a**2), 0, 0)) + return Aph, dAph + + +def field_B(r, th, ph): + B = B0*(iota0*(r-r**3/a**2)/(R0+r*np.cos(th)) + (1-r*np.cos(th)/R0)) + #dB = np.array((iota0*(1-3*r**2/a**2)/(R0+r*np.cos(th)) - iota0*(r-r**3/a**2)/(R0+r*np.cos(th))**2 * np.cos(th) - np.cos(th)/R0, + # iota0*(r-r**3/a**2)/(R0+r*np.cos(th))**2 *r*np.sin(th) + r*np.sin(th), 0)) + return B + + +#Initial Conditions +Ath = np.empty(1) +Aph = np.empty(1) +dAth = np.empty(3) +dAph = np.empty(3) +B = np.empty(1) +#dB = np.empty(3) +Ath, dAth = field_Ath(r0, th0, ph0) +Aph, dAph = field_Aph(r0, th0, ph0) +B = field_B(r0, th0, ph0) + +g = metric(z0_cpp[:]) +ctrv = np.zeros(3) +ctrv[0] = np.sqrt(g['^11']*mu*2*B/3) +ctrv[1] = np.sqrt(g['^22']*mu*2*B/3) +ctrv[2] = np.sqrt(g['^33']*mu*2*B/3) +p = np.zeros(3) +p[0] = g['_11']*ctrv[0] +p[1] = g['_22']*ctrv[1] + qe/c * Ath +p[2] = g['_33']*ctrv[2] + qe/c * Aph +pold = p + +#Initial Conditions Runge-Kutta +z0_cpp = np.array([r0, th0, ph0, p[0], p[1], p[2]]) +print(z0_cpp) + +def momenta_cpp(t, x): + Ath = np.empty(1) + Aph = np.empty(1) + dAth = np.empty(3) + dAph = np.empty(3) + #B = np.empty(1) + #dB = np.empty(3) + Ath, dAth = field_Ath(x[0], x[1], x[2]) + Aph, dAph = field_Aph(x[0], x[1], x[2]) + #B = field_B(x[0], x[1], x[2]) + g = metric(x[:3]) + + ctrv = np.empty(3) + ctrv[0] = 1/m *g['^11'] *(x[3]) + ctrv[1] = 1/m *g['^22'] *(x[4] - qe/c*Ath) + ctrv[2] = 1/m *g['^33'] *(x[5] - qe/c*Aph) + + xdot = np.empty(6) + xdot[0] = 1/m * g['^11']*x[3] + xdot[1] = 1/m * g['^22']*(x[4] - qe/c*Ath) + xdot[2] = 1/m * g['^33']*(x[5] - qe/c*Aph) + xdot[3] = qe/c*(ctrv[1]*dAth[0] + ctrv[2]*dAph[0]) + m/2*(g['d_11'][0]*ctrv[0]**2 + g['d_22'][0]*ctrv[1]**2 + g['d_33'][0]*ctrv[2]**2) + xdot[4] = qe/c*(ctrv[1]*dAth[1] + ctrv[2]*dAph[1]) + m/2*(g['d_11'][1]*ctrv[0]**2 + g['d_22'][1]*ctrv[1]**2 + g['d_33'][1]*ctrv[2]**2) + xdot[5] = qe/c*(ctrv[1]*dAth[2] + ctrv[2]*dAph[2]) + m/2*(g['d_11'][2]*ctrv[0]**2 + g['d_22'][2]*ctrv[1]**2 + g['d_33'][2]*ctrv[2]**2) + + return xdot + +sol_cpp = solve_ivp( + momenta_cpp, [0, 7000], z0_cpp, method="RK45", rtol=1e-9, atol=1e-9, +) + + +nt = len(sol_cpp.t) +print(nt) # This will print the number of time steps the solver used +print(7000/nt) +z = np.zeros([3, nt]) +z[0,:] = sol_cpp.y[0] +z[1,:] = sol_cpp.y[1] +z[2,:] = sol_cpp.y[2] + + +plot_orbit(z) +plt.plot(z[0,:5]*np.cos(z[1,:5]), z[0,:5]*np.sin(z[1,:5]), "o") +plt.show() + + + +# %% diff --git a/python/cp_sym_midpoint.py b/python/cp_sym_midpoint.py new file mode 100644 index 00000000..baaf7b62 --- /dev/null +++ b/python/cp_sym_midpoint.py @@ -0,0 +1,221 @@ +# %% +import matplotlib.pyplot as plt +import numpy as np +from scipy.optimize import root +from scipy.interpolate import lagrange +import common +from common import (r0,th0,ph0,vpar0,timesteps,get_val,get_der,mu,qe,m,c,) +from field_correct_test import field +from plotale import plot_orbit, plot_mani, plot_cost_function + + +f = field() + +dt = 1 +nt = 1000 +print(dt*nt) +metric = lambda z: { + "_ii": np.array([1, z[0]**2, (1 + z[0]*np.cos(z[1]))**2]), + "^ii": np.array([1, 1/z[0]**2, 1/(1 + z[0]*np.cos(z[1]))**2]), + "d_11": np.array([0,0,0]), + "d_22": np.array([2*z[0], 0,0]), + "d_33": np.array([2*(1+z[0]*np.cos(z[1]))*np.cos(z[1]), -2*(1+z[0]*np.cos(z[1]))*np.sin(z[1]), 0]), +} + + +def dLdx(v, dAth, dAph, g): + ret = np.zeros(3) + ret = (m/2*(g['d_11']*v[0]**2 + g['d_22']*v[1]**2 + g['d_33']*v[2]**2) + qe/c*(dAth*v[1] + dAph*v[2])) + return ret + + +def F(x, xold, pold): + global p + ret = np.zeros(3) + + xmid = (x + xold)/2 + vmid = (x - xold)/dt + + f.evaluate(xmid[0], xmid[1], xmid[2]) + Amid = np.array([0, f.co_Ath, f.co_Aph]) + g = metric(xmid) + pmid = pold + dt/2 * dLdx(vmid, f.co_dAth, f.co_dAph, g) + p = pold + dt * dLdx(vmid, f.co_dAth, f.co_dAph, g) + # d/dt p = d/dt dL/dq. = dL/dq + ret = x - xold - dt/m*g['^ii']*(pmid - qe/c*Amid) + return ret + + +# Initial Conditions +z = np.zeros([3, nt + 1]) +z[:, 0] = [r0, th0, ph0] + +f.evaluate(r0, th0, ph0) +g = metric(z[:,0]) + +vel = np.zeros([3, nt + 1]) +vel[0,0] = np.sqrt(g['^ii'][0]*mu*2*f.B) +vel[1,0] = 0 #-np.sqrt(g['^22']*mu*2*f.B/3) +vel[2,0] = 0 #-np.sqrt(g['^33']*mu*2*f.B/3) +p = np.zeros(3) +p[0] = g['_ii'][0]*vel[0,0] +p[1] = g['_ii'][1]*vel[1,0] + qe/c * f.co_Ath +p[2] = g['_ii'][2]*vel[2,0] + qe/c * f.co_Aph + +vper = np.zeros(nt) +vpar = np.zeros(nt) +v = np.zeros(nt) + +''' +v = np.zeros(3) +p = np.zeros(3) +vel = np.zeros([3, nt + 1]) + +p[0] = vel[0,0] +p[1] = vel[1,0] + qe/c*f.co_Ath +p[2] = vel[2,0] + qe/c*f.co_Aph +''' + +from time import time +tic = time() +for kt in range(nt): + pold = p + + sol = root(F, z[:,kt], method='hybr',tol=1e-12,args=(z[:,kt], pold)) + z[:,kt+1] = sol.x + + f.evaluate(z[0,kt+1], z[1,kt+1], z[2,kt+1]) + g = metric(z[:,kt+1]) + vel[:,kt+1] = 1/m * g['^ii']*(p - qe/c*np.array([0, f.co_Ath, f.co_Aph])) + + vper[kt] = np.sqrt(g['^ii'][0]*mu*2*f.B) + vpar[kt] = np.sqrt(np.sum(vel[:,kt+1]**2)) - vper[kt] + v[kt] = np.sum(vel[:,kt+1]) + + +q_cp = z +v_cp = vel + + +''' +t = np.linspace(0, nt, nt) + +# k = nt +k = 100 + +plt.plot(t[:k], vpar[:k], linewidth=0.5, c='r', label='vpar') + +plt.plot(t[:k], vper[:k], linewidth=0.5, c='g', label='vper') + +plt.plot(t[:k], v[:k], linewidth=0.5, c='b', label='v') +plt.show() + + +plt.hist(v, bins=30, color='blue', edgecolor='black', alpha=0.7) +plt.hist(vpar, bins=30, color='red', edgecolor='black', alpha=0.7) +plt.hist(vper, bins=30, color='green', edgecolor='black', alpha=0.7) + +plt.xlabel("Value") +plt.ylabel("Frequency") +plt.title("Histogram of velocity") +plt.grid(True, linestyle="--", alpha=0.5) +plt.show() + + + +plot_mani(q_cp) +plt.show() + +fig, ax = plt.subplots() +plot_orbit(q_cp, ax = ax) +plt.plot(q_cp[0,:100]*np.cos(q_cp[1,:100]), q_cp[0,:100]*np.sin(q_cp[1,:100]), "o") +plt.show() +''' + + + + + + +''' +dt, nt = timesteps(steps_per_bounce=8, nbounce=100) +print(dt) +print(nt) + +f = field() +dt = 1 +nt = 10000 + +z = np.zeros([3, nt + 1]) +z[:, 0] = [r0, th0, ph0] + +metric = lambda z: { + "_11": 1, + "_22": z[0]**2, + "_33": (1 + z[0]*np.cos(z[1]))**2, + "^11": 1, + "^22": 1/z[0]**2, + + + + + + + + "^33": 1/(1 + z[0]*np.cos(z[1]))**2, + "d_11": [0,0,0], + "d_22": [2*z[0], 0,0], + "d_33": [2*(1+z[0]*np.cos(z[1]))*np.cos(z[1]), -2*(1+z[0]*np.cos(z[1]))*np.sin(z[1]), 0], +} + +def implicit_p(p, pold, Ath, Aph, dAth, dAph, g): + ret = np.zeros(3) + + ctrv = np.zeros(3) + ctrv[0] = 1/m *g['^11'] *(p[0]) + ctrv[1] = 1/m *g['^22'] *(p[1] - qe/c*Ath) + ctrv[2] = 1/m *g['^33'] *(p[2] - qe/c*Aph) + + ret[0] = p[0] - pold[0] - dt*(qe/c*(ctrv[1]*dAth[0] + ctrv[2]*dAph[0]) + m/2*(g['d_11'][0]*ctrv[0]**2 + g['d_22'][0]*ctrv[1]**2 + g['d_33'][0]*ctrv[2]**2)) + ret[1] = p[1] - pold[1] - dt*(qe/c*(ctrv[1]*dAth[1] + ctrv[2]*dAph[1]) + m/2*(g['d_11'][1]*ctrv[0]**2 + g['d_22'][1]*ctrv[1]**2 + g['d_33'][1]*ctrv[2]**2)) + ret[2] = p[2] - pold[2] - dt*(qe/c*(ctrv[1]*dAth[2] + ctrv[2]*dAph[2]) + m/2*(g['d_11'][2]*ctrv[0]**2 + g['d_22'][2]*ctrv[1]**2 + g['d_33'][2]*ctrv[2]**2)) + + return ret + +#Initial Conditions +f.evaluate(r0, th0, ph0) +g = metric(z[:,0]) +ctrv = np.zeros(3) +ctrv[0] = np.sqrt(g['^11']*mu*2*f.B) +ctrv[1] = 0 #-np.sqrt(g['^22']*mu*2*f.B/3) +ctrv[2] = 0 #-np.sqrt(g['^33']*mu*2*f.B/3) +p = np.zeros(3) +p[0] = g['_11']*ctrv[0] +p[1] = g['_22']*ctrv[1] + qe/c * f.co_Ath +p[2] = g['_33']*ctrv[2] + qe/c * f.co_Aph +pold = p + +from time import time +tic = time() +for kt in range(nt): + + sol = root(implicit_p, p, method='hybr',tol=1e-12,args=(pold, f.co_Ath, f.co_Aph, f.co_dAth, f.co_dAph, g)) + p = sol.x + pold = p + + z[0,kt+1] = z[0,kt] + dt/m * g['^11']*(p[0]) + z[1,kt+1] = z[1,kt] + dt/m * g['^22']*(p[1] - qe/c*f.co_Ath) + z[2,kt+1] = z[2,kt] + dt/m * g['^33']*(p[2] - qe/c*f.co_Aph) + + f.evaluate(z[0, kt+1], z[1, kt+1], z[2, kt+1]) + g = metric(z[:,kt+1]) + + p = np.zeros(3) + p[0] = g['_11']*m*(z[0,kt+1]-z[0,kt])/dt + p[1] = g['_22']*m*(z[1,kt+1]-z[1,kt])/dt + qe/c*f.co_Ath + p[2] = g['_33']*m*(z[2,kt+1]-z[2,kt])/dt + qe/c*f.co_Aph +''' + + + +# %% diff --git a/python/cpp_expl_impl.py b/python/cpp_expl_impl.py new file mode 100644 index 00000000..50dc92df --- /dev/null +++ b/python/cpp_expl_impl.py @@ -0,0 +1,84 @@ +# %% +import matplotlib.pyplot as plt +import numpy as np +from scipy.optimize import root +from scipy.interpolate import lagrange +import common +from common import (r0,th0,ph0,pph0,timesteps,get_val,get_der,mu,qe,m,c,) +from field_correct_test import field +from plotting import plot_orbit, plot_cost_function + +f = field() +dt, nt = timesteps(steps_per_bounce=8, nbounce=100) + +print(dt) +print(nt) +dt = 1.9 +nt = 4000 + +z = np.zeros([3, nt + 1]) +z[:, 0] = [r0, th0, ph0] + +metric = lambda z: { + "_11": 1, + "_22": z[0]**2, + "_33": (1 + z[0]*np.cos(z[1]))**2, + "^11": 1, + "^22": 1/z[0]**2, + "^33": 1/(1 + z[0]*np.cos(z[1]))**2, + "d_11": [0,0,0], + "d_22": [2*z[0], 0,0], + "d_33": [2*(1+z[0]*np.cos(z[1]))*np.cos(z[1]), -2*(1+z[0]*np.cos(z[1]))*np.sin(z[1]), 0], +} + +def implicit_p(p, pold, Ath, Aph, dB, dAth, dAph, g): + ret = np.zeros(3) + + ctrv = np.zeros(3) + ctrv[0] = 1/m *g['^11'] *(p[0]) + ctrv[1] = 1/m *g['^22'] *(p[1] - qe/c*Ath) + ctrv[2] = 1/m *g['^33'] *(p[2] - qe/c*Aph) + + ret[0] = p[0] - pold[0] - dt*(qe/c*(ctrv[1]*dAth[0] + ctrv[2]*dAph[0]) - mu*dB[0] + m/2*(g['d_11'][0]*ctrv[0]**2 + g['d_22'][0]*ctrv[1]**2 + g['d_33'][0]*ctrv[2]**2)) + ret[1] = p[1] - pold[1] - dt*(qe/c*(ctrv[1]*dAth[1] + ctrv[2]*dAph[1]) - mu*dB[1] + m/2*(g['d_11'][1]*ctrv[0]**2 + g['d_22'][1]*ctrv[1]**2 + g['d_33'][1]*ctrv[2]**2)) + ret[2] = p[2] - pold[2] - dt*(qe/c*(ctrv[1]*dAth[2] + ctrv[2]*dAph[2]) - mu*dB[2] + m/2*(g['d_11'][2]*ctrv[0]**2 + g['d_22'][2]*ctrv[1]**2 + g['d_33'][2]*ctrv[2]**2)) + + return ret + +#Initial Conditions +f.evaluate(r0, th0, ph0) +g = metric(z[:,0]) +p = np.zeros(3) +p[0] = 0 +p[1] = qe/c * f.co_Ath +p[2] = qe/c * f.co_Aph +pold = p + +from time import time +tic = time() +for kt in range(nt): + sol = root(implicit_p, p, method='hybr',tol=1e-12,args=(pold, f.co_Ath, f.co_Aph, f.dB, f.co_dAth, f.co_dAph, g)) + p = sol.x + pold = p + + z[0,kt+1] = z[0,kt] + dt/m * g['^11']*(p[0]) + z[1,kt+1] = z[1,kt] + dt/m * g['^22']*(p[1] - qe/c*f.co_Ath) + z[2,kt+1] = z[2,kt] + dt/m * g['^33']*(p[2] - qe/c*f.co_Aph) + + f.evaluate(z[0, kt+1], z[1, kt+1], z[2, kt+1]) + g = metric(z[:,kt+1]) + + #p = np.zeros(3) + #p[0] = g['_11']*m*(z[0,kt+1]-z[0,kt])/dt + #p[1] = g['_22']*m*(z[1,kt+1]-z[1,kt])/dt + qe/c*f.co_Ath + #p[2] = g['_33']*m*(z[2,kt+1]-z[2,kt])/dt + qe/c*f.co_Aph + + +plot_orbit(z) +plt.plot(z[0,:5]*np.cos(z[1,:5]), z[0,:5]*np.sin(z[1,:5]), "o") +plt.show() + + + + +# %% diff --git a/python/cpp_lagrange.py b/python/cpp_lagrange.py new file mode 100644 index 00000000..9b83d682 --- /dev/null +++ b/python/cpp_lagrange.py @@ -0,0 +1,97 @@ +# %% +import matplotlib.pyplot as plt +import numpy as np +from scipy.optimize import root +from scipy.interpolate import lagrange +import common +from common import (f,r0,th0,ph0,pph0,timesteps,get_val,get_der,mu,qe,m,c,) +from plotting import plot_orbit, plot_cost_function + +dt, nt = timesteps(steps_per_bounce=8, nbounce=100) + +print(dt) +print(nt) +dt = 8 +nt = 5000 + +z = np.zeros([3, nt + 1]) +z[:, 0] = [r0, th0, ph0] + +f.evaluate(r0, th0, ph0) +v = np.zeros(3) +vold = v + +metric = lambda z: { + "_11": 1, + "_22": z[0]**2, + "_33": (1 + z[0]*np.cos(z[1]))**2, + "^11": 1, + "^22": 1/z[0]**2, + "^33": 1/(1 + z[0]*np.cos(z[1]))**2, + "d_11": [0,0,0], + "d_22": [2*z[0], 0,0], + "d_33": [2*(1+z[0]*np.cos(z[1]))*np.cos(z[1]), -2*(1+z[0]*np.cos(z[1]))*np.sin(z[1]), 0], +} +''' +metric = lambda z: { + "_11": 1, + "_22": z[0]**2, + "_33": z[0]**2 * np.sin(z[1])**2, + "^11": 1, + "^22": 1/z[0]**2, + "^33": 1/(z[0]**2 * np.sin(z[1])**2), + "d_11": [0,0,0], + "d_22": [2*z[0], 0,0], + "d_33": [2*z[0]*np.sin(z[1])**2, 2*z[0]**2 * np.sin(z[1])*np.cos(z[1]), 0], +} +''' + + +def implicit_v(v, vold, dAth, dAph, dB, g): + ret = np.zeros(3) + + + ret[0] = g['_11']*(v[0] - vold[0]) - dt*m*0.5*(v[0]**2 *g['d_11'][0] + v[1]**2 *g['d_22'][0] + v[2]**2 *g['d_33'][0]) + dt*m*v[0]*(g['d_11'][0]*v[0] + g['d_11'][1]*v[1] + g['d_11'][2]*v[2]) - dt*qe/c *(v[1]*(dAth[0]) + v[2]*(dAph[0])) + dt*mu*dB[0] + + ret[1] = g['_22']*(v[1] - vold[1]) - dt*m*0.5*(v[0]**2 *g['d_11'][1] + v[1]**2 *g['d_22'][1] + v[2]**2 *g['d_33'][1]) + dt*m*v[1]*(g['d_22'][0]*v[0] + g['d_22'][1]*v[1] + g['d_22'][2]*v[2]) - dt*qe/c *(v[0]*(-dAth[0]) + v[2]*(dAph[1] - dAth[2])) + dt*mu*dB[1] + + ret[2] = g['_33']*(v[2] - vold[2]) - dt*m*0.5*(v[0]**2 *g['d_11'][2] + v[1]**2 *g['d_22'][2] + v[2]**2 *g['d_33'][2]) + dt*m*v[2]*(g['d_33'][0]*v[0] + g['d_33'][1]*v[1] + g['d_33'][2]*v[2]) - dt*qe/c *(v[0]*(-dAph[0]) + v[1]*(dAth[2] - dAph[1])) + dt*mu*dB[2] + + return ret + + + +from time import time +tic = time() +for kt in range(nt): + + f.evaluate(z[0, kt], z[1, kt], z[2, kt]) + g = metric(z[:,kt]) + + + + sol = root(implicit_v, v, method='hybr',tol=1e-12,args=(vold, f.dAth, f.dAph, f.dB, g)) + vold = v + v = sol.x + + z[0,kt+1] = z[0,kt] + dt*v[0] + z[1,kt+1] = z[1,kt] + dt*v[1] + z[2,kt+1] = z[2,kt] + dt*v[2] + + + +plot_orbit(z) +plt.plot(z[0,:5]*np.cos(z[1,:5]), z[0,:5]*np.sin(z[1,:5]), "o") +plt.show() + + + + +#implicit + + + + + + +# %% diff --git a/python/cpp_runge_kutta.py b/python/cpp_runge_kutta.py new file mode 100644 index 00000000..e40d4d56 --- /dev/null +++ b/python/cpp_runge_kutta.py @@ -0,0 +1,113 @@ +# %% +from scipy.integrate import solve_ivp +import matplotlib.pyplot as plt +import numpy as np +from scipy.optimize import root +from scipy.interpolate import lagrange +import common +from common import (f,r0,th0,ph0,pph0,timesteps,get_val,get_der,mu,qe,m,c,) +from plotting import plot_orbit, plot_cost_function + +B0 = 1.0 # magnetic field modulus normalization +iota0 = 1.0 # constant part of rotational transform +a = 0.5 # (equivalent) minor radius +R0 = 1.0 # (equivalent) major radius + +z0_cpp = np.array([r0, th0, ph0]) + +metric = lambda z: { + "_11": 1, + "_22": z[0]**2, + "_33": (1 + z[0]*np.cos(z[1]))**2, + "^11": 1, + "^22": 1/z[0]**2, + "^33": 1/(1 + z[0]*np.cos(z[1]))**2, + "d_11": [0,0,0], + "d_22": [2*z[0], 0,0], + "d_33": [2*(1+z[0]*np.cos(z[1]))*np.cos(z[1]), -2*(1+z[0]*np.cos(z[1]))*np.sin(z[1]), 0], +} + +def field_Ath(r, th, ph): + Ath = B0*(r**2/2 - r**3/(3*R0)*np.cos(th)) + dAth = np.array((B0*(r - r**2/R0*np.cos(th)), B0*r**3*np.sin(th)/(3*R0), 0)) + return Ath, dAth + +def field_Aph(r, th, ph): + Aph = -B0*iota0*(r**2/2-r**4/(4*a**2)) + dAph = np.array((-B0*iota0*(r-r**3/a**2), 0, 0)) + return Aph, dAph + +def field_dB(r, th, ph): + dB = np.array((iota0*(1-3*r**2/a**2)/(R0+r*np.cos(th)) - iota0*(r-r**3/a**2)/(R0+r*np.cos(th))**2 * np.cos(th) - np.cos(th)/R0, + iota0*(r-r**3/a**2)/(R0+r*np.cos(th))**2 *r*np.sin(th) + r*np.sin(th), 0)) + return dB + + +#Initial Conditions +Ath = np.empty(1) +Aph = np.empty(1) +dAth = np.empty(3) +dAph = np.empty(3) +dB = np.empty(3) +Ath, dAth = field_Ath(r0, th0, ph0) +Aph, dAph = field_Aph(r0, th0, ph0) +dB = field_dB(r0, th0, ph0) + +g = metric(z0_cpp[:]) +p = np.zeros(3) +p[0] = 0 +p[1] = qe/c * Ath +p[2] = qe/c * Aph +pold = p + +#Initial Conditions Runge-Kutta +z0_cpp = np.array([r0, th0, ph0, p[0], p[1], p[2]]) +print(z0_cpp) + +def momenta_cpp(t, x): + Ath = np.empty(1) + Aph = np.empty(1) + dAth = np.empty(3) + dAph = np.empty(3) + dB = np.empty(3) + Ath, dAth = field_Ath(x[0], x[1], x[2]) + Aph, dAph = field_Aph(x[0], x[1], x[2]) + dB = field_dB(x[0], x[1], x[2]) + g = metric(x[:3]) + + ctrv = np.empty(3) + ctrv[0] = 1/m *g['^11'] *(x[3]) + ctrv[1] = 1/m *g['^22'] *(x[4] - qe/c*Ath) + ctrv[2] = 1/m *g['^33'] *(x[5] - qe/c*Aph) + + xdot = np.empty(6) + xdot[0] = 1/m * g['^11']*x[3] + xdot[1] = 1/m * g['^22']*(x[4] - qe/c*Ath) + xdot[2] = 1/m * g['^33']*(x[5] - qe/c*Aph) + xdot[3] = qe/c*(ctrv[1]*dAth[0] + ctrv[2]*dAph[0]) - mu*dB[0] + m/2*(g['d_11'][0]*ctrv[0]**2 + g['d_22'][0]*ctrv[1]**2 + g['d_33'][0]*ctrv[2]**2) + xdot[4] = qe/c*(ctrv[1]*dAth[1] + ctrv[2]*dAph[1]) - mu*dB[1] + m/2*(g['d_11'][1]*ctrv[0]**2 + g['d_22'][1]*ctrv[1]**2 + g['d_33'][1]*ctrv[2]**2) + xdot[5] = qe/c*(ctrv[1]*dAth[2] + ctrv[2]*dAph[2]) - mu*dB[2] + m/2*(g['d_11'][2]*ctrv[0]**2 + g['d_22'][2]*ctrv[1]**2 + g['d_33'][2]*ctrv[2]**2) + + return xdot + +sol_cpp = solve_ivp( + momenta_cpp, [0, 6000], z0_cpp, method="RK45", rtol=1e-6, atol=1e-6, +) + + +nt = len(sol_cpp.t) +print(nt) # This will print the number of time steps the solver used +print(6000/nt) +z = np.zeros([3, nt]) +z[0,:] = sol_cpp.y[0] +z[1,:] = sol_cpp.y[1] +z[2,:] = sol_cpp.y[2] + + +plot_orbit(z) +plt.plot(z[0,:5]*np.cos(z[1,:5]), z[0,:5]*np.sin(z[1,:5]), "o") +plt.show() + + + +# %% diff --git a/python/cpp_sym_euler.py b/python/cpp_sym_euler.py new file mode 100644 index 00000000..7f54b08a --- /dev/null +++ b/python/cpp_sym_euler.py @@ -0,0 +1,79 @@ +# %% +import matplotlib.pyplot as plt +import numpy as np +from scipy.optimize import root +from scipy.interpolate import lagrange +import common +from common import (r0,th0,ph0,pph0,timesteps,get_val,get_der,mu,qe,m,c,) +from field_correct_test import field +from plotale import plot_orbit, plot_mani, plot_cost_function + +f = field() + +dt = 1.9 + +nt = 4000 + +metric = lambda z: { + "_ii": np.array([1, z[0]**2, (1 + z[0]*np.cos(z[1]))**2]), + "^ii": np.array([1, 1/z[0]**2, 1/(1 + z[0]*np.cos(z[1]))**2]), + "d_11": np.array([0,0,0]), + "d_22": np.array([2*z[0], 0,0]), + "d_33": np.array([2*(1+z[0]*np.cos(z[1]))*np.cos(z[1]), -2*(1+z[0]*np.cos(z[1]))*np.sin(z[1]), 0]), +} + +def dLdx(v, dAth, dAph, dB, g): + ret = np.zeros(3) + ret = (m/2*(g['d_11']*v[0]**2 + g['d_22']*v[1]**2 + g['d_33']*v[2]**2) + qe/c*(dAth*v[1] + dAph*v[2]) - mu*dB) + return ret + + +def F(p, xold, pold): + global A + + ret = np.zeros(3) + f.evaluate(xold[0], xold[1], xold[2]) + g = metric(xold) + A = np.array([0, f.co_Ath, f.co_Aph]) + v = 1/m * g['^ii'] *(p - qe/c * A) + + ret = p - pold - dt * dLdx(v, f.co_dAth, f.co_dAph, f.dB, g) + return ret + + + +# Initial Conditions +z = np.zeros([3, nt + 1]) +z[:, 0] = [r0, th0, ph0] + +f.evaluate(r0, th0, ph0) +g = metric(z[:,0]) +v = np.zeros(3) +p = np.zeros(3) +p[0] = v[0] +p[1] = v[1] + qe/c*f.co_Ath +p[2] = v[2] + qe/c*f.co_Aph + + +from time import time +tic = time() +for kt in range(nt): + pold = p + + sol = root(F, p, method='hybr',tol=1e-12,args=(z[:,kt], pold)) + p = sol.x + + z[:,kt+1] = z[:,kt] + dt/m * g['^ii']*(p - qe/c * A) + + + + +plot_orbit(z) +plt.plot(z[0,:5]*np.cos(z[1,:5]), z[0,:5]*np.sin(z[1,:5]), "o") +plt.show() + +plot_mani(z) +plt.show() + + +# %% \ No newline at end of file diff --git a/python/cpp_sym_midpoint.py b/python/cpp_sym_midpoint.py new file mode 100644 index 00000000..646bf887 --- /dev/null +++ b/python/cpp_sym_midpoint.py @@ -0,0 +1,121 @@ +# %% +import matplotlib.pyplot as plt +import numpy as np +from scipy.optimize import root +from scipy.interpolate import lagrange +import common +from common import (r0,th0,ph0,vpar0,timesteps,get_val,get_der,mu,qe,m,c,) +from field_correct_test import field +from plotale import plot_orbit, plot_mani, plot_cost_function + +f = field() + +dt = 80 +nt = 1000 + +metric = lambda z: { + "_ii": np.array([1, z[0]**2, (1 + z[0]*np.cos(z[1]))**2]), + "^ii": np.array([1, 1/z[0]**2, 1/(1 + z[0]*np.cos(z[1]))**2]), + "d_11": np.array([0,0,0]), + "d_22": np.array([2*z[0], 0,0]), + "d_33": np.array([2*(1+z[0]*np.cos(z[1]))*np.cos(z[1]), -2*(1+z[0]*np.cos(z[1]))*np.sin(z[1]), 0]), +} + +def dLdx(v, dAth, dAph, dB, g): + ret = np.zeros(3) + ret = (m/2*(g['d_11']*v[0]**2 + g['d_22']*v[1]**2 + g['d_33']*v[2]**2) + qe/c*(dAth*v[1] + dAph*v[2]) - mu*dB) + return ret + + +def F(x, xold, pold): + global p + ret = np.zeros(3) + + xmid = (x + xold)/2 + vmid = (x - xold)/dt + + f.evaluate(xmid[0], xmid[1], xmid[2]) + Amid = np.array([0, f.co_Ath, f.co_Aph]) + g = metric(xmid) + pmid = pold + dt/2 * dLdx(vmid, f.co_dAth, f.co_dAph, f.dB, g) + p = pold + dt * dLdx(vmid, f.co_dAth, f.co_dAph, f.dB, g) + # d/dt p = d/dt dL/dq. = dL/dq + ret = x - xold - dt/m*g['^ii']*(pmid - qe/c*Amid) + return ret + + +# Initial Conditions +z = np.zeros([3, nt + 1]) +z[:, 0] = [r0, th0, ph0] + +f.evaluate(r0, th0, ph0) +g = metric(z[:,0]) +v = np.zeros(3) +p = np.zeros(3) +vel = np.zeros([3, nt + 1]) +vel[:,0] = [vpar0*f.co_hr, vpar0*f.co_hth, vpar0*f.co_hph] + +p[0] = vel[0,0] +p[1] = vel[1,0] + qe/c*f.co_Ath +p[2] = vel[2,0] + qe/c*f.co_Aph + +vper = np.zeros(nt) +vpar = np.zeros(nt) +v = np.zeros(nt) + +from time import time +tic = time() +for kt in range(nt): + pold = p + + sol = root(F, z[:,kt], method='hybr',tol=1e-12,args=(z[:,kt], pold)) + z[:,kt+1] = sol.x + + f.evaluate(z[0,kt+1], z[1,kt+1], z[2,kt+1]) + g = metric(z[:,kt+1]) + vel[:,kt+1] = 1/m * g['^ii']*(p - qe/c*np.array([0, f.co_Ath, f.co_Aph])) + + + vper[kt] = np.sqrt(g['^ii'][0]*mu*2*f.B) + vpar[kt] = np.sqrt(np.sum(vel[:,kt+1]**2)) - vper[kt] + v[kt] = np.sum(vel[:,kt+1]) + + + + + + +''' +t = np.linspace(0, nt, nt) + +# k = nt +k = 100 + +plt.plot(t[:k], vpar[:k], linewidth=0.5, c='r', label='vpar') + +plt.plot(t[:k], vper[:k], linewidth=0.5, c='g', label='vper') + +plt.plot(t[:k], v[:k], linewidth=0.5, c='b', label='v') +plt.show() + + +plt.hist(v, bins=30, color='blue', edgecolor='black', alpha=0.7) +plt.hist(vpar, bins=30, color='red', edgecolor='black', alpha=0.7) +plt.hist(vper, bins=30, color='green', edgecolor='black', alpha=0.7) + +plt.xlabel("Value") +plt.ylabel("Frequency") +plt.title("Histogram of velocity") +plt.grid(True, linestyle="--", alpha=0.5) +plt.show() + +plt.plot +plot_mani(z) +#plt.plot(z[0,:]*np.cos(z[1,:]), z[0,:]*np.sin(z[1,:]), ".") +plt.show() + +fig, ax = plt.subplots() +plot_orbit(z, ax = ax) +plt.plot(z[0,:5]*np.cos(z[1,:5]), z[0,:5]*np.sin(z[1,:5]), "o") +plt.show() +''' \ No newline at end of file diff --git a/python/cpp_var_midpoint.py b/python/cpp_var_midpoint.py new file mode 100644 index 00000000..9f081c68 --- /dev/null +++ b/python/cpp_var_midpoint.py @@ -0,0 +1,111 @@ +# %% +import matplotlib.pyplot as plt +import numpy as np +from scipy.optimize import root +from scipy.interpolate import lagrange +import common +from common import (r0,th0,ph0,pph0,timesteps,get_val,get_der,mu,qe,m,c,) +from field_correct_test import field +from plotting import plot_orbit, plot_cost_function + +f = field() + +qe = 1.0; m = 1.0; c = 1.0 # particle charge, mass and speed of light +mu = 1e-5 # magnetic moment + +# Initial conditions +r0 = 0.1 +th0 = 1.5 +ph0 = 1.0 +vpar0 = 0.0 + +dt = 800 +nt = 2000 + +metric = lambda z: { + "_ii": np.array([1, z[0]**2, (1 + z[0]*np.cos(z[1]))**2]), + "^ii": np.array([1, 1/z[0]**2, 1/(1 + z[0]*np.cos(z[1]))**2]), + "d_11": np.array([0,0,0]), + "d_22": np.array([2*z[0], 0,0]), + "d_33": np.array([2*(1+z[0]*np.cos(z[1]))*np.cos(z[1]), -2*(1+z[0]*np.cos(z[1]))*np.sin(z[1]), 0]), +} + + + +def F(x, xold, dLdxold, dLdxdotold): + global p, dpdt, vmid + ret = np.zeros(3) + + xmid = (x + xold)/2 + vmid = (x - xold)/dt + + f.evaluate(xmid[0], xmid[1], xmid[2]) + Amid = np.array([0, f.co_Ath, f.co_Aph]) + g = metric(xmid) + + dpdt = m/2*(g['d_11']*vmid[0]**2 + g['d_22']*vmid[1]**2 + g['d_33']*vmid[2]**2) + qe/c*(f.co_dAth*vmid[1] + f.co_dAph*vmid[2]) - mu*f.dB + p = m*g['_ii']*vmid + qe/c*Amid + + ret = (dpdt + dLdxold)*dt/2 - (p - dLdxdotold) + return ret + +def action(x, xold): + X = np.array([x, xold]) + + for i in 2: + print(i) + f.evaluate(X[1], X[2], X[3]) + g = metric(X) + + + + + xmid = (x + xold)/2 + vmid = (x - xold)/dt + f.evaluate(xmid[0], xmid[1], xmid[2]) + g = metric(xmid) + + Amid = np.array([0, f.co_Ath, f.co_Aph]) + hmid = np.array([f.co_hr, f.co_hth, f.co_hph]) + pmid = m*g['_ii']*vmid + qe/c*Amid + vpar = m* np.sqrt(np.sum(pmid**2)) + + S = np.sum((qe/c * Amid + m*vpar*hmid)*vmid) - (m/2*vpar**2 - mu*f.B)* dt + h = 6.6260755*1e-34 + + return S + + +# Initial Conditions +z = np.zeros([3, nt + 1]) +z[:, 0] = [r0, th0, ph0] +L = np.zeros(nt) + +f.evaluate(r0, th0, ph0) +g = metric(z[:,0]) +Amid = np.array([0, f.co_Ath, f.co_Aph]) +vmid = np.zeros(3) +p = m*g['_ii']*vmid + qe/c*Amid +dpdt = m/2*(g['d_11']*vmid[0]**2 + g['d_22']*vmid[1]**2 + g['d_33']*vmid[2]**2) + qe/c*(f.co_dAth*vmid[1] + f.co_dAph*vmid[2]) - mu*f.dB + + +from time import time +tic = time() +for kt in range(nt): + + sol = root(F, z[:,kt], method='hybr',tol=1e-12,args=(z[:,kt], dpdt, p)) + z[:,kt+1] = sol.x + + + + +plot_orbit(z) +plt.plot(z[0,:5]*np.cos(z[1,:5]), z[0,:5]*np.sin(z[1,:5]), "o") +plt.show() + + + + + + +# %% diff --git a/python/discrete_variational.py b/python/discrete_variational.py index cca6d2a9..412e1462 100644 --- a/python/discrete_variational.py +++ b/python/discrete_variational.py @@ -7,8 +7,10 @@ from common import (f,r0,th0,ph0,pph0,timesteps,get_val,get_der,mu,qe,m,c,) from plotting import plot_orbit, plot_cost_function -dt, nt = timesteps(steps_per_bounce=8, nbounce=100) + +dt, nt = timesteps(steps_per_bounce=8, nbounce=100) +print(dt) z = np.zeros([3, nt + 1]) z[:, 0] = [r0, th0, ph0] [H, pth, vpar] = get_val(np.array([r0, th0, ph0, pph0])) diff --git a/python/field_correct_test.py b/python/field_correct_test.py new file mode 100644 index 00000000..1975f721 --- /dev/null +++ b/python/field_correct_test.py @@ -0,0 +1,57 @@ +from numpy import sin, cos, zeros_like, array, sqrt + +B0 = 1.0 # magnetic field modulus normalization +iota0 = 1.0 # constant part of rotational transform +a = 0.5 # (equivalent) minor radius +R0 = 1.0 # (equivalent) major radius + +class field: + """ Model tokamak field with exact B = B0*(iota0*(r-r^3/a^2)/(R0+r*cos(th)) + (1-r/R0*cos(th)) """ + def evaluate(self, r, th, ph): + cth = cos(th) + sth = sin(th) + zer = zeros_like(r) + J = r*(R0+r*cth) + + self.co_Ath = B0*(r**2/2.0 - r**3/(3.0*R0)*cth) + self.co_dAth = array((B0*(r - r**2/R0*cth), B0*r**3*sth/(3.0*R0), zer)) + self.co_d2Ath = array((B0*(1.0 - 2.0*r/R0*cth), + B0*r**3*cth/(3.0*R0), + zer, + B0*r**2/R0*sth, + zer, + zer)) + + self.co_Aph = -B0*iota0*(r**2/2.0-r**4/(4.0*a**2)) + self.co_dAph = array((-B0*iota0*(r-r**3/a**2), zer, zer)) + self.co_d2Aph = array((-B0*iota0*(1.0-3.0*r**2/a**2), zer, zer, zer, zer, zer)) + + self.ctr_Br = 1/J *(self.co_dAph[1] - self.co_dAth[2]) + self.ctr_Bth = 1/J *(-self.co_dAph[0]) + self.ctr_Bph = 1/J *(self.co_dAth[0]) + + self.B = sqrt(self.ctr_Br**2 + r**2 *self.ctr_Bth**2 + (R0+r*cth)**2 *self.ctr_Bph**2) + self.dB = array(((r*self.ctr_Bth**2 + r**2*self.ctr_Bth*(-1/J*self.co_d2Aph[0] + 1/J**2*(R0+2*r*cth)*self.co_dAph[0]) + (R0+r*cth)*cth*self.ctr_Bph**2 + (R0+r*cth)**2*self.ctr_Bph*(1/J*self.co_d2Ath[0] - 1/J**2*(R0+2*r*cth)*self.co_dAth[0]))/self.B, + (-(R0+r*cth)*r*sth*self.ctr_Bph**2 + (R0+r*cth)**2*self.ctr_Bph*(1/J*self.co_d2Ath[3] + 1/J**2*r**2*sth*self.co_dAth[0]))/self.B, + zer)) + + self.B1 = B0*(iota0*(r-r**3/a**2)/(R0+r*cth) + (1.0 - r/R0*cth)) + + self.co_hr = self.ctr_Br/self.B + self.co_hth = r**2 *self.ctr_Bth/self.B + self.co_hph = (R0+r*cth)**2 *self.ctr_Bph/self.B + + + + self.Phie = zer + self.dPhie = array((zer, zer, zer)) + self.d2Phie = array((zer, zer, zer, zer, zer, zer)) + + + #self.B1 = B0*(iota0*(r-r**3/a**2)/(R0+r*cth) + (1.0 - r/R0*cth)) + #self.dB = array((B0*(iota0*(1-3*r**2/a**2)/(R0+r*cth) - iota0*(r-r**3/a**2)/(R0+r*cth)**2 *cth - cth/R0), + # B0*(iota0*(r-r**3/a**2)/(R0+r*cth)**2 *r*sth + r*sth/R0), + # zer)) + #self.d2B = array((B0*(-iota0*6*r/a**2/(R0+r*cth) - iota0*(1-3*r**2/a**2)/(R0+r*cth)**2*cth ................... ))) + #self.dB = array((-B0/R0*cth, B0*r/R0*sth, zer)) + #self.d2B = array((zer, B0*r/R0*cth, zer, B0/R0*sth, zer, zer)) diff --git a/python/manifolds.py b/python/manifolds.py new file mode 100644 index 00000000..09c9543c --- /dev/null +++ b/python/manifolds.py @@ -0,0 +1,431 @@ +# %% +import matplotlib.pyplot as plt +import numpy as np +from itertools import product +from scipy.optimize import root +from scipy.interpolate import lagrange +import common +from common import (r0,th0,ph0,pph0,timesteps,get_val,get_der,mu,qe,m,c,) +from field_correct_test import field +from plotale import plot_orbit, plot_mani, plot_cost_function +from cpp_sym_midpoint import z, vel +from cp_sym_midpoint import q_cp, v_cp + +''' +def position(M, z0): + + distances = np.linalg.norm(M - z0, axis=1) + + # Find index of the minimum distance + idx_min = np.argmin(distances) +''' +traj_v = vel +traj_q = z + +traj_v_cp = v_cp +traj_q_cp = q_cp + +f = field() + +R0 = 1 +z0 = 1 +print(r0, th0, ph0) + +#Resolution +n = 100#200 # not to be odd !! +nh = 100#100 # n/2 +l = n**2 * nh # Number of Pixels + + + +# Grid in flux coordinates + +gridR = np.linspace(0,0.5,n) +gridTh = np.linspace(0,2*np.pi,n) +gridPh = np.linspace(0,2*np.pi,n) + +a = 0.5 +th, ph = np.meshgrid(gridTh, gridPh, indexing='ij') + +X = (R0 + a*np.cos(th))*np.cos(ph) +Y = (R0 + a*np.cos(th))*np.sin(ph) +Z = z0 + a*np.sin(th) + + + +#fig = plt.figure() + +#ax = fig.add_subplot(111, projection='3d') + +#sc = ax.scatter(X, Y, Z, cmap='viridis', s=0.1) + +#plt.show() + + + + +# Cartesian grid +# Components x^i +X = np.linspace(-3,3,n) +Y = np.linspace(-3,3,n) +Z = np.linspace(0,2,nh) + +# Canonical coordinate Transformation into flux coordinates q^i +q = np.zeros([l,3]) +c = 0 + +for x in X: + for y in Y: + count = -1 + R = np.sqrt((R0 - np.sqrt(x**2 + y**2))**2 + (Z - z0)**2) + + for z in Z: + count += 1 + + + if R[count] <= 0.5: + + r = R[count] + + a = z-z0 + b = np.sqrt(r**2 - a**2) + theta = np.arctan(a/b) + cond = np.sqrt(x**2 + y**2) - R0 + + if a >= 0 and cond >= 0: + th = theta + elif a >= 0 and cond < 0: + th = np.pi - theta + elif a < 0 and cond < 0: + th = np.pi - theta + elif a < 0 and cond >= 0: + th = 2*np.pi + theta + + phi = np.arctan(y/x) + if x > 0 and y > 0: + ph = phi + elif x < 0 and y > 0: + ph = np.pi + phi + #elif x < 0 and y < 0: + # ph = np.pi + phi + #elif x > 0 and y < 0: + # ph = 2*np.pi + phi + else: break + + q[c,:] = [r, th, ph] + c +=1 + +q = q[:c,:] +print(c) + +# x^i = x^i(q) +x = (R0 + q[:,0] * np.cos(q[:,1])) * np.cos(q[:,2]) +y = (R0 + q[:,0] * np.cos(q[:,1])) * np.sin(q[:,2]) +z = z0 + q[:,0] * np.sin(q[:,1]) +### +# +# +# +# +# +# +# +# + + + #vx5[j] = x_dot[0] + #vy5[j] = x_dot[1] + #vz5[j] = x_dot[2] + + + +#velx5 = (R0 + traj_v[0,:500] * np.cos(traj_v[1,:500])) * np.cos(traj_v[2,:500]) +#vely5 = (R0 + traj_v[0,:500] * np.cos(traj_v[1,:500])) * np.sin(traj_v[2,:500]) +#velz5 = z0 + traj_v[0,:500] * np.sin(traj_v[1,:500]) + +# cp: + +vcp_x = np.zeros([3,len(traj_q_cp[0,:])]) + +for j in range(len(traj_q_cp[0,:])): + + A = np.array([ [np.cos(traj_q_cp[1,j])*np.cos(traj_q_cp[2,j]), -r*np.sin(traj_q_cp[1,j])*np.cos(traj_q_cp[2,j]), -np.sin(traj_q_cp[2,j])], + [np.cos(traj_q_cp[1,j])*np.sin(traj_q_cp[2,j]), -r*np.sin(traj_q_cp[1,j])*np.sin(traj_q_cp[2,j]), np.cos(traj_q_cp[2,j])], + [ np.sin(traj_q_cp[1,j]), r*np.cos(traj_q_cp[1,j]), 0] ]) + + q_dot = traj_v_cp[:,j] + + x_dot = A @ q_dot + + vcp_x[:,j] = x_dot + + +# Trajectory of particle in Cartesian coordinates + +qcpx = (R0 + traj_q_cp[0,:] * np.cos(traj_q_cp[1,:])) * np.cos(traj_q_cp[2,:]) +qcpy = (R0 + traj_q_cp[0,:] * np.cos(traj_q_cp[1,:])) * np.sin(traj_q_cp[2,:]) +qcpz = z0 + traj_q_cp[0,:] * np.sin(traj_q_cp[1,:]) + + +# Velocity in Cartesian coordinates +vcpx = (vcp_x[0,:]) # + qcpz +vcpy = (vcp_x[1,:]) # + qcpy +vcpz = (vcp_x[2,:]) # + qcpz + + + +# cpp: + +# v^x = dx/dt = dx/dq * vel^q +vel_x = np.zeros([3,len(traj_q[0,:])]) + +for j in range(len(traj_q[0,:])): + + A = np.array([ [np.cos(traj_q[1,j])*np.cos(traj_q[2,j]), -r*np.sin(traj_q[1,j])*np.cos(traj_q[2,j]), -np.sin(traj_q[2,j])], + [np.cos(traj_q[1,j])*np.sin(traj_q[2,j]), -r*np.sin(traj_q[1,j])*np.sin(traj_q[2,j]), np.cos(traj_q[2,j])], + [ np.sin(traj_q[1,j]), r*np.cos(traj_q[1,j]), 0] ]) + + q_dot = traj_v[:,j] + + x_dot = A @ q_dot + + vel_x[:,j] = x_dot + +# Trajectory of particle in Cartesian coordinates + +z1 = (R0 + traj_q[0,:] * np.cos(traj_q[1,:])) * np.cos(traj_q[2,:]) +z2 = (R0 + traj_q[0,:] * np.cos(traj_q[1,:])) * np.sin(traj_q[2,:]) +z3 = z0 + traj_q[0,:] * np.sin(traj_q[1,:]) + + +# first 50 evolutions of the trajectory + +x50 = (R0 + traj_q[0,:50] * np.cos(traj_q[1,:50])) * np.cos(traj_q[2,:50]) +y50 = (R0 + traj_q[0,:50] * np.cos(traj_q[1,:50])) * np.sin(traj_q[2,:50]) +z50 = z0 + traj_q[0,:50] * np.sin(traj_q[1,:50]) + + +# Velocity in Cartesian coordinates +velx = (vel_x[0,:])#*50 + z1 +vely = (vel_x[1,:])#*50 + z2 +velz = (vel_x[2,:])#*50 + z3 + + + + + + +# Calculate Surfaces and Curves at Point p0 +# +# +# +r0, th0, ph0 = traj_q[0,:50], traj_q[1,:50], traj_q[2,:50] +# +# +# +surf_R = [] +surf_Th = [] +surf_Phi = [] +curv_R = [] +curv_Th = [] +curv_Phi = [] + +for i in range(0,50): + print(i) + r0i = r0[i] + th0i = th0[i] + ph0i = ph0[i] + + surf_Ri = np.c_[ (R0 + r0i*np.cos(q[::100,1]))*np.cos(q[::100,2]), + (R0 + r0i*np.cos(q[::100,1]))*np.sin(q[::100,2]), + z0 + r0i*np.sin(q[::100,1]) ] + + surf_Thi = np.c_[ (R0 + q[:,0]*np.cos(th0i))*np.cos(q[:,2]), + (R0 + q[:,0]*np.cos(th0i))*np.sin(q[:,2]), + z0 + q[:,0]*np.sin(th0i) ] + + surf_Phii = np.c_[ (R0 + q[:,0]*np.cos(q[:,1]))*np.cos(ph0i), + (R0 + q[:,0]*np.cos(q[:,1]))*np.sin(ph0i), + z0 + q[:,0]*np.sin(q[:,1]) ] + + curv_Ri = np.c_[ (R0 + q[::100,0]*np.cos(th0i))*np.cos(ph0i), + (R0 + q[::100,0]*np.cos(th0i))*np.sin(ph0i), + z0 + q[::100,0]*np.sin(th0i) ] + + curv_Thi = np.c_[ (R0 + r0i*np.cos(q[::100,1]))*np.cos(ph0i), + (R0 + r0i*np.cos(q[::100,1]))*np.sin(ph0i), + z0 + r0i*np.sin(q[::100,1]) ] + + curv_Phii = np.c_[ (R0 + r0i*np.cos(th0i))*np.cos(q[::100,2]), + (R0 + r0i*np.cos(th0i))*np.sin(q[::100,2]), + z0 + r0i*np.sin(th0i)*np.ones(np.size(q[::100,2])) ] + + surf_R.append(surf_Ri) + surf_Th.append(surf_Thi) + surf_Phi.append(surf_Phii) + curv_R.append(curv_Ri) + curv_Th.append(curv_Thi) + curv_Phi.append(curv_Phii) + + + +# Interpolation of the field on the grid points + + +# Field evaluation on trajectory at grid points + +B = np.zeros(c) +k = 0 +zero = 1000 + + +for i in range(0,c): + r = q[i,0] + th = q[i,1] + ph = q[i,2] + f.evaluate(r,th,ph) + B[i] = f.B + + + if i == 0: + const_B = [] + f.evaluate(q[zero,0], q[zero,1], q[zero,2]) + const_B.append([q[zero,0], q[zero,1], q[zero,2], f.B]) + eps = abs(const_B[i][3] - B[0]) + print(eps) + + elif eps <= 1e-4: + const_B.append([q[i-1,0], q[i-1,1], q[i-1,2], B[i-1]]) + eps = abs(B[i] - const_B[0][3]) + k += 1 + + else: + eps = abs(const_B[0][3] - B[i]) + + + + + +print('c =' , c) +print('k =' , k) +# with k the number of points with almost the same B field value +# and c the total number of grid points inside the torus +arr = np.array(const_B) +print(np.shape(arr)) +rB = arr[:,0] +thB = arr[:,1] +phB = arr[:,2] +B_const = arr[:,3] +print(len(rB)) + +consBx = (R0 + rB * np.cos(thB)) * np.cos(phB) +consBy = (R0 + rB * np.cos(thB)) * np.sin(phB) +consBz = z0 + rB * np.sin(thB) +consB = np.array([consBx, consBy, consBz]) + +# Field evaluation on trajectory +B_traj = np.zeros(np.size(traj_q[0,:])) +for j in range(0,np.size(traj_q[0,:])): + r = traj_q[0,j] + th = traj_q[1,j] + ph = traj_q[2,j] + f.evaluate(r,th,ph) + B_traj[j] = f.B + + +# plot Banana orbit (Curie plot idk) + +fig, ax = plt.subplots() + +plot_orbit(traj_q, ax=ax) +plot_orbit(traj_q_cp, ax=ax) +plot_orbit(consB, ax=ax) +plt.show() + + +fig, ax = plt.subplots() + +plot_orbit(traj_q, ax=ax) +plot_orbit(traj_q_cp, ax=ax) +plt.show() + + + +# cp: Plot evverything + +fig = plt.figure() +ax = fig.add_subplot(111, projection='3d') + +sc = ax.scatter(qcpx, qcpy, qcpz, c='r', cmap='viridis', s=1) +sc = ax.scatter(vcpx+qcpx, vcpy+qcpy, vcpz+qcpz, color='green', marker='o', s=1, label="Velocity points") +#sc = ax.scatter(vcpx, vcpy, vcpz, color='green', marker='o', s=1, label="Velocity relative") +for (q1, th1, ph1, q2, th2, ph2) in zip(qcpx, qcpy, qcpz, qcpx+vcpx, qcpy+vcpy, qcpz+vcpz): + ax.plot([q1, q2], [th1, th2], [ph1, ph2], + color='gray', linestyle='--', linewidth=0.8) + +sc = ax.plot(qcpx[:50]+vcpx[:50], qcpy[:50]+vcpy[:50], qcpz[:50]+vcpz[:50], color='g', linewidth=1.5, label="Line through 50 points") + +plt.title('Classical Particle in Configuration Space Q') +#ax.set_xlim([0.95,1.05]) +#ax.set_ylim([-0.1,0]) +#ax.set_zlim([1.005,1.105]) +plt.show() + + + + +# cpp: Plot everything + +fig = plt.figure() + +ax = fig.add_subplot(111, projection='3d') + + +#sc = ax.scatter(x, y, z, c=B, cmap='viridis', s=0.1) + +#sc = ax.scatter(consBx, consBy, consBz, c=B_const, cmap='viridis', marker='o', s=50, label="Constant B") +sc = ax.plot(consBx, consBy, consBz, c='b', linewidth=1.5, label="Constant B") + +#sc = ax.scatter(surf_R[10][:,0], surf_R[10][:,1], surf_R[10][:,2], color='b', marker='o', s=20)#marker='o', s=200, cmap='viridis') +#sc = ax.scatter(surf_Th[10][:,0], surf_Th[10][:,1], surf_Th[10][:,2], c=B, cmap='viridis', s=0.1) +#sc = ax.scatter(surf_Phi[:,0], surf_Phi[:,1], surf_Phi[:,2], c=B, cmap='viridis', s=0.1) + +#sc = ax.scatter(curv_R[10][:,0], curv_R[10][:,1], curv_R[10][:,2], color='g', s=0.5) +#sc = ax.scatter(curv_Th[10][:,0], curv_Th[10][:,1], curv_Th[10][:,2], color='g', s=0.5) +#sc = ax.scatter(curv_Phi[10][:,0], curv_Phi[10][:,1], curv_Phi[10][:,2], color='g', s=0.5) + + + +### Trajectory of the particle: zi(t) = xi(q(t)) = f(T(t)), with q(t) being the curves on config manif Q. +# eqiv: zi(t) \in R <=> R -> Q -> R ---> 'Every value of 't' maps with a curve to a unique point in Q' ---> the coord. func. xi(T(t)) + +sc = ax.scatter(z1, z2, z3, c='r', cmap='viridis', s=1) + +#ax.plot(x50, y50, z50, color='red', linewidth=1.5, label="Line through 50 points") +#ax.scatter(x50[::10], y50[::10], z50[::10], color='red', marker='o') +#sc = ax.scatter(x50, y50, z50, c='red', cmap='viridis', s=1) + +### Plot velocity at each point of the trajectory + +# Scatter of velocity points +sc = ax.scatter(velx+z1, vely+z2, velz+z3, color='green', marker='o', s=1, label="Velocity points") +#sc = ax.scatter(velx, vely, velz, color='green', marker='o', s=1, label="Velocity relative") + +# Connect corresponding points from (x,y,z) to (velx,vely,velz) +for (x1, x2, x3, v1, v2, v3) in zip(z1, z2, z3, velx*50+z1, vely*50+z2, velz*50+z3): + ax.plot([x1, v1], [x2, v2], [x3, v3], + color='gray', linestyle='--', linewidth=0.8) + + +for (x1, y1, z1, x2, y2, z2) in zip(z1[::50], z2[::50], z3[::50], velx[::50], vely[::50], velz[::50]): + ax.plot([x1, x2], [y1, y2], [z1, z2], + color='gray', linestyle='--', linewidth=0.8) + + + +#ax.set_xlim([-1,1]) +#ax.set_ylim([0,2]) +#ax.set_zlim([0,2]) +plt.show() + + diff --git a/python/pauli_particle.py b/python/pauli_particle.py deleted file mode 100644 index 68d7354e..00000000 --- a/python/pauli_particle.py +++ /dev/null @@ -1,87 +0,0 @@ -# %% -import matplotlib.pyplot as plt -import numpy as np -from scipy.optimize import root -from scipy.interpolate import lagrange -import common -from common import (f,r0,th0,ph0,pph0,timesteps,get_val,get_der,mu,qe,m,c,) -from plotting import plot_orbit, plot_cost_function - -dt, nt = timesteps(steps_per_bounce=8, nbounce=100) - -z = np.zeros([3, nt + 1]) -z[:, 0] = [r0, th0, ph0] -[H, pth, vpar] = get_val(np.array([r0, th0, ph0, pph0])) - -''' -def newton(z, vpar, v): - r = z[0] - th = z[1] - e = 1 - - while e >= 1e-11: - f.evaluate(r,th,z[2]) - g = m*vpar*f.hth + qe/c*f.Ath - dg = m*vpar*f.dhth[0] + qe/c*f.dAth[0] - rnew = r - g/dg - r = rnew - e = g - return rnew -''' - -calc_fields = lambda vpar: { - "A2": m * vpar * f.hth + qe * f.Ath / c, - "dA2": m * vpar * f.dhth + qe * f.dAth / c, - "d2A2": m * vpar * f.d2hth + qe * f.d2Ath / c, - "A3": m * vpar * f.hph + qe * f.Aph / c, - "dA3": m * vpar * f.dhph + qe * f.dAph / c, - "d2A3": m * vpar * f.d2hph + qe * f.d2Aph / c, - "dB": mu * f.dB, - "d2B": mu * f.d2B, - "dPhie": f.dPhie, - "d2Phie": f.d2Phie, -} - -def F(z, vpar, v): - ret = np.zeros(2) - - f.evaluate(z[0],z[1],0) - - ret[0] = m*vpar*f.hth + qe/c *f.Ath - dt*v[0] - ret[1] = m*vpar*f.hph + qe/c *f.Aph - dt*v[1] - return ret - - - -from time import time -tic = time() -for kt in range(nt): - vparold = vpar - f.evaluate(z[0, kt], z[1, kt], z[2, kt]) - cov = calc_fields(vparold) - v1 = (qe*cov['dPhie'][2] - mu*cov['dB'][2] + vparold*(m * f.hph)) / (cov['dA2'][2] - cov['dA3'][1]) - v2 = (qe*cov['dPhie'][1] - mu*cov['dB'][1] + vparold*(m * f.hth)) / (cov['dA3'][1] - cov['dA2'][2]) - # v2 = (qe*f.dPhie[0] - f.dAph[0]*vparold - mu*f.dB[0]) / (f.dAth[0]) - # v1 = ((f.dAph[1] - f.dAth[2])*vparold - qe*f.dPhie[1] + mu*f.dB[1]) / (f.dAth[0]) - - #vpar = vparold + dt*(-f.dAph[0]*v1 + (f.dAth[2] - f.dAph[1])*v2) - vpar = vparold + dt*(m*f.hth*v1 + m*f.hph*v2) - v = np.array([v1,v2,vpar]) - - x0 = z[:2,kt] - sol = root(F, x0, method='hybr',tol=1e-12,args=(vpar, v)) - z[:2,kt+1] = sol.x - #z[0,kt+1] = newton(z[:,kt], vpar, v) - -plot_orbit(z) -plt.show() - - - - -#implicit - - - - - diff --git a/python/plotale.py b/python/plotale.py new file mode 100644 index 00000000..0a14041c --- /dev/null +++ b/python/plotale.py @@ -0,0 +1,98 @@ +from numpy import cos, sin, empty_like, linspace, roll +import matplotlib.pyplot as plt +import numpy as np +from matplotlib.pyplot import figure, plot, xlabel, ylabel, subplot, grid, legend + + + +def plot_orbit(z, ax): + + ax.plot(z[0, :] * cos(z[1, :]), + z[0, :] * sin(z[1, :]), + 'o', markersize=1) + + ax.set_xlabel(r'$R-R_0$') + ax.set_ylabel(r'$Z$') + return ax + + +def plot_mani(z): + # Torus parameters + R0 = 1 # Distance from the center of the tube to the center of the torus + r = 0.5 # Radius of the tube + + # Create meshgrid + theta = linspace(0, 2 * np.pi, 100) + phi = linspace(0.3, 1.4*np.pi, 100) + theta, phi = np.meshgrid(theta, phi) + + # Parametric equations for a torus + X = (R0 + r * cos(theta)) * cos(phi) + Y = (R0 + r * cos(theta)) * sin(phi) + Z = r * sin(theta) + + # Plotting + fig = plt.figure(figsize=(8, 6)) + ax = fig.add_subplot(111, projection='3d') + #ax.plot_surface(X, Y, Z, cmap='plasma', edgecolor='none') + + x = (R0 + z[0,:]*cos(z[1,:]))*cos(z[2,:]) + y = (R0 + z[0,:]*cos(z[1,:]))*sin(z[2,:]) + z = z[0,:]*sin(z[1,:]) + ax.scatter(x,y,z) + ax.set_zlim([-1,1]) + + + + + + + +def plot_cost_function(F, zlast, zold, pthold): + """Plot the cost function""" + r = linspace(-0.2,0.2,100) + fres = empty_like(r) + for kr in range(len(r)): + fres[kr] = F([r[kr]], zold[1:], pthold) + + figure() + subplot(2,1,1) + plot(r,fres) + out_conv = F([zlast[0]], zold[1:], pthold) + plot(zlast[0], out_conv, 'rx') + xlabel('r') + ylabel('F') + grid() + +def plot_cost_function_jac(F, zlast, zold, pthold): + """Plot the cost function with Jacobian""" + + r = linspace(-0.2,0.2,100) + fres = empty_like(r) + jac = empty_like(r) + for kr in range(len(r)): + out = F([r[kr]], zold[1:], pthold) + fres[kr] = out[0] + jac[kr] = out[1][0] + + jacnum = (roll(fres,-1)-roll(fres,1))/(2.*(r[1]-r[0])) + jacnum[0] = jacnum[1] + jacnum[-1] = jacnum[-2] + + figure() + subplot(2,1,1) + plot(r,fres) + out_conv = F([zlast[0]],zold[1:],pthold) + plot(zlast[0],out_conv[0],'rx') + xlabel('r') + ylabel('F') + grid() + + subplot(2,1,2) + plot(r,jac) + plot(r,jacnum,'r--') + plot(zlast[0],out_conv[1],'rx') + xlabel('r') + ylabel('dF/dr') + legend(['analytical', 'numerical']) + grid() diff --git a/python/test_delete_later/field.py b/python/test_delete_later/field.py new file mode 100644 index 00000000..ce6db16b --- /dev/null +++ b/python/test_delete_later/field.py @@ -0,0 +1,143 @@ +import numpy as np +from numba import jit, njit +from util import legendre_p, assoc_legendre_p, ellipk, ellipe + +EPS = 1e-13 + +# Fields are normalized to 1A*mu0/2pi + + +@njit +def afield_simple(X, Y, Z, A, dA3): + A[0] = 0.0 + A[1] = 0.0 + A[2] = 0.5 * (X**2 + Y**2) + dA3[:] = np.array([X,Y,0.0]) + +''' +@njit +def afield_simple(X, Y, Z, A): + A[0] = 0.0 + A[1] = 0.0 + A[2] = 0.5 * (X**2 + Y**2) +''' + +@njit +def bfield_simple(X, Y, Z, B): + B[0] = Y + B[1] = -X + B[2] = 0.0 + + +@njit +def bmod_simple(X, Y, Z, Bmod, dBmod): + Bmod[:] = np.sqrt(X**2 + Y**2) + dBmod[:] = np.array([X, Y, 0.0]) / Bmod + + +@njit +def afield_wire(X, Y, Z, A): + Acyl = np.zeros(3) + cyl_to_cart(afield_wire_cyl, X, Y, Z, A) + + +@njit +def bfield_wire(X, Y, Z, B): + Bcyl = np.zeros(3) + cyl_to_cart(bfield_wire_cyl, X, Y, Z, B) + + +@njit +def bfield_circle(X, Y, Z, B): + Bcyl = np.zeros(3) + cyl_to_cart(bfield_circle_cyl, X, Y, Z, B) + + +@njit +def afield_wire_cyl(R, phi, Z, A): + A[0] = 0.0 + A[1] = 0.0 + A[2] = -np.log(max(R, EPS)) + + +@njit +def bfield_wire_cyl(R, phi, Z, B): + B[0] = 0.0 + B[1] = 1.0 / max(R, EPS) + B[2] = 0.0 + + +@njit +def bfield_circle_cyl(R, phi, Z, B): + """ + Calculate the magnetic field of a circular loop in cylindrical coordinates. + http://plasmalab.pbworks.com/w/file/fetch/109941181/Bfield-calcs.pdf + https://tiggerntatie.github.io/emagnet-py/offaxis/off_axis_loop.html + """ + + a = 1.0 + r = max(np.sqrt(R**2 + Z**2), EPS) + + alpha = r / a + beta = Z / a + gamma = Z / r + Q = (1 + alpha) ** 2 + beta**2 + k2 = 4 * alpha / Q + + K = ellipk(k2) + E = ellipe(k2) + + B[0] = ( + gamma + / (np.pi * np.sqrt(Q)) + * (E * (1 + alpha**2 + beta**2) / (Q - 4 * alpha) - K) + ) + B[2] = ( + 1.0 + / (np.pi * np.sqrt(Q)) + * (E * (1 - alpha**2 - beta**2) / (Q - 4 * alpha) + K) + ) + + +@njit +def bfield_circle_cyl2(R, phi, Z, B): + """ + Calculate the magnetic field of a circular loop in cylindrical coordinates. + See https://farside.ph.utexas.edu/teaching/jk1/Electromagnetism/node52.html + Unstable at r=a + """ + LMAX = 128 + B[:] = 0.0 + + r = max(np.sqrt(R**2 + Z**2), EPS) + costh = Z / r + sinth = R / r + for l in np.arange(1, LMAX, 2): + Pl10 = assoc_legendre_p(l, 1, 0.0) + Br = Pl10 * legendre_p(l, costh) + Bth = Pl10 * assoc_legendre_p(l, 1, costh) + if r < 1: + Br *= r ** (l - 1.0) + Bth *= r ** (l - 1.0) / l + else: + Br *= r ** (-l - 2.0) + Bth *= -(r ** (-l - 2.0)) / (l + 1.0) + + B[0] += Br * sinth + Bth * costh + B[2] += Br * costh - Bth * sinth + + B[:] = -np.pi * B[:] + + +@njit +def cyl_to_cart(fun, X, Y, Z, B): + Bcyl = np.zeros(3) + R = np.sqrt(X**2 + Y**2) + phi = np.arctan2(Y, X) + fun(R, phi, Z, Bcyl) + + cosphi = X / max(R, EPS) + sinphi = Y / max(R, EPS) + B[0] = Bcyl[0] * cosphi - Bcyl[1] * sinphi + B[1] = Bcyl[0] * sinphi + Bcyl[1] * cosphi + B[2] = Bcyl[2] \ No newline at end of file diff --git a/python/test_delete_later/main.py b/python/test_delete_later/main.py new file mode 100644 index 00000000..f8de58bd --- /dev/null +++ b/python/test_delete_later/main.py @@ -0,0 +1,88 @@ +#%% +from numba import njit +from field import afield_simple, bfield_simple, bmod_simple +from scipy.integrate import solve_ivp +import matplotlib.pyplot as plt +from mpl_toolkits.mplot3d import Axes3D +import numpy as np + + +@njit +def velo(t, z): + B = np.empty(3) + bfield_simple(z[0], z[1], z[2], B) + + zdot = np.empty(6) + zdot[0] = z[3] + zdot[1] = z[4] + zdot[2] = z[5] + zdot[3] = z[4] * B[2] - z[5] * B[1] + zdot[4] = z[5] * B[0] - z[3] * B[2] + zdot[5] = z[3] * B[1] - z[4] * B[0] + + return zdot + + +@njit +def velo_cpp(t, z, mu0ovm): + B = np.empty(3) + Bmod = np.empty(1) + dBmod = np.empty(3) + bfield_simple(z[0], z[1], z[2], B) + bmod_simple(z[0], z[1], z[2], Bmod, dBmod) + + zdot = np.empty(6) + zdot[0] = z[3] + zdot[1] = z[4] + zdot[2] = z[5] + zdot[3] = z[4] * B[2] - z[5] * B[1] - mu0ovm * dBmod[0] + zdot[4] = z[5] * B[0] - z[3] * B[2] - mu0ovm * dBmod[1] + zdot[5] = z[3] * B[1] - z[4] * B[0] - mu0ovm * dBmod[2] + + return zdot + + +# %% +z0 = np.array([1.0, 2.0, 3.0, 0.1, 0.1, 0.1]) + +sol = solve_ivp(velo, [0, 500], z0, method="RK45", rtol=1e-9, atol=1e-9) + +z0_cpp = np.array([1.04, 2.04, 2.94, 0.1, 0.1, 0.1]) # Guessed guiding center + +# Set initial conditions for guiding center CPP +B = np.empty(3) +Bmod = np.empty(1) +dBmod = np.empty(3) +bfield_simple(z0_cpp[0], z0_cpp[1], z0_cpp[2], B) +bmod_simple(z0_cpp[0], z0_cpp[1], z0_cpp[2], Bmod, dBmod) + +vpar = np.dot(z0_cpp[3:], B/Bmod) +vperp = np.sqrt(np.sum(z0_cpp[3:]**2) - vpar**2) +mu0ovm = 0.5 * vperp**2 / Bmod +print(mu0ovm) + +z0_cpp[3:] = vpar*B/Bmod + + + + +sol_cpp = solve_ivp( + velo_cpp, [0, 500], z0_cpp, method="RK45", rtol=1e-9, atol=1e-9, args=(mu0ovm) +) + +# %% +plt.figure() +plt.plot(sol.y[0], sol.y[1]) +plt.plot(z0_cpp[0], z0_cpp[1], "ro") +plt.plot(sol_cpp.y[0], sol_cpp.y[1]) +plt.show() +# %% +plt.figure() +plt.plot(sol.y[1], sol.y[2]) +plt.plot(z0_cpp[1], z0_cpp[2], "ro") +plt.plot(sol_cpp.y[1], sol_cpp.y[2]) +plt.show() + +print(len(sol_cpp.t)) # This will print the number of time steps the solver used + +# %% \ No newline at end of file diff --git a/python/test_delete_later/mainHam.py b/python/test_delete_later/mainHam.py new file mode 100644 index 00000000..4bf0d4e7 --- /dev/null +++ b/python/test_delete_later/mainHam.py @@ -0,0 +1,88 @@ +# %% +from numba import njit +from field import afield_simple, bfield_simple, bmod_simple +from scipy.integrate import solve_ivp +from scipy.optimize import root +import matplotlib.pyplot as plt +from mpl_toolkits.mplot3d import Axes3D +import numpy as np + +taub = 7800 # estimated bounce time + +def timesteps(steps_per_bounce, nbounce): + """ compute step size and number of timesteps """ + return taub/steps_per_bounce, nbounce*steps_per_bounce + + +qe = 1.0; m = 1.0; c = 1.0 # particle charge, mass and speed of light +dt, nt = timesteps(steps_per_bounce = 8, nbounce = 100) + +dt = 0.01 +nt = 50000 +z0_cpp = np.array([1.04, 2.04, 2.94, 0.1, 0.1, 0.1]) # Guessed guiding center +p = np.zeros(3) + +# Set initial conditions for guiding center CPP +B = np.empty(3) +Bmod = np.empty(1) +dBmod = np.empty(3) +A = np.empty(3) +dA3 = np.empty(3) +bfield_simple(z0_cpp[0], z0_cpp[1], z0_cpp[2], B) +bmod_simple(z0_cpp[0], z0_cpp[1], z0_cpp[2], Bmod, dBmod) +afield_simple(z0_cpp[0], z0_cpp[1], z0_cpp[2], A, dA3) + +vpar = np.dot(z0_cpp[3:], B/Bmod) +vperp = np.sqrt(np.sum(z0_cpp[3:]**2) - vpar**2) +mu = 0.5 * vperp**2 / Bmod +mu0 = mu[0] + + +z = np.zeros([3,nt+1]) +z[:,0] = z0_cpp[0:3] +p = np.zeros(3) +p = vpar*B/Bmod + qe/c * A +pold = p + + + + +def implicit_p(p, pold, A, dA3, dBmod, mu0): + ret = np.zeros(3) + + ret[0] = p[0] - pold[0] - dt*(qe/c*(p[2] - qe/c * A[2]) * dA3[0] - mu0*dBmod[0]) + ret[1] = p[1] - pold[1] - dt*(qe/c*(p[2] - qe/c * A[2]) * dA3[1] - mu0*dBmod[1]) + ret[2] = p[2] - pold[2] - dt*(qe/c*(p[2] - qe/c * A[2]) * dA3[2] - mu0*dBmod[2]) + + return ret + + +for kt in range(nt): + + A = np.empty(3) + dA3 = np.empty(3) + Bmod = np.empty(1) + dBmod = np.empty(3) + + bmod_simple(z[0,kt], z[1,kt], z[2,kt], Bmod, dBmod) + afield_simple(z[0,kt], z[1,kt], z[2,kt], A, dA3) + + + sol = root(implicit_p, p, method='hybr',tol=1e-12,args=(pold, A, dA3, dBmod, mu0)) + pold = p + p = sol.x + + + z[0,kt+1] = z[0,kt] + dt*(p[0] - qe/c*A[0]) + z[1,kt+1] = z[1,kt] + dt*(p[1] - qe/c*A[1]) + z[2,kt+1] = z[2,kt] + dt*(p[2] - qe/c*A[2]) + + + +plt.figure() +plt.plot(z[0,:], z[1,:]) +plt.plot(z0_cpp[0], z0_cpp[1], "ro") +plt.show() + + +# %% diff --git a/python/test_delete_later/mainLa.py b/python/test_delete_later/mainLa.py new file mode 100644 index 00000000..5e18582e --- /dev/null +++ b/python/test_delete_later/mainLa.py @@ -0,0 +1,86 @@ +# %% +from numba import njit +from field import afield_simple, bfield_simple, bmod_simple +from scipy.integrate import solve_ivp +from scipy.optimize import root +import matplotlib.pyplot as plt +from mpl_toolkits.mplot3d import Axes3D +import numpy as np + +taub = 7800 # estimated bounce time + +def timesteps(steps_per_bounce, nbounce): + """ compute step size and number of timesteps """ + return taub/steps_per_bounce, nbounce*steps_per_bounce + + +qe = 1.0; m = 1.0; c = 1.0 # particle charge, mass and speed of light +dt, nt = timesteps(steps_per_bounce = 8, nbounce = 100) + +dt = 0.104 +nt = 5000 +z0_cpp = np.array([1.04, 2.04, 2.94, 0.1, 0.1, 0.1]) # Guessed guiding center +p = np.zeros(3) + +# Set initial conditions for guiding center CPP +B = np.empty(3) +Bmod = np.empty(1) +dBmod = np.empty(3) +A = np.empty(3) +dA3 = np.empty(3) +bfield_simple(z0_cpp[0], z0_cpp[1], z0_cpp[2], B) +bmod_simple(z0_cpp[0], z0_cpp[1], z0_cpp[2], Bmod, dBmod) +afield_simple(z0_cpp[0], z0_cpp[1], z0_cpp[2], A, dA3) + +vpar = np.dot(z0_cpp[3:], B/Bmod) +vperp = np.sqrt(np.sum(z0_cpp[3:]**2) - vpar**2) +mu = 0.5 * vperp**2 / Bmod +mu0 = mu[0] +print(mu0) + +z = np.zeros([3,nt+1]) +z[:,0] = z0_cpp[0:3] +vel = np.zeros([3,nt+1]) +vel[:,0] = vpar*B/Bmod +v = vel[:,0] +vold = v + + +def implicit_v(v, vold, B, dBmod, mu0): + ret = np.zeros(3) + + ret[0] = v[0] - vold[0] - dt*(v[1]*B[2] - v[2]*B[1] - mu0*dBmod[0]) + ret[1] = v[1] - vold[1] - dt*(v[2]*B[0] - v[0]*B[2] - mu0*dBmod[1]) + ret[2] = v[2] - vold[2] - dt*(v[0]*B[1] - v[1]*B[0] - mu0*dBmod[2]) + + return ret + + +for kt in range(nt): + + B = np.empty(3) + Bmod = np.empty(1) + dBmod = np.empty(3) + + bmod_simple(z[0,kt], z[1,kt], z[2,kt], Bmod, dBmod) + bfield_simple(z[0,kt], z[1,kt], z[2,kt], B) + + + sol = root(implicit_v, v, method='hybr',tol=1e-12,args=(vold, B, dBmod, mu0)) + vold = v + v = sol.x + + + z[0,kt+1] = z[0,kt] + dt*v[0] + z[1,kt+1] = z[1,kt] + dt*v[1] + z[2,kt+1] = z[2,kt] + dt*v[2] + + + +plt.figure() +plt.plot(z[0,:], z[1,:]) +plt.plot(z0_cpp[0], z0_cpp[1], "ro") +plt.show() + + +# %% diff --git a/python/test_delete_later/util.py b/python/test_delete_later/util.py new file mode 100644 index 00000000..8ad30d56 --- /dev/null +++ b/python/test_delete_later/util.py @@ -0,0 +1,197 @@ +import numpy as np +from numba import njit + + +@njit +def legendre_p(n, x): + """ " + Returns the value of the Legendre polynomial P_n(x). + + Parameters: + n (int): The degree of the polynomial. + x (float): The point at which to evaluate the polynomial. Must satisfy -1 <= x <= 1. + """ + if n == 0: + return 1.0 + if n == 1: + return x + + p_prev2 = 1.0 + p_prev1 = x + for n in range(2, n + 1): + p = ((2 * n - 1) * x * p_prev1 - (n - 1) * p_prev2) / n + p_prev2 = p_prev1 + p_prev1 = p + return p_prev1 + + +@njit +def assoc_legendre_p(n, m, x): + """ + Returns the value of the associated Legendre polynomial P_n^m(x). + + Parameters: + n (int): The degree of the polynomial. Must be greater than or equal to m. + m (int): The order of the polynomial. Must satisfy 0 <= m <= n. + x (float): The point at which to evaluate the polynomial. Must satisfy -1 <= x <= 1. + """ + if m < 0 or m > n: + return 0.0 + + if m == n: + fact = 1.0 + for k in range(n): + fact *= -(2 * k + 1) + return fact * (1 - x * x) ** (n / 2) + + pmm = 1.0 + fact = 1.0 + for k in range(m): + pmm *= -(2 * k + 1) * (1 - x * x) ** (0.5) + + if n == m: + return pmm + + pmmp1 = x * (2 * m + 1) * pmm + + if n == m + 1: + return pmmp1 + + for n in range(m + 2, n + 1): + pmn = (x * (2 * n - 1) * pmmp1 - (n + m - 1) * pmm) / (n - m) + pmm = pmmp1 + pmmp1 = pmn + + return pmmp1 + + +@njit +def polevl(x: float, coef: np.ndarray, N: int) -> float: + ans = coef[0] + i = N + p = 1 # Indexing in Python instead of pointer + + while i: + ans = ans * x + coef[p] + p += 1 + i -= 1 + + return ans + + +@njit +def ellipe(m): + """ + Compute the complete elliptic integral of the second kind. + Adapted from Cephes Math Library by Stephen L. Moshier + MIT License + + Parameters: + m (float): Input parameter (0 <= x <= 1) + + Returns: + float: Computed value of E(m) + """ + + P = np.array( + [ + 1.53552577301013293365e-4, + 2.50888492163602060990e-3, + 8.68786816565889628429e-3, + 1.07350949056076193403e-2, + 7.77395492516787092951e-3, + 7.58395289413514708519e-3, + 1.15688436810574127319e-2, + 2.18317996015557253103e-2, + 5.68051945617860553470e-2, + 4.43147180560990850618e-1, + 1.00000000000000000299e0, + ] + ) + + Q = np.array( + [ + 3.27954898576485872656e-5, + 1.00962792679356715133e-3, + 6.50609489976927491433e-3, + 1.68862163993311317300e-2, + 2.61769742454493659583e-2, + 3.34833904888224918614e-2, + 4.27180926518931511717e-2, + 5.85936634471101055642e-2, + 9.37499997197644278445e-2, + 2.49999999999888314361e-1, + ] + ) + + x = 1.0 - m + if x <= 0.0 or x > 1.0: + if x == 0.0: + return 1.0 + return 0.0 # Domain error handling + + return polevl(x, P, 10) - np.log(x) * (x * polevl(x, Q, 9)) + + +@njit +def ellipk(m): + """ + Compute the complete elliptic integral of the second kind. + Adapted from Cephes Math Library by Stephen L. Moshier + MIT License + + Parameters: + m (float): Input parameter (0 <= m <= 1) + + Returns: + float: Computed value of K(m) + """ + + MACHEP = 1.11022302462515654042e-16 + + P = np.array( + [ + 1.37982864606273237150e-4, + 2.28025724005875567385e-3, + 7.97404013220415179367e-3, + 9.85821379021226008714e-3, + 6.87489687449949877925e-3, + 6.18901033637687613229e-3, + 8.79078273952743772254e-3, + 1.49380448916805252718e-2, + 3.08851465246711995998e-2, + 9.65735902811690126535e-2, + 1.38629436111989062502e0, + ] + ) + + Q = np.array( + [ + 2.94078955048598507511e-5, + 9.14184723865917226571e-4, + 5.94058303753167793257e-3, + 1.54850516649762399335e-2, + 2.39089602715924892727e-2, + 3.01204715227604046988e-2, + 3.73774314173823228969e-2, + 4.88280347570998239232e-2, + 7.03124996963957469739e-2, + 1.24999999999870820058e-1, + 4.99999999999999999821e-1, + ] + ) + + C1 = 1.3862943611198906188e0 # log(4) + + x = 1.0 - m + + if x < 0.0 or x > 1.0: + return 0.0 + + if x > MACHEP: + return polevl(x, P, 10) - np.log(x) * polevl(x, Q, 10) + else: + if x == 0.0: + return np.inf + else: + return C1 - 0.5 * np.log(x) \ No newline at end of file