def example_solve():
##solves a simple ode using all three methods
##this is an example of how to use solve_ode
E1 = ode.solve_ode(
example_func, 0, 7, t_end=5
) ###Sets up a solver with the default step size h = 1e-3 from 0 to 5
###Solve with all three ode solvers
E1.Forward_Euler()
print(E1.fe_t[-1], E1.fe_y[-1])
E1.Heun()
print(E1.heun_t[-1], E1.heun_y[-1])
E1.RK4()
print(E1.rk4_t[-1], E1.rk4_y[-1])
###We can also solve this problem specifying a number of steps instead of a step size
E2 = ode.solve_ode(example_func, 0, 7, t_end=5,
n=1e4) ##Uses 10000 steps across the interval
E2.Forward_Euler()
E2.Heun()
E2.RK4()
def test_rk4(self):
Test = ode.solve_ode(example_func, 0, 7, t_end=5, n=2e4)
Test.RK4()
self.assertAlmostEqual(round(Test.rk4_y[-1], 3),
round(example_sltn(5), 3))
def problem_5():
'''
produces a mass vs. radius plot for our neutron star
runs with central densities from 1e4 to 1e6 in g / cm ^ 3
'''
###neutron stars
h = 1e4
rc = 1e14 * u.g / (u.cm ** 3)
Pc = NSPressure(rc)
Rad = []
M_total = []
while rc.value < 1e17:
P_c = NSPressure(rc).value ###Central Pressure in Ba
Neutron_Star = ode.solve_ode(NS_func , 1 , [P_c , 0] , h)
rend = 5e9
nsr , nsy = Neutron_Star.Forward_Euler(rend)
for i in range(len(nsr)):
if nsy[i][0] < 0:
R = nsr[i]
M = nsy[i][1]
Rad.append((R * u.cm).to(u.km).value)
M_total.append( (M * u.g).to(u.M_sun).value)
break
rc *= 1.2
plt.plot(M_total , Rad)
plt.xlabel("Mass (M_sun)")
plt.ylabel("Radius (km)")
plt.show()
def problem_2():
'''
solves problem 2 from the HW.
will show a plot comparing my solvers to scipy's odeint
'''
b = 0.25
c = 5.0
y0 = [np.pi - 0.1, 0.0]
t = np.linspace(0, 10, 101)
print (pend(y0 , t[0] , b , c))
sol = odeint(pend, y0, t, args=(b, c))
###Now we use my codes
A = ode.solve_ode(my_pend(b , c) , 0 , y0 , (t[1] - t[0]))
t_heun , y_heun = A.Heun(t[-1])
t_fe , y_fe = A.Forward_Euler(t[-1])
t_rk4 , y_rk4 = A.RK4(t[-1])
print (sol[0])
pt = []
ptheta = []
for i in range(len(t)):
pt.append(t[i])
ptheta.append(sol[i][0])
fe = []
for i in y_fe:
fe.append(i[0])
heun = []
for i in y_heun:
heun.append(i[0])
rk4 = []
for i in y_rk4:
rk4.append(i[0])
f , ax = plt.subplots(3 , 1 , sharex = True)
ax[0].plot(pt , ptheta , label = "scipy")
ax[0].plot(t_fe , fe , label = "Forward Euler")
#ax[0].set_xlabel("time")
ax[0].set_ylabel("theta")
ax[0].legend()
ax[1].plot(pt , ptheta , label = "scipy")
ax[1].plot(t_heun , heun , label = "Heun")
#ax[1].set_xlabel("time")
ax[1].set_ylabel("theta")
ax[1].legend()
ax[2].plot(pt , ptheta , label = "scipy")
ax[2].plot(t_rk4, rk4 , label = "RK4")
ax[2].set_xlabel("time")
ax[2].set_ylabel("theta")
ax[2].legend()
plt.show()
def evolve_rho(rho, t, dt, eta, dn, branch, theta0,
theta1=-pi / 2, theta2=pi / 2, pumpp=1., miss=0.0):
m, n = rho.shape
n = n // 2 - 1
gammas = gamma_pump(n, eta, theta0, branch, miss)
omegas = norm_omega_raman(n, eta, dn, theta1, theta2)
def f(t, rho):
return (calc_ddt_pump(rho, gammas) * pumpp +
calc_ddt_raman(rho, omegas))
return solve_ode(0, rho, f, t, dt)
def evolve_rho(rho,
t,
dt,
eta,
dn,
branch,
theta0,
theta1=-pi / 2,
theta2=pi / 2,
pumpp=1.,
miss=0.0):
m, n = rho.shape
n = n // 2 - 1
gammas = gamma_pump(n, eta, theta0, branch, miss)
omegas = norm_omega_raman(n, eta, dn, theta1, theta2)
def f(t, rho):
return (calc_ddt_pump(rho, gammas) * pumpp +
calc_ddt_raman(rho, omegas))
return solve_ode(0, rho, f, t, dt)
def problem_4():
'''
produces a mass vs. radius plot for our white dwarf
runs with central densities from 1e4 to 1e6 in g / cm ^ 3
'''
rc = 1e4
M_total = []
R = []
h = 5e6
while rc <= 1e6:
rho_c = rc * (u.g / (u.cm ** 3))
## we will solve in terms of P and M
P_c = WDPressure(rho_c).value ###Central Pressure in Ba
White_Dwarf = ode.solve_ode(WD_func , 1 , [P_c , 0] , h)
rearth = 6.378e+8 ##cm
r_end = 5 * rearth
wdr , wdy = White_Dwarf.RK4(r_end)
for i in range(len(wdr)):
if wdy[i][0] < 0:
R.append((wdr[i] * u.cm).to(u.R_sun).value)
M_total.append(((wdy[i][1] * u.g).to(u.M_sun).value))
break
rc *= 1.5
plt.plot(M_total , R)
plt.xlabel("Mass (M_sun)")
plt.ylabel("Radius (Solar Radii)")
plt.show()
def stiff(L): ''' solves our stiff ode given a value of lambda, L returns the maximum error of all three of our methors runs from 0 to 5 with aa step size of 1e-2 ''' h = 1e-2 stiff_solve = ode.solve_ode(stiff_ode(L) , 0 , [0] , h) tend = 5 t_h , y_h = stiff_solve.Heun(tend) t_fe , y_fe = stiff_solve.Forward_Euler(tend) t_rk4 , y_rk4 = stiff_solve.RK4(tend) rt = [] ry = [] for i in range(len(t_h)): rt.append(t_h[i]) ry.append(stiff_sltn(t_h[i] , L)) heun = [] rk4 = [] fe = [] t = 0 for i in y_h: heun.append(abs(i - stiff_sltn(t , L))) t += h t = 0 for i in y_fe: fe.append(abs(i - stiff_sltn(t , L))) t += h t = 0 for i in y_rk4: rk4.append(abs(i - stiff_sltn(t , L))) t += h return max(fe) , max(heun) , max(rk4)
def main_pump():
n = 100
gammas = gamma_pump(n, .8, pi / 2, 0.4, 0.02)
omegas = [omega_raman(n, .8, dn) for dn in range(1, 22)]
# omegas = omega_raman(n, .8, 10)
def f(t, rho):
dn = int(21 - t)
return calc_ddt_pump(rho, gammas) + calc_ddt_raman(rho, omegas[dn - 1])
ps0 = exp(-arange(n + 1, dtype=complex128) / 10) * (1 - exp(1 / 10))
rho0 = diag(r_[ps0, zeros(n + 1, dtype=complex128)])
print(sum(rho0))
print(calc_total_n(rho0))
ts, rhos = solve_ode(0, rho0, f, 20, .01)
print(abs(diag(rhos[0])))
print(abs(diag(rhos[-1])))
plot(abs(diag(rhos[0])))
plot(abs(diag(rhos[-1])))
figure()
imshow(abs(rhos[0]))
figure()
imshow(abs(rhos[-1]))
show()
def main_raman_sb_cooling():
# from pylab import plot, show, imshow, figure, colorbar, xlabel, ylabel
# from pylab import legend, title, savefig, close, grid
n = 100
nstart = 30
gammas = gamma_pump(n, .8, pi / 2, 0.4, 0.2)
omegas = [omega_raman(n, .8, dn, 0) for dn in range(1, n + 1)]
def get_f(dn):
def f(t, rho):
return (calc_ddt_pump(rho, gammas) +
calc_ddt_raman(rho, omegas[dn - 1]) * 10)
return f
ps0 = (exp(-arange(n + 1, dtype=complex128) / nstart) *
(1 - exp(-1 / nstart)))
rho0 = diag(r_[ps0, zeros(n + 1, dtype=complex128)])
dns = []
dnrange = range(1, 40)
for i in range(40):
print("\nstart iteration: %d" % i)
number = abs(sum(diag(rho0)))
ntotal = calc_total_n(rho0)
dnmax = 0
vmax = number**2 / ntotal
print("atom number: %f" % number)
print("total n: %f" % ntotal)
print("v: %f" % vmax)
for dn in dnrange:
print("dn: %d" % dn)
ts, rhos = solve_ode(0, rho0, get_f(dn), .45, 0.015)
print("dn: %d" % dn)
number = abs(sum(diag(rhos[-1])))
ntotal = calc_total_n(rhos[-1])
v = number**2 / ntotal
print("atom number: %f" % number)
print("total n: %f" % ntotal)
print("v: %f" % v)
if v > vmax:
print("use new dn: %d, v = %f" % (dn, v))
dnmax = dn
vmax = v
new_rho0 = rhos[-1]
# plot(abs(diag(rhos[0])), label='before')
# plot(abs(diag(rhos[-1])), label='after')
# legend()
# figure()
# plot(abs(diag(rhos[-1])) - abs(diag(rhos[0])))
# figure()
# imshow(abs(rhos[0]))
# figure()
# imshow(abs(rhos[-1]))
# show()
print('')
if dnmax == 0:
print("cooling stopped, abort")
break
dns.append(dnmax)
dnrange = range(max(dnmax - 15, 1), dnmax + 16)
rho0 = new_rho0
print('\n')
print("atom number: %f\n" % sum(diag(rho0)))
print("total n: %f\n" % calc_total_n(rho0))
print(dns)
def main_ode():
# from pylab import plot, show, imshow, figure, colorbar, xlabel, ylabel
# from pylab import legend, title, savefig, close, grid
ts, ys = solve_ode(0, [0, 1], lambda t, y: array([y[1], -y[0]]), 10000, .01)
def test_fe(self):
Test = ode.solve_ode(example_func, 0, 7, t_end=5, n=2e5)
Test.Forward_Euler()
self.assertAlmostEqual(round(Test.fe_y[-1], 3),
round(example_sltn(5), 3))