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
