#! pde The heat equation in two dimensions
#! equation u_t = u_{xx} + u_{yy}
#! kind parabolic
#! method alternating-direction implicit (Peaceman-Rachford): one tridiagonal solve per grid line, per half-step
#! checks two decaying modes against their exact rates, the measured order, and a time step 128 times the explicit limit
#! | In two dimensions an implicit step couples every point to its four neighbors, and the matrix
#! | is no longer tridiagonal. Peaceman and Rachford's way around it: split the step in half, and
#! | be implicit in one direction at a time.
#! |
#! | First half: implicit in $y$, explicit in $x$. Second half: the other way around. Each half is a
#! | set of independent *one-dimensional* problems -- one tridiagonal solve per grid line -- so a
#! | step costs $O(N)$ for $N$ unknowns, the same as an explicit step. It is still second order in
#! | time and stable for any time step, which the explicit scheme is not: that one needs
#! | $\Delta t \le h^2/4$.
#! |
#! | The second half-step is the first with the axes exchanged, so it is written as
#! | `transpose`, the same sweep, `transpose` back.
#! |
#! | The test field is two modes with different wavenumbers in $x$ and $y$,
#! | $\sin \pi x \sin 2\pi y + \tfrac12 \sin 3\pi x \sin \pi y$, decaying at $5\pi^2$ and $10\pi^2$. A
#! | field symmetric in $x$ and $y$ would not notice the two sweeps being mixed up; this one does.

pi = 3.141592653589793

#| The exact solution at time t on the padded (n+2) x (n+2) grid, boundary ring included.
def exact_field(n, t) {
    h = 1.0 / (n + 1.0)
    u = zeros(n + 2.0, n + 2.0)
    for i to n + 2.0 {
        for j to n + 2.0 {
            x = i * h
            y = j * h
            u[i, j] = exp(-5.0 * pi * pi * t) * sin(pi * x) * sin(2.0 * pi * y) + 0.5 * exp(-10.0 * pi * pi * t) * sin(3.0 * pi * x) * sin(pi * y)
        }
    }
    return u
}

#| Half a step: explicit along axis 0, implicit along axis 1. `r` is dt / h^2.
#| Every interior row is its own tridiagonal system; the boundary rows of each are the identity.
def half_step(u, r) {
    m = len(u)
    lo = zeros(m) - 0.5 * r
    di = zeros(m) + (1.0 + r)
    up = zeros(m) - 0.5 * r
    di[0] = 1.0
    up[0] = 0.0
    lo[m - 1.0] = 0.0
    di[m - 1.0] = 1.0
    rhs = u + 0.5 * r * (roll(u, 1.0, 0.0) - 2.0 * u + roll(u, -1.0, 0.0))
    out = zeros(m, m)
    for i to m - 2.0 {
        row = slab(rhs, i + 1.0)
        row[0] = 0.0
        row[m - 1.0] = 0.0
        setslab(out, i + 1.0, tridiag_solve(lo, di, up, row))
    }
    return out
}

#| One Peaceman-Rachford step.
def adi_step(u, r) {
    half = half_step(u, r)
    return transpose(half_step(transpose(half), r))
}

#| March to t_end and return the largest error against the exact field.
def adi_error(n, steps, t_end) {
    h = 1.0 / (n + 1.0)
    r = (t_end / steps) / (h * h)
    u = exact_field(n, 0.0)
    for s to steps {
        u = adi_step(u, r)
    }
    return take the largest difference between u and exact_field(n, t_end)
}

T_END = 0.02
e1 = adi_error(19.0, 10.0, T_END)
e2 = adi_error(39.0, 20.0, T_END)
order = log(e1 / e2) / log(2.0)
e3 = adi_error(79.0, 40.0, T_END)
amplitude = max(abs(exact_field(79.0, T_END)))
show e1
show e2
show e3
show order
show amplitude

# 4 steps on a 79 x 79 grid: dt / h^2 = 32, against an explicit limit of 1/4
big_r = (T_END / 4.0) * 80.0 * 80.0
big_step_error = adi_error(79.0, 4.0, T_END)
show big_r
show big_step_error

N = 39
H = 1.0 / (N + 1.0)
U = exact_field(N, 0.0)
FIG_STEPS = 20.0   # the figure's run: T_END in this many steps
for s to FIG_STEPS {
    U = adi_step(U, (T_END / FIG_STEPS) / (H * H))
}

# U[i, j] is the temperature at x = i h, y = j h: transposed, x runs across and y up;
# the field is signed, so the colors are centered on zero
figure("Two modes after t = 0.02: the faster one has nearly gone", "x", "y")
origin("lower")
colormap("diverging")
plot2d(transpose(U), "u")

test "both modes decay at their exact rates: the worst error is under a thousandth of the field" {
    assert_lt(e3, 1.0e-3 * amplitude)
}
test "halving h and the time step together cuts the error by four: second order in both" {
    assert_near(order, 2.0, 0.1)
}
test "a time step 128 times the explicit limit is still stable, and still accurate to a percent of the field" {
    assert_gt(big_r, 32.0 - 1.0e-9)
    assert_lt(big_step_error, 1.0e-2 * amplitude)
}

# SPDX-License-Identifier: MIT OR Apache-2.0
# SPDX-FileCopyrightText: 2026 The Viv authors