[PoC] : Add neural operator ROM proof-of-concept for GSoC 2026

Project.toml

@@ -35,7 +35,10 @@ StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 StatsBase = "2913bbd2-ae8a-5f71-8c99-4fb6c76f3a91"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
+Lux = "b2108857-7c20-44ae-9111-449ecde12c47"
+Optimisers = "3bd65402-5787-11e9-1adc-39752487f4e2"
 UnPack = "3a884ed6-31ef-47d7-9d2a-63182c4928ed"
+Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
 [compat]
 ArraysOfArrays = "0.6"
@@ -64,7 +67,10 @@ SparseMatricesCSR = "0.6"
 StaticArrays = "1"
 Statistics = "1"
 StatsBase = "0.34"
+Lux = "1"
+Optimisers = "0.3, 0.4"
 UnPack = "1"
+Zygote = "0.6"
 julia = "1.9"
 
 [extras]

examples/poisson_deeponet.jl

@@ -0,0 +1,139 @@
+# End-to-end PoC: Neural Operator ROM for Parametric Poisson Equation
+#
+# Problem: -∇·(κ(μ) ∇u) = f  on [0,1]²,  u = 0 on ∂Ω
+#
+# The diffusivity coefficient κ depends on parameters:
+#   κ(x; μ) = μ₁ + μ₂ · sin(π·x₁) · sin(π·x₂)
+#
+# We train a DeepONet to learn the map μ → u_h(μ) (the DOF vector),
+# then predict solutions for new parameters without running FEM.
+
+module PoissonDeepONet
+
+using Gridap
+using GridapNeuralROMs
+using LinearAlgebra
+using Random
+using Statistics
+
+function main(;n_train=60,n_test=10,epochs=300,seed=42)
+  rng = Random.MersenneTwister(seed)
+
+  println("="^60)
+  println("Neural Operator ROM for Parametric Poisson")
+  println("="^60)
+
+  # ── 1. Gridap FEM setup ──────────────────────────────────────────
+
+  domain = (0,1,0,1)
+  partition = (16,16)
+  model = CartesianDiscreteModel(domain,partition)
+
+  order = 1
+  reffe = ReferenceFE(lagrangian,Float64,order)
+  V = TestFESpace(model,reffe;conformity=:H1,dirichlet_tags="boundary")
+  U = TrialFESpace(V,0.0)
+
+  Ω = Triangulation(model)
+  dΩ = Measure(Ω,2*order)
+
+  N_dofs = num_free_dofs(V)
+  println("\nMesh: $(partition[1])×$(partition[2]),  Free DOFs: $N_dofs")
+
+  # ── 2. Parametric solver ─────────────────────────────────────────
+  # Black-box function: μ → FEFunction
+  # This is what makes neural operator ROMs non-intrusive —
+  # we only need input/output pairs, not the PDE operators.
+
+  function solve_poisson(μ)
+    κ(x) = μ[1] + μ[2] * sin(π*x[1]) * sin(π*x[2])
+    f(x) = 1.0
+    a(u,v) = ∫( κ ⊙ (∇(u)⋅∇(v)) )dΩ
+    l(v)   = ∫( f*v )dΩ
+    op = AffineFEOperator(a,l,U,V)
+    return Gridap.solve(op)
+  end
+
+  # ── 3. Collect snapshots ─────────────────────────────────────────
+
+  param_bounds = [(0.1,5.0),(0.0,4.0)]  # μ₁ ∈ [0.1,5], μ₂ ∈ [0,4]
+
+  println("\nGenerating $n_train training snapshots...")
+  train_params = sample_parameters(param_bounds,n_train)
+  t_snap = @elapsed train_data = collect_snapshots(
+    solve_poisson,train_params;trial=U
+  )
+  println("  Snapshot generation: $(round(t_snap,digits=2))s")
+  println("  Solution matrix: $(size(train_data.solutions))")
+  println("  Coordinate matrix: $(size(train_data.coordinates))")
+
+  # ── 4. Build and train DeepONet ──────────────────────────────────
+
+  d_param = length(first(train_params))
+  spatial_dim = size(train_data.coordinates,1)
+
+  deeponet = build_deeponet(;
+    param_dim=d_param,
+    n_dofs=N_dofs,
+    spatial_dim=spatial_dim,
+    latent_dim=32,
+    branch_width=64,
+    trunk_width=64,
+    n_branch_layers=3,
+    n_trunk_layers=3,
+  )
+
+  println("\nTraining DeepONet (latent_dim=32, width=64, 3 layers)...")
+  config = TrainingConfig(;epochs=epochs,lr=1e-3,batch_size=16,
+                           patience=80,verbose=true)
+  t_train = @elapsed result = train_operator(train_data,deeponet;config,rng)
+  println("  Training time: $(round(t_train,digits=2))s")
+  println("  Best epoch: $(result.best_epoch)")
+  println("  Final val loss: $(round(result.val_losses[result.best_epoch],digits=6))")
+
+  # ── 5. Evaluate on test parameters ──────────────────────────────
+
+  println("\nEvaluating on $n_test test parameters...")
+  test_params = sample_parameters(param_bounds,n_test)
+
+  relative_errors = Float64[]
+  speedups = Float64[]
+
+  for (i,μ) in enumerate(test_params)
+    # Ground truth via FEM
+    t_fem = @elapsed uh_true = solve_poisson(μ)
+    dofs_true = get_free_dof_values(uh_true)
+
+    # Neural operator prediction
+    t_rom = @elapsed uh_pred = reconstruct_fe_function(result,μ,U)
+    dofs_pred = get_free_dof_values(uh_pred)
+
+    # Relative L2 error on DOF vector
+    err = norm(dofs_pred .- dofs_true) / norm(dofs_true)
+    push!(relative_errors,err)
+    push!(speedups,t_fem / max(t_rom,1e-10))
+  end
+
+  # ── 6. Report ───────────────────────────────────────────────────
+
+  println("\n" * "="^60)
+  println("RESULTS")
+  println("="^60)
+  println("  Mean relative L2 error: $(round(mean(relative_errors),digits=6))")
+  println("  Max  relative L2 error: $(round(maximum(relative_errors),digits=6))")
+  println("  Min  relative L2 error: $(round(minimum(relative_errors),digits=6))")
+  println("  Mean speedup (FEM/ROM): $(round(mean(speedups),digits=1))×")
+  println("  Training samples:       $n_train")
+  println("  Test samples:           $n_test")
+  println("  FEM DOFs:               $N_dofs")
+  println("="^60)
+
+  return (;relative_errors,speedups,result)
+end
+
+end # module
+
+# Run if executed directly
+if abspath(PROGRAM_FILE) == @__FILE__
+  PoissonDeepONet.main()
+end

src/Exports.jl

@@ -146,3 +146,17 @@ using GridapROMs.Extensions: ⊕; export ⊕
 @publish RBTransient TransientHyperReduction
 @publish RBTransient TransientProjection
 @publish RBTransient TransientRBOperator
+
+@publish NeuralOperatorROMs SnapshotData
+@publish NeuralOperatorROMs collect_snapshots
+@publish NeuralOperatorROMs extract_coordinates
+@publish NeuralOperatorROMs sample_parameters
+@publish NeuralOperatorROMs DeepONetLayer
+@publish NeuralOperatorROMs build_deeponet
+@publish NeuralOperatorROMs precompute_trunk_matrix
+@publish NeuralOperatorROMs NormalizationStats
+@publish NeuralOperatorROMs TrainingConfig
+@publish NeuralOperatorROMs TrainingResult
+@publish NeuralOperatorROMs train_operator
+@publish NeuralOperatorROMs reconstruct_fe_function
+@publish NeuralOperatorROMs evaluate_rom

src/GridapROMs.jl

@@ -26,6 +26,8 @@ include("RB/RBTransient/RBTransient.jl")
 
 include("Distributed/Distributed.jl")
 
+include("NeuralOperatorROMs/NeuralOperatorROMs.jl")
+
 include("Exports.jl")
 
 end

src/NeuralOperatorROMs/DeepONet.jl

@@ -0,0 +1,100 @@
+# DeepONet implementation for parametric PDE ROMs using Lux.jl.
+#
+# Key insight for PDE ROMs: the trunk network evaluates at fixed DOF
+# coordinates that don't change between queries. So we precompute the
+# trunk matrix T ∈ R^{N_dofs × p} once, and online prediction reduces to:
+#
+#   û(μ) = T · b(μ) + bias
+#
+# where b(μ) = branch(μ) ∈ R^p. This is an O(N·p) matrix-vector product,
+# independent of the FEM assembly cost.
+
+"""
+    DeepONetLayer{B,T} <: Lux.AbstractLuxContainerLayer{(:branch,:trunk)}
+
+Custom Lux layer implementing DeepONet for parametric PDE surrogate modelling.
+
+The branch network encodes the parameter vector μ into a latent code b(μ).
+The trunk network encodes spatial coordinates x into basis functions t(x).
+The output at DOF location xᵢ is: û_i = Σₖ bₖ(μ) · tₖ(xᵢ) + bias_i.
+"""
+struct DeepONetLayer{B,T} <: Lux.AbstractLuxContainerLayer{(:branch,:trunk)}
+  branch::B
+  trunk::T
+  latent_dim::Int
+  n_dofs::Int
+end
+
+"""
+    build_deeponet(;
+      param_dim, n_dofs, spatial_dim,
+      latent_dim=32, branch_width=64, trunk_width=64,
+      n_branch_layers=2, n_trunk_layers=2,
+      activation=Lux.gelu
+    ) -> DeepONetLayer
+
+Construct a DeepONet for mapping parameter vectors to FEM DOF vectors.
+
+- `param_dim`:   dimension of the parameter space (branch input)
+- `n_dofs`:      number of free DOFs in the FEM discretization (output dim)
+- `spatial_dim`: spatial dimension D of the mesh (trunk input)
+- `latent_dim`:  dimension p of the shared latent space
+"""
+function build_deeponet(;
+  param_dim::Int,
+  n_dofs::Int,
+  spatial_dim::Int,
+  latent_dim::Int=32,
+  branch_width::Int=64,
+  trunk_width::Int=64,
+  n_branch_layers::Int=2,
+  n_trunk_layers::Int=2,
+  activation=Lux.gelu
+)
+  # Branch: μ ∈ R^d → b(μ) ∈ R^p
+  branch_layers = Any[Dense(param_dim,branch_width,activation)]
+  for _ in 2:n_branch_layers
+    push!(branch_layers,Dense(branch_width,branch_width,activation))
+  end
+  push!(branch_layers,Dense(branch_width,latent_dim))
+  branch = Chain(branch_layers...)
+
+  # Trunk: x ∈ R^D → t(x) ∈ R^p
+  trunk_layers = Any[Dense(spatial_dim,trunk_width,activation)]
+  for _ in 2:n_trunk_layers
+    push!(trunk_layers,Dense(trunk_width,trunk_width,activation))
+  end
+  push!(trunk_layers,Dense(trunk_width,latent_dim))
+  trunk = Chain(trunk_layers...)
+
+  return DeepONetLayer(branch,trunk,latent_dim,n_dofs)
+end
+
+"""
+    precompute_trunk_matrix(model, coord_matrix, ps, st) -> Matrix{Float32}
+
+Evaluate the trunk network at all DOF coordinates once.
+Returns T ∈ R^{N_dofs × latent_dim} so that prediction is T * b(μ) + bias.
+"""
+function precompute_trunk_matrix(model::DeepONetLayer,coord_matrix::AbstractMatrix,ps,st)
+  # coord_matrix: D × N_nodes (Float32)
+  trunk_out,_ = Lux.apply(model.trunk,coord_matrix,ps.trunk,st.trunk)
+  # trunk_out: latent_dim × N_nodes
+  return Matrix(transpose(trunk_out))  # N_nodes × latent_dim
+end
+
+# Forward pass: branch encodes μ, then multiply with precomputed trunk matrix.
+# Input:  x = (mu, trunk_matrix) packed as a tuple
+# Output: û ∈ R^{N_dofs × batch}
+function (l::DeepONetLayer)(x::Tuple,ps,st)
+  mu,trunk_matrix = x
+  # mu: d_param × batch
+  b,new_st_branch = Lux.apply(l.branch,mu,ps.branch,st.branch)
+  # b: latent_dim × batch
+
+  # û = T * b + bias  →  (N_dofs × batch)
+  u_hat = trunk_matrix * b
+
+  new_st = (branch=new_st_branch,trunk=st.trunk)
+  return u_hat,new_st
+end

src/NeuralOperatorROMs/NeuralOperatorROMs.jl

@@ -0,0 +1,41 @@
+module NeuralOperatorROMs
+
+using Gridap
+using Gridap.FESpaces
+using Gridap.Geometry
+using Gridap.CellData
+using Gridap.ReferenceFEs
+
+using Lux
+using Optimisers
+using Zygote
+using Random
+using LinearAlgebra
+using Statistics
+
+include("Snapshots.jl")
+
+include("DeepONet.jl")
+
+include("Training.jl")
+
+include("Reconstruction.jl")
+
+export SnapshotData
+export collect_snapshots
+export extract_coordinates
+export sample_parameters
+
+export DeepONetLayer
+export build_deeponet
+export precompute_trunk_matrix
+
+export NormalizationStats
+export TrainingConfig
+export TrainingResult
+export train_operator
+
+export reconstruct_fe_function
+export evaluate_rom
+
+end # module

src/NeuralOperatorROMs/Reconstruction.jl

@@ -0,0 +1,47 @@
+# Reconstruct Gridap FEFunction from neural operator predictions.
+#
+# This closes the loop: the neural operator outputs a raw DOF vector,
+# and we wrap it back into a Gridap FEFunction so it can be used for
+# visualization (writevtk), error computation, and postprocessing
+# with the full Gridap ecosystem.
+
+"""
+    evaluate_rom(result::TrainingResult, μ::AbstractVector) -> Vector{Float64}
+
+Predict the DOF vector for a new parameter μ using the trained neural operator.
+Returns denormalized DOF values ready for FEFunction reconstruction.
+"""
+function evaluate_rom(result::TrainingResult,μ::AbstractVector)
+  mu_col = Float32.(reshape(μ,length(μ),1))
+  mu_n = normalize(result.input_norm,mu_col)
+
+  u_hat_n,_ = Lux.apply(
+    result.model,(mu_n,result.trunk_matrix),result.params,result.state
+  )
+  u_hat = denormalize(result.output_norm,vec(u_hat_n))
+  return Float64.(u_hat)
+end
+
+"""
+    reconstruct_fe_function(
+      result::TrainingResult,
+      μ::AbstractVector,
+      trial::FESpace
+    ) -> FEFunction
+
+Evaluate the neural operator at parameter μ and reconstruct a Gridap
+FEFunction. The trial space provides the Dirichlet DOF values and the
+mesh/basis function information needed to build a proper FEFunction.
+
+This is the key integration point with Gridap: predicted DOFs go in via
+`FEFunction(trial, free_values)`, the same constructor Gridap uses
+internally after solving a linear system.
+"""
+function reconstruct_fe_function(
+  result::TrainingResult,
+  μ::AbstractVector,
+  trial::FESpace
+)
+  predicted_dofs = evaluate_rom(result,μ)
+  return FEFunction(trial,predicted_dofs)
+end