Covariance regularization¶

def regularize_covariance(
    covariance: npt.ArrayLike,
    /,
    *,
    method: CovarianceRegularizer = "diagonal_loading",
    batched: bool | None = None,
    copy: bool = True,
) -> Covariance | Covariances:
    """Regularize a covariance matrix or a batch of covariance matrices.

    The covariances are validated and regularized according to the specified
    method. The input can be either a single covariance matrix with shape ``(d,
    d)` or a batch of covariance matrices with shape `(n, d, d)``. If the
    batch flag is not explicitly provided, the shape of the input is used to infer whether it is
    batched or not.

    Parameters
    ----------
    covariance : array-like
        A covariance matrix with shape `(d, d)` or a batch with shape
        `(n, d, d)`.
    method : str or covariance regularizer configuration, default="diagonal_loading"
        Regularization method to apply. Passing a string uses the method's
        default configuration. `DiagonalLoading` and `LowRank` also accept
        epsilon heuristic strategy objects through their `eps` field.
    batched : bool or None, default=None
        Whether to interpret the input as batched. If `None`, the shape is
        inferred from the array rank.
    copy : bool, default=True
        Whether to regularize a copy of the input.

    Returns
    -------
    Covariance or batch of covariances
        Regularized covariance matrix or matrices.
    """
    spec, options = _REGISTRY.resolve(method)

    match spec.name:
        case "diagonal_loading":
            options = cast_options(options, DiagonalLoading)
            return diagonal_loading(covariance, eps=options.eps, batched=batched, copy=copy)
        case "linear_shrinkage":
            options = cast_options(options, LinearShrinkage)
            return linear_shrinkage(covariance, alpha=options.alpha, batched=batched, copy=copy)
        case "diagonal_shrinkage":
            options = cast_options(options, DiagonalShrinkage)
            return diagonal_shrinkage(covariance, alpha=options.alpha, batched=batched, copy=copy)
        case "eigenvalue_clipping":
            options = cast_options(options, EigenvalueClipping)
            return eigenvalue_clipping(
                covariance, min_eigenvalue=options.min_eigenvalue, batched=batched, copy=copy
            )
        case "lowrank":
            options = cast_options(options, LowRank)
            return lowrank(
                covariance, rank=options.rank, eps=options.eps, batched=batched, copy=copy
            )
        case _:
            msg = "Unhandled covariance regularizer registry entry."
            raise AssertionError(msg)

def estimate_epsilon(
    covariance: npt.ArrayLike,
    /,
    *,
    method: EpsilonMethod = "relative_trace",
    batched: bool | None = None,
) -> float | FloatArray:
    """Estimate a diagonal-loading epsilon from covariance scale or spectrum.

    Parameters
    ----------
    covariance : array-like
        Covariance matrix with shape `(d, d)` or batch of matrices with shape
        `(n, d, d)`.
    method : str or epsilon heuristic configuration, default="relative_trace"
        Heuristic used to estimate the epsilon value.
    batched : bool or None, default=None
        Whether to interpret the input as batched. If `None`, the shape is
        inferred from the input rank.

    Returns
    -------
    float or array
        Estimated epsilon value(s). If the input is a single covariance, a float
        is returned. If the input is a batch of covariances, an array of shape
        `(n,)` is returned with one epsilon per covariance.
    """
    covariance_arr: FloatArray = np.asarray(covariance, dtype=np.float64)
    shape_kind = check_covariance_shape(covariance_arr, batched=batched)
    spec, options = _EPSILON_REGISTRY.resolve(method)

    match spec.name:
        case "relative_trace":
            options = cast_options(options, RelativeToTrace)
            return _relative_trace(covariance_arr, c=options.c, batched=shape_kind)
        case "target_condition_number":
            options = cast_options(options, TargetConditionNumber)
            return _target_condition_number(covariance_arr, kappa=options.kappa, batched=shape_kind)
        case "residual_variance":
            options = cast_options(options, ResidualVariance)
            return _residual_variance(
                covariance_arr, c=options.c, rank=options.r, batched=shape_kind
            )
        case _:
            msg = "Unhandled covariance epsilon heuristic registry entry."
            raise AssertionError(msg)