import numpy as np
import pyoscode
import matplotlib.pyplot as plt
plt.ion()
import scipy.special as sp
import scipy.integrate
from scipy.optimize import minimize_scalar
import warnings
warnings.filterwarnings('ignore')
from ipywidgets import interact, IntSlider, FloatSlider,FloatLogSlider
%config InlineBackend.figure_formats = ['svg']
%matplotlib inline
from IPython.core.display import HTML
HTML("""
<style>
.slides {
width: 100% !important;
}
.middle > * {
vertical-align: middle;
}
</style>
""")
In /home/fruzsina/.config/matplotlib/stylelib/paper.mplstyle: The text.latex.preview rcparam was deprecated in Matplotlib 3.3 and will be removed two minor releases later. In /home/fruzsina/.config/matplotlib/stylelib/paper-elsevier.mplstyle: The text.latex.preview rcparam was deprecated in Matplotlib 3.3 and will be removed two minor releases later. Missing colon in file PosixPath('/home/fruzsina/.config/matplotlib/stylelib/kicc.mplstyle'), line 61 ('patch.linewidth = 0.5')
def wkb_step_airy(t,x,dx,h,order,n):
times = np.linspace(t,t+h,n)
ds0 = 1j*(np.sqrt(t))
ds1 = -1./4*(t**(-1))
ds2 = 1j*5./32*(t**(-5./2))
ds3 = 15./64*(t**(-4.))
s0 = 1j*2./3*((times)**(3./2)-t**(3./2))
s1 = -1./4*np.log((times)/t)
s2 = -1j*5./48*((times)**(-3./2)-t**(-3./2))
s3 = -5./64*((times)**(-3)-t**(-3))
alls = [s0,s1,s2,s3]
allds = [ds0,ds1,ds2,ds3]
fp = np.exp(sum(alls[:order+1]))
fm = np.conj(fp)
dfp = sum(allds[:order+1])
dfm = np.conj(dfp)
Ap = (dx - x*dfm)/(dfp - dfm)
Am = (dx - x*dfp)/(dfm - dfp)
y = Ap*fp + Am*fm
return times,y
To run these slides in jupyter,
head to https://github.com/fruzsinaagocs/oscode and
def f(t,x):
return np.array([x[1],-t*x[0]])
ti = 1.0
tf = 1e3
x0 = np.array([sp.airy(-ti)[0], -sp.airy(-ti)[1]])
%%time
sol = scipy.integrate.solve_ivp(f,(ti,tf),x0,rtol=1e-3)
print(sol.success)
print("x(t={})={:.4f}, dx/dt={:.4f}".format(tf,sol.y[0][-1],sol.y[1][-1]))
print("analytic solution: x={:.4f}, dx/dt={:.4f}".format(sp.airy(-tf)[0],-sp.airy(-tf)[1]))
True x(t=1000.0)=0.0302, dx/dt=-2.3852 analytic solution: x=0.0560, dx/dt=-2.6331 CPU times: user 1.63 s, sys: 0 ns, total: 1.63 s Wall time: 1.64 s
scipy.integrate.solve_ivp
's default method is 4(5) order Runge-KuttaBased on arxiv:1906.01421
def airy_order(t0, order):
t1 = t0 + 30.0
x0 = sp.airy(-t0)[0]
dx0 = -sp.airy(-t0)[1]
times, sol = wkb_step_airy(t0,x0,dx0,t1-t0,order,1000)
times_dense = np.linspace(t0,t1,1000)
sol_dense = np.array([sp.airy(-t)[0] for t in times_dense])
fig, ax = plt.subplots(2,1,figsize=(10,5),sharex=True)
plt.tight_layout()
plt.title('WKB approximation for the Airy equation, $\ddot{x} + tx = 0$')
ax[0].plot(times_dense,sol_dense,color='black',label='analytic solution')
ax[0].plot(times,sol,color='C1',label='WKB, order: {}'.format(order))
ax[0].set_ylim(-0.6,0.6)
ax[0].set_ylabel('$x$')
ax[0].legend()
ax[1].semilogy(times_dense,abs(sol-sol_dense),color='black')
ax[1].set_ylim((1e-7,1e-0))
ax[1].set_xlabel('t')
ax[1].set_ylabel('$|\Delta x|$')
plt.show()
interact(airy_order,
t0=FloatSlider(min=1.0, max=10.0, step=1.0), order=IntSlider(min=0, max=3,
step=1), continuous_update=False);
interactive(children=(FloatSlider(value=1.0, description='t0', max=10.0, min=1.0, step=1.0), IntSlider(value=0…
n = 40.
ts1 = np.linspace(-n/2,n/2,100000)
ws1 = np.log(np.sqrt((n**2-1.)))-np.log(1+ts1**2)
gs1 = np.zeros_like(ts1)
ti1 = ts1[0]
tf_dense = ts1[-1]
xi1 = np.sqrt(1.+ti1**2)/n*np.cos(n*np.arctan(ti1))
dxi1 = 1./(n*np.sqrt(1+ti1**2))*(ti1*np.cos(n*np.arctan(ti1))-n*np.sin(n*np.arctan(ti1)))
xs_ana = np.sqrt(1.+ts1**2)/n*(np.cos(n*np.arctan(ts1)))
solution1 = pyoscode.solve(ts1,ws1,gs1,ti1,tf_dense,xi1,dxi1,logw=True,rtol=1e-3)
def burst_f(t,y):
return np.array([y[1],-(n**2-1.)/(1.+t**2)**2*y[0]])
def burst_rk(tf):
sol = scipy.integrate.solve_ivp(burst_f,(ti1,tf),np.array([xi1,dxi1]),rtol=1e-3)
t_rk = sol.t
final_ind = np.argmax(t_rk >= tf)
t_rk = t_rk[:final_ind+1]
x_rk = sol.y[0][:final_ind+1]
plt.figure(figsize=(10,2.5))
plt.tight_layout()
plt.title("scipy.integrate.solve_ivp solving the equation $\ddot{x}+\\frac{n^2-1}{(1+t^2)^2}x=0$, with $n=40$")
plt.plot(ts1,xs_ana,label='analytic solution',color='black',lw=0.7)
plt.xlabel('t')
plt.ylabel('x')
plt.plot(t_rk,x_rk ,'.',label='RK',ms=9.0,color='C1')
plt.legend()
def burst(tf):
t_osc = np.array(solution1['t'])
final_ind = np.argmax(t_osc >= tf)
t_osc = t_osc[:final_ind+1]
x_osc = np.array(solution1['sol'])[:final_ind+1]
types = np.array(solution1['types'])[:final_ind+1]
t_rk = t_osc[types==0]
x_rk = x_osc[types==0]
t_wkb = t_osc[types==1]
x_wkb = x_osc[types==1]
plt.figure(figsize=(10,2.5))
plt.tight_layout()
plt.title("oscode solving the equation $\ddot{x}+\\frac{n^2-1}{(1+t^2)^2}x=0$, with $n=40$")
plt.plot(ts1,xs_ana,label='analytic solution',color='black',lw=0.7)
plt.xlabel('t')
plt.ylabel('x')
plt.plot(t_rk,x_rk ,'.',label='RK',ms=9.0,color='C1')
plt.plot(t_wkb,x_wkb,'.',label='WKB',ms=9.0,color='C0')
plt.legend()
interact(burst,
tf=FloatSlider(min=ti1, max=tf_dense,
step=1.0, continuous_update=False, value=ti1));
interact(burst_rk,
tf=FloatSlider(min=ti1, max=tf_dense,
step=1.0, continuous_update=False, value=ti1));
interactive(children=(FloatSlider(value=-20.0, continuous_update=False, description='tf', max=20.0, min=-20.0,…
interactive(children=(FloatSlider(value=-20.0, continuous_update=False, description='tf', max=20.0, min=-20.0,…
def burst2(n2):
ts2 = np.linspace(-n2,n2,500000)
ws2 = np.log(np.sqrt((n2**2-1.)))-np.log(1+ts2**2)
gs2 = np.zeros_like(ts2)
ti2 = ts2[0]
tf_dense = ts2[-1]
xi2 = np.sqrt(1.+ti2**2)/n2*np.cos(n2*np.arctan(ti2))
dxi2 = 1./(n2*np.sqrt(1+ti2**2))*(ti2*np.cos(n2*np.arctan(ti2))-n2*np.sin(n2*np.arctan(ti2)))
xs_ana2 = np.sqrt(1.+ts2**2)/n2*(np.cos(n2*np.arctan(ts2)))
solution2 = pyoscode.solve(ts2,ws2,gs2,ti2,tf_dense,xi2,dxi2,logw=True,rtol=1e-3)
t_osc = np.array(solution2['t'])
oscs = np.zeros(len(t_osc))
for i,t in enumerate(t_osc):
if i!=0:
oscs[i] = ((n2*np.arctan(t)/(2*np.pi))-(n2*np.arctan(t_osc[i-1])/(2*np.pi)))
else:
oscs[i] = None
types = np.array(solution2['types'])
t_rk = t_osc[types==0]
t_wkb = t_osc[types==1]
plt.figure(figsize=(10,5))
plt.tight_layout()
plt.title("Number of oscillations traversed per step, $\ddot{{x}}+\\frac{{n^2-1}}{{(1+t^2)^2}}x=0$, n={}".format(n2))
plt.xlabel('t')
plt.ylabel('$N_{\mathrm{osc}}$')
plt.semilogy(t_osc,oscs,color='black')
plt.semilogy(t_rk,oscs[types==0] ,'.',label='RK',ms=9.0,color='C1')
plt.semilogy(t_wkb,oscs[types==1],'.',label='WKB',ms=9.0,color='C0')
plt.ylim((1e-2,1e4))
plt.legend()
interact(burst2,
n2=FloatLogSlider(base=10, min=2, max=5,
step=1), continuous_update=False);
interactive(children=(FloatLogSlider(value=100.0, description='n2', max=5.0, min=2.0, step=1.0), Output()), _d…
scipy
):ts = np.linspace(0,1e3,5000)
ws = np.sqrt(ts)
gs = np.zeros_like(ts)
# Initial conditions
ti = 1.0
tf = ts[-1]
x0 = sp.airy(-ti)[0]
dx0 = -sp.airy(-ti)[1]
%%time
pyoscode_sol = pyoscode.solve(ts,ws,gs,ti,tf,x0,dx0,rtol=1e-3)
pyosc_y = pyoscode_sol['sol']
pyosc_dy = pyoscode_sol['dsol']
print("x(t={})={:.4f}, dx/dt={:.4f}".format(tf,pyosc_y[-1],pyosc_dy[-1]))
print("analytic solution: x={:.4f}, dx/dt={:.4f}".format(sp.airy(-tf)[0],-sp.airy(-tf)[1]))
x(t=1000.0)=0.0560+0.0000j, dx/dt=-2.6292+0.0000j analytic solution: x=0.0560, dx/dt=-2.6331 CPU times: user 646 µs, sys: 132 µs, total: 778 µs Wall time: 638 µs
import re
pair = re.compile(r'\(([^,\)]+),([^,\)]+)\)')
def parse_pair(s):
s = s.decode('utf-8')
return complex(*map(float, pair.match(s).groups()))
f1 = "images/bingo-singlek-k1e-5.txt"
f1ref = "images/bingo-singlek-k1e-5-ref.txt"
f1wkb = "images/rkwkb-single-k1e-5.txt"
d1 = np.genfromtxt(f1)
d1ref = np.genfromtxt(f1ref)
d1wkb = np.genfromtxt(f1wkb,dtype=complex,converters={1:parse_pair, 2:parse_pair})
n1 = d1[:,0]
n1ref = d1ref[:,0]
n1wkb = d1wkb[:,0]
rk1 = d1[:,1]
rk1ref = d1ref[:,1]
rk1wkb = d1wkb[:,1]
rk1steps = d1wkb[:,3]
plt.figure(figsize=(10,5))
plt.title('Evolution of a single primordial perturbation $\mathcal{R}_k$ of Fourier wavenumber $k$')
plt.plot(n1ref,rk1ref,color='black',lw=0.7)
plt.plot(n1wkb[rk1steps==0],rk1wkb[rk1steps==0],'.',color='C1',ms=9.0,label='RK step')
plt.plot(n1wkb[rk1steps==1],rk1wkb[rk1steps==1],'.',color='C0',ms=9.0,label='WKB step')
plt.ticklabel_format(axis='y',style='sci',scilimits=(-2,2))
plt.legend()
plt.xlabel('N')
plt.ylabel('$\Re{\{ \mathcal{R}_k \}} $')
plt.show()
fs1 = "images/singlek-kd-bd-2a.txt"
fs2 = "images/singlek-kd-bd-2a.txt"
fref = "images/singlek-kd-bg-2a-ref.txt"
d1 = np.genfromtxt(fs1,dtype=complex,converters={1:parse_pair,2:parse_pair,4:parse_pair})
d2 = np.genfromtxt(fs2,dtype=complex,converters={1:parse_pair,2:parse_pair,4:parse_pair})
dref = np.genfromtxt(fref)
n1 = d1[:,0]
n2 = d2[:,0]
nref = dref[:,0]
rk1 = d1[:,1]
rk2 = d2[:,1]
rkref = dref[:,1]
wkb1 = d1[:,3]
wkb2 = d2[:,3]
plt.figure(figsize=(10,5))
plt.title('Evolution of a single primordial perturbation $\mathcal{R}_k$ of Fourier wavenumber $k$')
plt.plot(nref,rkref,color='black',lw=0.7)
plt.plot(n1[wkb1==1],(rk1[wkb1==1]),'.',color='C0',label='WKB step',ms=9.)
plt.plot(n1[wkb1==0],(rk1[wkb1==0]),'.',color='C1',label='RK step',ms=9.)
plt.ticklabel_format(axis='y',style='sci',scilimits=(-2,2))
plt.xlabel('N')
plt.ylabel('$\Re{\{ \mathcal{R}_k \}} $')
plt.legend()
plt.show()
l=1.0
m=0.5
def Ei(n):
"""
Energy eigenvalues (if analytic solution available)
"""
return np.sqrt(2.0)*(n-0.5)
def V(t):
"""
Potential well
"""
return t**2 + l*t**4
def w(t,E):
"""
Frequency term in the Schrodinger equation
"""
return np.sqrt(2*m*(complex(E)-V(t)));
def f(E):
"""
Function to minimize wrt E to give the energy eigenvalues
"""
# Boundaries of integration
tl = -((E)**(0.25))-2.0
tr = -tl
tm = 0.5
# Grid of w, g
t = np.linspace(tl.real,tr.real,30000)
ws = np.log(w(t,E))
g = np.zeros(t.shape)
sol_l = pyoscode.solve(t,ws,g,tl,tm,0,1e-3,logw=True,rtol=1e-5)
sol_r = pyoscode.solve(t,ws,g,tr,tm,0,1e-3,h=-1,logw=True,rtol=1e-5)
psi_l = sol_l["sol"][-1]
psi_r = sol_r["sol"][-1]
dpsi_l = sol_l["dsol"][-1]
dpsi_r = sol_r["dsol"][-1]
try:
return abs(dpsi_l/psi_l - dpsi_r/psi_r)
except ZeroDivisionError:
return 1000.0
def plot(ns,Es):
plt.figure(figsize=(10,5))
t_v = np.linspace(-6,6,500)
plt.plot(t_v,V(t_v),color='black',label='V(x)')
for j,n,E in zip(range(len(ns)),ns,Es):
# Boundaries of integration
tl = -((E)**(0.25))-1.0
tr = -tl
tm = 0.0
# Grid of w, g
t = np.linspace(tl.real,tr.real,30000)
ws = np.log(w(t,E))
g = np.zeros(t.shape)
sol_l = pyoscode.solve(t,ws,g,tl,tr/2.,0,1e-3,logw=True,rtol=1e-5)
x_l = sol_l['sol']
ts_l = sol_l['t']
types_l = sol_l['types']
for i,typ in enumerate(types_l):
if typ==1 and 0 not in types_l[i:]:
firstwkb = i
break;
t_eval = np.linspace(ts_l[firstwkb],tm,2000)
sol = pyoscode.solve(t,ws,g,tl,tr,0,1e-3,logw=True,rtol=1e-5,t_eval=t_eval)
x_eval = sol['x_eval']
x_l = x_l[:firstwkb]
ts_l = ts_l[:firstwkb]
Ts_l = np.concatenate((np.array(ts_l),t_eval))
Xs_l = np.concatenate((np.array(x_l),x_eval))
maxx = max(np.real(Xs_l))
Xs_l = Xs_l/maxx*4*np.sqrt(E)
plt.plot(Ts_l,Xs_l+E,color='C{}'.format(j),label='$\Psi_n(x)$, n={}, $E_n$={:.4f}'.format(n,E))
plt.plot(-1*Ts_l,Xs_l+E,color='C{}'.format(j))
plt.xlabel('x')
plt.legend(loc='lower left')
plt.show()
bounds = [ (416.5,417.5),(1035,1037),(21930,21940),(471100,471110)]
ress = []
for bound in bounds:
res = minimize_scalar(f,bounds=bound,method='bounded')
print("Eigenenergy found: {}".format(res.x))
ress.append(res.x)
plot([50,100],ress[:2])
Eigenenergy found: 416.60834468082714 Eigenenergy found: 1035.3995965810213 Eigenenergy found: 21939.980473898962 Eigenenergy found: 471103.8196601125
scipy.integrate.solve_ivp(fun, t_span, y0, ...)
fun
: defines differential equationt_span
: integration rangey0
: initial conditionsimport numpy as np
import scipy.special as sp
import pyoscode
# Set up differential equation: x'' + tx = 0.
ts = np.linspace(0,40.0,5000)
ws = np.sqrt(ts)
gs = np.zeros_like(ts)
# Initial conditions
ti = 1.0
tf = 40.0
x0 = sp.airy(-ti)[0]
dx0 = -sp.airy(-ti)[1]
solution = pyoscode.solve(ts, ws, gs, ti, tf, x0, dx0)
solution.keys()
dict_keys(['sol', 'dsol', 't', 'types', 'x_eval'])
# Compute analytic solution for comparison
analytic = np.array([sp.airy(-t)[0] for t in ts])
plt.figure(figsize=(10,5))
plt.plot(ts,analytic,label='analytic solution',color='black',lw=0.7)
plt.plot(solution['t'],solution['sol'],'.',color='C1',label='pyoscode',ms=9.0)
plt.show()
# We'll get the solution at these points@
t_eval = np.linspace(15,35,600)
dense_solution = pyoscode.solve(ts, ws, gs, ti, tf, x0, dx0, t_eval=t_eval)
plt.figure(figsize=(10,5))
plt.plot(ts,analytic,label='analytic solution',color='black',lw=0.7)
plt.plot(dense_solution['t'],dense_solution['sol'],'.',color='C1',label='pyoscode internal step',ms=9.0)
plt.plot(t_eval, dense_solution['x_eval'], color='C1', label='pyoscode dense output')
plt.xlabel('x')
plt.ylabel('y(x)')
plt.legend()
plt.show()