SciFEM_Schratter/ex4/ex_4A.py

76 lines
No EOL
1.2 KiB
Python

import numpy as np
import matplotlib.pyplot as plt
np.set_printoptions(precision=2)
N = 10 # number of elements
a = 1
alpha = 1
g_b = 1
f_const = 1
x = np.linspace(0, 1, N + 1)
A = np.zeros((N + 1, N + 1))
f_vec = np.zeros(N + 1)
for i in range(1, N + 1):
h = x[i] - x[i - 1]
a_11 = 1./h + a*h/3.
a_12 = -1./h + a*h/6.
a_21 = -1./h + a*h/6.
a_22 = 1./h + a*h/3
A[i - 1, i - 1] += a_11
A[i - 1, i] += a_12
A[i, i - 1] += a_21
A[i, i] += a_22
f_vec[i] = f_const*h
print("A =\n", A)
# take Neumann data into account
A[N, N] += alpha
f_vec[N] += alpha*g_b
# take Dirichlet data into account
u_g = np.zeros(N + 1)
u_g[0] = 0
print("u_g =\n", u_g)
# remove first row of A
A_g = A[1:N+1, :]
#print("A_g =\n", A_g)
# remove first row of f_vec
f_g = f_vec[1:N+1]
# assemble RHS with dirichlet data
f_g -= A_g.dot(u_g)
#print(f_g)
# matrix for the inner nodes (excluding nodes with dirichlet bcs)
A_0 = A[1:N+1, 1:N+1]
#print(A_0)
# solve for u_0 (free dofs)
u_0 = np.linalg.solve(A_0, f_g)
# assemble "u = u_0 + u_g"
u = np.concatenate([[0], u_0])
print("u =\n", u)
plt.plot(x, u, '-')
plt.xlabel('x')
plt.ylabel('u_h(x)')
plt.grid()
plt.show()