Upload app.py

app.py

@@ -0,0 +1,468 @@
+import gradio as gr
+import torch
+import numpy as np
+import matplotlib.pyplot as plt
+from test_functions.Ackley10D import *
+from test_functions.Ackley2D import *
+from test_functions.Ackley6D import *
+from test_functions.HeatExchanger import *
+from test_functions.CantileverBeam import *
+from test_functions.Car import *
+from test_functions.CompressionSpring import *
+from test_functions.GKXWC1 import *
+from test_functions.GKXWC2 import *
+from test_functions.HeatExchanger import *
+from test_functions.JLH1 import *
+from test_functions.JLH2 import *
+from test_functions.KeaneBump import *
+from test_functions.GKXWC1 import *
+from test_functions.GKXWC2 import *
+from test_functions.PressureVessel import *
+from test_functions.ReinforcedConcreteBeam import *
+from test_functions.SpeedReducer import *
+from test_functions.ThreeTruss import *
+from test_functions.WeldedBeam import *
+# Import other objective functions as needed
+import time
+
+from Rosen_PFN4BO import *
+from PIL import Image
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+def s(input_string):
+    return input_string
+
+
+
+
+def optimize(objective_function, iteration_input, progress=gr.Progress()):
+
+    print(objective_function)
+
+    # Variable setup
+    Current_BEST = torch.tensor( -1e10 )   # Some arbitrary very small number
+    Prev_BEST = torch.tensor( -1e10 )
+
+    if objective_function=="CantileverBeam.png":
+        Current_BEST = torch.tensor( -82500  )  # Some arbitrary very small number
+        Prev_BEST = torch.tensor( -82500 )
+    elif objective_function=="CompressionSpring.png":
+        Current_BEST = torch.tensor( -8  )  # Some arbitrary very small number
+        Prev_BEST = torch.tensor( -8 )
+    elif objective_function=="HeatExchanger.png":
+        Current_BEST = torch.tensor( -30000 )   # Some arbitrary very small number
+        Prev_BEST = torch.tensor( -30000 )
+    elif objective_function=="ThreeTruss.png":
+        Current_BEST = torch.tensor( -300  )  # Some arbitrary very small number
+        Prev_BEST = torch.tensor( -300 )
+    elif objective_function=="Reinforcement.png":
+        Current_BEST = torch.tensor( -440   ) # Some arbitrary very small number
+        Prev_BEST = torch.tensor( -440 )
+    elif objective_function=="PressureVessel.png":
+        Current_BEST = torch.tensor( -40000  )  # Some arbitrary very small number
+        Prev_BEST = torch.tensor( -40000    ) 
+    elif objective_function=="SpeedReducer.png":
+        Current_BEST = torch.tensor( -3200  )  # Some arbitrary very small number
+        Prev_BEST = torch.tensor( -3200  )   
+    elif objective_function=="WeldedBeam.png":
+        Current_BEST = torch.tensor( -35  )  # Some arbitrary very small number
+        Prev_BEST = torch.tensor( -35   )
+    elif objective_function=="Car.png":
+        Current_BEST = torch.tensor( -35  )  # Some arbitrary very small number
+        Prev_BEST = torch.tensor( -35   )
+
+    # Initial random samples
+    # print(objective_functions)
+    trained_X = torch.rand(20, objective_functions[objective_function]['dim'])
+
+    # Scale it to the domain of interest using the selected function
+    # print(objective_function)
+    X_Scaled = objective_functions[objective_function]['scaling'](trained_X)
+
+    # Get the constraints and objective
+    trained_gx, trained_Y = objective_functions[objective_function]['function'](X_Scaled)
+
+    # Convergence list to store best values
+    convergence = []
+    time_conv = []
+
+    START_TIME = time.time()
+
+
+# with gr.Progress(track_tqdm=True) as progress:
+
+
+    # Optimization Loop
+    for ii in progress.tqdm(range(iteration_input)):  # Example with 100 iterations
+
+        # (0) Get the updated data for this iteration
+        X_scaled = objective_functions[objective_function]['scaling'](trained_X)
+        trained_gx, trained_Y = objective_functions[objective_function]['function'](X_scaled)
+
+        # (1) Randomly sample Xpen 
+        X_pen = torch.rand(1000,trained_X.shape[1])
+
+        # (2) PFN inference phase with EI
+        default_model = 'final_models/model_hebo_morebudget_9_unused_features_3.pt'
+        
+        ei, p_feas = Rosen_PFN_Parallel(default_model,
+                                           trained_X, 
+                                           trained_Y, 
+                                           trained_gx,
+                                           X_pen,
+                                           'power',
+                                           'ei'
+                                          )
+
+        # Calculating CEI
+        CEI = ei
+        for jj in range(p_feas.shape[1]):
+            CEI = CEI*p_feas[:,jj]
+
+        # (4) Get the next search value
+        rec_idx = torch.argmax(CEI)
+        best_candidate = X_pen[rec_idx,:].unsqueeze(0)
+
+        # (5) Append the next search point
+        trained_X = torch.cat([trained_X, best_candidate])
+
+
+        ################################################################################
+        # This is just for visualizing the best value. 
+        # This section can be remove for pure optimization purpose
+        Current_X = objective_functions[objective_function]['scaling'](trained_X)
+        Current_GX, Current_Y = objective_functions[objective_function]['function'](Current_X)
+        if ((Current_GX<=0).all(dim=1)).any():
+            Current_BEST = torch.max(Current_Y[(Current_GX<=0).all(dim=1)])
+        else:
+            Current_BEST = Prev_BEST
+        ################################################################################
+        
+        # (ii) Convergence tracking (assuming the best Y is to be maximized)
+        # if Current_BEST != -1e10:
+        print(Current_BEST)
+        print(convergence)
+        convergence.append(Current_BEST.abs())
+        time_conv.append(time.time() - START_TIME)
+
+    # Timing
+    END_TIME = time.time()
+    TOTAL_TIME = END_TIME - START_TIME
+    
+    # Website visualization
+    # (i) Radar chart for trained_X
+    radar_chart = None
+    # radar_chart = create_radar_chart(X_scaled)
+    # (ii) Convergence tracking (assuming the best Y is to be maximized)
+    convergence_plot = create_convergence_plot(objective_function, iteration_input, 
+                                               time_conv, 
+                                               convergence, TOTAL_TIME)
+
+
+    return convergence_plot
+    # return radar_chart, convergence_plot
+
+
+
+
+
+
+
+
+
+def create_radar_chart(X_scaled):
+    fig, ax = plt.subplots(figsize=(6, 6), subplot_kw=dict(polar=True))
+    labels = [f'x{i+1}' for i in range(X_scaled.shape[1])]
+    values = X_scaled.mean(dim=0).numpy()
+    
+    num_vars = len(labels)
+    angles = np.linspace(0, 2 * np.pi, num_vars, endpoint=False).tolist()
+    values = np.concatenate((values, [values[0]]))
+    angles += angles[:1]
+
+    ax.fill(angles, values, color='green', alpha=0.25)
+    ax.plot(angles, values, color='green', linewidth=2)
+    ax.set_yticklabels([])
+    ax.set_xticks(angles[:-1])
+    # ax.set_xticklabels(labels)
+    ax.set_xticklabels([f'{label}\n({value:.2f})' for label, value in zip(labels, values[:-1])])  # Show values
+    ax.set_title("Selected Design", size=15, color='black', y=1.1)
+    
+    plt.close(fig)
+    return fig
+
+
+
+
+
+
+
+def create_convergence_plot(objective_function, iteration_input, time_conv, convergence, TOTAL_TIME):
+    fig, ax = plt.subplots()
+    
+    # Realtime optimization data
+    ax.plot(time_conv, convergence, '^-', label='PFN-CBO (Realtime)' )
+
+    # Stored GP data
+    if objective_function=="CantileverBeam.png":
+        GP_TIME = torch.load('CantileverBeam_CEI_Avg_Time.pt')
+        GP_OBJ = torch.load('CantileverBeam_CEI_Avg_Obj.pt')
+        
+    elif objective_function=="CompressionSpring.png":
+        GP_TIME = torch.load('CompressionSpring_CEI_Avg_Time.pt')
+        GP_OBJ = torch.load('CompressionSpring_CEI_Avg_Obj.pt')
+
+    elif objective_function=="HeatExchanger.png":
+        GP_TIME = torch.load('HeatExchanger_CEI_Avg_Time.pt')
+        GP_OBJ = torch.load('HeatExchanger_CEI_Avg_Obj.pt')
+        
+    elif objective_function=="ThreeTruss.png":
+        GP_TIME = torch.load('ThreeTruss_CEI_Avg_Time.pt')
+        GP_OBJ = torch.load('ThreeTruss_CEI_Avg_Obj.pt')
+        
+    elif objective_function=="Reinforcement.png":
+        GP_TIME = torch.load('ReinforcedConcreteBeam_CEI_Avg_Time.pt')
+        GP_OBJ = torch.load('ReinforcedConcreteBeam_CEI_Avg_Obj.pt')
+        
+    elif objective_function=="PressureVessel.png":
+        GP_TIME = torch.load('PressureVessel_CEI_Avg_Time.pt')
+        GP_OBJ = torch.load('PressureVessel_CEI_Avg_Obj.pt')
+        
+    elif objective_function=="SpeedReducer.png":
+        GP_TIME = torch.load('SpeedReducer_CEI_Avg_Time.pt')
+        GP_OBJ = torch.load('SpeedReducer_CEI_Avg_Obj.pt')
+        
+    elif objective_function=="WeldedBeam.png":
+        GP_TIME = torch.load('WeldedBeam_CEI_Avg_Time.pt')
+        GP_OBJ = torch.load('WeldedBeam_CEI_Avg_Obj.pt')  
+
+    elif objective_function=="Car.png":
+        GP_TIME = torch.load('Car_CEI_Avg_Time.pt')
+        GP_OBJ = torch.load('Car_CEI_Avg_Obj.pt')    
+        
+    # Plot GP data    
+    ax.plot(GP_TIME[:iteration_input], GP_OBJ[:iteration_input], '^-', label='GP-CBO (Data)' )
+
+    
+    ax.set_xlabel('Time (seconds)')
+    ax.set_ylabel('Objective Value')
+    ax.set_title('Convergence Plot for {t} iterations'.format(t=iteration_input))
+    # ax.legend()
+
+    if objective_function=="CantileverBeam.png":
+        ax.axhline(y=50000, color='red', linestyle='--', label='Optimal Value')
+
+    elif objective_function=="CompressionSpring.png":
+        ax.axhline(y=0, color='red', linestyle='--', label='Optimal Value')
+
+    elif objective_function=="HeatExchanger.png":
+        ax.axhline(y=4700, color='red', linestyle='--', label='Optimal Value')
+        
+    elif objective_function=="ThreeTruss.png":
+        ax.axhline(y=262, color='red', linestyle='--', label='Optimal Value')
+        
+    elif objective_function=="Reinforcement.png":
+        ax.axhline(y=355, color='red', linestyle='--', label='Optimal Value')
+        
+    elif objective_function=="PressureVessel.png":
+        ax.axhline(y=5000, color='red', linestyle='--', label='Optimal Value')
+        
+    elif objective_function=="SpeedReducer.png":
+        ax.axhline(y=2650, color='red', linestyle='--', label='Optimal Value')
+        
+    elif objective_function=="WeldedBeam.png":
+        ax.axhline(y=6, color='red', linestyle='--', label='Optimal Value')  
+
+    elif objective_function=="Car.png":
+        ax.axhline(y=25, color='red', linestyle='--', label='Optimal Value')  
+
+    
+    ax.legend(loc='best')
+    # ax.legend(loc='lower left')
+        
+
+    # Add text to the top right corner of the plot
+    if len(convergence) == 0:
+        ax.text(0.5, 0.5, 'No Feasible Design Found', transform=ax.transAxes, fontsize=12,
+                verticalalignment='top', horizontalalignment='right')
+        
+    
+    plt.close(fig)
+    return fig
+
+
+
+
+
+
+# Define available objective functions
+objective_functions = {
+    # "ThreeTruss.png": {"image": "ThreeTruss.png", 
+    #                     "function": ThreeTruss, 
+    #                     "scaling": ThreeTruss_Scaling, 
+    #                     "dim": 2},
+    "CompressionSpring.png": {"image": "CompressionSpring.png", 
+                               "function": CompressionSpring, 
+                               "scaling": CompressionSpring_Scaling, 
+                               "dim": 3},
+    "Reinforcement.png": {"image": "Reinforcement.png", "function": ReinforcedConcreteBeam, "scaling": ReinforcedConcreteBeam_Scaling, "dim": 3},
+    "PressureVessel.png": {"image": "PressureVessel.png", "function": PressureVessel, "scaling": PressureVessel_Scaling, "dim": 4},
+    "SpeedReducer.png": {"image": "SpeedReducer.png", "function": SpeedReducer, "scaling": SpeedReducer_Scaling, "dim": 7},
+    "WeldedBeam.png": {"image": "WeldedBeam.png", "function": WeldedBeam, "scaling": WeldedBeam_Scaling, "dim": 4},
+    "HeatExchanger.png": {"image": "HeatExchanger.png", "function": HeatExchanger, "scaling": HeatExchanger_Scaling, "dim": 8},
+    "CantileverBeam.png": {"image": "CantileverBeam.png", "function": CantileverBeam, "scaling": CantileverBeam_Scaling, "dim": 10},
+    "Car.png": {"image": "Car.png", "function": Car, "scaling": Car_Scaling, "dim": 11},
+}
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+# Extract just the image paths for the gallery
+image_paths = [key for key in objective_functions]
+
+
+def submit_action(objective_function_choices, iteration_input):
+    # print(iteration_input)
+    # print(len(objective_function_choices))
+    # print(objective_functions[objective_function_choices]['function'])
+    if len(objective_function_choices)>0:
+        selected_function = objective_functions[objective_function_choices]['function']
+        return  optimize(objective_function_choices, iteration_input)
+    return None
+
+# Function to clear the output
+def clear_output():
+    # print(gallery.selected_index)
+    
+    return gr.update(value=[], selected=None),  None, 15, gr.Markdown(""), 'Test_formulation_default.png'
+
+def reset_gallery():
+    return gr.update(value=image_paths)
+
+
+with gr.Blocks() as demo:
+    # Centered Title and Description using gr.HTML
+    gr.HTML(
+        """
+        <div style="text-align: center;">
+            <h1>Pre-trained Transformer for Constrained Bayesian Optimization</h1>
+            <h4>Paper: <a href="https://arxiv.org/abs/2404.04495">
+            Fast and Accurate Bayesian Optimization with Pre-trained Transformers for Constrained Engineering Problems</a> 
+            </h4>
+
+            <p style="text-align: left;">This is a demo for Bayesian Optimization using PFN (Prior-Data Fitted Networks). 
+            Select your objective function by clicking on one of the check boxes below, then enter the iteration number to run the optimization process. 
+            The results will be visualized in the radar chart and convergence plot.</p>
+            
+            
+            
+
+        </div>
+        """
+    )
+
+    
+    with gr.Row():
+        
+        
+        with gr.Column(variant='compact'):
+            # gr.Markdown("# Inputs: ")
+            
+            with gr.Row():
+                gr.Markdown("## Select a problem (objective): ")
+                img_key = gr.Markdown(value="", visible=False)
+            
+            gallery = gr.Gallery(value=image_paths, label="Objective Functions", 
+                                 # height = 450, 
+                                 object_fit='contain',
+                                 columns=3, rows=3, elem_id="gallery")
+            
+            gr.Markdown("## Enter iteration Number: ")
+            iteration_input = gr.Slider(label="Iterations:", minimum=15, maximum=50, step=1, value=15)
+        
+
+            # Row for the Clear and Submit buttons
+            with gr.Row():
+                clear_button = gr.Button("Clear")
+                submit_button = gr.Button("Submit", variant="primary")
+
+        with gr.Column():
+            # gr.Markdown("# Outputs: ")
+            gr.Markdown("## Problem Formulation: ")
+            formulation = gr.Image(value='Formulation_default.png', height=150)
+            gr.Markdown("## Results: ")
+            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.")
+            convergence_plot = gr.Plot(label="Convergence Plot")
+
+
+
+    def handle_select(evt: gr.SelectData):
+        selected_image = evt.value
+        key = evt.value['image']['orig_name']
+        formulation = 'Test_formulation.png'
+        print('here')
+        print(key)
+
+        return key, formulation
+        
+    gallery.select(fn=handle_select, inputs=None, outputs=[img_key, formulation])
+
+
+    
+    submit_button.click(
+        submit_action,
+        inputs=[img_key, iteration_input],
+        # outputs= [radar_plot, convergence_plot],
+        outputs= convergence_plot,
+        
+        # progress=True  # Enable progress tracking
+        
+    )
+
+    clear_button.click(
+        clear_output,
+        inputs=None,
+        outputs=[gallery, convergence_plot, iteration_input, img_key, formulation]
+    ).then(
+        # Step 2: Reset the gallery to the original list
+        reset_gallery,
+        inputs=None,
+        outputs=gallery
+    )
+
+    
+
+demo.launch()