adding more plot types

app.py

@@ -27,8 +27,20 @@ def read_result_from_hub(filename):
     return local_path
 
 def make_visual(boundary):
-    vis = visualization.plot_surface(boundary)
-    return vis
+    boundary_plot = visualization.plot_boundary(boundary)
+    interactive_plot = visualization.plot_surface(boundary)
+    settings = forward_model.ConstellarationSettings.default_high_fidelity_skip_qi()
+    boundary_metrics, boundary_equilibrium = forward_model.forward_model(
+        boundary, settings=settings
+    )
+    boozer_settings = boozer.BoozerSettings(normalized_toroidal_flux=[0.0, 0.25, 1])
+    boozer_plot = visualization.plot_boozer_surfaces(
+        boundary_equilibrium, settings=boozer_settings
+    )
+    flux_surface_plot = visualization.plot_flux_surfaces(
+        boundary_equilibrium, boundary
+    )
+    return boundary_plot, interactive_plot, boozer_plot, flux_surface_plot
 
 def gradio_interface() -> gr.Blocks:
     with gr.Blocks() as demo:
@@ -62,8 +74,10 @@ def gradio_interface() -> gr.Blocks:
                 n_field_periods = gr.Number(label="Number of Period Fields", value=3)
                 generate_btn = gr.Button(value="Generate")
 
-            
-            plot = gr.Plot()
+            boundary_plot = gr.Plot()
+            interactive_plot = gr.Plot()
+            boozer_plot = gr.Plot()
+            flux_surface_plot = gr.Plot()
 
         def update_ui(mode):
             return (
@@ -77,15 +91,15 @@ def gradio_interface() -> gr.Blocks:
         def get_boundary_from_leaderboard(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!")
+                raise gr.Error("Sorry this isn't implemented for files with multiple boundaries yet!")
             else:
                 boundary = load_boundary(row['boundary_json'])
 
             vis = make_visual(boundary)
             return vis
 
-        dropdown.change(get_boundary_from_leaderboard, dropdown, plot)
-        rld_btn.click(get_boundary_from_leaderboard, dropdown, plot)
+        dropdown.change(get_boundary_from_leaderboard, dropdown, [boundary_plot, interactive_plot, boozer_plot, flux_surface_plot])
+        rld_btn.click(get_boundary_from_leaderboard, dropdown, [boundary_plot, interactive_plot, boozer_plot, flux_surface_plot])
 
         def get_boundary_vis_from_upload(uploaded_file):
             if uploaded_file is None:
@@ -95,7 +109,7 @@ def gradio_interface() -> gr.Blocks:
             boundary = load_boundary(data)
             return make_visual(boundary)
         
-        upload_box.change(get_boundary_vis_from_upload, inputs=[upload_box], outputs=[plot])
+        upload_box.change(get_boundary_vis_from_upload, inputs=[upload_box], outputs=[boundary_plot, interactive_plot, boozer_plot, flux_surface_plot])
 
         def generate_random_boundary(aspect_ratio, elongation, rotational_transform, n_field_periods):
             boundary = initial_guess.generate_rotating_ellipse(
@@ -104,7 +118,7 @@ def gradio_interface() -> gr.Blocks:
             vis = make_visual(boundary)
             return vis
         
-        generate_btn.click(generate_random_boundary, [aspect_ratio, elongation, rotational_transform, n_field_periods], plot)
+        generate_btn.click(generate_random_boundary, [aspect_ratio, elongation, rotational_transform, n_field_periods], [boundary_plot, interactive_plot, boozer_plot, flux_surface_plot])
 
     return demo