def keplerdirect(x, y, dx): nbodies = y.size // 4 # per body, we have four variables par = globalvar.get_odepar() npar = par.size gnewton = par[0] masses = par[1:npar] dydx = np.zeros(4 * nbodies) indx = 2 * np.arange(nbodies) indy = 2 * np.arange(nbodies) + 1 indvx = 2 * np.arange(nbodies) + 2 * nbodies indvy = 2 * np.arange(nbodies) + 2 * nbodies + 1 dydx[indx] = y[indvx] # x'=v dydx[indy] = y[indvy] for k in range(nbodies): gravx = 0.0 gravy = 0.0 for j in range(nbodies): if (k != j): dx = y[indx[k]] - y[indx[j]] dy = y[indy[k]] - y[indy[j]] R3 = np.power(dx * dx + dy * dy, 1.5) gravx = gravx - gnewton * masses[j] * dx / R3 gravy = gravy - gnewton * masses[j] * dy / R3 dydx[indvx[k]] = gravx dydx[indvy[k]] = gravy return dydx
def lunarlanding(x,y,dx): par = globalvar.get_odepar() dydx = np.zeros(4) thrmax = par[0] mode = par[1] g = par[2] Vnozz = par[3] mship = par[4] z = y[0] vz = y[1] fuel = y[2] thrfrac = y[3] # limit throttle value. This has no effect on dthrdt, # but is necessary for mode == 1, and for RHS of ODEs. thrfrac = np.max(np.array([0.0,np.min(np.array([1.0,thrfrac]))])) # make sure fuel can't go negative if (fuel <= 0.0): fuel = 0.0 thrfrac = 0.0 mtot = mship+fuel # calculate RHS for all ODEs, bc we'll need them if mode == 1 dzdt = vz dvzdt = thrmax*thrfrac/mtot - g dfdt = -thrmax*thrfrac/Vnozz if (int(mode)==0): # constant throttle dthrdt = 0.0 elif (int(mode)==1): # PID controller. This is a controller with beta = 0, since drdt == 0 # Note: dthrdt contains change of accelerations. kp = 0.10 ki = 0.60 kd = 0.30 vref = -0.2 dthrdt = (-kp*dvzdt+ki*(vref-vz)+kd*thrmax*thrfrac*dfdt/(mtot*mtot))/(1.0+kd*thrmax/mtot) if (dthrdt < 0.0): # limit adjustment of valve to prevent overshoots dthrdt = np.max(np.array([dthrdt,-thrfrac/dx])) else: dthrdt = np.min(np.array([dthrdt,(1.0-thrfrac)/dx])) # need to check for fuel again to make sure throttle change is 0. if (fuel <= 0.0): dthrdt = 0.0 elif (int(mode)==2): # off/on throttle. Note that we need to calculate the throttle change! # this just via guessing if (x > 11.5): dthrdt = (1.0-thrfrac)/dx else: dthrdt = 0.0 if (fuel <= 0.0): dthrdt = 0.0 dydx[0] = dzdt dydx[1] = dvzdt dydx[2] = dfdt dydx[3] = dthrdt return dydx
def ode_check(x, y, it): n = x.size par = globalvar.get_odepar() z = y[0, :] # altitude above ground vz = y[1, :] # vertical velocity f = y[2, :] # fuel mass thr = y[3, :] # throttle accG = np.zeros(n) # acceleration in units of g for k in range(n): accG[k] = ( (dydx.lunarlanding(x[k], y[:, k], n / (x[n - 1] - x[0])))[1]) / 9.81 # acceleration ftsz = 10 plt.figure(num=1, figsize=(8, 8), dpi=100, facecolor='white') plt.subplot(321) plt.plot(x, z, linestyle='-', color='black', linewidth=1.0) plt.xlabel('t [s]', fontsize=ftsz) plt.ylabel('z(t) [m]', fontsize=ftsz) util.rescaleplot(x, z, plt, 0.05) plt.tick_params(labelsize=ftsz) plt.subplot(322) plt.plot(x, vz, linestyle='-', color='black', linewidth=1.0) plt.xlabel('t [s]', fontsize=ftsz) plt.ylabel('v$_z$ [m s$^{-1}$]', fontsize=ftsz) util.rescaleplot(x, vz, plt, 0.05) plt.tick_params(labelsize=ftsz) plt.subplot(323) plt.plot(x, f, linestyle='-', color='black', linewidth=1.0) plt.xlabel('t [s]', fontsize=ftsz) plt.ylabel('fuel [kg]', fontsize=ftsz) util.rescaleplot(x, f, plt, 0.05) plt.tick_params(labelsize=ftsz) plt.subplot(324) plt.plot(x, thr, linestyle='-', color='black', linewidth=1.0) plt.xlabel('t [s]', fontsize=ftsz) plt.ylabel('throttle', fontsize=ftsz) util.rescaleplot(x, thr, plt, 0.05) plt.tick_params(labelsize=ftsz) plt.subplot(325) plt.plot(x, accG, linestyle='-', color='black', linewidth=1.0) plt.xlabel('t [s]', fontsize=ftsz) plt.ylabel('acc/G', fontsize=ftsz) util.rescaleplot(x, accG, plt, 0.05) plt.tick_params(labelsize=ftsz) plt.tight_layout() plt.show()
def kepler(x, y, dx): par = globalvar.get_odepar() dydx = np.zeros(4) R = np.sqrt(y[0] * y[0] + y[1] * y[1]) R3 = np.power(R, 3) dydx[0] = y[2] # x' = v_x dydx[1] = y[3] # y' = v_y dydx[2] = -par[0] * y[0] / R3 dydx[3] = -par[0] * y[1] / R3 return dydx
def h2formation(x,y): par = globalvar.get_odepar() dfdx = np.zeros(3) dfdy = np.zeros((3,3)) dfdy[0,0] = -4.0*par[0]*y[0]-par[2] dfdy[0,1] = 2.0*par[1] dfdy[0,2] = 2.0*par[3]*y[2] dfdy[1,0] = 2.0*par[0]*y[0] dfdy[1,1] = -par[1] dfdy[1,2] = 0.0 dfdy[2,0] = par[2] dfdy[2,1] = 0.0 dfdy[2,2] = -2.0*par[3]*y[2] return dfdy,dfdx
def h2formation(x, y, dx): nspec = y.size par = globalvar.get_odepar() A = y[0] # not necessary, but more readable B = y[1] C = y[2] #print 'A,B,C,sum = ',A,B,C,A+B+C # This is the stoichiometry matrix S S = np.array([[-2.0, 2.0, -1.0, 1.0], [1.0, -1.0, 0.0, 0.0], [0.0, 0.0, 1.0, -1.0]]) # This is the reaction rate vector. # Note that the stoichiometry would give par[3]*C, not par[3]*C*C. # However, we need to account for the electrons, whose density is that of C. v = np.array([par[0] * A * A, par[1] * B, par[2] * A, par[3] * C * C]) Myr = 1e6 * 3.65e2 * 3.6e3 * 2.4e1 dydx = S.dot(v) return dydx
def keplerdirect_symp1(x, y, dx): nbodies = y.size / 4 # per body, we have four variables par = globalvar.get_odepar() npar = par.size gnewton = par[0] masses = par[1:npar] dydx = np.zeros(4 * nbodies) pHpq = np.zeros(2 * nbodies) pHpp = np.zeros(2 * nbodies) indx = 2 * np.arange(nbodies) indy = 2 * np.arange(nbodies) + 1 indvx = 2 * np.arange(nbodies) + 2 * nbodies indvy = 2 * np.arange(nbodies) + 2 * nbodies + 1 px = y[indvx] * masses # this is more cumbersome than necessary, py = y[indvy] * masses # but for consistency with derivation. qx = y[indx] qy = y[indy] for i in range(nbodies): for j in range(0, i): ddx = qx[i] - qx[j] ddy = qy[i] - qy[j] R3 = np.power(ddx * ddx + ddy * ddy, 1.5) pHpq[indx[i]] = pHpq[indx[i]] + masses[j] * ddx / R3 pHpq[indy[i]] = pHpq[indy[i]] + masses[j] * ddy / R3 for j in range(i + 1, nbodies): ddx = qx[i] - qx[j] ddy = qy[i] - qy[j] R3 = np.power(ddx * ddx + ddy * ddy, 1.5) pHpq[indx[i]] = pHpq[indx[i]] + masses[j] * ddx / R3 pHpq[indy[i]] = pHpq[indy[i]] + masses[j] * ddy / R3 pHpq[indx] = pHpq[indx] * gnewton * masses pHpq[indy] = pHpq[indy] * gnewton * masses pHpp[indx] = (px - dx * pHpq[indx]) / masses pHpp[indy] = (py - dx * pHpq[indy]) / masses dydx[indx] = pHpp[indx] dydx[indy] = pHpp[indy] dydx[indvx] = -pHpq[indx] / masses dydx[indvy] = -pHpq[indy] / masses #for i in range(nbodies): # print('t=%10.2e mass=%10.2e: x,y=%13.5e %13.5e vx,vy=%13.5e %13.5e dx,dy=%13.5e %13.5e dvx,dvy= %13.5e %13.5e' # % (x,masses[i],qx[i],qy[i],px[i]/masses[i],py[i]/masses[i],dx*dydx[indx[i]],dx*dydx[indy[i]],dx*dydx[indvx[i]]/masses[i],dx*dydx[indvy[i]]/masses[i])) return dydx
def keplerdirect_symp2(x, y, dx): nbodies = y.size // 4 # per body, we have four variables par = globalvar.get_odepar() npar = par.size gnewton = par[0] masses = par[1:npar] dydx = np.zeros(4 * nbodies) pHpq = np.zeros(2 * nbodies) pHpq2 = np.zeros(2 * nbodies) pHpp = np.zeros(2 * nbodies) indx = 2 * np.arange(nbodies) indy = 2 * np.arange(nbodies) + 1 indvx = 2 * np.arange(nbodies) + 2 * nbodies indvy = 2 * np.arange(nbodies) + 2 * nbodies + 1 px = y[indvx] * masses # this is more cumbersome than necessary, py = y[indvy] * masses # but for consistency with derivation. qx = y[indx] qy = y[indy] # first step: constructing p(n+1/2) for i in range(nbodies): for j in range(0, i): ddx = qx[i] - qx[j] ddy = qy[i] - qy[j] R3 = np.power(ddx * ddx + ddy * ddy, 1.5) pHpq[indx[i]] = pHpq[indx[i]] + masses[j] * ddx / R3 pHpq[indy[i]] = pHpq[indy[i]] + masses[j] * ddy / R3 for j in range(i + 1, nbodies): ddx = qx[i] - qx[j] ddy = qy[i] - qy[j] R3 = np.power(ddx * ddx + ddy * ddy, 1.5) pHpq[indx[i]] = pHpq[indx[i]] + masses[j] * ddx / R3 pHpq[indy[i]] = pHpq[indy[i]] + masses[j] * ddy / R3 pHpq[indx] = pHpq[indx] * gnewton * masses pHpq[indy] = pHpq[indy] * gnewton * masses px2 = px - 0.5 * dx * pHpq[indx] py2 = py - 0.5 * dx * pHpq[indy] pHpp[indx] = px2 / masses pHpp[indy] = py2 / masses qx2 = qx + dx * pHpp[indx] qy2 = qy + dx * pHpp[indy] for i in range(nbodies): for j in range(0, i): ddx = qx2[i] - qx2[j] ddy = qy2[i] - qy2[j] R3 = np.power(ddx * ddx + ddy * ddy, 1.5) pHpq2[indx[i]] = pHpq2[indx[i]] + masses[j] * ddx / R3 pHpq2[indy[i]] = pHpq2[indy[i]] + masses[j] * ddy / R3 for j in range(i + 1, nbodies): ddx = qx2[i] - qx2[j] ddy = qy2[i] - qy2[j] R3 = np.power(ddx * ddx + ddy * ddy, 1.5) pHpq2[indx[i]] = pHpq2[indx[i]] + masses[j] * ddx / R3 pHpq2[indy[i]] = pHpq2[indy[i]] + masses[j] * ddy / R3 pHpq2[indx] = pHpq2[indx] * gnewton * masses pHpq2[indy] = pHpq2[indy] * gnewton * masses dydx[indx] = pHpp[indx] dydx[indy] = pHpp[indy] dydx[indvx] = -0.5 * (pHpq[indx] + pHpq2[indx]) / masses dydx[indvy] = -0.5 * (pHpq[indy] + pHpq2[indy]) / masses return dydx
def lunarlanding(x, y, dx): par = globalvar.get_odepar() dydx = np.zeros(4) thrmax = par[0] mode = par[1] g = par[2] Vnozz = par[3] mship = par[4] z = y[0] vz = y[1] fuel = y[2] thrfrac = y[3] # limit throttle value. This has no effect on dthrdt, # but is necessary for mode == 1, and for RHS of ODEs. thrfrac = np.max(np.array([0.0, np.min(np.array([1.0, thrfrac]))])) # make sure fuel can't go negative if (fuel <= 0.0): fuel = 0.0 thrfrac = 0.0 mtot = mship + fuel # calculate RHS for all ODEs, bc we'll need them if mode == 1 dzdt = vz dvzdt = thrmax * thrfrac / mtot - g dfdt = -thrmax * thrfrac / Vnozz if (int(mode) == 0): # constant throttle dthrdt = 0.0 elif (int(mode) == 1): # PID controller: constant reference velocity vref = par[5] kp = par[6] ki = par[7] kd = par[8] vintdt = z - 500.0 denom = (1.0 + kd * thrmax / mtot) numer = kp * (vref - vz) + ki * (vref * x - vintdt) + kd * g ddenom = -kd * thrmax * dfdt / mtot**2 dnumer = -kp * dvzdt + ki * (vref - vz) dthrdt = (denom * dnumer - numer * ddenom) / denom**2 if (dthrdt < 0.0): # limit adjustment of valve to prevent overshoots dthrdt = np.max(np.array([dthrdt, -thrfrac / dx])) else: dthrdt = np.min(np.array([dthrdt, (1.0 - thrfrac) / dx])) # need to check for fuel again to make sure throttle change is 0. if (fuel <= 0.0): dthrdt = 0.0 elif (int(mode) == 2): # PID controller with variable velocity kp = par[6] ki = par[7] kd = par[8] vref = -5.0 + (1.0 - z / 500.0) * 4.9 vintdt = z - 500.0 denom = (1.0 + kd * thrmax / mtot) numer = kp * (vref - vz) + ki * (vref * x - vintdt) + kd * (-4.9 / 500.0 * vz + g) ddenom = -kd * thrmax * dfdt / mtot**2 dnumer = -kp * dvzdt + ki * (vref - vz) dthrdt = (denom * dnumer - numer * ddenom) / denom**2 if (dthrdt < 0.0): # limit adjustment of valve to prevent overshoots dthrdt = np.max(np.array([dthrdt, -thrfrac / dx])) else: dthrdt = np.min(np.array([dthrdt, (1.0 - thrfrac) / dx])) # need to check for fuel again to make sure throttle change is 0. if (fuel <= 0.0): dthrdt = 0.0 else: raise Exception("[dydx.lunarlanding]: invalid imode %i" % (i)) dydx[0] = dzdt dydx[1] = dvzdt dydx[2] = dfdt dydx[3] = dthrdt return dydx
def ode_check(x, y, it, iprob): if (iprob == "lunarlander"): # lunar lander problem n = x.size par = globalvar.get_odepar() z = y[0, :] # altitude above ground vz = y[1, :] # vertical velocity f = y[2, :] # fuel mass thr = y[3, :] # throttle accG = np.zeros(n) # acceleration in units of g for k in range(n): accG[k] = ((dydx.lunarlanding( x[k], y[:, k], n / (x[n - 1] - x[0])))[1]) / 9.81 # acceleration ftsz = 10 plt.figure(num=1, figsize=(8, 8), dpi=100, facecolor='white') plt.subplot(321) plt.plot(x, z, linestyle='-', color='black', linewidth=1.0) plt.xlabel('t [s]', fontsize=ftsz) plt.ylabel('z(t) [m]', fontsize=ftsz) util.rescaleplot(x, z, plt, 0.05) plt.tick_params(labelsize=ftsz) plt.subplot(322) plt.plot(x, vz, linestyle='-', color='black', linewidth=1.0) plt.xlabel('t [s]', fontsize=ftsz) plt.ylabel('v$_z$ [m s$^{-1}$]', fontsize=ftsz) util.rescaleplot(x, vz, plt, 0.05) plt.tick_params(labelsize=ftsz) plt.subplot(323) plt.plot(x, f, linestyle='-', color='black', linewidth=1.0) plt.xlabel('t [s]', fontsize=ftsz) plt.ylabel('fuel [kg]', fontsize=ftsz) util.rescaleplot(x, f, plt, 0.05) plt.tick_params(labelsize=ftsz) plt.subplot(324) plt.plot(x, thr, linestyle='-', color='black', linewidth=1.0) plt.xlabel('t [s]', fontsize=ftsz) plt.ylabel('throttle', fontsize=ftsz) util.rescaleplot(x, thr, plt, 0.05) plt.tick_params(labelsize=ftsz) plt.subplot(325) plt.plot(x, accG, linestyle='-', color='black', linewidth=1.0) plt.xlabel('t [s]', fontsize=ftsz) plt.ylabel('acc/G', fontsize=ftsz) util.rescaleplot(x, accG, plt, 0.05) plt.tick_params(labelsize=ftsz) plt.tight_layout() plt.show() elif (iprob == "kepler"): # for the direct Kepler problem, we check for energy and angular momentum conservation, # and for the center-of-mass position and velocity color = [ 'black', 'green', 'cyan', 'blue', 'red', 'black', 'black', 'black', 'black' ] n = x.size par = globalvar.get_odepar() npar = par.size nbodies = par.size - 1 gnewton = par[0] masses = par[1:npar] Egrav = np.zeros(n) indx = 2 * np.arange(nbodies) indy = 2 * np.arange(nbodies) + 1 indvx = 2 * np.arange(nbodies) + 2 * nbodies indvy = 2 * np.arange(nbodies) + 2 * nbodies + 1 E = np.zeros(n) # total energy Lphi = np.zeros(n) # angular momentum R = np.sqrt( np.power(y[indx[0], :] - y[indx[1], :], 2) + np.power(y[indy[0], :] - y[indy[1], :], 2)) Rs = np.zeros(n) # center of mass position vs = np.zeros(n) # center of mass velocity for k in range(n): E[k] = 0.5 * np.sum( masses * (np.power(y[indvx, k], 2) + np.power(y[indvy, k], 2))) Lphi[k] = np.sum( masses * (y[indx, k] * y[indvy, k] - y[indy, k] * y[indvx, k])) Rsx = np.sum(masses * y[indx, k]) / np.sum(masses) Rsy = np.sum(masses * y[indy, k]) / np.sum(masses) vsx = np.sum(masses * y[indvx, k]) / np.sum(masses) vsy = np.sum(masses * y[indvy, k]) / np.sum(masses) Rs[k] = np.sqrt(Rsx * Rsx + Rsy * Rsy) vs[k] = np.sqrt(vsx * vsx + vsy * vsy) for j in range(nbodies): for i in range( j): # preventing double summation. Still O(N^2) though. dx = y[indx[j], :] - y[indx[i], :] dy = y[indy[j], :] - y[indy[i], :] Rt = np.sqrt(dx * dx + dy * dy) Egrav = Egrav - gnewton * masses[i] * masses[j] / Rt E = E + Egrav E = E / E[0] Lphi = Lphi / Lphi[0] for k in range(n): print( 'k=%7i t=%13.5e E/E0=%20.12e L/L0=%20.12e Rs=%10.2e vs=%10.2e' % (k, x[k], E[k], Lphi[k], Rs[k], vs[k])) Eplot = E - 1.0 Lplot = Lphi - 1.0 # try Poincare cuts p1tot = np.zeros(n) p2tot = np.zeros(n) q1tot = np.zeros(n) q2tot = np.zeros(n) for k in range(n): p1tot[k] = np.sum(masses * y[indvx, k]) / np.sum(masses) p2tot[k] = np.sum(masses * y[indvy, k]) / np.sum(masses) q1tot[k] = np.sum(y[indx, k]) q2tot[k] = np.sum(y[indy, k]) thq = 1e-1 thp = 1e-1 print( '[ode_test]: min/max(q1), min/max(p1): %13.5e %13.5e %13.5e %13.5e' % (np.min(q1tot), np.max(q1tot), np.min(p1tot), np.max(p1tot))) tarr = (np.abs(p1tot) <= thp) & (np.abs(q1tot) <= thq) if (len(tarr) == 0): raise Exception('[ode_test]: no indices found at these thresholds') indp = np.where(tarr)[0] nind = indp.size print('[ode_test]: found %i elements for Poincare section\n' % (i)) # now plot everything # (1) the orbits xmin = np.min(y[indx, :]) xmax = np.max(y[indx, :]) ymin = np.min(y[indy, :]) ymax = np.max(y[indy, :]) qmin = np.min(q2tot[indp]) qmax = np.max(q2tot[indp]) pmin = np.min(p2tot[indp]) pmax = np.max(p2tot[indp]) plt.figure(num=1, figsize=(8, 8), dpi=100, facecolor='white') plt.subplot(111) plt.xlim(1.05 * xmin, 1.05 * xmax) plt.ylim(1.05 * ymin, 1.05 * ymax) for k in range(nbodies): plt.plot(y[indx[k], :], y[indy[k], :], color=color[k], linewidth=1.0, linestyle='-') plt.axes().set_aspect('equal') plt.xlabel('x [AU]') plt.ylabel('y [AU]') # (2) the checks (total energy and angular momentum) plt.figure(num=2, dpi=100, facecolor='white') plt.subplot(311) plt.plot(x, Eplot, linestyle='-', color='black', linewidth=1.0) plt.xlabel('t [yr]') plt.ylabel('$\Delta$E/E') #plt.legend() plt.subplot(312) plt.plot(x, Lplot, linestyle='-', color='black', linewidth=1.0) plt.xlabel('t [yr]') plt.ylabel('$\Delta$L/L') #plt.legend() plt.subplot(313) plt.plot(q2tot[indp], p2tot[indp], '.', color='black') plt.xlabel('q$_2$') plt.ylabel('p$_2$') plt.tight_layout() plt.show() elif (iprob == "h2formation"): # H2 formation color = ['black', 'blue', 'red'] label = [u"H", u"H\u2082", u"H\u207A"] tlabel = ['k1', 'k2', 'k3', 'k4'] Myr = 1e6 * 3.6e3 * 3.65e2 * 2.4e1 nspec = (y.shape)[0] n = x.size par = globalvar.get_odepar() Ds = par[4] ys = get_h2eq(np.log10(Ds)) print( 'Equilibrium fractions: D=%13.5e A=%13.5e B=%13.5e C=%13.5e total=%13.5e' % (Ds, ys[0], ys[1], ys[2], np.sum(ys))) ys = ys * Ds th2 = np.zeros((4, n)) for i in range(n): th2[:, i] = get_h2times(y[:, i]) th2[2, :] = 0.0 # don't show CR ionization (really slow) t = x / Myr # convert time in seconds to Myr tmin = np.min(t) tmax = np.max(t) nmin = np.min(np.array([np.nanmin(y), np.min(ys)])) nmax = np.max(np.array([np.nanmax(y), np.max(ys)])) ly = np.log10(y) lnmin = np.min(np.array([np.nanmin(ly), np.min(np.log10(ys))])) lnmax = np.max(np.array([np.nanmax(ly), np.max(np.log10(ys))])) # Note that the first call to subplot is slightly different, because we need the # axes object to reset the y-labels. ftsz = 10 fig = plt.figure(num=1, figsize=(8, 8), dpi=100, facecolor='white') ax = fig.add_subplot(321) for i in range(nspec): plt.plot(t, y[i, :], linestyle='-', color=color[i], linewidth=1.0, label=label[i]) plt.plot([tmin, tmax], [ys[i], ys[i]], linestyle='--', color=color[i], linewidth=1.0) plt.xlabel('t [Myr]', fontsize=ftsz) plt.ylabel('n$_\mathrm{\mathsf{H}}$', fontsize=ftsz) plt.legend(fontsize=10) util.rescaleplot(t, y, plt, 0.05) ylabels = ax.yaxis.get_ticklocs() ylabels1 = ylabels[1:ylabels.size - 1] ylabelstext = ['%4.2f' % (lb / np.max(ylabels1)) for lb in ylabels1] plt.yticks(ylabels1, ylabelstext) plt.tick_params(labelsize=ftsz) plt.subplot(322) for i in range(nspec): plt.plot(t, ly[i, :], linestyle='-', color=color[i], linewidth=1.0, label=label[i]) plt.plot(np.array([tmin, tmax]), np.log10(np.array([ys[i], ys[i]])), linestyle='--', color=color[i], linewidth=1.0) plt.xlabel('t [Myr]', fontsize=ftsz) plt.ylabel('log n [cm$^{-3}$]', fontsize=ftsz) util.rescaleplot(t, ly, plt, 0.05) plt.tick_params(labelsize=ftsz) plt.subplot(323) for i in range(nspec): plt.plot(t, y[i, :] / ys[i], linestyle='-', color=color[i], linewidth=1.0, label=label[i]) plt.xlabel('t [Myr]', fontsize=ftsz) plt.ylabel('n/n$_{eq}$', fontsize=ftsz) util.rescaleplot(t, np.array([0, 2]), plt, 0.05) plt.tick_params(labelsize=ftsz) plt.subplot(324) for i in range(nspec): plt.plot(t, ly[i, :] - np.log10(ys[i]), linestyle='-', color=color[i], linewidth=1.0, label=label[i]) plt.xlabel('t [Myr]', fontsize=ftsz) plt.ylabel('log n/n$_{eq}$', fontsize=ftsz) util.rescaleplot(t, np.array([-1, 1]), plt, 0.05) plt.tick_params(labelsize=ftsz) plt.subplot(325) wplot = np.array([0, 1, 3]) for i in range(wplot.size): plt.plot(t, np.log10(th2[wplot[i], :]), linestyle='-', linewidth=1.0, label=tlabel[wplot[i]]) plt.xlabel('t [Myr]', fontsize=ftsz) plt.ylabel('log reaction time [Myr]', fontsize=ftsz) plt.legend(fontsize=ftsz) util.rescaleplot(t, np.log10(th2[wplot, :]), plt, 0.05) plt.tick_params(labelsize=ftsz) plt.subplot(326) plt.plot(t[1:t.size], np.log10(it[1:t.size]), linestyle='-', color='black', linewidth=1.0, label='iterations') plt.xlabel('t [Myr]', fontsize=ftsz) plt.ylabel('log #', fontsize=ftsz) plt.legend(fontsize=ftsz) util.rescaleplot(t, np.log10(it[1:t.size]), plt, 0.05) plt.tick_params(labelsize=ftsz) plt.tight_layout() plt.show() elif (iprob == "stiff1"): # Your code should show three plots in one window: # (1) the dependent variables y[0,:] and y[1,:] against x # (2) the residual y[0,:]-u and y[1,:]-v against x, see homework sheet # (3) the number of iterations against x. xmin = np.min(x) xmax = np.max(x) ymin = np.min(y) ymax = np.max(y) ftsz = 10 fig = plt.figure(num=1, figsize=(6, 8), dpi=100, facecolor='white') u = 2.0 * np.exp(-x) - np.exp(-1e3 * x) v = -np.exp(-x) + np.exp(-1e3 * x) diff = np.zeros((2, x.size)) diff[0, :] = y[0] - u diff[1, :] = y[1] - v plt.subplot(311) plt.plot(x, y[0, :], linestyle='-', color='blue', linewidth=1.0, label='u') plt.plot(x, u, linestyle='--', color='blue', linewidth=1.0, label='u [analytic]') plt.plot(x, y[1, :], linestyle='-', color='red', linewidth=1.0, label='v') plt.plot(x, v, linestyle='--', color='red', linewidth=1.0, label='v [analytic]') plt.xlabel('x', fontsize=ftsz) plt.ylabel('y', fontsize=ftsz) plt.legend(fontsize=ftsz) util.rescaleplot(x, y, plt, 0.05) plt.tick_params(labelsize=ftsz) plt.subplot(312) plt.plot(x, diff[0, :], linestyle='-', color='blue', linewidth=1.0, label='residual u') plt.plot(x, diff[1, :], linestyle='-', color='red', linewidth=1.0, label='residual v') plt.xlabel('x', fontsize=ftsz) plt.ylabel('y-y[analytic]', fontsize=ftsz) plt.legend(fontsize=ftsz) util.rescaleplot(x, diff, plt, 0.05) plt.tick_params(labelsize=ftsz) plt.subplot(313) plt.plot(x[1:x.size], np.log10(it[1:x.size]), linestyle='-', color='black', linewidth=1.0, label='iterations') plt.xlabel('x', fontsize=ftsz) plt.ylabel('log #', fontsize=ftsz) plt.legend(fontsize=ftsz) util.rescaleplot(x[1:x.size], np.log10(it[1:x.size]), plt, 0.05) plt.tick_params(labelsize=ftsz) plt.tight_layout() plt.show() elif (iprob == "stiff2"): xmin = np.min(x) xmax = np.max(x) ymin = np.min(y) ymax = np.max(y) ftsz = 10 fig = plt.figure(num=1, figsize=(6, 8), dpi=100, facecolor='white') plt.subplot(211) for i in range(3): plt.plot(x, y[i, :], linestyle='-', linewidth=1.0, label='y[%i]' % (i)) plt.xlabel('x', fontsize=ftsz) plt.ylabel('y', fontsize=ftsz) plt.legend(fontsize=ftsz) util.rescaleplot(x, y, plt, 0.05) plt.tick_params(labelsize=ftsz) plt.subplot(212) plt.plot(x[1:x.size], np.log10(it[1:x.size]), linestyle='-', color='black', linewidth=1.0, label='iterations') plt.xlabel('x', fontsize=ftsz) plt.ylabel('log #', fontsize=ftsz) plt.legend(fontsize=ftsz) util.rescaleplot(x[1:x.size], np.log10(it[1:x.size]), plt, 0.05) plt.tick_params(labelsize=ftsz) plt.tight_layout() plt.show() elif ((iprob == "exponential") or (iprob == "gaussian") or (iprob == "tanh")): if (iprob == "gaussian"): f = dydx.gaussian(x, y, 1.0) sol = np.zeros(x.size) for i in range(x.size): sol[i] = 0.5 * math.erf(x[i] / np.sqrt(2.0)) elif (iprob == "exponential"): f = dydx.exp(x, y, 1.0) sol = np.exp(x) elif (iprob == "tanh"): f = dydx.tanh(x, y, 1.0) sol = np.log(np.cosh(x)) res = y[0, :] - sol ftsz = 10 plt.figure(num=1, figsize=(8, 8), dpi=100, facecolor='white') plt.subplot(221) plt.plot(x, f, linestyle='-', color='black', linewidth=1.0) plt.xlabel('x', fontsize=ftsz) plt.ylabel("y'(x)", fontsize=ftsz) util.rescaleplot(x, f, plt, 0.05) plt.tick_params(labelsize=ftsz) plt.subplot(222) plt.plot(x, y[0, :], linestyle='-', color='black', linewidth=1.0, label='y(x)') plt.plot(x, sol, linestyle='-', color='red', linewidth=1.0, label='analytic') plt.xlabel('x', fontsize=ftsz) plt.ylabel("y(x)", fontsize=ftsz) plt.legend(fontsize=ftsz) util.rescaleplot(x, y, plt, 0.05) plt.tick_params(labelsize=ftsz) plt.subplot(223) plt.plot(x, res, linestyle='-', color='black', linewidth=1.0, label='residual') plt.xlabel('x', fontsize=ftsz) plt.ylabel("y(x)", fontsize=ftsz) plt.legend(fontsize=ftsz) util.rescaleplot(x, res, plt, 0.05) plt.tick_params(labelsize=ftsz) plt.subplot(224) plt.plot(x[1:x.size], it[1:x.size], linestyle='-', color='black', linewidth=1.0, label='iterations') plt.xlabel('x', fontsize=ftsz) plt.ylabel("#", fontsize=ftsz) plt.legend(fontsize=ftsz) util.rescaleplot(x, it, plt, 0.05) plt.tick_params(labelsize=ftsz) plt.tight_layout() plt.show() else: raise Exception('[ode_check]: invalid iprob %i\n' % (iprob))
def ode_check(x, y, it, iprob): if (iprob == "h2formation"): # H2 formation color = ['black', 'blue', 'red'] label = [u"H", u"H\u2082", u"H\u207A"] tlabel = ['k1', 'k2', 'k3', 'k4'] Myr = 1e6 * 3.6e3 * 3.65e2 * 2.4e1 nspec = (y.shape)[0] n = x.size par = globalvar.get_odepar() Ds = par[4] ys = get_h2eq(np.log10(Ds)) print( 'Equilibrium fractions: D=%13.5e A=%13.5e B=%13.5e C=%13.5e total=%13.5e' % (Ds, ys[0], ys[1], ys[2], np.sum(ys))) ys = ys * Ds th2 = np.zeros((4, n)) for i in range(n): th2[:, i] = get_h2times(y[:, i]) th2[2, :] = 0.0 # don't show CR ionization (really slow) t = x / Myr # convert time in seconds to Myr tmin = np.min(t) tmax = np.max(t) nmin = np.min(np.array([np.nanmin(y), np.min(ys)])) nmax = np.max(np.array([np.nanmax(y), np.max(ys)])) ly = np.log10(y) lnmin = np.min(np.array([np.nanmin(ly), np.min(np.log10(ys))])) lnmax = np.max(np.array([np.nanmax(ly), np.max(np.log10(ys))])) # Note that the first call to subplot is slightly different, because we need the # axes object to reset the y-labels. ftsz = 10 fig = plt.figure(num=1, figsize=(8, 8), dpi=100, facecolor='white') ax = fig.add_subplot(321) for i in range(nspec): plt.plot(t, y[i, :], linestyle='-', color=color[i], linewidth=1.0, label=label[i]) plt.plot([tmin, tmax], [ys[i], ys[i]], linestyle='--', color=color[i], linewidth=1.0) plt.xlabel('t [Myr]', fontsize=ftsz) plt.ylabel('n$_\mathrm{\mathsf{H}}$', fontsize=ftsz) plt.legend(fontsize=10) util.rescaleplot(t, y, plt, 0.05) ylabels = ax.yaxis.get_ticklocs() ylabels1 = ylabels[1:ylabels.size - 1] ylabelstext = ['%4.2f' % (lb / np.max(ylabels1)) for lb in ylabels1] plt.yticks(ylabels1, ylabelstext) plt.tick_params(labelsize=ftsz) plt.subplot(322) for i in range(nspec): plt.plot(t, ly[i, :], linestyle='-', color=color[i], linewidth=1.0, label=label[i]) plt.plot(np.array([tmin, tmax]), np.log10(np.array([ys[i], ys[i]])), linestyle='--', color=color[i], linewidth=1.0) plt.xlabel('t [Myr]', fontsize=ftsz) plt.ylabel('log n [cm$^{-3}$]', fontsize=ftsz) util.rescaleplot(t, ly, plt, 0.05) plt.tick_params(labelsize=ftsz) plt.subplot(323) for i in range(nspec): plt.plot(t, y[i, :] / ys[i], linestyle='-', color=color[i], linewidth=1.0, label=label[i]) plt.xlabel('t [Myr]', fontsize=ftsz) plt.ylabel('n/n$_{eq}$', fontsize=ftsz) util.rescaleplot(t, np.array([0, 2]), plt, 0.05) plt.tick_params(labelsize=ftsz) plt.subplot(324) for i in range(nspec): plt.plot(t, ly[i, :] - np.log10(ys[i]), linestyle='-', color=color[i], linewidth=1.0, label=label[i]) plt.xlabel('t [Myr]', fontsize=ftsz) plt.ylabel('log n/n$_{eq}$', fontsize=ftsz) util.rescaleplot(t, np.array([-1, 1]), plt, 0.05) plt.tick_params(labelsize=ftsz) plt.subplot(325) wplot = np.array([0, 1, 3]) for i in range(wplot.size): plt.plot(t, np.log10(th2[wplot[i], :]), linestyle='-', linewidth=1.0, label=tlabel[wplot[i]]) plt.xlabel('t [Myr]', fontsize=ftsz) plt.ylabel('log reaction time [Myr]', fontsize=ftsz) plt.legend(fontsize=ftsz) util.rescaleplot(t, np.log10(th2[wplot, :]), plt, 0.05) plt.tick_params(labelsize=ftsz) plt.subplot(326) plt.plot(t[1:t.size], np.log10(it[1:t.size]), linestyle='-', color='black', linewidth=1.0, label='iterations') plt.xlabel('t [Myr]', fontsize=ftsz) plt.ylabel('log #', fontsize=ftsz) plt.legend(fontsize=ftsz) util.rescaleplot(t, np.log10(it[1:t.size]), plt, 0.05) plt.tick_params(labelsize=ftsz) print("[ode_test]: integration took %6i iterations" % (np.sum(it))) plt.tight_layout() plt.show() elif (iprob == "doubleexp"): xmin = np.min(x) xmax = np.max(x) ymin = np.min(y) ymax = np.max(y) ftsz = 10 fig = plt.figure(num=1, figsize=(6, 8), dpi=100, facecolor='white') u = 2.0 * np.exp(-x) - np.exp(-1e3 * x) v = -np.exp(-x) + np.exp(-1e3 * x) diff = np.zeros((2, x.size)) diff[0, :] = y[0] - u diff[1, :] = y[1] - v plt.subplot(311) plt.plot(x, y[0, :], linestyle='-', color='blue', linewidth=1.0, label='u') plt.plot(x, u, linestyle='--', color='blue', linewidth=1.0, label='u [analytic]') plt.plot(x, y[1, :], linestyle='-', color='red', linewidth=1.0, label='v') plt.plot(x, v, linestyle='--', color='red', linewidth=1.0, label='v [analytic]') plt.xlabel('x', fontsize=ftsz) plt.ylabel('y', fontsize=ftsz) plt.legend(fontsize=ftsz) util.rescaleplot(x, y, plt, 0.05) plt.tick_params(labelsize=ftsz) plt.subplot(312) plt.plot(x, diff[0, :], linestyle='-', color='blue', linewidth=1.0, label='residual u') plt.plot(x, diff[1, :], linestyle='-', color='red', linewidth=1.0, label='residual v') plt.xlabel('x', fontsize=ftsz) plt.ylabel('y-y[analytic]', fontsize=ftsz) plt.legend(fontsize=ftsz) util.rescaleplot(x, diff, plt, 0.05) plt.tick_params(labelsize=ftsz) plt.subplot(313) plt.plot(x[1:x.size], np.log10(it[1:x.size]), linestyle='-', color='black', linewidth=1.0, label='iterations') plt.xlabel('x', fontsize=ftsz) plt.ylabel('log #', fontsize=ftsz) plt.legend(fontsize=ftsz) util.rescaleplot(x[1:x.size], np.log10(it[1:x.size]), plt, 0.05) plt.tick_params(labelsize=ftsz) print("[ode_test]: integration took %6i iterations" % (np.sum(it))) plt.tight_layout() plt.show() elif (iprob == "enrightpryce"): xmin = np.min(x) xmax = np.max(x) ymin = np.min(y) ymax = np.max(y) ftsz = 10 fig = plt.figure(num=1, figsize=(6, 8), dpi=100, facecolor='white') plt.subplot(211) for i in range(3): plt.plot(x, y[i, :], linestyle='-', linewidth=1.0, label='y[%i]' % (i)) plt.xlabel('x', fontsize=ftsz) plt.ylabel('y', fontsize=ftsz) plt.legend(fontsize=ftsz) util.rescaleplot(x, y, plt, 0.05) plt.tick_params(labelsize=ftsz) plt.subplot(212) plt.plot(x[1:x.size], np.log10(it[1:x.size]), linestyle='-', color='black', linewidth=1.0, label='iterations') plt.xlabel('x', fontsize=ftsz) plt.ylabel('log #', fontsize=ftsz) plt.legend(fontsize=ftsz) util.rescaleplot(x[1:x.size], np.log10(it[1:x.size]), plt, 0.05) plt.tick_params(labelsize=ftsz) print("[ode_test]: integration took %6i iterations" % (np.sum(it))) plt.tight_layout() plt.show() else: raise Exception('[ode_check]: invalid iprob %i\n' % (iprob))
def doubleexp(x, y, dx): par = globalvar.get_odepar() dydx = np.zeros(2) dydx[0] = 998.0 * y[0] + 1998.0 * y[1] dydx[1] = -999.0 * y[0] - 1999.0 * y[1] return dydx
def ode_check(x,y,it): # for the direct Kepler problem, we check for energy and angular momentum conservation, # and for the center-of-mass position and velocity color = ['black','green','cyan','blue','red','black','black','black','black'] n = x.size par = globalvar.get_odepar() npar = par.size nbodies = par.size-1 gnewton = par[0] masses = par[1:npar] Egrav = np.zeros(n) indx = 2*np.arange(nbodies) indy = 2*np.arange(nbodies)+1 indvx = 2*np.arange(nbodies)+2*nbodies indvy = 2*np.arange(nbodies)+2*nbodies+1 E = np.zeros(n) # total energy Lphi = np.zeros(n) # angular momentum R = np.sqrt(np.power(y[indx[0],:]-y[indx[1],:],2)+np.power(y[indy[0],:]-y[indy[1],:],2)) Rs = np.zeros(n) # center of mass position vs = np.zeros(n) # center of mass velocity for k in range(n): E[k] = 0.5*np.sum(masses*(np.power(y[indvx,k],2)+np.power(y[indvy,k],2))) Lphi[k] = np.sum(masses*(y[indx,k]*y[indvy,k]-y[indy,k]*y[indvx,k])) Rsx = np.sum(masses*y[indx,k])/np.sum(masses) Rsy = np.sum(masses*y[indy,k])/np.sum(masses) vsx = np.sum(masses*y[indvx,k])/np.sum(masses) vsy = np.sum(masses*y[indvy,k])/np.sum(masses) Rs[k] = np.sqrt(Rsx*Rsx+Rsy*Rsy) vs[k] = np.sqrt(vsx*vsx+vsy*vsy) for j in range(nbodies): for i in range(j): # preventing double summation. Still O(N^2) though. dx = y[indx[j],:]-y[indx[i],:] dy = y[indy[j],:]-y[indy[i],:] Rt = np.sqrt(dx*dx+dy*dy) Egrav = Egrav - gnewton*masses[i]*masses[j]/Rt E = E + Egrav E = E/E[0] Lphi = Lphi/Lphi[0] for k in range(n): print('k=%7i t=%13.5e E/E0=%20.12e L/L0=%20.12e Rs=%10.2e vs=%10.2e' % (k,x[k],E[k],Lphi[k],Rs[k],vs[k])) Eplot = E-1.0 Lplot = Lphi-1.0 # now plot everything # (1) the orbits xmin = np.min(y[indx,:]) xmax = np.max(y[indx,:]) ymin = np.min(y[indy,:]) ymax = np.max(y[indy,:]) #plt.figure(num=1,figsize=(8,8),dpi=100,facecolor='white') fig1,ax1 = plt.subplots(1,1) ax1.set_xlim(1.05*xmin,1.05*xmax) ax1.set_ylim(1.05*ymin,1.05*ymax) for k in range(nbodies): ax1.plot(y[indx[k],:],y[indy[k],:],color=color[k],linewidth=1.0,linestyle='-') ax1.set_aspect('equal') ax1.set_xlabel('x [AU]') ax1.set_ylabel('y [AU]') # (2) the checks (total energy and angular momentum) #plt.figure(num=2,figsize=(8,8),dpi=100,facecolor='white') fig2,(ax2,ax3) = plt.subplots(2,1) ax2.plot(x,Eplot,linestyle='-',color='black',linewidth=1.0) ax2.set_xlabel('t [yr]') ax2.set_ylabel('$\Delta$E/E') #plt.legend() #ax3.subplot(212) ax3.plot(x,Lplot,linestyle='-',color='black',linewidth=1.0) ax3.set_xlabel('t [yr]') ax3.set_ylabel('$\Delta$L/L') plt.tight_layout() plt.show()
def exponential(x, y, dx): par = globalvar.get_odepar() dydx = np.exp(y[0] - 2.0 * x) + 2.0 return dydx