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);
//! ```

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 rand::SeedableRng;
use rand_distr::{Distribution, StandardNormal};

// ---------------------------------------------------------------------------
// 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-3`, no fixed random seed.
    #[must_use]
    pub fn new(n_components: usize) -> Self {
        Self {
            n_components,
            max_iter: 1000,
            tol: 1e-3,
            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
    }
}

// ---------------------------------------------------------------------------
// 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 log-likelihood under the factor analysis model.
///
/// `log p(X) = -n/2 * [p*log(2π) + log|Σ| + tr(Σ⁻¹ S)]`
/// where `Σ = W W^T + diag(ψ)` and `S = X_c^T X_c / n`.
fn compute_log_likelihood<F: Float + Send + Sync + 'static>(
    x_centered: &Array2<F>,
    w: &Array2<F>,
    psi: &Array1<F>,
) -> F {
    let (n, p) = x_centered.dim();
    let k = w.ncols();
    // Σ = W W^T + diag(ψ)
    // We use the Woodbury identity for the log-det and trace.
    // log|Σ| = log|I_k + W^T Ψ⁻¹ W| + Σ_j log ψ_j
    let two_pi = F::from(2.0 * std::f64::consts::PI).unwrap();
    let n_f = F::from(n).unwrap();
    let p_f = F::from(p).unwrap();

    // W^T Ψ⁻¹ W: k × k
    let mut wtpsiw = Array2::<F>::zeros((k, k));
    for i in 0..k {
        for j in 0..k {
            let mut s = F::zero();
            for d in 0..p {
                s = s + w[[d, i]] * w[[d, j]] / psi[d];
            }
            wtpsiw[[i, j]] = s;
        }
    }
    // Add identity.
    for i in 0..k {
        wtpsiw[[i, i]] = wtpsiw[[i, i]] + F::one();
    }
    // log det of (I + W^T Ψ⁻¹ W) via Cholesky.
    let mut log_det_inner = F::zero();
    {
        let mut l = Array2::<F>::zeros((k, k));
        for i in 0..k {
            for j in 0..=i {
                let mut s = wtpsiw[[i, j]];
                for kk in 0..j {
                    s = s - l[[i, kk]] * l[[j, kk]];
                }
                if i == j {
                    s = if s > F::zero() {
                        s
                    } else {
                        F::from(1e-30).unwrap()
                    };
                    l[[i, j]] = s.sqrt();
                    log_det_inner = log_det_inner + l[[i, j]].ln();
                } else {
                    l[[i, j]] = s / l[[j, j]];
                }
            }
        }
        log_det_inner = log_det_inner * F::from(2.0).unwrap();
    }
    let log_det_psi: F = psi
        .iter()
        .copied()
        .map(|v| {
            let v_clamped = if v > F::zero() {
                v
            } else {
                F::from(1e-30).unwrap()
            };
            v_clamped.ln()
        })
        .fold(F::zero(), |a, b| a + b);
    let log_det_sigma = log_det_inner + log_det_psi;

    // Sample covariance S = X_c^T X_c / n.
    // tr(Σ⁻¹ S) using Woodbury: Σ⁻¹ = Ψ⁻¹ - Ψ⁻¹ W M⁻¹ W^T Ψ⁻¹
    // where M = I + W^T Ψ⁻¹ W.
    // tr(Σ⁻¹ S) = (1/n) Σ_i x_i^T Σ⁻¹ x_i
    // We compute it directly sample-by-sample for simplicity.
    // For efficiency, we use the factored form:
    // x^T Σ⁻¹ x = x^T Ψ⁻¹ x - (Ψ⁻¹ W m)^T M⁻¹ (W^T Ψ⁻¹ x)
    // where m = W^T Ψ⁻¹ x.

    // Invert M = I + W^T Ψ⁻¹ W.
    let m_inv = match cholesky_inv(&wtpsiw) {
        Ok(inv) => inv,
        Err(_) => return F::neg_infinity(),
    };

    let mut trace_sum = F::zero();
    for i in 0..n {
        // Ψ⁻¹ x_i
        let mut psi_inv_x = Array1::<F>::zeros(p);
        let mut xpsiinvx = F::zero();
        for d in 0..p {
            psi_inv_x[d] = x_centered[[i, d]] / psi[d];
            xpsiinvx = xpsiinvx + x_centered[[i, d]] * psi_inv_x[d];
        }
        // W^T Ψ⁻¹ x_i  (k-vector)
        let mut wtpx = Array1::<F>::zeros(k);
        for kk in 0..k {
            let mut s = F::zero();
            for d in 0..p {
                s = s + w[[d, kk]] * psi_inv_x[d];
            }
            wtpx[kk] = s;
        }
        // (W^T Ψ⁻¹ x)^T M⁻¹ (W^T Ψ⁻¹ x)
        let mut quad = F::zero();
        for ii in 0..k {
            let mut s = F::zero();
            for jj in 0..k {
                s = s + m_inv[[ii, jj]] * wtpx[jj];
            }
            quad = quad + wtpx[ii] * s;
        }
        trace_sum = trace_sum + xpsiinvx - quad;
    }
    let trace_term = trace_sum / n_f;

    // log p = -n/2 * [p*log(2π) + log|Σ| + tr(Σ⁻¹ S)]
    let half = F::from(0.5).unwrap();
    -n_f * half * (p_f * two_pi.ln() + log_det_sigma + trace_term)
}

// ---------------------------------------------------------------------------
// 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(),
            });
        }

        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;
            }
        }

        // Initialise W randomly, ψ = 1.
        let seed = self.random_state.unwrap_or(42);
        let mut rng = rand_xoshiro::Xoshiro256PlusPlus::seed_from_u64(seed);
        let std_normal = StandardNormal;
        let mut w = Array2::<F>::zeros((p, k));
        let scale = F::from(0.01).unwrap();
        for i in 0..p {
            for j in 0..k {
                let v: f64 = std_normal.sample(&mut rng);
                w[[i, j]] = F::from(v).unwrap() * scale;
            }
        }
        let mut psi = Array1::<F>::from_elem(p, F::one());

        let mut prev_ll = F::neg_infinity();
        let mut n_iter = 0usize;
        let tol_f = F::from(self.tol).unwrap();

        for iter in 0..self.max_iter {
            // --- E-step --------------------------------------------------------
            // Σ_z = (I_k + W^T Ψ⁻¹ W)⁻¹   shape k × k
            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..p {
                        s = s + w[[d, i]] * w[[d, j]] / psi[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: "FactorAnalysis: (I + W^T Ψ⁻¹ W) is singular".into(),
            })?;

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

            // E[Z | X] = β X_c^T   shape k × n
            let ez = beta.dot(&xc.t()); // k × n

            // E[Z Z^T | X] summed over samples = n * Σ_z + Σ_i e_i e_i^T
            // We keep the average: E_zzt = Σ_z + (1/n) Σ_i e_i e_i^T
            // shape k × k
            let ezz_t_sum = sigma_z.mapv(|v| v * n_f) + ez.dot(&ez.t()); // k × k

            // --- M-step --------------------------------------------------------
            // W_new = (Σ_i x_i e_i^T) (Σ_i e_i e_i^T)⁻¹
            //       = X_c^T E[Z|X]^T * (n Σ_z + E[Z|X] E[Z|X]^T)⁻¹

            // X_c^T E[Z|X]^T: xc^T is p×n, ez^T is n×k → result is p×k
            let xc_ez_t = xc.t().dot(&ez.t()); // p × k

            // ezz_t_sum is k × k
            let ezz_t_inv =
                cholesky_inv(&ezz_t_sum).map_err(|_| FerroError::NumericalInstability {
                    message: "FactorAnalysis: E[ZZ^T] is singular in M-step".into(),
                })?;

            let w_new = xc_ez_t.dot(&ezz_t_inv); // p × k

            // ψ_new[d] = (1/n) Σ_i (x_id² - w_new[d,:] e_i x_id)
            //          = (1/n) [Σ_i x_id² - w_new[d,:] Σ_i e_i x_id^T]
            //          = S[d,d] - (w_new[d,:] @ (1/n) Σ_i e_i x_id^T)
            // (1/n) Σ_i e_i x_id = (1/n) ez[:,i] * x_i[d] = (1/n) ez @ x_c[:,d]
            // = (1/n) ez @ xc[:, d]

            let mut psi_new = Array1::<F>::zeros(p);
            for d in 0..p {
                // Sample variance of feature d.
                let var_d = xc
                    .column(d)
                    .iter()
                    .copied()
                    .map(|v| v * v)
                    .fold(F::zero(), |a, b| a + b)
                    / n_f;
                // w_new[d,:] @ (1/n) ez @ xc[:,d]
                // (1/n) ez @ xc[:,d] is (1/n) Σ_i ez[:,i] * xc[i,d] — k-vector
                let mut ez_xd = Array1::<F>::zeros(k);
                for kk in 0..k {
                    let s = (0..n_samples)
                        .map(|i| ez[[kk, i]] * xc[[i, d]])
                        .fold(F::zero(), |a, b| a + b);
                    ez_xd[kk] = s / n_f;
                }
                let wd = w_new.row(d);
                let corr = wd
                    .iter()
                    .zip(ez_xd.iter())
                    .map(|(&wi, &ei)| wi * ei)
                    .fold(F::zero(), |a, b| a + b);
                let psi_d = var_d - corr;
                psi_new[d] = if psi_d > F::from(1e-6).unwrap() {
                    psi_d
                } else {
                    F::from(1e-6).unwrap()
                };
            }

            w = w_new;
            psi = psi_new;

            // --- Convergence check ------------------------------------------
            let ll = compute_log_likelihood(&xc, &w, &psi);
            let ll_change = (ll - prev_ll).abs();
            n_iter = iter + 1;
            if ll_change < tol_f && iter > 0 {
                prev_ll = ll;
                break;
            }
            prev_ll = ll;
        }

        Ok(FittedFactorAnalysis {
            components: w,
            noise_variance: psi,
            mean,
            n_iter,
            log_likelihood: prev_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(),
            });
        }
        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().iter() {
            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");
    }
}