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

import os
import sys
from subprocess import call

class secexcept(RuntimeError):
    def __init__(_self, msg):
        _self.message = msg

class selfextractorcreator:

    def __init__(_self, in_inst, in_pack, outfile):

        _self.input_inst = os.path.abspath(in_inst)
        _self.input_pack = os.path.abspath(in_pack)
        _self.outputfn = os.path.abspath(outfile)

        if not os.path.exists(_self.input_inst):
            raise secexcept("Input instructions file does not exist")

        if not os.path.exists(_self.input_pack):
            raise secexcept("Input package file does not exist")

        if os.path.exists(_self.outputfn):
            raise secexcept("%s already exists and I refuse to overwrite" % outfile)

    def _filecontents(_self, fname):
        ret = None
        with open(fname, "rb") as f:
            ret = f.read()
        return ret

    def inst_contents(_self):
        return _self._filecontents(_self.input_inst)

    def pack_contents(_self):
        return _self._filecontents(_self.input_pack)

    def template_contents(_self, inst, pack):

        template_edit = """#!/bin/bash

echo "Extracting..."

NAME_SELF=`basename $0`

OUT_INST=nop
OUT_PACK=nop

SIZE_INST=nop
SIZE_PACK=nop

# get byte offset of all magic string matches
FINDSENT=`grep -a "SENTINEL" --byte-offset --only-matching $NAME_SELF`
OFFSET_INST=""

# retrieves the last ocurrence of the magic string
for sents in $FINDSENT; do
    OFFSET_INST=${sents%%":SENTINEL"}
done

# check whether we did find the inst offset or not at all
if [ -z $OFFSET_INST ]; then
    echo "Unable to find instructions offset. Bailing out."
    exit 1
fi

# calc offsets
OFFSET_INST=$((OFFSET_INST + 9))
OFFSET_PACK=$((OFFSET_INST + SIZE_INST))

# extracts instruction script
dd skip=$OFFSET_INST count=$SIZE_INST if=$NAME_SELF of=$OUT_INST bs=1 &>/dev/null
chmod +x $OUT_INST

# extracts package file
dd skip=$OFFSET_PACK count=$SIZE_PACK if=$NAME_SELF of=$OUT_PACK bs=1 &>/dev/null

echo "Done. Will now run the embedded script."

# runs instruction script
source $OUT_INST

rm -rf ./$OUT_INST

exit 0 #SENTINEL\n"""

        # write names
        template_edit = template_edit.replace("OUT_INST=nop", "OUT_INST=%s" % os.path.basename(_self.input_inst))
        template_edit = template_edit.replace("OUT_PACK=nop", "OUT_PACK=%s" % os.path.basename(_self.input_pack))

        # write sizes
        template_edit = template_edit.replace("SIZE_INST=nop", "SIZE_INST=%s" % os.path.getsize(_self.input_inst))
        template_edit = template_edit.replace("SIZE_PACK=nop", "SIZE_PACK=%s" % os.path.getsize(_self.input_pack))

        return template_edit

    def produce(_self):
        inst_ct = _self.inst_contents()
        pack_ct = _self.pack_contents()
        templ_ct = _self.template_contents(inst_ct, pack_ct)

        with open(_self.outputfn, 'wb') as f:
            f.write(templ_ct.encode("ascii"))
            f.write(inst_ct)
            f.write(pack_ct)

        call(["chmod", "+x", _self.outputfn])

if __name__ == "__main__":

    if len(sys.argv) != 4:
        print("Usage: %s input_instructionsfile.sh input_packagefile.tar.gz output_filename.sh" % os.path.basename(__file__))
        sys.exit(1)

    infile_inst = sys.argv[1]
    infile_pack = sys.argv[2]
    outfile = sys.argv[3]

    try:
        sec = selfextractorcreator(infile_inst, infile_pack, outfile)
        sec.produce()
    except secexcept as sx:
        print("Failed: %s" % sx.message)

<commit_msg>use puaq pattern on autoselfext<commit_after>#!/usr/bin/env python3

import os
import sys
from subprocess import call

class secexcept(RuntimeError):
    def __init__(_self, msg):
        _self.message = msg

class selfextractorcreator:

    def __init__(_self, in_inst, in_pack, outfile):

        _self.input_inst = os.path.abspath(in_inst)
        _self.input_pack = os.path.abspath(in_pack)
        _self.outputfn = os.path.abspath(outfile)

        if not os.path.exists(_self.input_inst):
            raise secexcept("Input instructions file does not exist")

        if not os.path.exists(_self.input_pack):
            raise secexcept("Input package file does not exist")

        if os.path.exists(_self.outputfn):
            raise secexcept("%s already exists and I refuse to overwrite" % outfile)

    def _filecontents(_self, fname):
        ret = None
        with open(fname, "rb") as f:
            ret = f.read()
        return ret

    def inst_contents(_self):
        return _self._filecontents(_self.input_inst)

    def pack_contents(_self):
        return _self._filecontents(_self.input_pack)

    def template_contents(_self, inst, pack):

        template_edit = """#!/bin/bash

echo "Extracting..."

NAME_SELF=`basename $0`

OUT_INST=nop
OUT_PACK=nop

SIZE_INST=nop
SIZE_PACK=nop

# get byte offset of all magic string matches
FINDSENT=`grep -a "SENTINEL" --byte-offset --only-matching $NAME_SELF`
OFFSET_INST=""

# retrieves the last ocurrence of the magic string
for sents in $FINDSENT; do
    OFFSET_INST=${sents%%":SENTINEL"}
done

# check whether we did find the inst offset or not at all
if [ -z $OFFSET_INST ]; then
    echo "Unable to find instructions offset. Bailing out."
    exit 1
fi

# calc offsets
OFFSET_INST=$((OFFSET_INST + 9))
OFFSET_PACK=$((OFFSET_INST + SIZE_INST))

# extracts instruction script
dd skip=$OFFSET_INST count=$SIZE_INST if=$NAME_SELF of=$OUT_INST bs=1 &>/dev/null
chmod +x $OUT_INST

# extracts package file
dd skip=$OFFSET_PACK count=$SIZE_PACK if=$NAME_SELF of=$OUT_PACK bs=1 &>/dev/null

echo "Done. Will now run the embedded script."

# runs instruction script
source $OUT_INST

rm -rf ./$OUT_INST

exit 0 #SENTINEL\n"""

        # write names
        template_edit = template_edit.replace("OUT_INST=nop", "OUT_INST=%s" % os.path.basename(_self.input_inst))
        template_edit = template_edit.replace("OUT_PACK=nop", "OUT_PACK=%s" % os.path.basename(_self.input_pack))

        # write sizes
        template_edit = template_edit.replace("SIZE_INST=nop", "SIZE_INST=%s" % os.path.getsize(_self.input_inst))
        template_edit = template_edit.replace("SIZE_PACK=nop", "SIZE_PACK=%s" % os.path.getsize(_self.input_pack))

        return template_edit

    def produce(_self):
        inst_ct = _self.inst_contents()
        pack_ct = _self.pack_contents()
        templ_ct = _self.template_contents(inst_ct, pack_ct)

        with open(_self.outputfn, 'wb') as f:
            f.write(templ_ct.encode("ascii"))
            f.write(inst_ct)
            f.write(pack_ct)

        call(["chmod", "+x", _self.outputfn])

def puaq():
    print("Usage: %s input_instructionsfile.sh input_packagefile.tar.gz output_filename.sh" % os.path.basename(__file__))
    sys.exit(1)

if __name__ == "__main__":

    if len(sys.argv) != 4:
        puaq()

    infile_inst = sys.argv[1]
    infile_pack = sys.argv[2]
    outfile = sys.argv[3]

    try:
        sec = selfextractorcreator(infile_inst, infile_pack, outfile)
        sec.produce()
    except secexcept as sx:
        print("Failed: %s" % sx.message)

<|endoftext|>


<commit_before>from __future__ import absolute_import
from __future__ import print_function

import numpy as np
from itertools import tee
import socket
from multiprocessing import Process
import six.moves.cPickle as pickle
from six.moves import range
from flask import Flask, request
try:
    import urllib.request as urllib2
except ImportError:
    import urllib2

from pyspark.mllib.linalg import Matrix, Vector

from keras.models import model_from_yaml, slice_X

from .utils.rwlock import RWLock
from .utils.functional_utils import subtract_params
from .utils.rdd_utils import lp_to_simple_rdd
from .mllib.adapter import to_matrix, from_matrix, to_vector, from_vector
from .optimizers import SGD as default_optimizer

def get_server_weights(master_url='localhost:5000'):
    '''
    Retrieve master weights from parameter server
    '''
    request = urllib2.Request('http://{0}/parameters'.format(master_url),
                              headers={'Content-Type': 'application/elephas'})
    return pickle.loads(urllib2.urlopen(request).read())

def put_deltas_to_server(delta, master_url='localhost:5000'):
    '''
    Update master parameters with deltas from training process
    '''
    request = urllib2.Request('http://{0}/update'.format(master_url),
                              pickle.dumps(delta, -1), headers={'Content-Type': 'application/elephas'})
    return urllib2.urlopen(request).read()

class SparkModel(object):
    '''
    SparkModel is the main abstraction of elephas. Every other model
    should inherit from it.
    '''
    def __init__(self, sc, master_network, optimizer=None, mode='asynchronous', frequency='epoch',
                 num_workers=4, *args, **kwargs):
        self.spark_context = sc
        self._master_network = master_network
        if optimizer is None:
            self.optimizer = default_optimizer()
        else:
            self.optimizer = optimizer
        self.mode = mode
        self.frequency = frequency
        self.num_workers = num_workers
        self.weights = master_network.get_weights()
        self.pickled_weights = None
        self.lock = RWLock()

    @staticmethod
    def determine_master():
        '''
        Get URL of parameter server, running on master
        '''
        master_url = socket.gethostbyname(socket.gethostname()) + ':5000'
        return master_url

    def get_train_config(self, nb_epoch, batch_size,
                         verbose, validation_split):
        '''
        Get configuration of training parameters
        '''
        train_config = {}
        train_config['nb_epoch'] = nb_epoch
        train_config['batch_size'] = batch_size
        train_config['verbose'] = verbose
        train_config['validation_split'] = validation_split
        return train_config

    def get_config(self):
        '''
        Get configuration of model parameters
        '''
        model_config = {}
        model_config['model'] = self.master_network.get_config()
        model_config['optimizer'] = self.optimizer.get_config()
        model_config['mode'] = self.mode
        return model_config

    @property
    def master_network(self):
        ''' Get master network '''
        return self._master_network

    @master_network.setter
    def master_network(self, network):
        ''' Set master network '''
        self._master_network = network

    def start_server(self):
        ''' Start parameter server'''
        self.server = Process(target=self.start_service)
        self.server.start()

    def stop_server(self):
        ''' Terminate parameter server'''
        self.server.terminate()
        self.server.join()

    def start_service(self):
        ''' Define service and run flask app'''
        app = Flask(__name__)
        self.app = app

        @app.route('/')
        def home():
            return 'Elephas'

        @app.route('/parameters', methods=['GET'])
        def get_parameters():
            if self.mode == 'asynchronous':
                self.lock.acquire_read()
            self.pickled_weights = pickle.dumps(self.weights, -1)
            pickled_weights = self.pickled_weights
            if self.mode == 'asynchronous':
                self.lock.release()
            return pickled_weights

        @app.route('/update', methods=['POST'])
        def update_parameters():
            delta = pickle.loads(request.data)
            if self.mode == 'asynchronous':
                self.lock.acquire_write()
            constraints = self.master_network.constraints

            self.weights = self.optimizer.get_updates(self.weights, constraints, delta)
            if self.mode == 'asynchronous':
                self.lock.release()
            return 'Update done'

        self.app.run(host='0.0.0.0', debug=True,
                     threaded=True, use_reloader=False)

    def predict(self, data):
        '''
        Get prediction probabilities for a numpy array of features
        '''
        return self.master_network.predict(data)

    def predict_classes(self, data):
        '''
        Predict classes for a numpy array of features
        '''
        return self.master_network.predict_classes(data)

    def train(self, rdd, nb_epoch=10, batch_size=32,
              verbose=0, validation_split=0.1):
        '''
        Train an elephas model.
        '''
        rdd = rdd.repartition(self.num_workers)
        master_url = self.determine_master()

        if self.mode in ['asynchronous', 'synchronous', 'hogwild']:
            self._train(rdd, nb_epoch, batch_size, verbose, validation_split, master_url)
        else:
            print("""Choose from one of the modes: asynchronous, synchronous or hogwild""")

    def _train(self, rdd, nb_epoch=10, batch_size=32, verbose=0,
               validation_split=0.1, master_url='localhost:5000'):
        '''
        Protected train method to make wrapping of modes easier
        '''
        if self.mode in ['asynchronous', 'hogwild']:
            self.start_server()
        yaml = self.master_network.to_yaml()
        train_config = self.get_train_config(nb_epoch, batch_size,
                                             verbose, validation_split)
        if self.mode in ['asynchronous', 'hogwild']:
            worker = AsynchronousSparkWorker(yaml, train_config, self.frequency, master_url)
            rdd.mapPartitions(worker.train).collect()
            new_parameters = get_server_weights(master_url)
        elif self.mode == 'synchronous':
            init = self.master_network.get_weights()
            parameters = self.spark_context.broadcast(init)
            worker = SparkWorker(yaml, parameters, train_config)
            deltas = rdd.mapPartitions(worker.train).collect()
            new_parameters = self.master_network.get_weights()
            for delta in deltas:
                constraints = self.master_network.constraints
                new_parameters = self.optimizer.get_updates(self.weights, constraints, delta)
        self.master_network.set_weights(new_parameters)
        if self.mode in ['asynchronous', 'hogwild']:
            self.stop_server()

class SparkWorker(object):
    '''
    Synchronous Spark worker. This code will be executed on workers.
    '''
    def __init__(self, yaml, parameters, train_config):
        self.yaml = yaml
        self.parameters = parameters
        self.train_config = train_config

    def train(self, data_iterator):
        '''
        Train a keras model on a worker
        '''
        feature_iterator, label_iterator = tee(data_iterator, 2)
        x_train = np.asarray([x for x, y in feature_iterator])
        y_train = np.asarray([y for x, y in label_iterator])

        model = model_from_yaml(self.yaml)
        model.set_weights(self.parameters.value)
        weights_before_training = model.get_weights()
        if x_train.shape[0] > self.train_config.get('batch_size'):
            model.fit(x_train, y_train, show_accuracy=True, **self.train_config)
        weights_after_training = model.get_weights()
        deltas = subtract_params(weights_before_training, weights_after_training)
        yield deltas

class AsynchronousSparkWorker(object):
    '''
    Asynchronous Spark worker. This code will be executed on workers.
    '''
    def __init__(self, yaml, train_config, frequency, master_url):
        self.yaml = yaml
        self.train_config = train_config
        self.frequency = frequency
        self.master_url = master_url

    def train(self, data_iterator):
        '''
        Train a keras model on a worker and send asynchronous updates
        to parameter server
        '''
        feature_iterator, label_iterator = tee(data_iterator, 2)
        x_train = np.asarray([x for x, y in feature_iterator])
        y_train = np.asarray([y for x, y in label_iterator])

        if x_train.size == 0:
            return
        model = model_from_yaml(self.yaml)

        nb_epoch = self.train_config['nb_epoch']
        batch_size = self.train_config.get('batch_size')
        nb_train_sample = len(x_train[0])
        nb_batch = int(np.ceil(nb_train_sample/float(batch_size)))
        index_array = np.arange(nb_train_sample)
        batches = [(i*batch_size, min(nb_train_sample, (i+1)*batch_size)) for i in range(0, nb_batch)]

        if self.frequency == 'epoch':
            for epoch in range(nb_epoch):
                weights_before_training = get_server_weights(self.master_url)
                model.set_weights(weights_before_training)
                self.train_config['nb_epoch'] = 1
                if x_train.shape[0] > batch_size:
                    model.fit(x_train, y_train, show_accuracy=True, **self.train_config)
                weights_after_training = model.get_weights()
                deltas = subtract_params(weights_before_training, weights_after_training)
                put_deltas_to_server(deltas, self.master_url)
        elif self.frequency == 'batch':
            for epoch in range(nb_epoch):
                if x_train.shape[0] > batch_size:
                    for (batch_start, batch_end) in batches:
                        weights_before_training = get_server_weights(self.master_url)
                        model.set_weights(weights_before_training)
                        batch_ids = index_array[batch_start:batch_end]
                        X = slice_X(x_train, batch_ids)
                        y = slice_X(y_train, batch_ids)
                        model.train_on_batch(X, y)
                        weights_after_training = model.get_weights()
                        deltas = subtract_params(weights_before_training, weights_after_training)
                        put_deltas_to_server(deltas, self.master_url)
        else:
            print('Choose frequency to be either batch or epoch')
        yield []

class SparkMLlibModel(SparkModel):
    '''
    MLlib model takes RDDs of LabeledPoints. Internally we just convert
    back to plain old pair RDDs and continue as in SparkModel
    '''
    def __init__(self, **kwargs):
        super(SparkModel, self).__init__(**kwargs)

    def train(self, labeled_points, nb_epoch=10, batch_size=32, verbose=0, validation_split=0.1,
              categorical=False, nb_classes=None):
        '''
        Train an elephas model on an RDD of LabeledPoints
        '''
        rdd = lp_to_simple_rdd(labeled_points, categorical, nb_classes)
        rdd = rdd.repartition(self.num_workers)
        self._train(rdd, nb_epoch, batch_size, verbose, validation_split)

    def predict(self, mllib_data):
        '''
        Predict probabilities for an RDD of features
        '''
        if isinstance(mllib_data, Matrix):
            return to_matrix(self.master_network.predict(from_matrix(mllib_data)))
        elif isinstance(mllib_data, Vector):
            return to_vector(self.master_network.predict(from_vector(mllib_data)))
        else:
            print('Provide either an MLLib matrix or vector')
<commit_msg>Fix of mllib model init<commit_after>from __future__ import absolute_import
from __future__ import print_function

import numpy as np
from itertools import tee
import socket
from multiprocessing import Process
import six.moves.cPickle as pickle
from six.moves import range
from flask import Flask, request
try:
    import urllib.request as urllib2
except ImportError:
    import urllib2

from pyspark.mllib.linalg import Matrix, Vector

from keras.models import model_from_yaml, slice_X

from .utils.rwlock import RWLock
from .utils.functional_utils import subtract_params
from .utils.rdd_utils import lp_to_simple_rdd
from .mllib.adapter import to_matrix, from_matrix, to_vector, from_vector
from .optimizers import SGD as default_optimizer

def get_server_weights(master_url='localhost:5000'):
    '''
    Retrieve master weights from parameter server
    '''
    request = urllib2.Request('http://{0}/parameters'.format(master_url),
                              headers={'Content-Type': 'application/elephas'})
    return pickle.loads(urllib2.urlopen(request).read())

def put_deltas_to_server(delta, master_url='localhost:5000'):
    '''
    Update master parameters with deltas from training process
    '''
    request = urllib2.Request('http://{0}/update'.format(master_url),
                              pickle.dumps(delta, -1), headers={'Content-Type': 'application/elephas'})
    return urllib2.urlopen(request).read()

class SparkModel(object):
    '''
    SparkModel is the main abstraction of elephas. Every other model
    should inherit from it.
    '''
    def __init__(self, sc, master_network, optimizer=None, mode='asynchronous', frequency='epoch',
                 num_workers=4, *args, **kwargs):
        self.spark_context = sc
        self._master_network = master_network
        if optimizer is None:
            self.optimizer = default_optimizer()
        else:
            self.optimizer = optimizer
        self.mode = mode
        self.frequency = frequency
        self.num_workers = num_workers
        self.weights = master_network.get_weights()
        self.pickled_weights = None
        self.lock = RWLock()

    @staticmethod
    def determine_master():
        '''
        Get URL of parameter server, running on master
        '''
        master_url = socket.gethostbyname(socket.gethostname()) + ':5000'
        return master_url

    def get_train_config(self, nb_epoch, batch_size,
                         verbose, validation_split):
        '''
        Get configuration of training parameters
        '''
        train_config = {}
        train_config['nb_epoch'] = nb_epoch
        train_config['batch_size'] = batch_size
        train_config['verbose'] = verbose
        train_config['validation_split'] = validation_split
        return train_config

    def get_config(self):
        '''
        Get configuration of model parameters
        '''
        model_config = {}
        model_config['model'] = self.master_network.get_config()
        model_config['optimizer'] = self.optimizer.get_config()
        model_config['mode'] = self.mode
        return model_config

    @property
    def master_network(self):
        ''' Get master network '''
        return self._master_network

    @master_network.setter
    def master_network(self, network):
        ''' Set master network '''
        self._master_network = network

    def start_server(self):
        ''' Start parameter server'''
        self.server = Process(target=self.start_service)
        self.server.start()

    def stop_server(self):
        ''' Terminate parameter server'''
        self.server.terminate()
        self.server.join()

    def start_service(self):
        ''' Define service and run flask app'''
        app = Flask(__name__)
        self.app = app

        @app.route('/')
        def home():
            return 'Elephas'

        @app.route('/parameters', methods=['GET'])
        def get_parameters():
            if self.mode == 'asynchronous':
                self.lock.acquire_read()
            self.pickled_weights = pickle.dumps(self.weights, -1)
            pickled_weights = self.pickled_weights
            if self.mode == 'asynchronous':
                self.lock.release()
            return pickled_weights

        @app.route('/update', methods=['POST'])
        def update_parameters():
            delta = pickle.loads(request.data)
            if self.mode == 'asynchronous':
                self.lock.acquire_write()
            constraints = self.master_network.constraints

            self.weights = self.optimizer.get_updates(self.weights, constraints, delta)
            if self.mode == 'asynchronous':
                self.lock.release()
            return 'Update done'

        self.app.run(host='0.0.0.0', debug=True,
                     threaded=True, use_reloader=False)

    def predict(self, data):
        '''
        Get prediction probabilities for a numpy array of features
        '''
        return self.master_network.predict(data)

    def predict_classes(self, data):
        '''
        Predict classes for a numpy array of features
        '''
        return self.master_network.predict_classes(data)

    def train(self, rdd, nb_epoch=10, batch_size=32,
              verbose=0, validation_split=0.1):
        '''
        Train an elephas model.
        '''
        rdd = rdd.repartition(self.num_workers)
        master_url = self.determine_master()

        if self.mode in ['asynchronous', 'synchronous', 'hogwild']:
            self._train(rdd, nb_epoch, batch_size, verbose, validation_split, master_url)
        else:
            print("""Choose from one of the modes: asynchronous, synchronous or hogwild""")

    def _train(self, rdd, nb_epoch=10, batch_size=32, verbose=0,
               validation_split=0.1, master_url='localhost:5000'):
        '''
        Protected train method to make wrapping of modes easier
        '''
        if self.mode in ['asynchronous', 'hogwild']:
            self.start_server()
        yaml = self.master_network.to_yaml()
        train_config = self.get_train_config(nb_epoch, batch_size,
                                             verbose, validation_split)
        if self.mode in ['asynchronous', 'hogwild']:
            worker = AsynchronousSparkWorker(yaml, train_config, self.frequency, master_url)
            rdd.mapPartitions(worker.train).collect()
            new_parameters = get_server_weights(master_url)
        elif self.mode == 'synchronous':
            init = self.master_network.get_weights()
            parameters = self.spark_context.broadcast(init)
            worker = SparkWorker(yaml, parameters, train_config)
            deltas = rdd.mapPartitions(worker.train).collect()
            new_parameters = self.master_network.get_weights()
            for delta in deltas:
                constraints = self.master_network.constraints
                new_parameters = self.optimizer.get_updates(self.weights, constraints, delta)
        self.master_network.set_weights(new_parameters)
        if self.mode in ['asynchronous', 'hogwild']:
            self.stop_server()

class SparkWorker(object):
    '''
    Synchronous Spark worker. This code will be executed on workers.
    '''
    def __init__(self, yaml, parameters, train_config):
        self.yaml = yaml
        self.parameters = parameters
        self.train_config = train_config

    def train(self, data_iterator):
        '''
        Train a keras model on a worker
        '''
        feature_iterator, label_iterator = tee(data_iterator, 2)
        x_train = np.asarray([x for x, y in feature_iterator])
        y_train = np.asarray([y for x, y in label_iterator])

        model = model_from_yaml(self.yaml)
        model.set_weights(self.parameters.value)
        weights_before_training = model.get_weights()
        if x_train.shape[0] > self.train_config.get('batch_size'):
            model.fit(x_train, y_train, show_accuracy=True, **self.train_config)
        weights_after_training = model.get_weights()
        deltas = subtract_params(weights_before_training, weights_after_training)
        yield deltas

class AsynchronousSparkWorker(object):
    '''
    Asynchronous Spark worker. This code will be executed on workers.
    '''
    def __init__(self, yaml, train_config, frequency, master_url):
        self.yaml = yaml
        self.train_config = train_config
        self.frequency = frequency
        self.master_url = master_url

    def train(self, data_iterator):
        '''
        Train a keras model on a worker and send asynchronous updates
        to parameter server
        '''
        feature_iterator, label_iterator = tee(data_iterator, 2)
        x_train = np.asarray([x for x, y in feature_iterator])
        y_train = np.asarray([y for x, y in label_iterator])

        if x_train.size == 0:
            return
        model = model_from_yaml(self.yaml)

        nb_epoch = self.train_config['nb_epoch']
        batch_size = self.train_config.get('batch_size')
        nb_train_sample = len(x_train[0])
        nb_batch = int(np.ceil(nb_train_sample/float(batch_size)))
        index_array = np.arange(nb_train_sample)
        batches = [(i*batch_size, min(nb_train_sample, (i+1)*batch_size)) for i in range(0, nb_batch)]

        if self.frequency == 'epoch':
            for epoch in range(nb_epoch):
                weights_before_training = get_server_weights(self.master_url)
                model.set_weights(weights_before_training)
                self.train_config['nb_epoch'] = 1
                if x_train.shape[0] > batch_size:
                    model.fit(x_train, y_train, show_accuracy=True, **self.train_config)
                weights_after_training = model.get_weights()
                deltas = subtract_params(weights_before_training, weights_after_training)
                put_deltas_to_server(deltas, self.master_url)
        elif self.frequency == 'batch':
            for epoch in range(nb_epoch):
                if x_train.shape[0] > batch_size:
                    for (batch_start, batch_end) in batches:
                        weights_before_training = get_server_weights(self.master_url)
                        model.set_weights(weights_before_training)
                        batch_ids = index_array[batch_start:batch_end]
                        X = slice_X(x_train, batch_ids)
                        y = slice_X(y_train, batch_ids)
                        model.train_on_batch(X, y)
                        weights_after_training = model.get_weights()
                        deltas = subtract_params(weights_before_training, weights_after_training)
                        put_deltas_to_server(deltas, self.master_url)
        else:
            print('Choose frequency to be either batch or epoch')
        yield []

class SparkMLlibModel(SparkModel):
    '''
    MLlib model takes RDDs of LabeledPoints. Internally we just convert
    back to plain old pair RDDs and continue as in SparkModel
    '''
    def __init__(self, sc, master_network, optimizer=None, mode='asynchronous', frequency='epoch', num_workers=4):
        SparkModel.__init__(self, sc, master_network, optimizer, mode, frequency, num_workers)

    def train(self, labeled_points, nb_epoch=10, batch_size=32, verbose=0, validation_split=0.1,
              categorical=False, nb_classes=None):
        '''
        Train an elephas model on an RDD of LabeledPoints
        '''
        rdd = lp_to_simple_rdd(labeled_points, categorical, nb_classes)
        rdd = rdd.repartition(self.num_workers)
        self._train(rdd, nb_epoch, batch_size, verbose, validation_split)

    def predict(self, mllib_data):
        '''
        Predict probabilities for an RDD of features
        '''
        if isinstance(mllib_data, Matrix):
            return to_matrix(self.master_network.predict(from_matrix(mllib_data)))
        elif isinstance(mllib_data, Vector):
            return to_vector(self.master_network.predict(from_vector(mllib_data)))
        else:
            print('Provide either an MLLib matrix or vector')
<|endoftext|>


<commit_before>"""
Command to send weekly reminder mails.

This tool sends reminder mails to all tenders with a list of events with
missing availability for the tender.

Based on the reminder script of the previous Event management system by
Wietse Smid.

Author: Jelte Zeilstra
"""

from django.conf import settings
from django.core.management.base import BaseCommand, CommandError
from django.utils import timezone

from apps.organization.models import Organization
from apps.scheduling.models import MailTemplate
from utils.mail import mail

class Command(BaseCommand):
    help = 'Send reminder mails for events with missing availabilities'

    def handle(self, *args, **options):
        # Filter onganizations with an active reminder mail
        organizations = Organization.objects.filter(mailtemplate__name='reminder', mailtemplate__is_active=True)

        for organization in organizations:
            memberships = organization.membership_set.filter(is_tender=True)

            now = timezone.now()
            events = organization.participates.filter(starts_at__gte=now, is_closed=False).order_by('starts_at', )

            # Load templates
            try:
                mailtemplate = MailTemplate.objects.get(organization=organization, name="reminder", is_active=True)
            except MailTemplate.DoesNotExist:
                raise CommandError('MailTemplate "reminder" does not exist for %s' % organization)

            for membership in memberships:
                user = membership.user
                if not user.email:
                    self.stderr.write('%s heeft geen e-mailadres.\n' % (user,))
                    continue

                missing_events = events.exclude(pk__in=user.event_set.values_list('id', flat=True))
                addressees = [user]

                mail(settings.EMAIL_FROM, addressees, mailtemplate.subject, mailtemplate.template,
                     extraattrs={'missing_events': missing_events, 'now': now})
<commit_msg>scheduling: Exclude events too far in the future from the reminder mail<commit_after>"""
Command to send weekly reminder mails.

This tool sends reminder mails to all tenders with a list of events with
missing availability for the tender.

Based on the reminder script of the previous Event management system by
Wietse Smid.

Author: Jelte Zeilstra
"""

import datetime

from django.conf import settings
from django.core.management.base import BaseCommand, CommandError
from django.utils import timezone

from apps.organization.models import Organization
from apps.scheduling.models import MailTemplate
from utils.mail import mail

class Command(BaseCommand):
    help = 'Send reminder mails for events with missing availabilities'

    def handle(self, *args, **options):
        # Filter onganizations with an active reminder mail
        organizations = Organization.objects.filter(mailtemplate__name='reminder', mailtemplate__is_active=True)

        for organization in organizations:
            memberships = organization.membership_set.filter(is_tender=True)

            now = timezone.now()
            events = organization.participates.filter(starts_at__gte=now, is_closed=False).order_by('starts_at', )

            # Load template and settings
            try:
                mailtemplate = MailTemplate.objects.get(organization=organization, name="reminder", is_active=True)
                starts_before = now + datetime.timedelta(minutes=mailtemplate.send_at) if mailtemplate.send_at else None
            except MailTemplate.DoesNotExist:
                raise CommandError('MailTemplate "reminder" does not exist for %s' % organization)

            for membership in memberships:
                user = membership.user
                if not user.email:
                    self.stderr.write('%s heeft geen e-mailadres.\n' % (user,))
                    continue

                missing_events = events.exclude(pk__in=user.event_set.values_list('id', flat=True))
                if starts_before is not None:
                    missing_events = missing_events.exclude(starts_at__gte=starts_before)
                addressees = [user]

                if len(missing_events) > 0:
                    mail(settings.EMAIL_FROM, addressees, mailtemplate.subject, mailtemplate.template,
                         extraattrs={'missing_events': missing_events, 'now': now})
<|endoftext|>


<commit_before>#!/usr/bin/env python
"""EPP script to perform basic calculations on UDF:s in Clarity LIMS
Command to trigger this script:
bash -c "PATH/TO/INSTALLED/SCRIPT
--username {username}
--password {password}
--baseuri YOUR.URI
--pid {processLuid}
--log {compoundOutputFileLuidN}"
--output_files {outputFileLuids}
--udf1 NameOfUDF 1
--operator Mathematical operator
--udf2 NameOfUDF 2
--result_udf NameOfUDF to store result in
"

Johannes Alneberg, Science for Life Laboratory, Stockholm, Sweden
"""
from argparse import ArgumentParser

from genologics.lims import Lims
from genologics.epp import configure_logging,setup_standard_parser

def apply_calculations(lims,output_artifacts,udf1,op,udf2,result_udf):
    for artifact in output_artifacts:
        artifact[result_udf] = eval('{0}{1}{2}'.format(artifact.udf[udf1],op,artifact.udf[udf2]))
        artifact.put()

def main(lims,args):
    outputs = map(lambda id: Artifact(lims,id=id),args.output_files)
    apply_calculations(lims,outputs,args.udf1,args.operator,args.udf2,arrgs.result_udf)

if __name__ == "__main__":
    # Initialize parser with standard arguments and description
    desc = """EPP script to perform basic calculations on UDF:s in Clarity LIMS.
    result_udf=udf1 *operator* udf2"""
    parser = setup_standard_parser(description=desc)

    # Additional arguments
    parser.add_argument('--output_files',nargs='*',
                        help='Lims unique ids for each output file artifact')
    parser.add_argument('--udf1',
                        help='The first udf in the formula')
    parser.add_argument('--operator', choices =['+','-'],
                        help='operator to apply')
    parser.add_argument('--udf2',
                        help='The second udf in the formula')
    parser.add_argument('--result_udf',
                        help='Udf to store the result')
    args = parser.parse_args()

    # Start logging
    if args.log:
        configure_logging(args.log)
    lims = Lims(args.baseuri,args.username,args.password)
    lims.check_version()

    main(lims, args)

<commit_msg>Bugfix in udf calculations, import Artifact<commit_after>#!/usr/bin/env python
"""EPP script to perform basic calculations on UDF:s in Clarity LIMS
Command to trigger this script:
bash -c "PATH/TO/INSTALLED/SCRIPT
--username {username}
--password {password}
--baseuri YOUR.URI
--pid {processLuid}
--log {compoundOutputFileLuidN}"
--output_files {outputFileLuids}
--udf1 NameOfUDF 1
--operator Mathematical operator
--udf2 NameOfUDF 2
--result_udf NameOfUDF to store result in
"

Johannes Alneberg, Science for Life Laboratory, Stockholm, Sweden
"""
from argparse import ArgumentParser

from genologics.lims import Lims
from genologics.entities import Artifact
from genologics.epp import configure_logging,setup_standard_parser

def apply_calculations(lims,output_artifacts,udf1,op,udf2,result_udf):
    for artifact in output_artifacts:
        artifact[result_udf] = eval('{0}{1}{2}'.format(artifact.udf[udf1],op,artifact.udf[udf2]))
        artifact.put()

def main(lims,args):
    outputs = map(lambda id: Artifact(lims,id=id),args.output_files)
    apply_calculations(lims,outputs,args.udf1,args.operator,args.udf2,arrgs.result_udf)

if __name__ == "__main__":
    # Initialize parser with standard arguments and description
    desc = """EPP script to perform basic calculations on UDF:s in Clarity LIMS.
    result_udf=udf1 *operator* udf2"""
    parser = setup_standard_parser(description=desc)

    # Additional arguments
    parser.add_argument('--output_files',nargs='*',
                        help='Lims unique ids for each output file artifact')
    parser.add_argument('--udf1',
                        help='The first udf in the formula')
    parser.add_argument('--operator', choices =['+','-'],
                        help='operator to apply')
    parser.add_argument('--udf2',
                        help='The second udf in the formula')
    parser.add_argument('--result_udf',
                        help='Udf to store the result')
    args = parser.parse_args()

    # Start logging
    if args.log:
        configure_logging(args.log)
    lims = Lims(args.baseuri,args.username,args.password)
    lims.check_version()

    main(lims, args)

<|endoftext|>


<commit_before>import tensorflow as tf
import numpy as np

datax = 10000
testx = 5000
data = np.random.randint(0,2,(datax, 100))
labels = np.zeros((datax,2))
test_data = np.random.randint(0,2,(testx, 100))
test_labels = np.zeros((testx,2))
for i in range(datax):
	if (data[i,0]) == 1:
		labels[i,0] = 1
		labels[i,1] = 0
	else:
		labels[i,0] = 0
		labels[i,1] = 1
for i in range(testx):
	if (test_data[i,0]) == 1:
		test_labels[i,0] = 1
		test_labels[i,1] = 0
	else:
		test_labels[i,0] = 0
		test_labels[i,1] = 1

#Parameters
training_epochs = 51
learning_rate = .01
batch_size = 150

#Network Parameters
num_hidden = 200
num_classes = 2 #male or female
num_examples = data.shape[0]
num_attributes = data.shape[1]

#Inputs
x = tf.placeholder(tf.float32, shape=(None, num_attributes))
y = tf.placeholder(tf.float32, shape=(None, num_classes))

def multilayer_perceptron(train_data, weights, biases):
	hidden_layer = tf.add(tf.matmul(x, weights['hidden']), biases['hidden'])
	hidden_layer = tf.sigmoid(hidden_layer)

	out_layer = tf.matmul(hidden_layer, weights['out']) + biases['out']
	return out_layer

weights = {
	'hidden': tf.Variable(tf.random_normal([num_attributes, num_hidden])),
	'out': tf.Variable(tf.random_normal([num_hidden, num_classes]))
}

biases = {
	'hidden': tf.Variable(tf.random_normal([num_hidden])),
	'out': tf.Variable(tf.random_normal([num_classes]))
}

model = multilayer_perceptron(x, weights, biases)

cost = tf.reduce_mean(tf.nn.sigmoid_cross_entropy_with_logits(logits=model, labels=y))
optimizer = tf.train.AdamOptimizer(learning_rate=learning_rate).minimize(cost)

init = tf.global_variables_initializer()

with tf.Session() as sess:
	sess.run(init)

	#Training
	for epoch in range(training_epochs):
		average_cost = 0.
		total_batch = int(num_examples/batch_size)
		for i in range(total_batch):
			batch_x = data[i*batch_size:(i+1)*batch_size,:]
			batch_y = labels[i*batch_size:(i+1)*batch_size,:]

			_, c = sess.run([optimizer, cost],
				feed_dict = {x: batch_x, y: batch_y})

			# Compute loss
			average_cost += c / total_batch

		if epoch % 10 == 0:
			print("Epoch:" + str(epoch+1) + " cost = " + str(average_cost))
	print("Optimization Finished!")

	correct = tf.equal(tf.argmax(model, 1), tf.argmax(y, 1))

	accuracy = tf.reduce_mean(tf.cast(correct, "float"))
	print("Accuracy: " + str(accuracy.eval({x:test_data, y:test_labels})))<commit_msg>load test and train data from .npy file<commit_after>import tensorflow as tf
import numpy as np

# datax = 10000
# testx = 5000
# data = np.random.randint(0,2,(datax, 100))
# labels = np.zeros((datax,2))
# test_data = np.random.randint(0,2,(testx, 100))
# test_labels = np.zeros((testx,2))
# for i in range(datax):
# 	if (data[i,0]) == 1:
# 		labels[i,0] = 1
# 		labels[i,1] = 0
# 	else:
# 		labels[i,0] = 0
# 		labels[i,1] = 1
# for i in range(testx):
# 	if (test_data[i,0]) == 1:
# 		test_labels[i,0] = 1
# 		test_labels[i,1] = 0
# 	else:
# 		test_labels[i,0] = 0
# 		test_labels[i,1] = 1

data = np.load('data/training_data.npy')
labels = np.load('data/training_labels.npy')
test_data = np.load('data/testing_data.npy')
test_labels = np.load('data/testing_labels.npy')

#Parameters
training_epochs = 51
learning_rate = .01
batch_size = 150

#Network Parameters
num_hidden = 200
num_classes = 2 #male or female
num_examples = data.shape[0]
num_attributes = data.shape[1]

#Inputs
x = tf.placeholder(tf.float32, shape=(None, num_attributes))
y = tf.placeholder(tf.float32, shape=(None, num_classes))

def multilayer_perceptron(train_data, weights, biases):
	hidden_layer = tf.add(tf.matmul(x, weights['hidden']), biases['hidden'])
	hidden_layer = tf.sigmoid(hidden_layer)

	out_layer = tf.matmul(hidden_layer, weights['out']) + biases['out']
	return out_layer

weights = {
	'hidden': tf.Variable(tf.random_normal([num_attributes, num_hidden])),
	'out': tf.Variable(tf.random_normal([num_hidden, num_classes]))
}

biases = {
	'hidden': tf.Variable(tf.random_normal([num_hidden])),
	'out': tf.Variable(tf.random_normal([num_classes]))
}

model = multilayer_perceptron(x, weights, biases)

cost = tf.reduce_mean(tf.nn.sigmoid_cross_entropy_with_logits(logits=model, labels=y))
optimizer = tf.train.AdamOptimizer(learning_rate=learning_rate).minimize(cost)

init = tf.global_variables_initializer()

with tf.Session() as sess:
	sess.run(init)

	#Training
	for epoch in range(training_epochs):
		average_cost = 0.
		total_batch = int(num_examples/batch_size)
		for i in range(total_batch):
			batch_x = data[i*batch_size:(i+1)*batch_size,:]
			batch_y = labels[i*batch_size:(i+1)*batch_size,:]

			_, c = sess.run([optimizer, cost],
				feed_dict = {x: batch_x, y: batch_y})

			# Compute loss
			average_cost += c / total_batch

		if epoch % 10 == 0:
			print("Epoch:" + str(epoch+1) + " cost = " + str(average_cost))
	print("Optimization Finished!")

	correct = tf.equal(tf.argmax(model, 1), tf.argmax(y, 1))

	accuracy = tf.reduce_mean(tf.cast(correct, "float"))
	print("Accuracy: " + str(accuracy.eval({x:test_data, y:test_labels})))<|endoftext|>


<commit_before>#!/usr/bin/env python
"""EPP script to perform basic calculations on UDF:s in Clarity LIMS
Command to trigger this script:
bash -c "PATH/TO/INSTALLED/SCRIPT
--username {username}
--password {password}
--baseuri YOUR.URI
--pid {processLuid}
--log {compoundOutputFileLuidN}"
--output_files {outputFileLuids}
--udf1 NameOfUDF 1
--operator Mathematical operator
--udf2 NameOfUDF 2
--result_udf NameOfUDF to store result in
"

Johannes Alneberg, Science for Life Laboratory, Stockholm, Sweden
"""
from argparse import ArgumentParser

from genologics.lims import Lims
from genologics.entities import Artifact, Process
from genologics.epp import configure_logging,setup_standard_parser

def apply_calculations(lims,input_artifacts,udf1,op,udf2,result_udf):
    print 'result_udf: {0}, udf1: {1}, operator: {2}, udf2: {3}'.format(
        result_udf,udf1,op,udf2)
    for artifact in input_artifacts:
        try:
            result_udf = artifact.udf[result_udf]
            udf1=artifact.udf[udf1]
            udf2 = artifact.udf[udf2]
        except:
            print artifact.udf.items()
            raise
        print 'Updating: Artifact id: {0}, result_udf: {1}, udf1: {2}, operator: {3}, udf2: {4}'.format(
            artifact.id, artifact.udf[result_udf],artifact.udf[udf1],operator,artifact.udf[udf2])
        artifact.udf[result_udf] = eval('{0}{1}{2}'.format(artifact.udf[udf1],op,artifact.udf[udf2]))
        artifact.put()
        print 'Updated {0} to {1}.'.format(result_udf,artifact.udf[result_udf])

def main(lims,args):
    p = Process(lims,id = args.pid)
    input_ids = map(lambda io: io[0]['limsid'],p.input_output_maps)
    input_set = frozenset(input_ids)
    inputs = map(lambda id: Artifact(lims,id=id),list(input_set))

    apply_calculations(lims,inputs,args.udf1,args.operator,args.udf2,args.result_udf)

if __name__ == "__main__":
    # Initialize parser with standard arguments and description
    desc = """EPP script to perform basic calculations on UDF:s in Clarity LIMS.
    result_udf=udf1 *operator* udf2"""
    parser = setup_standard_parser(description=desc)

    # Additional arguments
    parser.add_argument('--output_files',nargs='*',
                        help='Lims unique ids for each output file artifact')
    parser.add_argument('--udf1',
                        help='The first udf in the formula')
    parser.add_argument('--operator', choices =['+','-','*'],
                        help='operator to apply')
    parser.add_argument('--udf2',
                        help='The second udf in the formula')
    parser.add_argument('--result_udf',
                        help='Udf to store the result')
    args = parser.parse_args()

    # Start logging
    if args.log:
        configure_logging(args.log)
    lims = Lims(args.baseuri,args.username,args.password)
    lims.check_version()

    main(lims, args)

<commit_msg>No overwriting in debuggin<commit_after>#!/usr/bin/env python
"""EPP script to perform basic calculations on UDF:s in Clarity LIMS
Command to trigger this script:
bash -c "PATH/TO/INSTALLED/SCRIPT
--username {username}
--password {password}
--baseuri YOUR.URI
--pid {processLuid}
--log {compoundOutputFileLuidN}"
--output_files {outputFileLuids}
--udf1 NameOfUDF 1
--operator Mathematical operator
--udf2 NameOfUDF 2
--result_udf NameOfUDF to store result in
"

Johannes Alneberg, Science for Life Laboratory, Stockholm, Sweden
"""
from argparse import ArgumentParser

from genologics.lims import Lims
from genologics.entities import Artifact, Process
from genologics.epp import configure_logging,setup_standard_parser

def apply_calculations(lims,input_artifacts,udf1,op,udf2,result_udf):
    print 'result_udf: {0}, udf1: {1}, operator: {2}, udf2: {3}'.format(
        result_udf,udf1,op,udf2)
    for artifact in input_artifacts:
        try:
            result_udf_val = artifact.udf[result_udf]
            udf1_val=artifact.udf[udf1]
            udf2_val = artifact.udf[udf2]
        except:
            print artifact.udf.items()
            raise
        print 'Updating: Artifact id: {0}, result_udf: {1}, udf1: {2}, operator: {3}, udf2: {4}'.format(
            artifact.id, artifact.udf[result_udf],artifact.udf[udf1],operator,artifact.udf[udf2])
        artifact.udf[result_udf] = eval('{0}{1}{2}'.format(artifact.udf[udf1],op,artifact.udf[udf2]))
        artifact.put()
        print 'Updated {0} to {1}.'.format(result_udf,artifact.udf[result_udf])

def main(lims,args):
    p = Process(lims,id = args.pid)
    input_ids = map(lambda io: io[0]['limsid'],p.input_output_maps)
    input_set = frozenset(input_ids)
    inputs = map(lambda id: Artifact(lims,id=id),list(input_set))

    apply_calculations(lims,inputs,args.udf1,args.operator,args.udf2,args.result_udf)

if __name__ == "__main__":
    # Initialize parser with standard arguments and description
    desc = """EPP script to perform basic calculations on UDF:s in Clarity LIMS.
    result_udf=udf1 *operator* udf2"""
    parser = setup_standard_parser(description=desc)

    # Additional arguments
    parser.add_argument('--output_files',nargs='*',
                        help='Lims unique ids for each output file artifact')
    parser.add_argument('--udf1',
                        help='The first udf in the formula')
    parser.add_argument('--operator', choices =['+','-','*'],
                        help='operator to apply')
    parser.add_argument('--udf2',
                        help='The second udf in the formula')
    parser.add_argument('--result_udf',
                        help='Udf to store the result')
    args = parser.parse_args()

    # Start logging
    if args.log:
        configure_logging(args.log)
    lims = Lims(args.baseuri,args.username,args.password)
    lims.check_version()

    main(lims, args)

<|endoftext|>


<commit_before>#! /usr/bin/env python
#
# Protein Engineering Analysis Tool Structure Analysis (PEATSA)
# Copyright (C) 2010 Michael Johnston & Jens Erik Nielsen
#
# Author: Michael Johnston
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.
#
# Contact information:
# Email: Jens.Nielsen_at_gmail.com
# Normal mail:
# Jens Nielsen
# SBBS, Conway Institute
# University College Dublin
# Dublin 4, Ireland
#

'''WebScript - Wrapper around ProteinDesignTool.py containing web related extensions.

	Written by Michael Johnston
	This programs copyright - Michael Johnston and Jens Nielsen 2008.
	For copyright of consituent programs check each

	NOTE: For detailed instructions see the programs README file.
	This tool should only be run from the JobSubmission.RunJob function or the RemoteLaunch script.
	It requires access to the PDT web-apps MySQL database

	Usage: WebScript -p [pdbname] -w [directory] -j [jobid] [optional arguments]

	-p - The name of the pdb the tool is to be run on

	-w - The directory where the program will be run. It MUST contain the pdb called [pdbname]
	     For certain calculations this directory must contain various other data files.
	     See README for more.

	-j - The JobID for the web-app job

	All other options are the same as for ProteinDesignTool.
	Except that the directory containing the output will be outputDir/jobid.
	The value for outputDir is determined by the entry in the webApplication.conf file
	and is not the directory from which the program is run.'''

import getopt, sys, os, traceback, math
import Data
import UtilityFunctions, Exceptions
import PEATSA.Core as Core

class CommandLineParser(Core.Utilities.CommandLineParser):

	'''Class for parsing the WebScripts command line and providing access to the data

	Adds handling of a couple of special features of the WebScript cl'''

	def parseCommandLine(self):

		'''Parses the command line

		Execeptions:
			Raises Core.Exceptions.ArgumentError if there was a problem with the arguments supplied'''

		try:
			self.opt = getopt.getopt(sys.argv[1:],
					"vw:p:o:j:h",
					["pKa", "scan", "path", "modes", "stability",
					"mutationList=", "ionisableGroups=","mutation=", "mutants=", "ligand=", "configurationFile=", "create-configuration"])
			self.opt = dict(self.opt[0])
		except getopt.GetoptError, data:
			raise Core.Exceptions.ArgumentError, data

		#Create configuration instance
		#This is needed for reading the default data and output directories
		if self.opt.has_key("--configurationFile"):
			configurationFile = self.opt["--configurationFile"]
		else:
			#Use default configuration for the web application
			configurationFile = os.path.join(UtilityFunctions.ResourceDirectory(), "webApplication.conf")
			self.opt["--configurationFile"] = configurationFile

		configuration = Core.Environment.Configuration(filename=configurationFile)

		#
		#Check required args
		#

		if not self.opt.has_key("-w") and not self.createConfigurationRequested():
			raise Core.Exceptions.ArgumentError, "Working directory (-w) must be specified"

		if not self.opt.has_key("-p") and not self.createConfigurationRequested():
			raise Core.Exceptions.ArgumentError, "PDB (-p) must be specified"

		if not self.opt.has_key("-j"):
			raise Core.Exceptions.ArgumentError, "A job id must be specified"

		if not self.opt.has_key("-o"):
			try:
				 self.opt["-o"] = os.path.join(os.getcwd(), self.opt["-j"])
				 self.environment.output('[PEAT-SA] No output directory specified - Defaulting to %s' % self.opt["-o"])
			except:
				Core.Exceptions.ArgumentError, "Output directory (-w) not specified and default value not available in configuration file"
		else:
			self.environment.output('[PEAT-SA] Output directory is %s' % self.opt["-o"])

		#
		#Check if any of the optional flags were specified
		#

		counts = [self.opt.keys().count(option) for option in self.optionalArguments]
		value = reduce(lambda x,y: x+y, counts)
		if value == 0 and  not self.createConfigurationRequested():
			self.environment.output('[PEAT-SA] Performing pka scan, stability analysis, ENM analysis and transition path finding')
			#Add the options arguments
			for option in self.optionalArguments:
				if option != '--mutationList':
					self.opt[option] = ""

		#Make any supplied paths absolute
		self.__fixPaths__()

	def jobID(self):

		if not self.opt.has_key('-j'):
			return None

		return self.opt['-j']

class WebScript:

	def __init__(self):

		#Get environment
		self.environment = Core.Environment.Environment()

		#Parse command line
		self.parser = CommandLineParser()
		self.parser.parseCommandLine()

		#Read Configuration, connect to database and create Job object.
		self.configuration = Core.Environment.Configuration(filename=self.parser.configurationFile())
		self.connection = None
		self.job = None
		if self.environment.isRoot():
			self.connection = UtilityFunctions.ConnectionFromConfiguration(self.configuration)
			print '[WEBAPP] Connection %s' % self.connection
			print '[WEBAPP] Retrieving job data'
			self.job = Data.Job(jobID=self.parser.jobID(), connection=self.connection)
			if not self.job.exists():
				raise Exceptions.DatabaseRetrievalError, 'Data for job %d is not present in database' % self.job.identification
			else:
				print '[WEBAPP] Obtained job data'
				self.job.setState('Running')
				self.job.setQueueStatusMessage('Your job is running')

		self.environment.wait()

		#Set the $HOME environment variable for apache to the path to the .WHAT_IF file
		#specifying the path to the whatif executable (used by pKarun.WHATIF.py)
		os.environ['HOME'] = self.configuration.get('PATHS', 'whatif')

		#Create the output directory
		if self.environment.isRoot() and self.parser.outputDirectory() is not None:
			if not os.path.exists(self.parser.outputDirectory()):
				os.mkdir(self.parser.outputDirectory())

		self.environment.wait()

	def main(self):

		#Set all job states to waiting
		if self.environment.isRoot():
			calculationStates = self.job.calculationStates()
			for calculation in calculationStates.keys():
				if calculationStates[calculation] == 'Queued':
					self.job.setCalculationState(calculation, 'Waiting')

		#If a delta-pka calculation is requested check that the directory exists
		#If it doesn't - raise an error
		if self.parser.scan() is True and not os.path.exists(self.parser.workingDirectory()):
			raise Core.Exceptions.MissingPKADataError, "Delta pKa scan requested but specified working dir does not exist"

		#Create the ProteinDesignTool instance - this does all cleaning etc.
		tool = Core.ProteinDesignTool.ProteinDesignTool(self.parser.configurationFile(),
				workingDirectory=self.parser.workingDirectory(),
				pdbFile=self.parser.pdbFile(),
				outputDirectory=self.parser.outputDirectory())

		#Check if a previously created MutantCollection was specified
		mutantCollection = self.parser.mutants()
		if mutantCollection is None:
			self.environment.output('[PEAT-SA] No mutants available from previous run.\n')
			#Note we use the ProteinDesignTool instances pdbFile attribute.
			#This points to the cleaned pdbFile the tool operates on.
			#Therefore we must use this to generate the mutants.
			mutantCollection = Core.ProteinDesignTool.CreateMutants(tool.pdbFile, tool.configuration, self.parser, self.environment)

		if self.environment.isRoot():
			self.job.addResults(matrix=mutantCollection.mutationScores, name='ModellingResults')

		#Check if any mutants were succesfully created
		if len(mutantCollection.mutantFiles()) == 0:
			raise Core.Exceptions.MutantCollectionError, 'None of the specified mutants could be modelled due to excessive clashes'
		#Do a scan if requested
		if self.parser.scan() is True:
			if self.environment.isRoot():
				self.job.setCalculationState('Scan', 'Running')
				#Add PDB file to job data since it wasn't uploaded by the user
				self.job.setStructureFromFile(self.parser.pdbFile())

			tool.runScan(mutantCollection=mutantCollection, verbose=self.parser.verbose())
			if self.environment.isRoot():
				self.job.addResults(matrix=tool.dataDirectory.scanResults, name='ScanResults')
				self.job.setCalculationState('Scan', 'Finished')
			self.environment.wait()

		#Get the mutants for the stability and binding calculations if these were requested
		if self.parser.stability() is True:
			if self.environment.isRoot():
				self.job.setCalculationState('Stability', 'Running')

			tool.runStabilityCalculation(mutantFiles=mutantCollection.mutantFiles(), verbose=self.parser.verbose())
			if self.environment.isRoot():
				self.job.addResults(matrix=tool.dataDirectory.stabilityResults, name='StabilityResults')
				#Create graph
				points = len(tool.dataDirectory.stabilityResults.total)
				data = Core.Utilities.CreateBarChartData(matrix=tool.dataDirectory.stabilityResults, column="Total",
						xlabel="Mutant Index", ylabel="KJ/mol", title="Stability",
						xticIncrement=int(math.ceil(points/20.0)))
				self.job.addImageData(data=data, name="StabilityImage")
				self.job.setCalculationState('Stability', 'Finished')
			self.environment.wait()

		if self.parser.binding() is True:
			if self.environment.isRoot():
				self.job.setCalculationState('Binding', 'Running')

			tool.runBindingCalculation(mutantFiles=mutantCollection.mutantFiles(),
					ligandFile=self.parser.ligandFile(),
					verbose=self.parser.verbose())
			if self.environment.isRoot():
				self.job.addResults(matrix=tool.dataDirectory.deltaBindingResults, name='BindingResults')
				#Create graph
				points = len(tool.dataDirectory.deltaBindingResults.total)
				data = Core.Utilities.CreateBarChartData(matrix=tool.dataDirectory.deltaBindingResults, column="Total",
						xlabel="Mutant Index", ylabel="KJ/mol", title="Binding",
						xticIncrement=int(math.ceil(points/20.0)))
				self.job.addImageData(data=data, name="BindingImage")
				self.job.setCalculationState('Binding', 'Finished')

			self.environment.wait()

		#FIXME: Delete temp data??
		if self.environment.isRoot():
			self.job.setState('Finished')
			self.connection.close()

		#Necessary in case something goes wrong in above root only section
		#to force all processes to quit
		self.environment.wait()

if __name__ == "__main__":

	tool = None
	environment = None

	try:
		environment = Core.Environment.Environment()
		tool = WebScript()
		tool.main()
	except Core.Exceptions.MissingPKADataError, data:
		print "\nError - %s" % data.errorString
		print "Description - %s\n" % data
		if tool is not None and environment.isRoot():
			code = os.path.split(tool.parser.workingDirectory())[1]
			tool.job.setError("No pKa data corresponding to supplied code %s is present" % code,
				"You must have run a previous pKa calculation on this pdb using the pKD server")
			tool.job.setState("Finished")
		#In serial this does nothing
		#In parallel if all processes are here they exit
		#In addition if only one process raised the exception and is here,
		#the others will be waiting elsewhere in the code
		#In this case, passing True here causes all the processes to exit.
		environment.wait(error=True)
		#This is only called if we are in serial
		environment.exit(1)
	except Core.Exceptions.ProteinDesignToolException, data:
		print "\nError - %s" % data.errorString
		print "Description - %s\n" % data
		if tool is not None and environment.isRoot():
			tool.job.setError(data.errorString,data.message)
			tool.job.setState("Finished")

		if data.__class__ == Core.Exceptions.ArgumentError:
			print __doc__

		environment.wait(error=True)
		environment.exit(1)
	except SystemExit, data:
		#This is raised by sys.exit - check was the arg 0
		if data.message != 0:
			if tool is not None and environment.isRoot():
				tool.job.setError("Encountered an unexpected error - This is likely a bug",
						"Program exited with exit code %d" % data.message)
				tool.job.setState("Finished")

			print "\nEncountered an unexpected error - This is likely a bug"
			print "Please contact the developers with the following traceback information",
			print "along with details on the calculation being run (arguments, configuration file etc.)\n"
			print "Program exited with exit code %d" % data.message

		environment.wait(error=True)
		environment.exit(1)
	except BaseException, data:
		#Get the traceback as a string = convoluted ....
		trace = traceback.format_list(traceback.extract_tb(sys.exc_info()[2]))
		trace = reduce(lambda x, y: x + y, trace)
		#Need to remove single quotes from the traceback
		trace = trace.replace("'", " ")
		if tool is not None and environment.isRoot():
			tool.job.setError("Encountered an unexpected error - This is likely a bug",
					"Check error logs for more information.")
			tool.job.setState("Finished")

		print "\nEncountered an unexpected error - This is likely a bug"
		print "Please contact the developers with the following traceback information",
		print "along with details on the calculation being run (arguments, configuration file etc.)\n"
		print "%s" % data
		environment.wait(error=True)
		environment.exit(1)

<commit_msg>Print traceback<commit_after>#! /usr/bin/env python
#
# Protein Engineering Analysis Tool Structure Analysis (PEATSA)
# Copyright (C) 2010 Michael Johnston & Jens Erik Nielsen
#
# Author: Michael Johnston
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.
#
# Contact information:
# Email: Jens.Nielsen_at_gmail.com
# Normal mail:
# Jens Nielsen
# SBBS, Conway Institute
# University College Dublin
# Dublin 4, Ireland
#

'''WebScript - Wrapper around ProteinDesignTool.py containing web related extensions.

	Written by Michael Johnston
	This programs copyright - Michael Johnston and Jens Nielsen 2008.
	For copyright of consituent programs check each

	NOTE: For detailed instructions see the programs README file.
	This tool should only be run from the JobSubmission.RunJob function or the RemoteLaunch script.
	It requires access to the PDT web-apps MySQL database

	Usage: WebScript -p [pdbname] -w [directory] -j [jobid] [optional arguments]

	-p - The name of the pdb the tool is to be run on

	-w - The directory where the program will be run. It MUST contain the pdb called [pdbname]
	     For certain calculations this directory must contain various other data files.
	     See README for more.

	-j - The JobID for the web-app job

	All other options are the same as for ProteinDesignTool.
	Except that the directory containing the output will be outputDir/jobid.
	The value for outputDir is determined by the entry in the webApplication.conf file
	and is not the directory from which the program is run.'''

import getopt, sys, os, traceback, math
import Data
import UtilityFunctions, Exceptions
import PEATSA.Core as Core

class CommandLineParser(Core.Utilities.CommandLineParser):

	'''Class for parsing the WebScripts command line and providing access to the data

	Adds handling of a couple of special features of the WebScript cl'''

	def parseCommandLine(self):

		'''Parses the command line

		Execeptions:
			Raises Core.Exceptions.ArgumentError if there was a problem with the arguments supplied'''

		try:
			self.opt = getopt.getopt(sys.argv[1:],
					"vw:p:o:j:h",
					["pKa", "scan", "path", "modes", "stability",
					"mutationList=", "ionisableGroups=","mutation=", "mutants=", "ligand=", "configurationFile=", "create-configuration"])
			self.opt = dict(self.opt[0])
		except getopt.GetoptError, data:
			raise Core.Exceptions.ArgumentError, data

		#Create configuration instance
		#This is needed for reading the default data and output directories
		if self.opt.has_key("--configurationFile"):
			configurationFile = self.opt["--configurationFile"]
		else:
			#Use default configuration for the web application
			configurationFile = os.path.join(UtilityFunctions.ResourceDirectory(), "webApplication.conf")
			self.opt["--configurationFile"] = configurationFile

		configuration = Core.Environment.Configuration(filename=configurationFile)

		#
		#Check required args
		#

		if not self.opt.has_key("-w") and not self.createConfigurationRequested():
			raise Core.Exceptions.ArgumentError, "Working directory (-w) must be specified"

		if not self.opt.has_key("-p") and not self.createConfigurationRequested():
			raise Core.Exceptions.ArgumentError, "PDB (-p) must be specified"

		if not self.opt.has_key("-j"):
			raise Core.Exceptions.ArgumentError, "A job id must be specified"

		if not self.opt.has_key("-o"):
			try:
				 self.opt["-o"] = os.path.join(os.getcwd(), self.opt["-j"])
				 self.environment.output('[PEAT-SA] No output directory specified - Defaulting to %s' % self.opt["-o"])
			except:
				Core.Exceptions.ArgumentError, "Output directory (-w) not specified and default value not available in configuration file"
		else:
			self.environment.output('[PEAT-SA] Output directory is %s' % self.opt["-o"])

		#
		#Check if any of the optional flags were specified
		#

		counts = [self.opt.keys().count(option) for option in self.optionalArguments]
		value = reduce(lambda x,y: x+y, counts)
		if value == 0 and  not self.createConfigurationRequested():
			self.environment.output('[PEAT-SA] Performing pka scan, stability analysis, ENM analysis and transition path finding')
			#Add the options arguments
			for option in self.optionalArguments:
				if option != '--mutationList':
					self.opt[option] = ""

		#Make any supplied paths absolute
		self.__fixPaths__()

	def jobID(self):

		if not self.opt.has_key('-j'):
			return None

		return self.opt['-j']

class WebScript:

	def __init__(self):

		#Get environment
		self.environment = Core.Environment.Environment()

		#Parse command line
		self.parser = CommandLineParser()
		self.parser.parseCommandLine()

		#Read Configuration, connect to database and create Job object.
		self.configuration = Core.Environment.Configuration(filename=self.parser.configurationFile())
		self.connection = None
		self.job = None
		if self.environment.isRoot():
			self.connection = UtilityFunctions.ConnectionFromConfiguration(self.configuration)
			print '[WEBAPP] Connection %s' % self.connection
			print '[WEBAPP] Retrieving job data'
			self.job = Data.Job(jobID=self.parser.jobID(), connection=self.connection)
			if not self.job.exists():
				raise Exceptions.DatabaseRetrievalError, 'Data for job %d is not present in database' % self.job.identification
			else:
				print '[WEBAPP] Obtained job data'
				self.job.setState('Running')
				self.job.setQueueStatusMessage('Your job is running')

		self.environment.wait()

		#Set the $HOME environment variable for apache to the path to the .WHAT_IF file
		#specifying the path to the whatif executable (used by pKarun.WHATIF.py)
		os.environ['HOME'] = self.configuration.get('PATHS', 'whatif')

		#Create the output directory
		if self.environment.isRoot() and self.parser.outputDirectory() is not None:
			if not os.path.exists(self.parser.outputDirectory()):
				os.mkdir(self.parser.outputDirectory())

		self.environment.wait()

	def main(self):

		#Set all job states to waiting
		if self.environment.isRoot():
			calculationStates = self.job.calculationStates()
			for calculation in calculationStates.keys():
				if calculationStates[calculation] == 'Queued':
					self.job.setCalculationState(calculation, 'Waiting')

		#If a delta-pka calculation is requested check that the directory exists
		#If it doesn't - raise an error
		if self.parser.scan() is True and not os.path.exists(self.parser.workingDirectory()):
			raise Core.Exceptions.MissingPKADataError, "Delta pKa scan requested but specified working dir does not exist"

		#Create the ProteinDesignTool instance - this does all cleaning etc.
		tool = Core.ProteinDesignTool.ProteinDesignTool(self.parser.configurationFile(),
				workingDirectory=self.parser.workingDirectory(),
				pdbFile=self.parser.pdbFile(),
				outputDirectory=self.parser.outputDirectory())

		#Check if a previously created MutantCollection was specified
		mutantCollection = self.parser.mutants()
		if mutantCollection is None:
			self.environment.output('[PEAT-SA] No mutants available from previous run.\n')
			#Note we use the ProteinDesignTool instances pdbFile attribute.
			#This points to the cleaned pdbFile the tool operates on.
			#Therefore we must use this to generate the mutants.
			mutantCollection = Core.ProteinDesignTool.CreateMutants(tool.pdbFile, tool.configuration, self.parser, self.environment)

		if self.environment.isRoot():
			self.job.addResults(matrix=mutantCollection.mutationScores, name='ModellingResults')

		#Check if any mutants were succesfully created
		if len(mutantCollection.mutantFiles()) == 0:
			raise Core.Exceptions.MutantCollectionError, 'None of the specified mutants could be modelled due to excessive clashes'
		#Do a scan if requested
		if self.parser.scan() is True:
			if self.environment.isRoot():
				self.job.setCalculationState('Scan', 'Running')
				#Add PDB file to job data since it wasn't uploaded by the user
				self.job.setStructureFromFile(self.parser.pdbFile())

			tool.runScan(mutantCollection=mutantCollection, verbose=self.parser.verbose())
			if self.environment.isRoot():
				self.job.addResults(matrix=tool.dataDirectory.scanResults, name='ScanResults')
				self.job.setCalculationState('Scan', 'Finished')
			self.environment.wait()

		#Get the mutants for the stability and binding calculations if these were requested
		if self.parser.stability() is True:
			if self.environment.isRoot():
				self.job.setCalculationState('Stability', 'Running')

			tool.runStabilityCalculation(mutantFiles=mutantCollection.mutantFiles(), verbose=self.parser.verbose())
			if self.environment.isRoot():
				self.job.addResults(matrix=tool.dataDirectory.stabilityResults, name='StabilityResults')
				#Create graph
				points = len(tool.dataDirectory.stabilityResults.total)
				data = Core.Utilities.CreateBarChartData(matrix=tool.dataDirectory.stabilityResults, column="Total",
						xlabel="Mutant Index", ylabel="KJ/mol", title="Stability",
						xticIncrement=int(math.ceil(points/20.0)))
				self.job.addImageData(data=data, name="StabilityImage")
				self.job.setCalculationState('Stability', 'Finished')
			self.environment.wait()

		if self.parser.binding() is True:
			if self.environment.isRoot():
				self.job.setCalculationState('Binding', 'Running')

			tool.runBindingCalculation(mutantFiles=mutantCollection.mutantFiles(),
					ligandFile=self.parser.ligandFile(),
					verbose=self.parser.verbose())
			if self.environment.isRoot():
				self.job.addResults(matrix=tool.dataDirectory.deltaBindingResults, name='BindingResults')
				#Create graph
				points = len(tool.dataDirectory.deltaBindingResults.total)
				data = Core.Utilities.CreateBarChartData(matrix=tool.dataDirectory.deltaBindingResults, column="Total",
						xlabel="Mutant Index", ylabel="KJ/mol", title="Binding",
						xticIncrement=int(math.ceil(points/20.0)))
				self.job.addImageData(data=data, name="BindingImage")
				self.job.setCalculationState('Binding', 'Finished')

			self.environment.wait()

		#FIXME: Delete temp data??
		if self.environment.isRoot():
			self.job.setState('Finished')
			self.connection.close()

		#Necessary in case something goes wrong in above root only section
		#to force all processes to quit
		self.environment.wait()

if __name__ == "__main__":

	tool = None
	environment = None

	try:
		environment = Core.Environment.Environment()
		tool = WebScript()
		tool.main()
	except Core.Exceptions.MissingPKADataError, data:
		print "\nError - %s" % data.errorString
		print "Description - %s\n" % data
		if tool is not None and environment.isRoot():
			code = os.path.split(tool.parser.workingDirectory())[1]
			tool.job.setError("No pKa data corresponding to supplied code %s is present" % code,
				"You must have run a previous pKa calculation on this pdb using the pKD server")
			tool.job.setState("Finished")
		#In serial this does nothing
		#In parallel if all processes are here they exit
		#In addition if only one process raised the exception and is here,
		#the others will be waiting elsewhere in the code
		#In this case, passing True here causes all the processes to exit.
		environment.wait(error=True)
		#This is only called if we are in serial
		environment.exit(1)
	except Core.Exceptions.ProteinDesignToolException, data:
		print "\nError - %s" % data.errorString
		print "Description - %s\n" % data
		if tool is not None and environment.isRoot():
			tool.job.setError(data.errorString,data.message)
			tool.job.setState("Finished")

		if data.__class__ == Core.Exceptions.ArgumentError:
			print __doc__

		environment.wait(error=True)
		environment.exit(1)
	except SystemExit, data:
		#This is raised by sys.exit - check was the arg 0
		if data.message != 0:
			if tool is not None and environment.isRoot():
				tool.job.setError("Encountered an unexpected error - This is likely a bug",
						"Program exited with exit code %d" % data.message)
				tool.job.setState("Finished")

			print "\nEncountered an unexpected error - This is likely a bug"
			print "Please contact the developers with the following traceback information",
			print "along with details on the calculation being run (arguments, configuration file etc.)\n"
			print "Program exited with exit code %d" % data.message

		environment.wait(error=True)
		environment.exit(1)
	except BaseException, data:
		#Get the traceback as a string = convoluted ....
		trace = traceback.format_list(traceback.extract_tb(sys.exc_info()[2]))
		trace = reduce(lambda x, y: x + y, trace)
		#Need to remove single quotes from the traceback
		trace = trace.replace("'", " ")
		if tool is not None and environment.isRoot():
			tool.job.setError("Encountered an unexpected error - This is likely a bug",
					"Check error logs for more information.")
			tool.job.setState("Finished")

		print "\nEncountered an unexpected error - This is likely a bug"
		print "Please contact the developers with the following traceback information",
		print "along with details on the calculation being run (arguments, configuration file etc.)\n"
		print "%s" % data
		print "%s" % trace
		environment.wait(error=True)
		environment.exit(1)

<|endoftext|>


<commit_before>"""A script for fetching public utility data from reporting agency servers."""

import os
import sys
import argparse

# This is a hack to make the pudl package importable from within this script,
# even though it isn't in one of the normal site-packages directories where
# Python typically searches.  When we have some real installation/packaging
# happening, this will no longer be necessary.
sys.path.append(os.path.abspath('..'))

def parse_command_line(argv):
    """
    Parse command line arguments. See the -h option.

    :param argv: arguments on the command line must include caller file name.
    """
    from pudl import settings
    from pudl import constants
    parser = argparse.ArgumentParser()

    parser.add_argument(
        '-q',
        '--quiet',
        dest='verbose',
        action='store_false',
        help="Quiet mode. Suppress download progress indicators and warnings.",
        default=True
    )
    parser.add_argument(
        '-z',
        '--zip',
        dest='unzip',
        action='store_false',
        help="Do not unzip downloaded data files.",
        default=True
    )
    parser.add_argument(
        '-c',
        '--clobber',
        action='store_true',
        help="Clobber existing zipfiles in the datastore if they exist.",
        default=False
    )
    parser.add_argument(
        '-d',
        '--datadir',
        type=str,
        help="""Path to the top level datastore directory. (default:
        %(default)s).""",
        default=settings.DATA_DIR
    )
    parser.add_argument(
        '-s',
        '--sources',
        nargs='+',
        choices=constants.data_sources,
        help="""List of data sources which should be downloaded.
        (default: %(default)s).""",
        default=constants.data_sources
    )
    parser.add_argument(
        '-y',
        '--years',
        dest='year',
        nargs='+',
        help="""List of years for which data should be downloaded. Different
        data sources have differet valid years. If data is not available for a
        specified year and data source, it will be ignored. If no years are
        specified, all available data will be downloaded for all requested data
        sources.""",
        default=[]
    )
    parser.add_argument(
        '-n',
        '--no-download',
        dest='no_download',
        action='store_true',
        help="Do not download data files, only unzip ones that are already present.",
        default=False
    )
    parser.add_argument(
        '-t',
        '--states',
        nargs='+',
        choices=constants.cems_states.keys(),
        help="""List of two letter US state abbreviations indicating which
        states data should be downloaded. Currently only applicable to the EPA's
        CEMS dataset.""",
        default=constants.cems_states.keys()
    )

    arguments = parser.parse_args(argv[1:])
    return arguments

def main():
    """Main function controlling flow of the script."""
    from pudl import datastore
    from pudl import constants
    import concurrent.futures

    args = parse_command_line(sys.argv)

    # Generate a list of valid years of data to download for each data source.
    # If no years were specified, use the full set of valid years.
    # If years were specified, keep only the years which are valid for that
    # data source, and optionally output a message saying which years are
    # being ignored because they aren't valid.
    yrs_by_src = {}
    for src in args.sources:
        if len(args.year) == 0:
            yrs_by_src[src] = constants.data_years[src]
        else:
            yrs_by_src[src] = [int(yr) for yr in args.year
                               if int(yr) in constants.data_years[src]]
            bad_yrs = [int(yr) for yr in args.year
                       if int(yr) not in constants.data_years[src]]
            if args.verbose and len(bad_yrs) > 0:
                print("Invalid {} years ignored: {}.".format(src, bad_yrs))

    with concurrent.futures.ThreadPoolExecutor() as executor:
        for src in args.sources:
            for yr in yrs_by_src[src]:
                executor.submit(datastore.update, src, yr, args.states,
                                clobber=args.clobber,
                                unzip=args.unzip,
                                verbose=args.verbose,
                                datadir=args.datadir,
                                no_download=args.no_download)

if __name__ == '__main__':
    sys.exit(main())
<commit_msg>Pointing the datastore script at the new config file<commit_after>"""A script for fetching public utility data from reporting agency servers."""

import os
import sys
import argparse

# This is a hack to make the pudl package importable from within this script,
# even though it isn't in one of the normal site-packages directories where
# Python typically searches.  When we have some real installation/packaging
# happening, this will no longer be necessary.
sys.path.append(os.path.abspath('..'))

def parse_command_line(argv):
    """
    Parse command line arguments. See the -h option.

    :param argv: arguments on the command line must include caller file name.
    """
    from config import SETTINGS
    from pudl import constants
    parser = argparse.ArgumentParser()

    parser.add_argument(
        '-q',
        '--quiet',
        dest='verbose',
        action='store_false',
        help="Quiet mode. Suppress download progress indicators and warnings.",
        default=True
    )
    parser.add_argument(
        '-z',
        '--zip',
        dest='unzip',
        action='store_false',
        help="Do not unzip downloaded data files.",
        default=True
    )
    parser.add_argument(
        '-c',
        '--clobber',
        action='store_true',
        help="Clobber existing zipfiles in the datastore if they exist.",
        default=False
    )
    parser.add_argument(
        '-d',
        '--datadir',
        type=str,
        help="""Path to the top level datastore directory. (default:
        %(default)s).""",
        default=SETTINGS['data_dir']
    )
    parser.add_argument(
        '-s',
        '--sources',
        nargs='+',
        choices=constants.data_sources,
        help="""List of data sources which should be downloaded.
        (default: %(default)s).""",
        default=constants.data_sources
    )
    parser.add_argument(
        '-y',
        '--years',
        dest='year',
        nargs='+',
        help="""List of years for which data should be downloaded. Different
        data sources have differet valid years. If data is not available for a
        specified year and data source, it will be ignored. If no years are
        specified, all available data will be downloaded for all requested data
        sources.""",
        default=[]
    )
    parser.add_argument(
        '-n',
        '--no-download',
        dest='no_download',
        action='store_true',
        help="Do not download data files, only unzip ones that are already present.",
        default=False
    )
    parser.add_argument(
        '-t',
        '--states',
        nargs='+',
        choices=constants.cems_states.keys(),
        help="""List of two letter US state abbreviations indicating which
        states data should be downloaded. Currently only applicable to the EPA's
        CEMS dataset.""",
        default=constants.cems_states.keys()
    )

    arguments = parser.parse_args(argv[1:])
    return arguments

def main():
    """Main function controlling flow of the script."""
    from pudl import datastore
    from pudl import constants
    import concurrent.futures

    args = parse_command_line(sys.argv)

    # Generate a list of valid years of data to download for each data source.
    # If no years were specified, use the full set of valid years.
    # If years were specified, keep only the years which are valid for that
    # data source, and optionally output a message saying which years are
    # being ignored because they aren't valid.
    yrs_by_src = {}
    for src in args.sources:
        if len(args.year) == 0:
            yrs_by_src[src] = constants.data_years[src]
        else:
            yrs_by_src[src] = [int(yr) for yr in args.year
                               if int(yr) in constants.data_years[src]]
            bad_yrs = [int(yr) for yr in args.year
                       if int(yr) not in constants.data_years[src]]
            if args.verbose and len(bad_yrs) > 0:
                print("Invalid {} years ignored: {}.".format(src, bad_yrs))

    with concurrent.futures.ThreadPoolExecutor() as executor:
        for src in args.sources:
            for yr in yrs_by_src[src]:
                executor.submit(datastore.update, src, yr, args.states,
                                clobber=args.clobber,
                                unzip=args.unzip,
                                verbose=args.verbose,
                                datadir=args.datadir,
                                no_download=args.no_download)

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