1_Fit_posterior.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "f4bc4ecc",
   "metadata": {},
   "source": [
    "In this tutorial, we show how to use the fSBI method proposed in the companion paper.  \n",
    "We will be performing an additional fSBI round as follows:\n",
    "- we start from the final posterior in the paper (the \"plausible candidate rules\", π<sub>3</sub>, either for MLP or polynomial plasticity rules).\n",
    "- we will generate additional rules with more specific excitatory rates (between 5 and 10Hz, the initial filtering done in the paper includes rules with rates between 1 and 50Hz).\n",
    "- we will train a new posterior p(θ|r<sub>exc</sub>)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "789acd2c",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2022-10-12T15:59:24.057282Z",
     "start_time": "2022-10-12T15:59:22.622463Z"
    }
   },
   "outputs": [],
   "source": [
    "from fsbi.density_estimator import MakePosterior\n",
    "from fsbi.utils import get_density_estim_data, resample_old_data, load_and_merge, apply_n_conditions, _make_unique_samples\n",
    "from fsbi.prior import RestrictedPrior as Prior\n",
    "from torch import FloatTensor as FT\n",
    "import matplotlib.pyplot as plt\n",
    "import torch\n",
    "import numpy as np\n",
    "import h5py\n",
    "import argparse\n",
    "import yaml\n",
    "import pickle\n",
    "from time import time"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "60436cb6",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true
   },
   "source": [
    "#### 1/ Choose samples from the polynomial search space (See Fig 2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c8c34726",
   "metadata": {},
   "outputs": [],
   "source": [
    "round_name = \"pi3_r5to10Hz\"\n",
    "\n",
    "# directory where all the rules simulated in the paper are stored, with all metrics computed.\n",
    "data_path = \"data_synapsesbi/bg_IF_EEEIIEII_6pPol/\"\n",
    "# name: bg: stability task (background inputs), IF: integrate and fire neuron model in Auryn,EE EI IE II recurrent synapses plastic, 6pPol: polynomial parmeterization with 6 params. \n",
    "\n",
    "# where to store the posterior object\n",
    "runs_path = \"runs_synapsesbi/bg_IF_EEEIIEII_6pPol/\"\n",
    "\n",
    "# simulation parameters for the auryn simulations, later down the road\n",
    "path_to_sim_config = \"tasks_configs/bg_IF_EEEIIEII_6pPol.yaml\"\n",
    "with open(path_to_sim_config, \"r\") as f:\n",
    "    sim_params = yaml.load(f, Loader=yaml.Loader)\n",
    "    \n",
    "# the bounds for the plasticity parameters\n",
    "# time constants are between 10ms and 100ms, other params between -2 and 2\n",
    "lower_lim=torch.tensor([0.01, 0.01, -2., -2., -2., -2.,\n",
    "            0.01, 0.01, -2., -2., -2., -2.,\n",
    "            0.01, 0.01, -2., -2., -2., -2.,\n",
    "            0.01, 0.01, -2., -2., -2., -2.])\n",
    "upper_lim=torch.tensor([.1, .1, 2., 2., 2., 2.,\n",
    "        .1, .1, 2., 2., 2., 2.,\n",
    "        .1, .1, 2., 2., 2., 2.,\n",
    "        .1, .1, 2., 2., 2., 2.])\n",
    "\n",
    "# which metrics to train the posterior on\n",
    "# fitting a posterior on too many metrics at the same time (especially similar/correlated metrics) generates worse quality posteriors, in our hands at least.\n",
    "metrics = [\"rate\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "711708ea",
   "metadata": {},
   "outputs": [],
   "source": [
    "# here we are building the training set to fit the posterior on.\n",
    "# there is a trade-off between 1/ including as many samples for training as possible and \n",
    "# 2/ allocating resources to fit the posterior to uninteresting regions of the parameter space (e.g. here we do not care about rules with high firing rates so need to include too many of them)\n",
    "\n",
    "# first, we load all the rules simulated in the paper are stored, with all metrics computed. (pi0 -> pi3).\n",
    "dataset_aux = load_and_merge(data_path, (\"bg_IF_EEEIIEII_6pPol_all.npy\",))\n",
    "\n",
    "# then we define constrains to only keep rules relevant for our round of fSBI.\n",
    "\n",
    "# all the conditions for plausibility from the paper\n",
    "cond_cv = (\"cv_isi\", 0.7, 1000)\n",
    "cond_sf = (\"spatial_Fano\", 0.5, 2.5)\n",
    "cond_tf = (\"temporal_Fano\", 0.5, 2.5)\n",
    "cond_ac = (\"auto_cov\", 0, 0.1)\n",
    "cond_fft = (\"fft\", 0, 1)\n",
    "cond_wb = (\"w_blow\", 0, 0.1)\n",
    "cond_srt = (\"std_rate_temporal\", 0, 0.5)\n",
    "cond_srs = (\"std_rate_spatial\", 0, 5)\n",
    "cond_scv = (\"std_cv\", 0, 0.2)\n",
    "cond_wc = (\"w_creep\", 0, 0.05)\n",
    "cond_ri = (\"rate_i\", 1, 50)\n",
    "cond_weef =(\"weef\", 0 ,0.5)\n",
    "cond_weif =(\"weif\", 0 ,0.5)\n",
    "cond_wief =(\"wief\", 0 ,5)\n",
    "cond_wiif =(\"wiif\", 0 ,5)\n",
    "\n",
    "# only keep the rules with exc rates between 3 and 15Hz for the training ([5,10]Hz + some leeway)\n",
    "cond_r = (\"rate\", 3, 15)\n",
    "\n",
    "condition = apply_n_conditions(dataset_aux, (cond_r,cond_ri,\n",
    "            cond_wb,cond_wc,cond_weef,cond_weif, cond_wief, cond_wiif,\n",
    "            cond_ac,cond_cv,cond_fft,cond_srt,cond_srs,cond_sf,cond_tf))\n",
    "\n",
    "dataset = dataset_aux[condition]\n",
    "print(str(np.sum(condition)) + \"/\" + str(len(condition)), \"samples kept for training\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "72a60cfb",
   "metadata": {
    "heading_collapsed": true,
    "jp-MarkdownHeadingCollapsed": true
   },
   "source": [
    "#### 1 bis/ Choose samples from the MLP search space (See Fig 4) \n",
    "\n",
    "This section is an alternative to the polynomial one above in case you want to use MLP rules."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "19c0c863",
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": [
    "round_name = \"pi3_r5to10Hz\"\n",
    "data_path = \"data_synapsesbi/bg_CVAIF_EEIE_T4wvceciMLP/\"\n",
    "runs_path = \"runs_synapsesbi/bg_CVAIF_EEIE_T4wvceciMLP/\"\n",
    "path_to_sim_config = \"./tasks_configs/bg_CVAIF_EEIE_T4wvceciMLP.yaml\"\n",
    "with open(path_to_sim_config, \"r\") as f:\n",
    "    sim_params = yaml.load(f, Loader=yaml.Loader)\n",
    "    \n",
    "lower_lim=torch.tensor([0., 0., -1., -1., -1., -1., #etaEE, etaIE, WpreEE, WpostEE, WpreIE, WpostIE\n",
    "            -1., -1., -1., -1., -1., -1.,\n",
    "            -1., -1., -1., -1., -1., -1.,\n",
    "            -1., -1., -1., -1.])\n",
    "upper_lim=torch.tensor([1., 1., 1., 1., 1., 1., #etaEE, etaIE, WpreEE, WpostEE, WpreIE, WpostIE\n",
    "            1., 1., 1., 1., 1., 1.,\n",
    "            1., 1., 1., 1., 1., 1.,\n",
    "            1., 1., 1., 1.])\n",
    "\n",
    "metrics = [\"rate\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "eb2e29eb",
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": [
    "## Choose which samples to use for the fit\n",
    "dataset = load_and_merge(data_path, (\"bg_CVAIF_EEIE_T4wvceciMLP_all.npy\",))\n",
    "\n",
    "cond_r = (\"rate\", 3, 15)\n",
    "\n",
    "cond_cv = (\"cv_isi\", 0.7, 1000)\n",
    "cond_sf = (\"spatial_Fano\", 0.5, 2.5)\n",
    "cond_tf = (\"temporal_Fano\", 0.5, 2.5)\n",
    "cond_ac = (\"auto_cov\", 0, 0.1)\n",
    "cond_fft = (\"fft\", 0, 1)\n",
    "cond_wb = (\"w_blow\", 0, 0.1)\n",
    "cond_srt = (\"std_rate_temporal\", 0, 0.5)\n",
    "cond_srs = (\"std_rate_spatial\", 0, 5)\n",
    "cond_scv = (\"std_cv\", 0, 0.2)\n",
    "cond_wc = (\"w_creep\", 0, 0.05)\n",
    "cond_ri = (\"rate_i\", 1, 50)\n",
    "cond_weef =(\"weef\", 0 ,0.5)\n",
    "cond_wief =(\"wief\", 0 ,5)\n",
    "\n",
    "condition = apply_n_conditions(dataset_aux, (cond_r,cond_ri,\n",
    "            cond_wb,cond_wc,cond_weef,cond_wief,\n",
    "            cond_ac,cond_cv,cond_fft,cond_srt,cond_srs,cond_sf,cond_tf))\n",
    "\n",
    "dataset = dataset_aux[condition]\n",
    "print(str(np.sum(condition)) + \"/\" + str(len(condition)), \"samples kept for training\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9c112094",
   "metadata": {},
   "source": [
    "#### 2/ Fit a posterior with the sbi package"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6694cb6d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# prepare the training set for the posterior: [theta, xs]\n",
    "thetas = torch.tensor(dataset['theta'][:,:-1], dtype=torch.float32) #remove nuisance parameter from training (input rate)\n",
    "xs = torch.tensor([[dataset[i][j] for i in metrics] for j in range(len(dataset))], dtype=torch.float32)\n",
    "\n",
    "# some requirements for using the sbi package, not useful in our specific case\n",
    "bounds = {'low':lower_lim,\n",
    "          'high':upper_lim}\n",
    "prior = torch.distributions.Uniform(low=lower_lim, high=upper_lim)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d73e6cc0",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2022-10-12T16:00:20.023663Z",
     "start_time": "2022-10-12T16:00:00.071564Z"
    }
   },
   "outputs": [],
   "source": [
    "# train and save the posterior: this can take time (~1h) depending on your hardware. If too long you can include fewer samples in the training set (at the cost of a potentially worse posterior)\n",
    "mk_post = MakePosterior(**sim_params[\"prior_params\"])\n",
    "mk_post.get_ensemble_posterior(\n",
    "    thetas,\n",
    "    xs,\n",
    "    save_path=runs_path + \"posterior_\" + round_name + \".pkl\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e7a28901",
   "metadata": {},
   "source": [
    "Congratulations, you now have a posterior you can sample new rules from!     \n",
    "\n",
    "Head to the Sample_from_posterior jupyter notebook to do that."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7c2b63b4",
   "metadata": {},
   "source": [
    "#### Bonus: initial round (starting fSBI from scratch)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d342daba",
   "metadata": {},
   "source": [
    "###### If round 0: draw samples from prior, WITHOUT classifier"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "715b4412",
   "metadata": {},
   "outputs": [],
   "source": [
    "data_path = \"data_synapsesbi/bg_IF_EEEIIEII_6pPol/\"\n",
    "\n",
    "### POLY\n",
    "prior_cls = \"_Prior_bg_IF_EEEIIEII_6pPol\"\n",
    "prior_cls_args = dict(lower_lim=[0.01, 0.01, -2., -2., -2., -2.,\n",
    "                         0.01, 0.01, -2., -2., -2., -2.,\n",
    "                         0.01, 0.01, -2., -2., -2., -2.,\n",
    "                         0.01, 0.01, -2., -2., -2., -2.,\n",
    "                         5],\n",
    "              upper_lim=[.1, .1, 2., 2., 2., 2.,\n",
    "                         .1, .1, 2., 2., 2., 2.,\n",
    "                         .1, .1, 2., 2., 2., 2.,\n",
    "                         .1, .1, 2., 2., 2., 2.,\n",
    "                         15]\n",
    "                     )\n",
    "\n",
    "## MLP\n",
    "# prior_cls = \"_Prior_bg_CVAIF_EEIE_T4wvceciMLP\"\n",
    "# prior_cls_args = dict(lower_lim=[0., 0., -1., -1., -1., -1., #etaEE, etaIE, WpreEE, WpostEE, WpreIE, WpostIE\n",
    "#                 -1., -1., -1., -1., -1., -1.,\n",
    "#                 -1., -1., -1., -1., -1., -1.,\n",
    "#                 -1., -1., -1., -1., 5],\n",
    "#     upper_lim= [1., 1., 1., 1., 1., 1., #etaEE, etaIE, WpreEE, WpostEE, WpreIE, WpostIE\n",
    "#                 1., 1., 1., 1., 1., 1.,\n",
    "#                 1., 1., 1., 1., 1., 1.,\n",
    "#                 1., 1., 1., 1., 15])\n",
    "\n",
    "\n",
    "\n",
    "prior = Prior(prior_class=prior_cls,\n",
    "              prior_class_args=prior_cls_args,\n",
    "              return_numpy=False,\n",
    "              restrict_prior=False,\n",
    "              )\n",
    "\n",
    "num_samples = 10\n",
    "thetas, seeds = _make_unique_samples((num_samples,),\n",
    "                                     prior=prior,\n",
    "                                     saved_seeds=[])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8d6634e7",
   "metadata": {},
   "source": [
    "###### If round 0+: draw samples from prior + classifier\n",
    "\n",
    "We train a classifier on the training set to accept or reject new samples (for sample efficiency).  \n",
    "This is done only during the inner iterations of round 0 for sample efficiency.  \n",
    "\n",
    "See this paper for more info: https://projecteuclid.org/ebooks/institute-of-mathematical-statistics-lecture-notes-monograph-series/A-Festschrift-for-Herman-Rubin/chapter/Generalized-Accept-Reject-sampling-schemes/10.1214/lnms/1196285403"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "593b379e",
   "metadata": {},
   "outputs": [],
   "source": [
    "## load training data for classifier\n",
    "\n",
    "## POLY\n",
    "for_classifier = np.load(data_path + \"box_prior_metrics.npy\")\n",
    "\n",
    "## MLP\n",
    "# for_classifier = np.load(save_dir + \"box_prior_metrics.npy\")\n",
    "\n",
    "theta_classifier, x_classifier = FT(for_classifier['theta']), FT(for_classifier['rate'])\n",
    "sum(~torch.isfinite(x_classifier)) / len(x_classifier)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.7"
  },
  "latex_envs": {
   "LaTeX_envs_menu_present": true,
   "autoclose": false,
   "autocomplete": true,
   "bibliofile": "biblio.bib",
   "cite_by": "apalike",
   "current_citInitial": 1,
   "eqLabelWithNumbers": true,
   "eqNumInitial": 1,
   "hotkeys": {
    "equation": "Ctrl-E",
    "itemize": "Ctrl-I"
   },
   "labels_anchors": false,
   "latex_user_defs": false,
   "report_style_numbering": false,
   "user_envs_cfg": false
  },
  "vscode": {
   "interpreter": {
    "hash": "61711b8cf7d2775f177b340d6eca0bb8d2b4f06be813726212ddbc6469b2ca7b"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}