run_bandit_parallel.py

import os
import sys
import time
import json
import pickle as pkl
import argparse
from functools import partial
from concurrent.futures import ProcessPoolExecutor
from typing import List, Tuple

import numpy as np
from tqdm import tqdm

from MWBM.bandits import UCB1
from MWBM.mwbm_algs import greedy_max_weighted_matching
from UserStudyAnalysis.utils import (
    human_matching_to_index,
    get_matching,
    sample_valid_pool,
)

# ============== WORKER FUNCTION ======================================== #

def run_single_sim(
    sim_seed: int,
    arms: np.ndarray,
    valid_sub: np.ndarray,
    allowed: np.ndarray,
    POOLS_FOLDER: str,
    SUBMISSION_FOLDER: str,
    HORIZON: int,
    min_arm: int,
    RANDOM_ASSIGN: bool,
    STRATEGY: bool,
) -> Tuple[list, list, list, list, list]:
    """
    Execute one full run of the bandit algorithm.

    Parameters
    ----------
    sim_seed : int
        Unique seed created by the parent for this worker; guarantees
        independent yet reproducible randomness.
    Other parameters
        Immutable objects produced once in the parent and *read-only*
        in every worker.

    Returns
    -------
    (chosen_arms, rewards, cum_rewards, conf_radii, arm_means)
        Five lists whose length == HORIZON.
    """
    rng = np.random.default_rng(sim_seed)

    # Instantiate the bandit algorithm
    algo = UCB1(arms)
    algo.initialize()

    # Pre-allocate results
    chosen_arms      = []
    rewards          = []
    cum_rewards      = []
    conf_radii       = []
    arm_means        = []

    # Track which (pool, arm, userIdx) triples are still unused
    available = np.full_like(valid_sub, True, dtype=bool)
    available[valid_sub == ""] = False

    for t in range(HORIZON):
        arm = algo.select_arm()

        # Draw a pool / user for this arm
        pool_id = sample_valid_pool(
            arm,
            valid_sub,
            available,
            allowed,
            min_human_arm=min_arm,
            rng=rng,           
        )

        #  Load pool object and load the matching problem
        user_id = None # prolific id
        if arm >= min_arm:
            pool_id, user_idx = pool_id # unpack tuple if user does the problem
            if STRATEGY == "human":
                user_id = valid_sub[pool_id, arm, user_idx]
                available[pool_id, arm, user_idx] = False

        pool_path   = os.path.join(POOLS_FOLDER, f"pool{pool_id}")
        problem_dir = os.path.join(pool_path, f"B{arm}")

        # Load pool object
        with open(os.path.join(pool_path, "pool.pkl"), "rb") as fh:
            pool = pkl.load(fh)

        # Algorithmic matching
        if arm != pool.N:
            alg_matching = np.loadtxt(
                os.path.join(problem_dir, "alg_matching.csv"), delimiter=","
            )
            if alg_matching.ndim == 1:
                alg_matching = alg_matching.reshape(-1, 2)
            pool.set_matching_pairs(alg_matching.astype(int))

        # Match the remaining users either at random or using human answers
        if arm > 0:
            if STRATEGY == 'random': # Match the remaining as random 
                remaining = np.repeat(np.arange(pool.M),
                                        pool.remaining_capacities)
                rng.shuffle(remaining)
                rand_match = np.array(list(zip(pool.remaining_people,
                                                remaining)))
                pool.set_matching_pairs(rand_match)
                #print(f'b={arm}, remaining shape = {remaining.shape}, rem_capacities = {pool.remaining_capacities}, rem_people = {pool.remaining_people.shape}, util= {pool.get_utility_current_matching()}')
            
            elif STRATEGY == 'greedy': # Greedy matching
                # Build the sub-matrix: rows = still-unmatched people,
                # columns = every location in the pool.
                remaining_people  = pool.remaining_people           # (n_rem, )
                remaining_cap     = pool.remaining_capacities       # (L, )

                # `pool.job_probs` (or similar) should be the full N×L matrix of
                # success probabilities.  Adapt the attribute name if needed.
                job_probs_sub = pool.g[remaining_people, :]

                greedy_dest = greedy_max_weighted_matching(job_probs_sub, remaining_cap)

                greedy_pairs = [(p, j) for p, j in zip(remaining_people, greedy_dest)
                                        if j != -1]
                #if greedy_pairs:          # only call set_matching_pairs if something to add
                pool.set_matching_pairs(np.asarray(greedy_pairs, dtype=int))

            elif STRATEGY == 'human': # Human matching
                if arm < min_arm:
                    human_matching = np.loadtxt(
                        os.path.join(problem_dir, "opt_rem_match.csv"),
                        delimiter=",",
                    ).astype(int).reshape(-1, 2)
                    pool.set_matching_pairs(human_matching)
                else:
                    hm = get_matching(f"pool{pool_id}/B{arm}", user_id,
                                    SUBMISSION_FOLDER)["matching"]
                    hm = human_matching_to_index(hm)
                    valid_hm = hm[hm[:, 1] != None].astype(int)
                    pool.set_matching_pairs(valid_hm)

                    if np.any(hm == None) and RANDOM_ASSIGN: # Randomly assign the individuals that user left unmatched
                        remaining = np.repeat(np.arange(pool.M),
                                            pool.remaining_capacities)
                        rng.shuffle(remaining)
                        rand_match = np.array(list(zip(pool.remaining_people,
                                                    remaining)))
                        pool.set_matching_pairs(rand_match)
            else:
                raise ValueError(f"Unknown assignment strategy: {STRATEGY}")

        if (RANDOM_ASSIGN and STRATEGY) or STRATEGY in ['human', 'greedy']:
            assert pool.check_matching_is_valid(complete=True)

        reward = pool.get_utility_current_matching()

        # -------- bookkeeping -------------------------------------------- #
        rewards.append(reward)
        cum_rewards.append(reward if t == 0 else cum_rewards[-1] + reward)

        algo.update(arm, reward)
        chosen_arms.append(arm)
        conf_radii.append(algo.conf_radius.copy())
        arm_means.append(algo.values.copy())

    return chosen_arms, rewards, cum_rewards, conf_radii, arm_means


# ============== Wrapper for paralellization ====================================== #

def main() -> None:
    """Parse args, load data once, fan out workers, save results."""
    # --------------------------------------------------------------------- #
    #  Argument parsing
    # --------------------------------------------------------------------- #
    parser = argparse.ArgumentParser(
        description="Run Bandit System with Human Data (parallel version)"
    )
    parser.add_argument("--pools-folder", required=True, type=str)
    parser.add_argument("--results-folder", required=True, type=str)
    parser.add_argument("--problems", default="pool_assignment", type=str)
    parser.add_argument("--num_sims", required=True, type=int)
    parser.add_argument("--horizon", required=True, type=int)
    parser.add_argument("--random-assign", action="store_true", default=False)
    parser.add_argument("--assignment-strategy",  default='human', type=str,
                        choices=['human', 'random', 'greedy']),
    parser.add_argument("--min-arm", required=True, type=int)
    parser.add_argument("--allowed-users", default="all", type=str)
    parser.add_argument("--title", default="exp", type=str)
    # NEW flags
    parser.add_argument("--workers", type=int, default=os.cpu_count(),
                        help="Number of parallel processes (default: all cores)")
    parser.add_argument("--seed", type=int,
                        help="Base RNG seed for reproducibility")

    args = parser.parse_args()

    # --------------------------------------------------------------------- #
    #  Derived paths & folders
    # --------------------------------------------------------------------- #
    POOLS_FOLDER    = args.pools_folder
    RESULTS_FOLDER  = args.results_folder
    SUBMISSION_FOLD = os.path.join(RESULTS_FOLDER, "submissions")
    STATS_FOLDER    = os.path.join(RESULTS_FOLDER, "stats")
    BANDIT_FOLDER   = os.path.join(RESULTS_FOLDER, "bandit")
    os.makedirs(BANDIT_FOLDER, exist_ok=True)

    # --------------------------------------------------------------------- #
    #  Load heavy data *once*
    # --------------------------------------------------------------------- #
    count_assignment = np.loadtxt(
        os.path.join(STATS_FOLDER, "pool_user_assignment", "count_assignment.csv"),
        dtype=int,
        delimiter=",",
    )

    valid_list = np.load(
        os.path.join(STATS_FOLDER, "pool_user_assignment",
                     f"{args.problems}.npy"),
        allow_pickle=True,
    )
    n_pools, n_arms = valid_list.shape
    valid_submissions = np.empty(
        (n_pools, n_arms, np.max(count_assignment).astype(int)),
        dtype="<U64",
    )
    for p in range(n_pools):
        for a in range(n_arms):
            for idx, usr in enumerate(valid_list[p, a]):
                valid_submissions[p, a, idx] = usr

    # Allowed users
    if args.allowed_users == "all":
        allowed_users = np.unique(valid_submissions[valid_submissions != ""])
    else:
        with open(os.path.join(BANDIT_FOLDER, "users", args.allowed_users)) as fh:
            allowed_users = np.array(json.load(fh), dtype="<U64")

    arms = np.arange(n_arms)

    # --------------------------------------------------------------------- #
    #  Prepare seeds for each worker
    # --------------------------------------------------------------------- #
    base_ss   = np.random.SeedSequence(args.seed)
    sim_seeds = [int(s.entropy) for s in base_ss.spawn(args.num_sims)]

    # --------------------------------------------------------------------- #
    #  Partial-bind the worker so map() sees a one-argument function
    # --------------------------------------------------------------------- #
    worker_fn = partial(
        run_single_sim,
        arms=arms,
        valid_sub=valid_submissions,
        allowed=allowed_users,
        POOLS_FOLDER=POOLS_FOLDER,
        SUBMISSION_FOLDER=SUBMISSION_FOLD,
        HORIZON=args.horizon,
        min_arm=args.min_arm,
        RANDOM_ASSIGN=args.random_assign,
        STRATEGY=args.assignment_strategy,
    )

    # --------------------------------------------------------------------- #
    #  Run the pool
    # --------------------------------------------------------------------- #
    chosen_arms, rewards, cum_rewards, conf_radii, arm_means = (
        [] for _ in range(5)
    )

    with ProcessPoolExecutor(max_workers=args.workers) as pool:
        for ca, rw, cr, cf, am in tqdm(
            pool.map(worker_fn, sim_seeds),
            total=args.num_sims,
            desc="simulations",
        ):
            chosen_arms.append(ca)
            rewards.append(rw)
            cum_rewards.append(cr)
            conf_radii.append(cf)
            arm_means.append(am)

    # --------------------------------------------------------------------- #
    #  Persist results *once*
    # --------------------------------------------------------------------- #
    outfile = os.path.join(
        BANDIT_FOLDER,
        f"bandit-results-h{args.horizon}"
        f"-s{args.num_sims}"
        f"-{args.title}"
        f"-{time.strftime('%Y%m%d-%H%M%S')}.npz",
    )
    np.savez(
        outfile,
        horizon=args.horizon,
        num_sims=args.num_sims,
        arms=arms,
        chosen_arms=chosen_arms,
        rewards=rewards,
        cum_rewards=cum_rewards,
        arm_means=arm_means,
        conf_radius=conf_radii,
    )
    print(f"\nSaved → {outfile}")


# ============== 3.  MAIN ENTRY POINT ====================================== #
if __name__ == "__main__":
    # Windows needs this guard, and it’s good practice everywhere.
    main()