falling body 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

def fall(state):
    v_prev = state[1]
    a = -9.8
    return np.array([v_prev, a])

##def falling_body():
dt = 0.0195  # seconds
g = 9.8  # m/s^2
y_0 = 10  # m
v_0 = 0

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

elapsed_time = 0
while states[-1][0] > 0:  # above ground
    states.append(euler(states[-1], fall, 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)
ys_an = -0.5 * g * an_times**2 + y_0

# When does the object hit the ground?
t_ground = abs(math.sqrt( (y_0 - 0) * 2.0 / g ))

# How close is the simulation?
#  Find the time closest to t_ground
sim_t_ground = sim_times[ sim_times >= t_ground][0]  # take earliest
##sim_value_at_approx_ground = positions[ sim_times >= t_ground ][0]  # take first
percent_diff_times = 100.0 * abs(sim_t_ground - t_ground) / t_ground


# Find the energy
mass = 1  # arbitrary
energy = mass*g*positions + 0.5*mass*velocities**2
an_energy = 98 # at t=0, v=0 so E = m*g*10 = 98 J

## Plotting
pl.subplot(2,1,1)
pl.title(r'Falling body ($dt=%.4f$, $diff=%.1f\%%$)' % (dt, percent_diff_times))
sim_pt_size = 7
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, ys_an, color='red', label='analytical')
# Vert line for time when ground is hit
pl.plot([t_ground, t_ground], [-1, 10], color='grey', linewidth=0.5, linestyle='--')
# Show point at which analytic hits ground and its value
pl.scatter([t_ground,], [0,], 50, color='red')
pl.annotate('$t={:.2f}$'.format(t_ground),
         xy=(t_ground, 0), xycoords='data',
         xytext=(-60, 40), textcoords='offset points', fontsize=16,
         arrowprops=dict(arrowstyle="->", connectionstyle="arc3,rad=-.5"))

pl.legend(loc='upper right')
pl.xlim(0, 1.6)
# 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.xaxis.set_ticks([0.5, 1, 1.5])
ax.yaxis.set_ticks([0, 10])
##ax.yaxis.set_major_locator(pl.MaxNLocator(nbins=1, prune='lower'))
##ax.xaxis.set_major_locator(pl.MaxNLocator(nbins=3, prune='lower'))

pl.xlabel('Time (s)')  # needs to be after messing with axes or won't show up
pl.ylabel('Height (m)')

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

ax = pl.gca()
pl.xlim(0, 1.6)
ax.xaxis.set_ticks([0.5, 1, 1.5])
pl.ylim(np.floor(energy.min())-1, np.ceil(energy.max()))
ax.yaxis.set_major_locator(pl.MaxNLocator(nbins=2))

pl.show()