adding more plot types

app.py

@@ -26,9 +26,15 @@ def read_result_from_hub(filename):
     )
     return local_path
 
-def make_visual(boundary):
+def make_boundary_plot(boundary):
     boundary_plot = visualization.plot_boundary(boundary)
+    return boundary_plot
+
+def make_interactive_plot(boundary):
     interactive_plot = visualization.plot_surface(boundary)
+    return interactive_plot
+
+def make_boozer_plot(boundary):
     settings = forward_model.ConstellarationSettings.default_high_fidelity_skip_qi()
     boundary_metrics, boundary_equilibrium = forward_model.forward_model(
         boundary, settings=settings
@@ -37,10 +43,17 @@ def make_visual(boundary):
     boozer_plot = visualization.plot_boozer_surfaces(
         boundary_equilibrium, settings=boozer_settings
     )
+    return
+
+def make_flux_surface_plot(boundary):
+    settings = forward_model.ConstellarationSettings.default_high_fidelity_skip_qi()
+    boundary_metrics, boundary_equilibrium = forward_model.forward_model(
+        boundary, settings=settings
+    )
     flux_surface_plot = visualization.plot_flux_surfaces(
         boundary_equilibrium, boundary
     )
-    return boundary_plot, interactive_plot, boozer_plot, flux_surface_plot
+    return flux_surface_plot
 
 def gradio_interface() -> gr.Blocks:
     with gr.Blocks() as demo:
@@ -54,6 +67,8 @@ def gradio_interface() -> gr.Blocks:
         full_df = pd.DataFrame(ds)
         filenames = full_df['result_filename'].to_list()
 
+        boundary = gr.State()
+
         mode_selector = gr.Radio(choices=["Leaderboard", "Upload", "Generate"],
                                  label="Select input method:",
                                  value="Leaderboard")
@@ -95,11 +110,15 @@ def gradio_interface() -> gr.Blocks:
             else:
                 boundary = load_boundary(row['boundary_json'])
 
-            vis = make_visual(boundary)
-            return vis
+            return boundary
+
+        # 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) \
+            .then(make_boundary_plot, boundary, boundary_plot) \
+            .then(make_interactive_plot, boundary, interactive_plot) \
+            .then(make_boozer_plot, boundary, boozer_plot) \
+            .then(make_flux_surface_plot, boundary, flux_surface_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:
@@ -107,18 +126,25 @@ def gradio_interface() -> gr.Blocks:
             with open(uploaded_file.name, 'r') as f:
                 data = f.read()
             boundary = load_boundary(data)
-            return make_visual(boundary)
+            return boundary
         
-        upload_box.change(get_boundary_vis_from_upload, inputs=[upload_box], outputs=[boundary_plot, interactive_plot, boozer_plot, flux_surface_plot])
+        upload_box.change(get_boundary_vis_from_upload, inputs=[upload_box], outputs=boundary) \
+            .then(make_boundary_plot, boundary, boundary_plot) \
+            .then(make_interactive_plot, boundary, interactive_plot) \
+            .then(make_boozer_plot, boundary, boozer_plot) \
+            .then(make_flux_surface_plot, boundary, flux_surface_plot)
 
         def generate_random_boundary(aspect_ratio, elongation, rotational_transform, n_field_periods):
             boundary = initial_guess.generate_rotating_ellipse(
                 aspect_ratio=aspect_ratio, elongation=elongation, rotational_transform=rotational_transform, n_field_periods=n_field_periods
             )
-            vis = make_visual(boundary)
-            return vis
+            return boundary
         
-        generate_btn.click(generate_random_boundary, [aspect_ratio, elongation, rotational_transform, n_field_periods], [boundary_plot, interactive_plot, boozer_plot, flux_surface_plot])
+        generate_btn.click(generate_random_boundary, [aspect_ratio, elongation, rotational_transform, n_field_periods], boundary) \
+            .then(make_boundary_plot, boundary, boundary_plot) \
+            .then(make_interactive_plot, boundary, interactive_plot) \
+            .then(make_boozer_plot, boundary, boozer_plot) \
+            .then(make_flux_surface_plot, boundary, flux_surface_plot)
 
     return demo