foxhunt/crates/ml-supervised/src/diffusion/noise.rs

//! Noise scheduler for the diffusion process.
//!
//! Precomputes `alpha_bar_t` for all timesteps and provides
//! forward process (add noise) operations.

use std::sync::Arc;

use cudarc::driver::CudaStream;
use ml_core::MLError;

use crate::gpu_tensor::{gpu_add, gpu_affine, GpuTensor};

use super::config::NoiseSchedule;

/// Precomputed noise schedule for the diffusion process.
///
/// Stores `alpha_bar_t` (cumulative product of (1 - `beta_t`)) for
/// all T timesteps, enabling efficient forward-process noise addition.
pub struct NoiseScheduler {
    /// Cumulative alpha products: `alpha_bar_t` for each timestep.
    alpha_bars: Vec<f32>,
    num_timesteps: usize,
    stream: Arc<CudaStream>,
}

impl std::fmt::Debug for NoiseScheduler {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("NoiseScheduler").finish_non_exhaustive()
    }
}

impl NoiseScheduler {
    /// Create a new noise scheduler with precomputed schedule.
    pub fn new(
        num_timesteps: usize,
        schedule: &NoiseSchedule,
        stream: &Arc<CudaStream>,
    ) -> Result<Self, MLError> {
        if num_timesteps == 0 {
            return Err(MLError::ConfigError("num_timesteps must be > 0".to_owned()));
        }

        let alpha_bars = match schedule {
            NoiseSchedule::Linear => Self::linear_schedule(num_timesteps),
            NoiseSchedule::Cosine => Self::cosine_schedule(num_timesteps),
        };

        Ok(Self {
            alpha_bars,
            num_timesteps,
            stream: Arc::clone(stream),
        })
    }

    /// Linear beta schedule: beta linearly from 1e-4 to 0.02.
    fn linear_schedule(t_max: usize) -> Vec<f32> {
        let beta_start = 1e-4_f64;
        let beta_end = 0.02_f64;
        let mut alpha_bars = Vec::with_capacity(t_max);
        let mut cumulative = 1.0_f64;
        for i in 0..t_max {
            let beta =
                beta_start + (beta_end - beta_start) * (i as f64 / (t_max - 1).max(1) as f64);
            cumulative *= 1.0 - beta;
            alpha_bars.push(cumulative as f32);
        }
        alpha_bars
    }

    /// Cosine schedule (Nichol & Dhariwal 2021).
    fn cosine_schedule(t_max: usize) -> Vec<f32> {
        let s = 0.008_f64;
        let mut alpha_bars = Vec::with_capacity(t_max);
        for i in 0..t_max {
            let t_frac = i as f64 / t_max as f64;
            let angle = (t_frac + s) / (1.0 + s) * std::f64::consts::FRAC_PI_2;
            let alpha_bar = angle.cos().powi(2);
            alpha_bars.push(alpha_bar as f32);
        }
        let first = alpha_bars.first().copied().unwrap_or(1.0);
        if first > 0.0 {
            for ab in &mut alpha_bars {
                *ab /= first;
            }
        }
        alpha_bars
    }

    /// Get `alpha_bar` for a specific timestep.
    pub fn get_alpha_bar(&self, t: usize) -> Result<f32, MLError> {
        self.alpha_bars.get(t).copied().ok_or_else(|| {
            MLError::ConfigError(format!(
                "Timestep {} out of range (max {})",
                t, self.num_timesteps
            ))
        })
    }

    /// Forward process: add noise to clean data x0 at timestep t.
    ///
    /// Returns (`noisy_x`, noise) where noise is the sampled epsilon.
    pub fn add_noise(&self, x0: &GpuTensor, t: usize) -> Result<(GpuTensor, GpuTensor), MLError> {
        let alpha_bar = self.get_alpha_bar(t)?;
        let sqrt_alpha_bar = alpha_bar.sqrt();
        let sqrt_one_minus_alpha_bar = (1.0_f32 - alpha_bar).max(0.0_f32).sqrt();

        // Sample noise epsilon ~ N(0, I)
        let noise = GpuTensor::randn(x0.shape.as_slice(), 1.0, &self.stream)?;

        // x_t = sqrt(alpha_bar_t) * x_0 + sqrt(1 - alpha_bar_t) * epsilon
        let signal = gpu_affine(x0, sqrt_alpha_bar as f64, 0.0)?;
        let noise_scaled = gpu_affine(&noise, sqrt_one_minus_alpha_bar as f64, 0.0)?;
        let noisy = gpu_add(&signal, &noise_scaled)?;

        Ok((noisy, noise))
    }

    /// Number of timesteps.
    pub const fn num_timesteps(&self) -> usize {
        self.num_timesteps
    }

    /// Get `alpha_bars` slice for DDIM sampling.
    pub fn alpha_bars(&self) -> &[f32] {
        &self.alpha_bars
    }
}

#[cfg(test)]
#[allow(clippy::assertions_on_result_states)]
mod tests {
    use super::*;
    use cudarc::driver::CudaContext;

    fn test_stream() -> Arc<CudaStream> {
        let ctx = CudaContext::new(0).expect("CUDA required");
        ctx.new_stream().expect("Failed to create stream")
    }

    #[test]
    fn test_linear_schedule_decreasing() {
        let stream = test_stream();
        let sched = NoiseScheduler::new(1000, &NoiseSchedule::Linear, &stream).unwrap();
        let a0 = sched.get_alpha_bar(0).unwrap();
        let a500 = sched.get_alpha_bar(500).unwrap();
        let a999 = sched.get_alpha_bar(999).unwrap();
        assert!(a0 > a500, "a0={a0} should be > a500={a500}");
        assert!(a500 > a999, "a500={a500} should be > a999={a999}");
    }

    #[test]
    fn test_cosine_schedule_decreasing() {
        let stream = test_stream();
        let sched = NoiseScheduler::new(1000, &NoiseSchedule::Cosine, &stream).unwrap();
        let a0 = sched.get_alpha_bar(0).unwrap();
        let a500 = sched.get_alpha_bar(500).unwrap();
        let a999 = sched.get_alpha_bar(999).unwrap();
        assert!(a0 > a500, "a0={a0} should be > a500={a500}");
        assert!(a500 > a999, "a500={a500} should be > a999={a999}");
    }

    #[test]
    fn test_alpha_bar_first_near_one() {
        let stream = test_stream();
        let sched = NoiseScheduler::new(1000, &NoiseSchedule::Cosine, &stream).unwrap();
        let a0 = sched.get_alpha_bar(0).unwrap();
        assert!(
            (a0 - 1.0).abs() < 0.05,
            "First alpha_bar should be near 1.0, got {a0}"
        );
    }

    #[test]
    fn test_add_noise_preserves_shape() {
        let stream = test_stream();
        let sched = NoiseScheduler::new(1000, &NoiseSchedule::Cosine, &stream).unwrap();
        let x = GpuTensor::from_vec(vec![1.0; 256], &[4, 64], &stream).unwrap();
        let (noisy, noise) = sched.add_noise(&x, 100).unwrap();
        assert_eq!(noisy.shape, vec![4, 64]);
        assert_eq!(noise.shape, vec![4, 64]);
    }

    #[test]
    fn test_out_of_range_errors() {
        let stream = test_stream();
        let sched = NoiseScheduler::new(100, &NoiseSchedule::Linear, &stream).unwrap();
        assert!(sched.get_alpha_bar(100).is_err());
    }

    #[test]
    fn test_zero_timesteps_errors() {
        let stream = test_stream();
        let result = NoiseScheduler::new(0, &NoiseSchedule::Linear, &stream);
        assert!(result.is_err());
    }
}