<commit_before>#-----------------------------------------------------------
# Multicorn, multiprocessing inspired by gunicorn
# TODO rename class: Multicorn -> Arbiter ?
#-----------------------------------------------------------
import errno
import fcntl
import logging
import os
import psutil
import random
import resource
import select
import signal
import socket
import sys
import time

import werkzeug.serving
try:
    from setproctitle import setproctitle
except ImportError:
    setproctitle = lambda x: None

import openerp
import openerp.tools.config as config

_logger = logging.getLogger(__name__)

class Multicorn(object):
    """ Multiprocessing inspired by (g)unicorn.
    Multicorn currently uses accept(2) as dispatching method between workers
    but we plan to replace it by a more intelligent dispatcher to will parse
    the first HTTP request line.
    """
    def __init__(self, app):
        # config
        self.address = (config['xmlrpc_interface'] or '0.0.0.0', config['xmlrpc_port'])
        self.population = config['workers']
        self.timeout = config['limit_time_real']
        self.limit_request = config['limit_request']
        # working vars
        self.beat = 4
        self.app = app
        self.pid = os.getpid()
        self.socket = None
        self.workers_http = {}
        self.workers_cron = {}
        self.workers = {}
        self.generation = 0
        self.queue = []

    def pipe_new(self):
        pipe = os.pipe()
        for fd in pipe:
            # non_blocking
            flags = fcntl.fcntl(fd, fcntl.F_GETFL) | os.O_NONBLOCK
            fcntl.fcntl(fd, fcntl.F_SETFL, flags)
            # close_on_exec
            flags = fcntl.fcntl(fd, fcntl.F_GETFD) | fcntl.FD_CLOEXEC
            fcntl.fcntl(fd, fcntl.F_SETFD, flags)
        return pipe

    def pipe_ping(self, pipe):
        try:
            os.write(pipe[1], '.')
        except IOError, e:
            if e.errno not in [errno.EAGAIN, errno.EINTR]:
                raise

    def signal_handler(self, sig, frame):
        if len(self.queue) < 5 or sig == signal.SIGCHLD:
            self.queue.append(sig)
            self.pipe_ping(self.pipe)
        else:
            _logger.warn("Dropping signal: %s", sig)

    def worker_spawn(self, klass, workers_registry):
        self.generation += 1
        worker = klass(self)
        pid = os.fork()
        if pid != 0:
            worker.pid = pid
            self.workers[pid] = worker
            workers_registry[pid] = worker
            return worker
        else:
            worker.run()
            sys.exit(0)

    def worker_pop(self, pid):
        if pid in self.workers:
            _logger.debug("Worker (%s) unregistered",pid)
            try:
                self.workers_http.pop(pid,None)
                self.workers_cron.pop(pid,None)
                u = self.workers.pop(pid)
                u.close()
            except OSError:
                return

    def worker_kill(self, pid, sig):
        try:
            os.kill(pid, sig)
        except OSError, e:
            if e.errno == errno.ESRCH:
                self.worker_pop(pid)

    def process_signals(self):
        while len(self.queue):
            sig = self.queue.pop(0)
            if sig in [signal.SIGINT,signal.SIGTERM]:
                raise KeyboardInterrupt

    def process_zombie(self):
        # reap dead workers
        while 1:
            try:
                wpid, status = os.waitpid(-1, os.WNOHANG)
                if not wpid:
                    break
                if (status >> 8) == 3:
                    msg = "Critial worker error (%s)"
                    _logger.critical(msg, wpid)
                    raise Exception(msg % wpid)
                self.worker_pop(wpid)
            except OSError, e:
                if e.errno == errno.ECHILD:
                    break
                raise

    def process_timeout(self):
        now = time.time()
        for (pid, worker) in self.workers.items():
            if now - worker.watchdog_time >= worker.watchdog_timeout:
                _logger.error("Worker (%s) timeout", pid)
                self.worker_kill(pid, signal.SIGKILL)

    def process_spawn(self):
        while len(self.workers_http) < self.population:
            self.worker_spawn(WorkerHTTP, self.workers_http)
        while len(self.workers_cron) < config['max_cron_threads']:
            self.worker_spawn(WorkerCron, self.workers_cron)

    def sleep(self):
        try:
            # map of fd -> worker
            fds = dict([(w.watchdog_pipe[0],w) for k,w in self.workers.items()])
            fd_in = fds.keys() + [self.pipe[0]]
            # check for ping or internal wakeups
            ready = select.select(fd_in, [], [], self.beat)
            # update worker watchdogs
            for fd in ready[0]:
                if fd in fds:
                    fds[fd].watchdog_time = time.time()
                try:
                    # empty pipe
                    while os.read(fd, 1):
                        pass
                except OSError, e:
                    if e.errno not in [errno.EAGAIN]:
                        raise
        except select.error, e:
            if e[0] not in [errno.EINTR]:
                raise

    def start(self):
        # wakeup pipe, python doesnt throw EINTR when a syscall is interrupted
        # by a signal simulating a pseudo SA_RESTART. We write to a pipe in the
        # signal handler to overcome this behaviour
        self.pipe = self.pipe_new()
        # set signal
        signal.signal(signal.SIGINT, self.signal_handler)
        signal.signal(signal.SIGTERM, self.signal_handler)
        signal.signal(signal.SIGCHLD, self.signal_handler)
        # listen to socket
        self.socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        self.socket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
        self.socket.setblocking(0)
        self.socket.bind(self.address)
        self.socket.listen(8)

    def stop(self, graceful=True):
        if graceful:
            _logger.info("Stopping gracefully")
            limit = time.time() + self.timeout
            for pid in self.workers.keys():
                self.worker_kill(pid, signal.SIGTERM)
            while self.workers and time.time() < limit:
                self.process_zombie()
                time.sleep(0.1)
        else:
            _logger.info("Stopping forcefully")
        for pid in self.workers.keys():
            self.worker_kill(pid, signal.SIGTERM)
        self.socket.close()
        openerp.cli.server.quit_signals_received = 1

    def run(self):
        self.start()
        _logger.debug("Multiprocess starting")
        while 1:
            try:
                #_logger.debug("Multiprocess beat (%s)",time.time())
                self.process_signals()
                self.process_zombie()
                self.process_timeout()
                self.process_spawn()
                self.sleep()
            except KeyboardInterrupt:
                _logger.debug("Multiprocess clean stop")
                self.stop()
                break
            except Exception,e:
                _logger.exception(e)
                self.stop(False)
                sys.exit(-1)

class Worker(object):
    """ Workers """
    def __init__(self, multi):
        self.multi = multi
        self.watchdog_time = time.time()
        self.watchdog_pipe = multi.pipe_new()
        self.watchdog_timeout = multi.timeout
        self.ppid = os.getpid()
        self.pid = None
        self.alive = True
        # should we rename into lifetime ?
        self.request_max = multi.limit_request
        self.request_count = 0

    def setproctitle(self, title=""):
        setproctitle('openerp: %s %s %s' % (self.__class__.__name__, self.pid, title))

    def close(self):
        os.close(self.watchdog_pipe[0])
        os.close(self.watchdog_pipe[1])

    def signal_handler(self, sig, frame):
        self.alive = False

    def sleep(self):
        try:
            ret = select.select([self.multi.socket], [], [], self.multi.beat)
        except select.error, e:
            if e[0] not in [errno.EINTR]:
                raise

    def process_limit(self):
        # If our parent changed sucide
        if self.ppid != os.getppid():
            _logger.info("Worker (%s) Parent changed", self.pid)
            self.alive = False
        # check for lifetime
        if self.request_count >= self.request_max:
            _logger.info("Worker (%d) max request (%s) reached.", self.pid, self.request_count)
            self.alive = False
        # Reset the worker if it consumes too much memory (e.g. caused by a memory leak).
        rss, vms = psutil.Process(os.getpid()).get_memory_info()
        if vms > config['limit_memory_soft']:
            _logger.info('Worker (%d) virtual memory limit (%s) reached.', self.pid, vms)
            self.alive = False # Commit suicide after the request.

        # VMS and RLIMIT_AS are the same thing: virtual memory, a.k.a. address space
        soft, hard = resource.getrlimit(resource.RLIMIT_AS)
        resource.setrlimit(resource.RLIMIT_AS, (config['limit_memory_hard'], hard))

        # SIGXCPU (exceeded CPU time) signal handler will raise an exception.
        r = resource.getrusage(resource.RUSAGE_SELF)
        cpu_time = r.ru_utime + r.ru_stime
        def time_expired(n, stack):
            _logger.info('Worker (%d) CPU time limit (%s) reached.', config['limit_time_cpu'])
            # We dont suicide in such case
            raise Exception('CPU time limit exceeded.')
        signal.signal(signal.SIGXCPU, time_expired)
        soft, hard = resource.getrlimit(resource.RLIMIT_CPU)
        resource.setrlimit(resource.RLIMIT_CPU, (cpu_time + config['limit_time_cpu'], hard))

    def process_work(self):
        pass

    def start(self):
        self.pid = os.getpid()
        self.setproctitle()
        _logger.info("Worker %s (%s) alive", self.__class__.__name__, self.pid)
        # Reseed the random number generator
        random.seed()
        # Prevent fd inherientence close_on_exec
        flags = fcntl.fcntl(self.multi.socket, fcntl.F_GETFD) | fcntl.FD_CLOEXEC
        fcntl.fcntl(self.multi.socket, fcntl.F_SETFD, flags)
        # reset blocking status
        self.multi.socket.setblocking(0)
        signal.signal(signal.SIGINT, self.signal_handler)
        signal.signal(signal.SIGTERM, signal.SIG_DFL)
        signal.signal(signal.SIGCHLD, signal.SIG_DFL)

    def stop(self):
        pass

    def run(self):
        try:
            self.start()
            while self.alive:
                self.process_limit()
                self.multi.pipe_ping(self.watchdog_pipe)
                self.sleep()
                self.process_work()
            _logger.info("Worker (%s) exiting. request_count: %s.", self.pid, self.request_count)
            self.stop()
        except Exception,e:
            _logger.exception("Worker (%s) Exception occured, exiting..." % self.pid)
            # should we use 3 to abort everything ?
            sys.exit(1)

class WorkerHTTP(Worker):
    """ HTTP Request workers """
    def process_request(self, client, addr):
        client.setblocking(1)
        client.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)
        # Prevent fd inherientence close_on_exec
        flags = fcntl.fcntl(client, fcntl.F_GETFD) | fcntl.FD_CLOEXEC
        fcntl.fcntl(client, fcntl.F_SETFD, flags)
        # do request using WorkerBaseWSGIServer monkey patched with socket
        self.server.socket = client
        # tolerate broken pipe when the http client closes the socket before
        # receiving the full reply
        try:
            self.server.process_request(client,addr)
        except IOError, e:
            if e.errno != errno.EPIPE:
                raise
        self.request_count += 1

    def process_work(self):
        try:
            client, addr = self.multi.socket.accept()
            self.process_request(client, addr)
        except socket.error, e:
            if e[0] not in (errno.EAGAIN, errno.ECONNABORTED):
                raise

    def start(self):
        Worker.start(self)
        self.server = WorkerBaseWSGIServer(self.multi.app)

class WorkerBaseWSGIServer(werkzeug.serving.BaseWSGIServer):
    """ werkzeug WSGI Server patched to allow using an external listen socket
    """
    def __init__(self, app):
        werkzeug.serving.BaseWSGIServer.__init__(self, "1", "1", app)
    def server_bind(self):
        # we dont bind beause we use the listen socket of Multicorn#socket
        # instead we close the socket
        if self.socket:
            self.socket.close()
    def server_activate(self):
        # dont listen as we use Multicorn#socket
        pass

class WorkerCron(Worker):
    """ Cron workers """

    def __init__(self, multi):
        super(WorkerCron, self).__init__(multi)
        # process_work() below process a single database per call.
        # The variable db_index is keeping track of the next database to
        # process.
        self.db_index = 0

    def sleep(self):
        # Really sleep once all the databases have been processed.
        if self.db_index == 0:
            interval = 60 + self.pid % 10 # chorus effect
            time.sleep(interval)

    def process_work(self):
        rpc_request = logging.getLogger('openerp.netsvc.rpc.request')
        rpc_request_flag = rpc_request.isEnabledFor(logging.DEBUG)
        _logger.debug("WorkerCron (%s) polling for jobs", self.pid)
        if config['db_name']:
            db_names = config['db_name'].split(',')
        else:
            db_names = openerp.netsvc.ExportService._services['db'].exp_list(True)
        if len(db_names):
            self.db_index = (self.db_index + 1) % len(db_names)
            db_name = db_names[self.db_index]
            self.setproctitle(db_name)
            if rpc_request_flag:
                start_time = time.time()
                start_rss, start_vms = psutil.Process(os.getpid()).get_memory_info()
            while True:
                # acquired = openerp.addons.base.ir.ir_cron.ir_cron._acquire_job(db_name)
                # TODO why isnt openerp.addons.base defined ?
                import base
                acquired = base.ir.ir_cron.ir_cron._acquire_job(db_name)
                if not acquired:
                    openerp.modules.registry.RegistryManager.delete(db_name)
                    break
            # dont keep cursors in multi database mode
            if len(db_names) > 1:
                openerp.sql_db.close_db(db_name)
            if rpc_request_flag:
                end_time = time.time()
                end_rss, end_vms = psutil.Process(os.getpid()).get_memory_info()
                logline = '%s time:%.3fs mem: %sk -> %sk (diff: %sk)' % (db_name, end_time - start_time, start_vms / 1024, end_vms / 1024, (end_vms - start_vms)/1024)
                _logger.debug("WorkerCron (%s) %s", self.pid, logline)

            self.request_count += 1
            if self.request_count >= self.request_max and self.request_max < len(db_names):
                _logger.error("There are more dabatases to process than allowed "
                    "by the `limit_request` configuration variable: %s more.",
                    len(db_names) - self.request_max)
        else:
            self.db_index = 0

    def start(self):
        Worker.start(self)
        openerp.service.start_internal()

# vim:expandtab:smartindent:tabstop=4:softtabstop=4:shiftwidth=4:
<commit_msg>[FIX] workers, listen backlog based on the number of workers<commit_after>#-----------------------------------------------------------
# Multicorn, multiprocessing inspired by gunicorn
# TODO rename class: Multicorn -> Arbiter ?
#-----------------------------------------------------------
import errno
import fcntl
import logging
import os
import psutil
import random
import resource
import select
import signal
import socket
import sys
import time

import werkzeug.serving
try:
    from setproctitle import setproctitle
except ImportError:
    setproctitle = lambda x: None

import openerp
import openerp.tools.config as config

_logger = logging.getLogger(__name__)

class Multicorn(object):
    """ Multiprocessing inspired by (g)unicorn.
    Multicorn currently uses accept(2) as dispatching method between workers
    but we plan to replace it by a more intelligent dispatcher to will parse
    the first HTTP request line.
    """
    def __init__(self, app):
        # config
        self.address = (config['xmlrpc_interface'] or '0.0.0.0', config['xmlrpc_port'])
        self.population = config['workers']
        self.timeout = config['limit_time_real']
        self.limit_request = config['limit_request']
        # working vars
        self.beat = 4
        self.app = app
        self.pid = os.getpid()
        self.socket = None
        self.workers_http = {}
        self.workers_cron = {}
        self.workers = {}
        self.generation = 0
        self.queue = []

    def pipe_new(self):
        pipe = os.pipe()
        for fd in pipe:
            # non_blocking
            flags = fcntl.fcntl(fd, fcntl.F_GETFL) | os.O_NONBLOCK
            fcntl.fcntl(fd, fcntl.F_SETFL, flags)
            # close_on_exec
            flags = fcntl.fcntl(fd, fcntl.F_GETFD) | fcntl.FD_CLOEXEC
            fcntl.fcntl(fd, fcntl.F_SETFD, flags)
        return pipe

    def pipe_ping(self, pipe):
        try:
            os.write(pipe[1], '.')
        except IOError, e:
            if e.errno not in [errno.EAGAIN, errno.EINTR]:
                raise

    def signal_handler(self, sig, frame):
        if len(self.queue) < 5 or sig == signal.SIGCHLD:
            self.queue.append(sig)
            self.pipe_ping(self.pipe)
        else:
            _logger.warn("Dropping signal: %s", sig)

    def worker_spawn(self, klass, workers_registry):
        self.generation += 1
        worker = klass(self)
        pid = os.fork()
        if pid != 0:
            worker.pid = pid
            self.workers[pid] = worker
            workers_registry[pid] = worker
            return worker
        else:
            worker.run()
            sys.exit(0)

    def worker_pop(self, pid):
        if pid in self.workers:
            _logger.debug("Worker (%s) unregistered",pid)
            try:
                self.workers_http.pop(pid,None)
                self.workers_cron.pop(pid,None)
                u = self.workers.pop(pid)
                u.close()
            except OSError:
                return

    def worker_kill(self, pid, sig):
        try:
            os.kill(pid, sig)
        except OSError, e:
            if e.errno == errno.ESRCH:
                self.worker_pop(pid)

    def process_signals(self):
        while len(self.queue):
            sig = self.queue.pop(0)
            if sig in [signal.SIGINT,signal.SIGTERM]:
                raise KeyboardInterrupt

    def process_zombie(self):
        # reap dead workers
        while 1:
            try:
                wpid, status = os.waitpid(-1, os.WNOHANG)
                if not wpid:
                    break
                if (status >> 8) == 3:
                    msg = "Critial worker error (%s)"
                    _logger.critical(msg, wpid)
                    raise Exception(msg % wpid)
                self.worker_pop(wpid)
            except OSError, e:
                if e.errno == errno.ECHILD:
                    break
                raise

    def process_timeout(self):
        now = time.time()
        for (pid, worker) in self.workers.items():
            if now - worker.watchdog_time >= worker.watchdog_timeout:
                _logger.error("Worker (%s) timeout", pid)
                self.worker_kill(pid, signal.SIGKILL)

    def process_spawn(self):
        while len(self.workers_http) < self.population:
            self.worker_spawn(WorkerHTTP, self.workers_http)
        while len(self.workers_cron) < config['max_cron_threads']:
            self.worker_spawn(WorkerCron, self.workers_cron)

    def sleep(self):
        try:
            # map of fd -> worker
            fds = dict([(w.watchdog_pipe[0],w) for k,w in self.workers.items()])
            fd_in = fds.keys() + [self.pipe[0]]
            # check for ping or internal wakeups
            ready = select.select(fd_in, [], [], self.beat)
            # update worker watchdogs
            for fd in ready[0]:
                if fd in fds:
                    fds[fd].watchdog_time = time.time()
                try:
                    # empty pipe
                    while os.read(fd, 1):
                        pass
                except OSError, e:
                    if e.errno not in [errno.EAGAIN]:
                        raise
        except select.error, e:
            if e[0] not in [errno.EINTR]:
                raise

    def start(self):
        # wakeup pipe, python doesnt throw EINTR when a syscall is interrupted
        # by a signal simulating a pseudo SA_RESTART. We write to a pipe in the
        # signal handler to overcome this behaviour
        self.pipe = self.pipe_new()
        # set signal
        signal.signal(signal.SIGINT, self.signal_handler)
        signal.signal(signal.SIGTERM, self.signal_handler)
        signal.signal(signal.SIGCHLD, self.signal_handler)
        # listen to socket
        self.socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        self.socket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
        self.socket.setblocking(0)
        self.socket.bind(self.address)
        self.socket.listen(8*self.population)

    def stop(self, graceful=True):
        if graceful:
            _logger.info("Stopping gracefully")
            limit = time.time() + self.timeout
            for pid in self.workers.keys():
                self.worker_kill(pid, signal.SIGTERM)
            while self.workers and time.time() < limit:
                self.process_zombie()
                time.sleep(0.1)
        else:
            _logger.info("Stopping forcefully")
        for pid in self.workers.keys():
            self.worker_kill(pid, signal.SIGTERM)
        self.socket.close()
        openerp.cli.server.quit_signals_received = 1

    def run(self):
        self.start()
        _logger.debug("Multiprocess starting")
        while 1:
            try:
                #_logger.debug("Multiprocess beat (%s)",time.time())
                self.process_signals()
                self.process_zombie()
                self.process_timeout()
                self.process_spawn()
                self.sleep()
            except KeyboardInterrupt:
                _logger.debug("Multiprocess clean stop")
                self.stop()
                break
            except Exception,e:
                _logger.exception(e)
                self.stop(False)
                sys.exit(-1)

class Worker(object):
    """ Workers """
    def __init__(self, multi):
        self.multi = multi
        self.watchdog_time = time.time()
        self.watchdog_pipe = multi.pipe_new()
        self.watchdog_timeout = multi.timeout
        self.ppid = os.getpid()
        self.pid = None
        self.alive = True
        # should we rename into lifetime ?
        self.request_max = multi.limit_request
        self.request_count = 0

    def setproctitle(self, title=""):
        setproctitle('openerp: %s %s %s' % (self.__class__.__name__, self.pid, title))

    def close(self):
        os.close(self.watchdog_pipe[0])
        os.close(self.watchdog_pipe[1])

    def signal_handler(self, sig, frame):
        self.alive = False

    def sleep(self):
        try:
            ret = select.select([self.multi.socket], [], [], self.multi.beat)
        except select.error, e:
            if e[0] not in [errno.EINTR]:
                raise

    def process_limit(self):
        # If our parent changed sucide
        if self.ppid != os.getppid():
            _logger.info("Worker (%s) Parent changed", self.pid)
            self.alive = False
        # check for lifetime
        if self.request_count >= self.request_max:
            _logger.info("Worker (%d) max request (%s) reached.", self.pid, self.request_count)
            self.alive = False
        # Reset the worker if it consumes too much memory (e.g. caused by a memory leak).
        rss, vms = psutil.Process(os.getpid()).get_memory_info()
        if vms > config['limit_memory_soft']:
            _logger.info('Worker (%d) virtual memory limit (%s) reached.', self.pid, vms)
            self.alive = False # Commit suicide after the request.

        # VMS and RLIMIT_AS are the same thing: virtual memory, a.k.a. address space
        soft, hard = resource.getrlimit(resource.RLIMIT_AS)
        resource.setrlimit(resource.RLIMIT_AS, (config['limit_memory_hard'], hard))

        # SIGXCPU (exceeded CPU time) signal handler will raise an exception.
        r = resource.getrusage(resource.RUSAGE_SELF)
        cpu_time = r.ru_utime + r.ru_stime
        def time_expired(n, stack):
            _logger.info('Worker (%d) CPU time limit (%s) reached.', config['limit_time_cpu'])
            # We dont suicide in such case
            raise Exception('CPU time limit exceeded.')
        signal.signal(signal.SIGXCPU, time_expired)
        soft, hard = resource.getrlimit(resource.RLIMIT_CPU)
        resource.setrlimit(resource.RLIMIT_CPU, (cpu_time + config['limit_time_cpu'], hard))

    def process_work(self):
        pass

    def start(self):
        self.pid = os.getpid()
        self.setproctitle()
        _logger.info("Worker %s (%s) alive", self.__class__.__name__, self.pid)
        # Reseed the random number generator
        random.seed()
        # Prevent fd inherientence close_on_exec
        flags = fcntl.fcntl(self.multi.socket, fcntl.F_GETFD) | fcntl.FD_CLOEXEC
        fcntl.fcntl(self.multi.socket, fcntl.F_SETFD, flags)
        # reset blocking status
        self.multi.socket.setblocking(0)
        signal.signal(signal.SIGINT, self.signal_handler)
        signal.signal(signal.SIGTERM, signal.SIG_DFL)
        signal.signal(signal.SIGCHLD, signal.SIG_DFL)

    def stop(self):
        pass

    def run(self):
        try:
            self.start()
            while self.alive:
                self.process_limit()
                self.multi.pipe_ping(self.watchdog_pipe)
                self.sleep()
                self.process_work()
            _logger.info("Worker (%s) exiting. request_count: %s.", self.pid, self.request_count)
            self.stop()
        except Exception,e:
            _logger.exception("Worker (%s) Exception occured, exiting..." % self.pid)
            # should we use 3 to abort everything ?
            sys.exit(1)

class WorkerHTTP(Worker):
    """ HTTP Request workers """
    def process_request(self, client, addr):
        client.setblocking(1)
        client.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)
        # Prevent fd inherientence close_on_exec
        flags = fcntl.fcntl(client, fcntl.F_GETFD) | fcntl.FD_CLOEXEC
        fcntl.fcntl(client, fcntl.F_SETFD, flags)
        # do request using WorkerBaseWSGIServer monkey patched with socket
        self.server.socket = client
        # tolerate broken pipe when the http client closes the socket before
        # receiving the full reply
        try:
            self.server.process_request(client,addr)
        except IOError, e:
            if e.errno != errno.EPIPE:
                raise
        self.request_count += 1

    def process_work(self):
        try:
            client, addr = self.multi.socket.accept()
            self.process_request(client, addr)
        except socket.error, e:
            if e[0] not in (errno.EAGAIN, errno.ECONNABORTED):
                raise

    def start(self):
        Worker.start(self)
        self.server = WorkerBaseWSGIServer(self.multi.app)

class WorkerBaseWSGIServer(werkzeug.serving.BaseWSGIServer):
    """ werkzeug WSGI Server patched to allow using an external listen socket
    """
    def __init__(self, app):
        werkzeug.serving.BaseWSGIServer.__init__(self, "1", "1", app)
    def server_bind(self):
        # we dont bind beause we use the listen socket of Multicorn#socket
        # instead we close the socket
        if self.socket:
            self.socket.close()
    def server_activate(self):
        # dont listen as we use Multicorn#socket
        pass

class WorkerCron(Worker):
    """ Cron workers """

    def __init__(self, multi):
        super(WorkerCron, self).__init__(multi)
        # process_work() below process a single database per call.
        # The variable db_index is keeping track of the next database to
        # process.
        self.db_index = 0

    def sleep(self):
        # Really sleep once all the databases have been processed.
        if self.db_index == 0:
            interval = 60 + self.pid % 10 # chorus effect
            time.sleep(interval)

    def process_work(self):
        rpc_request = logging.getLogger('openerp.netsvc.rpc.request')
        rpc_request_flag = rpc_request.isEnabledFor(logging.DEBUG)
        _logger.debug("WorkerCron (%s) polling for jobs", self.pid)
        if config['db_name']:
            db_names = config['db_name'].split(',')
        else:
            db_names = openerp.netsvc.ExportService._services['db'].exp_list(True)
        if len(db_names):
            self.db_index = (self.db_index + 1) % len(db_names)
            db_name = db_names[self.db_index]
            self.setproctitle(db_name)
            if rpc_request_flag:
                start_time = time.time()
                start_rss, start_vms = psutil.Process(os.getpid()).get_memory_info()
            while True:
                # acquired = openerp.addons.base.ir.ir_cron.ir_cron._acquire_job(db_name)
                # TODO why isnt openerp.addons.base defined ?
                import base
                acquired = base.ir.ir_cron.ir_cron._acquire_job(db_name)
                if not acquired:
                    openerp.modules.registry.RegistryManager.delete(db_name)
                    break
            # dont keep cursors in multi database mode
            if len(db_names) > 1:
                openerp.sql_db.close_db(db_name)
            if rpc_request_flag:
                end_time = time.time()
                end_rss, end_vms = psutil.Process(os.getpid()).get_memory_info()
                logline = '%s time:%.3fs mem: %sk -> %sk (diff: %sk)' % (db_name, end_time - start_time, start_vms / 1024, end_vms / 1024, (end_vms - start_vms)/1024)
                _logger.debug("WorkerCron (%s) %s", self.pid, logline)

            self.request_count += 1
            if self.request_count >= self.request_max and self.request_max < len(db_names):
                _logger.error("There are more dabatases to process than allowed "
                    "by the `limit_request` configuration variable: %s more.",
                    len(db_names) - self.request_max)
        else:
            self.db_index = 0

    def start(self):
        Worker.start(self)
        openerp.service.start_internal()

# vim:expandtab:smartindent:tabstop=4:softtabstop=4:shiftwidth=4:
<|endoftext|>


<commit_before>from __future__ import division, print_function

import os
import cStringIO as StringIO

from pymatgen.util.io_utils import which
from pymatgen.util.string_utils import list_strings

import logging
logger = logging.getLogger(__name__)

__author__ = "Matteo Giantomassi"
__copyright__ = "Copyright 2013, The Materials Project"
__version__ = "0.1"
__maintainer__ = "Matteo Giantomassi"
__email__ = "gmatteo at gmail.com"
__status__ = "Development"
__date__ = "$Feb 21, 2013M$"

__all__ = [
    "Mrgscr",
    "Mrggkk",
    "Mrgddb",
    "Anaddb",
]

class ExecWrapper(object):
    """This class runs an executable in a subprocess."""
    def __init__(self, executable=None, verbose=0):
        """
        Args:
            executable:
                path to the executable.
            verbose:
                Verbosity level.
        """
        if executable is None:
            executable = self.name

        self.executable = which(executable)

        self.verbose = int(verbose)

        if self.executable is None:
            msg = "Cannot find executable %s is PATH\n Use export PATH=/dir_with_exec:$PATH" % executable
            raise self.Error(msg)

        assert os.path.basename(self.executable) == self._name

    def __str__(self):
        return "%s" % self.executable

    def set_mpi_runner(self, mpi_runner="mpirun"):
        # TODO better treatment of mpirunner syntax.
        self._mpi_runner = mpi_runner

    @property
    def mpi_runner(self):
        try:
            return self._mpi_runner

        except AttributeError:
            return ""

    @property
    def name(self):
        return self._name

    def execute(self, cwd=None):

        # Try to execute binary without and with mpirun.
        try:
            self._execute(cwd=cwd, with_mpirun=False)

        except self.Error:
            self._execute(cwd=cwd, with_mpirun=True)

    def _execute(self, cwd=None, with_mpirun=False):
        """
        Execute the executable in a subprocess.
        """
        args = [self.executable, "<", self.stdin_fname, ">", self.stdout_fname, "2>", self.stderr_fname]

        if self.mpi_runner and with_mpirun:
            args.insert(0, self.mpi_runner)

        self.cmd_str = " ".join(args)

        p = Popen(self.cmd_str, shell=True, stdout=PIPE, stderr=PIPE, cwd=cwd)

        (self.stdout_data, self.stderr_data) = p.communicate()

        self.returncode = p.returncode

        if self.returncode != 0:
            with open(self.stdout_fname, "r") as out, open(self.stderr_fname, "r") as err:
                self.stdout_data = out.read()
                self.stderr_data = err.read()

            if self.verbose:
                print("*** stdout: ***\n", self.stdout_data)
                print("*** stderr  ***\n", self.stderr_data)

            raise self.Error("%s returned %s\n cmd_str: %s" % (self, self.returncode, self.cmd_str))

class MrgscrError(Exception):
    """Error class for Mrgscr"""

class Mrgscr(ExecWrapper):
    _name = "mrgscr"

    Error = MrgscrError

    def merge_qpoints(self, files_to_merge, out_prefix, cwd=None):
        """
        Execute mrgscr in a subprocess to merge files_to_merge. Produce new file with prefix out_prefix

        If cwd is not None, the child's current directory will be changed to cwd before it is executed.
        """
        # We work with absolute paths.
        files_to_merge = [os.path.abspath(s) for s in list_strings(files_to_merge)]
        nfiles = len(files_to_merge)

        if self.verbose:
            print("Will merge %d files with output_prefix %s" % (nfiles, out_prefix))
            for (i, f) in enumerate(files_to_merge):
                print(" [%d] %s" % (i, f))

        if nfiles == 1:
            raise self.Error("merge_qpoints does not support nfiles == 1")

        self.stdin_fname, self.stdout_fname, self.stderr_fname = (
            "mrgscr.stdin", "mrgscr.stdout", "mrgscr.stderr")

        if cwd is not None:
            self.stdin_fname, self.stdout_fname, self.stderr_fname = \
                map(os.path.join, 3 * [cwd], [self.stdin_fname, self.stdout_fname, self.stderr_fname])

        inp = StringIO.StringIO()

        inp.write(str(nfiles) + "\n")     # Number of files to merge.
        inp.write(out_prefix + "\n")      # Prefix for the final output file:

        for filename in files_to_merge:
            inp.write(filename + "\n")   # List with the files to merge.

        inp.write("1\n")                 # Option for merging q-points.

        inp.seek(0)
        self.stdin_data = [s for s in inp]

        with open(self.stdin_fname, "w") as fh:
            fh.writelines(self.stdin_data)

        try:
            self.execute(cwd=cwd)
        except self.Error:
            raise

class MrggkkError(Exception):
    """Error class for Mrggkk."""

class Mrggkk(ExecWrapper):
    _name = "mrggkk"

    Error = MrggkkError

    def merge(self, gswfk_file, dfpt_files, gkk_files, out_gkk, binascii=0, cwd=None):
        """
        Merge GGK files, return the absolute path of the new database.

        Args:
            gswfk_file:
                Ground-state WFK filename
            dfpt_files:
                List of 1WFK files to merge.
            gkk_files:
                List of GKK files to merge.
            out_gkk:
                Name of the output GKK file
            binascii:
                Integer flat. 0 --> binary output, 1 --> ascii formatted output
            cwd:
                Directory where the subprocess will be executed.
        """
        raise NotImplementedError("This method should be tested")

        out_gkk = out_gkk if cwd is None else os.path.join(os.path.abspath(cwd), out_gkk)

        # We work with absolute paths.
        gswfk_file = absath(gswfk_file)
        dfpt_files = [os.path.abspath(s) for s in list_strings(dfpt_files)]
        gkk_files = [os.path.abspath(s) for s in list_strings(gkk_files)]

        if self.verbose:
            print("Will merge %d 1WF files, %d GKK file in output %s" %
                  (len(dfpt_nfiles), (len_gkk_files), out_gkk))

            for (i, f) in enumerate(dfpt_files):
                print(" [%d] 1WF %s" % (i, f))

            for (i, f) in enumerate(gkk_files):
                print(" [%d] GKK %s" % (i, f))

        self.stdin_fname, self.stdout_fname, self.stderr_fname = (
            "mrggkk.stdin", "mrggkk.stdout", "mrggkk.stderr")

        if cwd is not None:
            self.stdin_fname, self.stdout_fname, self.stderr_fname = \
                map(os.path.join, 3 * [cwd], [self.stdin_fname, self.stdout_fname, self.stderr_fname])

        inp = StringIO.StringIO()

        inp.write(out_gkk + "\n")        # Name of the output file
        inp.write(str(binascii) + "\n")  # Integer flag: 0 --> binary output, 1 --> ascii formatted output
        inp.write(gswfk_file + "\n")     # Name of the groud state wavefunction file WF

        #dims = len(dfpt_files, gkk_files, ?)
        dims = " ".join([str(d) for d in dims])
        inp.write(dims + "\n")             # Number of 1WF, of GKK files, and number of 1WF files in all the GKK files

        # Names of the 1WF files...
        for fname in dfpt_files:
            inp.write(fname + "\n")

        # Names of the GKK files...
        for fname in gkk_files:
            inp.write(fname + "\n")

        inp.seek(0)
        self.stdin_data = [s for s in inp]

        with open(self.stdin_fname, "w") as fh:
            fh.writelines(self.stdin_data)

        try:
            self.execute(cwd=cwd)
        except self.Error:
            raise

        return out_gkk

class MrgddbError(Exception):
    """Error class for Mrgddb."""

class Mrgddb(ExecWrapper):
    _name = "mrgddb"

    Error = MrgddbError

    def merge(self, ddb_files, out_ddb, description, cwd=None):
        """Merge DDB file, return the absolute path of the new database."""

        # We work with absolute paths.
        ddb_files = [os.path.abspath(s) for s in list_strings(ddb_files)]

        out_ddb = out_ddb if cwd is None else os.path.join(os.path.abspath(cwd), out_ddb)

        if self.verbose:
            print("Will merge %d files into output DDB %s" % (len(ddb_files), out_ddb))
            for (i, f) in enumerate(ddb_files):
                print(" [%d] %s" % (i, f))

        # Handle the case of a single file since mrgddb uses 1 to denote GS files!
        if len(ddb_files) == 1:
            with open(ddb_files[0], "r") as inh, open(out_ddb, "w") as out:
                for line in inh:
                    out.write(line)
            return out_ddb

        self.stdin_fname, self.stdout_fname, self.stderr_fname = (
            "mrgddb.stdin", "mrgddb.stdout", "mrgddb.stderr")

        if cwd is not None:
            self.stdin_fname, self.stdout_fname, self.stderr_fname = \
                map(os.path.join, 3 * [cwd], [self.stdin_fname, self.stdout_fname, self.stderr_fname])

        inp = StringIO.StringIO()

        inp.write(out_ddb + "\n")            # Name of the output file.
        inp.write(str(description) + "\n")     # Description.
        inp.write(str(len(ddb_files)) + "\n")  # Number of input DDBs.

        # Names of the DDB files.
        for fname in ddb_files:
            inp.write(fname + "\n")

        inp.seek(0)
        self.stdin_data = [s for s in inp]

        with open(self.stdin_fname, "w") as fh:
            fh.writelines(self.stdin_data)

        try:
            self.execute(cwd=cwd)
        except self.Error:
            raise

        return out_ddb

class AnaddbError(Exception):
    """Error class for Anaddb."""

class Anaddb(ExecWrapper):
    _name = "anaddb"

    Error = AnaddbError

    #def make_stdin(self):
    #    # Files file
    #    inp = StringIO.StringIO()
    #    inp.write(self.input_fname + "\n")     # Input file.
    #    inp.write(self.stdout_fname + "\n")    # Output file.
    #    inp.write(ddb_file + "\n")             # DDB file
    #    inp.write("dummy_band2eps" + "\n")
    #    inp.write("dummy1" + "\n")
    #    inp.write("dummy2" + "\n")
    #    inp.write("dummy3" + "\n")
    #    inp.seek(0)
    #    self.stdin_data = [s for s in inp]

    def diagonalize_1q(self, ddb_file, cwd=None):

        # We work with absolute paths.
        ddb_file = os.path.abspath(ddb_file)

        self.stdin_fname, self.input_fname, self.stdout_fname, self.stderr_fname = (
            "anaddb.stdin", "anaddb.input", "anaddb.stdout", "anaddb.stderr")

        if cwd is not None:
            self.stdin_fname, self.input_fname, self.stdout_fname, self.stderr_fname = \
                map(os.path.join, 3 * [cwd], [self.stdin_fname, self.inp_fname, self.stdout_fname, self.stderr_fname])

        # Files file
        inp = StringIO.StringIO()

        inp.write(self.input_fname + "\n")     # Input file.
        inp.write(self.stdout_fname + "\n")    # Output file.
        inp.write(ddb_file + "\n")             # DDB file
        inp.write("dummy_band2eps" + "\n")
        inp.write("dummy1" + "\n")
        inp.write("dummy2" + "\n")
        inp.write("dummy3" + "\n")

        inp.seek(0)
        self.stdin_data = [s for s in inp]

        with open(self.stdin_fname, "w") as fh:
            fh.writelines(self.stdin_data)

        # Get the q-point from the DDB file
        with open(ddb_file, "r") as fh:
            nfound = 0
            tag = " qpt "
            for line in fh:
                print(line)
                if line.startswith(tag):
                    nfound += 1
                    # Coordinates of the q-points.
                    qcoords_str = line.split()[1:4]
                    #qcoords_str = [ s.replace("D", "E") for s in qcoords_str]
                    qpoint = map(float, qcoords_str)

        if nfound != 1:
            raise self.Error("Found %s occurrences of tag %s in file %s" % (nfound, tag, ddb_file))

        # Write simple input file for the anaddb code.
        with open(self.input_fname, "w") as inp:
            inp.write('# Flags\n')
            inp.write(' ifcflag   1 # Interatomic force constant flag\n\n')
            inp.write('# Wavevector grid number 1 (coarse grid, from DDB)\n\n')
            inp.write(' brav    1         # Bravais Lattice : 1-S.C., 2-F.C., 3-B.C., 4-Hex.\n')
            inp.write(' ngqpt   1  1  1   # Q-mesh\n')
            inp.write(' nqshft  1         # number of q-shifts\n')
            inp.write(' q1shft  %f %f %f' % tuple(qpoint))

            #inp.write('# Effective charges
            #inp.write('     asr   1     ! Acoustic Sum Rule. 1 => imposed asymetrically
            #inp.write('  chneut   1     ! Charge neutrality requirement for effective charges.
            #inp.write('# Interatomic force constant info
            #inp.write('  dipdip  1      ! Dipole-dipole interaction treatment
            #inp.write('  ifcana  1      ! Analysis of the IFCs
            #inp.write('  ifcout 20      ! Number of IFC's written in the output, per atom
            #inp.write('  natifc  1      ! Number of atoms in the cell for which ifc's are analysed
            #inp.write('   atifc  1      ! List of atoms
            #inp.write('
            #inp.write('# This line added when defaults were changed (v5.3) to keep the previous, old behaviour
            #inp.write('#  symdynmat 0

        if self.verbose:
            print("Will diagonalize DDB file : %s" % ddb_file)

        try:
            self.execute(cwd=cwd)
        except self.Error:
            raise

        # Get frequencies from the output file
        # TODO
        #with open(self.stdout_fname, "r") as out:
            #print(out.readlines())
            #for line in out:
            #    if line: raise
            #return frequencies
<commit_msg>  Popen import was erroneously removed in the last commit<commit_after>from __future__ import division, print_function

import os
import cStringIO as StringIO

from subprocess import Popen, PIPE
from pymatgen.util.io_utils import which
from pymatgen.util.string_utils import list_strings

import logging
logger = logging.getLogger(__name__)

__author__ = "Matteo Giantomassi"
__copyright__ = "Copyright 2013, The Materials Project"
__version__ = "0.1"
__maintainer__ = "Matteo Giantomassi"
__email__ = "gmatteo at gmail.com"
__status__ = "Development"
__date__ = "$Feb 21, 2013M$"

__all__ = [
    "Mrgscr",
    "Mrggkk",
    "Mrgddb",
    "Anaddb",
]

class ExecWrapper(object):
    """This class runs an executable in a subprocess."""
    def __init__(self, executable=None, verbose=0):
        """
        Args:
            executable:
                path to the executable.
            verbose:
                Verbosity level.
        """
        if executable is None:
            executable = self.name

        self.executable = which(executable)

        self.verbose = int(verbose)

        if self.executable is None:
            msg = "Cannot find executable %s is PATH\n Use export PATH=/dir_with_exec:$PATH" % executable
            raise self.Error(msg)

        assert os.path.basename(self.executable) == self._name

    def __str__(self):
        return "%s" % self.executable

    def set_mpi_runner(self, mpi_runner="mpirun"):
        # TODO better treatment of mpirunner syntax.
        self._mpi_runner = mpi_runner

    @property
    def mpi_runner(self):
        try:
            return self._mpi_runner

        except AttributeError:
            return ""

    @property
    def name(self):
        return self._name

    def execute(self, cwd=None):

        # Try to execute binary without and with mpirun.
        try:
            self._execute(cwd=cwd, with_mpirun=False)

        except self.Error:
            self._execute(cwd=cwd, with_mpirun=True)

    def _execute(self, cwd=None, with_mpirun=False):
        """
        Execute the executable in a subprocess.
        """
        args = [self.executable, "<", self.stdin_fname, ">", self.stdout_fname, "2>", self.stderr_fname]

        if self.mpi_runner and with_mpirun:
            args.insert(0, self.mpi_runner)

        self.cmd_str = " ".join(args)

        p = Popen(self.cmd_str, shell=True, stdout=PIPE, stderr=PIPE, cwd=cwd)

        (self.stdout_data, self.stderr_data) = p.communicate()

        self.returncode = p.returncode

        if self.returncode != 0:
            with open(self.stdout_fname, "r") as out, open(self.stderr_fname, "r") as err:
                self.stdout_data = out.read()
                self.stderr_data = err.read()

            if self.verbose:
                print("*** stdout: ***\n", self.stdout_data)
                print("*** stderr  ***\n", self.stderr_data)

            raise self.Error("%s returned %s\n cmd_str: %s" % (self, self.returncode, self.cmd_str))

class MrgscrError(Exception):
    """Error class for Mrgscr"""

class Mrgscr(ExecWrapper):
    _name = "mrgscr"

    Error = MrgscrError

    def merge_qpoints(self, files_to_merge, out_prefix, cwd=None):
        """
        Execute mrgscr in a subprocess to merge files_to_merge. Produce new file with prefix out_prefix

        If cwd is not None, the child's current directory will be changed to cwd before it is executed.
        """
        # We work with absolute paths.
        files_to_merge = [os.path.abspath(s) for s in list_strings(files_to_merge)]
        nfiles = len(files_to_merge)

        if self.verbose:
            print("Will merge %d files with output_prefix %s" % (nfiles, out_prefix))
            for (i, f) in enumerate(files_to_merge):
                print(" [%d] %s" % (i, f))

        if nfiles == 1:
            raise self.Error("merge_qpoints does not support nfiles == 1")

        self.stdin_fname, self.stdout_fname, self.stderr_fname = (
            "mrgscr.stdin", "mrgscr.stdout", "mrgscr.stderr")

        if cwd is not None:
            self.stdin_fname, self.stdout_fname, self.stderr_fname = \
                map(os.path.join, 3 * [cwd], [self.stdin_fname, self.stdout_fname, self.stderr_fname])

        inp = StringIO.StringIO()

        inp.write(str(nfiles) + "\n")     # Number of files to merge.
        inp.write(out_prefix + "\n")      # Prefix for the final output file:

        for filename in files_to_merge:
            inp.write(filename + "\n")   # List with the files to merge.

        inp.write("1\n")                 # Option for merging q-points.

        inp.seek(0)
        self.stdin_data = [s for s in inp]

        with open(self.stdin_fname, "w") as fh:
            fh.writelines(self.stdin_data)

        try:
            self.execute(cwd=cwd)
        except self.Error:
            raise

class MrggkkError(Exception):
    """Error class for Mrggkk."""

class Mrggkk(ExecWrapper):
    _name = "mrggkk"

    Error = MrggkkError

    def merge(self, gswfk_file, dfpt_files, gkk_files, out_gkk, binascii=0, cwd=None):
        """
        Merge GGK files, return the absolute path of the new database.

        Args:
            gswfk_file:
                Ground-state WFK filename
            dfpt_files:
                List of 1WFK files to merge.
            gkk_files:
                List of GKK files to merge.
            out_gkk:
                Name of the output GKK file
            binascii:
                Integer flat. 0 --> binary output, 1 --> ascii formatted output
            cwd:
                Directory where the subprocess will be executed.
        """
        raise NotImplementedError("This method should be tested")

        out_gkk = out_gkk if cwd is None else os.path.join(os.path.abspath(cwd), out_gkk)

        # We work with absolute paths.
        gswfk_file = absath(gswfk_file)
        dfpt_files = [os.path.abspath(s) for s in list_strings(dfpt_files)]
        gkk_files = [os.path.abspath(s) for s in list_strings(gkk_files)]

        if self.verbose:
            print("Will merge %d 1WF files, %d GKK file in output %s" %
                  (len(dfpt_nfiles), (len_gkk_files), out_gkk))

            for (i, f) in enumerate(dfpt_files):
                print(" [%d] 1WF %s" % (i, f))

            for (i, f) in enumerate(gkk_files):
                print(" [%d] GKK %s" % (i, f))

        self.stdin_fname, self.stdout_fname, self.stderr_fname = (
            "mrggkk.stdin", "mrggkk.stdout", "mrggkk.stderr")

        if cwd is not None:
            self.stdin_fname, self.stdout_fname, self.stderr_fname = \
                map(os.path.join, 3 * [cwd], [self.stdin_fname, self.stdout_fname, self.stderr_fname])

        inp = StringIO.StringIO()

        inp.write(out_gkk + "\n")        # Name of the output file
        inp.write(str(binascii) + "\n")  # Integer flag: 0 --> binary output, 1 --> ascii formatted output
        inp.write(gswfk_file + "\n")     # Name of the groud state wavefunction file WF

        #dims = len(dfpt_files, gkk_files, ?)
        dims = " ".join([str(d) for d in dims])
        inp.write(dims + "\n")             # Number of 1WF, of GKK files, and number of 1WF files in all the GKK files

        # Names of the 1WF files...
        for fname in dfpt_files:
            inp.write(fname + "\n")

        # Names of the GKK files...
        for fname in gkk_files:
            inp.write(fname + "\n")

        inp.seek(0)
        self.stdin_data = [s for s in inp]

        with open(self.stdin_fname, "w") as fh:
            fh.writelines(self.stdin_data)

        try:
            self.execute(cwd=cwd)
        except self.Error:
            raise

        return out_gkk

class MrgddbError(Exception):
    """Error class for Mrgddb."""

class Mrgddb(ExecWrapper):
    _name = "mrgddb"

    Error = MrgddbError

    def merge(self, ddb_files, out_ddb, description, cwd=None):
        """Merge DDB file, return the absolute path of the new database."""

        # We work with absolute paths.
        ddb_files = [os.path.abspath(s) for s in list_strings(ddb_files)]

        out_ddb = out_ddb if cwd is None else os.path.join(os.path.abspath(cwd), out_ddb)

        if self.verbose:
            print("Will merge %d files into output DDB %s" % (len(ddb_files), out_ddb))
            for (i, f) in enumerate(ddb_files):
                print(" [%d] %s" % (i, f))

        # Handle the case of a single file since mrgddb uses 1 to denote GS files!
        if len(ddb_files) == 1:
            with open(ddb_files[0], "r") as inh, open(out_ddb, "w") as out:
                for line in inh:
                    out.write(line)
            return out_ddb

        self.stdin_fname, self.stdout_fname, self.stderr_fname = (
            "mrgddb.stdin", "mrgddb.stdout", "mrgddb.stderr")

        if cwd is not None:
            self.stdin_fname, self.stdout_fname, self.stderr_fname = \
                map(os.path.join, 3 * [cwd], [self.stdin_fname, self.stdout_fname, self.stderr_fname])

        inp = StringIO.StringIO()

        inp.write(out_ddb + "\n")            # Name of the output file.
        inp.write(str(description) + "\n")     # Description.
        inp.write(str(len(ddb_files)) + "\n")  # Number of input DDBs.

        # Names of the DDB files.
        for fname in ddb_files:
            inp.write(fname + "\n")

        inp.seek(0)
        self.stdin_data = [s for s in inp]

        with open(self.stdin_fname, "w") as fh:
            fh.writelines(self.stdin_data)

        try:
            self.execute(cwd=cwd)
        except self.Error:
            raise

        return out_ddb

class AnaddbError(Exception):
    """Error class for Anaddb."""

class Anaddb(ExecWrapper):
    _name = "anaddb"

    Error = AnaddbError

    #def make_stdin(self):
    #    # Files file
    #    inp = StringIO.StringIO()
    #    inp.write(self.input_fname + "\n")     # Input file.
    #    inp.write(self.stdout_fname + "\n")    # Output file.
    #    inp.write(ddb_file + "\n")             # DDB file
    #    inp.write("dummy_band2eps" + "\n")
    #    inp.write("dummy1" + "\n")
    #    inp.write("dummy2" + "\n")
    #    inp.write("dummy3" + "\n")
    #    inp.seek(0)
    #    self.stdin_data = [s for s in inp]

    def diagonalize_1q(self, ddb_file, cwd=None):

        # We work with absolute paths.
        ddb_file = os.path.abspath(ddb_file)

        self.stdin_fname, self.input_fname, self.stdout_fname, self.stderr_fname = (
            "anaddb.stdin", "anaddb.input", "anaddb.stdout", "anaddb.stderr")

        if cwd is not None:
            self.stdin_fname, self.input_fname, self.stdout_fname, self.stderr_fname = \
                map(os.path.join, 3 * [cwd], [self.stdin_fname, self.inp_fname, self.stdout_fname, self.stderr_fname])

        # Files file
        inp = StringIO.StringIO()

        inp.write(self.input_fname + "\n")     # Input file.
        inp.write(self.stdout_fname + "\n")    # Output file.
        inp.write(ddb_file + "\n")             # DDB file
        inp.write("dummy_band2eps" + "\n")
        inp.write("dummy1" + "\n")
        inp.write("dummy2" + "\n")
        inp.write("dummy3" + "\n")

        inp.seek(0)
        self.stdin_data = [s for s in inp]

        with open(self.stdin_fname, "w") as fh:
            fh.writelines(self.stdin_data)

        # Get the q-point from the DDB file
        with open(ddb_file, "r") as fh:
            nfound = 0
            tag = " qpt "
            for line in fh:
                print(line)
                if line.startswith(tag):
                    nfound += 1
                    # Coordinates of the q-points.
                    qcoords_str = line.split()[1:4]
                    #qcoords_str = [ s.replace("D", "E") for s in qcoords_str]
                    qpoint = map(float, qcoords_str)

        if nfound != 1:
            raise self.Error("Found %s occurrences of tag %s in file %s" % (nfound, tag, ddb_file))

        # Write simple input file for the anaddb code.
        with open(self.input_fname, "w") as inp:
            inp.write('# Flags\n')
            inp.write(' ifcflag   1 # Interatomic force constant flag\n\n')
            inp.write('# Wavevector grid number 1 (coarse grid, from DDB)\n\n')
            inp.write(' brav    1         # Bravais Lattice : 1-S.C., 2-F.C., 3-B.C., 4-Hex.\n')
            inp.write(' ngqpt   1  1  1   # Q-mesh\n')
            inp.write(' nqshft  1         # number of q-shifts\n')
            inp.write(' q1shft  %f %f %f' % tuple(qpoint))

            #inp.write('# Effective charges
            #inp.write('     asr   1     ! Acoustic Sum Rule. 1 => imposed asymetrically
            #inp.write('  chneut   1     ! Charge neutrality requirement for effective charges.
            #inp.write('# Interatomic force constant info
            #inp.write('  dipdip  1      ! Dipole-dipole interaction treatment
            #inp.write('  ifcana  1      ! Analysis of the IFCs
            #inp.write('  ifcout 20      ! Number of IFC's written in the output, per atom
            #inp.write('  natifc  1      ! Number of atoms in the cell for which ifc's are analysed
            #inp.write('   atifc  1      ! List of atoms
            #inp.write('
            #inp.write('# This line added when defaults were changed (v5.3) to keep the previous, old behaviour
            #inp.write('#  symdynmat 0

        if self.verbose:
            print("Will diagonalize DDB file : %s" % ddb_file)

        try:
            self.execute(cwd=cwd)
        except self.Error:
            raise

        # Get frequencies from the output file
        # TODO
        #with open(self.stdout_fname, "r") as out:
            #print(out.readlines())
            #for line in out:
            #    if line: raise
            #return frequencies
<|endoftext|>


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

"""
This module provides classes to create phase diagrams.
"""

__author__ = "Shyue Ping Ong"
__copyright__ = "Copyright 2011, The Materials Project"
__version__ = "1.0"
__maintainer__ = "Shyue Ping Ong"
__email__ = "shyue@mit.edu"
__status__ = "Production"
__date__ = "Sep 23, 2011"

import numpy as np
import logging
import itertools

from scipy.spatial import Delaunay

from pymatgen.core.structure import Composition
from pymatgen.command_line.qhull_caller import qconvex
from pymatgen.phasediagram.entries import GrandPotPDEntry, TransformedPDEntry
from pymatgen.core.periodic_table import Element
from pymatgen.analysis.reaction_calculator import Reaction

logger = logging.getLogger(__name__)

class PhaseDiagram (object):
    '''
    Simple phase diagram class taking in elements and entries as inputs.
    '''
    FORMATION_ENERGY_TOLERANCE = 1e-11

    def __init__(self, entries, elements=None, use_external_qhull=False):
        """
        Standard constructor for phase diagram.

        Args:
            entries:
                A list of PDEntry-like objects having an energy, energy_per_atom and composition.
            elements:
                Optional list of elements in the phase diagram. If set to None, the elements are determined from the
                the entries themselves.
        """
        if elements == None:
            elements = set()
            map(elements.update, [entry.composition.elements for entry in entries])
        self._all_entries = entries
        self._elements = tuple(elements)
        self._qhull_data = None
        self._facets = None
        self._qhull_entries = None
        self._stable_entries = None
        self._all_entries_hulldata = None
        self._use_external_qhull = use_external_qhull
        self._make_phasediagram()

    @property
    def all_entries(self):
        """
        All entries submitted to construct the PhaseDiagram.
        Note that this does not mean that all these entries are actually used in the phase diagram.
        """
        return self._all_entries

    @property
    def dim(self):
        return len(self._elements)

    @property
    def elements(self):
        """
        Elements in the phase diagram.
        """
        return self._elements

    @property
    def facets(self):
        """
        Facets of the phase diagram in the form of  [[1,2,3],[4,5,6]...]
        """
        return self._facets

    @property
    def qhull_data(self):
        """
        Data used in the convex hull operation. This is essentially a matrix of
        composition data and energy per atom values created from the qhull_entries.
        """
        return self._qhull_data

    @property
    def qhull_entries(self):
        """
        Actual entries used in convex hull. Excludes all positive formation energy entries.
        """
        return self._qhull_entries

    @property
    def unstable_entries(self):
        """
        Entries that are unstable in the phase diagram. Includes positive formation energy entries.
        """
        return [e for e in self.all_entries if e not in self.stable_entries]

    @property
    def stable_entries(self):
        '''
        Returns the stable entries.
        '''
        return self._stable_entries

    @property
    def all_entries_hulldata(self):
        """
        Same as qhull_data, but for all entries rather than just positive formation energy ones.
        """
        return self._all_entries_hulldata

    @property
    def el_refs(self):
        """
        List of elemental references for the phase diagrams.
        These are typically entries corresponding to the lowest energy element entries.
        """
        return self._el_refs

    def get_form_energy(self, entry):
        '''
        Returns the formation energy for an entry (NOT normalized) from the
        elemental references.

        Args:
            entry:
                A PDEntry

        Returns:
            Formation energy from the elementals references.
        '''
        comp = entry.composition
        energy = entry.energy - sum([comp[el] * self._el_refs[el].energy_per_atom for el in comp.elements])
        return energy

    def get_form_energy_per_atom(self, entry):
        '''
        Returns the formation energy per atom for an entry from the
        elemental references.
        '''
        comp = entry.composition
        return self.get_form_energy(entry) / comp.num_atoms

    def _process_entries_qhulldata(self, entries_to_process):
        data = list()
        for entry in entries_to_process:
            comp = entry.composition
            energy_per_atom = entry.energy_per_atom
            row = list()
            for i in xrange(1, len(self._elements)):
                row.append(comp.get_atomic_fraction(self._elements[i]))
            row.append(energy_per_atom)
            data.append(row)
        return data

    def _create_convhull_data(self):
        '''
        Make data suitable for convex hull procedure from the list of entries.
        '''
        logger.debug("Creating convex hull data...")
        #Determine the elemental references based on lowest energy for each.
        self._el_refs = dict()
        for entry in self._all_entries:
            if entry.composition.is_element:
                for el in entry.composition.elements:
                    if entry.composition[el] > Composition.amount_tolerance:
                        break
                e_per_atom = entry.energy_per_atom
                if el not in self._el_refs:
                    self._el_refs[el] = entry
                elif self._el_refs[el].energy_per_atom > e_per_atom:
                    self._el_refs[el] = entry
        # Remove positive formation energy entries
        entries_to_process = list()
        for entry in self._all_entries:
            if self.get_form_energy(entry) <= self.FORMATION_ENERGY_TOLERANCE:
                entries_to_process.append(entry)
            else:
                logger.debug("Removing positive formation energy entry {}".format(entry))

        self._qhull_entries = entries_to_process
        return self._process_entries_qhulldata(entries_to_process)

    def _make_phasediagram(self):
        """
        Make the phase diagram.
        """
        stable_entries = set()
        dim = len(self._elements)
        self._qhull_data = self._create_convhull_data()
        if len(self._qhull_data) == dim:
            self._facets = [range(len(self._elements))]
        else:
            if self._use_external_qhull:
                logger.debug("> 4D hull encountered. Computing hull using external qconvex call.")
                self._facets = qconvex(self._qhull_data)
            else:
                logger.debug("Computing hull using scipy.spatial.delaunay")
                delau = Delaunay(self._qhull_data)
                self._facets = delau.convex_hull
            logger.debug("Final facets are\n{}".format(self._facets))

            logger.debug("Removing vertical facets...")
            finalfacets = list()
            for facet in self._facets:
                facetmatrix = np.zeros((len(facet), len(facet)))
                count = 0
                is_element_facet = True
                for vertex in facet:
                    facetmatrix[count] = np.array(self._qhull_data[vertex])
                    facetmatrix[count, dim - 1] = 1
                    count += 1
                    if len(self._qhull_entries[vertex].composition) > 1:
                        is_element_facet = False
                if abs(np.linalg.det(facetmatrix)) > 1e-8 and (not is_element_facet):
                    finalfacets.append(facet)
                else:
                    logger.debug("Removing vertical facet : {}".format(facet))
            self._facets = finalfacets

        for facet in self._facets:
            for vertex in facet:
                stable_entries.add(self._qhull_entries[vertex])
        self._stable_entries = stable_entries
        self._all_entries_hulldata = self._process_entries_qhulldata(self._all_entries)

class GrandPotentialPhaseDiagram (PhaseDiagram):
    '''
    Grand potential phase diagram class taking in elements and entries as inputs.
    '''

    def __init__(self, entries, chempots, elements=None, use_external_qhull=False):
        """
        Standard constructor for grand potential phase diagram.

        Args:
            entries:
                A list of PDEntry-like objects having an energy, energy_per_atom
                and composition.
            chempots:
                A dict of {element: float} to specify the chemical potentials
                of the open elements.
            elements:
                Optional list of elements in the phase diagram. If set to None,
                the elements are determined from the entries themselves.
        """
        if elements == None:
            elements = set()
            map(elements.update, [entry.composition.elements for entry in entries])
        allentries = list()
        for entry in entries:
            if not (entry.is_element and (entry.composition.elements[0] in chempots)):
                allentries.append(GrandPotPDEntry(entry, chempots))
        self.chempots = chempots
        filteredels = list()
        for el in elements:
            if el not in chempots:
                filteredels.append(el)
        elements = sorted(filteredels)
        super(GrandPotentialPhaseDiagram, self).__init__(allentries, elements, use_external_qhull)

class CompoundPhaseDiagram(PhaseDiagram):
    """
    Experimental feature. Generates phase diagrams from compounds as termninations
    instead of elements.
    """

    def __init__(self, entries, terminal_compositions, use_external_qhull=False):
        entries = get_entries_within_compositional_space(entries, terminal_compositions)
        elset = get_non_coplanar_element_set(entries)
        els = list(elset)
        pentries = get_transformed_entries(entries, els)
        super(CompoundPhaseDiagram, self).__init__(pentries, use_external_qhull=use_external_qhull)

def get_comp_matrix_from_comp(compositions, elements, normalize_row=True):
    """
    Helper function to generates a normalized composition matrix from a list of
    composition.
    """
    comp_matrix = np.array([[comp.get_atomic_fraction(el) for el in elements] for comp in compositions])
    if not normalize_row:
        return comp_matrix
    factor = np.tile(np.sum(comp_matrix, 1), (len(elements), 1)).transpose()
    return comp_matrix / factor

def get_comp_matrix(entries, elements, normalize_row=True):
    """
    Helper function to generates a normalized composition matrix from a list of
    composition.
    """
    return get_comp_matrix_from_comp([entry.composition for entry in entries], elements, normalize_row)

def is_coplanar(entries, elements):
    comp_matrix = get_comp_matrix(entries, elements)
    for submatrix in itertools.combinations(comp_matrix, min(len(elements), len(entries))):
        if abs(np.linalg.det(submatrix)) > 1e-5:
            return False
    return True

def get_non_coplanar_element_set(entries):
    elements = set()
    map(elements.update, [entry.composition.elements for entry in entries])
    for i in xrange(len(elements), 1, -1):
        for elset in itertools.combinations(elements, i):
            if not is_coplanar(entries, elset):
                return elset
    return None

def get_transformed_entries(entries, elements):
    comp_matrix = get_comp_matrix(entries, elements)
    newmat = []
    energies = []
    for i in xrange(len(elements)):
        col = comp_matrix[:, i]
        maxval = max(col)
        maxind = list(col).index(maxval)
        newmat.append(comp_matrix[maxind])
        energies.append(entries[i].energy_per_atom)
    invm = np.linalg.inv(np.array(newmat).transpose())
    newentries = []
    for i in xrange(len(entries)):
        entry = entries[i]
        lincomp = np.dot(invm, comp_matrix[i])
        lincomp = np.around(lincomp, 5)
        comp = Composition({Element.from_Z(j + 1):lincomp[j] for j in xrange(len(elements))})
        scaled_energy = entry.energy_per_atom - sum(lincomp * energies)
        newentries.append(TransformedPDEntry(comp, scaled_energy, entry))
    return newentries

def get_entries_within_compositional_space(entries, terminal_compositions):
    newentries = []
    for entry in entries:
        try:
            rxn = Reaction(terminal_compositions, [entry.composition])
            newentries.append(entry)
        except:
            pass
    return newentries

<commit_msg>Bug fix for tolerance of phase diagram showing positive formation energy entries.<commit_after>#!/usr/bin/env python

"""
This module provides classes to create phase diagrams.
"""

__author__ = "Shyue Ping Ong"
__copyright__ = "Copyright 2011, The Materials Project"
__version__ = "1.0"
__maintainer__ = "Shyue Ping Ong"
__email__ = "shyue@mit.edu"
__status__ = "Production"
__date__ = "Sep 23, 2011"

import numpy as np
import logging
import itertools

from scipy.spatial import Delaunay

from pymatgen.core.structure import Composition
from pymatgen.command_line.qhull_caller import qconvex
from pymatgen.phasediagram.entries import GrandPotPDEntry, TransformedPDEntry
from pymatgen.core.periodic_table import Element
from pymatgen.analysis.reaction_calculator import Reaction

logger = logging.getLogger(__name__)

class PhaseDiagram (object):
    '''
    Simple phase diagram class taking in elements and entries as inputs.
    '''
    FORMATION_ENERGY_TOLERANCE = 1e-11

    def __init__(self, entries, elements=None, use_external_qhull=False):
        """
        Standard constructor for phase diagram.

        Args:
            entries:
                A list of PDEntry-like objects having an energy, energy_per_atom and composition.
            elements:
                Optional list of elements in the phase diagram. If set to None, the elements are determined from the
                the entries themselves.
        """
        if elements == None:
            elements = set()
            map(elements.update, [entry.composition.elements for entry in entries])
        self._all_entries = entries
        self._elements = tuple(elements)
        self._qhull_data = None
        self._facets = None
        self._qhull_entries = None
        self._stable_entries = None
        self._all_entries_hulldata = None
        self._use_external_qhull = use_external_qhull
        self._make_phasediagram()

    @property
    def all_entries(self):
        """
        All entries submitted to construct the PhaseDiagram.
        Note that this does not mean that all these entries are actually used in the phase diagram.
        """
        return self._all_entries

    @property
    def dim(self):
        return len(self._elements)

    @property
    def elements(self):
        """
        Elements in the phase diagram.
        """
        return self._elements

    @property
    def facets(self):
        """
        Facets of the phase diagram in the form of  [[1,2,3],[4,5,6]...]
        """
        return self._facets

    @property
    def qhull_data(self):
        """
        Data used in the convex hull operation. This is essentially a matrix of
        composition data and energy per atom values created from the qhull_entries.
        """
        return self._qhull_data

    @property
    def qhull_entries(self):
        """
        Actual entries used in convex hull. Excludes all positive formation energy entries.
        """
        return self._qhull_entries

    @property
    def unstable_entries(self):
        """
        Entries that are unstable in the phase diagram. Includes positive formation energy entries.
        """
        return [e for e in self.all_entries if e not in self.stable_entries]

    @property
    def stable_entries(self):
        '''
        Returns the stable entries.
        '''
        return self._stable_entries

    @property
    def all_entries_hulldata(self):
        """
        Same as qhull_data, but for all entries rather than just positive formation energy ones.
        """
        return self._all_entries_hulldata

    @property
    def el_refs(self):
        """
        List of elemental references for the phase diagrams.
        These are typically entries corresponding to the lowest energy element entries.
        """
        return self._el_refs

    def get_form_energy(self, entry):
        '''
        Returns the formation energy for an entry (NOT normalized) from the
        elemental references.

        Args:
            entry:
                A PDEntry

        Returns:
            Formation energy from the elementals references.
        '''
        comp = entry.composition
        energy = entry.energy - sum([comp[el] * self._el_refs[el].energy_per_atom for el in comp.elements])
        return energy

    def get_form_energy_per_atom(self, entry):
        '''
        Returns the formation energy per atom for an entry from the
        elemental references.
        '''
        comp = entry.composition
        return self.get_form_energy(entry) / comp.num_atoms

    def _process_entries_qhulldata(self, entries_to_process):
        data = list()
        for entry in entries_to_process:
            comp = entry.composition
            energy_per_atom = entry.energy_per_atom
            row = list()
            for i in xrange(1, len(self._elements)):
                row.append(comp.get_atomic_fraction(self._elements[i]))
            row.append(energy_per_atom)
            data.append(row)
        return data

    def _create_convhull_data(self):
        '''
        Make data suitable for convex hull procedure from the list of entries.
        '''
        logger.debug("Creating convex hull data...")
        #Determine the elemental references based on lowest energy for each.
        self._el_refs = dict()
        for entry in self._all_entries:
            if entry.composition.is_element:
                for el in entry.composition.elements:
                    if entry.composition[el] > Composition.amount_tolerance:
                        break
                e_per_atom = entry.energy_per_atom
                if el not in self._el_refs:
                    self._el_refs[el] = entry
                elif self._el_refs[el].energy_per_atom > e_per_atom:
                    self._el_refs[el] = entry
        # Remove positive formation energy entries
        entries_to_process = list()
        for entry in self._all_entries:
            if self.get_form_energy(entry) <= -self.FORMATION_ENERGY_TOLERANCE:
                entries_to_process.append(entry)
            else:
                logger.debug("Removing positive formation energy entry {}".format(entry))
        entries_to_process.extend([entry for entry in self._el_refs.values()])

        self._qhull_entries = entries_to_process
        return self._process_entries_qhulldata(entries_to_process)

    def _make_phasediagram(self):
        """
        Make the phase diagram.
        """
        stable_entries = set()
        dim = len(self._elements)
        self._qhull_data = self._create_convhull_data()
        if len(self._qhull_data) == dim:
            self._facets = [range(len(self._elements))]
        else:
            if self._use_external_qhull:
                logger.debug("> 4D hull encountered. Computing hull using external qconvex call.")
                self._facets = qconvex(self._qhull_data)
            else:
                logger.debug("Computing hull using scipy.spatial.delaunay")
                delau = Delaunay(self._qhull_data)
                self._facets = delau.convex_hull
            logger.debug("Final facets are\n{}".format(self._facets))

            logger.debug("Removing vertical facets...")
            finalfacets = list()
            for facet in self._facets:
                facetmatrix = np.zeros((len(facet), len(facet)))
                count = 0
                is_element_facet = True
                for vertex in facet:
                    facetmatrix[count] = np.array(self._qhull_data[vertex])
                    facetmatrix[count, dim - 1] = 1
                    count += 1
                    if len(self._qhull_entries[vertex].composition) > 1:
                        is_element_facet = False
                if abs(np.linalg.det(facetmatrix)) > 1e-8 and (not is_element_facet):
                    finalfacets.append(facet)
                else:
                    logger.debug("Removing vertical facet : {}".format(facet))
            self._facets = finalfacets

        for facet in self._facets:
            for vertex in facet:
                stable_entries.add(self._qhull_entries[vertex])
        self._stable_entries = stable_entries
        self._all_entries_hulldata = self._process_entries_qhulldata(self._all_entries)

class GrandPotentialPhaseDiagram (PhaseDiagram):
    '''
    Grand potential phase diagram class taking in elements and entries as inputs.
    '''

    def __init__(self, entries, chempots, elements=None, use_external_qhull=False):
        """
        Standard constructor for grand potential phase diagram.

        Args:
            entries:
                A list of PDEntry-like objects having an energy, energy_per_atom
                and composition.
            chempots:
                A dict of {element: float} to specify the chemical potentials
                of the open elements.
            elements:
                Optional list of elements in the phase diagram. If set to None,
                the elements are determined from the entries themselves.
        """
        if elements == None:
            elements = set()
            map(elements.update, [entry.composition.elements for entry in entries])
        allentries = list()
        for entry in entries:
            if not (entry.is_element and (entry.composition.elements[0] in chempots)):
                allentries.append(GrandPotPDEntry(entry, chempots))
        self.chempots = chempots
        filteredels = list()
        for el in elements:
            if el not in chempots:
                filteredels.append(el)
        elements = sorted(filteredels)
        super(GrandPotentialPhaseDiagram, self).__init__(allentries, elements, use_external_qhull)

class CompoundPhaseDiagram(PhaseDiagram):
    """
    Experimental feature. Generates phase diagrams from compounds as termninations
    instead of elements.
    """

    def __init__(self, entries, terminal_compositions, use_external_qhull=False):
        entries = get_entries_within_compositional_space(entries, terminal_compositions)
        elset = get_non_coplanar_element_set(entries)
        els = list(elset)
        pentries = get_transformed_entries(entries, els)
        super(CompoundPhaseDiagram, self).__init__(pentries, use_external_qhull=use_external_qhull)

def get_comp_matrix_from_comp(compositions, elements, normalize_row=True):
    """
    Helper function to generates a normalized composition matrix from a list of
    composition.
    """
    comp_matrix = np.array([[comp.get_atomic_fraction(el) for el in elements] for comp in compositions])
    if not normalize_row:
        return comp_matrix
    factor = np.tile(np.sum(comp_matrix, 1), (len(elements), 1)).transpose()
    return comp_matrix / factor

def get_comp_matrix(entries, elements, normalize_row=True):
    """
    Helper function to generates a normalized composition matrix from a list of
    composition.
    """
    return get_comp_matrix_from_comp([entry.composition for entry in entries], elements, normalize_row)

def is_coplanar(entries, elements):
    comp_matrix = get_comp_matrix(entries, elements)
    for submatrix in itertools.combinations(comp_matrix, min(len(elements), len(entries))):
        if abs(np.linalg.det(submatrix)) > 1e-5:
            return False
    return True

def get_non_coplanar_element_set(entries):
    elements = set()
    map(elements.update, [entry.composition.elements for entry in entries])
    for i in xrange(len(elements), 1, -1):
        for elset in itertools.combinations(elements, i):
            if not is_coplanar(entries, elset):
                return elset
    return None

def get_transformed_entries(entries, elements):
    comp_matrix = get_comp_matrix(entries, elements)
    newmat = []
    energies = []
    for i in xrange(len(elements)):
        col = comp_matrix[:, i]
        maxval = max(col)
        maxind = list(col).index(maxval)
        newmat.append(comp_matrix[maxind])
        energies.append(entries[i].energy_per_atom)
    invm = np.linalg.inv(np.array(newmat).transpose())
    newentries = []
    for i in xrange(len(entries)):
        entry = entries[i]
        lincomp = np.dot(invm, comp_matrix[i])
        lincomp = np.around(lincomp, 5)
        comp = Composition({Element.from_Z(j + 1):lincomp[j] for j in xrange(len(elements))})
        scaled_energy = entry.energy_per_atom - sum(lincomp * energies)
        newentries.append(TransformedPDEntry(comp, scaled_energy, entry))
    return newentries

def get_entries_within_compositional_space(entries, terminal_compositions):
    newentries = []
    for entry in entries:
        try:
            rxn = Reaction(terminal_compositions, [entry.composition])
            newentries.append(entry)
        except:
            pass
    return newentries

<|endoftext|>


<commit_before>"""
Backend for NCS VPN module.

Author: Henrik Thostrup Jensen <htj at nordu.net>
Copyright: NORDUnet(2011-2013)
"""

import base64
import random

from twisted.python import log

from opennsa import constants as cnt, config
from opennsa.backends.common import genericbackend
from opennsa.protocols.shared import httpclient

# basic payload
#
#<service xmlns="http://tail-f.com/ns/ncs" >
#  <object-id>nsi-vpn</object-id>
#  <type>
#    <vpn xmlns="http://nordu.net/ns/ncs/vpn">
#      <side-a>
#        <router>routerA</router>
#        <interface>interface1</interface>
#      </side-a>
#      <side-b>
#        <router></router>
#        <interface>ge-1/0/1</interface>
#      </side-b>
#      <encapsulation-type>ethernet-vlan</encapsulation-type>
#      <vlan>1720</vlan>
#    </vpn>
#  </type>
#</service>
#
# encapsulation type can be ethernet or ethernet-vlan
# vlan must be specified if encapsulation-type is ethernet-vlan, otherwise not
#
# the payload must be posted to the services url, e.g.,:
# http://localhost:8080/api/running/services
#
# To tear down the VPN, do a DELETE against
# "http://localhost:8080/api/running/services/service/nsi-vpn"
#
# The connection id -> object-id mapping is hence rather important to remember, but it can be the

ETHERNET_VPN_PAYLOAD_BASE = """
<service xmlns="http://tail-f.com/ns/ncs" >
  <object-id>%(service_name)s</object-id>
  <type>
    <vpn xmlns="http://nordu.net/ns/ncs/vpn">
      <side-a>
        <router>%(router_a)s</router>
        <interface>%(interface_a)s</interface>
      </side-a>
      <side-b>
        <router>%(router_b)s</router>
        <interface>%(interface_b)s</interface>
      </side-b>
      <encapsulation-type>ethernet</encapsulation-type>
    </vpn>
  </type>
</service>
"""

ETHERNET_VLAN_VPN_PAYLOAD_BASE = """
<service xmlns="http://tail-f.com/ns/ncs" >
  <object-id>%(service_name)s</object-id>
  <type>
    <vpn xmlns="http://nordu.net/ns/ncs/vpn">
      <side-a>
        <router>%(router_a)s</router>
        <interface>%(interface_a)s</interface>
      </side-a>
      <side-b>
        <router>%(router_b)s</router>
        <interface>%(interface_b)s</interface>
      </side-b>
      <encapsulation-type>ethernet-vlan</encapsulation-type>
      <vlan>%(vlan)i</vlan>
    </vpn>
  </type>
</service>
"""

LOG_SYSTEM = 'opennsa.ncsvpn'

class NCSVPNTarget(object):

    def __init__(self, router, interface, vlan=None):
        self.router = router
        self.interface = interface
        self.vlan = vlan

    def __str__(self):
        if self.vlan:
            return '<NCSVPNTarget %s/%s#%i>' % (self.router, self.interface, self.vlan)
        else:
            return '<NCSVPNTarget %s/%s>' % (self.router, self.interface)

def createVPNPayload(service_name, source_target, dest_target):

    intps = {
        'service_name'  : service_name,
        'router_a'      : source_target.router,
        'interface_a'   : source_target.interface,
        'router_b'      : dest_target.router,
        'interface_b'   : dest_target.interface
    }

    if source_target.vlan and dest_target.vlan:
        assert source_target.vlan == dest_target.vlan, 'VLANs must match (until we get rewrite in place)'
        intps['vlan'] = source_target.vlan
        payload = ETHERNET_VLAN_VPN_PAYLOAD_BASE % intps
    else:
        payload = ETHERNET_VPN_PAYLOAD_BASE % intps

    return payload

class NCSVPNConnectionManager:

    def __init__(self, ncs_services_url, user, password, log_system):
        self.ncs_services_url = ncs_services_url
        self.user             = user
        self.password         = password
        self.log_system       = log_system

    def getResource(self, port, label_type, label_value):
        assert label_type in (None, cnt.ETHERNET_VLAN), 'Label must be None or VLAN'
        return port # this contains router and port

    def getTarget(self, port, label_type, label_value):
        assert label_type in (None, cnt.ETHERNET_VLAN), 'Label must be None or VLAN'
        if label_type == cnt.ETHERNET_VLAN:
            vlan = int(label_value)
            assert 1 <= vlan <= 4095, 'Invalid label value for vlan: %s' % label_value

        router, interface = port.split(':')
        return NCSVPNTarget(router, interface, vlan)

    def createConnectionId(self, source_target, dest_target):
        return 'ON-' + str(random.randint(100000,999999))

    def canSwapLabel(self, label_type):
        return False
        #return label_type == cnt.ETHERNET_VLAN:

    def _createAuthzHeader(self):
        return 'Basic ' + base64.b64encode( self.user + ':' + self.password)

    def _createHeaders(self):
        headers = {}
        headers['Content-Type'] = 'application/vnd.yang.data+xml'
        headers['Authorization'] = self._createAuthzHeader()
        return headers

    def setupLink(self, connection_id, source_target, dest_target, bandwidth):
        payload = createVPNPayload(connection_id, source_target, dest_target)
        headers = self._createHeaders()

        def linkUp(_):
            log.msg('Link %s -> %s up' % (source_target, dest_target), system=self.log_system)

        d = httpclient.httpRequest(self.ncs_services_url, payload, headers, method='POST')
        d.addCallback(linkUp)
        return d

    def teardownLink(self, connection_id, source_target, dest_target, bandwidth):
        service_url = self.ncs_services_url + '/service/' + connection_id
        headers = self._createHeaders()

        def linkDown(_):
            log.msg('Link %s -> %s down' % (source_target, dest_target), system=self.log_system)

        d = httpclient.httpRequest(service_url, None, headers, method='DELETE')
        d.addCallback(linkDown)
        return d

def NCSVPNBackend(network_name, network_topology, parent_requester, port_map, cfg):

    name = 'NCS VPN (%s)' % network_name

    # extract config items
    ncs_services_url = str(cfg[config.NCS_SERVICES_URL]) # convert from unicode
    user             = cfg[config.NCS_USER]
    password         = cfg[config.NCS_PASSWORD]

    cm = NCSVPNConnectionManager(ncs_services_url, user, password, name)
    return genericbackend.GenericBackend(network_name, network_topology, cm, parent_requester, name)

<commit_msg>slightly shorter logging name for ncs vpn backend<commit_after>"""
Backend for NCS VPN module.

Author: Henrik Thostrup Jensen <htj at nordu.net>
Copyright: NORDUnet(2011-2013)
"""

import base64
import random

from twisted.python import log

from opennsa import constants as cnt, config
from opennsa.backends.common import genericbackend
from opennsa.protocols.shared import httpclient

# basic payload
#
#<service xmlns="http://tail-f.com/ns/ncs" >
#  <object-id>nsi-vpn</object-id>
#  <type>
#    <vpn xmlns="http://nordu.net/ns/ncs/vpn">
#      <side-a>
#        <router>routerA</router>
#        <interface>interface1</interface>
#      </side-a>
#      <side-b>
#        <router></router>
#        <interface>ge-1/0/1</interface>
#      </side-b>
#      <encapsulation-type>ethernet-vlan</encapsulation-type>
#      <vlan>1720</vlan>
#    </vpn>
#  </type>
#</service>
#
# encapsulation type can be ethernet or ethernet-vlan
# vlan must be specified if encapsulation-type is ethernet-vlan, otherwise not
#
# the payload must be posted to the services url, e.g.,:
# http://localhost:8080/api/running/services
#
# To tear down the VPN, do a DELETE against
# "http://localhost:8080/api/running/services/service/nsi-vpn"
#
# The connection id -> object-id mapping is hence rather important to remember, but it can be the

ETHERNET_VPN_PAYLOAD_BASE = """
<service xmlns="http://tail-f.com/ns/ncs" >
  <object-id>%(service_name)s</object-id>
  <type>
    <vpn xmlns="http://nordu.net/ns/ncs/vpn">
      <side-a>
        <router>%(router_a)s</router>
        <interface>%(interface_a)s</interface>
      </side-a>
      <side-b>
        <router>%(router_b)s</router>
        <interface>%(interface_b)s</interface>
      </side-b>
      <encapsulation-type>ethernet</encapsulation-type>
    </vpn>
  </type>
</service>
"""

ETHERNET_VLAN_VPN_PAYLOAD_BASE = """
<service xmlns="http://tail-f.com/ns/ncs" >
  <object-id>%(service_name)s</object-id>
  <type>
    <vpn xmlns="http://nordu.net/ns/ncs/vpn">
      <side-a>
        <router>%(router_a)s</router>
        <interface>%(interface_a)s</interface>
      </side-a>
      <side-b>
        <router>%(router_b)s</router>
        <interface>%(interface_b)s</interface>
      </side-b>
      <encapsulation-type>ethernet-vlan</encapsulation-type>
      <vlan>%(vlan)i</vlan>
    </vpn>
  </type>
</service>
"""

LOG_SYSTEM = 'opennsa.ncsvpn'

class NCSVPNTarget(object):

    def __init__(self, router, interface, vlan=None):
        self.router = router
        self.interface = interface
        self.vlan = vlan

    def __str__(self):
        if self.vlan:
            return '<NCSVPNTarget %s/%s#%i>' % (self.router, self.interface, self.vlan)
        else:
            return '<NCSVPNTarget %s/%s>' % (self.router, self.interface)

def createVPNPayload(service_name, source_target, dest_target):

    intps = {
        'service_name'  : service_name,
        'router_a'      : source_target.router,
        'interface_a'   : source_target.interface,
        'router_b'      : dest_target.router,
        'interface_b'   : dest_target.interface
    }

    if source_target.vlan and dest_target.vlan:
        assert source_target.vlan == dest_target.vlan, 'VLANs must match (until we get rewrite in place)'
        intps['vlan'] = source_target.vlan
        payload = ETHERNET_VLAN_VPN_PAYLOAD_BASE % intps
    else:
        payload = ETHERNET_VPN_PAYLOAD_BASE % intps

    return payload

class NCSVPNConnectionManager:

    def __init__(self, ncs_services_url, user, password, log_system):
        self.ncs_services_url = ncs_services_url
        self.user             = user
        self.password         = password
        self.log_system       = log_system

    def getResource(self, port, label_type, label_value):
        assert label_type in (None, cnt.ETHERNET_VLAN), 'Label must be None or VLAN'
        return port # this contains router and port

    def getTarget(self, port, label_type, label_value):
        assert label_type in (None, cnt.ETHERNET_VLAN), 'Label must be None or VLAN'
        if label_type == cnt.ETHERNET_VLAN:
            vlan = int(label_value)
            assert 1 <= vlan <= 4095, 'Invalid label value for vlan: %s' % label_value

        router, interface = port.split(':')
        return NCSVPNTarget(router, interface, vlan)

    def createConnectionId(self, source_target, dest_target):
        return 'ON-' + str(random.randint(100000,999999))

    def canSwapLabel(self, label_type):
        return False
        #return label_type == cnt.ETHERNET_VLAN:

    def _createAuthzHeader(self):
        return 'Basic ' + base64.b64encode( self.user + ':' + self.password)

    def _createHeaders(self):
        headers = {}
        headers['Content-Type'] = 'application/vnd.yang.data+xml'
        headers['Authorization'] = self._createAuthzHeader()
        return headers

    def setupLink(self, connection_id, source_target, dest_target, bandwidth):
        payload = createVPNPayload(connection_id, source_target, dest_target)
        headers = self._createHeaders()

        def linkUp(_):
            log.msg('Link %s -> %s up' % (source_target, dest_target), system=self.log_system)

        d = httpclient.httpRequest(self.ncs_services_url, payload, headers, method='POST')
        d.addCallback(linkUp)
        return d

    def teardownLink(self, connection_id, source_target, dest_target, bandwidth):
        service_url = self.ncs_services_url + '/service/' + connection_id
        headers = self._createHeaders()

        def linkDown(_):
            log.msg('Link %s -> %s down' % (source_target, dest_target), system=self.log_system)

        d = httpclient.httpRequest(service_url, None, headers, method='DELETE')
        d.addCallback(linkDown)
        return d

def NCSVPNBackend(network_name, network_topology, parent_requester, port_map, cfg):

    name = 'NCS VPN %s' % network_name

    # extract config items
    ncs_services_url = str(cfg[config.NCS_SERVICES_URL]) # convert from unicode
    user             = cfg[config.NCS_USER]
    password         = cfg[config.NCS_PASSWORD]

    cm = NCSVPNConnectionManager(ncs_services_url, user, password, name)
    return genericbackend.GenericBackend(network_name, network_topology, cm, parent_requester, name)

<|endoftext|>