Blame for src/stiefel.rs

//! Stiefel/RCO — porta do `quanta/gsq_rco.py` (docs 26–27; Teos A.1/A.4/D.4).
//!
//! - `projeta_tangente_stiefel(X, G) = G − X·sym(XᵀG)` (eq. 26.2, Teo 26.2);
//! - retração: **polar Newton–Schulz** (sem LAPACK; equivalente ao QR/SVD pelo
//!   Teo 26.8-b — Y = fator polar de X, YᵀY = I até δ ≤ ~1e-6 em f32);
//! - backward = gradiente riemanniano projetado no tangente (Teo 26.3: adjunto
//!   exato no tangente);
//! - `erro_ortogonalidade δ(X) = ‖XᵀX − I‖_F` (eq. 26.7) — monitor V9 (≤ 3e-6).

use candle_core::{CpuStorage, CustomOp1, Device, Layout, Result, Shape, Tensor};

/// Projeção no espaço tangente de St(n,p) na métrica euclidiana embarcada.
/// P_X(G) = G − X·sym(XᵀG) com sym(A) = (A+Aᵀ)/2. (eq. 26.2)
pub fn projeta_tangente_stiefel(x: &Tensor, g: &Tensor) -> Result<Tensor> {
    let xtg = x.transpose(0, 1)?.matmul(g)?;
    let sim = (&xtg + &xtg.transpose(0, 1)?)?.affine(0.5, 0.0)?;
    g.sub(&x.matmul(&sim)?)
}

/// Uma iteração Newton–Schulz: Y ← ½·Y·(3I − YᵀY).
fn passo_ns(y: &Tensor) -> Result<Tensor> {
    let n = y.dim(1)?;
    let yty = y.transpose(0, 1)?.matmul(y)?;
    let ident = Tensor::eye(n, y.dtype(), y.device())?;
    let inner = ident.affine(3.0, 0.0)?.sub(&yty)?.affine(0.5, 0.0)?;
    y.matmul(&inner)
}

/// Norma espectral aproximada por iteração de potência (para escalar X antes
/// do Newton–Schulz — condição de convergência σ_max < √3).
fn norma_espectral_aprox(x: &Tensor, iters: usize) -> Result<f64> {
    // iteração de potência SIMÉTRICA: u = Xv; v = Xᵀu; σ² ≈ ‖Xᵀu‖²
    // (mantém v em (p,1) — o produto alternado preserva a dimensão)
    let p = x.dims()[1];
    let xt = x.transpose(0, 1)?;
    let mut v = Tensor::randn(0f32, 1f32, (p, 1), x.device())?;
    let mut lam = 1f64;
    for _ in 0..iters {
        let nv = v.sqr()?.sum_all()?.to_scalar::<f32>()?.sqrt();
        if nv < 1e-12 {
            return Ok(1.0);
        }
        let u = x.matmul(&v.affine(1.0 / nv as f64, 0.0)?)?; // (n,1)
        let nu = u.sqr()?.sum_all()?.to_scalar::<f32>()?.sqrt();
        if nu < 1e-12 {
            return Ok(1.0);
        }
        let w = xt.matmul(&u.affine(1.0 / nu as f64, 0.0)?)?; // (p,1)
        lam = w.sqr()?.sum_all()?.to_scalar::<f32>()? as f64;
        let nw = lam.sqrt().max(1e-12);
        v = w.affine(1.0 / nw, 0.0)?;
    }
    Ok(lam.max(1.0))
}

/// Retração polar: Y = polar(X) com colunas ortonormais (St(n,p)).
/// Newton–Schulz escalado: X₀ = X/σ_max, X_{k+1} = ½X_k(3I − X_kᵀX_k);
/// converge quadrática perto do polar; ~30 iterações dão δ_ort ≤ ~1e-6 (f32).
pub fn retrair_polar(x: &Tensor, iters: usize) -> Result<Tensor> {
    let sigma = norma_espectral_aprox(x, 12)?.sqrt();
    let mut y = x.affine(1.0 / sigma, 0.0)?;
    for _ in 0..iters {
        y = passo_ns(&y)?;
    }
    Ok(y)
}

/// Retração polar DETERMINÍSTICA (v0.7.0): escala pela norma de FROBENIUS
/// (σ_max ≤ ‖X‖_F sempre — sem iteração de potência aleatória) ⇒ σ_max(Y₀)
/// ≤ 1 < √3 ⇒ Newton–Schulz converge globalmente; a escala mais
/// conservadora pode exigir mais iterações (use 60–100).
/// Motivação: a versão aleatória deixa resíduo ~1e-8 em ~3% dos casos
/// quando a entrada é EXATAMENTE a identidade (quebra a reprodutibilidade
/// bit-a-bit dos gates de nascimento neutro — r-28 §1, v070-1).
pub fn retrair_polar_frobenius(x: &Tensor, iters: usize) -> Result<Tensor> {
    let fro = x.sqr()?.sum_all()?.to_scalar::<f32>()? as f64;
    let escala = fro.sqrt().max(1e-12);
    let mut y = x.affine(1.0 / escala, 0.0)?;
    for _ in 0..iters {
        y = passo_ns(&y)?;
    }
    Ok(y)
}

/// δ(X) = ‖XᵀX − I‖_F (eq. 26.7).
pub fn erro_ortogonalidade(x: &Tensor) -> Result<f64> {
    let n = x.dim(1)?;
    let xtx = x.transpose(0, 1)?.matmul(x)?;
    let ident = Tensor::eye(n, x.dtype(), x.device())?;
    let d = xtx.sub(&ident)?;
    let fro = d.sqr()?.sum_all()?.to_scalar::<f32>()? as f64;
    Ok(fro.sqrt())
}

/// CustomOp diferenciável: forward = retração polar; backward = projeção da
/// tangente no ponto de saída (gradiente riemanniano — Teo 26.3).
pub struct RetracaoStiefel {
    pub iters: usize,
}

impl CustomOp1 for RetracaoStiefel {
    fn name(&self) -> &'static str { "retracao_stiefel_polar" }

    fn cpu_fwd(&self, storage: &CpuStorage, layout: &Layout) -> Result<(CpuStorage, Shape)> {
        if !layout.is_contiguous() {
            return Err(candle_core::Error::Msg(
                "retracao_stiefel: entrada deve ser contígua (chame .contiguous())".into(),
            ));
        }
        let dados = storage.as_slice::<f32>()?;
        let (n, p) = (layout.shape().dims()[0], layout.shape().dims()[1]);
        // converte para candle, usa o caminho polar (CPU), devolve CpuStorage
        let dev = Device::Cpu;
        let x = Tensor::from_vec(dados.to_vec(), (n, p), &dev)?;
        let y = retrair_polar(&x, self.iters)?;
        let v = y.flatten_all()?.to_vec1::<f32>()?;
        let storage_out = CpuStorage::F32(v);
        Ok((storage_out, (n, p).into()))
    }

    fn bwd(&self, _arg: &Tensor, res: &Tensor, grad_res: &Tensor) -> Result<Option<Tensor>> {
        // adjunto: gradiente riemanniano no ponto Y = R(X) (Teo 26.3)
        Ok(Some(projeta_tangente_stiefel(res, grad_res)?))
    }
}

/// Matriz de codebook com linhas ortonormais: C = R(Wᵀ)ᵀ (port de
/// `QuantizadorGSQRF::codebook`). W ∈ R^{K×d} é o parâmetro bruto.
pub fn codebook_stiefel(w: &Tensor, iters: usize) -> Result<Tensor> {
    // a custom op exige layout padrão — materializa a transposta
    let wt = w.transpose(0, 1)?.contiguous()?;
    let q = wt.apply_op1(RetracaoStiefel { iters })?;
    Ok(q.transpose(0, 1)?)
}

/// Inicialização: linhas ortonormais via polar de gaussiana (nasce em St(d,K)).
pub fn init_stiefel(k: usize, d: usize, dev: &Device) -> Result<Tensor> {
    let g = Tensor::randn(0f32, 1f32, (k, d), dev)?;
    let wt = g.transpose(0, 1)?;
    let q = retrair_polar(&wt, 40)?;
    Ok(q.transpose(0, 1)?)
}

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

    #[test]
    fn polar_ortogonaliza() -> Result<()> {
        let dev = Device::Cpu;
        // C = R(Wᵀ)ᵀ tem LINHAS ortonormais; erro_ortogonalidade espera
        // colunas — transpõe (mesma convenção do gsq_rco.py: δ(R(Wᵀ)))
        let c = init_stiefel(16, 32, &dev)?;
        let delta = erro_ortogonalidade(&c.transpose(0, 1)?)?;
        assert!(delta < 3e-6, "δ_ort = {delta}");
        // projeção da tangente no ponto de colunas ortonormais (cᵀ)
        let g = Tensor::randn(0f32, 1f32, (32, 16), &dev)?;
        let p = projeta_tangente_stiefel(&c.transpose(0, 1)?, &g)?;
        let _ = p;
        Ok(())
    }

    #[test]
    fn gradiente_riemanniano_e_tangente() -> Result<()> {
        let dev = Device::Cpu;
        // X (n,p) com colunas ortonormais
        let x = retrair_polar(&Tensor::randn(0f32, 1f32, (24, 8), &dev)?, 40)?;
        let g = Tensor::randn(0f32, 1f32, (24, 8), &dev)?;
        let p = projeta_tangente_stiefel(&x, &g)?;
        // XᵀP + PᵀX ≈ 0 (Teo 26.2)
        let a = x.transpose(0, 1)?.matmul(&p)?;
        let soma = (&a + &a.transpose(0, 1)?)?;
        let m = soma.abs()?.max_all()?.to_scalar::<f32>()?;
        assert!(m < 1e-4, "resíduo tangente = {m}");
        Ok(())
    }
}