finite difference for harmonic oscillator

fd_harmonic.py

import matplotlib.pyplot as plt


def hom1(f, x0, dx0, delta, n):
    sol = [x0]
    x1, dx = x0 + dx0 * delta, dx0
    for _ in range(n-1):
        x0, x1, dx = x1, f(x0) * delta + x0, 1.0 * (x1 - x0) / delta        
        sol.append(x0)
    return sol

def hom2(g, x0, dx0, delta, n):
    sol = [x0]
    x1, dx = x0 + dx0 * delta, dx0
    for _ in range(n-1):
        x0, x1 = x1, g(x1, dx) * delta ** 2 + 2 * x1 - x0
        dx = 1.0 * (x1 - x0) / delta
        sol.append(x0)
    return sol

def harmonic_potential(x, dx):
    omega = 2.
    return -1. * omega ** 2 * x


def solve_harmonic():
    pi = 3.14159265
    period = pi
    omega = 2 * pi / period
    x0, dx0 = 0, omega
    delta = 0.01
    tmax = 6 * pi
    n = int(tmax / delta)
    harmonic_sol = hom2(harmonic_potential, x0, dx0, delta, n)
    ts = [delta * i for i in range(n)]
    plt.plot(ts, harmonic_sol)
    plt.show()

if __name__ == "__main__":
    solve_harmonic()