<commit_before>#!/usr/bin/env python
from __future__ import print_function

# standard lib
import glob
import itertools
import json
import math
import os
import random
import sys

# third party
import numpy as np
from scipy.spatial.distance import squareform

# treeCl
from .alignment import Alignment
from .concatenation import Concatenation
from .constants import SORT_KEY, PLL_RANDOM_SEED
from .distance_matrix import DistanceMatrix
from .errors import optioncheck, directorycheck
from . import tasks
from .parameters import PartitionParameters
from .partition import Partition
from .parutils import get_client, parallel_map, sequential_map
from .tree import Tree
from .utils import fileIO, setup_progressbar, model_translate
from .utils.decorators import lazyprop


def gapmask(simseqs, origseqs):
    """
    :param sims: list of (header, sequence) tuples of simulated sequences [no gaps]
    :param aln: list of (header, sequence) tuples of original sequences
    :return:
    """
    import numpy as np
    simdict = dict(simseqs)
    origdict = dict(origseqs)
    for k in origdict:
        origseq = np.array(list(origdict[k]))
        gap_pos = np.where(origseq=='-')
        simseq = np.array(list(simdict[k]))
        simseq[gap_pos] = '-'
        simdict[k] = ''.join(simseq)
    return list(simdict.items())


class NoRecordsError(Exception):
    def __init__(self, file_format, input_dir, compression):
        self.file_format = file_format
        self.input_dir = input_dir
        self.compression = compression

    def __str__(self):
        msg = ('No records were found in {0} matching\n'
               '\tfile_format = {1}\n'
               '\tcompression = {2}'.format(self.input_dir,
                                            self.file_format, self.compression))
        return msg


class Collection(object):
    """ Call:

    c = Collection(input_dir, file_format, datatype, tmpdir ...)
    c.calc_distances(), c.calc_TC_trees(), ...
    dm = c.distance_matrix('geo')
    cl = Clustering(dm)
    k = cl.spectral(4, prune='estimate', local_scale=7)
    p = Partition(k) """

    def __init__(
            self,
            records=None,
            input_dir=None,
            param_dir=None,
            file_format='fasta',
            compression=None,
            header_grep=None,
            show_progressbars=True,
    ):

        self._records = None
        self._input_files = None
        self.show_progressbars=show_progressbars

        if records is not None:
            self.records = records

        elif input_dir is not None:
            input_dir = os.path.abspath(input_dir)
            directorycheck(input_dir)
            optioncheck(file_format, ['fasta', 'phylip'])
            self.records = self.read_alignments(input_dir,
                                                file_format,
                                                header_grep,
                                                compression)

        else:
            raise Exception('Provide a list of records, '
                            'or the path to a set of alignments')

        if param_dir is not None:
            self.read_parameters(param_dir)

        if not self.records:
            raise NoRecordsError(file_format, input_dir, compression)


    def __len__(self):
        if hasattr(self, 'records'):
            return len(self.records)
        return 0

    def __getitem__(self, i):
        if hasattr(self, 'records'):
            return self.records[i]

    @property
    def records(self):
        """ Returns a list of records in SORT_KEY order """
        return [self._records[i] for i in range(len(self._records))]

    @records.setter
    def records(self, records):
        """ Sets a dictionary of records keyed by SORT_KEY order """
        self._records = dict(enumerate(records))

    @lazyprop
    def trees(self):
        """ Returns a list of trees in SORT_KEY order """
        try:
            return [rec.tree for rec in self]
        except ValueError:
            return []

    @lazyprop
    def names(self):
        """
        Returns a list of sequence record names in SORT_KEY order
        """
        try:
            return [rec.name for rec in self]
        except ValueError:
            return []

    @lazyprop
    def distances(self):
        try:
            return [rec.parameters.partitions.distances for rec in self]
        except ValueError:
            return []

    @lazyprop
    def variances(self):
        try:
            return [rec.parameters.partitions.variances for rec in self]
        except ValueError:
            return []

    @lazyprop
    def frequencies(self):
        try:
            return [rec.parameters.partitions.frequencies for rec in self]
        except ValueError:
            return []

    @lazyprop
    def alphas(self):
        try:
            return [rec.parameters.partitions.alpha for rec in self]
        except ValueError:
            return []

    @lazyprop
    def datatypes(self):
        try:
            return ['dna' if rec.is_dna() else 'protein' for rec in self]
        except ValueError:
            return []

    @lazyprop
    def lengths(self):
        try:
            return [len(rec) for rec in self]
        except ValueError:
            return []

    @lazyprop
    def headers(self):
        try:
            return [rec.get_names() for rec in self]
        except ValueError:
            return []

    @lazyprop
    def mrp_tree(self):
        trees = [tree.newick if hasattr('newick', tree) else tree for tree in self.trees]
        return Alignment().get_mrp_supertree(trees)

    def num_species(self):
        """ Returns the number of species found over all records
        """
        all_headers = reduce(lambda x, y: set(x) | set(y),
                             (rec.get_names() for rec in self.records))
        return len(all_headers)

    def read_alignments(self, input_dir, file_format, header_grep=None, compression=None):
        """ Get list of alignment files from an input directory *.fa, *.fas and
        *.phy files only

        Stores in self.files """

        optioncheck(compression, [None, 'gz', 'bz2'])

        if file_format == 'fasta':
            extensions = ['fa', 'fas', 'fasta']

        elif file_format == 'phylip':
            extensions = ['phy']

        else:
            extensions = []

        if compression:
            extensions = ['.'.join([x, compression]) for x in extensions]

        files = fileIO.glob_by_extensions(input_dir, extensions)
        files.sort(key=SORT_KEY)
        self._input_files = files
        records = []

        if self.show_progressbars:
            pbar = setup_progressbar("Loading files", len(files), simple_progress=True)
            pbar.start()

        for i, f in enumerate(files):
            if compression is not None:
                with fileIO.TempFile() as tmpfile:
                    with fileIO.freader(f, compression) as reader, fileIO.fwriter(tmpfile) as writer:
                        for line in reader:
                            writer.write(line)
                    try:
                        record = Alignment(tmpfile, file_format, True)
                    except RuntimeError:
                        record = Alignment(tmpfile, file_format, False)

            else:
                try:
                    record = Alignment(f, file_format, True)
                except RuntimeError:
                    record = Alignment(f, file_format, False)

            if header_grep:
                try:
                    datatype = 'dna' if record.is_dna() else 'protein'

                    record = Alignment([(header_grep(x), y) for (x, y) in record.get_sequences()], datatype)

                except TypeError:
                    raise TypeError("Couldn't apply header_grep to header\n"
                                    "alignment number={}, name={}\n"
                                    "header_grep={}".format(i, fileIO.strip_extensions(f), header_grep))
                except RuntimeError:
                    print('RuntimeError occurred processing alignment number={}, name={}'
                          .format(i, fileIO.strip_extensions(f)))
                    raise

            record.name = (fileIO.strip_extensions(f))
            records.append(record)
            if self.show_progressbars:
                pbar.update(i)
        if self.show_progressbars:
            pbar.finish()
        return records

    def read_parameters(self, input_dir):
        """ Read a directory full of tree files, matching them up to the
        already loaded alignments """

        if self.show_progressbars:
            pbar = setup_progressbar("Loading parameters", len(self.records))
            pbar.start()
        for i, rec in enumerate(self.records):
            hook = os.path.join(input_dir, '{}.json*'.format(rec.name))
            filename = glob.glob(hook)
            try:
                with fileIO.freader(filename[0]) as infile:
                    d = json.load(infile, parse_int=True)

                rec.parameters.construct_from_dict(d)

            except IOError, IndexError:
                continue

            finally:
                if self.show_progressbars:
                    pbar.update(i)
        if self.show_progressbars:
            pbar.finish()

    def write_parameters(self, output_dir, gz=False):
        if not os.path.exists(output_dir):
            try:
                os.makedirs(output_dir)
            except IOError as err:
                sys.stderr.write(err.message)
                raise err

        for rec in self.records:
            with fileIO.fwriter(os.path.join(output_dir, '{}.json'.format(rec.name)), gz=True) as outfile:
                rec.parameters.write(outfile, indent=4)

    def permuted_copy(self, partition=None):
        """ Return a copy of the collection with all alignment columns permuted
        """
        def take(n, iterable):
            return [iterable.next() for _ in range(n)]

        if partition is None:
            partition = Partition([1] * len(self))

        index_tuples = partition.get_membership()

        alignments = []
        for ix in index_tuples:
            concat = Concatenation(self, ix)
            sites = concat.alignment.get_sites()
            random.shuffle(sites)
            d = dict(zip(concat.alignment.get_names(), [iter(x) for x in zip(*sites)]))
            new_seqs = [[(k, ''.join(take(l, d[k]))) for k in d] for l in concat.lengths]

            for seqs, datatype, name in zip(new_seqs, concat.datatypes, concat.names):
                alignment = Alignment(seqs, datatype)
                alignment.name = name
                alignments.append(alignment)

        return self.__class__(records=sorted(alignments, key=lambda x: SORT_KEY(x.name)))

    def fast_calc_distances(self, batchsize=1, background=False):
        """
        Calculate fast approximate intra-alignment pairwise distances and variances using
        Jukes-Cantor closed formulae.
        :return: None (all side effects)
        """
        # Assemble argument lists
        args = []
        to_delete = []
        for rec in self:
            filename, delete = rec.get_alignment_file(as_phylip=True)
            if delete:
                to_delete.append(filename)
            args.append((filename,))

        # Dispatch work (either sequentially or in parallel)
        msg = 'Calculating fast distances'
        with fileIO.TempFileList(to_delete):
            client = get_client()
            if client is None:
                map_result = sequential_map(tasks.fast_calc_distances_task, args, msg)
            else:
                map_result = parallel_map(client, tasks.fast_calc_distances_task, args, msg, batchsize, background)
                if background:
                    return map_result

        # Process results
        pbar = setup_progressbar('Processing results', len(map_result))
        j = 0
        pbar.start()
        for i, result in enumerate(map_result):
            rec = self[i]
            distances = result['distances']
            variances = result['variances']
            tree = result['tree']
            rec.parameters.nj_tree = tree
            params = rec.parameters.partitions
            if params is None:
                params = PartitionParameters()
                rec.parameters.partitions = [params]
            params.distances = distances
            params.variances = variances
            pbar.update(i)
        pbar.finish()

    def calc_distances(self, batchsize=1, background=False):
        """
        Calculate fast approximate intra-alignment pairwise distances and variances using
        ML (requires ML models to have been set up using `calc_trees`).
        :return: None (all side effects)
        """
        # Assemble argument lists
        args = []
        to_delete = []
        for rec in self:
            filename, delete = rec.get_alignment_file(as_phylip=True)
            if delete:
                to_delete.append(filename)
            # Get input dict
            model = {'partitions': {}}
            data = {'alpha': rec.parameters.partitions.alpha, 'frequencies': rec.parameters.partitions.frequencies}
            if rec.is_dna():
                data['rates'] = rec.parameters.partitions.rates
            model['partitions'][0] = data
            args.append((model, filename))

        # Dispatch
        msg = 'Calculating ML distances'
        client = get_client()
        if client is None:
            map_result = sequential_map(tasks.calc_distances_task, args, msg)
        else:
            map_result = parallel_map(client, tasks.calc_distances_task, args, msg, batchsize, background)
            if background:
                return map_result

        # Process results
        with fileIO.TempFileList(to_delete):
            pbar = setup_progressbar('Processing results', len(map_result))
            j = 0
            pbar.start()
            for i, result in enumerate(map_result):
                rec = self[i]
                rec.parameters.partitions.distances = result['partitions'][0]['distances']
                rec.parameters.partitions.variances = result['partitions'][0]['variances']
                rec.parameters.nj_tree = result['nj_tree']
                pbar.update(j+1)
                j += 1

    def calc_trees(self, model=None, threads=1, indices=None, batchsize=1, output_dir=None, background=False):
        """
        Use pllpy to calculate maximum-likelihood trees
        :return: None (all side effects)
        """
        # Assemble argument lists
        if indices is None:
            indices = list(range(len(self)))
        args = []
        to_delete = []
        for i in indices:
            rec = self[i]
            filename, delete = rec.get_alignment_file(as_phylip=True)
            if delete:
                to_delete.append(filename)
            if model is None:
                model = ('DNA' if rec.is_dna() else 'LGX')
            if model == 'AUTOX':
                model = 'AUTO'
            partition = '{}, {} = 1 - {}'.format(model, rec.name, len(rec))
            tree = rec.parameters.nj_tree if rec.parameters.nj_tree is not None else True
            if output_dir is not None and os.path.isdir(output_dir):
                output_file = os.path.join(output_dir, '{}.json'.format(rec.name))
                curr_args = (filename, partition, tree, threads, PLL_RANDOM_SEED, None, output_file)
            else:
                curr_args = (filename, partition, tree, threads, PLL_RANDOM_SEED)
            args.append(curr_args)

        # Dispatch work
        msg = 'Calculating ML trees'
        client = get_client()
        if client is None:
            map_result = sequential_map(tasks.pll_task, args, msg)
        else:
            map_result = parallel_map(client, tasks.pll_task, args, msg, batchsize, background)
            if background:
                return map_result

        # Process results
        with fileIO.TempFileList(to_delete):
            pbar = setup_progressbar('Processing results', len(map_result))
            j = 0
            pbar.start()
            for i, result in zip(indices, map_result):
                rec = self[i]
                rec.parameters.construct_from_dict(result)
                pbar.update(j+1)
                j += 1

    def concatenate(self, indices):<commit_msg>Added missing pbar.finish() calls to calc_trees and calc_distances<commit_after>#!/usr/bin/env python
from __future__ import print_function

# standard lib
import glob
import itertools
import json
import math
import os
import random
import sys

# third party
import numpy as np
from scipy.spatial.distance import squareform

# treeCl
from .alignment import Alignment
from .concatenation import Concatenation
from .constants import SORT_KEY, PLL_RANDOM_SEED
from .distance_matrix import DistanceMatrix
from .errors import optioncheck, directorycheck
from . import tasks
from .parameters import PartitionParameters
from .partition import Partition
from .parutils import get_client, parallel_map, sequential_map
from .tree import Tree
from .utils import fileIO, setup_progressbar, model_translate
from .utils.decorators import lazyprop


def gapmask(simseqs, origseqs):
    """
    :param sims: list of (header, sequence) tuples of simulated sequences [no gaps]
    :param aln: list of (header, sequence) tuples of original sequences
    :return:
    """
    import numpy as np
    simdict = dict(simseqs)
    origdict = dict(origseqs)
    for k in origdict:
        origseq = np.array(list(origdict[k]))
        gap_pos = np.where(origseq=='-')
        simseq = np.array(list(simdict[k]))
        simseq[gap_pos] = '-'
        simdict[k] = ''.join(simseq)
    return list(simdict.items())


class NoRecordsError(Exception):
    def __init__(self, file_format, input_dir, compression):
        self.file_format = file_format
        self.input_dir = input_dir
        self.compression = compression

    def __str__(self):
        msg = ('No records were found in {0} matching\n'
               '\tfile_format = {1}\n'
               '\tcompression = {2}'.format(self.input_dir,
                                            self.file_format, self.compression))
        return msg


class Collection(object):
    """ Call:

    c = Collection(input_dir, file_format, datatype, tmpdir ...)
    c.calc_distances(), c.calc_TC_trees(), ...
    dm = c.distance_matrix('geo')
    cl = Clustering(dm)
    k = cl.spectral(4, prune='estimate', local_scale=7)
    p = Partition(k) """

    def __init__(
            self,
            records=None,
            input_dir=None,
            param_dir=None,
            file_format='fasta',
            compression=None,
            header_grep=None,
            show_progressbars=True,
    ):

        self._records = None
        self._input_files = None
        self.show_progressbars=show_progressbars

        if records is not None:
            self.records = records

        elif input_dir is not None:
            input_dir = os.path.abspath(input_dir)
            directorycheck(input_dir)
            optioncheck(file_format, ['fasta', 'phylip'])
            self.records = self.read_alignments(input_dir,
                                                file_format,
                                                header_grep,
                                                compression)

        else:
            raise Exception('Provide a list of records, '
                            'or the path to a set of alignments')

        if param_dir is not None:
            self.read_parameters(param_dir)

        if not self.records:
            raise NoRecordsError(file_format, input_dir, compression)


    def __len__(self):
        if hasattr(self, 'records'):
            return len(self.records)
        return 0

    def __getitem__(self, i):
        if hasattr(self, 'records'):
            return self.records[i]

    @property
    def records(self):
        """ Returns a list of records in SORT_KEY order """
        return [self._records[i] for i in range(len(self._records))]

    @records.setter
    def records(self, records):
        """ Sets a dictionary of records keyed by SORT_KEY order """
        self._records = dict(enumerate(records))

    @lazyprop
    def trees(self):
        """ Returns a list of trees in SORT_KEY order """
        try:
            return [rec.tree for rec in self]
        except ValueError:
            return []

    @lazyprop
    def names(self):
        """
        Returns a list of sequence record names in SORT_KEY order
        """
        try:
            return [rec.name for rec in self]
        except ValueError:
            return []

    @lazyprop
    def distances(self):
        try:
            return [rec.parameters.partitions.distances for rec in self]
        except ValueError:
            return []

    @lazyprop
    def variances(self):
        try:
            return [rec.parameters.partitions.variances for rec in self]
        except ValueError:
            return []

    @lazyprop
    def frequencies(self):
        try:
            return [rec.parameters.partitions.frequencies for rec in self]
        except ValueError:
            return []

    @lazyprop
    def alphas(self):
        try:
            return [rec.parameters.partitions.alpha for rec in self]
        except ValueError:
            return []

    @lazyprop
    def datatypes(self):
        try:
            return ['dna' if rec.is_dna() else 'protein' for rec in self]
        except ValueError:
            return []

    @lazyprop
    def lengths(self):
        try:
            return [len(rec) for rec in self]
        except ValueError:
            return []

    @lazyprop
    def headers(self):
        try:
            return [rec.get_names() for rec in self]
        except ValueError:
            return []

    @lazyprop
    def mrp_tree(self):
        trees = [tree.newick if hasattr('newick', tree) else tree for tree in self.trees]
        return Alignment().get_mrp_supertree(trees)

    def num_species(self):
        """ Returns the number of species found over all records
        """
        all_headers = reduce(lambda x, y: set(x) | set(y),
                             (rec.get_names() for rec in self.records))
        return len(all_headers)

    def read_alignments(self, input_dir, file_format, header_grep=None, compression=None):
        """ Get list of alignment files from an input directory *.fa, *.fas and
        *.phy files only

        Stores in self.files """

        optioncheck(compression, [None, 'gz', 'bz2'])

        if file_format == 'fasta':
            extensions = ['fa', 'fas', 'fasta']

        elif file_format == 'phylip':
            extensions = ['phy']

        else:
            extensions = []

        if compression:
            extensions = ['.'.join([x, compression]) for x in extensions]

        files = fileIO.glob_by_extensions(input_dir, extensions)
        files.sort(key=SORT_KEY)
        self._input_files = files
        records = []

        if self.show_progressbars:
            pbar = setup_progressbar("Loading files", len(files), simple_progress=True)
            pbar.start()

        for i, f in enumerate(files):
            if compression is not None:
                with fileIO.TempFile() as tmpfile:
                    with fileIO.freader(f, compression) as reader, fileIO.fwriter(tmpfile) as writer:
                        for line in reader:
                            writer.write(line)
                    try:
                        record = Alignment(tmpfile, file_format, True)
                    except RuntimeError:
                        record = Alignment(tmpfile, file_format, False)

            else:
                try:
                    record = Alignment(f, file_format, True)
                except RuntimeError:
                    record = Alignment(f, file_format, False)

            if header_grep:
                try:
                    datatype = 'dna' if record.is_dna() else 'protein'

                    record = Alignment([(header_grep(x), y) for (x, y) in record.get_sequences()], datatype)

                except TypeError:
                    raise TypeError("Couldn't apply header_grep to header\n"
                                    "alignment number={}, name={}\n"
                                    "header_grep={}".format(i, fileIO.strip_extensions(f), header_grep))
                except RuntimeError:
                    print('RuntimeError occurred processing alignment number={}, name={}'
                          .format(i, fileIO.strip_extensions(f)))
                    raise

            record.name = (fileIO.strip_extensions(f))
            records.append(record)
            if self.show_progressbars:
                pbar.update(i)
        if self.show_progressbars:
            pbar.finish()
        return records

    def read_parameters(self, input_dir):
        """ Read a directory full of tree files, matching them up to the
        already loaded alignments """

        if self.show_progressbars:
            pbar = setup_progressbar("Loading parameters", len(self.records))
            pbar.start()
        for i, rec in enumerate(self.records):
            hook = os.path.join(input_dir, '{}.json*'.format(rec.name))
            filename = glob.glob(hook)
            try:
                with fileIO.freader(filename[0]) as infile:
                    d = json.load(infile, parse_int=True)

                rec.parameters.construct_from_dict(d)

            except IOError, IndexError:
                continue

            finally:
                if self.show_progressbars:
                    pbar.update(i)
        if self.show_progressbars:
            pbar.finish()

    def write_parameters(self, output_dir, gz=False):
        if not os.path.exists(output_dir):
            try:
                os.makedirs(output_dir)
            except IOError as err:
                sys.stderr.write(err.message)
                raise err

        for rec in self.records:
            with fileIO.fwriter(os.path.join(output_dir, '{}.json'.format(rec.name)), gz=True) as outfile:
                rec.parameters.write(outfile, indent=4)

    def permuted_copy(self, partition=None):
        """ Return a copy of the collection with all alignment columns permuted
        """
        def take(n, iterable):
            return [iterable.next() for _ in range(n)]

        if partition is None:
            partition = Partition([1] * len(self))

        index_tuples = partition.get_membership()

        alignments = []
        for ix in index_tuples:
            concat = Concatenation(self, ix)
            sites = concat.alignment.get_sites()
            random.shuffle(sites)
            d = dict(zip(concat.alignment.get_names(), [iter(x) for x in zip(*sites)]))
            new_seqs = [[(k, ''.join(take(l, d[k]))) for k in d] for l in concat.lengths]

            for seqs, datatype, name in zip(new_seqs, concat.datatypes, concat.names):
                alignment = Alignment(seqs, datatype)
                alignment.name = name
                alignments.append(alignment)

        return self.__class__(records=sorted(alignments, key=lambda x: SORT_KEY(x.name)))

    def fast_calc_distances(self, batchsize=1, background=False):
        """
        Calculate fast approximate intra-alignment pairwise distances and variances using
        Jukes-Cantor closed formulae.
        :return: None (all side effects)
        """
        # Assemble argument lists
        args = []
        to_delete = []
        for rec in self:
            filename, delete = rec.get_alignment_file(as_phylip=True)
            if delete:
                to_delete.append(filename)
            args.append((filename,))

        # Dispatch work (either sequentially or in parallel)
        msg = 'Calculating fast distances'
        with fileIO.TempFileList(to_delete):
            client = get_client()
            if client is None:
                map_result = sequential_map(tasks.fast_calc_distances_task, args, msg)
            else:
                map_result = parallel_map(client, tasks.fast_calc_distances_task, args, msg, batchsize, background)
                if background:
                    return map_result

        # Process results
        pbar = setup_progressbar('Processing results', len(map_result))
        j = 0
        pbar.start()
        for i, result in enumerate(map_result):
            rec = self[i]
            distances = result['distances']
            variances = result['variances']
            tree = result['tree']
            rec.parameters.nj_tree = tree
            params = rec.parameters.partitions
            if params is None:
                params = PartitionParameters()
                rec.parameters.partitions = [params]
            params.distances = distances
            params.variances = variances
            pbar.update(i)
        pbar.finish()

    def calc_distances(self, batchsize=1, background=False):
        """
        Calculate fast approximate intra-alignment pairwise distances and variances using
        ML (requires ML models to have been set up using `calc_trees`).
        :return: None (all side effects)
        """
        # Assemble argument lists
        args = []
        to_delete = []
        for rec in self:
            filename, delete = rec.get_alignment_file(as_phylip=True)
            if delete:
                to_delete.append(filename)
            # Get input dict
            model = {'partitions': {}}
            data = {'alpha': rec.parameters.partitions.alpha, 'frequencies': rec.parameters.partitions.frequencies}
            if rec.is_dna():
                data['rates'] = rec.parameters.partitions.rates
            model['partitions'][0] = data
            args.append((model, filename))

        # Dispatch
        msg = 'Calculating ML distances'
        client = get_client()
        if client is None:
            map_result = sequential_map(tasks.calc_distances_task, args, msg)
        else:
            map_result = parallel_map(client, tasks.calc_distances_task, args, msg, batchsize, background)
            if background:
                return map_result

        # Process results
        with fileIO.TempFileList(to_delete):
            pbar = setup_progressbar('Processing results', len(map_result))
            j = 0
            pbar.start()
            for i, result in enumerate(map_result):
                rec = self[i]
                rec.parameters.partitions.distances = result['partitions'][0]['distances']
                rec.parameters.partitions.variances = result['partitions'][0]['variances']
                rec.parameters.nj_tree = result['nj_tree']
                pbar.update(j+1)
                j += 1
            pbar.finish()

    def calc_trees(self, model=None, threads=1, indices=None, batchsize=1, output_dir=None, background=False):
        """
        Use pllpy to calculate maximum-likelihood trees
        :return: None (all side effects)
        """
        # Assemble argument lists
        if indices is None:
            indices = list(range(len(self)))
        args = []
        to_delete = []
        for i in indices:
            rec = self[i]
            filename, delete = rec.get_alignment_file(as_phylip=True)
            if delete:
                to_delete.append(filename)
            if model is None:
                model = ('DNA' if rec.is_dna() else 'LGX')
            if model == 'AUTOX':
                model = 'AUTO'
            partition = '{}, {} = 1 - {}'.format(model, rec.name, len(rec))
            tree = rec.parameters.nj_tree if rec.parameters.nj_tree is not None else True
            if output_dir is not None and os.path.isdir(output_dir):
                output_file = os.path.join(output_dir, '{}.json'.format(rec.name))
                curr_args = (filename, partition, tree, threads, PLL_RANDOM_SEED, None, output_file)
            else:
                curr_args = (filename, partition, tree, threads, PLL_RANDOM_SEED)
            args.append(curr_args)

        # Dispatch work
        msg = 'Calculating ML trees'
        client = get_client()
        if client is None:
            map_result = sequential_map(tasks.pll_task, args, msg)
        else:
            map_result = parallel_map(client, tasks.pll_task, args, msg, batchsize, background)
            if background:
                return map_result

        # Process results
        with fileIO.TempFileList(to_delete):
            pbar = setup_progressbar('Processing results', len(map_result))
            j = 0
            pbar.start()
            for i, result in zip(indices, map_result):
                rec = self[i]
                rec.parameters.construct_from_dict(result)
                pbar.update(j+1)
                j += 1
            pbar.finish()

    def concatenate(self, indices):
{'fraction_non_alphanumeric': 0.09283835666132592, 'fraction_numerical': 0.0015115329967653195, 'mean_word_length': 1.8329194142331078, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 22, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '6572429', 'n_tokens_mistral': 9188, 'n_tokens_neox': 8804, 'n_words': 2392}