utils.py

import numpy as np
import matplotlib.pyplot as plt
from scipy import stats
import pandas as pd
import torch
import pyrosetta
from envs import Surface, Grid, Molecule
from common_nets import Mlp
from policy import Policy

# Default device
DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# Get correct environment
def get_environment(env_name):
    if env_name == 'naive_surface':
        return Surface(naive=True)
    if env_name == 'surface':
        return Surface()
    if env_name == 'naive_grid':
        return Grid(naive=True)
    if env_name == 'grid':
        return Grid()
    pose = pyrosetta.pose_from_sequence('A'*8)
    if env_name == 'naive_molecule':
        return Molecule(pose=pose, naive=True)
    if env_name == 'molecule':
        return Molecule(pose=pose)

def from_str_to_2D_arr(s):
    tokens = s[2:-2].split("],[")
    ans = []
    for arr_str in tokens:
        arr = []
        elem_str = arr_str.split(",")
        for e in elem_str:
            arr.append(int(e))
        ans.append(arr)
    return ans

def str_to_list(s):
    tokens = s[1:-1].split(",")
    ans = []
    for token in tokens:
        ans.append(int(token))
    return ans

# Get neural net architecture
def get_architectures(env_name, zero_order, arch_file='arch.csv'):
    # Get architecture info from arch_file
    df = pd.read_csv(arch_file)
    net_info = df[df['env_name']==env_name]
    if zero_order:
        layer_dims = str_to_list(net_info['derivative_layer_dims'].values[0])
    else:
        layer_dims = str_to_list(net_info['val_layer_dims'].values[0])

    return layer_dims

# Get predefined training parameters from file for a specific environment. 
def get_train_params(env_name, param_file='params.csv'):
    # Get parameter info from param file
    df = pd.read_csv(param_file)
    info = df[df['env_name']==env_name]

    # Rate underlying hamiltonian dynamics formula.
    rate = float(info['rate'].values[0])

    # Number of trajectories per stage
    num_traj = int(info['num_traj'].values[0])
    
    # Step size for discretized ODE
    step_size = float(info['step_size'].values[0])
    
    # Optimization params: learning rate, batch size, how often logging
    # and number of optimization steps per each sampling stage.
    lr = float(info['lr'].values[0])
    batch_size = int(info['batch_size'].values[0])
    log_interval = int(info['log_interval'].values[0])

    return rate, num_traj, step_size, lr, batch_size, log_interval

def setup_main_net(env_name, zero_order, state_dim):
    layer_dims = get_architectures(env_name, zero_order)
    if zero_order:
        output_dim = 1
    else:
        output_dim = state_dim
    main_net = Mlp(input_dim=state_dim, output_dim=output_dim, 
                    layer_dims=layer_dims, activation='relu').to(DEVICE)
    
    return main_net


def setup_dpo_model(method, env, env_name):
    rate, _, step_size, _, _, _ = get_train_params(env_name)
    zero_order = method.endswith('zero_order')
    state_dim = env.state_dim
    main_net = setup_main_net(env_name, zero_order, state_dim)
    path = 'models/' + env_name + '_' + method + '.pth'
    main_net.load_state_dict(torch.load(path))
    main_net.to(DEVICE)
    model = Policy(zero_order, main_net, rate, step_size)

    return model

### Plotting
def _bootstrap(data, n_boot=2000, ci=68, scale=1.0):
    boot_dist = []
    for _ in range(int(n_boot)):
        resampler = np.random.randint(0, data.shape[0], data.shape[0])
        sample = data.take(resampler, axis=0)
        boot_dist.append(np.mean(sample, axis=0))
    b = np.array(boot_dist)

    if scale < 2.0:
        b = np.array(boot_dist)
        s1 = np.apply_along_axis(stats.scoreatpercentile, 0, b, 50.-ci/2.)
        s2 = np.apply_along_axis(stats.scoreatpercentile, 0, b, 50.+ci/2.)
        return (s1,s2)
    else:
        median = np.median(b, axis=0)
        lower = np.percentile(b, 50 - ci / 2, axis=0)
        upper = np.percentile(b, 50 + ci / 2, axis=0)
        delta_low = scale * (median - lower)
        delta_high = scale * (upper - median)

        # Enforce minimum band
        delta_low = np.minimum(delta_low, np.abs(median)/5.0)
        delta_high = np.minimum(delta_high, np.abs(median)/5.0)

        # Scale CI around the median
        s1 = median - delta_low
        s2 = median + delta_high
        return s1, s2

def _tsplot(ax, x, data, mode='bootstrap', std_scale=1.0, **kw):
    est = np.mean(data, axis=0)
    if mode == 'bootstrap':
        cis = _bootstrap(data, scale=std_scale)
    else:
        sd = np.std(data, axis=0)
        cis = (est - std_scale * sd, est + std_scale * sd)
    p2 = ax.fill_between(x, cis[0], cis[1], alpha=0.2, **kw)
    p1 = ax.plot(x, est, **kw)
    ax.margins(x=0)
    return p1, p2

def plot_eval_benchmarks(eval_dict, time_steps, title, mode='bootstrap',
                         colors=['red', 'blue', 'green', 'orange'],
                         plot_filename='tmp.png',
                         std_scale=1.0):
    methods = list(eval_dict.keys())
    ax = plt.gca()
    graphic_list = []
    for i, method in enumerate(methods):
        data = eval_dict[method]
        _, p2 = _tsplot(ax, np.array(time_steps), data, mode,
                        label=method, color=colors[i], std_scale=std_scale)
        graphic_list.append(p2)
    ax.legend(graphic_list, methods)
    ax.set_title(title)
    ax.set_xlabel('Timestamp')
    ax.set_ylabel('Evaluation cost')
    # plt.tight_layout()
    plt.savefig(plot_filename)
    plt.show()