NeuroDynamics.jl: Scalable generative modeling of neural dynamics

NeuroDynamics.jl is a Julia package for scalable and efficient generative modeling of neural dynamics [1]. Neural systems are modeld as a system of of stochastic differential equations with differentiable drift and diffusion functions.
The package provides a high-level interface for specifying and fitting these models to neural data using variational inference [2, 3] and gradient-based optimization.

Version 0.1.0-DEV

The pacakage is still in DEV mode but if you can't wait, feel free to use it and report any issues you encounter.

Installation

The dev version can be installed by running the following command in the Julia REPL:

using Pkg
Pkg.add(url="https://github.com/elgazzarr/NeuroDynamics.jl")

Tutorials, Examples, and Documentation

For more information, check out the documentation.

Example of a simple Latent SDE model

using NeuroDynamics, Lux, LuxCUDA

obs_dim = 100
ctrl_dim = 10
dev = gpu_device()

hp = Dict("n_states" => 10, "hidden_dim" => 64, "context_dim" => 32, "t_init" => 50)

obs_encoder = Recurrent_Encoder(obs_dim, hp["n_states"], hp["context_dim"],  hp["hidden_dim"], hp["t_init"])
drift =  ModernWilsonCowan(hp["n_states"], ctrl_dim)
drift_aug = Chain(Dense(hp["n_states"] + hp["context_dim"], hp["hidden_dim"], softplus), Dense(hp["hidden_dim"], hp["n_states"], tanh))
diffusion = Dense(hp["n_states"], hp["n_states"], sigmoid)
dynamics =  SDE(drift, drift_aug, diffusion, EulerHeun(), dt=0.1)
obs_decoder = MLP_Decoder(hp["n_states"], obs_dim,  hp["hidden_dim"], 1, "Poisson")   
ctrl_encoder, ctrl_decoder = NoOpLayer(), NoOpLayer()

model = LatentUDE(obs_encoder, ctrl_encoder, dynamics, obs_decoder, ctrl_decoder, dev)

Citation

If you use this package in your research, please cite the following paper:

@article{elgazzar2024universal,
  title={Universal Differential Equations as a Common Modeling Language for Neuroscience},
  author={ElGazzar, Ahmed and van Gerven, Marcel},
  journal={arXiv preprint arXiv:2403.14510},
  year={2024}
}