Package: vbjax
vbjaxvbjaxis a Jax-based package for working with virtual brain style models.InstallationInstalls withpip install "vbjax", but for the latest features,
you can use the source,gitclonehttps://github.com/ins-amu/vbjaxcdvbjax
pipinstall-e".[dev]"You're encouraged to have the source handy to consult and change, but you can also justpipinstallgit+https://github.com/ins-amu/vbjaxThe primary additional dependency of vbjax isJAX, which itself depends only on
NumPy, SciPy & opt-einsum, so it should be safe to add to your
existing projects.gee pee youCUDAIf you have a CUDA-enabled GPU, you install the requisite dependencies like sopipinstall--upgrade"jax[cuda11_pip]"-fhttps://storage.googleapis.com/jax-releases/jax_cuda_releases.htmlM1/M2 🍎On newer Apple machines w/ M1 or M2 GPUs, JAX supports using the GPU experimentally
by installing just two extra packages:pip install ml-dtypes==0.2.0 jax-metalAbout a third of vbjax tests fail due to absence of certain operations like n-dim
scatter/gather & FFTs, and it may not be faster because these CPUs already have
excellent memory bandwidth & latency hiding.CUDA 🐳BUTbecause GPU software stack versions make aligning stars look like child's play,
container images are available and auto-built w/GitHub Actions, so you can use w/ Dockerdockerrun--rm-itghcr.io/ins-amu/vbjax:mainpython3-c'import vbjax; print(vbjax.__version__)'The images are built on Nvidia runtime images, so--gpus allis enough
for Jax to discover the GPU(s).ExamplesHere are some examples of simulations which show what you can
do with the library.  Because they are implemented atop Jax, it
is easy to take gradients for optimization or MCMC, or do efficient
GPU parallel batching. Or both ಠ_ರೃSimple networkHere's the smallest simulation you might want to do:
an all-to-all connected network with Montbrio-Pazo-Roxin
mass model dynamics,importvbjaxasvbimportjax.numpyasnpdefnetwork(x,p):c=0.03*x.sum(axis=1)returnvb.mpr_dfun(x,c,p)_,loop=vb.make_sde(dt=0.01,dfun=network,gfun=0.1)zs=vb.randn(500,2,32)xs=loop(zs[0],zs[1:],vb.mpr_default_theta)vb.plot_states(xs,'rV',jpg='example1',show=True)While integrators and mass models tend to be the same across publications, but
the network model itself varies (regions vs surface, stimulus etc), vbjax allows
user to focus on defining thenetworkand then getting time series.Parallel parameter space explorationOne of the key tools in use of these models is to run parameter sweeps, which
can be done easily with vbjax and parallelized with builtin Jax tools.importjax,jax.numpyasnpimportvbjaxasvbLet's
look at the effect of coupling, noise and excitability in a small network model,
where thepparameter is a tuple of coupling scaling, noise scaling and the
base MPR parameters:defnet(x,p):r,v=xk,_,mpr_p=pc=k*r.sum(),k*v.sum()returnvb.mpr_dfun(x,c,mpr_p)to ensure noise is a parameter we can tune later, we'll define the SDE with
a dynamic noise scaled by the second element of the parameter tuplep,defnoise(_,p):_,sigma,_=preturnsigmaturn that into an SDE, choose network size (8 nodes here), with random initial
conditions and realize an appropriate sample of noise for SDE integration:_,loop=vb.make_sde(0.01,net,noise)n_nodes=8rv0=vb.randn(2,n_nodes)zs=vb.randn(1000,*rv0.shape)to run map parameters of interest to some metric of interest, we write a function,
where we've chosen the standard deviation after a certain time as a index on the steady
state of the network dynamics,defrun(pars,mpr_p=vb.mpr_default_theta):k,sig,eta=pars# explored parsp=k,sig,mpr_p._replace(eta=eta)# set mprxs=loop(rv0,zs,p)# run simstd=xs[400:,0].std()# eval metricreturnstd# doneall that is plain virtual brain equations, but to enable efficient & parallel
sweeps, we can pull in jax primitivesjax.pmapto parallelize over compute
devices (by default, cores of the CPU) andjax.vmapto vectorize the functionrun,orjust ajax.vmapif using a GPU,using_cpu=jax.local_devices()[0].platform=='cpu'ifusing_cpu:run_batches=jax.pmap(jax.vmap(run,in_axes=1),in_axes=0)else:run_batches=jax.vmap(run,in_axes=1)now we run this over a parameter space (adapting for number of cores)# sweep sigma but just a few values are enoughsigmas=[0.0,0.2,0.3,0.4]results=[]ng=vb.cores*4ifusing_cpuelse32fori,sig_iinenumerate(sigmas):# create grid of k (on logarithmic scale) and etalog_ks,etas=np.mgrid[-9.0:-2.0:1j*ng,-4.0:-6.0:1j*ng]# reshape grid to big batch of valuespars=np.c_[np.exp(log_ks.ravel()),np.ones(log_ks.size)*sig_i,etas.ravel()].T.copy()# cpu w/ pmap expects a chunk for each coreifusing_cpu:pars=pars.reshape((3,vb.cores,-1)).transpose((1,0,2))# now runresult=run_batches(pars).block_until_ready()results.append(result)and plot the results,importpylabasplpl.figure(figsize=(8,2))fori,(sig_i,result)inenumerate(zip(sigmas,results)):pl.subplot(1,4,i+1)pl.imshow(result.reshape(log_ks.shape),vmin=0.2,vmax=0.7)pl.ylabel('k')ifi==0else(),pl.xlabel('eta')pl.title(f'sig ={sig_i:0.1f}')pl.show()pl.savefig('example3.jpg')A full runnable script is inexamples/parsweep.py.Performance notesIs (vb)jax fast?  Single threaded C/C++ code typically reaches 50 to 200 thousand
iterations per second on a single CPU core, for network sizes of e.g. 164
(e.g. Destrieux atlas from FreeSurfer).  Let's check efficiency of the Jax code
for some hardware on my desk:Xeon W-2133 (2017 14 nm Skylake) uses 88 W to do 5.7 Miter/s = 65 Kiter/WQuadro RTX 5000 (2018 12 nm Turing) uses 200 W to do 26 Miter/s = 130 Kiter/WM1 Air (2020 5 nm M1) uses 18 W to do 3.7 Miter/s = 205 Kiter/WYou can tweakthis scriptto see what you might
expect on your hardware.DistributedIf you are using a system like Dask or Slurm, you can then invoke
thatrun_batchesfunction in a distributed setting as required,
without needing to manage a per core or per node for loop.Simplest neural fieldHere's a neural field,importjax.numpyasnpimportvbjaxasvb# setup local connectivitylmax,nlat,nlon=16,32,64lc=vb.make_shtdiff(lmax=lmax,nlat=nlat,nlon=nlon)# network dynamicsdefnet(x,p):c=lc(x[0]),0.0returnvb.mpr_dfun(x,c,p)# solution + plotx0=vb.randn(2,nlat,nlon)*0.5+np.r_[0.2,-2.0][:,None,None]_,loop=vb.make_sde(0.1,net,0.2)zs=vb.randn(500,2,nlat,nlon)xt=loop(x0,zs,vb.mpr_default_theta._replace(eta=-3.9,cr=5.0))vb.make_field_gif(xt[::10],'example2.gif')This example shows how the field forms patterns gradually despite the
noise in the simulation, due to the effect of local connectivityMCMC estimation of neural field activityFor MCMC estimates with NumPyro we define a function to compute
posterior log probabilityp(theta | x),deflogp(xt=None):x0h=numpyro.sample('x0h',dist.Normal(jnp.zeros((nlat,nlon)),1))xth_mu=loop(x0h,ts,k)numpyro.sample('xth',dist.Normal(xth_mu,1),obs=xt)run MCMC w/ NUTS,mcmc=MCMC(NUTS(logp),num_warmup=500,num_samples=500)mcmc.run(jax.random.PRNGKey(0),xt=xt)x0h=mcmc.get_samples()['x0h']check diagnostics like estimated sample size, shrinkage and z-score,ess=numpyro.diagnostics.effective_sample_size(x0h.reshape((1,500,-1)))assertess.min()>100shrinkage,zscore=vbjax.shrinkage_zscore(x0,x0h,1)assertshrinkage.min()>0.7assertzscore.max()<1.5Full code is in thetest suite, can
be runpytest -m slow, since it takes about 5 minutes to run on a GPU, 7 min on m1 CPU core and
12 minutes on an x86_64 CPU core.Fitting an autoregressive processHere's a 1-lag MVARimportjaximportjax.numpyasnpimportvbjaxasvbnn=8true_A=vb.randn(nn,nn)_,loop=vb.make_sde(1,lambdax,A:-x+(A*x).mean(axis=1),1)x0=vb.randn(nn)zs=vb.randn(1000,nn)xt=loop(x0,zs,true_A)xtandtrue_Aare the simulated time series and ground truth
interaction matrices.To fit anything we need a loss function & gradient descent,defloss(est_A):returnnp.sum(np.square(xt-loop(x0,zs,est_A)))grad_loss=jax.grad(loss)est_A=np.ones((nn,nn))*0.3# wrongforiinrange(51):est_A=est_A-0.01*grad_loss(est_A)ifi%10==0:print('step',i,'log loss',np.log(loss(est_A)))print('mean sq err',np.square(est_A-true_A).mean())which printsstep 0 log loss 5.8016257
step 10 log loss 3.687574
step 20 log loss 1.7174681
step 30 log loss -0.15798996
step 40 log loss -1.9851608
step 50 log loss -3.7805486
mean sq err 8.422789e-05This is a pretty simple example but it's meant to show that any model
you build with vbjax like this is usable with optimization or NumPyro's
MCMC algorithms.ƪ(ړײ)‎ƪ​​ moar examples‽More complex examples are in theexamplesfolder:high resolution connectome neural field simulation & inferenceparameter sweep examplemore examples  cooking🍩HPC usageWe use this on HPC systems, most easily with container images.  Open an issue if it doesn't work.CSCS Piz DaintUseful modulesmoduleloaddaint-gpu
moduleloadcudatoolkit/11.2.0_3.39-2.1__gf93aa1c
moduleloadTensorFlowthen install in some Python environment; the default works finepip3install"jax[cuda]==0.3.8"-fhttps://storage.googleapis.com/jax-releases/jax_cuda_releases.html
pip3install"jaxlib==0.3.8+cuda11.cudnn805"-U-fhttps://storage.googleapis.com/jax-releases/jax_cuda_releases.htmlThis provides an older version of JAX unfortunately.The Sarus runtime can be used to make use of latest versions of vbjax and jax:$moduleloaddaint-gpu
$moduleloadsarus
$saruspullghcr.io/ins-amu/vbjax:main
...
$srun-pdebug-Aich042-Cgpu--ptysarusrunghcr.io/ins-amu/vbjax:mainpython3-c'import jax; print(jax.numpy.zeros(32).device())'...
gpu:0JSC JUSUFA nice module is available to get CUDA libsmoduleloadcuDNN/8.6.0.163-CUDA-11.7then you might set up a conda env,wgethttps://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh
bashMiniconda3-latest-Linux-x86_64.sh-b-p~/conda
.~/conda/bin/activate
condacreate-njaxpython=3.9numpyscipysourceactivatejaxonce you have an env, install the CUDA-enabled JAXpip3install--upgrade"jax[cuda]"-fhttps://storage.googleapis.com/jax-releases/jax_cuda_releases.htmland check it works(jax)[woodman1@jsfl02~]$srun-Aicei-hbp-2021-0002-pdevelgpus--ptypython3-c'import jax.numpy as np ; print(np.zeros(32).device())'gpu:0JSC also makes Singularity available, so the prebuilt image can be usedTODOCEAThe prebuilt image is the best route:TODODevelopmentNew ideas or even documenting tricks (like how Jax works) should go
into the test suite, and there are some ideas floating there before
making it into the library itself.git clone https://github.com/ins-amu/vbjax
cd vbjax
pip install '.[dev]'
pytestInstalling SHTnsThis library is used for some testing.  It is impossible to install on
Windows natively, so WSLx is required.On macOS,brewinstallfftw
gitclonehttps://bitbucket.org/nschaeff/shtns
./configure--enable-python--disable-simd--prefix=/opt/homebrew
make-j&&makeinstall&&pythonsetup.pyinstallReleasesa release of versionv1.2.3requires following stepsgit checkout main: tag releases from main for nowedit_version.pyto have correct release numberpython -m vbjax._version tagto create and push new tagGitHub tests, builds and pushes tag release to PyPIuse GitHub UI to create new release