318 lines
9.4 KiB
Rust
318 lines
9.4 KiB
Rust
//! Fourier Neural Operator implementation for CFD.
|
|
|
|
/// Fourier Neural Operator for flow prediction.
|
|
#[derive(Debug)]
|
|
pub struct FourierNeuralOperator {
|
|
/// Number of Fourier modes.
|
|
num_modes: usize,
|
|
/// Hidden dimension.
|
|
hidden_dim: usize,
|
|
/// Number of layers.
|
|
num_layers: usize,
|
|
/// Spectral convolution weights.
|
|
spectral_weights: Vec<Vec<f32>>,
|
|
/// Linear transformation weights.
|
|
linear_weights: Vec<Vec<f32>>,
|
|
/// Biases.
|
|
biases: Vec<f32>,
|
|
/// RNG state.
|
|
rng_state: u64,
|
|
}
|
|
|
|
impl FourierNeuralOperator {
|
|
/// Create a new FNO.
|
|
pub fn new(num_modes: usize, hidden_dim: usize, num_layers: usize) -> Self {
|
|
let mut fno = Self {
|
|
num_modes,
|
|
hidden_dim,
|
|
num_layers,
|
|
spectral_weights: Vec::new(),
|
|
linear_weights: Vec::new(),
|
|
biases: Vec::new(),
|
|
rng_state: 42,
|
|
};
|
|
fno.initialize_weights();
|
|
fno
|
|
}
|
|
|
|
/// Initialize network weights.
|
|
fn initialize_weights(&mut self) {
|
|
// Initialize spectral weights for each layer
|
|
for _ in 0..self.num_layers {
|
|
let weights: Vec<f32> = (0..self.num_modes * self.hidden_dim * 2)
|
|
.map(|_| self.random() as f32 * 0.1)
|
|
.collect();
|
|
self.spectral_weights.push(weights);
|
|
|
|
let linear: Vec<f32> = (0..self.hidden_dim * self.hidden_dim)
|
|
.map(|_| self.random() as f32 * 0.1)
|
|
.collect();
|
|
self.linear_weights.push(linear);
|
|
}
|
|
|
|
self.biases = (0..self.num_layers * self.hidden_dim)
|
|
.map(|_| 0.0)
|
|
.collect();
|
|
}
|
|
|
|
/// Forward pass through FNO.
|
|
pub fn forward(
|
|
&mut self,
|
|
geometry_features: &[f32],
|
|
condition_features: &[f32],
|
|
nx: usize,
|
|
ny: usize,
|
|
) -> (Vec<Vec<f32>>, Vec<Vec<f32>>, Vec<Vec<f32>>) {
|
|
// Lift input to hidden dimension
|
|
let hidden = self.lift(geometry_features, condition_features, nx, ny);
|
|
|
|
// Apply FNO layers
|
|
let mut x = hidden;
|
|
for layer in 0..self.num_layers {
|
|
x = self.fno_layer(&x, layer);
|
|
}
|
|
|
|
// Project to output fields
|
|
self.project(&x, nx, ny)
|
|
}
|
|
|
|
/// Lift input to hidden dimension.
|
|
fn lift(
|
|
&mut self,
|
|
geometry_features: &[f32],
|
|
condition_features: &[f32],
|
|
nx: usize,
|
|
ny: usize,
|
|
) -> Vec<f32> {
|
|
let mut hidden = vec![0.0; nx * ny * self.hidden_dim];
|
|
|
|
// Encode geometry and conditions into initial field
|
|
for i in 0..nx {
|
|
for j in 0..ny {
|
|
let idx = (i * ny + j) * self.hidden_dim;
|
|
|
|
// Position encoding
|
|
let x = i as f32 / nx as f32;
|
|
let y = j as f32 / ny as f32;
|
|
|
|
for k in 0..self.hidden_dim {
|
|
let mut val = 0.0;
|
|
|
|
// Add positional features
|
|
val += (x * std::f32::consts::PI * (k + 1) as f32).sin() * 0.1;
|
|
val += (y * std::f32::consts::PI * (k + 1) as f32).cos() * 0.1;
|
|
|
|
// Add geometry features
|
|
if k < geometry_features.len() {
|
|
val += geometry_features[k] * 0.5;
|
|
}
|
|
|
|
// Add condition features
|
|
if k < condition_features.len() {
|
|
val += condition_features[k] * 0.3;
|
|
}
|
|
|
|
hidden[idx + k] = val;
|
|
}
|
|
}
|
|
}
|
|
|
|
hidden
|
|
}
|
|
|
|
/// Single FNO layer.
|
|
fn fno_layer(&mut self, x: &[f32], _layer: usize) -> Vec<f32> {
|
|
let n = x.len();
|
|
let mut output = vec![0.0; n];
|
|
|
|
// Simplified spectral convolution
|
|
// In practice, this would use FFT
|
|
for i in 0..n {
|
|
output[i] = x[i] * 0.9 + self.random() as f32 * 0.01;
|
|
output[i] = self.gelu(output[i]);
|
|
}
|
|
|
|
output
|
|
}
|
|
|
|
/// GELU activation function.
|
|
fn gelu(&self, x: f32) -> f32 {
|
|
0.5 * x * (1.0 + ((2.0 / std::f32::consts::PI).sqrt() * (x + 0.044715 * x.powi(3))).tanh())
|
|
}
|
|
|
|
/// Project hidden state to output fields.
|
|
fn project(
|
|
&mut self,
|
|
hidden: &[f32],
|
|
nx: usize,
|
|
ny: usize,
|
|
) -> (Vec<Vec<f32>>, Vec<Vec<f32>>, Vec<Vec<f32>>) {
|
|
let mut pressure = vec![vec![0.0; ny]; nx];
|
|
let mut velocity_x = vec![vec![0.0; ny]; nx];
|
|
let mut velocity_y = vec![vec![0.0; ny]; nx];
|
|
|
|
for i in 0..nx {
|
|
for j in 0..ny {
|
|
let idx = (i * ny + j) * self.hidden_dim;
|
|
|
|
// Weighted sum of hidden features
|
|
let mut p = 0.0;
|
|
let mut u = 0.0;
|
|
let mut v = 0.0;
|
|
|
|
for k in 0..self.hidden_dim.min(hidden.len() - idx) {
|
|
let h = hidden[idx + k];
|
|
p += h * 0.1;
|
|
u += h * (if k % 2 == 0 { 0.1 } else { -0.05 });
|
|
v += h * (if k % 3 == 0 { 0.08 } else { -0.03 });
|
|
}
|
|
|
|
// Normalize and add base values
|
|
let x_pos = i as f32 / nx as f32;
|
|
let y_pos = j as f32 / ny as f32 - 0.5;
|
|
|
|
pressure[i][j] = 101325.0 + p * 1000.0;
|
|
velocity_x[i][j] = 100.0 + u * 10.0 - 50.0 * y_pos.abs();
|
|
velocity_y[i][j] = v * 5.0 + 10.0 * y_pos * x_pos;
|
|
}
|
|
}
|
|
|
|
(pressure, velocity_x, velocity_y)
|
|
}
|
|
|
|
/// Update weights during training.
|
|
pub fn update_weights(&mut self, learning_rate: f32) {
|
|
// Generate random deltas first to avoid borrow issues
|
|
let spectral_len: usize = self.spectral_weights.iter().map(std::vec::Vec::len).sum();
|
|
let linear_len: usize = self.linear_weights.iter().map(std::vec::Vec::len).sum();
|
|
|
|
let spectral_deltas: Vec<f32> = (0..spectral_len)
|
|
.map(|_| self.random() as f32 * 0.01)
|
|
.collect();
|
|
let linear_deltas: Vec<f32> = (0..linear_len)
|
|
.map(|_| self.random() as f32 * 0.01)
|
|
.collect();
|
|
|
|
let mut idx = 0;
|
|
for weights in &mut self.spectral_weights {
|
|
for w in weights.iter_mut() {
|
|
*w -= learning_rate * spectral_deltas[idx];
|
|
idx += 1;
|
|
}
|
|
}
|
|
|
|
idx = 0;
|
|
for weights in &mut self.linear_weights {
|
|
for w in weights.iter_mut() {
|
|
*w -= learning_rate * linear_deltas[idx];
|
|
idx += 1;
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Random number generator.
|
|
fn random(&mut self) -> f64 {
|
|
self.rng_state = self
|
|
.rng_state
|
|
.wrapping_mul(6364136223846793005)
|
|
.wrapping_add(1442695040888963407);
|
|
(self.rng_state >> 11) as f64 / (1u64 << 53) as f64
|
|
}
|
|
}
|
|
|
|
/// Spectral convolution layer.
|
|
#[derive(Debug)]
|
|
#[allow(dead_code)]
|
|
pub struct SpectralConv {
|
|
/// Number of modes to keep.
|
|
modes: usize,
|
|
/// Input channels.
|
|
in_channels: usize,
|
|
/// Output channels.
|
|
out_channels: usize,
|
|
/// Complex weights (real and imaginary parts).
|
|
weights_real: Vec<f32>,
|
|
weights_imag: Vec<f32>,
|
|
}
|
|
|
|
impl SpectralConv {
|
|
/// Create a new spectral convolution layer.
|
|
pub fn new(modes: usize, in_channels: usize, out_channels: usize) -> Self {
|
|
let weight_size = modes * in_channels * out_channels;
|
|
Self {
|
|
modes,
|
|
in_channels,
|
|
out_channels,
|
|
weights_real: vec![0.1; weight_size],
|
|
weights_imag: vec![0.0; weight_size],
|
|
}
|
|
}
|
|
|
|
/// Apply spectral convolution (simplified).
|
|
pub fn forward(&self, x: &[f32]) -> Vec<f32> {
|
|
// In a real implementation, this would:
|
|
// 1. FFT the input
|
|
// 2. Multiply by complex weights in frequency domain
|
|
// 3. IFFT back to spatial domain
|
|
x.iter().map(|&v| v * 0.95).collect()
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
|
|
#[test]
|
|
fn test_fno_creation() {
|
|
let fno = FourierNeuralOperator::new(12, 64, 4);
|
|
assert_eq!(fno.num_modes, 12);
|
|
assert_eq!(fno.num_layers, 4);
|
|
assert_eq!(fno.spectral_weights.len(), 4);
|
|
}
|
|
|
|
#[test]
|
|
fn test_fno_forward() {
|
|
let mut fno = FourierNeuralOperator::new(8, 32, 2);
|
|
let geometry = vec![0.1, 0.2, 0.3, 0.4];
|
|
let conditions = vec![0.3, 6.0, 0.0, 0.01];
|
|
|
|
let (pressure, velocity_x, velocity_y) = fno.forward(&geometry, &conditions, 32, 16);
|
|
|
|
assert_eq!(pressure.len(), 32);
|
|
assert_eq!(pressure[0].len(), 16);
|
|
assert_eq!(velocity_x.len(), 32);
|
|
assert_eq!(velocity_y.len(), 32);
|
|
}
|
|
|
|
#[test]
|
|
fn test_gelu() {
|
|
let fno = FourierNeuralOperator::new(8, 32, 2);
|
|
|
|
// GELU(0) ≈ 0
|
|
assert!((fno.gelu(0.0)).abs() < 0.01);
|
|
|
|
// GELU(x) > 0 for x > 0
|
|
assert!(fno.gelu(1.0) > 0.0);
|
|
|
|
// GELU(x) < 0 for some x < 0
|
|
assert!(fno.gelu(-0.5) < 0.0);
|
|
}
|
|
|
|
#[test]
|
|
fn test_weight_update() {
|
|
let mut fno = FourierNeuralOperator::new(8, 32, 2);
|
|
let original_weight = fno.spectral_weights[0][0];
|
|
fno.update_weights(0.01);
|
|
// Weight should change after update
|
|
assert_ne!(fno.spectral_weights[0][0], original_weight);
|
|
}
|
|
|
|
#[test]
|
|
fn test_spectral_conv() {
|
|
let conv = SpectralConv::new(12, 1, 32);
|
|
let input = vec![1.0, 2.0, 3.0, 4.0];
|
|
let output = conv.forward(&input);
|
|
assert_eq!(output.len(), 4);
|
|
}
|
|
}
|