Skip to content

Instantly share code, notes, and snippets.

@pckujawa
Created February 6, 2013 01:01
Show Gist options
  • Save pckujawa/4719284 to your computer and use it in GitHub Desktop.
Save pckujawa/4719284 to your computer and use it in GitHub Desktop.
#-------------------------------------------------------------------------------
# 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()
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment