<commit_before>import os

import numpy as np
import tensorflow as tf

import resnet
from read_data import getDataset, getTestBatchWithNoSegmentation, getFilenamesForSegmentedCats, \
    getTrainingBatchForSegmentation, getProducerForFilenames
from resnet import _bn, block, stack

UPDATE_OPERATORS = 'update_operators'

BATCH_SIZE = 32
IMG_SIZE = 224

is_training = tf.placeholder(dtype=tf.bool,
                          name='is_training')

def checkpoint_fn(layers):
    return 'ResNet-L%d.ckpt' % layers

def meta_fn(layers):
    return 'ResNet-L%d.meta' % layers

def weight_variable(shape, scale=1.0, scope='weight'):
    with tf.variable_scope(scope + '_scale_'+str(scale)):
        s = 0.2 * scale / np.sqrt(np.prod(shape[:3]))
        initial = tf.truncated_normal(shape, stddev=s)
        return tf.Variable(initial)

def deconv(x, filters_out, ksize=3, stride=1):
    shape = tf.shape(x)
    shape_list = x.get_shape().as_list()
    filters_in = shape_list[-1]

    weights = tf.get_variable('weights_transpose', [3, 3, filters_out, filters_in],
                             initializer=tf.contrib.layers.xavier_initializer_conv2d(dtype=x.dtype))  # weight_variable(, scale=7.07)
    out_size = tf.pack([shape[0], shape[1] * stride, shape[2] * stride, filters_out])
    upscaled = tf.nn.conv2d_transpose(x, weights, out_size, [1, stride, stride, 1])
    upscaled.set_shape([shape_list[0], shape_list[1] * stride, shape_list[2] * stride, filters_out])
    return upscaled

def findRegions(x):
    with tf.variable_scope('size56'):
        x = stack(x, 1, 16, bottleneck=False, is_training=is_training, stride=1)
        x = block(x, 16, is_training, 2, bottleneck=False, _conv=deconv)
    with tf.variable_scope('size112'):
        x = stack(x, 1, 8, bottleneck=False, is_training=is_training, stride=1)
        x = block(x, 8, is_training, 2, bottleneck=False, _conv=deconv)
    with tf.variable_scope('size224'):
        x = stack(x, 1, 4, bottleneck=False, is_training=is_training, stride=1)
        x = block(x, 2, is_training, 1, bottleneck=False)
    return x

def getLastLayersFromGraph(graph):
    scale2 = graph.get_tensor_by_name('scale2/block3/Relu:0') / (10)
    scale3 = graph.get_tensor_by_name('scale3/block8/Relu:0') / (10 ** 5)
    scale4 = graph.get_tensor_by_name('scale4/block36/Relu:0') / (10 ** 20)
    scale5 = graph.get_tensor_by_name('scale5/block3/Relu:0') / (10 ** 30)
    return scale2, scale3, scale4, scale5

def segmentationLoss(softmax, seg):
    tf.image_summary('segmentation', seg)
    seg = tf.squeeze(seg, [3])
    return tf.reduce_mean(-seg * tf.log(softmax[:, :, :, 0]) + (seg - 1) * tf.log(softmax[:, :, :, 1]))

def depth_conv(x, filters_out):
    with tf.variable_scope('depth_conv'):
        x_conv = resnet._conv(x, filters_out, ksize=3)
        x_conv = _bn(x_conv, is_training)
        return resnet._relu(x_conv)

def stichMutipleLayers(scales):
    with tf.variable_scope('Joining_layers'):
        scale2, scale3, scale4, scale5 = scales
        with tf.variable_scope('up2'):
            scale2 = depth_conv(scale2, 32)
            up2 = block(scale2, 8, is_training, stride=1, bottleneck=False)
        with tf.variable_scope('up3'):
            scale3 = depth_conv(scale3, 32)
            up3 = block(scale3, 8, is_training, stride=2, bottleneck=False, _conv=deconv)
        with tf.variable_scope('up4'):
            with tf.variable_scope('a'):
                scale4 = depth_conv(scale4, 32)
                up4 = block(scale4, 8, is_training, stride=2, bottleneck=False, _conv=deconv)
            with tf.variable_scope('b'):
                up4 = block(up4, 8, is_training, stride=2, bottleneck=False, _conv=deconv)
        with tf.variable_scope('up5'):
            with tf.variable_scope('a'):
                scale5 = depth_conv(scale5, 32)
                up5 = block(scale5, 8, is_training, stride=2, bottleneck=False, _conv=deconv)
            with tf.variable_scope('b'):
                up5 = block(up5, 8, is_training, stride=2, bottleneck=False, _conv=deconv)
            with tf.variable_scope('c'):
                up5 = block(up5, 8, is_training, stride=2, bottleneck=False, _conv=deconv)
        return tf.concat(3, [up2, up3, up4, up5])# tf.concat(3, [up2, up3, up4, up5])

def restore_resnet_graph(sess):
    layers = 152
    new_saver = tf.train.import_meta_graph(meta_fn(layers))
    new_saver.restore(sess, checkpoint_fn(layers))
    graph = tf.get_default_graph()
    return graph, graph.get_tensor_by_name("images:0")

def accuracy(softmax, label):
    y = tf.cast(tf.expand_dims(tf.argmax(softmax, 3), 3), tf.float32)
    return 1 - tf.reduce_mean(tf.abs(label - (1 - y)))

def makeSummary(summary_op, step, sess):
    summary_str = sess.run(summary_op, feed_dict={is_training: True, images: i, segmented_label: l})
    summary_writer.add_summary(summary_str, step)

    i_val, l_val = sess.run([img_val, seg_val])

    loss_val2, acc_val, test_summary_str = sess.run([cross_entropy, acc, summary_op],
                                  feed_dict={images: i_val, segmented_label: l_val, is_training: False})
    test_writer.add_summary(test_summary_str, step)
    print 'Validate loss:', loss_val2, 'Validate acc', acc_val

def findAvgVal(sess, img_val, seg_val):
    from sklearn.metrics import average_precision_score
    avg_acc = 0
    sm_avg = 0
    for i in range(1, 101):
        i_val, l_val = sess.run([img_val, seg_val])

        sm, acc_val = sess.run([softmaxed, acc],
                                                    feed_dict={images: i_val, segmented_label: l_val,
                                                               is_training: False})
        sm_cat = sm[:, :, :, 0]
        precision = average_precision_score(l_val.ravel(), sm_cat.ravel())
        sm_avg += precision
        avg_acc += acc_val

        print acc_val,  avg_acc/i, precision, sm_avg
    print "AVG", avg_acc/100, sm_avg/100

def loadModelFromFile(load_dir):
    step = 1
    if tf.gfile.Exists(load_dir) and load_old:
        restorer = tf.train.Saver(tf.all_variables())
        ckpt = tf.train.get_checkpoint_state(load_dir)
        if os.path.isabs(ckpt.model_checkpoint_path):
            restorer.restore(sess, ckpt.model_checkpoint_path)
        else:
            restorer.restore(sess, os.path.join(load_dir,
                                                ckpt.model_checkpoint_path))

        step_str = ckpt.model_checkpoint_path.split('/')[-1].split('-')[-1]
        print('Succesfully loaded model from %s at step=%s.' %
              (ckpt.model_checkpoint_path, step_str))
        if train_dir == load_dir:
            step = int(step_str) + 1
        return step

def inference(graph):
    #Stich togheter different layers
    stiched = stichMutipleLayers(getLastLayersFromGraph(graph))

    #Predict and upscale
    pred_region = findRegions(stiched)
    old_shape = pred_region.get_shape().as_list()
    softmaxed = tf.reshape(tf.nn.softmax(tf.reshape(pred_region, (-1, 2))), [-1] + old_shape[1:])
    tf.histogram_summary('pred_region', pred_region)
    tf.image_summary('predicted regions', tf.expand_dims(softmaxed[:, :, :, 0], 3))
    return softmaxed

def getAllFilenames():
    global img_coco_names, seg_coco_names, img_coco_names2, seg_coco_names2, img_default_names, seg_default_names
    img_coco_names, seg_coco_names = getFilenamesForSegmentedCats('/usr/local/data/coco2014_cats/segmented_cats.txt',
                                                                  image_base_path='/usr/local/data/coco2014_cats/images/',
                                                                  segmentation_base_path='/usr/local/data/coco2014_cats/segmentations/')
    img_coco_names2, seg_coco_names2 = getFilenamesForSegmentedCats('/usr/local/data/coco2014_cats2/segmented_cats.txt',
                                                                    image_base_path='/usr/local/data/coco2014_cats2/images/',
                                                                    segmentation_base_path='/usr/local/data/coco2014_cats2/segmentations/')
    img_default_names, seg_default_names = getFilenamesForSegmentedCats('data/cats_segmented_mixed.txt')

    return (img_coco_names + img_coco_names2 + img_default_names), (seg_coco_names + seg_coco_names2 + seg_default_names)

if __name__=='__main__':
    load_old = True
    train_dir = '/tmp/models/catnet12'
    test_dir = '/tmp/models/catnet12_test'
    load_dir = train_dir#'/tmp/models/catnet12'

    #Setup data-collection graph
    img_filenames, seg_filenames = getAllFilenames()
    img, seg = getTrainingBatchForSegmentation(getProducerForFilenames(img_filenames, seg_filenames), distort=True)
    img_val, seg_val = getDataset('data/cats_segmented_val.txt', scope='retrieve_test_data', distort=False)

    #Load pretrained
    sess = tf.Session()
    graph,images = restore_resnet_graph(sess)
    tf.image_summary('images', images)

    #placeholder for labels
    segmented_label = tf.placeholder(tf.float32, (None, IMG_SIZE, IMG_SIZE, 1), 'segmented_label')

    softmaxed = inference(graph)

    cross_entropy = segmentationLoss(softmaxed, segmented_label)
    tf.scalar_summary('cross_entropy', cross_entropy)
    acc = accuracy(softmaxed, segmented_label)
    tf.scalar_summary('acc', acc)

    # loss_avg
    ema = tf.train.ExponentialMovingAverage(0.9)
    apply_avg_loss =  ema.apply([cross_entropy])
    loss_avg = ema.average(cross_entropy)
    tf.scalar_summary('loss_avg', loss_avg)

    optimizer = tf.train.AdamOptimizer(0.0001)
    grads = optimizer.compute_gradients(cross_entropy)
    apply_gradients = optimizer.apply_gradients(grads)

    batchnorm_updates = tf.get_collection(resnet.UPDATE_OPS_COLLECTION)
    batchnorm_updates_op = tf.group(*batchnorm_updates)
    train_step = tf.group(apply_gradients, batchnorm_updates_op, apply_avg_loss)

    for grad, var in grads:
        if grad is not None:
            tf.histogram_summary(var.op.name + '/gradients', grad)

    saver = tf.train.Saver(tf.all_variables(), max_to_keep=10, keep_checkpoint_every_n_hours=2)
    summary_op = tf.merge_all_summaries()

    init = tf.initialize_all_variables()
    sess.run(init)

    if load_old:
        step = loadModelFromFile(load_dir)
    else:
        step = 1

    tf.train.start_queue_runners(sess=sess)
    summary_writer = tf.train.SummaryWriter(train_dir,
                                            graph=sess.graph)
    test_writer = tf.train.SummaryWriter(test_dir,
                                            graph=sess.graph)

    avg_loss = 0
    cnt = 1
    for step in range(step, 100000):
        i, l = sess.run([img, seg])
        loss_val, _ = sess.run([cross_entropy, train_step], feed_dict={images: i, segmented_label: l, is_training: True})

        avg_loss += loss_val
        print step, ':', loss_val, '\t', avg_loss/cnt
        cnt += 1
        if step % 100 == 0:
            total_loss = 0
            cnt = 1
            print 'Avg loss:', avg_loss / 100
            avg_loss = 0
            makeSummary(summary_op, step, sess)
        if step % 1000 == 0:
            saver.save(sess, train_dir + '/model.ckpt', global_step=step)
<commit_msg>Update train_catnet.py<commit_after>import os

import numpy as np
import tensorflow as tf

import resnet
from read_data import getDataset, getTestBatchWithNoSegmentation, getFilenamesForSegmentedCats, \
    getTrainingBatchForSegmentation, getProducerForFilenames
from resnet import _bn, block, stack

UPDATE_OPERATORS = 'update_operators'

BATCH_SIZE = 32
IMG_SIZE = 224

is_training = tf.placeholder(dtype=tf.bool,
                          name='is_training')

def checkpoint_fn(layers):
    return 'ResNet-L%d.ckpt' % layers

def meta_fn(layers):
    return 'ResNet-L%d.meta' % layers

def weight_variable(shape, scale=1.0, scope='weight'):
    with tf.variable_scope(scope + '_scale_'+str(scale)):
        s = 0.2 * scale / np.sqrt(np.prod(shape[:3]))
        initial = tf.truncated_normal(shape, stddev=s)
        return tf.Variable(initial)

def deconv(x, filters_out, ksize=3, stride=1):
    shape = tf.shape(x)
    shape_list = x.get_shape().as_list()
    filters_in = shape_list[-1]

    weights = tf.get_variable('weights_transpose', [3, 3, filters_out, filters_in],
                             initializer=tf.contrib.layers.xavier_initializer_conv2d(dtype=x.dtype))  # weight_variable(, scale=7.07)
    out_size = tf.pack([shape[0], shape[1] * stride, shape[2] * stride, filters_out])
    upscaled = tf.nn.conv2d_transpose(x, weights, out_size, [1, stride, stride, 1])
    upscaled.set_shape([shape_list[0], shape_list[1] * stride, shape_list[2] * stride, filters_out])
    return upscaled

def findRegions(x):
    with tf.variable_scope('size56'):
        x = stack(x, 1, 16, bottleneck=False, is_training=is_training, stride=1)
        x = block(x, 16, is_training, 2, bottleneck=False, _conv=deconv)
    with tf.variable_scope('size112'):
        x = stack(x, 1, 8, bottleneck=False, is_training=is_training, stride=1)
        x = block(x, 8, is_training, 2, bottleneck=False, _conv=deconv)
    with tf.variable_scope('size224'):
        x = stack(x, 1, 4, bottleneck=False, is_training=is_training, stride=1)
        x = block(x, 2, is_training, 1, bottleneck=False)
    return x

def getLastLayersFromGraph(graph):
    scale2 = graph.get_tensor_by_name('scale2/block3/Relu:0') / (10)
    scale3 = graph.get_tensor_by_name('scale3/block8/Relu:0') / (10 ** 5)
    scale4 = graph.get_tensor_by_name('scale4/block36/Relu:0') / (10 ** 20)
    scale5 = graph.get_tensor_by_name('scale5/block3/Relu:0') / (10 ** 30)
    return scale2, scale3, scale4, scale5

def segmentationLoss(softmax, seg):
    tf.image_summary('segmentation', seg)
    seg = tf.squeeze(seg, [3])
    return tf.reduce_mean(-seg * tf.log(softmax[:, :, :, 0]) + (seg - 1) * tf.log(softmax[:, :, :, 1]))

def depth_conv(x, filters_out):
    with tf.variable_scope('depth_conv'):
        x_conv = resnet._conv(x, filters_out, ksize=3)
        x_conv = _bn(x_conv, is_training)
        return resnet._relu(x_conv)

def stichMutipleLayers(scales):
    with tf.variable_scope('Joining_layers'):
        scale2, scale3, scale4, scale5 = scales
        with tf.variable_scope('up2'):
            scale2 = depth_conv(scale2, 32)
            up2 = block(scale2, 8, is_training, stride=1, bottleneck=False)
        with tf.variable_scope('up3'):
            scale3 = depth_conv(scale3, 32)
            up3 = block(scale3, 8, is_training, stride=2, bottleneck=False, _conv=deconv)
        with tf.variable_scope('up4'):
            with tf.variable_scope('a'):
                scale4 = depth_conv(scale4, 32)
                up4 = block(scale4, 8, is_training, stride=2, bottleneck=False, _conv=deconv)
            with tf.variable_scope('b'):
                up4 = block(up4, 8, is_training, stride=2, bottleneck=False, _conv=deconv)
        with tf.variable_scope('up5'):
            with tf.variable_scope('a'):
                scale5 = depth_conv(scale5, 32)
                up5 = block(scale5, 8, is_training, stride=2, bottleneck=False, _conv=deconv)
            with tf.variable_scope('b'):
                up5 = block(up5, 8, is_training, stride=2, bottleneck=False, _conv=deconv)
            with tf.variable_scope('c'):
                up5 = block(up5, 8, is_training, stride=2, bottleneck=False, _conv=deconv)
        return tf.concat(3, [up2, up3, up4, up5])# tf.concat(3, [up2, up3, up4, up5])

def restore_resnet_graph(sess):
    layers = 152
    new_saver = tf.train.import_meta_graph(meta_fn(layers))
    new_saver.restore(sess, checkpoint_fn(layers))
    graph = tf.get_default_graph()
    return graph, graph.get_tensor_by_name("images:0")

def accuracy(softmax, label):
    y = tf.cast(tf.expand_dims(tf.argmax(softmax, 3), 3), tf.float32)
    return 1 - tf.reduce_mean(tf.abs(label - (1 - y)))

def makeSummary(summary_op, step, sess):
    summary_str = sess.run(summary_op, feed_dict={is_training: True, images: i, segmented_label: l})
    summary_writer.add_summary(summary_str, step)

    i_val, l_val = sess.run([img_val, seg_val])

    loss_val2, acc_val, test_summary_str = sess.run([cross_entropy, acc, summary_op],
                                  feed_dict={images: i_val, segmented_label: l_val, is_training: False})
    test_writer.add_summary(test_summary_str, step)
    print 'Validate loss:', loss_val2, 'Validate acc', acc_val

def findAvgVal(sess, img_val, seg_val):
    from sklearn.metrics import average_precision_score
    avg_acc = 0
    sm_avg = 0
    for i in range(1, 101):
        i_val, l_val = sess.run([img_val, seg_val])

        sm, acc_val = sess.run([softmaxed, acc],
                                                    feed_dict={images: i_val, segmented_label: l_val,
                                                               is_training: False})
        sm_cat = sm[:, :, :, 0]
        precision = average_precision_score(l_val.ravel(), sm_cat.ravel())
        sm_avg += precision
        avg_acc += acc_val

        print acc_val,  avg_acc/i, precision, sm_avg
    print "AVG", avg_acc/100, sm_avg/100

def loadModelFromFile(load_dir):
    step = 1
    if tf.gfile.Exists(load_dir) and load_old:
        restorer = tf.train.Saver(tf.all_variables())
        ckpt = tf.train.get_checkpoint_state(load_dir)
        if os.path.isabs(ckpt.model_checkpoint_path):
            restorer.restore(sess, ckpt.model_checkpoint_path)
        else:
            restorer.restore(sess, os.path.join(load_dir,
                                                ckpt.model_checkpoint_path))

        step_str = ckpt.model_checkpoint_path.split('/')[-1].split('-')[-1]
        print('Succesfully loaded model from %s at step=%s.' %
              (ckpt.model_checkpoint_path, step_str))
        if train_dir == load_dir:
            step = int(step_str) + 1
        return step

def inference(graph):
    #Stich togheter different layers
    stiched = stichMutipleLayers(getLastLayersFromGraph(graph))

    #Predict and upscale
    pred_region = findRegions(stiched)
    old_shape = pred_region.get_shape().as_list()
    softmaxed = tf.reshape(tf.nn.softmax(tf.reshape(pred_region, (-1, 2))), [-1] + old_shape[1:])
    tf.histogram_summary('pred_region', pred_region)
    tf.image_summary('predicted regions', tf.expand_dims(softmaxed[:, :, :, 0], 3))
    return softmaxed

def getAllFilenames():
    global img_coco_names, seg_coco_names, img_coco_names2, seg_coco_names2, img_default_names, seg_default_names
    img_coco_names, seg_coco_names = getFilenamesForSegmentedCats('/usr/local/data/coco2014_cats/segmented_cats.txt',
                                                                  image_base_path='/usr/local/data/coco2014_cats/images/',
                                                                  segmentation_base_path='/usr/local/data/coco2014_cats/segmentations/')
    img_coco_names2, seg_coco_names2 = getFilenamesForSegmentedCats('/usr/local/data/coco2014_cats2/segmented_cats.txt',
                                                                    image_base_path='/usr/local/data/coco2014_cats2/images/',
                                                                    segmentation_base_path='/usr/local/data/coco2014_cats2/segmentations/')
    img_default_names, seg_default_names = getFilenamesForSegmentedCats('data/cats_segmented_mixed.txt')

    return (img_coco_names + img_coco_names2 + img_default_names), (seg_coco_names + seg_coco_names2 + seg_default_names)

if __name__=='__main__':
    load_old = True
    train_dir = '/tmp/models/catnet12'
    test_dir = '/tmp/models/catnet12_test'
    load_dir = train_dir#'/tmp/models/catnet12'

    #Setup data-collection graph
    img_filenames, seg_filenames = getAllFilenames()
    img, seg = getTrainingBatchForSegmentation(getProducerForFilenames(img_filenames, seg_filenames), distort=True)
    img_val, seg_val = getDataset('data/cats_segmented_val.txt', scope='retrieve_test_data', distort=False)

    #Load pretrained
    sess = tf.Session()
    graph,images = restore_resnet_graph(sess)
    tf.image_summary('images', images)

    #placeholder for labels
    segmented_label = tf.placeholder(tf.float32, (None, IMG_SIZE, IMG_SIZE, 1), 'segmented_label')

    softmaxed = inference(graph)

    cross_entropy = segmentationLoss(softmaxed, segmented_label)
    tf.scalar_summary('cross_entropy', cross_entropy)
    acc = accuracy(softmaxed, segmented_label)
    tf.scalar_summary('acc', acc)

    # loss_avg
    ema = tf.train.ExponentialMovingAverage(0.9)
    apply_avg_loss =  ema.apply([cross_entropy])
    loss_avg = ema.average(cross_entropy)
    tf.scalar_summary('loss_avg', loss_avg)

    optimizer = tf.train.AdamOptimizer(0.0001)
    grads = optimizer.compute_gradients(cross_entropy)
    apply_gradients = optimizer.apply_gradients(grads)

    batchnorm_updates = tf.get_collection(resnet.UPDATE_OPS_COLLECTION)
    batchnorm_updates_op = tf.group(*batchnorm_updates)
    train_step = tf.group(apply_gradients, batchnorm_updates_op, apply_avg_loss)

    for grad, var in grads:
        if grad is not None:
            tf.histogram_summary(var.op.name + '/gradients', grad)

    saver = tf.train.Saver(tf.all_variables(), max_to_keep=10, keep_checkpoint_every_n_hours=2)
    summary_op = tf.merge_all_summaries()

    init = tf.initialize_all_variables()
    sess.run(init)

    if load_old:
        step = loadModelFromFile(load_dir)
    else:
        step = 1

    tf.train.start_queue_runners(sess=sess)
    summary_writer = tf.train.SummaryWriter(train_dir,
                                            graph=sess.graph)
    test_writer = tf.train.SummaryWriter(test_dir,
                                            graph=sess.graph)

    avg_loss = 0
    cnt = 1
    for step in range(step, 100000):
        i, l = sess.run([img, seg])
        loss_val, _ = sess.run([cross_entropy, train_step], feed_dict={images: i, segmented_label: l, is_training: True})

        avg_loss += loss_val
        print step, ':', loss_val, '\t', avg_loss/cnt
        cnt += 1
        if step % 100 == 0:
            cnt = 1
            print 'Avg loss:', avg_loss / 100
            avg_loss = 0
            makeSummary(summary_op, step, sess)
        if step % 1000 == 0:
            saver.save(sess, train_dir + '/model.ckpt', global_step=step)
<|endoftext|>


<commit_before>#!/usr/bin/env python
import argparse
import os
import logging

# PYTHON_ARGCOMPLETE_OK

from sympy import (nsolve, symbols, Mul, Add, chebyshevt, exp, simplify,
    chebyshevt_root, Tuple, diff, plot, N, solve, together, Poly)

from sympy.utilities.decorator import conserve_mpmath_dps

logger = logging.getLogger(__name__)
logger.addHandler(logging.StreamHandler())
# Change INFO to DEBUG for more output
logger.setLevel(logging.INFO)

def general_rat_func(d, x, chebyshev=False):
    """
    Return a general rational function with numerator and denominator degree d

    Returns a tuple, (rational function, numerator coefficients, denominator coefficients)

    The constant coefficient in the denominator is always 1.
    """
    num_coeffs = symbols('p:%s' % (d+1))
    den_coeffs = symbols('q:%s' % (d+1))[1:]
    if chebyshev:
        num = Add(*(Mul(c, chebyshevt(i, x)) for i, c in enumerate(num_coeffs)))
        den = Add(*(Mul(c, chebyshevt(i, x)) for i, c in enumerate([1, *den_coeffs])))
        rat_func = num/den
    else:
        rat_func = Poly(reversed(num_coeffs), x)/Poly([*reversed(den_coeffs), 1], x)
    return rat_func, num_coeffs, den_coeffs

def nsolve_intervals(expr, bounds, division=30, solver='bisect', **kwargs):
    """
    Divide bounds into division intervals and nsolve in each one
    """
    roots = []
    L = bounds[1] - bounds[0]
    for i in range(division):
        interval = [bounds[0] + i*L/division, bounds[0] + (i + 1)*L/division]
        try:
            logger.debug("Solving in interval %s", interval)
            root = nsolve(expr, interval, solver=solver, **kwargs)
        except ValueError:
            logger.debug("No solution found")
            continue
        else:
            if interval[0] < root < interval[1]:
                logger.debug("Solution found: %s", root)
                roots.append(root)
            else:
                logger.warn("%s is not in %s, discarding", root, interval)

    return roots

def plot_in_terminal(*args, logname=None, **kwargs):
    """
    Run plot() but show in terminal if possible
    """
    try:
        from iterm2_tools.images import display_image_bytes
    except ImportError:
        p = plot(*args, **kwargs)
        if logname:
            os.makedirs('plots', exist_ok=True)
            p.save('plots/%s.png' % logname)
    else:
        from sympy.plotting.plot import unset_show
        from io import BytesIO
        unset_show()
        p = plot(*args, **kwargs, show=False)
        b = BytesIO()
        p.save(b)
        if logname:
            os.makedirs('plots', exist_ok=True)
            p.save('plots/%s.png' % logname)
        print(display_image_bytes(b.getvalue()))
        p._backend.close()

def _get_log_file_name(locals_dict):
    d = locals_dict.copy()
    kwargs = d.pop('kwargs')
    d.update(kwargs)
    d.setdefault('maxsteps')
    d.setdefault('division')
    degree = d.pop('degree')
    prec = d.pop('prec')
    info = 'degree=%s prec=%s ' % (degree, prec)
    info += ' '.join('%s=%s' % (i, d[i]) for i in sorted(d))
    return info

# This decorator is actually not needed any more, but we leave it in as it
# will fail early if we are not running in SymPy master.
@conserve_mpmath_dps
def CRAM_exp(degree, prec=128, *, max_loops=10, c=None, maxsteps=None, **kwargs):
    """
    Compute the CRAM approximation of exp(-t) from t in [0, oo) of the given degree

    The Remez algorithm is used.

    degree is the degree of the numerator and denominator of the
    approximation.

    prec is the precision of the floats used in the calculation. Note that, as
    of now, the result may not be accurate to prec digits.

    max_loops is the number of loops to run the Remez algorithm before giving
    up (default 10).

    c is the coefficient of the transformation c*(t + 1)/(t - 1) that is
    applied to exp(-t) to translate it to the interval [-1, 1]. It should be
    chosen so that the maximal errors of the approximation are equally spaced
    in that interval. The default is 0.6*degree.

    maxsteps is the argument passed to the mpmath bisection solver. The
    default is 1.7*prec. See also
    https://github.com/fredrik-johansson/mpmath/issues/339.

    Additional keyword arguments are passed to nsolve_intervals, such as
    division.

    The SymPy master branch is required for this to work.

    """
    logger.info("CRAM_exp with arguments %s", locals())

    if kwargs.get('log', True):
        os.makedirs('logs', exist_ok=True)
        logname = _get_log_file_name(locals())
        logger.addHandler(logging.FileHandler('logs/%s.log' % logname))
        logger.info("Logging to file 'logs/%s.log'", logname)
    else:
        logname = None

    epsilon, t, i, y = symbols("epsilon t i y")

    c = c or 0.6*degree
    maxsteps = int(maxsteps or 1.7*prec)

    r, num_coeffs, den_coeffs = general_rat_func(degree, t, chebyshev=True)
    E = exp(c*(t + 1)/(t - 1)) - r
    expr = E + (-1)**i*epsilon
    expr = expr*r.as_numer_denom()[1]
    expr = simplify(expr)

    points = [chebyshevt_root(2*(degree + 1) + 1, 2*(degree + 1) - j) for j in range(1, 2*(degree + 1) + 1)]
    points = [i.evalf(prec) for i in points]
    for iteration in range(max_loops):
        logger.info('-'*80)
        logger.info("Iteration %s:", iteration)
        system = Tuple(*[expr.subs({i: j, t: points[j]}) for j in range(2*degree + 1)])
        system = system + Tuple(expr.replace(exp, lambda i: 0).subs({i: 2*degree + 1, t: 1}))
        logger.debug('system: %s', system)
        logger.debug('[*num_coeffs, *den_coeffs, epsilon]: %s', [*num_coeffs, *den_coeffs, epsilon])
        sol = dict(zip([*num_coeffs, *den_coeffs, epsilon], nsolve(system, [*num_coeffs, *den_coeffs, epsilon], [*[1]*(2*(degree + 1) - 1), 0], prec=prec)))
        logger.info('sol: %s', sol)
        logger.info('system.subs(sol): %s', [i.evalf() for i in system.subs(sol)])
        D = diff(E.subs(sol), t)
        plot_in_terminal(E.subs(sol), (t, -1, 0.999), adaptive=False,
            nb_of_points=1000, logname=logname + 'iteration=%s' % iteration)
        logger.info('E.subs(sol): %s', E.subs(sol))

        # we can't use 1 because of the singularity
        points = [-1, *nsolve_intervals(D, [-1, 0.999999], prec=prec, tol=10**-(2*prec), **kwargs), 1]
        logger.debug('points: %s', points)
        logger.info('D: %s', D)
        logger.info('[(i, D.subs(t, i)) for i in points]: %s', [(i, D.subs(t, i)) for i in points])
        if not len(points) == 2*(degree + 1):
            logger.error("ERROR: len(points) is (%s), not 2*(degree + 1) (%s)",
                len(points), 2*(degree + 1))
            raise RuntimeError
        Evals = [E.evalf(prec, subs={**sol, t: point}) for point in points[:-1]] + [-r.evalf(prec, subs={**sol, t: 1})]
        logger.info('Evals: %s', Evals)
        maxmin = N(max(map(abs, Evals)) - min(map(abs, Evals)))
        logger.info('max - min: %s', maxmin)
        logger.info('epsilon: %s', N(sol[epsilon]))
        if maxmin < 10**-prec:
            logger.info("Converged in %d iterations.", iteration + 1)
            break
    else:
        logger.warn("!!!WARNING: DID NOT CONVERGE AFTER %d ITERATIONS!!!", max_loops)

    inv = solve(-c*(t + 1)/(t - 1) - y, t)[0].subs(y, t)
    n, d = together(r.subs(sol).subs(t, inv)).as_numer_denom() # simplify/cancel here will add degree to the numerator and denominator
    rat_func = (Poly(n)/Poly(d).TC())/(Poly(d)/Poly(d).TC())
    return rat_func.evalf(prec)

def main():
    parser = argparse.ArgumentParser(description=__doc__)

    parser.add_argument('degree', type=int)
    parser.add_argument('prec', type=int)
    parser.add_argument('--division', type=int)
    parser.add_argument('--c', type=float)
    parser.add_argument('--maxsteps', type=int)
    parser.add_argument('--max-loops', type=int)
    parser.add_argument('--log-level', default=None, choices=['debug', 'info',
        'warning', 'error', 'critical'])
    try:
        import argcomplete
        argcomplete.autocomplete(parser)
    except ImportError:
        pass
    args = parser.parse_args()

    if args.log_level:
        logger.setLevel(getattr(logging, args.log_level.upper()))

    t = symbols('t')
    arguments = args.__dict__.copy()
    for i in arguments.copy():
        if not arguments[i]:
           del arguments[i]
    del arguments['log_level']

    rat_func = CRAM_exp(**arguments)
    logger.info('rat_func: %s', rat_func)
    # XXX: log these
    plot_in_terminal(rat_func - exp(-t), (t, 0, 100), adaptive=False, nb_of_points=1000)

if __name__ == '__main__':
    main()
<commit_msg>Fix command when --log-level is not specified<commit_after>#!/usr/bin/env python
import argparse
import os
import logging

# PYTHON_ARGCOMPLETE_OK

from sympy import (nsolve, symbols, Mul, Add, chebyshevt, exp, simplify,
    chebyshevt_root, Tuple, diff, plot, N, solve, together, Poly)

from sympy.utilities.decorator import conserve_mpmath_dps

logger = logging.getLogger(__name__)
logger.addHandler(logging.StreamHandler())
# Change INFO to DEBUG for more output
logger.setLevel(logging.INFO)

def general_rat_func(d, x, chebyshev=False):
    """
    Return a general rational function with numerator and denominator degree d

    Returns a tuple, (rational function, numerator coefficients, denominator coefficients)

    The constant coefficient in the denominator is always 1.
    """
    num_coeffs = symbols('p:%s' % (d+1))
    den_coeffs = symbols('q:%s' % (d+1))[1:]
    if chebyshev:
        num = Add(*(Mul(c, chebyshevt(i, x)) for i, c in enumerate(num_coeffs)))
        den = Add(*(Mul(c, chebyshevt(i, x)) for i, c in enumerate([1, *den_coeffs])))
        rat_func = num/den
    else:
        rat_func = Poly(reversed(num_coeffs), x)/Poly([*reversed(den_coeffs), 1], x)
    return rat_func, num_coeffs, den_coeffs

def nsolve_intervals(expr, bounds, division=30, solver='bisect', **kwargs):
    """
    Divide bounds into division intervals and nsolve in each one
    """
    roots = []
    L = bounds[1] - bounds[0]
    for i in range(division):
        interval = [bounds[0] + i*L/division, bounds[0] + (i + 1)*L/division]
        try:
            logger.debug("Solving in interval %s", interval)
            root = nsolve(expr, interval, solver=solver, **kwargs)
        except ValueError:
            logger.debug("No solution found")
            continue
        else:
            if interval[0] < root < interval[1]:
                logger.debug("Solution found: %s", root)
                roots.append(root)
            else:
                logger.warn("%s is not in %s, discarding", root, interval)

    return roots

def plot_in_terminal(*args, logname=None, **kwargs):
    """
    Run plot() but show in terminal if possible
    """
    try:
        from iterm2_tools.images import display_image_bytes
    except ImportError:
        p = plot(*args, **kwargs)
        if logname:
            os.makedirs('plots', exist_ok=True)
            p.save('plots/%s.png' % logname)
    else:
        from sympy.plotting.plot import unset_show
        from io import BytesIO
        unset_show()
        p = plot(*args, **kwargs, show=False)
        b = BytesIO()
        p.save(b)
        if logname:
            os.makedirs('plots', exist_ok=True)
            p.save('plots/%s.png' % logname)
        print(display_image_bytes(b.getvalue()))
        p._backend.close()

def _get_log_file_name(locals_dict):
    d = locals_dict.copy()
    kwargs = d.pop('kwargs')
    d.update(kwargs)
    d.setdefault('maxsteps')
    d.setdefault('division')
    degree = d.pop('degree')
    prec = d.pop('prec')
    info = 'degree=%s prec=%s ' % (degree, prec)
    info += ' '.join('%s=%s' % (i, d[i]) for i in sorted(d))
    return info

# This decorator is actually not needed any more, but we leave it in as it
# will fail early if we are not running in SymPy master.
@conserve_mpmath_dps
def CRAM_exp(degree, prec=128, *, max_loops=10, c=None, maxsteps=None, **kwargs):
    """
    Compute the CRAM approximation of exp(-t) from t in [0, oo) of the given degree

    The Remez algorithm is used.

    degree is the degree of the numerator and denominator of the
    approximation.

    prec is the precision of the floats used in the calculation. Note that, as
    of now, the result may not be accurate to prec digits.

    max_loops is the number of loops to run the Remez algorithm before giving
    up (default 10).

    c is the coefficient of the transformation c*(t + 1)/(t - 1) that is
    applied to exp(-t) to translate it to the interval [-1, 1]. It should be
    chosen so that the maximal errors of the approximation are equally spaced
    in that interval. The default is 0.6*degree.

    maxsteps is the argument passed to the mpmath bisection solver. The
    default is 1.7*prec. See also
    https://github.com/fredrik-johansson/mpmath/issues/339.

    Additional keyword arguments are passed to nsolve_intervals, such as
    division.

    The SymPy master branch is required for this to work.

    """
    logger.info("CRAM_exp with arguments %s", locals())

    if kwargs.get('log', True):
        os.makedirs('logs', exist_ok=True)
        logname = _get_log_file_name(locals())
        logger.addHandler(logging.FileHandler('logs/%s.log' % logname))
        logger.info("Logging to file 'logs/%s.log'", logname)
    else:
        logname = None

    epsilon, t, i, y = symbols("epsilon t i y")

    c = c or 0.6*degree
    maxsteps = int(maxsteps or 1.7*prec)

    r, num_coeffs, den_coeffs = general_rat_func(degree, t, chebyshev=True)
    E = exp(c*(t + 1)/(t - 1)) - r
    expr = E + (-1)**i*epsilon
    expr = expr*r.as_numer_denom()[1]
    expr = simplify(expr)

    points = [chebyshevt_root(2*(degree + 1) + 1, 2*(degree + 1) - j) for j in range(1, 2*(degree + 1) + 1)]
    points = [i.evalf(prec) for i in points]
    for iteration in range(max_loops):
        logger.info('-'*80)
        logger.info("Iteration %s:", iteration)
        system = Tuple(*[expr.subs({i: j, t: points[j]}) for j in range(2*degree + 1)])
        system = system + Tuple(expr.replace(exp, lambda i: 0).subs({i: 2*degree + 1, t: 1}))
        logger.debug('system: %s', system)
        logger.debug('[*num_coeffs, *den_coeffs, epsilon]: %s', [*num_coeffs, *den_coeffs, epsilon])
        sol = dict(zip([*num_coeffs, *den_coeffs, epsilon], nsolve(system, [*num_coeffs, *den_coeffs, epsilon], [*[1]*(2*(degree + 1) - 1), 0], prec=prec)))
        logger.info('sol: %s', sol)
        logger.info('system.subs(sol): %s', [i.evalf() for i in system.subs(sol)])
        D = diff(E.subs(sol), t)
        plot_in_terminal(E.subs(sol), (t, -1, 0.999), adaptive=False,
            nb_of_points=1000, logname=logname + 'iteration=%s' % iteration)
        logger.info('E.subs(sol): %s', E.subs(sol))

        # we can't use 1 because of the singularity
        points = [-1, *nsolve_intervals(D, [-1, 0.999999], prec=prec, tol=10**-(2*prec), **kwargs), 1]
        logger.debug('points: %s', points)
        logger.info('D: %s', D)
        logger.info('[(i, D.subs(t, i)) for i in points]: %s', [(i, D.subs(t, i)) for i in points])
        if not len(points) == 2*(degree + 1):
            logger.error("ERROR: len(points) is (%s), not 2*(degree + 1) (%s)",
                len(points), 2*(degree + 1))
            raise RuntimeError
        Evals = [E.evalf(prec, subs={**sol, t: point}) for point in points[:-1]] + [-r.evalf(prec, subs={**sol, t: 1})]
        logger.info('Evals: %s', Evals)
        maxmin = N(max(map(abs, Evals)) - min(map(abs, Evals)))
        logger.info('max - min: %s', maxmin)
        logger.info('epsilon: %s', N(sol[epsilon]))
        if maxmin < 10**-prec:
            logger.info("Converged in %d iterations.", iteration + 1)
            break
    else:
        logger.warn("!!!WARNING: DID NOT CONVERGE AFTER %d ITERATIONS!!!", max_loops)

    inv = solve(-c*(t + 1)/(t - 1) - y, t)[0].subs(y, t)
    n, d = together(r.subs(sol).subs(t, inv)).as_numer_denom() # simplify/cancel here will add degree to the numerator and denominator
    rat_func = (Poly(n)/Poly(d).TC())/(Poly(d)/Poly(d).TC())
    return rat_func.evalf(prec)

def main():
    parser = argparse.ArgumentParser(description=__doc__)

    parser.add_argument('degree', type=int)
    parser.add_argument('prec', type=int)
    parser.add_argument('--division', type=int)
    parser.add_argument('--c', type=float)
    parser.add_argument('--maxsteps', type=int)
    parser.add_argument('--max-loops', type=int)
    parser.add_argument('--log-level', default=None, choices=['debug', 'info',
        'warning', 'error', 'critical'])
    try:
        import argcomplete
        argcomplete.autocomplete(parser)
    except ImportError:
        pass
    args = parser.parse_args()

    arguments = args.__dict__.copy()
    for i in arguments.copy():
        if not arguments[i]:
           del arguments[i]
    if args.log_level:
        logger.setLevel(getattr(logging, args.log_level.upper()))
        del arguments['log_level']

    rat_func = CRAM_exp(**arguments)
    logger.info('rat_func: %s', rat_func)
    # TODO: log this plot
    t = symbols('t')
    plot_in_terminal(rat_func - exp(-t), (t, 0, 100), adaptive=False, nb_of_points=1000)

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


<commit_before>from __future__ import print_function
import argparse
import os, sys, subprocess, json, re

def run(*args):
    ret = subprocess.call(args,  stdout=sys.stdout, stderr=sys.stderr)
    if ret != 0:
        exit(ret)

def run_filter(filter, *args):
    replacement = 'X' * len(filter)
    try:
        output = subprocess.check_output(args,
                                         stderr = subprocess.STDOUT)
    except subprocess.CalledProcessError as e:
        print(e.output.decode('utf-8').replace(filter, replacement))
        exit(e.returncode)
    print(output.decode('utf-8').replace(filter, replacement))

def run_output(*args):
    try:
        output = subprocess.check_output(args,
                                         stderr=sys.stderr)
    except subprocess.CalledProcessError as e:
        print(e.output.decode('utf-8'))
        exit(e.returncode)
    return output.decode('utf-8')

def target_binary_name(target):
    return target['name'].replace('-', '_') + target['metadata']['extra_filename']

class Manifest(object):
    def __init__(self, dir, version):
        # the --manifest-path behaviour changed in
        # https://github.com/rust-lang/cargo/pull/1955, so we need to
        # be careful to handle both
        path_file = os.path.join(dir, 'Cargo.toml')
        path_dir = dir

        try:
            output = subprocess.check_output(['cargo', 'read-manifest', '--manifest-path',
                                              path_file],
                                             stderr=subprocess.STDOUT)
        except subprocess.CalledProcessError as e:
            try:
                output = subprocess.check_output(['cargo', 'read-manifest', '--manifest-path',
                                                  path_dir],
                                                 stderr=subprocess.STDOUT)
            except subprocess.CalledProcessError as e2:
                print('Cargo failed to read `--manifest-path {}`'
                      ' and `--manifest-path {}`:'.format(path_file, path_dir),
                      file = sys.stderr)
                print(e.output.decode('utf-8'))
                print(e2.output.decode('utf-8'))
                exit(e.returncode)

        self.manifest = json.loads(output.decode('utf-8'))

    def targets(self):
        return self.manifest['targets']
    def lib_name(self):
        for target in self.targets():
            if 'lib' in target['kind']:
                return target['name'].replace('-', '_')
        return None

def add_features(cargo_args, version):
    nightly_feature = os.environ.get('TRAVIS_CARGO_NIGHTLY_FEATURE', 'unstable')
    if version == 'nightly' and nightly_feature != '':
        # only touch feature arguments when we are actually going to
        # add something non-trivial, avoids problems like that in
        # issue #14 (can't use -p ... on nightly even with an empty
        # nightly feature).
        for i in range(0, len(cargo_args)):
            if cargo_args[i] == '--features':
                cargo_args[i + 1] += ' ' + nightly_feature
                break
            elif cargo_args[i].startswith('--features='):
                cargo_args[i] += ' ' + nightly_feature
                break
        else:
            cargo_args += ['--features', nightly_feature]

def cargo_raw(feature, version, manifest, args):
    cargo_args = args.cargo_args
    subcommand = args.subcommand
    if subcommand == 'bench' and version != 'nightly':
        print('skipping `cargo bench` on non-nightly version')
        return

    if feature:
        add_features(cargo_args, version)

    if not args.quiet and '--verbose' not in args and '-v' not in cargo_args:
        cargo_args.append('--verbose')

    run('cargo', subcommand, *cargo_args)

def cargo_feature(version, manifest, args):
    cargo_raw(True, version, manifest, args)
def cargo_no_feature(version, manifest, args):
    cargo_raw(False, version, manifest, args)

def doc_upload(version, manifest, args):
    branch = os.environ.get('APPVEYOR_REPO_BRANCH') or os.environ['TRAVIS_BRANCH']
    repo = os.environ.get('APPVEYOR_REPO_NAME') or os.environ['TRAVIS_REPO_SLUG']
    if os.environ.get('APPVEYOR_PULL_REQUEST_NUMBER'):
        pr = 'true'
    else:
        pr = os.environ.get('TRAVIS_PULL_REQUEST', 'false')

    lib_name = manifest.lib_name()
    if lib_name is None:
        sys.stderr.write('error: uploading docs for package with no library')
        exit(1)

    if branch == args.branch and pr == 'false':
        # only load the token when we're sure we're uploading (travis
        # won't decrypt secret keys for PRs, so loading this with the
        # other vars causes problems with tests)
        token = os.environ['GH_TOKEN']

        print('uploading docs...')
        sys.stdout.flush()
        with open('target/doc/index.html', 'w') as f:
            f.write('<meta http-equiv=refresh content=0;url=%s/index.html>' % lib_name)

        run('git', 'clone', 'https://github.com/davisp/ghp-import')
        run(sys.executable, './ghp-import/ghp-import', '-n', 'target/doc')
        run_filter(token, 'git', 'push', '-fq', 'https://%s@github.com/%s.git' % (token, repo), 'gh-pages')

def build_kcov(use_sudo, verify):
    deps = ''
    if use_sudo:
        deps = 'sudo apt-get install libcurl4-openssl-dev libelf-dev libdw-dev cmake'
        if verify:
            deps += ' binutils-dev'
    init = deps + '''
    wget https://github.com/SimonKagstrom/kcov/archive/master.zip
    unzip master.zip
    mv kcov-master kcov
    mkdir kcov/build
    '''
    for line in init.split('\n'):
        line = line.strip()
        if line:
            print('Running: %s' % line)
            run(*line.split())
    current = os.getcwd()
    os.chdir('kcov/build')
    build = '''
    cmake ..
    make
    '''
    for line in build.split('\n'):
        line = line.strip()
        if line:
            print('Running: %s' % line)
            run(*line.split())
    os.chdir(current)
    return os.path.join(current, 'kcov/build/src/kcov')

def raw_coverage(use_sudo, verify, test_args, merge_msg, kcov_merge_args, kcov_merge_dir,
                 exclude_pattern):
    kcov = build_kcov(use_sudo, verify)

    test_binaries = []
    # look through the output of `cargo test` to find the test
    # binaries.
    # FIXME: the information cargo feeds us is
    # inconsistent/inaccurate, so using the output of read-manifest is
    # far too much trouble.

    output = run_output('cargo', 'test', *test_args)
    running = re.compile('^     Running target/debug/(.*)$', re.M)
    for line in running.finditer(output):
        test_binaries.append(line.group(1))

    # record coverage for each binary
    for binary in test_binaries:
        print('Recording %s' % binary)
        kcov_args = [kcov]
        if verify:
            kcov_args += ['--verify']
        exclude_pattern_arg = '--exclude-pattern=/.cargo'
        if exclude_pattern:
            exclude_pattern_arg += ',{}'.format(exclude_pattern)
        kcov_args += [exclude_pattern_arg, 'target/kcov-' + binary,
                      'target/debug/' + binary]
        print('Running: {}'.format(' '.join(kcov_args)))
        run(*kcov_args)
    # merge all the coverages and upload in one go
    print(merge_msg)
    kcov_args = [kcov, '--merge'] + kcov_merge_args + [kcov_merge_dir]
    kcov_args += ('target/kcov-' + b for b in test_binaries)
    run(*kcov_args)

def coverage(version, manifest, args):
    cargo_args = args.cargo_args
    add_features(cargo_args, version)

    kcov_merge_dir = args.merge_into
    raw_coverage(not args.no_sudo, args.verify, cargo_args, 'Merging coverage', [], kcov_merge_dir,
                 args.exclude_pattern)

def coveralls(version, manifest, args):
    job_id = os.environ['TRAVIS_JOB_ID']

    cargo_args = args.cargo_args
    add_features(cargo_args, version)

    raw_coverage(not args.no_sudo, args.verify, cargo_args, 'Uploading coverage',
                 ['--coveralls-id=' + job_id], 'target/kcov', args.exclude_pattern)

# user interface
class ScInfo(object):
    def __init__(self, func, description, arguments, help_=None, is_cargo=False):
        self.func = func
        self.help = help_
        self.description = description
        self.arguments = arguments
        self.is_cargo = is_cargo
        for _name, options in arguments:
            assert isinstance(options, dict)

EXCLUDE_PATTERN = (['--exclude-pattern'], {
    'metavar': 'PATTERN',
    'help': 'pass additional exclusionary patterns to kcov'
})

NO_SUDO = (['--no-sudo'], {
    'action': 'store_true',
    'default': False,
    'help': 'don\'t use `sudo` to install kcov\'s deps. Requires that '
    'libcurl4-openssl-dev, libelf-dev and libdw-dev are installed (e.g. via '
    '`addons: apt: packages:`)'
})
VERIFY = (['--verify'], {
    'action': 'store_true',
    'default': False,
    'help': 'pass `--verify` to kcov, to avoid some crashes. See '
    '<https://github.com/huonw/travis-cargo/issues/12>. This requires '
    'installing the `binutils-dev` package.'
})

SC_INFO = {
    'doc-upload': ScInfo(func = doc_upload,
                         description = 'Use ghp-import to upload cargo-rendered '
                         'docs to Github Pages, from the master branch.',
                         help_ = 'upload documentation to Github pages.',
                         arguments = [(['--branch'], {
                             'default': 'master',
                             'help': 'upload docs when on this branch, defaults to master',
                         })]),
    'coveralls': ScInfo(func = coveralls,
                        description = 'Record coverage of `cargo test` and upload to '
                        'coveralls.io with kcov, this runs all binaries that `cargo test` runs '
                        'but not doc tests. Merged kcov results can be accessed in `target/kcov`.',
                        help_ = 'record and upload code coverage to coveralls.io',
                        arguments = [(['cargo_args'], {
                            'metavar': 'ARGS',
                            'nargs': '*',
                            'help': 'arguments to pass to `cargo test`'
                        }),
                                     EXCLUDE_PATTERN,
                                     NO_SUDO,
                                     VERIFY]),
    'coverage': ScInfo(func = coverage,
                       description = 'Record coverage of `cargo test`, this runs all '
                       'binaries that `cargo test` runs but not doc tests. The results '
                       'of all tests are merged into a single directory.',
                       help_ = 'record code coverage',
                       arguments = [(['cargo_args'], {
                           'metavar': 'ARGS',
                           'nargs': '*',
                           'help': 'arguments to pass to `cargo test`'
                       }),
                                    (['-m', '--merge-into'], {
                                        'metavar': 'DIR',
                                        'default': 'target/kcov',
                                        'help': 'the directory to put the final merged kcov '
                                        'result into (default `target/kcov`)'
                                    }),
                                    EXCLUDE_PATTERN,
                                    NO_SUDO,
                                    VERIFY])
}

def cargo_sc(name, features):
    return ScInfo(func = cargo_feature if features else cargo_no_feature,
                  description = 'Run `cargo %s`' % name,
                  is_cargo = True,
                  arguments = [(['cargo_args'], {
                      'metavar': 'ARGS',
                      'nargs': '*',
                      'help': 'arguments to pass to `cargo %s`' % name
                  })])

NO_FEATURE_CARGO = [
    'clean', 'fetch', 'generate-lockfile', 'git-checkout', 'help', 'locate-project',
    'login', 'new', 'owner', 'package', 'pkgid', 'publish', 'read-manifest', 'search',
    'update', 'verify-project', 'version', 'yank'
]
FEATURE_CARGO = [
    'build', 'bench', 'test', 'doc', 'run', 'rustc', 'rustdoc',
]
SC_INFO.update((c, cargo_sc(c, False)) for c in NO_FEATURE_CARGO)
SC_INFO.update((c, cargo_sc(c, True)) for c in FEATURE_CARGO)

def main():
    parser = argparse.ArgumentParser(description = '''
Manages interactions between Travis and Cargo and common tooling tasks.
''')

    parser.add_argument('-q','--quiet', action='store_true', default=False,
                        help='don\'t pass --verbose to cargo subcommands')
    parser.add_argument('--only', metavar='VERSION',
                        help='only run the given command if the specified version matches '
                        '`TRAVIS_RUST_VERSION`')
    parser.add_argument('--skip', metavar='VERSION',
                        help='only run the given command if the specified version does not match '
                        '`TRAVIS_RUST_VERSION`')

    sb = parser.add_subparsers(metavar = '{coverage,coveralls,doc-upload,...}',
                               description = '''

travis-cargo supports all cargo subcommands, and selected others
(listed below).

Cargo subcommands have `--verbose` added to their invocation by
default, and, when running with a nightly compiler, `--features
unstable` (or `--features $TRAVIS_CARGO_NIGHTLY_FEATURE` if that
environment variable is defined) if `--features` is a valid argument.

''',
                               dest = 'subcommand')
    for _, name, sc in sorted((sc.is_cargo, n, sc) for n, sc in SC_INFO.items()):
        extra_args = {'help': sc.help} if sc.help is not None else {}
        sub_parser = sb.add_parser(name,
                                   description = sc.description,
                                   **extra_args)
        sub_parser.set_defaults(func = sc.func)
        for name, options in sc.arguments:
            sub_parser.add_argument(*name, **options)

    args = parser.parse_args()

    version = os.environ.get('TRAVIS_RUST_VERSION', None)
    if version is None:
        # fill in the version based on the compiler's version output.
        output = run_output('rustc', '-V')
        phrases = ['nightly', ('dev', 'nightly'), 'beta']
        for phrase in phrases:
            if isinstance(phrase, tuple):
                alias = phrase[1]
                phrase = phrase[0]
            else:
                alias = phrase
            if phrase in output:
                version = alias
                break

    if args.only and args.only != version:
        return

    if args.skip and args.skip == version:
      return

    manifest = Manifest(os.getcwd(), version)
    args.func(version, manifest, args)
<commit_msg>Add flag for arbitrary kcov options<commit_after>from __future__ import print_function
import argparse
import os, sys, subprocess, json, re

def run(*args):
    ret = subprocess.call(args,  stdout=sys.stdout, stderr=sys.stderr)
    if ret != 0:
        exit(ret)

def run_filter(filter, *args):
    replacement = 'X' * len(filter)
    try:
        output = subprocess.check_output(args,
                                         stderr = subprocess.STDOUT)
    except subprocess.CalledProcessError as e:
        print(e.output.decode('utf-8').replace(filter, replacement))
        exit(e.returncode)
    print(output.decode('utf-8').replace(filter, replacement))

def run_output(*args):
    try:
        output = subprocess.check_output(args,
                                         stderr=sys.stderr)
    except subprocess.CalledProcessError as e:
        print(e.output.decode('utf-8'))
        exit(e.returncode)
    return output.decode('utf-8')

def target_binary_name(target):
    return target['name'].replace('-', '_') + target['metadata']['extra_filename']

class Manifest(object):
    def __init__(self, dir, version):
        # the --manifest-path behaviour changed in
        # https://github.com/rust-lang/cargo/pull/1955, so we need to
        # be careful to handle both
        path_file = os.path.join(dir, 'Cargo.toml')
        path_dir = dir

        try:
            output = subprocess.check_output(['cargo', 'read-manifest', '--manifest-path',
                                              path_file],
                                             stderr=subprocess.STDOUT)
        except subprocess.CalledProcessError as e:
            try:
                output = subprocess.check_output(['cargo', 'read-manifest', '--manifest-path',
                                                  path_dir],
                                                 stderr=subprocess.STDOUT)
            except subprocess.CalledProcessError as e2:
                print('Cargo failed to read `--manifest-path {}`'
                      ' and `--manifest-path {}`:'.format(path_file, path_dir),
                      file = sys.stderr)
                print(e.output.decode('utf-8'))
                print(e2.output.decode('utf-8'))
                exit(e.returncode)

        self.manifest = json.loads(output.decode('utf-8'))

    def targets(self):
        return self.manifest['targets']
    def lib_name(self):
        for target in self.targets():
            if 'lib' in target['kind']:
                return target['name'].replace('-', '_')
        return None

def add_features(cargo_args, version):
    nightly_feature = os.environ.get('TRAVIS_CARGO_NIGHTLY_FEATURE', 'unstable')
    if version == 'nightly' and nightly_feature != '':
        # only touch feature arguments when we are actually going to
        # add something non-trivial, avoids problems like that in
        # issue #14 (can't use -p ... on nightly even with an empty
        # nightly feature).
        for i in range(0, len(cargo_args)):
            if cargo_args[i] == '--features':
                cargo_args[i + 1] += ' ' + nightly_feature
                break
            elif cargo_args[i].startswith('--features='):
                cargo_args[i] += ' ' + nightly_feature
                break
        else:
            cargo_args += ['--features', nightly_feature]

def cargo_raw(feature, version, manifest, args):
    cargo_args = args.cargo_args
    subcommand = args.subcommand
    if subcommand == 'bench' and version != 'nightly':
        print('skipping `cargo bench` on non-nightly version')
        return

    if feature:
        add_features(cargo_args, version)

    if not args.quiet and '--verbose' not in args and '-v' not in cargo_args:
        cargo_args.append('--verbose')

    run('cargo', subcommand, *cargo_args)

def cargo_feature(version, manifest, args):
    cargo_raw(True, version, manifest, args)
def cargo_no_feature(version, manifest, args):
    cargo_raw(False, version, manifest, args)

def doc_upload(version, manifest, args):
    branch = os.environ.get('APPVEYOR_REPO_BRANCH') or os.environ['TRAVIS_BRANCH']
    repo = os.environ.get('APPVEYOR_REPO_NAME') or os.environ['TRAVIS_REPO_SLUG']
    if os.environ.get('APPVEYOR_PULL_REQUEST_NUMBER'):
        pr = 'true'
    else:
        pr = os.environ.get('TRAVIS_PULL_REQUEST', 'false')

    lib_name = manifest.lib_name()
    if lib_name is None:
        sys.stderr.write('error: uploading docs for package with no library')
        exit(1)

    if branch == args.branch and pr == 'false':
        # only load the token when we're sure we're uploading (travis
        # won't decrypt secret keys for PRs, so loading this with the
        # other vars causes problems with tests)
        token = os.environ['GH_TOKEN']

        print('uploading docs...')
        sys.stdout.flush()
        with open('target/doc/index.html', 'w') as f:
            f.write('<meta http-equiv=refresh content=0;url=%s/index.html>' % lib_name)

        run('git', 'clone', 'https://github.com/davisp/ghp-import')
        run(sys.executable, './ghp-import/ghp-import', '-n', 'target/doc')
        run_filter(token, 'git', 'push', '-fq', 'https://%s@github.com/%s.git' % (token, repo), 'gh-pages')

def build_kcov(use_sudo, verify):
    deps = ''
    if use_sudo:
        deps = 'sudo apt-get install libcurl4-openssl-dev libelf-dev libdw-dev cmake'
        if verify:
            deps += ' binutils-dev'
    init = deps + '''
    wget https://github.com/SimonKagstrom/kcov/archive/master.zip
    unzip master.zip
    mv kcov-master kcov
    mkdir kcov/build
    '''
    for line in init.split('\n'):
        line = line.strip()
        if line:
            print('Running: %s' % line)
            run(*line.split())
    current = os.getcwd()
    os.chdir('kcov/build')
    build = '''
    cmake ..
    make
    '''
    for line in build.split('\n'):
        line = line.strip()
        if line:
            print('Running: %s' % line)
            run(*line.split())
    os.chdir(current)
    return os.path.join(current, 'kcov/build/src/kcov')

def raw_coverage(use_sudo, verify, test_args, merge_msg, kcov_merge_args, kcov_merge_dir,
                 exclude_pattern, extra_kcov_args):
    kcov = build_kcov(use_sudo, verify)

    test_binaries = []
    # look through the output of `cargo test` to find the test
    # binaries.
    # FIXME: the information cargo feeds us is
    # inconsistent/inaccurate, so using the output of read-manifest is
    # far too much trouble.

    output = run_output('cargo', 'test', *test_args)
    running = re.compile('^     Running target/debug/(.*)$', re.M)
    for line in running.finditer(output):
        test_binaries.append(line.group(1))

    # record coverage for each binary
    for binary in test_binaries:
        print('Recording %s' % binary)
        kcov_args = [kcov] + extra_kcov_args
        if verify:
            kcov_args += ['--verify']
        exclude_pattern_arg = '--exclude-pattern=/.cargo'
        if exclude_pattern:
            exclude_pattern_arg += ',{}'.format(exclude_pattern)
        kcov_args += [exclude_pattern_arg, 'target/kcov-' + binary,
                      'target/debug/' + binary]
        print('Running: {}'.format(' '.join(kcov_args)))
        run(*kcov_args)
    # merge all the coverages and upload in one go
    print(merge_msg)
    kcov_args = [kcov, '--merge'] + kcov_merge_args + [kcov_merge_dir]
    kcov_args += ('target/kcov-' + b for b in test_binaries)
    run(*kcov_args)

def coverage(version, manifest, args):
    cargo_args = args.cargo_args
    add_features(cargo_args, version)

    kcov_merge_dir = args.merge_into
    raw_coverage(not args.no_sudo, args.verify, cargo_args, 'Merging coverage', [], kcov_merge_dir,
                 args.exclude_pattern, args.kcov_args)

def coveralls(version, manifest, args):
    job_id = os.environ['TRAVIS_JOB_ID']

    cargo_args = args.cargo_args
    add_features(cargo_args, version)

    raw_coverage(not args.no_sudo, args.verify, cargo_args, 'Uploading coverage',
                 ['--coveralls-id=' + job_id], 'target/kcov', args.exclude_pattern, args.kcov_args)

# user interface
class ScInfo(object):
    def __init__(self, func, description, arguments, help_=None, is_cargo=False):
        self.func = func
        self.help = help_
        self.description = description
        self.arguments = arguments
        self.is_cargo = is_cargo
        for _name, options in arguments:
            assert isinstance(options, dict)

EXCLUDE_PATTERN = (['--exclude-pattern'], {
    'metavar': 'PATTERN',
    'help': 'pass additional exclusionary patterns to kcov'
})

KCOV_OPTIONS = (['--kcov-options'], {
    'action': 'append',
    'metavar': 'OPTION',
    'default': [],
    'help': 'pass additional arguments to kcov, apart from --verify '
    'and --exclude-pattern, when recording coverage. Specify multiple '
    'times for multiple arguments'
})

NO_SUDO = (['--no-sudo'], {
    'action': 'store_true',
    'default': False,
    'help': 'don\'t use `sudo` to install kcov\'s deps. Requires that '
    'libcurl4-openssl-dev, libelf-dev and libdw-dev are installed (e.g. via '
    '`addons: apt: packages:`)'
})
VERIFY = (['--verify'], {
    'action': 'store_true',
    'default': False,
    'help': 'pass `--verify` to kcov, to avoid some crashes. See '
    '<https://github.com/huonw/travis-cargo/issues/12>. This requires '
    'installing the `binutils-dev` package.'
})

SC_INFO = {
    'doc-upload': ScInfo(func = doc_upload,
                         description = 'Use ghp-import to upload cargo-rendered '
                         'docs to Github Pages, from the master branch.',
                         help_ = 'upload documentation to Github pages.',
                         arguments = [(['--branch'], {
                             'default': 'master',
                             'help': 'upload docs when on this branch, defaults to master',
                         })]),
    'coveralls': ScInfo(func = coveralls,
                        description = 'Record coverage of `cargo test` and upload to '
                        'coveralls.io with kcov, this runs all binaries that `cargo test` runs '
                        'but not doc tests. Merged kcov results can be accessed in `target/kcov`.',
                        help_ = 'record and upload code coverage to coveralls.io',
                        arguments = [(['cargo_args'], {
                            'metavar': 'ARGS',
                            'nargs': '*',
                            'help': 'arguments to pass to `cargo test`'
                        }),
                                     EXCLUDE_PATTERN,
                                     KCOV_OPTIONS,
                                     NO_SUDO,
                                     VERIFY]),
    'coverage': ScInfo(func = coverage,
                       description = 'Record coverage of `cargo test`, this runs all '
                       'binaries that `cargo test` runs but not doc tests. The results '
                       'of all tests are merged into a single directory.',
                       help_ = 'record code coverage',
                       arguments = [(['cargo_args'], {
                           'metavar': 'ARGS',
                           'nargs': '*',
                           'help': 'arguments to pass to `cargo test`'
                       }),
                                    (['-m', '--merge-into'], {
                                        'metavar': 'DIR',
                                        'default': 'target/kcov',
                                        'help': 'the directory to put the final merged kcov '
                                        'result into (default `target/kcov`)'
                                    }),
                                    EXCLUDE_PATTERN,
                                    KCOV_OPTIONS,
                                    NO_SUDO,
                                    VERIFY])
}

def cargo_sc(name, features):
    return ScInfo(func = cargo_feature if features else cargo_no_feature,
                  description = 'Run `cargo %s`' % name,
                  is_cargo = True,
                  arguments = [(['cargo_args'], {
                      'metavar': 'ARGS',
                      'nargs': '*',
                      'help': 'arguments to pass to `cargo %s`' % name
                  })])

NO_FEATURE_CARGO = [
    'clean', 'fetch', 'generate-lockfile', 'git-checkout', 'help', 'locate-project',
    'login', 'new', 'owner', 'package', 'pkgid', 'publish', 'read-manifest', 'search',
    'update', 'verify-project', 'version', 'yank'
]
FEATURE_CARGO = [
    'build', 'bench', 'test', 'doc', 'run', 'rustc', 'rustdoc',
]
SC_INFO.update((c, cargo_sc(c, False)) for c in NO_FEATURE_CARGO)
SC_INFO.update((c, cargo_sc(c, True)) for c in FEATURE_CARGO)

def main():
    parser = argparse.ArgumentParser(description = '''
Manages interactions between Travis and Cargo and common tooling tasks.
''')

    parser.add_argument('-q','--quiet', action='store_true', default=False,
                        help='don\'t pass --verbose to cargo subcommands')
    parser.add_argument('--only', metavar='VERSION',
                        help='only run the given command if the specified version matches '
                        '`TRAVIS_RUST_VERSION`')
    parser.add_argument('--skip', metavar='VERSION',
                        help='only run the given command if the specified version does not match '
                        '`TRAVIS_RUST_VERSION`')

    sb = parser.add_subparsers(metavar = '{coverage,coveralls,doc-upload,...}',
                               description = '''

travis-cargo supports all cargo subcommands, and selected others
(listed below).

Cargo subcommands have `--verbose` added to their invocation by
default, and, when running with a nightly compiler, `--features
unstable` (or `--features $TRAVIS_CARGO_NIGHTLY_FEATURE` if that
environment variable is defined) if `--features` is a valid argument.

''',
                               dest = 'subcommand')
    for _, name, sc in sorted((sc.is_cargo, n, sc) for n, sc in SC_INFO.items()):
        extra_args = {'help': sc.help} if sc.help is not None else {}
        sub_parser = sb.add_parser(name,
                                   description = sc.description,
                                   **extra_args)
        sub_parser.set_defaults(func = sc.func)
        for name, options in sc.arguments:
            sub_parser.add_argument(*name, **options)

    args = parser.parse_args()

    version = os.environ.get('TRAVIS_RUST_VERSION', None)
    if version is None:
        # fill in the version based on the compiler's version output.
        output = run_output('rustc', '-V')
        phrases = ['nightly', ('dev', 'nightly'), 'beta']
        for phrase in phrases:
            if isinstance(phrase, tuple):
                alias = phrase[1]
                phrase = phrase[0]
            else:
                alias = phrase
            if phrase in output:
                version = alias
                break

    if args.only and args.only != version:
        return

    if args.skip and args.skip == version:
      return

    manifest = Manifest(os.getcwd(), version)
    args.func(version, manifest, args)
<|endoftext|>


<commit_before>from __future__ import division, print_function
import numpy as np
from scipy import optimize as sciopt
import config as ttconf
from seq_utils import alphabets, profile_maps

class GTR(object):
    """
    Defines General-Time-Reversible model of character evolution.
    """
    def __init__(self, alphabet='nuc', prof_map=None):
        """
        Initialize empty evolutionary model.
        Args:
         - alphabet (numpy.array): alphabet of the sequence.
        """
        if type(alphabet)==str:
            if (alphabet not in alphabets):
                raise AttributeError("Unknown alphabet type specified")
            else:
                self.alphabet = alphabets[alphabet]
                self.profile_map = profile_maps[alphabet]
        else:
            self.alphabet = alphabet
            if prof_map is None:
                self.profile_map = {s:x for s,x in zip(self.alphabet, np.eye(len(self.alphabet)))}
            else:
                self.profile_map = prof_map
        n_states = len(self.alphabet)

        # general rate matrix
        self.W = np.zeros((n_states, n_states))
        # stationary states of the characters
        self.Pi = np.zeros((n_states, n_states))

        # mutation rate, scaling factor
        self.mu = 1.0
        # eigendecomposition of the GTR matrix
        # Pi.dot(W) = v.dot(eigenmat).dot(v_inv)
        self.v = np.zeros((n_states, n_states))
        self.v_inv = np.zeros((n_states, n_states))
        self.eigenmat = np.zeros(n_states)

        # distance matrix (needed for topology optimization and for NJ)
        self.dm = None

    def __str__(self):
        eq_freq_str = "Equilibrium frequencies (pi_i):\n"
        for a,p in zip(self.alphabet, np.diagonal(self.Pi)):
            eq_freq_str+=str(a)+': '+str(np.round(p,4))+'\n'

        W_str = "\nSymmetrized rates from j->i (W_ij):\n"
        W_str+='\t'+'\t'.join(map(str, self.alphabet))+'\n'
        for a,Wi in zip(self.alphabet, self.W):
            W_str+=str(a)+'\t'+'\t'.join([str(np.round(max(0,p),4)) for p in Wi])+'\n'

        Q_str = "\nActual rates from j->i (Q_ij):\n"
        Q_str+='\t'+'\t'.join(map(str, self.alphabet))+'\n'
        for a,Qi in zip(self.alphabet, self.Pi.dot(self.W)):
            Q_str+=str(a)+'\t'+'\t'.join([str(np.round(max(0,p),4)) for p in Qi])+'\n'

        return eq_freq_str + W_str + Q_str

    @classmethod
    def custom(cls, mu=1.0, pi=None, W=None, **kwargs):
        """
        Create a GTR model by specifying the matrix explicitly

        Args:
         - mu (float): mutation rate
         - W (nxn matrix): mutation matrix
         - pi (n vector): equilibrium frequencies

        KWargs:
         - alphabet(str): specify alphabet when applicable. If the alphabet specification
         is requred, but no alphabet specified, the nucleotide will be used as default.
        """
        gtr = cls(**kwargs)
        n = gtr.alphabet.shape[0]

        gtr.mu = mu
        if pi is not None and len(pi)==n:
            Pi = pi
        else:
            if pi is not None and len(pi)!=n:
                print("length of equilibrium frequency vector does not match alphabet length"
                      "Ignoring input equilibrium frequencies")
            Pi = np.ones(size=(n))
        Pi /= Pi.sum()
        gtr.Pi = np.diagflat(Pi)

        if W is None or W.shape!=(n,n):
            if W.shape!=(n,n):
                print("Mutation matrix size does not match alphabet size"
                      "Ignoring input mutation matrix")
            # flow matrix
            gtr.W = np.ones((a,a))
            np.fill_diagonal(gtr.W, - ((gtr.W).sum(axis=0) - 1))

        gtr.W = 0.5*(W+W.T)
        gtr._check_fix_Q()
        gtr._eig()
        return gtr

    @classmethod
    def standard(cls, model='Jukes-Cantor', **kwargs):
        """
        Create one of the standard GTR models.

        Args:

         - model (str): type of the model. Currently supported models are:
         Jukes-Cantor, random.

        KWargs:

         - alphabet(str): specify alphabet when applicable. If the alphabet specification
         is requred, but no alphabet specified, the nucleotide will be used as default.

         - mu(double): general mutation rate. **NOTE** that the mutation rate is the
         only object which sets the time-scale to the GTR model.
         In other words, the unit of branch length and the unit of time are
         connected through this variable. By default set to 1.
        """
        if 'alphabet' in kwargs and kwargs['alphabet'] in alphabets.keys():
            alphabet = kwargs['alphabet']
        else:
            print ("No alphabet specified. Using default nucleotide.")
            alphabet = 'nuc'
        if 'mu' in kwargs:
            mu = kwargs['mu']
        else:
            mu = 1.0

        if model=='Jukes-Cantor':

            gtr = cls('nuc')
            gtr.mu = mu
            a = gtr.alphabet.shape[0]

            # flow matrix
            gtr.W = np.ones((a,a))
            np.fill_diagonal(gtr.W, - ((gtr.W).sum(axis=0) - 1))

            # equilibrium concentrations matrix
            gtr.Pi = np.zeros(gtr.W.shape)
            np.fill_diagonal(gtr.Pi, 1.0/a)

            gtr._check_fix_Q() # make sure the main diagonal is correct
            gtr._eig() # eigendecompose the rate matrix
            return gtr

        elif model=='random':
            gtr = cls(alphabet)
            a = gtr.alphabet.shape[0]

            gtr.mu = mu

            Pi = 1.0*np.random.randint(0,100,size=(a))
            Pi /= Pi.sum()
            gtr.Pi = np.diagflat(Pi)

            W = 1.0*np.random.randint(0,100,size=(a,a)) # with gaps
            gtr.W = W+W.T

            gtr._check_fix_Q()
            gtr._eig()
            return gtr
        else:
            raise NotImplementedError("The specified evolutionary model is unsupported!")

    @classmethod
    def infer(cls, nij, Ti, root_state, pc=5.0, **kwargs):
        """
        Infer a GTR model by specifying the number of transitions and time spent in each
        character. The basic equation that is being solved is
            n_ij = pi_i W_ij T_j
        where n_ij are the transitions, pi_i are the equilibrium state frequencies,
        W_ij is the "mutation attempt matrix", while T_i is the time on the tree
        spent in character state i. To regularize the process, we add pseudocounts and
        also need to account for the fact that the root of the tree is in a particular
        state. the modified equation is
            n_ij + pc = pi_i W_ij (T_j+pc+root_state)

        Args:
         - nij (nxn matrix): the number of times a change in character state is observed
            between state i and j
         - Ti (n vector): the time spent in each character state
         - root_state( n vector): the number of characters in state i in the sequence
            of the root node.
         - pc (float): pseudocounts, this determines the lower cutoff on the rate when
            no mutation are observed
        KWargs:
         - alphabet(str): specify alphabet when applicable. If the alphabet specification
         is requred, but no alphabet specified, the nucleotide will be used as default.
        """
        gtr = cls(**kwargs)
        dp = 1e-5
        Nit = 20
        pc_mat = pc*np.ones_like(nij)
        pc_mat -= np.diag(np.diag(pc_mat))
        from scipy import linalg as LA
        count = 0
        pi_old = np.zeros_like(Ti)
        pi = np.ones_like(Ti)
        pi/=pi.sum()
        W_ij = np.ones_like(nij)
        mu = nij.sum()/Ti.sum()
        while LA.norm(pi_old-pi) > dp and count < Nit:
            print('GTR inference iteration ',count,'change:',LA.norm(pi_old-pi))
            count += 1
            pi_old = np.copy(pi)
            W_ij = (nij+nij.T+2*pc_mat)/mu/(np.outer(pi,Ti) + np.outer(Ti,pi)
                                                    + ttconf.TINY_NUMBER + 2*pc_mat)
            W_ij = W_ij/np.sum(W_ij)
            pi = (np.sum(nij+pc_mat,axis=1)+root_state)/(mu*np.dot(W_ij,Ti)+root_state.sum()+np.sum(pc_mat, axis=1))
            mu = nij.sum()/(ttconf.TINY_NUMBER + np.sum(pi * (W_ij.dot(Ti))))

        if count >= Nit:
            print ('WARNING: maximum number of iterations has been reached in GTR inference')
            np.min(pi.sum(axis=0)), np.max(pi.sum(axis=0))
            if LA.norm(pi_old-pi) > dp:
                print ('    the iterative scheme has not converged')
            elif np.abs(1-np.max(pi.sum(axis=0))) > dp:
                print ('    the iterative scheme has converged, but proper normalization was not reached')
        gtr.W = W_ij
        gtr.Pi = np.diag(pi)
        gtr._check_fix_Q()
        gtr._eig()
        return gtr

    def _check_fix_Q(self):
        """
        Check the main diagonal of Q and fix it in case it does not corresond
        the definition of the rate matrix. Should be run every time when creating
        custom GTR model.
        """
        # fix Q
        self.Pi /= self.Pi.sum() # correct the Pi manually
        break_degen = np.random.random(size=self.W.shape)*0.0001
        self.W += break_degen + break_degen.T
        # fix W
        np.fill_diagonal(self.W, 0)
        Wdiag = -((self.W.T*np.diagonal(self.Pi)).T).sum(axis=0)/ \
                np.diagonal(self.Pi)
        np.fill_diagonal(self.W, Wdiag)
        Q1 = self.Pi.dot(self.W)
        self.W /= -np.sum(np.diagonal(Q1*self.Pi))
        self.mu=1.0
        Q1 = self.Pi.dot(self.W)
        if (Q1.sum(axis=0) < 1e-10).sum() <  self.alphabet.shape[0]: # fix failed
            import ipdb; ipdb.set_trace()
            raise ArithmeticError("Cannot fix the diagonal of the GTR rate matrix.")

    def _eig(self):
        """
        Perform eigendecompositon of the rate matrix and stores the left- and right-
        matrices to convert the sequence profiles to the GTR matrix eigenspace
        and hence to speed-up the computations.
        """
        # eigendecomposition of the rate matrix
        eigvals, eigvecs = np.linalg.eig(self.Pi.dot(self.W))
        self.v = eigvecs
        self.v_inv = np.linalg.inv(self.v)
        self.eigenmat = eigvals
        return

    def prob_t(self, profile_p, profile_ch, t, rotated=False, return_log=False):
        """
        Compute the probability of the two profiles to be separated by the time t.
        Args:
         - profile_p(np.array): parent profile of shape (L, a), where
         L - length of the sequence, a - alpphabet size.

         - profile_ch(np.array): child profile of shape (L, a), where
         L - length of the sequence, a - alpphabet size.

         - t (double): time (branch len), separating the profiles.

         - rotated (bool, default False): if True, assume that the supplied
         profiles are already rotated.

         - return_log(bool, default False): whether return log-probability.

        Returns:
         - prob(np.array): resulting probability.
        """

        if t < 0:
            if return_log:
                return -BIG_NUMBER
            else:
                return 0.0

        L = profile_p.shape[0]
        if L != profile_ch.shape[0]:
            raise ValueError("Sequence lengths do not match!")
        eLambdaT = self._exp_lt(t)
        if not rotated: # we need to rotate first
            p1 = profile_p.dot(self.v) # (L x a).dot(a x a) = (L x a) - prof in eigenspace
            p2 = (self.v_inv.dot(profile_ch.T)).T # (L x a).dot(a x a) = (L x a) - prof in eigenspace
        else:
            p1 = profile_p
            p2 = profile_ch
            #prob = (profile_p*eLambdaT*profile_ch).sum(1) # sum over the alphabet

        prob = (p1*eLambdaT*p2).sum(1) # sum_i (p1_i * exp(l_i*t) * p_2_i) result = vector lenght L
        prob[prob<0] = 0.0 # avoid rounding instability

        if return_log:
            prob = (np.log(prob + ttconf.TINY_NUMBER)).sum() # sum all sites
        else:
            prob = prob.prod() # prod of all sites
        return prob

    def optimal_t(self, profile_p, profile_ch, rotated=False, return_log=False):
        """
        Find the optimal distance between the two profiles
        """

        def _neg_prob(t, parent, child):
            """
            Probability to observe child given the the parent state, transition
            matrix and the time of evolution (branch length).

            Args:
             - t(double): branch length (time between sequences)
             - parent (numpy.array): parent sequence
             - child(numpy.array): child sequence
             - tm (GTR): model of evolution

            Returns:
             - prob(double): negative probability of the two given sequences
               to be separated by the time t.
            """
            return -1*self.prob_t (parent, child, t, rotated=False, return_log=True)

        opt = sciopt.minimize_scalar(_neg_prob,
                bounds=[0,ttconf.MAX_BRANCH_LENGTH],
                method='Bounded',
                args=(profile_p, profile_ch))

        new_len = opt["x"]

        if new_len > .9 * ttconf.MAX_BRANCH_LENGTH or opt["success"] != True:
            if verbose > 0:
                print ("Cannot optimize branch length, minimization failed.")
            import ipdb; ipdb.set_trace()
            return -1.0
        else:
            return  new_len

    def propagate_profile(self, profile, t, rotated=False, return_log=False):
        """
        Compute the probability of the sequence state (profile) at time (t+t0),
        given the sequence state (profile) at time t0.
        Args:
         - profile(numpy.array): sequence profile. Shape = (L, a),
         where L - sequence length, a - alphabet size.

         - t(double): time to propagate

         - rotated(bool default False): whether the supplied profile is in the
         GTR matrix eigenspace

         - return log (bool, default False): whether to return log-probability

        Returns:
         - res(np.array): profile of the sequence after time t.
         Shape = (L, a), where L - sequence length, a - alphabet size.
        """
        eLambdaT = self._exp_lt(t) # vector lenght = a

        if not rotated:
            # rotate
            p = self.v_inv.dot(profile.T).T
        else:
            p = profile

        res = (self.v.dot((eLambdaT * p).T)).T

        if not return_log:
            return res
        else:
            return np.log(res)

    def _exp_lt(self, t):
        """
        Returns:
         - exp_lt(numpy.array): array of values exp(lambda(i) * t),
         where (i) - alphabet index (the eigenvalue number).
        """
        return np.exp(self.mu * t * self.eigenmat)

if __name__ == "__main__":
     pass
<commit_msg>localized scipy.optimize import<commit_after>from __future__ import division, print_function
import numpy as np
import config as ttconf
from seq_utils import alphabets, profile_maps

class GTR(object):
    """
    Defines General-Time-Reversible model of character evolution.
    """
    def __init__(self, alphabet='nuc', prof_map=None):
        """
        Initialize empty evolutionary model.
        Args:
         - alphabet (numpy.array): alphabet of the sequence.
        """
        if type(alphabet)==str:
            if (alphabet not in alphabets):
                raise AttributeError("Unknown alphabet type specified")
            else:
                self.alphabet = alphabets[alphabet]
                self.profile_map = profile_maps[alphabet]
        else:
            self.alphabet = alphabet
            if prof_map is None:
                self.profile_map = {s:x for s,x in zip(self.alphabet, np.eye(len(self.alphabet)))}
            else:
                self.profile_map = prof_map
        n_states = len(self.alphabet)

        # general rate matrix
        self.W = np.zeros((n_states, n_states))
        # stationary states of the characters
        self.Pi = np.zeros((n_states, n_states))

        # mutation rate, scaling factor
        self.mu = 1.0
        # eigendecomposition of the GTR matrix
        # Pi.dot(W) = v.dot(eigenmat).dot(v_inv)
        self.v = np.zeros((n_states, n_states))
        self.v_inv = np.zeros((n_states, n_states))
        self.eigenmat = np.zeros(n_states)

        # distance matrix (needed for topology optimization and for NJ)
        self.dm = None

    def __str__(self):
        eq_freq_str = "Equilibrium frequencies (pi_i):\n"
        for a,p in zip(self.alphabet, np.diagonal(self.Pi)):
            eq_freq_str+=str(a)+': '+str(np.round(p,4))+'\n'

        W_str = "\nSymmetrized rates from j->i (W_ij):\n"
        W_str+='\t'+'\t'.join(map(str, self.alphabet))+'\n'
        for a,Wi in zip(self.alphabet, self.W):
            W_str+=str(a)+'\t'+'\t'.join([str(np.round(max(0,p),4)) for p in Wi])+'\n'

        Q_str = "\nActual rates from j->i (Q_ij):\n"
        Q_str+='\t'+'\t'.join(map(str, self.alphabet))+'\n'
        for a,Qi in zip(self.alphabet, self.Pi.dot(self.W)):
            Q_str+=str(a)+'\t'+'\t'.join([str(np.round(max(0,p),4)) for p in Qi])+'\n'

        return eq_freq_str + W_str + Q_str

    @classmethod
    def custom(cls, mu=1.0, pi=None, W=None, **kwargs):
        """
        Create a GTR model by specifying the matrix explicitly

        Args:
         - mu (float): mutation rate
         - W (nxn matrix): mutation matrix
         - pi (n vector): equilibrium frequencies

        KWargs:
         - alphabet(str): specify alphabet when applicable. If the alphabet specification
         is requred, but no alphabet specified, the nucleotide will be used as default.
        """
        gtr = cls(**kwargs)
        n = gtr.alphabet.shape[0]

        gtr.mu = mu
        if pi is not None and len(pi)==n:
            Pi = pi
        else:
            if pi is not None and len(pi)!=n:
                print("length of equilibrium frequency vector does not match alphabet length"
                      "Ignoring input equilibrium frequencies")
            Pi = np.ones(size=(n))
        Pi /= Pi.sum()
        gtr.Pi = np.diagflat(Pi)

        if W is None or W.shape!=(n,n):
            if W.shape!=(n,n):
                print("Mutation matrix size does not match alphabet size"
                      "Ignoring input mutation matrix")
            # flow matrix
            gtr.W = np.ones((a,a))
            np.fill_diagonal(gtr.W, - ((gtr.W).sum(axis=0) - 1))

        gtr.W = 0.5*(W+W.T)
        gtr._check_fix_Q()
        gtr._eig()
        return gtr

    @classmethod
    def standard(cls, model='Jukes-Cantor', **kwargs):
        """
        Create one of the standard GTR models.

        Args:

         - model (str): type of the model. Currently supported models are:
         Jukes-Cantor, random.

        KWargs:

         - alphabet(str): specify alphabet when applicable. If the alphabet specification
         is requred, but no alphabet specified, the nucleotide will be used as default.

         - mu(double): general mutation rate. **NOTE** that the mutation rate is the
         only object which sets the time-scale to the GTR model.
         In other words, the unit of branch length and the unit of time are
         connected through this variable. By default set to 1.
        """
        if 'alphabet' in kwargs and kwargs['alphabet'] in alphabets.keys():
            alphabet = kwargs['alphabet']
        else:
            print ("No alphabet specified. Using default nucleotide.")
            alphabet = 'nuc'
        if 'mu' in kwargs:
            mu = kwargs['mu']
        else:
            mu = 1.0

        if model=='Jukes-Cantor':

            gtr = cls('nuc')
            gtr.mu = mu
            a = gtr.alphabet.shape[0]

            # flow matrix
            gtr.W = np.ones((a,a))
            np.fill_diagonal(gtr.W, - ((gtr.W).sum(axis=0) - 1))

            # equilibrium concentrations matrix
            gtr.Pi = np.zeros(gtr.W.shape)
            np.fill_diagonal(gtr.Pi, 1.0/a)

            gtr._check_fix_Q() # make sure the main diagonal is correct
            gtr._eig() # eigendecompose the rate matrix
            return gtr

        elif model=='random':
            gtr = cls(alphabet)
            a = gtr.alphabet.shape[0]

            gtr.mu = mu

            Pi = 1.0*np.random.randint(0,100,size=(a))
            Pi /= Pi.sum()
            gtr.Pi = np.diagflat(Pi)

            W = 1.0*np.random.randint(0,100,size=(a,a)) # with gaps
            gtr.W = W+W.T

            gtr._check_fix_Q()
            gtr._eig()
            return gtr
        else:
            raise NotImplementedError("The specified evolutionary model is unsupported!")

    @classmethod
    def infer(cls, nij, Ti, root_state, pc=5.0, **kwargs):
        """
        Infer a GTR model by specifying the number of transitions and time spent in each
        character. The basic equation that is being solved is
            n_ij = pi_i W_ij T_j
        where n_ij are the transitions, pi_i are the equilibrium state frequencies,
        W_ij is the "mutation attempt matrix", while T_i is the time on the tree
        spent in character state i. To regularize the process, we add pseudocounts and
        also need to account for the fact that the root of the tree is in a particular
        state. the modified equation is
            n_ij + pc = pi_i W_ij (T_j+pc+root_state)

        Args:
         - nij (nxn matrix): the number of times a change in character state is observed
            between state i and j
         - Ti (n vector): the time spent in each character state
         - root_state( n vector): the number of characters in state i in the sequence
            of the root node.
         - pc (float): pseudocounts, this determines the lower cutoff on the rate when
            no mutation are observed
        KWargs:
         - alphabet(str): specify alphabet when applicable. If the alphabet specification
         is requred, but no alphabet specified, the nucleotide will be used as default.
        """
        gtr = cls(**kwargs)
        dp = 1e-5
        Nit = 20
        pc_mat = pc*np.ones_like(nij)
        pc_mat -= np.diag(np.diag(pc_mat))
        from scipy import linalg as LA
        count = 0
        pi_old = np.zeros_like(Ti)
        pi = np.ones_like(Ti)
        pi/=pi.sum()
        W_ij = np.ones_like(nij)
        mu = nij.sum()/Ti.sum()
        while LA.norm(pi_old-pi) > dp and count < Nit:
            print('GTR inference iteration ',count,'change:',LA.norm(pi_old-pi))
            count += 1
            pi_old = np.copy(pi)
            W_ij = (nij+nij.T+2*pc_mat)/mu/(np.outer(pi,Ti) + np.outer(Ti,pi)
                                                    + ttconf.TINY_NUMBER + 2*pc_mat)
            W_ij = W_ij/np.sum(W_ij)
            pi = (np.sum(nij+pc_mat,axis=1)+root_state)/(mu*np.dot(W_ij,Ti)+root_state.sum()+np.sum(pc_mat, axis=1))
            mu = nij.sum()/(ttconf.TINY_NUMBER + np.sum(pi * (W_ij.dot(Ti))))

        if count >= Nit:
            print ('WARNING: maximum number of iterations has been reached in GTR inference')
            np.min(pi.sum(axis=0)), np.max(pi.sum(axis=0))
            if LA.norm(pi_old-pi) > dp:
                print ('    the iterative scheme has not converged')
            elif np.abs(1-np.max(pi.sum(axis=0))) > dp:
                print ('    the iterative scheme has converged, but proper normalization was not reached')
        gtr.W = W_ij
        gtr.Pi = np.diag(pi)
        gtr._check_fix_Q()
        gtr._eig()
        return gtr

    def _check_fix_Q(self):
        """
        Check the main diagonal of Q and fix it in case it does not corresond
        the definition of the rate matrix. Should be run every time when creating
        custom GTR model.
        """
        # fix Q
        self.Pi /= self.Pi.sum() # correct the Pi manually
        break_degen = np.random.random(size=self.W.shape)*0.0001
        self.W += break_degen + break_degen.T
        # fix W
        np.fill_diagonal(self.W, 0)
        Wdiag = -((self.W.T*np.diagonal(self.Pi)).T).sum(axis=0)/ \
                np.diagonal(self.Pi)
        np.fill_diagonal(self.W, Wdiag)
        Q1 = self.Pi.dot(self.W)
        self.W /= -np.sum(np.diagonal(Q1*self.Pi))
        self.mu=1.0
        Q1 = self.Pi.dot(self.W)
        if (Q1.sum(axis=0) < 1e-10).sum() <  self.alphabet.shape[0]: # fix failed
            import ipdb; ipdb.set_trace()
            raise ArithmeticError("Cannot fix the diagonal of the GTR rate matrix.")

    def _eig(self):
        """
        Perform eigendecompositon of the rate matrix and stores the left- and right-
        matrices to convert the sequence profiles to the GTR matrix eigenspace
        and hence to speed-up the computations.
        """
        # eigendecomposition of the rate matrix
        eigvals, eigvecs = np.linalg.eig(self.Pi.dot(self.W))
        self.v = eigvecs
        self.v_inv = np.linalg.inv(self.v)
        self.eigenmat = eigvals
        return

    def prob_t(self, profile_p, profile_ch, t, rotated=False, return_log=False):
        """
        Compute the probability of the two profiles to be separated by the time t.
        Args:
         - profile_p(np.array): parent profile of shape (L, a), where
         L - length of the sequence, a - alpphabet size.

         - profile_ch(np.array): child profile of shape (L, a), where
         L - length of the sequence, a - alpphabet size.

         - t (double): time (branch len), separating the profiles.

         - rotated (bool, default False): if True, assume that the supplied
         profiles are already rotated.

         - return_log(bool, default False): whether return log-probability.

        Returns:
         - prob(np.array): resulting probability.
        """

        if t < 0:
            if return_log:
                return -BIG_NUMBER
            else:
                return 0.0

        L = profile_p.shape[0]
        if L != profile_ch.shape[0]:
            raise ValueError("Sequence lengths do not match!")
        eLambdaT = self._exp_lt(t)
        if not rotated: # we need to rotate first
            p1 = profile_p.dot(self.v) # (L x a).dot(a x a) = (L x a) - prof in eigenspace
            p2 = (self.v_inv.dot(profile_ch.T)).T # (L x a).dot(a x a) = (L x a) - prof in eigenspace
        else:
            p1 = profile_p
            p2 = profile_ch
            #prob = (profile_p*eLambdaT*profile_ch).sum(1) # sum over the alphabet

        prob = (p1*eLambdaT*p2).sum(1) # sum_i (p1_i * exp(l_i*t) * p_2_i) result = vector lenght L
        prob[prob<0] = 0.0 # avoid rounding instability

        if return_log:
            prob = (np.log(prob + ttconf.TINY_NUMBER)).sum() # sum all sites
        else:
            prob = prob.prod() # prod of all sites
        return prob

    def optimal_t(self, profile_p, profile_ch, rotated=False, return_log=False):
        """
        Find the optimal distance between the two profiles
        """

        def _neg_prob(t, parent, child):
            """
            Probability to observe child given the the parent state, transition
            matrix and the time of evolution (branch length).

            Args:
             - t(double): branch length (time between sequences)
             - parent (numpy.array): parent sequence
             - child(numpy.array): child sequence
             - tm (GTR): model of evolution

            Returns:
             - prob(double): negative probability of the two given sequences
               to be separated by the time t.
            """
            return -1*self.prob_t (parent, child, t, rotated=False, return_log=True)

        from scipy.optimize import minimize_scalar
        opt = minimize_scalar(_neg_prob,
                bounds=[0,ttconf.MAX_BRANCH_LENGTH],
                method='Bounded',
                args=(profile_p, profile_ch))

        new_len = opt["x"]

        if new_len > .9 * ttconf.MAX_BRANCH_LENGTH or opt["success"] != True:
            if verbose > 0:
                print ("Cannot optimize branch length, minimization failed.")
            import ipdb; ipdb.set_trace()
            return -1.0
        else:
            return  new_len

    def propagate_profile(self, profile, t, rotated=False, return_log=False):
        """
        Compute the probability of the sequence state (profile) at time (t+t0),
        given the sequence state (profile) at time t0.
        Args:
         - profile(numpy.array): sequence profile. Shape = (L, a),
         where L - sequence length, a - alphabet size.

         - t(double): time to propagate

         - rotated(bool default False): whether the supplied profile is in the
         GTR matrix eigenspace

         - return log (bool, default False): whether to return log-probability

        Returns:
         - res(np.array): profile of the sequence after time t.
         Shape = (L, a), where L - sequence length, a - alphabet size.
        """
        eLambdaT = self._exp_lt(t) # vector lenght = a

        if not rotated:
            # rotate
            p = self.v_inv.dot(profile.T).T
        else:
            p = profile

        res = (self.v.dot((eLambdaT * p).T)).T

        if not return_log:
            return res
        else:
            return np.log(res)

    def _exp_lt(self, t):
        """
        Returns:
         - exp_lt(numpy.array): array of values exp(lambda(i) * t),
         where (i) - alphabet index (the eigenvalue number).
        """
        return np.exp(self.mu * t * self.eigenmat)

if __name__ == "__main__":
     pass
<|endoftext|>