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))