Time-dependent problems

TimeDependentPDEVariationalProblem solves an initial-condition or parameter inversion for a time-dependent PDE. The residual density takes four fields: the new time level, the previous one, the parameter and the adjoint.

def varf(u_new, u_old, m, p, x, t, dt):
    theta = 0.5                                   # Crank-Nicolson
    u = theta * u_new.val + (1 - theta) * u_old.val
    g = theta * u_new.grad + (1 - theta) * u_old.grad
    return ((u_new.val - u_old.val) / dt * p.val
            + jnp.exp(m.val) * jnp.dot(g, p.grad))

pde = hm.TimeDependentPDEVariationalProblem(
    [Vu, Vm, Vu], varf, bc, bc0, u0, t_init=0.0, t_final=1.0, dt=0.01)

where u0 is the initial state, a vector of Vu.

The initial state

The initial state may instead depend on the parameter through a weak equation of its own, given by an initial density init_varf(u0, m, p, x) in place of u0:

def init_varf(u0, m, p, x):                       # u(0) = m
    return (u0.val - m.val) * p.val

pde = hm.TimeDependentPDEVariationalProblem(
    [Vu, Vu, Vu], varf, bc, bc0, None, t_init=0.0, t_final=1.0, dt=0.01,
    init_varf=init_varf)

A stationary problem of its own, pde.init_problem, solves the initial equation with the state’s boundary conditions and differentiates it like the steps; its terms join theirs in the forward and adjoint marches, the gradient and every block of the Hessian. The parameter may therefore set the initial state through any map that JAX differentiates twice, (u0.val - g(m.val)) * p.val; init_linear=False declares a density that is nonlinear in u0, which Newton’s method then solves. Set the solvers of pde.init_problem as for any stationary problem. applications/ad_diff runs its initial-condition inversion with a hand-written wrapper or, with --initial-density, with this density. With essential conditions the two differ at the boundary: the wrapper zeroes the parameter’s boundary values, while the projection lets them reach the rows next to the boundary.

Why the previous level is a differentiation slot

The adjoint of a time-stepping scheme needs \(\partial r_{n+1}/\partial u_n\): the derivative of one step’s residual with respect to the previous state. Treating the old level as a frozen coefficient would force a finite difference there and make the gradient inexact. Making it a real slot keeps every block exact, and the stationary problem is then just the three-slot case.

The per-step blocks are A, B (the cross-time coupling), C, and the second-order blocks W_nn, W_no, W_oo, W_nm, W_om, W_mm. applyWuu accumulates into both time levels, which is what makes the time-dependent reduced Hessian symmetric; the test suite checks that W_um and W_mu are exact transposes and that the Hessian is symmetric to 1e-15 on a quasilinear problem where all four cross-time blocks are nonzero.

Operators along the trajectory

Nothing is assembled, and no multigrid hierarchy set up, inside a solve loop. The step operators of a trajectory (A_n, its transpose and B_n) are built once when the adjoint is first solved at that state and parameter, kept while the point does not change, and completed with C_n and the W blocks by setLinearizationPoint; every block of a step comes from one differentiation pass. The forward solve of a residual that is linear in the new level keeps its Jacobian across calls at the same parameter.

The common case of a step Jacobian that does not change in time (a linear PDE with coefficients independent of t and a fixed dt) is detected rather than assumed: the operators of every step are compared with those of the step before through their action on a random vector, and steps that agree share one matrix, one transpose and one solver with its hierarchy. When all steps agree, one of each serves the whole march. A step that differs gets its own, cloned from the solver you set (solver_fwd_inc, solver_adj_inc; solver_adj defaults to a clone of solver). On the 16-step advection-diffusion test problem this takes a forward or adjoint solve from 16 BoomerAMG setups to one, and a Hessian action from 32 to none. The comparison needs the operators of every step, because equal first steps say nothing about the later ones (a state at rest until a source switches on, a coefficient that changes in time): they are assembled at every step of a trajectory, and at every step of a linear forward solve at a new parameter. The comparison tolerance is TimeDependentPDEVariationalProblem.INVARIANCE_PROBE_TOL (1e-13); -1 treats every step as different.

Scalar parameters

t and dt arrive as traced values rather than constants, so changing the step size between steps reuses the compiled kernel instead of triggering a recompile. Declare them with nparams when building a kernel directly.

Time-dependent vectors

TimeDependentVector holds a snapshot per time level with the arithmetic the algorithms need, and MisfitTD compares snapshots against data at chosen times.

Note

apply_ijk, third derivatives, is not implemented for the time-dependent problem, because Newton-CG does not use it. A time-dependent Taylor approximation of a QoI would need it.

The time-dependent problem takes a domain density and a boundary density, bdr_varf(u, u_old, m, p, x, n, t, dt) over the boundary elements of bdr_attributes, which is where a Robin condition or a prescribed flux goes (test_timedependent.py checks a Robin term against MFEM’s boundary mass matrix and through an inversion). It takes no facet_varf: DG formulations are for the stationary class.