Skip to content

PID analysis pipeline: Gaussian MMI-PID across CTRNN/Elman ensembles, GPU speedup, and cleanup#43

Merged
jgendra merged 8 commits into
mainfrom
jgendra
Jul 4, 2026
Merged

PID analysis pipeline: Gaussian MMI-PID across CTRNN/Elman ensembles, GPU speedup, and cleanup#43
jgendra merged 8 commits into
mainfrom
jgendra

Conversation

@jgendra

@jgendra jgendra commented Jul 4, 2026

Copy link
Copy Markdown
Owner

TL;DR: Adds a full, verified Partial Information Decomposition (PID) analysis pipeline (notebook 06) comparing perceptual vs. context-dependent decision-making RNNs, plus a hidden-size sweep study and repo cleanup (removed an unrelated lecture example, reorganized results/figures directories).

Description:
This branch builds out the PID analysis workflow end-to-end: computing Gaussian MMI-PID over pre-trained CTRNN/Elman ensembles to test whether synergy between stimulus and context is higher for context-dependent decision-making than for purely perceptual decision-making. It includes a GPU-batched PID implementation for tractable runtime, time-resolved and decision-time PID figures across all seeds, accuracy/loss comparison plots, and a separate hidden-size sweep exploring how PID components scale with RNN capacity. Supporting data (trained model weights, activations, stimulus coherences, PID outputs) and output directories were added/reorganized alongside the code. An unrelated old lecture exercise (Jan_example_rnn/) was removed from the repo.

Changes:

  • notebooks/06_Full_PID_analysis_pipeline.ipynb: new full pipeline — loads the 10-seed × {perceptual, context} CTRNN ensemble, computes Gaussian MMI-PID (total MI, redundancy, unique×2, synergy) against signed stimulus coherence, produces accuracy/loss bar plots with Mann-Whitney U control and time-resolved PID plots for all seeds; iterated across several commits (markdown/comments pass, figure generation, final verification pass).
  • src/analysis/gaussian_pid.py: new GPU-batched PID implementation (gpu_time_pid) that computes the joint covariance once per timestep and batches sub-block slogdets, ~40x faster than the reference per-seed PID computation, validated to match the reference implementation to 3e-10 bits.
  • figures/: new accuracy_loss/ (test accuracy/loss bar plot) and all_time_pid/ (time-resolved PID plots per CTRNN seed) directories; added .gitkeep placeholders for learning_curves/{context,perceptual}.
  • results/: new output directories for accuracies/losses, model activations, model weights, PID outputs (including all_time_PID.npy + metadata for both tasks), and stimulus coherences (test coherence .npz files for context and perceptual tasks, 10 seeds each).
  • size_comparison/: new hidden-size sweep study for the Context Decision-Making (CDM) and Perceptual Decision-Making (PDM) tasks — sweep script (ctrnn_pid_sweep_cdm.py), PID-vs-hidden-size plots, per-hidden-size PID timeseries plots, and a comparison notebook (plot_pid.ipynb), plus a data-loading verification notebook/script.
  • notebooks/04_minimal_product.ipynb, notebooks/05_Full_train_pipeline.ipynb: corrected data path references.
  • Removed notebooks/Jan_example_rnn/ (unrelated Scientific Computing lecture exercise and its saved RNN weights).
  • .gitignore: updated to ignore Claude-related files and coherence results directories.

Copilot AI review requested due to automatic review settings July 4, 2026 21:33
@jgendra jgendra self-assigned this Jul 4, 2026
@jgendra
jgendra merged commit e8a44fb into main Jul 4, 2026
1 check passed

Copilot AI left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pull request overview

Adds an end-to-end PID analysis workflow to the project, including a GPU-batched Gaussian MMI-PID implementation and supporting scripts/artifacts for comparing perceptual vs. context-dependent RNNs and running hidden-size sweeps.

Changes:

  • Added a GPU-accelerated, time-resolved Gaussian MMI-PID path for RNN activations.
  • Added size-comparison / sweep scripts and reorganized result/figure output directories via placeholder folders and metadata.
  • Cleaned up notebooks/paths and removed an unrelated lecture example; added pre-trained weight artifacts for comparisons.

Reviewed changes

Copilot reviewed 13 out of 135 changed files in this pull request and generated 6 comments.

Show a summary per file
File Description
src/training/.gitkeep Keeps src/training/ tracked in git.
src/models/.gitkeep Keeps src/models/ tracked in git.
src/analysis/gaussian_pid.py Adds Torch-based GPU-batched time-resolved Gaussian MMI-PID implementation and bipartition helper.
size_comparison/data_loading_ver/ctrnn_pid_both_tasks.py New CTRNN training + PID analysis script for both tasks with plotting/saving.
size_comparison/CDM/results/ctrnn_pid_sweep_cdm.py New CDM hidden-size sweep script with PID analysis and plotting.
results/stimulus_coherences/perceptual/.gitkeep Tracks perceptual coherences output directory.
results/stimulus_coherences/context/.gitkeep Tracks context coherences output directory.
results/pid_outputs/perceptual/all_time_PID_meta.json Adds metadata describing saved perceptual all-time PID array.
results/pid_outputs/perceptual/.gitkeep Tracks perceptual PID outputs directory.
results/pid_outputs/context/all_time_PID_meta.json Adds metadata describing saved context all-time PID array.
results/pid_outputs/context/.gitkeep Tracks context PID outputs directory.
results/pid_outputs/.gitkeep Tracks PID outputs root directory.
results/model_weights/perceptual/.gitkeep Tracks perceptual model weights output directory.
results/model_weights/context/.gitkeep Tracks context model weights output directory.
results/model_weights/.gitkeep Tracks model weights root directory.
results/model_activations/perceptual/.gitkeep Tracks perceptual activations output directory.
results/model_activations/context/.gitkeep Tracks context activations output directory.
results/model_activations/.gitkeep Tracks activations root directory.
results/accuracies_n_losses/perceptual/.gitkeep Tracks perceptual accuracy/loss output directory.
results/accuracies_n_losses/context/.gitkeep Tracks context accuracy/loss output directory.
notebooks/Jan_example_rnn/true_rnn/utils.py Removes old lecture/example utility code.
notebooks/Jan_example_rnn/true_rnn/trueRNN.py Removes old lecture/example model code.
notebooks/Jan_example_rnn/true_rnn/main.py Removes old lecture/example training script.
notebooks/Jan_example_rnn/dyn_rnn/utils.py Removes old lecture/example utility code.
notebooks/Jan_example_rnn/dyn_rnn/main.py Removes old lecture/example training script.
notebooks/Jan_example_rnn/dyn_rnn/dynRNN.py Removes old lecture/example model code.
notebooks/05_Full_train_pipeline.ipynb Updates base data paths and clears execution counts / kernel metadata.
notebooks/04_minimal_product.ipynb Updates base data paths, clears execution counts, and adjusts weights output path.
figures/learning_curves/perceptual/.gitkeep Tracks perceptual learning-curve figure directory.
figures/learning_curves/context/.gitkeep Tracks context learning-curve figure directory.
figures/learning_curves/.gitkeep Tracks learning-curves root directory.
elman_vs_ctrnn_comparison/model_weights/Elman_PerceptualDecisionMaking-v0_h80.pt Adds pre-trained Elman weights artifact.
elman_vs_ctrnn_comparison/model_weights/Elman_PerceptualDecisionMaking-v0_h20.pt Adds pre-trained Elman weights artifact.
elman_vs_ctrnn_comparison/model_weights/Elman_ContextDecisionMaking-v0_h20.pt Adds pre-trained Elman weights artifact.
elman_vs_ctrnn_comparison/model_weights/CTRNN_PerceptualDecisionMaking-v0_h20.pt Adds pre-trained CTRNN weights artifact.
elman_vs_ctrnn_comparison/model_weights/CTRNN_ContextDecisionMaking-v0_h20.pt Adds pre-trained CTRNN weights artifact.
.gitignore Adjusts ignore/unignore rules for results coherences and ignores .claude/.

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

Comment on lines +425 to +427
Fast drop-in for `gaussian_pid_rnn(..., timestep=None, bipartitions="random")`:
it returns the same bipartition-averaged PID atoms at every timestep, but is
~40x faster. The trick is that at each timestep the joint covariance over *all*



def gaussian_pid_rnn_gpu(H, Y, n_bip=200, seed=0, reg=1e-5, device=None):
Comment on lines +83 to +87
RESULTS_DIR = Path(r"E:\TUM\3sem\NeuroAI\project\github_neuroai\results")

WEIGHTS_DIR = RESULTS_DIR / 'weights'
WEIGHTS_DIR.mkdir(parents=True, exist_ok=True)

Comment on lines +806 to +808
test_acc = compute_decision_accuracy(model, loaders[task_key]['test'], device)
history['test_acc'] = test_acc
print(f" Test decision accuracy: {test_acc:.3f}")
Comment on lines +958 to +973
for h in args.hidden_sizes:
print(f"\n{'#'*60}")
print(f" HIDDEN SIZE = {h}")
print(f"{'#'*60}")

main(
hidden_size = h,
pdm_dir = args.pdm_dir,
cdm_dir = args.cdm_dir,
num_epochs = args.num_epochs,
patience = args.patience,
lr = args.lr,
batch_size = args.batch_size,
n_test_trials = args.n_test_trials,
n_bipartitions = args.n_bipartitions,
) No newline at end of file
Comment on lines +98 to +101
'fixation': 200, # 10 steps @ dt=20
'stimulus': 1000, # 50 steps (rounded to 38 at dt=20)
'delay': 0, # 0 steps — MUST fix stochastic delay
'decision': 100, # 5 steps
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants