rustytorch/crates/specialized/rtx-neural-operator/src/layers.rs

//! Helper layers for neural operators.
//!
//! This module provides common layers used in neural operator architectures:
//! - **GridPositionalEncoding**: Adds 2D grid coordinates to input
//! - **Lifting**: Projects input to a higher-dimensional latent space (simple)
//! - **LiftingMLP**: Two-layer MLP for lifting (neuraloperator v2.0 compatible)
//! - **Projection**: Projects latent representation back to output space

use rtx_backend::Backend;
use rtx_nn::generic::{GenericLinear, GenericModule, GenericModule4D};
use rtx_tensor::generic::GenericTensor;
use std::fmt::Debug;

/// Grid positional encoding: adds 2D spatial coordinates to input.
///
/// Creates a meshgrid of x and y coordinates in [-1, 1] and concatenates
/// them with the input tensor. This is used in neuraloperator v2.0 to provide
/// spatial information to the network.
///
/// Input: `[B, C, H, W]` → Output: `[B, C+2, H, W]`
#[derive(Debug, Clone)]
pub struct GridPositionalEncoding<B: Backend<FloatElem = f32>> {
    device: B::Device,
}

impl<B: Backend<FloatElem = f32>> GridPositionalEncoding<B> {
    /// Create a new grid positional encoding layer.
    pub fn new(device: &B::Device) -> Self {
        Self {
            device: device.clone(),
        }
    }

    /// Apply positional encoding to input.
    ///
    /// Creates x and y coordinate grids normalized to [-1, 1] and concatenates
    /// them with the input along the channel dimension.
    pub fn forward_4d(&self, input: &GenericTensor<B, 4>) -> GenericTensor<B, 4> {
        let shape = input.shape();
        let batch = shape[0];
        let channels = shape[1];
        let height = shape[2];
        let width = shape[3];

        // Create x coordinates: [-1, 1] across width
        // Shape: [1, 1, 1, width] -> broadcast to [batch, 1, height, width]
        let mut x_coords = vec![0.0f32; batch * height * width];
        for b in 0..batch {
            for h in 0..height {
                for w in 0..width {
                    let x = 2.0 * (w as f32) / (width as f32 - 1.0).max(1.0) - 1.0;
                    x_coords[b * height * width + h * width + w] = x;
                }
            }
        }
        // Create y coordinates: [-1, 1] across height
        // Shape: [1, 1, height, 1] -> broadcast to [batch, 1, height, width]
        let mut y_coords = vec![0.0f32; batch * height * width];
        for b in 0..batch {
            for h in 0..height {
                let y = 2.0 * (h as f32) / (height as f32 - 1.0).max(1.0) - 1.0;
                for w in 0..width {
                    y_coords[b * height * width + h * width + w] = y;
                }
            }
        }

        // Concatenate: [B, C, H, W] + [B, 1, H, W] + [B, 1, H, W] = [B, C+2, H, W]
        // Manual concatenation by creating output tensor
        let out_channels = channels + 2;
        let mut output_data = vec![0.0f32; batch * out_channels * height * width];

        // Copy input channels
        let input_data = input.to_vec();
        for b in 0..batch {
            for c in 0..channels {
                for h in 0..height {
                    for w in 0..width {
                        let in_idx =
                            b * channels * height * width + c * height * width + h * width + w;
                        let out_idx =
                            b * out_channels * height * width + c * height * width + h * width + w;
                        output_data[out_idx] = input_data[in_idx];
                    }
                }
            }
        }

        // Copy x coordinates at channel index 'channels'
        for b in 0..batch {
            for h in 0..height {
                for w in 0..width {
                    let coord_idx = b * height * width + h * width + w;
                    let out_idx = b * out_channels * height * width
                        + channels * height * width
                        + h * width
                        + w;
                    output_data[out_idx] = x_coords[coord_idx];
                }
            }
        }

        // Copy y coordinates at channel index 'channels + 1'
        for b in 0..batch {
            for h in 0..height {
                for w in 0..width {
                    let coord_idx = b * height * width + h * width + w;
                    let out_idx = b * out_channels * height * width
                        + (channels + 1) * height * width
                        + h * width
                        + w;
                    output_data[out_idx] = y_coords[coord_idx];
                }
            }
        }

        GenericTensor::from_slice(
            &output_data,
            [batch, out_channels, height, width],
            &self.device,
        )
    }

    /// Get the device this layer is on.
    pub fn device(&self) -> &B::Device {
        &self.device
    }
}

/// Two-layer MLP for lifting (neuraloperator v2.0 compatible).
///
/// This layer matches the lifting architecture in neuraloperator v2.0:
/// - First layer: `in_channels → hidden_dim` with GELU activation
/// - Second layer: `hidden_dim → out_channels`
///
/// Where `hidden_dim = 2 * out_channels` (expansion factor of 2).
#[derive(Debug)]
pub struct LiftingMLP<B: Backend<FloatElem = f32>> {
    fc1: GenericLinear<B>,
    fc2: GenericLinear<B>,
    device: B::Device,
}

impl<B: Backend<FloatElem = f32>> LiftingMLP<B> {
    /// Create a new lifting MLP.
    ///
    /// # Arguments
    /// * `in_channels` - Number of input channels (should include positional encoding)
    /// * `out_channels` - Number of output channels (hidden dimension / width)
    /// * `device` - Device to create the layer on
    pub fn new(in_channels: usize, out_channels: usize, device: &B::Device) -> Self {
        let hidden_dim = 2 * out_channels;
        Self {
            fc1: GenericLinear::new(in_channels, hidden_dim, true, device),
            fc2: GenericLinear::new(hidden_dim, out_channels, true, device),
            device: device.clone(),
        }
    }

    /// Get weight and bias data from the first layer (fc1).
    ///
    /// Returns `(weight, bias)` as `(Vec<f32>, Vec<f32>)`.
    pub fn fc1_weights(&self) -> (Vec<f32>, Vec<f32>) {
        let weight = self.fc1.weight().to_vec();
        let bias = self
            .fc1
            .bias()
            .map(rtx_tensor::GenericTensor::to_vec)
            .unwrap_or_default();
        (weight, bias)
    }

    /// Get weight and bias data from the second layer (fc2).
    ///
    /// Returns `(weight, bias)` as `(Vec<f32>, Vec<f32>)`.
    pub fn fc2_weights(&self) -> (Vec<f32>, Vec<f32>) {
        let weight = self.fc2.weight().to_vec();
        let bias = self
            .fc2
            .bias()
            .map(rtx_tensor::GenericTensor::to_vec)
            .unwrap_or_default();
        (weight, bias)
    }

    /// Get the input dimension (in_channels).
    pub fn in_features(&self) -> usize {
        self.fc1.in_features()
    }

    /// Get the hidden dimension.
    pub fn hidden_dim(&self) -> usize {
        self.fc1.out_features()
    }

    /// Get the output dimension (out_channels / width).
    pub fn out_features(&self) -> usize {
        self.fc2.out_features()
    }

    /// Create a lifting MLP with pre-defined weights.
    ///
    /// # Arguments
    /// * `fc1_weight` - First layer weight [hidden_dim, in_channels]
    /// * `fc1_bias` - First layer bias [hidden_dim]
    /// * `fc2_weight` - Second layer weight [out_channels, hidden_dim]
    /// * `fc2_bias` - Second layer bias [out_channels]
    /// * `in_channels` - Number of input channels
    /// * `out_channels` - Number of output channels
    /// * `device` - Device to create the layer on
    pub fn from_weights(
        fc1_weight: &[f32],
        fc1_bias: &[f32],
        fc2_weight: &[f32],
        fc2_bias: &[f32],
        in_channels: usize,
        out_channels: usize,
        device: &B::Device,
    ) -> Self {
        let hidden_dim = fc1_bias.len();
        Self {
            fc1: GenericLinear::from_weights(
                fc1_weight,
                Some(fc1_bias),
                in_channels,
                hidden_dim,
                device,
            ),
            fc2: GenericLinear::from_weights(
                fc2_weight,
                Some(fc2_bias),
                hidden_dim,
                out_channels,
                device,
            ),
            device: device.clone(),
        }
    }
}

impl<B: Backend<FloatElem = f32>> GenericModule<B> for LiftingMLP<B> {
    fn forward(&self, _input: &GenericTensor<B, 2>) -> GenericTensor<B, 2> {
        panic!("Use forward_4d for spatial data");
    }

    fn device(&self) -> &B::Device {
        &self.device
    }
}

impl<B: Backend<FloatElem = f32>> GenericModule4D<B> for LiftingMLP<B> {
    fn forward_4d(&self, input: &GenericTensor<B, 4>) -> GenericTensor<B, 4> {
        // Input: [batch, in_channels, height, width]
        let shape = input.shape();
        let batch = shape[0];
        let in_ch = shape[1];
        let h = shape[2];
        let w = shape[3];

        // [B, C, H, W] -> [B, H, W, C]
        let step1 = input.swap_dims(1, 2);
        let permuted = step1.swap_dims(2, 3);

        // Reshape to [batch * height * width, in_channels]
        let reshaped = permuted.reshape([batch * h * w, in_ch]);

        // First linear + GELU
        let hidden = self.fc1.forward(&reshaped);
        let activated = hidden.gelu();

        // Second linear
        let output = self.fc2.forward(&activated);

        // Get output channels
        let out_ch = output.shape()[1];

        // Reshape back to [batch, height, width, out_channels]
        let output_nhwc = output.reshape([batch, h, w, out_ch]);

        // [B, H, W, C] -> [B, C, H, W]
        let step1 = output_nhwc.swap_dims(2, 3);
        step1.swap_dims(1, 2)
    }
}

/// Lifting layer: projects input channels to a higher-dimensional space.
///
/// In neural operators, the lifting layer embeds the input function into a
/// high-dimensional feature space where spectral operations are applied.
#[derive(Debug)]
pub struct Lifting<B: Backend<FloatElem = f32>> {
    linear: GenericLinear<B>,
    device: B::Device,
}

impl<B: Backend<FloatElem = f32>> Lifting<B> {
    /// Create a new lifting layer.
    ///
    /// # Arguments
    /// * `in_channels` - Number of input channels
    /// * `out_channels` - Number of output channels (latent dimension)
    /// * `device` - Device to create the layer on
    pub fn new(in_channels: usize, out_channels: usize, device: &B::Device) -> Self {
        Self {
            linear: GenericLinear::new(in_channels, out_channels, true, device),
            device: device.clone(),
        }
    }

    /// Create a lifting layer with pre-defined weights.
    ///
    /// # Arguments
    /// * `weight_data` - Weight matrix [out_channels, in_channels] in row-major
    /// * `bias_data` - Bias vector [out_channels]
    /// * `in_channels` - Number of input channels
    /// * `out_channels` - Number of output channels
    /// * `device` - Device to create the layer on
    pub fn from_weights(
        weight_data: &[f32],
        bias_data: &[f32],
        in_channels: usize,
        out_channels: usize,
        device: &B::Device,
    ) -> Self {
        Self {
            linear: GenericLinear::from_weights(
                weight_data,
                Some(bias_data),
                in_channels,
                out_channels,
                device,
            ),
            device: device.clone(),
        }
    }
}

impl<B: Backend<FloatElem = f32>> GenericModule<B> for Lifting<B> {
    fn forward(&self, _input: &GenericTensor<B, 2>) -> GenericTensor<B, 2> {
        panic!("Use forward_4d for spatial data");
    }

    fn device(&self) -> &B::Device {
        &self.device
    }
}

impl<B: Backend<FloatElem = f32>> GenericModule4D<B> for Lifting<B> {
    fn forward_4d(&self, input: &GenericTensor<B, 4>) -> GenericTensor<B, 4> {
        // Input: [batch, in_channels, height, width]
        let shape = input.shape();
        let batch = shape[0];
        let in_ch = shape[1];
        let h = shape[2];
        let w = shape[3];

        // [B, C, H, W] -> [B, H, W, C] using swap_dims
        // First swap C and H: [B, C, H, W] -> [B, H, C, W]
        let step1 = input.swap_dims(1, 2);
        // Then swap C and W: [B, H, C, W] -> [B, H, W, C]
        let permuted = step1.swap_dims(2, 3);

        // Reshape to [batch * height * width, in_channels]
        let reshaped = permuted.reshape([batch * h * w, in_ch]);

        // Apply linear transformation
        let transformed = self.linear.forward(&reshaped);

        // Get output channels
        let out_ch = transformed.shape()[1];

        // Reshape back to [batch, height, width, out_channels]
        let output_nhwc = transformed.reshape([batch, h, w, out_ch]);

        // [B, H, W, C] -> [B, C, H, W] using swap_dims
        // First swap W and C: [B, H, W, C] -> [B, H, C, W]
        let step1 = output_nhwc.swap_dims(2, 3);
        // Then swap H and C: [B, H, C, W] -> [B, C, H, W]
        step1.swap_dims(1, 2)
    }
}

/// Projection layer: projects latent representation back to output space.
///
/// The projection layer (also called "Q" in FNO papers) maps the high-dimensional
/// latent features back to the desired output channels.
#[derive(Debug)]
pub struct Projection<B: Backend<FloatElem = f32>> {
    linear: GenericLinear<B>,
    device: B::Device,
}

impl<B: Backend<FloatElem = f32>> Projection<B> {
    /// Create a new projection layer.
    ///
    /// # Arguments
    /// * `in_channels` - Number of input channels (latent dimension)
    /// * `out_channels` - Number of output channels
    /// * `device` - Device to create the layer on
    pub fn new(in_channels: usize, out_channels: usize, device: &B::Device) -> Self {
        Self {
            linear: GenericLinear::new(in_channels, out_channels, true, device),
            device: device.clone(),
        }
    }

    /// Get weight and bias data from the projection layer.
    ///
    /// Returns `(weight, bias)` as `(Vec<f32>, Vec<f32>)`.
    pub fn weights(&self) -> (Vec<f32>, Vec<f32>) {
        let weight = self.linear.weight().to_vec();
        let bias = self
            .linear
            .bias()
            .map(rtx_tensor::GenericTensor::to_vec)
            .unwrap_or_default();
        (weight, bias)
    }

    /// Get the input dimension.
    pub fn in_features(&self) -> usize {
        self.linear.in_features()
    }

    /// Get the output dimension.
    pub fn out_features(&self) -> usize {
        self.linear.out_features()
    }

    /// Create a projection layer with pre-defined weights.
    ///
    /// # Arguments
    /// * `weight_data` - Weight matrix [out_channels, in_channels] in row-major
    /// * `bias_data` - Bias vector [out_channels]
    /// * `in_channels` - Number of input channels
    /// * `out_channels` - Number of output channels
    /// * `device` - Device to create the layer on
    pub fn from_weights(
        weight_data: &[f32],
        bias_data: &[f32],
        in_channels: usize,
        out_channels: usize,
        device: &B::Device,
    ) -> Self {
        Self {
            linear: GenericLinear::from_weights(
                weight_data,
                Some(bias_data),
                in_channels,
                out_channels,
                device,
            ),
            device: device.clone(),
        }
    }
}

impl<B: Backend<FloatElem = f32>> GenericModule<B> for Projection<B> {
    fn forward(&self, _input: &GenericTensor<B, 2>) -> GenericTensor<B, 2> {
        panic!("Use forward_4d for spatial data");
    }

    fn device(&self) -> &B::Device {
        &self.device
    }
}

impl<B: Backend<FloatElem = f32>> GenericModule4D<B> for Projection<B> {
    fn forward_4d(&self, input: &GenericTensor<B, 4>) -> GenericTensor<B, 4> {
        // Same logic as Lifting but different semantic meaning
        let shape = input.shape();
        let batch = shape[0];
        let in_ch = shape[1];
        let h = shape[2];
        let w = shape[3];

        // [B, C, H, W] -> [B, H, W, C]
        let step1 = input.swap_dims(1, 2);
        let permuted = step1.swap_dims(2, 3);

        let reshaped = permuted.reshape([batch * h * w, in_ch]);
        let transformed = self.linear.forward(&reshaped);
        let out_ch = transformed.shape()[1];
        let output_nhwc = transformed.reshape([batch, h, w, out_ch]);

        // [B, H, W, C] -> [B, C, H, W]
        let step1 = output_nhwc.swap_dims(2, 3);
        step1.swap_dims(1, 2)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use rtx_backend_cpu::{CpuBackend, CpuDevice};

    #[test]
    fn test_lifting_shape() {
        let device = CpuDevice::new();
        let lifting = Lifting::<CpuBackend>::new(3, 32, &device);

        let input = GenericTensor::randn([2, 3, 64, 64], &device);
        let output = lifting.forward_4d(&input);

        assert_eq!(output.shape(), [2, 32, 64, 64]);
    }

    #[test]
    fn test_projection_shape() {
        let device = CpuDevice::new();
        let projection = Projection::<CpuBackend>::new(32, 1, &device);

        let input = GenericTensor::randn([2, 32, 64, 64], &device);
        let output = projection.forward_4d(&input);

        assert_eq!(output.shape(), [2, 1, 64, 64]);
    }

    #[test]
    fn test_lifting_preserves_spatial_dimensions() {
        let device = CpuDevice::new();
        let lifting = Lifting::<CpuBackend>::new(1, 16, &device);

        let input = GenericTensor::randn([4, 1, 32, 32], &device);
        let output = lifting.forward_4d(&input);

        assert_eq!(output.shape()[0], 4); // batch
        assert_eq!(output.shape()[1], 16); // channels
        assert_eq!(output.shape()[2], 32); // height
        assert_eq!(output.shape()[3], 32); // width
    }

    #[test]
    fn test_round_trip_preserves_spatial() {
        let device = CpuDevice::new();
        let lifting = Lifting::<CpuBackend>::new(2, 64, &device);
        let projection = Projection::<CpuBackend>::new(64, 2, &device);

        let input = GenericTensor::randn([1, 2, 16, 16], &device);
        let lifted = lifting.forward_4d(&input);
        let projected = projection.forward_4d(&lifted);

        assert_eq!(input.shape(), projected.shape());
    }

    #[test]
    fn test_grid_positional_encoding_shape() {
        let device = CpuDevice::new();
        let pos_enc = GridPositionalEncoding::<CpuBackend>::new(&device);

        let input = GenericTensor::randn([2, 3, 16, 16], &device);
        let output = pos_enc.forward_4d(&input);

        // Should add 2 channels for x and y coordinates
        assert_eq!(output.shape(), [2, 5, 16, 16]);
    }

    #[test]
    fn test_grid_positional_encoding_values() {
        let device = CpuDevice::new();
        let pos_enc = GridPositionalEncoding::<CpuBackend>::new(&device);

        // Create a simple input
        let input = GenericTensor::zeros([1, 1, 3, 3], &device);
        let output = pos_enc.forward_4d(&input);

        let data = output.to_vec();
        // Output shape: [1, 3, 3, 3] (1 batch, 3 channels: original + x + y)

        // Channel 0 should be zeros (original input)
        // Channel 1 should be x coordinates: [-1, 0, 1] for each row
        // Channel 2 should be y coordinates: [-1, 0, 1] for each column

        // Check x coordinates (channel 1)
        let x_channel_start = 1 * 3 * 3; // offset for channel 1
        assert!((data[x_channel_start] - (-1.0)).abs() < 0.01); // x=0 -> -1
        assert!((data[x_channel_start + 1] - 0.0).abs() < 0.01); // x=1 -> 0
        assert!((data[x_channel_start + 2] - 1.0).abs() < 0.01); // x=2 -> 1

        // Check y coordinates (channel 2)
        let y_channel_start = 2 * 3 * 3; // offset for channel 2
        assert!((data[y_channel_start] - (-1.0)).abs() < 0.01); // y=0 -> -1
        assert!((data[y_channel_start + 3] - 0.0).abs() < 0.01); // y=1 -> 0
        assert!((data[y_channel_start + 6] - 1.0).abs() < 0.01); // y=2 -> 1
    }

    #[test]
    fn test_lifting_mlp_shape() {
        let device = CpuDevice::new();
        // Input has 3 channels (1 data + 2 positional encoding)
        let lifting = LiftingMLP::<CpuBackend>::new(3, 32, &device);

        let input = GenericTensor::randn([2, 3, 16, 16], &device);
        let output = lifting.forward_4d(&input);

        assert_eq!(output.shape(), [2, 32, 16, 16]);
    }

    #[test]
    fn test_lifting_mlp_non_zero_output() {
        let device = CpuDevice::new();
        let lifting = LiftingMLP::<CpuBackend>::new(3, 32, &device);

        let input = GenericTensor::randn([1, 3, 8, 8], &device);
        let output = lifting.forward_4d(&input);

        let data = output.to_vec();
        let non_zero = data.iter().filter(|&&x| x.abs() > 1e-10).count();
        assert!(non_zero > 0, "LiftingMLP output should not be all zeros");
    }

    #[test]
    fn test_positional_encoding_with_lifting_mlp() {
        let device = CpuDevice::new();

        // Simulate neuraloperator v2.0 pipeline
        let pos_enc = GridPositionalEncoding::<CpuBackend>::new(&device);
        let lifting = LiftingMLP::<CpuBackend>::new(3, 32, &device); // 1 + 2 = 3 input channels

        let input = GenericTensor::randn([1, 1, 16, 16], &device);
        let with_pos = pos_enc.forward_4d(&input); // [1, 3, 16, 16]
        let lifted = lifting.forward_4d(&with_pos); // [1, 32, 16, 16]

        assert_eq!(with_pos.shape(), [1, 3, 16, 16]);
        assert_eq!(lifted.shape(), [1, 32, 16, 16]);
    }
}