app.py

import gradio as gr
import plotly.graph_objects as go
from huggingface_hub import hf_hub_download
from datasets import load_dataset
import pathlib
import json
import pandas as pd
from evaluation import load_boundary, load_boundaries
from constellaration import forward_model, initial_guess
from constellaration.boozer import boozer
from constellaration.utils import (
    file_exporter,
    visualization,
    visualization_utils,
)


organization = 'proxima-fusion'
results_repo = f'{organization}/constellaration-bench-results'

def read_result_from_hub(filename):
    local_path = hf_hub_download(
        repo_id=results_repo,
        repo_type="dataset",
        filename=filename,
    )
    return local_path

def make_visual(boundary):
    vis = visualization.plot_surface(boundary)
    return vis

def gradio_interface() -> gr.Blocks:
    with gr.Blocks() as demo:
        gr.Markdown("""
                    # Welcome to the ConStellaration Boundary Explorer!

                    ### Here, you can visualize submissions to the ConStellaration Leaderboard, generate and visualize new random boundaries, or upload and visualize your own!
                    """)
        
        ds = load_dataset(results_repo, split='train')
        full_df = pd.DataFrame(ds)
        filenames = full_df['result_filename'].to_list()
        with gr.Row():
            with gr.Column():
                dropdown = gr.Dropdown(choices=filenames, label="Choose a leaderboard entry", value=filenames[0])
                rld_btn = gr.Button(value="Reload")

            with gr.Column():
                plot = gr.Plot()

        def get_boundary_vis(selected_file):

            row = full_df[full_df['result_filename'] == selected_file].iloc[0]

            if row['problem_type'] == 'mhd_stable':
                raise gr.Error("Sorry this isn't implemented for mhd_stable submissions yet!")
            else:
                boundary = load_boundary(row['boundary_json'])

            vis = make_visual(boundary)
            return vis

        demo.load(get_boundary_vis, dropdown, plot)
        rld_btn.click(get_boundary_vis, dropdown, plot)
        

    return demo


if __name__ == "__main__":
    gradio_interface().launch()