Model API¶

def _fit_loop(self, steps: Iterator[tuple[np.ndarray, Step]]) -> SRF:
    """Drive the outer iterations; steps yields one measurement each."""
    history = defaultdict(list)

    recent = deque(maxlen=PROGRESS_WINDOW)
    pbar = self._progress_bar()
    for i, (w, step) in zip(pbar, steps):
        earlier = recent[0] if recent else None
        metrics = step.metrics(self._total_ss)
        metrics["converged"] = step.converged(earlier, self.tol)
        for key, value in metrics.items():
            history[key].append(value)

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

        if metrics["converged"]:
            logger.info("Converged at iteration %d", i)
            break
        recent.append(step)
    else:
        warnings.warn(
            f"Maximum number of outer iterations ({self.max_outer}) "
            "reached without convergence. Increase max_outer or tol.",
            ConvergenceWarning,
        )

    self.w_ = w
    self.components_ = w
    self.n_iter_ = i
    self.history_ = history

    return self

def _validate_input(self, x: np.ndarray) -> np.ndarray:
    """Validate the similarity matrix and derive the observation stats.

    Allocates only the one contract copy of x and the observation
    mask; all statistics are computed without materializing the
    observed entries.
    """
    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._observed_mask = observation_mask(x, self.missing_values)
    if not np.any(self._observed_mask):
        raise ValueError(
            "No observed entries found in the data. All values are missing."
        )

    if is_nan_marker(self.missing_values) or self.missing_values is None:
        np.nan_to_num(x, copy=False, nan=0.0, posinf=0.0, neginf=0.0)
    else:
        x[~self._observed_mask] = 0.0
    if self.check_input:
        check_symmetric(self._observed_mask, raise_exception=True)
        x = check_symmetric(x, raise_exception=True, tol=1e-10)

    # Zero-filled missing entries drop out of the sums, so the
    # variance identity avoids materializing the observed values
    n_observed = np.count_nonzero(self._observed_mask)
    observed_mean = np.sum(x) / n_observed
    observed_sumsq = np.einsum("ij,ij->", x, x)
    self._total_ss = observed_sumsq - n_observed * observed_mean**2

    if self.verbose:
        n_threads = os.environ.get("OMP_NUM_THREADS", "1")
        logger.info(
            "Using %s BLAS thread(s). For faster fitting, set "
            "OMP_NUM_THREADS before importing pysrf "
            "(e.g. export OMP_NUM_THREADS=4)",
            n_threads,
        )

    return x

def fit(self, x: np.ndarray, y: np.ndarray | None = None) -> SRF:
    """Fit the model to a similarity matrix.

    Parameters
    ----------
    x : array-like of shape (n_samples, n_samples)
        Symmetric similarity matrix. Missing values should be marked
        according to the ``missing_values`` parameter.
    y : Ignored
        Not used, present for sklearn API compatibility.

    Returns
    -------
    self : SRF
        Fitted estimator.
    """
    x = self._validate_input(x)
    w = _initialize_w(x, self.rank, self.init, self.random_state)

    if self._observed_mask.all():
        steps = _complete_steps(x, w, self.max_inner)
    else:
        steps = _missing_steps(
            x, self._observed_mask, w, self.rho, self.bounds, self.max_inner
        )
    return self._fit_loop(steps)

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

    Parameters
    ----------
    x : array-like of shape (n_samples, n_samples)
        Symmetric similarity matrix
    y : Ignored
        Not used, present for sklearn API compatibility.

    Returns
    -------
    w : ndarray of shape (n_samples, rank)
        Embedding matrix
    """
    return self.fit(x, y).transform(x)

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

    Parameters
    ----------
    w : ndarray of shape (n_samples, rank) or None
        Embedding matrix. If None, uses the fitted embedding.

    Returns
    -------
    s_hat : ndarray 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:
    """Negative MSE on observed entries (higher is better).

    Parameters
    ----------
    x : array-like of shape (n_samples, n_samples)
        Symmetric similarity matrix. Missing values should be marked
        according to the ``missing_values`` parameter.
    y : Ignored
        Not used, present for sklearn API compatibility.

    Returns
    -------
    score : float
        Negative mean squared reconstruction error 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 is_nan_marker(self.missing_values)
        else True,
    )
    observed_mask = observation_mask(x, self.missing_values)
    s_hat = self.reconstruct()
    mse = np.mean((x[observed_mask] - s_hat[observed_mask]) ** 2)
    return -mse

def transform(self, x: np.ndarray) -> np.ndarray:
    """Return the learned embedding matrix W.

    Parameters
    ----------
    x : array-like of shape (n_samples, n_samples)
        Ignored, present for sklearn API compatibility.

    Returns
    -------
    w : ndarray of shape (n_samples, rank)
        Embedding matrix
    """
    check_is_fitted(self)
    return self.w_