Feature Parity with NeuralOperators.jl

[avik-pal]

• [x] Fourier NO
• [ ] Markov NO
• [x] DeepONets
• [ ] Graph Kernel Networks (needs https://github.com/CarloLucibello/GraphNeuralNetworks.jl/issues/372)
• [x] NOMAD

[ayushinav]

IIUC, MNOs are not really a different architecture. Its using mostly another NO, FNO being used in the work, and defining different loss functions. We can add a number of other architectures though, eg., UNO, variants of DeepONets such as DeepM&Mnets and others. This work provides a good review of operators we can add. @avik-pal @ChrisRackauckas

[ayushinav]

Also, the continuous form would be another good addition. https://github.com/SciML/NeuralOperators.jl/issues/70

[ayushinav]

I'm struggling a bit to implement the continuous variant and would appreciate some discussion here. The general form of NO is

\mathcal{G}_{\theta} = \mathcal{Q} \circ \sigma_T(W_{T-1} + \mathcal{K}_{T-1}+ b_{T-1}) \circ \dots \circ \sigma_1(W_0 + \mathcal{K}_0+ b_0) \circ \mathcal{P}

$\mathcal{P}$ and $\mathcal{Q}$ are local operators that lift and project. IIUC, we would not need them as networks but as functions (?). Similarly, $W_t$'s should be functions too?

Coming to the kernel part $\mathcal{K}_t$'s, one version of the same is defined as:

(\mathcal{K}_t(v_t))(x) = \int_{D_t} \kappa^{(t)}(x,y) v_t(y) dy \quad \forall x \in D_t$

Now, $\kappa^{(t)}$ should be another function that takes in $x,y$. We'd need to have the domains defined for computing the integrals (the PDFs would be better but for now we can start with uniform distributions). The first kernel will have the domains defined by the input but will change for the subsequent layers, so we need to propagate the domains as well? Also, I'm not clear where the networks come in here. We can obviously have networks as function approximators for $\mathcal{P}, \mathcal{Q}, W_t, \kappa^{(t)}$ but I wanted to make sure if that's the case. @avik-pal @ChrisRackauckas

[avik-pal]

With the benchmarks looking positive should we try and get NOMAD and Markov NO? The FNO benchmarks are something I will have to handle on LuxLib end.

With those done, we can take a slight feature hit for the Graph NO (though GNNLux is a thing now so we can implement it soon enough after that) and tag a release for this version of Neural Operators replacing the old one. We can move this repo over to SciML or keep it here, don't have any significant preference on my end.

Thoughts @ChrisRackauckas @ayushinav .

[ayushinav]

IIUC, Markov NO is not much of a different architecture. It uses an existing NO, usually FNO (because it's readily available and popular). Even in the pervious NeuralOperators.jl, an FNO is returned https://github.com/SciML/NeuralOperators.jl/blob/main/src/FNO/FNO.jl#L196-L212.