cms.diffeq2nd_planet_verlet のソースコード

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

"""
  Planet simulator: Solve simulataneous second order diffrential equations by Verlet 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_Verlet.csv'
x0, vx0 = 0.0, 1.0
y0, vy0 = 2.0, 0.0
dt = 0.01
nt = 501
iprint_interval = 20
nprint_planets = 3

#===================
# functions
#===================


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[i] * 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
    return fxi, fyi, fzi




def initialize(planets, M, x, y, z, vx, vy, vz, fx, fy, fz):
    for i in range(0, len(planets)):
        M.append(planets[i]['Mass'])
        x.append([planets[i]['Revolution Radius']])
        y.append([0.0])
        z.append([0.0])
        vx.append([0.0])
        vy.append([planets[i]['Revolution Velocity']])
        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 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


#===================
# main routine
#===================


def main():
    global nt
    global dt

    print("Planet simulator: Solve simulataneous second order diffrential equations by Verlet method")

    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)

# 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):
        name = planets[i]['Name']
        print("  {:^12}  {:^12}".format(labellist[i], labellist[i+1]), end = '')
    print("")


"""
# Solve the 1st data by Heun method
    kvx0, kvy0 = force(0.0, x0, y0)
    kvx1, kvy1 = force(0.0+dt, x0+kvy0, y0+kvy0)
    vx1 = vx0 + dt * (kvx0 + kvx1) / 2.0
    vy1 = vy0 + dt * (kvy0 + kvy1) / 2.0
    kx0 = dt * vx0
    kx1 = dt * vx1
    ky0 = dt * vy0
    ky1 = dt * vy1
    x1  = x0 + (kx0 + kx1) / 2.0
    y1  = y0 + (ky0 + ky1) / 2.0
    xexact, yexact = xysolution(0.0)
    print("t={:5.2f}  {:12.6f}  {:12.6f}  {:12.6f}  {:12.6f}".format(0.0, x0, xexact, y0, yexact))
    fout.writerow([0.0, x0, xexact, y0, yexact])
    
    for i in range(1, nt):
        t1 = i * dt
        fx1, fy1 = force(t1, x1, y1)
        x2 = 2.0 * x1 - x0 + dt * dt * fx1
        y2 = 2.0 * y1 - y0 + dt * dt * fy1
        vx1 = (x2 - x0) / 2.0 / dt
        vy1 = (y2 - y0) / 2.0 / dt
        xexact, yexact = xysolution(t1)
        if i == 1 or i % iprint_interval == 0:
            print("t={:5.2f}  {:12.6f}  {:12.6f}  {:12.6f}  {:12.6f}".format(t1, x1, xexact, y1, yexact))
        fout.writerow([t1, x1, xexact, y1, yexact])
        x0  = x1
        vx0 = vx1
        y0  = y1
        vy0 = vy1
        x1  = x2
        y1  = y2

    f.close()
"""

if __name__ == '__main__':
    main()