Repository: rmarkello/abagen

Issue:
Donor-specific atlases and tolerance=0 fails for two donors
## The issue

When using donor-specific atlases and providing `tolerance=0` to the `abagen.get_expression_data()` workflow, two of the donors (`'9861'` and `'10021'`) have _no matched samples_. This is wrong and is due entirely to this line:

https://github.com/rmarkello/abagen/blob/ed4ce09d7e628b82a5f52d4c03f3fa61250a87e0/abagen/matching.py#L316

This line is important for when we're using a group-level atlas and trying to determine whether a sample (using its MNI coordinates) falls directly within a voxel of the provided atlas; unfortunately, the x-offset for the two aforementioned donors is a decimal (0.5), so the call to `np.floor` changes the coordinates such that they no longer align with the donor-specific atlases.

To reproduce:

```python
>>> import abagen
>>> atlases = abagen.fetch_desikan_killiany(native=True)['image']
>>> expression = abagen.get_expression_data(atlases, tolerance=0, return_donors=True)
>>> np.all(np.isnan(expression[0]))
True
>>> np.all(np.isnan(expression[1]))
True
```

## Proposed solution

The call to `np.floor` is a bit too basic. It assumes that all the provided atlases will always be on a regular grid (i.e., 1mm isotropic) and that the offsets will always be integer numbers—and while this may _generally_ be the case, it is impossible to say that it will always be true.

The `np.floor` line should be removed entirely, and, perhaps somewhere in `samples_.update_coords()` where we first handle the coordinates, we should update the MNI coordinates such that they align with the provided atlas grid. This will require maintaining some information about the image grid for the relevant atlas(es) (or at the very least voxel sizes).

(This does raise another point about whether we should be rounding the MNI coordinates down to the nearest atlas voxel at all. I maintain that this is reasonable to do, but it is a design choice that could be argued against.)

Discussion:

Patch:
diff --git a/abagen/allen.py b/abagen/allen.py
index e22338d..fa8ce03 100644
--- a/abagen/allen.py
+++ b/abagen/allen.py
@@ -13,7 +13,6 @@ import pandas as pd
 
 from . import (correct, datasets, images, io, matching, probes_, reporting,
                samples_, utils)
-from .transforms import xyz_to_ijk
 from .utils import first_entry, flatten_dict
 
 logging.basicConfig(format='%(levelname)s: %(message)s', level=logging.WARNING)
@@ -516,14 +515,14 @@ def get_expression_data(atlas,
 
     # if we don't want to aggregate over regions return sample-level results
     if region_agg is None:
-        # don't return samples that aren't matched to a region in the `atlas`
-        mask = {d: m.index != 0 for d, m in microarray.items()}
-        microarray = pd.concat([m[mask[d]] for d, m in microarray.items()])
-        # set index to well_id for all remaining tissue samples
-        well_ids = pd.concat([a[mask[d]] for d, a in annotation.items()])
-        microarray = microarray.set_index(well_ids['well_id'], append=True)
-        # return expression data (remove NaNs)
-        return microarray.dropna(axis=1, how='any')
+        microarray = {
+            donor: micro.set_index(annotation[donor]['well_id'], append=True)
+                        .dropna(axis=1, how='any')
+            for donor, micro in microarray.items()
+        }
+        if not return_donors:
+            microarray = pd.concat(microarray.values())
+        return microarray
 
     if missing == 'centroids':
         # labels that are missing across all donors
@@ -544,7 +543,7 @@ def get_expression_data(atlas,
             for subj in microarray:
                 microarray[subj] = microarray[subj].append(exp)
 
-    microarray = samples_.aggregate_samples(microarray.values(),
+    microarray = samples_.aggregate_samples(microarray,
                                             labels=all_labels,
                                             region_agg=region_agg,
                                             agg_metric=agg_metric,
@@ -567,7 +566,7 @@ def get_expression_data(atlas,
                                   return_donors=return_donors,
                                   data_dir=data_dir, counts=counts,
                                   n_probes=len(probe_info),
-                                  n_genes=(microarray[0].shape[1]
+                                  n_genes=(first_entry(microarray).shape[1]
                                            if return_donors
                                            else microarray.shape[1])).body
 
@@ -625,12 +624,10 @@ def get_samples_in_mask(mask=None, **kwargs):
 
     # get updated coordinates
     for donor, data in files.items():
-        annot, ontol = data['annotation'], data['ontology']
+        annot = data['annotation']
         if kwargs.get('corrected_mni', True):
             annot = samples_.update_mni_coords(annot)
-        if kwargs.get('lr_mirror', False):
-            annot = samples_.mirror_samples(annot, ontol)
-        data['annotation'] = annot
+        files[donor]['annotation'] = annot
     cols = ['well_id', 'mni_x', 'mni_y', 'mni_z']
     coords = np.asarray(pd.concat(flatten_dict(files, 'annotation'))[cols])
     well_id, coords = np.asarray(coords[:, 0], 'int'), coords[:, 1:]
@@ -640,16 +637,10 @@ def get_samples_in_mask(mask=None, **kwargs):
     if kwargs.get('atlas') is not None and mask is None:
         mask = kwargs['atlas']
     elif mask is None:
-        # create affine for "full" mask
+        # generate atlas image where each voxel with sample has value of 1
         affine = np.eye(4)
-        affine[:-1, -1] = np.floor(coords).min(axis=0)
-
-        # downsample coordinates to specified resolution and convert to ijk
-        ijk = np.unique(xyz_to_ijk(coords, affine), axis=0)
-
-        # generate atlas image where each voxel has
-        img = np.zeros(ijk.max(axis=0) + 2, dtype='int')
-        img[tuple(map(tuple, ijk.T))] = 1
+        affine[:-1, -1] = -1 * np.floor(np.max(np.abs(coords), axis=0))
+        img = np.ones(np.asarray(-2 * affine[:-1, -1], dtype=int))
         mask = nib.Nifti1Image(img, affine=affine)
 
     # reset these parameters
@@ -660,10 +651,23 @@ def get_samples_in_mask(mask=None, **kwargs):
     kwargs.setdefault('norm_matched', False)
 
     # get expression data + drop sample coordinates that weren't in atlas
-    exp = get_expression_data(**kwargs).droplevel('label')
-    keep = np.isin(coords.index, exp.index)
-
-    return exp, coords.loc[keep]
+    exp = get_expression_data(**kwargs)
+    if kwargs.get('return_donors'):
+        exp = {
+            donor: micro.drop(index=[0], level='label', errors='ignore')
+                        .droplevel('label')
+            for donor, micro in exp.items()
+        }
+        coords = {
+            donor: coords.loc[micro.index]
+            for donor, micro in exp.items()
+        }
+    else:
+        exp = exp.drop(index=[0], level='label', errors='ignore') \
+                 .droplevel('label')
+        coords = coords.loc[exp.index]
+
+    return exp, coords
 
 
 def get_interpolated_map(genes, mask, n_neighbors=10, **kwargs):
diff --git a/abagen/cli/run.py b/abagen/cli/run.py
index 4108ca9..aad2410 100644
--- a/abagen/cli/run.py
+++ b/abagen/cli/run.py
@@ -353,7 +353,6 @@ def main(args=None):
     """
 
     from ..allen import get_expression_data
-    from ..datasets import WELL_KNOWN_IDS as donors
 
     opts = get_parser().parse_args(args)
 
@@ -403,13 +402,8 @@ def main(args=None):
 
     # determine how best to save expression output files
     if opts.save_donors:
-        if opts.donors == 'all':
-            donors = list(donors.value_set('subj'))
-        else:
-            donors = [donors[f] for f in opts.donors]
-
         # save each donor dataframe as a separate file
-        for donor, exp in zip(donors, expression):
+        for donor, exp in expression.items():
             exp_fname = os.path.join(output_path,
                                      fname_pref + '_{}.csv'.format(donor))
             LGR.info('Saving donor {} info to {}'.format(donor, exp_fname))
diff --git a/abagen/images.py b/abagen/images.py
index 0fae903..3c08cff 100644
--- a/abagen/images.py
+++ b/abagen/images.py
@@ -544,4 +544,8 @@ def coerce_atlas_to_dict(atlas, donors, atlas_info=None, data_dir=None):
         LGR.info('Group-level atlas provided; using MNI coords for '
                  'tissue samples')
 
+    # update group atlas status based on what was decided / derived
+    for atl in atlas.values():
+        atl.group_atlas = group_atlas
+
     return atlas, group_atlas
diff --git a/abagen/matching.py b/abagen/matching.py
index daee9d2..0992a94 100644
--- a/abagen/matching.py
+++ b/abagen/matching.py
@@ -5,6 +5,7 @@ Structures and functions used for matching samples to atlas
 
 import warnings
 
+import nibabel as nib
 import numpy as np
 import pandas as pd
 from scipy.spatial import cKDTree, distance_matrix
@@ -35,9 +36,14 @@ class AtlasTree:
         information mapping `atlas` IDs to hemisphere (i.e., "L" or "R") and
         broad structural class (i.e.., "cortex", "subcortex/brainstem",
         "cerebellum", "white matter", or "other"). Default: None
+    group_atlas : bool, optional
+        Whether the provided `atlas` is a group atlas (in MNI space) or a
+        donor-level atlas (in native space). This will have an impact on how
+        provided sample coordinates are handled. Default: True
     """
 
-    def __init__(self, atlas, coords=None, *, triangles=None, atlas_info=None):
+    def __init__(self, atlas, coords=None, *, triangles=None, atlas_info=None,
+                 group_atlas=True):
         from .images import check_img
 
         self._full_coords = self._graph = None
@@ -49,7 +55,7 @@ class AtlasTree:
                               'constructor but `coords` is not None. Ignoring '
                               'supplied `coords` and using coordinates '
                               'derived from image.')
-            self._volumetric = True
+            self._volumetric = nib.affines.voxel_sizes(affine)
             self._shape = atlas.shape
             nz = atlas.nonzero()
             atlas, coords = atlas[nz], transforms.ijk_to_xyz(np.c_[nz], affine)
@@ -62,7 +68,7 @@ class AtlasTree:
             if len(atlas) != len(coords):
                 raise ValueError('Provided `atlas` and `coords` are of '
                                  'differing length.')
-            self._volumetric = False
+            self._volumetric = None
             self._shape = atlas.shape
             nz = atlas.nonzero()
             atlas, coords = atlas[nz], coords[nz]
@@ -73,12 +79,13 @@ class AtlasTree:
         self._labels = np.unique(self.atlas).astype(int)
         self._centroids = get_centroids(self.atlas, coords)
         # if not volumetric tree then centroid should be _on_ surface
-        if not self._volumetric:
+        if self._volumetric is None:
             centroids = np.r_[list(self._centroids.values())]
             _, idx = self.tree.query(centroids, k=1)
             self._centroids = dict(zip(self.labels, self.coords[idx]))
         self.atlas_info = atlas_info
         self.triangles = triangles
+        self.group_atlas = group_atlas
 
     def __repr__(self):
         if self.volumetric:
@@ -104,7 +111,7 @@ class AtlasTree:
     def volumetric(self):
         """ Return whether `self.atlas` is derived from a volumetric image
         """
-        return self._volumetric
+        return self._volumetric is not None
 
     @property
     def labels(self):
@@ -238,6 +245,10 @@ class AtlasTree:
                 samples['structure'] = 'cortex'
 
         if self.volumetric:
+            if self.group_atlas:
+                cols = ['mni_x', 'mni_y', 'mni_z']
+                vox_size = 1 / self._volumetric
+                samples[cols] = np.floor(samples[cols] * vox_size) / vox_size
             labels = self._match_volume(samples, abs(tolerance))
         else:
             cortex = samples['structure'] == 'cortex'
@@ -313,7 +324,6 @@ class AtlasTree:
         """
 
         cols = ['mni_x', 'mni_y', 'mni_z']
-        samples[cols] = np.floor(samples[cols])
         tol, labels = 0, np.zeros(len(samples))
         idx = np.ones(len(samples), dtype=bool)
         while tol <= tolerance and np.sum(idx) > 0:
diff --git a/abagen/samples_.py b/abagen/samples_.py
index b6588bd..42164f3 100644
--- a/abagen/samples_.py
+++ b/abagen/samples_.py
@@ -465,6 +465,10 @@ def aggregate_samples(microarray, labels=None, region_agg='donors',
     LGR.info('Aggregating samples to regions with provided region_agg: {}'
              .format(region_agg))
 
+    donors = None
+    if isinstance(microarray, dict):
+        donors, microarray = list(microarray.keys()), list(microarray.values())
+
     if region_agg == 'donors':
         microarray = [
             groupby_index(e, labels=labels, metric=metric) for e in microarray
@@ -482,4 +486,7 @@ def aggregate_samples(microarray, labels=None, region_agg='donors',
                  'poor normalization'.format(drop.sum()))
         microarray = microarray.drop(drop[drop].index, axis=1)
 
+    if donors is not None and return_donors:
+        microarray = dict(zip(donors, microarray))
+
     return microarray


Test patch:
diff --git a/abagen/tests/test_correct.py b/abagen/tests/test_correct.py
index 273a513..d76fbfc 100644
--- a/abagen/tests/test_correct.py
+++ b/abagen/tests/test_correct.py
@@ -16,9 +16,9 @@ from abagen.utils import flatten_dict
 
 @pytest.fixture(scope='module')
 def donor_expression(testfiles, atlas):
-    return allen.get_expression_data(atlas['image'], missing='centroids',
-                                     return_donors=True,
-                                     donors=['12876', '15496'])
+    return list(allen.get_expression_data(atlas['image'], missing='centroids',
+                                          return_donors=True,
+                                          donors=['12876', '15496']).values())
 
 
 @pytest.fixture(scope='module')