<commit_before># -*- coding:utf8 -*-
# File   : desc_a3c_atari.py
# Author : Jiayuan Mao
#          Honghua Dong
# Email  : maojiayuan@gmail.com
#          dhh19951@gmail,com
# Date   : 3/18/17
#
# This file is part of TensorArtist

"""This reproduction of A3C is based on ppwwyyxx's reproduction in his TensorPack framework.
Credit to : https://github.com/ppwwyyxx/tensorpack/tree/master/examples/A3C-Gym"""

import time
import queue
import threading
import functools
import collections
import numpy as np
import tensorflow as tf
import os

from tartist.core import get_env, get_logger, EnvBox
from tartist.core.utils.cache import cached_result
from tartist.core.utils.naming import get_dump_directory, get_data_directory
from tartist.data import flow
from tartist.data.rflow import QueryReqPipe
from tartist.nn import opr as O, optimizer, summary
from tartist import rl, random, image

logger = get_logger(__file__)

__envs__ = {
    'dir': {
        'root': get_dump_directory(__file__),
    },

    'a3c': {
        'env_name': 'Breakout-v0',
        'frame_history': 4,
        'limit_length': 40000,

        # gamma and acc_step in future_reward
        'gamma': 0.99,
        'acc_step': 5,

        'nr_players': 50,
        'nr_predictors': 2,

        'predictor': {
            'batch_size': 16,
            'outputs_name': ['value', 'policy_explore']
        },
        'inference': {
            'nr_players': 20,
            'max_stuck_repeat': 30
        },
        'demo': {
            'nr_plays': 5
        }
    },

    'dataset': {
        'input_shape': (84, 84)
    },

    'trainer': {
        'learning_rate': 0.001,

        'batch_size': 128,
        'epoch_size': 1000,
        'nr_epochs': 200,

        'gamma': 0.99,

        'env_flags': {
            'log_device_placement': False
        }
    },
    'demo': {
        'customized': True
    }
}

__trainer_cls__ = rl.train.A3CTrainer
__trainer_env_cls__ = rl.train.A3CTrainerEnv

def make_network(env):
    is_train = env.phase is env.Phase.TRAIN
    if is_train:
        slave_devices = env.slave_devices
        env.set_slave_devices([])

    with env.create_network() as net:
        h, w, c = get_input_shape()

        dpc = env.create_dpcontroller()
        with dpc.activate():
            def inputs():
                state = O.placeholder('state', shape=(None, h, w, c))
                return [state]

            def forward(x):
                _ = x / 255.0
                with O.argscope(O.conv2d, nonlin=O.relu):
                    _ = O.conv2d('conv0', _, 32, 5)
                    _ = O.max_pooling2d('pool0', _, 2)
                    _ = O.conv2d('conv1', _, 32, 5)
                    _ = O.max_pooling2d('pool1', _, 2)
                    _ = O.conv2d('conv2', _, 64, 4)
                    _ = O.max_pooling2d('pool2', _, 2)
                    _ = O.conv2d('conv3', _, 64, 3)

                dpc.add_output(_, name='feature')

            dpc.set_input_maker(inputs).set_forward_func(forward)

        _ = dpc.outputs['feature']
        _ = O.fc('fc0', _, 512, nonlin=O.p_relu)
        policy = O.fc('fc_policy', _, get_player_nr_actions())
        value = O.fc('fc_value', _, 1)

        expf = O.scalar('explore_factor', 1, trainable=False)
        policy_explore = O.softmax(policy * expf, name='policy_explore')

        policy = O.softmax(policy, name='policy')
        value = value.remove_axis(1, name='value')

        net.add_output(policy_explore, name='policy_explore')
        net.add_output(policy, name='policy')
        net.add_output(value, name='value')

        if env.phase is env.Phase.TRAIN:
            action = O.placeholder('action', shape=(None, ), dtype=tf.int64)
            future_reward = O.placeholder('future_reward', shape=(None, ))

            log_policy = O.log(policy + 1e-6)
            log_pi_a_given_s = (log_policy * tf.one_hot(action, get_player_nr_actions())).sum(axis=1)
            advantage = future_reward - O.zero_grad(value, name='advantage')
            policy_cost = (log_pi_a_given_s * advantage).mean(name='policy_cost')
            xentropy_cost = (-policy * log_policy).sum(axis=1).mean(name='xentropy_cost')
            value_loss = O.raw_l2_loss('raw_value_loss', future_reward, value).mean(name='value_loss')
            # value_loss = O.truediv(value_loss, future_reward.shape[0].astype('float32'), name='value_loss')
            entropy_beta = O.scalar('entropy_beta', 0.01, trainable=False)
            loss = tf.add_n([-policy_cost, -xentropy_cost * entropy_beta, value_loss], name='loss')

            net.set_loss(loss)

            for v in [policy_cost, xentropy_cost, value_loss,
                      value.mean(name='predict_value'), advantage.rms(name='rms_advantage'), loss]:
                summary.scalar(v)

    if is_train:
        env.set_slave_devices(slave_devices)

def make_player(is_train=True, dump_dir=None):
    def resize_state(s):
        return image.resize(s, get_env('dataset.input_shape'))

    p = rl.GymRLEnviron(get_env('a3c.env_name'), dump_dir=dump_dir)
    p = rl.MapStateProxyRLEnviron(p, resize_state)
    p = rl.GymHistoryProxyRLEnviron(p, get_env('a3c.frame_history'))

    p = rl.LimitLengthProxyRLEnviron(p, get_env('a3c.limit_length'))
    if is_train:
        p = rl.AutoRestartProxyRLEnviron(p)
    else:
        p = rl.GymPreventStuckProxyRLEnviron(p, get_env('a3c.inference.max_stuck_repeat'), 1)
    return p

def make_optimizer(env):
    lr = optimizer.base.make_optimizer_variable('learning_rate', get_env('trainer.learning_rate'))

    wrapper = optimizer.OptimizerWrapper()
    wrapper.set_base_optimizer(optimizer.base.AdamOptimizer(lr, epsilon=1e-3))
    wrapper.append_grad_modifier(optimizer.grad_modifier.LearningRateMultiplier([
        ('*/b', 2.0),
    ]))
    wrapper.append_grad_modifier(optimizer.grad_modifier.GlobalGradClipByAvgNorm(0.1))
    env.set_optimizer(wrapper)

def make_dataflow_train(env):
    batch_size = get_env('trainer.batch_size')

    df = flow.QueueDataFlow(env.data_queue)
    df = flow.BatchDataFlow(df, batch_size, sample_dict={
        'state': np.empty((batch_size, ) + get_input_shape(), dtype='float32'),
        'action': np.empty((batch_size, ), dtype='int32'),
        'future_reward': np.empty((batch_size, ), dtype='float32')
    })
    return df

@cached_result
def get_player_nr_actions():
    p = make_player()
    n = p.action_space.nr_actions
    del p
    return n

@cached_result
def get_input_shape():
    input_shape = get_env('dataset.input_shape')
    frame_history = get_env('a3c.frame_history')
    h, w, c = input_shape[0], input_shape[1], 3 * frame_history
    return h, w, c

def player_func(pid, requester):
    player = make_player()
    player.restart()
    state = player.current_state
    reward = 0
    is_over = False
    with requester.activate():
        while True:
            action = requester.query('data', (state, reward, is_over))
            reward, is_over = player.action(action)

            if len(player.stats['score']) > 0:
                score = player.stats['score'][-1]
                requester.query('stat', {'async/score': score}, do_recv=False)
                player.clear_stats()
            state = player.current_state

def inference_player_func(pid, requester):
    player = make_player(is_train=False)

    def get_action(o):
        action = requester.query('data', (o, ))
        return action

    with requester.activate():
        player.play_one_episode(get_action)
        score = player.stats['score'][-1]
        requester.query('stat', {'async/inference/score': score}, do_recv=False)

def _predictor_func(pid, router, task_queue, func, is_inference=False):
    batch_size = get_env('a3c.predictor.batch_size')
    batched_state = np.empty((batch_size, ) + get_input_shape(), dtype='float32')

    while True:
        callbacks = []
        nr_total = 0
        for i in range(batch_size):
            if i == 0 or not is_inference:
                identifier, inp, callback = task_queue.get()
            else:
                try:
                    identifier, inp, callback = task_queue.get_nowait()
                except queue.Empty:
                    break

            batched_state[i] = inp[0]
            callbacks.append(callback)
            nr_total += 1

        out = func(state=batched_state)
        for i in range(nr_total):
            policy = out['policy_explore'][i]
            if is_inference:
                # during inference, policy should be out['policy'][i]
                # but these two are equivalent under argmax operation
                # and we can only compile 'policy_explore' in output
                action = policy.argmax()
            else:
                action = random.choice(len(policy), p=policy)

            callbacks[i](action, out['value'][i])

predictor_func = functools.partial(_predictor_func, is_inference=False)
inference_predictor_func = functools.partial(_predictor_func, is_inference=True)

def make_a3c_configs(env):
    from common_a3c import on_data_func, on_stat_func, inference_on_data_func, inference_on_stat_func

    env.player_master.player_func = player_func
    env.player_master.predictor_func = predictor_func
    env.player_master.on_data_func = on_data_func
    env.player_master.on_stat_func = on_stat_func

    # currently we don't use multi-proc inference, so these settings are not used at all
    env.inference_player_master.player_func = inference_player_func
    env.inference_player_master.predictor_func = inference_predictor_func
    env.inference_player_master.on_data_func = inference_on_data_func
    env.inference_player_master.on_stat_func = inference_on_stat_func

    env.players_history = collections.defaultdict(list)
    env.players_history_lock = threading.Lock()

def main_train(trainer):
    from tartist.plugins.trainer_enhancer import summary
    summary.enable_summary_history(trainer, extra_summary_types={
        'async/score': 'async_scalar',
        'async/inference/score': 'async_scalar'
    })
    summary.enable_echo_summary_scalar(trainer, summary_spec={
        'async/score': ['avg', 'max'],
        'async/inference/score': ['avg', 'max']
    })

    from tartist.plugins.trainer_enhancer import progress
    progress.enable_epoch_progress(trainer)

    from tartist.plugins.trainer_enhancer import snapshot
    snapshot.enable_snapshot_saver(trainer)

    from tartist.core import register_event
    from common_a3c import main_inference_play_multithread

    def on_epoch_after(trainer):
        if trainer.epoch > 0 and trainer.epoch % 2 == 0:
            # main_inference_play(trainer, epoch=trainer.epoch)
            main_inference_play_multithread(trainer, make_player = make_player)

    # this one should run before monitor
    register_event(trainer, 'epoch:after', on_epoch_after, priority=5)

    trainer.train()

def main_demo(env, func):
    dump_dir = get_env('dir.demo', os.path.join(get_env('dir.root'), 'demo'))
    logger.info('demo dump dir: {}'.format(dump_dir))
    player = make_player(is_train=False, dump_dir=dump_dir)
    repeat_time = get_env('a3c.demo.nr_plays', 1)

    def get_action(inp, func=func):
        action = func(**{'state':[[inp]]})['policy'][0].argmax()
        return action

    for i in range(repeat_time):
        player.play_one_episode(get_action)
        logger.info('#{} play score={}'.format(i, player.stats['score'][-1]))

<commit_msg>[dev-dhh] fix another wrong email address<commit_after># -*- coding:utf8 -*-
# File   : desc_a3c_atari.py
# Author : Jiayuan Mao
#          Honghua Dong
# Email  : maojiayuan@gmail.com
#          dhh19951@gmail.com
# Date   : 3/18/17
#
# This file is part of TensorArtist

"""This reproduction of A3C is based on ppwwyyxx's reproduction in his TensorPack framework.
Credit to : https://github.com/ppwwyyxx/tensorpack/tree/master/examples/A3C-Gym"""

import time
import queue
import threading
import functools
import collections
import numpy as np
import tensorflow as tf
import os

from tartist.core import get_env, get_logger, EnvBox
from tartist.core.utils.cache import cached_result
from tartist.core.utils.naming import get_dump_directory, get_data_directory
from tartist.data import flow
from tartist.data.rflow import QueryReqPipe
from tartist.nn import opr as O, optimizer, summary
from tartist import rl, random, image

logger = get_logger(__file__)

__envs__ = {
    'dir': {
        'root': get_dump_directory(__file__),
    },

    'a3c': {
        'env_name': 'Breakout-v0',
        'frame_history': 4,
        'limit_length': 40000,

        # gamma and acc_step in future_reward
        'gamma': 0.99,
        'acc_step': 5,

        'nr_players': 50,
        'nr_predictors': 2,

        'predictor': {
            'batch_size': 16,
            'outputs_name': ['value', 'policy_explore']
        },
        'inference': {
            'nr_players': 20,
            'max_stuck_repeat': 30
        },
        'demo': {
            'nr_plays': 5
        }
    },

    'dataset': {
        'input_shape': (84, 84)
    },

    'trainer': {
        'learning_rate': 0.001,

        'batch_size': 128,
        'epoch_size': 1000,
        'nr_epochs': 200,

        'gamma': 0.99,

        'env_flags': {
            'log_device_placement': False
        }
    },
    'demo': {
        'customized': True
    }
}

__trainer_cls__ = rl.train.A3CTrainer
__trainer_env_cls__ = rl.train.A3CTrainerEnv

def make_network(env):
    is_train = env.phase is env.Phase.TRAIN
    if is_train:
        slave_devices = env.slave_devices
        env.set_slave_devices([])

    with env.create_network() as net:
        h, w, c = get_input_shape()

        dpc = env.create_dpcontroller()
        with dpc.activate():
            def inputs():
                state = O.placeholder('state', shape=(None, h, w, c))
                return [state]

            def forward(x):
                _ = x / 255.0
                with O.argscope(O.conv2d, nonlin=O.relu):
                    _ = O.conv2d('conv0', _, 32, 5)
                    _ = O.max_pooling2d('pool0', _, 2)
                    _ = O.conv2d('conv1', _, 32, 5)
                    _ = O.max_pooling2d('pool1', _, 2)
                    _ = O.conv2d('conv2', _, 64, 4)
                    _ = O.max_pooling2d('pool2', _, 2)
                    _ = O.conv2d('conv3', _, 64, 3)

                dpc.add_output(_, name='feature')

            dpc.set_input_maker(inputs).set_forward_func(forward)

        _ = dpc.outputs['feature']
        _ = O.fc('fc0', _, 512, nonlin=O.p_relu)
        policy = O.fc('fc_policy', _, get_player_nr_actions())
        value = O.fc('fc_value', _, 1)

        expf = O.scalar('explore_factor', 1, trainable=False)
        policy_explore = O.softmax(policy * expf, name='policy_explore')

        policy = O.softmax(policy, name='policy')
        value = value.remove_axis(1, name='value')

        net.add_output(policy_explore, name='policy_explore')
        net.add_output(policy, name='policy')
        net.add_output(value, name='value')

        if env.phase is env.Phase.TRAIN:
            action = O.placeholder('action', shape=(None, ), dtype=tf.int64)
            future_reward = O.placeholder('future_reward', shape=(None, ))

            log_policy = O.log(policy + 1e-6)
            log_pi_a_given_s = (log_policy * tf.one_hot(action, get_player_nr_actions())).sum(axis=1)
            advantage = future_reward - O.zero_grad(value, name='advantage')
            policy_cost = (log_pi_a_given_s * advantage).mean(name='policy_cost')
            xentropy_cost = (-policy * log_policy).sum(axis=1).mean(name='xentropy_cost')
            value_loss = O.raw_l2_loss('raw_value_loss', future_reward, value).mean(name='value_loss')
            # value_loss = O.truediv(value_loss, future_reward.shape[0].astype('float32'), name='value_loss')
            entropy_beta = O.scalar('entropy_beta', 0.01, trainable=False)
            loss = tf.add_n([-policy_cost, -xentropy_cost * entropy_beta, value_loss], name='loss')

            net.set_loss(loss)

            for v in [policy_cost, xentropy_cost, value_loss,
                      value.mean(name='predict_value'), advantage.rms(name='rms_advantage'), loss]:
                summary.scalar(v)

    if is_train:
        env.set_slave_devices(slave_devices)

def make_player(is_train=True, dump_dir=None):
    def resize_state(s):
        return image.resize(s, get_env('dataset.input_shape'))

    p = rl.GymRLEnviron(get_env('a3c.env_name'), dump_dir=dump_dir)
    p = rl.MapStateProxyRLEnviron(p, resize_state)
    p = rl.GymHistoryProxyRLEnviron(p, get_env('a3c.frame_history'))

    p = rl.LimitLengthProxyRLEnviron(p, get_env('a3c.limit_length'))
    if is_train:
        p = rl.AutoRestartProxyRLEnviron(p)
    else:
        p = rl.GymPreventStuckProxyRLEnviron(p, get_env('a3c.inference.max_stuck_repeat'), 1)
    return p

def make_optimizer(env):
    lr = optimizer.base.make_optimizer_variable('learning_rate', get_env('trainer.learning_rate'))

    wrapper = optimizer.OptimizerWrapper()
    wrapper.set_base_optimizer(optimizer.base.AdamOptimizer(lr, epsilon=1e-3))
    wrapper.append_grad_modifier(optimizer.grad_modifier.LearningRateMultiplier([
        ('*/b', 2.0),
    ]))
    wrapper.append_grad_modifier(optimizer.grad_modifier.GlobalGradClipByAvgNorm(0.1))
    env.set_optimizer(wrapper)

def make_dataflow_train(env):
    batch_size = get_env('trainer.batch_size')

    df = flow.QueueDataFlow(env.data_queue)
    df = flow.BatchDataFlow(df, batch_size, sample_dict={
        'state': np.empty((batch_size, ) + get_input_shape(), dtype='float32'),
        'action': np.empty((batch_size, ), dtype='int32'),
        'future_reward': np.empty((batch_size, ), dtype='float32')
    })
    return df

@cached_result
def get_player_nr_actions():
    p = make_player()
    n = p.action_space.nr_actions
    del p
    return n

@cached_result
def get_input_shape():
    input_shape = get_env('dataset.input_shape')
    frame_history = get_env('a3c.frame_history')
    h, w, c = input_shape[0], input_shape[1], 3 * frame_history
    return h, w, c

def player_func(pid, requester):
    player = make_player()
    player.restart()
    state = player.current_state
    reward = 0
    is_over = False
    with requester.activate():
        while True:
            action = requester.query('data', (state, reward, is_over))
            reward, is_over = player.action(action)

            if len(player.stats['score']) > 0:
                score = player.stats['score'][-1]
                requester.query('stat', {'async/score': score}, do_recv=False)
                player.clear_stats()
            state = player.current_state

def inference_player_func(pid, requester):
    player = make_player(is_train=False)

    def get_action(o):
        action = requester.query('data', (o, ))
        return action

    with requester.activate():
        player.play_one_episode(get_action)
        score = player.stats['score'][-1]
        requester.query('stat', {'async/inference/score': score}, do_recv=False)

def _predictor_func(pid, router, task_queue, func, is_inference=False):
    batch_size = get_env('a3c.predictor.batch_size')
    batched_state = np.empty((batch_size, ) + get_input_shape(), dtype='float32')

    while True:
        callbacks = []
        nr_total = 0
        for i in range(batch_size):
            if i == 0 or not is_inference:
                identifier, inp, callback = task_queue.get()
            else:
                try:
                    identifier, inp, callback = task_queue.get_nowait()
                except queue.Empty:
                    break

            batched_state[i] = inp[0]
            callbacks.append(callback)
            nr_total += 1

        out = func(state=batched_state)
        for i in range(nr_total):
            policy = out['policy_explore'][i]
            if is_inference:
                # during inference, policy should be out['policy'][i]
                # but these two are equivalent under argmax operation
                # and we can only compile 'policy_explore' in output
                action = policy.argmax()
            else:
                action = random.choice(len(policy), p=policy)

            callbacks[i](action, out['value'][i])

predictor_func = functools.partial(_predictor_func, is_inference=False)
inference_predictor_func = functools.partial(_predictor_func, is_inference=True)

def make_a3c_configs(env):
    from common_a3c import on_data_func, on_stat_func, inference_on_data_func, inference_on_stat_func

    env.player_master.player_func = player_func
    env.player_master.predictor_func = predictor_func
    env.player_master.on_data_func = on_data_func
    env.player_master.on_stat_func = on_stat_func

    # currently we don't use multi-proc inference, so these settings are not used at all
    env.inference_player_master.player_func = inference_player_func
    env.inference_player_master.predictor_func = inference_predictor_func
    env.inference_player_master.on_data_func = inference_on_data_func
    env.inference_player_master.on_stat_func = inference_on_stat_func

    env.players_history = collections.defaultdict(list)
    env.players_history_lock = threading.Lock()

def main_train(trainer):
    from tartist.plugins.trainer_enhancer import summary
    summary.enable_summary_history(trainer, extra_summary_types={
        'async/score': 'async_scalar',
        'async/inference/score': 'async_scalar'
    })
    summary.enable_echo_summary_scalar(trainer, summary_spec={
        'async/score': ['avg', 'max'],
        'async/inference/score': ['avg', 'max']
    })

    from tartist.plugins.trainer_enhancer import progress
    progress.enable_epoch_progress(trainer)

    from tartist.plugins.trainer_enhancer import snapshot
    snapshot.enable_snapshot_saver(trainer)

    from tartist.core import register_event
    from common_a3c import main_inference_play_multithread

    def on_epoch_after(trainer):
        if trainer.epoch > 0 and trainer.epoch % 2 == 0:
            # main_inference_play(trainer, epoch=trainer.epoch)
            main_inference_play_multithread(trainer, make_player = make_player)

    # this one should run before monitor
    register_event(trainer, 'epoch:after', on_epoch_after, priority=5)

    trainer.train()

def main_demo(env, func):
    dump_dir = get_env('dir.demo', os.path.join(get_env('dir.root'), 'demo'))
    logger.info('demo dump dir: {}'.format(dump_dir))
    player = make_player(is_train=False, dump_dir=dump_dir)
    repeat_time = get_env('a3c.demo.nr_plays', 1)

    def get_action(inp, func=func):
        action = func(**{'state':[[inp]]})['policy'][0].argmax()
        return action

    for i in range(repeat_time):
        player.play_one_episode(get_action)
        logger.info('#{} play score={}'.format(i, player.stats['score'][-1]))

<|endoftext|>


<commit_before>import unittest

from sublist import check_lists, SUBLIST, SUPERLIST, EQUAL, UNEQUAL

class SublistTest(unittest.TestCase):
    def test_unique_return_vals(self):
        self.assertEqual(len(set([SUBLIST, SUPERLIST, EQUAL, UNEQUAL])), 4)

    def test_empty_lists(self):
        self.assertEqual(check_lists([], []), EQUAL)

    def test_empty_list_within(self):
        self.assertEqual(check_lists([], [1, 2, 3]), SUBLIST)

    def test_within_empty_list(self):
        self.assertEqual(check_lists([1], []), SUPERLIST)

    def test_equal_lists(self):
        l1 = [0, 1, 2]
        l2 = [0, 1, 2]
        self.assertEqual(check_lists(l1, l2), EQUAL)

    def test_different_lists(self):
        l1 = list(range(1000000))
        l2 = list(range(1, 1000001))
        self.assertEqual(check_lists(l1, l2), UNEQUAL)

    def test_false_start(self):
        l1 = [1, 2, 5]
        l2 = [0, 1, 2, 3, 1, 2, 5, 6]
        self.assertEqual(check_lists(l1, l2), SUBLIST)

    def test_consecutive(self):
        l1 = [1, 1, 2]
        l2 = [0, 1, 1, 1, 2, 1, 2]
        self.assertEqual(check_lists(l1, l2), SUBLIST)

    def test_sublist_at_start(self):
        l1 = [0, 1, 2]
        l2 = [0, 1, 2, 3, 4, 5]
        self.assertEqual(check_lists(l1, l2), SUBLIST)

    def test_sublist_in_middle(self):
        l1 = [2, 3, 4]
        l2 = [0, 1, 2, 3, 4, 5]
        self.assertEqual(check_lists(l1, l2), SUBLIST)

    def test_sublist_at_end(self):
        l1 = [3, 4, 5]
        l2 = [0, 1, 2, 3, 4, 5]
        self.assertEqual(check_lists(l1, l2), SUBLIST)

    def test_at_start_of_superlist(self):
        l1 = [0, 1, 2, 3, 4, 5]
        l2 = [0, 1, 2]
        self.assertEqual(check_lists(l1, l2), SUPERLIST)

    def test_in_middle_of_superlist(self):
        l1 = [0, 1, 2, 3, 4, 5]
        l2 = [2, 3]
        self.assertEqual(check_lists(l1, l2), SUPERLIST)

    def test_at_end_of_superlist(self):
        l1 = [0, 1, 2, 3, 4, 5]
        l2 = [3, 4, 5]
        self.assertEqual(check_lists(l1, l2), SUPERLIST)

    def test_large_lists(self):
        l1 = list(range(1000)) * 1000 + list(range(1000, 1100))
        l2 = list(range(900, 1050))
        self.assertEqual(check_lists(l1, l2), SUPERLIST)

    def test_spread_sublist(self):
        multiples_of_3 = list(range(3, 200, 3))
        multiples_of_15 = list(range(15, 200, 15))
        self.assertEqual(check_lists(multiples_of_15, multiples_of_3), UNEQUAL)

    def test_double_digits(self):
        l1 = [1, 0, 1]
        l2 = [10, 1]
        self.assertEqual(check_lists(l1, l2), UNEQUAL)

    def test_inner_spaces(self):
        l1 = ['a c']
        l2 = ['a', 'c']
        self.assertEqual(check_lists(l1, l2), UNEQUAL)

    def test_avoid_sets(self):
        self.assertEqual(check_lists([1, 3], [1, 2, 3]), UNEQUAL)
        self.assertEqual(check_lists([1, 2, 3], [1, 3]), UNEQUAL)
        self.assertEqual(check_lists([1, 2, 3], [3, 2, 1]), UNEQUAL)

if __name__ == '__main__':
    unittest.main()
<commit_msg>sublist: Update test cases<commit_after>import unittest

from sublist import check_lists, SUBLIST, SUPERLIST, EQUAL, UNEQUAL

# test cases adapted from `x-common//canonical-data.json` @ version: 1.0.0

class SublistTest(unittest.TestCase):
    def test_unique_return_values(self):
        self.assertEqual(len(set([SUBLIST, SUPERLIST, EQUAL, UNEQUAL])), 4)

    def test_empty_lists(self):
        self.assertEqual(
            check_lists([], []),
            EQUAL
        )

    def test_empty_list_within_non_empty_list(self):
        self.assertEqual(
            check_lists([], [1, 2, 3]),
            SUBLIST
        )

    def test_non_empty_list_contains_empty_list(self):
        self.assertEqual(
            check_lists([1, 2, 3], []),
            SUPERLIST
        )

    def test_list_equals_itself(self):
        self.assertEqual(
            check_lists([1, 2, 3], [1, 2, 3]),
            EQUAL
        )

    def test_different_lists(self):
        self.assertEqual(
            check_lists([1, 2, 3], [2, 3, 4]),
            UNEQUAL
        )

    def test_false_start(self):
        self.assertEqual(
            check_lists([1, 2, 5], [0, 1, 2, 3, 1, 2, 5, 6]),
            SUBLIST
        )

    def test_consecutive(self):
        self.assertEqual(
            check_lists([1, 1, 2], [0, 1, 1, 1, 2, 1, 2]),
            SUBLIST
        )

    def test_sublist_at_start(self):
        self.assertEqual(
            check_lists([0, 1, 2], [0, 1, 2, 3, 4, 5]),
            SUBLIST
        )

    def test_sublist_in_middle(self):
        self.assertEqual(
            check_lists([2, 3, 4], [0, 1, 2, 3, 4, 5]),
            SUBLIST
        )

    def test_sublist_at_end(self):
        self.assertEqual(
            check_lists([3, 4, 5], [0, 1, 2, 3, 4, 5]),
            SUBLIST
        )

    def test_at_start_of_superlist(self):
        self.assertEqual(
            check_lists([0, 1, 2, 3, 4, 5], [0, 1, 2]),
            SUPERLIST
        )

    def test_in_middle_of_superlist(self):
        self.assertEqual(
            check_lists([0, 1, 2, 3, 4, 5], [2, 3]),
            SUPERLIST
        )

    def test_at_end_of_superlist(self):
        self.assertEqual(
            check_lists([0, 1, 2, 3, 4, 5], [3, 4, 5]),
            SUPERLIST
        )

    def test_first_list_missing_element_from_second_list(self):
        self.assertEqual(
            check_lists([1, 3], [1, 2, 3]),
            UNEQUAL
        )

    def test_second_list_missing_element_from_first_list(self):
        self.assertEqual(
            check_lists([1, 2, 3], [1, 3]),
            UNEQUAL
        )

    def test_order_matters_to_a_list(self):
        self.assertEqual(
            check_lists([1, 2, 3], [3, 2, 1]),
            UNEQUAL
        )

    def test_same_digits_but_different_numbers(self):
        self.assertEqual(
            check_lists([1, 0, 1], [10, 1]),
            UNEQUAL
        )

    # additional track specific test
    def test_inner_spaces(self):
        self.assertEqual(
            check_lists(['a c'], ['a', 'c']),
            UNEQUAL
        )

    # additional track specific test
    def test_large_lists(self):
        l1 = list(range(1000)) * 1000 + list(range(1000, 1100))
        l2 = list(range(900, 1050))
        self.assertEqual(check_lists(l1, l2), SUPERLIST)

    # additional track specific test
    def test_spread_sublist(self):
        multiples_of_3 = list(range(3, 200, 3))
        multiples_of_15 = list(range(15, 200, 15))
        self.assertEqual(check_lists(multiples_of_15, multiples_of_3), UNEQUAL)

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


<commit_before>#!/usr/bin/env python
import os
import sys
import platform
import tempfile

def cleanup(workingdir):
    print('+-' * 40)
    print('Cleanup Time...')
    print('Removing temporary data...')
    try:
        os.removedirs(workingdir)
    except:
        raise SystemError('Cleanup Failed!')

def main(**kwargs):
    """
Master Script that calls subscripts to be deployed to new Linux systems
    """

    scriptname = sys.argv[0]
<<<<<<< HEAD
    workingdir = '/usr/tmp/systemprep'
    readyfile = '/var/run/system-is-ready'
    scriptstart = '+' * 80
    scriptend = '-' * 80

    print ( str(scriptstart) )
    print ( 'Entering script -- ' + str(scriptname) )
    print ( 'Printing parameters...' )

    print 'kwargs = ', kwargs
    for key,value in kwargs.items() :
        print ( str(key) + ' = ' + str(value) )

=======
    #Logic to define variables based on OS
    if 'Windows' in platform.system():
        #workingdir = 'C:\WINDOWS\TEMP\SYSTEMPREP'
        workingdir = 'C:\WINDOWS\TEMP'
        readyfile = 'C:\WINDOWS\TEMP\system-is-ready'
    elif 'Linux' in platform.system():
        #workingdir = '/usr/tmp/systemprep'
        workingdir = '/usr/tmp/'
        readyfile = '/var/run/system-is-ready'
    else:
        raise SystemError('platform.system() is not recognized?')

    print('+' * 80)
    print('Entering script -- ' + scriptname)
    if not os.path.exists(readyfile):
        raise SystemError(readyfile + ' does not exist! System is not ready?')

    if os.path.exists(workingdir):
        print('Creating working directory.')
        workingdir = tempfile.mkdtemp(prefix='systemprep-', dir=workingdir)
        print(workingdir)
    else:
        raise SystemError('Defined workingdir does not exist!')

    print('Printing parameters...')
    print('kwargs = ' + str(kwargs))

    for key, value in kwargs.items():
        print(str(key) + ' = ' + str(value))

    cleanup(workingdir)
    print('-' * 80)
>>>>>>> origin/master
#    raw_input("\n\nPress the enter key to exit.")

if __name__ == "__main__":
    kwargs = dict(x.split('=', 1) for x in sys.argv[1:])
    main(**kwargs)
<commit_msg>Revert "-Added cleanup function and call before end"<commit_after>#!/usr/bin/env python
import os
import sys
import platform
import tempfile

def main(**kwargs):
    """
Master Script that calls subscripts to be deployed to new Linux systems
    """

    scriptname = sys.argv[0]
<<<<<<< HEAD
    workingdir = '/usr/tmp/systemprep'
    readyfile = '/var/run/system-is-ready'
    scriptstart = '+' * 80
    scriptend = '-' * 80

    print ( str(scriptstart) )
    print ( 'Entering script -- ' + str(scriptname) )
    print ( 'Printing parameters...' )

    print 'kwargs = ', kwargs
    for key,value in kwargs.items() :
        print ( str(key) + ' = ' + str(value) )

=======
    #Logic to define variables based on OS
    if 'Windows' in platform.system():
        #workingdir = 'C:\WINDOWS\TEMP\SYSTEMPREP'
        workingdir = 'C:\WINDOWS\TEMP'
        readyfile = 'C:\WINDOWS\TEMP\system-is-ready'
    elif 'Linux' in platform.system():
        #workingdir = '/usr/tmp/systemprep'
        workingdir = '/usr/tmp/'
        readyfile = '/var/run/system-is-ready'
    else:
        raise SystemError('platform.system() is not recognized?')

    print('+' * 80)
    print('Entering script -- ' + scriptname)
    if not os.path.exists(readyfile):
        raise SystemError(readyfile + ' does not exist! System is not ready?')

    if os.path.exists(workingdir):
        print('Creating working directory.')
        workingdir = tempfile.mkdtemp(prefix='systemprep-', dir=workingdir)
        print(workingdir)
    else:
        raise SystemError('Defined workingdir does not exist!')

    print('Printing parameters...')
    print('kwargs = ' + str(kwargs))

    for key, value in kwargs.items():
        print(str(key) + ' = ' + str(value))

    #TODO: Add Cleanup function and replace cleanup time. This does not cleanup if SystemError is raised after creation.
    #todo Cleanup should be a function that is called during a SystemError so files are never left behind unless needed.

    print('+-' * 40)
    print('Cleanup Time...')
    print('Removing temporary data...')
    os.removedirs(workingdir)
    print('-' * 80)
>>>>>>> origin/master
#    raw_input("\n\nPress the enter key to exit.")

if __name__ == "__main__":
    kwargs = dict(x.split('=', 1) for x in sys.argv[1:])
    main(**kwargs)
<|endoftext|>


<commit_before>"""
Expose top-level symbols that are safe for import *
"""

import platform
import re
import sys
import warnings

from ._version import get_versions
from numba.misc.init_utils import generate_version_info

__version__ = get_versions()["version"]
version_info = generate_version_info(__version__)
del get_versions
del generate_version_info

from numba.core import config
from numba.core import types, errors

# Re-export typeof
from numba.misc.special import (
    typeof,
    prange,
    pndindex,
    gdb,
    gdb_breakpoint,
    gdb_init,
    literally,
    literal_unroll,
)

# Re-export error classes
from numba.core.errors import *

# Re-export types itself
import numba.core.types as types

# Re-export all type names
from numba.core.types import *

# Re-export decorators
from numba.core.decorators import cfunc, generated_jit, jit, njit, stencil, jit_module

# Re-export vectorize decorators and the thread layer querying function
from numba.np.ufunc import (
    vectorize,
    guvectorize,
    threading_layer,
    get_num_threads,
    set_num_threads,
    set_parallel_chunksize,
    get_parallel_chunksize,
    get_thread_id,
)

# Re-export Numpy helpers
from numba.np.numpy_support import carray, farray, from_dtype

# Re-export experimental
from numba import experimental

# Initialize withcontexts
import numba.core.withcontexts
from numba.core.withcontexts import objmode_context as objmode
from numba.core.withcontexts import parallel_chunksize

# Initialize target extensions
import numba.core.target_extension

# Initialize typed containers
import numba.typed

# Keep this for backward compatibility.
def test(argv, **kwds):
    # To speed up the import time, avoid importing `unittest` and other test
    # dependencies unless the user is actually trying to run tests.
    from numba.testing import _runtests as runtests

    return runtests.main(argv, **kwds)

__all__ = (
    """
    cfunc
    from_dtype
    guvectorize
    jit
    experimental
    njit
    stencil
    jit_module
    typeof
    prange
    gdb
    gdb_breakpoint
    gdb_init
    vectorize
    objmode
    literal_unroll
    get_num_threads
    set_num_threads
    set_parallel_chunksize
    get_parallel_chunksize
    parallel_chunksize
    """.split()
    + types.__all__
    + errors.__all__
)

_min_llvmlite_version = (0, 40, 0)
_min_llvm_version = (11, 0, 0)

def _ensure_llvm():
    """
    Make sure llvmlite is operational.
    """
    import warnings
    import llvmlite

    # Only look at the the major, minor and bugfix version numbers.
    # Ignore other stuffs
    regex = re.compile(r"(\d+)\.(\d+).(\d+)")
    m = regex.match(llvmlite.__version__)
    if m:
        ver = tuple(map(int, m.groups()))
        if ver < _min_llvmlite_version:
            msg = (
                "Numba requires at least version %d.%d.%d of llvmlite.\n"
                "Installed version is %s.\n"
                "Please update llvmlite."
                % (_min_llvmlite_version + (llvmlite.__version__,))
            )
            raise ImportError(msg)
    else:
        # Not matching?
        warnings.warn("llvmlite version format not recognized!")

    from llvmlite.binding import llvm_version_info, check_jit_execution

    if llvm_version_info < _min_llvm_version:
        msg = (
            "Numba requires at least version %d.%d.%d of LLVM.\n"
            "Installed llvmlite is built against version %d.%d.%d.\n"
            "Please update llvmlite." % (_min_llvm_version + llvm_version_info)
        )
        raise ImportError(msg)

    check_jit_execution()

def _ensure_critical_deps():
    """
    Make sure Python, NumPy and SciPy have supported versions.
    """
    from numba.np.numpy_support import numpy_version
    from numba.core.utils import PYVERSION

    if PYVERSION < (3, 7):
        raise ImportError(f"Numba needs Python 3.7 or greater. Got {PYVERSION}")

    if numpy_version < (1, 18):
        raise ImportError(f"Numba needs NumPy 1.18 or greater. Got {numpy_version}")

    try:
        import scipy
    except ImportError:
        pass
    else:
        sp_version = tuple(map(int, scipy.__version__.split(".")[:2]))
        if sp_version < (1, 0):
            raise ImportError(
                f"Numba requires SciPy version 1.0 or greater. Got {sp_version}"
            )

def _try_enable_svml():
    """
    Tries to enable SVML if configuration permits use and the library is found.
    """
    if not config.DISABLE_INTEL_SVML:
        try:
            if sys.platform.startswith("linux"):
                llvmlite.binding.load_library_permanently("libsvml.so")
            elif sys.platform.startswith("darwin"):
                llvmlite.binding.load_library_permanently("libsvml.dylib")
            elif sys.platform.startswith("win"):
                llvmlite.binding.load_library_permanently("svml_dispmd")
            else:
                return False
            # The SVML library is loaded, therefore SVML *could* be supported.
            # Now see if LLVM has been compiled with the SVML support patch.
            # If llvmlite has the checking function `has_svml` and it returns
            # True, then LLVM was compiled with SVML support and the the setup
            # for SVML can proceed. We err on the side of caution and if the
            # checking function is missing, regardless of that being fine for
            # most 0.23.{0,1} llvmlite instances (i.e. conda or pip installed),
            # we assume that SVML was not compiled in. llvmlite 0.23.2 is a
            # bugfix release with the checking function present that will always
            # produce correct behaviour. For context see: #3006.
            try:
                if not getattr(llvmlite.binding.targets, "has_svml")():
                    # has detection function, but no svml compiled in, therefore
                    # disable SVML
                    return False
            except AttributeError:
                if platform.machine() == "x86_64" and config.DEBUG:
                    msg = (
                        "SVML was found but llvmlite >= 0.23.2 is "
                        "needed to support it."
                    )
                    warnings.warn(msg)
                # does not have detection function, cannot detect reliably,
                # disable SVML.
                return False

            # All is well, detection function present and reports SVML is
            # compiled in, set the vector library to SVML.
            llvmlite.binding.set_option("SVML", "-vector-library=SVML")
            return True
        except:
            if platform.machine() == "x86_64" and config.DEBUG:
                warnings.warn("SVML was not found/could not be loaded.")
    return False

_ensure_llvm()
_ensure_critical_deps()

# we know llvmlite is working as the above tests passed, import it now as SVML
# needs to mutate runtime options (sets the `-vector-library`).
import llvmlite

"""
Is set to True if Intel SVML is in use.
"""
config.USING_SVML = _try_enable_svml()

# ---------------------- WARNING WARNING WARNING ----------------------------
# The following imports occur below here (SVML init) because somewhere in their
# import sequence they have a `@njit` wrapped function. This triggers too early
# a bind to the underlying LLVM libraries which then irretrievably sets the LLVM
# SVML state to "no SVML". See https://github.com/numba/numba/issues/4689 for
# context.
# ---------------------- WARNING WARNING WARNING ----------------------------
<commit_msg>add .<commit_after>"""
Expose top-level symbols that are safe for import *
"""

import platform
import re
import sys
import warnings

from ._version import get_versions
from numba.misc.init_utils import generate_version_info

__version__ = get_versions()["version"]
version_info = generate_version_info(__version__)
del get_versions
del generate_version_info

from numba.core import config
from numba.core import types, errors

# Re-export typeof
from numba.misc.special import (
    typeof,
    prange,
    pndindex,
    gdb,
    gdb_breakpoint,
    gdb_init,
    literally,
    literal_unroll,
)

# Re-export error classes
from numba.core.errors import *

# Re-export types itself
import numba.core.types as types

# Re-export all type names
from numba.core.types import *

# Re-export decorators
from numba.core.decorators import cfunc, generated_jit, jit, njit, stencil, jit_module

# Re-export vectorize decorators and the thread layer querying function
from numba.np.ufunc import (
    vectorize,
    guvectorize,
    threading_layer,
    get_num_threads,
    set_num_threads,
    set_parallel_chunksize,
    get_parallel_chunksize,
    get_thread_id,
)

# Re-export Numpy helpers
from numba.np.numpy_support import carray, farray, from_dtype

# Re-export experimental
from numba import experimental

# Initialize withcontexts
import numba.core.withcontexts
from numba.core.withcontexts import objmode_context as objmode
from numba.core.withcontexts import parallel_chunksize

# Initialize target extensions
import numba.core.target_extension

# Initialize typed containers
import numba.typed

# Keep this for backward compatibility.
def test(argv, **kwds):
    # To speed up the import time, avoid importing `unittest` and other test
    # dependencies unless the user is actually trying to run tests.
    from numba.testing import _runtests as runtests

    return runtests.main(argv, **kwds)

__all__ = (
    """
    cfunc
    from_dtype
    guvectorize
    jit
    experimental
    njit
    stencil
    jit_module
    typeof
    prange
    gdb
    gdb_breakpoint
    gdb_init
    vectorize
    objmode
    literal_unroll
    get_num_threads
    set_num_threads
    set_parallel_chunksize
    get_parallel_chunksize
    parallel_chunksize
    """.split()
    + types.__all__
    + errors.__all__
)

_min_llvmlite_version = (0, 40, 0)
_min_llvm_version = (11, 0, 0)

def _ensure_llvm():
    """
    Make sure llvmlite is operational.
    """
    import warnings
    import llvmlite

    # Only look at the the major, minor and bugfix version numbers.
    # Ignore other stuffs
    regex = re.compile(r"(\d+)\.(\d+).(\d+)")
    m = regex.match(llvmlite.__version__)
    if m:
        ver = tuple(map(int, m.groups()))
        if ver < _min_llvmlite_version:
            msg = (
                "Numba requires at least version %d.%d.%d of llvmlite.\n"
                "Installed version is %s.\n"
                "Please update llvmlite."
                % (_min_llvmlite_version + (llvmlite.__version__,))
            )
            raise ImportError(msg)
    else:
        # Not matching?
        warnings.warn("llvmlite version format not recognized!")

    from llvmlite.binding import llvm_version_info, check_jit_execution

    if llvm_version_info < _min_llvm_version:
        msg = (
            "Numba requires at least version %d.%d.%d of LLVM.\n"
            "Installed llvmlite is built against version %d.%d.%d.\n"
            "Please update llvmlite." % (_min_llvm_version + llvm_version_info)
        )
        raise ImportError(msg)

    check_jit_execution()

def _ensure_critical_deps():
    """
    Make sure Python, NumPy and SciPy have supported versions.
    """
    from numba.np.numpy_support import numpy_version
    from numba.core.utils import PYVERSION

    if PYVERSION < (3, 7):
        raise ImportError(f"Numba needs Python 3.7 or greater. Got {PYVERSION}.")

    if numpy_version < (1, 18):
        raise ImportError(f"Numba needs NumPy 1.18 or greater. Got {numpy_version}.")

    try:
        import scipy
    except ImportError:
        pass
    else:
        sp_version = tuple(map(int, scipy.__version__.split(".")[:2]))
        if sp_version < (1, 0):
            raise ImportError(
                f"Numba requires SciPy version 1.0 or greater. Got {sp_version}."
            )

def _try_enable_svml():
    """
    Tries to enable SVML if configuration permits use and the library is found.
    """
    if not config.DISABLE_INTEL_SVML:
        try:
            if sys.platform.startswith("linux"):
                llvmlite.binding.load_library_permanently("libsvml.so")
            elif sys.platform.startswith("darwin"):
                llvmlite.binding.load_library_permanently("libsvml.dylib")
            elif sys.platform.startswith("win"):
                llvmlite.binding.load_library_permanently("svml_dispmd")
            else:
                return False
            # The SVML library is loaded, therefore SVML *could* be supported.
            # Now see if LLVM has been compiled with the SVML support patch.
            # If llvmlite has the checking function `has_svml` and it returns
            # True, then LLVM was compiled with SVML support and the the setup
            # for SVML can proceed. We err on the side of caution and if the
            # checking function is missing, regardless of that being fine for
            # most 0.23.{0,1} llvmlite instances (i.e. conda or pip installed),
            # we assume that SVML was not compiled in. llvmlite 0.23.2 is a
            # bugfix release with the checking function present that will always
            # produce correct behaviour. For context see: #3006.
            try:
                if not getattr(llvmlite.binding.targets, "has_svml")():
                    # has detection function, but no svml compiled in, therefore
                    # disable SVML
                    return False
            except AttributeError:
                if platform.machine() == "x86_64" and config.DEBUG:
                    msg = (
                        "SVML was found but llvmlite >= 0.23.2 is "
                        "needed to support it."
                    )
                    warnings.warn(msg)
                # does not have detection function, cannot detect reliably,
                # disable SVML.
                return False

            # All is well, detection function present and reports SVML is
            # compiled in, set the vector library to SVML.
            llvmlite.binding.set_option("SVML", "-vector-library=SVML")
            return True
        except:
            if platform.machine() == "x86_64" and config.DEBUG:
                warnings.warn("SVML was not found/could not be loaded.")
    return False

_ensure_llvm()
_ensure_critical_deps()

# we know llvmlite is working as the above tests passed, import it now as SVML
# needs to mutate runtime options (sets the `-vector-library`).
import llvmlite

"""
Is set to True if Intel SVML is in use.
"""
config.USING_SVML = _try_enable_svml()

# ---------------------- WARNING WARNING WARNING ----------------------------
# The following imports occur below here (SVML init) because somewhere in their
# import sequence they have a `@njit` wrapped function. This triggers too early
# a bind to the underlying LLVM libraries which then irretrievably sets the LLVM
# SVML state to "no SVML". See https://github.com/numba/numba/issues/4689 for
# context.
# ---------------------- WARNING WARNING WARNING ----------------------------
<|endoftext|>


<commit_before>import sys

import corner
import emcee
import matplotlib.pyplot as pl
import numpy as np
import scipy.optimize as op
from matplotlib.ticker import MaxNLocator

from model import MMNModel

class MMNFitter:
    """
    This class is used to fit a certain Multi Miyamoto Nagai model (with a predefined number of disks) to a datafile.
    """
    def __init__(self, n_walkers=100, n_steps=1000, n_threads=1, random_seed=120, fit_type='potential', check_positive_definite=False, verbose=True):
        """
        Constructor for the MultiMiyamotoNagai fitter. The fitter is based on emcee.

        :param n_walkers: emcee parameter to indicate how many parallel walkers the MCMC will use to fit the data
        :param n_steps: How many steps should the MCMC method proceed before stopping
        :param random_seed: A seed used to generate the initial positions of the walkers
        """
        self.n_walkers = n_walkers
        self.n_steps = n_steps
        self.n_threads = n_threads

        # The fitted models
        self.sampler = None
        self.samples = ()
        self.quantiles = None
        self.models = None
        self.axes = None
        self.ndim = 0
        self.fit_type = fit_type

        # The data samples we are fitting again :
        self.data = None
        self.n_values = 0
        self.yerr = None

        # Flags
        self.verbose = verbose
        self.check_DP = check_positive_definite
        if self.check_DP and self.verbose:
            print('Warning : Checking for definite-positiveness at every walker step. This ensures that the end model will be definite positive but' +
                  'might take a very long time to compute !')

        np.random.seed(random_seed)

    def set_model_type(self, nx=0, ny=0, nz=1):
        """
        Defines the type of model we are trying to fit
        :param nx: Number of disks on the yz plane
        :param ny: Number of disks on the xz plane
        :param nz: Number of disks on the xy plane
        """
        self.ndim = (nx+ny+nz)*3
        self.models = np.random.rand(self.ndim)
        self.axes = ['x']*nx + ['y']*ny + ['z']*nz

    def load_data(self, filename):
        """
        Loads the data that needs to be fitted. The data should be in an ascii file with four columns : X Y Z potential
        :param filename: The filename to open
        """
        self.data = np.loadtxt(filename)
        self.n_values = self.data.shape[0]
        self.yerr = 0.01*np.random.rand(self.n_values)

    def sMN(self, models=(), fit_type=None):
        """
        Sums the values of the different Miyamoto on a set of points
        The type of value summed depends on the fit we want to realize
        :param models: the list of models we are summing. If none the models of the instance are taken
        :param fit_type: the type of fit we are doing. Can be "potential", "density", "force" or None. If None then
        the default fit_type of the instance is taken
        :return: a scalar, the total sum of all the models on all points
        """
        if len(models) == 0:
            models = self.models

        if not fit_type:
            fit_type = self.fit_type

        # We pick the function to apply according to the fit model
        if fit_type == 'density':
            eval_func = MMNModel.mn_density
        elif fit_type == 'potential':
            eval_func = MMNModel.mn_potential
        else:
            eval_func = MMNModel.mn_force

        # The positions of the points
        x = self.data[:, 0]
        y = self.data[:, 1]
        z = self.data[:, 2]

        # Radius on each plane
        rxy = np.sqrt(x**2+y**2)
        rxz = np.sqrt(x**2+z**2)
        ryz = np.sqrt(y**2+z**2)

        total_sum = 0.0

        # Summing on each model
        for id_mod, axis in enumerate(self.axes):
            a, b, M = models[id_mod*3:(id_mod+1)*3]
            if axis == "x":
                value = eval_func(ryz, x, a, b, M)
            elif axis == "y":
                value = eval_func(rxz, y, a, b, M)
            else:
                value = eval_func(rxy, z, a, b, M)

            total_sum += value
        return total_sum

    def loglikelihood(self, models):
        """
        This function computes the loglikelihood of the
        :param models: the list of models
        :return: the loglikelihood of the sum of models
        """

        tmp_model = MMNModel()

        # Checking that a+b > 0 for every model :
        for id_mod, axis in enumerate(self.axes):
            a, b, M = models[id_mod*3:(id_mod+1)*3]
            if a+b < 0:
                return -np.inf

            # If we are checking for positive-definiteness we add the disk to the model
            if self.check_DP:
                tmp_model.add_model(axis, a, b, M)

        # Now checking for positive-definiteness:
        if self.check_DP:
            if not tmp_model.is_positive_definite():
                return -np.inf

        # Everything ok, we proceed with the likelihood :
        p = self.data[:, 3]
        model = self.sMN(models)
        inv_sigma2 = 1.0/(self.yerr**2)
        return -0.5*(np.sum((p-model)**2*inv_sigma2-np.log(inv_sigma2)))

    def maximum_likelihood(self):
        """
        Computation of the maximum of likelihood of the models
        """
        if self.verbose:
            print("Computing maximum of likelihood")

        # Optimizing the parameters of the models to minimize the loglikelihood
        models = self.models
        chi2 = lambda m: -2 * self.loglikelihood(m)
        result = op.minimize(chi2, models)
        values = result["x"]

        if self.verbose:
            print("Maximum of likelihood results :")

            axis_stat = {"x": [1, "yz"], "y": [1, "xz"], "z": [1, "xy"]}
            for id_mod, axis in enumerate(self.axes):
                stat = axis_stat[axis]
                axis_name = "{0}{1}".format(stat[1], stat[0])

                print("a{0} = {1} (initial : {2})".format(axis_name, values[id_mod*3], models[id_mod*3]))
                print("b{0} = {1} (initial : {2})".format(axis_name, values[id_mod*3+1], models[id_mod*3+1]))
                print("M{0} = {1} (initial : {2})".format(axis_name, values[id_mod*3+2], models[id_mod*3+2]))

                stat[0] += 1

        # Storing the best values as current models
        self.models = values

    def fit_data(self, burnin=100):
        """
        This function finds the parameters of the models using emcee
        :param burnin: the number of timesteps to keep after running emcee
        :returns: A list of all the samples truncated to give only from the burning timestep
        """

        # We initialize the positions of the walkers by adding a small random component to each parameter
        pos = [self.models + 1e-4*np.random.randn(self.ndim) for i in range(self.n_walkers)]

        # Running the MCMC to get the parameters
        if self.verbose:
            print("Running emcee ...")

        self.sampler = emcee.EnsembleSampler(self.n_walkers, self.ndim, self.loglikelihood, threads=self.n_threads)
        self.sampler.run_mcmc(pos, self.n_steps, rstate0=np.random.get_state())

        # Storing the last burnin results
        #nstp= 0.5*self.n_steps
        print(self.sampler.chain.shape)
        self.samples = self.sampler.chain[:, burnin:, :].reshape((-1, self.ndim))

        if self.verbose:
            print("Done.")

        # Checking for positive-definiteness
        everything_dp = True
        for sample in self.samples:
            tmp_model = MMNModel()
            for id_mod, axis in enumerate(self.axes):
                  a, b, M = sample[id_mod*3:(id_mod+1)*3]
                  tmp_model.add_model(axis, a, b, M)

            if not tmp_model.is_positive_definite:
                  everything_dp = False
                  break

        if not everything_dp:
            print('Warning : Some sample results are not positive definite ! You can end up with negative densities')
            print('To ensure a positive definite model, consider setting the parameter "check_positive_definite" to True in the fitter !')

        return self.samples

    def plot_disk_walkers(self, id_mod):
        """
        Plotting the walkers on each parameter of a certain model
        :param id_mod: the id of the disk parameters you want to plot
        """
        axis_name = {"x": "yz", "y": "xz", "z": "xy"}[self.axes[id_mod]]
        fig, axes = pl.subplots(3, 1, sharex=True, figsize=(8, 9))
        axes[0].plot(self.sampler.chain[:, :, 0].T, color="k", alpha=0.4)
        axes[0].yaxis.set_major_locator(MaxNLocator(5))
        axes[0].axhline(self.models[id_mod*3], color="#888888", lw=2)
        axes[0].set_ylabel("$a{0}$".format(axis_name))

        axes[1].plot(self.sampler.chain[:, :, 1].T, color="k", alpha=0.4)
        axes[1].yaxis.set_major_locator(MaxNLocator(5))
        axes[1].axhline(self.models[id_mod*3+1], color="#888888", lw=2)
        axes[1].set_ylabel("$b{0}$".format(axis_name))

        axes[2].plot(self.sampler.chain[:, :, 2].T, color="k", alpha=0.4)
        axes[2].yaxis.set_major_locator(MaxNLocator(5))
        axes[2].axhline(self.models[id_mod*3+2], color="#888888", lw=2)
        axes[2].set_ylabel("$M{0}$".format(axis_name))
        fig.savefig("Time.png")

    def corner_plot(self):
        """
        Draws the corner plot of the fitted data
        """
        labels = []
        axis_stat = {"x": [1, "yz"], "y": [1, "xz"], "z": [1, "xy"]}

        for id_mod, axis in enumerate(self.axes):
            stat = axis_stat[axis]
            axis_name = "{0}{1}".format(stat[1], stat[0])
            labels += ["a{0}".format(axis_name), "b{0}".format(axis_name), "M{0}".format(axis_name)]
            stat[0] += 1

        if self.verbose:
            print("Computing corner plot ...")

        figt = corner.corner(self.samples, labels=labels, truths=self.models)
        figt.savefig("Triangle.png")

    def compute_quantiles(self, quantiles=(16, 50, 84)):
        """
        Finds the quantiles values on the whole sample kept after emcee run. The results are stored in the quantiles
        attribute of the class.
        :param quantiles: a tuple indicating what quantiles are to be computed
        """
        if len(quantiles) != 3:
            sys.stderr.write('Warning : The quantile list should always be a triplet')
            return

        if len(self.samples) == 0:
            sys.stderr.write('Warning : You should not run compute_quantiles before fit_data ! Trying to compute quantiles.')

        k = lambda v: (v[1], v[2]-v[1], v[1]-v[0])
        qarray = np.array(np.percentile(self.samples, quantiles, axis=0))
        self.quantiles = np.array((qarray[1], qarray[2]-qarray[1], qarray[1]-qarray[0])).T

        if self.verbose:
            print("MCMC results :")
            axis_stat = {"x": [1, "yz"], "y": [1, "xz"], "z": [1, "xy"]}
            for id_mod, axis in enumerate(self.axes):
                stat = axis_stat[axis]
                axis_name = "{0}{1}".format(stat[1], stat[0])
                base_format = "{0} = {1[0]} +: {1[1]} -: {1[2]}"
                print("a"+base_format.format(axis_name, self.quantiles[id_mod*3], self.quantiles[id_mod*3]))
                print("b"+base_format.format(axis_name, self.quantiles[id_mod*3+1], self.quantiles[id_mod*3+1]))
                print("M"+base_format.format(axis_name, self.quantiles[id_mod*3+2], self.quantiles[id_mod*3+2]))
                stat[0] += 1

        return self.quantiles

<commit_msg>Added pickle functions to make multiprocessing inside the fitter class working<commit_after>import sys

import corner
import emcee
import matplotlib.pyplot as pl
import numpy as np
import scipy.optimize as op
from matplotlib.ticker import MaxNLocator

from model import MMNModel

# Thanks to Steven Bethard for this nice trick, found on :
# https://bytes.com/topic/python/answers/552476-why-cant-you-pickle-instancemethods

# Allows the methods of MMNFitter to be pickled for multiprocessing
import copy_reg
import types

def _pickle_method(method):
    func_name = method.im_func.__name__
    obj = method.im_self
    cls = method.im_class
    return _unpickle_method, (func_name, obj, cls)

def _unpickle_method(func_name, obj, cls):
    #for cls in cls.mro():
    try:
        func = cls.__dict__[func_name]
        return func.__get__(obj, cls)
    except KeyError:
        return None

copy_reg.pickle(types.MethodType, _pickle_method, _unpickle_method)

sampler = None

class MMNFitter:
    """
    This class is used to fit a certain Multi Miyamoto Nagai model (with a predefined number of disks) to a datafile.
    """
    def __init__(self, n_walkers=100, n_steps=1000, n_threads=1, random_seed=120, fit_type='potential', check_positive_definite=False, verbose=True):
        """
        Constructor for the MultiMiyamotoNagai fitter. The fitter is based on emcee.

        :param n_walkers: emcee parameter to indicate how many parallel walkers the MCMC will use to fit the data
        :param n_steps: How many steps should the MCMC method proceed before stopping
        :param random_seed: A seed used to generate the initial positions of the walkers
        """
        self.n_walkers = n_walkers
        self.n_steps = n_steps
        self.n_threads = n_threads

        # The fitted models
        self.samples = ()
        self.quantiles = None
        self.models = None
        self.axes = None
        self.ndim = 0
        self.fit_type = fit_type

        # The data samples we are fitting again :
        self.data = None
        self.n_values = 0
        self.yerr = None

        # Flags
        self.verbose = verbose
        self.check_DP = check_positive_definite
        if self.check_DP and self.verbose:
            print('Warning : Checking for definite-positiveness at every walker step. This ensures that the end model will be definite positive but' +
                  'might take a very long time to compute !')

        np.random.seed(random_seed)

    def set_model_type(self, nx=0, ny=0, nz=1):
        """
        Defines the type of model we are trying to fit
        :param nx: Number of disks on the yz plane
        :param ny: Number of disks on the xz plane
        :param nz: Number of disks on the xy plane
        """
        self.ndim = (nx+ny+nz)*3
        self.models = np.random.rand(self.ndim)
        self.axes = ['x']*nx + ['y']*ny + ['z']*nz

    def load_data(self, filename):
        """
        Loads the data that needs to be fitted. The data should be in an ascii file with four columns : X Y Z potential
        :param filename: The filename to open
        """
        self.data = np.loadtxt(filename)
        self.n_values = self.data.shape[0]
        self.yerr = 0.01*np.random.rand(self.n_values)

    def sMN(self, models=(), fit_type=None):
        """
        Sums the values of the different Miyamoto on a set of points
        The type of value summed depends on the fit we want to realize
        :param models: the list of models we are summing. If none the models of the instance are taken
        :param fit_type: the type of fit we are doing. Can be "potential", "density", "force" or None. If None then
        the default fit_type of the instance is taken
        :return: a scalar, the total sum of all the models on all points
        """
        if len(models) == 0:
            models = self.models

        if not fit_type:
            fit_type = self.fit_type

        # We pick the function to apply according to the fit model
        if fit_type == 'density':
            eval_func = MMNModel.mn_density
        elif fit_type == 'potential':
            eval_func = MMNModel.mn_potential
        else:
            eval_func = MMNModel.mn_force

        # The positions of the points
        x = self.data[:, 0]
        y = self.data[:, 1]
        z = self.data[:, 2]

        # Radius on each plane
        rxy = np.sqrt(x**2+y**2)
        rxz = np.sqrt(x**2+z**2)
        ryz = np.sqrt(y**2+z**2)

        total_sum = 0.0

        # Summing on each model
        for id_mod, axis in enumerate(self.axes):
            a, b, M = models[id_mod*3:(id_mod+1)*3]
            if axis == "x":
                value = eval_func(ryz, x, a, b, M)
            elif axis == "y":
                value = eval_func(rxz, y, a, b, M)
            else:
                value = eval_func(rxy, z, a, b, M)

            total_sum += value
        return total_sum

    def loglikelihood(self, models):
        """
        This function computes the loglikelihood of the
        :param models: the list of models
        :return: the loglikelihood of the sum of models
        """

        tmp_model = MMNModel()

        # Checking that a+b > 0 for every model :
        for id_mod, axis in enumerate(self.axes):
            a, b, M = models[id_mod*3:(id_mod+1)*3]
            if a+b < 0:
                return -np.inf

            # If we are checking for positive-definiteness we add the disk to the model
            if self.check_DP:
                tmp_model.add_model(axis, a, b, M)

        # Now checking for positive-definiteness:
        if self.check_DP:
            if not tmp_model.is_positive_definite():
                return -np.inf

        # Everything ok, we proceed with the likelihood :
        p = self.data[:, 3]
        model = self.sMN(models)
        inv_sigma2 = 1.0/(self.yerr**2)
        return -0.5*(np.sum((p-model)**2*inv_sigma2-np.log(inv_sigma2)))

    def maximum_likelihood(self):
        """
        Computation of the maximum of likelihood of the models
        """
        if self.verbose:
            print("Computing maximum of likelihood")

        # Optimizing the parameters of the models to minimize the loglikelihood
        models = self.models
        chi2 = lambda m: -2 * self.loglikelihood(m)
        result = op.minimize(chi2, models)
        values = result["x"]

        if self.verbose:
            print("Maximum of likelihood results :")

            axis_stat = {"x": [1, "yz"], "y": [1, "xz"], "z": [1, "xy"]}
            for id_mod, axis in enumerate(self.axes):
                stat = axis_stat[axis]
                axis_name = "{0}{1}".format(stat[1], stat[0])

                print("a{0} = {1} (initial : {2})".format(axis_name, values[id_mod*3], models[id_mod*3]))
                print("b{0} = {1} (initial : {2})".format(axis_name, values[id_mod*3+1], models[id_mod*3+1]))
                print("M{0} = {1} (initial : {2})".format(axis_name, values[id_mod*3+2], models[id_mod*3+2]))

                stat[0] += 1

        # Storing the best values as current models
        self.models = values

    def fit_data(self, burnin=100):
        """
        This function finds the parameters of the models using emcee
        :param burnin: the number of timesteps to keep after running emcee
        :returns: A list of all the samples truncated to give only from the burning timestep
        """

        # We initialize the positions of the walkers by adding a small random component to each parameter
        pos = [self.models + 1e-4*np.random.randn(self.ndim) for i in range(self.n_walkers)]

        # Running the MCMC to get the parameters
        if self.verbose:
            print("Running emcee ...")

        global sampler
        sampler = emcee.EnsembleSampler(self.n_walkers, self.ndim, self.loglikelihood, threads=self.n_threads)
        sampler.run_mcmc(pos, self.n_steps, rstate0=np.random.get_state())

        # Storing the last burnin results
        #nstp= 0.5*self.n_steps
        print(sampler.chain.shape)
        self.samples = sampler.chain[:, burnin:, :].reshape((-1, self.ndim))

        if self.verbose:
            print("Done.")

        # Checking for positive-definiteness
        everything_dp = True
        for sample in self.samples:
            tmp_model = MMNModel()
            for id_mod, axis in enumerate(self.axes):
                  a, b, M = sample[id_mod*3:(id_mod+1)*3]
                  tmp_model.add_model(axis, a, b, M)

            if not tmp_model.is_positive_definite:
                  everything_dp = False
                  break

        if not everything_dp:
            print('Warning : Some sample results are not positive definite ! You can end up with negative densities')
            print('To ensure a positive definite model, consider setting the parameter "check_positive_definite" to True in the fitter !')

        return self.samples

    def plot_disk_walkers(self, id_mod):
        """
        Plotting the walkers on each parameter of a certain model
        :param id_mod: the id of the disk parameters you want to plot
        """
        axis_name = {"x": "yz", "y": "xz", "z": "xy"}[self.axes[id_mod]]
        fig, axes = pl.subplots(3, 1, sharex=True, figsize=(8, 9))
        axes[0].plot(sampler.chain[:, :, 0].T, color="k", alpha=0.4)
        axes[0].yaxis.set_major_locator(MaxNLocator(5))
        axes[0].axhline(self.models[id_mod*3], color="#888888", lw=2)
        axes[0].set_ylabel("$a{0}$".format(axis_name))

        axes[1].plot(sampler.chain[:, :, 1].T, color="k", alpha=0.4)
        axes[1].yaxis.set_major_locator(MaxNLocator(5))
        axes[1].axhline(self.models[id_mod*3+1], color="#888888", lw=2)
        axes[1].set_ylabel("$b{0}$".format(axis_name))

        axes[2].plot(sampler.chain[:, :, 2].T, color="k", alpha=0.4)
        axes[2].yaxis.set_major_locator(MaxNLocator(5))
        axes[2].axhline(self.models[id_mod*3+2], color="#888888", lw=2)
        axes[2].set_ylabel("$M{0}$".format(axis_name))
        fig.savefig("Time.png")

    def corner_plot(self):
        """
        Draws the corner plot of the fitted data
        """
        labels = []
        axis_stat = {"x": [1, "yz"], "y": [1, "xz"], "z": [1, "xy"]}

        for id_mod, axis in enumerate(self.axes):
            stat = axis_stat[axis]
            axis_name = "{0}{1}".format(stat[1], stat[0])
            labels += ["a{0}".format(axis_name), "b{0}".format(axis_name), "M{0}".format(axis_name)]
            stat[0] += 1

        if self.verbose:
            print("Computing corner plot ...")

        figt = corner.corner(self.samples, labels=labels, truths=self.models)
        figt.savefig("Triangle.png")

    def compute_quantiles(self, quantiles=(16, 50, 84)):
        """
        Finds the quantiles values on the whole sample kept after emcee run. The results are stored in the quantiles
        attribute of the class.
        :param quantiles: a tuple indicating what quantiles are to be computed
        """
        if len(quantiles) != 3:
            sys.stderr.write('Warning : The quantile list should always be a triplet')
            return

        if len(self.samples) == 0:
            sys.stderr.write('Warning : You should not run compute_quantiles before fit_data ! Trying to compute quantiles.')

        k = lambda v: (v[1], v[2]-v[1], v[1]-v[0])
        qarray = np.array(np.percentile(self.samples, quantiles, axis=0))
        self.quantiles = np.array((qarray[1], qarray[2]-qarray[1], qarray[1]-qarray[0])).T

        if self.verbose:
            print("MCMC results :")
            axis_stat = {"x": [1, "yz"], "y": [1, "xz"], "z": [1, "xy"]}
            for id_mod, axis in enumerate(self.axes):
                stat = axis_stat[axis]
                axis_name = "{0}{1}".format(stat[1], stat[0])
                base_format = "{0} = {1[0]} +: {1[1]} -: {1[2]}"
                print("a"+base_format.format(axis_name, self.quantiles[id_mod*3], self.quantiles[id_mod*3]))
                print("b"+base_format.format(axis_name, self.quantiles[id_mod*3+1], self.quantiles[id_mod*3+1]))
                print("M"+base_format.format(axis_name, self.quantiles[id_mod*3+2], self.quantiles[id_mod*3+2]))
                stat[0] += 1

        return self.quantiles

<|endoftext|>