From 95d280cf3050a6aab1fa5e43607ab6118a104f53 Mon Sep 17 00:00:00 2001 From: SergeantChris Date: Sat, 18 Apr 2026 17:06:47 +0200 Subject: [PATCH] Add wrapper class CT-RSA for RSA over multiple neural and model timepoints + max + baseline-correction --- net2brain/evaluations/ct_rsa.py | 156 ++++++++++++++++++++++++++++++++ 1 file changed, 156 insertions(+) create mode 100644 net2brain/evaluations/ct_rsa.py diff --git a/net2brain/evaluations/ct_rsa.py b/net2brain/evaluations/ct_rsa.py new file mode 100644 index 0000000..09de238 --- /dev/null +++ b/net2brain/evaluations/ct_rsa.py @@ -0,0 +1,156 @@ +import numpy as np +import pandas as pd +from .rsa import RSA + + +class CTRSA(): + def __init__(self, model_rdms_path, brain_rdms_path, model_name, layer_skips=(), squared=False): + """ + Cross-Temporal RSA (CT-RSA) (Sartzetaki, Zonneveld et al., 2026). + Evaluation comparing model RDMs to brain RDMs across all model timepoints and layers, + and identifying the maximum RSA score for each brain timepoint across all model timepoints and layers. + Brain RDMs should be an array of shape (n_subjects, n_timepoints, n_stimuli, n_stimuli), + and model RDMs should be computed with the option multi_timepoint_rdms set in RDMCreator. + + Sartzetaki, C., Zonneveld, A.W., Oyarzo, P., Gifford, A.T., Cichy, R.M., Mettes, P. and Groen, I.I.A., 2026. + The Human Brain as a Dynamic Mixture of Expert Models in Video Understanding. In The Fourteenth + International Conference on Learning Representations. https://openreview.net/forum?id=bSsNSfyj8m + + Args: + model_rdms_path (str): Path to the folder containing the model RDMs. + brain_rdms_path (str): Path to the folder containing the brain RDMs. + model_name (str): Name of the model. + layer_skips (tuple, optional): Names of the model layers to skip. Use '_' instead of '.' in the names. + squared (bool): Whether to square the correlation values. + """ + + self.rsa = RSA( + model_rdms_path=model_rdms_path, + brain_rdms_path=brain_rdms_path, + model_name=model_name, + layer_skips=layer_skips, + squared=squared, + timepoint_agg=False, + model_timepoints=True, + ) + + def evaluate(self, layer_order=None, baseline_correction=True, sfreq=50, before_onset=0.2): + """ + Args: + layer_order: Provide the layer names in the correct depth order if this isn't the case. + baseline_correction: Whether to do baseline correction on the max scores by subtracting the mean of the baseline window from all scores. + sfreq: Sampling frequency of the sequential brain data, used to convert timepoint indices to seconds. + before_onset: Time in seconds before stimulus onset, used to convert timepoint indices to seconds and to define the baseline window for baseline correction. + + Returns: + final_df: A DataFrame with one row per brain timepoint, containing the maximum RSA score across all model timepoints and layers for that brain timepoint, + as well as the corresponding layer and model timepoint information. The DataFrame is baseline-corrected if baseline_correction is True. + """ + + df = self.rsa.evaluate() + df.drop(columns=["Significance"], inplace=True) + + df["Layer"] = ( + df["Layer"] + .apply(lambda lay: lay.split("RDM_")[1] if "RDM_" in lay else lay) + .apply(lambda lay: lay.split(".npz")[0] if ".npz" in lay else lay) + ) + df["Layer_FullNames"] = df["Layer"] + if not layer_order: + layer_order = df["Layer"].unique().tolist() + if self.rsa.layer_skips: + layer_order = [lay for lay in layer_order if lay not in self.rsa.layer_skips] + df["Layer"] = df["Layer"].map({lay: i for i, lay in enumerate(layer_order)}) + df.sort_values(by=["Layer"], inplace=True) + + df["R"] = df["R"].apply(np.array) + max_len = df["R"].apply(len).max() + + def repeat_elements_to_max_length(x): + repeat_count = max_len // len(x) + return [item for item in x for _ in range(repeat_count)] + + # Get original n_timepoints for each layer + all_og_n_time_points = [] + for layer in df["Layer"].unique(): + df_layer = df[df["Layer"] == layer] + og_n_time_points = len(df_layer['R'].values[0]) + all_og_n_time_points.append(og_n_time_points) + + # make all layers have max timepoints by repeating each timepoint max_timepoints // layer_timepoints times + df["R"] = df["R"].apply(repeat_elements_to_max_length) + + # Step 1: max layer for all model timepoints and brain timepoints + max_R_all_model_timepoints = [] + max_layers_all_model_timepoints = [] + max_layers_all_model_timepoints_idx = [] + max_layer_all_model_timepoints_names = [] + for model_timepoint in range(max_len): + R_matrix = np.stack(df["R"].apply(lambda x: x[model_timepoint]).values) + max_R = np.max(R_matrix, axis=0) + max_indices = np.argmax(R_matrix, axis=0) + max_layer_names = [df.iloc[idx]["Layer_FullNames"] for idx in max_indices] + layer_values = df["Layer"].values / df["Layer"].max() + max_layers = [layer_values[idx] for idx in max_indices] + max_R_all_model_timepoints.append(max_R) + max_layers_all_model_timepoints.append(max_layers) + max_layers_all_model_timepoints_idx.append(max_indices) + max_layer_all_model_timepoints_names.append(max_layer_names) + row = { + "Time": np.arange(len(max_R)) / sfreq - before_onset, + "R": max_R_all_model_timepoints, + "Layer_at_max_R": max_layers_all_model_timepoints, + "Layer_at_max_R_idx": max_layers_all_model_timepoints_idx, + "Layer_at_max_R_names": max_layer_all_model_timepoints_names, + } + metadata = df.drop(columns=["Layer", "Layer_FullNames", "R", "%R", "R_array"]).iloc[0].to_dict() + + # Step 2: flatten list-of-lists into long format + rows = [] + for rep_idx, (R_values_rep, Layer_values_rep) in enumerate(zip(row["R"], row["Layer_at_max_R"])): + # inner loop over timepoints + t_idx = 0 + for t, (R_val, Layer_val) in enumerate(zip(R_values_rep, Layer_values_rep)): + new_row = { + "rep_idx_abs": rep_idx, + "Time": row["Time"][t], # row["Time"] is the same for all reps + "R": R_val, + "Layer_at_max_R": Layer_val, + "Layer_at_max_R_idx": row["Layer_at_max_R_idx"][rep_idx][t_idx], + "Layer_at_max_R_names": row["Layer_at_max_R_names"][rep_idx][t_idx] + } + for k, v in metadata.items(): + new_row[k] = v + rows.append(new_row) + t_idx = t_idx + 1 + long_df = pd.DataFrame(rows) + + # Step 3: max model timepoint for each brain timepoint + max_dfs = [] + for t in long_df["Time"].unique(): + time_df = long_df[long_df["Time"] == t] + R_matrix = np.stack(time_df["R"].values) + max_index = np.argmax(R_matrix, axis=0) + max_layer = time_df.iloc[max_index]["Layer_at_max_R_idx"] + max_index_rescaled = int(max_index // (R_matrix.shape[0] / all_og_n_time_points[max_layer])) + max_slice = time_df.iloc[max_index].copy() + max_slice['rep_idx_abs'] = max_index_rescaled + try: # ratio of best timepoint to total number of timepoints in the original layer (before repeating to max_len) + max_slice['rep_idx_ratio'] = max_index_rescaled / (all_og_n_time_points[max_layer] - 1) + except ZeroDivisionError: + max_slice['rep_idx_ratio'] = 0 # more correct 0.5 to represent the averaging? + max_dfs.append(max_slice) + final_df = pd.concat(max_dfs, axis=1).T.reset_index(drop=True) + + # Step 4: do baseline correction on max scores + if baseline_correction: + r_values = final_df["R"].values + times = final_df["Time"].values + baseline_mask = (times >= -before_onset) & (times <= 0) + if not np.any(baseline_mask): + raise ValueError("No timepoints found in the specified baseline window.") + baseline_mean = np.mean(r_values[baseline_mask]) + corrected_r_values = np.array([r - baseline_mean for r in r_values]) + final_df["R"] = corrected_r_values + + return final_df