run_comparisons.py

#!/usr/bin/env python3

"""
Refactored script to run different Bayesian Optimization algorithms:
    - Single-task with constraints
    - Single-task without constraints
    - Multi-task without constraints

Usage:
    python run_comparisons.py <function_name> <algorithm_type>

Arguments:
    function_name (str): The name of the function to optimize. Options:
        - MTBranin
        - LbSync
        - MTPowell
    algorithm_type (str): The type of algorithm to run. Options:
        - "st_constraints" (Single-task with constraints)
        - "st_no_constraints" (Single-task without constraints)
        - "mt_no_constraints" (Multi-task without constraints)
"""

import os
import pickle
import random
import sys
from math import ceil

import numpy
import torch
from botorch.utils.transforms import normalize, unnormalize
from numpy import load

from bo.bo_loop import MultiTaskBayesianOptimization, SingleTaskBayesianOptimization
from utils.functions import LbSync, MTBranin, MTPowell
from utils.get_robust_gp import beta_bayes
from utils.optim import optimize_gp
from utils.utils import (
    build_mtgp,
    build_stgp,
    concat_data,
    sample_from_task,
    standardize,
)

torch.set_default_dtype(torch.float64)

# Global constants
NRUNS = 50
DELTA_MAX = 0.05
TAU = 0.001
DIST = 0.3


def initialize_function(function_name):
    """Initialize the objective function and related parameters."""
    if function_name == "MTBranin":
        obj = MTBranin(num_tsks=2)
    elif function_name == "LbSync":
        obj = LbSync(Ktyp="PI", num_lasers=5, num_tsks=2, disturbance=DIST)
    elif function_name == "MTPowell":
        obj = MTPowell(num_tsks=2)
    else:
        raise ValueError(f"Unknown function name: {function_name}")
    return obj, obj.dim, obj.bounds


def load_initial_data(function_name):
    """Load initial data points and thresholds."""
    data = load(f"data/X_init_{function_name}.npy", allow_pickle=True).item()
    return torch.tensor(data["X_init"]), data["threshold"]


def setup_directories(folder):
    """Create output directory if it doesn't exist."""
    if not os.path.exists(f"data/{folder}"):
        os.mkdir(f"data/{folder}")


def run_single_task_with_constraints(function_name, X_init, T):
    """Run single-task Bayesian Optimization with constraints."""
    folder = "Bayes_ST"
    setup_directories(folder)
    return run_bayesian_optimization(
        function_name, X_init, T, folder, single_task=True, constraints=True
    )


def run_single_task_without_constraints(function_name, X_init, T):
    """Run single-task Bayesian Optimization without constraints."""
    folder = "vanilla_bo"
    setup_directories(folder)
    return run_bayesian_optimization(
        function_name, X_init, T, folder, single_task=True, constraints=False
    )


def run_multi_task_without_constraints(function_name, X_init, T):
    """Run multi-task Bayesian Optimization without constraints."""
    folder = "Bayes_MT"
    setup_directories(folder)
    return run_bayesian_optimization(
        function_name, X_init, T, folder, single_task=False, constraints=False
    )


def build_gp(train_inputs, train_task, train_targets, single_task=True):
    """Build a Gaussian Process model."""
    if single_task:
        return build_stgp(train_inputs, train_targets)
    else:
        return build_mtgp((train_inputs, train_task), train_targets)


def run_bayesian_optimization(
    function_name, X_init, T, folder, single_task, constraints
):
    """Generalized Bayesian Optimization loop."""

    data_sets = []
    bests = []

    for i in range(X_init.size(0)):
        seeds = seeds_dict[function_name]
        torch.manual_seed(seeds[i])
        numpy.random.seed(seeds[i])
        random.seed(seeds[i])

        obj, d, bounds = initialize_function(function_name)
        num_tsks = 1 if single_task else obj.num_tsks
        norm_bounds = torch.vstack((torch.zeros(1, d), torch.ones(1, d)))
        num_sup_task_samples = 1 if single_task else ceil(2 * d / (num_tsks - 1))
        num_acq_samps = [1] + [num_sup_task_samples] * (num_tsks - 1)
        print(f"Round: {i + 1}")

        x0 = X_init[i, ...].view(1, bounds.size(-1))
        norm_x0 = normalize(x0, bounds)

        # Evaluate initial point for all tasks
        train_targets, train_tasks, norm_train_inputs = initialize_training_data(
            obj, norm_x0, num_tsks
        )

        # Evaluate supplementary tasks if multi-task
        if not single_task:
            norm_train_inputs, train_tasks, train_targets = (
                evaluate_supplementary_tasks(
                    obj,
                    num_tsks,
                    norm_bounds,
                    num_sup_task_samples,
                    norm_train_inputs,
                    train_tasks,
                    train_targets,
                )
            )

        norm_train_targets = standardize(train_targets, T=T)
        T_stdizd = standardize(T, T)
        if single_task:
            bo = SingleTaskBayesianOptimization(
                obj, norm_bounds, T_stdizd, T, num_acq_samps, constraints=constraints
            )
        else:
            bo = MultiTaskBayesianOptimization(
                obj,
                list(range(num_tsks)),
                norm_bounds,
                T_stdizd,
                T,
                num_acq_samps,
                constraints=constraints,
            )
        gp = build_gp(
            norm_train_inputs, train_tasks, norm_train_targets, single_task=single_task
        )

        bo = run_bo_iterations(
            bo,
            gp,
            norm_bounds,
            norm_train_inputs,
            train_tasks,
            norm_train_targets,
            single_task,
            constraints,
        )

        train_inputs = unnormalize(bo.train_inputs, bounds)
        train_targets = bo.unstd_train_targets
        data_sets.append(
            [
                train_inputs,
                torch.zeros(1, 1) if single_task else bo.train_tasks,
                train_targets,
            ]
        )
        bests.append([bo.best_x, bo.best_y])
        print(
            f"Best value: {round(bo.best_y[-1], 3)} at input: {unnormalize(bo.best_x[-1], bounds).round(decimals=3)}"
        )

    # Save data
    save_results(function_name, folder, data_sets, bests)


def initialize_training_data(obj, norm_x0, num_tsks):
    """Initialize training data with the initial point."""
    train_targets = torch.zeros(num_tsks, 1)
    for j in range(num_tsks):
        train_targets[j, ...] = obj.f(norm_x0, j)
    train_tasks = torch.arange(num_tsks).unsqueeze(-1)
    norm_train_inputs = norm_x0.repeat(num_tsks, 1)
    return train_targets, train_tasks, norm_train_inputs


def evaluate_supplementary_tasks(
    obj,
    num_tsks,
    norm_bounds,
    num_sup_task_samples,
    norm_train_inputs,
    train_tasks,
    train_targets,
):
    """Evaluate supplementary tasks for multi-task optimization."""
    for k in range(1, num_tsks):
        x, t, y = sample_from_task(obj, [k], norm_bounds, n=2 * num_sup_task_samples)
        norm_train_inputs, train_tasks, train_targets = concat_data(
            (x, t, y), (norm_train_inputs, train_tasks, train_targets)
        )
    return norm_train_inputs, train_tasks, train_targets


def run_bo_iterations(
    bo,
    gp,
    norm_bounds,
    norm_train_inputs,
    train_tasks,
    norm_train_targets,
    single_task,
    constraints,
):
    """Run the Bayesian Optimization iterations."""
    for _ in range(NRUNS):
        if constraints or _ % 2 == 0:
            robust_gp = build_gp(
                norm_train_inputs,
                train_tasks,
                norm_train_targets,
                single_task=single_task,
            )
        else:
            gp = build_gp(
                norm_train_inputs,
                train_tasks,
                norm_train_targets,
                single_task=single_task,
            )
            gp, _, _ = optimize_gp(
                gp, mode=1, max_iter=200, train_covar=not single_task
            )
            robust_gp = gp
        sqrtbeta = torch.sqrt(beta_bayes(norm_bounds, TAU, DELTA_MAX))
        bo.update_gp(robust_gp, sqrtbeta)
        norm_train_inputs, train_tasks, norm_train_targets = bo.step()
    return bo


def save_results(function_name, folder, data_sets, bests):
    """Save the results to a file."""
    with open(f"data/{folder}/{function_name}_12.obj", "wb") as file:
        pickle.dump({"data_sets": data_sets, "bests": bests}, file)


# Define a dictionary of seeds for each function_name
seeds_dict = {
    "MTBranin": [
        4001630359,
        1919308292,
        1427243158,
        2443879631,
        2368281548,
        1373693466,
        1878128376,
        2416680340,
        2577055479,
        473042899,
        2789764286,
        2245849798,
        3162095579,
        3697226237,
        3466549482,
    ],
    "LbSync": [
        378460615,
        2204309691,
        3714415149,
        3247551807,
        1394259410,
        2593874679,
        1338135722,
        229394925,
        1148606685,
        738629800,
        3400498408,
        3913446390,
        2823833501,
        2483689362,
        1656936394,
    ],
    "MTPowell": [
        2788530576,
        588811540,
        1073174780,
        1979726901,
        237964780,
        1720245674,
        371781323,
        946294099,
        69072920,
        1340726846,
        1378218114,
        1976946218,
        453164431,
        977118581,
        825415680,
    ],
}


if __name__ == "__main__":
    if len(sys.argv) != 3:
        print("Usage: python run_comparisons.py <function_name> <algorithm_type>")
        sys.exit(1)

    function_name = sys.argv[1]
    algorithm_type = sys.argv[2]

    print(algorithm_type)

    X_init, T = load_initial_data(function_name)

    if algorithm_type == "st_constraints":
        run_single_task_with_constraints(function_name, X_init, T)
    elif algorithm_type == "st_no_constraints":
        run_single_task_without_constraints(function_name, X_init, T)
    elif algorithm_type == "mt_no_constraints":
        run_multi_task_without_constraints(function_name, X_init, T)
    else:
        print(f"Unknown algorithm type: {algorithm_type}")
        sys.exit(1)