cms.diffeq2nd_planet_euler のソースコード

import csv
import numpy as np
from math import exp, sqrt, sin, cos, pi

"""
  Planet simulator: Solve simulataneous second order diffrential equations by Euler method
"""

# d2x/dt2 = force(t, x, y), d2y/dt2 = force(t, x, y)
# define function to be integrated
[ドキュメント] def force(t, x, y): return -x, -y
# solution: x(0.0) = 0.0, y(0.0) = 2.0, vx(0.0) = 1.0, vy(0.0) = 0.0
[ドキュメント] def xysolution(t): return sin(t), 2.0 * cos(t)
#=================== # constants #=================== G = 6.67259e-11 #Nm2/kg2 DayToSecond = 60 * 60 * 24 #s SecondToDay = 1.0 / DayToSecond AstronomicalUnit = 1.49597870e11 #m AU = AstronomicalUnit G1 = G * DayToSecond * DayToSecond / AU / AU / AU #=================== # parameters #=================== dbfile = 'planet_db.csv' outfile = 'diffeq2nd_Planet_Euler.csv' x0, vx0 = 0.0, 1.0 y0, vy0 = 2.0, 0.0 z0, vz0 = 0.0, 0.0 dt = 1.0 nt = 1000 iprint_interval = 100 nprint_planets = 4 CSVOutputInterval = 10 CriteriaR2ForIdenticalAtom = 0.001 DifferentialEquationSolver = "Verlet" #=================== # functions #===================
[ドキュメント] def readdb(dbfile): planets = [] f = open(dbfile, "r"); reader = csv.DictReader(f) for row in reader: planets.append(row) keys = list(planets[0].keys()) for d in planets: for key in keys: if key != 'Name': d[key] = float(d[key]) return planets
[ドキュメント] def sum(array): sum = 0.0 for e in array: sum += e return sum
[ドキュメント] def total_momentum(it, M, vx, vy, vz): Px, Py, Pz = 0.0, 0.0, 0.0 for i in range(0, len(M)): Px += M[i] * vx[i][it]; Py += M[i] * vy[i][it]; Pz += M[i] * vz[i][it]; return Px, Py, Pz
# Normalize total momentum to zero
[ドキュメント] def normalize_momentum(it, M, x, y, z, vx, vy, vz, fx, fy, fz): Mtot = sum(M) Px, Py, Pz = total_momentum(it, M, vx, vy, vz); for ip in range(0, len(M)): vx[ip][it] -= Px / Mtot; vy[ip][it] -= Py / Mtot; vz[ip][it] -= Pz / Mtot; return Px, Py, Pz
[ドキュメント] def initialize(planets, M, x, y, z, vx, vy, vz, fx, fy, fz): global AU global DayToSecond for i in range(0, len(planets)): M.append(planets[i]['Mass']) x.append([planets[i]['Revolution Radius'] / AU]) y.append([0.0]) z.append([0.0]) vx.append([0.0]) vy.append([planets[i]['Revolution Velocity'] * DayToSecond / AU]) vz.append([0.0]) for i in range(0, len(planets)): fxi, fyi, fzi = Fi(0, i, M, x, y, z) fx.append([fxi]) fy.append([fyi]) fz.append([fzi])
[ドキュメント] def Fij(istep, i, j, M, x, y, z): dx = x[j][istep] - x[i][istep] dy = y[j][istep] - y[i][istep] dz = z[j][istep] - z[i][istep] r2 = dx*dx + dy*dy + dz*dz r = sqrt(r2) g = G1 * M[j] f = g / r2 fx = f * dx / r fy = f * dy / r fz = f * dz / r return fx, fy, fz
[ドキュメント] def Fi(istep, i, M, x, y, z): fxi = 0.0 fyi = 0.0 fzi = 0.0 for j in range(0, len(M)): if i == j: continue fxj, fyj, fzj = Fij(istep, i, j, M, x, y, z) fxi += fxj fyi += fyj fzi += fzj # print("f={}, {}, {}".format(fxi, fyi, fzi)) return fxi, fyi, fzi
#=================== # main routine #===================
[ドキュメント] def main(): global nt global dt print("Planet simulator: Solve simulataneous second order diffrential equations by Euler method") print("G = {} Nm2/kg2".format(G)) print("AU = {:e} m".format(AU)) print("G1 = {}".format(G1)) print("") print("Planets:") planets = readdb(dbfile) keys = list(planets[0].keys()) for d in planets: print(" ", d['Name']) for key in keys: if key != 'Name': print(" {}: {}".format(key, d[key])) print("") M = [] x = [] y = [] z = [] vx = [] vy = [] vz = [] fx = [] fy = [] fz = [] initialize(planets, M, x, y, z, vx, vy, vz, fx, fy, fz) Px, Py, Pz = normalize_momentum(0, M, x, y, z, vx, vy, vz, fx, fy, fz) print("") # make label list for display / csv output labellist = ['t'] for i in range(0, len(planets)): labellist.append("x({})".format(planets[i]['Name'])) labellist.append("y({})".format(planets[i]['Name'])) # open outfile to write a csv file print("Write to [{}]".format(outfile)) f = open(outfile, 'w') fout = csv.writer(f, lineterminator='\n') fout.writerow(labellist) print("{:^5}".format('t'), end = '') for i in range(1, nprint_planets*2, 2): print(" {:^12} {:^12}".format(labellist[i], labellist[i+1]), end = '') print("") # Solve the 1st data by Heun method datalist = [0.0] print("{:^5}".format(0.0), end = '') for i in range(0, len(planets)): fx0, fy0, fz0 = Fi(0, i, M, x, y, z) vx1 = vx[i][0] + dt * fx0 vy1 = vy[i][0] + dt * fy0 vz1 = vz[i][0] + dt * fz0 x1 = x[i][0] + dt * vx[i][0] y1 = y[i][0] + dt * vy[i][0] z1 = z[i][0] + dt * vz[i][0] datalist.append(x[i][0]) datalist.append(y[i][0]) x[i].append(x1) y[i].append(y1) z[i].append(z1) vx[i].append(vx1) vy[i].append(vy1) vz[i].append(vz1) for i in range(1, nprint_planets*2, 2): print(" {:^12.4f} {:^12.4f}".format(x[i][0], y[i][0]), end = '') print("") fout.writerow(datalist) for it in range(1, nt): t = it * dt # print("it={} t={}".format(it, t)) datalist = [t] if it % iprint_interval == 0: print("{:^5}".format(t), end = '') for i in range(0, len(planets)): fx0, fy0, fz0 = Fi(it, i, M, x, y, z) vx1 = vx[i][it] + dt * fx0 vy1 = vy[i][it] + dt * fy0 vz1 = vz[i][it] + dt * fz0 x1 = x[i][it] + dt * vx[i][it] y1 = y[i][it] + dt * vy[i][it] z1 = z[i][it] + dt * vz[i][it] datalist.append(x[i][it]) datalist.append(y[i][it]) x[i].append(x1) y[i].append(y1) z[i].append(z1) vx[i].append(vx1) vy[i].append(vy1) vz[i].append(vz1) if it % iprint_interval == 0: for i in range(1, nprint_planets*2, 2): print(" {:^12.4g} {:^12.4g}".format(x[i][it], y[i][it]), end = '') print("") fout.writerow(datalist) f.close()
if __name__ == '__main__': main()