Make figures (simulations)_.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "d13254ad",
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys\n",
    "import importlib\n",
    "ARBOR_LOCATION = '/specialhome/lennart/Repos/arbor-sim/arbor/build/python/arbor/__init__.py'\n",
    "spec = importlib.util.spec_from_file_location('arbor', ARBOR_LOCATION)\n",
    "module = importlib.util.module_from_spec(spec)\n",
    "sys.modules[spec.name] = module\n",
    "spec.loader.exec_module(module)\n",
    "import iopublic\n",
    "iopublic.ARBOR_BUILD_CATALOGUE = '/specialhome/lennart/Repos/arbor-sim/arbor/build/arbor-build-catalogue'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "669f01b8",
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "import iopublic\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import gzip\n",
    "sys.path.append('../io/ArborNetworks')\n",
    "import lib_single\n",
    "import arbor\n",
    "import copy\n",
    "import h5py\n",
    "import random\n",
    "\n",
    "%matplotlib inline\n",
    "\n",
    "def get_tuning(selected):\n",
    "    fn_tuned = f'tuned_networks/{selected}'\n",
    "    tuning = json.load(open(fn_tuned))\n",
    "    for mod in tuning['mods']:\n",
    "        cal = mod.pop('scal_cal', None) # oops misnamed this one\n",
    "        if cal is not None:\n",
    "            mod['cal'] = cal\n",
    "    return tuning\n",
    "\n",
    "def run(network, tuning, connected=True):\n",
    "    recipe = iopublic.TunedIOModel(network, tuning, spikes=())\n",
    "    if not connected:\n",
    "        recipe.ggap = 0\n",
    "    context = arbor.context(threads=8, gpu_id=0)\n",
    "    domains = arbor.partition_load_balance(recipe, context)\n",
    "    sim = arbor.simulation(recipe, domains, context)\n",
    "    handles = [sim.sample((gid, 0), arbor.regular_schedule(1)) for gid in range(recipe.num_cells())]\n",
    "    sim.run(tfinal=8000, dt=0.025)\n",
    "    traces = [sim.samples(handle)[0][0].T for handle in handles]\n",
    "    vsall = np.array([vs for t, vs in traces])\n",
    "    return vsall[:, -3000:]\n",
    "\n",
    "def impulse_response(tuning):\n",
    "    tuning = copy.deepcopy(tuning)\n",
    "    for mod in tuning['mods']:\n",
    "        mod.clear()\n",
    "        mod.update(dict(\n",
    "                scal_na_s=0, scal_kdr=0, scal_ks=0, scal_cal=0, scal_cah=0,\n",
    "                scal_kca=0, scal_h=0, scal_cacc=0, scal_na_a=0, scal_k=0,\n",
    "                Vm=0, Vdend=10, cal=0))\n",
    "    recipe = iopublic.TunedIOModel(network, tuning, spikes=())\n",
    "    recipe.ggap = 0\n",
    "    context = arbor.context(threads=128)\n",
    "    domains = arbor.partition_load_balance(recipe, context)\n",
    "    sim = arbor.simulation(recipe, domains, context)\n",
    "    handles = [sim.sample((gid, 0), arbor.regular_schedule(0.01)) for gid in range(recipe.num_cells())]\n",
    "    sim.run(tfinal=10, dt=0.01/2)\n",
    "    traces = [sim.samples(handle)[0][0].T for handle in handles]\n",
    "    vs = np.array([vs for t, vs in traces])\n",
    "    x1 = (vs[:,10]-vs[:,0])/10\n",
    "    x2 = 10-vs[:,-1]\n",
    "    return vs, x1, x2\n",
    "\n",
    "# from sim_titan.py in llandsmeer/io\n",
    "random.seed(1234) # constant pool of randoms\n",
    "pick_samples = [random.randint(0,999) for _i in range(10000)]\n",
    "names='scal_ks scal_cal scal_cah scal_kca scal_h scal_cacc'.split()\n",
    "with gzip.open('../io/ArborNetworks/mech_config_version7.json.gz', 'rt') as f:\n",
    "    mech_config = json.load(f)\n",
    "def get_sc_mod(nid, x1, x2, mech_config):\n",
    "    b1, b2 = x1[nid]*100-.5, x2[nid]*100-.5\n",
    "    b = min(mech_config, key=lambda ex: (b1 - ex['bucket'][0])**2 + (b2 - ex['bucket'][1])**2)\n",
    "    opt = b['opt'][pick_samples[nid] % len(b['opt'])] # consistant same sample across sims\n",
    "    config = list(opt['ind'])\n",
    "    mod = dict(zip(names, config))\n",
    "    return mod\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 59,
   "id": "a5fd49d8",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 3min 29s, sys: 4min 16s, total: 7min 45s\n",
      "Wall time: 7min 42s\n",
      "CPU times: user 11min 4s, sys: 4min 22s, total: 15min 26s\n",
      "Wall time: 7min 44s\n",
      "CPU times: user 10min 31s, sys: 4min 7s, total: 14min 39s\n",
      "Wall time: 7min 20s\n",
      "CPU times: user 3min 18s, sys: 4min 5s, total: 7min 23s\n",
      "Wall time: 7min 20s\n",
      "CPU times: user 11min 15s, sys: 4min 24s, total: 15min 40s\n",
      "Wall time: 7min 51s\n",
      "CPU times: user 11min 19s, sys: 4min 29s, total: 15min 48s\n",
      "Wall time: 7min 55s\n",
      "CPU times: user 54min 14s, sys: 25min 45s, total: 1h 20min\n",
      "Wall time: 47min 34s\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "\n",
    "selected = '2021-12-08-shadow_averages_0.001_0.8_d1666304-c6fc-4346-a55d-a99b3aad55be'\n",
    "\n",
    "network = iopublic.get_network_for_tuning(selected)\n",
    "\n",
    "C = np.array([neuron.cluster for neuron in network.neurons])\n",
    "\n",
    "tuning = get_tuning(selected)\n",
    "%time vsall_conn_tune = run(network, tuning, connected=True)\n",
    "%time vsall_unconn_tune = run(network, tuning, connected=False)\n",
    "\n",
    "for mod in tuning['mods']:\n",
    "    mod.clear()\n",
    "\n",
    "%time vsall_conn_none = run(network, tuning, connected=True)\n",
    "%time vsall_unconn_none = run(network, tuning, connected=False)\n",
    "\n",
    "vstun, x1, x2 = impulse_response(tuning)\n",
    "for gid, mod in enumerate(tuning['mods']):\n",
    "    mod.clear()\n",
    "    mod.update(get_sc_mod(gid, x1, x2, mech_config))\n",
    "\n",
    "%time vsall_conn_sc = run(network, tuning, connected=True)\n",
    "%time vsall_unconn_sc = run(network, tuning, connected=False)\n",
    "\n",
    "I = 20 \n",
    "vstun, x1, x2 = impulse_response(tuning)\n",
    "for gid, mod in enumerate(tuning['mods']):\n",
    "    mod.clear()\n",
    "    mod.update(get_sc_mod(I, x1, x2, mech_config))\n",
    "    \n",
    "%time vsall20 = run(network, tuning, connected=False)\n",
    "\n",
    "\n",
    "import h5py\n",
    "with h5py.File('./imgdata/control_tuning_v2.h5', 'a') as f:\n",
    "    f.create_dataset(f'{selected}/tune/conn', data=vsall_conn_tune)\n",
    "    f.create_dataset(f'{selected}/tune/unconn', data=vsall_unconn_tune)\n",
    "    f.create_dataset(f'{selected}/sc/conn', data=vsall_conn_sc)\n",
    "    f.create_dataset(f'{selected}/sc/unconn', data=vsall_unconn_sc)\n",
    "    f.create_dataset(f'{selected}/none/conn', data=vsall_conn_none)\n",
    "    f.create_dataset(f'{selected}/none/unconn', data=vsall_unconn_none)\n",
    "    f.create_dataset(f'{selected}/gid20/unconn', data=vsall20)\n",
    "    f.create_dataset(f'{selected}/x1', data=x1)\n",
    "    f.create_dataset(f'{selected}/x2', data=x2)\n",
    "    f.create_dataset(f'{selected}/vstune', data=vstun)\n",
    "    f.create_dataset(f'{selected}/C', data=C)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e0915a95",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.8.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}