<commit_before>""" Base classes for representing and manipulating data. """

import logging
import re
import os
import sys
import numpy as np
import pandas as pd
from scipy import sparse
from neurosynth.base import mappable
from neurosynth.base import mask, imageutils, transformations
from neurosynth.base import lexparser as lp
from neurosynth.utils import deprecated

# For Python 2/3 compatibility
from six import string_types
from functools import reduce
try:
    from urllib2 import urlopen
except ImportError:
    from urllib.request import urlopen
try:
    import cPickle as pickle
except:
    import pickle

logger = logging.getLogger('neurosynth.dataset')


def download(path='.', url=None, unpack=False):
    """ Download the latest data files.
    Args:
        path (str): Location to save the retrieved data files. Defaults to
            current directory.
        unpack (bool): If True, unzips the data file post-download.
    """

    if url is None:
        url = 'https://github.com/neurosynth/neurosynth-data/blob/master/current_data.tar.gz?raw=true'
    if os.path.exists(path) and os.path.isdir(path):
        basename = os.path.basename(url).split('?')[0]
        filename = os.path.join(path, basename)
    else:
        filename = path

    f = open(filename, 'wb')

    u = urlopen(url)
    file_size = int(u.headers["Content-Length"][0])
    print("Downloading the latest Neurosynth files: {0} bytes: {1}".format(
        url, file_size))

    bytes_dl = 0
    block_size = 8192
    while True:
        buffer = u.read(block_size)
        if not buffer:
            break
        bytes_dl += len(buffer)
        f.write(buffer)
        p = float(bytes_dl) / file_size
        status = r"{0}  [{1:.2%}]".format(bytes_dl, p)
        status = status + chr(8) * (len(status) + 1)
        sys.stdout.write(status)

    f.close()

    if unpack:
        import tarfile
        tarfile.open(filename, 'r:gz').extractall(os.path.dirname(filename))


class Dataset(object):

    """ Base Dataset class.

    The core data-representing object in Neurosynth. Internally stores
    information about Mappables, including both reported activations and tagged
    features. Provides a variety of methods for manipulating and retrieving
    various kinds of data.

    The Dataset is typically initialized by passing in a database file as the
    first argument. At minimum, the input file must contain tab-delimited
    columns named x, y, z, id, and space (case-insensitive). The x/y/z columns
    indicate the coordinates of the activation center or peak, the id column is
    used to group multiple activations from a single Mappable (e.g. an
    article). Typically the id should be a uniquely identifying field
    accessible to others, e.g., a doi in the case of entire articles. The space
    column indicates the nominal atlas used to produce each activation.
    Currently all values except 'TAL' (Talairach) will be ignored. If space ==
    TAL and the transform argument is True, all activations reported in 
    Talairach space will be converted to MNI space using the Lancaster et al
    transform.

    Args:
        filename (str): The name of a database file containing a list of
            activations.
        feature_filename (str): An optional filename to construct a
            FeatureTable from.
        masker (str): An optional Nifti/Analyze image name defining the space
            to use for all operations. If no image is passed, defaults to the
            MNI152 2 mm template packaged with FSL.
        r (int): An optional integer specifying the radius of the smoothing
            kernel, in mm. Defaults to 6 mm.
        transform (bool, dict): Optional argument specifying how to handle
            transformation between coordinates reported in different
            stereotactic spaces. When True (default), activations in Talairach
            (T88) space will be converted to MNI space using the Lancaster et
            al (2007) transform; no other transformations will be applied. When
            False, no transformation will be applied. Alternatively, the user
            can pass their own dictionary of named transformations to apply, in
            which case each activation will be checked against the dictionary
            as it is read in and the specified transformation will be applied
            if found (for further details, see transformations.Transformer).
        target (str): The name of the target space within which activation
            coordinates are represented. By default, MNI.
        kwargs (dict): Additional optional arguments passed to add_features().
    """

    def __init__(
            self, filename, feature_filename=None, masker=None, r=6,
            transform=True, target='MNI', **kwargs):

        # Instance properties
        self.r = r

        # Set up transformations between different image spaces
        if transform:
            if not isinstance(transform, dict):
                transform = {'T88': transformations.t88_to_mni(),
                             'TAL': transformations.t88_to_mni()
                             }
            self.transformer = transformations.Transformer(transform, target)
        else:
            self.transformer = None

        # Load mappables
        self.mappables = self._load_mappables_from_txt(filename)

        # Load the volume into a new Masker
        if masker is None:
            resource_dir = os.path.join(os.path.dirname(__file__),
                                        os.path.pardir,
                                        'resources')
            masker = os.path.join(
                resource_dir, 'MNI152_T1_2mm_brain.nii.gz')
        self.masker = mask.Masker(masker)

        # Create supporting tables for images and features
        self.create_image_table()
        if feature_filename is not None:
            self.add_features(feature_filename, **kwargs)

    def _load_mappables_from_txt(self, filename):
        """ Load mappables from a text file.

        Args:
            filename (str): a string pointing to the location of the txt file
                to read from.
        """
        logger.info("Loading mappables from %s..." % filename)

        # Read in with pandas
        contents = pd.read_csv(filename, sep='\t')
        contents.columns = [col.lower() for col in list(contents.columns)]

        # Make sure all mandatory columns exist
        mc = ['x', 'y', 'z', 'id', 'space']
        if (set(mc) - set(list(contents.columns))):
            logger.error(
                "At least one of mandatory columns (x, y, z, id, and space) "
                "is missing from input file.")
            return

        # Initialize all mappables--for now, assume Articles are passed
        logger.info("Loading study data from database file...")
        return list(contents.groupby('id', as_index=False).apply(
            lambda x: mappable.Article(x, self.transformer)))

    def create_image_table(self, r=None):
        """ Create and store a new ImageTable instance based on the current
        Dataset. Will generally be called privately, but may be useful as a
        convenience method in cases where the user wants to re-generate the
        table with a new smoothing kernel of different radius.

        Args:
            r (int): An optional integer indicating the radius of the smoothing
                kernel. By default, this is None, which will keep whatever
                value is currently set in the Dataset instance.
        """
        logger.info("Creating image table...")
        if r is not None:
            self.r = r
        self.image_table = ImageTable(self)

    def add_mappables(self, filename=None, mappables=None, remap=True):
        """ Append new Mappable objects to the end of the list.

        Either a filename or a list of mappables must be passed.

        Args:
            filename (str): The name of the file containing new mappables.
            mappables (list): Mappable instances to append to existing list.
            remap (bool): Optional flag indicating whether to regenerate the
                entire ImageTable after appending the new Mappables.
        """
        # TODO: (i) it would be more effiicent to only map the new Mappables
        # into the ImageTable instead of redoing everything. (ii) we should
        # check for duplicates and prompt whether to overwrite or update in
        # cases where conflicts occur.
        if filename != None:
            self.mappables.extend(self._load_mappables_from_txt(filename))
        elif mappables != None:
            self.mappables.extend(mappables)
        if remap:
            self.create_image_table()

    def delete_mappables(self, ids, remap=True):
        """ Delete specific Mappables from the Dataset.

        Note that 'ids' is a list of unique identifiers of the Mappables (e.g.,
        doi's), and not indices in the current instance's mappables list.

        Args:
            ids (list): ids corresponding to the Mappables to delete.
            remap (bool): Optional flag indicating whether to regenerate the
                entire ImageTable after deleting undesired Mappables.
        """
        self.mappables = [m for m in self.mappables if m.id not in ids]
        if remap:
            self.create_image_table()

    def get_studies(self, features=None, expression=None, mask=None,
                    peaks=None, frequency_threshold=0.001,
                    activation_threshold=0.0, func=np.sum, return_type='ids',
                    r=6
                    ):
        """ Get IDs or data for studies that meet specific criteria.

        If multiple criteria are passed, the set intersection is returned. For
        example, passing expression='emotion' and mask='my_mask.nii.gz' would
        return only those studies that are associated with emotion AND report
        activation within the voxels indicated in the passed image.

        Args:
            ids (list): A list of IDs of studies to retrieve.
            features (list or str): The name of a feature, or a list of
                features, to use for selecting studies.
            expression (str): A string expression to pass to the PEG for study
                retrieval.
            mask: the mask image (see Masker documentation for valid data
                types).
            peaks (ndarray or list): Either an n x 3 numpy array, or a list of
                lists (e.g., [[-10, 22, 14]]) specifying the world (x/y/z)
                coordinates of the target location(s).
            frequency_threshold (float): For feature-based or expression-based
                selection, the threshold for selecting studies--i.e., the
                cut-off for a study to be included. Must be a float in range
                [0, 1].
            activation_threshold (int or float): For mask-based or peak-based
                selection, threshold for a study to be included based on
                amount of activation displayed. If an integer, represents the
                absolute number of voxels that must be active within the mask
                (or generated ROIs) in order for a study to be selected. If a
                float, it represents the proportion of voxels that must be
                active.
            func (Callable): The function to use when aggregating over the list
                of features. See documentation in FeatureTable.get_ids() for a
                full explanation. Only used for feature- or expression-based
                selection.
            return_type (str): A string specifying what data to return. Valid
                options are:
                'ids': returns a list of IDs of selected studies.
                'images': returns a voxel x mappable matrix of data for all
                selected studies.
                'weights': returns a dict where the keys are study IDs and the
                values are the computed weights. Only valid when performing
                feature-based selection.
            r (int): For peak-based selection, the radius in millimeters of the
                sphere to grow around each peak.

        Returns:
            When return_type is 'ids' (default), returns a list of IDs of the
            selected studies. When return_type is 'data', returns a 2D numpy
            array, with voxels in rows and studies in columns. When return_type
            is 'weights' (valid only for expression-based selection), returns
            a dict, where the keys are study IDs, and the values are the
            computed weights.

        Examples
        --------
        Select all studies tagged with the feature 'emotion':

            >>> ids = dataset.get_studies(features='emotion')

        Select all studies that activate at least 20% of voxels in an amygdala
        mask, and retrieve activation data rather than IDs:

            >>> data = dataset.get_studies(mask='amygdala_mask.nii.gz',
                threshold=0.2, return_type='images')

        Select studies that activate at least 5% of all voxels within 12 mm of
        three specific foci:

            >>> ids = dataset.get_studies(peaks=[[12, -20, 30], [-26, 22, 22],
                                                [0, 36, -20]], r=12)

        """
        results = []

        # Feature-based selection
        if features is not None:
            # Need to handle weights as a special case, because we can't
            # retrieve the weights later using just the IDs.
            if return_type == 'weights':
                if expression is not None or mask is not None or \
                        peaks is not None:
                    raise ValueError(
                        "return_type cannot be 'weights' when feature-based "
                        "search is used in conjunction with other search "
                        "modes.")
                return self.feature_table.get_ids(
                    features, frequency_threshold, func, get_weights=True)
            else:
                results.append(self.feature_table.get_ids(
                    features, frequency_threshold, func))

        # Logical expression-based selection
        if expression is not None:
            _ids = self.feature_table.get_ids_by_expression(
                expression, frequency_threshold, func)
            results.append(list(_ids))

        # Mask-based selection
        if mask is not None:
            mask = self.masker.mask(mask, in_global_mask=True).astype(bool)
            num_vox = np.sum(mask)
            prop_mask_active = self.image_table.data.T.dot(mask).astype(float)
            if isinstance(activation_threshold, float):
                prop_mask_active /= num_vox
            indices = np.where(prop_mask_active > activation_threshold)[0]
            results.append([self.image_table.ids[ind] for ind in indices])

        # Peak-based selection
        if peaks is not None:
            peaks = np.array(peaks)  # Make sure we have a numpy array
            peaks = transformations.xyz_to_mat(peaks)
            m = self.masker
            img = imageutils.map_peaks_to_image(
                peaks, r, vox_dims=m.vox_dims, dims=m.dims,
                header=m.get_header())
            results.append(self.get_studies(
                mask=img, activation_threshold=activation_threshold))

        # Get intersection of all sets
        ids = list(reduce(lambda x, y: set(x) & set(y), results))

        if return_type == 'ids':
            return ids
        elif return_type == 'data':
            return self.get_image_data(ids)

    @deprecated("get_mappables() is deprecated and will be removed in 0.5. "
                "Please use get_studies().")
    def get_mappables(self, ids, get_image_data=False):
        """ Takes a list of unique ids and returns corresponding Mappables.

        Args:
            ids (list): A list of ids of the mappables to return.
            get_image_data (bool): When True, returns a voxel x mappable matrix
            of image data rather than the Mappable instances themselves.

        Returns:
            If get_image_data is True, a 2D numpy array of voxels x Mappables.
            Otherwise, a list of Mappables.
        """
        if get_image_data:
            return self.get_image_data(ids)
        else:
            return [m for m in self.mappables if m.id in ids]

    @deprecated("get_ids_by_features() is deprecated and will be removed in "
                "0.5. Please use get_studies(features=...).")
    def get_ids_by_features(self, features, threshold=0.001, func=np.sum,
                            get_image_data=False, get_weights=False,
                            dense=True):
        """ A wrapper for FeatureTable.get_ids().

        Args:
            features (list): features to use when selecting Mappables.
            threshold (float): Float in range 0-1. Threshold used to select
                mappables.
            func (Callable): The function to use when aggregating over the list
                of features. See documentation in FeatureTable.get_ids() for a
                full explanation.
            get_image_data (bool): When True, returns a voxel x mappable matrix
            of image data rather than the Mappable instances themselves.
        """
        ids = self.feature_table.get_ids(
            features, threshold, func, get_weights)
        return self.get_image_data(ids, dense=dense) if get_image_data else ids

    @deprecated("get_ids_by_expression() is deprecated and will be removed in "
                "0.5. Please use get_studies(expression=...).")
    def get_ids_by_expression(self, expression, threshold=0.001, func=np.sum,
                              get_image_data=False):
        ids = self.feature_table.get_ids_by_expression(
            expression, threshold, func)
        return self.get_image_data(ids) if get_image_data else ids

    @deprecated("get_ids_by_mask() is deprecated and will be removed in "
                "0.5. Please use get_studies(mask=...).")
    def get_ids_by_mask(self, mask, threshold=0.0, get_image_data=False):
        """ Return all mappable objects that activate within the bounds
        defined by the mask image. 
        Args:
            mask: the mask image (see Masker documentation for valid data 
                types).
            threshold (int, float): an integer or float. If an integer, the 
                absolute number of voxels that must be active within the mask
                for a study to be retained. When a float, proportion of voxels
                that must be active.
            get_image_data (bool): When True, returns a voxel x mappable matrix
            of image data rather than the Mappable instances themselves.
        """
        mask = self.masker.mask(mask).astype(bool)
        num_vox = np.sum(mask)
        prop_mask_active = self.image_table.data.T.dot(mask).astype(float)
        if isinstance(threshold, float):
            prop_mask_active /= num_vox
        indices = np.where(prop_mask_active > threshold)[0]
        if get_image_data:
            return self.get_image_data(indices)
        else:
            return [self.image_table.ids[ind] for ind in indices]

    @deprecated("get_ids_by_peaks() is deprecated and will be removed in "
                "0.5. Please use get_studies(peaks=...).")
    def get_ids_by_peaks(self, peaks, r=10, threshold=0.0,
                         get_image_data=False):
        """ A wrapper for get_ids_by_mask. Takes a set of xyz coordinates and
        generates a new Nifti1Image to use as a mask.

        Args:
            peaks (ndarray, list): Either an n x 3 numpy array, or a list of 
                lists (e.g., [[-10, 22, 14]]) specifying the world (x/y/z)
                coordinates of the target location(s).
            r (int): Radius in millimeters of the sphere to grow around each
                location.
          threshold (float): Optional float indicating the proportion of voxels
            that must be active in order for a Mappable to be considered
            active.
            get_image_data (bool): When True, returns a voxel x mappable matrix
            of image data rather than the Mappable instances themselves.

        Returns:
            Either a list of ids (if get_image_data = False) or a numpy array
            of image data.

        """
        peaks = np.array(peaks)  # Make sure we have a numpy array
        peaks = transformations.xyz_to_mat(peaks)
        img = imageutils.map_peaks_to_image(
            peaks, r, vox_dims=self.masker.vox_dims,
            dims=self.masker.dims, header=self.masker.get_header())
        return self.get_ids_by_mask(img, threshold,
                                    get_image_data=get_image_data)

    def add_features(self, features, append=True, merge='outer',
                     duplicates='ignore', min_studies=0.0, threshold=0.001):
        """ Construct a new FeatureTable from file.

        Args:
            features: Feature data to add. Can be:
                (a) A text file containing the feature data, where each row is
                a study in the database, with features in columns. The first
                column must contain the IDs of the studies to match up with the
                image data.
                (b) A pandas DataFrame, where studies are in rows, features are 
                in columns, and the index provides the study IDs.
            append (bool): If True, adds new features to existing ones
                incrementally. If False, replaces old features.
            merge, duplicates, min_studies, threshold: Additional arguments
                passed to FeatureTable.add_features().
         """
        if (not append) or not hasattr(self, 'feature_table'):
            self.feature_table = FeatureTable(self)

        self.feature_table.add_features(features, merge=merge,
                                        duplicates=duplicates,
                                        min_studies=min_studies,
                                        threshold=threshold)

    def get_image_data(self, ids=None, voxels=None, dense=True):
        """ A convenience wrapper for ImageTable.get_image_data(). 

        Args:
            ids (list, array): A list or 1D numpy array of Mappable ids to 
                return. If None, returnsdata for all Mappables.
            voxels (list, array): A list or 1D numpy array of voxel indices
                (i.e., rows) to return. If None, returns data for all voxels.
        """
        return self.image_table.get_image_data(ids, voxels=voxels, dense=dense)

    def get_feature_data(self, ids=None, **kwargs):
        """ A convenience wrapper for FeatureTable.get_image_data(). """
        return self.feature_table.get_feature_data(ids, **kwargs)

    def get_feature_names(self, features=None):
        """ Returns names of features. If features is None, returns all 
        features. Otherwise assumes the user is trying to find the order of the
        features.  """
        if features:
            return self.feature_table.get_ordered_names(features)
        else:
            return self.feature_table.feature_names

    def get_feature_counts(self, threshold=0.001):
        """ Returns a dictionary, where the keys are the feature names
        and the values are the number of studies tagged with the feature. """
        counts = np.sum(self.get_feature_data() >= threshold, 0)
        return dict(zip(self.get_feature_names(), list(counts)))

    @classmethod
    def load(cls, filename):
        """ Load a pickled Dataset instance from file. """
        dataset = pickle.load(open(filename, 'rb'))
        if hasattr(dataset, 'feature_table'):
            dataset.feature_table._csr_to_sdf()
        return dataset

    def save(self, filename, keep_mappables=False):
        """ Pickle the Dataset instance to the provided file.

        If keep_mappables = False (default), will delete the Mappable objects
        themselves before pickling. This will save a good deal of space and
        is generally advisable once a stable Dataset is created, as the
        Mappables are rarely used after the ImageTable is generated.
        """
        if not keep_mappables:
            self.mappables = []

        if hasattr(self, 'feature_table'):
            self.feature_table._sdf_to_csr()

        pickle.dump(self, open(filename, 'wb'), -1)

        if hasattr(self, 'feature_table'):
            self.feature_table._csr_to_sdf()

    def to_json(self, filename=None):
        """ Save the Dataset to file in JSON format.

        This is not recommended, as the resulting file will typically be
        several GB in size. If no filename is provided, returns the JSON
        string.
        """
        import json
        mappables = [m.to_json() for m in self.mappables]<commit_msg>drop keep_mappables<commit_after>""" Base classes for representing and manipulating data. """

import logging
import re
import os
import sys
import numpy as np
import pandas as pd
from scipy import sparse
from neurosynth.base import mappable
from neurosynth.base import mask, imageutils, transformations
from neurosynth.base import lexparser as lp
from neurosynth.utils import deprecated

# For Python 2/3 compatibility
from six import string_types
from functools import reduce
try:
    from urllib2 import urlopen
except ImportError:
    from urllib.request import urlopen
try:
    import cPickle as pickle
except:
    import pickle

logger = logging.getLogger('neurosynth.dataset')


def download(path='.', url=None, unpack=False):
    """ Download the latest data files.
    Args:
        path (str): Location to save the retrieved data files. Defaults to
            current directory.
        unpack (bool): If True, unzips the data file post-download.
    """

    if url is None:
        url = 'https://github.com/neurosynth/neurosynth-data/blob/master/current_data.tar.gz?raw=true'
    if os.path.exists(path) and os.path.isdir(path):
        basename = os.path.basename(url).split('?')[0]
        filename = os.path.join(path, basename)
    else:
        filename = path

    f = open(filename, 'wb')

    u = urlopen(url)
    file_size = int(u.headers["Content-Length"][0])
    print("Downloading the latest Neurosynth files: {0} bytes: {1}".format(
        url, file_size))

    bytes_dl = 0
    block_size = 8192
    while True:
        buffer = u.read(block_size)
        if not buffer:
            break
        bytes_dl += len(buffer)
        f.write(buffer)
        p = float(bytes_dl) / file_size
        status = r"{0}  [{1:.2%}]".format(bytes_dl, p)
        status = status + chr(8) * (len(status) + 1)
        sys.stdout.write(status)

    f.close()

    if unpack:
        import tarfile
        tarfile.open(filename, 'r:gz').extractall(os.path.dirname(filename))


class Dataset(object):

    """ Base Dataset class.

    The core data-representing object in Neurosynth. Internally stores
    information about Mappables, including both reported activations and tagged
    features. Provides a variety of methods for manipulating and retrieving
    various kinds of data.

    The Dataset is typically initialized by passing in a database file as the
    first argument. At minimum, the input file must contain tab-delimited
    columns named x, y, z, id, and space (case-insensitive). The x/y/z columns
    indicate the coordinates of the activation center or peak, the id column is
    used to group multiple activations from a single Mappable (e.g. an
    article). Typically the id should be a uniquely identifying field
    accessible to others, e.g., a doi in the case of entire articles. The space
    column indicates the nominal atlas used to produce each activation.
    Currently all values except 'TAL' (Talairach) will be ignored. If space ==
    TAL and the transform argument is True, all activations reported in 
    Talairach space will be converted to MNI space using the Lancaster et al
    transform.

    Args:
        filename (str): The name of a database file containing a list of
            activations.
        feature_filename (str): An optional filename to construct a
            FeatureTable from.
        masker (str): An optional Nifti/Analyze image name defining the space
            to use for all operations. If no image is passed, defaults to the
            MNI152 2 mm template packaged with FSL.
        r (int): An optional integer specifying the radius of the smoothing
            kernel, in mm. Defaults to 6 mm.
        transform (bool, dict): Optional argument specifying how to handle
            transformation between coordinates reported in different
            stereotactic spaces. When True (default), activations in Talairach
            (T88) space will be converted to MNI space using the Lancaster et
            al (2007) transform; no other transformations will be applied. When
            False, no transformation will be applied. Alternatively, the user
            can pass their own dictionary of named transformations to apply, in
            which case each activation will be checked against the dictionary
            as it is read in and the specified transformation will be applied
            if found (for further details, see transformations.Transformer).
        target (str): The name of the target space within which activation
            coordinates are represented. By default, MNI.
        kwargs (dict): Additional optional arguments passed to add_features().
    """

    def __init__(
            self, filename, feature_filename=None, masker=None, r=6,
            transform=True, target='MNI', **kwargs):

        # Instance properties
        self.r = r

        # Set up transformations between different image spaces
        if transform:
            if not isinstance(transform, dict):
                transform = {'T88': transformations.t88_to_mni(),
                             'TAL': transformations.t88_to_mni()
                             }
            self.transformer = transformations.Transformer(transform, target)
        else:
            self.transformer = None

        # Load mappables
        self.mappables = self._load_mappables_from_txt(filename)

        # Load the volume into a new Masker
        if masker is None:
            resource_dir = os.path.join(os.path.dirname(__file__),
                                        os.path.pardir,
                                        'resources')
            masker = os.path.join(
                resource_dir, 'MNI152_T1_2mm_brain.nii.gz')
        self.masker = mask.Masker(masker)

        # Create supporting tables for images and features
        self.create_image_table()
        if feature_filename is not None:
            self.add_features(feature_filename, **kwargs)

    def _load_mappables_from_txt(self, filename):
        """ Load mappables from a text file.

        Args:
            filename (str): a string pointing to the location of the txt file
                to read from.
        """
        logger.info("Loading mappables from %s..." % filename)

        # Read in with pandas
        contents = pd.read_csv(filename, sep='\t')
        contents.columns = [col.lower() for col in list(contents.columns)]

        # Make sure all mandatory columns exist
        mc = ['x', 'y', 'z', 'id', 'space']
        if (set(mc) - set(list(contents.columns))):
            logger.error(
                "At least one of mandatory columns (x, y, z, id, and space) "
                "is missing from input file.")
            return

        # Initialize all mappables--for now, assume Articles are passed
        logger.info("Loading study data from database file...")
        return list(contents.groupby('id', as_index=False).apply(
            lambda x: mappable.Article(x, self.transformer)))

    def create_image_table(self, r=None):
        """ Create and store a new ImageTable instance based on the current
        Dataset. Will generally be called privately, but may be useful as a
        convenience method in cases where the user wants to re-generate the
        table with a new smoothing kernel of different radius.

        Args:
            r (int): An optional integer indicating the radius of the smoothing
                kernel. By default, this is None, which will keep whatever
                value is currently set in the Dataset instance.
        """
        logger.info("Creating image table...")
        if r is not None:
            self.r = r
        self.image_table = ImageTable(self)

    def add_mappables(self, filename=None, mappables=None, remap=True):
        """ Append new Mappable objects to the end of the list.

        Either a filename or a list of mappables must be passed.

        Args:
            filename (str): The name of the file containing new mappables.
            mappables (list): Mappable instances to append to existing list.
            remap (bool): Optional flag indicating whether to regenerate the
                entire ImageTable after appending the new Mappables.
        """
        # TODO: (i) it would be more effiicent to only map the new Mappables
        # into the ImageTable instead of redoing everything. (ii) we should
        # check for duplicates and prompt whether to overwrite or update in
        # cases where conflicts occur.
        if filename != None:
            self.mappables.extend(self._load_mappables_from_txt(filename))
        elif mappables != None:
            self.mappables.extend(mappables)
        if remap:
            self.create_image_table()

    def delete_mappables(self, ids, remap=True):
        """ Delete specific Mappables from the Dataset.

        Note that 'ids' is a list of unique identifiers of the Mappables (e.g.,
        doi's), and not indices in the current instance's mappables list.

        Args:
            ids (list): ids corresponding to the Mappables to delete.
            remap (bool): Optional flag indicating whether to regenerate the
                entire ImageTable after deleting undesired Mappables.
        """
        self.mappables = [m for m in self.mappables if m.id not in ids]
        if remap:
            self.create_image_table()

    def get_studies(self, features=None, expression=None, mask=None,
                    peaks=None, frequency_threshold=0.001,
                    activation_threshold=0.0, func=np.sum, return_type='ids',
                    r=6
                    ):
        """ Get IDs or data for studies that meet specific criteria.

        If multiple criteria are passed, the set intersection is returned. For
        example, passing expression='emotion' and mask='my_mask.nii.gz' would
        return only those studies that are associated with emotion AND report
        activation within the voxels indicated in the passed image.

        Args:
            ids (list): A list of IDs of studies to retrieve.
            features (list or str): The name of a feature, or a list of
                features, to use for selecting studies.
            expression (str): A string expression to pass to the PEG for study
                retrieval.
            mask: the mask image (see Masker documentation for valid data
                types).
            peaks (ndarray or list): Either an n x 3 numpy array, or a list of
                lists (e.g., [[-10, 22, 14]]) specifying the world (x/y/z)
                coordinates of the target location(s).
            frequency_threshold (float): For feature-based or expression-based
                selection, the threshold for selecting studies--i.e., the
                cut-off for a study to be included. Must be a float in range
                [0, 1].
            activation_threshold (int or float): For mask-based or peak-based
                selection, threshold for a study to be included based on
                amount of activation displayed. If an integer, represents the
                absolute number of voxels that must be active within the mask
                (or generated ROIs) in order for a study to be selected. If a
                float, it represents the proportion of voxels that must be
                active.
            func (Callable): The function to use when aggregating over the list
                of features. See documentation in FeatureTable.get_ids() for a
                full explanation. Only used for feature- or expression-based
                selection.
            return_type (str): A string specifying what data to return. Valid
                options are:
                'ids': returns a list of IDs of selected studies.
                'images': returns a voxel x mappable matrix of data for all
                selected studies.
                'weights': returns a dict where the keys are study IDs and the
                values are the computed weights. Only valid when performing
                feature-based selection.
            r (int): For peak-based selection, the radius in millimeters of the
                sphere to grow around each peak.

        Returns:
            When return_type is 'ids' (default), returns a list of IDs of the
            selected studies. When return_type is 'data', returns a 2D numpy
            array, with voxels in rows and studies in columns. When return_type
            is 'weights' (valid only for expression-based selection), returns
            a dict, where the keys are study IDs, and the values are the
            computed weights.

        Examples
        --------
        Select all studies tagged with the feature 'emotion':

            >>> ids = dataset.get_studies(features='emotion')

        Select all studies that activate at least 20% of voxels in an amygdala
        mask, and retrieve activation data rather than IDs:

            >>> data = dataset.get_studies(mask='amygdala_mask.nii.gz',
                threshold=0.2, return_type='images')

        Select studies that activate at least 5% of all voxels within 12 mm of
        three specific foci:

            >>> ids = dataset.get_studies(peaks=[[12, -20, 30], [-26, 22, 22],
                                                [0, 36, -20]], r=12)

        """
        results = []

        # Feature-based selection
        if features is not None:
            # Need to handle weights as a special case, because we can't
            # retrieve the weights later using just the IDs.
            if return_type == 'weights':
                if expression is not None or mask is not None or \
                        peaks is not None:
                    raise ValueError(
                        "return_type cannot be 'weights' when feature-based "
                        "search is used in conjunction with other search "
                        "modes.")
                return self.feature_table.get_ids(
                    features, frequency_threshold, func, get_weights=True)
            else:
                results.append(self.feature_table.get_ids(
                    features, frequency_threshold, func))

        # Logical expression-based selection
        if expression is not None:
            _ids = self.feature_table.get_ids_by_expression(
                expression, frequency_threshold, func)
            results.append(list(_ids))

        # Mask-based selection
        if mask is not None:
            mask = self.masker.mask(mask, in_global_mask=True).astype(bool)
            num_vox = np.sum(mask)
            prop_mask_active = self.image_table.data.T.dot(mask).astype(float)
            if isinstance(activation_threshold, float):
                prop_mask_active /= num_vox
            indices = np.where(prop_mask_active > activation_threshold)[0]
            results.append([self.image_table.ids[ind] for ind in indices])

        # Peak-based selection
        if peaks is not None:
            peaks = np.array(peaks)  # Make sure we have a numpy array
            peaks = transformations.xyz_to_mat(peaks)
            m = self.masker
            img = imageutils.map_peaks_to_image(
                peaks, r, vox_dims=m.vox_dims, dims=m.dims,
                header=m.get_header())
            results.append(self.get_studies(
                mask=img, activation_threshold=activation_threshold))

        # Get intersection of all sets
        ids = list(reduce(lambda x, y: set(x) & set(y), results))

        if return_type == 'ids':
            return ids
        elif return_type == 'data':
            return self.get_image_data(ids)

    @deprecated("get_mappables() is deprecated and will be removed in 0.5. "
                "Please use get_studies().")
    def get_mappables(self, ids, get_image_data=False):
        """ Takes a list of unique ids and returns corresponding Mappables.

        Args:
            ids (list): A list of ids of the mappables to return.
            get_image_data (bool): When True, returns a voxel x mappable matrix
            of image data rather than the Mappable instances themselves.

        Returns:
            If get_image_data is True, a 2D numpy array of voxels x Mappables.
            Otherwise, a list of Mappables.
        """
        if get_image_data:
            return self.get_image_data(ids)
        else:
            return [m for m in self.mappables if m.id in ids]

    @deprecated("get_ids_by_features() is deprecated and will be removed in "
                "0.5. Please use get_studies(features=...).")
    def get_ids_by_features(self, features, threshold=0.001, func=np.sum,
                            get_image_data=False, get_weights=False,
                            dense=True):
        """ A wrapper for FeatureTable.get_ids().

        Args:
            features (list): features to use when selecting Mappables.
            threshold (float): Float in range 0-1. Threshold used to select
                mappables.
            func (Callable): The function to use when aggregating over the list
                of features. See documentation in FeatureTable.get_ids() for a
                full explanation.
            get_image_data (bool): When True, returns a voxel x mappable matrix
            of image data rather than the Mappable instances themselves.
        """
        ids = self.feature_table.get_ids(
            features, threshold, func, get_weights)
        return self.get_image_data(ids, dense=dense) if get_image_data else ids

    @deprecated("get_ids_by_expression() is deprecated and will be removed in "
                "0.5. Please use get_studies(expression=...).")
    def get_ids_by_expression(self, expression, threshold=0.001, func=np.sum,
                              get_image_data=False):
        ids = self.feature_table.get_ids_by_expression(
            expression, threshold, func)
        return self.get_image_data(ids) if get_image_data else ids

    @deprecated("get_ids_by_mask() is deprecated and will be removed in "
                "0.5. Please use get_studies(mask=...).")
    def get_ids_by_mask(self, mask, threshold=0.0, get_image_data=False):
        """ Return all mappable objects that activate within the bounds
        defined by the mask image. 
        Args:
            mask: the mask image (see Masker documentation for valid data 
                types).
            threshold (int, float): an integer or float. If an integer, the 
                absolute number of voxels that must be active within the mask
                for a study to be retained. When a float, proportion of voxels
                that must be active.
            get_image_data (bool): When True, returns a voxel x mappable matrix
            of image data rather than the Mappable instances themselves.
        """
        mask = self.masker.mask(mask).astype(bool)
        num_vox = np.sum(mask)
        prop_mask_active = self.image_table.data.T.dot(mask).astype(float)
        if isinstance(threshold, float):
            prop_mask_active /= num_vox
        indices = np.where(prop_mask_active > threshold)[0]
        if get_image_data:
            return self.get_image_data(indices)
        else:
            return [self.image_table.ids[ind] for ind in indices]

    @deprecated("get_ids_by_peaks() is deprecated and will be removed in "
                "0.5. Please use get_studies(peaks=...).")
    def get_ids_by_peaks(self, peaks, r=10, threshold=0.0,
                         get_image_data=False):
        """ A wrapper for get_ids_by_mask. Takes a set of xyz coordinates and
        generates a new Nifti1Image to use as a mask.

        Args:
            peaks (ndarray, list): Either an n x 3 numpy array, or a list of 
                lists (e.g., [[-10, 22, 14]]) specifying the world (x/y/z)
                coordinates of the target location(s).
            r (int): Radius in millimeters of the sphere to grow around each
                location.
          threshold (float): Optional float indicating the proportion of voxels
            that must be active in order for a Mappable to be considered
            active.
            get_image_data (bool): When True, returns a voxel x mappable matrix
            of image data rather than the Mappable instances themselves.

        Returns:
            Either a list of ids (if get_image_data = False) or a numpy array
            of image data.

        """
        peaks = np.array(peaks)  # Make sure we have a numpy array
        peaks = transformations.xyz_to_mat(peaks)
        img = imageutils.map_peaks_to_image(
            peaks, r, vox_dims=self.masker.vox_dims,
            dims=self.masker.dims, header=self.masker.get_header())
        return self.get_ids_by_mask(img, threshold,
                                    get_image_data=get_image_data)

    def add_features(self, features, append=True, merge='outer',
                     duplicates='ignore', min_studies=0.0, threshold=0.001):
        """ Construct a new FeatureTable from file.

        Args:
            features: Feature data to add. Can be:
                (a) A text file containing the feature data, where each row is
                a study in the database, with features in columns. The first
                column must contain the IDs of the studies to match up with the
                image data.
                (b) A pandas DataFrame, where studies are in rows, features are 
                in columns, and the index provides the study IDs.
            append (bool): If True, adds new features to existing ones
                incrementally. If False, replaces old features.
            merge, duplicates, min_studies, threshold: Additional arguments
                passed to FeatureTable.add_features().
         """
        if (not append) or not hasattr(self, 'feature_table'):
            self.feature_table = FeatureTable(self)

        self.feature_table.add_features(features, merge=merge,
                                        duplicates=duplicates,
                                        min_studies=min_studies,
                                        threshold=threshold)

    def get_image_data(self, ids=None, voxels=None, dense=True):
        """ A convenience wrapper for ImageTable.get_image_data(). 

        Args:
            ids (list, array): A list or 1D numpy array of Mappable ids to 
                return. If None, returnsdata for all Mappables.
            voxels (list, array): A list or 1D numpy array of voxel indices
                (i.e., rows) to return. If None, returns data for all voxels.
        """
        return self.image_table.get_image_data(ids, voxels=voxels, dense=dense)

    def get_feature_data(self, ids=None, **kwargs):
        """ A convenience wrapper for FeatureTable.get_image_data(). """
        return self.feature_table.get_feature_data(ids, **kwargs)

    def get_feature_names(self, features=None):
        """ Returns names of features. If features is None, returns all 
        features. Otherwise assumes the user is trying to find the order of the
        features.  """
        if features:
            return self.feature_table.get_ordered_names(features)
        else:
            return self.feature_table.feature_names

    def get_feature_counts(self, threshold=0.001):
        """ Returns a dictionary, where the keys are the feature names
        and the values are the number of studies tagged with the feature. """
        counts = np.sum(self.get_feature_data() >= threshold, 0)
        return dict(zip(self.get_feature_names(), list(counts)))

    @classmethod
    def load(cls, filename):
        """ Load a pickled Dataset instance from file. """
        dataset = pickle.load(open(filename, 'rb'))
        if hasattr(dataset, 'feature_table'):
            dataset.feature_table._csr_to_sdf()
        return dataset

    def save(self, filename):
        """ Pickle the Dataset instance to the provided file.
        """
        if hasattr(self, 'feature_table'):
            self.feature_table._sdf_to_csr()

        pickle.dump(self, open(filename, 'wb'), -1)

        if hasattr(self, 'feature_table'):
            self.feature_table._csr_to_sdf()

    def to_json(self, filename=None):
        """ Save the Dataset to file in JSON format.

        This is not recommended, as the resulting file will typically be
        several GB in size. If no filename is provided, returns the JSON
        string.
        """
        import json
        mappables = [m.to_json() for m in self.mappables]
{'fraction_non_alphanumeric': 0.07279770792139138, 'fraction_numerical': 0.0054073685484847265, 'mean_word_length': 2.0731026785714284, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 27, 'https://': 2, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 24, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '6929207', 'n_tokens_mistral': 13183, 'n_tokens_neox': 12529, 'n_words': 4942}