stress_at_work_analysis/machine_learning/model.py

from sklearn.model_selection import LeaveOneGroupOut, cross_val_score


class ModelValidation:
    def __init__(self, X, y, group_variable=None, cv_name="loso"):
        self.model = None
        self.cv = None

        idx_common = X.index.intersection(y.index)
        self.y = y.loc[idx_common, "NA"]
        # TODO Handle the case of multiple labels.
        self.X = X.loc[idx_common]
        self.groups = self.y.index.get_level_values(group_variable)

        self.cv_name = cv_name
        print("ModelValidation initialized.")

    def set_cv_method(self):
        if self.cv_name == "loso":
            self.cv = LeaveOneGroupOut()
            self.cv.get_n_splits(X=self.X, y=self.y, groups=self.groups)
        print("Validation method set.")

    def cross_validate(self):
        print("Running cross validation ...")
        if self.model is None:
            raise TypeError(
                "Please, specify a machine learning model first, by setting the .model attribute. "
                "E.g. self.model = sklearn.linear_model.LinearRegression()"
            )
        if self.cv is None:
            raise TypeError(
                "Please, specify a cross validation method first, by using set_cv_method() first."
            )
        if self.X.isna().any().any() or self.y.isna().any().any():
            raise ValueError(
                "NaNs were found in either X or y. Please, check your data before continuing."
            )
        return cross_val_score(
            estimator=self.model,
            X=self.X,
            y=self.y,
            groups=self.groups,
            cv=self.cv,
            n_jobs=-1,
            scoring="r2",
        )