Autodownload monkey reaching data (#59)

* add two datasets to monkey-reaching

* add passive_all.jl

* add rest of urls + train/test/valid split

* fix typo

* add urls to test

---------

Co-authored-by: Mackenzie Mathis <[email protected]>

Author: gonlairo

cebra/datasets/monkey_reaching.py

@@ -132,6 +132,82 @@ def _get_info(trial_info, data):
     return data_dic
 
 
+monkey_reaching_urls = {
+    "all_all.jl": {
+        "url":
+            "https://figshare.com/ndownloader/files/41668764?private_link=6fa4ee74a8f465ec7914",
+        "checksum":
+            "dea556301fa4fafa86e28cf8621cab5a"
+    },
+    "all_train.jl": {
+        "url":
+            "https://figshare.com/ndownloader/files/41668752?private_link=6fa4ee74a8f465ec7914",
+        "checksum":
+            "e280e4cd86969e6fd8bfd3a8f402b2fe"
+    },
+    "all_test.jl": {
+        "url":
+            "https://figshare.com/ndownloader/files/41668761?private_link=6fa4ee74a8f465ec7914",
+        "checksum":
+            "25d3ff2c15014db8b8bf2543482ae881"
+    },
+    "all_valid.jl": {
+        "url":
+            "https://figshare.com/ndownloader/files/41668755?private_link=6fa4ee74a8f465ec7914",
+        "checksum":
+            "8cd25169d31f83ae01b03f7b1b939723"
+    },
+    "active_all.jl": {
+        "url":
+            "https://figshare.com/ndownloader/files/41668776?private_link=6fa4ee74a8f465ec7914",
+        "checksum":
+            "c626acea5062122f5a68ef18d3e45e51"
+    },
+    "active_train.jl": {
+        "url":
+            "https://figshare.com/ndownloader/files/41668770?private_link=6fa4ee74a8f465ec7914",
+        "checksum":
+            "72a48056691078eee22c36c1992b1d37"
+    },
+    "active_test.jl": {
+        "url":
+            "https://figshare.com/ndownloader/files/41668773?private_link=6fa4ee74a8f465ec7914",
+        "checksum":
+            "35b7e060008a8722c536584c4748f2ea"
+    },
+    "active_valid.jl": {
+        "url":
+            "https://figshare.com/ndownloader/files/41668767?private_link=6fa4ee74a8f465ec7914",
+        "checksum":
+            "dd58eb1e589361b4132f34b22af56b79"
+    },
+    "passive_all.jl": {
+        "url":
+            "https://figshare.com/ndownloader/files/41668758?private_link=6fa4ee74a8f465ec7914",
+        "checksum":
+            "bbb1bc9d8eec583a46f6673470fc98ad"
+    },
+    "passive_train.jl": {
+        "url":
+            "https://figshare.com/ndownloader/files/41668743?private_link=6fa4ee74a8f465ec7914",
+        "checksum":
+            "f22e05a69f70e18ba823a0a89162a45c"
+    },
+    "passive_test.jl": {
+        "url":
+            "https://figshare.com/ndownloader/files/41668746?private_link=6fa4ee74a8f465ec7914",
+        "checksum":
+            "42453ae3e4fd27d82d297f78c13cd6b7"
+    },
+    "passive_valid.jl": {
+        "url":
+            "https://figshare.com/ndownloader/files/41668749?private_link=6fa4ee74a8f465ec7914",
+        "checksum":
+            "2dcc10c27631b95a075eaa2d2297bb4a"
+    }
+}
+
+
 @register("area2-bump")
 class Area2BumpDataset(cebra.data.SingleSessionDataset):
     """Base dataclass to generate monkey reaching datasets.
@@ -151,19 +227,30 @@ class Area2BumpDataset(cebra.data.SingleSessionDataset):
 
     """
 
-    def __init__(
-        self,
-        path: str = get_datapath("monkey_reaching_preload_smth_40/"),
-        session: str = "active",
-    ):
+    def __init__(self,
+                 path: str = get_datapath("monkey_reaching_preload_smth_40/"),
+                 session: str = "active",
+                 download=True):
         super().__init__()
         self.path = path
+        self.download = download
         self.session = session
         if session == "active-passive":
             self.load_session = "all"
         else:
             self.load_session = session
-        self.data = jl.load(os.path.join(path, f"{self.load_session}_all.jl"))
+
+        super().__init__(
+            download=self.download,
+            data_url=monkey_reaching_urls[f"{self.load_session}_all.jl"]["url"],
+            data_checksum=monkey_reaching_urls[f"{self.load_session}_all.jl"]
+            ["checksum"],
+            location=self.path,
+            file_name=f"{self.load_session}_all.jl",
+        )
+
+        self.data = jl.load(
+            os.path.join(self.path, f"{self.load_session}_all.jl"))
         self._post_load()
 
     def split(self, split):
@@ -179,6 +266,15 @@ def split(self, split):
 
         """
 
+        super().__init__(
+            download=self.download,
+            data_url=monkey_reaching_urls[f"{self.load_session}_{split}.jl"]
+            ["url"],
+            data_checksum=monkey_reaching_urls[
+                f"{self.load_session}_{split}.jl"]["checksum"],
+            location=self.path,
+            file_name=f"{self.load_session}_{split}.jl",
+        )
         self.data = jl.load(
             os.path.join(self.path, f"{self.load_session}_{split}.jl"))
         self._post_load()

tests/test_datasets.py

@@ -299,7 +299,16 @@ def parametrize_data(function):
          "a83b02dbdc884fdd7e53df362499d42f"),
         ("gatsby.jl",
          "https://figshare.com/ndownloader/files/40849454?private_link=9f91576cbbcc8b0d8828",
-         "2b889da48178b3155011c12555342813")
+         "2b889da48178b3155011c12555342813"),
+        ("all_all.jl",
+         "https://figshare.com/ndownloader/files/41668764?private_link=6fa4ee74a8f465ec7914",
+         "dea556301fa4fafa86e28cf8621cab5a"),
+        ("active_all.jl",
+         "https://figshare.com/ndownloader/files/41668776?private_link=6fa4ee74a8f465ec7914",
+         "c626acea5062122f5a68ef18d3e45e51"),
+        ("passive_all.jl",
+         "https://figshare.com/ndownloader/files/41668758?private_link=6fa4ee74a8f465ec7914",
+         "bbb1bc9d8eec583a46f6673470fc98ad"),
     ])(function)