<commit_before>from __future__ import print_function
from __future__ import unicode_literals
from builtins import str
import os
import argparse
import inspect
import sys
from nipype.interfaces.base import Interface, InputMultiPath, traits
from nipype.utils.misc import str2bool

def listClasses(module=None):
    if module:
        __import__(module)
        pkg = sys.modules[module]
        print("Available Interfaces:")
        for k,v in sorted(list(pkg.__dict__.items())):
            if inspect.isclass(v) and issubclass(v, Interface):
                print("\t%s"%k)

def add_options(parser=None, module=None, function=None):
    interface = None
    if parser and module and function:
        __import__(module)
        interface = getattr(sys.modules[module],function)()

        inputs = interface.input_spec()
        for name, spec in sorted(interface.inputs.traits(transient=None).items()):
            desc = "\n".join(interface._get_trait_desc(inputs, name, spec))[len(name)+2:]
            args = {}

            if spec.is_trait_type(traits.Bool):
                args["action"] = 'store_true'

            if hasattr(spec, "mandatory") and spec.mandatory:
                if spec.is_trait_type(InputMultiPath):
                    args["nargs"]="+"
                parser.add_argument(name, help=desc, **args)
            else:
                if spec.is_trait_type(InputMultiPath):
                    args["nargs"]="*"
                parser.add_argument("--%s"%name, dest=name,
                                    help=desc, **args)
    return parser, interface

def run_instance(interface, options):
    if interface:
        print("setting function inputs")

        for input_name, _ in list(interface.inputs.items()):
            if getattr(options, input_name) != None:
                value = getattr(options, input_name)
                if not isinstance(value, bool):
                    #traits cannot cast from string to float or int
                    try:
                        value = float(value)
                    except:
                        pass
                    #try to cast string input to boolean
                    try:
                        value = str2bool(value)
                    except:
                        pass
                try:
                    setattr(interface.inputs, input_name,
                            value)
                except ValueError as e:
                    print("Error when setting the value of %s: '%s'"%(input_name, str(e)))

        print(interface.inputs)
        res = interface.run()
        print(res.outputs)

def main(argv):

    if len(argv) == 2 and not argv[1].startswith("-"):
        listClasses(argv[1])
        sys.exit(0)

    parser = argparse.ArgumentParser(description='Nipype interface runner', prog=argv[0])
    parser.add_argument("module", type=str, help="Module name")
    parser.add_argument("interface", type=str, help="Interface name")
    parsed = parser.parse_args(args=argv[1:3])

    _, prog = os.path.split(argv[0])
    interface_parser = argparse.ArgumentParser(description="Run %s"%parsed.interface, prog=" ".join([prog] + argv[1:3]))
    interface_parser, interface  = add_options(interface_parser, parsed.module, parsed.interface)
    args = interface_parser.parse_args(args=argv[3:])
    run_instance(interface, args)
<commit_msg>fix: utils/nipype_cmd: cosmetics<commit_after>from __future__ import print_function
from __future__ import unicode_literals
from builtins import str
import os
import argparse
import inspect
import sys
from nipype.interfaces.base import Interface, InputMultiPath, traits
from nipype.utils.misc import str2bool

def listClasses(module=None):
    if module:
        __import__(module)
        pkg = sys.modules[module]
        print("Available Interfaces:")
        for k, v in sorted(list(pkg.__dict__.items())):
            if inspect.isclass(v) and issubclass(v, Interface):
                print("\t%s"%k)

def add_options(parser=None, module=None, function=None):
    interface = None
    if parser and module and function:
        __import__(module)
        interface = getattr(sys.modules[module], function)()

        inputs = interface.input_spec()
        for name, spec in sorted(interface.inputs.traits(transient=None).items()):
            desc = "\n".join(interface._get_trait_desc(inputs, name, spec))[len(name)+2:]
            args = {}

            if spec.is_trait_type(traits.Bool):
                args["action"] = 'store_true'

            if hasattr(spec, "mandatory") and spec.mandatory:
                if spec.is_trait_type(InputMultiPath):
                    args["nargs"] = "+"
                parser.add_argument(name, help=desc, **args)
            else:
                if spec.is_trait_type(InputMultiPath):
                    args["nargs"] = "*"
                parser.add_argument("--%s"%name, dest=name,
                                    help=desc, **args)
    return parser, interface

def run_instance(interface, options):
    if interface:
        print("setting function inputs")

        for input_name, _ in list(interface.inputs.items()):
            if getattr(options, input_name) != None:
                value = getattr(options, input_name)
                if not isinstance(value, bool):
                    #traits cannot cast from string to float or int
                    try:
                        value = float(value)
                    except:
                        pass
                    #try to cast string input to boolean
                    try:
                        value = str2bool(value)
                    except:
                        pass
                try:
                    setattr(interface.inputs, input_name,
                            value)
                except ValueError as e:
                    print("Error when setting the value of %s: '%s'"%(input_name, str(e)))

        print(interface.inputs)
        res = interface.run()
        print(res.outputs)

def main(argv):

    if len(argv) == 2 and not argv[1].startswith("-"):
        listClasses(argv[1])
        sys.exit(0)

    parser = argparse.ArgumentParser(description='Nipype interface runner', prog=argv[0])
    parser.add_argument("module", type=str, help="Module name")
    parser.add_argument("interface", type=str, help="Interface name")
    parsed = parser.parse_args(args=argv[1:3])

    _, prog = os.path.split(argv[0])
    interface_parser = argparse.ArgumentParser(description="Run %s"%parsed.interface, prog=" ".join([prog] + argv[1:3]))
    interface_parser, interface  = add_options(interface_parser, parsed.module, parsed.interface)
    args = interface_parser.parse_args(args=argv[3:])
    run_instance(interface, args)
<|endoftext|>


<commit_before><commit_msg>unlink mx_graal.py<commit_after><|endoftext|>


<commit_before># Adapted from https://tensorflow.googlesource.com/tensorflow/+/master/tensorflow/models/image/cifar10/cifar10.py

from datetime import datetime
import math
import time
import numpy as np
import dataset
import tensorflow.python.platform
import tensorflow as tf

batch_size = 8

def conv_op(input_op, name, kw, kh, n_in, n_out, dw, dh):
    with tf.name_scope(name) as scope:
        kernel_init_val = tf.truncated_normal([kh, kw, n_in, n_out], dtype=tf.float32, stddev=1e-1)
        kernel = tf.Variable(kernel_init_val, trainable=True, name='w')
        conv = tf.nn.conv2d(input_op, kernel, (1, dh, dw, 1), padding='SAME')
        bias_init_val = tf.constant(0.0, shape=[n_out], dtype=tf.float32)
        biases = tf.Variable(bias_init_val, trainable=True, name='b')
        z = tf.reshape(tf.nn.bias_add(conv, biases), conv.get_shape())
        activation = tf.nn.relu(z, name=scope)
        return activation

def affine_op(input_op, name, n_in, n_out):
    with tf.name_scope(name) as scope:
        kernel = tf.Variable(tf.truncated_normal([n_in, n_out], dtype=tf.float32, stddev=1e-1), trainable=True, name='weights')
        biases = tf.Variable(tf.constant(0.0, shape=[n_out], dtype=tf.float32), trainable=True, name='biases')
        affine1 = tf.nn.relu_layer(input_op, kernel, biases, name=name)
        return affine1

def mpool_op(input_op, name, kh, kw, dh, dw):
    return tf.nn.max_pool(input_op,
                          ksize=[1, kh, kw, 1],
                          strides=[1, dh, dw, 1],
                          padding='VALID',
                          name=name)

def loss(logits, labels):
    labels = tf.expand_dims(labels, 1)
    indices = tf.expand_dims(tf.range(0, batch_size, 1), 1)
    concated = tf.concat(1, [indices, labels])
    onehot_labels = tf.sparse_to_dense(concated, tf.pack([batch_size, 10]), 1.0, 0.0)
    cross_entropy = tf.nn.softmax_cross_entropy_with_logits(logits, onehot_labels, name='xentropy')
    loss = tf.reduce_mean(cross_entropy, name='xentropy_mean')
    return loss

def inference(input_op):
    conv1 = conv_op(input_op, name="conv1", kh=3, kw=3, n_in=3, n_out=64, dh=1, dw=1)
    pool1 = mpool_op(conv1,   name="pool1", kh=2, kw=2, dw=2, dh=2)
    conv2 = conv_op(pool1,    name="conv2", kh=3, kw=3, n_in=64, n_out=128, dh=1, dw=1)
    pool2 = mpool_op(conv2,   name="pool2", kh=2, kw=2, dh=2, dw=2)
    conv3 = conv_op(pool2,    name="conv3", kh=3, kw=3, n_in=128, n_out=512, dh=1, dw=1)
    conv4 = conv_op(conv3,    name="conv4", kh=3, kw=3, n_in=512, n_out=512, dh=1, dw=1)
    pool4 = mpool_op(conv4,   name="pool4", kh=2, kw=2, dh=2, dw=2)
    conv5 = conv_op(pool4,    name="conv5", kh=3, kw=3, n_in=512, n_out=512, dh=1, dw=1)
    conv6 = conv_op(conv5,    name="conv6", kh=3, kw=3, n_in=512, n_out=512, dh=1, dw=1)
    pool6 = mpool_op(conv6,   name="pool6", kh=2, kw=2, dh=2, dw=2)
    conv7 = conv_op(pool6,    name="conv7", kh=3, kw=3, n_in=512, n_out=512, dh=1, dw=1)
    conv8 = conv_op(conv7,    name="conv8", kh=3, kw=3, n_in=512, n_out=512, dh=1, dw=1)
    pool8 = mpool_op(conv8,   name="pool8", kh=2, kw=2, dw=2, dh=2)
    resh1 = tf.reshape(pool8, [-1,512*7*7], name="resh1")
    affn1 = affine_op(resh1,  name="affn1", n_in=512*7*7, n_out=4096)
    affn2 = affine_op(affn1,  name="affn2", n_in=4096, n_out=10)
    affn3 = affine_op(affn2,  name="affn3", n_in=10, n_out=10)
    return affn3
    # return affn1

def train(lr=0.0001, max_step=1000):
    """
    train model

    :param lr:  This is the learning rate
    """
    with tf.Graph().as_default():
        # Generate some dummy images.
        # image_size = 224
        # images = tf.Variable(tf.random_normal([batch_size, image_size, image_size, 3],
        #                                       dtype=tf.float32,
        #                                       stddev=1))
        #
        # labels = tf.Variable(tf.ones([batch_size], dtype=tf.int32))

        in_images = tf.placeholder("float", [batch_size, 32, 32, 3])
        images = tf.image.resize_images(in_images, 224, 224)
        labels = tf.placeholder("int32", [batch_size])

        # Build a Graph that computes the logits predictions from the
        # inference model.
        last_layer = inference(images)

        # Add a simple objective so we can calculate the backward pass.
        objective = loss(last_layer, labels)
        optimizer = tf.train.AdagradOptimizer(lr)
        global_step = tf.Variable(0, name="global_step", trainable=False)
        train_step = optimizer.minimize(objective)

        # grab variables we want to log
        tf.scalar_summary("loss", objective)
        tf.image_summary("images", images)

        summaries = tf.merge_all_summaries()

        # Build an initialization operation.
        initializer = tf.initialize_all_variables()

        # Start running operations on the Graph.
        with tf.Session() as sess:
            sess.run(initializer)
            writer = tf.train.SummaryWriter("train_logs", graph_def=sess.graph_def)
            for i in range(max_step):
                trn, tst = dataset.get_cifar10(batch_size)
                for batch in trn:
                    X = np.vstack(batch[0]).reshape(-1, 3, 32, 32).transpose(0, 2, 3, 1)
                    Y = np.array(batch[1])
                    result = sess.run(
                        [train_step, summaries, objective],
                        feed_dict = {
                            in_images: X,
                            labels: Y
                        }
                    )
                    writer.add_summary(result[1], i)
                    print i, result[2]
                    if result[2] is np.NaN:
                        return

if __name__ == '__main__':
    train()
<commit_msg>adjust training parameters<commit_after># Adapted from https://tensorflow.googlesource.com/tensorflow/+/master/tensorflow/models/image/cifar10/cifar10.py

from datetime import datetime
import math
import time
import numpy as np
import dataset
import tensorflow.python.platform
import tensorflow as tf

batch_size = 8

def conv_op(input_op, name, kw, kh, n_in, n_out, dw, dh):
    with tf.name_scope(name) as scope:
        kernel_init_val = tf.truncated_normal([kh, kw, n_in, n_out], dtype=tf.float32, stddev=1e-1)
        kernel = tf.Variable(kernel_init_val, trainable=True, name='w')
        conv = tf.nn.conv2d(input_op, kernel, (1, dh, dw, 1), padding='SAME')
        bias_init_val = tf.constant(0.0, shape=[n_out], dtype=tf.float32)
        biases = tf.Variable(bias_init_val, trainable=True, name='b')
        z = tf.reshape(tf.nn.bias_add(conv, biases), conv.get_shape())
        activation = tf.nn.relu(z, name=scope)
        return activation

def affine_op(input_op, name, n_in, n_out):
    with tf.name_scope(name) as scope:
        kernel = tf.Variable(tf.truncated_normal([n_in, n_out], dtype=tf.float32, stddev=1e-1), trainable=True, name='weights')
        biases = tf.Variable(tf.constant(0.0, shape=[n_out], dtype=tf.float32), trainable=True, name='biases')
        affine1 = tf.nn.relu_layer(input_op, kernel, biases, name=name)
        return affine1

def mpool_op(input_op, name, kh, kw, dh, dw):
    return tf.nn.max_pool(input_op,
                          ksize=[1, kh, kw, 1],
                          strides=[1, dh, dw, 1],
                          padding='VALID',
                          name=name)

def loss(logits, labels):
    labels = tf.expand_dims(labels, 1)
    indices = tf.expand_dims(tf.range(0, batch_size, 1), 1)
    concated = tf.concat(1, [indices, labels])
    onehot_labels = tf.sparse_to_dense(concated, tf.pack([batch_size, 10]), 1.0, 0.0)
    cross_entropy = tf.nn.softmax_cross_entropy_with_logits(logits, onehot_labels, name='xentropy')
    loss = tf.reduce_mean(cross_entropy, name='xentropy_mean')
    return loss

def inference(input_op):
    conv1 = conv_op(input_op, name="conv1", kh=3, kw=3, n_in=3, n_out=64, dh=1, dw=1)
    pool1 = mpool_op(conv1,   name="pool1", kh=2, kw=2, dw=2, dh=2)
    conv2 = conv_op(pool1,    name="conv2", kh=3, kw=3, n_in=64, n_out=128, dh=1, dw=1)
    pool2 = mpool_op(conv2,   name="pool2", kh=2, kw=2, dh=2, dw=2)
    conv3 = conv_op(pool2,    name="conv3", kh=3, kw=3, n_in=128, n_out=512, dh=1, dw=1)
    conv4 = conv_op(conv3,    name="conv4", kh=3, kw=3, n_in=512, n_out=512, dh=1, dw=1)
    pool4 = mpool_op(conv4,   name="pool4", kh=2, kw=2, dh=2, dw=2)
    conv5 = conv_op(pool4,    name="conv5", kh=3, kw=3, n_in=512, n_out=512, dh=1, dw=1)
    conv6 = conv_op(conv5,    name="conv6", kh=3, kw=3, n_in=512, n_out=512, dh=1, dw=1)
    pool6 = mpool_op(conv6,   name="pool6", kh=2, kw=2, dh=2, dw=2)
    conv7 = conv_op(pool6,    name="conv7", kh=3, kw=3, n_in=512, n_out=512, dh=1, dw=1)
    conv8 = conv_op(conv7,    name="conv8", kh=3, kw=3, n_in=512, n_out=512, dh=1, dw=1)
    pool8 = mpool_op(conv8,   name="pool8", kh=2, kw=2, dw=2, dh=2)
    resh1 = tf.reshape(pool8, [-1,512*7*7], name="resh1")
    affn1 = affine_op(resh1,  name="affn1", n_in=512*7*7, n_out=4096)
    affn2 = affine_op(affn1,  name="affn2", n_in=4096, n_out=10)
    affn3 = affine_op(affn2,  name="affn3", n_in=10, n_out=10)
    return affn3
    # return affn1

def train(lr=0.00001, max_step=1000):
    """
    train model

    :param lr:  This is the learning rate
    """
    with tf.Graph().as_default():
        # Generate some dummy images.
        # image_size = 224
        # images = tf.Variable(tf.random_normal([batch_size, image_size, image_size, 3],
        #                                       dtype=tf.float32,
        #                                       stddev=1))
        #
        # labels = tf.Variable(tf.ones([batch_size], dtype=tf.int32))

        in_images = tf.placeholder("float", [batch_size, 32, 32, 3])
        images = tf.image.resize_images(in_images, 224, 224)
        labels = tf.placeholder("int32", [batch_size])

        # Build a Graph that computes the logits predictions from the
        # inference model.
        last_layer = inference(images)

        # Add a simple objective so we can calculate the backward pass.
        objective = loss(last_layer, labels)
        optimizer = tf.train.AdagradOptimizer(lr)
        global_step = tf.Variable(0, name="global_step", trainable=False)
        train_step = optimizer.minimize(objective)

        # grab variables we want to log
        tf.scalar_summary("loss", objective)
        tf.image_summary("images", images)

        summaries = tf.merge_all_summaries()

        # Build an initialization operation.
        initializer = tf.initialize_all_variables()

        # Start running operations on the Graph.
        with tf.Session() as sess:
            sess.run(initializer)
            writer = tf.train.SummaryWriter("train_logs", graph_def=sess.graph_def)
            for i in range(max_step):
                trn, tst = dataset.get_cifar10(batch_size)
                for batch in trn:
                    X = np.vstack(batch[0]).reshape(-1, 3, 32, 32).transpose(0, 2, 3, 1)
                    Y = np.array(batch[1])
                    result = sess.run(
                        [train_step, summaries, objective],
                        feed_dict = {
                            in_images: X,
                            labels: Y
                        }
                    )
                    writer.add_summary(result[1], i)
                    print i, result[2]
                    if result[2] is np.NaN:
                        return

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


<commit_before>'''
The spec util
Handles the Lab experiment spec: reading, writing(evolution), validation and default setting
Expands the spec and params into consumable inputs in info space for lab units.
'''
from slm_lab.lib import logger, util
from string import Template
import itertools
import json
import numpy as np
import os
import pydash as ps

SPEC_DIR = 'slm_lab/spec'
'''
All spec values are already param, inferred automatically.
To change from a value into param range, e.g.
- single: "explore_anneal_epi": 50
- continuous param: "explore_anneal_epi": {"min": 50, "max": 100, "dist": "uniform"}
- discrete range: "explore_anneal_epi": {"values": [50, 75, 100]}
'''
SPEC_FORMAT = {
    "agent": [{
        "name": str,
        "algorithm": dict,
        "memory": dict,
        "net": dict,
    }],
    "env": [{
        "name": str,
        "max_t": (type(None), int, float),
        "max_tick": (int, float),
    }],
    "body": {
        "product": ["outer", "inner", "custom"],
        "num": (int, list),
    },
    "meta": {
        "distributed": bool,
        "eval_frequency": (int, float),
        "max_tick_unit": str,
        "max_session": int,
        "max_trial": (type(None), int),
    },
    "name": str,
}
logger = logger.get_logger(__name__)

def check_comp_spec(comp_spec, comp_spec_format):
    '''Base method to check component spec'''
    for spec_k, spec_format_v in comp_spec_format.items():
        comp_spec_v = comp_spec[spec_k]
        if ps.is_list(spec_format_v):
            v_set = spec_format_v
            assert comp_spec_v in v_set, f'Component spec value {ps.pick(comp_spec, spec_k)} needs to be one of {util.to_json(v_set)}'
        else:
            v_type = spec_format_v
            assert isinstance(comp_spec_v, v_type), f'Component spec {ps.pick(comp_spec, spec_k)} needs to be of type: {v_type}'
            if isinstance(v_type, tuple) and int in v_type and isinstance(comp_spec_v, float):
                # cast if it can be int
                comp_spec[spec_k] = int(comp_spec_v)

def check_body_spec(spec):
    '''Base method to check body spec for AEB space resolution'''
    ae_product = ps.get(spec, 'body.product')
    body_num = ps.get(spec, 'body.num')
    if ae_product == 'outer':
        pass
    elif ae_product == 'inner':
        agent_num = len(spec['agent'])
        env_num = len(spec['env'])
        assert agent_num == env_num, 'Agent and Env spec length must be equal for body `inner` product. Given {agent_num}, {env_num}'
    else:  # custom AEB
        assert ps.is_list(body_num)

def check(spec):
    '''Check a single spec for validity'''
    try:
        spec_name = spec.get('name')
        assert set(spec.keys()) >= set(SPEC_FORMAT.keys()), f'Spec needs to follow spec.SPEC_FORMAT. Given \n {spec_name}: {util.to_json(spec)}'
        for agent_spec in spec['agent']:
            check_comp_spec(agent_spec, SPEC_FORMAT['agent'][0])
        for env_spec in spec['env']:
            check_comp_spec(env_spec, SPEC_FORMAT['env'][0])
        check_comp_spec(spec['body'], SPEC_FORMAT['body'])
        check_comp_spec(spec['meta'], SPEC_FORMAT['meta'])
        check_body_spec(spec)
    except Exception as e:
        logger.exception(f'spec {spec_name} fails spec check')
        raise e
    return True

def check_all():
    '''Check all spec files, all specs.'''
    spec_files = ps.filter_(os.listdir(SPEC_DIR), lambda f: f.endswith('.json') and not f.startswith('_'))
    for spec_file in spec_files:
        spec_dict = util.read(f'{SPEC_DIR}/{spec_file}')
        for spec_name, spec in spec_dict.items():
            try:
                spec['name'] = spec_name
                spec['git_SHA'] = util.get_git_sha()
                check(spec)
            except Exception as e:
                logger.exception(f'spec_file {spec_file} fails spec check')
                raise e
    logger.info(f'Checked all specs from: {ps.join(spec_files, ",")}')
    return True

def get(spec_file, spec_name):
    '''
    Get an experiment spec from spec_file, spec_name.
    Auto-check spec.
    @example

    spec = spec_util.get('base.json', 'base_case_openai')
    '''
    spec_file = spec_file.replace(SPEC_DIR, '')  # cleanup
    if 'data/' in spec_file:
        assert spec_name in spec_file, 'spec_file in data/ must be lab-generated and contains spec_name'
        spec = util.read(spec_file)
    else:
        spec_file = f'{SPEC_DIR}/{spec_file}'  # allow direct filename
        spec_dict = util.read(spec_file)
        assert spec_name in spec_dict, f'spec_name {spec_name} is not in spec_file {spec_file}. Choose from:\n {ps.join(spec_dict.keys(), ",")}'
        spec = spec_dict[spec_name]
        spec['name'] = spec_name
        spec['git_SHA'] = util.get_git_sha()
    check(spec)
    return spec

def get_param_specs(spec):
    '''Return a list of specs with substituted spec_params'''
    assert 'spec_params' in spec, 'Parametrized spec needs a spec_params key'
    spec_params = spec.pop('spec_params')
    spec_template = Template(json.dumps(spec))
    keys = spec_params.keys()
    specs = []
    for vals in itertools.product(*spec_params.values()):
        spec_str = spec_template.substitute(dict(zip(keys, vals)))
        spec = json.loads(spec_str)
        spec['name'] += f'_{"_".join(vals)}'
        specs.append(spec)
    return specs

def is_aeb_compact(aeb_list):
    '''
    Check if aeb space (aeb_list) is compact; uniq count must equal shape in each of a,e axes. For b, per unique a,e hash, uniq must equal shape.'''
    aeb_shape = util.get_aeb_shape(aeb_list)
    aeb_uniq = [len(np.unique(col)) for col in np.transpose(aeb_list)]
    ae_compact = np.array_equal(aeb_shape, aeb_uniq)
    b_compact = True
    for ae, ae_b_list in ps.group_by(aeb_list, lambda aeb: f'{aeb[0]}{aeb[1]}').items():
        b_shape = util.get_aeb_shape(ae_b_list)[2]
        b_uniq = [len(np.unique(col)) for col in np.transpose(ae_b_list)][2]
        b_compact = b_compact and np.array_equal(b_shape, b_uniq)
    aeb_compact = ae_compact and b_compact
    return aeb_compact

def is_singleton(spec):
    '''Check if spec uses a singleton Session'''
    return len(spec['agent']) == 1 and len(spec['env']) == 1 and spec['body']['num'] == 1

def override_dev_spec(spec):
    spec['meta']['max_session'] = 1
    spec['meta']['max_trial'] = 2
    return spec

def override_enjoy_spec(spec):
    spec['meta']['max_session'] = 1
    return spec

def override_eval_spec(spec):
    for agent_spec in spec['agent']:
        if 'max_size' in agent_spec['memory']:
            agent_spec['memory']['max_size'] = 100
    # evaluate by episode is set in env clock init in env/base.py
    return spec

def override_test_spec(spec):
    for agent_spec in spec['agent']:
        # covers episodic and timestep
        agent_spec['algorithm']['training_frequency'] = 1
        agent_spec['algorithm']['training_start_step'] = 1
        agent_spec['algorithm']['training_epoch'] = 1
        agent_spec['algorithm']['training_batch_epoch'] = 1
    for env_spec in spec['env']:
        env_spec['max_t'] = 20
        env_spec['max_tick'] = 3
    spec['meta']['eval_frequency'] = 1000
    spec['meta']['max_tick_unit'] = 'epi'
    spec['meta']['max_session'] = 1
    spec['meta']['max_trial'] = 2
    return spec

def resolve_aeb(spec):
    '''
    Resolve an experiment spec into the full list of points (coordinates) in AEB space.
    @param {dict} spec An experiment spec.
    @returns {list} aeb_list Resolved array of points in AEB space.
    @example

    spec = spec_util.get('base.json', 'general_inner')
    aeb_list = spec_util.resolve_aeb(spec)
    # => [(0, 0, 0), (0, 0, 1), (1, 1, 0), (1, 1, 1)]
    '''
    agent_num = len(spec['agent']) if ps.is_list(spec['agent']) else 1
    env_num = len(spec['env']) if ps.is_list(spec['env']) else 1
    ae_product = ps.get(spec, 'body.product')
    body_num = ps.get(spec, 'body.num')
    body_num_list = body_num if ps.is_list(body_num) else [body_num] * env_num

    aeb_list = []
    if ae_product == 'outer':
        for e in range(env_num):
            sub_aeb_list = list(itertools.product(range(agent_num), [e], range(body_num_list[e])))
            aeb_list.extend(sub_aeb_list)
    elif ae_product == 'inner':
        for a, e in zip(range(agent_num), range(env_num)):
            sub_aeb_list = list(itertools.product([a], [e], range(body_num_list[e])))
            aeb_list.extend(sub_aeb_list)
    else:  # custom AEB, body_num is a aeb_list
        aeb_list = [tuple(aeb) for aeb in body_num]
    aeb_list.sort()
    assert is_aeb_compact(aeb_list), 'Failed check: for a, e, uniq count == len (shape), and for each a,e hash, b uniq count == b len (shape)'
    return aeb_list
<commit_msg>put git_SHA into meta as git_sha<commit_after>'''
The spec util
Handles the Lab experiment spec: reading, writing(evolution), validation and default setting
Expands the spec and params into consumable inputs in info space for lab units.
'''
from slm_lab.lib import logger, util
from string import Template
import itertools
import json
import numpy as np
import os
import pydash as ps

SPEC_DIR = 'slm_lab/spec'
'''
All spec values are already param, inferred automatically.
To change from a value into param range, e.g.
- single: "explore_anneal_epi": 50
- continuous param: "explore_anneal_epi": {"min": 50, "max": 100, "dist": "uniform"}
- discrete range: "explore_anneal_epi": {"values": [50, 75, 100]}
'''
SPEC_FORMAT = {
    "agent": [{
        "name": str,
        "algorithm": dict,
        "memory": dict,
        "net": dict,
    }],
    "env": [{
        "name": str,
        "max_t": (type(None), int, float),
        "max_tick": (int, float),
    }],
    "body": {
        "product": ["outer", "inner", "custom"],
        "num": (int, list),
    },
    "meta": {
        "distributed": bool,
        "eval_frequency": (int, float),
        "max_tick_unit": str,
        "max_session": int,
        "max_trial": (type(None), int),
    },
    "name": str,
}
logger = logger.get_logger(__name__)

def check_comp_spec(comp_spec, comp_spec_format):
    '''Base method to check component spec'''
    for spec_k, spec_format_v in comp_spec_format.items():
        comp_spec_v = comp_spec[spec_k]
        if ps.is_list(spec_format_v):
            v_set = spec_format_v
            assert comp_spec_v in v_set, f'Component spec value {ps.pick(comp_spec, spec_k)} needs to be one of {util.to_json(v_set)}'
        else:
            v_type = spec_format_v
            assert isinstance(comp_spec_v, v_type), f'Component spec {ps.pick(comp_spec, spec_k)} needs to be of type: {v_type}'
            if isinstance(v_type, tuple) and int in v_type and isinstance(comp_spec_v, float):
                # cast if it can be int
                comp_spec[spec_k] = int(comp_spec_v)

def check_body_spec(spec):
    '''Base method to check body spec for AEB space resolution'''
    ae_product = ps.get(spec, 'body.product')
    body_num = ps.get(spec, 'body.num')
    if ae_product == 'outer':
        pass
    elif ae_product == 'inner':
        agent_num = len(spec['agent'])
        env_num = len(spec['env'])
        assert agent_num == env_num, 'Agent and Env spec length must be equal for body `inner` product. Given {agent_num}, {env_num}'
    else:  # custom AEB
        assert ps.is_list(body_num)

def check(spec):
    '''Check a single spec for validity'''
    try:
        spec_name = spec.get('name')
        assert set(spec.keys()) >= set(SPEC_FORMAT.keys()), f'Spec needs to follow spec.SPEC_FORMAT. Given \n {spec_name}: {util.to_json(spec)}'
        for agent_spec in spec['agent']:
            check_comp_spec(agent_spec, SPEC_FORMAT['agent'][0])
        for env_spec in spec['env']:
            check_comp_spec(env_spec, SPEC_FORMAT['env'][0])
        check_comp_spec(spec['body'], SPEC_FORMAT['body'])
        check_comp_spec(spec['meta'], SPEC_FORMAT['meta'])
        check_body_spec(spec)
    except Exception as e:
        logger.exception(f'spec {spec_name} fails spec check')
        raise e
    return True

def check_all():
    '''Check all spec files, all specs.'''
    spec_files = ps.filter_(os.listdir(SPEC_DIR), lambda f: f.endswith('.json') and not f.startswith('_'))
    for spec_file in spec_files:
        spec_dict = util.read(f'{SPEC_DIR}/{spec_file}')
        for spec_name, spec in spec_dict.items():
            try:
                check(spec)
            except Exception as e:
                logger.exception(f'spec_file {spec_file} fails spec check')
                raise e
    logger.info(f'Checked all specs from: {ps.join(spec_files, ",")}')
    return True

def get(spec_file, spec_name):
    '''
    Get an experiment spec from spec_file, spec_name.
    Auto-check spec.
    @example

    spec = spec_util.get('base.json', 'base_case_openai')
    '''
    spec_file = spec_file.replace(SPEC_DIR, '')  # cleanup
    if 'data/' in spec_file:
        assert spec_name in spec_file, 'spec_file in data/ must be lab-generated and contains spec_name'
        spec = util.read(spec_file)
    else:
        spec_file = f'{SPEC_DIR}/{spec_file}'  # allow direct filename
        spec_dict = util.read(spec_file)
        assert spec_name in spec_dict, f'spec_name {spec_name} is not in spec_file {spec_file}. Choose from:\n {ps.join(spec_dict.keys(), ",")}'
        spec = spec_dict[spec_name]
        spec['name'] = spec_name
        spec['meta']['git_sha'] = util.get_git_sha()
    check(spec)
    return spec

def get_param_specs(spec):
    '''Return a list of specs with substituted spec_params'''
    assert 'spec_params' in spec, 'Parametrized spec needs a spec_params key'
    spec_params = spec.pop('spec_params')
    spec_template = Template(json.dumps(spec))
    keys = spec_params.keys()
    specs = []
    for vals in itertools.product(*spec_params.values()):
        spec_str = spec_template.substitute(dict(zip(keys, vals)))
        spec = json.loads(spec_str)
        spec['name'] += f'_{"_".join(vals)}'
        specs.append(spec)
    return specs

def is_aeb_compact(aeb_list):
    '''
    Check if aeb space (aeb_list) is compact; uniq count must equal shape in each of a,e axes. For b, per unique a,e hash, uniq must equal shape.'''
    aeb_shape = util.get_aeb_shape(aeb_list)
    aeb_uniq = [len(np.unique(col)) for col in np.transpose(aeb_list)]
    ae_compact = np.array_equal(aeb_shape, aeb_uniq)
    b_compact = True
    for ae, ae_b_list in ps.group_by(aeb_list, lambda aeb: f'{aeb[0]}{aeb[1]}').items():
        b_shape = util.get_aeb_shape(ae_b_list)[2]
        b_uniq = [len(np.unique(col)) for col in np.transpose(ae_b_list)][2]
        b_compact = b_compact and np.array_equal(b_shape, b_uniq)
    aeb_compact = ae_compact and b_compact
    return aeb_compact

def is_singleton(spec):
    '''Check if spec uses a singleton Session'''
    return len(spec['agent']) == 1 and len(spec['env']) == 1 and spec['body']['num'] == 1

def override_dev_spec(spec):
    spec['meta']['max_session'] = 1
    spec['meta']['max_trial'] = 2
    return spec

def override_enjoy_spec(spec):
    spec['meta']['max_session'] = 1
    return spec

def override_eval_spec(spec):
    for agent_spec in spec['agent']:
        if 'max_size' in agent_spec['memory']:
            agent_spec['memory']['max_size'] = 100
    # evaluate by episode is set in env clock init in env/base.py
    return spec

def override_test_spec(spec):
    for agent_spec in spec['agent']:
        # covers episodic and timestep
        agent_spec['algorithm']['training_frequency'] = 1
        agent_spec['algorithm']['training_start_step'] = 1
        agent_spec['algorithm']['training_epoch'] = 1
        agent_spec['algorithm']['training_batch_epoch'] = 1
    for env_spec in spec['env']:
        env_spec['max_t'] = 20
        env_spec['max_tick'] = 3
    spec['meta']['eval_frequency'] = 1000
    spec['meta']['max_tick_unit'] = 'epi'
    spec['meta']['max_session'] = 1
    spec['meta']['max_trial'] = 2
    return spec

def resolve_aeb(spec):
    '''
    Resolve an experiment spec into the full list of points (coordinates) in AEB space.
    @param {dict} spec An experiment spec.
    @returns {list} aeb_list Resolved array of points in AEB space.
    @example

    spec = spec_util.get('base.json', 'general_inner')
    aeb_list = spec_util.resolve_aeb(spec)
    # => [(0, 0, 0), (0, 0, 1), (1, 1, 0), (1, 1, 1)]
    '''
    agent_num = len(spec['agent']) if ps.is_list(spec['agent']) else 1
    env_num = len(spec['env']) if ps.is_list(spec['env']) else 1
    ae_product = ps.get(spec, 'body.product')
    body_num = ps.get(spec, 'body.num')
    body_num_list = body_num if ps.is_list(body_num) else [body_num] * env_num

    aeb_list = []
    if ae_product == 'outer':
        for e in range(env_num):
            sub_aeb_list = list(itertools.product(range(agent_num), [e], range(body_num_list[e])))
            aeb_list.extend(sub_aeb_list)
    elif ae_product == 'inner':
        for a, e in zip(range(agent_num), range(env_num)):
            sub_aeb_list = list(itertools.product([a], [e], range(body_num_list[e])))
            aeb_list.extend(sub_aeb_list)
    else:  # custom AEB, body_num is a aeb_list
        aeb_list = [tuple(aeb) for aeb in body_num]
    aeb_list.sort()
    assert is_aeb_compact(aeb_list), 'Failed check: for a, e, uniq count == len (shape), and for each a,e hash, b uniq count == b len (shape)'
    return aeb_list
<|endoftext|>


<commit_before>import logging
import os

import docker
from docker.utils import kwargs_from_env

from vdist import defaults

class BuildMachine(object):

    def __init__(self, machine_logs=True, image=None, insecure_registry=False):
        self.logger = logging.getLogger('BuildMachine')

        self.dockerclient = docker.Client(version='1.15', **kwargs_from_env())
        self.container = None

        self.machine_logs = machine_logs
        self.image = image

        self.insecure_registry = insecure_registry

    def _image_exists(self, image):
        for img in self.dockerclient.images():
            if image in img['RepoTags']:
                return True
        return False

    def _pull_image(self, image):
        self.logger.info('are we using an insecure registry? %s' % self.insecure_registry)
        self.dockerclient.pull(
            image,
            insecure_registry=self.insecure_registry)

    def launch(self, build_dir, extra_binds=None):
        if not self._image_exists(self.image):
            self.logger.info(
                'Image does not exist: %s, pulling from repo..' % self.image)
            self._pull_image(self.image)

        binds = {build_dir: defaults.SHARED_DIR}
        if extra_binds:
            binds = binds.items() + extra_binds.items()

        path_to_command = os.path.join(
            defaults.SHARED_DIR,
            defaults.SCRATCH_DIR,
            defaults.SCRATCH_BUILDSCRIPT_NAME
        )
        self.logger.info('Starting container: %s' % self.image)
        self.container = self.dockerclient.create_container(
            image=self.image,
            command=path_to_command)
        self.dockerclient.start(
            container=self.container.get('Id'),
            binds=binds)
        lines = self.dockerclient.logs(container=self.container.get('Id'),
                                       stdout=True, stderr=True, stream=True)
        for line in lines:
            if self.machine_logs:
                self.logger.info(line.strip())

    def shutdown(self):
        self.dockerclient.stop(container=self.container.get('Id'))
        self.dockerclient.remove_container(container=self.container.get('Id'))
<commit_msg>formatting fixes<commit_after>import logging
import os

import docker
from docker.utils import kwargs_from_env

from vdist import defaults

class BuildMachine(object):

    def __init__(self, machine_logs=True, image=None, insecure_registry=False):
        self.logger = logging.getLogger('BuildMachine')

        self.dockerclient = docker.Client(version='1.15', **kwargs_from_env())
        self.container = None

        self.machine_logs = machine_logs
        self.image = image

        self.insecure_registry = insecure_registry

    def _image_exists(self, image):
        for img in self.dockerclient.images():
            if image in img['RepoTags']:
                return True
        return False

    def _pull_image(self, image):
        self.logger.info('Insecure registry: %s' % self.insecure_registry)

        self.dockerclient.pull(image, insecure_registry=self.insecure_registry)

    def launch(self, build_dir, extra_binds=None):
        if not self._image_exists(self.image):
            self.logger.info(
                'Image does not exist: %s, pulling from repo..' % self.image)
            self._pull_image(self.image)

        binds = {build_dir: defaults.SHARED_DIR}
        if extra_binds:
            binds = binds.items() + extra_binds.items()

        path_to_command = os.path.join(
            defaults.SHARED_DIR,
            defaults.SCRATCH_DIR,
            defaults.SCRATCH_BUILDSCRIPT_NAME
        )
        self.logger.info('Starting container: %s' % self.image)
        self.container = self.dockerclient.create_container(
            image=self.image,
            command=path_to_command)

        self.dockerclient.start(
            container=self.container.get('Id'),
            binds=binds)

        lines = self.dockerclient.logs(container=self.container.get('Id'),
                                       stdout=True, stderr=True, stream=True)
        for line in lines:
            if self.machine_logs:
                self.logger.info(line.strip())

    def shutdown(self):
        self.dockerclient.stop(container=self.container.get('Id'))
        self.dockerclient.remove_container(container=self.container.get('Id'))
<|endoftext|>


<commit_before>
from bisect import insort
import collections
import hashlib
import logging
import re

from ..uri import Path
from . import core

log = logging.getLogger(__name__)

class FormatError(ValueError):
    pass

class Pattern(object):
    """
    >>> r = Pattern(r'/{controller}/{action}/{id:\d+}')
    >>> r
    <Pattern:r'/{controller}/{action}/{id:\d+}'>

    >>> r.match('/gallery/photo/12')[0]
    {'action': 'photo', 'controller': 'gallery', 'id': '12'}

    >>> r.match('/controller/action/12extra')
    >>> r.match('/controller/action/not_digits')

    >>> r.match('/controller/action/12/extra')
    ({'action': 'action', 'controller': 'controller', 'id': '12'}, '/extra')

    >>> r.format(controller='news', action='archive', id='24')
    '/news/archive/24'

    >>> r = Pattern('/gallery/{action}', controller='gallery')
    >>> m = r.match('/gallery/view')
    >>> m[0]
    {'action': 'view', 'controller': 'gallery'}
    >>> r.format(**m[0])
    '/gallery/view'

    >>> r = Pattern('/{id}', _requirements={'id': r'\d+'})
    >>> r.match('/12')[0]
    {'id': '12'}
    >>> r.match('/hi')

    >>> r = Pattern('/{id:\d+}', _parsers={'id':int})
    >>> r.match('/12')[0]
    {'id': 12}

    >>> r = Pattern('/{method:[A-Z]+}', _formatters={'method': str.upper})
    >>> r.match('/GET')[0]
    {'method': 'GET'}
    >>> r.format(method='post')
    '/POST'

    >>> r = Pattern('/{id:\d+}', _parsers=dict(id=int),
    ...     _formatters=dict(id='%04d'))
    >>> r.match('/123')[0]
    {'id': 123}
    >>> r.format(id=456)
    '/0456'

    >>> r = Pattern('/{action:get}/{id:\d+}', _formatters=dict(id=int))
    >>> r.match('/get/12')[0]
    {'action': 'get', 'id': '12'}
    >>> r.format(action='test', id=4)
    Traceback (most recent call last):
    ...
    FormatError: cannot match against output

    """

    default_pattern = '[^/]+'
    token_re = re.compile(r'''
        {                           # start of keyword match
        ([a-zA-Z_]\w*)              # group 1: name
        (?::(                       # colon and group 2: pattern
        [^}\\]*(?:\\.[^}\\]*)*      # zero or more chars. } can be escaped.
        ))?                         # the colon and pattern are optional
        }
    ''', re.X)

    def __init__(self, raw, **kwargs):

        self._raw = raw
        self._constants = kwargs
        self._keys = set()

        self._requirements = kwargs.pop('_requirements', {})
        self._requirements = dict((k, re.compile(v + '$'))
            for k, v in self._requirements.items())

        self._parsers = kwargs.pop('_parsers', {})
        self._formatters = kwargs.pop('_formatters', {})

        self._compile()

    def __repr__(self):
        return '<%s:r%s>' % (self.__class__.__name__,
            repr(self._raw).replace('\\\\', '\\'))

    def _compile(self):

        self._hash_to_key = {}
        self._hash_to_pattern = {}

        format = self.token_re.sub(self._compile_sub, self._raw)

        pattern = re.escape(format)
        for hash, patt in self._hash_to_pattern.items():
            pattern = pattern.replace(hash, patt, 1)

        for hash, key in self._hash_to_key.items():
            format = format.replace(hash, '%%(%s)s' % key, 1)

        self._format = format
        self._compiled = re.compile(pattern + r'(?=/|$)')

        del self._hash_to_key
        del self._hash_to_pattern

    def _compile_sub(self, match):
        name = match.group(1)
        self._keys.add(name)
        patt = match.group(2) or self.default_pattern
        hash = 'x%s' % hashlib.md5(name).hexdigest()
        self._hash_to_key[hash] = name
        self._hash_to_pattern[hash] = '(?P<%s>%s)' % (name, patt)
        return hash

    def match(self, value):
        """Match this pattern against some text. Returns the matched data, and
        the unmatched string, or None if there is no match.
        """

        m = self._compiled.match(value)
        if not m:
            return None

        result = self._constants.copy()
        result.update(m.groupdict())

        for key, pattern in self._requirements.items():
            if not key in result or not pattern.match(result[key]):
                return None

        self._parse_data(result)

        return result, value[m.end():]

    def _parse_data(self, data):
        for key, callback in self._parsers.items():
            if key in data:
                data[key] = callback(data[key])

    def _format_data(self, data):
        for key, formatter in self._formatters.items():
            if key in data:
                if isinstance(formatter, basestring):
                    data[key] = formatter % data[key]
                else:
                    data[key] = formatter(data[key])

    def format(self, **kwargs):
        data = self._constants.copy()
        data.update(kwargs)
        self._format_data(data)

        out = self._format % data

        x = self.match(out)
        if x is None:
            raise FormatError('cannot match against output')
        m, d = x
        if d:
            raise FormatError('did not match all output')

        self._parse_data(data)

        for k, v in m.iteritems():
            if k in data and data[k] != v:
                raise FormatError('got different value for %r: got %r, expected %r' % (k, v, data[k]))

        return out

class Match(collections.Mapping):

    def __init__(self, pattern, data, unmatched):
        self.pattern = pattern
        self.data = data
        self.unmatched = unmatched

    def __getitem__(self, key):
        return self.data[key]

    def __getattr__(self, name):
        return getattr(self.data, name)

    def __iter__(self):
        return iter(self.data)

    def __len__(self):
        return len(self.data)

    def format(self, **kwargs):
        data = self.data.copy()
        data.update(kwargs)
        return self.pattern.format(**data)

class ReRouter(core.Router):

    def __init__(self):
        self._apps = []

    def register(self, pattern, app=None, **kwargs):
        """Register directly, or use as a decorator.

        Params:
            pattern -- The pattern to match with. Should start with a '/'.
            app -- The app to register. If not provided this method returns
                a decorator which can be used to register with.

        """

        # We are being used directly here.
        if app:
            # We are creating a key here that will first respect the requested
            # priority of apps relative to each other, but in the case of
            # multiple apps at the same priority, respect the registration
            # order.
            priority = (-kwargs.pop('_priority', 0), len(self._apps))
            insort(self._apps, (priority, Pattern(pattern, **kwargs), app))
            return app

        # We are not being used directly, so return a decorator to do the
        # work later.
        def ReRouter_register(app):
            self.register(pattern, app, **kwargs)
            return app
        return ReRouter_register

    def route_step(self, path):
        for _, pattern, app in self._apps:
            m = pattern.match(path)
            if m:
                kwargs, path = m
                return core.RoutingStep(next=app, path=path, data=kwargs)

    def generate_step(self, data):
        log.debug('generate_step(%r, %r)' % (self, data))
        for _, pattern, app in self._apps:
            if any(k in data and data[k] != v for k, v in
                pattern._constants.iteritems()):
                continue
            try:
                return core.GenerationStep(segment=pattern.format(**data), next=app)
            except FormatError:
                pass
                # log.exception('FormatError while generating')
            except KeyError:
                pass
                # log.exception('KeyError while generating')

<commit_msg>ReRouter will not generate to anonymous paths.<commit_after>
from bisect import insort
import collections
import hashlib
import logging
import re

from ..uri import Path
from . import core

log = logging.getLogger(__name__)

class FormatError(ValueError):
    pass

class Pattern(object):
    """
    >>> r = Pattern(r'/{controller}/{action}/{id:\d+}')
    >>> r
    <Pattern:r'/{controller}/{action}/{id:\d+}'>

    >>> r.match('/gallery/photo/12')[0]
    {'action': 'photo', 'controller': 'gallery', 'id': '12'}

    >>> r.match('/controller/action/12extra')
    >>> r.match('/controller/action/not_digits')

    >>> r.match('/controller/action/12/extra')
    ({'action': 'action', 'controller': 'controller', 'id': '12'}, '/extra')

    >>> r.format(controller='news', action='archive', id='24')
    '/news/archive/24'

    >>> r = Pattern('/gallery/{action}', controller='gallery')
    >>> m = r.match('/gallery/view')
    >>> m[0]
    {'action': 'view', 'controller': 'gallery'}
    >>> r.format(**m[0])
    '/gallery/view'

    >>> r = Pattern('/{id}', _requirements={'id': r'\d+'})
    >>> r.match('/12')[0]
    {'id': '12'}
    >>> r.match('/hi')

    >>> r = Pattern('/{id:\d+}', _parsers={'id':int})
    >>> r.match('/12')[0]
    {'id': 12}

    >>> r = Pattern('/{method:[A-Z]+}', _formatters={'method': str.upper})
    >>> r.match('/GET')[0]
    {'method': 'GET'}
    >>> r.format(method='post')
    '/POST'

    >>> r = Pattern('/{id:\d+}', _parsers=dict(id=int),
    ...     _formatters=dict(id='%04d'))
    >>> r.match('/123')[0]
    {'id': 123}
    >>> r.format(id=456)
    '/0456'

    >>> r = Pattern('/{action:get}/{id:\d+}', _formatters=dict(id=int))
    >>> r.match('/get/12')[0]
    {'action': 'get', 'id': '12'}
    >>> r.format(action='test', id=4)
    Traceback (most recent call last):
    ...
    FormatError: cannot match against output

    """

    default_pattern = '[^/]+'
    token_re = re.compile(r'''
        {                           # start of keyword match
        ([a-zA-Z_]\w*)              # group 1: name
        (?::(                       # colon and group 2: pattern
        [^}\\]*(?:\\.[^}\\]*)*      # zero or more chars. } can be escaped.
        ))?                         # the colon and pattern are optional
        }
    ''', re.X)

    def __init__(self, raw, **kwargs):

        self._raw = raw
        self._constants = kwargs
        self._keys = set()

        self._requirements = kwargs.pop('_requirements', {})
        self._requirements = dict((k, re.compile(v + '$'))
            for k, v in self._requirements.items())

        self._parsers = kwargs.pop('_parsers', {})
        self._formatters = kwargs.pop('_formatters', {})

        self._compile()

    def __repr__(self):
        return '<%s:r%s>' % (self.__class__.__name__,
            repr(self._raw).replace('\\\\', '\\'))

    def _compile(self):

        self._hash_to_key = {}
        self._hash_to_pattern = {}

        format = self.token_re.sub(self._compile_sub, self._raw)

        pattern = re.escape(format)
        for hash, patt in self._hash_to_pattern.items():
            pattern = pattern.replace(hash, patt, 1)

        for hash, key in self._hash_to_key.items():
            format = format.replace(hash, '%%(%s)s' % key, 1)

        self._format = format
        self._compiled = re.compile(pattern + r'(?=/|$)')

        del self._hash_to_key
        del self._hash_to_pattern

    def _compile_sub(self, match):
        name = match.group(1)
        self._keys.add(name)
        patt = match.group(2) or self.default_pattern
        hash = 'x%s' % hashlib.md5(name).hexdigest()
        self._hash_to_key[hash] = name
        self._hash_to_pattern[hash] = '(?P<%s>%s)' % (name, patt)
        return hash

    def match(self, value):
        """Match this pattern against some text. Returns the matched data, and
        the unmatched string, or None if there is no match.
        """

        m = self._compiled.match(value)
        if not m:
            return None

        result = self._constants.copy()
        result.update(m.groupdict())

        for key, pattern in self._requirements.items():
            if not key in result or not pattern.match(result[key]):
                return None

        self._parse_data(result)

        return result, value[m.end():]

    def _parse_data(self, data):
        for key, callback in self._parsers.items():
            if key in data:
                data[key] = callback(data[key])

    def _format_data(self, data):
        for key, formatter in self._formatters.items():
            if key in data:
                if isinstance(formatter, basestring):
                    data[key] = formatter % data[key]
                else:
                    data[key] = formatter(data[key])

    def format(self, **kwargs):
        data = self._constants.copy()
        data.update(kwargs)
        self._format_data(data)

        out = self._format % data

        x = self.match(out)
        if x is None:
            raise FormatError('cannot match against output')
        m, d = x
        if d:
            raise FormatError('did not match all output')

        self._parse_data(data)

        for k, v in m.iteritems():
            if k in data and data[k] != v:
                raise FormatError('got different value for %r: got %r, expected %r' % (k, v, data[k]))

        return out

class Match(collections.Mapping):

    def __init__(self, pattern, data, unmatched):
        self.pattern = pattern
        self.data = data
        self.unmatched = unmatched

    def __getitem__(self, key):
        return self.data[key]

    def __getattr__(self, name):
        return getattr(self.data, name)

    def __iter__(self):
        return iter(self.data)

    def __len__(self):
        return len(self.data)

    def format(self, **kwargs):
        data = self.data.copy()
        data.update(kwargs)
        return self.pattern.format(**data)

class ReRouter(core.Router):

    def __init__(self):
        self._apps = []

    def register(self, pattern, app=None, **kwargs):
        """Register directly, or use as a decorator.

        Params:
            pattern -- The pattern to match with. Should start with a '/'.
            app -- The app to register. If not provided this method returns
                a decorator which can be used to register with.

        """

        # We are being used directly here.
        if app:
            # We are creating a key here that will first respect the requested
            # priority of apps relative to each other, but in the case of
            # multiple apps at the same priority, respect the registration
            # order.
            priority = (-kwargs.pop('_priority', 0), len(self._apps))
            insort(self._apps, (priority, Pattern(pattern, **kwargs), app))
            return app

        # We are not being used directly, so return a decorator to do the
        # work later.
        def ReRouter_register(app):
            self.register(pattern, app, **kwargs)
            return app
        return ReRouter_register

    def route_step(self, path):
        for _, pattern, app in self._apps:
            m = pattern.match(path)
            if m:
                kwargs, path = m
                return core.RoutingStep(next=app, path=path, data=kwargs)

    def generate_step(self, data):
        log.debug('generate_step(%r, %r)' % (self, data))
        for _, pattern, app in self._apps:
            # Skip patterns that are not identifiable.
            if not (pattern._keys or pattern._constants):
                continue
            # Filter out unmatching constant data.
            if any(k in data and data[k] != v for k, v in
                pattern._constants.iteritems()):
                continue
            try:
                return core.GenerationStep(segment=pattern.format(**data), next=app)
            except FormatError:
                pass
                # log.exception('FormatError while generating')
            except KeyError:
                pass
                # log.exception('KeyError while generating')

<|endoftext|>


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

"""
gifshare - Share Gifs via Amazon S3
"""

FOOTER = """
Copyright (c) 2014 by Mark Smith.
MIT Licensed, see LICENSE.txt for more details.
"""

__version__ = '0.0.1'

import argparse
from ConfigParser import SafeConfigParser
import logging
from os.path import expanduser, isfile, basename, splitext
import random
import re
from StringIO import StringIO
import sys

from boto.s3.key import Key
from boto.s3.connection import S3Connection

import magic
import progressbar
import requests

class UserException(Exception):
    """
    Indicates that the user should not see a stack-trace - just the
    exception message.
    """
    pass

class UnknownFileType(UserException):
    pass

class FileAlreadyExists(UserException):
    pass

URL_RE = re.compile(r'^http.*')
CONTENT_TYPE_MAP = {
    'gif': 'image/gif',
    'jpeg': 'image/jpeg',
    'png': 'image/png',
}
LOG = logging.getLogger('gifshare')

def correct_ext(data, is_buffer=False):
    magic_output = magic.from_buffer(data) if is_buffer else magic.from_file(
        data)
    match = re.search(r'JPEG|GIF|PNG', magic_output)
    if match:
        return match.group(0).lower()
    else:
        raise UnknownFileType("Unknown file type: {}".format(magic_output))

def load_config():
    config = SafeConfigParser()
    config.read([expanduser('~/.gifshare'), '.gifshare'])
    return config

def download_file(url):
    LOG.debug("Downloading image ...")
    response = requests.get(url, stream=True)
    length = int(response.headers['content-length'])
    print 'Content length:', length
    content = StringIO()
    i = 0
    widgets = ['Downloading image ', progressbar.Bar(), progressbar.Percentage()]
    pbar = progressbar.ProgressBar(widgets=widgets, maxval=length).start()
    for chunk in response.iter_content(64):
        i += len(chunk)
        content.write(chunk)
        pbar.update(i)
    pbar.finish()

    return content.getvalue()

def get_name_from_url(url):
    return re.match(r'.*/([^/\.]+)', url).group(1)

def upload_callback():
    pbar = [None]
    widgets = ['Uploading image ', progressbar.Bar(), progressbar.Percentage()]

    def callback(update, total):
        if pbar[0] is None:
            pbar[0] = progressbar.ProgressBar(widgets=widgets, maxval=total)
            pbar[0].start()
        else:
            pbar[0].update(update)
        if update == total:
            pbar[0].finish()

    return callback

def upload_url(config, url, name=None, force=False):
    LOG.debug('Uploading URL ...')
    data = download_file(url)
    ext = correct_ext(data, True)
    filename = (name or get_name_from_url(url)) + '.' + ext
    dest_url = config.get('default', 'web_root') + filename
    bucket = Bucket(config)
    key = bucket.key_for(filename, CONTENT_TYPE_MAP[ext])
    if key.exists() and not force:
        raise FileAlreadyExists("File at {} already exists!".format(dest_url))
    LOG.debug("Uploading image ...")

    key.set_contents_from_string(data, cb=upload_callback())

    return dest_url

def upload_file(config, path, name=None, force=False):
    LOG.debug("Uploading file ...")
    ext = correct_ext(path)
    filename = (name or splitext(basename(path))[0]) + '.' + ext
    url = config.get('default', 'web_root') + filename
    bucket = Bucket(config)
    key = bucket.key_for(filename, CONTENT_TYPE_MAP[ext])
    if key.exists() and not force:
        raise FileAlreadyExists("File at {} already exists!".format(url))
    key.set_contents_from_filename(path, cb=upload_callback())

    return url

class Bucket(object):
    def __init__(self, config):
        self._bucket = None
        self._key_id = config.get('default', 'aws_access_id')
        self._access_key = config.get('default', 'aws_secret_access_key')
        self._bucket_name = config.get('default', 'bucket')
        self._web_root = config.get('default', 'web_root')

    @property
    def bucket(self):
        if not self._bucket:
            conn = S3Connection(self._key_id, self._access_key)
            self._bucket = conn.get_bucket(self._bucket_name)
        return self._bucket

    def key_for(self, filename, content_type):
        k = Key(self.bucket, filename)
        k.content_type = content_type
        return k

    def list(self):
        bucket = self.bucket
        for key in bucket.list():
            url = self._web_root + key.name
            yield url

def command_upload(arguments, config):
    path = arguments.path
    if not URL_RE.match(path):
        if isfile(path):
            print upload_file(
                config, path, arguments.key, force=arguments.force)
        else:
            raise IOError(
                '{} does not exist or is not a file!'.format(path))
    else:
        print upload_url(
            config, path, arguments.key, force=arguments.force)

def command_list(arguments, config):
    bucket = Bucket(config)
    if not arguments.random:
        for item in bucket.list():
            print item
    else:
        print random.choice(list(bucket.list()))

def main(argv=sys.argv[1:]):
    try:
        a_parser = argparse.ArgumentParser(description=__doc__, epilog=FOOTER)
        a_parser.add_argument(
            '--version',
            action='version',
            version='%(prog)s ' + __version__)

        a_parser.add_argument(
            '-v', '--verbose',
            action='store_true',
            help='print out more stuff')

        subparsers = a_parser.add_subparsers()

        upload_parser = subparsers.add_parser(
            "upload",
            help="Upload an image to your bucket."
        )
        upload_parser.set_defaults(target=command_upload)

        upload_parser.add_argument(
            '--force', '-f',
            action='store_true',
            default=False,
            help='Overwrite any existing files if necessary.')

        upload_parser.add_argument(
            'path',
            help='The path to a file to upload')

        upload_parser.add_argument(
            'key',
            nargs='?',
            help='A nice filename for the gif.')

        list_parser = subparsers.add_parser(
            "list",
            help="List images stored in your bucket."
        )
        list_parser.add_argument(
            '-r', '--random',
            action='store_true',
            help='Display a single random image URL.'
        )
        list_parser.set_defaults(target=command_list)

        arguments = a_parser.parse_args(argv)
        config = load_config()

        logging.basicConfig()
        LOG.setLevel(
            level=logging.DEBUG if arguments.verbose else logging.WARN)

        arguments.target(arguments, config)
    except UserException as ue:
        print >> sys.stderr, ue
        return 1

if __name__ == '__main__':
    sys.exit(main())
<commit_msg>print should only be used for output.<commit_after>#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""
gifshare - Share Gifs via Amazon S3
"""

FOOTER = """
Copyright (c) 2014 by Mark Smith.
MIT Licensed, see LICENSE.txt for more details.
"""

__version__ = '0.0.1'

import argparse
from ConfigParser import SafeConfigParser
import logging
from os.path import expanduser, isfile, basename, splitext
import random
import re
from StringIO import StringIO
import sys

from boto.s3.key import Key
from boto.s3.connection import S3Connection

import magic
import progressbar
import requests

class UserException(Exception):
    """
    Indicates that the user should not see a stack-trace - just the
    exception message.
    """
    pass

class UnknownFileType(UserException):
    pass

class FileAlreadyExists(UserException):
    pass

URL_RE = re.compile(r'^http.*')
CONTENT_TYPE_MAP = {
    'gif': 'image/gif',
    'jpeg': 'image/jpeg',
    'png': 'image/png',
}
LOG = logging.getLogger('gifshare')

def correct_ext(data, is_buffer=False):
    magic_output = magic.from_buffer(data) if is_buffer else magic.from_file(
        data)
    match = re.search(r'JPEG|GIF|PNG', magic_output)
    if match:
        return match.group(0).lower()
    else:
        raise UnknownFileType("Unknown file type: {}".format(magic_output))

def load_config():
    config = SafeConfigParser()
    config.read([expanduser('~/.gifshare'), '.gifshare'])
    return config

def download_file(url):
    LOG.debug("Downloading image ...")
    response = requests.get(url, stream=True)
    length = int(response.headers['content-length'])
    LOG.debug('Content length: %d', length)
    content = StringIO()
    i = 0
    widgets = ['Downloading image ', progressbar.Bar(), progressbar.Percentage()]
    pbar = progressbar.ProgressBar(widgets=widgets, maxval=length).start()
    for chunk in response.iter_content(64):
        i += len(chunk)
        LOG.debug('Update: %d', i)
        content.write(chunk)
        pbar.update(i)
    pbar.finish()

    return content.getvalue()

def get_name_from_url(url):
    return re.match(r'.*/([^/\.]+)', url).group(1)

def upload_callback():
    pbar = [None]
    widgets = ['Uploading image ', progressbar.Bar(), progressbar.Percentage()]

    def callback(update, total):
        if pbar[0] is None:
            pbar[0] = progressbar.ProgressBar(widgets=widgets, maxval=total)
            pbar[0].start()
        else:
            pbar[0].update(update)
        if update == total:
            pbar[0].finish()

    return callback

def upload_url(config, url, name=None, force=False):
    LOG.debug('Uploading URL ...')
    data = download_file(url)
    ext = correct_ext(data, True)
    filename = (name or get_name_from_url(url)) + '.' + ext
    dest_url = config.get('default', 'web_root') + filename
    bucket = Bucket(config)
    key = bucket.key_for(filename, CONTENT_TYPE_MAP[ext])
    if key.exists() and not force:
        raise FileAlreadyExists("File at {} already exists!".format(dest_url))
    LOG.debug("Uploading image ...")

    key.set_contents_from_string(data, cb=upload_callback())

    return dest_url

def upload_file(config, path, name=None, force=False):
    LOG.debug("Uploading file ...")
    ext = correct_ext(path)
    filename = (name or splitext(basename(path))[0]) + '.' + ext
    url = config.get('default', 'web_root') + filename
    bucket = Bucket(config)
    key = bucket.key_for(filename, CONTENT_TYPE_MAP[ext])
    if key.exists() and not force:
        raise FileAlreadyExists("File at {} already exists!".format(url))
    key.set_contents_from_filename(path, cb=upload_callback())

    return url

class Bucket(object):
    def __init__(self, config):
        self._bucket = None
        self._key_id = config.get('default', 'aws_access_id')
        self._access_key = config.get('default', 'aws_secret_access_key')
        self._bucket_name = config.get('default', 'bucket')
        self._web_root = config.get('default', 'web_root')

    @property
    def bucket(self):
        if not self._bucket:
            conn = S3Connection(self._key_id, self._access_key)
            self._bucket = conn.get_bucket(self._bucket_name)
        return self._bucket

    def key_for(self, filename, content_type):
        k = Key(self.bucket, filename)
        k.content_type = content_type
        return k

    def list(self):
        bucket = self.bucket
        for key in bucket.list():
            url = self._web_root + key.name
            yield url

def command_upload(arguments, config):
    path = arguments.path
    if not URL_RE.match(path):
        if isfile(path):
            print upload_file(
                config, path, arguments.key, force=arguments.force)
        else:
            raise IOError(
                '{} does not exist or is not a file!'.format(path))
    else:
        print upload_url(
            config, path, arguments.key, force=arguments.force)

def command_list(arguments, config):
    bucket = Bucket(config)
    if not arguments.random:
        for item in bucket.list():
            print item
    else:
        print random.choice(list(bucket.list()))

def main(argv=sys.argv[1:]):
    try:
        a_parser = argparse.ArgumentParser(description=__doc__, epilog=FOOTER)
        a_parser.add_argument(
            '--version',
            action='version',
            version='%(prog)s ' + __version__)

        a_parser.add_argument(
            '-v', '--verbose',
            action='store_true',
            help='print out more stuff')

        subparsers = a_parser.add_subparsers()

        upload_parser = subparsers.add_parser(
            "upload",
            help="Upload an image to your bucket."
        )
        upload_parser.set_defaults(target=command_upload)

        upload_parser.add_argument(
            '--force', '-f',
            action='store_true',
            default=False,
            help='Overwrite any existing files if necessary.')

        upload_parser.add_argument(
            'path',
            help='The path to a file to upload')

        upload_parser.add_argument(
            'key',
            nargs='?',
            help='A nice filename for the gif.')

        list_parser = subparsers.add_parser(
            "list",
            help="List images stored in your bucket."
        )
        list_parser.add_argument(
            '-r', '--random',
            action='store_true',
            help='Display a single random image URL.'
        )
        list_parser.set_defaults(target=command_list)

        arguments = a_parser.parse_args(argv)
        config = load_config()

        logging.basicConfig()
        LOG.setLevel(
            level=logging.DEBUG if arguments.verbose else logging.WARN)

        arguments.target(arguments, config)
    except UserException as ue:
        print >> sys.stderr, ue
        return 1

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


<commit_before>"""
Text types
"""
from scrapely.extractors import safehtml, htmlregion

def escape_html(text):
    """escape text for use in an html page"""
    # import here so jinja2 is only required when displaying html
    from jinja2.utils import escape
    return escape(text)

class _BaseTextProcessor(object):
    """basic text processor, defines identity functions, some of which
    are overridden in subclasses
    """
    def extract(self, text):
        """Matches and extracts any string, as it is"""
        return text

    def adapt(self, text, htmlpage):
        return text

    def render(self, field_name, field_value, item):
        return escape_html(field_value)

class RawFieldTypeProcessor(_BaseTextProcessor):
    """Extracts the raw data, without processing. Data is escaped for presentation

    >>> r = RawFieldTypeProcessor()
    >>> html = htmlregion(u'
test
')
    >>> r.extract(html)
    u'
test
'
    >>> r.adapt(html, None)
    u'
test
'
    >>> str(r.render(None, '
test
', {}))
    '<p>test</p>'
    """
    name = 'raw html'
    description = 'raw html as it appears in the page'

class TextFieldTypeProcessor(_BaseTextProcessor):
    """Extracts strings, removing all HTML markup

    Extraction and rendering does nothing, but adapt converts html into text
    >>> p = TextFieldTypeProcessor()
    >>> html = htmlregion(u'
test
!')
    >>> extracted = p.extract(html)
    >>> extracted
    u'test !'
    >>> p.adapt(extracted, None)
    u'test !'
    """
    name = 'text'
    description = 'extracts text from web pages, cleaning all markup'

    def extract(self, htmlregion):
        return htmlregion.text_content

class SafeHtmlFieldTypeProcessor(_BaseTextProcessor):
    """Extracts strings, with only a safe subset of HTML remaining

    Extraction and rendering does nothing, but adapt transforms the HTML
    >>> p = SafeHtmlFieldTypeProcessor()
    >>> html = htmlregion(u'
test
 <blink>foo')
    >>> p.extract(html)
    u'
test
 <blink>foo'
    >>> p.adapt(html)
    u'
test
 foo'
    """
    name = 'safe html'
    description = 'removes all but a small subset of html tags'

    def adapt(self, text, htmlpage=None):
        """Remove html markup"""
        return safehtml(text)

    def render(self, field_name, field_value, item):
        return field_value
<commit_msg>fix text extraction to avoid to match whitespace chars<commit_after>"""
Text types
"""
from scrapely.extractors import text as extract_text, safehtml, htmlregion

def escape_html(text):
    """escape text for use in an html page"""
    # import here so jinja2 is only required when displaying html
    from jinja2.utils import escape
    return escape(text)

class _BaseTextProcessor(object):
    """basic text processor, defines identity functions, some of which
    are overridden in subclasses
    """
    def extract(self, text):
        """Matches and extracts any string, as it is"""
        return text

    def adapt(self, text, htmlpage):
        return text

    def render(self, field_name, field_value, item):
        return escape_html(field_value)

class RawFieldTypeProcessor(_BaseTextProcessor):
    """Extracts the raw data, without processing. Data is escaped for presentation

    >>> r = RawFieldTypeProcessor()
    >>> html = htmlregion(u'
test
')
    >>> r.extract(html)
    u'
test
'
    >>> r.adapt(html, None)
    u'
test
'
    >>> str(r.render(None, '
test
', {}))
    '<p>test</p>'
    """
    name = 'raw html'
    description = 'raw html as it appears in the page'

class TextFieldTypeProcessor(_BaseTextProcessor):
    """Extracts strings, removing all HTML markup

    Extraction and rendering does nothing, but adapt converts html into text
    >>> p = TextFieldTypeProcessor()
    >>> html = htmlregion(u'
test
!')
    >>> extracted = p.extract(html)
    >>> extracted
    u'test !'
    >>> p.adapt(extracted, None)
    u'test !'
    >>> html = htmlregion(u'
 \\n
')
    >>> p.extract(html)
    u''
    """
    name = 'text'
    description = 'extracts text from web pages, cleaning all markup'

    def extract(self, htmlregion):
        return extract_text(htmlregion.text_content)

class SafeHtmlFieldTypeProcessor(_BaseTextProcessor):
    """Extracts strings, with only a safe subset of HTML remaining

    Extraction and rendering does nothing, but adapt transforms the HTML
    >>> p = SafeHtmlFieldTypeProcessor()
    >>> html = htmlregion(u'
test
 <blink>foo')
    >>> p.extract(html)
    u'
test
 <blink>foo'
    >>> p.adapt(html)
    u'
test
 foo'
    """
    name = 'safe html'
    description = 'removes all but a small subset of html tags'

    def adapt(self, text, htmlpage=None):
        """Remove html markup"""
        return safehtml(text)

    def render(self, field_name, field_value, item):
        return field_value
<|endoftext|>


<commit_before>'''
Main file for this module.
Takes the gadget base as input and an output name.
Calculates the matter correlation function in a gadget output.
Downsamples the particles randomly by a certain fraction.
'''

#TODO wrap this in a main()
import os
from argparse import ArgumentParser
from glob import glob
import numpy as np
#import Corrfunc
from .readGadgetSnapshot import *

desc = '''Main file for this module.
Takes the gadget base as input and an output name.
Calculates the matter correlation function in a gadget output.
Downsamples the particles randomly by a certain fraction.'''

parser = ArgumentParser(description=desc)

parser.add_argument('inbase', type=str, help= \
    'Path and file base for the gadget files.')
# May want to rename to tmp
parser.add_argument('outputfile', type=str, help= \
    'Filename to write the final correlation function.')

#TODO find a good value for p
p0 = 1e-3
parser.add_argument('-p', '--sample_frac', default = p0,type = float, help='Fraction to sample. Default is %d.'%p0)

args = vars(parser.parse_args())

inbase = args['inbase']
outputfile= args['outputfile']
p = args['sample_frac']

assert os.path.exists(inbase+ '.0')#check that the path is valid
assert os.path.exists(outputfile) or os.access(os.path.dirname(outputfile), os.W_OK) #check if it exists or we have access
assert 0 < p < 1 #p must be a valid fraction!

all_pos = np.array([])

for file in glob(inbase+'*'):
    pos = readGadgetSnapshot(file, read_pos=True)[1] #Think this returns some type of tuple; should check
    pos = pos[np.random.rand(pos.shape[0]) < p] #downsample
    all_pos = np.resize(all_pos, (all_pos.shape[0]+pos.shape[0], 3))
    all_pos[-pos.shape[0]:, :] = pos
    break

print all_pos

#all pos have been collected. now run corrFunc.<commit_msg>Checking that my loading of snaps works as expected.<commit_after>'''
Main file for this module.
Takes the gadget base as input and an output name.
Calculates the matter correlation function in a gadget output.
Downsamples the particles randomly by a certain fraction.
'''

#TODO wrap this in a main()
import os
from argparse import ArgumentParser
from glob import glob
import numpy as np
#import Corrfunc
from readGadgetSnapshot import *#TODO fix relative imports

desc = '''Main file for this module.
Takes the gadget base as input and an output name.
Calculates the matter correlation function in a gadget output.
Downsamples the particles randomly by a certain fraction.'''

parser = ArgumentParser(description=desc)

parser.add_argument('inbase', type=str, help= \
    'Path and file base for the gadget files.')
# May want to rename to tmp
parser.add_argument('outputfile', type=str, help= \
    'Filename to write the final correlation function.')

#TODO find a good value for p
p0 = 1e-3
parser.add_argument('-p', '--sample_frac', default = p0,type = float, help='Fraction to sample. Default is %f.'%p0)

args = vars(parser.parse_args())

inbase = args['inbase']
outputfile= args['outputfile']
p = args['sample_frac']

assert os.path.exists(inbase+ '.0'), "%s is not a valid path!"%(inbase+'.0') #check that the path is valid
assert 0 < p < 1 , "%f is not a valid fraction between 0 and 1"%p #p must be a valid fraction!

#check we can write to the outputfile
try:
    open(outputfile, 'w')
except OSError:
    raise IOError('Outputfile %s is not writable!'%outputfile)

all_pos = np.array([])

for file in glob(inbase+'*'):
    pos = readGadgetSnapshot(file, read_pos=True)[1] #Think this returns some type of tuple; should check
    pos = pos[np.random.rand(pos.shape[0]) < p] #downsample
    all_pos = np.resize(all_pos, (all_pos.shape[0]+pos.shape[0], 3))
    all_pos[-pos.shape[0]:, :] = pos
    break

print all_pos

#all pos have been collected. now run corrFunc.<|endoftext|>