simple harmonic oscillator sim.py

#-------------------------------------------------------------------------------
# Purpose:
# Author:      Pat
#-------------------------------------------------------------------------------
#!/usr/bin/env python

import math
import numpy as np
import pylab as pl
from pprint import pprint, pformat

def euler(x,f,dt):
    return x + f(x) * dt

k = 1.0
m = 1.0

def SHO(state):
    v_prev = state[1]
    x_prev = state[0]
    return np.array([v_prev, -k*x_prev/m])

dt = 0.00065  # seconds
x_0 = 1  # m
v_0 = 0  # m/s

# Timing
#  Run for five complete cycles
t_end = 5 * (2 * math.pi * k/m)

# Let us call the state of the system 'state'
states = [np.array([x_0, v_0])]

elapsed_time = 0
while elapsed_time < t_end:
    states.append(euler(states[-1], SHO, dt)) # This is what does the hard work
    elapsed_time += dt

# The trouble with the above, is that what you really want is all the positions
# Turn the list to an array:
state = np.array(states)

# Now array slices get all positions and velocities:
positions = state[:,0]
velocities = state[:,1]
sim_times = pl.frange(0, elapsed_time, dt)

## Analytical
an_times = np.linspace(0, elapsed_time, 100)
positions_an = x_0 * np.cos(math.sqrt(k/m) * an_times)
# Note: could also use `np.vectorize(math.cos)()`

# Find the energy
energies = 0.5*k*positions**2 + 0.5*m*velocities**2
an_energy = 0.5  # at t=0, v=0 so E = 0.5*k*x^2=0.5 J

##sim_t_end = sim_times[ sim_times >= t_end][0]  # take earliest
sim_value_at_approx_end = positions[ sim_times >= t_end][0]  # take first
an_value_at_end = positions_an[-1]
percent_diff = 100.0 * abs(sim_value_at_approx_end - an_value_at_end) / an_value_at_end


## Plotting
pl.subplot(2,1,1)
pl.title(r'Simple Harmonic Osc. ($dt=%.5f$, $diff=%.1f\%%$)' % (dt, percent_diff))
sim_pt_size = 3
pl.scatter(sim_times, positions, sim_pt_size, color='black', label='simulation')
pl.plot(sim_times, positions, color='black', alpha=0.5)
pl.plot(an_times, positions_an, color='red', label='analytical')
pl.legend(loc='lower left')
# Move axes to be on the plot
ax = pl.gca()  # get current axes
#  Remove right and top axes
ax.spines['right'].set_color('none')
ax.spines['top'].set_color('none')
ax.spines['bottom'].set_position(('data',0))
ax.xaxis.set_ticks_position('bottom')
ax.yaxis.set_major_locator(pl.MaxNLocator(nbins=3))
pl.xlabel('Time (s)', bbox=dict(facecolor='white', alpha=0.5))  # needs to be after messing with axes or won't show up
pl.ylabel('Position (m)')
pl.xlim(-1, sim_times.max()+1)

pl.subplot(2,1,2)
pl.ylabel(r'Energy (J)')
pl.plot(sim_times, energies, color='black', label='simulation')
pl.plot([an_times.min(), an_times.max()], [an_energy, an_energy], color='black', alpha=0.5, label='analytical')
pl.annotate('analytical',
         xy=(30, an_energy), xycoords='data',
         xytext=(-60, -50), textcoords='offset points', fontsize=14,
         arrowprops=dict(arrowstyle="->", connectionstyle="arc3,rad=.5"))

ax = pl.gca()
ax.yaxis.set_major_locator(pl.MaxNLocator(nbins=3))
pl.ylim(energies.min()*0.9, energies.max()*1.1)
pl.xlim(-1, sim_times.max()+1)

pl.show()