foxhunt/crates/ml-ppo/src/symlog.rs

//! Symlog value transform (`DreamerV3`)
//!
//! Compresses large values while preserving sign. Near-identity for small values.
//! Critical for financial returns spanning multiple orders of magnitude.
//!
//! # References
//! - Hafner et al., "Mastering Diverse Domains through World Models" (`DreamerV3`, 2023)
//!
//! # Properties
//! - `symlog(0) = 0`
//! - Near-identity for `|x| << 1`
//! - Logarithmic compression for `|x| >> 1`
//! - Preserves sign: `sign(symlog(x)) == sign(x)`
//! - Invertible: `symexp(symlog(x)) == x`

use ml_core::MLError;

/// Symlog transform: `sign(x) * ln(|x| + 1)`
///
/// Compresses the magnitude of large values logarithmically while
/// leaving small values approximately unchanged.
///
/// # Examples
/// - `symlog(0.0)   = 0.0`
/// - `symlog(1.0)   = ln(2) ≈ 0.693`
/// - `symlog(100.0)  = ln(101) ≈ 4.615`
/// - `symlog(-5.0)  = -ln(6) ≈ -1.792`
#[inline]
pub fn symlog(x: f64) -> f64 {
    x.signum() * (x.abs() + 1.0).ln()
}

/// Inverse of symlog: `sign(x) * (exp(|x|) - 1)`
///
/// Recovers the original value from its symlog representation.
#[inline]
pub fn symexp(x: f64) -> f64 {
    x.signum() * (x.abs().exp() - 1.0)
}

/// Element-wise symlog for a batch of f32 values.
///
/// Applies `sign(x) * ln(|x| + 1)` to each element.
pub fn symlog_vec(data: &[f32]) -> Vec<f32> {
    data.iter()
        .map(|&x| {
            let xf = x as f64;
            (xf.signum() * (xf.abs() + 1.0).ln()) as f32
        })
        .collect()
}

/// Element-wise symexp for a batch of f32 values.
///
/// Applies `sign(x) * (exp(|x|) - 1)` to each element.
pub fn symexp_vec(data: &[f32]) -> Vec<f32> {
    data.iter()
        .map(|&x| {
            let xf = x as f64;
            (xf.signum() * (xf.abs().exp() - 1.0)) as f32
        })
        .collect()
}

/// GPU-accelerated symlog via CUDA kernel (operates on `CudaVec` data).
///
/// For GPU data, download to host, apply symlog, and re-upload. For large
/// batches, a fused CUDA kernel would be more efficient but the download
/// path is correct and used only at batch boundaries.
#[cfg(feature = "cuda")]
pub fn symlog_gpu(
    data: &crate::cuda_nn::CudaVec,
    stream: &std::sync::Arc<cudarc::driver::CudaStream>,
) -> Result<crate::cuda_nn::CudaVec, MLError> {
    let host = data.to_vec(stream)?;
    let transformed = symlog_vec(&host);
    crate::cuda_nn::cuda_from_slice(stream, &transformed)
}

/// GPU-accelerated symexp via CUDA kernel (operates on `CudaVec` data).
#[cfg(feature = "cuda")]
pub fn symexp_gpu(
    data: &crate::cuda_nn::CudaVec,
    stream: &std::sync::Arc<cudarc::driver::CudaStream>,
) -> Result<crate::cuda_nn::CudaVec, MLError> {
    let host = data.to_vec(stream)?;
    let transformed = symexp_vec(&host);
    crate::cuda_nn::cuda_from_slice(stream, &transformed)
}

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

    const EPSILON: f64 = 1e-9;

    #[test]
    fn test_symlog_zero() {
        let result = symlog(0.0);
        assert!(
            result.abs() < EPSILON,
            "symlog(0) should be 0, got {result}"
        );
    }

    #[test]
    fn test_symlog_positive() {
        let result = symlog(1.0);
        let expected = 2.0_f64.ln(); // ln(|1| + 1) = ln(2)
        assert!(
            (result - expected).abs() < EPSILON,
            "symlog(1.0) should be ln(2) ≈ {expected}, got {result}"
        );
    }

    #[test]
    fn test_symlog_negative() {
        let result = symlog(-1.0);
        let expected = -(2.0_f64.ln());
        assert!(
            (result - expected).abs() < EPSILON,
            "symlog(-1.0) should be -ln(2) ≈ {expected}, got {result}"
        );
    }

    #[test]
    fn test_symlog_large() {
        let result = symlog(1000.0);
        let expected = (1001.0_f64).ln(); // ≈ 6.908
        assert!(
            (result - expected).abs() < EPSILON,
            "symlog(1000.0) should be ln(1001) ≈ {expected}, got {result}"
        );
        // Verify significant compression
        assert!(
            result < 7.0,
            "symlog(1000) should compress to < 7, got {result}"
        );
    }

    #[test]
    fn test_symlog_negative_large() {
        let result = symlog(-500.0);
        let expected = -(501.0_f64.ln());
        assert!(
            (result - expected).abs() < EPSILON,
            "symlog(-500.0) should be -ln(501) ≈ {expected}, got {result}"
        );
    }

    #[test]
    fn test_symexp_inverse() {
        // symexp(symlog(x)) should equal x for various values
        let test_values = [
            0.0, 1.0, -1.0, 0.5, -0.5, 10.0, -10.0, 100.0, -100.0, 0.001, -0.001,
        ];
        for &x in &test_values {
            let roundtrip = symexp(symlog(x));
            assert!(
                (roundtrip - x).abs() < 1e-6,
                "symexp(symlog({x})) should be {x}, got {roundtrip}"
            );
        }
    }

    #[test]
    fn test_symlog_near_identity() {
        // For |x| < 0.5, symlog(x) should be close to x
        // Because ln(|x|+1) ≈ |x| for small |x| (first-order Taylor)
        let small_values = [-0.4, -0.2, -0.1, -0.01, 0.01, 0.1, 0.2, 0.4];
        for &x in &small_values {
            let result = symlog(x);
            let diff = (result - x).abs();
            assert!(
                diff < 0.1,
                "|symlog({x}) - {x}| = {diff} should be < 0.1"
            );
        }
    }

    #[test]
    fn test_symlog_financial_returns() {
        // Realistic trading return magnitudes
        let returns = [-0.05, -0.001, 0.001, 0.05, 5.0];

        // Small returns should be nearly unchanged
        let small_return = symlog(0.001);
        assert!(
            (small_return - 0.001).abs() < 0.001,
            "tiny return should be near-identity, got {small_return}"
        );

        // Large returns should be compressed
        let large_return = symlog(5.0);
        assert!(
            large_return < 2.0,
            "5.0 return should compress below 2.0, got {large_return}"
        );

        // All returns should preserve sign
        for &r in &returns {
            let s = symlog(r);
            if r > 0.0 {
                assert!(s > 0.0, "positive return {r} should yield positive symlog");
            } else if r < 0.0 {
                assert!(s < 0.0, "negative return {r} should yield negative symlog");
            }
        }
    }

    #[test]
    fn test_symlog_vec() {
        let data = vec![-10.0_f32, -1.0, 0.0, 1.0, 10.0];
        let result = symlog_vec(&data);

        // Check each element matches scalar symlog
        for (i, (&input, &output)) in data.iter().zip(result.iter()).enumerate() {
            let expected = symlog(f64::from(input)) as f32;
            assert!(
                (output - expected).abs() < 1e-4,
                "element {i}: symlog_vec({input}) = {output}, expected {expected}"
            );
        }
    }

    #[test]
    fn test_symexp_vec() {
        let data = vec![-2.0_f32, -0.5, 0.0, 0.5, 2.0];
        let result = symexp_vec(&data);

        for (i, (&input, &output)) in data.iter().zip(result.iter()).enumerate() {
            let expected = symexp(f64::from(input)) as f32;
            assert!(
                (output - expected).abs() < 1e-3,
                "element {i}: symexp_vec({input}) = {output}, expected {expected}"
            );
        }
    }

    #[test]
    fn test_symlog_symexp_vec_roundtrip() {
        let data = vec![-5.0_f32, -1.0, -0.1, 0.0, 0.1, 1.0, 5.0];
        let encoded = symlog_vec(&data);
        let decoded = symexp_vec(&encoded);

        for (i, (&original, &roundtrip)) in data.iter().zip(decoded.iter()).enumerate() {
            assert!(
                (roundtrip - original).abs() < 1e-4,
                "element {i}: roundtrip of {original} = {roundtrip}"
            );
        }
    }
}