<filename>example_opt_root/util.py
import os
from torch import Tensor
import torch
from typing import Any
import json
from collections import namedtuple
import numpy as np
import re

def set_seed(SEED):
    torch.manual_seed(SEED)
    torch.cuda.manual_seed(SEED)

def log_check(file, logdir = './log/'):
    if not os.path.exists(os.path.join(logdir, file)):
        return False
    with open(os.path.join(logdir, file), 'r') as f:
        for line in f:
            if re.match('Traceback+', line):
                # print('error!', file)
                return False
        f.close()
    return True

def clear_config(configdir = './config/'):
    files = os.listdir(configdir)
    for f in files:
        loc = f.find('_')
        outname = f[loc + 1:] + '.out'
        if log_check(outname):
            print('remove', os.path.join(configdir, f))
            os.remove(os.path.join(configdir, f))
    

def generate_initial_data(n, fun, device, all_value=False):
    R_MIN = 50
    n = n * 10
    # generate training data
    DIM = fun.fun._dimensions
    n = 10 * DIM  # fixed
    bounds = torch.tensor([(0, 1) for _ in range(DIM)],
                          device=device, dtype=torch.double).T

    train_x = torch.rand(n, DIM, device=device, dtype=torch.double) * \
        (bounds[1] - bounds[0]) + bounds[0]
    train_obj = fun(train_x).unsqueeze(-1)
    for _ in range(R_MIN-1):
        train_obj = torch.cat((train_obj, fun(train_x).unsqueeze(-1)), dim=-1)
    var_obj = torch.var(train_obj, dim=-1)
    mean_obj = torch.mean(train_obj, dim=-1)
    best_observed_value = fun.exact(train_x[mean_obj.argmax().item()])
    if all_value:
        return n * R_MIN, train_x, mean_obj.unsqueeze(-1), best_observed_value, var_obj.unsqueeze(-1), train_obj
    else:
        return n * R_MIN, train_x, mean_obj.unsqueeze(-1), best_observed_value, var_obj.unsqueeze(-1)

def get_config(config_path):
    with open(config_path) as config_file:
        base_config = json.load(config_file)

    if os.path.exists('job_parameters.json'):
        with open('job_parameters.json') as param_config_file:
            param_config = json.load(param_config_file)
    else:
        param_config = {}

    config = base_config
    for section, d in param_config.items():
        for k in d:
            assert k in config[section]
        config[section].update(d)
    return config


def config_to_namedtuple(obj):
    if isinstance(obj, dict):
        for key, value in obj.items():
            obj[key] = config_to_namedtuple(value)
        return namedtuple('GenericDict', obj.keys())(**obj)
    elif isinstance(obj, list):
        return [config_to_namedtuple(item) for item in obj]
    else:
        return obj


class linear_noisy_function():
    def __init__(self, fun, var) -> None: 
        self.fun = fun
        self.var = np.sqrt(var)
        self.bounds = torch.tensor(self.fun.bounds, dtype=torch.double).T
        optimum_x = torch.tensor(fun.global_optimum).squeeze()
        if fun._dimensions > 1:
            if optimum_x.dim() >= 2:
                optimum_x = optimum_x[0]
        else:
            if optimum_x.dim() >= 1:
                optimum_x = optimum_x[0]
            optimum_x = optimum_x.unsqueeze(0)

        optimum_x = optimum_x.numpy()
        # optimum_x = torch.tensor(fun.global_optimum[0]).numpy()
        self.opt = torch.tensor(self.fun.fun(optimum_x)).squeeze()
        
        if self.opt.dim() != 0:
            self.opt = self.opt[0]
        self.opt = self.opt.item()
        train_x = torch.rand( 1000, fun._dimensions, dtype=torch.double) * (self.bounds[1] - self.bounds[0]) + self.bounds[0]
        train_x = train_x.cpu().numpy()
        train_obj = torch.tensor(list(map(fun.fun, train_x))).unsqueeze(-1)
        self.min = train_obj.min().item()
        self.max = train_obj.max().item()
        print(self.min, self.max, self.opt)
        assert self.opt <= self.min + 1e-5
    def exact(self, X, *args: Any) -> Tensor:
        cuda = X.device
        X = X.cpu()
        X = X * (self.bounds[1] - self.bounds[0]) + self.bounds[0]
        X = torch.clamp(X, min=self.bounds[0], max=self.bounds[1])
        X = X.numpy()
        if X.ndim >= 2:
            v = map(self.fun.fun, X)
            u = np.array(list(v))
            exact = torch.tensor(u)
        else:
            exact = torch.tensor(self.fun.fun(X))
        exact = (self.max - exact) / (self.max - self.opt)
        # print(exact, self.opt, self.min, self.max)
        return exact.to(cuda)


    def __call__(self, X, *args: Any) -> Tensor:
        exact = self.exact(X)
        return exact + torch.minimum(exact * self.var * torch.randn_like(exact), torch.ones_like(exact))



class homo_noisy_function(linear_noisy_function):

    def __call__(self, X, *args: Any) -> Tensor:
        exact = self.exact(X)
        return exact + torch.minimum(self.var * torch.randn_like(exact), torch.ones_like(exact))


class sin_noisy_function(linear_noisy_function):

    def __call__(self, X, *args: Any) -> Tensor:
        exact = self.exact(X)
        return exact + torch.minimum((torch.sin(exact) * self.var * torch.randn_like(exact)), torch.ones_like(exact))


if __name__ == '__main__':
    clear_config()