JEDI: Jointly Embedded Inference of Neural Dynamics

Under development

The current research code remains work-under-progress and could use more documentation and examples. Please reach out to anirudh.jamkhandi@umontreal.ca if you have any questions.

Installation

Tested python version : 3.11 Tested Pip version : 23.2.1

python -m venv .venv #creates a virtual environment
git clone https://github.com/sinthlab/JEDI.git # clones the repository 
cd JEDI/ 
pip install -r requirements.txt # installs the required packages
source .venv/bin/activate

Data setup

1. To generate the synthetic teacher data, run the notebook :

JEDI/notebooks/multi_task_teacher.ipynb

2. To generate data from task trained RNN : Follow the instructions on Computation Through Dynamics Benckmark Github

3. The motor cortex recordings was obtained from the Perich et al., 2018, Link. To preprocess the data, use the notebook :

JEDI/notebooks/preprocess_motor_data.ipynb

Training JEDI

To train the model on :

1. Synthetic data(multi-task and multi-frequency) locally,

cd JEDI/scripts/local/
./train_teacher_data.sh

2. Task Trained RNN activations locally,

cd JEDI/scripts/local/
./train_ctd.sh

3. Motor Cortex data locally,

cd JEDI/scripts/local/
./train_motor_data.sh

Extracting fixed points

To extract the fixed poins for :

1. Task trained RNN activations on MemoryPro Task on cluster,

cd JEDI/scripts/slurm/sarray_get_fp_per_condition.sh

Visualize them here,

JEDI/notebooks/fp_memory_pro.ipynb

2. Motor Cortex recordings on cluster,

cd JEDI/scripts/slurm/sarray_get_fp_per_condition.sh

Visualize them here,

JEDI/notebooks/fp_motor_cortex.ipynb

Eigen Spectrum Analysis of JEDI weights

Visualize the spectrum in these notebooks,

JEDI/notebooks/eigen_spectrum_motorcortex.ipynb

and

JEDI/notebooks/eigen_spectrum_motorcortex.ipynb

Cite

Please cite our paper if you use this code in your own work:

@article{jamkhandi2026jedi,
  title={JEDI: Jointly Embedded Inference of Neural Dynamics},
  author={Jamkhandi, Anirudh and Korojy, Ali and Codol, Olivier and Lajoie, Guillaume and Perich, Matthew G},
  journal={arXiv preprint arXiv:2603.10489},
  year={2026}
}