Extend experiments to use set-valued predictors using the SAPS scores

Author: Eleni Straitouri

config.py

@@ -10,6 +10,7 @@
 N_RUNS = 50
 DEBUG = False
 N_LABELS = 16
+SAPS_WEIGHTS = [0.02, 0.05, 0.1, 0.15, 0.2, 0.25, 0.3, 0.35]
 
 parser = argparse.ArgumentParser()
 parser.add_argument("--model_name", type=str, \
@@ -35,6 +36,10 @@
                     type=bool, help="Set to True to run experiments\
                     with a mixture of MNLs using the ImageNet16H-PS data",
                     default=False)
+parser.add_argument("--score_type", choices=["vanilla", "SAPS"], \
+                    default="vanilla", type=str,
+                    help="Choose the label score function type (non-conformity score)\
+                    to select the labels in the prediction set.")
 
 args,unknown = parser.parse_known_args()
 
@@ -62,12 +67,27 @@
 # Initialize random generators
 entropy = 0x3034c61a9ae04ff8cb62ab8ec2c4b501
 numpy_rng = np.random.default_rng(entropy)
+tune_rng = np.random.default_rng(entropy)
+saps_rng = np.random.default_rng(entropy)
 
 # Fraction of the dataset to use as calibration set
 calibration_split = 0.1
 
+# Fraction of the dataset to use as a tuning set for SAPS
+tuning_split = 0.1
+
+# Name of score function (non-conformity score)
+score_type = args.score_type
+
+# Range of lambda values according to score type
+lambda_min = 0. 
+lambda_max = 1. if score_type == 'vanilla' else (1. + SAPS_WEIGHTS[-1]*(N_LABELS-1))
+
 # The granularity for the lambda grid
-lambda_step = 0.001
+lambda_step = (lambda_max - lambda_min) / 1000
+
+# Decimals to round up spurious lambda values due to np.arange
+lamda_dec = 3 if score_type == 'vanilla' else 5
 
 # Set up paths for results and plots
 # Path to store results

counterfactual_harm.py

@@ -2,23 +2,38 @@
 import numpy as np
 import os
 import pandas as pd
-
+import json
+from utils import saps, saps_batch, create_path
+from copy import deepcopy
 """Main module to control and evaluate counterfactual harm (bound)"""
 
 class CounterfactualHarm:
-    def __init__(self, model, human, data) -> None:
+    def __init__(self, model, human, data, score_type='vanilla') -> None:
         self.model = model
         self.human = human
         self.data = data
         self.emp_risks = {}
+        self.score_type = score_type
         self.H = self.fn_per_set("h")
         self.G = self.fn_per_set("g")
         self.alphas = np.arange(0,1+config.alpha_step, config.alpha_step)
 
     def _empirical_risk(self, lamda):
         """Computes hat H(\lambda)"""
         # Set is valid or corner case of empty set
-        is_valid_set = (self.H_data_set['model_score_true_label'] >= (1 - lamda)) | (self.H_data_set['model_score_max'] == self.H_data_set['model_score_true_label'])
+        if self.score_type == 'vanilla':
+            is_valid_set = (self.H_data_set['model_score_true_label'] >= (1 - lamda)) | (self.H_data_set['model_score_max'] == self.H_data_set['model_score_true_label'])
+        else:
+            lamda_k = str(np.around(lamda, decimals=config.lamda_dec))
+            assert lamda_k in self.weights, self.weights.keys()
+            is_valid_set = (saps(
+                self.weights[lamda_k], 
+                self.H_data_set['true_label_rank'],
+                self.H_data_set['model_score_max'],
+                config.saps_rng
+                ) <= lamda) | \
+            (self.H_data_set['true_label_rank'] == 1)
+            
         if config.mode == 'control':
             emp_risk = self.H_data_set.where(is_valid_set, self.H_data_set['h_invalid_sets'], axis=0)['h_valid_sets'].mean()
             return emp_risk
@@ -32,7 +47,18 @@ def _empirical_risk(self, lamda):
     def _empirical_benefit(self, lamda):
         """Computes hat G(\lambda)"""
         # Set is valid or corner case of empty set
-        is_valid_set = (self.G_data_set['model_score_true_label'] >= (1 - lamda)) | (self.H_data_set['model_score_max'] == self.H_data_set['model_score_true_label'])
+        if self.score_type == 'vanilla':
+            is_valid_set = (self.G_data_set['model_score_true_label'] >= (1 - lamda)) | (self.H_data_set['model_score_max'] == self.H_data_set['model_score_true_label'])
+        else:
+            lamda_k = str(np.around(lamda, decimals=config.lamda_dec))
+            is_valid_set = (saps(
+                self.weights[lamda_k], 
+                self.G_data_set['true_label_rank'],
+                self.G_data_set['model_score_max'],
+                config.saps_rng
+                ) <= lamda) | \
+            (self.G_data_set['true_label_rank'] == 1)
+           
 
         if config.mode == 'control':
             emp_ben = self.G_data_set.where(is_valid_set, self.G_data_set['g_invalid_sets'], axis=0)['g_valid_sets'].mean()
@@ -46,39 +72,38 @@ def _empirical_benefit(self, lamda):
 
     def fn_per_set(self, fn_name):
         """Reads/computes the h/g function for each prediction set"""
-        data_path = f"{config.ROOT_DIR}/data/{fn_name}/noise{config.noise_level}"
+        data_path = f"{config.ROOT_DIR}/data/{fn_name}_{self.score_type}/noise{config.noise_level}"
         file_path = f"{data_path}/{config.model_name}{'_pred_set' if config.HUMAN_DATASET=='PS' else ''}.csv"
         if not os.path.exists(file_path):
             if not os.path.exists(data_path):
                 os.makedirs(data_path)
             fn_per_set = []
+            
             # Compute the h/g value given each human prediction 
             for image_name, participant_id, human_correct in self.human.itertuples(index=True):
                 true_label = self.data.loc[image_name]["category"]
                 model_score_true_label = self.model.loc[image_name][true_label]
-                
+                label_ranks = config.N_LABELS - self.model.drop(columns=['correct']).loc[image_name].argsort().argsort()
+                true_label_rank = label_ranks[true_label]
                 model_score_max = self.model.drop(columns=['correct']).loc[image_name].max()
 
-                if fn_name == "h":
+                if "h" in fn_name:
                     fn_value_valid = 0
                     fn_value_invalid = human_correct
                 else:
                     fn_value_valid = 1 - human_correct
                     fn_value_invalid = 0
 
-                fn_per_set.append((image_name, participant_id, model_score_true_label, model_score_max, fn_value_valid, fn_value_invalid))
+                fn_per_set.append((image_name, participant_id, model_score_true_label, model_score_max, true_label_rank, fn_value_valid, fn_value_invalid))
 
-            columns = ["image_name", "participant_id", "model_score_true_label", "model_score_max", f"{fn_name}_valid_sets", f"{fn_name}_invalid_sets"]
-            
-            fn_df = pd.DataFrame(fn_per_set, columns=columns).set_index('image_name')
-            
+            columns = ["image_name", "participant_id", "model_score_true_label", "model_score_max", "true_label_rank",f"{fn_name}_valid_sets", f"{fn_name}_invalid_sets"]            
+            fn_df = pd.DataFrame(fn_per_set, columns=columns).set_index('image_name')  
             fn_df.to_csv(file_path)
         else:
             fn_df = pd.read_csv(file_path, index_col='image_name')
 
         return fn_df
 
-        
     def set_data_set(self, data_set):
         self.data_set = data_set
         self.data_set_size = len(data_set)
@@ -90,34 +115,86 @@ def control(self):
         """Min control level per lambda for h and g"""
         n = self.data_set_size
         thresholds = (((n+1)*self.alphas - 1)/n)
-        lamdas_dict = { np.around(lamda, decimals=3):{} for lamda in np.arange(0,1+config.lambda_step,config.lambda_step)}
+        lamdas_dict = {np.around(lamda, decimals=config.lamda_dec):{} for lamda in np.arange(config.lambda_min,config.lambda_max+config.lambda_step,config.lambda_step)}
 
-        for lamda in np.arange(0,1+config.lambda_step, config.lambda_step):
+        for lamda in np.arange(config.lambda_min,config.lambda_max+config.lambda_step, config.lambda_step):
             emp_risk_lamda = self._empirical_risk(lamda)
 
             # Min alpha such that lambda is harm controlling under CF (Counterfactual) monotonicity
             min_alpha_idx_CF = np.searchsorted(thresholds, emp_risk_lamda, side='left')
 
             # For each lambda keep the min level of control under CF monotonicity
-            lamdas_dict[np.around(lamda, decimals=3)]['CF'] = np.round(self.alphas[min_alpha_idx_CF], decimals=2)
+            lamdas_dict[np.around(lamda, decimals=config.lamda_dec)]['CF'] = np.round(self.alphas[min_alpha_idx_CF], decimals=2)
 
             # Empirical benefit (\hat G)
-            emp_benefit_lamda = self._empirical_benefit(1.-lamda)
+            emp_benefit_lamda = self._empirical_benefit(config.lambda_max-lamda)
 
             # Select smallest alpha that for which lambda is g controlling under cI (Interventional) monotonicity
             min_alpha_idx_cI = np.searchsorted(thresholds, emp_benefit_lamda, side='left')
-            lamdas_dict[np.around(1 - lamda, decimals=3)]['cI'] = np.round(self.alphas[min_alpha_idx_cI], decimals=2)
+            lamdas_dict[np.around(config.lambda_max - lamda, decimals=config.lamda_dec)]['cI'] = np.round(self.alphas[min_alpha_idx_cI], decimals=2)
 
         return lamdas_dict
 
     def compute(self):
         """Evaluate the counterfactual harm (bound)"""
-        lamdas_dict = { lamda:{} for lamda in np.arange(0,1+config.lambda_step,config.lambda_step)}
+        lamdas_dict = {np.around(lamda, decimals=config.lamda_dec):{} for lamda in np.arange(config.lambda_min,config.lambda_max+config.lambda_step,config.lambda_step)}
 
-        for lamda in np.arange(0,1+config.lambda_step, config.lambda_step):
+        for lamda in lamdas_dict.keys():
             harm_sum, harm_count = self._empirical_risk(lamda)
             g_harm, g_count = self._empirical_benefit(lamda)
             lamdas_dict[lamda]['hat_H'] = (harm_sum, harm_count)
             lamdas_dict[lamda]['hat_G'] = (g_harm, g_count)
 
-        return lamdas_dict
+        return lamdas_dict
+    
+    def tune_saps(self, x_y, run):
+        """Find the optimal saps weight parameter value for each lamda value in a given run"""
+        self.weights = {}
+        weight_per_run_path = f"{config.results_path}/{config.model_name}/noise{config.noise_level}/saps/saps_weights_tune{config.tuning_split}_run{run}.json" 
+        # Read the weights if already computed 
+        if os.path.exists(weight_per_run_path):
+            with open(weight_per_run_path, 'rt') as f:
+                self.weights = json.load(f)
+        else:
+            self.weight_per_run = {}
+            lamdas = np.array([lamda for lamda in np.arange(config.lambda_min,config.lambda_max+config.lambda_step,config.lambda_step)])
+            min_avg_size_per_lambda = np.ones_like(lamdas)*config.N_LABELS
+            self.weight_per_lamda = np.ones_like(lamdas)*config.SAPS_WEIGHTS[0] 
+            tune_set_probs = self.model.drop(columns=['correct']).loc[x_y.index]
+            tune_set_label_ranks = deepcopy(tune_set_probs)
+            for idx in x_y.index:
+                tune_set_label_ranks.loc[idx] = tune_set_probs.loc[idx].argsort()
+            max_scores = tune_set_probs.max(axis=1)
+            
+            for weight in config.SAPS_WEIGHTS:
+                # Get the saps scores in the tune set
+                saps_scores_per_label = saps_batch(
+                    weight,
+                    tune_set_probs.loc[x_y.index.to_numpy()],
+                    tune_set_label_ranks,
+                    max_scores
+                )
+                # Compute the average prediction set size given a weight value
+                set_sizes_per_lambda = []
+                for idx in saps_scores_per_label.index:
+                    set_sizes_per_lambda.append(np.searchsorted(
+                        saps_scores_per_label.loc[idx],
+                        lamdas,
+                        side='right',
+                        sorter=saps_scores_per_label.loc[idx].argsort()
+                    ))
+                set_sizes_per_lambda = np.stack(set_sizes_per_lambda, ).T
+                set_sizes_per_lambda[set_sizes_per_lambda == 0] = 1
+                avg_set_sizes_per_lambda = set_sizes_per_lambda.mean(axis=1)
+                # Keep track of the weight value with the minimum average set size 
+                self.weight_per_lamda[
+                    avg_set_sizes_per_lambda < 
+                    min_avg_size_per_lambda] = weight
+                min_avg_size_per_lambda = np.minimum(avg_set_sizes_per_lambda, min_avg_size_per_lambda)
+            # Save the optimal weight for each lambda 
+            for i,lamda in enumerate(lamdas):
+                lam_key = str(np.around(lamda, decimals=config.lamda_dec))
+                self.weights[lam_key] = self.weight_per_lamda[i]
+            create_path(f"{config.results_path}/{config.model_name}/noise{config.noise_level}/saps/")
+            with open(weight_per_run_path, 'w') as f:
+                json.dump(self.weights,f)