refactor: generate curriculum plot from df_session

Author: hanhou

code/Home.py

@@ -348,37 +348,12 @@ def init(if_load_bpod_data_override=None, if_load_docDB_override=None, if_load_s
         st.session_state[f'df_selected_from_{source}'] = pd.DataFrame(columns=['subject_id', 'session'])
 
     # Load autotrain
-    auto_train_manager, curriculum_manager = load_auto_train()
-    st.session_state.auto_train_manager = auto_train_manager
+    _, curriculum_manager = load_auto_train()
     st.session_state.curriculum_manager = curriculum_manager
 
     # Some ad-hoc modifications on df_sessions
     _df = st.session_state.df['sessions_main'].copy()
 
-    # -- overwrite the `if_stage_overriden_by_trainer`
-    # Previously it was set to True if the trainer changes stage during a session.
-    # But it is more informative to define it as whether the trainer has overridden the curriculum.
-    # In other words, it is set to True only when stage_suggested ~= stage_actual, as defined in the autotrain curriculum.
-    _df.drop(columns=['if_overriden_by_trainer'], inplace=True)
-    tmp_auto_train = (
-        auto_train_manager.df_manager.query("if_closed_loop == True")[
-            [
-                "subject_id",
-                "session_date",
-                "current_stage_suggested",
-                "if_stage_overriden_by_trainer",
-            ]
-        ]
-        .copy()
-        .drop_duplicates(subset=["subject_id", "session_date"], keep="first")
-    )
-    tmp_auto_train["session_date"] = pd.to_datetime(tmp_auto_train["session_date"])
-    _df = _df.merge(
-        tmp_auto_train,
-        on=["subject_id", "session_date"],
-        how='left',
-    )
-
     # --- Load data from docDB ---
     if_load_docDb = if_load_docDB_override if if_load_docDB_override is not None else (
         st.query_params['if_load_docDB'].lower() == 'true'

code/util/plot_autotrain_manager.py

@@ -32,42 +32,17 @@ def plot_manager_all_progress_bokeh_source(
     filtered_session_ids=None,
 ):
     # --- Prepare data ---
-    manager = st.session_state.auto_train_manager
-    df_manager = manager.df_manager.sort_values(
+    # Now we already merged full curriculum info in the master df_session by the result access api
+    # manager = st.session_state.auto_train_manager
+
+    df_to_draw = st.session_state.df['sessions_main'].sort_values(
         by=["subject_id", "session"],
         ascending=[sort_order == "ascending", False]
-    )
+    ).copy()
 
-    if not len(df_manager):
+    if not len(df_to_draw):
         return None
 
-    # Metadata merge from df_master
-    df_tmp_rig_trainer = st.session_state.df["sessions_main"][
-        ["subject_id", "session_date", "session", "rig", "trainer", "PI", "nwb_suffix", 
-         "foraging_eff_random_seed", "finished_trials", "finished_rate", 
-         "task", "curriculum_name", "curriculum_version", "current_stage_actual"]
-    ]
-    df_tmp_rig_trainer["session_date"] = df_tmp_rig_trainer["session_date"].astype(str)
-
-    df_to_draw = (
-        df_manager.drop_duplicates(
-            subset=["subject_id", "session_date"], keep="last"
-        )  # Duplicte sessions in the autotrain due to pipeline issues
-        .drop(
-            columns=[
-                "session",
-                "task",
-                "foraging_efficiency", 
-                "finished_trials", 
-            ]
-        )  # df_master has higher priority in session numbers
-        .merge(
-            df_tmp_rig_trainer.query(f"current_stage_actual != 'None'"),
-            on=["subject_id", "session_date"],
-            how="right",
-        )
-    )
-
     # If use_filtered_data, filter the data
     if if_use_filtered_data:
         df_to_draw = df_to_draw.merge(
@@ -76,11 +51,6 @@ def plot_manager_all_progress_bokeh_source(
             how="inner",
         )
 
-    # Correct df_manager missing sessions (df_manager has higher priority in curriculum-related fields)
-    df_to_draw["curriculum_name"] = df_to_draw["curriculum_name_x"].fillna(df_to_draw["curriculum_name_y"])
-    df_to_draw["curriculum_version"] = df_to_draw["curriculum_version_x"].fillna(df_to_draw["curriculum_version_y"])
-    df_to_draw["current_stage_actual"] = df_to_draw["current_stage_actual_x"].fillna(df_to_draw["current_stage_actual_y"])
-
     df_to_draw["color"] = df_to_draw["current_stage_actual"].map(stage_color_mapper)
     df_to_draw["edge_color"] = (  # Use grey edge to indicate stage without suggestion 
         df_to_draw["current_stage_suggested"].map(stage_color_mapper).fillna("#d3d3d3")

code/util/streamlit.py

@@ -774,8 +774,6 @@ def _add_download_filtered_session():
 
 def add_auto_train_manager():
 
-    df_training_manager = st.session_state.auto_train_manager.df_manager
-
     # -- Show plotly chart --
     cols = st.columns([1, 1, 1, 0.7, 0.7, 1.5, 1.5, 2])
     options = ["date", "session", "relative_date"]
@@ -856,7 +854,7 @@ def add_auto_train_manager():
             else None
         ),
     )
-    
+
     if fig_auto_train is None:
         st.markdown("### In the filtered sessions, no AutoTrain history to show!")
         return
@@ -878,20 +876,29 @@ def add_auto_train_manager():
             st.write(data_df.iloc[indices])
 
     # -- Show dataframe --
+    df_training_manager = st.session_state.df_session_filtered[
+        [
+            "subject_id",
+            "session_date",
+            "session",
+            "curriculum_name",
+            "curriculum_version",
+            "curriculum_schema_version",
+            "current_stage_suggested",
+            "current_stage_actual",
+            "session_at_current_stage",
+            "if_overriden_by_trainer",
+            "foraging_efficiency",
+            "finished_trials",
+            "decision",
+            "next_stage_suggested",
+        ]
+    ]
+
     # only show filtered subject
     df_training_manager = df_training_manager[df_training_manager['subject_id'].isin(
         st.session_state.df_session_filtered['subject_id'].unique().astype(str))]
 
-    # reorder columns
-    df_training_manager = df_training_manager[['subject_id', 'session_date', 'session', 
-                                                'curriculum_name', 'curriculum_version', 'curriculum_schema_version',
-                                                'current_stage_suggested', 'current_stage_actual',
-                                                'session_at_current_stage',
-                                                'if_closed_loop', 'if_stage_overriden_by_trainer',
-                                                'foraging_efficiency', 'finished_trials', 
-                                                'decision', 'next_stage_suggested'
-                                                ]]
-
     with st.expander('Automatic training manager', expanded=False):
         st.dataframe(df_training_manager, height=3000)
 
@@ -1195,4 +1202,4 @@ def add_download_plotly_as_svg(fig, file_name="plot.svg"):
         data=svg_file,
         file_name=file_name.replace(".svg", "") + ".svg",
         mime="image/svg+xml"
-    )
+    )