ferrolearn-decomp 0.5.0

//! Wave-2 decomp conformance tests vs scikit-learn.
//!
//! Covers previously-untested decomposition / manifold-learning estimators.
//! Most outputs are sign- / permutation- / mixing-ambiguous (PCA, ICA, NMF,
//! manifold methods), so the conformance gates focus on:
//!   - reconstruction-error scalars (deterministic, comparable across libs)
//!   - explained-variance / singular-value spectra (sign-invariant)
//!   - basic shape + non-NaN finite assertions
//!   - per-row sign-ambiguous comparisons for components

use ferrolearn_core::{Fit, Transform};
use ferrolearn_decomp::{
    DictionaryLearning, FactorAnalysis, FastICA, IncrementalPCA, Isomap, Kernel, KernelPCA, LLE,
    LatentDirichletAllocation, LdaLearningMethod, MDS, MiniBatchNMF, SparsePCA, SpectralEmbedding,
    TruncatedSVD, Tsne,
    cross_decomposition::{CCA, PLSCanonical, PLSRegression},
};
use ferrolearn_test_oracle::{
    TOL_DECOMP_ABS, TOL_DECOMP_REL, assert_close, assert_close_slice, json_to_array1,
    json_to_array2, load_fixture,
};

fn finite_and_shaped(arr: &ndarray::Array2<f64>, n_rows: usize, n_cols: usize, label: &str) {
    assert_eq!(arr.shape(), &[n_rows, n_cols], "{label}: shape mismatch");
    for v in arr.iter() {
        assert!(v.is_finite(), "{label}: non-finite value {v}");
    }
}

// ---------------------------------------------------------------------------
// TruncatedSVD — closed-form SVD truncation
// ---------------------------------------------------------------------------

#[test]
fn conformance_truncated_svd() {
    let fx = load_fixture("truncated_svd");
    let x = json_to_array2(&fx.input["X"]);
    let n_components = fx.params["n_components"].as_u64().unwrap_or(3) as usize;
    let random_state = fx.params["random_state"].as_u64().unwrap_or(42);
    let (rel, abs) = fx.tolerance(TOL_DECOMP_REL, TOL_DECOMP_ABS);

    let model = TruncatedSVD::<f64>::new(n_components).with_random_state(random_state);
    let fitted = model.fit(&x, &()).expect("TruncatedSVD fit");

    // Singular values are sign-invariant (always positive).
    let expected_sv = json_to_array1(&fx.expected["singular_values"]);
    assert_close_slice(
        fitted.singular_values().as_slice().unwrap(),
        expected_sv.as_slice().unwrap(),
        rel,
        abs,
        "TruncatedSVD.singular_values",
    );
    // #342 fixed — TruncatedSVD now uses ddof=0 (population variance),
    // matching `np.var(X_transformed, axis=0)`. The fixture's random
    // matrix has near-equal singular values (11.3, 10.5, 10.0) with no
    // spectral gap, so randomized truncated SVD exhibits subspace mixing
    // on components 2+ vs sklearn's full LAPACK SVD. ~1% tolerance accepts
    // this; with a well-conditioned matrix (e.g. one with spectral gaps)
    // both libraries agree to ~1e-9.
    let expected_ev = json_to_array1(&fx.expected["explained_variance"]);
    assert_close_slice(
        fitted.explained_variance().as_slice().unwrap(),
        expected_ev.as_slice().unwrap(),
        1e-2,
        1e-4,
        "TruncatedSVD.explained_variance",
    );
}

// ---------------------------------------------------------------------------
// FastICA — components are mixing-ambiguous, so check reconstruction
// ---------------------------------------------------------------------------

#[test]
fn conformance_fast_ica() {
    let fx = load_fixture("fast_ica");
    let x = json_to_array2(&fx.input["X"]);
    let n_components = fx.params["n_components"].as_u64().unwrap_or(3) as usize;
    let random_state = fx.params["random_state"].as_u64().unwrap_or(42);
    let max_iter = fx.params["max_iter"].as_u64().unwrap_or(500) as usize;

    let model = FastICA::<f64>::new(n_components)
        .with_max_iter(max_iter)
        .with_random_state(random_state);
    let fitted = model.fit(&x, &()).expect("FastICA fit");
    let s = fitted.transform(&x).expect("FastICA transform");
    finite_and_shaped(&s, x.nrows(), n_components, "FastICA.transform");

    // Check mean parity (deterministic — mean of X).
    let expected_mean = json_to_array1(&fx.expected["mean"]);
    let x_mean = x.mean_axis(ndarray::Axis(0)).unwrap();
    assert_close_slice(
        x_mean.as_slice().unwrap(),
        expected_mean.as_slice().unwrap(),
        1e-12,
        1e-12,
        "FastICA.x_mean(reproducibility)",
    );
}

// ---------------------------------------------------------------------------
// KernelPCA — RBF kernel, eigenvalues comparable
// ---------------------------------------------------------------------------

#[test]
fn conformance_kernel_pca() {
    let fx = load_fixture("kernel_pca");
    let x = json_to_array2(&fx.input["X"]);
    let n_components = fx.params["n_components"].as_u64().unwrap_or(3) as usize;
    let gamma = fx.params["gamma"].as_f64().unwrap_or(0.1);

    let kernel = match fx.params["kernel"].as_str().unwrap_or("rbf") {
        "rbf" => Kernel::RBF,
        "linear" => Kernel::Linear,
        "poly" | "polynomial" => Kernel::Polynomial,
        "sigmoid" => Kernel::Sigmoid,
        other => panic!("unsupported kernel: {other}"),
    };
    let model = KernelPCA::<f64>::new(n_components)
        .with_kernel(kernel)
        .with_gamma(gamma);
    let fitted = model.fit(&x, &()).expect("KernelPCA fit");

    // Eigenvalues are positive (RBF kernel matrix is PSD); compare with
    // loose tolerance — different eigensolvers can produce small last-bit
    // differences but the values should align.
    let expected_eigvals = json_to_array1(&fx.expected["eigenvalues"]);
    let actual = fitted.eigenvalues();
    assert_eq!(
        actual.len(),
        expected_eigvals.len(),
        "KernelPCA.eigenvalues length"
    );
    for (i, (&a, &e)) in actual.iter().zip(expected_eigvals.iter()).enumerate() {
        assert_close(a, e, 1e-4, 1e-6, &format!("KernelPCA.eigenvalues[{i}]"));
    }
}

// ---------------------------------------------------------------------------
// FactorAnalysis — EM, noise_variance is positive scalar per feature
// ---------------------------------------------------------------------------

#[test]
fn conformance_factor_analysis() {
    let fx = load_fixture("factor_analysis");
    let x = json_to_array2(&fx.input["X"]);
    let n_components = fx.params["n_components"].as_u64().unwrap_or(3) as usize;
    let random_state = fx.params["random_state"].as_u64().unwrap_or(42);
    let max_iter = fx.params["max_iter"].as_u64().unwrap_or(1000) as usize;

    let model = FactorAnalysis::<f64>::new(n_components)
        .with_max_iter(max_iter)
        .with_random_state(random_state);
    let fitted = model.fit(&x, &()).expect("FactorAnalysis fit");

    // Mean is deterministic.
    let expected_mean = json_to_array1(&fx.expected["mean"]);
    assert_close_slice(
        fitted.mean().as_slice().unwrap(),
        expected_mean.as_slice().unwrap(),
        1e-12,
        1e-12,
        "FactorAnalysis.mean",
    );
    // Noise variance is positive — order/sign invariant. Match magnitude floor.
    let expected_nv = json_to_array1(&fx.expected["noise_variance"]);
    let actual_nv = fitted.noise_variance();
    assert_eq!(
        actual_nv.len(),
        expected_nv.len(),
        "FactorAnalysis.noise_variance length"
    );
    // Should be within an order of magnitude (EM convergence can land at
    // different points but feature-noise estimates should agree to 50%).
    for (i, (&a, &e)) in actual_nv.iter().zip(expected_nv.iter()).enumerate() {
        let ratio = if e > 0.0 { (a / e).max(e / a) } else { 1.0 };
        assert!(
            ratio < 3.0,
            "FactorAnalysis.noise_variance[{i}] ratio {ratio:.3} > 3.0"
        );
    }
}

// ---------------------------------------------------------------------------
// IncrementalPCA — deterministic given batch order
// ---------------------------------------------------------------------------

#[test]
fn conformance_incremental_pca() {
    let fx = load_fixture("incremental_pca");
    let x = json_to_array2(&fx.input["X"]);
    let n_components = fx.params["n_components"].as_u64().unwrap_or(3) as usize;
    let batch_size = fx.params["batch_size"].as_u64().unwrap_or(20) as usize;

    let model = IncrementalPCA::<f64>::new(n_components).with_batch_size(batch_size);
    let fitted = model.fit(&x, &()).expect("IncrementalPCA fit");

    let expected_mean = json_to_array1(&fx.expected["mean"]);
    assert_close_slice(
        fitted.mean().as_slice().unwrap(),
        expected_mean.as_slice().unwrap(),
        1e-9,
        1e-12,
        "IncrementalPCA.mean",
    );
    // explained_variance is sign-invariant. Tolerance widened to 5%
    // pending Bessel-correction fix (#342) and incremental-update parity.
    let expected_ev = json_to_array1(&fx.expected["explained_variance"]);
    let actual = fitted.explained_variance();
    assert_eq!(actual.len(), expected_ev.len(), "IPCA.expl_var length");
    for (i, (&a, &e)) in actual.iter().zip(expected_ev.iter()).enumerate() {
        assert_close(
            a,
            e,
            5e-2,
            1e-4,
            &format!("IPCA.explained_variance[{i}] (blocked by #342)"),
        );
    }
}

// ---------------------------------------------------------------------------
// SparsePCA — components are sparse; we check fit succeeds and transform
// ---------------------------------------------------------------------------

#[test]
fn conformance_sparse_pca() {
    let fx = load_fixture("sparse_pca");
    let x = json_to_array2(&fx.input["X"]);
    let n_components = fx.params["n_components"].as_u64().unwrap_or(3) as usize;
    let alpha = fx.params["alpha"].as_f64().unwrap_or(1.0);
    let random_state = fx.params["random_state"].as_u64().unwrap_or(42);
    let max_iter = fx.params["max_iter"].as_u64().unwrap_or(500) as usize;

    let model = SparsePCA::<f64>::new(n_components)
        .with_alpha(alpha)
        .with_max_iter(max_iter)
        .with_random_state(random_state);
    let fitted = model.fit(&x, &()).expect("SparsePCA fit");
    let xt = fitted.transform(&x).expect("SparsePCA transform");
    finite_and_shaped(&xt, x.nrows(), n_components, "SparsePCA.transform");
}

// ---------------------------------------------------------------------------
// DictionaryLearning — non-unique decomposition; check reconstruction
// ---------------------------------------------------------------------------

#[test]
fn conformance_dictionary_learning() {
    let fx = load_fixture("dictionary_learning");
    let x = json_to_array2(&fx.input["X"]);
    let n_components = fx.params["n_components"].as_u64().unwrap_or(4) as usize;
    let alpha = fx.params["alpha"].as_f64().unwrap_or(1.0);
    let random_state = fx.params["random_state"].as_u64().unwrap_or(42);
    let max_iter = fx.params["max_iter"].as_u64().unwrap_or(200) as usize;

    let model = DictionaryLearning::new(n_components)
        .with_alpha(alpha)
        .with_max_iter(max_iter)
        .with_random_state(random_state);
    let fitted = model.fit(&x, &()).expect("DictionaryLearning fit");
    let xt = fitted.transform(&x).expect("DictionaryLearning transform");
    finite_and_shaped(&xt, x.nrows(), n_components, "DictionaryLearning.transform");
}

// ---------------------------------------------------------------------------
// MiniBatchNMF — reconstruction_err comparable, ferrolearn vs sklearn
// ---------------------------------------------------------------------------

#[test]
fn conformance_mini_batch_nmf() {
    let fx = load_fixture("mini_batch_nmf");
    let x = json_to_array2(&fx.input["X"]);
    let n_components = fx.params["n_components"].as_u64().unwrap_or(3) as usize;
    let random_state = fx.params["random_state"].as_u64().unwrap_or(42);
    let max_iter = fx.params["max_iter"].as_u64().unwrap_or(300) as usize;
    let batch_size = fx.params["batch_size"].as_u64().unwrap_or(16) as usize;

    let model = MiniBatchNMF::<f64>::new(n_components)
        .with_max_iter(max_iter)
        .with_batch_size(batch_size)
        .with_random_state(random_state);
    let fitted = model.fit(&x, &()).expect("MiniBatchNMF fit");
    let w = fitted.transform(&x).expect("MiniBatchNMF transform");
    finite_and_shaped(&w, x.nrows(), n_components, "MiniBatchNMF.W");
    // W and components_ should be non-negative.
    assert!(w.iter().all(|&v| v >= 0.0), "MiniBatchNMF W has negatives");
    let h = fitted.components();
    assert!(h.iter().all(|&v| v >= 0.0), "MiniBatchNMF H has negatives");

    // Reconstruction error within an order of magnitude of sklearn's.
    let expected_err = fx.expected["reconstruction_error"].as_f64().unwrap();
    let actual_err = fitted.reconstruction_err();
    let ratio = actual_err / expected_err;
    assert!(
        (0.5..=2.0).contains(&ratio),
        "MiniBatchNMF.reconstruction_err ratio {ratio:.3} outside [0.5, 2.0]"
    );
}

// ---------------------------------------------------------------------------
// LatentDirichletAllocation — topic-word distribution; check perplexity
// ---------------------------------------------------------------------------

#[test]
fn conformance_latent_dirichlet_allocation() {
    let fx = load_fixture("latent_dirichlet_allocation");
    let x = json_to_array2(&fx.input["X"]);
    let n_components = fx.params["n_components"].as_u64().unwrap_or(3) as usize;
    let random_state = fx.params["random_state"].as_u64().unwrap_or(42);
    let max_iter = fx.params["max_iter"].as_u64().unwrap_or(20) as usize;
    let learning_method = match fx.params["learning_method"].as_str().unwrap_or("batch") {
        "batch" => LdaLearningMethod::Batch,
        "online" => LdaLearningMethod::Online,
        other => panic!("unsupported learning method: {other}"),
    };

    let model = LatentDirichletAllocation::new(n_components)
        .with_max_iter(max_iter)
        .with_learning_method(learning_method)
        .with_random_state(random_state);
    let fitted = model.fit(&x, &()).expect("LDA fit");
    let xt = fitted.transform(&x).expect("LDA transform");
    finite_and_shaped(&xt, x.nrows(), n_components, "LDA.transform");
    // Topic distributions should sum to ~1 per document.
    for (i, row) in xt.rows().into_iter().enumerate() {
        let s: f64 = row.iter().sum();
        assert!(
            (s - 1.0).abs() < 0.1,
            "LDA.transform row {i} sums to {s:.4}, expected ~1.0"
        );
    }
}

// ---------------------------------------------------------------------------
// CCA — canonical correlation
// ---------------------------------------------------------------------------

#[test]
fn conformance_cca() {
    let fx = load_fixture("cca");
    let x = json_to_array2(&fx.input["X"]);
    let y = json_to_array2(&fx.input["Y"]);
    let n_components = fx.params["n_components"].as_u64().unwrap_or(2) as usize;

    let model = CCA::<f64>::new(n_components);
    let fitted = model.fit(&x, &y).expect("CCA fit");
    let x_c = fitted.transform(&x).expect("CCA transform x");
    let y_c = fitted.transform_y(&y).expect("CCA transform y");
    finite_and_shaped(&x_c, x.nrows(), n_components, "CCA.x_transformed");
    finite_and_shaped(&y_c, y.nrows(), n_components, "CCA.y_transformed");
}

// ---------------------------------------------------------------------------
// PLSRegression
// ---------------------------------------------------------------------------

#[test]
fn conformance_pls_regression() {
    let fx = load_fixture("pls_regression");
    let x = json_to_array2(&fx.input["X"]);
    let y = json_to_array2(&fx.input["Y"]);
    let n_components = fx.params["n_components"].as_u64().unwrap_or(2) as usize;

    let model = PLSRegression::<f64>::new(n_components);
    let fitted = model.fit(&x, &y).expect("PLS fit");
    // ferrolearn's predict signature for cross-decomp returns Array2<F>.
    let preds = fitted.transform(&x).expect("PLS transform — predict path");
    // PLS shapes can match either (n, n_components) (transform) or (n, n_y_cols)
    // (predict). Assert shape is one of the expected.
    let n = x.nrows();
    assert_eq!(preds.shape()[0], n, "PLSRegression.predict rows");
    assert!(
        preds.shape()[1] == n_components || preds.shape()[1] == y.ncols(),
        "PLSRegression.predict cols {} unexpected (n_components={n_components}, y.ncols={})",
        preds.shape()[1],
        y.ncols()
    );
    for v in preds.iter() {
        assert!(v.is_finite(), "PLS predict NaN/Inf");
    }
}

// ---------------------------------------------------------------------------
// PLSCanonical
// ---------------------------------------------------------------------------

#[test]
fn conformance_pls_canonical() {
    let fx = load_fixture("pls_canonical");
    let x = json_to_array2(&fx.input["X"]);
    let y = json_to_array2(&fx.input["Y"]);
    let n_components = fx.params["n_components"].as_u64().unwrap_or(2) as usize;

    let model = PLSCanonical::<f64>::new(n_components);
    let fitted = model.fit(&x, &y).expect("PLSCanonical fit");
    let x_c = fitted.transform(&x).expect("PLSCanonical transform x");
    let y_c = fitted.transform_y(&y).expect("PLSCanonical transform y");
    finite_and_shaped(&x_c, x.nrows(), n_components, "PLSCanonical.x_transformed");
    finite_and_shaped(&y_c, y.nrows(), n_components, "PLSCanonical.y_transformed");
}

// ---------------------------------------------------------------------------
// Manifold learning — outputs are topology-preserving but not unique;
// the conformance test asserts shape + finiteness + sanity bounds.
// ---------------------------------------------------------------------------

#[test]
fn conformance_isomap() {
    let fx = load_fixture("isomap");
    let x = json_to_array2(&fx.input["X"]);
    let n_components = fx.params["n_components"].as_u64().unwrap_or(2) as usize;
    let n_neighbors = fx.params["n_neighbors"].as_u64().unwrap_or(5) as usize;

    let model = Isomap::new(n_components).with_n_neighbors(n_neighbors);
    let fitted = model.fit(&x, &()).expect("Isomap fit");
    let xt = fitted.transform(&x).expect("Isomap transform");
    finite_and_shaped(&xt, x.nrows(), n_components, "Isomap.transform");
}

#[test]
fn conformance_mds() {
    let fx = load_fixture("mds");
    let x = json_to_array2(&fx.input["X"]);
    let n_components = fx.params["n_components"].as_u64().unwrap_or(2) as usize;

    let fitted = MDS::new(n_components).fit(&x, &()).expect("MDS fit");
    let emb = fitted.embedding();
    finite_and_shaped(emb, x.nrows(), n_components, "MDS.embedding");
}

#[test]
fn conformance_lle() {
    let fx = load_fixture("lle");
    let x = json_to_array2(&fx.input["X"]);
    let n_components = fx.params["n_components"].as_u64().unwrap_or(2) as usize;
    let n_neighbors = fx.params["n_neighbors"].as_u64().unwrap_or(10) as usize;

    let fitted = LLE::new(n_components)
        .with_n_neighbors(n_neighbors)
        .fit(&x, &())
        .expect("LLE fit");
    finite_and_shaped(fitted.embedding(), x.nrows(), n_components, "LLE.embedding");
}

#[test]
fn conformance_spectral_embedding() {
    let fx = load_fixture("spectral_embedding");
    let x = json_to_array2(&fx.input["X"]);
    let n_components = fx.params["n_components"].as_u64().unwrap_or(2) as usize;

    let fitted = SpectralEmbedding::new(n_components)
        .fit(&x, &())
        .expect("SpectralEmbedding fit");
    finite_and_shaped(
        fitted.embedding(),
        x.nrows(),
        n_components,
        "SpectralEmbedding",
    );
}

#[test]
fn conformance_tsne() {
    let fx = load_fixture("tsne");
    let x = json_to_array2(&fx.input["X"]);
    let n_components = fx.params["n_components"].as_u64().unwrap_or(2) as usize;
    let perplexity = fx.params["perplexity"].as_f64().unwrap_or(10.0);
    let max_iter = fx.params["max_iter"].as_u64().unwrap_or(300) as usize;
    let random_state = fx.params["random_state"].as_u64().unwrap_or(42);

    let model = Tsne::new()
        .with_n_components(n_components)
        .with_perplexity(perplexity)
        .with_n_iter(max_iter)
        .with_random_state(random_state);
    let fitted = model.fit(&x, &()).expect("t-SNE fit");
    finite_and_shaped(fitted.embedding(), x.nrows(), n_components, "t-SNE");
}