PETSc[06] petsc4py中的TS

TS

TS（Time Stepper）是 PETSc 中用于求解时间依赖问题的接口，其不负责空间离散或线性/非线性求解，而是管理时间推进的整体流程，即：时间区间，时间步长，时间离散格式，以及每一步中该调用哪些算子。

在TS.pyx源文件和petscts.pxi源文件中，可以看到更完整的函数列表。

时间显格式/隐格式

PETSc 支持的时间依赖问题，通常写成下面两类形式之一。

显式（或半显式）形式：
$$
\frac{d u}{d t} = G(t, u)
$$
隐式形式：
$$
F(t, u, \dot u) = 0
$$
在TS中，通过RHSFunction, IFunction 区分是哪种形式。

显格式推进：RHSFunction

当问题可以写成 $\dot u = G(t, u)$ 时，可以使用显式/半隐式时间格式，在这种情况下，TS 只需要你提供一个函数，用于计算右端项 $G(t, u)$，其代码为：

ts.setRHSFunction(rhs_function, F)

其中， rhs_function(ts, t, U, F) 包含了：

• 输入：当前时间 t，当前解向量 U
• 输出：右端项 F = G(t, U)

在这种模式下，不需要 SNES 和 IJacobian，每个时间步中只是若干次向量更新和RHS计算。

隐格式推进：IFunction

当时间推进格式是隐式的，TS 在每一个时间步都会形成一个非线性方程，以 Backward Euler 为例：
$$
\frac{u^{n+1} - u^n}{\Delta t} - G(t^{n+1}, u^{n+1}) = 0
$$
这符合隐格式的形式 $F(t, u, \dot u) = 0$，而可以通过下面代码实现：

ts.setIFunction(ifunction, F)

其中 ifunction(ts, t, U, Udot, F) 计算了 $F = F(t, U, \dot U)$。

只要使用了隐式方法，TS 在内部就会调用 SNES，因此就需要 Jacobian，即需要计算IJacobian函数。

对于隐格式 $F(t, u, \dot u) = 0$，其Jacobian实际上是：
$$
J = \frac{\partial F}{\partial u} + \alpha \frac{\partial F}{\partial \dot u}
$$
其中系数 $\alpha$ 由时间格式和步长给出，如 Backward Euler 中，有 $\alpha = \frac{1}{\Delta t}$，代码实现为：

ts.setIJacobian(ijacobian, J)

回调函数ijacobian(ts, t, U, Udot, shift, J, P)中的shift为上述$\alpha$。

TS流程

从执行流程上看，一个隐式 TS 步如下：

1. TS 选择时间步长 $\Delta t$
2. TS 构造当前时间步的非线性问题
3. TS 调用 snes.solve()
4. SNES 在每一步 Newton 迭代中：
   • 调用 IFunction
   • 调用 IJacobian
   • 使用 KSP 解线性系统
5. TS 判断是否接受该时间步，必要时调整步长

可以使用下面的参数列表进行设置内部求解器细节：

-snes_type newtonls
-ksp_type gmres
-pc_type ilu

TS 提供了一套与 KSP / SNES 风格一致的监控机制，可以通过 -ts_monitor 打开，TS会在每一个时间步输出当前时间步的信息。和KSP与SNES类似，也可以通过ts.setMonitor(monitor)设置自定义的monitor，从而记录能量、误差或某个分量的时间演化。

1D Heat equation

考虑1D热传导问题：
$$
u_t = \kappa u_{xx},\quad x\in[0,1]
$$
边界条件选择$u(0,t)=0, u(1,t)=0$，初值为 $u(x,0)=\sin(\pi x)$，其解析解为：
$$
u(x,t)=\exp(-\pi^2 t) \cdot \sin(\pi x)
$$

用 TS 做隐格式时间推进，空间用二阶中心差分，代码如下：

import sys, petsc4py
petsc4py.init(sys.argv)
from petsc4py import PETSc
import numpy as np
import matplotlib.pyplot as plt

comm  = PETSc.COMM_WORLD
rank  = comm.getRank()
opt   = PETSc.Options()
kappa = opt.getReal("kappa", 1.0)
nx    = opt.getInt("nx", 100)
tmax  = opt.getReal("tmax", 1.0)
dt    = opt.getReal("dt", 1e-3)
hx    = 1.0 / (nx - 1)
save_interval = opt.getReal("save_interval", 1.0)

# create 1D DMDA
dm = PETSc.DMDA().create(
    dim=1,
    sizes=(nx,),
    dof=1,
    stencil_width=1,
    stencil_type=PETSc.DMDA.StencilType.STAR,
    comm=comm,
)
dm.setUp()

# create vectors and matrix
F = dm.createGlobalVec()
J = dm.createMatrix()
U = dm.createGlobalVec()
with U as u0:
    (xs, xe), = dm.getRanges()
    u0[:] = np.sin(np.pi * hx * np.arange(xs, xe))  # initial condition: u(x,0) = sin(pi x)


def ifunction(ts, t, Uin, Udot, Fout):
    # compute F(t, U, Udot) = Udot - kappa * u_xx = 0
    local_U = dm.getLocalVec()
    dm.globalToLocal(Uin, local_U)
    
    with dm.getVecArray(local_U) as u, \
         dm.getVecArray(Udot) as udot, \
         dm.getVecArray(Fout) as f:
        
        (xs, xe), = dm.getRanges()
        if xs == 0:
            f[0] = u[0] - 0.0        # u(0, t) = 0
            xs += 1
        if xe == nx:
            f[nx-1] = u[nx-1] - 0.0  # u(1, t) = 0
            xe -= 1
        if xs < xe:
            # u_t = kappa * u_xx
            uxx = (u[xs-1:xe-1] - 2.0 * u[xs:xe] + u[xs+1:xe+1]) / (hx * hx)
            f[xs:xe] = udot[xs:xe] - kappa * uxx

    dm.restoreLocalVec(local_U)


def ijacobian(ts, t, Uin, Udot, shift, J, P):
    # compute: J = dF/dU + shift * dF/dUdot
    #    F = Udot - kappa * u_xx
    # 
    #    F[i] = Udot[i] - kappa*(U[i-1]-2U[i]+U[i+1]) / (hx * hx)
    # => dF[i]/dU[i]   = 2*kappa / (hx * hx)
    # => dF[i]/dU[i+1] =  -kappa / (hx * hx)
    # => dF[i]/dU[i-1] =  -kappa / (hx * hx)
    # 
    #    dF[i]/dUdot[i] = 1
    # => J[i, i]   = dF[i]/dU[i] + shift * dF[i]/dUdot[i]
    # => J[i, i+1] = dF[i]/dU[i+1]
    # => J[i, i-1] = dF[i]/dU[i-1]
    # 
    #    Boundary conditions: F = U - g
    # => dF/dU = 1
    
    J.zeroEntries()
    (xs, xe), = dm.getRanges()
    for i in range(xs, xe):
        if i == 0 or i == nx - 1:
            # d(U - g)/dU = 1
            J.setValue(i, i, 1.0)
        else:
            off = -kappa / (hx * hx)
            diag = shift + 2.0 * kappa / (hx * hx)
            if i - 1 != 0: J.setValue(i, i - 1, off)      # BC
            J.setValue(i, i, diag)
            if i + 1 != nx - 1: J.setValue(i, i + 1, off) # BC
    J.assemble() # P=J, so `P.assemble()` is not needed


u_history = [] # store solution history for output
last_save = -np.inf
def ts_monitor(ts, step, time, u):    
    global u_history, last_save, rank
    flag = time - last_save - save_interval
    if flag >= -1e-8 or step == -1:
        scatter, u_seq = PETSc.Scatter.toZero(u)
        scatter.scatter(u, u_seq, addv=PETSc.InsertMode.INSERT, mode=PETSc.ScatterMode.FORWARD)
        if rank == 0:
            u_arr = u_seq.getArray(readonly=True).copy()
            u_history.append((time, u_arr))
        last_save = time

    if step == -1 and rank == 0 and u_history:
        plt.figure(figsize=(10, 6))
        x_grid = np.linspace(0, 1, nx)
        for (t_val, u_val) in u_history:
            plt.plot(x_grid, u_val, label=f"t={t_val:.2f}")
        plt.xlabel("x")
        plt.ylabel("u")
        plt.legend()
        plt.grid(True)
        plt.savefig("heat.png")
        plt.close()
        PETSc.Sys.Print("Result saved to 'heat.png'")


# Create TS
ts = PETSc.TS().create(comm=comm)
ts.setDM(dm)
ts.setType(ts.Type.BEULER) # Backward Euler
ts.setProblemType(ts.ProblemType.NONLINEAR)
ts.setMonitor(ts_monitor)
ts.setIFunction(ifunction, F)
ts.setIJacobian(ijacobian, J)

ts.setTime(0.0)
ts.setTimeStep(dt)
ts.setMaxTime(tmax)
ts.setMaxSteps(1e9)

ts.setFromOptions()
ts.solve(U)

运行方法为：

mpirun -n 4 python test.py \
  -nx 100 -kappa 0.1 \
  -dt 0.1 -tmax 10.0 -save_interval 1.0 \
  -ts_type beuler -ksp_type cg -pc_type jacobi