osobhandClaude Opus 4.6
02d382d5f6
Consistent formatting pass: line wrapping, import sorting, trailing whitespace removal, let-chain indentation, merged derive attributes, and unsafe block reformatting. Co-Authored-By: Claude Opus 4.6 (1M context) <[email protected]>
547 lines
19 KiB
Rust
547 lines
19 KiB
Rust
//! Fourier Neural Operator (FNO) architectures.
|
|
//!
|
|
//! FNO is a neural operator architecture that learns mappings between function
|
|
//! spaces using spectral convolutions in the Fourier domain.
|
|
//!
|
|
//! ## Architecture
|
|
//!
|
|
//! FNO consists of:
|
|
//! 1. Lifting layer: P (input → latent space)
|
|
//! 2. Multiple Fourier layers: spectral conv + residual + activation
|
|
//! 3. Projection layer: Q (latent space → output)
|
|
//!
|
|
//! ## References
|
|
//!
|
|
//! Li, Z., et al. (2020). "Fourier Neural Operator for Parametric Partial
|
|
//! Differential Equations." arXiv:2010.08895
|
|
|
|
use rtx_backend::Backend;
|
|
use rtx_nn::generic::{GenericLinear, GenericModule, GenericModule4D};
|
|
use rtx_tensor::generic::GenericTensor;
|
|
use std::fmt::Debug;
|
|
|
|
use crate::Result;
|
|
use crate::layers::{GridPositionalEncoding, LiftingMLP, Projection};
|
|
use crate::spectral::SpectralConv2d;
|
|
use crate::weights::FNO2dWeights;
|
|
|
|
/// 1D Fourier Neural Operator.
|
|
#[derive(Debug)]
|
|
pub struct FNO1d<B: Backend<FloatElem = f32>> {
|
|
_device: B::Device,
|
|
}
|
|
|
|
impl<B: Backend<FloatElem = f32>> FNO1d<B> {
|
|
/// Create a new 1D FNO.
|
|
///
|
|
/// # Arguments
|
|
/// * `in_channels` - Number of input channels
|
|
/// * `out_channels` - Number of output channels
|
|
/// * `width` - Width of the latent space
|
|
/// * `n_modes` - Number of Fourier modes
|
|
/// * `device` - Device to create the model on
|
|
pub fn new(
|
|
_in_channels: usize,
|
|
_out_channels: usize,
|
|
_width: usize,
|
|
_n_modes: usize,
|
|
device: &B::Device,
|
|
) -> Result<Self> {
|
|
Ok(Self {
|
|
_device: device.clone(),
|
|
})
|
|
}
|
|
}
|
|
|
|
/// 2D Fourier Neural Operator.
|
|
///
|
|
/// FNO2d learns mappings between 2D function spaces by alternating between
|
|
/// spectral convolutions (in Fourier space) and pointwise operations (in physical space).
|
|
///
|
|
/// ## Architecture (neuraloperator v2.0 compatible)
|
|
///
|
|
/// 1. Positional encoding: Adds 2D grid coordinates to input
|
|
/// 2. Lifting MLP (P): Two-layer MLP projects (in_channels + 2) to width
|
|
/// 3. N Fourier blocks: Each has spectral_conv + 1x1 conv + residual + GELU
|
|
/// 4. Projection layers (Q): Projects width -> 128 -> output channels
|
|
#[derive(Debug)]
|
|
pub struct FNO2d<B: Backend<FloatElem = f32>> {
|
|
/// Grid positional encoding layer
|
|
positional_encoding: GridPositionalEncoding<B>,
|
|
/// Two-layer MLP lifting (neuraloperator v2.0 compatible)
|
|
lifting: LiftingMLP<B>,
|
|
/// Spectral convolution layers
|
|
spectral_convs: Vec<SpectralConv2d<B>>,
|
|
/// 1x1 convolution layers (skip connections)
|
|
convs: Vec<GenericLinear<B>>,
|
|
/// First projection layer (width -> hidden)
|
|
projection1: Projection<B>,
|
|
/// Second projection layer (hidden -> out_channels)
|
|
projection2: Projection<B>,
|
|
device: B::Device,
|
|
}
|
|
|
|
impl<B: Backend<FloatElem = f32>> FNO2d<B> {
|
|
/// Create a new 2D FNO with default number of layers (4).
|
|
///
|
|
/// # Arguments
|
|
/// * `in_channels` - Number of input channels
|
|
/// * `out_channels` - Number of output channels
|
|
/// * `width` - Width of the latent space
|
|
/// * `n_modes` - Number of Fourier modes (same for both dimensions)
|
|
/// * `device` - Device to create the model on
|
|
pub fn new(
|
|
in_channels: usize,
|
|
out_channels: usize,
|
|
width: usize,
|
|
n_modes: usize,
|
|
device: &B::Device,
|
|
) -> Result<Self> {
|
|
Self::new_with_layers(in_channels, out_channels, width, n_modes, 4, device)
|
|
}
|
|
|
|
/// Create a new 2D FNO with custom number of Fourier blocks.
|
|
///
|
|
/// # Arguments
|
|
/// * `in_channels` - Number of input channels
|
|
/// * `out_channels` - Number of output channels
|
|
/// * `width` - Width of the latent space
|
|
/// * `n_modes` - Number of Fourier modes (same for both dimensions)
|
|
/// * `n_layers` - Number of Fourier blocks
|
|
/// * `device` - Device to create the model on
|
|
/// Hidden dimension for projection layers (matches neuraloperator)
|
|
const PROJECTION_HIDDEN: usize = 128;
|
|
|
|
/// Creates a new FNO2d model with a specified number of Fourier layers.
|
|
///
|
|
/// # Arguments
|
|
///
|
|
/// * `in_channels` - Number of input channels
|
|
/// * `out_channels` - Number of output channels
|
|
/// * `width` - Width of hidden channels in Fourier blocks
|
|
/// * `n_modes` - Number of Fourier modes to keep
|
|
/// * `n_layers` - Number of Fourier blocks
|
|
/// * `device` - Device to create the model on
|
|
pub fn new_with_layers(
|
|
in_channels: usize,
|
|
out_channels: usize,
|
|
width: usize,
|
|
n_modes: usize,
|
|
n_layers: usize,
|
|
device: &B::Device,
|
|
) -> Result<Self> {
|
|
// 1. Create positional encoding layer
|
|
let positional_encoding = GridPositionalEncoding::new(device);
|
|
|
|
// 2. Create lifting MLP: (in_channels + 2) -> width
|
|
// The +2 accounts for x and y positional encoding channels
|
|
let lifting = LiftingMLP::new(in_channels + 2, width, device);
|
|
|
|
// 3. Create n_layers spectral convolution layers: width -> width
|
|
let mut spectral_convs = Vec::with_capacity(n_layers);
|
|
for _ in 0..n_layers {
|
|
spectral_convs.push(SpectralConv2d::new(width, width, n_modes, n_modes, device));
|
|
}
|
|
|
|
// 4. Create n_layers 1x1 convolution layers (skip connections): width -> width
|
|
// Using GenericLinear as 1x1 conv (pointwise linear transformation)
|
|
let mut convs = Vec::with_capacity(n_layers);
|
|
for _ in 0..n_layers {
|
|
convs.push(GenericLinear::new(width, width, true, device));
|
|
}
|
|
|
|
// 5. Create two-stage projection: width -> 128 -> out_channels
|
|
let projection1 = Projection::new(width, Self::PROJECTION_HIDDEN, device);
|
|
let projection2 = Projection::new(Self::PROJECTION_HIDDEN, out_channels, device);
|
|
|
|
Ok(Self {
|
|
positional_encoding,
|
|
lifting,
|
|
spectral_convs,
|
|
convs,
|
|
projection1,
|
|
projection2,
|
|
device: device.clone(),
|
|
})
|
|
}
|
|
|
|
/// Create a new FNO2d model from pre-trained weights.
|
|
///
|
|
/// # Arguments
|
|
/// * `weights` - Pre-trained weights loaded from SafeTensors file
|
|
/// * `device` - Device to create the model on
|
|
///
|
|
/// # Example
|
|
/// ```rust,ignore
|
|
/// use rtx_neural_operator::{FNO2d, load_fno2d_weights};
|
|
///
|
|
/// let weights = load_fno2d_weights("./weights/fno2d_darcy/model.safetensors")?;
|
|
/// let fno = FNO2d::<CpuBackend>::from_weights(&weights, &device)?;
|
|
/// ```
|
|
pub fn from_weights(weights: &FNO2dWeights, device: &B::Device) -> Result<Self> {
|
|
let config = &weights.config;
|
|
|
|
// Create positional encoding layer
|
|
let positional_encoding = GridPositionalEncoding::new(device);
|
|
|
|
// Create lifting MLP from weights
|
|
let lifting = LiftingMLP::from_weights(
|
|
&weights.lifting_fc1_weight,
|
|
&weights.lifting_fc1_bias,
|
|
&weights.lifting_fc2_weight,
|
|
&weights.lifting_fc2_bias,
|
|
config.in_channels + 2, // +2 for positional encoding
|
|
config.width,
|
|
device,
|
|
);
|
|
|
|
// Create spectral convolution layers from weights
|
|
let mut spectral_convs = Vec::with_capacity(config.n_layers);
|
|
for sw in &weights.spectral_weights {
|
|
spectral_convs.push(SpectralConv2d::from_weights_dual(
|
|
&sw.weights1_real,
|
|
&sw.weights1_imag,
|
|
&sw.weights2_real,
|
|
&sw.weights2_imag,
|
|
config.width,
|
|
config.width,
|
|
config.n_modes.0,
|
|
config.n_modes.1,
|
|
device,
|
|
));
|
|
}
|
|
|
|
// Create 1x1 convolution layers (skip connections) from weights
|
|
let mut convs = Vec::with_capacity(config.n_layers);
|
|
for (weight, bias) in &weights.conv_weights {
|
|
convs.push(GenericLinear::from_weights(
|
|
weight,
|
|
Some(bias),
|
|
config.width,
|
|
config.width,
|
|
device,
|
|
));
|
|
}
|
|
|
|
// Create two-stage projection layers: width -> 128 -> out_channels
|
|
let (proj0_weight, proj0_bias) = &weights.projection_weights[0];
|
|
let projection1 = Projection::from_weights(
|
|
proj0_weight,
|
|
proj0_bias,
|
|
config.width,
|
|
Self::PROJECTION_HIDDEN,
|
|
device,
|
|
);
|
|
|
|
let (proj1_weight, proj1_bias) = &weights.projection_weights[1];
|
|
let projection2 = Projection::from_weights(
|
|
proj1_weight,
|
|
proj1_bias,
|
|
Self::PROJECTION_HIDDEN,
|
|
config.out_channels,
|
|
device,
|
|
);
|
|
|
|
Ok(Self {
|
|
positional_encoding,
|
|
lifting,
|
|
spectral_convs,
|
|
convs,
|
|
projection1,
|
|
projection2,
|
|
device: device.clone(),
|
|
})
|
|
}
|
|
}
|
|
|
|
impl<B: Backend<FloatElem = f32>> FNO2d<B> {
|
|
/// Extract all model weights into a serializable format.
|
|
///
|
|
/// This is the inverse of `from_weights()` and can be used to save
|
|
/// trained models to SafeTensors format.
|
|
pub fn to_weights(&self, pde_type: &str) -> FNO2dWeights {
|
|
// Extract lifting MLP weights
|
|
let (lifting_fc1_weight, lifting_fc1_bias) = self.lifting.fc1_weights();
|
|
let (lifting_fc2_weight, lifting_fc2_bias) = self.lifting.fc2_weights();
|
|
|
|
// Extract spectral convolution weights
|
|
let spectral_weights: Vec<crate::weights::SpectralConvWeights> = self
|
|
.spectral_convs
|
|
.iter()
|
|
.map(|sc| {
|
|
let (w1r, w1i) = sc.weights1_real_imag();
|
|
let (w2r, w2i) = sc.weights2_real_imag();
|
|
crate::weights::SpectralConvWeights {
|
|
weights1_real: w1r,
|
|
weights1_imag: w1i,
|
|
weights2_real: w2r,
|
|
weights2_imag: w2i,
|
|
}
|
|
})
|
|
.collect();
|
|
|
|
// Extract 1x1 convolution (skip connection) weights
|
|
let conv_weights: Vec<(Vec<f32>, Vec<f32>)> = self
|
|
.convs
|
|
.iter()
|
|
.map(|linear| {
|
|
let weight = linear.weight().to_vec();
|
|
let bias = linear
|
|
.bias()
|
|
.map(rtx_tensor::GenericTensor::to_vec)
|
|
.unwrap_or_default();
|
|
(weight, bias)
|
|
})
|
|
.collect();
|
|
|
|
// Extract projection weights
|
|
let proj1 = self.projection1.weights();
|
|
let proj2 = self.projection2.weights();
|
|
let projection_weights = vec![proj1, proj2];
|
|
|
|
// Determine config from model structure
|
|
let width = self.lifting.out_features();
|
|
let n_layers = self.spectral_convs.len();
|
|
let n_modes = if !self.spectral_convs.is_empty() {
|
|
(
|
|
self.spectral_convs[0].n_modes_h(),
|
|
self.spectral_convs[0].n_modes_w(),
|
|
)
|
|
} else {
|
|
(12, 12)
|
|
};
|
|
let in_channels = self.lifting.in_features().saturating_sub(2); // Remove positional encoding
|
|
let out_channels = self.projection2.out_features();
|
|
|
|
let config = crate::weights::FNO2dConfig {
|
|
in_channels,
|
|
out_channels,
|
|
width,
|
|
n_modes,
|
|
n_layers,
|
|
pde_type: pde_type.to_string(),
|
|
};
|
|
|
|
FNO2dWeights {
|
|
lifting_fc1_weight,
|
|
lifting_fc1_bias,
|
|
lifting_fc2_weight,
|
|
lifting_fc2_bias,
|
|
spectral_weights,
|
|
conv_weights,
|
|
projection_weights,
|
|
config,
|
|
}
|
|
}
|
|
}
|
|
|
|
impl<B: Backend<FloatElem = f32>> GenericModule<B> for FNO2d<B> {
|
|
fn forward(&self, _input: &GenericTensor<B, 2>) -> GenericTensor<B, 2> {
|
|
panic!("Use forward_4d for FNO2d");
|
|
}
|
|
|
|
fn device(&self) -> &B::Device {
|
|
&self.device
|
|
}
|
|
}
|
|
|
|
impl<B: Backend<FloatElem = f32>> GenericModule4D<B> for FNO2d<B> {
|
|
fn forward_4d(&self, input: &GenericTensor<B, 4>) -> GenericTensor<B, 4> {
|
|
// Step 1: Add positional encoding
|
|
// input: [batch, in_channels, height, width]
|
|
// with_pos: [batch, in_channels + 2, height, width]
|
|
let with_pos = self.positional_encoding.forward_4d(input);
|
|
|
|
// Step 2: Lift input to latent space using MLP
|
|
// with_pos: [batch, in_channels + 2, height, width]
|
|
// x: [batch, width, height, width_spatial]
|
|
let mut x = self.lifting.forward_4d(&with_pos);
|
|
|
|
// Step 3: Apply Fourier blocks
|
|
let n_layers = self.spectral_convs.len();
|
|
for i in 0..n_layers {
|
|
// Apply spectral convolution in Fourier domain
|
|
let spectral_out = self.spectral_convs[i].forward_4d(&x);
|
|
|
|
// Apply 1x1 convolution (skip connection) in physical space
|
|
// For 1x1 conv, we need to apply GenericLinear pointwise across spatial dimensions
|
|
let conv_out = self.apply_pointwise_linear(&x, &self.convs[i]);
|
|
|
|
// Residual connection: combine spectral and conv paths
|
|
let combined = spectral_out.add(&conv_out);
|
|
|
|
// Apply GELU activation (except on last layer, matching neuraloperator v2.0)
|
|
if i < n_layers - 1 {
|
|
x = combined.gelu();
|
|
} else {
|
|
x = combined;
|
|
}
|
|
}
|
|
|
|
// Step 4: Project back to output space
|
|
// Two-stage projection: width -> 128 -> out_channels with GELU
|
|
let x = self.projection1.forward_4d(&x);
|
|
let x = x.gelu();
|
|
self.projection2.forward_4d(&x)
|
|
}
|
|
}
|
|
|
|
impl<B: Backend<FloatElem = f32>> FNO2d<B> {
|
|
/// Apply pointwise linear transformation (1x1 convolution) to 4D tensor.
|
|
///
|
|
/// This mimics a 1x1 convolution by applying the linear layer independently
|
|
/// at each spatial location.
|
|
fn apply_pointwise_linear(
|
|
&self,
|
|
input: &GenericTensor<B, 4>,
|
|
linear: &GenericLinear<B>,
|
|
) -> GenericTensor<B, 4> {
|
|
let shape = input.shape();
|
|
let batch = shape[0];
|
|
let channels = shape[1];
|
|
let height = shape[2];
|
|
let width = shape[3];
|
|
|
|
// Reshape: [batch, channels, height, width] -> [batch, height, width, channels]
|
|
// First swap channels and height: [B, C, H, W] -> [B, H, C, W]
|
|
let step1 = input.swap_dims(1, 2);
|
|
// Then swap channels and width: [B, H, C, W] -> [B, H, W, C]
|
|
let permuted = step1.swap_dims(2, 3);
|
|
|
|
// Reshape to [batch * height * width, channels]
|
|
let reshaped = permuted.reshape([batch * height * width, channels]);
|
|
|
|
// Apply linear transformation
|
|
let transformed = linear.forward(&reshaped);
|
|
|
|
// Get output channels
|
|
let out_channels = transformed.shape()[1];
|
|
|
|
// Reshape back to [batch, height, width, out_channels]
|
|
let output_nhwc = transformed.reshape([batch, height, width, out_channels]);
|
|
|
|
// Permute back to [batch, out_channels, height, width]
|
|
// First swap width and channels: [B, H, W, C] -> [B, H, C, W]
|
|
let step1 = output_nhwc.swap_dims(2, 3);
|
|
// Then swap height and channels: [B, H, C, W] -> [B, C, H, W]
|
|
step1.swap_dims(1, 2)
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
use rtx_backend_cpu::{CpuBackend, CpuDevice};
|
|
|
|
#[test]
|
|
fn test_fno2d_creation() {
|
|
let device = CpuDevice::new();
|
|
let result = FNO2d::<CpuBackend>::new(1, 1, 32, 12, &device);
|
|
assert!(result.is_ok());
|
|
}
|
|
|
|
#[test]
|
|
fn test_fno2d_forward_shape() {
|
|
let device = CpuDevice::new();
|
|
let fno = FNO2d::<CpuBackend>::new(2, 3, 64, 16, &device).unwrap();
|
|
|
|
let input = GenericTensor::randn([4, 2, 32, 32], &device);
|
|
let output = fno.forward_4d(&input);
|
|
|
|
// For now, output has same shape as input (placeholder)
|
|
assert_eq!(output.shape()[0], 4); // batch
|
|
assert_eq!(output.shape()[2], 32); // height
|
|
assert_eq!(output.shape()[3], 32); // width
|
|
}
|
|
|
|
// ==================== TDD Tests for Full FNO2d Implementation ====================
|
|
|
|
#[test]
|
|
fn test_fno2d_output_channels() {
|
|
let device = CpuDevice::new();
|
|
let fno = FNO2d::<CpuBackend>::new(3, 8, 64, 12, &device).unwrap();
|
|
|
|
let input = GenericTensor::randn([2, 3, 32, 32], &device);
|
|
let output = fno.forward_4d(&input);
|
|
|
|
// Output should transform from in_channels to out_channels
|
|
assert_eq!(output.shape()[0], 2); // batch
|
|
assert_eq!(output.shape()[1], 8); // out_channels
|
|
assert_eq!(output.shape()[2], 32); // height preserved
|
|
assert_eq!(output.shape()[3], 32); // width preserved
|
|
}
|
|
|
|
#[test]
|
|
fn test_fno2d_non_zero_output() {
|
|
let device = CpuDevice::new();
|
|
let fno = FNO2d::<CpuBackend>::new(2, 4, 32, 8, &device).unwrap();
|
|
|
|
// Create non-zero input
|
|
let input = GenericTensor::ones([1, 2, 16, 16], &device);
|
|
let output = fno.forward_4d(&input);
|
|
|
|
// FNO should produce meaningful non-zero output
|
|
let output_data = output.to_vec();
|
|
let non_zero_count = output_data.iter().filter(|&&x| x.abs() > 1e-6).count();
|
|
|
|
assert!(
|
|
non_zero_count > 0,
|
|
"FNO should produce non-zero outputs for non-zero inputs"
|
|
);
|
|
|
|
// Output should not be all zeros
|
|
let all_zeros = output_data.iter().all(|&x| x.abs() < 1e-6);
|
|
assert!(!all_zeros, "FNO output should not be all zeros");
|
|
}
|
|
|
|
#[test]
|
|
fn test_fno2d_different_widths() {
|
|
let device = CpuDevice::new();
|
|
|
|
// Test with different latent widths
|
|
let fno_small = FNO2d::<CpuBackend>::new(1, 1, 16, 4, &device).unwrap();
|
|
let fno_large = FNO2d::<CpuBackend>::new(1, 1, 128, 4, &device).unwrap();
|
|
|
|
let input = GenericTensor::randn([1, 1, 32, 32], &device);
|
|
|
|
let output_small = fno_small.forward_4d(&input);
|
|
let output_large = fno_large.forward_4d(&input);
|
|
|
|
// Both should produce valid outputs with correct shape
|
|
assert_eq!(output_small.shape(), [1, 1, 32, 32]);
|
|
assert_eq!(output_large.shape(), [1, 1, 32, 32]);
|
|
|
|
// Outputs should be non-zero
|
|
let data_small = output_small.to_vec();
|
|
let data_large = output_large.to_vec();
|
|
|
|
assert!(data_small.iter().any(|&x| x.abs() > 1e-6));
|
|
assert!(data_large.iter().any(|&x| x.abs() > 1e-6));
|
|
}
|
|
|
|
#[test]
|
|
fn test_fno2d_custom_layers() {
|
|
let device = CpuDevice::new();
|
|
|
|
// Test with different number of Fourier blocks
|
|
let fno_2layers = FNO2d::<CpuBackend>::new_with_layers(2, 3, 32, 8, 2, &device).unwrap();
|
|
let fno_6layers = FNO2d::<CpuBackend>::new_with_layers(2, 3, 32, 8, 6, &device).unwrap();
|
|
|
|
let input = GenericTensor::randn([1, 2, 24, 24], &device);
|
|
|
|
let output_2 = fno_2layers.forward_4d(&input);
|
|
let output_6 = fno_6layers.forward_4d(&input);
|
|
|
|
// Both should produce valid outputs
|
|
assert_eq!(output_2.shape(), [1, 3, 24, 24]);
|
|
assert_eq!(output_6.shape(), [1, 3, 24, 24]);
|
|
|
|
// Outputs should be different (different network depths)
|
|
let data_2 = output_2.to_vec();
|
|
let data_6 = output_6.to_vec();
|
|
|
|
assert!(data_2.iter().any(|&x| x.abs() > 1e-6));
|
|
assert!(data_6.iter().any(|&x| x.abs() > 1e-6));
|
|
}
|
|
}
|