{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Wong - Wang DMF\n",
    "\n",
    "\n",
    "Kong, X., Kong, R., Orban, C., Wang, P., Zhang, S., Anderson, K., Holmes, A., Murray, J.D., Deco, G., van den Heuvel, M. and Yeo, B.T., 2021. Sensory-motor cortices shape functional connectivity dynamics in the human brain. Nature communications, 12(1), p.6373.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os \n",
    "import sbi.utils as utils\n",
    "from copy import deepcopy\n",
    "from vbi.sbi_inference import Inference\n",
    "import matplotlib.pyplot as plt\n",
    "from vbi import report_cfg, update_cfg\n",
    "from vbi import extract_features_list, extract_features\n",
    "from vbi import get_features_by_domain, get_features_by_given_names\n",
    "from vbi.models.pytorch. ww_sde_kong import *"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "seed = 42\n",
    "torch.manual_seed(seed)\n",
    "np.random.seed(seed)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "# engine = \"gpu\"  # \"cpu\"\n",
    "dtype = torch.float64\n",
    "path = \"output\"\n",
    "os.makedirs(path, exist_ok=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "dict_keys(['G', 'J', 'w', 's', 'I0', 'a', 'b', 'd', 'tau_s', 'gamma_s', 't_end', 't_cut', 'tr', 'dt', 'n_sim', 'weights', 'engine', 'device', 'dtype'])"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "obj = WW_SDE_KONG()\n",
    "params = obj.get_default_params()\n",
    "params.keys()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "data = obj.run()\n",
    "t = data[\"t\"]\n",
    "x = data[\"x\"]\n",
    "\n",
    "fig, ax = plt.subplots(1, 1, figsize=(12, 2.5))\n",
    "ax.plot(t, x[:,0,:].T, lw=1, alpha=0.7)\n",
    "ax.set_xlabel(\"Time [s]\")\n",
    "ax.set_ylabel(\"BOLD\")\n",
    "plt.tight_layout()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "obj = WW_SDE_KONG()\n",
    "params = obj.get_default_params()\n",
    "par = deepcopy(params)\n",
    "par['G'] = np.linspace(0,7,10)\n",
    "par['n_sim'] = len(par['G'])\n",
    "obj = WW_SDE_KONG(par)\n",
    "data = obj.run()\n",
    "t = data[\"t\"]\n",
    "x = data[\"x\"]\n",
    "\n",
    "print(t.shape, x.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "# fig, ax = plt.subplots(1, 1, figsize=(12, 2.5))\n",
    "# ax.plot(t, x[:,9,:].T, lw=0.5, alpha=0.8)\n",
    "# ax.set_xlabel(\"Time [s]\")\n",
    "# ax.set_ylabel(\"BOLD\")\n",
    "# plt.tight_layout()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}