Upload 529 files

Ackley10D_CEI_Avg_Obj.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ef078fab3310090b39fa175558b54f1a6819ea07cb0b9e19f38b39bd4c27c12b
+size 2968

Ackley10D_CEI_Avg_Time.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:0e43f57c340f37f7e912b2143872910e14e797a16f1c16c1b3088cf3d550c64a
+size 3484

Car_CEI_Avg_Obj.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:2703f7d2083dfc5a340b082b9b16406467443a82ced26ac7202f7440f68c9854
+size 3008

Car_CEI_Avg_Time.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:acaa29100dfafb78d40a7391feae3d042716159955ec46772eb2f0c017830d1a
+size 3396

Gradio_important.ipynb

@@ -0,0 +1,588 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "id": "0823043e-8451-4dc8-968c-ca066003f4a7",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Running on local URL:  http://127.0.0.1:7958\n",
+      "\n",
+      "To create a public link, set `share=True` in `launch()`.\n"
+     ]
+    },
+    {
+     "data": {
+      "text/html": [
+       "<div><iframe src=\"http://127.0.0.1:7958/\" width=\"100%\" height=\"500\" allow=\"autoplay; camera; microphone; clipboard-read; clipboard-write;\" frameborder=\"0\" allowfullscreen></iframe></div>"
+      ],
+      "text/plain": [
+       "<IPython.core.display.HTML object>"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "text/plain": []
+     },
+     "execution_count": 1,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "import gradio as gr\n",
+    "import torch\n",
+    "import numpy as np\n",
+    "import matplotlib.pyplot as plt\n",
+    "from test_functions.Ackley10D import *\n",
+    "from test_functions.Ackley2D import *\n",
+    "from test_functions.Ackley6D import *\n",
+    "from test_functions.HeatExchanger import *\n",
+    "from test_functions.CantileverBeam import *\n",
+    "from test_functions.Car import *\n",
+    "from test_functions.CompressionSpring import *\n",
+    "from test_functions.GKXWC1 import *\n",
+    "from test_functions.GKXWC2 import *\n",
+    "from test_functions.HeatExchanger import *\n",
+    "from test_functions.JLH1 import *\n",
+    "from test_functions.JLH2 import *\n",
+    "from test_functions.KeaneBump import *\n",
+    "from test_functions.GKXWC1 import *\n",
+    "from test_functions.GKXWC2 import *\n",
+    "from test_functions.PressureVessel import *\n",
+    "from test_functions.ReinforcedConcreteBeam import *\n",
+    "from test_functions.SpeedReducer import *\n",
+    "from test_functions.ThreeTruss import *\n",
+    "from test_functions.WeldedBeam import *\n",
+    "# Import other objective functions as needed\n",
+    "import time\n",
+    "\n",
+    "from Rosen_PFN4BO import *\n",
+    "from PIL import Image\n",
+    "\n",
+    "\n",
+    "\n",
+    "\n",
+    "\n",
+    "\n",
+    "\n",
+    "\n",
+    "\n",
+    "\n",
+    "\n",
+    "\n",
+    "\n",
+    "\n",
+    "\n",
+    "def s(input_string):\n",
+    "    return input_string\n",
+    "\n",
+    "\n",
+    "\n",
+    "\n",
+    "def optimize(objective_function, iteration_input, progress=gr.Progress()):\n",
+    "\n",
+    "    print(objective_function)\n",
+    "\n",
+    "    # Variable setup\n",
+    "    Current_BEST = torch.tensor( -1e10 )   # Some arbitrary very small number\n",
+    "    Prev_BEST = torch.tensor( -1e10 )\n",
+    "\n",
+    "    if objective_function==\"CantileverBeam.png\":\n",
+    "        Current_BEST = torch.tensor( -82500  )  # Some arbitrary very small number\n",
+    "        Prev_BEST = torch.tensor( -82500 )\n",
+    "    elif objective_function==\"CompressionSpring.png\":\n",
+    "        Current_BEST = torch.tensor( -8  )  # Some arbitrary very small number\n",
+    "        Prev_BEST = torch.tensor( -8 )\n",
+    "    elif objective_function==\"HeatExchanger.png\":\n",
+    "        Current_BEST = torch.tensor( -30000 )   # Some arbitrary very small number\n",
+    "        Prev_BEST = torch.tensor( -30000 )\n",
+    "    elif objective_function==\"ThreeTruss.png\":\n",
+    "        Current_BEST = torch.tensor( -300  )  # Some arbitrary very small number\n",
+    "        Prev_BEST = torch.tensor( -300 )\n",
+    "    elif objective_function==\"Reinforcement.png\":\n",
+    "        Current_BEST = torch.tensor( -440   ) # Some arbitrary very small number\n",
+    "        Prev_BEST = torch.tensor( -440 )\n",
+    "    elif objective_function==\"PressureVessel.png\":\n",
+    "        Current_BEST = torch.tensor( -40000  )  # Some arbitrary very small number\n",
+    "        Prev_BEST = torch.tensor( -40000    ) \n",
+    "    elif objective_function==\"SpeedReducer.png\":\n",
+    "        Current_BEST = torch.tensor( -3200  )  # Some arbitrary very small number\n",
+    "        Prev_BEST = torch.tensor( -3200  )   \n",
+    "    elif objective_function==\"WeldedBeam.png\":\n",
+    "        Current_BEST = torch.tensor( -35  )  # Some arbitrary very small number\n",
+    "        Prev_BEST = torch.tensor( -35   )\n",
+    "    elif objective_function==\"Car.png\":\n",
+    "        Current_BEST = torch.tensor( -35  )  # Some arbitrary very small number\n",
+    "        Prev_BEST = torch.tensor( -35   )\n",
+    "\n",
+    "    # Initial random samples\n",
+    "    # print(objective_functions)\n",
+    "    trained_X = torch.rand(20, objective_functions[objective_function]['dim'])\n",
+    "\n",
+    "    # Scale it to the domain of interest using the selected function\n",
+    "    # print(objective_function)\n",
+    "    X_Scaled = objective_functions[objective_function]['scaling'](trained_X)\n",
+    "\n",
+    "    # Get the constraints and objective\n",
+    "    trained_gx, trained_Y = objective_functions[objective_function]['function'](X_Scaled)\n",
+    "\n",
+    "    # Convergence list to store best values\n",
+    "    convergence = []\n",
+    "    time_conv = []\n",
+    "\n",
+    "    START_TIME = time.time()\n",
+    "\n",
+    "\n",
+    "# with gr.Progress(track_tqdm=True) as progress:\n",
+    "\n",
+    "\n",
+    "    # Optimization Loop\n",
+    "    for ii in progress.tqdm(range(iteration_input)):  # Example with 100 iterations\n",
+    "\n",
+    "        # (0) Get the updated data for this iteration\n",
+    "        X_scaled = objective_functions[objective_function]['scaling'](trained_X)\n",
+    "        trained_gx, trained_Y = objective_functions[objective_function]['function'](X_scaled)\n",
+    "\n",
+    "        # (1) Randomly sample Xpen \n",
+    "        X_pen = torch.rand(1000,trained_X.shape[1])\n",
+    "\n",
+    "        # (2) PFN inference phase with EI\n",
+    "        default_model = 'final_models/model_hebo_morebudget_9_unused_features_3.pt'\n",
+    "        \n",
+    "        ei, p_feas = Rosen_PFN_Parallel(default_model,\n",
+    "                                           trained_X, \n",
+    "                                           trained_Y, \n",
+    "                                           trained_gx,\n",
+    "                                           X_pen,\n",
+    "                                           'power',\n",
+    "                                           'ei'\n",
+    "                                          )\n",
+    "\n",
+    "        # Calculating CEI\n",
+    "        CEI = ei\n",
+    "        for jj in range(p_feas.shape[1]):\n",
+    "            CEI = CEI*p_feas[:,jj]\n",
+    "\n",
+    "        # (4) Get the next search value\n",
+    "        rec_idx = torch.argmax(CEI)\n",
+    "        best_candidate = X_pen[rec_idx,:].unsqueeze(0)\n",
+    "\n",
+    "        # (5) Append the next search point\n",
+    "        trained_X = torch.cat([trained_X, best_candidate])\n",
+    "\n",
+    "\n",
+    "        ################################################################################\n",
+    "        # This is just for visualizing the best value. \n",
+    "        # This section can be remove for pure optimization purpose\n",
+    "        Current_X = objective_functions[objective_function]['scaling'](trained_X)\n",
+    "        Current_GX, Current_Y = objective_functions[objective_function]['function'](Current_X)\n",
+    "        if ((Current_GX<=0).all(dim=1)).any():\n",
+    "            Current_BEST = torch.max(Current_Y[(Current_GX<=0).all(dim=1)])\n",
+    "        else:\n",
+    "            Current_BEST = Prev_BEST\n",
+    "        ################################################################################\n",
+    "        \n",
+    "        # (ii) Convergence tracking (assuming the best Y is to be maximized)\n",
+    "        # if Current_BEST != -1e10:\n",
+    "        print(Current_BEST)\n",
+    "        print(convergence)\n",
+    "        convergence.append(Current_BEST.abs())\n",
+    "        time_conv.append(time.time() - START_TIME)\n",
+    "\n",
+    "    # Timing\n",
+    "    END_TIME = time.time()\n",
+    "    TOTAL_TIME = END_TIME - START_TIME\n",
+    "    \n",
+    "    # Website visualization\n",
+    "    # (i) Radar chart for trained_X\n",
+    "    radar_chart = None\n",
+    "    # radar_chart = create_radar_chart(X_scaled)\n",
+    "    # (ii) Convergence tracking (assuming the best Y is to be maximized)\n",
+    "    convergence_plot = create_convergence_plot(objective_function, iteration_input, \n",
+    "                                               time_conv, \n",
+    "                                               convergence, TOTAL_TIME)\n",
+    "\n",
+    "\n",
+    "    return convergence_plot\n",
+    "    # return radar_chart, convergence_plot\n",
+    "\n",
+    "\n",
+    "\n",
+    "\n",
+    "\n",
+    "\n",
+    "\n",
+    "\n",
+    "\n",
+    "def create_radar_chart(X_scaled):\n",
+    "    fig, ax = plt.subplots(figsize=(6, 6), subplot_kw=dict(polar=True))\n",
+    "    labels = [f'x{i+1}' for i in range(X_scaled.shape[1])]\n",
+    "    values = X_scaled.mean(dim=0).numpy()\n",
+    "    \n",
+    "    num_vars = len(labels)\n",
+    "    angles = np.linspace(0, 2 * np.pi, num_vars, endpoint=False).tolist()\n",
+    "    values = np.concatenate((values, [values[0]]))\n",
+    "    angles += angles[:1]\n",
+    "\n",
+    "    ax.fill(angles, values, color='green', alpha=0.25)\n",
+    "    ax.plot(angles, values, color='green', linewidth=2)\n",
+    "    ax.set_yticklabels([])\n",
+    "    ax.set_xticks(angles[:-1])\n",
+    "    # ax.set_xticklabels(labels)\n",
+    "    ax.set_xticklabels([f'{label}\\n({value:.2f})' for label, value in zip(labels, values[:-1])])  # Show values\n",
+    "    ax.set_title(\"Selected Design\", size=15, color='black', y=1.1)\n",
+    "    \n",
+    "    plt.close(fig)\n",
+    "    return fig\n",
+    "\n",
+    "\n",
+    "\n",
+    "\n",
+    "\n",
+    "\n",
+    "\n",
+    "def create_convergence_plot(objective_function, iteration_input, time_conv, convergence, TOTAL_TIME):\n",
+    "    fig, ax = plt.subplots()\n",
+    "    \n",
+    "    # Realtime optimization data\n",
+    "    ax.plot(time_conv, convergence, '^-', label='PFN-CBO (Realtime)' )\n",
+    "\n",
+    "    # Stored GP data\n",
+    "    if objective_function==\"CantileverBeam.png\":\n",
+    "        GP_TIME = torch.load('CantileverBeam_CEI_Avg_Time.pt')\n",
+    "        GP_OBJ = torch.load('CantileverBeam_CEI_Avg_Obj.pt')\n",
+    "        \n",
+    "    elif objective_function==\"CompressionSpring.png\":\n",
+    "        GP_TIME = torch.load('CompressionSpring_CEI_Avg_Time.pt')\n",
+    "        GP_OBJ = torch.load('CompressionSpring_CEI_Avg_Obj.pt')\n",
+    "\n",
+    "    elif objective_function==\"HeatExchanger.png\":\n",
+    "        GP_TIME = torch.load('HeatExchanger_CEI_Avg_Time.pt')\n",
+    "        GP_OBJ = torch.load('HeatExchanger_CEI_Avg_Obj.pt')\n",
+    "        \n",
+    "    elif objective_function==\"ThreeTruss.png\":\n",
+    "        GP_TIME = torch.load('ThreeTruss_CEI_Avg_Time.pt')\n",
+    "        GP_OBJ = torch.load('ThreeTruss_CEI_Avg_Obj.pt')\n",
+    "        \n",
+    "    elif objective_function==\"Reinforcement.png\":\n",
+    "        GP_TIME = torch.load('ReinforcedConcreteBeam_CEI_Avg_Time.pt')\n",
+    "        GP_OBJ = torch.load('ReinforcedConcreteBeam_CEI_Avg_Obj.pt')\n",
+    "        \n",
+    "    elif objective_function==\"PressureVessel.png\":\n",
+    "        GP_TIME = torch.load('PressureVessel_CEI_Avg_Time.pt')\n",
+    "        GP_OBJ = torch.load('PressureVessel_CEI_Avg_Obj.pt')\n",
+    "        \n",
+    "    elif objective_function==\"SpeedReducer.png\":\n",
+    "        GP_TIME = torch.load('SpeedReducer_CEI_Avg_Time.pt')\n",
+    "        GP_OBJ = torch.load('SpeedReducer_CEI_Avg_Obj.pt')\n",
+    "        \n",
+    "    elif objective_function==\"WeldedBeam.png\":\n",
+    "        GP_TIME = torch.load('WeldedBeam_CEI_Avg_Time.pt')\n",
+    "        GP_OBJ = torch.load('WeldedBeam_CEI_Avg_Obj.pt')  \n",
+    "\n",
+    "    elif objective_function==\"Car.png\":\n",
+    "        GP_TIME = torch.load('Car_CEI_Avg_Time.pt')\n",
+    "        GP_OBJ = torch.load('Car_CEI_Avg_Obj.pt')    \n",
+    "        \n",
+    "    # Plot GP data    \n",
+    "    ax.plot(GP_TIME[:iteration_input], GP_OBJ[:iteration_input], '^-', label='GP-CBO (Data)' )\n",
+    "\n",
+    "    \n",
+    "    ax.set_xlabel('Time (seconds)')\n",
+    "    ax.set_ylabel('Objective Value')\n",
+    "    ax.set_title('Convergence Plot for {t} iterations'.format(t=iteration_input))\n",
+    "    # ax.legend()\n",
+    "\n",
+    "    if objective_function==\"CantileverBeam.png\":\n",
+    "        ax.axhline(y=50000, color='red', linestyle='--', label='Optimal Value')\n",
+    "\n",
+    "    elif objective_function==\"CompressionSpring.png\":\n",
+    "        ax.axhline(y=0, color='red', linestyle='--', label='Optimal Value')\n",
+    "\n",
+    "    elif objective_function==\"HeatExchanger.png\":\n",
+    "        ax.axhline(y=4700, color='red', linestyle='--', label='Optimal Value')\n",
+    "        \n",
+    "    elif objective_function==\"ThreeTruss.png\":\n",
+    "        ax.axhline(y=262, color='red', linestyle='--', label='Optimal Value')\n",
+    "        \n",
+    "    elif objective_function==\"Reinforcement.png\":\n",
+    "        ax.axhline(y=355, color='red', linestyle='--', label='Optimal Value')\n",
+    "        \n",
+    "    elif objective_function==\"PressureVessel.png\":\n",
+    "        ax.axhline(y=5000, color='red', linestyle='--', label='Optimal Value')\n",
+    "        \n",
+    "    elif objective_function==\"SpeedReducer.png\":\n",
+    "        ax.axhline(y=2650, color='red', linestyle='--', label='Optimal Value')\n",
+    "        \n",
+    "    elif objective_function==\"WeldedBeam.png\":\n",
+    "        ax.axhline(y=6, color='red', linestyle='--', label='Optimal Value')  \n",
+    "\n",
+    "    elif objective_function==\"Car.png\":\n",
+    "        ax.axhline(y=25, color='red', linestyle='--', label='Optimal Value')  \n",
+    "\n",
+    "    \n",
+    "    ax.legend(loc='best')\n",
+    "    # ax.legend(loc='lower left')\n",
+    "        \n",
+    "\n",
+    "    # Add text to the top right corner of the plot\n",
+    "    if len(convergence) == 0:\n",
+    "        ax.text(0.5, 0.5, 'No Feasible Design Found', transform=ax.transAxes, fontsize=12,\n",
+    "                verticalalignment='top', horizontalalignment='right')\n",
+    "        \n",
+    "    \n",
+    "    plt.close(fig)\n",
+    "    return fig\n",
+    "\n",
+    "\n",
+    "\n",
+    "\n",
+    "\n",
+    "\n",
+    "# Define available objective functions\n",
+    "objective_functions = {\n",
+    "    # \"ThreeTruss.png\": {\"image\": \"ThreeTruss.png\", \n",
+    "    #                     \"function\": ThreeTruss, \n",
+    "    #                     \"scaling\": ThreeTruss_Scaling, \n",
+    "    #                     \"dim\": 2},\n",
+    "    \"CompressionSpring.png\": {\"image\": \"CompressionSpring.png\", \n",
+    "                               \"function\": CompressionSpring, \n",
+    "                               \"scaling\": CompressionSpring_Scaling, \n",
+    "                               \"dim\": 3},\n",
+    "    \"Reinforcement.png\": {\"image\": \"Reinforcement.png\", \"function\": ReinforcedConcreteBeam, \"scaling\": ReinforcedConcreteBeam_Scaling, \"dim\": 3},\n",
+    "    \"PressureVessel.png\": {\"image\": \"PressureVessel.png\", \"function\": PressureVessel, \"scaling\": PressureVessel_Scaling, \"dim\": 4},\n",
+    "    \"SpeedReducer.png\": {\"image\": \"SpeedReducer.png\", \"function\": SpeedReducer, \"scaling\": SpeedReducer_Scaling, \"dim\": 7},\n",
+    "    \"WeldedBeam.png\": {\"image\": \"WeldedBeam.png\", \"function\": WeldedBeam, \"scaling\": WeldedBeam_Scaling, \"dim\": 4},\n",
+    "    \"HeatExchanger.png\": {\"image\": \"HeatExchanger.png\", \"function\": HeatExchanger, \"scaling\": HeatExchanger_Scaling, \"dim\": 8},\n",
+    "    \"CantileverBeam.png\": {\"image\": \"CantileverBeam.png\", \"function\": CantileverBeam, \"scaling\": CantileverBeam_Scaling, \"dim\": 10},\n",
+    "    \"Car.png\": {\"image\": \"Car.png\", \"function\": Car, \"scaling\": Car_Scaling, \"dim\": 11},\n",
+    "}\n",
+    "\n",
+    "\n",
+    "\n",
+    "\n",
+    "\n",
+    "\n",
+    "\n",
+    "\n",
+    "\n",
+    "\n",
+    "\n",
+    "\n",
+    "\n",
+    "\n",
+    "\n",
+    "\n",
+    "\n",
+    "\n",
+    "\n",
+    "\n",
+    "\n",
+    "\n",
+    "\n",
+    "\n",
+    "# Extract just the image paths for the gallery\n",
+    "image_paths = [key for key in objective_functions]\n",
+    "\n",
+    "\n",
+    "def submit_action(objective_function_choices, iteration_input):\n",
+    "    # print(iteration_input)\n",
+    "    # print(len(objective_function_choices))\n",
+    "    # print(objective_functions[objective_function_choices]['function'])\n",
+    "    if len(objective_function_choices)>0:\n",
+    "        selected_function = objective_functions[objective_function_choices]['function']\n",
+    "        return  optimize(objective_function_choices, iteration_input)\n",
+    "    return None\n",
+    "\n",
+    "# Function to clear the output\n",
+    "def clear_output():\n",
+    "    # print(gallery.selected_index)\n",
+    "    \n",
+    "    return gr.update(value=[], selected=None),  None, 15, gr.Markdown(\"\"), 'Test_formulation_default.png'\n",
+    "\n",
+    "def reset_gallery():\n",
+    "    return gr.update(value=image_paths)\n",
+    "\n",
+    "\n",
+    "with gr.Blocks() as demo:\n",
+    "    # Centered Title and Description using gr.HTML\n",
+    "    gr.HTML(\n",
+    "        \"\"\"\n",
+    "        <div style=\"text-align: center;\">\n",
+    "            <h1>Pre-trained Transformer for Constrained Bayesian Optimization</h1>\n",
+    "            <h4>Paper: <a href=\"https://arxiv.org/abs/2404.04495\">\n",
+    "            Fast and Accurate Bayesian Optimization with Pre-trained Transformers for Constrained Engineering Problems</a> \n",
+    "            </h4>\n",
+    "\n",
+    "            <p style=\"text-align: left;\">This is a demo for Bayesian Optimization using PFN (Prior-Data Fitted Networks). \n",
+    "            Select your objective function by clicking on one of the check boxes below, then enter the iteration number to run the optimization process. \n",
+    "            The results will be visualized in the radar chart and convergence plot.</p>\n",
+    "            \n",
+    "            \n",
+    "            \n",
+    "\n",
+    "        </div>\n",
+    "        \"\"\"\n",
+    "    )\n",
+    "\n",
+    "    \n",
+    "    with gr.Row():\n",
+    "        \n",
+    "        \n",
+    "        with gr.Column(variant='compact'):\n",
+    "            # gr.Markdown(\"# Inputs: \")\n",
+    "            \n",
+    "            with gr.Row():\n",
+    "                gr.Markdown(\"## Select a problem (objective): \")\n",
+    "                img_key = gr.Markdown(value=\"\", visible=False)\n",
+    "            \n",
+    "            gallery = gr.Gallery(value=image_paths, label=\"Objective Functions\", \n",
+    "                                 # height = 450, \n",
+    "                                 object_fit='contain',\n",
+    "                                 columns=3, rows=3, elem_id=\"gallery\")\n",
+    "            \n",
+    "            gr.Markdown(\"## Enter iteration Number: \")\n",
+    "            iteration_input = gr.Slider(label=\"Iterations:\", minimum=15, maximum=50, step=1, value=15)\n",
+    "        \n",
+    "\n",
+    "            # Row for the Clear and Submit buttons\n",
+    "            with gr.Row():\n",
+    "                clear_button = gr.Button(\"Clear\")\n",
+    "                submit_button = gr.Button(\"Submit\", variant=\"primary\")\n",
+    "\n",
+    "        with gr.Column():\n",
+    "            # gr.Markdown(\"# Outputs: \")\n",
+    "            gr.Markdown(\"## Problem Formulation: \")\n",
+    "            formulation = gr.Image(value='Formulation_default.png', height=150)\n",
+    "            gr.Markdown(\"## Results: \")\n",
+    "            gr.Markdown(\"The graph will plot the best observed data v.s. the time for the algorithm to run up until the iteration. The PFN-CBO shows the result of the realtime optimization running in the backend while the GP-CBO shows the stored data from our previous experiments since running GP-CBO will take longer time.\")\n",
+    "            convergence_plot = gr.Plot(label=\"Convergence Plot\")\n",
+    "\n",
+    "\n",
+    "\n",
+    "    def handle_select(evt: gr.SelectData):\n",
+    "        selected_image = evt.value\n",
+    "        key = evt.value['image']['orig_name']\n",
+    "        formulation = 'Test_formulation.png'\n",
+    "        print('here')\n",
+    "        print(key)\n",
+    "\n",
+    "        return key, formulation\n",
+    "        \n",
+    "    gallery.select(fn=handle_select, inputs=None, outputs=[img_key, formulation])\n",
+    "\n",
+    "\n",
+    "    \n",
+    "    submit_button.click(\n",
+    "        submit_action,\n",
+    "        inputs=[img_key, iteration_input],\n",
+    "        # outputs= [radar_plot, convergence_plot],\n",
+    "        outputs= convergence_plot,\n",
+    "        \n",
+    "        # progress=True  # Enable progress tracking\n",
+    "        \n",
+    "    )\n",
+    "\n",
+    "    clear_button.click(\n",
+    "        clear_output,\n",
+    "        inputs=None,\n",
+    "        outputs=[gallery, convergence_plot, iteration_input, img_key, formulation]\n",
+    "    ).then(\n",
+    "        # Step 2: Reset the gallery to the original list\n",
+    "        reset_gallery,\n",
+    "        inputs=None,\n",
+    "        outputs=gallery\n",
+    "    )\n",
+    "\n",
+    "    \n",
+    "\n",
+    "demo.launch()\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "776c7ab2-96a1-4e22-9b4b-daf69960e3c4",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "9d2c7c58-43b1-4e5b-9135-17683dac1788",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "5d33a24c-818c-4023-bbbd-495f992a9d1a",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "764d0258-ec88-41d5-b5b5-e0bcb39ff313",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "d2c35245-543c-4b82-8d12-04f3dda1468b",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "e3663adc-3e95-418b-bf50-0a372615cdd6",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "30886262-bd87-4760-a585-7872e071663f",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3 (ipykernel)",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.10.14"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

Gradio_test.ipynb

@@ -0,0 +1,569 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "id": "4453c5ad-ec87-42e0-a6d5-e3fd3593aec2",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Running on local URL:  http://127.0.0.1:7891\n",
+      "Running on public URL: https://f714b6f956fb581264.gradio.live\n",
+      "\n",
+      "This share link expires in 72 hours. For free permanent hosting and GPU upgrades, run `gradio deploy` from Terminal to deploy to Spaces (https://huggingface.co/spaces)\n"
+     ]
+    },
+    {
+     "data": {
+      "text/html": [
+       "<div><iframe src=\"https://f714b6f956fb581264.gradio.live\" width=\"100%\" height=\"500\" allow=\"autoplay; camera; microphone; clipboard-read; clipboard-write;\" frameborder=\"0\" allowfullscreen></iframe></div>"
+      ],
+      "text/plain": [
+       "<IPython.core.display.HTML object>"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "text/plain": []
+     },
+     "execution_count": 1,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "import gradio as gr\n",
+    "import torch\n",
+    "import numpy as np\n",
+    "import matplotlib.pyplot as plt\n",
+    "from test_functions.Ackley10D import *\n",
+    "from test_functions.Ackley2D import *\n",
+    "from test_functions.Ackley6D import *\n",
+    "from test_functions.HeatExchanger import *\n",
+    "from test_functions.CantileverBeam import *\n",
+    "from test_functions.Car import *\n",
+    "from test_functions.CompressionSpring import *\n",
+    "from test_functions.GKXWC1 import *\n",
+    "from test_functions.GKXWC2 import *\n",
+    "from test_functions.HeatExchanger import *\n",
+    "from test_functions.JLH1 import *\n",
+    "from test_functions.JLH2 import *\n",
+    "from test_functions.KeaneBump import *\n",
+    "from test_functions.GKXWC1 import *\n",
+    "from test_functions.GKXWC2 import *\n",
+    "from test_functions.PressureVessel import *\n",
+    "from test_functions.ReinforcedConcreteBeam import *\n",
+    "from test_functions.SpeedReducer import *\n",
+    "from test_functions.ThreeTruss import *\n",
+    "from test_functions.WeldedBeam import *\n",
+    "# Import other objective functions as needed\n",
+    "import time\n",
+    "\n",
+    "from Rosen_PFN4BO import *\n",
+    "\n",
+    "def optimize(objective_function, iteration_input):\n",
+    "\n",
+    "    # Variable setup\n",
+    "    Current_BEST = -1e10   # Some arbitrary very small number\n",
+    "    Prev_BEST = -1e10\n",
+    "\n",
+    "    # Initial random samples\n",
+    "    # print(objective_functions)\n",
+    "    trained_X = torch.rand(20, objective_functions[objective_function]['dim'])\n",
+    "\n",
+    "    # Scale it to the domain of interest using the selected function\n",
+    "    # print(objective_function)\n",
+    "    X_Scaled = objective_functions[objective_function]['scaling'](trained_X)\n",
+    "\n",
+    "    # Get the constraints and objective\n",
+    "    trained_gx, trained_Y = objective_functions[objective_function]['function'](X_Scaled)\n",
+    "\n",
+    "    # Convergence list to store best values\n",
+    "    convergence = []\n",
+    "\n",
+    "    START_TIME = time.time()\n",
+    "\n",
+    "    # Optimization Loop\n",
+    "    for ii in range(iteration_input):  # Example with 100 iterations\n",
+    "\n",
+    "        # (0) Get the updated data for this iteration\n",
+    "        X_scaled = objective_functions[objective_function]['scaling'](trained_X)\n",
+    "        trained_gx, trained_Y = objective_functions[objective_function]['function'](X_scaled)\n",
+    "\n",
+    "        # (1) Randomly sample Xpen \n",
+    "        X_pen = torch.rand(1000,trained_X.shape[1])\n",
+    "\n",
+    "        # (2) PFN inference phase with EI\n",
+    "        default_model = 'final_models/Cyril_500features_800epoch_cpu.pt'\n",
+    "        \n",
+    "        ei, p_feas = Rosen_PFN_Parallel(default_model,\n",
+    "                                           trained_X, \n",
+    "                                           trained_Y, \n",
+    "                                           trained_gx,\n",
+    "                                           X_pen,\n",
+    "                                           'power',\n",
+    "                                           'ei'\n",
+    "                                          )\n",
+    "\n",
+    "        # Calculating CEI\n",
+    "        CEI = ei\n",
+    "        for jj in range(p_feas.shape[1]):\n",
+    "            CEI = CEI*p_feas[:,jj]\n",
+    "\n",
+    "        # (4) Get the next search value\n",
+    "        rec_idx = torch.argmax(CEI)\n",
+    "        best_candidate = X_pen[rec_idx,:].unsqueeze(0)\n",
+    "\n",
+    "        # (5) Append the next search point\n",
+    "        trained_X = torch.cat([trained_X, best_candidate])\n",
+    "\n",
+    "\n",
+    "        ################################################################################\n",
+    "        # This is just for visualizing the best value. \n",
+    "        # This section can be remove for pure optimization purpose\n",
+    "        Current_X = objective_functions[objective_function]['scaling'](trained_X)\n",
+    "        Current_GX, Current_Y = objective_functions[objective_function]['function'](Current_X)\n",
+    "        if ((Current_GX<=0).all(dim=1)).any():\n",
+    "            Current_BEST = torch.max(Current_Y[(Current_GX<=0).all(dim=1)])\n",
+    "        else:\n",
+    "            Current_BEST = Prev_BEST\n",
+    "        ################################################################################\n",
+    "        \n",
+    "        # (ii) Convergence tracking (assuming the best Y is to be maximized)\n",
+    "        if Current_BEST != -1e10:\n",
+    "            convergence.append(Current_BEST.abs())\n",
+    "\n",
+    "    # Timing\n",
+    "    END_TIME = time.time()\n",
+    "    TOTAL_TIME = END_TIME - START_TIME\n",
+    "    \n",
+    "    # Website visualization\n",
+    "    # (i) Radar chart for trained_X\n",
+    "    radar_chart = create_radar_chart(X_scaled)\n",
+    "    # (ii) Convergence tracking (assuming the best Y is to be maximized)\n",
+    "    convergence_plot = create_convergence_plot(convergence, TOTAL_TIME)\n",
+    "    \n",
+    "    return radar_chart, convergence_plot\n",
+    "\n",
+    "def create_radar_chart(X_scaled):\n",
+    "    fig, ax = plt.subplots(figsize=(6, 6), subplot_kw=dict(polar=True))\n",
+    "    labels = [f'x{i+1}' for i in range(X_scaled.shape[1])]\n",
+    "    values = X_scaled.mean(dim=0).numpy()\n",
+    "    \n",
+    "    num_vars = len(labels)\n",
+    "    angles = np.linspace(0, 2 * np.pi, num_vars, endpoint=False).tolist()\n",
+    "    values = np.concatenate((values, [values[0]]))\n",
+    "    angles += angles[:1]\n",
+    "\n",
+    "    ax.fill(angles, values, color='green', alpha=0.25)\n",
+    "    ax.plot(angles, values, color='green', linewidth=2)\n",
+    "    ax.set_yticklabels([])\n",
+    "    ax.set_xticks(angles[:-1])\n",
+    "    # ax.set_xticklabels(labels)\n",
+    "    ax.set_xticklabels([f'{label}\\n({value:.2f})' for label, value in zip(labels, values[:-1])])  # Show values\n",
+    "    ax.set_title(\"Selected Design\", size=15, color='black', y=1.1)\n",
+    "    \n",
+    "    plt.close(fig)\n",
+    "    return fig\n",
+    "\n",
+    "def create_convergence_plot(convergence, TOTAL_TIME):\n",
+    "    fig, ax = plt.subplots()\n",
+    "    # print(len(convergence))\n",
+    "    ax.plot(convergence, label='Best Objective Value')\n",
+    "    ax.set_xlabel('Iteration')\n",
+    "    ax.set_ylabel('Objective Value')\n",
+    "    ax.set_title('Convergence Plot (Opt Runtime: {t} sec)'.format(t=round(TOTAL_TIME, 2)))\n",
+    "    ax.legend()\n",
+    "\n",
+    "    # Add text to the top right corner of the plot\n",
+    "    if len(convergence) == 0:\n",
+    "        ax.text(0.5, 0.5, 'No Feasible Design Found', transform=ax.transAxes, fontsize=12,\n",
+    "                verticalalignment='top', horizontalalignment='right')\n",
+    "    \n",
+    "    plt.close(fig)\n",
+    "    return fig\n",
+    "\n",
+    "# Define available objective functions\n",
+    "objective_functions = {\n",
+    "    \"Ackley2D\": {\"function\": Ackley2D, \"scaling\": Ackley2D_Scaling, \"dim\": 2},\n",
+    "    \"Ackley6D\": {\"function\": Ackley6D, \"scaling\": Ackley6D_Scaling, \"dim\": 6},\n",
+    "    \"Ackley10D\": {\"function\": Ackley10D, \"scaling\": Ackley10D_Scaling, \"dim\": 10},\n",
+    "    \"GKXWC1\": {\"function\": GKXWC1, \"scaling\": GKXWC1_Scaling, \"dim\": 2},\n",
+    "    \"GKXWC2\": {\"function\": GKXWC2, \"scaling\": GKXWC2_Scaling, \"dim\": 2},\n",
+    "    \"JLH1\": {\"function\": JLH1, \"scaling\": JLH1_Scaling, \"dim\": 2},\n",
+    "    \"JLH2\": {\"function\": JLH2, \"scaling\": JLH2_Scaling, \"dim\": 2},\n",
+    "    \"Keane Bump\": {\"function\": KeaneBump, \"scaling\": KeaneBump_Scaling, \"dim\": 18},\n",
+    "    \"Three Truss\": {\"function\": ThreeTruss, \"scaling\": ThreeTruss_Scaling, \"dim\": 2},\n",
+    "    \"Compression Spring\": {\"function\": CompressionSpring, \"scaling\": CompressionSpring_Scaling, \"dim\": 3},\n",
+    "    \"Reinforced Concrete Beam\": {\"function\": ReinforcedConcreteBeam, \"scaling\": ReinforcedConcreteBeam_Scaling, \"dim\": 3},\n",
+    "    \"Pressure Vessel\": {\"function\": PressureVessel, \"scaling\": PressureVessel_Scaling, \"dim\": 4},\n",
+    "    \"Speed Reducer\": {\"function\": SpeedReducer, \"scaling\": SpeedReducer_Scaling, \"dim\": 4},\n",
+    "    \"Welded Beam\": {\"function\": WeldedBeam, \"scaling\": WeldedBeam_Scaling, \"dim\": 4},\n",
+    "    \"Heat Exchanger\": {\"function\": HeatExchanger, \"scaling\": HeatExchanger_Scaling, \"dim\": 8},\n",
+    "    \"Cantilever Beam\": {\"function\": CantileverBeam, \"scaling\": CantileverBeam_Scaling, \"dim\": 10},\n",
+    "    \"Car\": {\"function\": Car, \"scaling\": Car_Scaling, \"dim\": 11},\n",
+    "    \n",
+    "    # Add more functions here\n",
+    "}\n",
+    "\n",
+    "\n",
+    "\n",
+    "\n",
+    "\n",
+    "with gr.Blocks(theme=gr.themes.Default()) as demo:\n",
+    "    # Centered Title and Description using gr.HTML\n",
+    "    gr.HTML(\n",
+    "        \"\"\"\n",
+    "        <div style=\"text-align: center;\">\n",
+    "            <h1>Pre-trained Transformer for Constrained Bayesian Optimization</h1>\n",
+    "            <p>This is a demo for Bayesian Optimization using PFN (Prior-Data Fitted Networks). \n",
+    "            Select your objective function by clicking on one of the check boxes below, then enter the iteration number to run the optimization process. \n",
+    "            The results will be visualized in the radar chart and convergence plot.</p>\n",
+    "            <img src=\"https://github.com/rosenyu304/BOEngineeringBenchmark/blob/main/Icons.png?raw=true\" \n",
+    "                 alt=\"Example Image\" \n",
+    "                 style=\"width: 800px; height: auto; margin-top: 20px; display: block; margin-left: auto; margin-right: auto;\">\n",
+    "\n",
+    "        </div>\n",
+    "        \"\"\"\n",
+    "    )\n",
+    "\n",
+    "    selected_objective = gr.State(None)  # To store the selected objective function\n",
+    "\n",
+    "\n",
+    "\n",
+    "\n",
+    "\n",
+    "    \n",
+    "    with gr.Row():\n",
+    "        \n",
+    "        objective_checkbox_group = gr.CheckboxGroup(\n",
+    "            choices=[\"JLH1\", \"JLH2\", \"GKXWC1\", \"GKXWC2\", \"Ackley2D\", \"Ackley6D\", \"Ackley10D\", \"Keane Bump\", \"Three Truss\", \"Reinforced Concrete Beam\", \"Pressure Vessel\", \"Welded Beam\", \"Speed Reducer\", \"Car\"],\n",
+    "            label=\"Select the design problem:\"\n",
+    "        )\n",
+    "    with gr.Row():\n",
+    "        iteration_input = gr.Number(label=\"Enter Iteration Number:\", value=10)\n",
+    "            \n",
+    "\n",
+    "    # Row for the Clear and Submit buttons\n",
+    "    with gr.Row():\n",
+    "        clear_button = gr.Button(\"Clear\")\n",
+    "        submit_button = gr.Button(\"Submit\", variant=\"primary\")\n",
+    "        \n",
+    "    \n",
+    "    with gr.Row():\n",
+    "        with gr.Column():\n",
+    "            radar_plot = gr.Plot(label=\"Resulting Design\")\n",
+    "        with gr.Column():\n",
+    "            convergence_plot = gr.Plot(label=\"Convergence Plot\")\n",
+    "\n",
+    "\n",
+    "\n",
+    "    # Define actions for buttons\n",
+    "    def clear_action():\n",
+    "        return None, None, None\n",
+    "\n",
+    "    def submit_action(objective_function_choices, iteration_input):\n",
+    "        # Handle the case where multiple choices are selected\n",
+    "        if len(objective_function_choices) > 0:\n",
+    "            selected_function = objective_function_choices[0]  # Assuming using the first selected function\n",
+    "            return optimize(selected_function, iteration_input)\n",
+    "        return None, None\n",
+    "\n",
+    "    # Button click actions\n",
+    "    clear_button.click(clear_action, outputs=[objective_checkbox_group, radar_plot, convergence_plot])\n",
+    "    submit_button.click(\n",
+    "        submit_action, \n",
+    "        inputs=[objective_checkbox_group, iteration_input], \n",
+    "        outputs=[radar_plot, convergence_plot]\n",
+    "    )\n",
+    "\n",
+    "demo.launch(share=True)\n",
+    "\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "352d0291-93b4-43eb-b683-3d48776dc670",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "92ecbbe6-dea6-4e7f-aae1-f0d442dbda3b",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "ba69b5f9-c52c-4c23-8645-c81c27f7a815",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "id": "05789fba-2099-46b7-8675-64b7969427a1",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Running on local URL:  http://127.0.0.1:7899\n",
+      "\n",
+      "To create a public link, set `share=True` in `launch()`.\n"
+     ]
+    },
+    {
+     "data": {
+      "text/html": [
+       "<div><iframe src=\"http://127.0.0.1:7899/\" width=\"100%\" height=\"500\" allow=\"autoplay; camera; microphone; clipboard-read; clipboard-write;\" frameborder=\"0\" allowfullscreen></iframe></div>"
+      ],
+      "text/plain": [
+       "<IPython.core.display.HTML object>"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    }
+   ],
+   "source": [
+    "import gradio as gr\n",
+    "\n",
+    "def calculator(num1, operation, num2):\n",
+    "    if operation == \"add\":\n",
+    "        return num1 + num2\n",
+    "    elif operation == \"subtract\":\n",
+    "        return num1 - num2\n",
+    "    elif operation == \"multiply\":\n",
+    "        return num1 * num2\n",
+    "    elif operation == \"divide\":\n",
+    "        return num1 / num2\n",
+    "\n",
+    "with gr.Blocks() as demo:\n",
+    "    with gr.Row():\n",
+    "        with gr.Column():\n",
+    "            num_1 = gr.Number(value=4)\n",
+    "            operation = gr.Radio([\"add\", \"subtract\", \"multiply\", \"divide\"])\n",
+    "            num_2 = gr.Number(value=0)\n",
+    "            submit_btn = gr.Button(value=\"Calculate\")\n",
+    "        with gr.Column():\n",
+    "            result = gr.Number()\n",
+    "\n",
+    "    submit_btn.click(\n",
+    "        calculator, inputs=[num_1, operation, num_2], outputs=[result], api_name=False\n",
+    "    )\n",
+    "    examples = gr.Examples(\n",
+    "        examples=[\n",
+    "            [5, \"add\", 3],\n",
+    "            [4, \"divide\", 2],\n",
+    "            [-4, \"multiply\", 2.5],\n",
+    "            [0, \"subtract\", 1.2],\n",
+    "        ],\n",
+    "        inputs=[num_1, operation, num_2],\n",
+    "    )\n",
+    "\n",
+    "if __name__ == \"__main__\":\n",
+    "    demo.launch(show_api=False)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "a4bf709a-ff0a-4aac-a4b4-fd98cd5948bb",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "679f7647-ca68-46f9-a1da-81d6c96267c9",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "ea40bfac-e090-4cd5-9caa-99b06db3ea8d",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 50,
+   "id": "928ac99a-af8f-401c-8c0b-ef83cfef5ba9",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Running on local URL:  http://127.0.0.1:7890\n",
+      "\n",
+      "To create a public link, set `share=True` in `launch()`.\n"
+     ]
+    },
+    {
+     "data": {
+      "text/html": [
+       "<div><iframe src=\"http://127.0.0.1:7890/\" width=\"100%\" height=\"500\" allow=\"autoplay; camera; microphone; clipboard-read; clipboard-write;\" frameborder=\"0\" allowfullscreen></iframe></div>"
+      ],
+      "text/plain": [
+       "<IPython.core.display.HTML object>"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    }
+   ],
+   "source": [
+    "import gradio as gr\n",
+    "\n",
+    "def calculator(num1, operation, num2):\n",
+    "    if operation == \"add\":\n",
+    "        return num1 + num2\n",
+    "    elif operation == \"subtract\":\n",
+    "        return num1 - num2\n",
+    "    elif operation == \"multiply\":\n",
+    "        return num1 * num2\n",
+    "    elif operation == \"divide\":\n",
+    "        return num1 / num2\n",
+    "\n",
+    "with gr.Blocks() as demo:\n",
+    "    with gr.Row():\n",
+    "        with gr.Column():\n",
+    "            num_1 = gr.Number(value=4)\n",
+    "            operation = gr.Radio([\"add\", \"subtract\", \"multiply\", \"divide\"])\n",
+    "            num_2 = gr.Number(value=0)\n",
+    "            submit_btn = gr.Button(value=\"Calculate\")\n",
+    "        with gr.Column():\n",
+    "            result = gr.Number()\n",
+    "\n",
+    "    submit_btn.click(\n",
+    "        calculator, inputs=[num_1, operation, num_2], outputs=[result], api_name=False\n",
+    "    )\n",
+    "    examples = gr.Examples(\n",
+    "        examples=[\n",
+    "            [5, \"add\", 3],\n",
+    "            [4, \"divide\", 2],\n",
+    "            [-4, \"multiply\", 2.5],\n",
+    "            [0, \"subtract\", 1.2],\n",
+    "        ],\n",
+    "        inputs=[num_1, operation, num_2],\n",
+    "    )\n",
+    "\n",
+    "if __name__ == \"__main__\":\n",
+    "    demo.launch(show_api=False)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 36,
+   "id": "09a251df-4076-4925-8799-9a2a59cb8246",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# import gradio as gr\n",
+    "\n",
+    "# def greet(selected_options):\n",
+    "#     return f\"You selected: {', '.join(selected_options)}\"\n",
+    "\n",
+    "# with gr.Blocks() as demo:\n",
+    "#     with gr.Row():\n",
+    "#         checkbox_group = gr.CheckboxGroup(\n",
+    "#             choices=[\"Option 1\", \"Option 2\"],\n",
+    "#             label=\"Select your options\",\n",
+    "#             elem_id=\"custom_checkbox_group\"\n",
+    "#         )\n",
+    "#         output = gr.Textbox(label=\"Output\")\n",
+    "    \n",
+    "#     checkbox_group.change(greet, checkbox_group, output)\n",
+    "\n",
+    "#     gr.HTML(\n",
+    "#         f\"\"\"\n",
+    "#         <style>\n",
+    "#         #custom_checkbox_group label {\n",
+    "#             display: block;\n",
+    "#             width: 200pt;\n",
+    "#             height: 200pt;\n",
+    "#             border: 1px solid #ccc;\n",
+    "#             margin-bottom: 10pt;\n",
+    "#             padding: 10pt;\n",
+    "#             box-sizing: border-box;\n",
+    "#             position: relative;\n",
+    "#         }\n",
+    "#         #custom_checkbox_group label input {\n",
+    "#             position: absolute;\n",
+    "#             top: 10pt;\n",
+    "#             left: 10pt;\n",
+    "#         }\n",
+    "#         #custom_checkbox_group label span {\n",
+    "#             position: absolute;\n",
+    "#             top: 10pt;\n",
+    "#             left: 40pt; /* Adjust this value to control the distance between the checkbox and the label */\n",
+    "#         }\n",
+    "#         #custom_checkbox_group label img {\n",
+    "#             position: absolute;\n",
+    "#             bottom: 10pt;\n",
+    "#             left: 10pt;\n",
+    "#             width: 180pt; /* Adjust the size of the image if needed */\n",
+    "#             height: auto;\n",
+    "#         }\n",
+    "#         </style>\n",
+    "#         <label>\n",
+    "#             <input type=\"checkbox\" />\n",
+    "#             <span>Option 1</span>\n",
+    "#             <img src=\"https://images.pexels.com/photos/1108099/pexels-photo-1108099.jpeg\" alt=\"Dog image\"/>\n",
+    "#         </label>\n",
+    "#         \"\"\"\n",
+    "#     )\n",
+    "\n",
+    "# demo.launch()\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "f52549d5-4be0-4672-be6d-df462957cb56",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3 (ipykernel)",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.10.14"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

HeatExchanger_CEI_Avg_Obj.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:6fc95574634750d3dc892076b26e55c6f79d4dbb128d5b65e6832e83783c89a8
+size 3432

HeatExchanger_CEI_Avg_Time.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:0ed8e01768b9cc8bf82c51f523c9ea46c4f3e7e3e9e6c8e04edb0d615032f1e9
+size 3500

PressureVessel_CEI_Avg_Obj.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f7c81ea242bdcb45cb644cd5f18b941ff8ebbcbbb81b9965eea251c01f9f6c78
+size 3628

PressureVessel_CEI_Avg_Time.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:eb204723d0523baebbfda6e4f1fdbc7506c66bfc0ed0cbc7ec5ea485451660a7
+size 3504

ReinforcedConcreteBeam_CEI_Avg_Obj.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:5d2e3212a28eb9cb59212d876c1ddae2f1b37950974eed01683c7d4180206c7e
+size 3532

ReinforcedConcreteBeam_CEI_Avg_Time.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:aae5672638e75081965450635cc15310f90ee167c3264399bee07afc2ad3a58d
+size 3472

Rosen_PFN4BO.py

@@ -0,0 +1,442 @@
+import contextlib
+import torch
+import scipy
+import math
+from sklearn.preprocessing import power_transform, PowerTransformer, StandardScaler
+
+from torchvision.transforms.functional import to_tensor
+from pfns4bo import transformer
+from pfns4bo import bar_distribution
+
+import torch
+import numpy as np
+
+import pfns4bo
+from pfns4bo.scripts.acquisition_functions import TransformerBOMethod
+
+
+import warnings
+warnings.filterwarnings('ignore')
+
+device = torch.device("cpu")
+# device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+dtype = torch.float32
+
+
+from sklearn.utils import resample
+
+@torch.enable_grad()
+def Rosen_PFN(model_name,
+              trained_X, 
+              trained_Y, 
+              X_pen,
+              trasform_type,
+              what_do_you_want
+             ):
+    
+    PFN = TransformerBOMethod(torch.load(model_name).requires_grad_(False), device=device)
+
+    # X_pen.requires_grad_(True)
+
+    # with torch.no_grad():
+
+
+    dim = trained_X.shape[1]
+
+    x_given = trained_X
+    x_eval = X_pen
+    x_predict = torch.cat([x_given, x_eval], dim=0)
+    x_full_feed = torch.cat([x_given, x_given, x_eval], dim=0).unsqueeze(1)
+
+
+
+    if trasform_type== 'std':
+        pt = StandardScaler()
+        pt.fit(trained_Y)
+        PT_trained_Y = pt.transform(trained_Y)
+        trained_Y = to_tensor(PT_trained_Y).to(torch.float32).reshape(trained_Y.shape)
+    elif trasform_type== 'power':
+        pt = PowerTransformer(method="yeo-johnson")
+        pt.fit(trained_Y.detach().numpy())
+        # PT_trained_Y = pt.transform(trained_Y.detach().numpy())
+        # trained_Y = to_tensor(PT_trained_Y).to(torch.float32).reshape(trained_Y.shape)
+        # print(trained_Y.shape)
+
+        # print(trained_Y)
+        trained_Y, _ = general_power_transform(trained_Y, 
+                                            trained_Y, 
+                                            .0, 
+                                            less_safe=False) #.squeeze(1)
+        # print(trained_Y.shape)
+        # .squeeze(1)
+
+
+        # y_given = general_power_transform(y_given.unsqueeze(1), 
+        #                                   y_given.unsqueeze(1), 
+        #                                   .0, 
+        #                                   less_safe=False).squeeze(1)
+
+    y_given = trained_Y
+
+    y_given = y_given.reshape(-1)
+    y_full_feed = y_given.unsqueeze(1)
+
+    criterion: bar_distribution.BarDistribution = PFN.model.criterion
+
+    style = None
+    logits = PFN.model(
+                (style,
+                 x_full_feed.repeat_interleave(dim=1, repeats=y_full_feed.shape[1]),
+                 y_full_feed.repeat(1,x_full_feed.shape[1])),
+                single_eval_pos=len(x_given)
+            )
+
+    # logits = logits.softmax(-1).log_()
+    logits = logits.softmax(-1).log()
+
+    logits_given = logits[:len(x_given)]
+    logits_eval = logits[len(x_given):]
+
+    best_f = torch.max(y_given)
+
+    if what_do_you_want == 'mean':
+        output = criterion.mean(logits_eval)
+
+        
+        if trasform_type== 'std' or trasform_type== 'power':
+            
+            if pt.standardize:
+                XX = output.clone()
+                scale = torch.from_numpy(pt._scaler.scale_)
+                std_mean = torch.from_numpy(pt._scaler.mean_)
+                XX = torch_std_inverse_transform(XX, scale, std_mean)
+            
+            for i, lmbda in enumerate(pt.lambdas_):
+                with np.errstate(invalid="ignore"):  # hide NaN warnings
+                    XX = torch_power_inverse_transform(XX, lmbda)
+                    # print(XX)
+                    return XX
+            
+
+
+
+            
+            # output = pt.inverse_transform(output)
+            # output = torch.from_numpy(output)
+            
+            
+    elif what_do_you_want == 'ei':
+        output = criterion.ei(logits_eval, best_f)
+        
+    elif what_do_you_want == 'ucb':
+        acq_function = criterion.ucb
+        ucb_rest_prob = .05
+        if ucb_rest_prob is not None:
+            acq_function = lambda *args: criterion.ucb(*args, rest_prob=ucb_rest_prob)
+        output = acq_ensembling(acq_function(logits_eval, best_f))
+        
+    elif what_do_you_want == 'variance':
+        output = criterion.variance(logits_eval)
+        
+    elif what_do_you_want == 'mode':
+        output = criterion.mode(logits_eval)
+        
+    elif what_do_you_want == 'ts':
+        mn = criterion.mean(logits_eval)
+        
+        
+        if trasform_type== 'std' or trasform_type== 'power':
+            
+            if pt.standardize:
+                XX = mn.clone()
+                scale = torch.from_numpy(pt._scaler.scale_)
+                std_mean = torch.from_numpy(pt._scaler.mean_)
+                XX = torch_std_inverse_transform(XX, scale, std_mean)
+            
+            for i, lmbda in enumerate(pt.lambdas_):
+                with np.errstate(invalid="ignore"):  # hide NaN warnings
+                    XX = torch_power_inverse_transform(XX, lmbda)
+
+        var = criterion.variance(logits_eval)
+
+        return XX, var
+        
+    return output
+
+
+
+
+
+
+
+
+
+
+
+def Rosen_PFN_Parallel(model_name,
+                       trained_X, 
+                       trained_Y, 
+                       GX,
+                       X_pen,
+                       trasform_type,
+                       what_do_you_want
+                      ):
+    
+    PFN = TransformerBOMethod(torch.load(model_name), device=device)
+
+    with torch.no_grad():
+
+
+        dim = trained_X.shape[1]
+
+        x_given = trained_X
+        x_eval = X_pen
+        x_predict = torch.cat([x_given, x_eval], dim=0)
+        x_full_feed = torch.cat([x_given, x_given, x_eval], dim=0).unsqueeze(1)
+        
+        
+        
+        y_given = trained_Y
+        y_given = y_given.reshape(-1)
+        
+        ######################################################################
+        # Objective Power Transform
+        y_given, pt_y = general_power_transform(y_given.unsqueeze(1), 
+                                          y_given.unsqueeze(1), 
+                                          .0, 
+                                          less_safe=False)
+        y_given = y_given.squeeze(1)
+        ######################################################################
+        
+
+        ######################################################################
+        # Constraints Power Transform
+        # Changes for Parallel:
+        GX = -GX
+        GX_t, pt_GX = general_power_transform(GX, GX, .0, less_safe=False)
+        G_thres, _ = general_power_transform(GX, 
+                                          torch.zeros((1, GX.shape[1])).to(GX.device), 
+                                          .0, 
+                                          less_safe=False)
+        GX = GX_t
+        ######################################################################
+        
+
+        
+        y_full_feed = y_given.unsqueeze(1)
+
+        criterion: bar_distribution.BarDistribution = PFN.model.criterion
+
+        style = None
+        logits = PFN.model(
+                    (style,
+                     x_full_feed.repeat_interleave(dim=1, repeats=y_full_feed.shape[1]+GX.shape[1]),
+                     torch.cat([y_full_feed, GX], dim=1).unsqueeze(2) ),
+                    single_eval_pos=len(x_given)
+                )
+        
+        logits = logits.softmax(-1).log_()
+
+        logits_given = logits[:len(x_given)]
+        logits_eval = logits[len(x_given):]
+        
+        best_f = torch.max(y_given)
+        
+        objective_given = logits_given[:,0,:].unsqueeze(1)
+        objective_eval  = logits_eval[:,0,:].unsqueeze(1)
+        constraint_given = logits_given[:,1:,:]
+        constraint_eval  = logits_eval[:,1:,:]
+        
+        
+        
+        if what_do_you_want == 'mean':
+            obj_output = criterion.mean(objective_eval)
+            con_output = criterion.mean(constraint_eval)
+            
+        elif what_do_you_want == 'ei':
+            # Changes for CEI
+
+            # Objective
+            tau = torch.max(y_given)
+            objective_acq_value = acq_ensembling(criterion.ei(objective_eval, tau))
+
+            # Constraints
+            constraints_acq_value = acq_ensembling(criterion.pi(constraint_eval[:,0,:].unsqueeze(1), G_thres[0, 0].item()))
+            constraints_acq_value = constraints_acq_value.unsqueeze(1)
+
+
+            for jj in range(1,constraint_eval.shape[1]):
+                next_constraints_acq_value = acq_ensembling(criterion.pi(constraint_eval[:,jj,:].unsqueeze(1), G_thres[0, jj].item()))
+                next_constraints_acq_value = next_constraints_acq_value.unsqueeze(1)
+                constraints_acq_value = torch.cat([constraints_acq_value,next_constraints_acq_value], dim=1)
+
+            return objective_acq_value, constraints_acq_value
+
+            
+        elif what_do_you_want == 'variance':
+            output = criterion.variance(logits_eval)
+        elif what_do_you_want == 'mode':
+            output = criterion.mode(logits_eval)
+        elif what_do_you_want == 'cts':
+            obj_mnn = criterion.mean(objective_eval)
+            obj_mnn = pt_y.inverse_transform(obj_mnn)
+            obj_mnn = torch.from_numpy(obj_mnn)
+            
+            
+            con_mnn = criterion.mean(constraint_eval)
+            con_mnn = pt_GX.inverse_transform(con_mnn)
+            con_mnn = torch.from_numpy(-con_mnn)
+            
+            obj_varr = criterion.variance(objective_eval)
+            con_varr = criterion.variance(constraint_eval)
+            
+            return obj_mnn, obj_varr, con_mnn, con_varr
+            
+            
+            
+    return output
+
+
+
+
+def acq_ensembling(acq_values): # (points, ensemble dim)
+        return acq_values.max(1).values
+
+
+
+
+
+
+
+def torch_std_inverse_transform(X, scale, mean):
+    X *= scale
+    X += mean
+    return X
+
+
+def torch_power_inverse_transform(x, lmbda):
+    out = torch.zeros_like(x)
+    pos = x >= 0
+
+    # when x >= 0
+    if abs(lmbda) < np.spacing(1.0):
+        out[pos] = torch.exp(x[pos])-1
+    else:  # lmbda != 0
+        out[pos] = torch.pow(x[pos] * lmbda + 1, 1 / lmbda) - 1
+
+    # when x < 0
+    if abs(lmbda - 2) > np.spacing(1.0):
+        out[~pos] = 1 - torch.pow(-(2 - lmbda) * x[~pos] + 1, 1 / (2 - lmbda))
+    else:  # lmbda == 2
+        out[~pos] = 1 - torch.exp(-x[~pos])
+
+    return out
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+################################################################################
+## PFN defined functions
+################################################################################
+
+
+def log01(x, eps=.0000001, input_between_zero_and_one=False):
+    logx = torch.log(x + eps)
+    if input_between_zero_and_one:
+        return (logx - math.log(eps)) / (math.log(1 + eps) - math.log(eps))
+    return (logx - logx.min(0)[0]) / (logx.max(0)[0] - logx.min(0)[0])
+
+def log01_batch(x, eps=.0000001, input_between_zero_and_one=False):
+    x = x.repeat(1, x.shape[-1] + 1, 1)
+    for b in range(x.shape[-1]):
+        x[:, b, b] = log01(x[:, b, b], eps=eps, input_between_zero_and_one=input_between_zero_and_one)
+    return x
+
+def lognormed_batch(x, eval_pos, eps=.0000001):
+    x = x.repeat(1, x.shape[-1] + 1, 1)
+    for b in range(x.shape[-1]):
+        logx = torch.log(x[:, b, b]+eps)
+        x[:, b, b] = (logx - logx[:eval_pos].mean(0))/logx[:eval_pos].std(0)
+    return x
+
+def _rank_transform(x_train, x):
+    assert len(x_train.shape) == len(x.shape) == 1
+    relative_to = torch.cat((torch.zeros_like(x_train[:1]),x_train.unique(sorted=True,), torch.ones_like(x_train[-1:])),-1)
+    higher_comparison = (relative_to < x[...,None]).sum(-1).clamp(min=1)
+    pos_inside_interval = (x - relative_to[higher_comparison-1])/(relative_to[higher_comparison] - relative_to[higher_comparison-1])
+    x_transformed = higher_comparison - 1 + pos_inside_interval
+    return x_transformed/(len(relative_to)-1.)
+
+def rank_transform(x_train, x):
+    assert x.shape[1] == x_train.shape[1], f"{x.shape=} and {x_train.shape=}"
+    # make sure everything is between 0 and 1
+    assert (x_train >= 0.).all() and (x_train <= 1.).all(), f"{x_train=}"
+    assert (x >= 0.).all() and (x <= 1.).all(), f"{x=}"
+    return_x = x.clone()
+    for feature_dim in range(x.shape[1]):
+        return_x[:, feature_dim] = _rank_transform(x_train[:, feature_dim], x[:, feature_dim])
+    return return_x
+
+
+
+def general_power_transform(x_train, x_apply, eps, less_safe=False):
+
+    # print('in function')
+    # print(x_train)
+    # print(x_apply)
+    # print('in function')
+    
+    if eps > 0:
+        try:
+            pt = PowerTransformer(method='box-cox')
+            pt.fit(x_train.cpu()+eps)
+            x_out = torch.tensor(pt.transform(x_apply.cpu()+eps), dtype=x_apply.dtype, device=x_apply.device)
+        except Exception as e:
+            print(e)
+            x_out = x_apply - x_train.mean(0)
+            print(x_train)
+            print(x_out)
+    else:
+        pt = PowerTransformer(method='yeo-johnson')
+        if not less_safe and (x_train.std() > 1_000 or x_train.mean().abs() > 1_000):
+            x_apply = (x_apply - x_train.mean(0)) / x_train.std(0)
+            x_train = (x_train - x_train.mean(0)) / x_train.std(0)
+            # print('inputs are LAARGEe, normalizing them')
+        try:
+            pt.fit(x_train.cpu().double())
+        # except ValueError as e:
+        except Exception as e:
+            # print(x_train)
+            # print('caught this errrr', e)
+            if less_safe:
+                x_train = (x_train - x_train.mean(0)) / x_train.std(0)
+                x_apply = (x_apply - x_train.mean(0)) / x_train.std(0)
+            else:
+                x_train = x_train - x_train.mean(0)
+                x_apply = x_apply - x_train.mean(0)
+            # print(x_train)
+            pt.fit(x_train.cpu().double())
+            # print(x_train)
+        x_out = torch.tensor(pt.transform(x_apply.cpu()), dtype=x_apply.dtype, device=x_apply.device)
+    if torch.isnan(x_out).any() or torch.isinf(x_out).any():
+        print('WARNING: power transform failed')
+        print(f"{x_train=} and {x_apply=}")
+        x_out = x_apply - x_train.mean(0)
+    return x_out, pt

SpeedReducer_CEI_Avg_Obj.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:1567f1e9557cb6d701605a2ec74c6e294c42a85c88ddf3c0f33e307bf7f9a07f
+size 3684

SpeedReducer_CEI_Avg_Time.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:9cfa75441b15a3e1b8dbdc1e4a074e7a3682c41c9a85924793c67a41bec86acd
+size 3496

ThreeTruss_CEI_Avg_Obj.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b02d1c95ef3aee57fee8804a82119d9b68453e182184cf47970779742d059bed
+size 2844

ThreeTruss_CEI_Avg_Time.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:7e9b30e0f99096ab84b8a545d2c1f24b80cd2d0bce1df6bc7f268b32c88a5b4f
+size 2912

WeldedBeam_CEI_Avg_Obj.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:efba0c05f0ac9803ee75caa3396983535bc3a104b47db2a3e463b1497ab5a93b
+size 3164

WeldedBeam_CEI_Avg_Time.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:06225c2be0d11a0cb563f1eb525401d5f5536401694d7ed7e3f7179a1f51352b
+size 3552

final_models/Cyril_500features.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:43218823860a5ca71657fd25c50bcc1209c3c570bcbb9df9ed2822bbb9f6f9c8
+size 239411934

final_models/Cyril_500features_800epoch_cpu.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:25884174687cfbde831badc4f1d05e94f860711dc3a07f4dde09930860e63603
+size 239408346

final_models/Cyril_50features.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d5b680c4b72e72a33a21896885de7e8fba52c42612a6165a7cf60afede2e425d
+size 107333480

final_models/hebo_morebudget_9_unused_features_3_userpriorperdim2_8.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ff2a4aa60feeca59e80f3b272d7b2ab521e1e82189469db494068de33dcaba17
+size 107378616

final_models/heboplus_500features_retrain_epoch800_cpu.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:25884174687cfbde831badc4f1d05e94f860711dc3a07f4dde09930860e63603
+size 239408346

final_models/model_hebo_morebudget_9_unused_features_3.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:cc7f6f4b9b06e59987e42845b7b0d31ffa5b414b9eddfe14d88b25120e3cd4f8
+size 107262245

final_models/model_sampled_warp_simple_mlp_for_hpob_46.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ebb2d4d7f419ca4617fdf85c663a69b3b0285fef91712e0dd69d5ab2d61754fd
+size 56761718

pfns4bo/.ipynb_checkpoints/__init__-checkpoint.py

@@ -0,0 +1,50 @@
+import os
+
+model_path = 'final_models'
+
+def prepare_models():
+    pfns4bo_dir = os.path.dirname(__file__)
+    model_names = ['hebo_morebudget_9_unused_features_3_userpriorperdim2_8.pt',
+                   'model_sampled_warp_simple_mlp_for_hpob_46.pt',
+                   'model_hebo_morebudget_9_unused_features_3.pt',]
+
+    for name in model_names:
+        weights_path = os.path.join(pfns4bo_dir, model_path, name)
+        compressed_weights_path = os.path.join(pfns4bo_dir, model_path, name + '.gz')
+        if not os.path.exists(weights_path):
+            if not os.path.exists(compressed_weights_path):
+                print("Downloading", os.path.abspath(compressed_weights_path))
+                import requests
+                url = f'https://github.com/automl/PFNs4BO/raw/main/pfns4bo/final_models/{name + ".gz"}'
+                r = requests.get(url, allow_redirects=True)
+                os.makedirs(os.path.dirname(compressed_weights_path), exist_ok=True)
+                with open(compressed_weights_path, 'wb') as f:
+                    f.write(r.content)
+            if os.path.exists(compressed_weights_path):
+                print("Unzipping", name)
+                os.system(f"gzip -dk {compressed_weights_path}")
+            else:
+                print("Failed to find", compressed_weights_path)
+                print("Make sure you have an internet connection to download the model automatically..")
+        if os.path.exists(weights_path):
+            print("Successfully located model at", weights_path)
+
+
+model_dict = {
+    'hebo_plus_userprior_model': os.path.join(os.path.dirname(__file__),model_path,
+                                              'hebo_morebudget_9_unused_features_3_userpriorperdim2_8.pt'),
+    'hebo_plus_model': os.path.join(os.path.dirname(__file__),model_path,
+                                    'model_hebo_morebudget_9_unused_features_3.pt'),
+    'bnn_model': os.path.join(os.path.dirname(__file__),model_path,'model_sampled_warp_simple_mlp_for_hpob_46.pt')
+}
+
+
+def __getattr__(name):
+    if name in model_dict:
+        if not os.path.exists(model_dict[name]):
+            print("Can't find", os.path.abspath(model_dict[name]), "thus unzipping/downloading models now.")
+            print("This might take a while..")
+            prepare_models()
+        return model_dict[name]
+    raise AttributeError(f"module '{__name__}' has no attribute '{name}'")
+

pfns4bo/.ipynb_checkpoints/bar_distribution-checkpoint.py

@@ -0,0 +1,410 @@
+from .utils import print_once
+
+import torch
+from torch import nn
+
+
+class BarDistribution(nn.Module):
+    def __init__(self, borders: torch.Tensor, smoothing=.0, ignore_nan_targets=True): # here borders should start with min and end with max, where all values lie in (min,max) and are sorted
+        '''
+        :param borders:
+        :param smoothing:
+        :param append_mean_pred: Whether to predict the mean of the other positions as a last output in forward,
+        is enabled when additionally y has a sequence length 1 shorter than logits, i.e. len(logits) == 1 + len(y)
+        '''
+        super().__init__()
+        assert len(borders.shape) == 1
+        self.register_buffer('borders', borders)
+        self.register_buffer('smoothing', torch.tensor(smoothing))
+        self.register_buffer('bucket_widths', self.borders[1:] - self.borders[:-1])
+        full_width = self.bucket_widths.sum()
+
+        assert (1 - (full_width / (self.borders[-1] - self.borders[0]))).abs() < 1e-2, f'diff: {full_width - (self.borders[-1] - self.borders[0])} with {full_width} {self.borders[-1]} {self.borders[0]}'
+        assert (self.bucket_widths >= 0.0).all() , "Please provide sorted borders!" # This also allows size zero buckets
+        self.num_bars = len(borders) - 1
+        self.ignore_nan_targets = ignore_nan_targets
+        self.to(borders.device)
+
+    def __setstate__(self, state):
+        super().__setstate__(state)
+        self.__dict__.setdefault('append_mean_pred', False)
+
+    def map_to_bucket_idx(self, y):
+        target_sample = torch.searchsorted(self.borders, y) - 1
+        target_sample[y == self.borders[0]] = 0
+        target_sample[y == self.borders[-1]] = self.num_bars - 1
+        return target_sample
+
+    def ignore_init(self, y):
+        ignore_loss_mask = torch.isnan(y)
+        if ignore_loss_mask.any():
+            if not self.ignore_nan_targets: raise ValueError(f'Found NaN in target {y}')
+            print_once("A loss was ignored because there was nan target.")
+        y[ignore_loss_mask] = self.borders[0] # this is just a default value, it will be ignored anyway
+        return ignore_loss_mask
+
+    def compute_scaled_log_probs(self, logits):
+        # this is equivalent to log(p(y)) of the density p
+        bucket_log_probs = torch.log_softmax(logits, -1)
+        scaled_bucket_log_probs = bucket_log_probs - torch.log(self.bucket_widths)
+        return scaled_bucket_log_probs
+
+    def forward(self, logits, y, mean_prediction_logits=None): # gives the negative log density (the _loss_), y: T x B, logits: T x B x self.num_bars
+        y = y.clone().view(*logits.shape[:-1]) # no trailing one dimension
+        ignore_loss_mask = self.ignore_init(y)
+        target_sample = self.map_to_bucket_idx(y)
+        assert (target_sample >= 0).all() and (target_sample < self.num_bars).all(), f'y {y} not in support set for borders (min_y, max_y) {self.borders}'
+        assert logits.shape[-1] == self.num_bars, f'{logits.shape[-1]} vs {self.num_bars}'
+
+        scaled_bucket_log_probs = self.compute_scaled_log_probs(logits)
+        nll_loss = -scaled_bucket_log_probs.gather(-1,target_sample[..., None]).squeeze(-1) # T x B
+
+        if mean_prediction_logits is not None:
+            if not self.training:
+                print('Calculating loss incl mean prediction loss for nonmyopic BO.')
+            scaled_mean_log_probs = self.compute_scaled_log_probs(mean_prediction_logits)
+            nll_loss = torch.cat((nll_loss, self.mean_loss(logits, scaled_mean_log_probs)), 0)
+
+        smooth_loss = -scaled_bucket_log_probs.mean(dim=-1)
+        smoothing = self.smoothing if self.training else 0.
+        loss = (1. - smoothing) * nll_loss + smoothing * smooth_loss
+        loss[ignore_loss_mask] = 0.
+        return loss
+
+    def mean_loss(self, logits, scaled_mean_logits):
+        assert (len(logits.shape) == 3) and (len(scaled_mean_logits.shape) == 2), \
+            (len(logits.shape), len(scaled_mean_logits.shape))
+        means = self.mean(logits).detach()  # T x B
+        target_mean = self.map_to_bucket_idx(means).clamp_(0, self.num_bars - 1)  # T x B
+        return -scaled_mean_logits.gather(1, target_mean.T).mean(1).unsqueeze(0)  # 1 x B
+
+    def mean(self, logits):
+        bucket_means = self.borders[:-1] + self.bucket_widths/2
+        p = torch.softmax(logits, -1)
+        return p @ bucket_means
+
+    def median(self, logits):
+        return self.icdf(logits, 0.5)
+
+    def icdf(self, logits, left_prob):
+        """
+        Implementation of the quantile function
+        :param logits: Tensor of any shape, with the last dimension being logits
+        :param left_prob: float: The probability mass to the left of the result.
+        :return: Position with `left_prob` probability weight to the left.
+        """
+        probs = logits.softmax(-1)
+        cumprobs = torch.cumsum(probs, -1)
+        idx = torch.searchsorted(cumprobs, left_prob * torch.ones(*cumprobs.shape[:-1], 1, device=logits.device))\
+            .squeeze(-1).clamp(0, cumprobs.shape[-1] - 1)  # this might not do the right for outliers
+        cumprobs = torch.cat([torch.zeros(*cumprobs.shape[:-1], 1, device=logits.device), cumprobs], -1)
+
+        rest_prob = left_prob - cumprobs.gather(-1, idx[..., None]).squeeze(-1)
+        left_border = self.borders[idx]
+        right_border = self.borders[idx+1]
+        return left_border + (right_border - left_border) * rest_prob / probs.gather(-1, idx[..., None]).squeeze(-1)
+
+    def quantile(self, logits, center_prob=.682):
+        side_probs = (1.-center_prob)/2
+        return torch.stack((self.icdf(logits, side_probs), self.icdf(logits, 1.-side_probs)),-1)
+
+    def ucb(self, logits, best_f, rest_prob=(1-.682)/2, maximize=True):
+        """
+        UCB utility. Rest Prob is the amount of utility above (below) the confidence interval that is ignored.
+        Higher rest_prob is equivalent to lower beta in the standard GP-UCB formulation.
+        :param logits: Logits, as returned by the Transformer.
+        :param rest_prob: The amount of utility above (below) the confidence interval that is ignored.
+        The default is equivalent to using GP-UCB with `beta=1`.
+        To get the corresponding `beta`, where `beta` is from
+        the standard GP definition of UCB `ucb_utility = mean + beta * std`,
+        you can use this computation: `beta = math.sqrt(2)*torch.erfinv(torch.tensor(2*(1-rest_prob)-1))`.
+        :param maximize:
+        :return: utility
+        """
+        if maximize:
+            rest_prob = 1 - rest_prob
+        return self.icdf(logits, rest_prob)
+
+    def mode(self, logits):
+        mode_inds = logits.argmax(-1)
+        bucket_means = self.borders[:-1] + self.bucket_widths/2
+        return bucket_means[mode_inds]
+
+    def ei(self, logits, best_f, maximize=True): # logits: evaluation_points x batch x feature_dim
+        bucket_diffs = self.borders[1:] - self.borders[:-1]
+        assert maximize
+        if not torch.is_tensor(best_f) or not len(best_f.shape):
+            best_f = torch.full(logits[...,0].shape, best_f, device=logits.device)
+
+        best_f = best_f[..., None].repeat(*[1]*len(best_f.shape), logits.shape[-1])
+        clamped_best_f = best_f.clamp(self.borders[:-1], self.borders[1:])
+
+        #bucket_contributions = (best_f[...,None] < self.borders[:-1]).float() * bucket_means
+        # true bucket contributions
+        bucket_contributions = ((self.borders[1:]**2-clamped_best_f**2)/2 - best_f*(self.borders[1:] - clamped_best_f))/bucket_diffs
+
+        p = torch.softmax(logits, -1)
+        return torch.einsum("...b,...b->...", p, bucket_contributions)
+
+    def pi(self, logits, best_f, maximize=True):# logits: evaluation_points x batch x feature_dim
+        """
+        Acquisition Function: Probability of Improvement
+        :param logits: as returned by Transformer
+        :param best_f: best evaluation so far (the incumbent)
+        :param maximize: whether to maximize
+        :return: utility
+        """
+        assert maximize is True
+        if not torch.is_tensor(best_f) or not len(best_f.shape):
+            best_f = torch.full(logits[...,0].shape, best_f, device=logits.device)
+        p = torch.softmax(logits, -1)
+        border_widths = self.borders[1:] - self.borders[:-1]
+        factor = 1. - ((best_f[...,None] - self.borders[:-1]) / border_widths).clamp(0., 1.)
+        return (p * factor).sum(-1)
+
+
+    def mean_of_square(self, logits):
+        """
+        Computes E[x^2].
+        :param logits: Output of the model.
+        """
+        left_borders = self.borders[:-1]
+        right_borders = self.borders[1:]
+        bucket_mean_of_square = (left_borders.square() + right_borders.square() + left_borders*right_borders)/3.
+        p = torch.softmax(logits, -1)
+        return p @ bucket_mean_of_square
+
+    def variance(self, logits):
+        return self.mean_of_square(logits) - self.mean(logits).square()
+
+    def pi(self, logits, best_f, maximize=True):# logits: evaluation_points x batch x feature_dim
+        """
+        Acquisition Function: Probability of Improvement
+        :param logits: as returned by Transformer
+        :param best_f: best evaluation so far (the incumbent)
+        :param maximize: whether to maximize
+        :return: utility
+        """
+        assert maximize is True
+        p = torch.softmax(logits, -1)
+        border_widths = self.borders[1:] - self.borders[:-1]
+        factor = 1. - ((best_f - self.borders[:-1]) / border_widths).clamp(0., 1.)
+        return (p * factor).sum(-1)
+
+
+    def mean_of_square(self, logits):
+        """
+        Computes E[x^2].
+        :param logits: Output of the model.
+        """
+        left_borders = self.borders[:-1]
+        right_borders = self.borders[1:]
+        bucket_mean_of_square = (left_borders.square() + right_borders.square() + left_borders*right_borders)/3.
+        p = torch.softmax(logits, -1)
+        return p @ bucket_mean_of_square
+
+    def variance(self, logits):
+        return self.mean_of_square(logits) - self.mean(logits).square()
+
+
+class FullSupportBarDistribution(BarDistribution):
+    @staticmethod
+    def halfnormal_with_p_weight_before(range_max,p=.5):
+        s = range_max / torch.distributions.HalfNormal(torch.tensor(1.)).icdf(torch.tensor(p))
+        return torch.distributions.HalfNormal(s)
+
+
+    def forward(self, logits, y, mean_prediction_logits=None): # gives the negative log density (the _loss_), y: T x B, logits: T x B x self.num_bars
+        assert self.num_bars > 1
+        y = y.clone().view(len(y),-1) # no trailing one dimension
+        ignore_loss_mask = self.ignore_init(y) # alters y
+        target_sample = self.map_to_bucket_idx(y) # shape: T x B (same as y)
+        target_sample.clamp_(0, self.num_bars - 1)
+
+        assert logits.shape[-1] == self.num_bars, f'{logits.shape[-1]} vs {self.num_bars}'
+        assert (target_sample >= 0).all() and (target_sample < self.num_bars).all(), \
+            f'y {y} not in support set for borders (min_y, max_y) {self.borders}'
+        assert logits.shape[-1] == self.num_bars, f'{logits.shape[-1]} vs {self.num_bars}'
+        # ignore all position with nan values
+
+
+        scaled_bucket_log_probs = self.compute_scaled_log_probs(logits)
+
+        assert len(scaled_bucket_log_probs) == len(target_sample), (len(scaled_bucket_log_probs), len(target_sample))
+        log_probs = scaled_bucket_log_probs.gather(-1, target_sample.unsqueeze(-1)).squeeze(-1)
+
+        side_normals = (self.halfnormal_with_p_weight_before(self.bucket_widths[0]), self.halfnormal_with_p_weight_before(self.bucket_widths[-1]))
+
+
+        log_probs[target_sample == 0] += side_normals[0].log_prob((self.borders[1]-y[target_sample == 0]).clamp(min=.00000001)) + torch.log(self.bucket_widths[0])
+        log_probs[target_sample == self.num_bars-1] += side_normals[1].log_prob((y[target_sample == self.num_bars-1]-self.borders[-2]).clamp(min=.00000001)) + torch.log(self.bucket_widths[-1])
+
+        nll_loss = -log_probs
+
+        if mean_prediction_logits is not None:
+            assert not ignore_loss_mask.any(), "Ignoring examples is not implemented with mean pred."
+            if not self.training:
+                print('Calculating loss incl mean prediction loss for nonmyopic BO.')
+            if not torch.is_grad_enabled():
+                print("Warning: loss is not correct in absolute terms, only the gradient is right, when using `append_mean_pred`.")
+            scaled_mean_log_probs = self.compute_scaled_log_probs(mean_prediction_logits)
+            nll_loss = torch.cat((nll_loss, self.mean_loss(logits, scaled_mean_log_probs)), 0)
+            #ignore_loss_mask = torch.zeros_like(nll_loss, dtype=torch.bool)
+
+        if self.smoothing:
+            smooth_loss = -scaled_bucket_log_probs.mean(dim=-1)
+            smoothing = self.smoothing if self.training else 0.
+            nll_loss = (1. - smoothing) * nll_loss + smoothing * smooth_loss
+
+        if ignore_loss_mask.any():
+            nll_loss[ignore_loss_mask] = 0.
+
+        return nll_loss
+
+    def mean(self, logits):
+        bucket_means = self.borders[:-1] + self.bucket_widths / 2
+        p = torch.softmax(logits, -1)
+        side_normals = (self.halfnormal_with_p_weight_before(self.bucket_widths[0]),
+                        self.halfnormal_with_p_weight_before(self.bucket_widths[-1]))
+        bucket_means[0] = -side_normals[0].mean + self.borders[1]
+        bucket_means[-1] = side_normals[1].mean + self.borders[-2]
+        return p @ bucket_means.to(logits.device)
+
+    def mean_of_square(self, logits):
+        """
+        Computes E[x^2].
+        :param logits: Output of the model.
+        """
+        left_borders = self.borders[:-1]
+        right_borders = self.borders[1:]
+        bucket_mean_of_square = (left_borders.square() + right_borders.square() + left_borders*right_borders)/3.
+        side_normals = (self.halfnormal_with_p_weight_before(self.bucket_widths[0]),
+                        self.halfnormal_with_p_weight_before(self.bucket_widths[-1]))
+        bucket_mean_of_square[0] = side_normals[0].variance + (-side_normals[0].mean + self.borders[1]).square()
+        bucket_mean_of_square[-1] = side_normals[1].variance + (side_normals[1].variance + self.borders[-2]).square()
+        p = torch.softmax(logits, -1)
+        return p @ bucket_mean_of_square
+
+    def pi(self, logits, best_f, maximize=True):# logits: evaluation_points x batch x feature_dim
+        """
+        Acquisition Function: Probability of Improvement
+        :param logits: as returned by Transformer (evaluation_points x batch x feature_dim)
+        :param best_f: best evaluation so far (the incumbent)
+        :param maximize: whether to maximize
+        :return: utility
+        """
+        assert maximize is True
+        if not torch.is_tensor(best_f) or not len(best_f.shape):
+            best_f = torch.full(logits[...,0].shape, best_f, device=logits.device) # evaluation_points x batch
+        assert best_f.shape == logits[...,0].shape, f"best_f.shape: {best_f.shape}, logits.shape: {logits.shape}"
+        p = torch.softmax(logits, -1) # evaluation_points x batch
+        border_widths = self.borders[1:] - self.borders[:-1]
+        factor = 1. - ((best_f[...,None] - self.borders[:-1]) / border_widths).clamp(0., 1.) # evaluation_points x batch x num_bars
+
+        side_normals = (self.halfnormal_with_p_weight_before(self.bucket_widths[0]),
+                        self.halfnormal_with_p_weight_before(self.bucket_widths[-1]))
+        position_in_side_normals = (-(best_f - self.borders[1]).clamp(max=0.), (best_f - self.borders[-2]).clamp(min=0.)) # evaluation_points x batch
+        factor[...,0] = 0.
+        factor[...,0][position_in_side_normals[0] > 0.] = side_normals[0].cdf(position_in_side_normals[0][position_in_side_normals[0] > 0.])
+        factor[...,-1] = 1.
+        factor[...,-1][position_in_side_normals[1] > 0.] = 1. - side_normals[1].cdf(position_in_side_normals[1][position_in_side_normals[1] > 0.])
+        return (p * factor).sum(-1)
+
+
+    def ei_for_halfnormal(self, scale, best_f, maximize=True):
+        """
+        This is the EI for a standard normal distribution with mean 0 and variance `scale` times 2.
+        Which is the same as the half normal EI.
+        I tested this with MC approximation:
+        ei_for_halfnormal = lambda scale, best_f: (torch.distributions.HalfNormal(torch.tensor(scale)).sample((10_000_000,))- best_f ).clamp(min=0.).mean()
+        print([(ei_for_halfnormal(scale,best_f), FullSupportBarDistribution().ei_for_halfnormal(scale,best_f)) for scale in [0.1,1.,10.] for best_f in [.1,10.,4.]])
+        :param scale:
+        :param best_f:
+        :param maximize:
+        :return:
+        """
+        assert maximize
+        mean = torch.tensor(0.)
+        u = (mean - best_f) / scale
+        normal = torch.distributions.Normal(torch.zeros_like(u), torch.ones_like(u))
+        try:
+            ucdf = normal.cdf(u)
+        except ValueError:
+            print(f"u: {u}, best_f: {best_f}, scale: {scale}")
+            raise
+        updf = torch.exp(normal.log_prob(u))
+        normal_ei = scale * (updf + u * ucdf)
+        return 2*normal_ei
+
+    def ei(self, logits, best_f, maximize=True): # logits: evaluation_points x batch x feature_dim
+        if torch.isnan(logits).any():
+            raise ValueError(f"logits contains NaNs: {logits}")
+        bucket_diffs = self.borders[1:] - self.borders[:-1]
+        assert maximize
+        if not torch.is_tensor(best_f) or not len(best_f.shape):
+            best_f = torch.full(logits[...,0].shape, best_f, device=logits.device)
+        assert best_f.shape == logits[...,0].shape, f"best_f.shape: {best_f.shape}, logits.shape: {logits.shape}"
+
+
+        best_f_per_logit = best_f[..., None].repeat(*[1]*len(best_f.shape), logits.shape[-1])
+        clamped_best_f = best_f_per_logit.clamp(self.borders[:-1], self.borders[1:])
+
+        # true bucket contributions
+        bucket_contributions = ((self.borders[1:]**2-clamped_best_f**2)/2 - best_f_per_logit*(self.borders[1:] - clamped_best_f))/bucket_diffs
+
+        # extra stuff for continuous
+        side_normals = (self.halfnormal_with_p_weight_before(self.bucket_widths[0]),
+                        self.halfnormal_with_p_weight_before(self.bucket_widths[-1]))
+        position_in_side_normals = (-(best_f - self.borders[1]).clamp(max=0.),
+                                    (best_f - self.borders[-2]).clamp(min=0.))  # evaluation_points x batch
+
+        bucket_contributions[...,-1] = self.ei_for_halfnormal(side_normals[1].scale, position_in_side_normals[1])
+
+        bucket_contributions[...,0] = self.ei_for_halfnormal(side_normals[0].scale, torch.zeros_like(position_in_side_normals[0])) \
+                                  - self.ei_for_halfnormal(side_normals[0].scale, position_in_side_normals[0])
+
+        p = torch.softmax(logits, -1)
+        return torch.einsum("...b,...b->...", p, bucket_contributions)
+
+
+def get_bucket_limits(num_outputs:int, full_range:tuple=None, ys:torch.Tensor=None, verbose:bool=False):
+    assert (ys is None) != (full_range is None), 'Either full_range or ys must be passed.'
+
+    if ys is not None:
+        ys = ys.flatten()
+        ys = ys[~torch.isnan(ys)]
+        if len(ys) % num_outputs: ys = ys[:-(len(ys) % num_outputs)]
+        print(f'Using {len(ys)} y evals to estimate {num_outputs} buckets. Cut off the last {len(ys) % num_outputs} ys.')
+        ys_per_bucket = len(ys) // num_outputs
+        if full_range is None:
+            full_range = (ys.min(), ys.max())
+        else:
+            assert full_range[0] <= ys.min() and full_range[1] >= ys.max(), f'full_range {full_range} not in range of ys {ys.min(), ys.max()}'
+            full_range = torch.tensor(full_range)
+        ys_sorted, ys_order = ys.sort(0)
+        bucket_limits = (ys_sorted[ys_per_bucket-1::ys_per_bucket][:-1]+ys_sorted[ys_per_bucket::ys_per_bucket])/2
+        if verbose:
+            print(f'Using {len(ys)} y evals to estimate {num_outputs} buckets. Cut off the last {len(ys) % num_outputs} ys.')
+            print(full_range)
+        bucket_limits = torch.cat([full_range[0].unsqueeze(0), bucket_limits, full_range[1].unsqueeze(0)],0)
+
+    else:
+        class_width = (full_range[1] - full_range[0]) / num_outputs
+        bucket_limits = torch.cat([full_range[0] + torch.arange(num_outputs).float()*class_width, torch.tensor(full_range[1]).unsqueeze(0)], 0)
+
+    assert len(bucket_limits) - 1 == num_outputs, f'len(bucket_limits) - 1 == {len(bucket_limits) - 1} != {num_outputs} == num_outputs'
+    assert full_range[0] == bucket_limits[0], f'{full_range[0]} != {bucket_limits[0]}'
+    assert full_range[-1] == bucket_limits[-1], f'{full_range[-1]} != {bucket_limits[-1]}'
+
+    return bucket_limits
+
+
+def get_custom_bar_dist(borders, criterion):
+    # Tested that a bar_dist with borders 0.54 (-> softplus 1.0) yields the same bar distribution as the passed one.
+    borders_ = torch.nn.functional.softplus(borders) + 0.001
+    borders_ = (torch.cumsum(torch.cat([criterion.borders[0:1], criterion.bucket_widths]) * borders_, 0))
+    criterion_ = criterion.__class__(borders=borders_, handle_nans=criterion.handle_nans)
+    return criterion_
+
+
+

pfns4bo/.ipynb_checkpoints/lost_functions-checkpoint.py

@@ -0,0 +1,177 @@
+#!/usr/bin/env python3
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+
+r"""
+A converter that simplifies using numpy-based optimizers with generic torch
+`nn.Module` classes. This enables using a `scipy.optim.minimize` optimizer
+for optimizing module parameters.
+"""
+
+from __future__ import annotations
+
+from collections import OrderedDict
+from math import inf
+from numbers import Number
+from typing import Dict, List, Optional, Set, Tuple
+from warnings import warn
+
+import numpy as np
+import torch
+from botorch.optim.utils import (
+    _get_extra_mll_args,
+    _handle_numerical_errors,
+    get_name_filter,
+    get_parameters_and_bounds,
+    TorchAttr,
+)
+from gpytorch.mlls import MarginalLogLikelihood
+from torch.nn import Module
+
+
+def module_to_array(
+    module: Module,
+    bounds: Optional[Dict[str, Tuple[Optional[float], Optional[float]]]] = None,
+    exclude: Optional[Set[str]] = None,
+) -> Tuple[np.ndarray, Dict[str, TorchAttr], Optional[np.ndarray]]:
+    r"""Extract named parameters from a module into a numpy array.
+
+    Only extracts parameters with requires_grad, since it is meant for optimizing.
+
+    Args:
+        module: A module with parameters. May specify parameter constraints in
+            a `named_parameters_and_constraints` method.
+        bounds: A dictionary mapping parameter names t lower and upper bounds.
+            of lower and upper bounds. Bounds specified here take precedence
+            over bounds on the same parameters specified in the constraints
+            registered with the module.
+        exclude: A list of parameter names that are to be excluded from extraction.
+
+    Returns:
+        3-element tuple containing
+        - The parameter values as a numpy array.
+        - An ordered dictionary with the name and tensor attributes of each
+        parameter.
+        - A `2 x n_params` numpy array with lower and upper bounds if at least
+        one constraint is finite, and None otherwise.
+
+    Example:
+        >>> mll = ExactMarginalLogLikelihood(model.likelihood, model)
+        >>> parameter_array, property_dict, bounds_out = module_to_array(mll)
+    """
+    warn(
+        "`module_to_array` is marked for deprecation, consider using "
+        "`get_parameters_and_bounds`, `get_parameters_as_ndarray_1d`, or "
+        "`get_bounds_as_ndarray` instead.",
+        DeprecationWarning,
+    )
+    param_dict, bounds_dict = get_parameters_and_bounds(
+        module=module,
+        name_filter=None if exclude is None else get_name_filter(exclude),
+        requires_grad=True,
+    )
+    if bounds is not None:
+        bounds_dict.update(bounds)
+
+    # Record tensor metadata and read parameter values to the tape
+    param_tape: List[Number] = []
+    property_dict = OrderedDict()
+    with torch.no_grad():
+        for name, param in param_dict.items():
+            property_dict[name] = TorchAttr(param.shape, param.dtype, param.device)
+            param_tape.extend(param.view(-1).cpu().double().tolist())
+
+    # Extract lower and upper bounds
+    start = 0
+    bounds_np = None
+    params_np = np.asarray(param_tape)
+    for name, param in param_dict.items():
+        numel = param.numel()
+        if name in bounds_dict:
+            for row, bound in enumerate(bounds_dict[name]):
+                if bound is None:
+                    continue
+
+                if torch.is_tensor(bound):
+                    if (bound == (2 * row - 1) * inf).all():
+                        continue
+                    bound = bound.detach().cpu()
+
+                elif bound == (2 * row - 1) * inf:
+                    continue
+
+                if bounds_np is None:
+                    bounds_np = np.full((2, len(params_np)), ((-inf,), (inf,)))
+
+                bounds_np[row, start : start + numel] = bound
+        start += numel
+
+    return params_np, property_dict, bounds_np
+
+
+
+
+def set_params_with_array(
+    module: Module, x: np.ndarray, property_dict: Dict[str, TorchAttr]
+) -> Module:
+    r"""Set module parameters with values from numpy array.
+
+    Args:
+        module: Module with parameters to be set
+        x: Numpy array with parameter values
+        property_dict: Dictionary of parameter names and torch attributes as
+            returned by module_to_array.
+
+    Returns:
+        Module: module with parameters updated in-place.
+
+    Example:
+        >>> mll = ExactMarginalLogLikelihood(model.likelihood, model)
+        >>> parameter_array, property_dict, bounds_out = module_to_array(mll)
+        >>> parameter_array += 0.1  # perturb parameters (for example only)
+        >>> mll = set_params_with_array(mll, parameter_array,  property_dict)
+    """
+    warn(
+        "`_set_params_with_array` is marked for deprecation, consider using "
+        "`set_parameters_from_ndarray_1d` instead.",
+        DeprecationWarning,
+    )
+    param_dict = OrderedDict(module.named_parameters())
+    start_idx = 0
+    for p_name, attrs in property_dict.items():
+        # Construct the new tensor
+        if len(attrs.shape) == 0:  # deal with scalar tensors
+            end_idx = start_idx + 1
+            new_data = torch.tensor(
+                x[start_idx], dtype=attrs.dtype, device=attrs.device
+            )
+        else:
+            end_idx = start_idx + np.prod(attrs.shape)
+            new_data = torch.tensor(
+                x[start_idx:end_idx], dtype=attrs.dtype, device=attrs.device
+            ).view(*attrs.shape)
+        start_idx = end_idx
+        # Update corresponding parameter in-place. Disable autograd to update.
+        param_dict[p_name].requires_grad_(False)
+        param_dict[p_name].copy_(new_data)
+        param_dict[p_name].requires_grad_(True)
+    return module
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+

pfns4bo/.ipynb_checkpoints/transformer-checkpoint.py

@@ -0,0 +1,327 @@
+import math
+from typing import Optional
+
+import torch
+import torch.nn as nn
+from torch import Tensor
+from torch.nn import Module, TransformerEncoder
+
+from .layer import TransformerEncoderLayer, _get_activation_fn
+from .utils import SeqBN, bool_mask_to_att_mask
+
+
+
+class TransformerModel(nn.Module):
+    def __init__(self, encoder, ninp, nhead, nhid, nlayers, dropout=0.0, style_encoder=None, y_encoder=None,
+                 pos_encoder=None, decoder_dict=None, input_normalization=False, init_method=None, pre_norm=False,
+                 activation='gelu', recompute_attn=False, num_global_att_tokens=0, full_attention=False,
+                 all_layers_same_init=False, efficient_eval_masking=True, decoder_once_dict=None, return_all_outputs=False,
+                 save_trainingset_representations=False):
+        super().__init__()
+        self.model_type = 'Transformer'
+        encoder_layer_creator = lambda: TransformerEncoderLayer(ninp, nhead, nhid, dropout, activation=activation,
+                                                                pre_norm=pre_norm, recompute_attn=recompute_attn,
+                                                                save_trainingset_representations=save_trainingset_representations)
+        self.transformer_encoder = TransformerEncoder(encoder_layer_creator(), nlayers)\
+            if all_layers_same_init else TransformerEncoderDiffInit(encoder_layer_creator, nlayers)
+        self.ninp = ninp
+        self.encoder = encoder
+        self.y_encoder = y_encoder
+        self.pos_encoder = pos_encoder
+        self.return_all_outputs = return_all_outputs
+
+        def make_decoder_dict(decoder_description_dict):
+            if decoder_description_dict is None or len(decoder_description_dict) == 0:
+                return None
+            initialized_decoder_dict = {}
+            for decoder_key in decoder_description_dict:
+                decoder_model, decoder_n_out = decoder_description_dict[decoder_key]
+                if decoder_model is None:
+                    initialized_decoder_dict[decoder_key] = nn.Sequential(nn.Linear(ninp, nhid), nn.GELU(), nn.Linear(nhid, decoder_n_out))
+                else:
+                    initialized_decoder_dict[decoder_key] = decoder_model(ninp, nhid, decoder_n_out)
+                print('Initialized decoder for', decoder_key, 'with', decoder_description_dict[decoder_key], ' and nout', decoder_n_out)
+            return torch.nn.ModuleDict(initialized_decoder_dict)
+
+        self.decoder_dict = make_decoder_dict(decoder_dict)
+        self.decoder_dict_once = make_decoder_dict(decoder_once_dict)
+
+        # N(0,1) is the initialization as the default of nn.Embedding
+        self.decoder_dict_once_embeddings = torch.nn.Parameter(torch.randn((len(self.decoder_dict_once), 1, ninp))) if self.decoder_dict_once is not None else None
+            #nn.Embedding(len(self.decoder_dict.keys()), nhid)
+        self.input_ln = SeqBN(ninp) if input_normalization else None
+        self.style_encoder = style_encoder
+        self.init_method = init_method
+        if num_global_att_tokens is not None:
+            assert not full_attention
+        self.global_att_embeddings = nn.Embedding(num_global_att_tokens, ninp) if num_global_att_tokens else None
+        self.full_attention = full_attention
+        self.efficient_eval_masking = efficient_eval_masking
+
+        self.nhid = nhid
+
+        self.init_weights()
+
+    def __setstate__(self, state):
+        super().__setstate__(state)
+        self.__dict__.setdefault('efficient_eval_masking', False)
+        if not hasattr(self, 'decoder_dict_once'):
+            self.__dict__.setdefault('decoder_dict_once', None)
+        if hasattr(self, 'decoder') and not hasattr(self, 'decoder_dict'):
+            self.add_module('decoder_dict', nn.ModuleDict({'standard': self.decoder}))
+        self.__dict__.setdefault('return_all_outputs', False)
+
+        def add_approximate_false(module):
+            if isinstance(module, nn.GELU):
+                module.__dict__.setdefault('approximate', 'none')
+
+        self.apply(add_approximate_false)
+
+    @staticmethod
+    def generate_square_subsequent_mask(sz):
+        mask = (torch.triu(torch.ones(sz, sz)) == 1).transpose(0, 1)
+        return bool_mask_to_att_mask(mask)
+
+    @staticmethod
+    def generate_D_q_matrix(sz, query_size):
+        train_size = sz-query_size
+        mask = torch.zeros(sz,sz) == 0
+        mask[:,train_size:].zero_()
+        mask |= torch.eye(sz) == 1
+        return bool_mask_to_att_mask(mask)
+
+    @staticmethod
+    def generate_global_att_query_matrix(num_global_att_tokens, seq_len, num_query_tokens):
+        train_size = seq_len + num_global_att_tokens - num_query_tokens
+        sz = seq_len + num_global_att_tokens
+        mask = torch.zeros(num_query_tokens, sz) == 0
+        mask[:,train_size:].zero_()
+        mask[:,train_size:] |= torch.eye(num_query_tokens) == 1
+        return bool_mask_to_att_mask(mask)
+
+    @staticmethod
+    def generate_global_att_trainset_matrix(num_global_att_tokens, seq_len, num_query_tokens):
+        train_size = seq_len + num_global_att_tokens - num_query_tokens
+        trainset_size = seq_len - num_query_tokens
+        mask = torch.zeros(trainset_size, num_global_att_tokens) == 0
+        #mask[:,num_global_att_tokens:].zero_()
+        #mask[:,num_global_att_tokens:] |= torch.eye(trainset_size) == 1
+        return bool_mask_to_att_mask(mask)
+
+    @staticmethod
+    def generate_global_att_globaltokens_matrix(num_global_att_tokens, seq_len, num_query_tokens):
+        mask = torch.zeros(num_global_att_tokens, num_global_att_tokens+seq_len-num_query_tokens) == 0
+        return bool_mask_to_att_mask(mask)
+
+    def init_weights(self):
+        initrange = 1.
+        # if isinstance(self.encoder,EmbeddingEncoder):
+        #    self.encoder.weight.data.uniform_(-initrange, initrange)
+        # self.decoder.bias.data.zero_()
+        # self.decoder.weight.data.uniform_(-initrange, initrange)
+        if self.init_method is not None:
+            self.apply(self.init_method)
+        for layer in self.transformer_encoder.layers:
+            nn.init.zeros_(layer.linear2.weight)
+            nn.init.zeros_(layer.linear2.bias)
+            attns = layer.self_attn if isinstance(layer.self_attn, nn.ModuleList) else [layer.self_attn]
+            for attn in attns:
+                nn.init.zeros_(attn.out_proj.weight)
+                nn.init.zeros_(attn.out_proj.bias)
+
+    def forward(self, *args, **kwargs):
+        """
+        This will perform a forward-pass (possibly recording gradients) of the model.
+        We have multiple interfaces we support with this model:
+
+        model(train_x, train_y, test_x, src_mask=None, style=None, only_return_standard_out=True)
+        model((x,y), src_mask=None, single_eval_pos=None, only_return_standard_out=True)
+        model((style,x,y), src_mask=None, single_eval_pos=None, only_return_standard_out=True)
+        """
+        if len(args) == 3:
+            # case model(train_x, train_y, test_x, src_mask=None, style=None, only_return_standard_out=True)
+            assert all(kwarg in {'src_mask', 'style', 'only_return_standard_out'} for kwarg in kwargs.keys()), \
+                f"Unrecognized keyword argument in kwargs: {set(kwargs.keys()) - {'src_mask', 'style', 'only_return_standard_out'}}"
+            x = args[0]
+            if args[2] is not None:
+                x = torch.cat((x, args[2]), dim=0)
+            style = kwargs.pop('style', None)
+            return self._forward((style, x, args[1]), single_eval_pos=len(args[0]), **kwargs)
+        elif len(args) == 1 and isinstance(args, tuple):
+            # case model((x,y), src_mask=None, single_eval_pos=None, only_return_standard_out=True)
+            # case model((style,x,y), src_mask=None, single_eval_pos=None, only_return_standard_out=True)
+            assert all(kwarg in {'src_mask', 'single_eval_pos', 'only_return_standard_out'} for kwarg in kwargs.keys()), \
+                f"Unrecognized keyword argument in kwargs: {set(kwargs.keys()) - {'src_mask', 'single_eval_pos', 'only_return_standard_out'}}"
+            return self._forward(*args, **kwargs)
+
+    def _forward(self, src, src_mask=None, single_eval_pos=None, only_return_standard_out=True):
+        assert isinstance(src, tuple), 'inputs (src) have to be given as (x,y) or (style,x,y) tuple'
+
+        if len(src) == 2: # (x,y) and no style
+            src = (None,) + src
+
+        style_src, x_src, y_src = src
+
+        if single_eval_pos is None:
+            single_eval_pos = x_src.shape[0]
+
+
+        x_src = self.encoder(x_src)
+
+        if self.decoder_dict_once is not None:
+            x_src = torch.cat([x_src, self.decoder_dict_once_embeddings.repeat(1, x_src.shape[1], 1)], dim=0)
+
+        y_src = self.y_encoder(y_src.unsqueeze(-1) if len(y_src.shape) < len(x_src.shape) else y_src) if y_src is not None else None
+        if self.style_encoder:
+            assert style_src is not None, 'style_src must be given if style_encoder is used'
+            style_src = self.style_encoder(style_src).unsqueeze(0)
+        else:
+            style_src = torch.tensor([], device=x_src.device)
+        global_src = torch.tensor([], device=x_src.device) if self.global_att_embeddings is None else \
+            self.global_att_embeddings.weight.unsqueeze(1).repeat(1, x_src.shape[1], 1)
+
+        if src_mask is not None:
+            assert self.global_att_embeddings is None or isinstance(src_mask, tuple)
+
+        if src_mask is None:
+            if self.global_att_embeddings is None:
+                full_len = len(x_src) + len(style_src)
+                if self.full_attention:
+                    src_mask = bool_mask_to_att_mask(torch.ones((full_len, full_len), dtype=torch.bool)).to(x_src.device)
+                elif self.efficient_eval_masking:
+                    src_mask = single_eval_pos + len(style_src)
+                else:
+                    src_mask = self.generate_D_q_matrix(full_len, len(x_src) - single_eval_pos).to(x_src.device)
+            else:
+                src_mask_args = (self.global_att_embeddings.num_embeddings,
+                                 len(x_src) + len(style_src),
+                                 len(x_src) + len(style_src) - single_eval_pos)
+                src_mask = (self.generate_global_att_globaltokens_matrix(*src_mask_args).to(x_src.device),
+                            self.generate_global_att_trainset_matrix(*src_mask_args).to(x_src.device),
+                            self.generate_global_att_query_matrix(*src_mask_args).to(x_src.device))
+
+        train_x = x_src[:single_eval_pos]
+        if y_src is not None:
+            train_x = train_x + y_src[:single_eval_pos]
+        src = torch.cat([global_src, style_src, train_x, x_src[single_eval_pos:]], 0)
+
+        if self.input_ln is not None:
+            src = self.input_ln(src)
+
+        if self.pos_encoder is not None:
+            src = self.pos_encoder(src)
+
+
+        output = self.transformer_encoder(src, src_mask)
+
+        num_prefix_positions = len(style_src)+(self.global_att_embeddings.num_embeddings if self.global_att_embeddings else 0)
+        if self.return_all_outputs:
+            out_range_start = num_prefix_positions
+        else:
+            out_range_start = single_eval_pos + num_prefix_positions
+
+        # In the line below, we use the indexing feature, that we have `x[i:None] == x[i:]`
+        out_range_end = -len(self.decoder_dict_once_embeddings) if self.decoder_dict_once is not None else None
+
+        # take care the output once are counted from the end
+        output_once = {k: v(output[-(i+1)]) for i, (k, v) in enumerate(self.decoder_dict_once.items())}\
+            if self.decoder_dict_once is not None else {}
+
+        output = {k: v(output[out_range_start:out_range_end]) for k,v in self.decoder_dict.items()}\
+            if self.decoder_dict is not None else {}
+
+        if only_return_standard_out:
+            return output['standard']
+
+        if output_once:
+            return output, output_once
+        return output
+
+    @torch.no_grad()
+    def init_from_small_model(self, small_model):
+        assert isinstance(self.decoder, nn.Linear) and isinstance(self.encoder, (nn.Linear, nn.Sequential)) \
+               and isinstance(self.y_encoder, (nn.Linear, nn.Sequential))
+
+        def set_encoder_weights(my_encoder, small_model_encoder):
+            my_encoder_linear, small_encoder_linear = (my_encoder, small_model_encoder) \
+                if isinstance(my_encoder, nn.Linear) else (my_encoder[-1], small_model_encoder[-1])
+            small_in_dim = small_encoder_linear.out_features
+            my_encoder_linear.weight.zero_()
+            my_encoder_linear.bias.zero_()
+            my_encoder_linear.weight[:small_in_dim] = small_encoder_linear.weight
+            my_encoder_linear.bias[:small_in_dim] = small_encoder_linear.bias
+
+        set_encoder_weights(self.encoder, small_model.encoder)
+        set_encoder_weights(self.y_encoder, small_model.y_encoder)
+
+        small_in_dim = small_model.decoder.in_features
+
+        self.decoder.weight[:, :small_in_dim] = small_model.decoder.weight
+        self.decoder.bias = small_model.decoder.bias
+
+        for my_layer, small_layer in zip(self.transformer_encoder.layers, small_model.transformer_encoder.layers):
+            small_hid_dim = small_layer.linear1.out_features
+            my_in_dim = my_layer.linear1.in_features
+
+            # packed along q,k,v order in first dim
+            my_in_proj_w = my_layer.self_attn.in_proj_weight
+            small_in_proj_w = small_layer.self_attn.in_proj_weight
+
+            my_in_proj_w.view(3, my_in_dim, my_in_dim)[:, :small_in_dim, :small_in_dim] = small_in_proj_w.view(3,
+                                                                                                               small_in_dim,
+                                                                                                               small_in_dim)
+            my_layer.self_attn.in_proj_bias.view(3, my_in_dim)[:,
+            :small_in_dim] = small_layer.self_attn.in_proj_bias.view(3, small_in_dim)
+
+            my_layer.self_attn.out_proj.weight[:small_in_dim, :small_in_dim] = small_layer.self_attn.out_proj.weight
+            my_layer.self_attn.out_proj.bias[:small_in_dim] = small_layer.self_attn.out_proj.bias
+
+            my_layer.linear1.weight[:small_hid_dim, :small_in_dim] = small_layer.linear1.weight
+            my_layer.linear1.bias[:small_hid_dim] = small_layer.linear1.bias
+
+            my_layer.linear2.weight[:small_in_dim, :small_hid_dim] = small_layer.linear2.weight
+            my_layer.linear2.bias[:small_in_dim] = small_layer.linear2.bias
+
+            my_layer.norm1.weight[:small_in_dim] = math.sqrt(small_in_dim / my_in_dim) * small_layer.norm1.weight
+            my_layer.norm2.weight[:small_in_dim] = math.sqrt(small_in_dim / my_in_dim) * small_layer.norm2.weight
+
+            my_layer.norm1.bias[:small_in_dim] = small_layer.norm1.bias
+            my_layer.norm2.bias[:small_in_dim] = small_layer.norm2.bias
+
+
+class TransformerEncoderDiffInit(Module):
+    r"""TransformerEncoder is a stack of N encoder layers
+
+    Args:
+        encoder_layer_creator: a function generating objects of TransformerEncoderLayer class without args (required).
+        num_layers: the number of sub-encoder-layers in the encoder (required).
+        norm: the layer normalization component (optional).
+    """
+    __constants__ = ['norm']
+
+    def __init__(self, encoder_layer_creator, num_layers, norm=None):
+        super().__init__()
+        self.layers = nn.ModuleList([encoder_layer_creator() for _ in range(num_layers)])
+        self.num_layers = num_layers
+        self.norm = norm
+
+    def forward(self, src: Tensor, mask: Optional[Tensor] = None, src_key_padding_mask: Optional[Tensor] = None) -> Tensor:
+        r"""Pass the input through the encoder layers in turn.
+
+        Args:
+            src: the sequence to the encoder (required).
+            mask: the mask for the src sequence (optional).
+            src_key_padding_mask: the mask for the src keys per batch (optional).
+
+        Shape:
+            see the docs in Transformer class.
+        """
+        output = src
+
+        for mod in self.layers:
+            output = mod(output, src_mask=mask, src_key_padding_mask=src_key_padding_mask)
+
+        if self.norm is not None:
+            output = self.norm(output)
+
+        return output

pfns4bo/__init__.py

@@ -0,0 +1,50 @@
+import os
+
+model_path = 'final_models'
+
+def prepare_models():
+    pfns4bo_dir = os.path.dirname(__file__)
+    model_names = ['hebo_morebudget_9_unused_features_3_userpriorperdim2_8.pt',
+                   'model_sampled_warp_simple_mlp_for_hpob_46.pt',
+                   'model_hebo_morebudget_9_unused_features_3.pt',]
+
+    for name in model_names:
+        weights_path = os.path.join(pfns4bo_dir, model_path, name)
+        compressed_weights_path = os.path.join(pfns4bo_dir, model_path, name + '.gz')
+        if not os.path.exists(weights_path):
+            if not os.path.exists(compressed_weights_path):
+                print("Downloading", os.path.abspath(compressed_weights_path))
+                import requests
+                url = f'https://github.com/automl/PFNs4BO/raw/main/pfns4bo/final_models/{name + ".gz"}'
+                r = requests.get(url, allow_redirects=True)
+                os.makedirs(os.path.dirname(compressed_weights_path), exist_ok=True)
+                with open(compressed_weights_path, 'wb') as f:
+                    f.write(r.content)
+            if os.path.exists(compressed_weights_path):
+                print("Unzipping", name)
+                os.system(f"gzip -dk {compressed_weights_path}")
+            else:
+                print("Failed to find", compressed_weights_path)
+                print("Make sure you have an internet connection to download the model automatically..")
+        if os.path.exists(weights_path):
+            print("Successfully located model at", weights_path)
+
+
+model_dict = {
+    'hebo_plus_userprior_model': os.path.join(os.path.dirname(__file__),model_path,
+                                              'hebo_morebudget_9_unused_features_3_userpriorperdim2_8.pt'),
+    'hebo_plus_model': os.path.join(os.path.dirname(__file__),model_path,
+                                    'model_hebo_morebudget_9_unused_features_3.pt'),
+    'bnn_model': os.path.join(os.path.dirname(__file__),model_path,'model_sampled_warp_simple_mlp_for_hpob_46.pt')
+}
+
+
+def __getattr__(name):
+    if name in model_dict:
+        if not os.path.exists(model_dict[name]):
+            print("Can't find", os.path.abspath(model_dict[name]), "thus unzipping/downloading models now.")
+            print("This might take a while..")
+            prepare_models()
+        return model_dict[name]
+    raise AttributeError(f"module '{__name__}' has no attribute '{name}'")
+