# 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