run.py

import numpy as np
import pandas as pd
import random
import time
import shap
from tqdm import tqdm
import matplotlib.pyplot as plt

from sklearn.model_selection import train_test_split
from sklearn.metrics import matthews_corrcoef
from sklearn.ensemble import RandomForestClassifier

import mlflow
from mlflow.models import infer_signature

import warnings
warnings.filterwarnings("ignore", module="mlflow.types.utils")

def get_columns_filter(column_csv: str, col: str, val) -> list[str]:
    """
    Filters a CSV file for rows where a specified column matches a given value and 
    returns a list of feature names.

    Args:
        column_csv (str): Path to the CSV file containing feature information.
        col (str): The column name to filter on.
        val: The value to match in the specified column.

    Returns:
        List[str]: A list of feature names ('name' column) where the specified 
        column equals the given value.
    """
    prostate_cols = pd.read_csv(column_csv)
    prostate_cols_filtered = prostate_cols[prostate_cols[col] == val]
    return prostate_cols_filtered['name'].to_list()

class Runner(object):
    """
    Runner is a utility class for managing the end-to-end process of training, 
    evaluating, and tracking machine learning models on tabular data, with a focus
    on privacy and explainability metrics.
    Attributes:
        data_link (str): Path to the CSV data file.
        params (dict): Parameters for the model and experiment.
        model (object): Scikit-learn compatible model instance.
        filter_columns (list): Columns to filter out for wide feature set.
        filter_narrow_columns (list): Columns to filter out for narrow feature set.
        categorical_columns (list): Columns to treat as categorical.
        meaning_columns (list): Columns with clinical meaning.
        random_state (int): Random seed for reproducibility.
        data_type (str): Type of data input ('pandas' or 'numpy').
        columns (list): Columns used in the narrow feature set.
        random_col_test (bool): flag to indicate whether to perform test of adding
        column with random values.
    Methods:
        get_data():
            Loads and preprocesses the data, splits into train/test sets, and 
            returns feature matrices and targets.
        train_track(data):
            Trains the model, evaluates performance, computes privacy/explainability
            metrics, and logs results to MLflow.
        PBI_acc(data, sensitive_attribute):
            Computes Privacy Breach Index (PBI) based on accuracy for a given 
            sensitive attribute.
        PBI_MCC_scale(baseline_MCC, data, subset_columns, sensitive_attribute):
            Computes scaled PBI using Matthews Correlation Coefficient (MCC).
        PBI_MCC_add(baseline_MCC, data, subset_columns, sensitive_attribute):
            Computes additive PBI using MCC.
        frac_rw_meaning():
            Calculates the fraction of features with real-world (clinical) meaning.
        expected_feature_importance(data):
            Estimates expected feature importance using KL divergence over 
            perturbations.
        weighted_KL(ref, dist, weights, eps=1e-14):
            Computes weighted KL divergence between reference and distribution 
            probabilities.
        get_importance_df(data, expected_imp):
            Constructs a DataFrame of feature importances using SHAP and expected 
            importance.
        monotonicity(importance_df):
            Computes Spearman correlation between expected and actual feature 
            importances.
        non_sensitivity(importance_df):
            Measures non-sensitivity as the symmetric difference between zero-importance
            features.
    """


    def __init__(
            self, 
            data_link: str, 
            params: dict, 
            model: object, 
            filter_columns_csv: str, 
            random_state: int = 42, 
            data_type: str = 'pandas',
            random_col_test: bool = False,
        ):
        """
        Initializes the runner class with data source, parameters, model, and column
        filters.

        Args:
            data_link (str): Path or link to the dataset.
            params (dict): Dictionary of parameters for model or processing.
            model (object): Machine learning model instance.
            filter_columns_csv (str): Path to CSV file specifying column filters.
            random_state (int, optional): Seed for random number generators. 
            Defaults to 42.
            data_type (str, optional): Type of data structure to use ('pandas' by 
            default).
            random_col_test (bool, optional): Whether to include a column with random
            noise to test validity.
        """
        self.data_link = data_link
        self.params = params
        self.model = model
        self.random_state = random_state
        self.data_type = data_type
        self.filter_columns = get_columns_filter(filter_columns_csv, "Keep", "no")
        self.filter_narrow_columns = get_columns_filter(filter_columns_csv, "Keep Narrow", "no")
        self.categorical_columns = get_columns_filter(filter_columns_csv, "Categorical", "yes")
        self.meaning_columns = get_columns_filter(filter_columns_csv, "Clinical Meaning", 1)
        self.random_col_test = random_col_test
        np.random.seed(random_state)
        random.seed(random_state)

    def get_data(self) -> tuple[pd.DataFrame]:
        """
        Loads and preprocesses the dataset for model training and evaluation.

        Reads the data from the specified CSV file, filters rows where 
        'fstcan_cancersite' equals 1 (prostate cancer), and prepares wide and narrow
        feature sets by dropping specified columns. Categorical columns are 
        converted to the 'category' dtype. The target variable `y` is defined as a 
        boolean indicating whether the difference between 'mortality_exitdays' and 
        'fstcan_exitdays' exceeds 4500.
        Splits the data into training and test sets for both wide and narrow feature
        sets, as well as the target variable.
        
        Returns:
            X_wide_train (pd.DataFrame): Training set with wide features.

            X_wide_test (pd.DataFrame): Test set with wide features.

            X_narrow_train (pd.DataFrame): Training set with narrow features.
            
            X_narrow_test (pd.DataFrame): Test set with narrow features.
            
            y_train (pd.Series): Training target variable.
            
            y_test (pd.Series): Test target variable.
        """
        df_raw = pd.read_csv(self.data_link, sep=',')
        # Filter for prostate cancer only
        df = df_raw[df_raw['fstcan_cancersite'] == 1].copy()

        if self.random_col_test:
            # Add a column with random values
            df['random_col'] = np.random.rand(len(df))

        # Filter on chosen columns
        X_wide = df.drop(columns=self.filter_columns)
        X_narrow = df.drop(columns=self.filter_narrow_columns)
        # Define target
        y = (df['mortality_exitdays'] - df['fstcan_exitdays']) > 4500
        
        # Change the type of the columns
        for col in self.categorical_columns:
            if col in X_wide.columns:
                X_wide[col] = X_wide[col].copy().astype('category')
            if col in X_narrow.columns:
                X_narrow[col] = X_narrow[col].copy().astype('category')

        self.columns = X_narrow.columns.copy()
        
        X_wide_train, X_wide_test, X_narrow_train, X_narrow_test, y_train, y_test = train_test_split(
            X_wide, 
            X_narrow, 
            y, 
            test_size=.2, 
            random_state=self.random_state
        )
        return X_wide_train, X_wide_test, X_narrow_train, X_narrow_test, y_train, y_test
    
    def train_track(self, data: tuple[pd.DataFrame]):
        """
        Trains the model on the provided dataset, evaluates performance, computes 
        metrics, and logs results to MLflow.
        Args:
            data (tuple): A tuple containing training and testing data splits in 
                the following order:
                (X_wide_train, X_wide_test, X_narrow_train, X_narrow_test, y_train, y_test).
        Performs:
            - Fits the model using the narrow training data.
            - Predicts outcomes for both training and testing sets.
            - Calculates training time, accuracy, overfitting, and Matthews 
              correlation coefficient (MMC) for train/test.
            - Computes privacy bias index (PBI) and fraction of features with 
              real-world meaning.
            - Logs all metrics, parameters, and input data to MLflow.
            - Saves the trained model to MLflow with input signature and example.
        Returns:
            None
        """
        X_wide_train, X_wide_test, X_narrow_train, X_narrow_test, y_train, y_test = data
        # Time training and making prediction by model
        start_t = time.time()
        self.model.fit(X_narrow_train, y_train)
        y_train_pred = self.model.predict(X_narrow_train)
        y_test_pred = self.model.predict(X_narrow_test)
        end_t = time.time()
        # Calculate time elapsed and accuracy metrics
        training_time = end_t - start_t
        train_acc = self.model.score(X_narrow_train, y_train)
        test_acc = self.model.score(X_narrow_test, y_test)
        overfitting = train_acc - test_acc

        # Calculate MCC metrics
        MMC_train = matthews_corrcoef(y_train, y_train_pred)
        MMC_test = matthews_corrcoef(y_test, y_test_pred)

        MCC_scaled = (MMC_test + 1)/2

        # Calculate privacy metric
        pbi = self.PBI_acc(data, sensitive_attribute=['race7'])
        
        # Calculate explainability metrics
        frac_meaning = self.frac_rw_meaning()

        # If experimenting comment these 5 lines out
        expected_imp = self.expected_feature_importance(data)
        importance_matrix = self.get_importance_df(data, expected_imp)
        monotonicity = self.monotonicity(importance_matrix)
        non_sensitivity = self.non_sensitivity(importance_matrix)
        explainability = (frac_meaning + monotonicity + (1 - non_sensitivity/len(self.columns))) / 3

        explainer = shap.TreeExplainer(self.model)
        shap_values = explainer(X_narrow_test)
        # Create a matplotlib figure and axis, and pass them to shap.plots.bar
        fig, ax = plt.subplots()
        plt.tight_layout()
        shap.plots.bar(shap_values, ax=ax, max_display=15)

        # Start tracking experiment
        mlflow.set_tracking_uri(uri="http://127.0.0.1:5000")
        mlflow.set_experiment("privacy_exp_prostate")

        if self.data_type == 'numpy':
            dataset = mlflow.data.from_numpy(X_narrow_train, name="PLCO")
        elif self.data_type == 'pandas':
            dataset = mlflow.data.from_pandas(X_narrow_train, name="PLCO")

        with mlflow.start_run():
            mlflow.log_input(dataset, context="training")
            mlflow.log_params(self.params)
            mlflow.log_param("model", self.model.__class__)
            mlflow.log_metric("train accuracy", train_acc)
            mlflow.log_metric("test accuracy", test_acc)
            mlflow.log_metric("training time", training_time)
            mlflow.log_metric("overfitting", overfitting)
            mlflow.log_metric("random_state", self.random_state)
            mlflow.log_metric("train MMC", MMC_train)
            mlflow.log_metric("test MMC", MMC_test)
            mlflow.log_metric("PBI", pbi)
            mlflow.log_metric('Fraction of Features with Real World Meaning', frac_meaning)
            # Comment out these 4 lines if experimenting
            mlflow.log_metric("Monotonicity", monotonicity)
            mlflow.log_metric("Non-sensitivity", non_sensitivity)
            mlflow.log_metric("Predictive performance", MCC_scaled)
            mlflow.log_metric("Explainability", explainability)


            fig.subplots_adjust(left=0.4, right=0.95, top=0.95, bottom=0.25)
            mlflow.log_figure(fig, "shap_vals.png")
            plt.close(fig)

            feature_importance = np.abs(shap_values.values).mean(axis=0)
            importance_df = pd.DataFrame({
                'feature': X_narrow_train.columns, 
                'importance': feature_importance
            })
            importance_df = importance_df.sort_values(by='importance', ascending=False)
            if self.random_col_test:
                print(importance_df[importance_df['feature']=='random_col'])
            mlflow.log_table(importance_df, "shap_vals.json")

            signature = infer_signature(X_narrow_train, self.model.predict(X_narrow_train))

            model_info = mlflow.sklearn.log_model(
                sk_model = self.model,
                artifact_path="model",
                signature=signature,
                input_example=X_narrow_train,
            )

    def PBI_acc(self, data: tuple[pd.DataFrame], sensitive_attribute: list[str]) -> float:
        """
        Calculates the Privacy-Breach-Increase (PBI) metric for a given dataset and 
        sensitive attribute.

        The PBI metric quantifies the change in accuracy when predicting the 
        sensitive attribute using a subset of features (narrow) compared to using a
        wider set of features, after removing the sensitive attribute itself from 
        the predictors.
        Args:
            data (tuple): A tuple containing the following elements:
                - X_wide_train (pd.DataFrame): Training features (wide set).
                - X_wide_test (pd.DataFrame): Test features (wide set).
                - X_narrow_train (pd.DataFrame): Training features (narrow subset).
                - X_narrow_test (pd.DataFrame): Test features (narrow subset).
                - y_train (pd.Series or np.ndarray): Training labels (not used in this method).
                - y_test (pd.Series or np.ndarray): Test labels (not used in this method).
            sensitive_attribute (list[str]): The name of the sensitive attribute to be predicted.
        Returns:
            float: The computed PBI value, representing the relative change in accuracy.
        """

        X_wide_train, X_wide_test, X_narrow_train, X_narrow_test, y_train, y_test = data
        
        # Remove the sensitive attribute from input, and make it the target.
        X_train_pbi = X_wide_train.drop(sensitive_attribute, axis=1)
        X_test_pbi = X_wide_test.drop(sensitive_attribute, axis=1)
        X_train_pbi_subset = X_narrow_train.drop(sensitive_attribute, axis=1)
        X_test_pbi_subset = X_narrow_test.drop(sensitive_attribute, axis=1)
        y_train_pbi = X_wide_train[sensitive_attribute[0]]
        y_test_pbi = X_wide_test[sensitive_attribute[0]]



        # Train two different models and compare their accuracy
        model = RandomForestClassifier(random_state=self.random_state)
        baseline_acc = model.fit(X_train_pbi,y_train_pbi).score(X_test_pbi,y_test_pbi)
        acc = model.fit(X_train_pbi_subset,y_train_pbi).score(X_test_pbi_subset,y_test_pbi)
        pbi = (acc / baseline_acc) - 1
        return pbi

    def frac_rw_meaning(self) -> float:
        """
        Calculates the fraction of columns used to train the model that have real
        world (clinincal) meaning.
        Returns:
            float: The ratio of the number of columns in `self.meaning_columns` that
                   exist in `self.columns`
                   to the total number of columns in `self.columns`.
        """

        sum = 0

        for col in self.meaning_columns:
            if col in self.columns:
                sum += 1
        return sum / len(self.columns)
                
    def expected_feature_importance(self, data: tuple[pd.DataFrame]) -> list[float]:
        """
        Computes the expected feature importance for each feature in the input data using a model's predictions.
        The method perturbs each feature in the training set, calculates the change in predicted probabilities,
        and aggregates the weighted Kullback-Leibler (KL) divergence to estimate feature importance.
        Args:
            data (tuple[pd.DataFrame]): A tuple containing the following elements:
                - X_wide_train (pd.DataFrame): Training data with wide features.
                - X_wide_test (pd.DataFrame): Test data with wide features.
                - X_narrow_train (pd.DataFrame): Training data with narrow features.
                - X_narrow_test (pd.DataFrame): Test data with narrow features.
                - y_train (pd.Series or pd.DataFrame): Training labels.
                - y_test (pd.Series or pd.DataFrame): Test labels.
        Returns:
            list[float]: A list of expected feature importance values, one for each feature in X_narrow_train.
        """
        X_wide_train, X_wide_test, X_narrow_train, X_narrow_test, y_train, y_test = data

        y_pred = self.model.predict(X_narrow_train)
        y_prob = self.model.predict_proba(X_narrow_train)

        expected_imp = []

        bar = tqdm(X_narrow_train.columns, desc="Processing column...")

        # For each column
        for col in bar:
            bar.set_description(f"Processing column {col}")
            total = 0

            # Check type of column and ensure the correct number of bins is set
            match X_narrow_test[col].dtype:
                case 'int64' | 'float64':
                    for bins in [10, 5, 3]:
                        try:
                            ints = pd.qcut(X_narrow_train[col], bins, retbins=True)[1]
                            break  # Exit the loop if successful
                        except ValueError as e:
                            pass
                    # Get the midpoints of the quantiles.
                    vals = np.array([(ints[i+1] + ints[i])/2 for i in range(len(ints)-1)])
                case 'category':
                    freqs = X_narrow_train[col].value_counts(normalize=True)
                    vals = freqs.index.values
            # For each row in the dataset predict likelihood of each class.
            for i in range(X_narrow_train.shape[0]):
                x_star = X_narrow_train.iloc[i]
                y_star = y_pred[i]

                y_prob_arr = np.repeat([y_prob[i]], len(vals), axis=0)

                # Get rows with the columns value perturbed.
                x_is = []
                for val in vals:
                    x_i = x_star.copy()
                    x_i[col] = val
                    x_is.append(x_i)

                x_is = pd.DataFrame(x_is)
                y_pert = self.model.predict_proba(x_is)

                # Assign weights to each possible value
                if X_narrow_test[col].dtype == 'category':
                    sample_weight = np.array([freqs[val] for val in vals])
                else:
                    sample_weight = np.array([1.0/len(vals)]*len(vals))

                # Find the overall prediction error
                total += self.weighted_KL(y_prob_arr, y_pert, sample_weight)
            expected_imp.append(total)
        return expected_imp

    def weighted_KL(self, ref: np.ndarray, dist: np.ndarray, weights: np.ndarray, eps: float=1e-14) -> float:
        """
        Computes the weighted Kullback-Leibler (KL) divergence between two distributions.
        Parameters
        ----------
        ref : np.ndarray
            Reference probability distribution(s), shape (n_samples, n_features).
        dist : np.ndarray
            Comparison probability distribution(s), shape (n_samples, n_features).
        weights : np.ndarray
            Weights for each sample, shape (n_samples,).
        eps : float, optional
            Small value to avoid numerical instability in log and division (default is 1e-14).
        Returns
        -------
        float
            The weighted sum of KL divergences across all samples.
        Notes
        -----
        The KL divergence is computed for each sample and then weighted by the provided weights.
        """
        ref_clip = np.clip(ref, eps, 1-eps)
        dist_clip = np.clip(dist, eps, 1-eps)
        kl = np.sum(ref_clip * np.log(ref_clip/dist_clip), axis=1)
        weighted = kl * weights
        return np.sum(weighted)
    
    def get_importance_df(self, data: tuple[pd.DataFrame], expected_imp: list) -> pd.DataFrame:
        """
        Collects expected and actual feature importance in a dataframe 
        Parameters
        ----------
        data : tuple
            A tuple containing training and test datasets in the following order:
            (X_wide_train, X_wide_test, X_narrow_train, X_narrow_test, y_train, y_test).
            X_narrow_train should be a pandas DataFrame of features used for SHAP analysis.
        expected_imp : array-like
            Expected importance values for each feature, to be included in the output DataFrame.
        Returns
        -------
        importance_df : pandas.DataFrame
            DataFrame containing columns:
                - 'feature': Feature names from X_narrow_train.
                - 'ex_importance': Expected importance values.
                - 'importance': Mean absolute SHAP value for each feature.
            The DataFrame is sorted by 'importance' in descending order.
        """
        X_wide_train, X_wide_test, X_narrow_train, X_narrow_test, y_train, y_test = data
        # Get explanations from shap
        explainer = shap.TreeExplainer(self.model)
        shap_values = explainer.shap_values(X_narrow_train)

        # Compile expected 
        feature_importance = np.abs(shap_values).mean(axis=0)
        importance_df = pd.DataFrame({
            'feature': X_narrow_train.columns, 
            'ex_importance': expected_imp, 
            'importance': feature_importance
        })
        importance_df = importance_df.sort_values(by='importance', ascending=False)
        return importance_df
    
    def monotonicity(self, importance_df: pd.DataFrame) -> float:
        """
        Calculates the Spearman rank correlation coefficient between the 'ex_importance' and 'importance' columns
        in the provided DataFrame to assess monotonicity.
        Args:
            importance_df (pd.DataFrame): DataFrame containing 'ex_importance' and 'importance' columns.
        Returns:
            float: Spearman correlation coefficient indicating the monotonic relationship between the two columns.
        """
        return importance_df['ex_importance'].corr(importance_df['importance'], method='spearman')
    
    def non_sensitivity(self, importance_df: pd.DataFrame) -> int:
        """
        Calculates the non-sensitivity metric between two sets of features based on their importance values.
        This method compares features with zero 'importance' and zero 'ex_importance' in the provided DataFrame.
        It computes the symmetric difference between these two sets, which represents features that are considered
        non-important by one metric but not the other. The non-sensitivity is defined as the number of such features.
        Args:
            importance_df (pd.DataFrame): DataFrame containing feature importance values with columns
                'feature', 'importance', and 'ex_importance'.
        Returns:
            int: The count of features that are non-important in one metric but not the other.
        """
        A0 = set(importance_df[importance_df['importance'] == 0]['feature'])
        X0 = set(importance_df[importance_df['ex_importance'] == 0]['feature'])

        symmetric_diff = A0.symmetric_difference(X0)
        non_sensitivity = len(symmetric_diff)

        return non_sensitivity