(六)、Fealpy 组装刚度（质量）矩阵和载荷向量

给定一个单纯形网格 $\mathcal{T}$, 其有 $NN$ 个节点， $NC$ 个单元。 定义在 $\mathcal{T}$ 的分片 $p$ 次连续有限元空间 $V_h$ 有 $gdof$ 个基函数， 其组成的函数的行向量为：
$$\varphi =[{\varphi }_0,{\varphi }_1,\cdots ,{\varphi }_{gdof-1}],$$
限制在每个网格单元 $\tau$ 上， 共有 $ldof$ 个基函数：
$${\varphi }_{\tau }=[{\varphi }_0,{\varphi }_1,\cdots ,{\varphi }_{ldof-1}].$$
此时 stiff matrix 为
$$A={\int }_{\Omega }(\nabla \varphi {)}^T\nabla \varphi \mathrm{d}\mathbf{x}.$$
注意 $A$ 是一个稀疏矩阵。 Fealpy 的想法是先在局部组装刚度矩阵，再以某种方式拼接在一起得到全局的刚度矩阵。 实际计算中， Fealpy 采用数值积分的方式来实现：

# 刚度矩阵的计算
# 假设已经定义好 mesh 和 space
qf = mesh.integrator(q, 'cell')  # 获得第q个积分公式
bcs, ws = qf.get_quadrature_points_and_weights()  # 拿到相应的重心坐标和权重
cellmeasure = mesh.entity_measure('cell')  # 得到单元的面积
gphi = space.grad_basis(bcs)  # 计算基函数在重心坐标处的梯度值(NQ, NC, ldof, GD)
# NC: 单元个数 ldof:局部基函数个数 # GD: 几何维数 NQ:积分点个数 
# 组装刚度矩阵 A.shape==(NC, ldof, ldof) <= (ldof, gd) * (gd, ldof)
A = np.einsum('i, ijkl, ijml, j -> jkl', ws, gphi, gphi, cellmeasure)

我们来说明一下上述代码中的最后一行：‘i’ 代表权重，‘l’代表gd, 由einsum知， 这是对权重和gd求和， 正好是我们要得到的式子。
有了单元刚度矩阵，我们便可以计算总体刚度矩阵。

import scipy.sparse import csr_matrix

gdof = space.number_of_global_dof()  # 全局自由度的个数
# (NC, ldof) cell2dof[i,j] 存储第i个单元上局部第j个自由度的全局编号
cell2dof = space.cell_to_dof()

# (NC, ldof) -> (Nc, ldof,1) -> (NC, ldof, ldof)
I = np.broadcast_to(cell2dof[:, :, None], shape=A.shape)
# (NC, ldof) -> (NC, 1, ldof) -> (NC, ldof, ldof)
J = np.broadcast_to(cell2dof[:, None, :], shape=A.shape)

A = csr_matrix((A.flat, (I.flat, J.flat)), shape=(gdodf, gdof))

类似地， 单元载荷向量的组装便很简单了：

实现的代码如下：

import numpy as np
from fealpy.mesh import MeshFactory
from fealpy.functionspace import LagrangeFiniteElementSpace
from matplotlib import pyplot as plt
from scipy.sparse import csr_matrix
mf = MeshFactory()
mesh = mf.boxmesh2d([0, 1, 0, 1], nx=1, ny=1, meshtype='tri')

fig = plt.figure()
axes = fig.gca()
mesh.add_plot(axes)

space = LagrangeFiniteElementSpace(mesh, p=1)
gdof = space.number_of_global_dofs()
print('gdof:', gdof)

cell2dof = space.cell_to_dof()
print('cell2dof:', cell2dof)
cell = mesh.entity('cell')
print('cell:', cell)

mesh.find_node(axes, showindex=True, fontsize=24)
mesh.find_cell(axes, showindex=True, fontsize=24)

cellmeasure = mesh.entity_measure('cell')
print('cellmeasure:', cellmeasure)

qf = mesh.integrator(1, 'cell')  # 积分公式
bcs, ws = qf.get_quadrature_points_and_weights()
print('bcs:', bcs, 'ws:', ws)  # 重心坐标与权重
gphi = space.grad_basis(bcs)  # [1, 2, 3, 2]  q * NC * gphi * dim 积分点*单元*基函数*导数分量(维数)

A = np.einsum('i, ijkl, ijml, j-> jkm', ws, gphi, gphi, cellmeasure)  # (2, 3, 3) 单元 基函数 基函数

I = np.broadcast_to(cell2dof[:, :, None], shape=A.shape)
J = np.broadcast_to(cell2dof[:, None, :], shape=A.shape)
# A, I, J  # (2, 3, 3)
for i, j, val in zip(I.flat, J.flat, A.flat):
    print('(', i, ',', j, '):', val)
A = csr_matrix((A.flat, (I.flat, J.flat)), shape=(gdof, gdof))  # (gdof, gdof)
A.toarray()

# mass matrix
phi = space.basis(bcs)  # (NQ, 1, ) 1:NC 节省内存
M = np.einsum('i, ijk, ijm, j->jkm', ws, phi, phi, cellmeasure)
M = csr_matrix((M.flat, (I.flat, J.flat)), shape=(gdof, gdof))  # (gdof, gdof)
M.toarray()

# 对流扩散方程
# 对流项组装
b = np.array([1, 1], dtype=np.float64)
B = np.einsum('i, q, ijkq, ijm, j-> jkm', ws, b, gphi, phi, cellmeasure)
B = csr_matrix((B.flat, (I.flat, J.flat)), shape=(gdof, gdof))  # (gdof, gdof)
B.toarray()
# 扩散项组装
D = np.array([[10.0, -1.0], [-1.0, 2.0]], dtype=np.float64)
B = np.einsum('i,sq, ijkq, ijls,j->jkl', ws, D, gphi, gphi, cellmeasure)
B = csr_matrix((B.flat, (I.flat, J.flat)), shape=(gdof, gdof))  # (gdof, gdof)
print('B', B.toarray())

B = space.stiff_matrix(c=D)
print(B.toarray())
# 组装向量
from fealpy.decorator import cartesian
@cartesian
def f(p):
    x = p[..., 0]
    y = p[..., 1]
    return np.exp(x**2 + y**2)
# print(f.coordtype)
qf= mesh.integrator(3, 'cell')
bcs, ws = qf.get_quadrature_point_and_weight()
phi = space.basis(bcs)  # (NQ, 1, 3) 基函数
ps = mesh.bc_to_point(bcs)  # (6, 2, 2) 积分点 单元 分量
val = f(ps)
bb = np.einsum('i, ij, ijk, j -> jk', ws, val, phi, cellmeasure)
F = np.zeros(gdof, dtype=np.float64)
np.add.at(F, cell2dof, bb)

plt.show()

笔者感悟： 我在魏老师上课代码中加了一个扩散项的组装，经过与fealpy中space.stiff_matrix(c=diffusion_coefficient) 比对，结果是对的。 我花了三小节的时间来整理、处理矩阵的组装， 是因为我觉得这是fealpy 的灵魂所在：包括cell_to_dof， csr_matrix, np.einsum 的引入等等。学习felapy, 我们要学到其背后的原理与编程的想法， 这对我我们解自己专业的问题至关重要。我这么说并不意味着用fealpy解PDE都要自己组装。 事实上， 魏老师写了很多PDE的接口， 下一节我们就来介绍这些接口， 和展示一些数值实验。