#(c) Rory Clune 03.24.2011
# Animation of a bending beam, considering only lateral displacements

from numpy import *
from numpy.linalg import inv
import pygame

elements = 100	#the number of elements for the beam
load = 1.0		#the vertical point load at the end of the beam
EI = 1.0		#the bending rigidity
L = 1.0			#the length of the beam
h = L/elements	#the length of an element
dof = 2*(elements+1) #the numbers of degrees of freedom for the problem

#1. set up the matrices and vectors (derived from Direct Stiffness Method)
A = zeros((dof,dof))
for i in range (2,dof-2,1):
	if i%2==0:
		A[i,i-2] = -12*EI/h**3
		A[i,i-1] = -6*EI/h**2	
		A[i,i] = 24*EI/h**3
		A[i,i+1] = 0
		A[i,i+2] = -12*EI/h**3
		A[i,i+3] = 6*EI/h**2
	else:		
		A[i,i-3] = 6*EI/h**2
		A[i,i-2] = 2*EI/h
		A[i,i-1] = 0	
		A[i,i] = 8*EI/h
		A[i,i+1] = -6*EI/h**2
		A[i,i+2] = 2*EI/h
A[dof-2,dof-4] = -12*EI/h**3
A[dof-2,dof-3] = -6*EI/h**2
A[dof-2,dof-2] = 12*EI/h**3
A[dof-2,dof-1] = -6*EI/h**2
A[dof-1,dof-4] = 6*EI/h**2
A[dof-1,dof-3] = 2*EI/h
A[dof-1,dof-2] = -6*EI/h**2
A[dof-1,dof-1] = 4*EI/h
#Apply Dirichlet boundary conditions (slope and deflection at left = 0):
for i in range (0,dof,1):
	for j in range (0,2,1):
		A[i,j] = 0.0
		A[j,i] = 0.0		
A[0,0] = 1.0
A[1,1] = 1.0
Am = mat(A)
b = zeros((dof,1))
bm = mat(b)

print Am
#2. Solve for response and animate using pygame
screen = pygame.display.set_mode((640, 480))
clock = pygame.time.Clock()    
running = 1
k = 0; dk = 1
while running:
	pygame.time.wait(10)
	event = pygame.event.poll()
	if event.type == pygame.QUIT:
		running = 0
	screen.fill((255, 255, 255))		
	load = k
	b[dof-2,0] = load #point load at end
	u = inv(Am)*bm
	pointList = []
	for i in range(0,dof-1,2):
		pointList.append([10+i*h*300,240+u[i,0]*50])
	# Draw the beam:
	pygame.draw.aalines(screen, (0, 0, 255), False, pointList, 1)
	if k == 10 or k == -10:
		dk *= -1
	k += dk
	pygame.display.flip()