rustytorch/demos/rtx-neural-operator-demo/src/data.rs

//! Synthetic data generation for FNO training
//!
//! Generates training pairs (input, output) for various PDE types.
//! The primary focus is Darcy flow, which is the standard benchmark
//! for Fourier Neural Operators.
//!
//! # Darcy Flow Equation
//!
//! The 2D Darcy equation describes flow through porous media:
//!
//! ```text
//! -∇·(k(x)∇p) = f
//! ```
//!
//! where:
//! - k(x): permeability field (input) - random log-normal field
//! - p(x): pressure field (output) - solved via finite differences
//! - f: source term (constant forcing)
//!
//! # Usage
//!
//! ```rust,ignore
//! use rtx_neural_operator_demo::data::DarcyDataGenerator;
//!
//! let generator = DarcyDataGenerator::new(64);
//! let (inputs, outputs) = generator.generate_batch(32);
//! ```

use rand::Rng;
use rand_distr::{Distribution, Normal};

/// Configuration for Darcy flow data generation
#[derive(Debug, Clone)]
pub struct DarcyDataConfig {
    /// Grid resolution (assumes square grid)
    pub resolution: usize,
    /// Mean of log-permeability field
    pub log_k_mean: f32,
    /// Standard deviation of log-permeability field
    pub log_k_std: f32,
    /// Correlation length for permeability field (as fraction of domain)
    pub correlation_length: f32,
    /// Source term magnitude
    pub source_magnitude: f32,
    /// Number of Jacobi iterations for PDE solver
    pub solver_iterations: usize,
    /// Convergence tolerance for solver
    pub solver_tolerance: f32,
}

impl Default for DarcyDataConfig {
    fn default() -> Self {
        Self {
            resolution: 64,
            log_k_mean: 0.0,
            log_k_std: 1.0,
            correlation_length: 0.1,
            source_magnitude: 1.0,
            solver_iterations: 1000,
            solver_tolerance: 1e-6,
        }
    }
}

impl DarcyDataConfig {
    /// Create config for quick training (lower resolution)
    #[must_use]
    pub fn quick() -> Self {
        Self {
            resolution: 32,
            solver_iterations: 500,
            ..Self::default()
        }
    }

    /// Create config for standard training
    #[must_use]
    pub fn standard() -> Self {
        Self::default()
    }

    /// Create config for high-resolution training
    #[must_use]
    pub fn high_resolution() -> Self {
        Self {
            resolution: 128,
            solver_iterations: 2000,
            ..Self::default()
        }
    }

    /// Set resolution
    #[must_use]
    pub const fn with_resolution(mut self, resolution: usize) -> Self {
        self.resolution = resolution;
        self
    }
}

/// A single training sample (input, output pair)
#[derive(Debug, Clone)]
pub struct DarcySample {
    /// Permeability field k(x) - flattened [H, W]
    pub permeability: Vec<f32>,
    /// Pressure field p(x) - flattened [H, W]
    pub pressure: Vec<f32>,
    /// Resolution of the grid
    pub resolution: usize,
}

/// Generator for synthetic Darcy flow training data
pub struct DarcyDataGenerator {
    config: DarcyDataConfig,
}

impl DarcyDataGenerator {
    /// Create a new data generator with given resolution
    #[must_use]
    pub fn new(resolution: usize) -> Self {
        Self {
            config: DarcyDataConfig::default().with_resolution(resolution),
        }
    }

    /// Create a data generator with custom configuration
    #[must_use]
    pub fn with_config(config: DarcyDataConfig) -> Self {
        Self { config }
    }

    /// Generate a single training sample
    ///
    /// Returns (`permeability_field`, `pressure_field`)
    #[must_use]
    pub fn generate_sample(&self) -> DarcySample {
        let n = self.config.resolution;

        // Generate random permeability field
        let permeability = self.generate_permeability_field();

        // Solve Darcy equation to get pressure field
        let pressure = self.solve_darcy(&permeability);

        DarcySample {
            permeability,
            pressure,
            resolution: n,
        }
    }

    /// Generate a batch of training samples
    ///
    /// Returns (inputs, outputs) where each is a vector of flattened fields
    #[must_use]
    pub fn generate_batch(&self, batch_size: usize) -> (Vec<Vec<f32>>, Vec<Vec<f32>>) {
        let mut inputs = Vec::with_capacity(batch_size);
        let mut outputs = Vec::with_capacity(batch_size);

        for _ in 0..batch_size {
            let sample = self.generate_sample();
            inputs.push(sample.permeability);
            outputs.push(sample.pressure);
        }

        (inputs, outputs)
    }

    /// Generate multiple samples with progress callback
    pub fn generate_samples_with_progress<F>(
        &self,
        n_samples: usize,
        mut progress_callback: F,
    ) -> Vec<DarcySample>
    where
        F: FnMut(usize, usize),
    {
        let mut samples = Vec::with_capacity(n_samples);

        for i in 0..n_samples {
            samples.push(self.generate_sample());
            progress_callback(i + 1, n_samples);
        }

        samples
    }

    /// Generate a random permeability field k(x)
    ///
    /// Uses a Gaussian random field with log-normal distribution
    /// to ensure positivity (permeability must be > 0).
    fn generate_permeability_field(&self) -> Vec<f32> {
        let n = self.config.resolution;
        let mut rng = rand::rng();

        // Generate correlated Gaussian random field
        let gaussian_field = self.generate_gaussian_random_field(&mut rng);

        // Apply spatial smoothing for correlation
        let smoothed = self.apply_gaussian_smoothing(&gaussian_field);

        // Convert to permeability values
        let mut permeability = vec![0.0; n * n];
        for i in 0..n * n {
            // Use smoothed field to modulate log-normal
            let z = smoothed[i].clamp(-3.0, 3.0);
            let k = (self.config.log_k_mean + z * self.config.log_k_std).exp();
            permeability[i] = k.clamp(0.1, 10.0); // Reasonable bounds
        }

        permeability
    }

    /// Generate a Gaussian random field
    fn generate_gaussian_random_field<R: Rng>(&self, rng: &mut R) -> Vec<f32> {
        let n = self.config.resolution;
        let normal = Normal::new(0.0f32, 1.0f32).unwrap();

        (0..n * n).map(|_| normal.sample(rng)).collect()
    }

    /// Apply Gaussian smoothing for spatial correlation
    fn apply_gaussian_smoothing(&self, field: &[f32]) -> Vec<f32> {
        let n = self.config.resolution;
        let sigma = (self.config.correlation_length * n as f32).max(1.0);
        let kernel_radius = (3.0 * sigma).ceil() as i32;

        let mut smoothed = vec![0.0; n * n];

        // Build Gaussian kernel
        let kernel_size = (2 * kernel_radius + 1) as usize;
        let mut kernel = vec![0.0f32; kernel_size * kernel_size];
        let mut kernel_sum = 0.0f32;

        for dy in -kernel_radius..=kernel_radius {
            for dx in -kernel_radius..=kernel_radius {
                let dist_sq = (dx * dx + dy * dy) as f32;
                let weight = (-dist_sq / (2.0 * sigma * sigma)).exp();
                let ki =
                    ((dy + kernel_radius) as usize) * kernel_size + (dx + kernel_radius) as usize;
                kernel[ki] = weight;
                kernel_sum += weight;
            }
        }

        // Normalize kernel
        for k in &mut kernel {
            *k /= kernel_sum;
        }

        // Apply convolution with periodic boundary conditions
        for j in 0..n {
            for i in 0..n {
                let mut sum = 0.0f32;

                for dy in -kernel_radius..=kernel_radius {
                    for dx in -kernel_radius..=kernel_radius {
                        let ni = ((i as i32 + dx).rem_euclid(n as i32)) as usize;
                        let nj = ((j as i32 + dy).rem_euclid(n as i32)) as usize;

                        let ki = ((dy + kernel_radius) as usize) * kernel_size
                            + (dx + kernel_radius) as usize;

                        sum += field[nj * n + ni] * kernel[ki];
                    }
                }

                smoothed[j * n + i] = sum;
            }
        }

        smoothed
    }

    /// Solve Darcy equation: -∇·(k∇p) = f
    ///
    /// Uses iterative Jacobi method with:
    /// - Dirichlet BC: p = 0 on boundary
    /// - Constant source term f
    fn solve_darcy(&self, permeability: &[f32]) -> Vec<f32> {
        let n = self.config.resolution;
        let h = 1.0 / (n - 1) as f32; // Grid spacing
        let h2 = h * h;
        let f = self.config.source_magnitude;

        // Initialize pressure field
        let mut p = vec![0.0f32; n * n];
        let mut p_new = vec![0.0f32; n * n];

        // Iterative solver (Jacobi method)
        for _iter in 0..self.config.solver_iterations {
            let mut max_diff = 0.0f32;

            for j in 1..n - 1 {
                for i in 1..n - 1 {
                    let idx = j * n + i;

                    // Get permeability at cell faces (harmonic mean for interface)
                    let k_c = permeability[idx];
                    let k_e = permeability[idx + 1];
                    let k_w = permeability[idx - 1];
                    let k_n = permeability[(j + 1) * n + i];
                    let k_s = permeability[(j - 1) * n + i];

                    // Harmonic means at interfaces
                    let k_pe = 2.0 * k_c * k_e / (k_c + k_e);
                    let k_pw = 2.0 * k_c * k_w / (k_c + k_w);
                    let k_pn = 2.0 * k_c * k_n / (k_c + k_n);
                    let k_ps = 2.0 * k_c * k_s / (k_c + k_s);

                    // Discrete Laplacian: -∇·(k∇p)
                    // Using central differences with variable coefficients
                    let sum_k = k_pe + k_pw + k_pn + k_ps;

                    let weighted_neighbors = k_pe * p[idx + 1]
                        + k_pw * p[idx - 1]
                        + k_pn * p[(j + 1) * n + i]
                        + k_ps * p[(j - 1) * n + i];

                    // Solve: sum_k * p_c - weighted_neighbors = f * h²
                    p_new[idx] = (weighted_neighbors + f * h2) / sum_k;

                    let diff = (p_new[idx] - p[idx]).abs();
                    max_diff = max_diff.max(diff);
                }
            }

            // Swap buffers
            std::mem::swap(&mut p, &mut p_new);

            // Check convergence
            if max_diff < self.config.solver_tolerance {
                break;
            }
        }

        // Normalize to [0, 1] range for neural network training
        let (min_p, max_p) = p.iter().fold((f32::MAX, f32::MIN), |(min, max), &v| {
            (min.min(v), max.max(v))
        });

        let range = (max_p - min_p).max(1e-8);
        for v in &mut p {
            *v = (*v - min_p) / range;
        }

        p
    }

    /// Get the configuration
    #[must_use]
    pub fn config(&self) -> &DarcyDataConfig {
        &self.config
    }
}

/// Batch training samples for efficient GPU/CPU processing
#[derive(Debug, Clone)]
pub struct TrainingBatch {
    /// Batch of input tensors [B, 1, H, W] flattened as [B * H * W]
    pub inputs: Vec<f32>,
    /// Batch of output tensors [B, 1, H, W] flattened as [B * H * W]
    pub outputs: Vec<f32>,
    /// Batch size
    pub batch_size: usize,
    /// Height/width of each sample
    pub resolution: usize,
}

impl TrainingBatch {
    /// Create a training batch from samples
    #[must_use]
    pub fn from_samples(samples: &[DarcySample]) -> Self {
        let batch_size = samples.len();
        let resolution = samples.first().map_or(64, |s| s.resolution);
        let sample_size = resolution * resolution;

        let mut inputs = Vec::with_capacity(batch_size * sample_size);
        let mut outputs = Vec::with_capacity(batch_size * sample_size);

        for sample in samples {
            inputs.extend(&sample.permeability);
            outputs.extend(&sample.pressure);
        }

        Self {
            inputs,
            outputs,
            batch_size,
            resolution,
        }
    }

    /// Get total number of elements in the batch
    #[must_use]
    pub fn total_elements(&self) -> usize {
        self.batch_size * self.resolution * self.resolution
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_data_config_defaults() {
        let config = DarcyDataConfig::default();
        assert_eq!(config.resolution, 64);
        assert_eq!(config.solver_iterations, 1000);
    }

    #[test]
    fn test_data_config_presets() {
        let quick = DarcyDataConfig::quick();
        assert_eq!(quick.resolution, 32);

        let standard = DarcyDataConfig::standard();
        assert_eq!(standard.resolution, 64);

        let high_res = DarcyDataConfig::high_resolution();
        assert_eq!(high_res.resolution, 128);
    }

    #[test]
    fn test_generator_creation() {
        let generator = DarcyDataGenerator::new(64);
        assert_eq!(generator.config().resolution, 64);
    }

    #[test]
    fn test_generate_sample() {
        let generator = DarcyDataGenerator::new(16); // Small for speed
        let sample = generator.generate_sample();

        assert_eq!(sample.permeability.len(), 16 * 16);
        assert_eq!(sample.pressure.len(), 16 * 16);
        assert_eq!(sample.resolution, 16);

        // Check permeability is positive
        for &k in &sample.permeability {
            assert!(k > 0.0, "Permeability must be positive");
        }

        // Check pressure is normalized to [0, 1]
        for &p in &sample.pressure {
            assert!(p >= 0.0 && p <= 1.0, "Pressure should be normalized");
        }
    }

    #[test]
    fn test_generate_batch() {
        let generator = DarcyDataGenerator::new(8); // Small for speed
        let (inputs, outputs) = generator.generate_batch(4);

        assert_eq!(inputs.len(), 4);
        assert_eq!(outputs.len(), 4);

        for input in &inputs {
            assert_eq!(input.len(), 8 * 8);
        }
    }

    #[test]
    fn test_training_batch() {
        let generator = DarcyDataGenerator::new(8);
        let samples: Vec<_> = (0..4).map(|_| generator.generate_sample()).collect();

        let batch = TrainingBatch::from_samples(&samples);

        assert_eq!(batch.batch_size, 4);
        assert_eq!(batch.resolution, 8);
        assert_eq!(batch.total_elements(), 4 * 8 * 8);
        assert_eq!(batch.inputs.len(), 4 * 8 * 8);
        assert_eq!(batch.outputs.len(), 4 * 8 * 8);
    }

    #[test]
    fn test_permeability_field_variation() {
        let generator = DarcyDataGenerator::new(16);
        let sample = generator.generate_sample();

        // Check there's actual variation in the field
        let min_k = sample.permeability.iter().fold(f32::MAX, |a, &b| a.min(b));
        let max_k = sample.permeability.iter().fold(f32::MIN, |a, &b| a.max(b));

        assert!(max_k > min_k, "Permeability field should have variation");
    }

    #[test]
    fn test_pressure_boundary_conditions() {
        let generator = DarcyDataGenerator::new(16);
        let sample = generator.generate_sample();

        // With normalization, we can't check exact boundary values,
        // but we verify the solution is well-behaved
        let n = sample.resolution;

        // Interior should generally have higher values than boundaries
        // (due to the Dirichlet BC and positive source)
        let interior_avg: f32 = {
            let mut sum = 0.0f32;
            let mut count = 0;
            for j in n / 4..3 * n / 4 {
                for i in n / 4..3 * n / 4 {
                    sum += sample.pressure[j * n + i];
                    count += 1;
                }
            }
            sum / count as f32
        };

        // Just verify it's a reasonable value
        assert!(interior_avg >= 0.0 && interior_avg <= 1.0);
    }

    #[test]
    fn test_progress_callback() {
        let generator = DarcyDataGenerator::new(8);
        let mut calls = Vec::new();

        generator.generate_samples_with_progress(3, |current, total| {
            calls.push((current, total));
        });

        assert_eq!(calls, vec![(1, 3), (2, 3), (3, 3)]);
    }

    #[test]
    fn test_different_resolutions() {
        for res in [8, 16, 32] {
            let generator = DarcyDataGenerator::new(res);
            let sample = generator.generate_sample();
            assert_eq!(sample.resolution, res);
            assert_eq!(sample.permeability.len(), res * res);
        }
    }
}