rustytorch/crates/specialized/rtx-neural-operator

RTX Neural Operator

Neural operators for learning solution operators to PDEs (Partial Differential Equations).

Overview

This crate provides specialized neural network architectures for scientific computing and physics-informed machine learning:

• Fourier Neural Operators (FNO): Learn mappings between function spaces using spectral convolutions
• DeepONet: Separate branch and trunk networks for operator learning
• Spectral Convolutions: Efficient convolutions in the Fourier domain

Architecture

Neural operators differ from traditional neural networks by learning mappings between infinite-dimensional function spaces rather than finite-dimensional vectors. They are particularly effective for solving PDEs and learning physical systems.

Key Components

SpectralConv2d

The core building block - performs convolution in the Fourier domain:

1. Apply 2D FFT to input → frequency representation
2. Truncate to n_modes (low-pass filter, keep low frequencies)
3. Multiply by learnable complex weights in frequency space
4. Apply inverse 2D FFT → return to physical space

This approach is dramatically more efficient than spatial convolutions for learning global patterns, as multiplication in Fourier space is O(n_modes) vs O(kernel_size²).

Lifting and Projection Layers

• Lifting: Projects input channels to high-dimensional latent space
• Projection: Projects latent representation back to output space

FNO Architecture

Input → Lifting → [SpectralConv + Residual + Activation]×L → Projection → Output

Where L is the number of Fourier layers.

Implementation Status

✅ Complete (with Tests)

• Error types and Result aliases
• Lifting and Projection layers (fully tested)
• SpectralConv1d/2d structure and weight initialization
• FNO1d/2d architecture scaffolding
• DeepONet architecture scaffolding
• 20 passing unit tests

🚧 TODO (Future Implementation)

1. FFT Integration: Implement actual spectral convolution using rtx-tensor's ComplexTensor API

   • Real-to-complex conversion
   • 2D FFT/IFFT operations
   • Mode truncation/padding
   • Complex weight multiplication

2. FNO Forward Pass: Complete implementation with:

   • Multiple Fourier layers
   • Residual connections
   • Activation functions (GELU)
   • Skip connections

3. DeepONet Implementation:

   • Branch network (MLP)
   • Trunk network (MLP)
   • Inner product aggregation

4. Training utilities:

   • Loss functions for operator learning
   • Relative L2 error metric

Testing

All tests follow strict TDD principles:

cargo test -p rtx-neural-operator

Current test coverage:

• Weight initialization (Xavier uniform)
• Shape preservation through layers
• Batch independence
• Parameter validation
• Panic conditions

Dependencies

• rtx-tensor: Tensor operations and FFT
• rtx-nn: Neural network layers
• rtx-autograd: Automatic differentiation
• rtx-backend: Backend abstraction (CPU, CUDA, etc.)

Design Principles

1. TDD First: All code has tests written before implementation
2. No External ML Frameworks: Pure RTX stack, no burn/candle/torch
3. Type Safety: Generic over Backend with compile-time dispatch
4. Production Ready: No unwrap(), proper error handling with Result<T, E>
5. Rust 2024 Edition: Uses latest stable features

File Organization

rtx-neural-operator/
├── src/
│   ├── lib.rs          (76 lines)   - Error types, exports
│   ├── layers.rs       (196 lines)  - Lifting, Projection
│   ├── spectral.rs     (366 lines)  - SpectralConv1d/2d
│   ├── fno.rs          (127 lines)  - FNO1d/2d architectures
│   └── deeponet.rs     (85 lines)   - DeepONet
└── Cargo.toml

Total: 850 lines (well under 1000-line limit per file)

References

• Li, Z., et al. (2020). "Fourier Neural Operator for Parametric Partial Differential Equations." arXiv:2010.08895
• Lu, L., et al. (2021). "Learning nonlinear operators via DeepONet based on the universal approximation theorem of operators." Nature Machine Intelligence.

License

MIT OR Apache-2.0