Merge pull request #117 from AllenNeuralDynamics/han_optimize_autotrain

hanhou · web-flow · commit 0631d520ba37 · 2025-06-10T19:34:31.000-07:00
Multiple Improvements on autotrain
diff --git a/code/Home.py b/code/Home.py
@@ -302,7 +302,7 @@ def show_curriculums():
 
 
 # ------- Layout starts here -------- #
-def init(if_load_bpod_data_override=None, if_load_docDB_override=None):
+def init(if_load_bpod_data_override=None, if_load_docDB_override=None, if_load_sessions_older_than_6_month_override=None):
 
     # Clear specific session state and all filters
     for key in st.session_state:
@@ -319,9 +319,25 @@ def init(if_load_bpod_data_override=None, if_load_docDB_override=None):
         if 'if_load_bpod_sessions' in st.session_state
         else False)
     st.session_state.bpod_loaded = _if_load_bpod
-    
+
+    _if_load_sessions_older_than_6_month = (
+        if_load_sessions_older_than_6_month_override
+        if if_load_sessions_older_than_6_month_override is not None
+        else (
+            st.query_params["if_load_sessions_older_than_6_month"].lower() == "true"
+            if "if_load_sessions_older_than_6_month" in st.query_params
+            else (
+                st.session_state.if_load_sessions_older_than_6_month
+                if "if_load_sessions_older_than_6_month" in st.session_state
+                else False
+            )
+        )
+    )
+
     # --- Load data using aind-analysis-arch-result-access ---
-    df_han = get_session_table(if_load_bpod=_if_load_bpod)
+    # Convert boolean to months: if True, load all sessions (None), if False, load only recent 6 months
+    only_recent_n_month = None if _if_load_sessions_older_than_6_month else 6
+    df_han = get_session_table(if_load_bpod=_if_load_bpod, only_recent_n_month=only_recent_n_month)
     df = {'sessions_main': df_han}  # put it in df['session_main'] for backward compatibility
 
     if not len(df):
@@ -332,38 +348,12 @@ def init(if_load_bpod_data_override=None, if_load_docDB_override=None):
         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'
@@ -439,9 +429,6 @@ def app():
         # -- 1. unit dataframe --
         
         cols = st.columns([4, 4, 4, 1])
-        cols[0].markdown(f'### Filter the sessions on the sidebar\n'
-                         f'#####  {len(st.session_state.df_session_filtered)} sessions, '
-                         f'{len(st.session_state.df_session_filtered.subject_id.unique())} mice filtered')
         
         with cols[0].expander(':bulb: Get the master session table by code', expanded=False):
             st.code(f'''
@@ -453,6 +440,12 @@ def app():
         
         with cols[1]:
             with st.form(key='load_settings', clear_on_submit=False):
+                if_load_sessions_older_than_6_month = checkbox_wrapper_for_url_query(
+                    st_prefix=st,
+                    label='Include sessions older than 6 months (reload after change)',
+                    key='if_load_sessions_older_than_6_month',
+                    default=False,
+                )
                 if_load_bpod_sessions = checkbox_wrapper_for_url_query(
                     st_prefix=st,
                     label='Include old Bpod sessions (reload after change)',
@@ -472,6 +465,12 @@ def app():
                     sync_session_state_to_URL()
                     init()
                     st.rerun()  # Reload the page to apply the changes
+                    
+        cols[0].markdown(f'### Filter the sessions on the sidebar\n' +
+                    f'#####  {len(st.session_state.df_session_filtered)} sessions, ' +
+                    f'{len(st.session_state.df_session_filtered.subject_id.unique())} mice filtered' +
+                    (f' (recent 6 months only)' if not st.session_state.if_load_sessions_older_than_6_month else '')
+                    )
               
         table_height = slider_wrapper_for_url_query(st_prefix=cols[-1],
                                                     label='Table height',
@@ -696,4 +695,4 @@ def add_main_tabs():
         st.markdown('####  1. Reload the page')
         st.markdown('####  2. Click this original URL https://foraging-behavior-browser.allenneuraldynamics-test.org/')
         st.markdown('####  3. Report your bug here: https://github.com/AllenNeuralDynamics/foraging-behavior-browser/issues (paste your URL and screenshoots)')
-        raise e
+        raise e
diff --git a/code/pages/0_Data inventory.py b/code/pages/0_Data inventory.py
@@ -559,6 +559,6 @@ def add_venn_diagrms(df_merged):
     # Share the same master df as the Home page
     if "df" not in st.session_state or "sessions_main" not in st.session_state.df.keys() or not st.session_state.bpod_loaded:
         st.spinner("Loading data from Han temp pipeline...")
-        init(if_load_docDB_override=False, if_load_bpod_data_override=True)
+        init(if_load_docDB_override=False, if_load_bpod_data_override=True, if_load_sessions_older_than_6_month_override=True)
 
     app()
diff --git a/code/util/aws_s3.py b/code/util/aws_s3.py
@@ -62,7 +62,7 @@ def load_auto_train():
     curriculum_manager = CurriculumManager(
         saved_curriculums_on_s3=dict(
             bucket='aind-behavior-data',
-            root='foraging_auto_training/saved_curriculums/'
+            root='foraging_nwb_bonsai_processed/foraging_auto_training/saved_curriculums/'
         ),
         saved_curriculums_local=os.path.expanduser('~/curriculum_manager/'),
     )
@@ -72,7 +72,7 @@ def load_auto_train():
                                 root='foraging_nwb_bonsai_processed/',
                                 file_name='df_sessions.pkl'),
         df_manager_root_on_s3=dict(bucket='aind-behavior-data',
-                                root='foraging_auto_training/')
+                                root='foraging_nwb_bonsai_processed/foraging_auto_training/')
     )
 
     _df = auto_train_manager.df_manager.copy().rename(
diff --git a/code/util/plot_autotrain_manager.py b/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,19 +51,17 @@ 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")
     )
+    
+    # Convert session_date to string for URL generation
+    df_to_draw["session_date_str"] = df_to_draw["session_date"].astype(str)
     df_to_draw["imgs_1"] = df_to_draw.apply(
         lambda x: get_s3_public_url(
             subject_id=x["subject_id"],
-            session_date=x["session_date"],
+            session_date=x["session_date_str"],
             nwb_suffix=x["nwb_suffix"],
             figure_suffix="choice_history.png",
         ),
@@ -97,13 +70,17 @@ def plot_manager_all_progress_bokeh_source(
     df_to_draw["imgs_2"] = df_to_draw.apply(
         lambda x: get_s3_public_url(
             subject_id=x["subject_id"],
-            session_date=x["session_date"],
+            session_date=x["session_date_str"],
             nwb_suffix=x["nwb_suffix"],
             figure_suffix="logistic_regression_Su2022.png",
         ),
         axis=1,
     )
     df_to_draw.round(3)
+    
+    # --- Remove rows with NaN in color or edge_color ---
+    # to fix a bug where non-normalized stages appears in the autotrain table
+    df_to_draw = df_to_draw.dropna(subset=["color", "edge_color"])
 
     # --- Filter recent days ---
     df_to_draw['session_date'] = pd.to_datetime(df_to_draw['session_date'])
diff --git a/code/util/streamlit.py b/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"]
@@ -818,7 +816,11 @@ def add_auto_train_manager():
         ),
         default=True,
     )
-    only_filtered_effective = only_filtered and len(st.session_state.df_session_filtered) < len(st.session_state.df["sessions_main"])
+    only_filtered_effective = (
+        only_filtered 
+        and len(st.session_state.df_session_filtered) < len(st.session_state.df["sessions_main"])
+        and len(st.session_state.df_session_filtered) < 500  # Only if filtered down to less than 500 sessions
+    )
 
     recent_months = slider_wrapper_for_url_query(cols[6],
                                                 label="only recent months",
@@ -856,7 +858,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 +880,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_eff",
+            "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 +1206,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"
-    )
+    )
diff --git a/environment/Dockerfile b/environment/Dockerfile
@@ -1,22 +1,14 @@
 # hash:sha256:51bda5f40316acb89ad85a82e996448f5a31d6f40b5b443e817e9b346eee2f67
 ARG REGISTRY_HOST
-FROM $REGISTRY_HOST/codeocean/jupyterlab:3.6.1-miniconda4.12.0-python3.9-ubuntu20.04
+FROM $REGISTRY_HOST/codeocean/mambaforge3:23.1.0-4-python3.10.12-ubuntu22.04
 
 ARG DEBIAN_FRONTEND=noninteractive
 
 ARG GIT_ASKPASS
 ARG GIT_ACCESS_TOKEN
 COPY git-askpass /
 
-RUN pip install -r https://raw.githubusercontent.com/AllenNeuralDynamics/foraging-behavior-browser/main/requirements.txt
-
-ADD "https://github.com/coder/code-server/releases/download/v4.21.1/code-server-4.21.1-linux-amd64.tar.gz" /.code-server/code-server.tar.gz
-
-RUN cd /.code-server \
-    && tar -xvf code-server.tar.gz \
-    && rm code-server.tar.gz \
-    && ln -s /.code-server/code-server-4.21.1-linux-amd64/bin/code-server  /usr/bin/code-server
-
+RUN pip install -r https://raw.githubusercontent.com/AllenNeuralDynamics/foraging-behavior-browser/main/requirements.txt --no-cache-dir
 
 COPY postInstall /
 RUN /postInstall
diff --git a/environment/postInstall b/environment/postInstall
