[ENH] Allow label exclusion in remove_distance()

abagen/correct.py

@@ -48,7 +48,8 @@ def remove_distance(coexpression, atlas, atlas_info=None, labels=None):
     # load atlas_info, if provided
     atlas = check_niimg_3d(atlas)
     if atlas_info is not None:
-        atlas_info = utils.check_atlas_info(atlas, atlas_info)
+        atlas_info = utils.check_atlas_info(atlas, atlas_info,
+                                            labels_of_interest=labels)
 
     # check that provided coexpression array is symmetric
     check_symmetric(coexpression, raise_exception=True)

abagen/tests/test_correct.py

@@ -22,6 +22,15 @@ def test_remove_distance(donor_expression):
         assert np.allclose(out, out.T)
         assert isinstance(out, np.ndarray)
 
+    # subset expression data + and atlas_info
+    coexpr = np.corrcoef(expr.iloc[:-1])
+    removed_label = pd.read_csv(atlas_info).iloc[:-1]
+    out = correct.remove_distance(coexpr, ATLAS.image, removed_label,
+                                  labels=removed_label.id)
+    assert np.allclose(out, out.T)
+    assert isinstance(out, np.ndarray)
+    assert len(out) == len(removed_label)
+
     with pytest.raises(ValueError):
         correct.remove_distance(expr, ATLAS.image, ATLAS.info)
 

abagen/utils.py

@@ -14,7 +14,7 @@
 )
 
 
-def check_atlas_info(atlas, atlas_info):
+def check_atlas_info(atlas, atlas_info, labels_of_interest=None):
     """
     Checks whether provided `info` on `atlas` is sufficient for processing
 
@@ -29,6 +29,9 @@ def check_atlas_info(atlas, atlas_info):
         'structure' containing information mapping atlas IDs to hemisphere and
         broad structural class (i.e., "cortex", "subcortex", "cerebellum").
         Default: None
+    labels_of_interest : array_like, optional
+        List of values containing labels to compare between `atlas` and
+        `atlas_info`, if they don't all match. Default: None
 
     Returns
     -------
@@ -54,7 +57,11 @@ def check_atlas_info(atlas, atlas_info):
     if 'id' in atlas_info.columns:
         atlas_info = atlas_info.set_index('id')
 
-    ids = get_unique_labels(atlas)
+    if labels_of_interest is None:
+        ids = get_unique_labels(atlas)
+    else:
+        ids = labels_of_interest
+
     cols = ['hemisphere', 'structure']
     try:
         assert all(c in atlas_info.columns for c in cols)