RasmahRasmah

Operator Learning

println("Julia ", VERSION)
Julia 1.12.6

Overview

A classical PDE solver maps one input to one solution. Operator learning instead learns the map itself — the operator G : a ↦ u that takes an input function (a coefficient field, a boundary condition, a source) straight to the solution. Once trained, a neural operator evaluates new inputs in a single forward pass, no solve required.

Rasmah provides the two families of neural operators from the literature: the Deep operator network (DeepONet), which encodes the input function at sensor points and the query point separately, and the Fourier neural operator (FNO), which mixes frequencies with learned spectral filters.

using Rasmah
using Random

DeepONet: the antiderivative operator

A DeepONet approximates G(u)(y) ≈ ⟨branch(u), trunk(y)⟩ + b. We teach it the antiderivative operator G(u)(y) = ∫₀ʸ u(s) ds, whose exact value for the input u(s) = a sin(s) is a(1 − cos y). Training pairs an input function (sampled at fixed sensor points) with a query point and the target value.

rng = MersenneTwister(4)
m = 12
sensors = collect(range(0.0, π; length = m))
branch = MLP([m, 32, 32])
trunk = MLP([1, 32, 32])
op = DeepONet(branch, trunk)
θ0 = deeponet_init(rng, branch.sizes, trunk.sizes)

N = 300
Us = zeros(m, N); Ys = zeros(1, N); ys = zeros(N)
for i in 1:N
    a = 0.5 + rand(rng); y = rand(rng) * π
    Us[:, i] = a .* sin.(sensors)
    Ys[1, i] = y
    ys[i] = a * (1 - cos(y))
end
θ = deeponet_fit(op, θ0, Us, Ys, ys; epochs = 800, lr = 2e-3, seed = 4)
2593-element Vector{Float64}:
 -0.062580881733357
  0.17521063082901525
 -0.2533196875365946
  0.04571569258007756
 -0.26434700373406617
 -0.13088575571200678
 -0.2472998058163845
 -0.006866751618390729
  0.17621900818964964
  0.16696273802557435
  ⋮
  0.011688303817861512
  0.01852118486314324
 -0.029728842113223235
  0.011149498538082907
 -0.02188392224050907
 -0.00533282843181668
  0.010720002509735753
 -0.005067793366333773
  0.009265624447125176

Check the learned operator against the exact antiderivative on unseen inputs:

err = 0.0
for _ in 1:50
    a = 0.5 + rand(rng); y = rand(rng) * π
    err = max(err, abs(apply(op, θ, a .* sin.(sensors), [y]) - a * (1 - cos(y))))
end
println("DeepONet max error = ", round(err; digits = 4))
DeepONet max error = 0.0781

FNO: a spectral solution operator

The FNO applies v = σ(idct(R ⊙ dct(u)) + w·u + b) per block — a learned filter in the cosine (DCT) basis plus a pointwise skip. For a 1-D low-pass operator the true map is diagonal in the DCT basis, so a linear FNO can represent it exactly. We fit FNO1D to paired input/output signals:

n = 16
K = [1.0, 0.8, 0.6, 0.4, 0.2, 0.1, 0.05, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]
target(u) = idct2(K .* dct2(u))

fno = FNO1D(8, 1)
θf0 = fno1d_init(rng, 8, 1)
Uin = randn(rng, n, 100)
Uout = hcat([target(Uin[:, i]) for i in 1:100]...)
θf = fno1d_fit(fno, θf0, Uin, Uout; act = identity, epochs = 500, lr = 1e-2, seed = 5)

errf = maximum(abs.(apply(fno, θf, randn(rng, n); act = identity) .- target(randn(rng, n))))
println("FNO1D error (single sample) = ", round(errf; digits = 4))
FNO1D error (single sample) = 0.4615

The 2-D Poisson solution operator, exactly

The multi-channel FNO2DChannel lifts the input to width channels, applies channel-mixing spectral layers, and projects back. For the 2-D Poisson problem −Δu = f (Neumann) the solution operator is diagonal in the DCT basis with eigenvalues λ_{k,l} = (πk)² + (πl)², so a full-mode linear FNO reproduces it exactly — a useful self-check:

n2 = 8
λmat = [(π * (k - 1))^2 + (π * (l - 1))^2 for k in 1:n2, l in 1:n2]
λmat[1, 1] = 1.0
G(f) = idct2d(dct2d(f) ./ λmat)

fno2 = FNO2D(n2, 1)
θpois = vcat(vec(1 ./ λmat), 0.0, 0.0)   # R = 1/λ, w = b = 0
f = randn(rng, n2, n2)
println("FNO2D Poisson exactness error = ", maximum(abs.(apply(fno2, θpois, f; act = identity) .- G(f))))
FNO2D Poisson exactness error = 2.7755575615628914e-17

Next steps