cms.diffeq2nd_2d_verlet のソースコード

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

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


#===================
# parameters
#===================
outfile = 'diffeq2nd_2D_Verlet.csv'
x0, vx0 = 0.0, 1.0
y0, vy0 = 2.0, 0.0
dt = 0.01
nt = 501
iprint_interval = 20

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


def main(x0, y0, vx0, vy0, dt, nt):
    print("Solve simulataneous second order diffrential equations by Verlet method")

    print("Write to [{}]".format(outfile))

# open outfile to write a csv file
    f = open(outfile, 'w')
    fout = csv.writer(f, lineterminator='\n')
    fout.writerow([
        't', 'x(cal)', 'x(exact)', 'y(cal)', 'y(exact)'
        ])
    print("{:^5}  {:^12}  {:^12}  {:^12}  {:^12}".format('t', 'x(cal)', 'x(exact)', 'y(cal)', 'y(exact)'))

# 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(x0, y0, vx0, vy0, dt, nt)