Google Colabで始めるFEniCSx入門：Poisson方程式，誤差評価，log-logプロット

# すでに dolfinx が入っているかを試します
try:
    import dolfinx

# 入っていなければ，Colab 用のインストールスクリプトを実行します
except ImportError:
    !wget "https://fem-on-colab.github.io/releases/fenicsx-install-release-real.sh" -O "/tmp/fenicsx-install.sh" \
        && bash "/tmp/fenicsx-install.sh"

# 最後に，ちゃんと import できることを確認します
import dolfinx
print("dolfinx is ready.")

# NumPy を使います
import numpy as np

# グラフ描画に Matplotlib を使います
import matplotlib.pyplot as plt

# 弱形式を書くために UFL を使います
import ufl

# FEniCSx は MPI ベースなので communicator を読み込みます
from mpi4py import MPI

# メッシュと有限要素空間のための機能を読み込みます
from dolfinx import mesh, fem

# 線形問題をまとめて解くためのクラスを読み込みます
from dolfinx.fem.petsc import LinearProblem

# 今回は 1 プロセスとして，communicator を 1 つだけ名前で持ちます
COMM = MPI.COMM_WORLD


# 厳密解 u(x,y)=sin(pi x)sin(pi y) を Python 関数として定義します
# これは境界値の補間に使います
def exact_solution(x):
    return np.sin(np.pi * x[0]) * np.sin(np.pi * x[1])


# N を受け取って，N×N 分割メッシュ上で計算し，誤差を返す関数を作ります
def solve_and_measure(N):

    # 対角線を交互に入れて，角の「死んだ要素」を避けます
    domain = mesh.create_unit_square(
        COMM,
        N,
        N,
        diagonal=mesh.DiagonalType.left_right,
    )

    # 1 次 Lagrange 要素による有限要素空間 V_h を作ります
    V = fem.functionspace(domain, ("Lagrange", 1))

    # 空間座標 x=(x_0,x_1) を UFL の記号として定義します
    x = ufl.SpatialCoordinate(domain)

    # 厳密解を UFL の式として定義します
    # 誤差評価ではこちらを使って「真の誤差」を測ります
    u_exact_expr = ufl.sin(np.pi * x[0]) * ufl.sin(np.pi * x[1])

    # 右辺 f = -Δu を UFL で定義します
    f = 2 * np.pi**2 * ufl.sin(np.pi * x[0]) * ufl.sin(np.pi * x[1])

    # 境界条件で使うために，有限要素空間上の関数 u_D を作ります
    u_D = fem.Function(V)

    # 厳密解を節点に補間して，Dirichlet 境界値として使える形にします
    u_D.interpolate(exact_solution)

    # 空間次元は 2 なので，境界は 1 次元の面です
    fdim = domain.topology.dim - 1

    # 境界上の全ての facet を見つけます
    boundary_facets = mesh.locate_entities_boundary(
        domain, fdim, lambda x: np.full(x.shape[1], True, dtype=bool)
    )

    # 境界上の自由度を取り出します
    boundary_dofs = fem.locate_dofs_topological(V, fdim, boundary_facets)

    # Dirichlet 境界条件 u=u_D を作ります
    bc = fem.dirichletbc(u_D, boundary_dofs)

    # TrialFunction は未知関数 u_h を表す記号です
    u = ufl.TrialFunction(V)

    # TestFunction は試験関数 v を表す記号です
    v = ufl.TestFunction(V)

    # 双線形形式 a(u,v)=∫ grad(u)・grad(v) dx を定義します
    a = ufl.inner(ufl.grad(u), ufl.grad(v)) * ufl.dx

    # 線形形式 L(v)=∫ f v dx を定義します
    L = f * v * ufl.dx

    # 線形問題 a(u_h,v)=L(v) を解く準備をします
    problem = LinearProblem(
        a,
        L,
        bcs=[bc],
        petsc_options_prefix=f"poisson_{N}_",
        petsc_options={
            "ksp_type": "preonly",
            "pc_type": "lu",
            "ksp_error_if_not_converged": True,
        },
    )

    # 実際に線形方程式を解いて数値解 u_h を得ます
    u_h = problem.solve()

    # L2 絶対誤差 ||u-u_h||_{L2} を計算します
    abs_L2 = np.sqrt(
        fem.assemble_scalar(
            fem.form((u_h - u_exact_expr) ** 2 * ufl.dx)
        )
    )

    # H1 半ノルム絶対誤差 |u-u_h|_{H1} を計算します
    # ここでは補間関数ではなく，厳密解そのものの勾配と比較しています
    abs_H1 = np.sqrt(
        fem.assemble_scalar(
            fem.form(
                ufl.inner(
                    ufl.grad(u_h - u_exact_expr),
                    ufl.grad(u_h - u_exact_expr),
                ) * ufl.dx
            )
        )
    )

    # 厳密解の L2 ノルム ||u||_{L2} を計算します
    norm_L2 = np.sqrt(
        fem.assemble_scalar(
            fem.form(u_exact_expr ** 2 * ufl.dx)
        )
    )

    # 厳密解の H1 半ノルム |u|_{H1} を計算します
    seminorm_H1 = np.sqrt(
        fem.assemble_scalar(
            fem.form(
                ufl.inner(ufl.grad(u_exact_expr), ufl.grad(u_exact_expr)) * ufl.dx
            )
        )
    )

    # 相対 L2 誤差を計算します
    rel_L2 = abs_L2 / norm_L2

    # 相対 H1 半ノルム誤差を計算します
    rel_H1 = abs_H1 / seminorm_H1

    # 一様メッシュでは代表的なメッシュ幅として h=1/N とみなせます
    h = 1.0 / N

    # 後で表やグラフに使うため，必要な値をまとめて返します
    return h, abs_L2, rel_L2, abs_H1, rel_H1


# 誤差を見るために，粗いメッシュから細かいメッシュまで並べます
Ns = [8, 16, 32, 64, 128]

# h を入れる空リストを用意します
hs = []

# L2 の絶対誤差を入れる空リストを用意します
abs_L2s = []

# L2 の相対誤差を入れる空リストを用意します
rel_L2s = []

# H1 半ノルムの絶対誤差を入れる空リストを用意します
abs_H1s = []

# H1 半ノルムの相対誤差を入れる空リストを用意します
rel_H1s = []

# 各メッシュサイズで solve_and_measure を実行します
for N in Ns:
    h, abs_L2, rel_L2, abs_H1, rel_H1 = solve_and_measure(N)
    hs.append(h)
    abs_L2s.append(abs_L2)
    rel_L2s.append(rel_L2)
    abs_H1s.append(abs_H1)
    rel_H1s.append(rel_H1)

# あとで計算しやすいように NumPy 配列に変えます
hs = np.array(hs)
abs_L2s = np.array(abs_L2s)
rel_L2s = np.array(rel_L2s)
abs_H1s = np.array(abs_H1s)
rel_H1s = np.array(rel_H1s)


# 連続する 2 点の log の差から実験的収束率を計算する関数を作ります
def rates(errors, hs):
    return np.log(errors[1:] / errors[:-1]) / np.log(hs[1:] / hs[:-1])


# 絶対誤差と相対誤差の収束率を表示します
print("L2 absolute rates  =", rates(abs_L2s, hs))
print("L2 relative rates  =", rates(rel_L2s, hs))
print("H1 absolute rates  =", rates(abs_H1s, hs))
print("H1 relative rates  =", rates(rel_H1s, hs))

# L2 誤差の log-log プロットを描きます
plt.figure(figsize=(6, 5))
plt.loglog(hs, abs_L2s, "o-", label="absolute L2 error")
plt.loglog(hs, rel_L2s, "s--", label="relative L2 error")
plt.xlabel("h")
plt.ylabel("error")
plt.title("L2 errors on a log-log plot")
plt.grid(True, which="both")
plt.legend()
plt.show()

# H1 半ノルム誤差の log-log プロットを描きます
plt.figure(figsize=(6, 5))
plt.loglog(hs, abs_H1s, "o-", label="absolute H1 seminorm error")
plt.loglog(hs, rel_H1s, "s--", label="relative H1 seminorm error")
plt.xlabel("h")
plt.ylabel("error")
plt.title("H1 seminorm errors on a log-log plot")
plt.grid(True, which="both")
plt.legend()
plt.show()

# 3D 描画を使うために読み込みます
from mpl_toolkits.mplot3d import Axes3D  # noqa: F401

# 可視化専用の分割数です
# 12〜24 程度にすると見やすいです
N_vis = 16

# -------------------------
# 1. 可視化用の問題をもう一度解きます
# -------------------------
domain_vis = mesh.create_unit_square(
    COMM,
    N_vis,
    N_vis,
    diagonal=mesh.DiagonalType.left_right,
)

V_vis = fem.functionspace(domain_vis, ("Lagrange", 1))

x_vis = ufl.SpatialCoordinate(domain_vis)
f_vis = 2 * np.pi**2 * ufl.sin(np.pi * x_vis[0]) * ufl.sin(np.pi * x_vis[1])

u_D_vis = fem.Function(V_vis)
u_D_vis.interpolate(exact_solution_numpy)

fdim_vis = domain_vis.topology.dim - 1
boundary_facets_vis = mesh.locate_entities_boundary(
    domain_vis,
    fdim_vis,
    lambda x: np.full(x.shape[1], True, dtype=bool),
)
boundary_dofs_vis = fem.locate_dofs_topological(V_vis, fdim_vis, boundary_facets_vis)
bc_vis = fem.dirichletbc(u_D_vis, boundary_dofs_vis)

u_vis = ufl.TrialFunction(V_vis)
v_vis = ufl.TestFunction(V_vis)

a_vis = ufl.inner(ufl.grad(u_vis), ufl.grad(v_vis)) * ufl.dx
L_vis = f_vis * v_vis * ufl.dx

problem_vis = LinearProblem(
    a_vis,
    L_vis,
    bcs=[bc_vis],
    petsc_options_prefix=f"visual_{N_vis}_",
    petsc_options={
        "ksp_type": "preonly",
        "pc_type": "lu",
        "ksp_error_if_not_converged": True,
    },
)

u_h_vis = problem_vis.solve()

# -------------------------
# 2. 2 次元のメッシュ図を描きます
# -------------------------
domain_vis.topology.create_connectivity(domain_vis.topology.dim, 0)
cells_vis = domain_vis.topology.connectivity(domain_vis.topology.dim, 0).array.reshape(-1, 3)
points_vis = domain_vis.geometry.x[:, :2]

plt.figure(figsize=(6, 6))
plt.triplot(points_vis[:, 0], points_vis[:, 1], cells_vis, linewidth=0.8)
plt.xlabel("x")
plt.ylabel("y")
plt.title(f"2D mesh (N = {N_vis})")
plt.gca().set_aspect("equal")
plt.show()

# -------------------------
# 3. 2 次元の等高線図を描きます
# -------------------------
num_cells_vis = domain_vis.topology.index_map(domain_vis.topology.dim).size_local
triangles_vis = np.vstack([V_vis.dofmap.cell_dofs(c) for c in range(num_cells_vis)])

# ここでは reshape せず，自由度座標をそのまま受け取ります
dof_coords_vis = V_vis.tabulate_dof_coordinates()

# 2 次元表示なので x, y の 2 成分だけを使います
xy_vis = dof_coords_vis[:, :2]

# 自由度に対応する数値解の値を取り出します
values_vis = u_h_vis.x.array.real

plt.figure(figsize=(6, 5))
plt.tricontourf(
    xy_vis[:, 0],
    xy_vis[:, 1],
    triangles_vis,
    values_vis,
    levels=20,
)
plt.colorbar(label="u_h")
plt.xlabel("x")
plt.ylabel("y")
plt.title(f"2D contour of FEM solution (N = {N_vis})")
plt.gca().set_aspect("equal")
plt.show()

# -------------------------
# 4. 近似解の 3 次元曲面図を描きます
# -------------------------
fig = plt.figure(figsize=(8, 6))
ax = fig.add_subplot(111, projection="3d")

ax.plot_trisurf(
    xy_vis[:, 0],
    xy_vis[:, 1],
    values_vis,
    triangles=triangles_vis,
    linewidth=0.2,
    antialiased=True,
)

ax.set_xlabel("x")
ax.set_ylabel("y")
ax.set_zlabel("u_h")
ax.set_title(f"3D surface of FEM solution (N = {N_vis})")
plt.show()