Adding FNO and DeepONets to Neuromancer

[Parv621]

This PR updates #263.

It introduces DeepONets and Fourier Neural Operators (FNOs) into Neuromancer, with integrations to the external libraries neuralop and DeepXDE

Integration is provided via lightweight wrappers that preserve Neuromancer’s abstractions while maintaining compatibility with upstream libraries.

---

Summary of Changes

Operator Integrations

1. FNO Wrapper
   src/neuromancer/modules/operators.py
   A thin wrapper enabling direct use of neuralop models inside Neuromancer. The wrapper attaches the required _metadata for robust checkpointing across PyTorch versions. The design generalizes to any operator provided by neuralop.

2. Native DeepONet (Cartesian Product)
   src/neuromancer/modules/operators.py
   Native implementation of DeepONet with Cartesian-product evaluation, consistent with standard operator-learning formulations.

3. DeepXDE Wrapper
   src/neuromancer/modules/operators.py
   Wrapper for DeepONet-style models implemented in DeepXDE. Supports both
   dde.nn.DeepONetCartesianProd and dde.nn.DeepONet.

4. Custom Operator Losses
   src/neuromancer/modules/operators.py
   Implementation of LpLoss and H1Loss, suitable for grid-based operator learning.

---

Training and Validation Infrastructure

5. Validation Dataset Support During Training
   Changes in trainer.py, problem.py, and callbacks.py to support validation datasets during training (inspired by DeepONet Example #168).
   A new LossHistoryCallback is added to callbacks.py to track and plot training and validation loss per epoch.

---

Implemented Notebooks

DeepONet Examples

1. Antiderivative (Aligned GRF Dataset)
   examples/neural_operators/Part_1_DeepONet_antiderivative_aligned.ipynb
   Inspiration: DeepXDE aligned antiderivative demo and Lu et al. (2021).

2. Antiderivative (Unaligned GRF Dataset)
   examples/neural_operators/Part_2_DeepONet_antiderivative_unaligned.ipynb
   Inspiration: DeepXDE unaligned antiderivative demo.

3. 1D Advection equation (Aligned initial conditions from GRF fields)
   examples/neural_operators/Part_6_PIDeepONet_Advection_aligned.ipynb
   Inspiration: DeepXDE aligned 1D Advection demo.

---

FNO Examples

1. 1D Allen–Cahn Equation (from-scratch FNO)
   examples/neural_operators/Part_4_FNO_1DAllenCahn.ipynb
   Inspiration: ETH “Operator Learning – FNO” tutorial and the original FNO paper.

2. Darcy Flow (FNO, Neuromancer pipeline)
   examples/neural_operators/Part_5_FNO_DarcyFlow.ipynb
   Inspiration: neuralop Darcy Flow example and NVIDIA’s FNO Darcy guide.

3. Diffusion Equation (FNO, Neuromancer pipeline)
   examples/neural_operators/Part_6_FNO_DiffusionEquation.ipynb
   Inspiration: examples/PDEs/Part_1_PINN_DiffusionEquation.ipynb.

---

Notebooks I could not implement (DeepONet)

The following notebooks are not included as they do not achieve accurate solutions using standard MLP DeepONets:

• Cartesian-Product DeepONet – Darcy Flow
• Cartesian-Product DeepONet – Diffusion Equation
• Vanilla DeepONet – Darcy Flow
• Vanilla DeepONet – Diffusion Equation

These are moved to another SOLARIS branch (feature/DeepONet_notebooks_debug); we could revisit them later.

Some papers try to solve DarcyFlow using NOs, might be useful for debugging. NVIDIA example on Darcy Flow using DeepONet uses a convnet-based model pix2pix in the branch net. Another paper (https://www.mdpi.com/2227-9717/13/9/2754) solves DarcyFlow using various methods

---

Remaining Work

• Complete a Spherical FNO (SFNO) example for the shallow-water equations, inspired by the neuralop SFNO SWE example:
  https://neuraloperator.github.io/dev/auto_examples/models/plot_SFNO_swe.html

[drgona]

we can add the LossHistoryCallback class to callbacks.py

Since users will start from the first example, it would be good if the Part 1 example has extensive method markdown, similar to Part 3.

[Parv621]

For ifft in PyTorch needed for FNOs

[Parv621]

Should I add dependencies for neuralop and deepxde to the environment? Or would you like to keep it flexible?

[drgona]

It depends how many and what dependencies there are. If it is a lot, we can add the necessary dependencies for neuralop and deepxde to examples dependencies in our toml file.

https://github.com/pnnl/neuromancer/blob/master/pyproject.toml#L69

[Parv621]

So, the deepxde library can be added like that, as we use it only in examples.
The neuraloperator library is being used here https://github.com/SOLARIS-JHU/neuromancer_NO/blob/feature/FNO/src/neuromancer/modules/operators.py#L6, so it will either need to be added as a main dependency. Or we could do something like this in pyproject.toml file, loading this only for operators.

examples = [
    "deepxde",
    "imageio",
    "neuraloperator",
]
[project.optional-dependencies]
operators = [
  "neuraloperator",
]

Let me know what you prefer!

[Parv621]

Formatting changes only (linter)

[drgona]

lets run all the unit tests on your on this PR to make sure we are not braking the core

[drgona]

Well done with the examples!
Few minor comments on notebooks below.

1. Lets make sure the examples are ordered withour overlapping Part numbers. Lets consolidate examples to make them more compact.
2. Each new example should add markdown to new features or architectures it introduces compared to previous examples.
3. Can we have antiderivative example to introduce simple FNO as well?

Part_1_DeepONet_antiderivative_aligned.ipynb
markdown cell with Operator evaluation: lets specify b for grid point evaluations, similarly we define m for sensor points
lets explain what is the attribute "Glorot normal" in DeepXDE DeepOnet

Part_2_DeepONet_antiderivative_unaligned.ipynb
it would be nice to have base DeepONet class in neuromancer and compare with DeepXDE in this example similarly as in part 1. Nice to have if easy to implement, but not critical.

Part_5_FNO_DarcyFlow.ipynb
If we don't introduce FNO on antiderivative, we can use this example as main intro to FNO and make it part 3 and potentially merge with Part_3_DeepONet_DarcyFlow.ipynb to compare FNO vs DeepOnet

Part_3_DeepONet_DarcyFlow.ipynb
this notebook introduces for the first time DeepONetBatchWrapper but it is not explained, what is the difference between the DeepXDEWrapper?

It is not apparent from the notebooks what is the main difference between:
Part_3_DeepONetCartesianProd_DarcyFlow.ipynb
Part_3_DeepONet_DarcyFlow.ipynb
Maybe we can merge these two into one notebook and compare cartesiaon prod with non-catesian DeepOnet and focus on explaining the main differences and their implications.
Currently both notebooks have large cells of code without markdown.

Similarly:
Part_6_FNO_DiffusionEquation.ipynb
Part_3_DeepONet_DiffusionEquation.ipynb
Part_3_DeepONetCartesianProd_DiffusionEquation.ipynb
These can be merged into single exmaple that compares FNO vs DeepONet on diffusion example and focus on explaining main differences.

Part_3_DeepONetCartesianProd_DiffusionEquation.ipynb this title does not match the content which shows FNO. It looks like deprecated file to be deleted.

Part_4_FNO_1DAllenCahn.ipynb
The role of convolution layers should be explained.

[drgona]

lets make sure this dependency is included in our toml file:
from pathlib import Path

[Parv621]

pathlib is a standard library module. No dependency entry is required.

[drgona]

since you introduce change to the core abstraction, we need to run all the tests to make sure we are not breaking the core

[Parv621]

[drgona]

similarly here, lets run all the tests to make sure these changes to the core abstraction are backward compatible

[Parv621]

similarly here, lets run all the tests to make sure these changes to the core abstraction are backward compatible

this passed with a few warnings

[drgona]

Well done with the examples! Few minor comments on notebooks below.

1. Lets make sure the examples are ordered withour overlapping Part numbers. Lets consolidate examples to make them more compact.
2. Each new example should add markdown to new features or architectures it introduces compared to previous examples.
3. Can we have antiderivative example to introduce simple FNO as well?

Part_1_DeepONet_antiderivative_aligned.ipynb markdown cell with Operator evaluation: lets specify b for grid point evaluations, similarly we define m for sensor points lets explain what is the attribute "Glorot normal" in DeepXDE DeepOnet

Part_2_DeepONet_antiderivative_unaligned.ipynb it would be nice to have base DeepONet class in neuromancer and compare with DeepXDE in this example similarly as in part 1. Nice to have if easy to implement, but not critical.

Part_5_FNO_DarcyFlow.ipynb If we don't introduce FNO on antiderivative, we can use this example as main intro to FNO and make it part 3 and potentially merge with Part_3_DeepONet_DarcyFlow.ipynb to compare FNO vs DeepOnet

Part_3_DeepONet_DarcyFlow.ipynb this notebook introduces for the first time DeepONetBatchWrapper but it is not explained, what is the difference between the DeepXDEWrapper?

It is not apparent from the notebooks what is the main difference between: Part_3_DeepONetCartesianProd_DarcyFlow.ipynb Part_3_DeepONet_DarcyFlow.ipynb Maybe we can merge these two into one notebook and compare cartesiaon prod with non-catesian DeepOnet and focus on explaining the main differences and their implications. Currently both notebooks have large cells of code without markdown.

Similarly: Part_6_FNO_DiffusionEquation.ipynb Part_3_DeepONet_DiffusionEquation.ipynb Part_3_DeepONetCartesianProd_DiffusionEquation.ipynb These can be merged into single exmaple that compares FNO vs DeepONet on diffusion example and focus on explaining main differences.

Part_3_DeepONetCartesianProd_DiffusionEquation.ipynb this title does not match the content which shows FNO. It looks like deprecated file to be deleted.

Part_4_FNO_1DAllenCahn.ipynb The role of convolution layers should be explained.

FNO examples:
Part_4_FNO_1DAllenCahn.ipynb -> Part_3_FNO_1DAllenCahn.ipynb
Part_4_FNO_DarcyFlow.ipynb -> Part_4_FNO_DarcyFlow.ipynb - explain the embeddings
Part_6_FNO_DiffusionEquation.ipynb -> Part_5_FNO_DiffusionEquation.ipynb explain the diff with Part_4_FNO_DarcyFlow.ipynb