Training a PINN on 2D PDE

In this tutorial we will go over using a PINN to solve 2D PDEs. We will be using the system from NeuralPDE Tutorials. However, we will be using our custom loss function and use nested AD capabilities of Lux.jl.

This is a demonstration of Lux.jl. For serious use cases of PINNs, please refer to the package: NeuralPDE.jl.

Package Imports

using Lux,
    Optimisers,
    Random,
    Printf,
    Statistics,
    MLUtils,
    OnlineStats,
    CairoMakie,
    Reactant,
    Enzyme

const xdev = reactant_device(; force=true)
const cdev = cpu_device()

Problem Definition

Since Lux supports efficient nested AD upto 2nd order, we will rewrite the problem with first order derivatives, so that we can compute the gradients of the loss using 2nd order AD.

Define the Neural Networks

All the networks take 3 input variables and output a scalar value. Here, we will define a wrapper over the 3 networks, so that we can train them using Training.TrainState.

struct PINN{M} <: AbstractLuxWrapperLayer{:model}
    model::M
end

function PINN(; hidden_dims::Int=32)
    return PINN(
        Chain(
            Dense(3 => hidden_dims, tanh),
            Dense(hidden_dims => hidden_dims, tanh),
            Dense(hidden_dims => hidden_dims, tanh),
            Dense(hidden_dims => 1),
        ),
    )
end

Define the Loss Functions

We will define a custom loss function to compute the loss using 2nd order AD. For that, first we'll need to define the derivatives of our model:

function ∂u_∂t(model::StatefulLuxLayer, xyt::AbstractArray)
    return Enzyme.gradient(Enzyme.Reverse, sum ∘ model, xyt)[1][3, :]
end

function ∂u_∂x(model::StatefulLuxLayer, xyt::AbstractArray)
    return Enzyme.gradient(Enzyme.Reverse, sum ∘ model, xyt)[1][1, :]
end

function ∂u_∂y(model::StatefulLuxLayer, xyt::AbstractArray)
    return Enzyme.gradient(Enzyme.Reverse, sum ∘ model, xyt)[1][2, :]
end

function ∂²u_∂x²(model::StatefulLuxLayer, xyt::AbstractArray)
    return Enzyme.gradient(Enzyme.Reverse, sum ∘ ∂u_∂x, Enzyme.Const(model), xyt)[2][1, :]
end

function ∂²u_∂y²(model::StatefulLuxLayer, xyt::AbstractArray)
    return Enzyme.gradient(Enzyme.Reverse, sum ∘ ∂u_∂y, Enzyme.Const(model), xyt)[2][2, :]
end

We will use the following loss function

function physics_informed_loss_function(model::StatefulLuxLayer, xyt::AbstractArray)
    return mean(abs2, ∂u_∂t(model, xyt) .- ∂²u_∂x²(model, xyt) .- ∂²u_∂y²(model, xyt))
end

Additionally, we need to compute the loss with respect to the boundary conditions.

function mse_loss_function(
    model::StatefulLuxLayer, target::AbstractArray, xyt::AbstractArray
)
    return MSELoss()(model(xyt), target)
end

function loss_function(model, ps, st, (xyt, target_data, xyt_bc, target_bc))
    smodel = StatefulLuxLayer(model, ps, st)
    physics_loss = physics_informed_loss_function(smodel, xyt)
    data_loss = mse_loss_function(smodel, target_data, xyt)
    bc_loss = mse_loss_function(smodel, target_bc, xyt_bc)
    loss = physics_loss + data_loss + bc_loss
    return loss, smodel.st, (; physics_loss, data_loss, bc_loss)
end

Generate the Data

We will generate some random data to train the model on. We will take data on a square spatial and temporal domain $x\in [0,2]$, $y\in [0,2]$, and $t\in [0,2]$. Typically, you want to be smarter about the sampling process, but for the sake of simplicity, we will skip that.

analytical_solution(x, y, t) = @. exp(x + y) * cos(x + y + 4t)
analytical_solution(xyt) = analytical_solution(xyt[1, :], xyt[2, :], xyt[3, :])

grid_len = 16

grid = range(0.0f0, 2.0f0; length=grid_len)
xyt = stack([[elem...] for elem in vec(collect(Iterators.product(grid, grid, grid)))])

target_data = reshape(analytical_solution(xyt), 1, :)

bc_len = 512

x = collect(range(0.0f0, 2.0f0; length=bc_len))
y = collect(range(0.0f0, 2.0f0; length=bc_len))
t = collect(range(0.0f0, 2.0f0; length=bc_len))

xyt_bc = hcat(
    stack((x, y, zeros(Float32, bc_len)); dims=1),
    stack((zeros(Float32, bc_len), y, t); dims=1),
    stack((ones(Float32, bc_len) .* 2, y, t); dims=1),
    stack((x, zeros(Float32, bc_len), t); dims=1),
    stack((x, ones(Float32, bc_len) .* 2, t); dims=1),
)
target_bc = reshape(analytical_solution(xyt_bc), 1, :)

min_target_bc, max_target_bc = extrema(target_bc)
min_data, max_data = extrema(target_data)
min_pde_val, max_pde_val = min(min_data, min_target_bc), max(max_data, max_target_bc)

xyt = (xyt .- minimum(xyt)) ./ (maximum(xyt) .- minimum(xyt))
xyt_bc = (xyt_bc .- minimum(xyt_bc)) ./ (maximum(xyt_bc) .- minimum(xyt_bc))
target_bc = (target_bc .- min_pde_val) ./ (max_pde_val - min_pde_val)
target_data = (target_data .- min_pde_val) ./ (max_pde_val - min_pde_val)

Training

function train_model(
    xyt,
    target_data,
    xyt_bc,
    target_bc;
    seed::Int=0,
    maxiters::Int=50000,
    hidden_dims::Int=128,
)
    rng = Random.default_rng()
    Random.seed!(rng, seed)

    pinn = PINN(; hidden_dims)
    ps, st = Lux.setup(rng, pinn) |> xdev

    bc_dataloader =
        DataLoader((xyt_bc, target_bc); batchsize=128, shuffle=true, partial=false) |> xdev
    pde_dataloader =
        DataLoader((xyt, target_data); batchsize=128, shuffle=true, partial=false) |> xdev

    train_state = Training.TrainState(pinn, ps, st, Adam(0.005f0))

    lr = i -> i < 5000 ? 0.005f0 : (i < 10000 ? 0.0005f0 : 0.00005f0)

    total_loss_tracker, physics_loss_tracker, data_loss_tracker, bc_loss_tracker = ntuple(
        _ -> OnlineStats.CircBuff(Float32, 32; rev=true), 4
    )

    iter = 1
    for ((xyt_batch, target_data_batch), (xyt_bc_batch, target_bc_batch)) in
        zip(Iterators.cycle(pde_dataloader), Iterators.cycle(bc_dataloader))
        Optimisers.adjust!(train_state, lr(iter))

        _, loss, stats, train_state = Training.single_train_step!(
            AutoEnzyme(),
            loss_function,
            (xyt_batch, target_data_batch, xyt_bc_batch, target_bc_batch),
            train_state;
            return_gradients=Val(false),
        )

        fit!(total_loss_tracker, Float32(loss))
        fit!(physics_loss_tracker, Float32(stats.physics_loss))
        fit!(data_loss_tracker, Float32(stats.data_loss))
        fit!(bc_loss_tracker, Float32(stats.bc_loss))

        mean_loss = mean(OnlineStats.value(total_loss_tracker))
        mean_physics_loss = mean(OnlineStats.value(physics_loss_tracker))
        mean_data_loss = mean(OnlineStats.value(data_loss_tracker))
        mean_bc_loss = mean(OnlineStats.value(bc_loss_tracker))

        isnan(loss) && throw(ArgumentError("NaN Loss Detected"))

        if iter % 1000 == 1 || iter == maxiters
            @printf(
                "Iteration: [%6d/%6d] \t Loss: %.9f (%.9f) \t Physics Loss: %.9f \
                 (%.9f) \t Data Loss: %.9f (%.9f) \t BC \
                 Loss: %.9f (%.9f)\n",
                iter,
                maxiters,
                loss,
                mean_loss,
                stats.physics_loss,
                mean_physics_loss,
                stats.data_loss,
                mean_data_loss,
                stats.bc_loss,
                mean_bc_loss
            )
        end

        iter += 1
        iter ≥ maxiters && break
    end

    return StatefulLuxLayer(pinn, cdev(train_state.parameters), cdev(train_state.states))
end

trained_model = train_model(xyt, target_data, xyt_bc, target_bc)

Iteration: [     1/ 50000] 	 Loss: 20.523933411 (20.523933411) 	 Physics Loss: 16.931318283 (16.931318283) 	 Data Loss: 2.007483244 (2.007483244) 	 BC Loss: 1.585133195 (1.585133195)
Iteration: [  1001/ 50000] 	 Loss: 0.017368639 (0.019241152) 	 Physics Loss: 0.000384365 (0.000523638) 	 Data Loss: 0.005318499 (0.007538576) 	 BC Loss: 0.011665775 (0.011178942)
Iteration: [  2001/ 50000] 	 Loss: 0.015431636 (0.018665709) 	 Physics Loss: 0.001248549 (0.001662073) 	 Data Loss: 0.004322524 (0.006408242) 	 BC Loss: 0.009860563 (0.010595393)
Iteration: [  3001/ 50000] 	 Loss: 0.015749739 (0.015216020) 	 Physics Loss: 0.000569920 (0.001279077) 	 Data Loss: 0.004014884 (0.004232429) 	 BC Loss: 0.011164935 (0.009704513)
Iteration: [  4001/ 50000] 	 Loss: 0.009719806 (0.008716607) 	 Physics Loss: 0.002388194 (0.003383097) 	 Data Loss: 0.003175713 (0.002104781) 	 BC Loss: 0.004155898 (0.003228731)
Iteration: [  5001/ 50000] 	 Loss: 0.005263248 (0.006891137) 	 Physics Loss: 0.001360453 (0.002432577) 	 Data Loss: 0.002217358 (0.001982781) 	 BC Loss: 0.001685437 (0.002475780)
Iteration: [  6001/ 50000] 	 Loss: 0.001140392 (0.001318317) 	 Physics Loss: 0.000323426 (0.000313914) 	 Data Loss: 0.000575230 (0.000728794) 	 BC Loss: 0.000241736 (0.000275610)
Iteration: [  7001/ 50000] 	 Loss: 0.001363131 (0.001061225) 	 Physics Loss: 0.000297997 (0.000400512) 	 Data Loss: 0.000963937 (0.000523300) 	 BC Loss: 0.000101198 (0.000137412)
Iteration: [  8001/ 50000] 	 Loss: 0.000873748 (0.000750067) 	 Physics Loss: 0.000501427 (0.000265209) 	 Data Loss: 0.000304499 (0.000403834) 	 BC Loss: 0.000067822 (0.000081024)
Iteration: [  9001/ 50000] 	 Loss: 0.001739428 (0.001934528) 	 Physics Loss: 0.000829241 (0.001233693) 	 Data Loss: 0.000761855 (0.000475479) 	 BC Loss: 0.000148332 (0.000225356)
Iteration: [ 10001/ 50000] 	 Loss: 0.000633724 (0.000649034) 	 Physics Loss: 0.000274444 (0.000288766) 	 Data Loss: 0.000306615 (0.000308623) 	 BC Loss: 0.000052666 (0.000051645)
Iteration: [ 11001/ 50000] 	 Loss: 0.000386282 (0.000382222) 	 Physics Loss: 0.000160045 (0.000067849) 	 Data Loss: 0.000185642 (0.000279561) 	 BC Loss: 0.000040596 (0.000034812)
Iteration: [ 12001/ 50000] 	 Loss: 0.000273355 (0.000358341) 	 Physics Loss: 0.000058881 (0.000067159) 	 Data Loss: 0.000175104 (0.000256167) 	 BC Loss: 0.000039371 (0.000035016)
Iteration: [ 13001/ 50000] 	 Loss: 0.000312572 (0.000335592) 	 Physics Loss: 0.000064870 (0.000069724) 	 Data Loss: 0.000214451 (0.000233520) 	 BC Loss: 0.000033251 (0.000032348)
Iteration: [ 14001/ 50000] 	 Loss: 0.000406529 (0.000346434) 	 Physics Loss: 0.000077713 (0.000076441) 	 Data Loss: 0.000297825 (0.000242032) 	 BC Loss: 0.000030991 (0.000027961)
Iteration: [ 15001/ 50000] 	 Loss: 0.000240908 (0.000292884) 	 Physics Loss: 0.000039256 (0.000056216) 	 Data Loss: 0.000173427 (0.000205937) 	 BC Loss: 0.000028225 (0.000030731)
Iteration: [ 16001/ 50000] 	 Loss: 0.000221967 (0.000299261) 	 Physics Loss: 0.000045093 (0.000063268) 	 Data Loss: 0.000144775 (0.000206898) 	 BC Loss: 0.000032098 (0.000029095)
Iteration: [ 17001/ 50000] 	 Loss: 0.000429942 (0.000302526) 	 Physics Loss: 0.000094877 (0.000067734) 	 Data Loss: 0.000310956 (0.000207792) 	 BC Loss: 0.000024110 (0.000027000)
Iteration: [ 18001/ 50000] 	 Loss: 0.000221261 (0.000289581) 	 Physics Loss: 0.000045943 (0.000061020) 	 Data Loss: 0.000142210 (0.000200792) 	 BC Loss: 0.000033109 (0.000027769)
Iteration: [ 19001/ 50000] 	 Loss: 0.000221184 (0.000282269) 	 Physics Loss: 0.000063333 (0.000057770) 	 Data Loss: 0.000139595 (0.000200582) 	 BC Loss: 0.000018256 (0.000023917)
Iteration: [ 20001/ 50000] 	 Loss: 0.000300043 (0.000259325) 	 Physics Loss: 0.000047858 (0.000051022) 	 Data Loss: 0.000233215 (0.000184559) 	 BC Loss: 0.000018970 (0.000023743)
Iteration: [ 21001/ 50000] 	 Loss: 0.000304431 (0.000253003) 	 Physics Loss: 0.000058605 (0.000055728) 	 Data Loss: 0.000223942 (0.000173880) 	 BC Loss: 0.000021883 (0.000023395)
Iteration: [ 22001/ 50000] 	 Loss: 0.000167038 (0.000236288) 	 Physics Loss: 0.000029658 (0.000044059) 	 Data Loss: 0.000110797 (0.000170320) 	 BC Loss: 0.000026583 (0.000021910)
Iteration: [ 23001/ 50000] 	 Loss: 0.000226478 (0.000254691) 	 Physics Loss: 0.000034662 (0.000053648) 	 Data Loss: 0.000170027 (0.000178980) 	 BC Loss: 0.000021789 (0.000022064)
Iteration: [ 24001/ 50000] 	 Loss: 0.000316798 (0.000266476) 	 Physics Loss: 0.000068699 (0.000066841) 	 Data Loss: 0.000227380 (0.000172998) 	 BC Loss: 0.000020719 (0.000026637)
Iteration: [ 25001/ 50000] 	 Loss: 0.000205806 (0.000228483) 	 Physics Loss: 0.000041027 (0.000037151) 	 Data Loss: 0.000145431 (0.000170033) 	 BC Loss: 0.000019348 (0.000021299)
Iteration: [ 26001/ 50000] 	 Loss: 0.000247680 (0.000264209) 	 Physics Loss: 0.000062845 (0.000073178) 	 Data Loss: 0.000161656 (0.000165484) 	 BC Loss: 0.000023179 (0.000025548)
Iteration: [ 27001/ 50000] 	 Loss: 0.000220185 (0.000241178) 	 Physics Loss: 0.000040774 (0.000052417) 	 Data Loss: 0.000152581 (0.000166542) 	 BC Loss: 0.000026830 (0.000022219)
Iteration: [ 28001/ 50000] 	 Loss: 0.000231527 (0.000230345) 	 Physics Loss: 0.000060593 (0.000049008) 	 Data Loss: 0.000153760 (0.000158954) 	 BC Loss: 0.000017174 (0.000022382)
Iteration: [ 29001/ 50000] 	 Loss: 0.000239694 (0.000235708) 	 Physics Loss: 0.000073950 (0.000056235) 	 Data Loss: 0.000137358 (0.000158777) 	 BC Loss: 0.000028386 (0.000020696)
Iteration: [ 30001/ 50000] 	 Loss: 0.000203338 (0.000223219) 	 Physics Loss: 0.000028516 (0.000040052) 	 Data Loss: 0.000154368 (0.000162479) 	 BC Loss: 0.000020454 (0.000020688)
Iteration: [ 31001/ 50000] 	 Loss: 0.000279015 (0.000222952) 	 Physics Loss: 0.000043878 (0.000042610) 	 Data Loss: 0.000213876 (0.000159681) 	 BC Loss: 0.000021261 (0.000020662)
Iteration: [ 32001/ 50000] 	 Loss: 0.000217566 (0.000204443) 	 Physics Loss: 0.000044902 (0.000031157) 	 Data Loss: 0.000150807 (0.000153447) 	 BC Loss: 0.000021857 (0.000019839)
Iteration: [ 33001/ 50000] 	 Loss: 0.000176795 (0.000204427) 	 Physics Loss: 0.000021615 (0.000034391) 	 Data Loss: 0.000131915 (0.000149062) 	 BC Loss: 0.000023265 (0.000020975)
Iteration: [ 34001/ 50000] 	 Loss: 0.000194323 (0.000201223) 	 Physics Loss: 0.000031174 (0.000035442) 	 Data Loss: 0.000142124 (0.000145737) 	 BC Loss: 0.000021025 (0.000020043)
Iteration: [ 35001/ 50000] 	 Loss: 0.000132700 (0.000211725) 	 Physics Loss: 0.000015713 (0.000044191) 	 Data Loss: 0.000095944 (0.000148370) 	 BC Loss: 0.000021043 (0.000019164)
Iteration: [ 36001/ 50000] 	 Loss: 0.000164996 (0.000200748) 	 Physics Loss: 0.000028600 (0.000034721) 	 Data Loss: 0.000116301 (0.000147227) 	 BC Loss: 0.000020095 (0.000018800)
Iteration: [ 37001/ 50000] 	 Loss: 0.000304521 (0.000191764) 	 Physics Loss: 0.000067376 (0.000026386) 	 Data Loss: 0.000218824 (0.000146148) 	 BC Loss: 0.000018321 (0.000019230)
Iteration: [ 38001/ 50000] 	 Loss: 0.000242138 (0.000197651) 	 Physics Loss: 0.000025409 (0.000028536) 	 Data Loss: 0.000192619 (0.000151911) 	 BC Loss: 0.000024111 (0.000017204)
Iteration: [ 39001/ 50000] 	 Loss: 0.000167793 (0.000194161) 	 Physics Loss: 0.000026857 (0.000029151) 	 Data Loss: 0.000124928 (0.000146219) 	 BC Loss: 0.000016008 (0.000018790)
Iteration: [ 40001/ 50000] 	 Loss: 0.000168046 (0.000191795) 	 Physics Loss: 0.000023108 (0.000027603) 	 Data Loss: 0.000125795 (0.000145107) 	 BC Loss: 0.000019142 (0.000019085)
Iteration: [ 41001/ 50000] 	 Loss: 0.000156071 (0.000189268) 	 Physics Loss: 0.000019337 (0.000024171) 	 Data Loss: 0.000120357 (0.000147406) 	 BC Loss: 0.000016377 (0.000017691)
Iteration: [ 42001/ 50000] 	 Loss: 0.000180218 (0.000193830) 	 Physics Loss: 0.000026169 (0.000027747) 	 Data Loss: 0.000138332 (0.000146765) 	 BC Loss: 0.000015716 (0.000019317)
Iteration: [ 43001/ 50000] 	 Loss: 0.000189987 (0.000193514) 	 Physics Loss: 0.000023531 (0.000027774) 	 Data Loss: 0.000151968 (0.000145500) 	 BC Loss: 0.000014488 (0.000020240)
Iteration: [ 44001/ 50000] 	 Loss: 0.000183824 (0.000215055) 	 Physics Loss: 0.000025106 (0.000047723) 	 Data Loss: 0.000142475 (0.000144129) 	 BC Loss: 0.000016243 (0.000023203)
Iteration: [ 45001/ 50000] 	 Loss: 0.000273151 (0.000213490) 	 Physics Loss: 0.000042826 (0.000038736) 	 Data Loss: 0.000204913 (0.000152898) 	 BC Loss: 0.000025412 (0.000021856)
Iteration: [ 46001/ 50000] 	 Loss: 0.000200716 (0.000179186) 	 Physics Loss: 0.000024293 (0.000022313) 	 Data Loss: 0.000157517 (0.000139066) 	 BC Loss: 0.000018906 (0.000017807)
Iteration: [ 47001/ 50000] 	 Loss: 0.000173622 (0.000179886) 	 Physics Loss: 0.000017607 (0.000022635) 	 Data Loss: 0.000137739 (0.000139571) 	 BC Loss: 0.000018277 (0.000017679)
Iteration: [ 48001/ 50000] 	 Loss: 0.000183665 (0.000180370) 	 Physics Loss: 0.000022360 (0.000023689) 	 Data Loss: 0.000142342 (0.000138544) 	 BC Loss: 0.000018963 (0.000018137)
Iteration: [ 49001/ 50000] 	 Loss: 0.000177444 (0.000207296) 	 Physics Loss: 0.000032091 (0.000046293) 	 Data Loss: 0.000125603 (0.000142999) 	 BC Loss: 0.000019750 (0.000018005)

Visualizing the Results

ts, xs, ys = 0.0f0:0.05f0:2.0f0, 0.0f0:0.02f0:2.0f0, 0.0f0:0.02f0:2.0f0
grid = stack([[elem...] for elem in vec(collect(Iterators.product(xs, ys, ts)))])

u_real = reshape(analytical_solution(grid), length(xs), length(ys), length(ts))

grid_normalized = (grid .- minimum(grid)) ./ (maximum(grid) .- minimum(grid))
u_pred = reshape(trained_model(grid_normalized), length(xs), length(ys), length(ts))
u_pred = u_pred .* (max_pde_val - min_pde_val) .+ min_pde_val

begin
    fig = Figure()
    ax = CairoMakie.Axis(fig[1, 1]; xlabel="x", ylabel="y")
    errs = [abs.(u_pred[:, :, i] .- u_real[:, :, i]) for i in 1:length(ts)]
    Colorbar(fig[1, 2]; limits=extrema(stack(errs)))

    CairoMakie.record(fig, "pinn_nested_ad.gif", 1:length(ts); framerate=10) do i
        ax.title = "Abs. Predictor Error | Time: $(ts[i])"
        err = errs[i]
        contour!(ax, xs, ys, err; levels=10, linewidth=2)
        heatmap!(ax, xs, ys, err)
        return fig
    end

    fig
end

Appendix

using InteractiveUtils
InteractiveUtils.versioninfo()

if @isdefined(MLDataDevices)
    if @isdefined(CUDA) && MLDataDevices.functional(CUDADevice)
        println()
        CUDA.versioninfo()
    end

    if @isdefined(AMDGPU) && MLDataDevices.functional(AMDGPUDevice)
        println()
        AMDGPU.versioninfo()
    end
end

Julia Version 1.11.7
Commit f2b3dbda30a (2025-09-08 12:10 UTC)
Build Info:
  Official https://julialang.org/ release
Platform Info:
  OS: Linux (x86_64-linux-gnu)
  CPU: 4 × AMD EPYC 7763 64-Core Processor
  WORD_SIZE: 64
  LLVM: libLLVM-16.0.6 (ORCJIT, znver3)
Threads: 4 default, 0 interactive, 2 GC (on 4 virtual cores)
Environment:
  JULIA_DEBUG = Literate
  LD_LIBRARY_PATH = 
  JULIA_NUM_THREADS = 4
  JULIA_CPU_HARD_MEMORY_LIMIT = 100%
  JULIA_PKG_PRECOMPILE_AUTO = 0

---

This page was generated using Literate.jl.