app.py

import numpy as np
import gradio as gr
import pickle
import plotly.graph_objects as go
import plotly.express as px
import pandas as pd

# Define styling for success/failure highlighting
def highlight_success(val):
    color = 'lightgreen' if val == 'Success' else 'lightcoral'
    return f'color:white;background-color: {color}'

# Simulation function for electrospraying
def sim_espray_constrained(x, noise_se=None):
    # Ensure x is a numpy array with float data type
    x = np.array(x, dtype=float)
    
    # Ensure x is a 2D array
    if x.ndim == 1:
        x = x.reshape(1, -1)
        
    # Define the equations
    conc = x[:, 0]
    flow_rate = x[:, 1]
    voltage = x[:, 2]
    solvent = x[:, 3]
    diameter = (np.sqrt(conc) * np.sqrt(flow_rate)) / np.log2(voltage) * 10 + 0.4 + solvent  # Diameter in micrometers
    if noise_se is not None:
        diameter = diameter + noise_se * np.random.randn(*diameter.shape)
    exp_con = (np.log(flow_rate) * (solvent - 0.5) + 1.40 >= 0).astype(float)
    return np.column_stack((diameter, exp_con))

# Initialize experiment data
X_init = np.array([[0.5, 15, 10, 0],
                   [0.5, 0.1, 10, 1],
                   [3, 20, 15, 0],
                   [1, 20, 10, 1],
                   [0.2, 0.02, 10, 1]])

Y_init = sim_espray_constrained(X_init)
exp_record_df = pd.DataFrame(X_init, columns=['Concentration (%w/v)', 'Flow Rate (mL/h)', 'Voltage (kV)', 'Solvent'])
exp_record_df['Size (um)'] = Y_init[:, 0]
exp_record_df['Solvent'] = ['DMAc' if x == 0 else 'CHCl3' for x in exp_record_df['Solvent']]
exp_record_df['Feasible?'] = ['Success' if x == 1 else 'Failed' for x in Y_init[:, 1]]
prior_experiments_display = exp_record_df.style.map(highlight_success, subset=['Feasible?']).format(precision=3)

# Functions for data processing and visualization
def import_results():
    strategies = ['qEI', 'qEI_vi_mixed_con', 'qEICF_vi_mixed_con', 'rnd']
    
    # Load results from pickle file
    with open('best_distances.pkl', 'rb') as f:
        best_distances = pickle.load(f)

    # vstack all values in best_distances
    best_distances_vstack = {k: np.vstack(best_distances[k]) for k in strategies}
    best_distances_all_trials = -np.vstack([best_distances_vstack[k] for k in strategies])
    best_distances_all_trials_df = pd.DataFrame(best_distances_all_trials)
    best_distances_all_trials_df['strategy'] = np.repeat(['Vanilla BO', 'Constrained BO', 'CCBO', 'Random'], 20)
    best_distances_all_trials_df['trial'] = list(range(20)) * len(strategies)

    best_distances_df_long = pd.melt(best_distances_all_trials_df, id_vars=['strategy', 'trial'], var_name='iteration', value_name='regret')
    return best_distances_df_long

def calc_human_performance(df):
    # Make a copy of the dataframe to avoid modifying the original
    df_copy = df.copy()
    
    # convert back solvent to 0 and 1
    df_copy['Solvent'] = [0 if x == 'DMAc' else 1 for x in df_copy['Solvent']]
    
    TARGET_SIZE = 3.0  # Example target size
    ROUNDS = len(df_copy) // 2

    # Ensure all values are numeric
    numeric_cols = ['Concentration (%w/v)', 'Flow Rate (mL/h)', 'Voltage (kV)', 'Solvent']
    for col in numeric_cols:
        df_copy[col] = pd.to_numeric(df_copy[col])
    
    X_human = df_copy[numeric_cols].values

    X_human_init = X_init.copy()
    Y_human_init = Y_init.copy()
    
    best_human_distance = []

    for iter in range(ROUNDS + 1):
        Y_distance = -np.abs(Y_human_init[:, 0] - TARGET_SIZE)
        best_human_distance.append(np.ma.masked_array(Y_distance, mask=~Y_human_init[:, 1].astype(bool)).max())
        
        # Check if we have more data for this iteration
        if 2 * iter < len(X_human):
            # Get the slice of new experiments
            new_x = X_human[2 * iter:min(2 * (iter + 1), len(X_human))]
            
            # Add the new experiments to our dataset
            X_human_init = np.vstack([X_human_init, new_x])
            Y_human_init = np.vstack([Y_human_init, sim_espray_constrained(new_x)])

    return -np.array(best_human_distance)

def plot_results(exp_data_df):
    # Extract human performance
    best_human_distance = calc_human_performance(exp_data_df)
    
    # Import results
    best_distances_df_long = import_results()
    
    fig = go.Figure()
    
    strategies = best_distances_df_long['strategy'].unique()
    
    for strategy in strategies:
        strategy_data = best_distances_df_long[best_distances_df_long['strategy'] == strategy]
        
        # Calculate mean and standard deviation
        mean_regret = strategy_data.groupby('iteration')['regret'].mean()
        std_regret = strategy_data.groupby('iteration')['regret'].std()
        
        iterations = mean_regret.index
        color = px.colors.qualitative.Set2[strategies.tolist().index(strategy)]
        
        # Add trace for mean line
        mean_trace = go.Scatter(
            x=iterations,
            y=mean_regret,
            mode='lines',
            name=strategy,
            line=dict(width=2, color=color)
        )
        fig.add_trace(mean_trace)
        
        # Add trace for shaded area (standard deviation)
        fig.add_trace(go.Scatter(
            x=list(iterations) + list(iterations[::-1]),
            y=list(mean_regret + std_regret) + list((mean_regret - std_regret)[::-1]),
            fill='toself',
            fillcolor=mean_trace.line.color.replace('rgb', 'rgba').replace(')', ',0.2)'),
            line=dict(color='rgba(255,255,255,0)'),
            showlegend=False,
            name=f'{strategy} (std dev)'
        ))
    # Add trace for human performance
    fig.add_trace(go.Scatter(
        x=list(range(len(best_human_distance))),
        y=best_human_distance,
        mode='lines+markers',
        name='Human',
        line=dict(width=2, color='brown')
    ))

    fig.update_layout(
        title='Performance Comparison',
        xaxis_title='Iteration',
        yaxis_title='Regret (μm)',
        legend_title='Strategy',
        template='plotly_white',
        legend=dict(
            x=0.01,
            y=0.01,
            bgcolor='rgba(255, 255, 255, 0.5)',
            bordercolor='rgba(0, 0, 0, 0.5)',
            borderwidth=1
        )
    )
    
    return fig

# Prediction function - simplified signature by removing unnecessary text params
def predict(state, conc1, flow_rate1, voltage1, solvent1, conc2, flow_rate2, voltage2, solvent2):
    # Get current results storage from state or initialize if None
    if state is None:
        results_storage = pd.DataFrame(columns=['Concentration (%w/v)', 'Flow Rate (mL/h)', 'Voltage (kV)', 'Solvent', 'Size (um)', 'Feasible?'])
    else:
        results_storage = state.copy()
        
    solvent_value1 = 0 if solvent1 == 'DMAc' else 1
    solvent_value2 = 0 if solvent2 == 'DMAc' else 1
    
    # Process inputs and get predictions
    inputs1 = np.array([[conc1, flow_rate1, voltage1, solvent_value1]])
    inputs2 = np.array([[conc2, flow_rate2, voltage2, solvent_value2]])
    results1 = sim_espray_constrained(inputs1)
    results2 = sim_espray_constrained(inputs2)
    
    # Format and store results
    results_df = pd.DataFrame([
        [conc1, flow_rate1, voltage1, solvent_value1, results1[0, 0], results1[0, 1]],
        [conc2, flow_rate2, voltage2, solvent_value2, results2[0, 0], results2[0, 1]]
    ], columns=['Concentration (%w/v)', 'Flow Rate (mL/h)', 'Voltage (kV)', 'Solvent', 'Size (um)', 'Feasible?'])

    results_df['Solvent'] = ['DMAc' if x == 0 else 'CHCl3' for x in results_df['Solvent']]
    results_df['Feasible?'] = ['Success' if x == 1 else 'Failed' for x in results_df['Feasible?']]
    
    results_storage = pd.concat([results_storage, results_df], ignore_index=True)
    results_display = results_storage.style.map(highlight_success, subset=['Feasible?']).format(precision=3)
   
    # Return updated state and UI updates
    return (
        results_storage, 
        gr.DataFrame(value=prior_experiments_display, label="Prior Experiments"),
        gr.DataFrame(value=results_display, label="Your Results"), 
        plot_results(results_storage)
    )

# Reset results function
def reset_results(state):
    results_storage = pd.DataFrame(columns=['Concentration (%w/v)', 'Flow Rate (mL/h)', 'Voltage (kV)', 'Solvent', 'Size (um)', 'Feasible?'])
    return (
        results_storage, 
        gr.DataFrame(value=prior_experiments_display, label="Prior Experiments"),
        gr.DataFrame(value=results_storage.style.map(highlight_success, subset=['Feasible?']).format(precision=3), label="Your Results"), 
        plot_results(results_storage)
    )

# Application description
description = "<h3>Welcome, challenger!</h3><p> If you think you may perform better than <strong>CCBO</strong>, try this interactive game to optimize electrospray! Rules are simple: <ul><li>Examine carefully the initial experiments you have on the right (or below if you're using your phone), remember, your target size is <u><i><strong>3.000 um</strong></i></u> ----></li><li>Select your parameters, you have <strong>2</strong> experiments (chances) in each round, use them wisely! </li><li>Click <strong>Submit</strong> to see the results, reflect and improve your selection!</li><li>Repeat the process for <strong>5</strong> rounds to see if you can beat CCBO!</li></ul></p><p>Your data will not be stored, so feel free to play again, good luck!</p>"

# Create Gradio interface
with gr.Blocks() as demo:
    gr.Markdown("## Human vs CCBO Campaign - Optimize Electrospray")
    gr.Markdown(description)
    
    # Add state component to store user-specific results
    results_state = gr.State()
    
    with gr.Row():
        # Input parameters column
        with gr.Column():
            gr.Markdown("### Experiment 1")
            conc1 = gr.Slider(minimum=0.05, maximum=5.0, value=1.2, step=0.001, label="Concentration (%w/v)")
            flow_rate1 = gr.Slider(minimum=0.01, maximum=60.0, value=20.0, step=0.001, label="Flow Rate (mL/h)")
            voltage1 = gr.Slider(minimum=10.0, maximum=18.0, value=15.0, step=0.001, label="Voltage (kV)")
            solvent1 = gr.Dropdown(['DMAc', 'CHCl3'], value='DMAc', label='Solvent')
            
            gr.Markdown("### Experiment 2")
            conc2 = gr.Slider(minimum=0.05, maximum=5.0, value=2.8, step=0.001, label="Concentration (%w/v)")
            flow_rate2 = gr.Slider(minimum=0.01, maximum=60.0, value=20.0, step=0.001, label="Flow Rate (mL/h)")
            voltage2 = gr.Slider(minimum=10.0, maximum=18.0, value=15.0, step=0.001, label="Voltage (kV)")
            solvent2 = gr.Dropdown(['DMAc', 'CHCl3'], value='CHCl3', label='Solvent')
        
        # Results display column
        with gr.Column():
            prior_experiments = gr.DataFrame(value=prior_experiments_display, label="Prior Experiments")
            results_df = gr.DataFrame(label="Your Results")
            perf_plot = gr.Plot(label="Performance Comparison")
    
    # Action buttons
    with gr.Row():
        submit_btn = gr.Button("Submit")
        reset_btn = gr.Button("Reset Results")
    
    # Connect the submit button to the predict function
    submit_btn.click(
        fn=predict,
        inputs=[
            results_state, 
            conc1, flow_rate1, voltage1, solvent1,
            conc2, flow_rate2, voltage2, solvent2
        ],
        outputs=[results_state, prior_experiments, results_df, perf_plot]
    )
    
    # Connect the reset button to the reset_results function
    reset_btn.click(
        fn=reset_results,
        inputs=[results_state],
        outputs=[results_state, prior_experiments, results_df, perf_plot]
    )

demo.launch()