#
# beam.py
# (c) Neil Gershenfeld  3/18/11
#

#
# import symbolic math
#
from sympy import *
#
# define symbols
#
x,h = symbols('xh')
#
# define and invert coefficient matrix
#
M = Matrix([[1,0,0,0],[0,1,0,0],[1,h,h**2,h**3],[0,1,2*h,3*h**2]])
Minv = M.inv()
#
# take derivatives
#
v = [1,x,x**2,x**3]
psi = zeros((1,4))
dpsi = zeros((1,4))
ddpsi = zeros((1,4))
for i in range(4):
   psi[i] = Minv[:,i].dot(v)
   dpsi[i] = diff(psi[i],x)
   ddpsi[i] = diff(dpsi[i],x)
#
# evaluate overlap integrals
#
I = zeros((4,4))
for i in range(4):
   for j in range(4):
      I[i,j] = integrate(ddpsi[i]*ddpsi[j],(x,0,h)).subs(h,1)
#
# plot interactively
#
from numpy import *
import matplotlib.pyplot as plt
plt.ion()
N = 100
x=linspace(0,1,N)
y0 = eval(str(psi[0].subs(h,1)))
y1 = eval(str(psi[1].subs(h,1)))
y2 = eval(str(psi[2].subs(h,1)))
y3 = eval(str(psi[3].subs(h,1)))
plt.plot(x,y0,'r',x,y1,'g',x,y2,'b',x,y3,'k')
plt.draw()
raw_input("functions (enter to continue) ")
plt.clf()
y0 = eval(str(dpsi[0].subs(h,1)))
y1 = eval(str(dpsi[1].subs(h,1)))
y2 = eval(str(dpsi[2].subs(h,1)))
y3 = eval(str(dpsi[3].subs(h,1)))
plt.plot(x,y0,'r',x,y1,'g',x,y2,'b',x,y3,'k')
plt.draw()
raw_input("first derivatives (enter to continue) ")
plt.clf()
y0 = eval(str(ddpsi[0].subs(h,1)))
y1 = eval(str(ddpsi[1].subs(h,1)))
y2 = eval(str(ddpsi[2].subs(h,1)))
y3 = eval(str(ddpsi[3].subs(h,1)))
plt.plot(x,y0,'r',x,y1,'g',x,y2,'b',x,y3,'k')
plt.draw()
raw_input("second derivatives (enter to continue) ")
plt.clf()
#
# construct overlap matrix
#
A = zeros((N,N))
A[0,0] = I[0,0]
A[0,1] = I[0,1]
A[0,2] = I[0,2]
A[0,3] = I[0,3]
A[1,0] = I[1,0]
A[1,1] = I[1,1]
A[1,2] = I[1,2]
A[1,3] = I[1,3]
for i in range(2,N-3,2):
   A[i,i-2] = I[2,0]
   A[i,i-1] = I[2,1]
   A[i,i] = I[2,2]+I[0,0]
   A[i,i+1] = I[2,3]+I[0,1]
   A[i,i+2] = I[0,2]
   A[i,i+3] = I[0,3]
   A[i+1,i-2] = I[3,0]
   A[i+1,i-1] = I[3,1]
   A[i+1,i] = I[3,2]+I[1,0]
   A[i+1,i+1] = I[3,3]+I[1,1]
   A[i+1,i+2] = I[1,2]
   A[i+1,i+3] = I[1,3]
A[N-2,N-4] = I[2,0]
A[N-2,N-3] = I[2,1]
A[N-2,N-2] = I[2,2]
A[N-2,N-1] = I[2,3]
A[N-1,N-4] = I[3,0]
A[N-1,N-3] = I[3,1]
A[N-1,N-2] = I[3,2]
A[N-1,N-1] = I[3,3]
#
# apply boundary condition at origin for 0 displacement, slope
#
A[:,0] = 0
A[0,:] = 0
A[:,1] = 0
A[1,:] = 0
A[0,0] = 1
A[1,1] = 1
#
# invert
#
B = linalg.inv(A)
#
# loop over forces at end and plot
#
b = zeros(N)
for i in linspace(0,1,10):
   b[-2] = i
   c = dot(B,b)
   u = c[::2]
   plt.plot(u,'k')
raw_input("beam bending (enter to continue) ")