<commit_before>#   Copyright (c) 2018 PaddlePaddle Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import core
import multiprocessing
import framework
import executor

__all__ = ['ParallelExecutor']

class ParallelExecutor(object):
    def __init__(self,
                 use_cuda,
                 loss_name=None,
                 main_program=None,
                 num_threads=None,
                 allow_op_delay=False,
                 share_vars_from=None):
        """
        ParallelExecutor can run program in parallel.

        Args:
            use_cuda(bool): Whether to use CUDA or not.
            loss_name(str, default None): The loss name must set in training.
            main_program(Program, default None): The program that need to run,
                if not provided, then default_main_program will be used.
            num_threads(int, default None): How many threads are used for
                training.
            allow_op_delay(bool, default False): Whether to delay and buffer
                some operators together for scheduling or not, which may
                improve performance in some cases, defalut False.
            share_vars_from(ParallelExecutor, default None): If provied,
                it will share variables from the specified ParallelExecutor.

        Returns:
            A ParallelExecutor object.

        Raises:
            TypeError: If share_vars_from is provided, but not ParallelExecutor
                object.

        Examples:
            .. code-block:: python

              train_exe = fluid.ParallelExecutor(
                  use_cuda=True, loss_name=loss.name)
              test_exe = fluid.ParallelExecutor(
                  use_cuda=True,
                  main_program=test_program,
                  share_vars_from=train_exe)

              train_loss, = train_exe.run([loss.name], feed_dict=feed_dict)
              test_loss, = test_exe.run([loss.name], feed_dict=feed_dict)
        """

        self._places = []
        self._act_places = []
        if use_cuda:
            for i in xrange(core.get_cuda_device_count()):
                p = core.Place()
                self._act_places.append(core.CUDAPlace(i))
                p.set_place(self._act_places[-1])
                self._places.append(p)
        else:
            for i in xrange(multiprocessing.cpu_count()):
                p = core.Place()
                self._act_places.append(core.CPUPlace(i))
                p.set_place(self._act_places[-1])
                self._places.append(p)
        assert self._places, "no place for execution"

        if num_threads is None:
            if use_cuda:
                # Experiments on se-resnext shows that too many threads hurt
                # performance. Worth tunning for other models in the future.
                num_threads = len(self._places)
            else:
                num_threads = min(
                    len(self._places) * 2, multiprocessing.cpu_count())

        main = main_program
        main = main if main else framework.default_main_program()
        scope = executor.global_scope()

        if share_vars_from and not isinstance(share_vars_from,
                                              ParallelExecutor):
            raise TypeError("share_vars_from must be ParallelExecutor.")
        local_scopes = share_vars_from.executor.local_scopes(
        ) if share_vars_from else []

        self.persistable_vars = [
            v.name
            for v in filter(lambda var: \
                var.persistable and var.type != core.VarDesc.VarType.RAW,
                main.list_vars())
        ]

        self.executor = core.ParallelExecutor(
            num_threads,
            True if use_cuda else False,  # use_event
            self._places,
            set([
                p.name for p in main.global_block().iter_parameters()
                if not p.stop_gradient
            ]),
            set(self.persistable_vars),
            main.desc,
            loss_name if loss_name else '',
            scope,
            local_scopes,
            allow_op_delay)
        self.scope = scope

    def run(self, fetch_list, feed=None, feed_dict=None):
        """

        Args:
            fetch_list(list): The fetched variable names
            feed(list|dict|None): The feed variables. If the feed is a dict,
                tensors in that dict will be splitted into each devices. If
                the feed is a list, each element of the list will be copied
                to each device.
            feed_dict: Alias for feed parameter, for backward compatibility.

        Returns: fetched result list.

        """
        if feed is None:
            feed = feed_dict

        if isinstance(feed, dict):
            feed_tensor_dict = dict()
            for feed_name in feed:
                feed_tensor = feed[feed_name]
                if not isinstance(feed_tensor, core.LoDTensor):
                    feed_tensor = core.LoDTensor()
                    # always set to CPU place, since the tensor need to be splitted
                    # it is fast in CPU
                    feed_tensor.set(feed[feed_name], core.CPUPlace())
                feed_tensor_dict[feed_name] = feed_tensor

            self.executor.feed_and_split_tensor_into_local_scopes(
                feed_tensor_dict)
        elif isinstance(feed, list) or isinstance(feed, tuple):
            if len(feed) != len(self._act_places):
                raise ValueError(
                    "Feed a list of tensor, the list should be the same size as places"
                )

            res = list()

            for i, each in enumerate(feed):
                if not isinstance(each, dict):
                    raise TypeError(
                        "Each element of feed list should be a dict")
                res_dict = dict()
                for feed_name in each:
                    tensor = each[feed_name]
                    if not isinstance(tensor, core.LoDTensor):
                        tmp = core.LoDTensor()
                        tmp.set(tensor, self._act_places[i])
                        tensor = tmp
                    res_dict[feed_name] = tensor
                res.append(res_dict)
            self.executor.feed_tensors_into_local_scopes(res)

        fetch_var_name = '@FETCHED_VAR_NAME@'
        self.executor.run(fetch_list, fetch_var_name)
        arr = self.scope.find_var(fetch_var_name).get_lod_tensor_array()
        return [arr[i] for i in range(len(arr))]

    def bcast_params(self):
        self.executor.bcast_params(set(self.persistable_vars))

    @property
    def device_count(self):
        return len(self._act_places)
<commit_msg>Add more comments<commit_after>#   Copyright (c) 2018 PaddlePaddle Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import core
import multiprocessing
import framework
import executor
import sys

__all__ = ['ParallelExecutor']

class ParallelExecutor(object):
    def __init__(self,
                 use_cuda,
                 loss_name=None,
                 main_program=None,
                 num_threads=None,
                 allow_op_delay=False,
                 share_vars_from=None):
        """
        ParallelExecutor can run program in parallel.

        Args:
            use_cuda(bool): Whether to use CUDA or not.
            loss_name(str, default None): The loss name must set in training.
            main_program(Program, default None): The program that need to run,
                if not provided, then default_main_program will be used.
            num_threads(int, default None): How many threads are used for
                training.
            allow_op_delay(bool, default False): Whether to delay and buffer
                some operators together for scheduling or not, which may
                improve performance in some cases, defalut False.
            share_vars_from(ParallelExecutor, default None): If provied,
                it will share variables from the specified ParallelExecutor.

        Returns:
            A ParallelExecutor object.

        Raises:
            TypeError: If share_vars_from is provided, but not ParallelExecutor
                object.

        Examples:
            .. code-block:: python

              train_exe = fluid.ParallelExecutor(
                  use_cuda=True, loss_name=loss.name)
              test_exe = fluid.ParallelExecutor(
                  use_cuda=True,
                  main_program=test_program,
                  share_vars_from=train_exe)

              train_loss, = train_exe.run([loss.name], feed_dict=feed_dict)
              test_loss, = test_exe.run([loss.name], feed_dict=feed_dict)
        """

        self._places = []
        self._act_places = []
        if use_cuda:
            for i in xrange(core.get_cuda_device_count()):
                p = core.Place()
                self._act_places.append(core.CUDAPlace(i))
                p.set_place(self._act_places[-1])
                self._places.append(p)
        else:
            for i in xrange(multiprocessing.cpu_count()):
                p = core.Place()
                self._act_places.append(core.CPUPlace(i))
                p.set_place(self._act_places[-1])
                self._places.append(p)
        assert self._places, "no place for execution"

        if num_threads is None:
            if use_cuda:
                # Experiments on se-resnext shows that too many threads hurt
                # performance. Worth tunning for other models in the future.
                num_threads = len(self._places)
            else:
                num_threads = min(
                    len(self._places) * 2, multiprocessing.cpu_count())

        main = main_program
        main = main if main else framework.default_main_program()
        scope = executor.global_scope()

        if share_vars_from and not isinstance(share_vars_from,
                                              ParallelExecutor):
            raise TypeError("share_vars_from must be ParallelExecutor.")
        local_scopes = share_vars_from.executor.local_scopes(
        ) if share_vars_from else []

        self.persistable_vars = [
            v.name
            for v in filter(lambda var: \
                var.persistable and var.type != core.VarDesc.VarType.RAW,
                main.list_vars())
        ]

        self.executor = core.ParallelExecutor(
            num_threads,
            True if use_cuda else False,  # use_event
            self._places,
            set([
                p.name for p in main.global_block().iter_parameters()
                if not p.stop_gradient
            ]),
            set(self.persistable_vars),
            main.desc,
            loss_name if loss_name else '',
            scope,
            local_scopes,
            allow_op_delay)
        self.scope = scope

    def run(self, fetch_list, feed=None, feed_dict=None):
        """
        Run a parallel executor with fetch_list.

        The feed parameter can be a dict or a list. If feed is a dict, the
        feed data will be split into multiple devices. If feed is a list, we
        assume the data has been splitted into multiple devices, the each
        element in the list will be copied to each device directly.

        For example, if the feed is a dict:
        >>> exe = ParallelExecutor()
        >>> # the image will be splitted into devices. If there is two devices
        >>> # each device will process an image with shape (24, 1, 28, 28)
        >>> exe.run(feed={'image': numpy.random.random(size=(48, 1, 28, 28))})

        For example, if the feed is a list:
        >>> exe = ParallelExecutor()
        >>> # each device will process each element in the list.
        >>> # the 1st device will process an image with shape (48, 1, 28, 28)
        >>> # the 2nd device will process an image with shape (32, 1, 28, 28)
        >>> #
        >>> # you can use exe.device_count to get the device number.
        >>> exe.run(feed=[{"image": numpy.random.random(size=(48, 1, 28, 28))},
        >>>               {"image": numpy.random.random(size=(32, 1, 28, 28))},
        >>>              ])

        Args:
            fetch_list(list): The fetched variable names
            feed(list|dict|None): The feed variables. If the feed is a dict,
                tensors in that dict will be splitted into each devices. If
                the feed is a list, each element of the list will be copied
                to each device.
            feed_dict: Alias for feed parameter, for backward compatibility.
                This parameter is deprecated.

        Returns: fetched result list.

        """
        if feed is None:
            feed = feed_dict
            print >> sys.stderr, "`feed_dict` is deprecated. Please use `feed=`"

        if isinstance(feed, dict):
            feed_tensor_dict = dict()
            for feed_name in feed:
                feed_tensor = feed[feed_name]
                if not isinstance(feed_tensor, core.LoDTensor):
                    feed_tensor = core.LoDTensor()
                    # always set to CPU place, since the tensor need to be splitted
                    # it is fast in CPU
                    feed_tensor.set(feed[feed_name], core.CPUPlace())
                feed_tensor_dict[feed_name] = feed_tensor

            self.executor.feed_and_split_tensor_into_local_scopes(
                feed_tensor_dict)
        elif isinstance(feed, list) or isinstance(feed, tuple):
            if len(feed) != len(self._act_places):
                raise ValueError(
                    "Feed a list of tensor, the list should be the same size as places"
                )

            res = list()

            for i, each in enumerate(feed):
                if not isinstance(each, dict):
                    raise TypeError(
                        "Each element of feed list should be a dict")
                res_dict = dict()
                for feed_name in each:
                    tensor = each[feed_name]
                    if not isinstance(tensor, core.LoDTensor):
                        tmp = core.LoDTensor()
                        tmp.set(tensor, self._act_places[i])
                        tensor = tmp
                    res_dict[feed_name] = tensor
                res.append(res_dict)
            self.executor.feed_tensors_into_local_scopes(res)

        fetch_var_name = '@FETCHED_VAR_NAME@'
        self.executor.run(fetch_list, fetch_var_name)
        arr = self.scope.find_var(fetch_var_name).get_lod_tensor_array()
        return [arr[i] for i in range(len(arr))]

    def bcast_params(self):
        self.executor.bcast_params(set(self.persistable_vars))

    @property
    def device_count(self):
        return len(self._act_places)
<|endoftext|>


<commit_before>from tota.utils import closest, distance, sort_by_distance, possible_moves
from tota import settings

def create():
    def simple_hero_logic(self, things, t):
        # some useful data about the enemies I can see in the map
        enemy_team = settings.ENEMY_TEAMS[self.team]
        enemies = [thing for thing in things.values()
                   if thing.team == enemy_team]
        closest_enemy = closest(self, enemies)
        closest_enemy_distance = distance(self, closest_enemy)

        if self.life < self.max_life and self.can('heal', t):
            # if I'm hurt and can heal, heal
            return 'heal', self.position
        else:
            # else, try to attack
            if closest_enemy:
                # there is an enemy
                if closest_enemy_distance < settings.STUN_DISTANCE and self.can('stun', t):
                    # try to stun him
                    return 'stun', closest_enemy.position
                elif closest_enemy_distance < settings.FIREBALL_DISTANCE and self.can('fireball', t):
                    # else try to fireball him
                    return 'firebal', closest_enemy.position
                elif closest_enemy_distance < settings.HERO_ATTACK_DISTANCE:
                    # else try to attack him
                    return 'attack', closest_enemy.position
                else:
                    # of finally just move to him
                    moves = sort_by_distance(closest_enemy,
                                             possible_moves(self, things))
                    for move in moves:
                        return 'move', move
            else:
                # no enemies, do nothing
                return None

    return simple_hero_logic
<commit_msg>Simple doesn't fireball himself anymore<commit_after>from tota.utils import closest, distance, sort_by_distance, possible_moves
from tota import settings

def create():
    def simple_hero_logic(self, things, t):
        # some useful data about the enemies I can see in the map
        enemy_team = settings.ENEMY_TEAMS[self.team]
        enemies = [thing for thing in things.values()
                   if thing.team == enemy_team]
        closest_enemy = closest(self, enemies)
        closest_enemy_distance = distance(self, closest_enemy)

        if self.life < self.max_life and self.can('heal', t):
            # if I'm hurt and can heal, heal
            return 'heal', self.position
        else:
            # else, try to attack
            if closest_enemy:
                # there is an enemy
                if closest_enemy_distance < settings.STUN_DISTANCE and self.can('stun', t):
                    # try to stun him
                    return 'stun', closest_enemy.position
                elif closest_enemy_distance < settings.FIREBALL_DISTANCE and self.can('fireball', t) and closest_enemy_distance > settings.FIREBALL_RADIUS:
                    # else try to fireball him, but only if I'm not in range
                    return 'firebal', closest_enemy.position
                elif closest_enemy_distance < settings.HERO_ATTACK_DISTANCE:
                    # else try to attack him
                    return 'attack', closest_enemy.position
                else:
                    # of finally just move to him
                    moves = sort_by_distance(closest_enemy,
                                             possible_moves(self, things))
                    for move in moves:
                        return 'move', move
            else:
                # no enemies, do nothing
                return None

    return simple_hero_logic
<|endoftext|>


<commit_before>import onmt
import argparse
import os
import torch
import torch.nn as nn
from torch.autograd import Variable
import math
import time

parser = argparse.ArgumentParser(description='train.lua')

parser.add_argument('-config', help="Read options from this file")

##
## **Data options**
##

parser.add_argument('-data', help="Path to the training *-train.pt file from preprocess.lua")
parser.add_argument('-save_model', help="""Model filename (the model will be saved as
                              <save_model>_epochN_PPL.pt where PPL is the validation perplexity""")
parser.add_argument('-train_from', help="If training from a checkpoint then this is the path to the pretrained model.")
# parser.add_argument('-cont', action="store_true", help="If training from a checkpoint, whether to continue the training in the same configuration or not.")

##
## **Model options**
##

parser.add_argument('-layers',        type=int, default=2,   help="Number of layers in the LSTM encoder/decoder")
parser.add_argument('-rnn_size',       type=int, default=500, help="Size of LSTM hidden states")
parser.add_argument('-word_vec_size', type=int, default=500, help="Word embedding sizes")
parser.add_argument('-feat_merge', default='concat', help="Merge action for the features embeddings: concat or sum")
parser.add_argument('-feat_vec_exponent', type=float, default=0.7, help="""When using concatenation, if the feature takes N values
                                                                then the embedding dimension will be set to N^exponent""")
parser.add_argument('-feat_vec_size', type=int, default=20, help="When using sum, the common embedding size of the features")
parser.add_argument('-input_feed', type=int,    default=1,  help="Feed the context vector at each time step as additional input (via concatenation with the word embeddings) to the decoder.")
# parser.add_argument('-residual',   action="store_true",     help="Add residual connections between RNN layers.")
parser.add_argument('-brnn',       action="store_true",     help="Use a bidirectional encoder")
parser.add_argument('-brnn_merge', default='concat',        help="Merge action for the bidirectional hidden states: concat or sum")

##
## **Optimization options**
##

parser.add_argument('-max_batch_size',  type=int, default=64,  help="Maximum batch size")
parser.add_argument('-max_generator_batches', type=int, default=16, help="""Maximum batches of words in a sequence to run the generator on in parallel.
                                                                           Higher is faster, but uses more memory.""")
parser.add_argument('-epochs',          type=int, default=13,  help="Number of training epochs")
parser.add_argument('-start_epoch',     type=int, default=0,   help="If loading from a checkpoint, the epoch from which to start")
parser.add_argument('-start_iteration', type=int, default=0,   help="If loading from a checkpoint, the iteration from which to start")
# this gives really bad initialization; Xavier better
parser.add_argument('-param_init',      type=float, default=0.1, help="Parameters are initialized over uniform distribution with support (-param_init, param_init)")
parser.add_argument('-optim', default='sgd', help="Optimization method. Possible options are: sgd, adagrad, adadelta, adam")
parser.add_argument('-learning_rate', type=float, default=1, help="""Starting learning rate. If adagrad/adadelta/adam is used,
                                then this is the global learning rate. Recommed settings. sgd = 1,
                                adagrad = 0.1, adadelta = 1, adam = 0.1""")
parser.add_argument('-max_grad_norm',       type=float, default=5,   help="If the norm of the gradient vector exceeds this renormalize it to have the norm equal to max_grad_norm")
parser.add_argument('-dropout',             type=float, default=0.3, help="Dropout probability. Dropout is applied between vertical LSTM stacks.")
parser.add_argument('-learning_rate_decay', type=float, default=0.5, help="""Decay learning rate by this much if (i) perplexity does not decrease
                                        on the validation set or (ii) epoch has gone past the start_decay_at_limit""")
parser.add_argument('-start_decay_at', default=8, help="Start decay after this epoch")
parser.add_argument('-curriculum', action="store_true", help="""For this many epochs, order the minibatches based on source
                             sequence length. Sometimes setting this to 1 will increase convergence speed.""")
parser.add_argument('-pre_word_vecs_enc', help="""If a valid path is specified, then this will load
                                     pretrained word embeddings on the encoder side.
                                     See README for specific formatting instructions.""")
parser.add_argument('-pre_word_vecs_dec', help="""If a valid path is specified, then this will load
                                     pretrained word embeddings on the decoder side.
                                     See README for specific formatting instructions.""")
# parser.add_argument('-fix_word_vecs_enc', action="store_true", help="Fix word embeddings on the encoder side")
# parser.add_argument('-fix_word_vecs_dec', action="store_true", help="Fix word embeddings on the decoder side")

##
## **Other options**
##

# GPU
parser.add_argument('-cuda', action='store_true', help="Use CUDA")
# parser.add_argument('-nparallel', type=int, default=1,  help="""When using GPUs, how many batches to execute in parallel.
#                             Note. this will technically change the final batch size to max_batch_size*nparallel.""")

# bookkeeping
# parser.add_argument('-save_every', type=int, default=0, help="""Save intermediate models every this many iterations within an epoch.
#                              If = 0, will not save models within an epoch. """)
parser.add_argument('-report_every', type=int, default=50,   help="Print stats every this many iterations within an epoch.")
# parser.add_argument('-seed',         type=int, default=3435, help="Seed for random initialization")
# parser.add_argument('-json_log', action="store_true", help="Outputs logs in JSON format.")

opt = parser.parse_args()
print(opt)

if torch.cuda.is_available() and not opt.cuda:
    print("WARNING: You have a CUDA device, so you should probably run with -cuda")

class NMTCriterion(nn.Container):
    def __init__(self, vocabSize, features):
        self.sub = []
        super(NMTCriterion, self).__init__()

        def makeOne(size):
            weight = torch.ones(vocabSize)
            weight[onmt.Constants.PAD] = 0
            crit = nn.NLLLoss(weight, size_average=False)
            if opt.cuda:
                crit.cuda()
            return crit

        self.sub += [makeOne(vocabSize)]
        for feature in features:
            self.sub += [makeOne(features.size())]

    def forward(self, inputs, targets):
        if len(self.sub) == 1:
            total_size = targets.nelement()
            loss = self.sub[0](inputs.view(total_size, -1), targets.view(total_size))
            return loss
        else:
            assert False, "FIXME: features"
            loss = Variable(inputs.new(1).zero_())
            for sub, input, target in zip(self.sub, inputs, targets):
                loss += sub(input, target)
            return loss

def memoryEfficientLoss(outputs, targets, generator, crit, eval=False):
    # compute generations one piece at a time
    loss = 0
    outputs_rewrapped = Variable(outputs.data, requires_grad=(not eval), volatile=eval)

    batch_size = outputs.size(1)
    chunks = int(math.ceil(targets.size(0) / opt.max_generator_batches))
    outputs_chunked = torch.chunk(outputs_rewrapped, chunks)
    targets_chunked = torch.chunk(targets, chunks)
    for out_t, targ_t in zip(outputs_chunked, targets_chunked):
        out_t = out_t.view(-1, out_t.size(2))
        pred_t = generator(out_t)
        loss_t = crit(pred_t, targ_t)
        loss += loss_t.data[0]
        if not eval:
            loss_t.div(batch_size).backward()

    return loss, outputs_rewrapped.grad

def eval(model, criterion, data):
    total_loss = 0
    total_words = 0

    model.eval()
    for i in range(len(data)):
        batch = data[i]
        outputs = model(batch)
        targets = batch[1][1:]  # exclude <s> from targets
        loss, _ = memoryEfficientLoss(outputs, targets, model.generator, criterion, eval=True)
        total_loss += loss
        total_words += targets.data.ne(onmt.Constants.PAD).sum()

    model.train()
    return total_loss / total_words

def trainModel(model, trainData, validData, dataset):
    print(model)
    model.train()
    for p in model.parameters():
        p.data.uniform_(-opt.param_init, opt.param_init)

    # define criterion of each GPU
    criterion = NMTCriterion(dataset['dicts']['tgt']['words'].size(),
                             dataset['dicts']['tgt']['features'])

    optim = onmt.Optim(
        model.parameters(), opt.optim, opt.learning_rate, opt.max_grad_norm,
        lr_decay=opt.learning_rate_decay,
        start_decay_at=opt.start_decay_at
    )

    # checkpoint = onmt.train.Checkpoint.new(opt, model, optim, dataset)

    def trainEpoch(epoch):

        startI = opt.start_iteration
        opt.start_iteration = 1

        # shuffle mini batch order
        batchOrder = torch.randperm(len(trainData))

        total_loss, report_loss = 0, 0
        total_words, report_words = 0, 0
        start = time.time()
        for i in range(startI, len(trainData)):

            batchIdx = batchOrder[i] if epoch >= opt.curriculum else i
            batch = trainData[batchIdx]
            # srcData, tgtData = batch
            # srcDict = dataset['dicts']['src']['words']
            # tgtDict = dataset['dicts']['tgt']['words']
            # print(srcData)
            # for i in range(srcData.size(1)):
            #     print(' '.join(srcDict.convertToLabels(srcData.data[:,i], onmt.Constants.EOS)))
            #     print(' '.join(tgtDict.convertToLabels(tgtData.data[:,i], onmt.Constants.EOS)))
            #     print()
            # for n, p in model.state_dict().items():
            #     if n not in norms:
            #         norms[n] = []
            #     norms[n] += [p.data.abs().norm()]
            #     print(n, p.data.abs().norm())

            model.zero_grad()
            outputs = model(batch)
            targets = batch[1][1:]  # exclude <s> from targets
            loss, gradOutput = memoryEfficientLoss(outputs, targets, model.generator, criterion)

            outputs.backward(gradOutput)

            # for module in model.modules():
            #     params = list(module.parameters())
            #     if not isinstance(module, nn.Container) and len(params) > 0:
            #         print(module)
            #         for p in params:
            #             print('p', p.data.nelement(), p.data.norm())
            #             print('gp', p.grad.nelement(), p.grad.norm())
            # assert False

            # update the parameters
            grad_norm = optim.step()

            report_loss += loss
            total_loss += loss
            num_words = targets.data.ne(onmt.Constants.PAD).sum()
            total_words += num_words
            report_words += num_words
            if i % opt.report_every == 0 and i > 0:
                print("Epoch %2d, %5d/%5d batches; perplexity: %6.2f; %3.0f tokens/s" %
                      (epoch+1, i, len(trainData), math.exp(report_loss / report_words), report_words/(time.time()-start)))
                report_loss = report_words = 0
                start = time.time()
                # for k,v in norms.items():
                #     print(k,v)

            # if opt.save_every > 0 and ii % opt.save_every == 0:
            #     checkpoint.saveIteration(ii, epochState, batchOrder, not opt.json_log)

        return total_loss / total_words

    for epoch in range(opt.start_epoch, opt.epochs):
        print('')
        train_loss = trainEpoch(epoch)
        print('Train loss: %g' % train_loss)
        print('Train perplexity: %g' % math.exp(min(train_loss, 100)))
        valid_loss = eval(model, criterion, validData)
        print('Validation perplexity: %g' % math.exp(min(valid_loss, 100)))
        if opt.optim == 'sgd':
            optim.updateLearningRate(valid_loss, epoch)

        # checkpoint.saveEpoch(validPpl, epochState, not opt.json_log)

def main():

    checkpoint = {}
    assert opt.train_from is None, "FIXME: Load from checkpoint"
    # if opt.train_from is not None:
    #     assert os.path.exists(opt.train_from), 'checkpoint path invalid'
    #
    #     if not opt.json_log:
    #       print('Loading checkpoint \'' + opt.train_from + '\'...')
    #
    #     checkpoint = torch.load(opt.train_from)
    #
    #     opt.layers = checkpoint.options.layers
    #     opt.rnn_size = checkpoint.options.rnn_size
    #     opt.brnn = checkpoint.options.brnn
    #     opt.brnn_merge = checkpoint.options.brnn_merge
    #     opt.input_feed = checkpoint.options.input_feed
    #
    #     # Resume training from checkpoint
    #     if opt.cont:
    #         opt.optim = checkpoint.options.optim
    #         opt.learning_rate_decay = checkpoint.options.learning_rate_decay
    #         opt.start_decay_at = checkpoint.options.start_decay_at
    #         opt.epochs = checkpoint.options.epochs
    #         opt.curriculum = checkpoint.options.curriculum
    #
    #         opt.learning_rate = checkpoint.info.learning_rate
    #         opt.optim_states = checkpoint.info.optim_states
    #         opt.start_epoch = checkpoint.info.epoch
    #         opt.start_iteration = checkpoint.info.iteration
    #     print('Resuming training from epoch %d at iteration %d...' % (opt.start_epoch, opt.start_iteration))

    print("Loading data from '%s'" % opt.data)

    dataset = torch.load(opt.data)

    trainData = onmt.Dataset(dataset['train']['src'], dataset['train']['tgt'], opt.max_batch_size, opt.cuda)
    validData = onmt.Dataset(dataset['valid']['src'], dataset['valid']['tgt'], opt.max_batch_size, opt.cuda)
    print(' * vocabulary size. source = %d; target = %d' %
            (dataset['dicts']['src']['words'].size(), dataset['dicts']['tgt']['words'].size()))
    print(' * additional features. source = %d; target = %d' %
            (len(dataset['dicts']['src']['features']), len(dataset['dicts']['tgt']['features'])))
    print(' * number of training sentences. %d' % len(dataset['train']['src']['words']))
    print(' * maximum batch size. %d' % opt.max_batch_size)

    print('Building model...')

    if 'model' in checkpoint:
        model = checkpoint['model']
    else:
        encoder = onmt.Models.Encoder(opt, dataset['dicts']['src'])
        decoder = onmt.Models.Decoder(opt, dataset['dicts']['tgt'])
        generator = nn.Sequential(
            nn.Linear(opt.rnn_size, dataset['dicts']['tgt']['words'].size()),
            nn.LogSoftmax())
        model = onmt.Models.NMTModel(encoder, decoder, generator)

    if opt.cuda:
        model.cuda()

    nParams = sum([p.nelement() for p in model.parameters()])
    print('* number of parameters: %d' % nParams)

    trainModel(model, trainData, validData, dataset)

if __name__ == "__main__":
    main()
<commit_msg>memory efficient decoding<commit_after>import onmt
import argparse
import os
import torch
import torch.nn as nn
from torch.autograd import Variable
import math
import time

parser = argparse.ArgumentParser(description='train.lua')

parser.add_argument('-config', help="Read options from this file")

##
## **Data options**
##

parser.add_argument('-data', help="Path to the training *-train.pt file from preprocess.lua")
parser.add_argument('-save_model', help="""Model filename (the model will be saved as
                              <save_model>_epochN_PPL.pt where PPL is the validation perplexity""")
parser.add_argument('-train_from', help="If training from a checkpoint then this is the path to the pretrained model.")
# parser.add_argument('-cont', action="store_true", help="If training from a checkpoint, whether to continue the training in the same configuration or not.")

##
## **Model options**
##

parser.add_argument('-layers',        type=int, default=2,   help="Number of layers in the LSTM encoder/decoder")
parser.add_argument('-rnn_size',       type=int, default=500, help="Size of LSTM hidden states")
parser.add_argument('-word_vec_size', type=int, default=500, help="Word embedding sizes")
parser.add_argument('-feat_merge', default='concat', help="Merge action for the features embeddings: concat or sum")
parser.add_argument('-feat_vec_exponent', type=float, default=0.7, help="""When using concatenation, if the feature takes N values
                                                                then the embedding dimension will be set to N^exponent""")
parser.add_argument('-feat_vec_size', type=int, default=20, help="When using sum, the common embedding size of the features")
parser.add_argument('-input_feed', type=int,    default=1,  help="Feed the context vector at each time step as additional input (via concatenation with the word embeddings) to the decoder.")
# parser.add_argument('-residual',   action="store_true",     help="Add residual connections between RNN layers.")
parser.add_argument('-brnn',       action="store_true",     help="Use a bidirectional encoder")
parser.add_argument('-brnn_merge', default='concat',        help="Merge action for the bidirectional hidden states: concat or sum")

##
## **Optimization options**
##

parser.add_argument('-max_batch_size',  type=int, default=64,  help="Maximum batch size")
parser.add_argument('-max_generator_batches', type=int, default=16, help="""Maximum batches of words in a sequence to run the generator on in parallel.
                                                                           Higher is faster, but uses more memory.""")
parser.add_argument('-epochs',          type=int, default=13,  help="Number of training epochs")
parser.add_argument('-start_epoch',     type=int, default=0,   help="If loading from a checkpoint, the epoch from which to start")
parser.add_argument('-start_iteration', type=int, default=0,   help="If loading from a checkpoint, the iteration from which to start")
# this gives really bad initialization; Xavier better
parser.add_argument('-param_init',      type=float, default=0.1, help="Parameters are initialized over uniform distribution with support (-param_init, param_init)")
parser.add_argument('-optim', default='sgd', help="Optimization method. Possible options are: sgd, adagrad, adadelta, adam")
parser.add_argument('-learning_rate', type=float, default=1, help="""Starting learning rate. If adagrad/adadelta/adam is used,
                                then this is the global learning rate. Recommed settings. sgd = 1,
                                adagrad = 0.1, adadelta = 1, adam = 0.1""")
parser.add_argument('-max_grad_norm',       type=float, default=5,   help="If the norm of the gradient vector exceeds this renormalize it to have the norm equal to max_grad_norm")
parser.add_argument('-dropout',             type=float, default=0.3, help="Dropout probability. Dropout is applied between vertical LSTM stacks.")
parser.add_argument('-learning_rate_decay', type=float, default=0.5, help="""Decay learning rate by this much if (i) perplexity does not decrease
                                        on the validation set or (ii) epoch has gone past the start_decay_at_limit""")
parser.add_argument('-start_decay_at', default=8, help="Start decay after this epoch")
parser.add_argument('-curriculum', action="store_true", help="""For this many epochs, order the minibatches based on source
                             sequence length. Sometimes setting this to 1 will increase convergence speed.""")
parser.add_argument('-pre_word_vecs_enc', help="""If a valid path is specified, then this will load
                                     pretrained word embeddings on the encoder side.
                                     See README for specific formatting instructions.""")
parser.add_argument('-pre_word_vecs_dec', help="""If a valid path is specified, then this will load
                                     pretrained word embeddings on the decoder side.
                                     See README for specific formatting instructions.""")
# parser.add_argument('-fix_word_vecs_enc', action="store_true", help="Fix word embeddings on the encoder side")
# parser.add_argument('-fix_word_vecs_dec', action="store_true", help="Fix word embeddings on the decoder side")

##
## **Other options**
##

# GPU
parser.add_argument('-cuda', action='store_true', help="Use CUDA")
# parser.add_argument('-nparallel', type=int, default=1,  help="""When using GPUs, how many batches to execute in parallel.
#                             Note. this will technically change the final batch size to max_batch_size*nparallel.""")

# bookkeeping
# parser.add_argument('-save_every', type=int, default=0, help="""Save intermediate models every this many iterations within an epoch.
#                              If = 0, will not save models within an epoch. """)
parser.add_argument('-report_every', type=int, default=50,   help="Print stats every this many iterations within an epoch.")
# parser.add_argument('-seed',         type=int, default=3435, help="Seed for random initialization")
# parser.add_argument('-json_log', action="store_true", help="Outputs logs in JSON format.")

opt = parser.parse_args()
print(opt)

if torch.cuda.is_available() and not opt.cuda:
    print("WARNING: You have a CUDA device, so you should probably run with -cuda")

class NMTCriterion(nn.Container):
    def __init__(self, vocabSize, features):
        self.sub = []
        super(NMTCriterion, self).__init__()

        def makeOne(size):
            weight = torch.ones(vocabSize)
            weight[onmt.Constants.PAD] = 0
            crit = nn.NLLLoss(weight, size_average=False)
            if opt.cuda:
                crit.cuda()
            return crit

        self.sub += [makeOne(vocabSize)]
        for feature in features:
            self.sub += [makeOne(features.size())]

    def forward(self, inputs, targets):
        if len(self.sub) == 1:
            total_size = targets.nelement()
            loss = self.sub[0](inputs.view(total_size, -1), targets.view(total_size))
            return loss
        else:
            assert False, "FIXME: features"
            loss = Variable(inputs.new(1).zero_())
            for sub, input, target in zip(self.sub, inputs, targets):
                loss += sub(input, target)
            return loss

def memoryEfficientLoss(outputs, targets, generator, crit, eval=False):
    # compute generations one piece at a time
    loss = 0
    outputs= Variable(outputs.data, requires_grad=(not eval), volatile=eval)

    batch_size = outputs.size(1)
    outputs_split = torch.split(outputs, opt.max_generator_batches)
    targets_split = torch.split(targets, opt.max_generator_batches)
    for out_t, targ_t in zip(outputs_split, targets_split):
        out_t = out_t.view(-1, out_t.size(2))
        pred_t = generator(out_t)
        loss_t = crit(pred_t, targ_t)
        loss += loss_t.data[0]
        if not eval:
            loss_t.div(batch_size).backward()

    return loss, outputs.grad

def eval(model, criterion, data):
    total_loss = 0
    total_words = 0

    model.eval()
    for i in range(len(data)):
        batch = data[i]
        outputs = model(batch)
        targets = batch[1][1:]  # exclude <s> from targets
        loss, _ = memoryEfficientLoss(outputs, targets, model.generator, criterion, eval=True)
        total_loss += loss
        total_words += targets.data.ne(onmt.Constants.PAD).sum()

    model.train()
    return total_loss / total_words

def trainModel(model, trainData, validData, dataset):
    print(model)
    model.train()
    for p in model.parameters():
        p.data.uniform_(-opt.param_init, opt.param_init)

    # define criterion of each GPU
    criterion = NMTCriterion(dataset['dicts']['tgt']['words'].size(),
                             dataset['dicts']['tgt']['features'])

    optim = onmt.Optim(
        model.parameters(), opt.optim, opt.learning_rate, opt.max_grad_norm,
        lr_decay=opt.learning_rate_decay,
        start_decay_at=opt.start_decay_at
    )

    # checkpoint = onmt.train.Checkpoint.new(opt, model, optim, dataset)

    def trainEpoch(epoch):

        startI = opt.start_iteration
        opt.start_iteration = 1

        # shuffle mini batch order
        batchOrder = torch.randperm(len(trainData))

        total_loss, report_loss = 0, 0
        total_words, report_words = 0, 0
        start = time.time()
        for i in range(startI, len(trainData)):

            batchIdx = batchOrder[i] if epoch >= opt.curriculum else i
            batch = trainData[batchIdx]
            # srcData, tgtData = batch
            # srcDict = dataset['dicts']['src']['words']
            # tgtDict = dataset['dicts']['tgt']['words']
            # print(srcData)
            # for i in range(srcData.size(1)):
            #     print(' '.join(srcDict.convertToLabels(srcData.data[:,i], onmt.Constants.EOS)))
            #     print(' '.join(tgtDict.convertToLabels(tgtData.data[:,i], onmt.Constants.EOS)))
            #     print()
            # for n, p in model.state_dict().items():
            #     if n not in norms:
            #         norms[n] = []
            #     norms[n] += [p.data.abs().norm()]
            #     print(n, p.data.abs().norm())

            model.zero_grad()
            outputs = model(batch)
            targets = batch[1][1:]  # exclude <s> from targets
            loss, gradOutput = memoryEfficientLoss(outputs, targets, model.generator, criterion)

            outputs.backward(gradOutput)

            # for module in model.modules():
            #     params = list(module.parameters())
            #     if not isinstance(module, nn.Container) and len(params) > 0:
            #         print(module)
            #         for p in params:
            #             print('p', p.data.nelement(), p.data.norm())
            #             print('gp', p.grad.nelement(), p.grad.norm())
            # assert False

            # update the parameters
            grad_norm = optim.step()

            report_loss += loss
            total_loss += loss
            num_words = targets.data.ne(onmt.Constants.PAD).sum()
            total_words += num_words
            report_words += num_words
            if i % opt.report_every == 0 and i > 0:
                print("Epoch %2d, %5d/%5d batches; perplexity: %6.2f; %3.0f tokens/s" %
                      (epoch+1, i, len(trainData), math.exp(report_loss / report_words), report_words/(time.time()-start)))
                report_loss = report_words = 0
                start = time.time()
                # for k,v in norms.items():
                #     print(k,v)

            # if opt.save_every > 0 and ii % opt.save_every == 0:
            #     checkpoint.saveIteration(ii, epochState, batchOrder, not opt.json_log)

        return total_loss / total_words

    for epoch in range(opt.start_epoch, opt.epochs):
        print('')
        train_loss = trainEpoch(epoch)
        print('Train loss: %g' % train_loss)
        print('Train perplexity: %g' % math.exp(min(train_loss, 100)))
        valid_loss = eval(model, criterion, validData)
        print('Validation perplexity: %g' % math.exp(min(valid_loss, 100)))
        if opt.optim == 'sgd':
            optim.updateLearningRate(valid_loss, epoch)

        # checkpoint.saveEpoch(validPpl, epochState, not opt.json_log)

def main():

    checkpoint = {}
    assert opt.train_from is None, "FIXME: Load from checkpoint"
    # if opt.train_from is not None:
    #     assert os.path.exists(opt.train_from), 'checkpoint path invalid'
    #
    #     if not opt.json_log:
    #       print('Loading checkpoint \'' + opt.train_from + '\'...')
    #
    #     checkpoint = torch.load(opt.train_from)
    #
    #     opt.layers = checkpoint.options.layers
    #     opt.rnn_size = checkpoint.options.rnn_size
    #     opt.brnn = checkpoint.options.brnn
    #     opt.brnn_merge = checkpoint.options.brnn_merge
    #     opt.input_feed = checkpoint.options.input_feed
    #
    #     # Resume training from checkpoint
    #     if opt.cont:
    #         opt.optim = checkpoint.options.optim
    #         opt.learning_rate_decay = checkpoint.options.learning_rate_decay
    #         opt.start_decay_at = checkpoint.options.start_decay_at
    #         opt.epochs = checkpoint.options.epochs
    #         opt.curriculum = checkpoint.options.curriculum
    #
    #         opt.learning_rate = checkpoint.info.learning_rate
    #         opt.optim_states = checkpoint.info.optim_states
    #         opt.start_epoch = checkpoint.info.epoch
    #         opt.start_iteration = checkpoint.info.iteration
    #     print('Resuming training from epoch %d at iteration %d...' % (opt.start_epoch, opt.start_iteration))

    print("Loading data from '%s'" % opt.data)

    dataset = torch.load(opt.data)

    trainData = onmt.Dataset(dataset['train']['src'], dataset['train']['tgt'], opt.max_batch_size, opt.cuda)
    validData = onmt.Dataset(dataset['valid']['src'], dataset['valid']['tgt'], opt.max_batch_size, opt.cuda)
    print(' * vocabulary size. source = %d; target = %d' %
            (dataset['dicts']['src']['words'].size(), dataset['dicts']['tgt']['words'].size()))
    print(' * additional features. source = %d; target = %d' %
            (len(dataset['dicts']['src']['features']), len(dataset['dicts']['tgt']['features'])))
    print(' * number of training sentences. %d' % len(dataset['train']['src']['words']))
    print(' * maximum batch size. %d' % opt.max_batch_size)

    print('Building model...')

    if 'model' in checkpoint:
        model = checkpoint['model']
    else:
        encoder = onmt.Models.Encoder(opt, dataset['dicts']['src'])
        decoder = onmt.Models.Decoder(opt, dataset['dicts']['tgt'])
        generator = nn.Sequential(
            nn.Linear(opt.rnn_size, dataset['dicts']['tgt']['words'].size()),
            nn.LogSoftmax())
        model = onmt.Models.NMTModel(encoder, decoder, generator)

    if opt.cuda:
        model.cuda()

    nParams = sum([p.nelement() for p in model.parameters()])
    print('* number of parameters: %d' % nParams)

    trainModel(model, trainData, validData, dataset)

if __name__ == "__main__":
    main()
<|endoftext|>


<commit_before>#!/usr/bin/env python

"""
Random Tree generation Class
"""

import numpy as np
import random
import toytree
from .utils import ToytreeError

class RandomTree(object):

    @staticmethod
    def unittree(ntips, treeheight=1.0, random_names=False, seed=None):
        """
        Returns a random tree ultrametric topology w/ N tips and a root
        height set to 1 or a user-entered treeheight value. Descendant
        nodes are evenly spaced between the root and time 0.

        Parameters
        -----------
        ntips (int):
            The number of tips in the randomly generated tree

        treeheight(float):
            Scale tree height (all edges) so that root is at this height.

        seed (int):
            Random number generator seed.
        """
        # seed generator
        random.seed(seed)

        # generate tree with N tips.
        tmptree = toytree.tree().treenode  # TreeNode()
        tmptree.populate(ntips)
        self = toytree.tree(newick=tmptree.write())

        # set tip names by labeling sequentially from 0
        self = (
            self
            .ladderize()
            .mod.make_ultrametric()
            .mod.node_scale_root_height(treeheight)
        )

        # set tipnames randomly (doesn't have to match idx)
        nidx = list(range(self.ntips))[::-1]
        if random_names:
            random.shuffle(nidx)
        for tidx, node in enumerate(self.treenode.get_leaves()):
            node.name = "r{}".format(nidx[tidx])

        # set all support values to 100 default
        for node in self.treenode.traverse():
            node.support = 100

        # fill internal node names and idxs
        self.treenode.ladderize()
        self._coords.update()
        return self

    @staticmethod
    def imbtree(ntips, treeheight=1.0, random_names=False):
        """
        Return an imbalanced (comb-like) tree topology.
        """
        node = toytree.TreeNode.TreeNode()
        node.add_child(name="r0")
        node.add_child(name="r1")

        for i in range(2, ntips):
            # empty node
            cherry = toytree.TreeNode.TreeNode()
            # add new child
            cherry.add_child(name="r" + str(i))
            # add old tree
            cherry.add_child(node)
            # update rtree
            node = cherry

        # get toytree from newick
        tre = toytree.tree(node)
        tre = tre.mod.make_ultrametric(nocopy=True)
        tre = tre.mod.node_scale_root_height(treeheight, nocopy=True)
        tre._coords.update()

        # randomize tip names
        nidx = list(range(tre.ntips))
        if random_names:
            random.shuffle(nidx)
        for idx, node in tre.idx_dict.items():
            if node.is_leaf():
                node.name = "r{}".format(nidx[idx])
        return tre

    @staticmethod
    def baltree(ntips, treeheight=1.0, random_names=False):
        """
        Returns a balanced tree topology.
        """
        # require even number of tips
        if ntips % 2:
            raise ToytreeError("balanced trees must have even number of tips.")

        # make first cherry
        rtree = toytree.tree()
        rtree.treenode.add_child(name="r0")
        rtree.treenode.add_child(name="r1")

        # add tips in a balanced way
        for i in range(2, ntips):

            # get node to split
            node = return_small_clade(rtree.treenode)

            # add two children
            node.add_child(name=node.name)
            node.add_child(name="r" + str(i))

            # rename ancestral node
            node.name = None

        # get toytree from newick
        tre = toytree.tree(rtree)  # .write(tree_format=9))
        tre = tre.mod.make_ultrametric().mod.node_scale_root_height(treeheight)
        tre._coords.update()

        # rename tips so names are in order
        nidx = list(range(tre.ntips))
        if random_names:
            random.shuffle(nidx)
        for idx, node in tre.idx_dict.items():
            if node.is_leaf():
                node.name = "r{}".format(nidx[idx])
        return tre

    @staticmethod
    def rtree(ntips, treeheight=1.0, random_names=False, seed=None):
        """
        Randomly assigns edge lengths between U(0-1) to edges
        and then scales total treeheight to 1.0 or entered value.
        """
        # get a unit tree
        self = toytree.rtree.unittree(ntips, treeheight, random_names, seed)

        # randomly assign node dists
        self = self.set_node_values(
            "dist", {i: random.random() for i in range(self.nnodes)}
        )

        # rescale total height to .
        self = self.mod.node_scale_root_height(treeheight, nocopy=True)
        return self

    @staticmethod
    def coaltree(ntips, ne=None, random_names=False, seed=None):
        """
        Returns a coalescent tree with ntips samples and waiting times
        between coalescent events drawn from the kingman coalescent:
        (4N)/(k*(k-1)), where N is population size and k is sample size.
        Edge lengths on the tree are in generations.

        If no Ne argument is entered then edge lengths are returned in units
        of 2*Ne, i.e., coalescent time units.
        """
        # seed generator
        random.seed(seed)

        # convert units
        coalunits = False
        if not ne:
            coalunits = True
            ne = 10000

        # build tree: generate N tips as separate Nodes then attach together
        # at internal nodes drawn randomly from coalescent waiting times.
        tips = [
            toytree.tree().treenode.add_child(name=str(i))
            for i in range(ntips)
        ]
        while len(tips) > 1:
            rtree = toytree.tree()
            tip1 = tips.pop(random.choice(range(len(tips))))
            tip2 = tips.pop(random.choice(range(len(tips))))
            kingman = (4. * ne) / float(ntips * (ntips - 1))
            dist = random.expovariate(1. / kingman)
            rtree.treenode.add_child(tip1, dist=tip2.height + dist)
            rtree.treenode.add_child(tip2, dist=tip1.height + dist)
            tips.append(rtree.treenode)

        # build new tree from the newick string
        self = toytree.tree(tips[0].write())
        self.treenode.ladderize()

        # make tree edges in units of 2N (then N doesn't matter!)
        if coalunits:
            for node in self.treenode.traverse():
                node.dist /= (2. * ne)

        # ensure tips are at zero (they sometime vary just slightly)
        for node in self.treenode.traverse():
            if node.is_leaf():
                node.dist += node.height

        # set tipnames to r{idx}
        nidx = list(range(self.ntips))
        if random_names:
            random.shuffle(nidx)
        for idx, node in self.idx_dict.items():
            if node.is_leaf():
                node.name = "r{}".format(nidx[idx])

        # decompose fills in internal node names and idx
        self._coords.update()
        return self

    @staticmethod
    def bdtree(ntips=10,
                time=4,
                b=1,
                d=0,
                stop="taxa",
                seed=None,
                retain_extinct=False,
                random_names=False,
                verbose=False):
        """
        Generate a classic birth/death tree.

        Parameters
        -----------
        ntips (int):
            Number of tips to generate for 'taxa' stopping criterion.

        time (float):
            Amount of time to simulate for 'time' stopping criterion.

        b (float):
            Birth rate per time unit

        d (float):
            Death rate per time unit (d=0 produces Yule trees)

        stop (str):
            Stopping critereon. Valid values are only 'taxa' or 'time'.

        seed (int):
            Random number generator seed.

        retain_extinct (bool):
            Whether to retain internal nodes leading to extinct tips.

        random_names (bool):
            Whether to randomize tip names or name them in order.

        verbose (bool):
            Print some useful information
        """

        # set random seed
        if seed:
            np.random.seed(seed)

        # start from random tree (idxs will be re-assigned at end)
        tre = toytree.tree()
        tre.treenode.idx = 0
        gidx = 1

        # counters for extinctions, total events, and time
        resets = 0
        ext = 0
        evnts = 0
        t = 0

        # continue until stop var
        while 1:

            # get current tips
            tips = tre.treenode.get_leaves()

            # sample time until next event, increment t and evnts
            dt = np.random.exponential(1 / (len(tips) * (b + d)))
            t = t + dt
            evnts += 1

            # sample a [0-1] to choose birth or death and sample a tip node
            r = np.random.sample()
            sp = np.random.choice(tips)

            # event is a birth
            if (r <= b / (b + d)):
                c1 = sp.add_child(name=str(t) + "-1", dist=0)
                c1.add_feature("tdiv", t)
                c1.idx = gidx
                gidx += 1

                c2 = sp.add_child(name=str(t) + "-2", dist=0)
                c2.add_feature("tdiv", t)
                c2.idx = gidx
                gidx += 1

            # else event is extinction
            else:
                # get parent node
                parent = sp.up

                # if no parent then reset to empty tree
                if parent is None:
                    resets += 1
                    tre = toytree.tree()
                    tre.treenode.idx = 0
                    gidx = 1
                    ext = 0
                    evnts = 0
                    t = 0

                # connect parent to sp' children
                else:
                    # if parent is None then reset
                    if sp.up is None:
                        tre = toytree.tree()
                        tre.treenode.idx = 0
                        gidx = 1
                        ext = 0
                        evnts = 0
                        t = 0

                    # if parent is root then sister is new root
                    elif sp.up is tre.treenode:
                        tre = [i for i in sp.up.children if i != sp][0]
                        tre = toytree.tree(tre)
                        tre.up = None

                    # if parent is a non-root node then connect sister to up.up
                    else:
                        # get sister
                        sister = [i for i in sp.up.children if i != sp][0]

                        # connect sister to grandparent
                        sister.up = sp.up.up

                        # drop parent from grandparent
                        sp.up.up.children.remove(sp.up)

                        # add sister to grandparent children
                        sp.up.up.children.append(sister)

                        # extend sisters dist to reach grandparent
                        sister.dist += sp.up.dist

                    ext += 1

            # update branch lengths so all tips end at time=present
            tips = tre.treenode.get_leaves()
            for x in tips:
                x.dist += dt

            # time stopping criterion... not quite working right.
            if maxtime:
                if t >= maxtime:
                    break

            # ntips stopping criterion.
            else:
                if len(tips) >= ntips:
                    break

        # report status
        if verbose:
            fill = (evnts - ext, ext, evnts / (evnts - ext), resets)
            print("b:\t{}\nd:\t{}\nb/d:\t{:.2f}\nreset:\t{}".format(*fill))

        # update coords and return
        tre.treenode.ladderize()
        tre._coords.update()

        # rename tips so names are in order else random
        nidx = list(range(tre.ntips))
        if random_names:
            random.shuffle(nidx)
        for idx, node in tre.idx_dict.items():
            if node.is_leaf():
                node.name = "r{}".format(nidx[idx])
        return tre

def _prune(tre, verbose=False):
    """
    Helper function for recursively pruning extinct branches in bd trees.
    Dynamic func!
    """
    ttree = tre.copy()
    tips = ttree.treenode.get_leaves()

    if np.any(np.array([x.height for x in tips]) > 0):
        for t in tips:
            if not np.isclose(t.height, 0):
                if verbose:
                    print("Removing node/height {}/{}".format(t.name, t.height))
                t.delete(prevent_nondicotomic=False)
                ttree = _prune(ttree)
    return ttree

def return_small_clade(treenode):
    "used to produce balanced trees, returns a tip node from the smaller clade"
    node = treenode
    while 1:
        if node.children:
            c1, c2 = node.children
            node = sorted([c1, c2], key=lambda x: len(x.get_leaves()))[0]
        else:
            return node
<commit_msg>Add more interpretable and consistent stopping criteria.<commit_after>#!/usr/bin/env python

"""
Random Tree generation Class
"""

import numpy as np
import random
import toytree
from .utils import ToytreeError

class RandomTree(object):

    @staticmethod
    def unittree(ntips, treeheight=1.0, random_names=False, seed=None):
        """
        Returns a random tree ultrametric topology w/ N tips and a root
        height set to 1 or a user-entered treeheight value. Descendant
        nodes are evenly spaced between the root and time 0.

        Parameters
        -----------
        ntips (int):
            The number of tips in the randomly generated tree

        treeheight(float):
            Scale tree height (all edges) so that root is at this height.

        seed (int):
            Random number generator seed.
        """
        # seed generator
        random.seed(seed)

        # generate tree with N tips.
        tmptree = toytree.tree().treenode  # TreeNode()
        tmptree.populate(ntips)
        self = toytree.tree(newick=tmptree.write())

        # set tip names by labeling sequentially from 0
        self = (
            self
            .ladderize()
            .mod.make_ultrametric()
            .mod.node_scale_root_height(treeheight)
        )

        # set tipnames randomly (doesn't have to match idx)
        nidx = list(range(self.ntips))[::-1]
        if random_names:
            random.shuffle(nidx)
        for tidx, node in enumerate(self.treenode.get_leaves()):
            node.name = "r{}".format(nidx[tidx])

        # set all support values to 100 default
        for node in self.treenode.traverse():
            node.support = 100

        # fill internal node names and idxs
        self.treenode.ladderize()
        self._coords.update()
        return self

    @staticmethod
    def imbtree(ntips, treeheight=1.0, random_names=False):
        """
        Return an imbalanced (comb-like) tree topology.
        """
        node = toytree.TreeNode.TreeNode()
        node.add_child(name="r0")
        node.add_child(name="r1")

        for i in range(2, ntips):
            # empty node
            cherry = toytree.TreeNode.TreeNode()
            # add new child
            cherry.add_child(name="r" + str(i))
            # add old tree
            cherry.add_child(node)
            # update rtree
            node = cherry

        # get toytree from newick
        tre = toytree.tree(node)
        tre = tre.mod.make_ultrametric(nocopy=True)
        tre = tre.mod.node_scale_root_height(treeheight, nocopy=True)
        tre._coords.update()

        # randomize tip names
        nidx = list(range(tre.ntips))
        if random_names:
            random.shuffle(nidx)
        for idx, node in tre.idx_dict.items():
            if node.is_leaf():
                node.name = "r{}".format(nidx[idx])
        return tre

    @staticmethod
    def baltree(ntips, treeheight=1.0, random_names=False):
        """
        Returns a balanced tree topology.
        """
        # require even number of tips
        if ntips % 2:
            raise ToytreeError("balanced trees must have even number of tips.")

        # make first cherry
        rtree = toytree.tree()
        rtree.treenode.add_child(name="r0")
        rtree.treenode.add_child(name="r1")

        # add tips in a balanced way
        for i in range(2, ntips):

            # get node to split
            node = return_small_clade(rtree.treenode)

            # add two children
            node.add_child(name=node.name)
            node.add_child(name="r" + str(i))

            # rename ancestral node
            node.name = None

        # get toytree from newick
        tre = toytree.tree(rtree)  # .write(tree_format=9))
        tre = tre.mod.make_ultrametric().mod.node_scale_root_height(treeheight)
        tre._coords.update()

        # rename tips so names are in order
        nidx = list(range(tre.ntips))
        if random_names:
            random.shuffle(nidx)
        for idx, node in tre.idx_dict.items():
            if node.is_leaf():
                node.name = "r{}".format(nidx[idx])
        return tre

    @staticmethod
    def rtree(ntips, treeheight=1.0, random_names=False, seed=None):
        """
        Randomly assigns edge lengths between U(0-1) to edges
        and then scales total treeheight to 1.0 or entered value.
        """
        # get a unit tree
        self = toytree.rtree.unittree(ntips, treeheight, random_names, seed)

        # randomly assign node dists
        self = self.set_node_values(
            "dist", {i: random.random() for i in range(self.nnodes)}
        )

        # rescale total height to .
        self = self.mod.node_scale_root_height(treeheight, nocopy=True)
        return self

    @staticmethod
    def coaltree(ntips, ne=None, random_names=False, seed=None):
        """
        Returns a coalescent tree with ntips samples and waiting times
        between coalescent events drawn from the kingman coalescent:
        (4N)/(k*(k-1)), where N is population size and k is sample size.
        Edge lengths on the tree are in generations.

        If no Ne argument is entered then edge lengths are returned in units
        of 2*Ne, i.e., coalescent time units.
        """
        # seed generator
        random.seed(seed)

        # convert units
        coalunits = False
        if not ne:
            coalunits = True
            ne = 10000

        # build tree: generate N tips as separate Nodes then attach together
        # at internal nodes drawn randomly from coalescent waiting times.
        tips = [
            toytree.tree().treenode.add_child(name=str(i))
            for i in range(ntips)
        ]
        while len(tips) > 1:
            rtree = toytree.tree()
            tip1 = tips.pop(random.choice(range(len(tips))))
            tip2 = tips.pop(random.choice(range(len(tips))))
            kingman = (4. * ne) / float(ntips * (ntips - 1))
            dist = random.expovariate(1. / kingman)
            rtree.treenode.add_child(tip1, dist=tip2.height + dist)
            rtree.treenode.add_child(tip2, dist=tip1.height + dist)
            tips.append(rtree.treenode)

        # build new tree from the newick string
        self = toytree.tree(tips[0].write())
        self.treenode.ladderize()

        # make tree edges in units of 2N (then N doesn't matter!)
        if coalunits:
            for node in self.treenode.traverse():
                node.dist /= (2. * ne)

        # ensure tips are at zero (they sometime vary just slightly)
        for node in self.treenode.traverse():
            if node.is_leaf():
                node.dist += node.height

        # set tipnames to r{idx}
        nidx = list(range(self.ntips))
        if random_names:
            random.shuffle(nidx)
        for idx, node in self.idx_dict.items():
            if node.is_leaf():
                node.name = "r{}".format(nidx[idx])

        # decompose fills in internal node names and idx
        self._coords.update()
        return self

    @staticmethod
    def bdtree(ntips=10,
                time=4,
                b=1,
                d=0,
                stop="taxa",
                seed=None,
                retain_extinct=False,
                random_names=False,
                verbose=False):
        """
        Generate a classic birth/death tree.

        Parameters
        -----------
        ntips (int):
            Number of tips to generate for 'taxa' stopping criterion.

        time (float):
            Amount of time to simulate for 'time' stopping criterion.

        b (float):
            Birth rate per time unit

        d (float):
            Death rate per time unit (d=0 produces Yule trees)

        stop (str):
            Stopping critereon. Valid values are only 'taxa' or 'time'.

        seed (int):
            Random number generator seed.

        retain_extinct (bool):
            Whether to retain internal nodes leading to extinct tips.

        random_names (bool):
            Whether to randomize tip names or name them in order.

        verbose (bool):
            Print some useful information
        """
        if stop not in ["taxa", "time"]:
            raise ToytreeError("stop must be either 'taxa' or 'time'")

        # set random seed
        if seed:
            np.random.seed(seed)

        # start from random tree (idxs will be re-assigned at end)
        tre = toytree.tree()
        tre.treenode.idx = 0
        gidx = 1

        # counters for extinctions, total events, and time
        resets = 0
        ext = 0
        evnts = 0
        t = 0

        # continue until stop var
        while 1:

            # get current tips
            tips = tre.treenode.get_leaves()

            # sample time until next event, increment t and evnts
            dt = np.random.exponential(1 / (len(tips) * (b + d)))
            t = t + dt
            evnts += 1

            # sample a [0-1] to choose birth or death and sample a tip node
            r = np.random.sample()
            sp = np.random.choice(tips)

            # event is a birth
            if (r <= b / (b + d)):
                c1 = sp.add_child(name=str(t) + "-1", dist=0)
                c1.add_feature("tdiv", t)
                c1.idx = gidx
                gidx += 1

                c2 = sp.add_child(name=str(t) + "-2", dist=0)
                c2.add_feature("tdiv", t)
                c2.idx = gidx
                gidx += 1

            # else event is extinction
            else:
                # get parent node
                parent = sp.up

                # if no parent then reset to empty tree
                if parent is None:
                    resets += 1
                    tre = toytree.tree()
                    tre.treenode.idx = 0
                    gidx = 1
                    ext = 0
                    evnts = 0
                    t = 0

                # connect parent to sp' children
                else:
                    # if parent is None then reset
                    if sp.up is None:
                        tre = toytree.tree()
                        tre.treenode.idx = 0
                        gidx = 1
                        ext = 0
                        evnts = 0
                        t = 0

                    # if parent is root then sister is new root
                    elif sp.up is tre.treenode:
                        tre = [i for i in sp.up.children if i != sp][0]
                        tre = toytree.tree(tre)
                        tre.up = None

                    # if parent is a non-root node then connect sister to up.up
                    else:
                        # get sister
                        sister = [i for i in sp.up.children if i != sp][0]

                        # connect sister to grandparent
                        sister.up = sp.up.up

                        # drop parent from grandparent
                        sp.up.up.children.remove(sp.up)

                        # add sister to grandparent children
                        sp.up.up.children.append(sister)

                        # extend sisters dist to reach grandparent
                        sister.dist += sp.up.dist

                    ext += 1

            # update branch lengths so all tips end at time=present
            tips = tre.treenode.get_leaves()
            for x in tips:
                x.dist += dt

            # check stopping criterion
            if stop == "taxa":
                if len(tips) >= taxa_stop:
                    break
            else:
                if t >= time_stop:
                    break

            # ntips stopping criterion.
            else:
                if len(tips) >= ntips:
                    break

        # report status
        if verbose:
            fill = (evnts - ext, ext, evnts / (evnts - ext), resets)
            print("b:\t{}\nd:\t{}\nb/d:\t{:.2f}\nreset:\t{}".format(*fill))

        # update coords and return
        tre.treenode.ladderize()
        tre._coords.update()

        # rename tips so names are in order else random
        nidx = list(range(tre.ntips))
        if random_names:
            random.shuffle(nidx)
        for idx, node in tre.idx_dict.items():
            if node.is_leaf():
                node.name = "r{}".format(nidx[idx])
        return tre

def _prune(tre, verbose=False):
    """
    Helper function for recursively pruning extinct branches in bd trees.
    Dynamic func!
    """
    ttree = tre.copy()
    tips = ttree.treenode.get_leaves()

    if np.any(np.array([x.height for x in tips]) > 0):
        for t in tips:
            if not np.isclose(t.height, 0):
                if verbose:
                    print("Removing node/height {}/{}".format(t.name, t.height))
                t.delete(prevent_nondicotomic=False)
                ttree = _prune(ttree)
    return ttree

def return_small_clade(treenode):
    "used to produce balanced trees, returns a tip node from the smaller clade"
    node = treenode
    while 1:
        if node.children:
            c1, c2 = node.children
            node = sorted([c1, c2], key=lambda x: len(x.get_leaves()))[0]
        else:
            return node
<|endoftext|>


<commit_before>#!/usr/bin/env python3

import sys
import os

import generic_run
import path_utils

def generic_parse(str_line, separator):
    if str_line is None:
        return None
    n = str_line.find(separator)
    if n == -1:
        return None
    return str_line[:n]

def get_prev_hash(str_line):
    return generic_parse(str_line, " ")

def is_non_generic(char_input, list_select):
    for c in list_select:
        if c == char_input:
            return False
    return True

def is_nonnumber(thechar):
    return is_non_generic(thechar, ["0", "1", "2", "3", "4", "5", "6", "7", "8", "9"])

def is_nonspaceortabs(thechar):
    return is_non_generic(thechar, [" ", "\t"])

def revision_filter_function(the_output):

    find_string = "Revision: "
    n = the_output.find(find_string)
    revision_left = the_output[n + len(find_string):]
    revision = ""
    for i in range(len(revision_left)):
        if is_nonnumber(revision_left[i]):
            revision = revision_left[0:i]
            break
    return revision

def status_filter_function(the_line):

    if the_line is None:
        return None
    if len(the_line) == 0:
        return None
    if the_line[0] != "?":
        return None
    the_line = the_line[1:]

    for i in range(len(the_line)):
        if is_nonspaceortabs(the_line[i]):
            return the_line[i:]

def detect_separator(the_string):

    if the_string is None:
        return None
    if len(the_string) == 0:
        return None

    rep_char = the_string[0]
    rep_line = ""
    for i in range(len(the_string)):
        if is_non_generic(the_string[i], [rep_char]):
            rep_line = the_string[0:i]
            return rep_line

def rev_single_entry_filter(log_entry):

    if log_entry is None:
        return None
    if len(log_entry) == 0:
        return None

    na = log_entry.find("r")
    if na == -1:
        return None
    nb = log_entry.find(" ", na)
    the_rev = log_entry[na:nb]
    return the_rev

def rev_entries_filter(log_entries_list):

    rev_list = []

    for le in log_entries_list:
        rle = rev_single_entry_filter(le)
        if rle is not None:
            rev_list.append(rle)

    return rev_list

def collect_svn_patch_cmd_generic(repo, storage_path, output_filename, log_title, cmd, filter_function=None):

    fullbasepath = path_utils.concat_path(storage_path, repo)
    output_filename_full = path_utils.concat_path(fullbasepath, output_filename)

    try:
        path_utils.guaranteefolder(fullbasepath)
    except path_utils.PathUtilsException as puex:
        return False, "Can't collect patch for %s: Failed guaranteeing folder [%s]" % (log_title, fullbasepath)

    if os.path.exists(output_filename_full):
        return False, "Can't collect patch for %s: %s already exists" % (log_title, output_filename_full)

    v, r = generic_run.run_cmd_simple(cmd, use_cwd=repo)
    if not v:
        return False, "Failed calling svn command for %s: %s. Repository: %s." % (log_title, r, repo)
    if filter_function is not None:
        final_output = filter_function(r)
    else:
        final_output = r

    with open(output_filename_full, "w") as f:
        f.write(final_output)

    return True, ""

def collect_svn_patch_head(repo, storage_path):
    return collect_svn_patch_cmd_generic(repo, storage_path, "head.patch", "head", ["svn", "diff", "--internal-diff"], None)

def collect_svn_patch_head_id(repo, storage_path):
    return collect_svn_patch_cmd_generic(repo, storage_path, "head_id.txt", "head-id", ["svn", "info"], revision_filter_function)

def collect_svn_patch_head_unversioned(repo, storage_path):
    v, r = generic_run.run_cmd_simple(["svn", "status"], use_cwd=repo)
    if not v:
        return False, "Failed calling svn command for head-unversioned: %s. Repository: %s." % (r, repo)

    unversioned_files = [x for x in r.split(os.linesep) if x != ""]
    for uf in unversioned_files:
        file_filtered = status_filter_function(uf)
        if file_filtered is None:
            continue
        target_base = path_utils.concat_path(storage_path, repo, "head_unversioned")
        target_file = path_utils.concat_path(target_base, os.sep)
        source_file = path_utils.concat_path(repo, file_filtered)
        try:
            path_utils.guaranteefolder( os.path.dirname(target_file) )
        except path_utils.PathUtilsException as puex:
            return False, "Can't collect patch for head-unversioned: Failed guaranteeing folder [%s]" % os.path.dirname(target_file)
        if os.path.exists( path_utils.concat_path(target_base, file_filtered) ):
            return False, "Can't collect patch for head-unversioned: %s already exists" % path_utils.concat_path(target_base, file_filtered)
        if not path_utils.copy_to( source_file, target_file ):
            return False, "Can't collect patch for head-unversioned: Cant copy %s to %s" % (source_file, target_file)

    return True, ""

def collect_svn_patch_previous(repo, storage_path, previous_number):

    if not previous_number > 0:
        return False, "Can't collect patch for previous: nothing to format"

    v, r = generic_run.run_cmd_simple(["svn", "log", "--limit", str(previous_number)], use_cwd=repo)
    if not v:
        return False, "Failed calling svn command for previous: %s. Repository: %s." % (r, repo)
    log_out = r

    sep_line = detect_separator(log_out)
    log_entries_pre = log_out.split(sep_line)
    log_entries = []
    for le in log_entries_pre:
        if le != "":
            log_entries.append(le)

    prev_list = rev_entries_filter(log_entries)

    if previous_number > len(prev_list):
        return False, "Can't collect patch for previous: requested %d commits, but there are only %d in total" % (previous_number, len(prev_list))

    for i in range(previous_number):
        v, r = generic_run.run_cmd_simple(["svn", "diff", "-c", prev_list[i]], use_cwd=repo)
        if not v:
            return False, "Failed calling svn command for previous: %s. Repository: %s. Revision: %s." % (r, repo, prev_list[i])

        previous_file_content = r
        previous_file_name = path_utils.concat_path(storage_path, repo, "previous_%d_%s.patch" % ((i+1), prev_list[i]))

        try:
            path_utils.guaranteefolder( os.path.dirname(previous_file_name) )
        except path_utils.PathUtilsException as puex:
            return False, "Can't collect patch for previous: Failed guaranteeing folder [%s]" % os.path.dirname(previous_file_name)

        if os.path.exists(previous_file_name):
            return False, "Can't collect patch for previous: %s already exists" % previous_file_name
        with open(previous_file_name, "w") as f:
            f.write(previous_file_content)

    return True, ""

def collect_svn_patch(repo, storage_path, head, head_id, head_unversioned, previous):

    repo = path_utils.filter_remove_trailing_sep(repo)
    repo = os.path.abspath(repo)
    storage_path = path_utils.filter_remove_trailing_sep(storage_path)
    storage_path = os.path.abspath(storage_path)

    if not os.path.exists(repo):
        return False, ["Repository %s does not exist" % repo]

    if not os.path.exists(storage_path):
        return False, ["Storage path %s does not exist" % storage_path]

    report = []

    # head
    if head:
        v, r = collect_svn_patch_head(repo, storage_path)
        if not v:
            report.append(r)

    # head-id
    if head_id:
        v, r = collect_svn_patch_head_id(repo, storage_path)
        if not v:
            report.append(r)

    # head_unversioned
    if head_unversioned:
        v, r = collect_svn_patch_head_unversioned(repo, storage_path)
        if not v:
            report.append(r)

    # previous
    if previous > 0:
        v, r = collect_svn_patch_previous(repo, storage_path, previous)
        if not v:
            report.append(r)

    return (len(report)==0), report

def puaq():
    print("Usage: %s repo [--storage-path the_storage_path] [--head] [--head-id] [--head-unversioned] [--previous X]" % os.path.basename(__file__))
    sys.exit(1)

if __name__ == "__main__":

    if len(sys.argv) < 2:
        puaq()

    repo = sys.argv[1]
    storage_path = None
    params = sys.argv[2:]

    # parse options
    storage_path_parse_next = False
    head = False
    head_id = False
    head_unversioned = False
    previous = 0
    previous_parse_next = False

    for p in params:

        if storage_path_parse_next:
            storage_path = p
            storage_path_parse_next = False
            continue

        if previous_parse_next:
            previous = int(p)
            previous_parse_next = False
            continue

        if p == "--storage-path":
            storage_path_parse_next = True
        elif p == "--head":
            head = True
        elif p == "--head-id":
            head_id = True
        elif p == "--head-unversioned":
            head_unversioned = True
        elif p == "--previous":
            previous_parse_next = True

    if storage_path is None:
        storage_path = os.getcwd()

    v, r = collect_svn_patch(repo, storage_path, head, head_id, head_unversioned, previous)
    if not v:
        for i in r:
            print("Failed: %s" % i)
        sys.exit(1)
    else:
        print("All succeeded")
<commit_msg>fix duplicate folders when copying head-unversioned, collect_svn_patch<commit_after>#!/usr/bin/env python3

import sys
import os

import generic_run
import path_utils

def generic_parse(str_line, separator):
    if str_line is None:
        return None
    n = str_line.find(separator)
    if n == -1:
        return None
    return str_line[:n]

def get_prev_hash(str_line):
    return generic_parse(str_line, " ")

def is_non_generic(char_input, list_select):
    for c in list_select:
        if c == char_input:
            return False
    return True

def is_nonnumber(thechar):
    return is_non_generic(thechar, ["0", "1", "2", "3", "4", "5", "6", "7", "8", "9"])

def is_nonspaceortabs(thechar):
    return is_non_generic(thechar, [" ", "\t"])

def revision_filter_function(the_output):

    find_string = "Revision: "
    n = the_output.find(find_string)
    revision_left = the_output[n + len(find_string):]
    revision = ""
    for i in range(len(revision_left)):
        if is_nonnumber(revision_left[i]):
            revision = revision_left[0:i]
            break
    return revision

def status_filter_function(the_line):

    if the_line is None:
        return None
    if len(the_line) == 0:
        return None
    if the_line[0] != "?":
        return None
    the_line = the_line[1:]

    for i in range(len(the_line)):
        if is_nonspaceortabs(the_line[i]):
            return the_line[i:]

def detect_separator(the_string):

    if the_string is None:
        return None
    if len(the_string) == 0:
        return None

    rep_char = the_string[0]
    rep_line = ""
    for i in range(len(the_string)):
        if is_non_generic(the_string[i], [rep_char]):
            rep_line = the_string[0:i]
            return rep_line

def rev_single_entry_filter(log_entry):

    if log_entry is None:
        return None
    if len(log_entry) == 0:
        return None

    na = log_entry.find("r")
    if na == -1:
        return None
    nb = log_entry.find(" ", na)
    the_rev = log_entry[na:nb]
    return the_rev

def rev_entries_filter(log_entries_list):

    rev_list = []

    for le in log_entries_list:
        rle = rev_single_entry_filter(le)
        if rle is not None:
            rev_list.append(rle)

    return rev_list

def collect_svn_patch_cmd_generic(repo, storage_path, output_filename, log_title, cmd, filter_function=None):

    fullbasepath = path_utils.concat_path(storage_path, repo)
    output_filename_full = path_utils.concat_path(fullbasepath, output_filename)

    try:
        path_utils.guaranteefolder(fullbasepath)
    except path_utils.PathUtilsException as puex:
        return False, "Can't collect patch for %s: Failed guaranteeing folder [%s]" % (log_title, fullbasepath)

    if os.path.exists(output_filename_full):
        return False, "Can't collect patch for %s: %s already exists" % (log_title, output_filename_full)

    v, r = generic_run.run_cmd_simple(cmd, use_cwd=repo)
    if not v:
        return False, "Failed calling svn command for %s: %s. Repository: %s." % (log_title, r, repo)
    if filter_function is not None:
        final_output = filter_function(r)
    else:
        final_output = r

    with open(output_filename_full, "w") as f:
        f.write(final_output)

    return True, ""

def collect_svn_patch_head(repo, storage_path):
    return collect_svn_patch_cmd_generic(repo, storage_path, "head.patch", "head", ["svn", "diff", "--internal-diff"], None)

def collect_svn_patch_head_id(repo, storage_path):
    return collect_svn_patch_cmd_generic(repo, storage_path, "head_id.txt", "head-id", ["svn", "info"], revision_filter_function)

def collect_svn_patch_head_unversioned(repo, storage_path):
    v, r = generic_run.run_cmd_simple(["svn", "status"], use_cwd=repo)
    if not v:
        return False, "Failed calling svn command for head-unversioned: %s. Repository: %s." % (r, repo)

    unversioned_files = [x for x in r.split(os.linesep) if x != ""]
    for uf in unversioned_files:
        file_filtered = status_filter_function(uf)
        if file_filtered is None:
            continue
        target_base = path_utils.concat_path(storage_path, repo, "head_unversioned")
        target_file = path_utils.concat_path(target_base, file_filtered)
        source_file = path_utils.concat_path(repo, file_filtered)

        try:
            path_utils.guaranteefolder( os.path.dirname(target_file) )
        except path_utils.PathUtilsException as puex:
            return False, "Can't collect patch for head-unversioned: Failed guaranteeing folder [%s]" % os.path.dirname(target_file)
        if os.path.exists( target_file ):
            return False, "Can't collect patch for head-unversioned: %s already exists" % target_file

        if os.path.isdir(source_file): # this is necessary to avoid duplicating the copied folder
            target_file = os.path.dirname(target_file)

        if not path_utils.copy_to( source_file, target_file ):
            return False, "Can't collect patch for head-unversioned: Cant copy %s to %s" % (source_file, target_file)

    return True, ""

def collect_svn_patch_previous(repo, storage_path, previous_number):

    if not previous_number > 0:
        return False, "Can't collect patch for previous: nothing to format"

    v, r = generic_run.run_cmd_simple(["svn", "log", "--limit", str(previous_number)], use_cwd=repo)
    if not v:
        return False, "Failed calling svn command for previous: %s. Repository: %s." % (r, repo)
    log_out = r

    sep_line = detect_separator(log_out)
    log_entries_pre = log_out.split(sep_line)
    log_entries = []
    for le in log_entries_pre:
        if le != "":
            log_entries.append(le)

    prev_list = rev_entries_filter(log_entries)

    if previous_number > len(prev_list):
        return False, "Can't collect patch for previous: requested %d commits, but there are only %d in total" % (previous_number, len(prev_list))

    for i in range(previous_number):
        v, r = generic_run.run_cmd_simple(["svn", "diff", "-c", prev_list[i]], use_cwd=repo)
        if not v:
            return False, "Failed calling svn command for previous: %s. Repository: %s. Revision: %s." % (r, repo, prev_list[i])

        previous_file_content = r
        previous_file_name = path_utils.concat_path(storage_path, repo, "previous_%d_%s.patch" % ((i+1), prev_list[i]))

        try:
            path_utils.guaranteefolder( os.path.dirname(previous_file_name) )
        except path_utils.PathUtilsException as puex:
            return False, "Can't collect patch for previous: Failed guaranteeing folder [%s]" % os.path.dirname(previous_file_name)

        if os.path.exists(previous_file_name):
            return False, "Can't collect patch for previous: %s already exists" % previous_file_name
        with open(previous_file_name, "w") as f:
            f.write(previous_file_content)

    return True, ""

def collect_svn_patch(repo, storage_path, head, head_id, head_unversioned, previous):

    repo = path_utils.filter_remove_trailing_sep(repo)
    repo = os.path.abspath(repo)
    storage_path = path_utils.filter_remove_trailing_sep(storage_path)
    storage_path = os.path.abspath(storage_path)

    if not os.path.exists(repo):
        return False, ["Repository %s does not exist" % repo]

    if not os.path.exists(storage_path):
        return False, ["Storage path %s does not exist" % storage_path]

    report = []

    # head
    if head:
        v, r = collect_svn_patch_head(repo, storage_path)
        if not v:
            report.append(r)

    # head-id
    if head_id:
        v, r = collect_svn_patch_head_id(repo, storage_path)
        if not v:
            report.append(r)

    # head_unversioned
    if head_unversioned:
        v, r = collect_svn_patch_head_unversioned(repo, storage_path)
        if not v:
            report.append(r)

    # previous
    if previous > 0:
        v, r = collect_svn_patch_previous(repo, storage_path, previous)
        if not v:
            report.append(r)

    return (len(report)==0), report

def puaq():
    print("Usage: %s repo [--storage-path the_storage_path] [--head] [--head-id] [--head-unversioned] [--previous X]" % os.path.basename(__file__))
    sys.exit(1)

if __name__ == "__main__":

    if len(sys.argv) < 2:
        puaq()

    repo = sys.argv[1]
    storage_path = None
    params = sys.argv[2:]

    # parse options
    storage_path_parse_next = False
    head = False
    head_id = False
    head_unversioned = False
    previous = 0
    previous_parse_next = False

    for p in params:

        if storage_path_parse_next:
            storage_path = p
            storage_path_parse_next = False
            continue

        if previous_parse_next:
            previous = int(p)
            previous_parse_next = False
            continue

        if p == "--storage-path":
            storage_path_parse_next = True
        elif p == "--head":
            head = True
        elif p == "--head-id":
            head_id = True
        elif p == "--head-unversioned":
            head_unversioned = True
        elif p == "--previous":
            previous_parse_next = True

    if storage_path is None:
        storage_path = os.getcwd()

    v, r = collect_svn_patch(repo, storage_path, head, head_id, head_unversioned, previous)
    if not v:
        for i in r:
            print("Failed: %s" % i)
        sys.exit(1)
    else:
        print("All succeeded")
<|endoftext|>


<commit_before>#!/usr/bin/env python
# -*- coding: utf-8 -*-

from runner.koan import *

class AboutNewStyleClasses(Koan):
    class OldStyleClass:
        "An old style class"
        # Original class style have been phased out in Python 3.

    class NewStyleClass(object):
        "A new style class"
        # Introduced in Python 2.2
        #
        # Aside from this set of tests, Python Koans sticks exclusively to this
        # kind of class
        pass

    def test_new_style_classes_inherit_from_object_base_class(self):
        self.assertEqual(____, issubclass(self.NewStyleClass, object))
        self.assertEqual(____, issubclass(self.OldStyleClass, object))

    def test_new_style_classes_have_more_attributes(self):
        self.assertEqual(__, len(dir(self.OldStyleClass)))
        self.assertEqual(__, self.OldStyleClass.__doc__)
        self.assertEqual(__, self.OldStyleClass.__module__)

        self.assertEqual(__, len(dir(self.NewStyleClass)))
        # To examine the available attributes, run
        # 'dir(<Class name goes here>)'
        # from a python console

    # ------------------------------------------------------------------

    def test_old_style_classes_have_type_but_no_class_attribute(self):
        self.assertEqual(__, self.OldStyleClass.__class__)

        try:
            cls = self.OldStyleClass.__class__
        except Exception as ex:
            pass

        self.assertMatch(__, ex[0])

    def test_new_style_classes_have_same_class_as_type(self):
        new_style = self.NewStyleClass()
        self.assertEqual(__, self.NewStyleClass.__class__)
        self.assertEqual(
            __,
            type(self.NewStyleClass) == self.NewStyleClass.__class__)

    # ------------------------------------------------------------------

    def test_in_old_style_instances_class_is_different_to_type(self):
        old_style = self.OldStyleClass()
        self.assertEqual(__, old_style.__class__)

    def test_new_style_instances_have_same_class_as_type(self):
        new_style = self.NewStyleClass()
        self.assertEqual(__, new_style.__class__)
        self.assertEqual(__, type(new_style) == new_style.__class__)
<commit_msg>New style classes koan<commit_after>#!/usr/bin/env python
# -*- coding: utf-8 -*-

from runner.koan import *

class AboutNewStyleClasses(Koan):
    class OldStyleClass:
        "An old style class"
        # Original class style have been phased out in Python 3.

    class NewStyleClass(object):
        "A new style class"
        # Introduced in Python 2.2
        #
        # Aside from this set of tests, Python Koans sticks exclusively to this
        # kind of class
        pass

    def test_new_style_classes_inherit_from_object_base_class(self):
        self.assertEqual(True, issubclass(self.NewStyleClass, object))
        self.assertEqual(False, issubclass(self.OldStyleClass, object))

    def test_new_style_classes_have_more_attributes(self):
        self.assertEqual(2, len(dir(self.OldStyleClass)))
        self.assertEqual('An old style class', self.OldStyleClass.__doc__)
        self.assertEqual('koans.about_new_style_classes', self.OldStyleClass.__module__)

        self.assertEqual(18, len(dir(self.NewStyleClass)))
        # To examine the available attributes, run
        # 'dir(<Class name goes here>)'
        # from a python console

    # ------------------------------------------------------------------

    def test_old_style_classes_have_type_but_no_class_attribute(self):
        # Problem....?????
        # self.assertEqual(raise AttributeError, self.OldStyleClass.__class__)

        try:
            cls = self.OldStyleClass.__class__
        except Exception as ex:
            pass

        self.assertMatch("class OldStyleClass has no attribute '__class__'", ex[0])

    def test_new_style_classes_have_same_class_as_type(self):
        new_style = self.NewStyleClass()
        self.assertEqual(type, self.NewStyleClass.__class__)
        self.assertEqual(
            True,
            type(self.NewStyleClass) == self.NewStyleClass.__class__)

    # ------------------------------------------------------------------

    def test_in_old_style_instances_class_is_different_to_type(self):
        old_style = self.OldStyleClass()
        self.assertEqual(self.OldStyleClass, old_style.__class__)

    def test_new_style_instances_have_same_class_as_type(self):
        new_style = self.NewStyleClass()
        self.assertEqual(self.NewStyleClass, new_style.__class__)
        self.assertEqual(True, type(new_style) == new_style.__class__)
<|endoftext|>