[GSoC 2026] Add Neural Operators (DeepONet and NOMAD)

[bitsarebits]

This pull request contains my final work submission for Google Summer of Code 2026.

Project Goals

The goal of this project was to integrate Neural Operators into GridapROMs.jl, providing surrogate models for high-fidelity parametric finite element solvers in both steady-state and transient settings.

What I Did

The implementation integrates two Neural Operator architectures into the GridapROMs.jl reduction and solver workflow:

• DeepONet (Deep Operator Network) for steady-state and transient problems.
• NOMAD (Non-linear Manifold Decoder) for steady-state and transient problems.

The implementation includes:

• Automatic architecture construction through AutoDeepONet and AutoNOMAD, together with support for explicitly specified network architectures.
• A configurable NeuralOpStrategy for model architecture, training parameters, spatial and temporal subsampling, verbosity, and learning-rate scheduling.
• Custom parameter and sensor preprocessing through branch_sampler, including transformations such as logarithmic scaling and construction of multi-sensor inputs.
• Spatial and temporal subsampling through step_x and step_t to control the size of the training data for large meshes and transient problems.
• Input normalization and output scaling, with normalization statistics stored in the resulting NeuralRBOperator for consistent online inference.
• Fine-tuning of pretrained neural operators, supporting both continued training with the original normalization statistics and transfer to a new data distribution through update_stats.
• Learning-rate scheduling through CosineAnnealing and ReduceLROnPlateau.
• Separate offline training and online inference paths for steady-state and transient problems.
• Integration with the existing GridapROMs solver and operator abstractions through multiple dispatch.
• Tests covering utility functions, normalization, learning-rate schedulers, architecture construction, coordinate extraction, model construction, steady and transient inference, fine-tuning, and selected error conditions.

The implementation is organized into dedicated training, solver, scheduler, and logging components in both RBSteady and RBTransient.

Current Status

The main GSoC scope has been implemented and exercised through integration tests and example workflows.

The implemented API covers the complete offline/online workflow: snapshots are used to train a neural operator, the resulting model and normalization data are stored in a NeuralRBOperator, and the trained operator can subsequently be used through the standard solve interface.

Both steady and transient workflows are supported, including transient inference over the full space-time grid.

The examples also include fine-tuning workflows for continual learning and transfer learning, with the corresponding normalization behavior documented in the API.

Follow-up Work

The current implementation provides a foundation for further development. Planned follow-up work includes:

• Adding and evaluating an additional Neural Operator architecture.
• Improving the extraction of physical coordinates from finite element spaces using a more efficient mechanism.
• Further evaluating the implementation across additional PDEs, meshes, and parameter distributions.

These items are intended as post-GSoC extensions and improvements rather than blockers for the current project scope.

AI Assistance

Gemini was used as a supplementary development tool during the project, primarily for code-related discussion, debugging support, and documentation assistance. The implementation, design decisions, testing, and final integration were carried out and reviewed by me.