%load_ext sympyprinting
from sympy import *
from numpy import *
from sympy.abc import x
from matplotlib import pyplot as plt
%matplotlib inline 

h=symbols('h')

#For A_{i,i}, linear hat functions completely overlap
f1 = x/h;
plot(f1.subs(h,.1),(f1*f1).subs(h,.1),(x,0,.1))
2*integrate(f1*f1,(x,0,h))

#For A_{i,i+1} and A_{i,i-1}, linear hat functions partially overlap
f1 = x/h;
f2 = 1-x/h;
plot(f1.subs(h,.1),f2.subs(h,.1),(f1*f2).subs(h,.1),(x,0,.1))
integrate(f1*f2,(x,0,h))

f1 = x/h;
integrate(f1,(x,0,h))

#Now for Hermite polynomials.  As above, we consider on a single interval, with varying overlaps.
a2f = lambda a: a[0]+a[1]*x+a[2]*x**2+a[3]*x**3
M = Matrix([[1,0,0,0],[0,1,0,0],[1,h,h**2,h**3],[0,1,2*h,3*h**2]]).inv()
f1 = a2f(M.dot(Matrix([1,0,0,0])));
f2 = a2f(M.dot(Matrix([0,1,0,0])));
f3 = a2f(M.dot(Matrix([0,0,1,0])));
f4 = a2f(M.dot(Matrix([0,0,0,1])));
hval = 1.6
plot(f1.subs(h,hval),f2.subs(h,hval),f3.subs(h,hval),f4.subs(h,hval),(x,0,hval))

#test
f2u = lambda f: [f.subs(x,0),diff(f,x).subs(x,0),f.subs(x,h),diff(f,x).subs(x,h)]
[f2u(f) for f in [f1,f2,f3,f4]]

#A_ii for i even
2*integrate(f1*f1,(x,0,h))

#A_ii for i odd
2*integrate(f2*f2,(x,0,h))

#A_{i,i+1} for i even
2*integrate(f1*f2,(x,0,h))

#A_{i,i+1} for i odd
integrate(f2*f3,(x,0,h))

#A_{i,i+2} for i even
integrate(f1*f3,(x,0,h))

#A_{i,i+2} for i odd
integrate(f2*f4,(x,0,h))

#A_{i,i+3} for i even
integrate(f1*f4,(x,0,h))

#To calculate B, we need shape function derivatives.
df1 = diff(f1,x);
df2 = diff(f2,x);
df3 = diff(f3,x);
df4 = diff(f4,x);
hval = 1.6
plot(df1.subs(h,hval),df2.subs(h,hval),df3.subs(h,hval),df4.subs(h,hval),(x,0,hval))

#B_ii for i even
2*integrate(df1*df1,(x,0,h))

#B_ii for i even
2*integrate(df2*df2,(x,0,h))

#B_{i,i+1} for i even
2*integrate(df1*df2,(x,0,h))

#B_{i,i+1} for i odd
integrate(df2*df3,(x,0,h))

#A_{i,i+2} for i even
integrate(df1*df3,(x,0,h))

#A_{i,i+2} for i odd
integrate(df2*df4,(x,0,h))

#A_{i,i+3} for i even
integrate(df1*df4,(x,0,h))

#For 9.2, we need second spatial derivatives.
ddf1 = diff(f1,x,2);
ddf2 = diff(f2,x,2);
ddf3 = diff(f3,x,2);
ddf4 = diff(f4,x,2);
hval = 1.6
plot(ddf1.subs(h,hval),ddf2.subs(h,hval),ddf3.subs(h,hval),ddf4.subs(h,hval),(x,0,hval))

def draw_from_a(a,hval,ax,color=None,label=None):
    #take a hermite solution and draw it
    u = dstack((a[:-2:2],a[1:-2:2],a[2::2],a[3::2]))[0]
    f = [a2f(M.dot(Matrix(xi))) for xi in u]
    seg_pts = linspace(0,hval,seg_res)
    for i,seg_f in enumerate(f):
        seg_f = lambdify((x),seg_f.subs(h,hval),'numpy')
        ax.plot(seg_pts+i*hval, [seg_f(xi) for xi in seg_pts],color=color if color else 'c')
        ax.plot(seg_pts[[0,-1]]+i*hval,[seg_f(xi) for xi in seg_pts[[0,-1]]],marker='.',linestyle='',color='b')

def symmetrize(a):
    #make a matrix symmetry, assuming you haven't set complementary elements
    #http://stackoverflow.com/questions/2572916/numpy-smart-symmetric-matrix
    return a + a.T - diag(a.diagonal())

from numpy.linalg import solve
def hermite_beam(n,E,I,f):
    #n: how many nodes, int
    #E: Young's Modulus, float
    #I: area moment of inertia, float
    #f: load on nodes, sympy expression
    seg_res = 20 #how many points to plot on each segment
    hval = 1./(n-1)

    integ = lambda func: integrate(func,(x,0,h)).subs(h,hval)
    an = arange(n); anm1 = arange(n-1)
    A = zeros((2*n,2*n))
    #A_i,i
    A[2*an,2*an] = 2*integ(ddf1*ddf1) 
    A[2*an+1,2*an+1] = 2*integ(ddf2*ddf2)
    A[2*n-2,2*n-2] /= 2
    A[2*n-1,2*n-1] /= 2
    #A_i,i+1
    A[2*an,2*an+1] = integ(ddf1*ddf2) + integ(ddf3*ddf4)
    A[2*n-2,2*n-1] = integ(ddf3*ddf4)
    A[2*anm1+1,2*anm1+2] = integ(ddf2*ddf3)
    #A_i,i+2
    A[2*anm1,2*anm1+2] = integ(ddf1*ddf3)
    A[2*anm1+1,2*anm1+3] = integ(ddf2*ddf4)
    #A_i,i+3
    A[2*anm1,2*anm1+3] = integ(ddf1*ddf4)

    A = symmetrize(A)
    A *= E*I

    #set up RHS
    B = zeros(2*n)
    B[2*n-2] = integ(f3*(f.subs(x,x+h*(n-1))))
    B[2*n-1] = integ(f4*(f.subs(x,x+h*(n-1))))
    for i in range(1,n-1):
        B[2*i] =   integ(f3*(f.subs(x,x+h*(i-1)))) + integ(f1*(f.subs(x,x+h*i)))
        B[2*i+1] = integ(f4*(f.subs(x,x+h*(i-1)))) + integ(f2*(f.subs(x,x+h*i)))

    # before solving we want to enforce boundary conditions a0 = a1 = 0
    def e(i): r = zeros(2*n); r[i]=1; return r
    A[0] = e(0); A[:,0] = e(0); B[0] = 0
    A[1] = e(1); B[1] = 0; A[:,1] = e(1)
    return solve(A,B)

f = -.8 + 0*x     #gravity load on nodes, use x to create Sympy expression
fig,ax = plt.subplots()
n=10
draw_from_a(hermite_beam(n,1.,1.,f),1./(n-1),ax)
plt.savefig('converge.png',dpi=200)

#f = x*((x-x**3/6+x**5/120-x**7/5040).subs(x,5*x))    
f = (1 - x**2 + x**4/2 - x**6/6 + x**8/24 - x**10/120 + x**12/720 - x**14/5040).subs(x,3*(x-.5))-.5
fig,ax = plt.subplots()
for n in [7,9,11,13]:
    draw_from_a(hermite_beam(n,.5,.5,f),1./(n-1),ax)
t = linspace(0,1,4*n)
f = lambdify((x),f,'numpy')
ax.plot(t,.1*f(t),color='r',label='load')
ax.legend(loc='upper left')
plt.savefig('converge.png',dpi=200)

exp(-x**2).series(x, 0, 20)