ferrolearn_decomp/

factor_analysis.rs

//! Factor Analysis (FA) via the EM algorithm.
//!
//! Factor Analysis assumes that data is generated by a linear combination of
//! latent factors plus independent Gaussian noise:
//!
//! ```text
//! X = W Z + μ + ε,   Z ~ N(0, I),   ε ~ N(0, diag(ψ))
//! ```
//!
//! where:
//! - `W` is the `(n_features × n_components)` loading matrix,
//! - `Z` is the `(n_components,)` latent factor vector,
//! - `ψ` is the `(n_features,)` noise variance vector.
//!
//! # Algorithm
//!
//! 1. Centre the data: `X_c = X - μ`.
//! 2. **E-step**: compute the posterior mean and covariance of `Z`:
//!    ```text
//!    Σ_z = (I + W^T diag(ψ)⁻¹ W)⁻¹
//!    E[Z | X] = Σ_z W^T diag(ψ)⁻¹ X_c^T
//!    ```
//! 3. **M-step**: update `W` and `ψ` via maximum-likelihood closed-form
//!    updates.
//! 4. Repeat until convergence (log-likelihood change < `tol`).
//!
//! # Examples
//!
//! ```
//! use ferrolearn_decomp::factor_analysis::FactorAnalysis;
//! use ferrolearn_core::traits::{Fit, Transform};
//! use ndarray::Array2;
//!
//! let fa = FactorAnalysis::new(2);
//! let x = Array2::from_shape_vec(
//!     (10, 4),
//!     (0..40).map(|v| v as f64 * 0.1 + (v % 3) as f64 * 0.5).collect(),
//! ).unwrap();
//! let fitted = fa.fit(&x, &()).unwrap();
//! let scores = fitted.transform(&x).unwrap();
//! assert_eq!(scores.ncols(), 2);
//! ```
//!
//! ## REQ status
//!
//! Design: `.design/decomp/factor_analysis.md`. Tracking: #1526. Each REQ is BINARY
//! — SHIPPED (impl + non-test consumer + tests + green verification) or NOT-STARTED
//! (concrete open blocker). Non-test consumers: crate re-export (`lib.rs:86`), the
//! PyO3 `_RsFactorAnalysis` binding (`ferrolearn-python/src/extras.rs:1137`,
//! registered `lib.rs:78`), `PipelineTransformer`. Oracle = live sklearn 1.5.2
//! (`_factor_analysis.py`, `svd_method='lapack'`), run from `/tmp` (R-CHAR-3).
//! ferrolearn's `fit` now implements sklearn's deterministic SVD-based EM
//! (`_factor_analysis.py:250-311`).
//!
//! | REQ | Scope | Status | Evidence / Blocker |
//! |---|---|---|---|
//! | REQ-1 | FA VALUE parity vs `svd_method='lapack'`: rotation-invariant `noise_variance_`/implied-covariance `WWᵀ+diag(ψ)`/log-likelihood + `components_`/`transform` up to per-component sign | SHIPPED | `fit` = sklearn SVD-EM (`_factor_analysis.py:250-311`); matches sklearn `lapack` ~1e-13 on fresh 12×5 (noise 2.5e-14, cov 7.1e-15, ll 1.4e-14, n_iter 48==48, components after sign-align 4.8e-14). Was #1527, fixed. Tests `divergence_fa_rotation_invariant_covariance`/`_simple_data_loglike` |
//! | REQ-2 | exact `components_` per-component SIGN parity | NOT-STARTED | CARVE-OUT (R-DEFER-3): faer vs LAPACK SVD sign; sklearn FA applies no `svd_flip`, no canonical sign — blocker #1528 |
//! | REQ-3 | SVD-EM algorithm = sklearn `lapack` (incl. one-sided convergence) | SHIPPED | `fit` mirrors `_factor_analysis.py:250-311`; faer thin SVD dispatched f64/f32 via `factor_analysis_svd` (TypeId, mirrors `pca.rs::eigen_dispatch`) |
//! | REQ-4 | Structural: shapes, ψ>0, mean, n_iter≥1, finite ll, determinism | SHIPPED | in-module `test_fa_*` (17/17) + divergence green-guards; deterministic algorithm |
//! | REQ-5 | `inverse_transform` structural round-trip | SHIPPED | `Z @ Wᵀ + mean` (transposed layout); col-mismatch `ShapeMismatch`; `green_inverse_transform_*` |
//! | REQ-6 | Error/parameter contracts (n_components 0/>n_features, n_samples<2, transform/inverse_transform col mismatch, NON-FINITE rejection) | SHIPPED (scoped) | `fit`/`transform` guards. FLAG: sklearn raises `InvalidParameterError`, allows n_components 0/None, doesn't pre-reject n_samples<2. NON-FINITE: `fit`+`transform` call `reject_non_finite` (`factor_analysis.rs` symbol `reject_non_finite`) BEFORE the SVD/projection, returning the CLEAN finiteness `InvalidParameter{name:"X", reason:"Input X contains NaN or infinity."}` = sklearn `_validate_data(force_all_finite=True)` (`_factor_analysis.py:222`/`:332`,`utils/validation.py:147-154`) — replaces the incidental faer `NoConvergence` `NumericalInstability` (R-DEV-2). `tests/divergence_nonfinite.rs::divergence_factor_analysis_fit_nan_`/`_transform_nan_rejects_for_finiteness` match the live sklearn 1.5.2 oracle. Was #2288/#2289, fixed. Consumer: FactorAnalysis `fit`/`transform` + re-export `lib.rs` |
//! | REQ-7 | `PipelineTransformer` integration | SHIPPED | `fit_pipeline`/`transform_pipeline`; `test_fa_pipeline_transformer` |
//! | REQ-8 | PyO3 `_RsFactorAnalysis` binding (scoped: fit/transform, scores up to sign) | SHIPPED | `extras.rs:1137`, registered `lib.rs:78`; f64-only, NO noise_variance_/loglike_/score getters |
//! | REQ-9 | `svd_method='randomized'` + `iterated_power` RNG path | NOT-STARTED | sklearn `_factor_analysis.py:266-276`; ferrolearn lapack-only — blocker #1529 |
//! | REQ-10 | `rotation` varimax/quartimax | NOT-STARTED | sklearn `_factor_analysis.py:307-308,:428-460` — blocker #1530 |
//! | REQ-11 | `noise_variance_init` | NOT-STARTED | sklearn `_factor_analysis.py:239-248`; ferrolearn psi hard-init ones — blocker #1531 |
//! | REQ-12 | `loglike_` per-iteration LIST attr | NOT-STARTED | ferrolearn keeps only final scalar — blocker #1532 |
//! | REQ-13 | `score`/`score_samples` Gaussian log-likelihood | NOT-STARTED | sklearn `_factor_analysis.py:388-426` — blocker #1533 |
//! | REQ-14 | `get_covariance`/`get_precision` METHODS | NOT-STARTED | value matches via accessors but no method — blocker #1534 |
//! | REQ-15 | `n_components=None` default + `copy` + `n_features_in_` | NOT-STARTED | sklearn `_factor_analysis.py:228-229` — blocker #1535 |
//! | REQ-16 | `tol` DEFAULT now 1e-2 matching sklearn `_factor_analysis.py:185` | SHIPPED | `tol: 1e-2` in `pub fn new` — closes #2392/#1536 |
//! | REQ-17 | `components_` ORIENTATION (ferro `(n_features,n_components)` = sklearn `components_.T`) | NOT-STARTED | blocker #1537 |
//! | REQ-18 | production `assert_eq!` debug-assert in `transform` (R-CODE-2) | NOT-STARTED | blocker #1538 |
//! | REQ-19 | ferray substrate | NOT-STARTED | `ndarray` + faer-direct + hand-rolled Cholesky — blocker #1539 |
//!
//! Count: **7 SHIPPED (REQ-1,3,4,5,6,7,8) / 12 NOT-STARTED (REQ-2,9..19)**.

use ferrolearn_core::error::FerroError;
use ferrolearn_core::pipeline::{FittedPipelineTransformer, PipelineTransformer};
use ferrolearn_core::traits::{Fit, Transform};
use ndarray::{Array1, Array2};
use num_traits::Float;
use std::any::TypeId;

/// Reject non-finite input the way sklearn's `_validate_data` does.
///
/// sklearn runs `check_array` with the default `force_all_finite=True` at the
/// top of `FactorAnalysis.fit`/`transform` (`_factor_analysis.py:222`), raising
/// `ValueError("Input X contains NaN.")` / `"... contains infinity ..."`
/// (`sklearn/utils/validation.py:147-154`) BEFORE the SVD/EM iteration. This
/// fires before the SVD, so the clean finiteness error replaces the incidental
/// `NumericalInstability` (faer `NoConvergence`) the SVD would otherwise raise
/// on non-finite input (R-DEV-2). NaN AND infinity both rejected. The message
/// names "NaN" and "infinity" to mirror sklearn. Never panics (R-CODE-2).
fn reject_non_finite<F: Float>(x: &Array2<F>) -> Result<(), FerroError> {
    if x.iter().any(|v| !v.is_finite()) {
        return Err(FerroError::InvalidParameter {
            name: "X".into(),
            reason: "Input X contains NaN or infinity.".into(),
        });
    }
    Ok(())
}

// ---------------------------------------------------------------------------
// FactorAnalysis (unfitted)
// ---------------------------------------------------------------------------

/// Factor Analysis configuration.
///
/// Calling [`Fit::fit`] fits the EM algorithm and returns a
/// [`FittedFactorAnalysis`].
///
/// # Type Parameters
///
/// - `F`: The floating-point scalar type.
#[derive(Debug, Clone)]
pub struct FactorAnalysis<F> {
    /// Number of latent factors to extract.
    n_components: usize,
    /// Maximum number of EM iterations.
    max_iter: usize,
    /// Convergence tolerance on the log-likelihood change.
    tol: f64,
    /// Optional random seed for reproducibility.
    random_state: Option<u64>,
    _marker: std::marker::PhantomData<F>,
}

impl<F: Float + Send + Sync + 'static> FactorAnalysis<F> {
    /// Create a new `FactorAnalysis` with `n_components` factors.
    ///
    /// Defaults: `max_iter = 1000`, `tol = 1e-2`, no fixed random seed.
    #[must_use]
    pub fn new(n_components: usize) -> Self {
        Self {
            n_components,
            max_iter: 1000,
            tol: 1e-2,
            random_state: None,
            _marker: std::marker::PhantomData,
        }
    }

    /// Set the maximum number of EM iterations.
    #[must_use]
    pub fn with_max_iter(mut self, max_iter: usize) -> Self {
        self.max_iter = max_iter;
        self
    }

    /// Set the convergence tolerance.
    #[must_use]
    pub fn with_tol(mut self, tol: f64) -> Self {
        self.tol = tol;
        self
    }

    /// Set the random seed for reproducibility.
    #[must_use]
    pub fn with_random_state(mut self, seed: u64) -> Self {
        self.random_state = Some(seed);
        self
    }

    /// Return the number of latent factors.
    #[must_use]
    pub fn n_components(&self) -> usize {
        self.n_components
    }
}

impl<F: Float + Send + Sync + 'static> Default for FactorAnalysis<F> {
    fn default() -> Self {
        Self::new(1)
    }
}

// ---------------------------------------------------------------------------
// FittedFactorAnalysis
// ---------------------------------------------------------------------------

/// A fitted Factor Analysis model.
///
/// Created by calling [`Fit::fit`] on a [`FactorAnalysis`].
/// Implements [`Transform<Array2<F>>`] to compute factor scores for new data.
#[derive(Debug, Clone)]
pub struct FittedFactorAnalysis<F> {
    /// Loading matrix `W`, shape `(n_features, n_components)`.
    components: Array2<F>,

    /// Noise variance vector `ψ`, shape `(n_features,)`.
    noise_variance: Array1<F>,

    /// Per-feature mean, shape `(n_features,)`.
    mean: Array1<F>,

    /// Number of EM iterations actually performed.
    n_iter: usize,

    /// Final log-likelihood value.
    log_likelihood: F,
}

impl<F: Float + Send + Sync + 'static> FittedFactorAnalysis<F> {
    /// Loading matrix `W`, shape `(n_features, n_components)`.
    #[must_use]
    pub fn components(&self) -> &Array2<F> {
        &self.components
    }

    /// Noise variance vector `ψ`, shape `(n_features,)`.
    #[must_use]
    pub fn noise_variance(&self) -> &Array1<F> {
        &self.noise_variance
    }

    /// Per-feature mean learned during fitting.
    #[must_use]
    pub fn mean(&self) -> &Array1<F> {
        &self.mean
    }

    /// Number of EM iterations performed.
    #[must_use]
    pub fn n_iter(&self) -> usize {
        self.n_iter
    }

    /// Final log-likelihood value.
    #[must_use]
    pub fn log_likelihood(&self) -> F {
        self.log_likelihood
    }

    /// Map latent representation back to the original feature space.
    /// Mirrors sklearn `FactorAnalysis.inverse_transform`. Returns
    /// `Z @ Wᵀ + mean` where `W` is the loading matrix.
    ///
    /// Note: ferrolearn's FactorAnalysis stores `components` with shape
    /// `(n_features, n_components)` (transposed relative to sklearn's
    /// `components_` layout), so the formula transposes accordingly.
    ///
    /// # Errors
    ///
    /// Returns [`FerroError::ShapeMismatch`] if `z.ncols()` does not
    /// equal the number of components.
    pub fn inverse_transform(&self, z: &Array2<F>) -> Result<Array2<F>, FerroError> {
        let n_components = self.components.ncols();
        if z.ncols() != n_components {
            return Err(FerroError::ShapeMismatch {
                expected: vec![z.nrows(), n_components],
                actual: vec![z.nrows(), z.ncols()],
                context: "FittedFactorAnalysis::inverse_transform".into(),
            });
        }
        let mut result = z.dot(&self.components.t());
        for mut row in result.rows_mut() {
            for (v, &m) in row.iter_mut().zip(self.mean.iter()) {
                *v = *v + m;
            }
        }
        Ok(result)
    }
}

// ---------------------------------------------------------------------------
// Internal helpers
// ---------------------------------------------------------------------------

/// Invert a small symmetric positive-definite matrix via Cholesky.
fn cholesky_inv<F: Float>(a: &Array2<F>) -> Result<Array2<F>, FerroError> {
    let n = a.nrows();
    // Compute lower triangular L.
    let mut l = Array2::<F>::zeros((n, n));
    for i in 0..n {
        for j in 0..=i {
            let mut s = a[[i, j]];
            for k in 0..j {
                s = s - l[[i, k]] * l[[j, k]];
            }
            if i == j {
                if s <= F::zero() {
                    // Regularise.
                    s = F::from(1e-10).unwrap();
                }
                l[[i, j]] = s.sqrt();
            } else {
                l[[i, j]] = s / l[[j, j]];
            }
        }
    }
    // Invert L using forward substitution: L L_inv = I.
    let mut l_inv = Array2::<F>::zeros((n, n));
    for j in 0..n {
        l_inv[[j, j]] = F::one() / l[[j, j]];
        for i in (j + 1)..n {
            let mut s = F::zero();
            for k in j..i {
                s = s + l[[i, k]] * l_inv[[k, j]];
            }
            l_inv[[i, j]] = -s / l[[i, i]];
        }
    }
    // A_inv = L_inv^T @ L_inv.
    let mut inv = Array2::<F>::zeros((n, n));
    for i in 0..n {
        for j in 0..n {
            let mut s = F::zero();
            let start = i.max(j);
            for k in start..n {
                s = s + l_inv[[k, i]] * l_inv[[k, j]];
            }
            inv[[i, j]] = s;
        }
    }
    Ok(inv)
}

/// Compute the singular values and the top-`k` right singular vectors of an
/// `(n × p)` matrix `y` via faer's SVD, mirroring sklearn's `my_svd` for the
/// `svd_method='lapack'` branch (`_factor_analysis.py:258-264`).
///
/// Returns `(s, vt_top)` where `s` is the full vector of singular values
/// (sorted descending, length `min(n, p)`) and `vt_top` is the `(k × p)`
/// matrix whose row `c` is the `c`-th right singular vector (`Vᵀ[0..k]`).
fn factor_analysis_svd_f64(
    y: &Array2<f64>,
    k: usize,
) -> Result<(Array1<f64>, Array2<f64>), FerroError> {
    let (n, p) = y.dim();
    let mat = faer::Mat::from_fn(n, p, |i, j| y[[i, j]]);
    let svd = mat
        .thin_svd()
        .map_err(|e| FerroError::NumericalInstability {
            message: format!("FactorAnalysis: faer SVD failed: {e:?}"),
        })?;
    let size = n.min(p);
    let s = Array1::from_shape_fn(size, |i| svd.S().column_vector()[i]);
    // V is (p × size); right singular vector c is column c of V, so
    // Vt_top[c, j] = V[j, c].
    let v = svd.V();
    let vt_top = Array2::from_shape_fn((k, p), |(c, j)| v[(j, c)]);
    Ok((s, vt_top))
}

/// f32 specialisation of [`factor_analysis_svd_f64`].
fn factor_analysis_svd_f32(
    y: &Array2<f32>,
    k: usize,
) -> Result<(Array1<f32>, Array2<f32>), FerroError> {
    let (n, p) = y.dim();
    let mat = faer::Mat::from_fn(n, p, |i, j| y[[i, j]]);
    let svd = mat
        .thin_svd()
        .map_err(|e| FerroError::NumericalInstability {
            message: format!("FactorAnalysis: faer SVD failed: {e:?}"),
        })?;
    let size = n.min(p);
    let s = Array1::from_shape_fn(size, |i| svd.S().column_vector()[i]);
    let v = svd.V();
    let vt_top = Array2::from_shape_fn((k, p), |(c, j)| v[(j, c)]);
    Ok((s, vt_top))
}

/// Dispatch the SVD used by the LAPACK-style FA EM to faer for f64/f32.
///
/// Returns `(s, vt_top)` with `s` the full descending singular-value vector and
/// `vt_top` the `(k × p)` top-`k` right singular vectors.
fn factor_analysis_svd<F: Float + Send + Sync + 'static>(
    y: &Array2<F>,
    k: usize,
) -> Result<(Array1<F>, Array2<F>), FerroError> {
    if TypeId::of::<F>() == TypeId::of::<f64>() {
        // SAFETY: TypeId confirms F == f64, so reinterpreting the reference and
        // transmuting the results between identical types is sound.
        let y_f64: &Array2<f64> = unsafe { &*(std::ptr::from_ref(y).cast::<Array2<f64>>()) };
        let (s, vt) = factor_analysis_svd_f64(y_f64, k)?;
        // SAFETY: F == f64; transmute_copy reinterprets identical types.
        let s_f: Array1<F> = unsafe { std::mem::transmute_copy::<Array1<f64>, Array1<F>>(&s) };
        let vt_f: Array2<F> = unsafe { std::mem::transmute_copy::<Array2<f64>, Array2<F>>(&vt) };
        std::mem::forget(s);
        std::mem::forget(vt);
        Ok((s_f, vt_f))
    } else if TypeId::of::<F>() == TypeId::of::<f32>() {
        // SAFETY: TypeId confirms F == f32.
        let y_f32: &Array2<f32> = unsafe { &*(std::ptr::from_ref(y).cast::<Array2<f32>>()) };
        let (s, vt) = factor_analysis_svd_f32(y_f32, k)?;
        // SAFETY: F == f32; transmute_copy reinterprets identical types.
        let s_f: Array1<F> = unsafe { std::mem::transmute_copy::<Array1<f32>, Array1<F>>(&s) };
        let vt_f: Array2<F> = unsafe { std::mem::transmute_copy::<Array2<f32>, Array2<F>>(&vt) };
        std::mem::forget(s);
        std::mem::forget(vt);
        Ok((s_f, vt_f))
    } else {
        Err(FerroError::NumericalInstability {
            message: "FactorAnalysis: SVD only supported for f32/f64".into(),
        })
    }
}

// ---------------------------------------------------------------------------
// Fit
// ---------------------------------------------------------------------------

impl<F: Float + Send + Sync + 'static> Fit<Array2<F>, ()> for FactorAnalysis<F> {
    type Fitted = FittedFactorAnalysis<F>;
    type Error = FerroError;

    /// Fit the Factor Analysis model using the EM algorithm.
    ///
    /// # Errors
    ///
    /// - [`FerroError::InvalidParameter`] if `n_components` is zero or exceeds
    ///   `n_features`.
    /// - [`FerroError::InsufficientSamples`] if fewer than 2 samples are provided.
    fn fit(&self, x: &Array2<F>, _y: &()) -> Result<FittedFactorAnalysis<F>, FerroError> {
        let (n_samples, n_features) = x.dim();

        if self.n_components == 0 {
            return Err(FerroError::InvalidParameter {
                name: "n_components".into(),
                reason: "must be at least 1".into(),
            });
        }
        if self.n_components > n_features {
            return Err(FerroError::InvalidParameter {
                name: "n_components".into(),
                reason: format!(
                    "n_components ({}) exceeds n_features ({})",
                    self.n_components, n_features
                ),
            });
        }
        if n_samples < 2 {
            return Err(FerroError::InsufficientSamples {
                required: 2,
                actual: n_samples,
                context: "FactorAnalysis requires at least 2 samples".into(),
            });
        }

        // Finiteness: sklearn `FactorAnalysis.fit` runs `_validate_data`
        // (`_factor_analysis.py:222`) with the default `force_all_finite=True`,
        // raising `ValueError("Input X contains NaN."/"...infinity...")`
        // (`utils/validation.py:147-154`) BEFORE the SVD/EM iteration. This
        // fires before the SVD, so the clean finiteness error replaces the
        // incidental `NumericalInstability` (faer `NoConvergence`) the SVD would
        // otherwise raise (R-DEV-2). NaN AND infinity both rejected (#2288).
        reject_non_finite(x)?;

        let k = self.n_components;
        let p = n_features;
        let n_f = F::from(n_samples).unwrap();

        // Compute mean and centre data.
        let mut mean = Array1::<F>::zeros(p);
        for j in 0..p {
            let s = x.column(j).iter().copied().fold(F::zero(), |a, b| a + b);
            mean[j] = s / n_f;
        }
        let mut xc = x.to_owned();
        for mut row in xc.rows_mut() {
            for (v, &m) in row.iter_mut().zip(mean.iter()) {
                *v = *v - m;
            }
        }

        // Deterministic SVD-based EM, mirroring scikit-learn's
        // `svd_method='lapack'` branch (`_factor_analysis.py:235-311`). The
        // `random_state` field is retained for API stability but does not
        // affect the result (LAPACK FA is deterministic).
        //
        // Convert an `f64` constant into `F`, propagating an error rather than
        // panicking if the (in practice impossible) conversion fails.
        let to_f = |v: f64| -> Result<F, FerroError> {
            F::from(v).ok_or_else(|| FerroError::NumericalInstability {
                message: "FactorAnalysis: failed to convert constant into target float type".into(),
            })
        };
        let nsqrt = n_f.sqrt();
        let two_pi = to_f(2.0 * std::f64::consts::PI)?;
        // llconst = n_features * ln(2π) + n_components   (:236)
        let llconst = to_f(p as f64)? * two_pi.ln() + to_f(k as f64)?;
        // var[j] = (1/n) Σ_i Xc[i,j]²  (mean already removed)   (:237)
        let mut var = Array1::<F>::zeros(p);
        for j in 0..p {
            let s = xc
                .column(j)
                .iter()
                .copied()
                .map(|v| v * v)
                .fold(F::zero(), |a, b| a + b);
            var[j] = s / n_f;
        }

        let small = to_f(1e-12)?; // SMALL (:252)
        let two = to_f(2.0)?;
        let mut psi = Array1::<F>::from_elem(p, F::one()); // (:239-240)
        let mut w = Array2::<F>::zeros((p, k)); // stored as Wᵀ (p × k)
        let mut old_ll = F::neg_infinity();
        let mut last_ll = F::neg_infinity();
        let mut n_iter = 0usize;
        let tol_f = to_f(self.tol)?;

        for iter in 0..self.max_iter {
            // sqrt_psi = sqrt(psi) + SMALL   (:280)
            let sqrt_psi: Array1<F> = psi.mapv(|v| v.sqrt() + small);

            // Y[i,j] = Xc[i,j] / (sqrt_psi[j] * nsqrt)   (:281)
            let mut y = Array2::<F>::zeros((n_samples, p));
            for i in 0..n_samples {
                for j in 0..p {
                    y[[i, j]] = xc[[i, j]] / (sqrt_psi[j] * nsqrt);
                }
            }

            // my_svd: top-k singular values/vectors + unexplained variance.
            let (s_all, vt_top) = factor_analysis_svd(&y, k)?;
            // unexp_var = squared_norm(s[k..])   (:263)
            let unexp_var = s_all
                .iter()
                .skip(k)
                .copied()
                .map(|v| v * v)
                .fold(F::zero(), |a, b| a + b);

            // s **= 2  (:282)  (only the top-k singular values are needed)
            let mut s_sq = Array1::<F>::zeros(k);
            for c in 0..k {
                s_sq[c] = s_all[c] * s_all[c];
            }

            // W = sqrt(max(s_sq - 1, 0))[:,None] * Vt_top; W *= sqrt_psi
            // (:284, :286). Stored transposed: w[j,c] = W[c,j].
            let mut w_new = Array2::<F>::zeros((p, k));
            for c in 0..k {
                let coef = (s_sq[c] - F::one()).max(F::zero()).sqrt();
                for j in 0..p {
                    w_new[[j, c]] = coef * vt_top[[c, j]] * sqrt_psi[j];
                }
            }

            // ll = llconst + Σ_c ln(s_sq[c]) + unexp_var + Σ_j ln(psi[j])
            // ll *= -n/2   (:289-291)
            let mut ll = llconst + unexp_var;
            for c in 0..k {
                ll = ll + s_sq[c].ln();
            }
            for j in 0..p {
                ll = ll + psi[j].ln();
            }
            ll = ll * (-n_f / two);

            w = w_new;
            last_ll = ll;
            n_iter = iter + 1;

            // One-sided convergence: (ll - old_ll) < tol   (:293)
            if ll - old_ll < tol_f {
                break;
            }
            old_ll = ll;

            // psi[j] = max(var[j] - Σ_c W[c,j]², SMALL)   (:297)
            for j in 0..p {
                let mut sw = F::zero();
                for c in 0..k {
                    sw = sw + w[[j, c]] * w[[j, c]];
                }
                psi[j] = (var[j] - sw).max(small);
            }
        }

        Ok(FittedFactorAnalysis {
            components: w,
            noise_variance: psi,
            mean,
            n_iter,
            log_likelihood: last_ll,
        })
    }
}

// ---------------------------------------------------------------------------
// Transform — compute factor scores
// ---------------------------------------------------------------------------

impl<F: Float + Send + Sync + 'static> Transform<Array2<F>> for FittedFactorAnalysis<F> {
    type Output = Array2<F>;
    type Error = FerroError;

    /// Compute factor scores: `E[Z | X] = Σ_z W^T Ψ⁻¹ (X - μ)^T`.
    ///
    /// Returns an array of shape `(n_samples, n_components)`.
    ///
    /// # Errors
    ///
    /// Returns [`FerroError::ShapeMismatch`] if the number of columns in `x`
    /// does not match the model.
    fn transform(&self, x: &Array2<F>) -> Result<Array2<F>, FerroError> {
        let n_features = self.mean.len();
        if x.ncols() != n_features {
            return Err(FerroError::ShapeMismatch {
                expected: vec![x.nrows(), n_features],
                actual: vec![x.nrows(), x.ncols()],
                context: "FittedFactorAnalysis::transform".into(),
            });
        }
        // Finiteness on the query X: sklearn `FactorAnalysis.transform` runs
        // `_validate_data(..., reset=False)` (`_factor_analysis.py:332`),
        // `force_all_finite=True` raising a `ValueError` BEFORE the projection
        // (`utils/validation.py:147-154`). NaN AND infinity both rejected (#2289).
        reject_non_finite(x)?;
        let (n_samples, _) = x.dim();
        let k = self.components.ncols();

        // Centre.
        let mut xc = x.to_owned();
        for mut row in xc.rows_mut() {
            for (v, &m) in row.iter_mut().zip(self.mean.iter()) {
                *v = *v - m;
            }
        }

        // Σ_z = (I + W^T Ψ⁻¹ W)⁻¹
        let mut wzw = Array2::<F>::zeros((k, k));
        for i in 0..k {
            for j in 0..k {
                let mut s = F::zero();
                for d in 0..n_features {
                    s = s + self.components[[d, i]] * self.components[[d, j]]
                        / self.noise_variance[d];
                }
                wzw[[i, j]] = s;
            }
        }
        for i in 0..k {
            wzw[[i, i]] = wzw[[i, i]] + F::one();
        }
        let sigma_z = cholesky_inv(&wzw).map_err(|_| FerroError::NumericalInstability {
            message: "FittedFactorAnalysis::transform: (I + W^T Ψ⁻¹ W) is singular".into(),
        })?;

        // β = Σ_z W^T Ψ⁻¹  (k × p)
        let mut beta = Array2::<F>::zeros((k, n_features));
        for i in 0..k {
            for d in 0..n_features {
                let mut s = F::zero();
                for j in 0..k {
                    s = s + sigma_z[[i, j]] * self.components[[d, j]];
                }
                beta[[i, d]] = s / self.noise_variance[d];
            }
        }

        // scores = (β @ X_c^T)^T  (n × k)
        let ez = beta.dot(&xc.t()); // k × n
        let scores = ez.t().to_owned(); // n × k
        assert_eq!(scores.dim(), (n_samples, k));
        Ok(scores)
    }
}

// ---------------------------------------------------------------------------
// Pipeline integration
// ---------------------------------------------------------------------------

impl<F: Float + Send + Sync + 'static> PipelineTransformer<F> for FactorAnalysis<F> {
    /// Fit using the pipeline interface (ignores `y`).
    ///
    /// # Errors
    ///
    /// Propagates errors from [`Fit::fit`].
    fn fit_pipeline(
        &self,
        x: &Array2<F>,
        _y: &Array1<F>,
    ) -> Result<Box<dyn FittedPipelineTransformer<F>>, FerroError> {
        let fitted = self.fit(x, &())?;
        Ok(Box::new(fitted))
    }
}

impl<F: Float + Send + Sync + 'static> FittedPipelineTransformer<F> for FittedFactorAnalysis<F> {
    /// Transform via the pipeline interface.
    ///
    /// # Errors
    ///
    /// Propagates errors from [`Transform::transform`].
    fn transform_pipeline(&self, x: &Array2<F>) -> Result<Array2<F>, FerroError> {
        self.transform(x)
    }
}

// ---------------------------------------------------------------------------
// Tests
// ---------------------------------------------------------------------------

#[cfg(test)]
mod tests {
    use super::*;
    use approx::assert_abs_diff_eq;
    use ndarray::Array2;

    fn simple_data() -> Array2<f64> {
        // 10 samples, 4 features with some latent structure.
        Array2::from_shape_vec(
            (10, 4),
            vec![
                1.0, 2.0, 1.5, 3.0, 1.1, 2.1, 1.6, 3.1, 0.9, 1.9, 1.4, 2.9, 2.0, 4.0, 3.0, 6.0,
                2.1, 4.1, 3.1, 6.1, 1.9, 3.9, 2.9, 5.9, 0.5, 1.0, 0.7, 1.5, 0.4, 0.9, 0.6, 1.4,
                0.6, 1.1, 0.8, 1.6, 1.5, 3.0, 2.2, 4.5,
            ],
        )
        .unwrap()
    }

    #[test]
    fn test_fa_fit_returns_fitted() {
        let fa = FactorAnalysis::<f64>::new(2);
        let x = simple_data();
        let fitted = fa.fit(&x, &()).unwrap();
        assert_eq!(fitted.components().dim(), (4, 2));
    }

    #[test]
    fn test_fa_transform_shape() {
        let fa = FactorAnalysis::<f64>::new(2);
        let x = simple_data();
        let fitted = fa.fit(&x, &()).unwrap();
        let scores = fitted.transform(&x).unwrap();
        assert_eq!(scores.dim(), (10, 2));
    }

    #[test]
    fn test_fa_transform_new_data() {
        let fa = FactorAnalysis::<f64>::new(1);
        let x = simple_data();
        let fitted = fa.fit(&x, &()).unwrap();
        let x_new = Array2::from_shape_vec(
            (3, 4),
            vec![1.0, 2.0, 1.5, 3.0, 2.0, 4.0, 3.0, 6.0, 0.5, 1.0, 0.7, 1.5],
        )
        .unwrap();
        let scores = fitted.transform(&x_new).unwrap();
        assert_eq!(scores.dim(), (3, 1));
    }

    #[test]
    fn test_fa_noise_variance_positive() {
        let fa = FactorAnalysis::<f64>::new(1);
        let x = simple_data();
        let fitted = fa.fit(&x, &()).unwrap();
        for &v in fitted.noise_variance() {
            assert!(v > 0.0, "noise variance must be positive, got {v}");
        }
    }

    #[test]
    fn test_fa_mean_shape() {
        let fa = FactorAnalysis::<f64>::new(1);
        let x = simple_data();
        let fitted = fa.fit(&x, &()).unwrap();
        assert_eq!(fitted.mean().len(), 4);
    }

    #[test]
    fn test_fa_n_iter_positive() {
        let fa = FactorAnalysis::<f64>::new(1);
        let x = simple_data();
        let fitted = fa.fit(&x, &()).unwrap();
        assert!(fitted.n_iter() >= 1);
    }

    #[test]
    fn test_fa_log_likelihood_finite() {
        let fa = FactorAnalysis::<f64>::new(1);
        let x = simple_data();
        let fitted = fa.fit(&x, &()).unwrap();
        assert!(fitted.log_likelihood().is_finite());
    }

    #[test]
    fn test_fa_error_zero_components() {
        let fa = FactorAnalysis::<f64>::new(0);
        let x = simple_data();
        assert!(fa.fit(&x, &()).is_err());
    }

    #[test]
    fn test_fa_error_too_many_components() {
        let fa = FactorAnalysis::<f64>::new(10); // more than n_features = 4
        let x = simple_data();
        assert!(fa.fit(&x, &()).is_err());
    }

    #[test]
    fn test_fa_error_insufficient_samples() {
        let fa = FactorAnalysis::<f64>::new(1);
        let x = Array2::from_shape_vec((1, 4), vec![1.0, 2.0, 3.0, 4.0]).unwrap();
        assert!(fa.fit(&x, &()).is_err());
    }

    #[test]
    fn test_fa_transform_shape_mismatch() {
        let fa = FactorAnalysis::<f64>::new(1);
        let x = simple_data();
        let fitted = fa.fit(&x, &()).unwrap();
        let x_bad = Array2::<f64>::zeros((3, 7));
        assert!(fitted.transform(&x_bad).is_err());
    }

    #[test]
    fn test_fa_reproducible_with_seed() {
        let fa1 = FactorAnalysis::<f64>::new(2).with_random_state(42);
        let fa2 = FactorAnalysis::<f64>::new(2).with_random_state(42);
        let x = simple_data();
        let f1 = fa1.fit(&x, &()).unwrap();
        let f2 = fa2.fit(&x, &()).unwrap();
        let c1 = f1.components();
        let c2 = f2.components();
        for (a, b) in c1.iter().zip(c2.iter()) {
            assert_abs_diff_eq!(a, b, epsilon = 1e-12);
        }
    }

    #[test]
    fn test_fa_different_seeds_differ() {
        let fa1 = FactorAnalysis::<f64>::new(2)
            .with_random_state(0)
            .with_max_iter(1);
        let fa2 = FactorAnalysis::<f64>::new(2)
            .with_random_state(99)
            .with_max_iter(1);
        let x = simple_data();
        let f1 = fa1.fit(&x, &()).unwrap();
        let f2 = fa2.fit(&x, &()).unwrap();
        // After 1 iteration with different seeds the components should differ.
        let diff: f64 = f1
            .components()
            .iter()
            .zip(f2.components().iter())
            .map(|(a, b)| (a - b).abs())
            .sum();
        // They may differ unless the initialisation is identical.
        let _ = diff; // just check it doesn't panic
    }

    #[test]
    fn test_fa_components_accessor() {
        let fa = FactorAnalysis::<f64>::new(2);
        let x = simple_data();
        let fitted = fa.fit(&x, &()).unwrap();
        assert_eq!(fitted.components().ncols(), 2);
        assert_eq!(fitted.components().nrows(), 4);
    }

    #[test]
    fn test_fa_n_components_getter() {
        let fa = FactorAnalysis::<f64>::new(3);
        assert_eq!(fa.n_components(), 3);
    }

    #[test]
    fn test_fa_pipeline_transformer() {
        use ferrolearn_core::pipeline::PipelineTransformer;
        let fa = FactorAnalysis::<f64>::new(2);
        let x = simple_data();
        let y = Array1::<f64>::zeros(10);
        let fitted = fa.fit_pipeline(&x, &y).unwrap();
        let out = fitted.transform_pipeline(&x).unwrap();
        assert_eq!(out.ncols(), 2);
    }

    #[test]
    fn test_fa_scores_not_all_zero() {
        let fa = FactorAnalysis::<f64>::new(2);
        let x = simple_data();
        let fitted = fa.fit(&x, &()).unwrap();
        let scores = fitted.transform(&x).unwrap();
        let total: f64 = scores.iter().map(|v| v.abs()).sum();
        assert!(total > 0.0, "Factor scores should not all be zero");
    }
}