stratified_cv¶

def stratified_cv_modeling(
    y: pd.DataFrame,
    X: pd.DataFrame,
    classes: np.ndarray,
    estimator: BaseEstimator = LassoCV(),
    skf: StratifiedKFold = StratifiedKFold(n_splits=4, shuffle=True, random_state=42),
    sample_weight: np.ndarray | None = None,
    **kwargs,
) -> BaseEstimator:
    """
    Fit a model using stratified cross-validation splits.

    This function wraps a scikit-learn estimator with user-defined stratified folds.
    While it defaults to `LassoCV`, any estimator with a `cv` attribute can be used.

    :param y: Response variable. Must be a single-column DataFrame.
    :param X: Predictor matrix. Must be a DataFrame with the same number of rows as `y`.
    :param classes: Array of class labels for stratification, typically generated by
        `stratification_classification()`.
    :param estimator: scikit-learn estimator to use for modeling. Must support `cv` as
        an attribute.
    :param skf: StratifiedKFold object to control how splits are generated.
    :param sample_weight: Optional array of per-sample weights for the estimator.
    :param kwargs: Additional arguments passed to the estimator's `fit()` method.

    :return: A fitted estimator with the best parameters determined via
        cross-validation.

    :raises ValueError: If inputs are misformatted or incompatible with the estimator.

    """
    # Validate data
    if not isinstance(y, pd.DataFrame):
        raise ValueError("The response variable y must be a DataFrame.")
    if y.shape[1] != 1:
        raise ValueError("The response variable y must be a single column DataFrame.")
    if not isinstance(X, pd.DataFrame):
        raise ValueError("The predictors X must be a DataFrame.")
    if X.shape[0] != y.shape[0]:
        raise ValueError("The number of rows in X must match the number of rows in y.")
    if classes.size == 0 or not isinstance(classes, np.ndarray):
        raise ValueError("The classes must be a non-empty numpy array.")

    # Verify estimator has a `cv` attribute
    if not hasattr(estimator, "cv"):
        raise ValueError("The estimator must support a `cv` parameter.")

    # Initialize StratifiedKFold for stratified splits
    with warnings.catch_warnings(record=True) as w:
        warnings.simplefilter("always")
        # default setting for shuffle is False, which means the partitioning is
        # deterministic and static. Recommendation for bootstrapping is to
        # set shuffle=True and use random_state = bootstrap_iteration in order to
        # have random, but reproducible, partitions
        folds = list(skf.split(X, classes))
        for warning in w:
            logger.debug(
                f"Warning encountered during stratified k-fold split: {warning.message}"
            )

    # Clone the estimator and set the `cv` attribute with predefined folds
    model = clone(estimator)
    model.cv = folds

    # Step 7: Fit the model using the custom cross-validation folds
    model.fit(
        X,
        y.values.ravel(),
        sample_weight=sample_weight,
    )

    return model