{
"cells": [
{
"cell_type": "markdown",
"metadata": {},
"source": [
"# [Montbrio SDE model using Cupy](https://github.com/Ziaeemehr/vbi_paper/blob/main/docs/examples/mpr_sde_cupy.ipynb)\n",
"\n",
"Estimation of global coupling $G$.\n",
"\n",
"![\"Open](\"https://colab.research.google.com/assets/colab-badge.svg\")
"
]
},
{
"cell_type": "code",
"execution_count": 18,
"metadata": {},
"outputs": [],
"source": [
"import os\n",
"import vbi\n",
"import torch\n",
"import numpy as np\n",
"import networkx as nx\n",
"from copy import deepcopy\n",
"import sbi.utils as utils\n",
"import matplotlib.pyplot as plt\n",
"from sbi.analysis import pairplot\n",
"from vbi.sbi_inference import Inference\n",
"from vbi.models.cupy.mpr import MPR_sde\n",
"\n",
"import warnings\n",
"warnings.filterwarnings(\"ignore\")"
]
},
{
"cell_type": "code",
"execution_count": 19,
"metadata": {},
"outputs": [],
"source": [
"seed = 42\n",
"np.random.seed(seed)\n",
"path = \"output/mpr_sde_cupy_\"\n",
"os.makedirs(path, exist_ok=True)\n",
"LOAD_DATA = True"
]
},
{
"cell_type": "code",
"execution_count": 20,
"metadata": {},
"outputs": [],
"source": [
"LABESSIZE = 10\n",
"plt.rcParams['axes.labelsize'] = LABESSIZE\n",
"plt.rcParams['xtick.labelsize'] = LABESSIZE\n",
"plt.rcParams['ytick.labelsize'] = LABESSIZE"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"loading connectivity matrix"
]
},
{
"cell_type": "code",
"execution_count": 21,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"number of nodes: 88\n"
]
}
],
"source": [
"D = vbi.LoadSample(nn=88)\n",
"weights = D.get_weights()\n",
"nn = weights.shape[0]\n",
"print(f\"number of nodes: {nn}\")"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Simulating BOLD single for a sample value of $G$"
]
},
{
"cell_type": "code",
"execution_count": 22,
"metadata": {},
"outputs": [],
"source": [
"TR = 300.0\n",
"fs = 1 / (TR / 1000)\n",
"t_cut = 20\n",
"par = {\n",
" \"G\": 0.506, # global coupling strength\n",
" \"weights\": weights, # connection matrix\n",
" \"method\": \"heun\", # integration method\n",
" \"dt\": 0.01,\n",
" \"t_cut\": 20_000,\n",
" \"t_end\": 100_000, # [ms]\n",
" \"num_sim\": 1, # number of simulations\n",
" \"tr\": TR,\n",
" \"rv_decimate\": 10,\n",
" \"engine\": \"cpu\", # cpu or gpu\n",
" \"seed\": seed, # seed for random number generator\n",
" \"RECORD_RV\": True,\n",
" \"RECORD_BOLD\": True,\n",
"}"
]
},
{
"cell_type": "code",
"execution_count": 54,
"metadata": {},
"outputs": [
{
"name": "stderr",
"output_type": "stream",
"text": [
"Integrating: 0%| | 0/999999 [00:00"
]
},
"metadata": {},
"output_type": "display_data"
}
],
"source": [
"if fmri_d.ndim == 3:\n",
" fig, ax = plt.subplots(3, figsize=(10, 5), sharex=True)\n",
" ax[0].set_ylabel(\"BOLD\")\n",
" ax[0].plot(fmri_t, fmri_d[:,:,0], lw=0.1)\n",
" ax[0].margins(0, 0.1)\n",
" ax[1].plot(rv_t, rv_d[:, :nn, 0], lw=0.1, alpha=0.1)\n",
" ax[2].plot(rv_t, rv_d[:, nn:, 0], lw=0.1, alpha=0.1)\n",
" ax[1].set_ylabel(\"r\")\n",
" ax[2].set_ylabel(\"v\")\n",
" ax[2].set_xlabel(\"Time [s]\")\n",
" ax[1].margins(0, 0.01)\n",
" plt.tight_layout()\n",
" plt.show()\n"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"**Training data**\n",
"\n",
"- Uniform prior for $G$ and sampling from prior;\n",
"- Selecting GPU as engine;\n",
"- Storing training BOLD signals;\n",
"- Extracting features from the simulated BOLD signals;\n",
"- Visualizing some of the features;\n",
"- Training NN and estimating parameter of G for given observed signal;\n",
"- Visualising the posterior distribution.\n"
]
},
{
"cell_type": "code",
"execution_count": 31,
"metadata": {},
"outputs": [],
"source": [
"num_sim = 512\n",
"G_min, G_max = 0.0, 1.0\n",
"\n",
"prior_min = [G_min]\n",
"prior_max = [G_max]\n",
"prior = utils.torchutils.BoxUniform(\n",
" low=torch.as_tensor(prior_min), high=torch.as_tensor(prior_max)\n",
")\n",
"\n",
"obj = Inference()\n",
"theta = obj.sample_prior(prior, num_sim, seed=seed)\n",
"\n",
"par_batch = deepcopy(par)\n",
"par_batch['G'] = theta.numpy().astype(np.float64).squeeze()\n",
"par_batch['num_sim'] = num_sim\n",
"par_batch['engine'] = 'gpu'\n",
"par_batch['RECORD_RV'] = False"
]
},
{
"cell_type": "code",
"execution_count": 43,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"['fmri_d', 'fmri_t', 'theta']\n"
]
}
],
"source": [
"if LOAD_DATA and os.path.exists(path + \"/bolds.npz\"):\n",
" sol = np.load(path + \"/bolds.npz\")\n",
" print(sol.files)\n",
"else:\n",
" obj = MPR_sde(par_batch)\n",
" sol = obj.run()"
]
},
{
"cell_type": "code",
"execution_count": 44,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"fmri_d.shape = (266, 88, 512)\n",
"fmri_t.shape = (266,)\n",
"theta.shape = torch.Size([512, 1])\n",
"bolds.shape = (512, 88, 266)\n"
]
}
],
"source": [
"fmri_d = sol[\"fmri_d\"]\n",
"fmri_t = sol[\"fmri_t\"]\n",
"fmri_d = fmri_d\n",
"fmri_t = fmri_t\n",
"bolds = np.transpose(fmri_d, (2, 1, 0))\n",
"\n",
"if not LOAD_DATA:\n",
" np.savez(path + \"/bolds.npz\", fmri_d=bolds, fmri_t=fmri_t, theta=theta.numpy().squeeze())\n",
"print(f\"fmri_d.shape = {fmri_d.shape}\")\n",
"print(f\"fmri_t.shape = {fmri_t.shape}\")\n",
"print(f\"theta.shape = {theta.shape}\")\n",
"print(f\"bolds.shape = {bolds.shape}\")"
]
},
{
"cell_type": "code",
"execution_count": 47,
"metadata": {},
"outputs": [
{
"data": {
"text/plain": [
"torch.Size([512])"
]
},
"execution_count": 47,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"theta = np.load(path + \"/bolds.npz\")[\"theta\"]\n",
"theta = torch.tensor(theta).float()\n",
"theta.shape"
]
},
{
"cell_type": "code",
"execution_count": 48,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"Selected features:\n",
"------------------\n",
"■ Domain: connectivity\n",
" ▢ Function: fcd_stat\n",
" ▫ description: Extracts features from dynamic functional connectivity (FCD)\n",
" ▫ function : vbi.feature_extraction.features.fcd_stat\n",
" ▫ parameters : {'TR': 1.0, 'win_len': 30, 'positive': False, 'eigenvalues': True, 'masks': None, 'verbose': False, 'pca_num_components': 3, 'quantiles': [0.05, 0.25, 0.5, 0.75, 0.95], 'k': None, 'features': ['sum', 'max', 'min', 'mean', 'std', 'skew', 'kurtosis']}\n",
" ▫ tag : ['fmri', 'eeg', 'meg']\n",
" ▫ use : yes\n"
]
}
],
"source": [
"from vbi import (\n",
" get_features_by_domain,\n",
" get_features_by_given_names,\n",
" report_cfg,\n",
" extract_features,\n",
")\n",
"\n",
"cfg = get_features_by_domain(\"connectivity\")\n",
"cfg = get_features_by_given_names(cfg, [\"fcd_stat\"])\n",
"report_cfg(cfg)"
]
},
{
"cell_type": "code",
"execution_count": 49,
"metadata": {},
"outputs": [
{
"name": "stderr",
"output_type": "stream",
"text": [
"100%|██████████| 512/512 [00:04<00:00, 108.12it/s]\n"
]
},
{
"data": {
"application/vnd.microsoft.datawrangler.viewer.v0+json": {
"columns": [
{
"name": "index",
"rawType": "int64",
"type": "integer"
},
{
"name": "fcd_full_sum",
"rawType": "float64",
"type": "float"
},
{
"name": "fcd_full_ut_std",
"rawType": "float64",
"type": "float"
},
{
"name": "G",
"rawType": "float32",
"type": "float"
}
],
"ref": "9e6c644e-daa6-4be4-b502-98cf2b3c4577",
"rows": [
[
"0",
"7322.4306640625",
"0.03303716331720352",
"0.4294045"
],
[
"1",
null,
null,
"0.8854429"
],
[
"2",
"7455.4052734375",
"0.032545045018196106",
"0.57390445"
],
[
"3",
"6776.82421875",
"0.025232113897800446",
"0.26658005"
],
[
"4",
"7109.46142578125",
"0.025810761377215385",
"0.62744915"
]
],
"shape": {
"columns": 3,
"rows": 5
}
},
"text/html": [
"\n",
"\n",
"
\n",
" \n",
" \n",
" | \n",
" fcd_full_sum | \n",
" fcd_full_ut_std | \n",
" G | \n",
"
\n",
" \n",
" \n",
" \n",
" | 0 | \n",
" 7322.430664 | \n",
" 0.033037 | \n",
" 0.429404 | \n",
"
\n",
" \n",
" | 1 | \n",
" NaN | \n",
" NaN | \n",
" 0.885443 | \n",
"
\n",
" \n",
" | 2 | \n",
" 7455.405273 | \n",
" 0.032545 | \n",
" 0.573904 | \n",
"
\n",
" \n",
" | 3 | \n",
" 6776.824219 | \n",
" 0.025232 | \n",
" 0.266580 | \n",
"
\n",
" \n",
" | 4 | \n",
" 7109.461426 | \n",
" 0.025811 | \n",
" 0.627449 | \n",
"
\n",
" \n",
"
\n",
"