Repository: rmarkello/abagen
Pull request #186 (2021-03-29 15:36:53)

[ENH] Add geometry/space option to check_atlas
Makes it much easier to specify non-standard atlases / surfaces, which benefits construction of `AtlasTree` objects and enables better use of graph representation of surfaces


Patch:
diff --git a/abagen/datasets/fetchers.py b/abagen/datasets/fetchers.py
index 049de56..fdc31fe 100644
--- a/abagen/datasets/fetchers.py
+++ b/abagen/datasets/fetchers.py
@@ -465,15 +465,21 @@ Brain = namedtuple('Brain', ('lh', 'rh'))
 Surface = namedtuple('Surface', ('vertices', 'faces'))
 
 
-def fetch_fsaverage5():
+def fetch_fsaverage5(load=True):
     """
-    Fetches and load fsaverage5 surface
+    Fetches and optionally loads fsaverage5 surface
+
+    Parameters
+    ----------
+    load : bool, optional
+        Whether to pre-load files. Default: True
 
     Returns
     -------
     brain : namedtuple ('lh', 'rh')
-        Where each entry in the tuple is a hemisphere, represented as a
-        namedtuple with fields ('vertices', 'faces')
+        If `load` is True, a namedtuple where each entry in the tuple is a
+        hemisphere, represented as a namedtuple with fields ('vertices',
+        'faces'). If `load` is False, a namedtuple where entries are filepaths.
     """
 
     hemispheres = []
@@ -481,14 +487,18 @@ def fetch_fsaverage5():
         fn = RESOURCE(
             os.path.join('data', f'fsaverage5-pial-{hemi}.surf.gii.gz')
         )
-        hemispheres.append(Surface(*load_gifti(fn).agg_data()))
+        if load:
+            hemispheres.append(Surface(*load_gifti(fn).agg_data()))
+        else:
+            hemispheres.append(fn)
 
     return Brain(*hemispheres)
 
 
-def fetch_fsnative(donors, surf='pial', data_dir=None, resume=True, verbose=1):
+def fetch_fsnative(donors, surf='pial', load=True, data_dir=None, resume=True,
+                   verbose=1):
     """
-    Fetches and load fsnative surface of `donor`
+    Fetches and optionally loads fsnative surface of `donor`
 
     Parameters
     ----------
@@ -497,6 +507,8 @@ def fetch_fsnative(donors, surf='pial', data_dir=None, resume=True, verbose=1):
         specify 'all' to download all available donors.
     surf : {'orig', 'white', 'pial', 'inflated', 'sphere'}, optional
         Which surface to load. Default: 'pial'
+    load : bool, optional
+        Whether to pre-load files. Default: True
     data_dir : str, optional
         Directory where data should be downloaded and unpacked. Default: $HOME/
         abagen-data
@@ -508,9 +520,11 @@ def fetch_fsnative(donors, surf='pial', data_dir=None, resume=True, verbose=1):
     Returns
     -------
     brain : namedtuple ('lh', 'rh')
-        Where each entry in the tuple is a hemisphere, represented as a
-        namedtuple with fields ('vertices', 'faces'). If multiple donors are
-        requested a dictionary is returned where keys are donor IDs.
+        If `load` is True, a namedtuple where each entry in the tuple is a
+        hemisphere, represented as a namedtuple with fields ('vertices',
+        'faces'). If `load` is False, a namedtuple where entries are filepaths.
+        If multiple donors are requested a dictionary is returned where keys
+        are donor IDs.
     """
 
     donors = check_donors(donors)
@@ -524,6 +538,9 @@ def fetch_fsnative(donors, surf='pial', data_dir=None, resume=True, verbose=1):
     hemispheres = []
     for hemi in ('lh', 'rh'):
         fn = os.path.join(fpath, 'surf', f'{hemi}.{surf}')
-        hemispheres.append(Surface(*nib.freesurfer.read_geometry(fn)))
+        if load:
+            hemispheres.append(Surface(*nib.freesurfer.read_geometry(fn)))
+        else:
+            hemispheres.append(fn)
 
     return Brain(*hemispheres)
diff --git a/abagen/images.py b/abagen/images.py
index 199c856..0fae903 100644
--- a/abagen/images.py
+++ b/abagen/images.py
@@ -256,7 +256,8 @@ def check_surface(atlas):
                 raise ValueError('Provided GIFTIs do not seem to be valid '
                                  'label.gii files')
         adata.append(data)
-        labs.append(hemi.labeltable.get_labels_as_dict())
+        ldict = hemi.labeltable.get_labels_as_dict()
+        labs.append({k: ldict.get(k) for k in np.unique(data)})
 
     # we need each hemisphere to have unique values so they don't get averaged
     # check to see if the two hemispheres have more than 1 overlapping value
@@ -273,7 +274,8 @@ def check_surface(atlas):
     return adata, atlas_info
 
 
-def check_atlas(atlas, atlas_info=None, donor=None, data_dir=None):
+def check_atlas(atlas, atlas_info=None, geometry=None, space=None, donor=None,
+                data_dir=None):
     """
     Checks that `atlas` is a valid atlas
 
@@ -285,9 +287,14 @@ def check_atlas(atlas, atlas_info=None, donor=None, data_dir=None):
     atlas_info : {os.PathLike, pandas.DataFrame, None}, optional
         Filepath or dataframe containing information about `atlas`. Must have
         at least columns ['id', 'hemisphere', 'structure'] containing
-        information mapping `atlas` IDs to hemisphere (i.e., "L" or "R") and
+        information mapping `atlas` IDs to hemisphere (i.e., "L", "R", "B") and
         broad structural class (i.e.., "cortex", "subcortex/brainstem",
         "cerebellum", "white matter", or "other"). Default: None
+    geometry : (2,) tuple-of-GIFTI, optional
+        Surfaces files defining geometry of `atlas`, if `atlas` is a tuple of
+        GIFTI images. Default: None
+    space : {'fsaverage', 'fsnative', 'fslr'}, optional
+        If `geometry` is supplied, what space files are in. Default: None
     donor : str, optional
         If specified, indicates which donor the specified `atlas` belongs to.
         Only relevant when `atlas` is surface-based, to ensure the correct
@@ -314,21 +321,18 @@ def check_atlas(atlas, atlas_info=None, donor=None, data_dir=None):
         coords = triangles = None
     except TypeError:
         atlas, info = check_surface(atlas)
-        if donor is None:
-            data = fetch_fsaverage5()
-            coords = transforms.fsaverage_to_mni152(
-                np.row_stack([hemi.vertices for hemi in data])
-            )
-        else:
-            data = fetch_fsnative(donor, data_dir=data_dir)
-            coords = transforms.fsnative_to_xyz(
-                np.row_stack([hemi.vertices for hemi in data]), donor
-            )
-        triangles, offset = [], 0
-        for hemi in data:
-            triangles.append(hemi.faces + offset)
-            offset += hemi.vertices.shape[0]
-        triangles = np.row_stack(triangles)
+        # backwards compatibility for `donor` keyword
+        if geometry is None and donor is None:
+            geometry = fetch_fsaverage5()
+            space = 'fsaverage5'
+        elif geometry is None and donor is not None:
+            geometry = fetch_fsnative(donor, data_dir=data_dir)
+            space = 'fsnative'
+        elif geometry is not None and space is None:
+            raise ValueError('If providing geometry files space parameter '
+                             'must be specified')
+        coords, triangles = check_geometry(geometry, space, donor=donor,
+                                           data_dir=data_dir)
         if atlas_info is None and info is not None:
             atlas_info = info
 
@@ -340,6 +344,63 @@ def check_atlas(atlas, atlas_info=None, donor=None, data_dir=None):
     return atlas
 
 
+def check_geometry(surface, space, donor=None, data_dir=None):
+    """
+    Loads geometry `surface` files and transforms coordinates in `space`
+
+    Parameters
+    ----------
+    surface : (2,) tuple-of-GIFTI
+        Surface geometry files in GIFTI format (lh, rh)
+    space : {'fsaverage', 'fsnative', 'fslr'}
+        What space `surface` files are in; used to apply appropriate transform
+        to MNI152 space. If 'fsnative' then `donor` must be supplied as well
+    donor : str, optional
+        If specified, indicates which donor the specified `surface` belongs to
+    data_dir : str, optional
+        Directory where donor-specific FreeSurfer data exists (or should be
+        downloaded and unpacked). Only used if provided `donor` is not None.
+        Default: $HOME/abagen-data
+
+    Returns
+    -------
+    coords : (N, 3) np.ndarray
+        Coordinates from `surface` files
+    triangles : (T, 3) np.ndarray
+        Triangles from `surface` files
+    """
+
+    if len(surface) != 2:
+        raise TypeError('Must provide a tuple of geometry files')
+
+    # fsaverage5, fsaverage6, etc
+    if 'fsaverage' in space and space != 'fsaverage':
+        space = 'fsaverage'
+    space_opts = ('fsaverage', 'fsnative', 'fslr')
+    if space not in space_opts:
+        raise ValueError(f'Provided space must be one of {space_opts}.')
+    if space == 'fsnative' and donor is None:
+        raise ValueError('Specified space is "fsnative" but no donor ID '
+                         'supplied')
+
+    try:
+        coords, triangles = map(list, zip(*[
+            load_gifti(img).agg_data() for img in surface
+        ]))
+    except TypeError:
+        coords, triangles = map(list, zip(*[i for i in surface]))
+
+    triangles[-1] += coords[0].shape[0]
+    coords, triangles = np.row_stack(coords), np.row_stack(triangles)
+
+    if space == 'fsaverage':
+        coords = transforms.fsaverage_to_mni152(coords)
+    elif space == 'fsnative':
+        coords = transforms.fsnative_to_xyz(coords, donor, data_dir=data_dir)
+
+    return coords, triangles
+
+
 def check_atlas_info(atlas_info, labels):
     """
     Checks whether provided `atlas_info` is correct format for processing
@@ -461,11 +522,11 @@ def coerce_atlas_to_dict(atlas, donors, atlas_info=None, data_dir=None):
     donors = check_donors(donors)
     group_atlas = True
 
-    # FIXME: so that we're not depending on type checks so much :grimacing:
-    if isinstance(atlas, dict):
+    try:
         atlas = {
             WELL_KNOWN_IDS.subj[donor]: check_atlas(atl, atlas_info,
-                                                    donor, data_dir)
+                                                    donor=donor,
+                                                    data_dir=data_dir)
             for donor, atl in atlas.items()
         }
         # if it's a group atlas they should all be the same object
@@ -477,7 +538,7 @@ def coerce_atlas_to_dict(atlas, donors, atlas_info=None, data_dir=None):
                              f'requested donors. Missing donors: {donors}.')
         LGR.info('Donor-specific atlases provided; using native coords for '
                  'tissue samples')
-    else:
+    except AttributeError:
         atlas = check_atlas(atlas, atlas_info)
         atlas = {donor: atlas for donor in donors}
         LGR.info('Group-level atlas provided; using MNI coords for '
diff --git a/abagen/matching.py b/abagen/matching.py
index 7b01993..daee9d2 100644
--- a/abagen/matching.py
+++ b/abagen/matching.py
@@ -40,38 +40,34 @@ class AtlasTree:
     def __init__(self, atlas, coords=None, *, triangles=None, atlas_info=None):
         from .images import check_img
 
-        graph = None
+        self._full_coords = self._graph = None
         try:  # let's first check if it's an image
             atlas = check_img(atlas)
-            data, affine = np.asarray(atlas.dataobj), atlas.affine
-            self._shape = atlas.shape
-            nz = data.nonzero()
+            atlas, affine = np.asarray(atlas.dataobj), atlas.affine
             if coords is not None:
                 warnings.warn('Volumetric image supplied to `AtlasTree` '
                               'constructor but `coords` is not None. Ignoring '
                               'supplied `coords` and using coordinates '
                               'derived from image.')
-            atlas, coords = data[nz], transforms.ijk_to_xyz(np.c_[nz], affine)
             self._volumetric = True
+            self._shape = atlas.shape
+            nz = atlas.nonzero()
+            atlas, coords = atlas[nz], transforms.ijk_to_xyz(np.c_[nz], affine)
         except TypeError:
             atlas = np.asarray(atlas)
+            self._full_coords = coords
             if coords is None:
                 raise ValueError('When providing a surface atlas you must '
                                  'also supply relevant geometry `coords`.')
             if len(atlas) != len(coords):
                 raise ValueError('Provided `atlas` and `coords` are of '
                                  'differing length.')
+            self._volumetric = False
             self._shape = atlas.shape
             nz = atlas.nonzero()
-            if triangles is not None:
-                graph = surfaces.make_surf_graph(coords, triangles,
-                                                 atlas == 0)[nz].T[nz].T
             atlas, coords = atlas[nz], coords[nz]
-            self._volumetric = False
 
-        self._triangles = triangles
         self._nz = nz
-        self._graph = graph
         self._tree = cKDTree(coords)
         self._atlas = np.asarray(atlas)
         self._labels = np.unique(self.atlas).astype(int)
@@ -82,6 +78,7 @@ class AtlasTree:
             _, idx = self.tree.query(centroids, k=1)
             self._centroids = dict(zip(self.labels, self.coords[idx]))
         self.atlas_info = atlas_info
+        self.triangles = triangles
 
     def __repr__(self):
         if self.volumetric:
@@ -121,6 +118,12 @@ class AtlasTree:
         """
         return self._centroids
 
+    @property
+    def graph(self):
+        """ Returns graph of underlying parcellation
+        """
+        return self._graph
+
     @property
     def coords(self):
         """ Returns coordinates of underlying cKDTree
@@ -131,13 +134,41 @@ class AtlasTree:
     def coords(self, pts):
         """ Sets underlying cKDTree to represent provided `pts`
         """
-        if len(pts) != len(self.atlas):
+        pts = np.asarray(pts)
+        if pts.shape[0] != self.atlas.shape[0]:
             raise ValueError('Provided coordinates do not match length of '
                              'current atlas. Expected {}. Received {}'
                              .format(len(self.atlas), len(pts)))
         if not np.allclose(pts, self.coords):
             self._tree = cKDTree(pts)
             self._centroids = get_centroids(self.atlas, pts)
+            # update graph with new coordinates (if relevant)
+            self.triangles = self.triangles
+
+    @property
+    def triangles(self):
+        """ Returns triangles of underlying graph (if applicable)
+        """
+        return self._triangles
+
+    @triangles.setter
+    def triangles(self, tris):
+        """ Sets triangles of underlying graph (if applicable)
+        """
+        if self.volumetric or tris is None:
+            self._triangles = None
+            return
+
+        tris = np.asarray(tris)
+        atlas = np.zeros(self._shape)
+        atlas[self._nz] = self.atlas
+        if np.any(tris.max(axis=0) >= self._full_coords.shape[0]):
+            raise ValueError('Cannot provide triangles with indices greater '
+                             'than tree coordinate array')
+        self._triangles = tris
+        self._graph = surfaces.make_surf_graph(
+            self._full_coords, self._triangles, atlas == 0
+        )[self._nz].T[self._nz].T
 
     @property
     def atlas_info(self):
diff --git a/abagen/utils.py b/abagen/utils.py
index 526578f..61dd02f 100644
--- a/abagen/utils.py
+++ b/abagen/utils.py
@@ -182,7 +182,7 @@ def labeltable_to_df(labels):
                 dict(id=ids, label=label, hemisphere=hemi, structure='cortex')
             ), ignore_index=True
         )
-    info = info.set_index('id').drop([0], axis=0)
+    info = info.set_index('id').drop([0], axis=0).sort_index()
 
     if len(info) != 0:
         return info
diff --git a/docs/user_guide/parcellations.rst b/docs/user_guide/parcellations.rst
index a9443fc..7dcdb5b 100644
--- a/docs/user_guide/parcellations.rst
+++ b/docs/user_guide/parcellations.rst
@@ -180,6 +180,10 @@ by default, be based on the geometry of the FreeSurfer surfaces provided with
 based on different geometry please refer to :ref:`usage_parcellations_special`,
 below.
 
+Finally, when in doubt we recommend simply using a standard-space, group-level
+atlas; however, we are actively investigating whether native-space atlases
+provide any measurable benefits to the ``abagen`` workflows.
+
 .. note::
 
     The donor-native volumetric versions of the DK parcellation shipped with
@@ -211,45 +215,25 @@ internally coerced to `AtlasTree` instances, which is then used to assign
 microarray tissue samples to parcels in the atlas.
 
 Take, for example, a surface atlas in fsaverage6 resolution (by default,
-surface atlases are assumed to be fsaverage5 resolution). In this case, you can
-load the surface atlas and relevant geometry files, create an `AtlasTree`
-object, and use that as your atlas moving forward in the same way you would
-with a pair of GIFTI files:
+surface atlases are assumed to be fsaverage5 resolution). In this case, you
+simply need to supply the relevant geometry files for the atlas and specify
+the space of the atlas:
 
 .. code::
 
-    >>> from abagen import matching, transforms
-    >>> atlas = ('/.../myatlas-fsaverage6-lh.gii', '/.../myatlas-fsaverage6-rh.gii')
-    >>> surf = ('/.../mysurf-fsaverage6-lh.gii', '/.../mysurf-fsaverage6-lh.gii')
-    >>> parcellation, atlas_info = images.check_surface(atlas)
-    >>> coords = np.row_stack([nib.load(fn).agg_data('NIFTI_INTENT_POINTSET') for fn in surf])
-    >>> mni152 = transforms.fsaverage_to_mni152(coords)
-    >>> atlas = matching.AtlasTree(parcellation, coords=mni152, atlas_info=atlas_info)
-
-(Note that you need to coerce the fsaverage coordinates---which are in MNI305
-space---to MNI152 space before using them in the :obj:`~.AtlasTree`
-constructor!)
+    >>> from abagen import images
+    >>> atlas = ('/.../fsaverage6-lh.label.gii', '/.../fsaverage6-rh.label.gii')
+    >>> surf = ('/.../fsaverage6-lh.surf.gii', '/.../fsaverage6-lh.surf.gii')
+    >>> atlas = images.check_atlas(atlas, geometry=surf, space='fsaverage6')
 
-The same procedure can be used for an atlas using fslr32k geometry:
+The same procedure can be used for an atlas using fsLR geometry:
 
 .. code::
 
-    >>> from abagen import matching, transforms
-    >>> atlas = ('/.../myatlas-fslr32k-lh.gii', '/.../myatlas-fslr32k-rh.gii')
-    >>> surf = ('/.../mysurf-fslr32k-lh.gii', '/.../mysurf-fslr32k-lh.gii')
-    >>> parcellation, atlas_info = images.check_surface(atlas)
-    >>> coords = np.row_stack([nib.load(fn).agg_data('NIFTI_INTENT_POINTSET') for fn in surf])
-    >>> atlas = matching.AtlasTree(parcellation, coords=coords, atlas_info=atlas_info)
-
-(Note that because fslr32k is already in MNI space you do not need to transform
-the coordinates before using them in the :obj:`~.matching.AtlasTree`
-constructor.)
-
-.. note::
-
-    When in doubt it is likely best to use a standard-space, group-level atlas;
-    however, we are actively investigating whether native-space atlases provide
-    any measurable benefits to the ``abagen`` workflow.
+    >>> from abagen import images
+    >>> atlas = ('/.../fslr32k-lh.label.gii', '/.../fslr32k-rh.label.gii')
+    >>> surf = ('/.../fslr32k-lh.surf.gii', '/.../fslr32k-lh.surf.gii')
+    >>> atlas = images.check_atlas(atlas, geometry=surf, space='fslr')
 
 .. _MNI space: https://www.ncbi.nlm.nih.gov/pmc/articles/PMC1088516/
 .. _fsaverage space: https://surfer.nmr.mgh.harvard.edu/


Test patch:
diff --git a/abagen/tests/datasets/test_fetchers.py b/abagen/tests/datasets/test_fetchers.py
index 7bc3352..fcfd36b 100644
--- a/abagen/tests/datasets/test_fetchers.py
+++ b/abagen/tests/datasets/test_fetchers.py
@@ -171,6 +171,12 @@ def test_fetch_fsaverage5():
             assert hasattr(hemi, attr)
             assert getattr(hemi, attr).shape == (exp, 3)
 
+    fs5 = fetchers.fetch_fsaverage5(load=False)
+    for hemi in ('lh', 'rh'):
+        assert hasattr(fs5, hemi)
+        hemi = getattr(fs5, hemi)
+        assert Path(hemi).is_file()
+
 
 def test_fetch_fsnative():
     fsn = fetchers.fetch_fsnative(donors=['12876'])
@@ -184,3 +190,9 @@ def test_fetch_fsnative():
         hemi = getattr(fsn, hemi)
         for attr in ('vertices', 'faces'):
             assert hasattr(hemi, attr)
+
+    fsn = fetchers.fetch_fsnative(donors=['12876'], load=False)
+    for hemi in ('lh', 'rh'):
+        assert hasattr(fsn, hemi)
+        hemi = getattr(fsn, hemi)
+        assert Path(hemi).is_file()
diff --git a/abagen/tests/test_images.py b/abagen/tests/test_images.py
index 07c4281..453ab48 100644
--- a/abagen/tests/test_images.py
+++ b/abagen/tests/test_images.py
@@ -32,6 +32,14 @@ def annotation(tmp_path_factory):
     return fname
 
 
+@pytest.fixture(scope='module')
+def fsgeometry():
+    return (
+        resource_filename('abagen', 'data/fsaverage5-pial-lh.surf.gii.gz'),
+        resource_filename('abagen', 'data/fsaverage5-pial-rh.surf.gii.gz'),
+    )
+
+
 def test_leftify_atlas(atlas):
     out = images.leftify_atlas(atlas['image'])
     assert len(np.unique(out.dataobj)) == 44
@@ -113,7 +121,7 @@ def test_check_surface(surface):
         images.check_surface(('lh.nii.gz', 'rh.nii.gz'))
 
 
-def test_check_atlas(atlas, surface):
+def test_check_atlas(atlas, surface, fsgeometry):
     # check loading volumetric atlas
     tree = images.check_atlas(atlas['image'])
     assert isinstance(tree, AtlasTree)
@@ -135,6 +143,16 @@ def test_check_atlas(atlas, surface):
     assert not tree.volumetric
     assert len(tree.coords) == 18426
 
+    tree = images.check_atlas(surface['image'], geometry=fsgeometry,
+                              space='fsaverage')
+    assert isinstance(tree, AtlasTree)
+    assert isinstance(tree.atlas_info, pd.DataFrame)
+    assert not tree.volumetric
+    assert len(tree.coords) == 18426
+
+    with pytest.raises(ValueError):
+        images.check_atlas(surface['image'], geometry=fsgeometry)
+
     # check loading donor-specific surface file
     fp = 'data/native_dk/12876/atlas-desikankilliany-{}.label.gii.gz'
     surf = [
@@ -147,6 +165,19 @@ def test_check_atlas(atlas, surface):
     assert len(tree.coords) == 386566
 
 
+def test_check_geometry(fsgeometry):
+    coords, triangles = images.check_geometry(fsgeometry, 'fsaverage')
+    assert len(coords) == 20484
+    assert len(triangles) == 40960
+
+    with pytest.raises(ValueError):
+        images.check_geometry(fsgeometry, 'notaspace')
+    with pytest.raises(ValueError):
+        images.check_geometry(fsgeometry, 'fsnative', donor=None)
+    with pytest.raises(TypeError):
+        images.check_geometry(fsgeometry[0], 'fsaverage')
+
+
 def test_check_atlas_info(atlas):
     labels = np.trim_zeros(np.unique(nib.load(atlas['image']).dataobj))
 
diff --git a/abagen/tests/test_matching.py b/abagen/tests/test_matching.py
index 9b95321..64ab194 100644
--- a/abagen/tests/test_matching.py
+++ b/abagen/tests/test_matching.py
@@ -142,11 +142,16 @@ def test_AtlasTree(atlas, surface, testfiles):
     tree = images.check_atlas(surface['image'])
     assert str(tree) == 'AtlasTree[n_rois=68, n_vertex=18426]'
     assert not tree.volumetric
+    assert tree.triangles is not None
+    assert tree.graph is not None
     assert isinstance(tree.atlas_info, pd.DataFrame)
     assert len(tree.centroids) == 68
     labels = tree.label_samples([tree.centroids[1], tree.centroids[2]])
     assert np.all(labels['label'] == [1, 2])
 
+    with pytest.raises(ValueError):
+        tree.triangles = [[512423, 512312, 4213215]]
+
     # check negative surface tolerance
     labels = tree.label_samples([-72, -25, -13], tolerance=-4)
     assert np.all(labels['label'] == 14)