Merge pull request #5 from AllenNeuralDynamics/detached_local_branch

Detached local branch got lost, adding it back in.

Author: rachelstephlee

src/rachel_analysis_utils/analysis_utils.py

@@ -23,9 +23,9 @@ def get_RPE_by_avg_signal_fit(data, avg_signal_col):
         slope = np.nan
     return (x_fit, y_fit, slope)
 
-output_col_name = lambda channel, data_column, alignment_event: f"avg_{data_column}_{channel[:3]}_{alignment_event.split("_in_")[0]}"
+output_col_name = lambda channel, data_column, alignment_event: f"avg_{data_column}_{channel.split('_dff')[0]}_{alignment_event.split('_in_')[0]}"
 
-def add_AUC_and_rpe_slope(nwbs_by_week, parameters, data_column = 'data_z_norm', 
+def add_AUC_and_rpe_slope(nwbs_by_week, all_channels, save_dfs, data_column = 'data_z_norm', 
                             alignment_event = 'choice_time_in_session',offsets = [0.33,1]):
     """
     Enrich NWB weeks with average signal windows and compute RPE slopes per session for each channel.
@@ -36,11 +36,7 @@ def add_AUC_and_rpe_slope(nwbs_by_week, parameters, data_column = 'data_z_norm',
     # Enrich each week with average signals for every channel
     for nwb_week in nwbs_by_week:
         nwb_week_enriched = copy.deepcopy(nwb_week)
-        for channel in list(parameters["channels"].keys()):
-            # build the channel name used for processing (append preprocessing suffix if present)
-            if parameters.get('preprocessing', 'raw') != 'raw':
-                channel = channel + '_' + parameters['preprocessing']
-
+        for channel in all_channels:
             avg_signal_col = output_col_name(channel, data_column, alignment_event)
 
             nwb_week_enriched = trial_metrics.get_average_signal_window_multi(
@@ -58,13 +54,11 @@ def add_AUC_and_rpe_slope(nwbs_by_week, parameters, data_column = 'data_z_norm',
     rpe_rows = []
     subject_id = str(nwbs_by_week_enriched[0][0]).split(' ')[1].split('_')[0]
 
-    for ch in list(parameters["channels"].keys()):
-        channel = ch
-        if parameters.get('preprocessing', 'raw') != 'raw':
-            channel = channel + '_' + parameters['preprocessing']
+    for channel in all_channels:
 
         avg_signal_col = output_col_name(channel, data_column, alignment_event)
-
+        if avg_signal_col not in df_trials_all.columns:
+            continue
 
         for ses_idx in sorted(df_trials_all['ses_idx'].unique()):
             data = df_trials_all[df_trials_all['ses_idx'] == ses_idx]
@@ -86,7 +80,7 @@ def add_AUC_and_rpe_slope(nwbs_by_week, parameters, data_column = 'data_z_norm',
 
 
 
-    if parameters.get("save_dfs", False) == True:
+    if save_dfs == True:
         combined_rpe_slope.to_csv(f"/results/data/{subject_id}/rpe_slope.csv")
 
     return nwbs_by_week_enriched, combined_rpe_slope
@@ -236,7 +230,9 @@ def add_sliding_window_corr(
     else:
         raise ValueError("Input 'nwb' must have a 'df_fip' attribute or be a DataFrame with fip data.")
 
-
+    if signal1name not in df_fip.event.unique() or signal2name not in df_fip.event.unique():
+        print("One or both signals not found in df_fip events. Skipping correlation.")
+        return nwb
    # Select and pivot the two signals directly (align by timestamp)
     df_sel = (
         df_fip.loc[df_fip['event'].isin([signal1name, signal2name]),
@@ -295,7 +291,7 @@ def add_sliding_window_corr(
     r = r_full.to_numpy(dtype=float)[centers]
 
     df_corr = pd.DataFrame({'data_z':r, 'timestamps':t_centers})
-    df_corr['event'] = f'{signal1name[:3]}:{signal2name[:3]}_pearsonR'
+    df_corr['event'] = f'{signal1name.split('_dff')[0]}:{signal2name.split('_dff')[0]}_pearsonR'
 
     # If the entire correlation vector is NaN, nothing to merge — return unchanged.
     if np.all(np.isnan(r)):

src/rachel_analysis_utils/nwb_utils.py

@@ -482,7 +482,7 @@ def load(cls, session_folder, load_fip = False):
 
         return obj
 
-def save_nwb_list(nwb_list, plot_loc, df_sess=None):
+def save_nwb_list(flat_dummy_nwbs, plot_loc, df_sess=None):
     """
     Save a list or list-of-lists of dummy_nwb objects.
 
@@ -494,12 +494,6 @@ def save_nwb_list(nwb_list, plot_loc, df_sess=None):
                     <attr>.parquet
     """
 
-    # flatten list or list-of-lists
-    flat_dummy_nwbs = [
-        nwb
-        for item in nwb_list
-        for nwb in (item if isinstance(item, list) else [item])
-    ]
 
     subject_ids = set()
 
@@ -633,6 +627,22 @@ def get_date_and_week_interval(df, start_date):
     week_interval_series = ((date_series - start_date).dt.days // 7) + 1
     return week_interval_series
 
+def split_nwbs_by_week(nwbs_all):
+    nwbs_by_week = []
+    nwb_week_i = []
+    curr_week = 1
+    for nwb in nwbs_all:
+        week_interval = nwb.df_trials.week_interval.unique()[0]
+        if week_interval == curr_week:
+            nwb_week_i.append(nwb)
+        else:
+            nwbs_by_week.append(nwb_week_i)
+            nwb_week_i = [nwb]
+            curr_week = week_interval
+    nwbs_by_week.append(nwb_week_i)
+
+    return nwbs_by_week
+
 def get_dummy_nwbs_by_week(df_sess,df_trials, df_events, df_fip):
     start_date = pd.to_datetime(df_sess['session_date'].min())