#51 Solution of Laplace equation using FEM

import os
import matplotlib.pyplot as plt
import numpy as np
from numpy import linspace
from scipy.spatial import Delaunay
from numpy import cross
from scipy.sparse import dok_matrix
import matplotlib.cm as cm
os.chdir(r'D:\projects\wordpress\ex51')
os.getcwd()
xmin = 0 ; xmax = 1 ; nXpoints = 10
ymin = 0 ; ymax = 1 ; nYpoints = 10
horizontal = linspace(xmin,xmax,nXpoints)
vertical = linspace(ymin,ymax,nYpoints)
y, x = np.meshgrid(horizontal, vertical)
vertices = np.array([x.flatten(),y.flatten()])
triangulation = Delaunay(vertices.T)
index2point = lambda index: triangulation.points[index]
all_centers = index2point(triangulation.vertices).mean(axis=1)
trngl_set=triangulation.vertices
plt.triplot(vertices[0],vertices[1],triangles=trngl_set)
plt.savefig("example51_1.png", dpi=100)
plt.show()
points=triangulation.points.shape[0]
stiff_matrix=dok_matrix((points,points))
for triangle in triangulation.vertices:
    helper_matrix=dok_matrix((points,points))
    pt1,pt2,pt3=index2point(triangle)
    area=abs(0.5*cross(pt2-pt1,pt3-pt1))
    coeffs=0.5*np.vstack((pt2-pt3,pt3-pt1,pt1-pt2))/area
    #helper_matrix[triangle,triangle] = \
    np.array(np.mat(coeffs)*np.mat(coeffs).T)
    u=None
    u=np.array(np.mat(coeffs)*np.mat(coeffs).T)
    for i in range(len(triangle)):
        for j in range(len(triangle)):
            helper_matrix[triangle[i],triangle[j]] = u[i,j]
    stiff_matrix=stiff_matrix+helper_matrix

allNodes = np.unique(trngl_set)
boundaryNodes = np.unique(triangulation.convex_hull)
NonBoundaryNodes = np.array([])
for x in allNodes:
    if x not in boundaryNodes:
        NonBoundaryNodes = np.append(NonBoundaryNodes,x)

NonBoundaryNodes = NonBoundaryNodes.astype(int)
nbnodes = len(boundaryNodes) # number of boundary nodes
FbVals=np.zeros([nbnodes,1]) # Values on the boundary
FbVals[(nbnodes-nXpoints+1):-1]=np.ones([nXpoints-2, 1])
totalNodes = len(allNodes)
stiff_matrixDense = stiff_matrix.todense()
stiffNonb = stiff_matrixDense[np.ix_(NonBoundaryNodes,NonBoundaryNodes)]
stiffAtb = stiff_matrixDense[np.ix_(NonBoundaryNodes,boundaryNodes)]
U=np.zeros([totalNodes, 1])
U[NonBoundaryNodes] = np.linalg.solve( - stiffNonb,stiffAtb * FbVals )
U[boundaryNodes] = FbVals
X = vertices[0]
Y = vertices[1]
Z = U.T.flatten()
from mpl_toolkits.mplot3d import axes3d
fig = plt.figure()
ax = fig.add_subplot(111, projection='3d')
surf = ax.plot_trisurf(X, Y, Z, cmap=cm.jet, linewidth=0)
fig.colorbar(surf)
ax.set_xlabel('X',fontsize=14)
ax.set_ylabel('Y',fontsize=14)
ax.set_zlabel(r"$\phi$",fontsize=16)
plt.savefig("example51_2.png", dpi=100)
plt.show()
from numpy import pi, sinh, sin, cos
def f(x,y):
    return sum( 2*(1.0/(n*pi) - cos(n*pi)/(n*pi))*(sinh(n*pi*x)/ \
                   sinh(n*pi))*sin(n*pi*y) for n in range(1,200))
Ze = f(X,Y)
ZdiffZe = Ze - Z
plt.plot(ZdiffZe)
plt.savefig("example51_3.png", dpi=100)
plt.show()
plt.close()