Source code for njab.sklearn.types

from dataclasses import dataclass
from collections import namedtuple

import pandas as pd
import sklearn

AucRocCurve = namedtuple("AucRocCurve", 'fpr tpr cutoffs')
PrecisionRecallCurve = namedtuple("PrecisionRecallCurve",
                                  'precision recall cutoffs')




@dataclass
class ResultsSplit:
    """Results for a single split of train and test data."""
    auc: float = None  # receiver operation curve area under the curve
    aps: float = None  # average precision score
    roc: AucRocCurve = None
    prc: PrecisionRecallCurve = None





@dataclass
class Results:
    """Results for a model trained on a single split of train and test data."""
    model: sklearn.base.BaseEstimator = None
    selected_features: list = None
    train: ResultsSplit = None
    test: ResultsSplit = None
    name: str = None





@dataclass
class Splits:
    """Splits of train and test data with targets."""
    X_train: pd.DataFrame
    X_test: pd.DataFrame
    y_train: pd.Series
    y_test: pd.Series