Repository: rmarkello/abagen

Issue:
Add new RNAseq-based probe selection option
## The issue

`abagen` currently offers a number of different methods for probe selection. Probe selection is required when there are several (>1) probes indexing the same gene; there exist several ways to either select a single "representative" probe, or collapse the data amongst the probes such that we obtain a single value, per brain region, for each gene (refer to the [online documentation](https://abagen.readthedocs.io/en/stable/user_guide/probes.html) for more information).

One of the methods used—and recommended—by Arnatkevičiūte et al., 2019 is to select the probe that has the highest spatial correlation with RNAseq data for the corresponding gene. Unfortunately, `abagen` does not currently support this option (primarily since we don't [yet!] have a fetcher / loader for the RNAseq data; see #98).

## Proposed solution

**tl;dr** Add a new `probe_selection` option: `'rnaseq'`.

Selecting this option would fetch + load RNAseq data (for the two donors that have RNAseq data) and then, similar to the procedure for `probe_selection='diff_stability'`, correlate the microarray expression for each probe across anatomical regions with the RNAseq expression for the corresponding gene. The microarray probe with the highest correlation to the RNAseq data would be selected for further consideration (above a certain threshold; below this threshold the probe would simply be removed).

## Open question

Users are allowed to specify which donors they want to use via the `donors` option to `abagen.get_expression_data()`. Unfortunately, only two of the donors have RNAseq data, and it probably only makes sense to allow `probe_selection='rnaseq'` when at least one of those two donors are included in the call. This can be a pretty simple check at the beginning of the workflow that raises a `ValueError` if the required donors aren't provided.

Unfortunately, this leads to a somewhat more complicated issue: beyond fetching the relevant files at the beginning of the main workflow, we're not maintaining any information about the donors being used. This becomes necessary when we get to the actual probe selection stage, because we only want to be correlating microarray and RNAseq data from the _same donors_. That is, even if a user wants to generate microarray expression data with all six donors, we would only use the two donors with RNAseq data to perform the probe selection. Critically, this requires us to be able to match microarray data with the RNAseq data, but beyond attempting to parse filepaths this is not feasible in the current setup.

### Potential fixes

#### Option one

Modify the dictionary that is returned by `fetch_microarray()` (and related functions) such that it is a dictionary of the form:

```json
{
    "donorXXXX": {
        "microarray": "/path/to/MicroarrayExpression.csv",
        "annotation": "/path/to/Annotation.csv",
        "probes": "/path/to/Probes.csv",
        "ontology": "/path/to/Ontology.csv",
        "pacall": "/path/to/PACall.csv"
    }
}
```

Instead of the current method (where the dictionary keys are "microarray," "annotation," etc. and the values are lists of filepaths). 

The **benefit** of this method is that we would be able to use the keys of the dictionary (the donor IDs) to match the microarray and RNAseq data in the probe selection step. The **drawback** to this method is that a significant portion of the current codebase relies on this list-of-filepaths format, so modifying the structure would require refactoring a _lot_ of code.

#### Option two

We could correlate ALL the donors' microarray expression data with the RNAseq data available from the two donors, and then average across these correlations (similar to the method for `probe_selection='diff_stability'`), picking the probe with the highest average correlation.

The **benefit** of this method is that we wouldn't have to take into account which data belongs to which donor.  The **downside** of this method is that this is almost certainly not as accurate as only correlating microarray and RNAseq data within the same donor.

#### Option three

_\<Insert option here>_  Very open to suggestions and discussion about solutions to this issue of donor matching for RNAseq-based probe selection!

Discussion:

Patch:
diff --git a/abagen/allen.py b/abagen/allen.py
index 337ccad..50b3e04 100644
--- a/abagen/allen.py
+++ b/abagen/allen.py
@@ -9,6 +9,7 @@ import numpy as np
 import pandas as pd
 
 from . import correct, datasets, io, probes_, samples_, utils
+from .utils import flatten_dict
 
 import logging
 logging.basicConfig(format='%(levelname)s: %(message)s', level=logging.INFO)
@@ -89,8 +90,8 @@ def get_expression_data(atlas, atlas_info=None, *,
         Selection method for subsetting (or collapsing across) probes that
         index the same gene. Must be one of 'average', 'max_intensity',
         'max_variance', 'pc_loading', 'corr_variance', 'corr_intensity', or
-        'diff_stability'; see Notes for more information on different options.
-        Default: 'diff_stability'
+        'diff_stability', 'rnaseq'; see Notes for more information on different
+        options. Default: 'diff_stability'
     lr_mirror : bool, optional
         Whether to mirror microarray expression samples across hemispheres to
         increase spatial coverage. This will duplicate samples across both
@@ -225,6 +226,11 @@ def get_expression_data(atlas, atlas_info=None, *,
         variation across donors (i.e., the highest average correlation across
         brain regions between all pairs of donors).
 
+        8. ``method='rnaseq'``
+
+        Selects probes with most consistent pattern of regional variation to
+        RNAseq data (across the two donors with RNAseq data).
+
     The following methods can be used for normalizing microarray expression
     values prior to aggregating:
 
@@ -275,15 +281,14 @@ def get_expression_data(atlas, atlas_info=None, *,
     # check probe_selection input
     if probe_selection not in probes_.SELECTION_METHODS:
         raise ValueError('Provided probe_selection method is invalid, must be '
-                         'one of {}. Received value: \'{}\''
-                         .format(list(probes_.SELECTION_METHODS),
-                                 probe_selection))
+                         f'one of {list(probes_.SELECTION_METHODS)}. Received '
+                         f'value: \'{probe_selection}\'')
 
     # fetch files (downloading if necessary) and unpack to variables
     files = datasets.fetch_microarray(data_dir=data_dir, donors=donors,
                                       verbose=verbose, n_proc=n_proc)
 
-    if probe_selection == 'diff_stability' and len(files['microarray']) == 1:
+    if probe_selection == 'diff_stability' and len(files) == 1:
         raise ValueError('Cannot use diff_stability for probe_selection with '
                          'only one donor. Please specify a different probe_'
                          'selection method or use more donors.')
@@ -294,26 +299,27 @@ def get_expression_data(atlas, atlas_info=None, *,
     # get some info on labels in `atlas_img`
     all_labels = utils.get_unique_labels(atlas)
     if not exact:
-        lgr.info('Pre-calculating centroids for {} regions in provided atlas'
-                 .format(len(all_labels)))
+        lgr.info(f'Pre-calculating centroids for {len(all_labels)} regions in '
+                 'provided atlas')
         centroids = utils.get_centroids(atlas, labels=all_labels,
                                         image_space=True)
 
     # update the annotation "files". this handles updating the MNI coordinates,
     # dropping mistmatched samples (where MNI coordinates don't match the
     # provided ontology), and mirroring samples across hemispheres, if desired
-    annotation = files['annotation']
-    ontology = files['ontology']
-    for n, (annot, ontol) in enumerate(zip(annotation, ontology)):
+    for donor, data in files.items():
+        annot, ontol = data['annotation'], data['ontology']
         if corrected_mni:
             annot = samples_.update_mni_coords(annot)
         annot = samples_.drop_mismatch_samples(annot, ontol)
         if lr_mirror:
             annot = samples_.mirror_samples(annot, ontol)
-        annotation[n] = annot
+        data['annotation'] = annot
+    annotation = flatten_dict(files, 'annotation')
 
     # get dataframe of probe information (reannotated or otherwise)
-    probe_info = io.read_probes(files['probes'][0])
+    # the Probes.csv files are the same for every donor so just grab the first
+    probe_info = io.read_probes([files[d]['probes'] for d in files][0])
     if reannotated:
         probe_info = probes_.reannotate_probes(probe_info)
 
@@ -321,20 +327,22 @@ def get_expression_data(atlas, atlas_info=None, *,
     probe_info = probe_info.dropna(subset=['entrez_id'])
 
     # intensity-based filtering of probes
-    probe_info = probes_.filter_probes(files['pacall'], annotation, probe_info,
+    probe_info = probes_.filter_probes(flatten_dict(files, 'pacall'),
+                                       annotation, probe_info,
                                        threshold=ibf_threshold)
 
     # get probe-reduced microarray expression data for all donors based on
     # selection method; this will be a list of gene x sample dataframes (one
     # for each donor)
-    microarray = probes_.collapse_probes(files['microarray'], annotation,
-                                         probe_info, method=probe_selection)
-
+    microarray = probes_.collapse_probes(flatten_dict(files, 'microarray'),
+                                         annotation, probe_info,
+                                         method=probe_selection)
     missing = []
     counts = pd.DataFrame(np.zeros((len(all_labels) + 1, len(microarray)),
                                    dtype=int),
-                          index=np.append([0], all_labels))
-    for subj in range(len(microarray)):
+                          index=np.append([0], all_labels),
+                          columns=microarray.keys())
+    for subj in microarray:
         if lr_mirror:  # reset index (duplicates will cause issues if we don't)
             # TODO: come up with alternative sample IDs for mirrored samples
             annotation[subj] = annotation[subj].reset_index(drop=True)
@@ -360,9 +368,8 @@ def get_expression_data(atlas, atlas_info=None, *,
         # get counts of samples collapsed into each ROI
         labs, num = np.unique(labels, return_counts=True)
         counts.loc[labs, subj] = num
-        lgr.info('{:>3} / {} samples matched to regions for donor #{}'
-                 .format(counts.iloc[1:, subj].sum(), len(annotation[subj]),
-                         datasets.WELL_KNOWN_IDS.value_set('subj')[subj]))
+        lgr.info(f'{counts.iloc[1:][subj].sum():>3} / {len(annotation[subj])} '
+                 f'samples matched to regions for donor #{subj}')
 
         # if we don't want to do exact matching then cache which parcels are
         # missing data and the expression data for the closest sample to that
@@ -385,19 +392,19 @@ def get_expression_data(atlas, atlas_info=None, *,
     if not exact:  # check for missing ROIs and fill in, as needed
         # labels that are missing across all donors
         empty = reduce(set.intersection, [set(f.index) for f, d in missing])
-        lgr.info('Matching {} regions with no data to nearest samples'
-                 .format(len(empty)))
+        lgr.info(f'Matching {len(empty)} regions w/no data to nearest samples')
         for roi in empty:
             # find donor with sample closest to centroid of empty parcel
             ind = np.argmin([dist.get(roi) for micro, dist in missing])
-            donor = datasets.WELL_KNOWN_IDS.value_set("subj")[ind]
-            lgr.debug('Assigning sample from donor {} to region id #{}'
-                      .format(donor, roi))
+            subj = list(microarray.keys())[ind]
+            lgr.debug(f'Assigning sample from donor {subj} to region #{roi}')
             # assign expression data from that sample and add to count
-            microarray[ind] = microarray[ind].append(missing[ind][0].loc[roi])
-            counts.loc[roi, ind] += 1
+            exp = missing[ind][0].loc[roi]
+            microarray[subj] = microarray[subj].append(exp)
+            counts.loc[roi, subj] += 1
 
-    microarray = samples_.aggregate_samples(microarray, labels=all_labels,
+    microarray = samples_.aggregate_samples(microarray.values(),
+                                            labels=all_labels,
                                             region_agg=region_agg,
                                             agg_metric=agg_metric,
                                             return_donors=return_donors)
diff --git a/abagen/correct.py b/abagen/correct.py
index 0d899a1..4e42b0e 100644
--- a/abagen/correct.py
+++ b/abagen/correct.py
@@ -387,8 +387,8 @@ def normalize_expression(expression, norm='srs', ignore_warn=False):
         normfunc = NORMALIZATION_METHODS[norm]
     except KeyError:
         raise ValueError('Provided value for `norm` not recognized. Must be '
-                         'one of {}. Received: {}'
-                         .format(list(NORMALIZATION_METHODS), norm))
+                         f'one of {list(NORMALIZATION_METHODS)}. Received: '
+                         f'{norm}')
 
     # FIXME: I hate having to do this...
     if isinstance(expression, pd.DataFrame):
@@ -457,10 +457,10 @@ def remove_distance(coexpression, atlas, atlas_info=None, labels=None):
     if atlas_info is not None:
         atlas_info = utils.check_atlas_info(atlas, atlas_info, labels=labels)
         if labels is not None and len(labels) != len(coexpression):
-            raise ValueError('Provided labels {} are a different length than '
-                             'provided coexpression matrix of size {}. Please '
-                             'confirm inputs and try again.'
-                             .format(labels, coexpression.shape))
+            raise ValueError(f'Provided labels {labels} are a different '
+                             'length than provided coexpression matrix of '
+                             f'size {coexpression.shape}. Please confirm '
+                             'inputs and try again.')
 
     # check that provided coexpression array is symmetric
     if not np.allclose(coexpression, coexpression.T, atol=1e-10):
diff --git a/abagen/datasets/fetchers.py b/abagen/datasets/fetchers.py
index f0d3cb3..661b8b5 100644
--- a/abagen/datasets/fetchers.py
+++ b/abagen/datasets/fetchers.py
@@ -12,6 +12,7 @@ import pandas as pd
 
 from .. import io
 from .utils import _get_dataset_dir, _fetch_files
+
 WELL_KNOWN_IDS = Recoder(
     (('9861', 'H0351.2001', '178238387', '157722636', '157722638'),
      ('10021', 'H0351.2002', '178238373', '157723301', '157723303'),
@@ -21,11 +22,49 @@ WELL_KNOWN_IDS = Recoder(
      ('15697', 'H0351.1016', '178236545', '157682966', '157682968')),
     fields=('subj', 'uid', 'url', 't1w', 't2w')
 )
-
 VALID_DONORS = sorted(WELL_KNOWN_IDS.value_set('subj')
                       | WELL_KNOWN_IDS.value_set('uid'))
 
 
+def check_donors(donors, default='12876', valid=VALID_DONORS):
+    """
+    Checks that provided `donors` are valid
+
+    Parameters
+    ----------
+    donors : list of str
+        List of donors to download; can be either donor number or UID. Can also
+        specify 'all' to download all available donors. If 'None' is provided
+        then `default` will be used.
+    default : str, optional
+        Default donor to use if `donors` is None. Default: '12876'
+    valid : list of str, optional
+        List of valid donnor numbers and UIDs. Default: :obj:`VALID_DONORS`
+
+    Returns
+    -------
+    donors : list of str
+        Donor subject IDs
+    """
+
+    if donors is None:
+        donors = [default]
+    elif donors == 'all':
+        donors = valid
+    elif isinstance(donors, str):
+        donors = [donors]
+
+    donors = list(donors).copy()
+    for n, sub_id in enumerate(donors):
+        if sub_id not in valid:
+            raise ValueError('Invalid subject id: {0}. Subjects must in: {1}.'
+                             .format(sub_id, valid))
+        donors[n] = WELL_KNOWN_IDS[sub_id]  # convert to ID system
+    donors = sorted(set(donors), key=lambda x: int(x))
+
+    return donors
+
+
 def fetch_microarray(data_dir=None, donors=None, resume=True, verbose=1,
                      convert=True, n_proc=1):
     """
@@ -109,13 +148,11 @@ def fetch_microarray(data_dir=None, donors=None, resume=True, verbose=1,
         for fn in files[0::n_files] + files[2::n_files]:
             io._make_parquet(fn, convert_only=True)
 
-    return dict(
-        microarray=files[0::n_files],
-        ontology=files[1::n_files],
-        pacall=files[2::n_files],
-        probes=files[3::n_files],
-        annotation=files[4::n_files]
-    )
+    keys = ['microarray', 'ontology', 'pacall', 'probes', 'annotation']
+    return {
+        donor: dict(zip(keys, files[k:k + n_files]))
+        for k, donor in zip(range(0, len(files), n_files), donors)
+    }
 
 
 def fetch_rnaseq(data_dir=None, donors=None, resume=True, verbose=1):
@@ -159,11 +196,12 @@ def fetch_rnaseq(data_dir=None, donors=None, resume=True, verbose=1):
     data_dir = _get_dataset_dir(dataset_name, data_dir=data_dir,
                                 verbose=verbose)
 
-    sub_files = ('Contents.txt', 'Genes.csv', 'Ontology.csv',
+    sub_files = ('Genes.csv', 'Ontology.csv',
                  'RNAseqCounts.csv', 'RNAseqTPM.csv', 'SampleAnnot.csv')
     n_files = len(sub_files)
     valid = ['9861', '10021', 'H0351.2001', 'H0351.2002']
-    donors = check_donors(donors, default=valid[0], valid=valid)
+    donors = sorted(set(check_donors(donors, default=valid[0])) & set(valid),
+                    key=lambda x: int(x))
 
     files = [
         [(os.path.join('rnaseq_donor{}'.format(sub), fname),
@@ -178,13 +216,11 @@ def fetch_rnaseq(data_dir=None, donors=None, resume=True, verbose=1):
     files = [l for f in files for l in f]
     files = _fetch_files(data_dir, files, resume=resume, verbose=verbose)
 
-    return dict(
-        genes=files[1::n_files],
-        ontology=files[2::n_files],
-        counts=files[3::n_files],
-        tpm=files[4::n_files],
-        annotation=files[5::n_files]
-    )
+    keys = ['genes', 'ontology', 'counts', 'tpm', 'annotation']
+    return {
+        donor: dict(zip(keys, files[k:k + n_files]))
+        for k, donor in zip(range(0, len(files), n_files), donors)
+    }
 
 
 def fetch_raw_mri(data_dir=None, donors=None, resume=True, verbose=1):
@@ -232,49 +268,10 @@ def fetch_raw_mri(data_dir=None, donors=None, resume=True, verbose=1):
 
     files = _fetch_files(data_dir, files, resume=resume, verbose=verbose)
 
-    return dict(
-        t1w=files[0::n_files],
-        t2w=files[1::n_files],
-    )
-
-
-def check_donors(donors, default='12876', valid=VALID_DONORS):
-    """
-    Checks that provided `donors` are valid
-
-    Parameters
-    ----------
-    donors : list of str
-        List of donors to download; can be either donor number or UID. Can also
-        specify 'all' to download all available donors. If 'None' is provided
-        then `default` will be used.
-    default : str, optional
-        Default donor to use if `donors` is None. Default: '12876'
-    valid : list of str, optional
-        List of valid donnor numbers and UIDs. Default: :obj:`VALID_DONORS`
-
-    Returns
-    -------
-    donors : list of str
-        Donor subject IDs
-    """
-
-    if donors is None:
-        donors = [default]
-    elif donors == 'all':
-        donors = valid
-    elif isinstance(donors, str):
-        donors = [donors]
-
-    donors = list(donors)
-    for n, sub_id in enumerate(donors):
-        if sub_id not in valid:
-            raise ValueError('Invalid subject id: {0}. Subjects must in: {1}.'
-                             .format(sub_id, valid))
-        donors[n] = WELL_KNOWN_IDS[sub_id]  # convert to ID system
-    donors = sorted(set(donors), key=lambda x: int(x))
-
-    return donors
+    return {
+        donor: dict(zip(sub_files.keys(), files[k:k + n_files]))
+        for k, donor in zip(range(0, len(files), n_files), donors)
+    }
 
 
 def fetch_freesurfer(data_dir=None, donors=None, resume=True, verbose=1):
@@ -310,24 +307,25 @@ def fetch_freesurfer(data_dir=None, donors=None, resume=True, verbose=1):
 
     url = "https://www.repository.cam.ac.uk/bitstream/handle/1810/265272/" \
           "donor{}.zip"
-    dataset_name = 'allenbrain'
+    dataset_name = 'freesurfer'
     data_dir = _get_dataset_dir(dataset_name, data_dir=data_dir,
                                 verbose=verbose)
 
     donors = check_donors(donors)
     files = [
-        (os.path.join('normalized_microarray_donor{}'.format(sub),
-                      'donor{}'.format(sub)),
+        ('donor{}'.format(sub),
          url.format(sub),
          dict(uncompress=True,
-              move=os.path.join('normalized_microarray_donor{}'.format(sub),
-                                'freesurfer.tar.gz')))
+              move=os.path.join('freesurfer.tar.gz')))
         for sub in donors
     ]
 
     files = _fetch_files(data_dir, files, resume=resume, verbose=verbose)
 
-    return files
+    return {
+        donor: files[k]
+        for k, donor in enumerate(donors)
+    }
 
 
 def fetch_desikan_killiany(*args, **kwargs):
diff --git a/abagen/io.py b/abagen/io.py
index 22b797e..e06e6f2 100644
--- a/abagen/io.py
+++ b/abagen/io.py
@@ -282,12 +282,16 @@ def read_annotation(fname, copy=False):
     `abagen.fetch_rnaseq()`), then the returned DataFrame will have the
     following columns: 'RNAseq_sample_name', 'replicate_sample', 'sample_name',
     'well_id', 'microarray_run_id', 'ontology_color', 'main_structure',
-    'sub_structure', 'ontology_structure_id', 'ontology_structure_acronym',
-    'hemisphere', 'brain', 'million_clusters', 'clip_percentage',
-    'RIN_RNA_squality', 'rnaseq_run_id', 'A.Pct', 'C.Pct', 'G.Pct', 'T.Pct',
-    'N.Pct'
+    'sub_structure', 'structure_id', 'structure_acronym', 'hemisphere',
+    'brain', 'million_clusters', 'clip_percentage', 'RIN_RNA_squality',
+    'rnaseq_run_id', 'A.Pct', 'C.Pct', 'G.Pct', 'T.Pct', 'N.Pct'
     """
 
+    mapper = dict(
+        ontology_structure_id='structure_id',
+        ontology_structure_acronym='structure_acronym'
+    )
+
     try:
         data = pd.read_csv(fname)
         data.index.name = 'sample_id'
@@ -296,6 +300,7 @@ def read_annotation(fname, copy=False):
             raise TypeError('Provided fname must be filepath to Annotation'
                             '.csv file from Allen Human Brain Atlas.')
         data = fname.copy() if copy else fname
+        data.rename(mapper, axis='columns', inplace=True, errors='ignore')
 
     return data
 
diff --git a/abagen/probes_.py b/abagen/probes_.py
index 25a62b9..cb1ca9b 100644
--- a/abagen/probes_.py
+++ b/abagen/probes_.py
@@ -13,7 +13,7 @@ from pkg_resources import resource_filename
 import numpy as np
 import pandas as pd
 
-from . import io, utils
+from . import datasets, io, utils
 
 lgr = logging.getLogger('abagen')
 
@@ -67,7 +67,7 @@ def reannotate_probes(probes):
     return reannotated
 
 
-def filter_probes(pacall, samples, probes, threshold=0.5):
+def filter_probes(pacall, annotation, probes, threshold=0.5):
     """
     Performs intensity based filtering (IBF) of expression probes
 
@@ -77,15 +77,16 @@ def filter_probes(pacall, samples, probes, threshold=0.5):
 
     Parameters
     ----------
-    pacall : list of str
-        List of filepaths to PAcall files from Allen Brain Institute (i.e.,
-        as obtained by calling :func:`abagen.fetch_microarray` and accessing
-        the `pacall` attribute on the resulting object).
-    samples : list of str or pandas.DataFrame
-
+    pacall : dict
+        Dictionary where keys are donor IDs and values are filepaths to (or
+        dataframes of) PACall.csv files from Allen Brain Institute
+    annotation : dict
+        Dictionary where keys are donor IDs and values are filepaths to (or
+        dataframes of) SampleAnnot.csv files from Allen Brain Institute
     probes : str or pandas.DataFrame
-        Dataframe containing information on microarray probes that should be
-        considered in filtering (probes not in this dataframe will be ignored)
+        Filepath to Probes.csv or dataframe containing information on
+        microarray probes that should be considered in filtering (probes not in
+        this will be ignored)
     threshold : (0, 1) float, optional
         Threshold for filtering probes. Specifies the proportion of samples for
         which a given probe must have expression levels above background noise.
@@ -100,14 +101,13 @@ def filter_probes(pacall, samples, probes, threshold=0.5):
 
     threshold = np.clip(threshold, 0.0, 1.0)
 
-    lgr.info('Filtering probes with intensity-based threshold of {}'
-             .format(threshold))
+    lgr.info(f'Filtering probes with intensity-based threshold of {threshold}')
 
     probes = io.read_probes(probes)
     signal, n_samp = np.zeros(len(probes), dtype=int), 0
-    for pa, annot in zip(pacall, samples):
-        data = io.read_pacall(pa).loc[probes.index,
-                                      io.read_annotation(annot).index]
+    for donor, pa in pacall.items():
+        annot = io.read_annotation(annotation[donor]).index
+        data = io.read_pacall(pa).loc[probes.index, annot]
         n_samp += data.shape[-1]
         # sum binary expression indicator across samples for current subject
         signal += np.asarray(data.sum(axis=1))
@@ -115,21 +115,46 @@ def filter_probes(pacall, samples, probes, threshold=0.5):
     # calculate proportion of signal to noise for given probe across samples
     keep = (signal / n_samp) >= threshold
 
-    lgr.info('{} probes survive intensity-based filtering'.format(keep.sum()))
+    lgr.info(f'{keep.sum()} probes survive intensity-based filtering')
 
     return probes[keep]
 
 
-def _groupby_and_apply(expression, probes, info, applyfunc):
+def _groupby_structure_id(microarray, annotation):
     """
-    Subsets `expression` based on most representative probe
+    Averages samples in `microarray` having identical structure IDs
 
     Parameters
     ----------
-    expression : list of (P, S) pandas.DataFrame
-        Each dataframe should have `P` rows representing probes and `S` columns
+    microarray : (P, S) pandas.DataFrame
+        Dataframe should have `P` rows representing probes and `S` columns
         representing distinct samples, with values indicating microarray
         expression levels
+    annotation : (S, A) pandas.DataFrame
+        Annotation dataframe, obtained by loading a SampleAnnot.csv file from
+        Allen Brain Institute
+
+    Returns
+    -------
+    expression : (P, R) pandas.DataFrame
+        Input `microarray` dataframe but with `S` samples averaged into `R`
+        regions
+    """
+
+    sid = io.read_annotation(annotation)['structure_id']
+
+    return io.read_microarray(microarray).groupby(sid, axis=1).mean()
+
+
+def _groupby_and_apply(expression, probes, info, applyfunc):
+    """
+    Subsets `expression` based on most representative probe
+
+    Parameters
+    ----------
+    expression : dict of (P, S) pandas.DataFrame
+        Dictionary where keys are donor IDs and values are dataframes with `P`
+        rows representing probes and `S` columns representing distinct samples
     probes : pandas.DataFrame
         Dataframe containing information on probes that should be considered in
         representative analysis. Generally, intensity-based-filtering (i.e.,
@@ -145,26 +170,153 @@ def _groupby_and_apply(expression, probes, info, applyfunc):
 
     Returns
     -------
-    representative : list of (S, G) pandas.DataFrame
-        Dataframes with rows representing `S` distinct samples and columns
-        representing `G` unique genes, with values indicating microarray
-        expression levels for each combination of samples and genes
+    representative : dict of (S, G) pandas.DataFrame
+        Dictionary where keys are donor IDs and values are dataframes with `S`
+        rows representing distinct samples and `G` columns representing unique
+        genes
     """
 
-    # group probes by gene and get probe corresponding to max avg correlationx
-    retained = info.groupby('gene_symbol').apply(applyfunc)
-    probes = probes.loc[sorted(np.squeeze(retained).astype(int))]
+    # group probes by gene and get probe corresponding to relevant feature
+    retained = info.groupby('gene_symbol').apply(applyfunc).dropna()
+    probes = probes.loc[sorted(np.squeeze(retained.astype(int)))]
 
     # subset expression dataframes to retain only desired probes and reassign
     # (and sort) index to gene symbols in lieu of probe IDs
-    representative = [
-        e.loc[probes.index].set_index(probes['gene_symbol']).sort_index()
-        for e in expression
-    ]
+    representative = {
+        d: e.loc[probes.index].set_index(probes['gene_symbol']).sort_index().T
+        for d, e in utils.check_dict(expression).items()
+    }
 
     return representative
 
 
+def _diff_stability(expression, probes, annotation, *args, **kwargs):
+    """
+    Picks one probe to represent `expression` data for each gene in `probes`
+
+    If there are multiple probes with expression data for the same gene, this
+    function will calculate the similarity of each probes' expression across
+    donors and select the probe with the most consistent pattern of regional
+    variation (i.e., "differential stability" or DS). Regions are defined by
+    the "structure_id" column in `annotation`; similarity is calculated by the
+    Spearman correlation coefficient (but see `rank` kwarg).
+
+    Parameters
+    ----------
+    expression : dict of (P, S) pandas.DataFrame
+        Dictionary where keys are donor IDs and values are dataframes with `P`
+        rows representing probes and `S` columns representing distinct samples
+    probes : pandas.DataFrame
+        Dataframe containing information on probes that should be considered in
+        representative analysis. Generally, intensity-based-filtering (i.e.,
+        `filter_probes()`) should have been used to reduce this list to only
+        those probes with good expression signal
+    annotation : list of str
+        List of filepaths to annotation files from Allen Brain Institute (i.e.,
+        as obtained by calling :func:`abagen.fetch_microarray` and accessing
+        the `annotation` attribute on the resulting object).
+
+    Returns
+    -------
+    representative : dict of (S, G) pandas.DataFrame
+        Dictionary where keys are donor IDs and values are dataframes with `S`
+        rows representing distinct samples and `G` columns representing unique
+        genes
+    """
+
+    # confirm inputs are expected dictionaries
+    expression = utils.check_dict(expression)
+    annotation = utils.check_dict(annotation)
+
+    # collapse (i.e., average) expression across AHBA anatomical regions
+    region_exp = [
+        _groupby_structure_id(microarray, annotation[donor])
+        for donor, microarray in expression.items()
+    ]
+
+    # get correlation of probe expression across samples for all donor pairs
+    probe_exp = np.zeros((len(probes), sum(range(len(expression)))))
+    for n, (exp1, exp2) in enumerate(itertools.combinations(region_exp, 2)):
+        # samples that current donor pair have in common
+        samples = np.intersect1d(exp1.columns, exp2.columns)
+
+        # the ranking process can take a few seconds on each loop
+        # unfortunately, we have to do it each time because `samples` changes
+        # based on which anatomical regions the two subjects have in common
+        exp1 = exp1.loc[:, samples].T.rank()
+        exp2 = exp2.loc[:, samples].T.rank()
+
+        probe_exp[:, n] = utils.efficient_corr(exp1, exp2)
+
+    info = pd.DataFrame(dict(gene_symbol=np.asarray(probes.gene_symbol),
+                             diff_stability=probe_exp.mean(axis=1)),
+                        index=probes.index)
+    applyfunc = functools.partial(_max_idx, column='diff_stability')
+
+    return _groupby_and_apply(expression, probes, info, applyfunc)
+
+
+def _rnaseq(expression, probes, annotation, *args, **kwargs):
+    """
+
+    Parameters
+    ----------
+    expression : dict of (P, S) pandas.DataFrame
+        Dictionary where keys are donor IDs and values are dataframes with `P`
+        rows representing probes and `S` columns representing distinct samples
+    probes : pandas.DataFrame
+        Dataframe containing information on probes that should be considered in
+        representative analysis. Generally, intensity-based-filtering (i.e.,
+        `filter_probes()`) should have been used to reduce this list to only
+        those probes with good expression signal
+    annotation : list of str
+        List of filepaths to annotation files from Allen Brain Institute (i.e.,
+        as obtained by calling :func:`abagen.fetch_microarray` and accessing
+        the `annotation` attribute on the resulting object).
+
+    Returns
+    -------
+    representative : dict of (S, G) pandas.DataFrame
+        Dictionary where keys are donor IDs and values are dataframes with `S`
+        rows representing distinct samples and `G` columns representing unique
+        genes
+    """
+    # confirm inputs are expected dictionaries
+    expression = utils.check_dict(expression)
+    annotation = utils.check_dict(annotation)
+
+    # fetch RNAseq data
+    rnaseq = datasets.fetch_rnaseq(donors=expression.keys())
+
+    probe_exp = np.ones((len(probes), len(rnaseq))) * np.nan
+    for n, (donor, data) in enumerate(rnaseq.items()):
+        # collapse (i.e., average) data across AHAB  anatomical regions
+        micro = _groupby_structure_id(expression[donor], annotation[donor])
+        rna = _groupby_structure_id(io.read_tpm(data['tpm']),
+                                    data['annotation'])
+
+        # get rid of "constant" RNAseq genes
+        rna = rna[np.logical_not(np.isclose(rna.std(axis=1, ddof=1), 0))]
+
+        # get matching genes + strcutres between microarray + RNAseq
+        regions = np.intersect1d(micro.columns, rna.columns)
+        mask = np.isin(np.asarray(probes.gene_symbol),
+                       np.intersect1d(probes.gene_symbol, rna.index))
+        genes = np.asarray(probes.loc[mask, 'gene_symbol'])
+        micro, rna = micro.loc[mask, regions].T, rna.loc[genes, regions].T
+
+        # correlate expression values across regions for each gene
+        probe_exp[mask, n] = utils.efficient_corr(micro.rank(), rna.rank())
+
+    mask = np.sum(np.isnan(probe_exp), axis=1) < len(rnaseq)
+    info = pd.DataFrame(dict(gene_symbol=np.asarray(probes[mask].gene_symbol),
+                             rna_corr=np.nanmean(probe_exp[mask], axis=1)),
+                        index=probes.index[mask])
+    applyfunc = functools.partial(_max_idx, column='rna_corr')
+
+    return _groupby_and_apply(expression, probes, info, applyfunc)
+
+
 def _max_idx(df, column=0):
     """
     Returns probe ID with max index in `df`
@@ -253,123 +405,35 @@ def _correlate(df, method):
     return df.index[xmax.argmax()]
 
 
-def _diff_stability(expression, probes, annotation, *args, **kwargs):
-    """
-    Picks one probe to represent `expression` data for each gene in `probes`
-
-    If there are multiple probes with expression data for the same gene, this
-    function will calculate the similarity of each probes' expression across
-    donors and select the probe with the most consistent pattern of regional
-    variation (i.e., "differential stability" or DS). Regions are defined by
-    the "structure_id" column in `annotation`; similarity is calculated by the
-    Spearman correlation coefficient (but see `rank` kwarg).
-
-    Parameters
-    ----------
-    expression : list of (P, S) pandas.DataFrame
-        Each dataframe should have `P` rows representing probes and `S` columns
-        representing distinct samples, with values indicating microarray
-        expression levels
-    probes : pandas.DataFrame
-        Dataframe containing information on probes that should be considered in
-        representative analysis. Generally, intensity-based-filtering (i.e.,
-        `filter_probes()`) should have been used to reduce this list to only
-        those probes with good expression signal
-    annotation : list of str
-        List of filepaths to annotation files from Allen Brain Institute (i.e.,
-        as obtained by calling :func:`abagen.fetch_microarray` and accessing
-        the `annotation` attribute on the resulting object).
-
-    Returns
-    -------
-    representative : list of (S, G) pandas.DataFrame
-        Dataframes with rows representing `S` distinct samples and columns
-        representing `G` unique genes, with values indicating microarray
-        expression levels for each combination of samples and genes
-    """
-
-    # collapse (i.e., average) expression across AHBA anatomical regions
-    region_exp = [_groupby_structure_id(exp, annot)
-                  for exp, annot in zip(expression, annotation)]
-
-    # get correlation of probe expression across samples for all donor pairs
-    probe_exp = np.zeros((len(probes), sum(range(len(expression)))))
-    for n, (exp1, exp2) in enumerate(itertools.combinations(region_exp, 2)):
-        # samples that current donor pair have in common
-        samples = np.intersect1d(exp1.columns, exp2.columns)
-
-        # the ranking process can take a few seconds on each loop
-        # unfortunately, we have to do it each time because `samples` changes
-        # based on which anatomical regions the two subjects have in common
-        exp1 = exp1.loc[:, samples].T.rank()
-        exp2 = exp2.loc[:, samples].T.rank()
-
-        probe_exp[:, n] = utils.efficient_corr(exp1, exp2)
-
-    info = pd.DataFrame(dict(gene_symbol=np.asarray(probes.gene_symbol),
-                             diff_stability=probe_exp.mean(axis=1)),
-                        index=probes.index)
-    applyfunc = functools.partial(_max_idx, column='diff_stability')
-
-    return _groupby_and_apply(expression, probes, info, applyfunc)
-
-
-def _groupby_structure_id(microarray, annotation):
-    """
-    Averages samples in `microarray` having identical structure IDs
-
-    Parameters
-    ----------
-    microarray : pandas.DataFrame
-        Dataframe should have `P` rows representing probes and `S` columns
-        representing distinct samples, with values indicating microarray
-        expression levels
-    annotation : pandas.DataFrame
-        Annotation dataframe, obtained by loading an annotation file from Allen
-        Brain Institute
-
-    Returns
-    -------
-    expression : pandas.DataFrame
-
-    """
-    sid = io.read_annotation(annotation)['structure_id']
-    microarray = io.read_microarray(microarray)
-    microarray = microarray.rename(dict(zip(microarray.columns, sid)), axis=1)
-
-    return microarray.groupby(microarray.columns, axis=1).mean()
-
-
 def _average(expression, probes, *args, **kwargs):
     """
     Averages expression data for probes representing the same gene
 
     Parameters
     ----------
-    expression : list of (P, S) pandas.DataFrame
-        Each dataframe should have `P` rows representing probes and `S` columns
-        representing distinct samples, with values indicating microarray
-        expression levels
+    expression : dict of (P, S) pandas.DataFrame
+        Dictionary where keys are donor IDs and values are dataframes with `P`
+        rows representing probes and `S` columns representing distinct samples
     probes : pandas.DataFrame
-        Dataframe containing information on probes that should be considered in
-        representative analysis. Generally, intensity-based-filtering (i.e.,
-        `filter_probes()`) should have been used to reduce this list to only
-        those probes with good expression signal
+        Dataframe containing information on microarray probes that should be
+        considered in representative analysis. Generally intensity-based
+        filtering (i.e., via :func:`filter_probes()`) should have been used to
+        reduce this list to only those probes with good expression signal
 
     Returns
     -------
-    representative : list of (S, G) pandas.DataFrame
-        Dataframes with rows representing `S` distinct samples and columns
-        representing `G` unique genes, with values indicating microarray
-        expression levels for each combination of samples and genes
+    representative : dict of (S, G) pandas.DataFrame
+        Dictionary where keys are donor IDs and values are dataframes with `S`
+        rows representing distinct samples and `G` columns representing unique
+        genes
     """
 
     def _avg(df):
         return df.join(probes['gene_symbol'], on='probe_id') \
                  .groupby('gene_symbol') \
-                 .mean()
+                 .mean().T
 
-    return [_avg(exp) for exp in expression]
+    return {d: _avg(exp) for d, exp in utils.check_dict(expression).items()}
 
 
 def _collapse(expression, probes, *args, method='max_variance', **kwargs):
@@ -378,15 +442,14 @@ def _collapse(expression, probes, *args, method='max_variance', **kwargs):
 
     Parameters
     ----------
-    expression : list of (P, S) pandas.DataFrame
-        Each dataframe should have `P` rows representing probes and `S` columns
-        representing distinct samples, with values indicating microarray
-        expression levels
+    expression : dict of (P, S) pandas.DataFrame
+        Dictionary where keys are donor IDs and values are dataframes with `P`
+        rows representing probes and `S` columns representing distinct samples
     probes : pandas.DataFrame
-        Dataframe containing information on probes that should be considered in
-        representative analysis. Generally, intensity-based-filtering (i.e.,
-        `filter_probes()`) should have been used to reduce this list to only
-        those probes with good expression signal
+        Dataframe containing information on microarray probes that should be
+        considered in representative analysis. Generally intensity-based
+        filtering (i.e., via :func:`filter_probes()`) should have been used to
+        reduce this list to only those probes with good expression signal
     method : str, optional
         Method by which to select represenative probes for each gene. Must be
         one of ['max_variance', 'max_intensity', 'pc_loading', 'corr_variance',
@@ -394,14 +457,15 @@ def _collapse(expression, probes, *args, method='max_variance', **kwargs):
 
     Returns
     -------
-    representative : list of (S, G) pandas.DataFrame
-        Dataframes with rows representing `S` distinct samples and columns
-        representing `G` unique genes, with values indicating microarray
-        expression levels for each combination of samples and genes
+    representative : dict of (S, G) pandas.DataFrame
+        Dictionary where keys are donor IDs and values are dataframes with `S`
+        rows representing distinct samples and `G` columns representing unique
+        genes
     """
 
     # concatenate all donors into giant probe x sample expression dataframe
-    probe_exp = pd.concat(expression, axis=1)
+    expression = utils.check_dict(expression)
+    probe_exp = pd.concat(expression.values(), axis=1)
 
     # determine aggregation function based on provided method; also reduce
     # probe expression if required (i.e., max_variance, max_intensity)
@@ -417,8 +481,8 @@ def _collapse(expression, probes, *args, method='max_variance', **kwargs):
     elif method in ['corr_variance', 'corr_intensity']:
         agg = functools.partial(_correlate, method=method[5:])
     else:
-        raise ValueError('Provided method {} invalid. Please check inputs '
-                         'and try again.'.format(method))
+        raise ValueError(f'Provided method {method} is invalid. Please check '
+                         'inputs and try again.')
 
     info = pd.merge(probes[['gene_symbol']], probe_exp, on='probe_id')
     return _groupby_and_apply(expression, probes, info, agg)
@@ -439,7 +503,8 @@ SELECTION_METHODS = dict(
     pc_loading=_pc_loading,
     corr_variance=_corr_variance,
     corr_intensity=_corr_intensity,
-    diff_stability=_diff_stability
+    diff_stability=_diff_stability,
+    rnaseq=_rnaseq
 )
 
 
@@ -454,33 +519,34 @@ def collapse_probes(microarray, annotation, probes, method='diff_stability'):
 
     Parameters
     ----------
-    microarray : list of str
-        List of filepaths to microarray expression files from Allen Brain
-        Institute (i.e., as obtained by calling :func:`abagen.fetch_microarray`
-        and accessing the `microarray` attribute on the resulting object).
-    annotation : list of str
-        List of filepaths to annotation files from Allen Brain Institute (i.e.,
-        as obtained by calling :func:`abagen.fetch_microarray` and accessing
-        the `annotation` attribute on the resulting object). Only used if
-        provided `method` is 'diff_stability'
-    probes : pandas.DataFrame
-        Dataframe containing information on probes that should be considered in
-        representative analysis. Generally, intensity-based-filtering (i.e.,
-        `filter_probes()`) should have been used to reduce this list to only
-        those probes with good expression signal
+    microarray : dict of str or pandas.DataFrame
+        Dictionary where keys are donor IDs and values are filepaths to (or
+        dataframes of) MicroarrayExpression.csv files from Allen Brain
+        Institute
+    annotation : dict of str or pandas.DataFrame
+        Dictionary where keys are donor IDs and values are filepaths to (or
+        dataframes of) SampleAnnot.csv files from Allen Brain Institute. Only
+        used if `method='diff_stability'`
+    probes : str or pandas.DataFrame
+        Filepath to Probes.csv or dataframe containing information on
+        microarray probes that should be considered in representative analysis.
+        Generally intensity-based filtering (i.e., via :func:`filter_probes()`)
+        should have been used to reduce this list to only those probes with
+        good expression signal
     method : str, optional
         Selection method for subsetting (or collapsing across) probes from the
-        same gene. Must be one of 'average', 'intensity', 'variance',
-        'pc_loading', 'corr_variance', 'corr_intensity', or 'diff_stability';
-        see Notes for more information. Default: 'diff_stability'
+        same gene. Must be one of 'average', 'max_intensity', 'max_variance',
+        'pc_loading', 'corr_intensity', 'corr_variance', 'diff_stability', or
+        'rnaseq'; see Notes for more information. Default: 'diff_stability'
 
     Returns
     -------
-    expression : list of (S, G) pandas.DataFrame
-        Dataframes with rows representing `S` distinct samples and columns
-        representing `G` unique genes, with values indicating microarray
-        expression levels for each combination of samples and genes. Columns
-        will be identical across all dataframes but `S` will vary by donor.
+    expression : dict of (S, G) pandas.DataFrame
+        Dictionary where keys are donor IDs and values are dataframes with `S`
+        rows representing distinct samples and `G` columns representing unique
+        genes. Entries of dataframe indicate microarray expression levels for
+        each combination of sample + gene. Columns will be identical across all
+        dataframes, but `S` will vary by donor.
 
     Notes
     -----
@@ -526,6 +592,11 @@ def collapse_probes(microarray, annotation, probes, method='diff_stability'):
     donors (i.e., highest average correlation across brain regions between all
     pairs of donors) as in [PR7]_ and [PR8]_.
 
+    8. ``method='rnaseq'``
+
+    Selects probes with most consistent pattern of regional variation to RNAseq
+    data (across the two donors with RNAseq data).
+
     References
     ----------
     .. [PR1] Anderson, K. M., Krienen, F. M., Choi, E. Y., Reinen, J. M., Yeo,
@@ -615,26 +686,24 @@ def collapse_probes(microarray, annotation, probes, method='diff_stability'):
     try:
         collfunc = SELECTION_METHODS[method]
     except KeyError:
-        raise ValueError('Provided method must be one of {}, not: {}'
-                         .format(list(SELECTION_METHODS), method))
+        raise ValueError('Provided method must be one of '
+                         f'{list(SELECTION_METHODS)}, not: {method}')
 
-    lgr.info('Reducing probes indexing same gene with method: "{}"'
-             .format(method))
+    lgr.info(f'Reducing probes indexing same gene with method: "{method}"')
 
-    # read in microarray data for all subjects; this can be quite slow...
+    # subset microarray data for pre-selected probes + samples
+    # this will also left/right mirror samples, if previously requested
     probes = io.read_probes(probes)
-    exp = [
-        # subset microarray data for pre-selected probes + samples
-        # this will also left/right mirror samples, if previously requested
-        io.read_microarray(micro).loc[probes.index,
-                                      io.read_annotation(annot).index]
-        for micro, annot in zip(microarray, annotation)
-    ]
-    exp = [
-        e.T for e in collfunc(exp, probes, annotation)
-    ]
+    microarray = utils.check_dict(microarray)
+    annotation = utils.check_dict(annotation)
+    for donor, micro in microarray.items():
+        samp = io.read_annotation(annotation[donor]).index
+        microarray[donor] = io.read_microarray(micro).loc[probes.index, samp]
+
+    # now, "collect" the probes based on the provided `method`
+    expression = collfunc(microarray, probes, annotation)
+    n_genes = utils.first_entry(expression).shape[-1]
 
-    lgr.info('{} genes remain after probe filtering and selection'
-             .format(exp[0].shape[-1]))
+    lgr.info(f'{n_genes} genes remain after probe filtering + selection')
 
-    return exp
+    return expression
diff --git a/abagen/samples_.py b/abagen/samples_.py
index 9c49aa7..a9cd5ff 100644
--- a/abagen/samples_.py
+++ b/abagen/samples_.py
@@ -471,8 +471,6 @@ def groupby_index(microarray, labels=None, metric='mean'):
     ----------
     microarray : (S, G) pandas.DataFrame
         Microarray expression data, where `S` is samples and `G` is genes
-    sample_labels : (S, 1) pandas.DataFrame
-        Parcel labels for `S` samples, as returned by e.g., `label_samples()`
     labels : (L,) array_like, optional
         All possible labels for parcellation (to account for possibility that
         some parcels have NO expression data). Default: None
@@ -522,7 +520,9 @@ def aggregate_samples(microarray, labels=None, region_agg='donors',
         `G` is genes. Index of dataframes should identify to which region each
         sample was assigned
     labels : array_like
-
+        Array containing IDs of all possible region labels. If no samples from
+        any donor were assigned to a region NaNs are inserted in place of
+        expression values
     region_agg : {'samples', 'donors'}, optional
         When multiple samples are identified as belonging to a region in
         `atlas` this determines how they are aggegated. If 'samples',
diff --git a/abagen/utils.py b/abagen/utils.py
index e8ea8b7..d674b9e 100644
--- a/abagen/utils.py
+++ b/abagen/utils.py
@@ -16,6 +16,80 @@ AGG_FUNCS = dict(
 )
 
 
+def check_dict(dictionary):
+    """
+    Checks that `dictionary` is a dict and makes it one, if not
+
+    Parameters
+    ----------
+    dictionary : dict
+        Presumptive dictionary
+
+    Returns
+    -------
+    dictionary : dict
+        Actual dictionary. If it wasn't one before, keys are chronological int,
+        starting at 0
+    """
+
+    if isinstance(dictionary, str):
+        dictionary = [dictionary]
+    if not isinstance(dictionary, dict):
+        return dict(zip(range(len(dictionary)), dictionary))
+    return dictionary
+
+
+def flatten_dict(dictionary, subkey):
+    """
+    Extracts `subkey` from entries of `dictionary` and creates new dictionary
+
+    Parameters
+    ----------
+    dictionary : dict
+        Input dictionary, where values are sub-dictionaries
+    subkey : str
+        Key to extract from `dictionary` entries
+
+    Returns
+    -------
+    flattened : dict
+        Flattened input `dictionary`
+
+    Examples
+    --------
+    >>> test = {'one': {'value': 1}, 'two': {'value': 2}}
+    >>> flatten_dict(test, 'value')
+    {'one': 1, 'two': 2}
+    """
+
+    return {k: v.get(subkey, None) for k, v in check_dict(dictionary).items()}
+
+
+def first_entry(dictionary, subkey=None):
+    """
+    Extracts `subkey` from the first entry in `dictionary`
+
+    Parameters
+    ----------
+     dictionary : dict
+        Input dictionary, where values are sub-dictionaries
+    subkey : str
+        Key to extract from `dictionary` entries
+
+    Returns
+    -------
+    entry
+        Extracted entry
+    """
+
+    dictionary = check_dict(dictionary)
+    entry = dictionary[list(dictionary)[0]]
+    if subkey is not None:
+        entry = entry.get(subkey, None)
+
+    return entry
+
+
 def leftify_atlas(atlas):
     """
     Zeroes out all ROIs in the right hemisphere of `atlas`
diff --git a/docs/api.rst b/docs/api.rst
index a1683a4..496a321 100644
--- a/docs/api.rst
+++ b/docs/api.rst
@@ -32,18 +32,19 @@ Reference API
    :no-members:
    :no-inherited-members:
 
-.. currentmodule:: abagen
+.. currentmodule:: abagen.datasets
 
 .. autosummary::
    :template: function.rst
    :toctree: generated/
 
-   abagen.fetch_microarray
-   abagen.fetch_rnaseq
-   abagen.fetch_raw_mri
-   abagen.fetch_freesurfer
-   abagen.fetch_desikan_killiany
-   abagen.fetch_gene_group
+   abagen.datasets.fetch_microarray
+   abagen.datasets.fetch_rnaseq
+   abagen.datasets.fetch_raw_mri
+   abagen.datasets.fetch_freesurfer
+   abagen.datasets.fetch_desikan_killiany
+   abagen.datasets.fetch_gene_group
+   abagen.datasets.fetch_donor_info
 
 .. _ref_io:
 
@@ -64,6 +65,9 @@ Reference API
    abagen.io.read_ontology
    abagen.io.read_pacall
    abagen.io.read_probes
+   abagen.io.read_genes
+   abagen.io.read_tpm
+   abagen.io.read_counts
 
 .. _ref_correct:
 
@@ -73,15 +77,15 @@ Reference API
    :no-members:
    :no-inherited-members:
 
-.. currentmodule:: abagen
+.. currentmodule:: abagen.correct
 
 .. autosummary::
    :template: function.rst
    :toctree: generated/
 
-   abagen.remove_distance
-   abagen.keep_stable_genes
-   abagen.normalize_expression
+   abagen.correct.remove_distance
+   abagen.correct.keep_stable_genes
+   abagen.correct.normalize_expression
 
 .. _ref_utils:
 
@@ -97,5 +101,6 @@ Reference API
    :template: function.rst
    :toctree: generated/
 
+   abagen.utils.leftify_atlas
    abagen.utils.check_atlas_info
    abagen.utils.check_metric
diff --git a/docs/user_guide/download.rst b/docs/user_guide/download.rst
index 5a86015..903cc8b 100644
--- a/docs/user_guide/download.rst
+++ b/docs/user_guide/download.rst
@@ -25,7 +25,7 @@ following command:
     (e.g., ``['9861', '10021']``) instead of passing ``'all'``.
 
 This command will download data from the specified donors into a folder called
-``allenbrain`` in the ``$HOME/abagen-data`` directory. If you have already
+``microarray`` in the ``$HOME/abagen-data`` directory. If you have already
 downloaded the data you can provide the ``data_dir`` argument to specify where
 the files have been stored:
 
@@ -55,20 +55,27 @@ the directory specified should have the following structure:
     ├── normalized_microarray_donor15697/
     └── normalized_microarray_donor9861/
 
-(Note the directory does not have to be named ``allenbrain`` for this to work.)
+(Note the directory does not have to be named ``microarray`` for this to work.)
 
 .. _usage_download_loading:
 
 Loading the AHBA data
 ---------------------
 
-The ``files`` object returned by :func:`abagen.fetch_microarray` is a
+The ``files`` object returned by :func:`abagen.fetch_microarray` is a nested
 dictionary with filepaths to the five different file types in the AHBA
-microarray dataset:
+microarray dataset. The keys are the donor IDs:
 
 .. doctest::
 
     >>> print(sorted(files))
+    ['10021', '12876', '14380', '15496', '15697', '9861']
+
+And the values for each entry are a sub-dictionary of the downloaded files:
+
+.. doctest::
+
+    >>> print(sorted(files['9861']))
     ['annotation', 'microarray', 'ontology', 'pacall', 'probes']
 
 You can load the data in these files using the :mod:`abagen.io` functions.
@@ -81,7 +88,8 @@ For example, you can load the annotation file for the first donor with:
 
 .. doctest::
 
-    >>> annotation = abagen.io.read_annotation(files['annotation'][0])
+    >>> data = files['9861']
+    >>> annotation = abagen.io.read_annotation(data['annotation'])
     >>> print(annotation)
                structure_id  slab_num  well_id  ... mni_x mni_y mni_z
     sample_id                                   ...
@@ -100,7 +108,7 @@ And you can do the same for, e.g., the probe file with:
 
 .. doctest::
 
-    >>> probes = abagen.io.read_probes(files['probes'][0])
+    >>> probes = abagen.io.read_probes(data['probes'])
     >>> print(probes)
                           probe_name  gene_id  ... entrez_id chromosome
     probe_id                                   ...


Test patch:
diff --git a/abagen/tests/datasets/test_fetchers.py b/abagen/tests/datasets/test_fetchers.py
index 9501f52..6c63cde 100644
--- a/abagen/tests/datasets/test_fetchers.py
+++ b/abagen/tests/datasets/test_fetchers.py
@@ -22,8 +22,9 @@ def test_fetch_microarray():
     # f5 = datasets.fetch_microarray(donors='all')
 
     assert f1 == f2 == f3 == f4
+    assert len(f1) == 1  # only one donor
     for k in ['microarray', 'annotation', 'pacall', 'probes', 'ontology']:
-        assert len(f1.get(k)) == 1
+        assert k in f1['12876']
 
     # check downloading incorrect donor
     with pytest.raises(ValueError):
@@ -41,8 +42,9 @@ def test_fetch_raw_mri():
     f4 = fetchers.fetch_raw_mri(donors=None)
 
     assert f1 == f2 == f3 == f4
+    assert len(f1) == 1
     for k in ['t1w', 't2w']:
-        assert len(f1.get(k)) == 1
+        assert k in f1['12876']
 
     # check downloading incorrect donor
     with pytest.raises(ValueError):
@@ -59,8 +61,9 @@ def test_fetch_rnaseq():
     f4 = fetchers.fetch_rnaseq(donors=None)
 
     assert f1 == f2 == f3 == f4
+    assert len(f1) == 1
     for k in ['genes', 'ontology', 'counts', 'tpm', 'annotation']:
-        assert len(f1.get(k)) == 1
+        assert k in f1['9861']
 
     with pytest.raises(ValueError):
         fetchers.fetch_rnaseq(donors='notadonor')
diff --git a/abagen/tests/test_allen.py b/abagen/tests/test_allen.py
index e71ff98..62b24e5 100644
--- a/abagen/tests/test_allen.py
+++ b/abagen/tests/test_allen.py
@@ -70,4 +70,4 @@ def test_missing_labels(testfiles, atlas):
     assert len(out) == len(info)
 
     assert isinstance(counts, pd.DataFrame)
-    assert counts.shape == (len(info), len(testfiles['probes']))
+    assert counts.shape == (len(info), len(testfiles))
diff --git a/abagen/tests/test_correct.py b/abagen/tests/test_correct.py
index 04f1e34..8124d0a 100644
--- a/abagen/tests/test_correct.py
+++ b/abagen/tests/test_correct.py
@@ -11,6 +11,7 @@ import pytest
 import scipy.stats as sstats
 
 from abagen import allen, correct, io
+from abagen.utils import flatten_dict
 
 
 @pytest.fixture(scope='module')
@@ -120,7 +121,10 @@ def test__srs(a):
 ])
 def test_normalize_expression_real(testfiles, method):
     # load in data and add some NaN values for "realness"
-    micro = [io.read_microarray(f).T for f in testfiles['microarray']]
+    micro = [
+        io.read_microarray(f).T
+        for f in flatten_dict(testfiles, 'microarray').values()
+    ]
     inds = [[5, 15, 25], [0, 10, 20]]
     for n, idx in enumerate(inds):
         micro[n].iloc[idx] = np.nan
diff --git a/abagen/tests/test_io.py b/abagen/tests/test_io.py
index aeb4b9c..601f581 100644
--- a/abagen/tests/test_io.py
+++ b/abagen/tests/test_io.py
@@ -10,6 +10,7 @@ import pandas as pd
 import pytest
 
 from abagen import io
+from abagen.utils import flatten_dict
 
 
 annot_cols = [
@@ -35,10 +36,9 @@ gene_cols = [
 rna_annot_cols = [
     'RNAseq_sample_name', 'replicate_sample', 'sample_name', 'well_id',
     'microarray_run_id', 'ontology_color', 'main_structure',
-    'sub_structure', 'ontology_structure_id', 'ontology_structure_acronym',
-    'hemisphere', 'brain', 'million_clusters', 'clip_percentage',
-    'RIN_RNA_quality', 'rnaseq_run_id', 'A.Pct', 'C.Pct', 'G.Pct', 'T.Pct',
-    'N.Pct'
+    'sub_structure', 'structure_id', 'structure_acronym', 'hemisphere',
+    'brain', 'million_clusters', 'clip_percentage', 'RIN_RNA_quality',
+    'rnaseq_run_id', 'A.Pct', 'C.Pct', 'G.Pct', 'T.Pct', 'N.Pct'
 ]
 
 
@@ -50,7 +50,7 @@ rna_annot_cols = [
     ('probes', False, probe_cols),
 ])
 def test_readfiles(testfiles, key, has_parq, columns):
-    for fn in testfiles.get(key):
+    for d, fn in flatten_dict(testfiles, key).items():
         func = getattr(io, 'read_{}'.format(key))
 
         # check file (CSV + parquet) exist
@@ -94,7 +94,7 @@ def test_readfiles(testfiles, key, has_parq, columns):
     ('annotation', rna_annot_cols)
 ])
 def test_readrnaseq(rnafiles, key, columns):
-    for fn in rnafiles.get(key):
+    for d, fn in flatten_dict(rnafiles, key).items():
         func = getattr(io, 'read_{}'.format(key))
 
         # check file exists
diff --git a/abagen/tests/test_probes.py b/abagen/tests/test_probes.py
index ed6037e..208b0ce 100644
--- a/abagen/tests/test_probes.py
+++ b/abagen/tests/test_probes.py
@@ -9,11 +9,12 @@ import pytest
 
 import abagen
 from abagen import probes_
+from abagen.utils import first_entry, flatten_dict
 
 
 def test_reannotate_probes(testfiles):
     # set up a few useful variables
-    probe_file = testfiles['probes'][0]
+    probe_file = first_entry(testfiles, 'probes')
 
     # should work with either a filename _or_ a dataframe
     reannot = probes_.reannotate_probes(probe_file)
@@ -42,9 +43,9 @@ def test_reannotate_probes(testfiles):
 ])
 def test_filter_probes(testfiles, threshold, expected_length):
     # set up a few useful variables
-    pacall = testfiles['pacall']
-    probe_file = testfiles['probes'][0]
-    samples = testfiles['annotation']
+    pacall = flatten_dict(testfiles, 'pacall')
+    probe_file = first_entry(testfiles, 'probes')
+    samples = flatten_dict(testfiles, 'annotation')
     probe_df = abagen.io.read_probes(probe_file)
 
     # should work with either a filename _or_ a dataframe
@@ -80,10 +81,10 @@ def test_groupby_and_apply(func, a_expected, b_expected):
     prb = pd.DataFrame(dict(gene_symbol=gene_symbol), index=index)
     exp = [info[['a', 'b']].copy()]
 
-    mi = probes_._groupby_and_apply(exp, prb, info, func)[0]
+    mi = list(probes_._groupby_and_apply(exp, prb, info, func).values())[0]
     expected = pd.DataFrame(dict(a=a_expected, b=b_expected),
                             index=pd.Series(['a', 'b', 'c'],
-                                            name='gene_symbol'))
+                                            name='gene_symbol')).T
     pd.testing.assert_frame_equal(mi, expected)
 
 
@@ -189,15 +190,15 @@ def test_diff_stability():
                       [4, 6, 6, 7],
                       [5, 7, 3, 4]],
                      index=pd.Series(['a', 'b', 'c'], name='gene_symbol'),
-                     columns=expression[0].columns),
+                     columns=expression[0].columns).T,
         pd.DataFrame([[1, 2, 3],
                       [1, 2, 3],
                       [1, 2, 3]],
                      index=pd.Series(['a', 'b', 'c'], name='gene_symbol'),
-                     columns=expression[1].columns),
+                     columns=expression[1].columns).T,
     ]
 
-    out = probes_._diff_stability(expression, prb, annotation)
+    out = list(probes_._diff_stability(expression, prb, annotation).values())
     assert len(out) == len(expected)
     for df, exp in zip(out, expected):
         pd.testing.assert_frame_equal(df, exp)
@@ -219,7 +220,7 @@ def test_average():
     expected = pd.DataFrame(dict(a=[1.0, 3.5, 5.0],
                                  b=[11.0, 8.5, 7.0]),
                             index=pd.Series(['a', 'b', 'c'],
-                                            name='gene_symbol'))
+                                            name='gene_symbol')).T
     pd.testing.assert_frame_equal(out, expected)
 
 
@@ -237,11 +238,12 @@ def test_collapse_probes(testfiles, method):
     # we've aleady tested the underlying methods so here we just want to do
     # some smoke tests to make sure the function returns what we expected
     # regardless of the provided method
-    out = probes_.collapse_probes(testfiles['microarray'],
-                                  testfiles['annotation'],
-                                  testfiles['probes'][0],
+    out = probes_.collapse_probes(flatten_dict(testfiles, 'microarray'),
+                                  flatten_dict(testfiles, 'annotation'),
+                                  first_entry(testfiles, 'probes'),
                                   method=method)
 
+    out = list(out.values())
     assert len(out) == 2  # number of donors
     assert np.all([len(exp) == n_samp for exp, n_samp in zip(out, [363, 470])])
     assert np.all([len(exp.columns) == 29131 for exp in out])
diff --git a/abagen/tests/test_samples.py b/abagen/tests/test_samples.py
index 2d0215a..a3dca88 100644
--- a/abagen/tests/test_samples.py
+++ b/abagen/tests/test_samples.py
@@ -8,7 +8,7 @@ import pandas as pd
 import pytest
 
 from abagen import samples_
-
+from abagen.utils import first_entry, flatten_dict
 
 # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #
 #  generate fake data (based largely on real data) so we know what to expect  #
@@ -275,25 +275,30 @@ def test_aggregate_samples():
 
 
 def test_update_mni_coords_real(testfiles):
-    for annotation in testfiles['annotation']:
+    for annotation in flatten_dict(testfiles, 'annotation').values():
         samples_.update_mni_coords(annotation)
 
 
 def test_label_samples_real(testfiles, atlas):
-    out = samples_.label_samples(testfiles['annotation'][0], atlas['image'])
+    out = samples_.label_samples(first_entry(testfiles, 'annotation'),
+                                 atlas['image'])
     assert isinstance(out, pd.DataFrame)
     assert out.index.name == 'sample_id'
     assert out.columns == ['label']
 
 
 def test_drop_mismatch_samples_real(testfiles):
-    for an, on in zip(testfiles['annotation'], testfiles['ontology']):
+    annotation = flatten_dict(testfiles, 'annotation').values()
+    ontology = flatten_dict(testfiles, 'ontology').values()
+    for an, on in zip(annotation, ontology):
         samples_.drop_mismatch_samples(an, on)
 
 
 def test_mirror_samples_real(testfiles):
+    annotation = flatten_dict(testfiles, 'annotation').values()
+    ontology = flatten_dict(testfiles, 'ontology').values()
     orig = [363, 470]
-    for an, on, o in zip(testfiles['annotation'], testfiles['ontology'], orig):
+    for an, on, o in zip(annotation, ontology, orig):
         out = samples_.mirror_samples(an, on)
         # there should be more than the original # of samples but less than or
         # equal to 2x that number (we can't MORE than duplicate)
@@ -301,7 +306,9 @@ def test_mirror_samples_real(testfiles):
 
 
 def test__mirror_ontology_real(testfiles):
+    annotation = flatten_dict(testfiles, 'annotation').values()
+    ontology = flatten_dict(testfiles, 'ontology').values()
     orig = [363, 470]
-    for a, o, l in zip(testfiles['annotation'], testfiles['ontology'], orig):
+    for a, o, l in zip(annotation, ontology, orig):
         annot = samples_._mirror_ontology(a, o)
         assert len(annot) == l
diff --git a/abagen/tests/test_utils.py b/abagen/tests/test_utils.py
index 815c56a..413f780 100644
--- a/abagen/tests/test_utils.py
+++ b/abagen/tests/test_utils.py
@@ -10,6 +10,21 @@ import pytest
 from abagen import utils
 
 
+@pytest.mark.xfail
+def test_check_dict():
+    assert False
+
+
+@pytest.mark.xfail
+def test_flatten_dict():
+    assert False
+
+
+@pytest.mark.xfail
+def test_first_entry():
+    assert False
+
+
 def test_leftify_atlas(atlas):
     out = utils.leftify_atlas(atlas['image'])
     assert len(np.unique(out.dataobj)) == 44