DeepExplainerDataset.py

import pandas as pd
import numpy as np 
import torch

class DeepExplainerDataset:
    def __init__(self, train_X, train_y, transform_func, 
                 test_X=None, test_y=None,
                 feature_names=None, outcome_name=None, 
                 categorical_variables=None):
        """
        Create an ExplainerDataset instance for the training and testing dataset. The class will check the dataset, identify the categorical variables, 
        and create mappings for categorical variables.
    
        Parameters
        ----------
        train_X : np.ndarray or pandas.dataframe
            Training dataset contains features only.
        train_y : np.ndarray or pandas.dataframe
            Label of training dataset.
        transform_func: function
            Transform the train_X into tensor format input for deep learning model. 
        test_X : np.ndarray or pandas.dataframe
            Test dataset contains features only.
        test_y : np.ndarray or pandas.dataframe
            Label of test dataset.
        feature_names : list or np.ndarray
            Feature names.
        outcome_name : str
            Label name.
        categorical_variables : list
            List of categorical variables. It can contains all string-like feature names, or all interger-like feature variable indices.
    
        Returns
        -------
        ExplainerDataset
            ExplainerDataset class instance.
    
        """                
        
        self.train_x = train_X 
        self.train_y = train_y
        self.test_x = test_X
        self.test_y = test_y 
        self.feature_names = feature_names 
        self.outcome_name = outcome_name
        self.transform_func = transform_func
        
        if isinstance(self.train_x, pd.DataFrame) or isinstance(self.train_x, np.ndarray):
            if isinstance(self.train_x, pd.DataFrame):
                if self.feature_names and isinstance(self.feature_names, list):
                    if len(self.feature_names) == self.train_x.shape[1]: 
                        self.train_x.columns = self.feature_names
                    else:
                        print('Warning: length of feature names is not consistent with the dataset shap. Ignore the input feature_names.')
                        self.feature_names = self.train_x.columns
                else:
                    self.feature_names = self.train_x.columns 
            elif isinstance(self.train_x, np.ndarray):
                if self.feature_names and isinstance(self.feature_names, list):
                    if len(self.feature_names) == self.train_x.shape[1]: 
                        self.train_x = pd.DataFrame(self.train_x, columns=self.feature_names)
                    else:
                        print('Warning: length of feature names is not consistent with the dataset shap. Ignore the input feature_names. Use the default feature names.')
                        self.feature_names = [f'Feature_{i}' for i in range(self.train_x.shape[1])]
                        self.train_x = pd.DataFrame(self.train_x, columns=self.feature_names) 
                else:
                    print('Warning: feature_names is not provided. Use the default feature names.')
                    self.feature_names = [f'Feature_{i}' for i in range(self.train_x.shape[1])]
                    self.train_x = pd.DataFrame(self.train_x, columns=self.feature_names)                
        else:
            print('Error: The input train_X should be numpy.array or pandas.DataFrame.')
            return

        if isinstance(self.test_x, pd.DataFrame) or isinstance(self.test_x, np.ndarray):
            if isinstance(self.test_x, np.ndarray): 
                self.test_x = pd.DataFrame(self.test_x)
            assert self.train_x.shape[1] == self.test_x.shape[1], "Error: train_x and test_x must have the same number of features."
            self.test_x.columns = self.feature_names   
        else:
            print('Warning: The input test_X should be numpy.array or pandas.DataFrame. Ignore test_X.')
            self.test_x = None                  
        
        if isinstance(self.train_y, pd.DataFrame) or isinstance(self.train_y, np.ndarray):
            if isinstance(self.train_y, np.ndarray):
                self.train_y = pd.DataFrame(self.train_y)
                assert self.train_y.shape[1] == 1, "Error: train_y must have one outcome."    
                
                if self.outcome_name:
                    if (type(self.outcome_name) is str) and len(self.outcome_name) > 0: 
                        self.train_y.columns = [self.outcome_name]
                    else:
                        print('Warning: support only one outcome, and the outcome name should be a string. Ignore outcome_name.')
                        self.outcome_name = self.train_y.columns[0]                    
                else:
                    print('Warning: Outcome name is not provided. Use the default outcome name.')
                    self.train_y.columns=['Outcome']   
            elif isinstance(self.train_y, pd.DataFrame): 
                assert self.train_y.shape[1] == 1, "Error: train_y must have one outcome."            
                if self.outcome_name:
                    if (type(self.outcome_name) is str) and len(self.outcome_name) > 0: 
                        self.train_y.columns = [self.outcome_name]
                    else:
                        print('Warning: support only one outcome, and the outcome name should be a string. Ignore outcome_name.')
                        self.outcome_name = self.train_y.columns[0]                    
                else:
                    self.outcome_name = self.train_y.columns[0]
        else:
            print('Error: The input train_y should be numpy.array or pandas.DataFrame')
            return        

        if isinstance(self.test_y, pd.DataFrame) or isinstance(self.test_y, np.ndarray):
            if isinstance(self.test_y, np.ndarray):
                self.test_y = pd.DataFrame(self.test_y)
            assert self.test_y.shape[1] == 1, "Error: test_y must have one outcome."
            self.test_y.columns = [self.outcome_name]
        else:
            print('Warning: The input test_y should be numpy.array or pandas.DataFrame. Ignore test_y.')
            self.test_y = None

        ###check if the train_X, train_y, test_X, test_y are complete
        assert self.train_x.isnull().values.any() == False, "Error: train_X has missing values."
        assert self.train_y.isnull().values.any() == False, "Error: train_y has missing values."
        assert self.train_x.shape[0] == self.train_y.shape[0], "Error: number of samples of train_X and train_y is not consistent."
        series = self.train_y[self.outcome_name]
        assert pd.api.types.is_integer_dtype(series) and set(series.dropna().unique()).issubset({0, 1}), "Error: outcome column in train_y should be a binary integer column."
        if test_X is not None and test_y is not None:
            assert self.test_x.isnull().values.any() == False, "Error: test_X has missing values."
            assert self.test_y.isnull().values.any() == False, "Error: test_y has missing values."
            assert self.test_x.shape[0] == self.test_y.shape[0], "Error: number of samples of test_X and test_y is not consistent."
            series = self.test_y[self.outcome_name]
            assert pd.api.types.is_integer_dtype(series) and set(series.dropna().unique()).issubset({0, 1}), "Error: outcome column in test_y should be a binary integer column."
        else:
            self.test_x = None 
            self.test_y = None

        self.categorical_variables = []
        self.categorical_variable_indices = []
        self.categorical_variable_level_map = {}
        self.features = {}
        self.meta_data = {}
        self.process_categorical_variables(categorical_variables)
        self.generate_metadata()
        self.transformer_input_type = None
        self.get_transformer_input_type()
        assert isinstance(self.transformer_input_type, list) == True, "Error: The provided transformer function does not support input data in format of numpy.array or pandas.DataFrame"
    
    def process_categorical_variables(self, categorical_variables):
        """
        Process the categorical features, called by init() function. The process will check the categorical variables, identify their names and feature column indices, 
        and create level mappings.
    
        Parameters
        ----------
        categorical_variables : list
            List of categorical variables. It can contains all string-like feature names, or all interger-like feature variable indices.
    
        Returns
        -------
    
        """ 
        
        
        if isinstance(categorical_variables, list) and len(categorical_variables) > 0:
            all_str = all(isinstance(item, str) for item in categorical_variables)
            all_int = all(isinstance(item, int) for item in categorical_variables)
            assert all_str or all_int, "Input categorical variables should contain column names or the indices of column names"
            
            if all_str:
                self.categorical_variables = categorical_variables 
                self.categorical_variable_indices = [] 
                for var in self.categorical_variables:
                    if var not in self.feature_names:
                        print(f'Error: Input categorical variable {var} does not exit in the feature list.')
                        return 
                    for i, feature in enumerate(self.feature_names):
                        if var == feature:
                            self.categorical_variable_indices.append(i)
                            break 
            else:
                self.categorical_variable_indices = categorical_variables 
            
            ###check all columns, if it is string type but not in categorical variable list
            for index, feature in enumerate(self.feature_names):
                if pd.api.types.is_string_dtype(self.train_x[feature]):
                    if index not in self.categorical_variable_indices:
                        print(f'Warning: The data type of column {feature} is string, but it is not in the categorical variable list. Add it to the list of categorical variables.')
                        self.categorical_variable_indices.append(index)  
                elif pd.api.types.is_bool_dtype(self.train_x[feature]):
                    if index not in self.categorical_variable_indices:
                        print(f'Warning: The data type of column {feature} is boolean, but it is not in the categorical variable list. Add it to the list of categorical variables.')
                        self.categorical_variable_indices.append(index)  
            self.categorical_variable_indices = sorted(set(self.categorical_variable_indices))
            self.categorical_variables = [self.feature_names[i] for i in self.categorical_variable_indices] 
            
            for i in self.categorical_variable_indices:
                feature = self.feature_names[i]
                unique_levels = self.train_x[feature].unique().tolist()
                self.categorical_variable_level_map[i] = np.asarray(unique_levels)
                
    def generate_metadata(self):  
        """
        Summarize the characteristics of the dataset including data type, value ranges for continuous variables, and levels for categorical variables. It is called by init() function.
    
        Parameters
        ----------
    
        Returns
        -------
    
        """ 
        numeric_columns = self.train_x.select_dtypes(include=['number']).columns.tolist()
        for index, feature in enumerate(self.feature_names):
            if index in self.categorical_variable_indices:
                self.features[feature] = self.train_x[feature].unique().tolist()
                self.meta_data[feature] = {'data_type' : 'categorical', 'levels': self.features[feature]}
            elif feature in numeric_columns:
                min_value = self.train_x[feature].min()
                max_value = self.train_x[feature].max()
                self.features[feature] = [min_value, max_value]
                self.meta_data[feature] = {'data_type' : 'continuous', 'min':min_value, 'max':max_value}
            else:
                assert feature in self.features, f"Error: Feature {feature} is not a continuous or categorical variable. Consider to remove the feature or add it to the categorical variable list."
        
    def get_metadata(self, print_summary=True):
        """
        Return and print (optional) the summary of the dataset.
    
        Parameters
        ----------
        print_summary : bool
            Whether to print out the summary.
    
        Returns
        -------
        dict
            Dictionary contains the summary of the dataset.
        """         
        
        if print_summary:
            print(self.meta_data)
        return self.meta_data
    
    def get_data(self):
        """
        Get the datasets, return as dataframes.
    
        Parameters
        ----------
    
        Returns
        -------
        pandas.dataframe
            Four dataframes for train_x, train_y, test_x, test_y.
        """ 
        
        return self.train_x, self.train_y, self.test_x, self.test_y
    
    def get_transformer_input_type(self):
        ###test if the input type is pandas dataframe
        input_types = []
        try: 
            _ = self.transform_func(self.train_x)
            input_types.append('DataFrame')
        except Exception:
            pass 
        
        try:
            _ = self.transform_func(self.train_x.values)
            input_types.append('Array')
        except Exception:
            pass 
        
        if len(input_types) > 0:
            self.transformer_input_type = input_types 
            
    def get_tensor_data(self, y_unit=1):
        """
        Get the datasets, return as tensors.
        We will use tansformer_func to convert the train_X, test_X into tensors.
        We will directly make train_y, test_y into tensors.
        
        Parameters
        ----------
        y_unit: int. 
             The number of outcome columns, can be 1 or 2. For binary classification, if y_unit is 1, the column is class 1; if y_unit is 2, one column is class 0, 
             another column is class 1. Default value is 1.
        Returns
        -------
        torch.tensor
            Four tensors for train_x, train_y, test_x, test_y.
        """
        if y_unit == 2:
            train_y_temp = self.train_y.copy()
            train_y_temp = pd.get_dummies(train_y_temp, columns=[self.outcome_name], drop_first=False)
            train_y_tensor = torch.tensor(train_y_temp.values, dtype=torch.float32)
        else:    
            train_y_tensor = torch.tensor(self.train_y.values, dtype=torch.float32)
        
        train_x_tensor = None
        if 'DataFrame' in self.transformer_input_type:
            train_x_tensor = self.transform_func(self.train_x)
        elif 'Array' in self.transformer_input_type:
            train_x_tensor = self.transform_func(self.train_x.values)
        
        test_x_tensor = None 
        test_y_tensor = None 
        if self.test_x is not None and self.test_y is not None:
            if y_unit == 2:
                test_y_temp = self.test_y.copy()
                test_y_temp = pd.get_dummies(test_y_temp, columns=[self.outcome_name], drop_first=False)
                test_y_tensor = torch.tensor(test_y_temp.values, dtype=torch.float32)
            else:    
                test_y_tensor = torch.tensor(self.test_y.values, dtype=torch.float32)
            
            if 'DataFrame' in self.transformer_input_type:
                test_x_tensor = self.transform_func(self.test_x)
            elif 'Array' in self.transformer_input_type:
                test_x_tensor = self.transform_func(self.test_x.values)
        
        return train_x_tensor, train_y_tensor, test_x_tensor, test_y_tensor
    
    def transform_data(self, input_data):
        if isinstance(input_data, pd.DataFrame) or isinstance(input_data, np.ndarray):
            if isinstance(input_data, pd.DataFrame):
                if 'DataFrame' in self.transformer_input_type:
                    transform_data = self.transform_func(input_data)
                    return transform_data
                elif 'Array' in self.transformer_input_type:
                    transform_data = self.transform_func(input_data.values)
                    return transform_data                    
            elif isinstance(input_data, np.ndarray):
                if 'Array' in self.transformer_input_type:
                    transform_data = self.transform_func(input_data)
                    return transform_data
                elif 'DataFrame' in self.transformer_input_type:
                    data = pd.DataFrame(input_data, columns=self.feature_names)
                    transform_data = self.transform_func(data)
                    return transform_data             
        else:
            data_type = type(input_data)
            print(f'The input data type is {data_type}, not numpy.array or pandas.DataFrame. Try ignoring data type and directly transforming.')
            try: 
                transform_data = self.transform_func(input_data)
                return transform_data
            except Exception:
                print(f'Error: The provided transformer function does not support data type: {data_type}.')