Merge pull request #32 from Dynamics-of-Neural-Systems-Lab/fix_rat_example

small fixes to rat example

Author: peach-lucien

examples/rat_hippocampus/rat_utils.py

@@ -14,7 +14,7 @@
 import cebra
 
 def convert_spikes_to_rates(spikes,
-                            labels,
+                            labels=None,
                             pca=None,
                             pca_n=10,
                             spiking_rates=True,
@@ -42,7 +42,8 @@ def convert_spikes_to_rates(spikes,
 
     rates_pca = rates_pca[:-1,:] # skip last
 
-    labels = labels[:rates_pca.shape[0]]
+    if labels is not None:
+        labels = labels[:rates_pca.shape[0]]
 
     data = MARBLE.construct_dataset(
         anchor=rates_pca,

examples/rat_hippocampus/run_marble_and_cebra.ipynb

@@ -46,10 +46,7 @@
     "!wget -nc https://dataverse.harvard.edu/api/access/datafile/7609512 -O data/rat_data.pkl\n",
     "\n",
     "with open('data/rat_data.pkl', 'rb') as handle:\n",
-    "    hippocampus_pos = pickle.load(handle)\n",
-    "    \n",
-    "rat = 'achilles'\n",
-    "hippocampus_pos = hippocampus_pos[rat]"
+    "    hippocampus_pos = pickle.load(handle)\n"
    ]
   },
   {
@@ -79,7 +76,8 @@
     "    \n",
     "    return neural_train.numpy(), neural_test.numpy(), label_train.numpy(), label_test.numpy()\n",
     "\n",
-    "neural_train, neural_test, label_train, label_test = split_data(hippocampus_pos, 0.2)"
+    "\n",
+    "neural_train, neural_test, label_train, label_test = split_data(hippocampus_pos['achilles'], 0.2)"
    ]
   },
   {
@@ -113,7 +111,17 @@
    "outputs": [],
    "source": [
     "max_iterations = 10000\n",
-    "output_dimension = 32 #set to 3 for embeddings and 32 for decoding\n",
+    "output_dimension = 32 #set to 3 for embeddings and 32 for decoding\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "scrolled": true
+   },
+   "outputs": [],
+   "source": [
     "\n",
     "for rat in list(hippocampus_pos.keys()):\n",
     "\n",
@@ -167,7 +175,9 @@
    "source": [
     "for rat in list(hippocampus_pos.keys()):\n",
     "    # build model \n",
-    "    data, labels, _ = convert_spikes_to_rates(hippocampus_pos[rat][\"neural\"], label_train, pca_n=10)\n",
+    "    data, labels, _ = convert_spikes_to_rates(hippocampus_pos[rat][\"neural\"],\n",
+    "                                              labels=hippocampus_pos[rat]['continuous_index'].numpy(),\n",
+    "                                              pca_n=10)\n",
     "    pickle.dump([data, labels], open(f'data/{rat}_preprocessed_data.pkl','wb'))\n",
     "    \n",
     "    # build model\n",
@@ -182,7 +192,7 @@
     "        \"include_positions\": True,\n",
     "      }\n",
     "    \n",
-    "    model = MARBLE.net(data, params=params) #define model\n",
+    "    model = MARBLE.net(data, params=params.copy()) #define model\n",
     "    model.fit(data, outdir=f\"data/hippocampus_{rat}\") # train model\n",
     "    data = model.transform(data) #evaluate model"
    ]
@@ -212,6 +222,7 @@
     "        \"order\": 1,  # order of derivatives\n",
     "        \"hidden_channels\": [64],  # number of internal dimensions in MLP\n",
     "        \"out_channels\": 32, \n",
+    "        \"inner_product_features\": False,\n",
     "        \"emb_norm\": True, # spherical output embedding\n",
     "        \"diffusion\": False,\n",
     "        \"include_positions\": True,\n",
@@ -221,20 +232,21 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "metadata": {},
+   "metadata": {
+    "scrolled": true
+   },
    "outputs": [],
    "source": [
     "rat = 'achilles'\n",
     "kernel_width = 10\n",
     "\n",
     "for pca_n in [3,5,10,20,30]:\n",
-    "    data_train, label_train_marble, pca = prepare_marble(neural_train.T, \n",
-    "                                                         label_train, \n",
-    "                                                         pca_n=pca_n, \n",
-    "                                                         kernel_width=kernel_width)\n",
     "\n",
-    "    model = MARBLE.net(data_train, params=params)\n",
-    "    model.fit(data_train, outdir=f\"data/hippocampus_{rat}_pca{pca_n}\")"
+    "    data, _, _ = convert_spikes_to_rates(hippocampus_pos[rat][\"neural\"],\n",
+    "                                        pca_n=pca_n,\n",
+    "                                        kernel_width=kernel_width)\n",
+    "    model = MARBLE.net(data, params=params.copy())\n",
+    "    model.fit(data, outdir=f\"data/hippocampus_{rat}_pca{pca_n}\")\n"
    ]
   },
   {
@@ -247,20 +259,21 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "metadata": {},
+   "metadata": {
+    "scrolled": true
+   },
    "outputs": [],
    "source": [
     "pca_n = 20\n",
     "rat = 'achilles'\n",
     "\n",
     "for kernel_width in [3,5,10,20,30,50,100]:\n",
-    "    data_train, label_train_marble, pca = convert_spikes_to_rates(neural_train.T, \n",
-    "                                                                  label_train, \n",
-    "                                                                  pca_n=pca_n, \n",
-    "                                                                  kernel_width=kernel_width)\n",
-    "    \n",
-    "    model = MARBLE.net(data_train, params=params)\n",
-    "    model.fit(data_train, outdir=f\"data/hippocampus_achilles_kw{k_width}\")"
+    "    data, _, _ = convert_spikes_to_rates(hippocampus_pos[rat][\"neural\"], \n",
+    "                                         pca_n=pca_n,\n",
+    "                                         kernel_width=kernel_width)\n",
+    "    model = MARBLE.net(data, params=params.copy())\n",
+    "    model.fit(data, outdir=f\"data/hippocampus_{rat}_kw{kernel_width}\")\n",
+    "    "
    ]
   }
  ],
@@ -280,7 +293,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.10.14"
+   "version": "3.9.18"
   }
  },
  "nbformat": 4,