Model API¶

def _check_convergence(
    self,
    primal_res: float,
    dual_res: float,
    v: np.ndarray,
    x_hat: np.ndarray,
    lam: np.ndarray,
) -> bool:
    """Check ADMM convergence using primal and dual residual norms.

    Parameters
    ----------
    primal_res : float
        Norm of the primal residual ||V - WW'||_F
    dual_res : float
        Norm of the dual residual rho * ||V - V_old||_F
    v : ndarray of shape (n, n)
        Current auxiliary variable
    x_hat : ndarray of shape (n, n)
        Current estimate WW'
    lam : ndarray of shape (n, n)
        Current Lagrange multipliers

    Returns
    -------
    converged : bool
        True if both primal and dual tolerances are satisfied
    """
    eps_abs = self.tol
    eps_rel = 1e-4
    n = v.size

    eps_pri = np.sqrt(n) * eps_abs + eps_rel * max(
        np.linalg.norm(v), np.linalg.norm(x_hat)
    )
    eps_dual = np.sqrt(n) * eps_abs + eps_rel * np.linalg.norm(lam)
    return primal_res <= eps_pri and dual_res <= eps_dual

def _compute_metrics(
    self,
    x: np.ndarray,
    v: np.ndarray,
    x_hat: np.ndarray,
    lam: np.ndarray,
    primal_residual: np.ndarray,
    v_old: np.ndarray | None = None,
) -> dict[str, float]:
    """Compute ADMM optimization metrics for monitoring and convergence.

    Parameters
    ----------
    x : ndarray of shape (n, n)
        Original similarity matrix
    v : ndarray of shape (n, n)
        Auxiliary variable
    x_hat : ndarray of shape (n, n)
        Current estimate WW'
    lam : ndarray of shape (n, n)
        Lagrange multipliers
    primal_residual : ndarray of shape (n, n)
        V - WW'
    v_old : ndarray of shape (n, n) or None
        Previous V, used to compute the dual residual

    Returns
    -------
    metrics : dict[str, float]
        Keys: total_objective, data_fit, penalty, lagrangian,
        rec_error, evar, primal_residual, dual_residual
    """
    data_residual = x - v

    observed_mask = self._observation_mask
    data_fit = np.sum(data_residual[observed_mask] ** 2)

    penalty_term = np.sum(primal_residual**2)
    penalty = (self.rho / 2.0) * penalty_term

    lagrangian = np.sum(lam * primal_residual)
    total_obj = data_fit + penalty + lagrangian

    rec_residual_obs = (data_residual + primal_residual)[observed_mask]
    rec_ss = np.sum(rec_residual_obs**2)
    rec_error = np.sqrt(rec_ss)

    if self._total_var > 0:
        evar = 1.0 - rec_ss / self._total_var
    else:
        evar = 0.0

    primal_res = np.sqrt(penalty_term)
    dual_res = self.rho * np.linalg.norm(v - v_old) if v_old is not None else np.inf

    return {
        "total_objective": total_obj,
        "data_fit": data_fit,
        "penalty": penalty,
        "lagrangian": lagrangian,
        "rec_error": rec_error,
        "evar": evar,
        "primal_residual": primal_res,
        "dual_residual": dual_res,
    }

def _fit_complete_data(self, x: np.ndarray) -> SRF:
    """Fit model with complete data (no missing values).

    Uses _frobenius_residual to compute ||X - WW'||_F without forming WW',
    so memory stays at O(n^2).
    """
    w = _initialize_w(x, self.rank, self.init, self.random_state)
    history = defaultdict(list)

    n = x.shape[0]
    x_norm = np.linalg.norm(x, "fro")
    total_var = np.var(x)
    eps_rel = 1e-4

    pbar = trange(
        1,
        self.max_outer + 1,
        disable=not self.verbose,
        desc="SRF",
        mininterval=10.0 if not sys.stderr.isatty() else 0.1,
    )
    for i in pbar:
        w = self._solver(x, w, max_iter=self.max_inner, tol=self.tol)

        residual_norm, xhat_norm = _frobenius_residual(x, w)

        rec_error = residual_norm
        mse = residual_norm**2 / (n * n)
        evar = 1.0 - mse / total_var if total_var > 0 else 0.0

        history["rec_error"].append(rec_error)
        history["evar"].append(evar)
        history["primal_residual"].append(residual_norm)

        pbar.set_postfix(
            rec_error=f"{rec_error:.3f}",
            evar=f"{evar:.3f}",
        )

        eps_pri = n * self.tol + eps_rel * max(x_norm, xhat_norm)
        if residual_norm <= eps_pri:
            logger.info("Converged at iteration %d", i)
            break

    self._store_results(w, i, history)

    return self

def _fit_missing_data(self, x: np.ndarray) -> SRF:
    """Fit model with missing data using ADMM."""
    bound_min, bound_max = self.bounds
    history = defaultdict(list)

    w = _initialize_w(x, self.rank, self.init, self.random_state)
    lam = np.zeros_like(x)

    v = x.copy()
    x_hat = w @ w.T

    v_old = np.empty_like(x)
    target = np.empty_like(x)
    primal_residual = np.empty_like(x)

    pbar = trange(
        1,
        self.max_outer + 1,
        disable=not self.verbose,
        desc="SRF",
        mininterval=10.0 if not sys.stderr.isatty() else 0.1,
    )
    for i in pbar:
        v_old[:] = v

        np.divide(lam, self.rho, out=target)
        np.add(v, target, out=target)
        w = self._solver(target, w, max_iter=self.max_inner, tol=self.tol)
        x_hat[:] = w @ w.T

        update_v_(
            self._observation_mask, x, x_hat, lam, self.rho, bound_min, bound_max, v
        )
        np.subtract(v, x_hat, out=primal_residual)
        lam += self.rho * primal_residual

        metrics = self._compute_metrics(x, v, x_hat, lam, primal_residual, v_old)
        for key, value in metrics.items():
            history[key].append(value)

        pbar.set_postfix(
            obj=f"{metrics['total_objective']:.3f}",
            rec_error=f"{metrics['rec_error']:.3f}",
            evar=f"{metrics['evar']:.3f}",
        )

        if i > 1 and self._check_convergence(
            metrics["primal_residual"],
            metrics["dual_residual"],
            v,
            x_hat,
            lam,
        ):
            logger.info("Converged at iteration %d", i)
            break

    self._store_results(w, i, history)

    return self

def _store_results(self, w: np.ndarray, n_iter: int, history: dict) -> None:
    """Set fitted attributes after optimization."""
    self.w_ = w
    self.components_ = w
    self.n_iter_ = n_iter
    self.history_ = history

def fit(self, x: np.ndarray, y: np.ndarray | None = None) -> SRF:
    """
    Fit the symmetric NMF model to the data.

    Parameters
    ----------
    x : array-like of shape (n_samples, n_samples)
        Symmetric similarity matrix. Missing values are allowed and should
        be marked according to the missing_values parameter.
    y : Ignored
        Not used, present here for API consistency by convention.

    Returns
    -------
    self : object
        Fitted estimator.
    """
    self._validate_params()
    _validate_bounds(self.bounds)

    x = validate_data(
        self,
        x,
        reset=True,
        ensure_all_finite=(
            "allow-nan" if _is_nan_marker(self.missing_values) else True
        ),
        ensure_2d=True,
        dtype=np.float64,
        copy=True,
    )

    self._missing_mask = _get_missing_mask(x, self.missing_values)
    if np.all(self._missing_mask):
        raise ValueError(
            "No observed entries found in the data. All values are missing."
        )

    check_symmetric(self._missing_mask, raise_exception=True)
    self._observation_mask = ~self._missing_mask
    x[self._missing_mask] = 0.0
    x = check_symmetric(x, raise_exception=True, tol=1e-10)

    observed = self._observation_mask
    observed_count = int(np.sum(observed))
    if observed_count > 0:
        observed_mean = np.sum(x[observed]) / observed_count
        self._total_var = np.sum((x[observed] - observed_mean) ** 2)
    else:
        self._total_var = 0.0

    if np.all(self._observation_mask):
        return self._fit_complete_data(x)
    else:
        return self._fit_missing_data(x)

def fit_transform(self, x: np.ndarray, y: np.ndarray | None = None) -> np.ndarray:
    """
    Fit the model and return the learned factors.

    Parameters
    ----------
    x : array-like of shape (n_samples, n_samples)
        Symmetric similarity matrix
    y : Ignored
        Not used, present here for API consistency by convention.

    Returns
    -------
    w : array-like of shape (n_samples, rank)
        Learned factor matrix
    """
    return self.fit(x, y).transform(x)

def reconstruct(self, w: np.ndarray | None = None) -> np.ndarray:
    """
    Reconstruct the similarity matrix from factors.

    Parameters
    ----------
    w : array-like of shape (n_samples, rank) or None
        Factor matrix to use for reconstruction.
        If None, uses the fitted factors.

    Returns
    -------
    s_hat : array-like of shape (n_samples, n_samples)
        Reconstructed similarity matrix
    """
    if w is None:
        check_is_fitted(self)
        w = self.w_

    return w @ w.T

def score(self, x: np.ndarray, y: np.ndarray | None = None) -> float:
    """
    Score the model using reconstruction error on observed entries only.

    Parameters
    ----------
    x : array-like of shape (n_samples, n_samples)
        Symmetric similarity matrix. Missing values are allowed and should
        be marked according to the missing_values parameter.
    y : Ignored
        Not used, present here for API consistency by convention.

    Returns
    -------
    mse : float
        Mean squared error of the reconstruction on observed entries.
    """
    check_is_fitted(self)

    x = validate_data(
        self,
        x,
        reset=False,
        ensure_2d=True,
        dtype=np.float64,
        ensure_all_finite="allow-nan" if self.missing_values is np.nan else True,
    )
    observation_mask = ~_get_missing_mask(x, self.missing_values)
    reconstruction = self.reconstruct()
    mse = np.mean((x[observation_mask] - reconstruction[observation_mask]) ** 2)
    return -mse

def transform(self, x: np.ndarray) -> np.ndarray:
    """
    Project data onto the learned factor space.

    Parameters
    ----------
    x : array-like of shape (n_samples, n_samples)
        Symmetric matrix to transform

    Returns
    -------
    w : array-like of shape (n_samples, rank)
        Transformed data
    """
    check_is_fitted(self)
    return self.w_