Initial commit
This commit is contained in:
redclawsystems
@@ -0,0 +1,317 @@
|
||||
//! 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);
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user