Initial commit
This commit is contained in:
redclawsystems
@@ -0,0 +1,779 @@
|
||||
//! Spectral convolution layers for Fourier Neural Operators.
|
||||
//!
|
||||
//! Spectral convolutions operate in the Fourier domain, enabling efficient
|
||||
//! learning of long-range dependencies in spatial data. These layers are the
|
||||
//! core building blocks of Fourier Neural Operators (FNO).
|
||||
//!
|
||||
//! ## Algorithm
|
||||
//!
|
||||
//! For SpectralConv2d:
|
||||
//! 1. Apply 2D FFT to input: x_freq = FFT2d(x)
|
||||
//! 2. Truncate to n_modes (keep only low frequencies)
|
||||
//! 3. Multiply by learnable weights in frequency domain
|
||||
//! 4. Apply inverse FFT: y = IFFT2d(x_freq * W)
|
||||
//!
|
||||
//! This is dramatically more efficient than spatial convolutions for learning
|
||||
//! global patterns, as multiplication in Fourier space is O(n_modes) vs O(kernel_size).
|
||||
|
||||
use rtx_backend::Backend;
|
||||
use rtx_nn::generic::{GenericModule, GenericModule4D};
|
||||
use rtx_tensor::generic::GenericTensor;
|
||||
use rtx_tensor::{ComplexTensor, Device, Tensor};
|
||||
use std::fmt::Debug;
|
||||
|
||||
/// 1D Spectral Convolution layer.
|
||||
///
|
||||
/// Performs convolution in the Fourier domain for 1D signals.
|
||||
#[derive(Debug)]
|
||||
pub struct SpectralConv1d<B: Backend<FloatElem = f32>> {
|
||||
in_channels: usize,
|
||||
out_channels: usize,
|
||||
n_modes: usize,
|
||||
device: B::Device,
|
||||
// Weights are stored as [in_channels, out_channels, n_modes, 2]
|
||||
// where the last dimension holds [real, imag] components
|
||||
weights: GenericTensor<B, 4>,
|
||||
}
|
||||
|
||||
impl<B: Backend<FloatElem = f32>> SpectralConv1d<B> {
|
||||
/// Create a new 1D spectral convolution layer.
|
||||
///
|
||||
/// # Arguments
|
||||
/// * `in_channels` - Number of input channels
|
||||
/// * `out_channels` - Number of output channels
|
||||
/// * `n_modes` - Number of Fourier modes to keep (frequency truncation)
|
||||
/// * `device` - Device to create the layer on
|
||||
pub fn new(
|
||||
in_channels: usize,
|
||||
out_channels: usize,
|
||||
n_modes: usize,
|
||||
device: &B::Device,
|
||||
) -> Self {
|
||||
// Initialize weights with Xavier uniform: scale * (2 * rand - 1)
|
||||
let scale = (2.0 / (in_channels + out_channels) as f32).sqrt();
|
||||
let weights = GenericTensor::rand([in_channels, out_channels, n_modes, 2], device)
|
||||
.mul_scalar(2.0 * scale)
|
||||
.add_scalar(-scale);
|
||||
|
||||
Self {
|
||||
in_channels,
|
||||
out_channels,
|
||||
n_modes,
|
||||
device: device.clone(),
|
||||
weights,
|
||||
}
|
||||
}
|
||||
|
||||
/// Get the number of Fourier modes.
|
||||
pub fn n_modes(&self) -> usize {
|
||||
self.n_modes
|
||||
}
|
||||
}
|
||||
|
||||
impl<B: Backend<FloatElem = f32>> GenericModule<B> for SpectralConv1d<B> {
|
||||
fn forward(&self, _input: &GenericTensor<B, 2>) -> GenericTensor<B, 2> {
|
||||
panic!("Use forward_3d for 1D spectral convolution");
|
||||
}
|
||||
|
||||
fn device(&self) -> &B::Device {
|
||||
&self.device
|
||||
}
|
||||
}
|
||||
|
||||
/// 2D Spectral Convolution layer.
|
||||
///
|
||||
/// Performs convolution in the Fourier domain for 2D spatial data.
|
||||
/// This is the core component of the Fourier Neural Operator (FNO).
|
||||
///
|
||||
/// # Algorithm (neuraloperator v2.0 compatible)
|
||||
///
|
||||
/// Given input x of shape [batch, in_channels, height, width]:
|
||||
/// 1. Apply 2D FFT to get frequency representation
|
||||
/// 2. Apply weights1 to upper portion of frequency domain (low frequencies)
|
||||
/// 3. Apply weights2 to lower portion of frequency domain (high frequencies in h)
|
||||
/// 4. Apply inverse 2D FFT to return to spatial domain
|
||||
///
|
||||
/// The learnable parameters are two sets of complex weights in frequency space,
|
||||
/// matching neuraloperator v2.0's dual-weight architecture.
|
||||
#[derive(Debug)]
|
||||
pub struct SpectralConv2d<B: Backend<FloatElem = f32>> {
|
||||
in_channels: usize,
|
||||
out_channels: usize,
|
||||
n_modes_h: usize,
|
||||
n_modes_w: usize,
|
||||
device: B::Device,
|
||||
/// Complex weights for upper frequency portion: [in_ch, out_ch, n_modes_h, n_modes_w, 2]
|
||||
weights1: GenericTensor<B, 5>,
|
||||
/// Complex weights for lower frequency portion: [in_ch, out_ch, n_modes_h, n_modes_w, 2]
|
||||
weights2: GenericTensor<B, 5>,
|
||||
}
|
||||
|
||||
impl<B: Backend<FloatElem = f32>> SpectralConv2d<B> {
|
||||
/// Create a new 2D spectral convolution layer.
|
||||
///
|
||||
/// # Arguments
|
||||
/// * `in_channels` - Number of input channels
|
||||
/// * `out_channels` - Number of output channels
|
||||
/// * `n_modes_h` - Number of Fourier modes along height dimension
|
||||
/// * `n_modes_w` - Number of Fourier modes along width dimension
|
||||
/// * `device` - Device to create the layer on
|
||||
pub fn new(
|
||||
in_channels: usize,
|
||||
out_channels: usize,
|
||||
n_modes_h: usize,
|
||||
n_modes_w: usize,
|
||||
device: &B::Device,
|
||||
) -> Self {
|
||||
// Xavier initialization for complex weights: scale * (2 * rand - 1)
|
||||
let scale = (2.0 / (in_channels + out_channels) as f32).sqrt();
|
||||
|
||||
// weights1: for upper portion of frequency domain
|
||||
let weights1 =
|
||||
GenericTensor::rand([in_channels, out_channels, n_modes_h, n_modes_w, 2], device)
|
||||
.mul_scalar(2.0 * scale)
|
||||
.add_scalar(-scale);
|
||||
|
||||
// weights2: for lower portion of frequency domain
|
||||
let weights2 =
|
||||
GenericTensor::rand([in_channels, out_channels, n_modes_h, n_modes_w, 2], device)
|
||||
.mul_scalar(2.0 * scale)
|
||||
.add_scalar(-scale);
|
||||
|
||||
Self {
|
||||
in_channels,
|
||||
out_channels,
|
||||
n_modes_h,
|
||||
n_modes_w,
|
||||
device: device.clone(),
|
||||
weights1,
|
||||
weights2,
|
||||
}
|
||||
}
|
||||
|
||||
/// Get the number of Fourier modes in height dimension.
|
||||
pub fn n_modes_h(&self) -> usize {
|
||||
self.n_modes_h
|
||||
}
|
||||
|
||||
/// Get the number of Fourier modes in width dimension.
|
||||
pub fn n_modes_w(&self) -> usize {
|
||||
self.n_modes_w
|
||||
}
|
||||
|
||||
/// Get a reference to weights1 tensor.
|
||||
pub fn weights1(&self) -> &GenericTensor<B, 5> {
|
||||
&self.weights1
|
||||
}
|
||||
|
||||
/// Get a reference to weights2 tensor.
|
||||
pub fn weights2(&self) -> &GenericTensor<B, 5> {
|
||||
&self.weights2
|
||||
}
|
||||
|
||||
/// Get a reference to the weights tensor (returns weights1 for compatibility).
|
||||
#[deprecated(note = "Use weights1() or weights2() instead")]
|
||||
pub fn weights(&self) -> &GenericTensor<B, 5> {
|
||||
&self.weights1
|
||||
}
|
||||
|
||||
/// Extract weights1 as separate real and imaginary parts.
|
||||
///
|
||||
/// Returns `(real, imag)` where each is a `Vec<f32>` of shape
|
||||
/// `[in_channels * out_channels * n_modes_h * n_modes_w]`.
|
||||
pub fn weights1_real_imag(&self) -> (Vec<f32>, Vec<f32>) {
|
||||
let data = self.weights1.to_vec();
|
||||
let n_complex = self.in_channels * self.out_channels * self.n_modes_h * self.n_modes_w;
|
||||
let mut real = Vec::with_capacity(n_complex);
|
||||
let mut imag = Vec::with_capacity(n_complex);
|
||||
|
||||
for i in 0..n_complex {
|
||||
real.push(data[i * 2]);
|
||||
imag.push(data[i * 2 + 1]);
|
||||
}
|
||||
|
||||
(real, imag)
|
||||
}
|
||||
|
||||
/// Extract weights2 as separate real and imaginary parts.
|
||||
///
|
||||
/// Returns `(real, imag)` where each is a `Vec<f32>` of shape
|
||||
/// `[in_channels * out_channels * n_modes_h * n_modes_w]`.
|
||||
pub fn weights2_real_imag(&self) -> (Vec<f32>, Vec<f32>) {
|
||||
let data = self.weights2.to_vec();
|
||||
let n_complex = self.in_channels * self.out_channels * self.n_modes_h * self.n_modes_w;
|
||||
let mut real = Vec::with_capacity(n_complex);
|
||||
let mut imag = Vec::with_capacity(n_complex);
|
||||
|
||||
for i in 0..n_complex {
|
||||
real.push(data[i * 2]);
|
||||
imag.push(data[i * 2 + 1]);
|
||||
}
|
||||
|
||||
(real, imag)
|
||||
}
|
||||
|
||||
/// Get the input channels.
|
||||
pub fn in_channels(&self) -> usize {
|
||||
self.in_channels
|
||||
}
|
||||
|
||||
/// Get the output channels.
|
||||
pub fn out_channels(&self) -> usize {
|
||||
self.out_channels
|
||||
}
|
||||
|
||||
/// Create a spectral convolution layer with pre-defined weights.
|
||||
///
|
||||
/// # Arguments
|
||||
/// * `weights1_real` - Real part of weights1 [in_ch, out_ch, n_modes_h, n_modes_w]
|
||||
/// * `weights1_imag` - Imaginary part of weights1
|
||||
/// * `in_channels` - Number of input channels
|
||||
/// * `out_channels` - Number of output channels
|
||||
/// * `n_modes_h` - Number of Fourier modes along height
|
||||
/// * `n_modes_w` - Number of Fourier modes along width
|
||||
/// * `device` - Device to create the layer on
|
||||
///
|
||||
/// Note: This creates weights2 as a copy of weights1 for backwards compatibility.
|
||||
/// Use `from_weights_dual()` for full neuraloperator v2.0 compatibility.
|
||||
pub fn from_weights(
|
||||
weights1_real: &[f32],
|
||||
weights1_imag: &[f32],
|
||||
in_channels: usize,
|
||||
out_channels: usize,
|
||||
n_modes_h: usize,
|
||||
n_modes_w: usize,
|
||||
device: &B::Device,
|
||||
) -> Self {
|
||||
// For backwards compatibility, use the same weights for both
|
||||
Self::from_weights_dual(
|
||||
weights1_real,
|
||||
weights1_imag,
|
||||
weights1_real,
|
||||
weights1_imag,
|
||||
in_channels,
|
||||
out_channels,
|
||||
n_modes_h,
|
||||
n_modes_w,
|
||||
device,
|
||||
)
|
||||
}
|
||||
|
||||
/// Create a spectral convolution layer with dual pre-defined weights.
|
||||
///
|
||||
/// This matches neuraloperator v2.0 architecture with two weight matrices.
|
||||
///
|
||||
/// # Arguments
|
||||
/// * `weights1_real` - Real part of weights1 [in_ch, out_ch, n_modes_h, n_modes_w]
|
||||
/// * `weights1_imag` - Imaginary part of weights1
|
||||
/// * `weights2_real` - Real part of weights2 [in_ch, out_ch, n_modes_h, n_modes_w]
|
||||
/// * `weights2_imag` - Imaginary part of weights2
|
||||
/// * `in_channels` - Number of input channels
|
||||
/// * `out_channels` - Number of output channels
|
||||
/// * `n_modes_h` - Number of Fourier modes along height
|
||||
/// * `n_modes_w` - Number of Fourier modes along width
|
||||
/// * `device` - Device to create the layer on
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
pub fn from_weights_dual(
|
||||
weights1_real: &[f32],
|
||||
weights1_imag: &[f32],
|
||||
weights2_real: &[f32],
|
||||
weights2_imag: &[f32],
|
||||
in_channels: usize,
|
||||
out_channels: usize,
|
||||
n_modes_h: usize,
|
||||
n_modes_w: usize,
|
||||
device: &B::Device,
|
||||
) -> Self {
|
||||
let total_complex = in_channels * out_channels * n_modes_h * n_modes_w;
|
||||
|
||||
// Interleave weights1
|
||||
let mut interleaved1 = Vec::with_capacity(total_complex * 2);
|
||||
for i in 0..total_complex {
|
||||
interleaved1.push(weights1_real[i]);
|
||||
interleaved1.push(weights1_imag[i]);
|
||||
}
|
||||
let weights1 = GenericTensor::from_slice(
|
||||
&interleaved1,
|
||||
[in_channels, out_channels, n_modes_h, n_modes_w, 2],
|
||||
device,
|
||||
);
|
||||
|
||||
// Interleave weights2
|
||||
let mut interleaved2 = Vec::with_capacity(total_complex * 2);
|
||||
for i in 0..total_complex {
|
||||
interleaved2.push(weights2_real[i]);
|
||||
interleaved2.push(weights2_imag[i]);
|
||||
}
|
||||
let weights2 = GenericTensor::from_slice(
|
||||
&interleaved2,
|
||||
[in_channels, out_channels, n_modes_h, n_modes_w, 2],
|
||||
device,
|
||||
);
|
||||
|
||||
Self {
|
||||
in_channels,
|
||||
out_channels,
|
||||
n_modes_h,
|
||||
n_modes_w,
|
||||
device: device.clone(),
|
||||
weights1,
|
||||
weights2,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl<B: Backend<FloatElem = f32>> GenericModule<B> for SpectralConv2d<B> {
|
||||
fn forward(&self, _input: &GenericTensor<B, 2>) -> GenericTensor<B, 2> {
|
||||
panic!("Use forward_4d for 2D spectral convolution");
|
||||
}
|
||||
|
||||
fn device(&self) -> &B::Device {
|
||||
&self.device
|
||||
}
|
||||
}
|
||||
|
||||
impl<B: Backend<FloatElem = f32>> GenericModule4D<B> for SpectralConv2d<B> {
|
||||
fn forward_4d(&self, input: &GenericTensor<B, 4>) -> GenericTensor<B, 4> {
|
||||
let shape = input.shape();
|
||||
let batch = shape[0];
|
||||
let in_ch = shape[1];
|
||||
let height = shape[2];
|
||||
let width = shape[3];
|
||||
|
||||
// Validate input channels match
|
||||
assert_eq!(
|
||||
in_ch, self.in_channels,
|
||||
"Input channels {} do not match layer in_channels {}",
|
||||
in_ch, self.in_channels
|
||||
);
|
||||
|
||||
// Convert GenericTensor to Vec<f32> for processing
|
||||
let input_data = input.to_vec();
|
||||
|
||||
// Allocate output buffer [batch, out_channels, height, width]
|
||||
let output_size = batch * self.out_channels * height * width;
|
||||
let mut output_data = vec![0.0f32; output_size];
|
||||
|
||||
// Use CPU device for ComplexTensor operations
|
||||
let cpu_device = Device::cpu();
|
||||
|
||||
// Get weight data for both weight matrices
|
||||
let weight1_data = self.weights1.to_vec();
|
||||
let weight2_data = self.weights2.to_vec();
|
||||
|
||||
// Process each batch
|
||||
for b in 0..batch {
|
||||
// Process each output channel
|
||||
for out_c in 0..self.out_channels {
|
||||
// Accumulate contributions from all input channels
|
||||
let mut channel_sum_real = Tensor::zeros(vec![height, width], &cpu_device)
|
||||
.expect("Failed to create zeros");
|
||||
let mut channel_sum_imag = Tensor::zeros(vec![height, width], &cpu_device)
|
||||
.expect("Failed to create zeros");
|
||||
|
||||
for in_c in 0..self.in_channels {
|
||||
// Extract 2D slice [height, width] for this batch and input channel
|
||||
let slice_start = (b * in_ch + in_c) * height * width;
|
||||
let slice_end = slice_start + (height * width);
|
||||
let slice_data = input_data[slice_start..slice_end].to_vec();
|
||||
|
||||
// Create real tensor and convert to ComplexTensor
|
||||
let real_tensor =
|
||||
Tensor::from_data(slice_data, vec![height, width], &cpu_device)
|
||||
.expect("Failed to create tensor from data");
|
||||
|
||||
let complex_input = ComplexTensor::<f32>::from_real(real_tensor)
|
||||
.expect("Failed to create complex tensor");
|
||||
|
||||
// Apply 2D FFT
|
||||
let fft_result = complex_input.fft2d().expect("Failed to perform FFT2D");
|
||||
|
||||
// Truncate to n_modes
|
||||
let modes_h = self.n_modes_h.min(height);
|
||||
let modes_w = self.n_modes_w.min(width);
|
||||
|
||||
// Extract FFT data
|
||||
let fft_real_data = fft_result
|
||||
.real()
|
||||
.to_cpu()
|
||||
.expect("Failed to get FFT real data");
|
||||
let fft_imag_data = fft_result
|
||||
.imag()
|
||||
.to_cpu()
|
||||
.expect("Failed to get FFT imag data");
|
||||
|
||||
// Weight offset for [in_c, out_c, :, :, :]
|
||||
let weight_slice_start =
|
||||
(in_c * self.out_channels + out_c) * self.n_modes_h * self.n_modes_w * 2;
|
||||
|
||||
// Perform complex multiplication in frequency domain
|
||||
let mut freq_result_real = vec![0.0f32; height * width];
|
||||
let mut freq_result_imag = vec![0.0f32; height * width];
|
||||
|
||||
// Apply weights1 to upper portion (low frequencies in h)
|
||||
// Positions [0..modes_h, 0..modes_w]
|
||||
for h in 0..modes_h {
|
||||
for w in 0..modes_w {
|
||||
let fft_idx = h * width + w;
|
||||
|
||||
// Get complex FFT value at this position
|
||||
let fft_r = fft_real_data[fft_idx];
|
||||
let fft_i = fft_imag_data[fft_idx];
|
||||
|
||||
// Get complex weight from weights1
|
||||
let weight_idx = weight_slice_start + (h * self.n_modes_w + w) * 2;
|
||||
let weight_r = weight1_data[weight_idx];
|
||||
let weight_i = weight1_data[weight_idx + 1];
|
||||
|
||||
// Complex multiplication: (a+bi) * (c+di) = (ac-bd) + (ad+bc)i
|
||||
let result_r = fft_r * weight_r - fft_i * weight_i;
|
||||
let result_i = fft_r * weight_i + fft_i * weight_r;
|
||||
|
||||
freq_result_real[fft_idx] = result_r;
|
||||
freq_result_imag[fft_idx] = result_i;
|
||||
}
|
||||
}
|
||||
|
||||
// Apply weights2 to lower portion (high frequencies in h)
|
||||
// Positions [height-modes_h..height, 0..modes_w]
|
||||
// This corresponds to negative frequencies in the FFT output
|
||||
if height > modes_h {
|
||||
for h in 0..modes_h {
|
||||
let fft_h = height - modes_h + h;
|
||||
for w in 0..modes_w {
|
||||
let fft_idx = fft_h * width + w;
|
||||
|
||||
// Get complex FFT value at this position
|
||||
let fft_r = fft_real_data[fft_idx];
|
||||
let fft_i = fft_imag_data[fft_idx];
|
||||
|
||||
// Get complex weight from weights2
|
||||
let weight_idx = weight_slice_start + (h * self.n_modes_w + w) * 2;
|
||||
let weight_r = weight2_data[weight_idx];
|
||||
let weight_i = weight2_data[weight_idx + 1];
|
||||
|
||||
// Complex multiplication
|
||||
let result_r = fft_r * weight_r - fft_i * weight_i;
|
||||
let result_i = fft_r * weight_i + fft_i * weight_r;
|
||||
|
||||
freq_result_real[fft_idx] = result_r;
|
||||
freq_result_imag[fft_idx] = result_i;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Create ComplexTensor from frequency domain result
|
||||
let freq_real_tensor =
|
||||
Tensor::from_data(freq_result_real, vec![height, width], &cpu_device)
|
||||
.expect("Failed to create frequency real tensor");
|
||||
|
||||
let freq_imag_tensor =
|
||||
Tensor::from_data(freq_result_imag, vec![height, width], &cpu_device)
|
||||
.expect("Failed to create frequency imag tensor");
|
||||
|
||||
let freq_complex =
|
||||
ComplexTensor::<f32>::from_real_imag(freq_real_tensor, freq_imag_tensor)
|
||||
.expect("Failed to create frequency complex tensor");
|
||||
|
||||
// Apply inverse 2D FFT
|
||||
let ifft_result = freq_complex.ifft2d().expect("Failed to perform IFFT2D");
|
||||
|
||||
// Accumulate to channel sum
|
||||
let ifft_real = ifft_result.real();
|
||||
let ifft_imag = ifft_result.imag();
|
||||
|
||||
channel_sum_real = channel_sum_real
|
||||
.add(ifft_real)
|
||||
.expect("Failed to add real part");
|
||||
channel_sum_imag = channel_sum_imag
|
||||
.add(ifft_imag)
|
||||
.expect("Failed to add imag part");
|
||||
}
|
||||
|
||||
// Extract real part and copy to output
|
||||
let final_real_data = channel_sum_real
|
||||
.to_cpu()
|
||||
.expect("Failed to get final real data");
|
||||
|
||||
let output_offset = (b * self.out_channels + out_c) * height * width;
|
||||
for i in 0..final_real_data.len() {
|
||||
output_data[output_offset + i] = final_real_data[i];
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Convert back to GenericTensor
|
||||
GenericTensor::from_slice(
|
||||
&output_data,
|
||||
[batch, self.out_channels, height, width],
|
||||
&self.device,
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use rtx_backend_cpu::{CpuBackend, CpuDevice};
|
||||
|
||||
// ==================== SpectralConv1d Tests ====================
|
||||
|
||||
#[test]
|
||||
fn test_spectral_conv1d_creation() {
|
||||
let device = CpuDevice::new();
|
||||
let conv = SpectralConv1d::<CpuBackend>::new(4, 8, 12, &device);
|
||||
|
||||
assert_eq!(conv.in_channels, 4);
|
||||
assert_eq!(conv.out_channels, 8);
|
||||
assert_eq!(conv.n_modes(), 12);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_spectral_conv1d_weight_shape() {
|
||||
let device = CpuDevice::new();
|
||||
let conv = SpectralConv1d::<CpuBackend>::new(3, 6, 8, &device);
|
||||
|
||||
// Weights should be [in_ch, out_ch, n_modes, 2]
|
||||
assert_eq!(conv.weights.shape(), [3, 6, 8, 2]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_spectral_conv1d_weight_initialization() {
|
||||
let device = CpuDevice::new();
|
||||
let conv = SpectralConv1d::<CpuBackend>::new(4, 8, 10, &device);
|
||||
|
||||
// Check Xavier initialization bounds
|
||||
let scale = (2.0_f32 / (4.0_f32 + 8.0_f32)).sqrt();
|
||||
let data = conv.weights.to_vec();
|
||||
|
||||
// All values should be in range [-scale, scale]
|
||||
for &val in &data {
|
||||
assert!(
|
||||
val >= -scale && val <= scale,
|
||||
"Value {} outside Xavier range [{}, {}]",
|
||||
val,
|
||||
-scale,
|
||||
scale
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
// ==================== SpectralConv2d Tests ====================
|
||||
|
||||
#[test]
|
||||
fn test_spectral_conv2d_creation() {
|
||||
let device = CpuDevice::new();
|
||||
let conv = SpectralConv2d::<CpuBackend>::new(3, 16, 12, 12, &device);
|
||||
|
||||
assert_eq!(conv.in_channels, 3);
|
||||
assert_eq!(conv.out_channels, 16);
|
||||
assert_eq!(conv.n_modes_h(), 12);
|
||||
assert_eq!(conv.n_modes_w(), 12);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_spectral_conv2d_weight_shape() {
|
||||
let device = CpuDevice::new();
|
||||
let conv = SpectralConv2d::<CpuBackend>::new(4, 8, 6, 10, &device);
|
||||
|
||||
// Both weight matrices should be [in_ch, out_ch, n_modes_h, n_modes_w, 2]
|
||||
assert_eq!(conv.weights1().shape(), [4, 8, 6, 10, 2]);
|
||||
assert_eq!(conv.weights2().shape(), [4, 8, 6, 10, 2]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_spectral_conv2d_weight_initialization() {
|
||||
let device = CpuDevice::new();
|
||||
let conv = SpectralConv2d::<CpuBackend>::new(3, 12, 8, 8, &device);
|
||||
|
||||
// Check Xavier initialization bounds for both weight matrices
|
||||
let scale = (2.0_f32 / (3.0_f32 + 12.0_f32)).sqrt();
|
||||
|
||||
for weights in [conv.weights1(), conv.weights2()] {
|
||||
let data = weights.to_vec();
|
||||
// All values should be in range [-scale, scale]
|
||||
for &val in &data {
|
||||
assert!(
|
||||
val >= -scale && val <= scale,
|
||||
"Value {} outside Xavier range [{}, {}]",
|
||||
val,
|
||||
-scale,
|
||||
scale
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_spectral_conv2d_forward_shape() {
|
||||
let device = CpuDevice::new();
|
||||
let conv = SpectralConv2d::<CpuBackend>::new(4, 8, 12, 12, &device);
|
||||
|
||||
let input = GenericTensor::randn([2, 4, 64, 64], &device);
|
||||
let output = conv.forward_4d(&input);
|
||||
|
||||
// Output should have shape [batch, out_channels, height, width]
|
||||
assert_eq!(output.shape(), [2, 8, 64, 64]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_spectral_conv2d_preserves_spatial_dims() {
|
||||
let device = CpuDevice::new();
|
||||
let conv = SpectralConv2d::<CpuBackend>::new(1, 32, 8, 8, &device);
|
||||
|
||||
let input = GenericTensor::randn([4, 1, 32, 48], &device);
|
||||
let output = conv.forward_4d(&input);
|
||||
|
||||
// Spatial dimensions should be preserved
|
||||
assert_eq!(output.shape()[2], 32); // height
|
||||
assert_eq!(output.shape()[3], 48); // width
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_spectral_conv2d_different_modes() {
|
||||
let device = CpuDevice::new();
|
||||
let conv = SpectralConv2d::<CpuBackend>::new(2, 4, 6, 10, &device);
|
||||
|
||||
let input = GenericTensor::randn([1, 2, 64, 64], &device);
|
||||
let output = conv.forward_4d(&input);
|
||||
|
||||
// Output channels should match conv.out_channels
|
||||
assert_eq!(output.shape(), [1, 4, 64, 64]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_spectral_conv2d_batch_independence() {
|
||||
let device = CpuDevice::new();
|
||||
let conv = SpectralConv2d::<CpuBackend>::new(3, 6, 8, 8, &device);
|
||||
|
||||
// Single sample
|
||||
let input1 = GenericTensor::randn([1, 3, 32, 32], &device);
|
||||
let output1 = conv.forward_4d(&input1);
|
||||
|
||||
// Batch of 4
|
||||
let input4 = GenericTensor::randn([4, 3, 32, 32], &device);
|
||||
let output4 = conv.forward_4d(&input4);
|
||||
|
||||
// Batch dimension should be the only difference
|
||||
assert_eq!(output1.shape()[1..], output4.shape()[1..]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_spectral_conv2d_mode_truncation() {
|
||||
let device = CpuDevice::new();
|
||||
|
||||
// Small mode count (strong low-pass filtering)
|
||||
let conv_small = SpectralConv2d::<CpuBackend>::new(1, 1, 4, 4, &device);
|
||||
|
||||
// Large mode count (weak filtering)
|
||||
let conv_large = SpectralConv2d::<CpuBackend>::new(1, 1, 30, 30, &device);
|
||||
|
||||
let input = GenericTensor::randn([2, 1, 64, 64], &device);
|
||||
|
||||
let output_small = conv_small.forward_4d(&input);
|
||||
let output_large = conv_large.forward_4d(&input);
|
||||
|
||||
// Both should preserve spatial dimensions
|
||||
assert_eq!(output_small.shape(), [2, 1, 64, 64]);
|
||||
assert_eq!(output_large.shape(), [2, 1, 64, 64]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
#[should_panic(expected = "Use forward_4d for 2D spectral convolution")]
|
||||
fn test_spectral_conv2d_forward_panics() {
|
||||
let device = CpuDevice::new();
|
||||
let conv = SpectralConv2d::<CpuBackend>::new(2, 4, 8, 8, &device);
|
||||
|
||||
let input = GenericTensor::randn([4, 2], &device);
|
||||
let _ = conv.forward(&input);
|
||||
}
|
||||
|
||||
// ==================== forward_4d Implementation Tests ====================
|
||||
|
||||
#[test]
|
||||
fn test_forward_4d_output_channels() {
|
||||
let device = CpuDevice::new();
|
||||
let conv = SpectralConv2d::<CpuBackend>::new(3, 8, 12, 12, &device);
|
||||
|
||||
let input = GenericTensor::randn([2, 3, 64, 64], &device);
|
||||
let output = conv.forward_4d(&input);
|
||||
|
||||
// Output should transform from in_channels to out_channels
|
||||
assert_eq!(output.shape(), [2, 8, 64, 64]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_forward_4d_non_zero_output() {
|
||||
let device = CpuDevice::new();
|
||||
let conv = SpectralConv2d::<CpuBackend>::new(2, 4, 8, 8, &device);
|
||||
|
||||
// Create non-zero input
|
||||
let input = GenericTensor::ones([1, 2, 32, 32], &device);
|
||||
let output = conv.forward_4d(&input);
|
||||
|
||||
// FFT-based transformation should produce 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,
|
||||
"FFT transformation should produce non-zero outputs"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_forward_4d_complex_multiplication() {
|
||||
let device = CpuDevice::new();
|
||||
|
||||
// Use small modes for easier verification
|
||||
let conv = SpectralConv2d::<CpuBackend>::new(1, 1, 4, 4, &device);
|
||||
|
||||
// Create a simple input pattern (constant)
|
||||
let input = GenericTensor::ones([1, 1, 16, 16], &device);
|
||||
let output = conv.forward_4d(&input);
|
||||
|
||||
// Output shape should be preserved
|
||||
assert_eq!(output.shape(), [1, 1, 16, 16]);
|
||||
|
||||
// The output should be different from input (due to complex multiplication)
|
||||
let input_data = input.to_vec();
|
||||
let output_data = output.to_vec();
|
||||
|
||||
let mut differences = 0;
|
||||
for (inp, out) in input_data.iter().zip(output_data.iter()) {
|
||||
if (inp - out).abs() > 1e-5 {
|
||||
differences += 1;
|
||||
}
|
||||
}
|
||||
|
||||
assert!(
|
||||
differences > 0,
|
||||
"Complex multiplication in frequency domain should modify values"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_forward_4d_multiple_batches() {
|
||||
let device = CpuDevice::new();
|
||||
let conv = SpectralConv2d::<CpuBackend>::new(2, 3, 8, 8, &device);
|
||||
|
||||
let input = GenericTensor::randn([5, 2, 32, 32], &device);
|
||||
let output = conv.forward_4d(&input);
|
||||
|
||||
// Should handle multiple batches correctly
|
||||
assert_eq!(output.shape(), [5, 3, 32, 32]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_forward_4d_different_spatial_sizes() {
|
||||
let device = CpuDevice::new();
|
||||
let conv = SpectralConv2d::<CpuBackend>::new(1, 2, 4, 6, &device);
|
||||
|
||||
// Non-square spatial dimensions
|
||||
let input = GenericTensor::randn([1, 1, 24, 36], &device);
|
||||
let output = conv.forward_4d(&input);
|
||||
|
||||
// Spatial dimensions should be preserved
|
||||
assert_eq!(output.shape(), [1, 2, 24, 36]);
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user