
# coding: utf-8

# ## Finite elements method
# This program is developed for solving equation: 
# 
# $$ -pu''(x) + qu(x) = f(x), \quad [0, \pi], \quad u(0)=u'(\pi)=0, 
# \quad p = 1, \quad q = 10, \quad f(x) = \sin(x) $$
# 
# with the third order hierarchical finite elements.

# #### Exact solution of the equation:
# $$ u(x) = \frac{ \exp{(kx)} - \exp{(-kx)} }{ 11k(\exp{(-k\pi)} +\exp{(k\pi)} )} + \frac{1}{11}\sin{(x)}, \quad k=\sqrt{10} $$

# #### Third order hierarchical basis:
# $$ N^1_{-1} = N^0 = \frac{1-\xi}{2} \\ $$ 
# $$ N^1_1 = N^1 = \frac{\xi+1}{2} \\ $$
# $$ N^2 = \frac{1}{2}(\xi^2 - 1) \\ $$
# $$ N^3 = \frac{1}{6}\xi(\xi^2-1) \\ $$

# #### K and M matricies of the problem:
# $ K = $    $ M = $

# In[44]:

import numpy as np
import matplotlib.pyplot as plt
np.set_printoptions(precision=5)
plt.rcParams['figure.dpi'] = 110


# In[45]:

def r(x):    # wolphram mathematica solution
    k = np.sqrt(10)
    c = 1/(110*(1 + np.exp(2*k*np.pi)))
    return c*(-k*np.exp(k*np.pi - k*x) + k*np.exp(k*np.pi + k*x) + 
              (k**2)*np.sin(x)*np.exp(k*x) + (k**2)*np.exp(2*k*np.pi)*np.sin(x))


# In[46]:

def u(x):    # my own solution
    k = np.sqrt(10)
    c = 1/(11*k)*(np.exp(-k*np.pi) + np.exp(k*np.pi))**(-1)
    return (-c*np.exp(-k*x) + c*np.exp(k*x) + 1/11*np.sin(x))


# In[47]:

t = np.linspace(0, np.pi, 100)
plt.plot(t, u(t), 'r', t, r(t), 'b--')
plt.xlim(0, np.pi)
plt.grid(True)
plt.show()


# #### Matricies construction:

# In[114]:

# parameters of the problem
p = 1.0; q = 10.0
def f(x): 
    return np.sin(x)

n = 10       # number of finite elements
N = 3*n        # size of global matricies = 4 + 3*(n-1) - 1
h = np.pi/n    # physical element length

# grid construction 
# example: 1 element on [0, 1] => Corresponding grid: [0, 1/2, 2/3, 1]
t = np.linspace(0, np.pi, N+1)
for i in np.arange(1, N+1, 3): t[i] += h/6
    
# grid for f(x) values in nodes    
tt = np.linspace(0, np.pi, N+1-n)

Kj = 2*p/h*np.array([ [ 1/2, 0  , 0   , -1/2], 
                      [   0, 2/3, 0   ,  0  ],
                      [   0, 0  , 2/45,  0  ],
                      [-1/2, 0  , 0   ,  1/2] ])

Mj = q*h/2*np.array([ [  2/3 , -1/3 ,  1/45 ,  1/3 ], 
                      [ -1/3 ,  4/15,  0    , -1/3 ],
                      [  1/45,  0   ,  4/945, -1/45],
                      [  1/3 , -1/3 , -1/45 ,  2/3 ] ])

lj = h/2*np.array([ [  2/3, -1/3,  1/45,   1/3], 
                    [ -1/3, 4/15,     0,  -1/3],
                    [  1/3, -1/3, -1/45,   2/3] ])


# In[121]:

K = np.zeros([N, N])
M = np.zeros([N, N])
l = np.zeros(N)

# fill matricies M and K
# boundary part
for i in range(3):
    for j in range(3):
        M[i, j] = Mj[i+1, j+1]
        K[i, j] = Kj[i+1, j+1]
# inner part
for s in range(3*n-3):
    if (s%3 == 0):
        M[2+s:6+s, 2+s:6+s] += Mj
        K[2+s:6+s, 2+s:6+s] += Kj

# fill vector l
# boundary part
tmp = np.array([f(tt[i]) for i in range(3)])
tmp = np.dot(tmp, lj)
l[0:3] = tmp[1:4]

#inner part
for s in range(3*n-3):
    if (s%3 == 0):
        index = int(s/3)*2 + 2
        tmp = np.array([f(tt[index + i]) for i in range(3)])
        tmp = np.dot(tmp, lj)
        l[2+s:6+s] += tmp

A = (K + M)
coeffs = np.linalg.solve(A, l)


# In[122]:

def N0(z):
    if (z>=-1 and z<=1): return (1 - z)/2
    else: return 0
    
def N1(z):
    if (z>=-1 and z<=1): return (z + 1)/2
    else: return 0
    
def N2(z):
    if (z>=-1 and z<=1): return (z**2 - 1)/2
    else: return 0

def N3(z):
    if (z>=-1 and z<=1): return z*(z**2 - 1)/6
    else: return 0
    
def mapping(z, xj_1, xj):
    return (1-z)/2*xj_1 + (1+z)/2*xj

def mapping_1(x, xj, xj_1):
    return -(2*x - (xj + xj_1))/(xj - xj_1)


# In[123]:

# solution construction
def solution(x, coeffs):
    
    if x==0:     return 0
    if x==np.pi: return coeffs[-1]
    
    fCoeffs = np.concatenate(([0], coeffs))
    
    t = np.linspace(0, np.pi, n+1)
    xj=0; xj_1=0; index=0
    for i in range(np.size(t)):
        if (x <= t[i]): 
            xj = t[i]; xj_1 = t[i-1]; index=i-1
            break
    
    c = [0, 0, 0, 0]
    for i in range(4):
        c[i] = fCoeffs[index*3 + i]
        
    res = (c[0]*N0(mapping_1(x, xj_1, xj)) + c[1]*N2(mapping_1(x, xj_1, xj)) + 
           c[2]*N3(mapping_1(x, xj_1, xj)) + c[3]*N1(mapping_1(x, xj_1, xj)))
    
    return res

t = np.linspace(0, np.pi, n+1)
fCoeffs = np.concatenate(([0], coeffs))
y = []
for i in range(np.size(fCoeffs)):
    if i%3 == 0: y.append(fCoeffs[i])
plt.plot(t, y, t, u(t), 'r', t, [solution(i, coeffs) for i in t], 'k')
plt.xlim(0, np.pi+0.1)
plt.show()
# In[124]:

t = np.linspace(0, np.pi, 2000)
plt.plot(t, u(t), 'r', t, [solution(i, coeffs) for i in t], 'k')
plt.xlim(0, np.pi)
plt.show()
d = u(np.pi)/solution(np.pi, coeffs)
print(d)
print()
print(coeffs)
print(solution(np.pi/4, coeffs))


# In[39]:

1.82675889088
1.57338794184
1.51788424441
1.50002837549
1.50000709381


# In[24]:

t = np.linspace(0, np.pi, 100)
plt.plot(t, list(map(solution, t)))
plt.show()


# In[ ]:




# In[117]:

6%3


# In[218]:

a = np.array([[0, 0, 0],
              [1, 2, 4]])
b = np.array([[1, 3], 
              [2, 3]])
a[0:2,1:3] += b
a


# In[60]:

a = np.array([0, 1, 2, 3, 4])
b = 2.4
for i in range(np.size(a)):
    if (b <= a[i]): 
        print(a[i])
        break


# In[70]:

t = np.linspace(0, np.pi, n+1)
t


# In[75]:

x = 0
t = np.linspace(0, np.pi, n+1)
xj=0; xj_1=0;
for i in range(np.size(t)):
    if (x <= t[i]): 
        xj = t[i]; xj_1 = t[i-1];
        break


# In[76]:

print(xj, xj_1)


# In[80]:

e = np.concatenate(([0], coeffs))


# In[81]:

e


# In[82]:

coeffs


# In[89]:

for i in range(0, np.size(coeffs), 3):
    print(i)


# In[ ]:



