feat: store model selection when switching pages (#96)

* feat: add `fetch_selected_models`

This function is not yet stable and requires some iteration.

* feat: remember selected models between pages

* feat: remove spurious title in sidebar + some missing data fixes

* feat: make the default all models

* fix: failing tests

* feat: final fix to model selection

* test: update model selection import

* feat: ensure model names in correct order

* fix: identical model names

Author: lwalew

src/mlipaudit/app.py

@@ -25,6 +25,7 @@
 from mlipaudit.io import load_benchmark_results_from_disk, load_scores_from_disk
 from mlipaudit.ui import leaderboard_page
 from mlipaudit.ui.page_wrapper import UIPageWrapper
+from mlipaudit.ui.utils import model_selection
 
 
 def _data_func_from_key(
@@ -156,6 +157,10 @@ def main() -> None:
     else:
         pages_to_show = [leaderboard] + page_categories[selected_category]
 
+    # Add model selection
+    with st.sidebar:
+        model_selection(unique_model_names=list(results.keys()))
+
     # Set up navigation in main area
     pg = st.navigation(pages_to_show)
 

src/mlipaudit/ui/bond_length_distribution.py

@@ -22,7 +22,7 @@
     BondLengthDistributionResult,
 )
 from mlipaudit.ui.page_wrapper import UIPageWrapper
-from mlipaudit.ui.utils import display_model_scores
+from mlipaudit.ui.utils import display_model_scores, fetch_selected_models
 
 ModelName: TypeAlias = str
 BenchmarkResultForMultipleModels: TypeAlias = dict[
@@ -53,7 +53,6 @@ def bond_length_distribution_page(
                    keys and the benchmark results objects as values.
     """
     st.markdown("# Bond length distribution")
-    st.sidebar.markdown("# Bond length distribution")
 
     st.markdown(
         "The benchmark runs short simulations of small molecules to check whether the "
@@ -84,11 +83,11 @@ def bond_length_distribution_page(
         st.markdown("**No results to display**.")
         return
 
-    unique_model_names = list(set(data.keys()))
-    model_select = st.sidebar.multiselect(
-        "Select model(s)", unique_model_names, default=unique_model_names
-    )
-    selected_models = model_select if model_select else unique_model_names
+    selected_models = fetch_selected_models(available_models=list(data.keys()))
+
+    if not selected_models:
+        st.markdown("**No results to display**.")
+        return
 
     distribution_data = [
         {

src/mlipaudit/ui/conformer_selection.py

@@ -22,7 +22,11 @@
 
 from mlipaudit.benchmarks import ConformerSelectionBenchmark, ConformerSelectionResult
 from mlipaudit.ui.page_wrapper import UIPageWrapper
-from mlipaudit.ui.utils import create_st_image, display_model_scores
+from mlipaudit.ui.utils import (
+    create_st_image,
+    display_model_scores,
+    fetch_selected_models,
+)
 
 APP_DATA_DIR = Path(__file__).parent.parent / "app_data"
 CONFORMER_IMG_DIR = APP_DATA_DIR / "conformer_selection" / "img"
@@ -35,6 +39,7 @@ def _process_data_into_dataframe(
     selected_models: list[str],
 ) -> pd.DataFrame:
     converted_data_scores = []
+    model_names = []
     for model_name, results in data.items():
         if model_name in selected_models:
             model_data_converted = {
@@ -46,8 +51,9 @@ def _process_data_into_dataframe(
                 ),
             }
             converted_data_scores.append(model_data_converted)
+            model_names.append(model_name)
 
-    return pd.DataFrame(converted_data_scores, index=selected_models)
+    return pd.DataFrame(converted_data_scores, index=model_names)
 
 
 def _molecule_stats_df(results: ConformerSelectionResult) -> pd.DataFrame:
@@ -76,7 +82,6 @@ def conformer_selection_page(
                    keys and the benchmark results objects as values.
     """
     st.markdown("# Conformer selection")
-    st.sidebar.markdown("# Conformer selection")
 
     st.markdown(
         "Organic molecules are flexible and able to adopt multiple conformations. "
@@ -125,17 +130,11 @@ def conformer_selection_page(
         st.markdown("**No results to display**.")
         return
 
-    model_names = list(data.keys())
-    model_select = st.sidebar.multiselect(
-        "Select model(s)", model_names, default=model_names
-    )
-    # with st.sidebar.container():
-    #     selected_energy_unit = st.selectbox(
-    #         "Select an energy unit:",
-    #         ["kcal/mol", "eV"],
-    #     )
+    selected_models = fetch_selected_models(available_models=list(data.keys()))
 
-    selected_models = model_select if model_select else model_names
+    if not selected_models:
+        st.markdown("**No results to display**.")
+        return
 
     df = _process_data_into_dataframe(data, selected_models)
     df_display = df.copy()

src/mlipaudit/ui/dihedral_scan.py

@@ -28,7 +28,11 @@
     DihedralScanResult,
 )
 from mlipaudit.ui.page_wrapper import UIPageWrapper
-from mlipaudit.ui.utils import create_st_image, display_model_scores
+from mlipaudit.ui.utils import (
+    create_st_image,
+    display_model_scores,
+    fetch_selected_models,
+)
 
 APP_DATA_DIR = Path(__file__).parent.parent / "app_data"
 DIHEDRAL_SCAN_DATA_DIR = APP_DATA_DIR / "dihedral_scan"
@@ -88,7 +92,6 @@ def dihedral_scan_page(
                    keys and the benchmark results objects as values.
     """
     st.markdown("# Dihedral scan")
-    st.sidebar.markdown("# Dihedral scan")
 
     with st.sidebar.container():
         selected_energy_unit = st.selectbox(
@@ -134,11 +137,11 @@ def dihedral_scan_page(
         st.markdown("**No results to display**.")
         return
 
-    unique_model_names = list(set(data.keys()))
-    model_select = st.sidebar.multiselect(
-        "Select model(s)", unique_model_names, default=unique_model_names
-    )
-    selected_models = model_select if model_select else unique_model_names
+    selected_models = fetch_selected_models(available_models=list(data.keys()))
+
+    if not selected_models:
+        st.markdown("**No results to display**.")
+        return
 
     conversion_factor = (
         1.0 if selected_energy_unit == "kcal/mol" else (units.kcal / units.mol)
@@ -153,6 +156,7 @@ def dihedral_scan_page(
             "Pearson Correlation": result.avg_pearson_r,
         }
         for model_name, result in data.items()
+        if model_name in selected_models
     ]
 
     # Create summary dataframe

src/mlipaudit/ui/folding_stability.py

@@ -20,46 +20,51 @@
 
 from mlipaudit.benchmarks import FoldingStabilityBenchmark, FoldingStabilityResult
 from mlipaudit.ui.page_wrapper import UIPageWrapper
-from mlipaudit.ui.utils import display_model_scores
+from mlipaudit.ui.utils import display_model_scores, fetch_selected_models
 
 ModelName: TypeAlias = str
 BenchmarkResultForMultipleModels: TypeAlias = dict[ModelName, FoldingStabilityResult]
 
 
 def _data_to_dataframes(
     data: BenchmarkResultForMultipleModels,
+    selected_models: list[str],
 ) -> tuple[pd.DataFrame, pd.DataFrame]:
     plot_data = []
     agg_data = []
 
     for model_name, result in data.items():
-        for molecule_result in result.molecules:
-            if not molecule_result.failed:
-                for idx in range(len(molecule_result.rmsd_trajectory)):  # type: ignore
-                    plot_data.append({
-                        "Model": model_name,
-                        "Structure": molecule_result.structure_name,
-                        "Frame": idx,
-                        "RMSD": molecule_result.rmsd_trajectory[idx],  # type: ignore
-                        "TM score": molecule_result.tm_score_trajectory[idx],  # type: ignore
-                        "Rad of Gyr Dev": molecule_result.radius_of_gyration_deviation[  # type: ignore
-                            idx
-                        ],
-                        "DSSP match": molecule_result.match_secondary_structure[idx],  # type: ignore
-                    })
-                    # Next line is to stay within max. line length below
-                    max_dev_rad_of_gyr = (
-                        molecule_result.max_abs_deviation_radius_of_gyration
-                    )
-                    agg_data.append({
-                        "Model": model_name,
-                        "Score": result.score,
-                        "Structure": molecule_result.structure_name,
-                        "avg. RMSD": molecule_result.avg_rmsd,
-                        "avg. TM score": molecule_result.avg_tm_score,
-                        "avg. DSSP match": molecule_result.avg_match,
-                        "max. abs. dev. Rad. of Gyr.": max_dev_rad_of_gyr,
-                    })
+        if model_name in selected_models:
+            for molecule_result in result.molecules:
+                if not molecule_result.failed:
+                    for idx in range(len(molecule_result.rmsd_trajectory)):  # type: ignore
+                        rad_of_gyr_dev = (
+                            molecule_result.radius_of_gyration_deviation[idx]  # type: ignore
+                        )
+                        plot_data.append({
+                            "Model": model_name,
+                            "Structure": molecule_result.structure_name,
+                            "Frame": idx,
+                            "RMSD": molecule_result.rmsd_trajectory[idx],  # type: ignore
+                            "TM score": molecule_result.tm_score_trajectory[idx],  # type: ignore
+                            "Rad of Gyr Dev": rad_of_gyr_dev,
+                            "DSSP match": molecule_result.match_secondary_structure[
+                                idx
+                            ],  # type: ignore
+                        })
+                        # Next line is to stay within max. line length below
+                        max_dev_rad_of_gyr = (
+                            molecule_result.max_abs_deviation_radius_of_gyration
+                        )
+                        agg_data.append({
+                            "Model": model_name,
+                            "Score": result.score,
+                            "Structure": molecule_result.structure_name,
+                            "avg. RMSD": molecule_result.avg_rmsd,
+                            "avg. TM score": molecule_result.avg_tm_score,
+                            "avg. DSSP match": molecule_result.avg_match,
+                            "max. abs. dev. Rad. of Gyr.": max_dev_rad_of_gyr,
+                        })
 
     df = pd.DataFrame(plot_data)
     df_agg = pd.DataFrame(agg_data)
@@ -181,7 +186,6 @@ def folding_stability_page(
                    keys and the benchmark results objects as values.
     """
     st.markdown("# Folding stability of trajectories")
-    st.sidebar.markdown("# Folding stability")
 
     st.markdown(
         "This module examines the folding stability trajectories of proteins in MLIP "
@@ -209,21 +213,13 @@ def folding_stability_page(
         st.markdown("**No results to display**.")
         return
 
-    df, df_agg = _data_to_dataframes(data)
-
-    unique_model_ids = list(data.keys())
-
-    # Add "Select All" option
-    all_models_option = st.sidebar.checkbox("Select all models", value=False)
+    selected_models = fetch_selected_models(available_models=list(data.keys()))
 
-    if all_models_option:
-        model_select = unique_model_ids
-    else:
-        model_select = st.sidebar.multiselect(
-            "Select model(s)", unique_model_ids, default=unique_model_ids
-        )
+    if not selected_models:
+        st.markdown("**No results to display**.")
+        return
 
-    selected_models = model_select if model_select else unique_model_ids
+    df, df_agg = _data_to_dataframes(data, selected_models)
 
     unique_structures = list(set(df["Structure"].unique()))