run_bandit.py

import numpy as np

import os
import pickle as pkl
import argparse
from tqdm import tqdm
import time
import json
import sys
#sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))

from MWBM.bandits import UCB1
from UserStudyAnalysis.utils import human_matching_to_index, get_matching, sample_valid_pool


parser = argparse.ArgumentParser(description='Run Bandit System with Prolific Data')
parser.add_argument('--pools-folder', type=str, required=True, help='Folder with the pools')
parser.add_argument('--results-folder', type=str, required=True, help='Folder to save data')
parser.add_argument('--problems', type=str, help='PROBLEMS considered AS VALID for Score Matrix of all Problem', default='pool_assignment')
parser.add_argument('--num_sims', type=int, help='Number of simulations')
parser.add_argument('--horizon', type=int, help='Number of iterations', default=None)
parser.add_argument('--random-assign', action='store_true', help='Assign people to random places if not all are assigned', default=False)
parser.add_argument('--min-arm', type=int, help='Minimum arm')
parser.add_argument('--allowed-users', type=str, help='Allowed users for the bandit algorithm', default='all')
parser.add_argument('--title', type=str, help='Title for the results file')
args = parser.parse_args()

POOLS_FOLDER = args.pools_folder
RESULTS_FOLDER = args.results_folder
NUM_SIMS = args.num_sims
HORIZON = args.horizon
RANDOM_ASSIGN = args.random_assign
problem_assignment_file = args.problems
min_arm = args.min_arm
allowed_users_file = args.allowed_users
title_experiment = args.title

SUBMISSION_FOLDER = os.path.join(RESULTS_FOLDER, 'submissions')
STATS_FOLDER = os.path.join(RESULTS_FOLDER, 'stats')
BANDIT_FOLDER = os.path.join(RESULTS_FOLDER, 'bandit')
if not os.path.exists(BANDIT_FOLDER):
    os.makedirs(BANDIT_FOLDER)

count_assignment = np.loadtxt(os.path.join(STATS_FOLDER,'pool_user_assignment','count_assignment.csv'), dtype=int, delimiter=',')

valid_submissions_list = np.load(os.path.join(STATS_FOLDER, 'pool_user_assignment', f'{problem_assignment_file}.npy'), allow_pickle=True)
n_pools, n_arms = valid_submissions_list.shape[0], valid_submissions_list.shape[1]
#convert to unmpy for easier indexing and vectorized funtions
valid_submissions = np.empty((n_pools, n_arms, np.max(count_assignment).astype(int)), dtype='<U64')
for pool_id in range(n_pools):
    for arm_id in range(n_arms):
        for person_count_idx in range(len(valid_submissions_list[pool_id, arm_id])):
            valid_submissions[pool_id, arm_id, person_count_idx] = valid_submissions_list[pool_id, arm_id][person_count_idx]


#read allowed users list (one file with one user id per line)
print(allowed_users_file)
if allowed_users_file == 'all':
    allowed_users = np.unique(valid_submissions[valid_submissions != ''])
else:
    with open(os.path.join(BANDIT_FOLDER, 'users', allowed_users_file), 'r') as f:
        allowed_users = np.array(json.load(f), dtype='<U64')

#*###################################################################################################
#* Define Bandit algorithm: arms are the ids of the arms. Here we assume that arms are 0,1,2,---,Bmax
arms = np.arange(n_arms)
algo = UCB1(arms)

BASE_SEED = 42         
rng = np.random.default_rng(BASE_SEED)

#* Initialize variables to store results
chosen_arms = []
rewards = []
cumulative_rewards = []
confidence_radius = []
arm_means = []


user_mask = np.isin(valid_submissions, allowed_users)

print(f'------- Selected users: {allowed_users_file}, n={allowed_users.shape}')
print(f'------- Number of valid submissions per arm: {user_mask.sum(axis=2).sum(axis=0)}')


for _ in tqdm(range(NUM_SIMS), desc='simulations'): # Loop through number of simulations

    #* Initialize the algorithm for each simulation
    algo.initialize()

    #* Initialize variables to store results for each simulation
    iter_chosen_arms = []
    iter_rewards = []
    iter_cumulative_rewards = []
    iter_confidence_radius = []
    iter_arm_means = []

    #* Initialize variables to store the selected pools: sample without replacement
    available_problems = np.full_like(valid_submissions, fill_value = True, dtype=bool)
    available_problems[valid_submissions==''] = False

    for t in tqdm(range(HORIZON), leave=False):

        #* UCB selects arm
        arm = algo.select_arm()
       
        #*Sample a pool and a user from that pool
        user_id = None
        pool_id = sample_valid_pool(arm, valid_submissions, available_problems, allowed_users, min_human_arm=min_arm, rng=rng)
        if arm>=min_arm: 
            pool_id, user_idx = pool_id
            user_id = valid_submissions[pool_id, arm, user_idx]
            available_problems[pool_id, arm, user_idx] = False  # Mark this user as used

        #*Load the pool object
        pool_path = os.path.join(POOLS_FOLDER, f'pool{pool_id}')
        problem_path = os.path.join(pool_path, f'B{arm}')
        with open(os.path.join(pool_path, 'pool.pkl'), 'rb') as f:
            pool = pkl.load(f)  #probs = np.load(os.path.join(pool_path, 'g.npy')), capacities = np.load(os.path.join(pool_path, 'capacities.npy'))
        #*Set algorithm matching
        if arm != pool.N:
            alg_matching = np.loadtxt(os.path.join(problem_path, 'alg_matching.csv'), delimiter=",")
            if alg_matching.ndim == 1: #if the algorithm only matches one person return [[person,loc]] instead of [person,loc]. 
                alg_matching = alg_matching.reshape(-1, 2)
            pool.set_matching_pairs(alg_matching.astype(int))

        if arm>0:
            if arm<min_arm: #assign the remaining person to the remaining place
                human_matching = np.loadtxt(os.path.join(os.path.join(pool_path, f'B{arm}'), 'opt_rem_match.csv'), delimiter=",").astype(int).reshape(-1,2)
                pool.set_matching_pairs(human_matching)
            else:
                #*Load the human partial matching
                human_matching = get_matching(f'pool{pool_id}/B{arm}', user_id, SUBMISSION_FOLDER)['matching']
                assert len(human_matching) == arm, f"Matching length is {len(human_matching)} instead of {arm}"
                human_matching = human_matching_to_index(human_matching)
                #*Set human matchings. If there are people assigned to None, assign them to random available places.
                valid_matchings = human_matching[human_matching[:,1] != None].astype(int)
                pool.set_matching_pairs(valid_matchings)
                #*check if there are people assigned to none place, assign them to random available places.
                if np.any(np.array(human_matching) == None) and RANDOM_ASSIGN:
                    rand_match = np.array([])
                    rem_slots = np.repeat(np.arange(pool.M), pool.remaining_capacities)
                    np.random.shuffle(rem_slots)
                    rand_match = np.array(list(zip(pool.remaining_people, rem_slots)))
                    pool.set_matching_pairs(rand_match)

        if RANDOM_ASSIGN:
            assert pool.check_matching_is_valid(complete=True) 

        reward = pool.get_utility_current_matching()
        iter_rewards.append(reward)
            
        if t==0:
            iter_cumulative_rewards.append(reward)
        else:
            iter_cumulative_rewards.append(iter_cumulative_rewards[t-1] + reward)
            
        algo.update(arm, reward)
        iter_chosen_arms.append(arm)
        iter_confidence_radius.append(algo.conf_radius.copy())
        iter_arm_means.append(algo.values.copy())

    # Append results of each simulation to the main variables
    chosen_arms.append(iter_chosen_arms)
    rewards.append(iter_rewards)
    cumulative_rewards.append(iter_cumulative_rewards)
    confidence_radius.append(iter_confidence_radius)
    arm_means.append(iter_arm_means)


np.savez(os.path.join(BANDIT_FOLDER, f'bandit-results-h{HORIZON}-s{NUM_SIMS}-{title_experiment}-{time.strftime("%Y%m%d-%H%M%S")}.npz'),
         horizon=HORIZON,
         num_sims=NUM_SIMS,
         arms=arms,
         chosen_arms=chosen_arms,
         rewards=rewards,
         cum_rewards=cumulative_rewards,
         arm_means=arm_means, 
         conf_radius=confidence_radius)