Merge pull request #186 from rmarkello/geometry

[ENH] Add geometry/space option to check_atlas

abagen/datasets/fetchers.py

@@ -465,30 +465,40 @@ def fetch_donor_info():
 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 = []
     for hemi in ('lh', 'rh'):
         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)

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 '

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 @@ def __init__(self, atlas, coords=None, *, triangles=None, atlas_info=None):
             _, 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 @@ def centroids(self):
         """
         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 @@ def coords(self):
     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):

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()