[ENH] Add fslr32k schaefer atlases

Author: rmarkello

netneurotools/data/osf.json

@@ -343,6 +343,13 @@
                 "5dbc8d7bcfc96c000ec6dca2"
             ],
             "md5": "8738daccab4648c3e891a1c8d3a9ec1f"
+        },
+        "fslr32k": {
+            "url": [
+                "mb37e",
+                "5e3086e4af75930094bdd507"
+            ],
+            "md5": "d8378f33107ed5d98c27e8070ebb5aa2"
         }
     }
 }

netneurotools/datasets/fetchers.py

@@ -489,7 +489,7 @@ def fetch_schaefer2018(version='fsaverage', data_dir=None, url=None,
 
     Parameters
     ----------
-    version : {'fsaverage', 'fsaverage5', 'fsaverage6'}
+    version : {'fsaverage', 'fsaverage5', 'fsaverage6', 'fslr32k'}
         Specifies which surface annotation files should be matched to. Default:
         'fsaverage'
     data_dir : str, optional
@@ -523,7 +523,7 @@ def fetch_schaefer2018(version='fsaverage', data_dir=None, url=None,
     License: https://github.com/ThomasYeoLab/CBIG/blob/master/LICENSE.md
     """
 
-    versions = ['fsaverage', 'fsaverage5', 'fsaverage6']
+    versions = ['fsaverage', 'fsaverage5', 'fsaverage6', 'fslr32k']
     if version not in versions:
         raise ValueError('The version of Schaefer et al., 2018 parcellation '
                          'requested "{}" does not exist. Must be one of {}'
@@ -546,16 +546,22 @@ def fetch_schaefer2018(version='fsaverage', data_dir=None, url=None,
         'move': '{}.tar.gz'.format(dataset_name)
     }
 
+    if version == 'fslr32k':
+        hemispheres, suffix = ['LR'], 'dlabel.nii'
+    else:
+        hemispheres, suffix = ['L', 'R'], 'annot'
     filenames = [
-        'atl-Schaefer2018_space-{}_hemi-{}_desc-{}_deterministic.annot'
-        .format(version, hemi, desc) for desc in keys for hemi in ['L', 'R']
+        'atl-Schaefer2018_space-{}_hemi-{}_desc-{}_deterministic.{}'
+        .format(version, hemi, desc, suffix)
+        for desc in keys for hemi in hemispheres
     ]
 
     files = [(op.join(dataset_name, version, f), url, opts)
              for f in filenames]
     data = _fetch_files(data_dir, files=files, resume=resume, verbose=verbose)
 
-    data = [data[i:i + 2] for i in range(0, len(keys) * 2, 2)]
+    if suffix == 'annot':
+        data = [data[i:i + 2] for i in range(0, len(keys) * 2, 2)]
 
     return Bunch(**dict(zip(keys, data)))
 

netneurotools/tests/test_datasets.py

@@ -125,7 +125,7 @@ def test_fetch_connectome(tmpdir, dataset, expected):
 
 
 @pytest.mark.parametrize('version', [
-    'fsaverage', 'fsaverage5', 'fsaverage6'
+    'fsaverage', 'fsaverage5', 'fsaverage6', 'fslr32k'
 ])
 def test_fetch_schaefer2018(tmpdir, version):
     keys = [
@@ -134,10 +134,13 @@ def test_fetch_schaefer2018(tmpdir, version):
     ]
     schaefer = datasets.fetch_schaefer2018(version, data_dir=tmpdir, verbose=0)
 
-    assert all(k in schaefer
-               and len(schaefer[k]) == 2
-               and all(os.path.isfile(hemi) for hemi in schaefer[k])
-               for k in keys)
+    if version == 'fslr32k':
+        assert all(k in schaefer and os.path.isfile(schaefer[k]) for k in keys)
+    else:
+        assert all(k in schaefer
+                   and len(schaefer[k]) == 2
+                   and all(os.path.isfile(hemi) for hemi in schaefer[k])
+                   for k in keys)
 
 
 def test_fetch_hcp_standards(tmpdir):