RasmahRasmah

Method of Lines: Space-Discretized PDEs and Their Gradients

Overview

Time-dependent partial differential equations (PDEs) — heat conduction, diffusion, reaction–diffusion — are often solved by the method of lines: discretize space with the finite-element method, turning the PDE into a large system of ordinary differential equations (ODEs), then march that system in time. Because the resulting ODE is just a function, you can then differentiate the solution with respect to the PDE’s inputs — the strength of a forcing term and, through the mesh and the assembled operators, the geometry itself.

This tutorial does exactly that for the heat equation. We build the semi-discrete ODE with Rasmah’s weak-form FEM, solve it with solve_ode, and use ode_gradient to pull gradients of the temperature field with respect to a heat source and the rod length — verifying every derivative against finite differences.

The heat equation

A rod of length L conducts heat with diffusivity α and a uniform volumetric source q. The temperature u(x, t) obeys

\[ \frac{\partial u}{\partial t} = \alpha \frac{\partial^2 u}{\partial x^2} + q, \qquad u(0,t) = u(L,t) = 0. \]

The steady state ∂u/∂t = 0 is the parabola u(x) = q x(L-x)/(2α), whose center value is q L²/(8α) — a clean reference we can check against.

Discretizing space

Space is discretized with continuous piecewise-linear finite elements on a thin box (one element thick in y and z, so the problem is effectively 1-D). The weak form gives a mass matrix M = ∫ φᵢφⱼ and a stiffness matrix K = α ∫ ∇φᵢ·∇φⱼ, assembled through the weak-form layer. The boundary nodes (x = 0, x = L) are fixed to 0.

using Rasmah
using LinearAlgebra

α = 1.0        # diffusivity
N = 10         # elements along the rod

# The semi-discrete operators as functions of the geometry L.
function operators(L)
    mesh = tetrahedralize_box(L, 1.0, 1.0, N, 1, 1)
    U = FESpace(mesh, ReferenceFE(3, 1))
    M = assemble_mass(U)
    K = assemble_matrix(heat_form(U, α))
    fixed = Int[]
    for i in 1:size(mesh.nodes, 2)
        x = mesh.nodes[1, i]
        (x < 1e-12 || x > L - 1e-12) && push!(fixed, i)
    end
    free = setdiff(1:size(mesh.nodes, 2), fixed)
    lump = vec(sum(M; dims = 2))[free]      # row-sum (lumped) mass
    return K[free, free], lump, free, mesh
end
operators (generic function with 1 method)

Mass lumping (replacing M by its diagonal row sums) is the standard choice for explicit time stepping: it keeps the inverse-mass operator local and positive, so the explicit integrator stays stable at a reasonable step size.

The semi-discrete ODE

Eliminating the boundary nodes turns the PDE into the ODE

\[ \frac{du}{dt} = -D^{-1} K\, u + q, \]

where D is the lumped mass and u now lives on the interior nodes. The ODE vector field closes over the geometry L (it re-assembles K and D from the mesh) and takes the parameters p = [q, L]:

function rhs(u, p, t)
    K_r, lump = operators(p[2])
    -(K_r * u) ./ lump .+ p[1] .* ones(length(u))
end
rhs (generic function with 1 method)

Solve it and check the steady-state centre temperature:

q, L = 2.0, 1.0
u0 = zeros(length(operators(L)[3]))
u_end = solve_ode(rhs, u0, (0.0, 2.0), [q, L]; method = RK4(), steps = 1000).u[end]

_, _, free, mesh = operators(L)
c = argmin(abs.(mesh.nodes[1, free] .- L / 2))
println("centre temperature u(L/2) = ", round(u_end[c]; digits = 4),
        "   (exact q L²/(8α) = ", q * L^2 / (8α), ")")
centre temperature u(L/2) = 0.25   (exact q L²/(8α) = 0.25)

The explicit RK4 step is chosen small enough to satisfy the diffusion CFL condition. Colour the whole mesh by the recovered temperature field (boundary nodes are 0):

T = zeros(size(mesh.nodes, 2))
T[free] .= u_end
vtk(mesh; field = T, fieldname = "temperature", scalarbar = true)

Gradients: forcing and geometry

We differentiate the total heat Σᵢ uᵢ(T) through the ODE — and therefore through the mesh, the mass/stiffness assembly, and the time integrator. p = [q, L] carries both the forcing strength and the geometry, so ode_gradient returns both partial derivatives at once:

dL_du = ones(length(u0))   # L = Σᵢ uᵢ(T)
for sensealg in (DiscreteAdjoint(), ContinuousAdjoint())
    du0, dp = ode_gradient(rhs, u0, (0.0, 2.0), [q, L], dL_du, nothing;
                           sensealg = sensealg, method = RK4(), steps = 1000)
    println(nameof(typeof(sensealg)), ":  ∂(Σu)/∂q = ", round(dp[1]; digits = 4),
            ",  ∂(Σu)/∂L = ", round(dp[2]; digits = 4))
end
DiscreteAdjoint:  ∂(Σu)/∂q = 3.3,  ∂(Σu)/∂L = 13.2
ContinuousAdjoint:  ∂(Σu)/∂q = 3.3,  ∂(Σu)/∂L = 13.2

Cross-check against forward-mode AD through the entire pipeline (mesh → FEM assembly → ODE solve → loss):

loss(p) = sum(solve_ode(rhs, u0, (0.0, 2.0), p; method = RK4(), steps = 1000).u[end])
g = gradient(loss, [q, L])
println("forward-mode AD:  ∂(Σu)/∂q = ", round(g[1]; digits = 4),
        ",  ∂(Σu)/∂L = ", round(g[2]; digits = 4))
forward-mode AD:  ∂(Σu)/∂q = 3.3,  ∂(Σu)/∂L = 13.2

The three agree. Physically, ∂(Σu)/∂q > 0 (a stronger source heats the rod more) and ∂(Σu)/∂L > 0 (a longer rod holds more heat in the parabolic steady state) — the gradient has already flowed back through the geometry.

Next steps