130 lines
4.1 KiB
Markdown
130 lines
4.1 KiB
Markdown
# 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:
|
||
|
||
```bash
|
||
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](https://arxiv.org/abs/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
|