Skip to content

Fixes to unified coefficient analysis pipeline #6

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
MartynaPlomecka wants to merge 1 commit into
base: main
Choose a base branch
from
Open

Fixes to unified coefficient analysis pipeline #6

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
178 changes: 135 additions & 43 deletions weinhardt2025/analysis/analysis_coefficients.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -8,20 +8,23 @@

Usage examples:

# Discrete (diagnosis-based, odds ratios):
# Discrete nominal (diagnosis-based, odds ratios):
python analysis_coefficients.py \
--model weinhardt2025/params/dezfouli2019/spice_dezfouli2019_a0_05.pkl \
--data weinhardt2025/data/dezfouli2019/dezfouli2019.csv \
--analysis discrete \
--analysis disc \
--criterion-type nominal \
--criterion diag \
--reference Healthy
--reference Healthy \
--participant-col session

# Continuous (age-based, effect sizes):
python analysis_coefficients.py \
--model weinhardt2025/params/eckstein2022/spice_eckstein2022.pkl \
--data weinhardt2025/data/eckstein2022/eckstein2022.csv \
--analysis continuous \
--criterion age
--analysis cont \
--criterion-type continuous \
--criterion age_years
"""

import argparse
Expand All @@ -41,7 +44,27 @@
from scipy.stats import spearmanr, kruskal, norm, chi2, f_oneway
from sklearn.linear_model import LogisticRegression
from sklearn.preprocessing import StandardScaler
from statsmodels.stats.multitest import multipletests
try:
from statsmodels.stats.multitest import multipletests
except ModuleNotFoundError:
def multipletests(pvals, method="fdr_bh"):
if method != "fdr_bh":
raise ValueError("Fallback multipletests only supports method='fdr_bh'.")
pvals = np.asarray(pvals, dtype=float)
n = len(pvals)
if n == 0:
return np.array([], dtype=bool), np.array([]), np.array([]), np.array([])

order = np.argsort(pvals)
ranked = pvals[order]
adjusted_ranked = ranked * n / np.arange(1, n + 1)
adjusted_ranked = np.minimum.accumulate(adjusted_ranked[::-1])[::-1]
adjusted_ranked = np.clip(adjusted_ranked, 0, 1)

adjusted = np.empty_like(adjusted_ranked)
adjusted[order] = adjusted_ranked
rejected = adjusted <= 0.05
return rejected, adjusted, np.full(n, np.nan), np.full(n, np.nan)

warnings.filterwarnings("ignore")

Expand Down Expand Up @@ -78,7 +101,16 @@ def clean_name(col):
# 1. Preparation – extract coefficients
# ---------------------------------------------------------------------------

def prepare(model_path: str, data_path: str, model_module: str, criterion_col, dataset_kwargs: dict = {}):
def prepare(
model_path: str,
data_path: str,
model_module: str,
criterion_col,
participant_col: str = "participant",
block_col: str = "block",
experiment_col: str = "experiment",
dataset_kwargs: dict = {},
):
"""Load a trained SPICE model, extract ensemble-averaged SINDy coefficients
per participant and merge with the data file.

Expand All @@ -91,11 +123,17 @@ def prepare(model_path: str, data_path: str, model_module: str, criterion_col, d
Names of the SINDy coefficient columns.
"""
# --- load data to infer dimensions ---
dataset = csv_to_dataset(file=data_path, **dataset_kwargs)
dataset = csv_to_dataset(
file=data_path,
df_participant_id=participant_col,
df_block=block_col,
df_experiment_id=experiment_col,
**dataset_kwargs,
)
raw_df = pd.read_csv(data_path)
n_actions = dataset.ys.shape[-1]
n_participants = len(dataset.xs[..., -1].unique())
unique_sessions = raw_df["participant"].unique().tolist()
unique_sessions = raw_df[participant_col].unique().tolist()

# --- load SPICE model via precoded module ---
mod = importlib.import_module(model_module)
Expand Down Expand Up @@ -153,8 +191,8 @@ def prepare(model_path: str, data_path: str, model_module: str, criterion_col, d
print(f"Extracted {len(sindy_cols)} SINDy coefficient columns for {len(sindy_df)} participants.")

# --- build criterion column per participant from raw data ---
crit_df = raw_df.groupby("participant").first().reset_index()
crit_df = crit_df.rename(columns={"participant": "participant_id"})
crit_df = raw_df.groupby(participant_col).first().reset_index()
crit_df = crit_df.rename(columns={participant_col: "participant_id"})
crit_df = crit_df[["participant_id", criterion_col]]

# Merge
Expand Down Expand Up @@ -383,6 +421,8 @@ def run_continuous(df, sindy_cols, criterion_col, output_dir):

scaler = StandardScaler()
crit_std = scaler.fit_transform(df_clean[[criterion_col]]).flatten()
crit_mean = float(scaler.mean_[0])
crit_scale = float(scaler.scale_[0]) if float(scaler.scale_[0]) != 0 else 1.0

results = []
skipped = []
Expand Down Expand Up @@ -483,7 +523,15 @@ def run_continuous(df, sindy_cols, criterion_col, output_dir):
# ---- plots ----
if not reg_df.empty:
_plot_beta_bars(reg_df, criterion_col, output_dir)
_plot_logistic_curves(reg_df, criterion_col, crit_min, crit_max, output_dir)
_plot_logistic_curves(
reg_df,
criterion_col,
crit_min,
crit_max,
crit_mean,
crit_scale,
output_dir,
)
return res_df


Expand All @@ -506,13 +554,13 @@ def _plot_beta_bars(df, criterion_col, output_dir):
print(f" β bar-plot saved.")


def _plot_logistic_curves(df, criterion_col, crit_min, crit_max, output_dir):
def _plot_logistic_curves(df, criterion_col, crit_min, crit_max, crit_mean, crit_scale, output_dir):
n = min(len(df), 12)
top = df.head(n)
ncols = min(4, n)
nrows = math.ceil(n / ncols)
xs = np.linspace(crit_min, crit_max, 200)
xs_std = (xs - xs.mean()) / xs.std()
xs_std = (xs - crit_mean) / crit_scale

fig, axes = plt.subplots(nrows, ncols, figsize=(4 * ncols, 3 * nrows))
axes_flat = np.array(axes).flatten() if n > 1 else [axes]
Expand Down Expand Up @@ -555,7 +603,7 @@ def jonckheere_terpstra(groups):
return z, 2 * (1 - norm.cdf(abs(z)))


def run_magnitude_analysis(df, sindy_cols, criterion_col, analysis_type,
def run_magnitude_analysis(df, sindy_cols, criterion_col, analysis_type, criterion_type,
output_dir, group_labels=None):
"""Non-parametric magnitude analysis (Spearman, Kruskal-Wallis,
Jonckheere-Terpstra) of SINDy coefficients vs the criterion.
Expand All @@ -566,25 +614,35 @@ def run_magnitude_analysis(df, sindy_cols, criterion_col, analysis_type,
out = Path(output_dir)
out.mkdir(parents=True, exist_ok=True)

if analysis_type == "disc":
if criterion_type == "nominal":
unique_vals = sorted(df[criterion_col].unique())
group_col = "__group__"
df[group_col] = df[criterion_col].map({v: i for i, v in enumerate(unique_vals)})
if group_labels is None:
group_labels = [str(v) for v in unique_vals]
n_groups = len(unique_vals)
else:
elif criterion_type in {"ordinal", "continuous"}:
if analysis_type == "disc":
unique_vals = sorted(df[criterion_col].unique())
group_col = "__group__"
df[group_col] = df[criterion_col].map({v: i for i, v in enumerate(unique_vals)})
if group_labels is None:
group_labels = [str(v) for v in unique_vals]
n_groups = len(unique_vals)
else:
# Bin continuous criterion into quantile-based groups
n_groups = min(6, len(df) // 10)
if n_groups < 3:
print("Not enough data for magnitude analysis grouping. Skipping.")
return None
group_col = "__group__"
df[group_col] = pd.qcut(df[criterion_col], q=n_groups, labels=False, duplicates="drop")
n_groups = df[group_col].nunique()
if group_labels is None:
bounds = pd.qcut(df[criterion_col], q=n_groups, duplicates="drop").cat.categories
group_labels = [f"{iv.left:.0f}-{iv.right:.0f}" for iv in bounds]
n_groups = min(6, len(df) // 10)
if n_groups < 3:
print("Not enough data for magnitude analysis grouping. Skipping.")
return None
group_col = "__group__"
df[group_col] = pd.qcut(df[criterion_col], q=n_groups, labels=False, duplicates="drop")
n_groups = df[group_col].nunique()
if group_labels is None:
bounds = pd.qcut(df[criterion_col], q=n_groups, duplicates="drop").cat.categories
group_labels = [f"{iv.left:.0f}-{iv.right:.0f}" for iv in bounds]
else:
raise ValueError(f"Unsupported criterion_type '{criterion_type}'.")

keep = [c for c in sindy_cols
if (nz := df.loc[df[c] != 0, c]).size > 10 and nz.std() > 1e-10]
Expand All @@ -601,12 +659,16 @@ def run_magnitude_analysis(df, sindy_cols, criterion_col, analysis_type,
groups_vals = df.loc[nz_mask, group_col]
raw = [vals[groups_vals == g].values for g in range(n_groups)]

rho, p_s = spearmanr(groups_vals, vals) if vals.size > 10 else (np.nan, np.nan)
valid = [x for x in raw if x.size and x.std() > 1e-10]
kw, p_kw = kruskal(*valid) if len(valid) >= 3 else (np.nan, np.nan)
z, p_jt = jonckheere_terpstra(raw)

trend = ("Increasing" if rho > 0 else "Decreasing") if not np.isnan(rho) else "Undetermined"
if criterion_type == "nominal":
rho, p_s = np.nan, np.nan
z, p_jt = np.nan, np.nan
trend = "Not applicable"
else:
rho, p_s = spearmanr(groups_vals, vals) if vals.size > 10 else (np.nan, np.nan)
z, p_jt = jonckheere_terpstra(raw)
trend = ("Increasing" if rho > 0 else "Decreasing") if not np.isnan(rho) else "Undetermined"

grp_stats = []
for g_idx in range(n_groups):
Expand All @@ -632,10 +694,12 @@ def run_magnitude_analysis(df, sindy_cols, criterion_col, analysis_type,
"group_stats": grp_stats,
})

res = pd.DataFrame(results).sort_values("jt_p")
sort_col = "kruskal_p" if criterion_type == "nominal" else "jt_p"
res = pd.DataFrame(results).sort_values(sort_col)

# FDR correction
for col in ["spearman_p", "kruskal_p", "jt_p"]:
pvalue_cols = ["kruskal_p"] if criterion_type == "nominal" else ["spearman_p", "kruskal_p", "jt_p"]
for col in pvalue_cols:
m = res[col].notna()
if m.sum() > 0:
res.loc[m, f"{col}_fdr"] = multipletests(res.loc[m, col], method="fdr_bh")[1]
Expand All @@ -662,18 +726,23 @@ def run_magnitude_analysis(df, sindy_cols, criterion_col, analysis_type,
plt.figure(figsize=(12, max(8, len(top15) * 0.5)))
sns.heatmap(heat, annot=True, cmap="RdBu_r", center=0, fmt=".3f",
cbar_kws={"label": "Mean Coefficient"})
plt.title(f"SINDy coefficient means by {criterion_col} group (top-{len(top15)} JT)")
title_suffix = "KW" if criterion_type == "nominal" else "JT"
plt.title(f"SINDy coefficient means by {criterion_col} group (top-{len(top15)} {title_suffix})")
plt.tight_layout()
plt.savefig(out / "magnitude_heatmap.png", dpi=300, bbox_inches="tight")
plt.close()

# Summary
print(f"Significant Spearman (p<0.05): {(res.spearman_p < 0.05).sum()}")
print(f"Significant Kruskal-Wallis (p<0.05): {(res.kruskal_p < 0.05).sum()}")
print(f"Significant Jonckheere-Terpstra: {(res.jt_p < 0.05).sum()}")
print("\nTop 10 by JT p-value:")
print(res[["coefficient", "spearman_rho", "spearman_p", "jt_p", "trend"]]
.head(10).to_string(index=False))
if criterion_type == "nominal":
print("\nTop 10 by Kruskal-Wallis p-value:")
print(res[["coefficient", "kruskal_p"]].head(10).to_string(index=False))
else:
print(f"Significant Spearman (p<0.05): {(res.spearman_p < 0.05).sum()}")
print(f"Significant Jonckheere-Terpstra: {(res.jt_p < 0.05).sum()}")
print("\nTop 10 by JT p-value:")
print(res[["coefficient", "spearman_rho", "spearman_p", "jt_p", "trend"]]
.head(10).to_string(index=False))

# Clean up temp column
if group_col in df.columns:
Expand All @@ -700,8 +769,11 @@ def parse_args():
help="Column name for the regression criterion "
"(e.g. 'Diagnosis', 'Age', 'diag')")
p.add_argument("--analysis", required=True, choices=["disc", "cont"],
help="'discrete' for odds-ratio analysis, "
"'continuous' for effect-size analysis")
help="'disc' for odds-ratio analysis, "
"'cont' for effect-size analysis")
p.add_argument("--criterion-type", required=True,
choices=["nominal", "ordinal", "continuous"],
help="Statistical type of the criterion")
p.add_argument("--model-name", default="spice.precoded.workingmemory_rewardbinary",
help="Name of the SPICE model module "
"(default: spice.precoded.workingmemory_rewardbinary)")
Expand All @@ -710,14 +782,24 @@ def parse_args():
"(e.g. 'Healthy'). Required for --analysis discrete.")
p.add_argument("--output", default=None,
help="Output directory (default: auto-generated next to data)")
p.add_argument("--participant-col", default="participant",
help="Participant/session ID column in the input CSV")
p.add_argument("--block-col", default="block",
help="Block column in the input CSV")
p.add_argument("--experiment-col", default="experiment",
help="Experiment column in the input CSV")
return p.parse_args()


def main():
args = parse_args()

if args.analysis == "disc" and args.reference is None:
raise ValueError("--reference-group is required for discrete analysis.")
raise ValueError("--reference is required for discrete analysis.")
if args.analysis == "disc" and args.criterion_type == "continuous":
raise ValueError("--criterion-type continuous is incompatible with --analysis disc.")
if args.analysis == "cont" and args.criterion_type != "continuous":
raise ValueError("--analysis cont requires --criterion-type continuous.")

output_dir = args.output or os.path.join(
os.path.dirname(args.data),
Expand All @@ -733,6 +815,9 @@ def main():
data_path=args.data,
model_module=args.model_name,
criterion_col=args.criterion,
participant_col=args.participant_col,
block_col=args.block_col,
experiment_col=args.experiment_col,
)

# 2. Regression analysis
Expand All @@ -749,7 +834,14 @@ def main():
print("\n" + "=" * 70)
print("STEP 3: Magnitude analysis")
print("=" * 70)
run_magnitude_analysis(df, sindy_cols, args.criterion, args.analysis, output_dir)
run_magnitude_analysis(
df,
sindy_cols,
args.criterion,
args.analysis,
args.criterion_type,
output_dir,
)

print(f"\nAll results saved to: {output_dir}")

Expand Down