integration_code/solutions/simulation.py

"""W05: velocity Verlet for x''=-omega^2*x, unit mass, fixed time step."""
import math


def verlet(x=1.,v=0.,omega=1.,dt=.1,steps=100):
    if not all(math.isfinite(z) for z in (x,v,omega,dt)) or omega<=0 or dt<=0:
        raise ValueError("finite state, positive frequency and step required")
    if type(steps) is not int or not 0<=steps<=100000:
        raise ValueError("bounded nonnegative step count required")
    energy0=.5*(v*v+omega*omega*x*x)
    if not math.isfinite(energy0):
        raise ValueError("initial energy is not finite; rescale the problem")
    rows=[(0.,x,v,energy0)]
    for k in range(1,steps+1):
        half=v-.5*dt*omega*omega*x
        new_x=x+dt*half
        new_v=half-.5*dt*omega*omega*new_x
        x,v=new_x,new_v
        energy=.5*(v*v+omega*omega*x*x)
        if not all(math.isfinite(z) for z in (x,v,energy,k*dt)):
            raise ValueError("nonfinite trajectory; check step and scale")
        rows.append((k*dt,x,v,energy))
    return rows


def exact(t,x0=1.,v0=0.,omega=1.):
    if not all(math.isfinite(z) for z in (t,x0,v0,omega)) or omega<=0:
        raise ValueError("finite inputs and positive frequency required")
    phase=omega*t
    if not math.isfinite(phase):
        raise ValueError("nonfinite phase; rescale the problem")
    result=(x0*math.cos(phase)+(v0/omega)*math.sin(phase),
            -omega*x0*math.sin(phase)+v0*math.cos(phase))
    if not all(math.isfinite(z) for z in result):
        raise ValueError("nonfinite analytic state; rescale the problem")
    return result


if __name__=='__main__':
    for dt in (.1,.05,.025):
        t,x,v,e=verlet(dt=dt,steps=round(1/dt))[-1]
        target=exact(t)
        print(dt,math.hypot(x-target[0],v-target[1]),e)