brainiak · mihaic · Apr 25, 2017 · Apr 11, 2017 · Apr 25, 2017
diff --git a/brainiak/fcma/preprocessing.py b/brainiak/fcma/preprocessing.py
@@ -23,7 +23,7 @@
 from scipy.stats.mstats import zscore
 from mpi4py import MPI
 
-from ..image import multimask_images
+from ..image import mask_images, multimask_images
 
 
 logger = logging.getLogger(__name__)
@@ -123,12 +123,13 @@ def prepare_fcma_data(images, conditions, mask1, mask2=None,
     raw_data2 = []
     if rank == 0:
         if mask2 is not None:
-            activity_data1, activity_data2 = multimask_images(images,
-                                                              [mask1, mask2],
-                                                              np.float32)
+            masks = (mask1, mask2)
+            activity_data1, activity_data2 = zip(*multimask_images(images,
+                                                                   masks,
+                                                                   np.float32))
             raw_data2, _ = _separate_epochs(activity_data2, conditions)
         else:
-            (activity_data1,) = multimask_images(images, (mask1,), np.float32)
+            activity_data1 = list(mask_images(images, mask1, np.float32))
         raw_data1, labels = _separate_epochs(activity_data1, conditions)
         time1 = time.time()
     raw_data_length = len(raw_data1)
@@ -216,7 +217,7 @@ def prepare_mvpa_data(images, conditions, mask):
     labels: 1D array
         contains labels of the data
     """
-    (activity_data,) = multimask_images(images, (mask,), np.float32)
+    activity_data = list(mask_images(images, mask, np.float32))
     epoch_info = generate_epochs_info(conditions)
     num_epochs = len(epoch_info)
     (d1, _) = activity_data[0].shape

diff --git a/brainiak/image.py b/brainiak/image.py
@@ -19,14 +19,63 @@
     "mask_image",
     "multimask_images",
 ]
-from typing import Iterable, Sequence
-from typing import List  # noqa F401 https://gitlab.com/pycqa/flake8/issues/118
+
+import itertools
+
+from typing import Iterable, Sequence, Type, TypeVar
 
 import numpy as np
 
 from nibabel.spatialimages import SpatialImage
 
 
+T = TypeVar("T", bound="MaskedMultiSubjectData")
+
+
+class MaskedMultiSubjectData(np.ndarray):
+    """Array in shape n_voxels, n_trs, n_subjects."""
+    @classmethod
+    def from_masked_images(cls: Type[T], masked_images: Iterable[np.ndarray],
+                           n_sub: int) -> T:
+        """Create a new instance from masked images.
+
+        Parameters
+        ----------
+        masked_images
+            Images to concatenate.
+        n_sub
+            Number of subjects. Must match the number of images.
+
+        Returns
+        -------
+        T
+            A new instance.
+
+        Raises
+        ------
+        ValueError
+            Images have different shapes.
+
+            The number of images differs from n_sub.
+        """
+        images_iterator = iter(masked_images)
+        first_image = next(images_iterator)
+        result = np.empty((first_image.shape[0], first_image.shape[1], n_sub))
+        for n_images, image in enumerate(itertools.chain([first_image],
+                                                         images_iterator)):
+            if image.shape != first_image.shape:
+                raise ValueError("Image {} has different shape from first "
+                                 "image: {} != {}".format(n_images,
+                                                          image.shape,
+                                                          first_image.shape))
+            result[:, :, n_images] = image
+        n_images += 1
+        if n_images != n_sub:
+            raise ValueError("n_sub != number of images: {} != {}"
+                             .format(n_sub, n_images))
+        return result.view(cls)
+
+
 class ConditionSpec(np.ndarray):
     """One-hot representation of conditions across epochs and TRs.
 
@@ -85,7 +134,7 @@ def mask_image(image: SpatialImage, mask: np.ndarray, data_type: type = None
 
 def multimask_images(images: Iterable[SpatialImage],
                      masks: Sequence[np.ndarray], image_type: type = None
-                     ) -> Sequence[Sequence[np.ndarray]]:
+                     ) -> Iterable[Sequence[np.ndarray]]:
     """Mask images with multiple masks.
 
     Parameters
@@ -97,14 +146,32 @@ def multimask_images(images: Iterable[SpatialImage],
     image_type:
         Type to cast images to.
 
-    Returns
-    -------
-    List[List[np.ndarray]]
-        For each mask, a list of masked images.
+    Yields
+    ------
+    Sequence[np.ndarray]
+        For each mask, a masked image.
     """
-    masked_images = [[] for _ in range(len(masks))
-                     ]  # type: List[List[np.ndarray]]
     for image in images:
-        for i, mask in enumerate(masks):
-            masked_images[i].append(mask_image(image, mask, image_type))
-    return masked_images
+        yield [mask_image(image, mask, image_type) for mask in masks]
+
+
+def mask_images(images: Iterable[SpatialImage], mask: np.ndarray,
+                image_type: type = None) -> Iterable[np.ndarray]:
+    """Mask images.
+
+    Parameters
+    ----------
+    images:
+        Images to mask.
+    mask:
+        Mask to apply.
+    image_type:
+        Type to cast images to.
+
+    Yields
+    ------
+    np.ndarray
+        Masked image.
+    """
+    for images in multimask_images(images, (mask,), image_type):
+        yield images[0]
diff --git a/brainiak/io.py b/brainiak/io.py
@@ -21,7 +21,7 @@
 ]
 
 from pathlib import Path
-from typing import Iterable, List, Union
+from typing import Callable, Iterable, List, Union
 
 import nibabel as nib
 import numpy as np
@@ -61,13 +61,46 @@ def load_images_from_dir(in_dir: Union[str, Path], suffix: str = "nii.gz",
         yield nib.load(str(f))
 
 
-def load_boolean_mask(path: Union[str, Path]) -> np.ndarray:
+def load_images(image_paths: Iterable[Union[str, Path]]
+                ) -> Iterable[SpatialImage]:
+    """Load images from paths.
+
+    For efficiency, returns an iterator, not a sequence, so the results cannot
+    be accessed by indexing.
+
+    For every new iteration through the images, load_images must be called
+    again.
+
+    Parameters
+    ----------
+    image_paths:
+        Paths to images.
+
+    Yields
+    ------
+    SpatialImage
+        Image.
+    """
+    for image_path in image_paths:
+        if isinstance(image_path, Path):
+            string_path = str(image_path)
+        else:
+            string_path = image_path
+        yield nib.load(string_path)
+
+
+def load_boolean_mask(path: Union[str, Path],
+                      predicate: Callable[[np.ndarray], np.ndarray] = None
+                      ) -> np.ndarray:
     """Load boolean nibabel.SpatialImage mask.
 
     Parameters
     ----------
     path
         Mask path.
+    predicate
+        Callable used to create boolean values, e.g. a threshold function
+        ``lambda x: x > 50``.
 
     Returns
     -------
@@ -76,7 +109,12 @@ def load_boolean_mask(path: Union[str, Path]) -> np.ndarray:
     """
     if not isinstance(path, str):
         path = str(path)
-    return nib.load(path).get_data().astype(np.bool)
+    data = nib.load(path).get_data()
+    if predicate is not None:
+        mask = predicate(data)
+    else:
+        mask = data.astype(np.bool)
+    return mask
 
 
 def load_labels(path: Union[str, Path]) -> List[SingleConditionSpec]:

diff --git a/brainiak/isfc.py b/brainiak/isfc.py
@@ -31,7 +31,6 @@
 # Princeton University, 2017
 
 from brainiak.fcma.util import compute_correlation
-import nibabel as nib
 import numpy as np
 from scipy import stats
 
@@ -115,53 +114,3 @@ def isfc(D, collapse_subj=True):
     if collapse_subj:
         ISFC = np.mean(ISFC, axis=2)
     return ISFC
-
-
-def load_subjects_nii(data_files, mask_file, mask_func=None):
-    """Loading masked nifti data into matrix
-
-    Given a list of subject data files and a mask, loads voxel timecourses
-    inside mask into a matrix, which is returned along with the voxel
-    coordinates of the mask. If mask_func is given, it specifies which mask
-    values should be included (e.g. for thresholding a continuous-valued mask)
-
-    Parameters
-    ----------
-    data_files : list of filenames of subject nii files
-        each nii file should be 4D (space+time), with all dimensions identical
-        for all subjects
-
-    mask_file : filename of mask nii
-        mask should be 3D (space), with dimensions identical to the first
-        three dimensions of the data files
-
-    mask_func : Callable[[ndarray], bool] : default x>0
-
-    Returns
-    -------
-    D : voxel by time by subject ndarray
-        all data within mask, from all subjects
-    coords : tuple of 3 ndarrays
-        x,y,z (as provided by nibabel) coordinates of mask voxel locations
-    """
-
-    mask_nii = nib.load(mask_file)
-    mask = mask_nii.get_data()
-    if mask_func is not None:
-        mask = mask_func(mask)
-    else:
-        mask = mask > 0
-    mask_shape = mask_nii.shape
-    coords = np.where(mask)
-
-    data_shape = nib.load(data_files[0]).shape
-    D = np.zeros((np.sum(mask), data_shape[3], len(data_files)))
-    for s in range(len(data_files)):
-        nii = nib.load(data_files[s])
-        if nii.shape != data_shape:
-            raise ValueError("Data has different shapes across subjects")
-        if nii.shape[:3] != mask_shape:
-            raise ValueError("Data and mask have different shapes")
-        D[:, :, s] = nii.get_data()[mask]
-
-    return (D, coords)
diff --git a/examples/isfc/isfc.py b/examples/isfc/isfc.py
@@ -21,6 +21,7 @@
 # Princeton University, 2017
 
 import brainiak.isfc
+from brainiak import image, io
 import nibabel as nib
 import numpy as np
 from matplotlib import pyplot as plt
@@ -36,8 +37,11 @@
           subj in np.arange(1, 5)]
 
 print('Loading data from ', len(fnames), ' subjects...')
-D, coords = brainiak.isfc.load_subjects_nii(fnames, brain_fname,
-                                            lambda x: x > 50)
+
+brain_mask = io.load_boolean_mask(brain_fname, lambda x: x > 50)
+masked_images = image.mask_images(io.load_images(fnames), brain_mask)
+coords = np.where(brain_mask)
+D = image.MaskedMultiSubjectData.from_masked_images(masked_images, len(fnames))
 
 print('Calculating ISC on ', D.shape[0], ' voxels')
 ISC = brainiak.isfc.isc(D)