plot_I_with_no_context.py

import importlib
import os

import hydra
from hydra.utils import get_original_cwd
from omegaconf import DictConfig
import numpy as np
from xgboost import XGBClassifier, XGBRegressor
from tqdm import tqdm
import matplotlib
import matplotlib.pyplot as plt

from utils import int_expectation, freq


@hydra.main(config_path='config', config_name='config')
def run_main(cfg: DictConfig) -> None:
    fontsize = cfg.plotting.importance_fontsize
    matplotlib.rcParams.update({
        "font.size": fontsize,          # base font size
        "axes.labelsize": fontsize,     # x and y labels
        "xtick.labelsize": fontsize,    # x tick labels
        "ytick.labelsize": fontsize,    # y tick labels
        "legend.fontsize": fontsize,   # legend
    })
    if cfg.data.problem_type == 'classification':
        classification = True
    else:
        classification = False
    if cfg.data.univariate_kwargs.feature_type == 'categorical':
        continuous = False
    else:
        continuous = True  
    results_path = os.path.join(get_original_cwd(), cfg.results_folder, cfg.data.results_subfolder)
    os.makedirs(results_path, exist_ok=True)
    data_module = importlib.import_module(cfg.data.data_module)
    dataset_loader = getattr(data_module, cfg.data.dataset_loader)
    n_train = cfg.data.univariate_kwargs.n_samples
    n_repetitions = cfg.data.univariate_kwargs.repetitions
    scm = dataset_loader()
    evaluations = {}
    features = cfg.data.univariate_kwargs.features
    for i, feature in enumerate(features):
        evaluations[feature] = []
        for rep in tqdm(range(n_repetitions), desc=f'Univariate importance for {feature}'):
            # sample
            train_sample = scm.marginal([feature, cfg.data.fit_kwargs.target], n_train).T

            # train model
            if classification:
                univariate_model = XGBClassifier()
            else:
                univariate_model = XGBRegressor()
            univariate_model.fit(train_sample[:,0:1], train_sample[:,-1])

            # fix range by first feature and first repetition
            if i==0 and rep==0:
                if continuous:
                    plotting_range = np.linspace(
                        np.quantile(train_sample[:,0], 0.05, axis=0),
                        np.quantile(train_sample[:,0], 0.95, axis=0),
                        cfg.data.plot_kwargs.resolution,
                    )[:,None] # add feature dimension
                else:
                    plotting_range = np.unique(
                        train_sample[:,0]
                    )[:,None] # add feature dimension

            # evaluate model
            if classification:
                # use specified class - by default last class
                red_ind = cfg.data.plot_kwargs.get('multivariate_reduction',-1)
                freq_in_data = freq(train_sample[:, -1])
                if red_ind is not None:
                    f_empty = freq_in_data[red_ind].item()
                    f_feature = univariate_model.predict_proba(
                        plotting_range
                    )[:,red_ind]
                else:
                    f_empty = int_expectation(
                        freq_in_data, axis=-1
                    ).item()
                    f_feature = int_expectation(
                        univariate_model.predict_proba(plotting_range),
                        axis=-1,
                        )
            else:
                # multivariate regression case not implemented currently
                assert train_sample.shape[1] == 1,\
                    "Multivariate regression not implemented currently"
                f_empty = train_sample[:, -1].mean().item()
                f_feature = univariate_model.predict(train_sample[:,0:1]).squeeze()
            evaluations[feature].append(f_feature- f_empty)
        # convert list of arrays to array of arrays
        evaluations[feature] = np.stack(evaluations[feature], axis=0)
    
    # plot
    plt.figure(1)
    # add horizontal line at y=0
    plt.axhline(0, color='k', linestyle='--', linewidth=1, alpha=1.0)
    # add univariate features importances
    for i, feature in enumerate(features):
        # mean feature importance across repetitions    
        mean_evaluation = evaluations[feature].mean(axis=0)
        # standard error of the mean
        std_evaluation = evaluations[feature].std(axis=0)  / np.sqrt(n_repetitions)
        # plot as curve
        plt.plot(plotting_range.squeeze(),
                mean_evaluation,  
                linestyle='-',
                linewidth=cfg.data.univariate_kwargs.linewidth,
                color=cfg.data.univariate_kwargs.colors[i],
                label=feature,
        )
        plt.fill_between(plotting_range.squeeze(),
            mean_evaluation-std_evaluation,
            mean_evaluation+std_evaluation,
            color=cfg.data.univariate_kwargs.colors[i],
            alpha=cfg.data.univariate_kwargs.alpha,
        )
        if not continuous:
            plt.plot(plotting_range,
                     mean_evaluation,
                     marker="o",
                     linestyle='None',
                     markersize=cfg.data.univariate_kwargs.markersize,
                     color=cfg.data.univariate_kwargs.colors[i]
            )
            plt.xticks(plotting_range.squeeze())
    if len(features) > 1:
        # plt.legend()
        plt.legend(loc='center left', bbox_to_anchor=(1, 0.5))
        plt.xlabel('feature value')
        plt.ylabel(r'$I_{\emptyset}$')
    else:
        plt.xlabel(feature)
        plt.ylabel(r'$I_{\emptyset}$'+f'({feature})')

    feature_string_list = features[0]
    for feature in features[1:]:
        feature_string_list += f'_{feature}'
    plt.tight_layout()
    plt.savefig(
        os.path.join(
            results_path,
            f'univariate_importance_{cfg.data.name}_'\
                f'{cfg.data.fit_kwargs.target}_features_'\
                    f'{feature_string_list}.png'),
    dpi=300)
    
if __name__ == "__main__":
    run_main()