experimental-data-ci.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "2d3f0113",
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "3fd2102b-9d9c-447e-a6bb-6bbb4527836c",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "plt.style.use(\"bioAI.mplstyle\")\n",
    "import scipy\n",
    "import astropy\n",
    "import tqdm"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "723ca03d",
   "metadata": {},
   "source": [
    "## Load data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "22a52e79",
   "metadata": {},
   "outputs": [],
   "source": [
    "def data_scope(arr, t_arr, t_start, t_stop):\n",
    "    idxs = np.where((t_arr >= t_start) & (t_arr <= t_stop))[0]\n",
    "    return arr[idxs]\n",
    "\n",
    "def spike_scope(spikes, t_arr):\n",
    "    hists = []\n",
    "    for spike_train in spikes.item().values():\n",
    "        hist, _ = np.histogram(spike_train, bins=np.append(t_arr, t_arr[-1] + t_arr[1] - t_arr[0]))\n",
    "        hists.append(hist)\n",
    "    hists = np.array(hists)\n",
    "    return hists"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "f2fb324b",
   "metadata": {},
   "outputs": [],
   "source": [
    "data_path = '/Users/vemundss/Desktop/Toroidal_topology_grid_cell_data/'\n",
    "# data_path = './Toroidal_topology_grid_cell_data/'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0826beca",
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(data_path + 'rat_q_sessions.txt') as file:\n",
    "#with open(data_path + 'rat_r_day1_sessions.txt') as file:\n",
    "#with open(data_path + 'rat_r_day2_sessions.txt') as file:\n",
    "#with open(data_path + 'rat_s_sessions.txt') as file:\n",
    "    sessions = file.readlines()\n",
    "sessions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "97ac41e8",
   "metadata": {},
   "outputs": [],
   "source": [
    "data = {}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "49f0b25d-e56d-45d6-9023-c36235fcd225",
   "metadata": {},
   "outputs": [],
   "source": [
    "f = np.load(data_path + 'rat_q_grid_modules_1_2.npz', allow_pickle=True)\n",
    "x = f['x']\n",
    "y = f['y']\n",
    "t = f['t']\n",
    "spikes1 = f['spikes_mod1']\n",
    "spikes2 = f['spikes_mod2']\n",
    "\n",
    "t_start, t_stop = 27826, 31223\n",
    "x = data_scope(f['x'], f['t'], t_start, t_stop)\n",
    "y = data_scope(f['y'], f['t'], t_start, t_stop)\n",
    "t = data_scope(f['t'], f['t'], t_start, t_stop)\n",
    "spikes1 = spike_scope(f['spikes_mod1'], t)\n",
    "spikes2 = spike_scope(f['spikes_mod2'], t)\n",
    "\n",
    "data['rat_q'] = {'x': x, 'y': y, 't': t, 'modules': {1: spikes1, 2: spikes2}}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "63d4e4ee",
   "metadata": {},
   "outputs": [],
   "source": [
    "f = np.load(data_path + 'rat_r_day1_grid_modules_1_2_3.npz', allow_pickle=True)\n",
    "x = f['x']\n",
    "y = f['y']\n",
    "t = f['t']\n",
    "spikes1 = f['spikes_mod1']\n",
    "spikes2 = f['spikes_mod2']\n",
    "spikes3 = f['spikes_mod3']\n",
    "\n",
    "t_start, t_stop = 7457, 14778\n",
    "x = data_scope(f['x'], f['t'], t_start, t_stop)\n",
    "y = data_scope(f['y'], f['t'], t_start, t_stop)\n",
    "t = data_scope(f['t'], f['t'], t_start, t_stop)\n",
    "spikes1 = spike_scope(f['spikes_mod1'], t)\n",
    "spikes2 = spike_scope(f['spikes_mod2'], t)\n",
    "spikes3 = spike_scope(f['spikes_mod3'], t)\n",
    "\n",
    "data['rat_r_day1'] = {'x': x, 'y': y, 't': t, 'modules': {1: spikes1, 2: spikes2, 3: spikes3}}\n",
    "\n",
    "# can add second slice (valid times) 14890, 16045 here"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "aa024643",
   "metadata": {},
   "outputs": [],
   "source": [
    "f = np.load(data_path + 'rat_r_day2_grid_modules_1_2_3.npz', allow_pickle=True)\n",
    "x = f['x']\n",
    "y = f['y']\n",
    "t = f['t']\n",
    "spikes1 = f['spikes_mod1']\n",
    "spikes2 = f['spikes_mod2']\n",
    "spikes3 = f['spikes_mod3']\n",
    "\n",
    "t_start, t_stop = 10617, 13004\n",
    "x = data_scope(f['x'], f['t'], t_start, t_stop)\n",
    "y = data_scope(f['y'], f['t'], t_start, t_stop)\n",
    "t = data_scope(f['t'], f['t'], t_start, t_stop)\n",
    "spikes1 = spike_scope(f['spikes_mod1'], t)\n",
    "spikes2 = spike_scope(f['spikes_mod2'], t)\n",
    "spikes3 = spike_scope(f['spikes_mod3'], t)\n",
    "\n",
    "data['rat_r_day2'] = {'x': x, 'y': y, 't': t, 'modules': {1: spikes1, 2: spikes2, 3: spikes3}}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "0e440953",
   "metadata": {},
   "outputs": [],
   "source": [
    "f = np.load(data_path + 'rat_s_grid_modules_1.npz', allow_pickle=True)\n",
    "x = f['x']\n",
    "y = f['y']\n",
    "t = f['t']\n",
    "spikes1 = f['spikes_mod1']\n",
    "\n",
    "t_start, t_stop = 9939, 12363\n",
    "x = data_scope(f['x'], f['t'], t_start, t_stop)\n",
    "y = data_scope(f['y'], f['t'], t_start, t_stop)\n",
    "t = data_scope(f['t'], f['t'], t_start, t_stop)\n",
    "spikes1 = spike_scope(f['spikes_mod1'], t)\n",
    "\n",
    "data['rat_s'] = {'x': x, 'y': y, 't': t, 'modules': {1: spikes1}}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e8c49cd0",
   "metadata": {},
   "source": [
    "## Create ratemaps"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "abf1f047",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, axs = plt.subplots(ncols=4, figsize=(20, 5))\n",
    "for i, (rat, rat_data) in enumerate(data.items()):\n",
    "    axs[i].scatter(rat_data['x'], rat_data['y'], s=0.01)\n",
    "    axs[i].set_title(rat)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2d385151",
   "metadata": {},
   "outputs": [],
   "source": [
    "res = 32\n",
    "all_ratemaps = []\n",
    "for rat_i, data_i in data.items():\n",
    "    x, y = data_i['x'], data_i['y']\n",
    "    for module_j, spikes_j in data_i['modules'].items():\n",
    "        ratemaps_ij = scipy.stats.binned_statistic_2d(x, y, spikes_j, bins=res).statistic\n",
    "        all_ratemaps.append(ratemaps_ij)\n",
    "        print(len(ratemaps_ij))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e63566b7",
   "metadata": {},
   "source": [
    "## Calculate grid statistics"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b911c468",
   "metadata": {},
   "outputs": [],
   "source": [
    "import spatial_maps as sm\n",
    "grid_threshold = 0.4 # <----------------- threshold for gridness score\n",
    "ratemaps = all_ratemaps[2] # <----------------- select grid module\n",
    "ratemaps = np.array([astropy.convolution.convolve(ratemaps[i], astropy.convolution.Gaussian2DKernel(2), boundary='wrap') for i in range(ratemaps.shape[0])])\n",
    "\n",
    "center_peaks = np.array([sm.find_peaks(ratemaps[i])[0] for i in range(ratemaps.shape[0])])\n",
    "grid_scores = np.array([sm.gridness(ratemap) for ratemap in ratemaps])\n",
    "acorrs = np.array([sm.autocorrelation(ratemap) for ratemap in ratemaps])\n",
    "peaks = [sm.find_peaks(acorr) for acorr in acorrs]\n",
    "thetas = np.array([sm.spacing_and_orientation(peak, acorrs.shape[-1])[1] for peak in peaks])\n",
    "spacings = np.array([sm.spacing_and_orientation(peak, acorrs.shape[-1])[0] for peak in peaks])\n",
    "# get the most common spacing and theta\n",
    "bins=25\n",
    "hist, bin_edges = np.histogram(spacings, bins=bins)\n",
    "spacing = bin_edges[np.argmax(hist)] + (bin_edges[1] - bin_edges[0])/2\n",
    "hist, bin_edges = np.histogram(thetas, bins=bins)\n",
    "theta = bin_edges[np.argmax(hist)] + (bin_edges[1] - bin_edges[0])/2\n",
    "\n",
    "from methods import Hexagon, HexagonalGCs\n",
    "f = 2/(np.sqrt(3)*spacing)\n",
    "#radius = spacing/np.sqrt(3) # spacing is peak-to-peak distance\n",
    "radius = 2/(3*f)\n",
    "orientation = theta*180/np.pi - 15 # theta is angle from x-axis\n",
    "hexagon = HexagonalGCs(f=f, init_rot=orientation)\n",
    "\n",
    "# mask out non-grid cells\n",
    "mask = grid_scores > grid_threshold\n",
    "mask = hexagon.unit_cell.is_in_hexagon(center_peaks-res/2) & mask\n",
    "print(\"n=\",(mask).sum(), \", N=\",len(mask))\n",
    "\n",
    "# remove stats based on phases outside hexagon and grid scores below threshold\n",
    "center_peaks = center_peaks[mask]\n",
    "ratemaps = ratemaps[mask]\n",
    "thetas = thetas[mask]\n",
    "spacings = spacings[mask]\n",
    "grid_scores = grid_scores[mask]\n",
    "\n",
    "# set phases to hexagon\n",
    "hexagon.set_phases(center_peaks-res/2) # WRAPPY?\n",
    "\n",
    "# create random hexagon baseline\n",
    "random_hexagon = HexagonalGCs(f=f, init_rot=orientation)\n",
    "random_hexagon.set_phases(random_hexagon.unit_cell.wrap(np.around(random_hexagon.unit_cell.sample(len(center_peaks)))))\n",
    "print(f\"Radius: {radius}, Orientation: {orientation}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d67d9a6a",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, axs = plt.subplots(ncols=4, figsize=(8, 2))\n",
    "axs[0].scatter(*(center_peaks-res/2).T,alpha=.3)\n",
    "axs[0].set_xlim(-res/2, res/2)\n",
    "axs[0].set_ylim(-res/2, res/2)\n",
    "hexagon.unit_cell.plot(fig, axs[0], color='black')#, center=np.zeros(2)+res/2-1)\n",
    "#ax.scatter(*(center_peaks-res/2).T, color='green', alpha=.2)\n",
    "axs[0].set_title('Phases')\n",
    "axs[0].set_xticks([])\n",
    "axs[0].set_yticks([])\n",
    "axs[0].set_aspect('equal')\n",
    "axs[0].set_xlabel('box size')\n",
    "axs[0].set_ylabel('box size')\n",
    "axs[1].hist(spacings, bins=bins)\n",
    "axs[1].axvline(spacing, color='r')\n",
    "axs[1].set_title('Spacing')\n",
    "axs[2].hist(thetas, bins=bins)\n",
    "axs[2].axvline(theta, color='r')\n",
    "axs[2].set_title('Orientation')\n",
    "axs[2].set_xlabel('Angle [rad]')\n",
    "axs[3].hist(grid_scores, bins=bins)\n",
    "axs[3].set_title('Gridness')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2f8f22d8",
   "metadata": {},
   "outputs": [],
   "source": [
    "from plotting_functions import multiimshow\n",
    "fig, axs = multiimshow(ratemaps[np.argsort(grid_scores)[::-1]].transpose(0,2,1), \n",
    "                       origin='lower', figsize=(8,8), normalize=True, add_colorbar=True);\n",
    "for i in range(len(center_peaks)):\n",
    "    axs[i].scatter(*center_peaks[np.argsort(grid_scores)[::-1]][i], color='red', s=3)\n",
    "    #hexagon.unit_cell.plot(fig, axs[i], center=center_peaks[np.argsort(grid_scores)[::-1]][i], color='black')\n",
    "    hexagon.unit_cell.plot(fig, axs[i], center=np.ones(2)*res/2, color='black')\n",
    "\n",
    "fig.savefig('./results/experimental-phases.png', dpi=300, bbox_inches='tight')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c5fbf742",
   "metadata": {},
   "source": [
    "## KDE, Ripley and Grid score of phases"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "648f2901",
   "metadata": {},
   "outputs": [],
   "source": [
    "bw_method = np.linspace(0, 0.2, 21)[1:]#hexagon.unit_cell.radius/2/10, 21)[1:]\n",
    "fig, axs = plt.subplots(nrows=2, ncols=len(bw_method)//2, figsize=(len(bw_method)//2*3, 3*2))\n",
    "for i, r in enumerate(bw_method[::2]):\n",
    "    kde, mesh, _, _ = hexagon.phase_kde(bw_method=r, res=64)\n",
    "    axs[0,i].scatter(*mesh.T, c=kde, s=10)\n",
    "    axs[0,i].set_title(f'{r:.2f}')\n",
    "\n",
    "\n",
    "for i, r in enumerate(bw_method[::2]):\n",
    "    kde, mesh, _, _ = random_hexagon.phase_kde(bw_method=r, res=64)\n",
    "    axs[1,i].scatter(*mesh.T, c=kde, s=10)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d96ee804",
   "metadata": {},
   "outputs": [],
   "source": [
    "from models import JacobianCI\n",
    "ci_model = JacobianCI(ncells=7,lr=1e-3, seed=3, f=2/(3*radius*1.4), init_rot=orientation)\n",
    "ci_model_small = JacobianCI(ncells=7,lr=1e-3, seed=3, f=f, init_rot=orientation)\n",
    "loss_history = ci_model.train_model(ntrain_steps=5000)\n",
    "_ = ci_model_small.train_model(ntrain_steps=5000)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0189a29f",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, axs = plt.subplots(ncols=2, figsize=(3, 1.5))\n",
    "axs[0].semilogy(loss_history)\n",
    "ci_model.unit_cell.plot(fig=fig, ax=axs[1], color='black')\n",
    "axs[1].scatter(*ci_model.phases.detach().numpy().T,s=1)\n",
    "axs[1].axis('equal')\n",
    "#ax.axis('off')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5e14fcef",
   "metadata": {},
   "outputs": [],
   "source": [
    "from utils import find_peaks\n",
    "import torch\n",
    "\n",
    "rel_box_size = radius * res / radius\n",
    "\n",
    "x = np.linspace(-rel_box_size/2,rel_box_size/2,res)\n",
    "y = np.linspace(-rel_box_size/2,rel_box_size/2,res)\n",
    "X,Y = np.meshgrid(x,y)\n",
    "square_mesh = np.stack([X,Y],axis=-1)\n",
    "square_mesh.shape\n",
    "\n",
    "# calculate ratemaps with given resolution and find peaks in the sqaure mesh\n",
    "ci_ratemaps = ci_model(torch.tensor(square_mesh.reshape(-1,2),\n",
    "                                    dtype=torch.float32)).detach().numpy().reshape(res,res,-1)\n",
    "ci_center_peaks = np.array([find_peaks(ci_ratemaps[:,:,i])[0] for i in range(ci_ratemaps.shape[-1])])\n",
    "# convert peaks to unit cell coordinates again\n",
    "ci_center_peaks = ci_center_peaks / res * rel_box_size - rel_box_size/2\n",
    "\n",
    "plt.scatter(*ci_center_peaks.T)\n",
    "plt.scatter(*ci_model.phases.detach().numpy().T, color='red')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d55a8ceb",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, axs = multiimshow(ci_ratemaps.transpose(2,1,0), origin='lower', figsize=(5,5), normalize=False, \n",
    "                       extent=(-rel_box_size/2,rel_box_size/2,-rel_box_size/2,rel_box_size/2));\n",
    "for i in range(len(ci_center_peaks)):\n",
    "    axs[i].scatter(*ci_center_peaks[i], color='red', s=3)\n",
    "    ci_model.unit_cell.plot(fig, axs[i], center=np.zeros(2), color='black')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f0b8fa18",
   "metadata": {},
   "outputs": [],
   "source": [
    "import copy\n",
    "#ci_center_peaks = ci_center_peaks[:,None] + np.random.normal(0, 10, size=(ratemaps.shape[0]//7,2))[None]\n",
    "#ci_center_peaks = ci_model.phases.detach().numpy()[:,None] +  ci_model.unit_cell.sample(ratemaps.shape[0]//7)[None]\n",
    "ci_center_peaks = ci_model.phases.detach().numpy()[:,None] + np.random.normal(0, 1.4, size=(ratemaps.shape[0]//7,2))[None]\n",
    "ci_center_peaks = ci_center_peaks.reshape(-1,2)\n",
    "ci_center_peaks = np.around(ci_model.unit_cell.wrap(ci_center_peaks))\n",
    "#ci_center_peaks = ci_model.unit_cell.wrap(ci_center_peaks)\n",
    "ci_center_peaks.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f79bdcda",
   "metadata": {},
   "outputs": [],
   "source": [
    "tmp_phase = ci_model_small.phases.detach().numpy()[:,None] + np.random.normal(0, 1.4, size=(ratemaps.shape[0]//7,2))[None]\n",
    "tmp_phase = tmp_phase.reshape(-1,2)\n",
    "tmp_phase = ci_model.unit_cell.wrap(tmp_phase)\n",
    "ci_model_small.set_phases(tmp_phase)\n",
    "ci_model_small.set_scale()\n",
    "ci_cpy_ratemaps = ci_model_small(torch.tensor(square_mesh.reshape(-1,2),\n",
    "                                    dtype=torch.float32)).detach().numpy().reshape(res,res,-1)\n",
    "\n",
    "random_ratemaps = random_hexagon(torch.tensor(square_mesh.reshape(-1,2),\n",
    "                                    dtype=torch.float32)).detach().numpy().reshape(res,res,-1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a0d75b79",
   "metadata": {},
   "outputs": [],
   "source": [
    "kde_data, mesh, _, _ = hexagon.phase_kde(bw_method=0.09, res=64)#,wrap=False,double_extension=False)\n",
    "kde_random, mesh, _, _ = random_hexagon.phase_kde(bw_method=0.09, res=64)#,wrap=False,double_extension=False)\n",
    "kde_ci, mesh_ci, _, _ = ci_model.phase_kde(bw_method=0.09, res=64)\n",
    "min_kde = min(kde_data.min(), kde_random.min(), kde_ci.min())\n",
    "max_kde = max(kde_data.max(), kde_random.max(), kde_ci.max())\n",
    "min_kde, max_kde"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "20b47406",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(figsize=(1,1))\n",
    "phase_grid_score = hexagon.grid_score(bw_method=0.11)\n",
    "ax.scatter(*mesh.T, c=kde_data, s=10, vmin=min_kde, vmax=max_kde)\n",
    "ax.scatter(*(center_peaks-res/2).T,alpha=0.5,color='grey')\n",
    "ax.set_title('data',fontsize='small')\n",
    "ax.axis('off')\n",
    "ax.axis('equal')\n",
    "# add colorbar\n",
    "#cbar = plt.colorbar(ax.collections[0], ax=ax)#, orientation='horizontal')\n",
    "\n",
    "fig.tight_layout(pad=0)\n",
    "fig.savefig('./results/hexagon_phase_kde.pdf')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5a39eb79",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(figsize=(1,1))\n",
    "ax.scatter(*mesh.T, c=kde_random, s=10, vmin=min_kde, vmax=max_kde)\n",
    "ax.scatter(*random_hexagon.phases.detach().numpy().T,alpha=0.5,color='grey')\n",
    "ax.set_title('random',fontsize='small')\n",
    "ax.axis('off')\n",
    "ax.axis('equal')\n",
    "\n",
    "fig.tight_layout(pad=0)\n",
    "fig.savefig('./results/random_hexagon_phase_kde.pdf')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1ea958cc",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig,ax = plt.subplots(figsize=(1*1.4,1*1.4))\n",
    "fig, ax = hexagon.unit_cell.plot(color='grey', fig=fig, ax=ax)\n",
    "ax.scatter(*mesh_ci.T, c=kde_ci, s=10, vmin=min_kde, vmax=max_kde)\n",
    "ax.scatter(*ci_center_peaks.T)\n",
    "ax.axis('equal')\n",
    "ax.axis('off')\n",
    "ax.set_title('CI+noise',fontsize='small')\n",
    "\n",
    "fig.tight_layout(pad=0)\n",
    "fig.savefig('./results/ci_phase_kde.pdf')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e3bac8a0",
   "metadata": {},
   "outputs": [],
   "source": [
    "radii = np.linspace(0, hexagon.unit_cell.radius-1e-6, 32)[1:]\n",
    "random_ripleys = []\n",
    "ci_ripleys = []\n",
    "for _ in tqdm.trange(100):\n",
    "    random_hexagon.set_phases(np.around(random_hexagon.unit_cell.wrap(random_hexagon.unit_cell.sample(len(center_peaks)))))\n",
    "    random_ripleys_i = [random_hexagon.ripleys(r) for r in radii]\n",
    "    random_ripleys.append(random_ripleys_i)\n",
    "    ci_center_peaks = ci_model.phases.detach().numpy()[:,None] + np.random.normal(0, 1.4, size=(ratemaps.shape[0]//7,2))[None]\n",
    "    ci_center_peaks = ci_center_peaks.reshape(-1,2)\n",
    "    ci_center_peaks = np.around(ci_model.unit_cell.wrap(ci_center_peaks))\n",
    "    ci_ripleys_i = [ci_model.ripleys(r,phases=ci_center_peaks) for r in radii]\n",
    "    ci_ripleys.append(ci_ripleys_i)\n",
    "random_ripleys = np.array(random_ripleys)\n",
    "ci_ripleys = np.array(ci_ripleys)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "0046b0e7",
   "metadata": {},
   "outputs": [],
   "source": [
    "ripleys = [hexagon.ripleys(r) for r in radii]\n",
    "#ci_ripleys = [ci_model.ripleys(r,phases=ci_center_peaks) for r in radii]\n",
    "ripleys = np.array(ripleys)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cb6ec654",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig=plt.figure(figsize = (1.8,1.5))\n",
    "\n",
    "# random baseline\n",
    "plt.plot(radii, np.mean(random_ripleys, axis=0), label = \"random\", alpha=0.5)\n",
    "mu_random = np.mean(random_ripleys, axis=0)\n",
    "sigma_random = np.std(random_ripleys, axis=0)\n",
    "plt.fill_between(radii, mu_random - 2*sigma_random, mu_random + 2*sigma_random, alpha=0.2)\n",
    "# hexagonal baseline\n",
    "#plt.plot(radii, ci_ripleys, label = \"Optimised\", alpha=0.5)\n",
    "# hexagonal baseline\n",
    "plt.plot(radii, np.mean(ci_ripleys, axis=0), label = \"CI+noise\", alpha=0.5)\n",
    "mu_optimised = np.mean(ci_ripleys, axis=0)\n",
    "sigma_optimised = np.std(ci_ripleys, axis=0)\n",
    "plt.fill_between(radii, mu_optimised - 2*sigma_optimised, mu_optimised + 2*sigma_optimised, alpha=0.2)\n",
    "# data\n",
    "plt.plot(radii, ripleys, label = \"data\",alpha=0.5)\n",
    "# plot star where significant\n",
    "plt.scatter(radii[ripleys > mu_random + 2*sigma_random], \n",
    "            ripleys[ripleys > mu_random + 2*sigma_random], color='red', marker='*', s=5)\n",
    "plt.scatter(radii[ripleys < mu_random -2*sigma_random], \n",
    "            ripleys[ripleys < mu_random-2*sigma_random], color='red', marker='*', s=5)\n",
    "\n",
    "#plt.axvline(2/(ci_model.f*np.sqrt(21)), color='grey', linestyle='--', label=r'$2/f\\sqrt{21}$')\n",
    "\n",
    "plt.xlabel('radius', fontsize='small')\n",
    "plt.ylabel('Ripley\\'s H', fontsize='small')\n",
    "plt.legend(frameon = False, fontsize='small')\n",
    "\n",
    "fig.tight_layout(pad=0)\n",
    "fig.savefig('./results/ripleys.pdf')#, dpi=300)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fc15aab4",
   "metadata": {},
   "outputs": [],
   "source": [
    "grid_scores = [hexagon.grid_score(bw_method=r) for r in bw_method]\n",
    "plt.plot(bw_method, grid_scores)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a44d03f7",
   "metadata": {},
   "source": [
    "**We want to** find the distance between two neighbouring small hexagons centers, i.e. $2a_p$. This gives us an expectation for when the ripleys would change from below random expectation to above random expectation of phases in a given ball radius, in the scenario where the phases where hexagonally arranged.\n",
    "\n",
    "We have the frequency of the generating pattern is $f=\\frac{2}{3R}$, where $R$ is the radius of the corresponding unit cell. The apothem $A_p$ of the unit cell can be calculated from the special triangle (30-60-90 triangle) with side lengths (from long to short) as $R, A_p, R/2$. This gives\n",
    "\n",
    "\\begin{align}\n",
    "R^2 &= A_p^2 + \\left(\\frac{R}{2}\\right)^2 \\\\\n",
    "A_p  &= \\frac{\\sqrt{3}}{2}R\n",
    "\\end{align}\n",
    "\n",
    "We further have that the unit cell can have a good cover with 7 regular hexagons. The area of each of these subhexagons are $a = \\frac{1}{7}A$ of the size of the area of the unit cell hexagon. The area of a hexagon is in general given by $A = \\frac{3\\sqrt(3)R^2}{2}$. So get we\n",
    "\n",
    "\\begin{align}\n",
    "a &= \\frac{1}{7}A \\\\\n",
    "\\frac{3\\sqrt(3)r^2}{2} &= \\frac{1}{7} \\cdot \\frac{3\\sqrt(3)R^2}{2} \\\\\n",
    "r &= \\frac{1}{\\sqrt{7}}R\n",
    "\\end{align}\n",
    "\n",
    "Since $A_p$ is proportional to $R$, we also have $a_p = \\frac{1}{\\sqrt{7}} A_p$ \n",
    "\n",
    "This gives $2a_p$, in terms of $f$\n",
    "\n",
    "\\begin{align}\n",
    "2a_p = 2\\cdot\\frac{1}{\\sqrt{7}} A_p = \\frac{2}{\\sqrt{7}}\\cdot \\frac{\\sqrt{3}}{2}R = \\frac{\\sqrt{3}}{\\sqrt{7}}\\cdot \\frac{2}{3f} = \\frac{2}{f\\sqrt{21}}\n",
    "\\end{align}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c09accce",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.imshow(random_ratemaps[:,:,0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6b68df93",
   "metadata": {},
   "outputs": [],
   "source": [
    "np.min(ci_ratemaps), np.max(ci_ratemaps), np.min(ratemaps), np.max(ratemaps)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 71,
   "id": "52727e61",
   "metadata": {},
   "outputs": [],
   "source": [
    "ncells = len(ratemaps)\n",
    "dist = np.linalg.norm(ratemaps.reshape(ncells,res**2)[:,None] - \n",
    "                      ratemaps.reshape(ncells,res**2)[...,None],axis=0)\n",
    "\n",
    "\n",
    "ci_cpy_dist = np.linalg.norm(ci_cpy_ratemaps.reshape(res**2,-1).T[:,None] -\n",
    "                            ci_cpy_ratemaps.reshape(res**2,-1).T[...,None],axis=0)\n",
    "\n",
    "random_dist = np.linalg.norm(random_ratemaps.reshape(res**2,-1).T[:,None] -\n",
    "                            random_ratemaps.reshape(res**2,-1).T[...,None],axis=0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "678be289",
   "metadata": {},
   "outputs": [],
   "source": [
    "row = 10\n",
    "col = 16\n",
    "a = row*32 + col\n",
    "fig,ax = plt.subplots(figsize = (1.25,1.25))\n",
    "im = ax.imshow(ci_cpy_dist[a].reshape(res,res))\n",
    "plt.plot(col,row, \"rx\", markersize = 5)\n",
    "# add colorbar that fits the image and doesnt stretch beyond\n",
    "cbar = fig.colorbar(im, ax=ax, fraction=0.046, pad=0.04)\n",
    "ax.axis('off')\n",
    "ax.axis('equal')\n",
    "ax.set_title('neural distance',fontsize='small')\n",
    "plt.tight_layout(pad=0)\n",
    "fig.savefig('./results/neural_distance.pdf')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "99d98ee3",
   "metadata": {},
   "outputs": [],
   "source": [
    "grads = np.array([np.gradient(im) for im in random_ratemaps.transpose(2,0,1)])\n",
    "grads = np.transpose(grads, (0, -2, -1, 1))\n",
    "grads = grads[:, 1:-1, 1:-1] # remove 1st order boundary term\n",
    "\n",
    "fig, axs = plt.subplots(1, 3, figsize = (3.9,1.3))\n",
    "\n",
    "m11 = np.sum(grads[...,0]**2, axis = 0)\n",
    "m22 = np.sum(grads[...,1]**2, axis = 0)\n",
    "m12 = np.sum(grads[...,0]*grads[...,1], axis = 0)\n",
    "\n",
    "vmin = min(m11.min(), m22.min(), m12.min())\n",
    "vmax = max(m11.max(), m22.max(), m12.max())\n",
    "\n",
    "im1 = axs[0].imshow(m11, vmin=vmin, vmax=vmax)\n",
    "im2 = axs[1].imshow(m12, vmin=vmin, vmax=vmax)\n",
    "im3 = axs[2].imshow(m22, vmin=vmin, vmax=vmax)\n",
    "[ax.axis(\"off\") for ax in axs]\n",
    "[ax.axis(\"equal\") for ax in axs]\n",
    "[ax.set_title(title,fontsize='small') for ax, title in zip(axs, [\"$G_{xx}$\", \"$G_{xy}$\", \"$G_{xz}$\"])]\n",
    "[fig.colorbar(im, ax = ax, fraction=0.046, pad=0.04) for ax, im in zip(axs, [im1, im2, im3])]\n",
    "plt.tight_layout(pad=0)\n",
    "fig.savefig('./results/metric_tensor_components.pdf')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "24a8d08c",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig,ax = plt.subplots(figsize = (0.85,1.25))\n",
    "im1 = ax.imshow(m11, vmin=vmin, vmax=vmax)\n",
    "ax.axis('off')\n",
    "ax.axis('equal')\n",
    "ax.set_title(\"$G_{xx}$\",fontsize='small',y=0.85)\n",
    "#fig.colorbar(im1, ax = ax, fraction=0.046, pad=0.04)\n",
    "fig.tight_layout(pad=0)\n",
    "fig.savefig('./results/metric_tensor_components11.pdf')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a282bdf7",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig,ax = plt.subplots(figsize = (0.85,1.25))\n",
    "im1 = ax.imshow(m12, vmin=vmin, vmax=vmax)\n",
    "ax.axis('off')\n",
    "ax.axis('equal')\n",
    "ax.set_title(\"$G_{xy}$\",fontsize='small',y=0.85)\n",
    "#fig.colorbar(im1, ax = ax, fraction=0.046, pad=0.04)\n",
    "fig.tight_layout(pad=0)\n",
    "fig.savefig('./results/metric_tensor_components12.pdf')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e70d9260",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig,ax = plt.subplots(figsize = (1.32,1.25))\n",
    "im1 = ax.imshow(m22, vmin=vmin, vmax=vmax)\n",
    "ax.axis('off')\n",
    "ax.axis('equal')\n",
    "ax.set_title(\"$G_{yy}$\",fontsize='small',y=0.85)\n",
    "fig.colorbar(im1, ax = ax, fraction=0.046, pad=0.04)\n",
    "fig.tight_layout(pad=0)\n",
    "fig.savefig('./results/metric_tensor_components22.pdf')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "id": "ae4c8576",
   "metadata": {},
   "outputs": [],
   "source": [
    "# shuffled distribution (space)\n",
    "n_shuffles = 1000\n",
    "\n",
    "a = []\n",
    "b = []\n",
    "c = []\n",
    "\n",
    "for i in range(n_shuffles):\n",
    "    abx = np.copy(ratemaps).reshape(-1, res*res) # shuffle ratemaps\n",
    "    abx = abx[:, np.random.choice(res*res, size = res*res, replace = False)]\n",
    "    abx = abx.reshape(-1, res, res)\n",
    "    grads = np.array([np.gradient(im) for im in abx])\n",
    "\n",
    "    grads = np.transpose(grads, (0, -2, -1, 1))\n",
    "    grads = grads[:, 1:-1, 1:-1] # remove 1st order boundary term\n",
    "\n",
    "    a.append(np.sum(grads[...,0]**2, axis=0))\n",
    "    b.append(np.sum(grads[...,1]**2, axis = 0))\n",
    "    c.append(np.sum(grads[...,1]*grads[...,0], axis = 0))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "eff8c4e5",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig,ax = plt.subplots(figsize = (2,1.5))\n",
    "bins = np.linspace(np.min([m11.min(), m22.min(), m12.min(), \n",
    "                           np.array(a).min(), np.array(b).min(), np.array(c).min()]), \n",
    "                           np.max([m11.max(), m22.max(), m12.max(), np.array(a).max(), \n",
    "                                   np.array(b).max(), np.array(c).max()]), 100)\n",
    "ax.hist(m11.ravel(), label = \"$G_{xx}$\", alpha = 0.5, density=True,bins=bins)\n",
    "ax.hist(m22.ravel(), label = \"$G_{yy}$\", alpha = 0.5, density=True, bins=bins)\n",
    "ax.hist(m12.ravel(), label = \"$G_{xy}$\", alpha = 0.5, density=True, bins=bins)\n",
    "\n",
    "ax.hist(np.array(a).ravel(), label = \"$G^{perm}_{xx}$\", alpha = 0.5, density=True, bins=bins)\n",
    "ax.hist(np.array(b).ravel(), label = \"$G^{perm}_{yy}$\", alpha = 0.5, density=True, bins=bins)\n",
    "ax.hist(np.array(c).ravel(), label = \"$G^{perm}_{xy}$\", alpha = 0.5, density=True, bins=bins)\n",
    "ax.legend(frameon = False)\n",
    "\n",
    "ax.set_xlabel('metric tensor components',fontsize='small')\n",
    "ax.set_ylabel('density',fontsize='small')\n",
    "\n",
    "fig.tight_layout(pad=0)\n",
    "fig.savefig('./results/metric_tensor_components_hist.pdf')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "845af9ad",
   "metadata": {},
   "outputs": [],
   "source": [
    "ci_ratemaps.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "63528964",
   "metadata": {},
   "outputs": [],
   "source": [
    "x = np.linspace(0, 1.5, res) # in meters\n",
    "y = np.copy(x)\n",
    "xx, yy = np.meshgrid(x, y)\n",
    "rr = np.stack((xx.ravel(), yy.ravel()), axis = -1)\n",
    "dr = scipy.spatial.distance.pdist(rr)\n",
    "\n",
    "flat_maps = ratemaps.reshape(-1, res*res)\n",
    "# dg = np.linalg.norm(flat_maps.T[:,None] - flat_maps.T[None], axis = -1)\n",
    "dg = scipy.spatial.distance.pdist(flat_maps.T)\n",
    "\n",
    "flat_maps_ci_cpy = ci_cpy_ratemaps.reshape(res**2, -1)\n",
    "dgi = scipy.spatial.distance.pdist(flat_maps_ci_cpy)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6d2b6347",
   "metadata": {},
   "outputs": [],
   "source": [
    "# rolled baseline\n",
    "rolled = np.array([np.roll(ratemaps[i], -(center_peaks[i] - res//2), axis = (-2, -1)) for i in range(len(ratemaps))])\n",
    "rolled = np.array([np.roll(ratemaps[i], -(center_peaks[i] - res//2), axis = (-2, -1)) for i in range(len(ratemaps))])\n",
    "multiimshow(rolled,figsize = (5,5));\n",
    "rolled = np.reshape(rolled, (-1, res*res)).T\n",
    "droll = scipy.spatial.distance.pdist(rolled)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "id": "69dec88c",
   "metadata": {},
   "outputs": [],
   "source": [
    "def shuffle_ratemap_distances(ratemaps, n_shuffles = 100):\n",
    "    \n",
    "    dns = []\n",
    "    res = ratemaps.shape[-1]\n",
    "\n",
    "    for i in range(n_shuffles):\n",
    "        shuffled = np.copy(ratemaps).reshape(-1, res*res) # shuffle ratemaps\n",
    "        inds = np.random.choice(res*res, size = res*res, replace = False)\n",
    "        shuffled = shuffled[:, inds]\n",
    "        shuffled = np.reshape(shuffled, (-1, res*res)).T\n",
    "        dn = scipy.spatial.distance.pdist(shuffled)\n",
    "        dns.append(dn)\n",
    "\n",
    "    return np.array(dns)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 56,
   "id": "aae39b22",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Data\n",
    "hist_g, xedges_g, _ = scipy.stats.binned_statistic(dr.ravel(), dg.ravel(), \"mean\", bins = 32)\n",
    "std_g, _, _ = scipy.stats.binned_statistic(dr.ravel(), dg.ravel(), \"std\", bins = xedges_g)\n",
    "\n",
    "# Shuffled\n",
    "dns = shuffle_ratemap_distances(ratemaps)\n",
    "drs = np.array([dr for i in range(len(dns))])\n",
    "\n",
    "hist_n, xedges_n, _ = scipy.stats.binned_statistic(drs.ravel(), dns.ravel(), \"mean\", bins = 32)\n",
    "std_n, _, _ = scipy.stats.binned_statistic(drs.ravel(), dns.ravel(), \"std\", bins = xedges_n)\n",
    "\n",
    "# Rolled\n",
    "hist_r, xedges_r, _ = scipy.stats.binned_statistic(dr.ravel(), droll.ravel(), \"mean\", bins = 32)\n",
    "std_r, _, _ = scipy.stats.binned_statistic(dr.ravel(), droll.ravel(), \"std\", bins = xedges_r)\n",
    "\n",
    "\n",
    "xg = xedges_g[:-1] + np.diff(xedges_g)/2 # plot vs. center of bin\n",
    "xn = xedges_n[:-1] + np.diff(xedges_n)/2 \n",
    "xr = xedges_r[:-1] + np.diff(xedges_r)/2 "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "776acca9",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure(figsize = (5,5))\n",
    "\n",
    "# shuffled\n",
    "plt.plot(xn, hist_n)\n",
    "plt.fill_between(xn, hist_n - 2*std_n, hist_n + 2*std_n, alpha = 0.25, label = \"Space Shuffled\")\n",
    "\n",
    "# rolled\n",
    "plt.plot(xr, hist_r)\n",
    "plt.fill_between(xr, hist_r - 2*std_r, hist_r + 2*std_r, alpha = 0.25, label = \"Rolled\")\n",
    "\n",
    "# data\n",
    "plt.plot(xg, hist_g)\n",
    "plt.fill_between(xg, hist_g - 2*std_g, hist_g + 2*std_g, alpha = 0.25, label = \"Data\")\n",
    "\n",
    "plt.legend(frameon = False)\n",
    "plt.xlabel(\"Physical Distance\")\n",
    "plt.ylabel(\"Neural Distance\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "ml",
   "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.12.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}