testbook.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "%reload_ext autoreload\n",
    "%autoreload 2\n",
    "\n",
    "# %matplotlib widget\n",
    "%matplotlib qt\n",
    "# %gui qt\n",
    "\n",
    "import time\n",
    "import random\n",
    "from copy import copy\n",
    "from tqdm import tqdm\n",
    "from pathlib import Path\n",
    "import numpy as np\n",
    "import scipy as sp\n",
    "from sklearn.decomposition import PCA\n",
    "# import torch\n",
    "import matplotlib as mpl\n",
    "import matplotlib.pyplot as plt\n",
    "import faststats as fs\n",
    "from vrAnalysis import helpers\n",
    "\n",
    "import allensdk"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "from vrAnalysis.simulations.simulator import Simulator"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 82,
   "metadata": {},
   "outputs": [
    {
     "ename": "ValueError",
     "evalue": "not enough values to unpack (expected 2, got 0)",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[1;31mValueError\u001b[0m                                Traceback (most recent call last)",
      "Cell \u001b[1;32mIn[82], line 2\u001b[0m\n\u001b[0;32m      1\u001b[0m box_length \u001b[38;5;241m=\u001b[39m \u001b[38;5;241m80\u001b[39m\n\u001b[1;32m----> 2\u001b[0m sim \u001b[38;5;241m=\u001b[39m \u001b[43mSimulator\u001b[49m\u001b[43m(\u001b[49m\n\u001b[0;32m      3\u001b[0m \u001b[43m    \u001b[49m\u001b[43mbox_length\u001b[49m\u001b[43m,\u001b[49m\n\u001b[0;32m      4\u001b[0m \u001b[43m    \u001b[49m\u001b[43mspacing\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;241;43m1\u001b[39;49m\u001b[43m,\u001b[49m\n\u001b[0;32m      5\u001b[0m \u001b[43m    \u001b[49m\u001b[43mdt\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;241;43m0.1\u001b[39;49m\u001b[43m,\u001b[49m\n\u001b[0;32m      6\u001b[0m \u001b[43m    \u001b[49m\u001b[43mspeed_mean\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;241;43m0.12\u001b[39;49m\u001b[43m,\u001b[49m\n\u001b[0;32m      7\u001b[0m \u001b[43m    \u001b[49m\u001b[43mnum_place_cells\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;241;43m50\u001b[39;49m\u001b[43m,\u001b[49m\n\u001b[0;32m      8\u001b[0m \u001b[43m    \u001b[49m\u001b[43mnum_grid_cells\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;241;43m0\u001b[39;49m\u001b[43m,\u001b[49m\n\u001b[0;32m      9\u001b[0m \u001b[43m    \u001b[49m\u001b[43mplace_width_mean\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;241;43m20\u001b[39;49m\u001b[43m,\u001b[49m\n\u001b[0;32m     10\u001b[0m \u001b[43m    \u001b[49m\u001b[43mplace_width_std\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;241;43m10.0\u001b[39;49m\u001b[43m,\u001b[49m\n\u001b[0;32m     11\u001b[0m \u001b[43m    \u001b[49m\u001b[43mnum_grid_modules\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;241;43m4\u001b[39;49m\u001b[43m,\u001b[49m\n\u001b[0;32m     12\u001b[0m \u001b[43m    \u001b[49m\u001b[43mgrid_expansion\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;241;43m1.3\u001b[39;49m\u001b[43m,\u001b[49m\n\u001b[0;32m     13\u001b[0m \u001b[43m    \u001b[49m\u001b[43mbase_grid_spacing\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;241;43m39.8\u001b[39;49m\u001b[43m,\u001b[49m\n\u001b[0;32m     14\u001b[0m \u001b[43m    \u001b[49m\u001b[43mbase_grid_width\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;241;43m27.4\u001b[39;49m\u001b[43m,\u001b[49m\n\u001b[0;32m     15\u001b[0m \u001b[43m    \u001b[49m\u001b[43mg_noise_amp\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;241;43m1\u001b[39;49m\u001b[43m \u001b[49m\u001b[38;5;241;43m/\u001b[39;49m\u001b[43m \u001b[49m\u001b[38;5;241;43m100000\u001b[39;49m\n\u001b[0;32m     16\u001b[0m \u001b[43m)\u001b[49m\n",
      "File \u001b[1;32mc:\\Users\\Andrew\\Documents\\GitHub\\vrAnalysis\\vrAnalysis\\simulations\\simulator.py:44\u001b[0m, in \u001b[0;36mSimulator.__init__\u001b[1;34m(self, box_length, spacing, dt, speed_mean, num_place_cells, num_grid_cells, place_width_mean, place_width_std, num_grid_modules, grid_expansion, base_grid_spacing, base_grid_width, g_noise_amp)\u001b[0m\n\u001b[0;32m     42\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mset_environment()\n\u001b[0;32m     43\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mcreate_agent()\n\u001b[1;32m---> 44\u001b[0m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mcreate_cell_library\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[1;32mc:\\Users\\Andrew\\Documents\\GitHub\\vrAnalysis\\vrAnalysis\\simulations\\simulator.py:65\u001b[0m, in \u001b[0;36mSimulator.create_cell_library\u001b[1;34m(self)\u001b[0m\n\u001b[0;32m     60\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mplace_library \u001b[38;5;241m=\u001b[39m np\u001b[38;5;241m.\u001b[39mstack(\n\u001b[0;32m     61\u001b[0m     [support\u001b[38;5;241m.\u001b[39mget_place_map(pxc, pyc, \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mxgrid, \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mygrid, pw) \u001b[38;5;28;01mfor\u001b[39;00m pxc, pyc, pw \u001b[38;5;129;01min\u001b[39;00m \u001b[38;5;28mzip\u001b[39m(\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mplace_xc, \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mplace_yc, \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mplace_width)]\n\u001b[0;32m     62\u001b[0m )\n\u001b[0;32m     64\u001b[0m \u001b[38;5;66;03m# Create grid cell tuning curves\u001b[39;00m\n\u001b[1;32m---> 65\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mgrid_xc, \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mgrid_yc \u001b[38;5;241m=\u001b[39m helpers\u001b[38;5;241m.\u001b[39mnamed_transpose([support\u001b[38;5;241m.\u001b[39mrand_centroid(\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mbox_length) \u001b[38;5;28;01mfor\u001b[39;00m _ \u001b[38;5;129;01min\u001b[39;00m \u001b[38;5;28mrange\u001b[39m(\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mnum_grid_cells)])\n\u001b[0;32m     66\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mgrid_spacing \u001b[38;5;241m=\u001b[39m np\u001b[38;5;241m.\u001b[39marray(\n\u001b[0;32m     67\u001b[0m     [\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mbase_grid_spacing \u001b[38;5;241m*\u001b[39m (\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mgrid_expansion \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39m np\u001b[38;5;241m.\u001b[39mrandom\u001b[38;5;241m.\u001b[39mrandint(\u001b[38;5;241m0\u001b[39m, \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mnum_grid_modules)) \u001b[38;5;28;01mfor\u001b[39;00m _ \u001b[38;5;129;01min\u001b[39;00m \u001b[38;5;28mrange\u001b[39m(\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mnum_grid_cells)]\n\u001b[0;32m     68\u001b[0m )\n\u001b[0;32m     69\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mgrid_angle \u001b[38;5;241m=\u001b[39m np\u001b[38;5;241m.\u001b[39marray([np\u001b[38;5;241m.\u001b[39mpi \u001b[38;5;241m/\u001b[39m \u001b[38;5;241m3\u001b[39m \u001b[38;5;241m*\u001b[39m np\u001b[38;5;241m.\u001b[39mrandom\u001b[38;5;241m.\u001b[39mrandom() \u001b[38;5;28;01mfor\u001b[39;00m _ \u001b[38;5;129;01min\u001b[39;00m \u001b[38;5;28mrange\u001b[39m(\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mnum_grid_cells)])\n",
      "\u001b[1;31mValueError\u001b[0m: not enough values to unpack (expected 2, got 0)"
     ]
    }
   ],
   "source": [
    "box_length = 80\n",
    "sim = Simulator(\n",
    "    box_length,\n",
    "    spacing=1,\n",
    "    dt=0.1,\n",
    "    speed_mean=0.12,\n",
    "    num_place_cells=50,\n",
    "    num_grid_cells=0,\n",
    "    place_width_mean=20,\n",
    "    place_width_std=10.0,\n",
    "    num_grid_modules=4,\n",
    "    grid_expansion=1.3,\n",
    "    base_grid_spacing=39.8,\n",
    "    base_grid_width=27.4,\n",
    "    g_noise_amp=1 / 100000\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Show Examples\n",
    "numExamplesEach = 6\n",
    "figdim = 2\n",
    "f, ax = plt.subplots(2,numExamplesEach,figsize=(figdim*numExamplesEach,figdim*2), layout='constrained')\n",
    "for n in range(numExamplesEach):\n",
    "    ax[0,n].imshow(sim.place_library[np.random.randint(sim.num_place_cells)])\n",
    "    ax[1,n].imshow(sim.grid_library[np.random.randint(sim.num_grid_cells)])\n",
    "    ax[0,n].axis('off')\n",
    "    ax[1,n].axis('off')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 81,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.hist(np.reshape(sim.grid_library, -1))\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 78,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Do a trajectory\n",
    "t, pos, posidx = sim.run_simulation(10000)\n",
    "\n",
    "# Return place and grid cell activity along trajectory, with noise if requested\n",
    "place_signal = sim.place_library[:,posidx[:,0], posidx[:,1]].T\n",
    "grid_signal = sim.grid_library[:,posidx[:,0], posidx[:,1]].T\n",
    "hippo_signal = np.concatenate((place_signal,grid_signal),axis=1)\n",
    "\n",
    "# Add Noise\n",
    "gNoiseAmp = 1/5\n",
    "gNoisePlace = gNoiseAmp * np.std(place_signal,axis=0)\n",
    "gNoiseHippo = gNoiseAmp * np.std(hippo_signal,axis=0)\n",
    "place_activity = place_signal + np.random.normal(0,gNoisePlace,place_signal.shape)\n",
    "hippo_activity = hippo_signal + np.random.normal(0,gNoiseHippo,hippo_signal.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 79,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Plot Trajectory and plot examples of place cell and grid cell activity overlaid on trajectory\n",
    "numExamplesEach = 2\n",
    "pexidx = np.random.randint(0,sim.num_place_cells,numExamplesEach)\n",
    "gexidx = np.random.randint(0,sim.num_grid_cells,numExamplesEach)\n",
    "\n",
    "figdim = 3\n",
    "f, ax = plt.subplots(1,numExamplesEach*2+1,figsize=((numExamplesEach*2+1)*figdim,figdim))\n",
    "# Ag.plot_trajectory(ax=ax[0])\n",
    "plt.set_cmap('cool')\n",
    "ax[0].scatter(pos[:,0],pos[:,1],s=15,c=range(pos.shape[0]),alpha=0.3)\n",
    "plt.set_cmap('jet')\n",
    "for n in range(numExamplesEach):\n",
    "    ax[n+1].scatter(pos[:,0],pos[:,1],s=5,c=place_activity[:,pexidx[n]])\n",
    "    ax[n+1+numExamplesEach].scatter(pos[:,0],pos[:,1],s=5,c=hippo_activity[:,sim.num_place_cells+gexidx[n]])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "vrAnalysis",
   "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.9.19"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}