Cross-Validation API¶

def cross_val_score(
    similarity_matrix: np.ndarray,
    estimator: BaseEstimator | None = None,
    param_grid: dict[str, list] | None = None,
    n_repeats: int = 5,
    sampling_fraction: float = 0.8,
    estimate_sampling_fraction: bool | dict = False,
    sampling_selection: str = "mean",
    random_state: int = 0,
    verbose: int = 1,
    n_jobs: int = -1,
    missing_values: float | None = np.nan,
    fit_final_estimator: bool = False,
) -> GridSearchCV:
    """
    Cross-validate any estimator for matrix completion.

    Generic cross-validation function that works with SRF or any sklearn-compatible
    estimator with a .reconstruct() method.

    Parameters
    ----------
    similarity_matrix : ndarray
        Symmetric similarity matrix to cross-validate
    estimator : BaseEstimator or None, default=None
        Estimator to cross-validate. If None, uses SRF(random_state=random_state).
        Can be a single estimator or a Pipeline. Must have a .reconstruct() method.
    param_grid : dict or None, default=None
        Dictionary with parameter names (str) as keys and lists of values to try
        as values. If None, uses default {'rank': [5, 10, 15, 20]} for SRF.
    n_repeats : int, default=5
        Number of times to repeat the cross-validation
    sampling_fraction : float, default=0.8
        Fraction of eligible entries to use for training in each split; must be in (0, 1).
        The remaining (1 - sampling_fraction) becomes validation.
        Note: Constant diagonal entries are excluded from both train and validation.
        Ignored when estimate_sampling_fraction is True or a dict; if both are provided,
        estimate_sampling_fraction takes precedence.
    estimate_sampling_fraction : bool or dict, default=False
        If True, automatically estimate optimal sampling fraction using sampling
        bound estimation from Random Matrix Theory. If dict, passed as kwargs to
        estimate_sampling_bounds_fast(). When enabled, overrides sampling_fraction.
    sampling_selection : str, default="mean"
        Selection method for the estimated sampling fraction; one of {"mean", "min", "max"}.
    random_state : int, default=0
        Random seed for reproducibility
    verbose : int, default=1
        Verbosity level
    n_jobs : int, default=-1
        Number of jobs to run in parallel (-1 uses all processors)
    missing_values : float or None, default=np.nan
        Value to consider as missing in original data
    fit_final_estimator : bool, default=False
        Whether to fit the final estimator on the best parameters

    Returns
    -------
    grid : GridSearchCV
        Fitted GridSearchCV object with best parameters and scores

    Examples
    --------
    >>> from pysrf.cross_validation import cross_val_score
    >>> result = cross_val_score(similarity_matrix, param_grid={'rank': [5, 10, 15]})
    """
    if estimator is None:
        estimator = SRF(random_state=random_state)

    if param_grid is None:
        param_grid = {"rank": [5, 10, 15, 20]}

    valid_selections = {"mean", "min", "max"}
    if sampling_selection not in valid_selections:
        raise ValueError(
            f"sampling_selection must be one of {sorted(valid_selections)}"
        )

    if estimate_sampling_fraction:
        from .bounds import estimate_sampling_bounds_fast

        kwargs = (
            estimate_sampling_fraction
            if isinstance(estimate_sampling_fraction, dict)
            else {}
        )
        if "random_state" not in kwargs:
            kwargs["random_state"] = random_state
        if "n_jobs" not in kwargs:
            kwargs["n_jobs"] = n_jobs
        kwargs.pop("verbose", None)

        pmin, pmax, s_noise = estimate_sampling_bounds_fast(similarity_matrix, **kwargs)
        sampling_fraction = {
            "mean": np.mean([pmin, pmax]),
            "min": pmin,
            "max": pmax,
        }[sampling_selection]

    else:
        _validate_sampling_fraction(sampling_fraction)

    cv = EntryMaskSplit(
        n_repeats=n_repeats,
        sampling_fraction=sampling_fraction,
        random_state=random_state,
        missing_values=missing_values,
    )
    grid = GridSearchCV(
        estimator=estimator,
        param_grid=param_grid,
        cv=cv,
        n_jobs=n_jobs,
        verbose=verbose,
        fit_final_estimator=fit_final_estimator,
    )
    grid.fit(similarity_matrix)

    return grid

def fit_and_score(
    estimator: BaseEstimator,
    x: np.ndarray,
    train_mask: np.ndarray,
    validation_mask: np.ndarray,
    fit_params: dict,
    split_idx: int | None = None,
) -> dict:
    """
    Fit estimator with parameters and return validation score.

    Parameters
    ----------
    estimator : BaseEstimator
        Model instance to fit (works with SRF or any estimator with .reconstruct())
    x : ndarray
        Full data matrix
    train_mask : ndarray of bool
        Boolean mask where True = training entry
    validation_mask : ndarray of bool
        Boolean mask where True = validation entry
    fit_params : dict
        Parameters to set on the estimator
    split_idx : int or None
        Index of the CV split

    Returns
    -------
    result : dict
        Dictionary with score, parameters, and fitted estimator
    """
    est = clone(estimator).set_params(**fit_params)

    # Set SRF-specific params if estimator supports them
    if hasattr(est, "missing_values"):
        est.set_params(missing_values=np.nan)

    if hasattr(est, "bounds"):
        if "bounds" not in fit_params or fit_params["bounds"] is None:
            original_bounds = (np.nanmin(x), np.nanmax(x))
            est.set_params(bounds=original_bounds)

    # Track which entries were already NaN in the original data
    originally_nan = np.isnan(x)

    # Create training data: keep only training entries, mask everything else
    x_train = np.full_like(x, np.nan)
    x_train[train_mask] = x[train_mask]

    # Fit model on training data only
    est.fit(x_train)

    # Get reconstruction
    if hasattr(est, "reconstruct"):
        reconstruction = est.reconstruct()
    else:
        raise ValueError(
            f"Estimator {type(est).__name__} must have a .reconstruct() method "
            "for matrix completion cross-validation"
        )

    # Evaluate only on validation entries that were originally observed
    valid_eval_mask = validation_mask & ~originally_nan

    if not valid_eval_mask.any():
        raise ValueError("No valid validation entries to evaluate")

    mse = np.mean((x[valid_eval_mask] - reconstruction[valid_eval_mask]) ** 2)

    result = {
        "score": mse,
        "split": split_idx if split_idx is not None else 0,
        "estimator": est,
        "params": fit_params,
    }

    # Include history if available (optional)
    if hasattr(est, "history_"):
        result["history"] = est.history_

    return result