<commit_before>#!/usr/bin/python
# Copyright 2009 The Tor Project, Inc.
# License at end of file.
#
# This tests connections to a list of Tor nodes in a given Tor consensus file
# while also recording the certificates - it's not a perfect tool but complete
# or even partial failure should raise alarms.
#
# This plugoo uses threads and as a result, it's not friendly to SIGINT signals.
#

import logging
import socket
import time
import random
import threading
import sys
import os
try:
    from ooni.plugooni import Plugoo
    from ooni.common import Storage
    from ooni import output
except:
    print "ERR: Not running from ooni"

ssl = OpenSSL = None

try:
    import ssl
except ImportError:
    pass

if ssl is None:
    try:
        import OpenSSL.SSL
        import OpenSSL.crypto
    except ImportError:
        pass

if ssl is None and OpenSSL is None:
    if socket.ssl:
        print """Your Python is too old to have the ssl module, and you haven't
installed pyOpenSSL.  I'll try to work with what you've got, but I can't
record certificates so well."""
    else:
        print """Your Python has no OpenSSL support.  Upgrade to 2.6, install
pyOpenSSL, or both."""
        sys.exit(1)

################################################################

# How many servers should we test in parallel?
N_THREADS = 16

# How long do we give individual socket operations to succeed or fail?
# (Seconds)
TIMEOUT = 10

################################################################

CONNECTING = "noconnect"
HANDSHAKING = "nohandshake"
OK = "ok"
ERROR = "err"

LOCK = threading.RLock()
socket.setdefaulttimeout(TIMEOUT)

def clean_pem_cert(cert):
    idx = cert.find('-----END')
    if idx > 1 and cert[idx-1] != '\n':
        cert = cert.replace('-----END','\n-----END')
    return cert

def record((addr,port), state, extra=None, cert=None):
    LOCK.acquire()
    try:
        OUT.append({'addr' : addr,
                         'port' : port,
                         'state' : state,
                         'extra' : extra})
        if cert:
            CERT_OUT.append({'addr' : addr,
                                  'port' : port,
                                  'clean_cert' : clean_pem_cert(cert)})
    finally:
        LOCK.release()

def probe(address,theCtx=None):
    sock = s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    logging.info("Opening socket to %s",address)
    try:
        s.connect(address)
    except IOError, e:
        logging.info("Error %s from socket connect.",e)
        record(address, CONNECTING, e)
        s.close()
        return
    logging.info("Socket to %s open.  Launching SSL handshake.",address)
    if ssl:
        try:
            s = ssl.wrap_socket(s,cert_reqs=ssl.CERT_NONE,ca_certs=None)
            # "MARCO!"
            s.do_handshake()
        except IOError, e:
            logging.info("Error %s from ssl handshake",e)
            record(address, HANDSHAKING, e)
            s.close()
            sock.close()
            return
        cert = s.getpeercert(True)
        if cert != None:
            cert = ssl.DER_cert_to_PEM_cert(cert)
    elif OpenSSL:
        try:
            s = OpenSSL.SSL.Connection(theCtx, s)
            s.set_connect_state()
            s.setblocking(True)
            s.do_handshake()
            cert = s.get_peer_certificate()
            if cert != None:
                cert = OpenSSL.crypto.dump_certificate(
                    OpenSSL.crypto.FILETYPE_PEM, cert)
        except IOError, e:
            logging.info("Error %s from OpenSSL handshake",e)
            record(address, HANDSHAKING, e)
            s.close()
            sock.close()
            return
    else:
        try:
            s = socket.ssl(s)
            s.write('a')
            cert = s.server()
        except IOError, e:
            logging.info("Error %s from socket.ssl handshake",e)
            record(address, HANDSHAKING, e)
            sock.close()
            return

    logging.info("SSL handshake with %s finished",address)
    # "POLO!"
    record(address,OK, cert=cert)
    if (ssl or OpenSSL):
        s.close()
    sock.close()

def parseNetworkstatus(ns):
    for line in ns:
        if line.startswith('r '):
            r = line.split()
            yield (r[-3],int(r[-2]))

def parseCachedDescs(cd):
    for line in cd:
        if line.startswith('router '):
            r = line.split()
            yield (r[2],int(r[3]))

def worker(addrList, origLength):
    done = False
    logging.info("Launching thread.")

    if OpenSSL is not None:
        context = OpenSSL.SSL.Context(OpenSSL.SSL.TLSv1_METHOD)
    else:
        context = None

    while True:
        LOCK.acquire()
        try:
            if addrList:
                print "Starting test %d/%d"%(
                    1+origLength-len(addrList),origLength)
                addr = addrList.pop()
            else:
                return
        finally:
            LOCK.release()

        try:
            logging.info("Launching probe for %s",addr)
            probe(addr, context)
        except Exception, e:
            logging.info("Unexpected error from %s",addr)
            record(addr, ERROR, e)

def runThreaded(addrList, nThreads):
    ts = []
    origLen = len(addrList)
    for num in xrange(nThreads):
        t = threading.Thread(target=worker, args=(addrList,origLen))
        t.setName("Th#%s"%num)
        ts.append(t)
        t.start()
    for t in ts:
        logging.info("Joining thread %s",t.getName())
        t.join()

def main(self, args):
    # BEGIN
    # This logic should be present in more or less all plugoos
    global OUT
    global CERT_OUT
    global OUT_DATA
    global CERT_OUT_DATA
    OUT_DATA = []
    CERT_OUT_DATA = []

    try:
        OUT = output.data(name=args.output.main) #open(args.output.main, 'w')
    except:
        print "No output file given. quitting..."
        return -1

    try:
        CERT_OUT = output.data(args.output.certificates) #open(args.output.certificates, 'w')
    except:
        print "No output cert file given. quitting..."
        return -1

    logging.basicConfig(format='%(asctime)s [%(levelname)s] [%(threadName)s] %(message)s',
                        datefmt="%b %d %H:%M:%S",
                        level=logging.INFO,
                        filename=args.log)
    logging.info("============== STARTING NEW LOG")
    # END

    if ssl is not None:
        methodName = "ssl"
    elif OpenSSL is not None:
        methodName = "OpenSSL"
    else:
        methodName = "socket"
    logging.info("Running marco with method '%s'", methodName)

    addresses = []

    if args.input.ips:
        for fn in input.file(args.input.ips).simple():
            addresses.append( (a,b) )

    elif args.input.consensus:
        for fn in args:
            print fn
            for a,b in parseNetworkstatus(open(args.input.consensus)):
                addresses.append( (a,b) )

    if args.input.randomize:
        # Take a random permutation of the set the knuth way!
        for i in range(0, len(addresses)):
            j = random.randint(0, i)
            addresses[i], addresses[j] = addresses[j], addresses[i]

    if len(addresses) == 0:
        logging.error("No input source given, quiting...")
        return -1

    addresses = list(addresses)

    if not args.input.randomize:
        addresses.sort()

    runThreaded(addresses, N_THREADS)

class MarcoPlugin(Plugoo):
  def __init__(self):
    self.name = ""

    self.modules = [ "logging", "socket", "time", "random", "threading", "sys",
                     "OpenSSL.SSL", "OpenSSL.crypto", "os" ]

    self.input = Storage()
    self.input.ip = None
    self.input.consensus = os.path.expanduser("~/.tor/cached-consensus")

    self.output = Storage()
    self.output.main = 'reports/marco.yamlooni'
    self.output.certificates = 'reports/marco_certs.out'

    # We require for Tor to already be running or have recently run
    self.args = Storage()
    self.args.input = self.input
    self.args.output = self.output
    self.args.log = 'reports/marco.log'

  def ooni_main(self, cmd):
    self.args.input.randomize = cmd.randomize
    self.args.input.ip = cmd.listfile
    main(self, self.args)

if __name__ == '__main__':
    if len(sys.argv) < 2:
        print >> sys.stderr, ("This script takes one or more networkstatus "
                              "files as an argument.")
    self = None
    main(self, sys.argv[1:])

# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are
# met:
#
#     * Redistributions of source code must retain the above copyright
# notice, this list of conditions and the following disclaimer.
#
#     * Redistributions in binary form must reproduce the above
# copyright notice, this list of conditions and the following disclaimer
# in the documentation and/or other materials provided with the
# distribution.
#
#     * Neither the names of the copyright owners nor the names of its
# contributors may be used to endorse or promote products derived from
# this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
# "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
# LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
# A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
# OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
# SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
# LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
# DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
# THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
<commit_msg>Update the path to the tor consensus file, might be better to take this as input<commit_after>#!/usr/bin/python
# Copyright 2009 The Tor Project, Inc.
# License at end of file.
#
# This tests connections to a list of Tor nodes in a given Tor consensus file
# while also recording the certificates - it's not a perfect tool but complete
# or even partial failure should raise alarms.
#
# This plugoo uses threads and as a result, it's not friendly to SIGINT signals.
#

import logging
import socket
import time
import random
import threading
import sys
import os
try:
    from ooni.plugooni import Plugoo
except:
    print "Error importing Plugoo"

try:
    from ooni.common import Storage
except:
    print "Error importing Storage"

try:
    from ooni import output
except:
    print "Error importing output"

ssl = OpenSSL = None

try:
    import ssl
except ImportError:
    pass

if ssl is None:
    try:
        import OpenSSL.SSL
        import OpenSSL.crypto
    except ImportError:
        pass

if ssl is None and OpenSSL is None:
    if socket.ssl:
        print """Your Python is too old to have the ssl module, and you haven't
installed pyOpenSSL.  I'll try to work with what you've got, but I can't
record certificates so well."""
    else:
        print """Your Python has no OpenSSL support.  Upgrade to 2.6, install
pyOpenSSL, or both."""
        sys.exit(1)

################################################################

# How many servers should we test in parallel?
N_THREADS = 16

# How long do we give individual socket operations to succeed or fail?
# (Seconds)
TIMEOUT = 10

################################################################

CONNECTING = "noconnect"
HANDSHAKING = "nohandshake"
OK = "ok"
ERROR = "err"

LOCK = threading.RLock()
socket.setdefaulttimeout(TIMEOUT)

def clean_pem_cert(cert):
    idx = cert.find('-----END')
    if idx > 1 and cert[idx-1] != '\n':
        cert = cert.replace('-----END','\n-----END')
    return cert

def record((addr,port), state, extra=None, cert=None):
    LOCK.acquire()
    try:
        OUT.append({'addr' : addr,
                         'port' : port,
                         'state' : state,
                         'extra' : extra})
        if cert:
            CERT_OUT.append({'addr' : addr,
                                  'port' : port,
                                  'clean_cert' : clean_pem_cert(cert)})
    finally:
        LOCK.release()

def probe(address,theCtx=None):
    sock = s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    logging.info("Opening socket to %s",address)
    try:
        s.connect(address)
    except IOError, e:
        logging.info("Error %s from socket connect.",e)
        record(address, CONNECTING, e)
        s.close()
        return
    logging.info("Socket to %s open.  Launching SSL handshake.",address)
    if ssl:
        try:
            s = ssl.wrap_socket(s,cert_reqs=ssl.CERT_NONE,ca_certs=None)
            # "MARCO!"
            s.do_handshake()
        except IOError, e:
            logging.info("Error %s from ssl handshake",e)
            record(address, HANDSHAKING, e)
            s.close()
            sock.close()
            return
        cert = s.getpeercert(True)
        if cert != None:
            cert = ssl.DER_cert_to_PEM_cert(cert)
    elif OpenSSL:
        try:
            s = OpenSSL.SSL.Connection(theCtx, s)
            s.set_connect_state()
            s.setblocking(True)
            s.do_handshake()
            cert = s.get_peer_certificate()
            if cert != None:
                cert = OpenSSL.crypto.dump_certificate(
                    OpenSSL.crypto.FILETYPE_PEM, cert)
        except IOError, e:
            logging.info("Error %s from OpenSSL handshake",e)
            record(address, HANDSHAKING, e)
            s.close()
            sock.close()
            return
    else:
        try:
            s = socket.ssl(s)
            s.write('a')
            cert = s.server()
        except IOError, e:
            logging.info("Error %s from socket.ssl handshake",e)
            record(address, HANDSHAKING, e)
            sock.close()
            return

    logging.info("SSL handshake with %s finished",address)
    # "POLO!"
    record(address,OK, cert=cert)
    if (ssl or OpenSSL):
        s.close()
    sock.close()

def parseNetworkstatus(ns):
    for line in ns:
        if line.startswith('r '):
            r = line.split()
            yield (r[-3],int(r[-2]))

def parseCachedDescs(cd):
    for line in cd:
        if line.startswith('router '):
            r = line.split()
            yield (r[2],int(r[3]))

def worker(addrList, origLength):
    done = False
    logging.info("Launching thread.")

    if OpenSSL is not None:
        context = OpenSSL.SSL.Context(OpenSSL.SSL.TLSv1_METHOD)
    else:
        context = None

    while True:
        LOCK.acquire()
        try:
            if addrList:
                print "Starting test %d/%d"%(
                    1+origLength-len(addrList),origLength)
                addr = addrList.pop()
            else:
                return
        finally:
            LOCK.release()

        try:
            logging.info("Launching probe for %s",addr)
            probe(addr, context)
        except Exception, e:
            logging.info("Unexpected error from %s",addr)
            record(addr, ERROR, e)

def runThreaded(addrList, nThreads):
    ts = []
    origLen = len(addrList)
    for num in xrange(nThreads):
        t = threading.Thread(target=worker, args=(addrList,origLen))
        t.setName("Th#%s"%num)
        ts.append(t)
        t.start()
    for t in ts:
        logging.info("Joining thread %s",t.getName())
        t.join()

def main(self, args):
    # BEGIN
    # This logic should be present in more or less all plugoos
    global OUT
    global CERT_OUT
    global OUT_DATA
    global CERT_OUT_DATA
    OUT_DATA = []
    CERT_OUT_DATA = []

    try:
        OUT = output.data(name=args.output.main) #open(args.output.main, 'w')
    except:
        print "No output file given. quitting..."
        return -1

    try:
        CERT_OUT = output.data(args.output.certificates) #open(args.output.certificates, 'w')
    except:
        print "No output cert file given. quitting..."
        return -1

    logging.basicConfig(format='%(asctime)s [%(levelname)s] [%(threadName)s] %(message)s',
                        datefmt="%b %d %H:%M:%S",
                        level=logging.INFO,
                        filename=args.log)
    logging.info("============== STARTING NEW LOG")
    # END

    if ssl is not None:
        methodName = "ssl"
    elif OpenSSL is not None:
        methodName = "OpenSSL"
    else:
        methodName = "socket"
    logging.info("Running marco with method '%s'", methodName)

    addresses = []

    if args.input.ips:
        for fn in input.file(args.input.ips).simple():
            addresses.append( (a,b) )

    elif args.input.consensus:
        for fn in args:
            print fn
            for a,b in parseNetworkstatus(open(args.input.consensus)):
                addresses.append( (a,b) )

    if args.input.randomize:
        # Take a random permutation of the set the knuth way!
        for i in range(0, len(addresses)):
            j = random.randint(0, i)
            addresses[i], addresses[j] = addresses[j], addresses[i]

    if len(addresses) == 0:
        logging.error("No input source given, quiting...")
        return -1

    addresses = list(addresses)

    if not args.input.randomize:
        addresses.sort()

    runThreaded(addresses, N_THREADS)

class MarcoPlugin(Plugoo):
  def __init__(self):
    self.name = ""

    self.modules = [ "logging", "socket", "time", "random", "threading", "sys",
                     "OpenSSL.SSL", "OpenSSL.crypto", "os" ]

    self.input = Storage()
    self.input.ip = None
    try:
        self.input.consensus = os.path.expanduser("~/.tor/cached-consensus")
    try:
        self.input.consensus = os.path.expanduser("~/tor/bundle/tor-browser_en-US/Data/Tor/cached-consensus")
    if not self.input.consensus:
        print "Error importing consensus file"
        return -2

    self.output = Storage()
    self.output.main = 'reports/marco.yamlooni'
    self.output.certificates = 'reports/marco_certs.out'

    # We require for Tor to already be running or have recently run
    self.args = Storage()
    self.args.input = self.input
    self.args.output = self.output
    self.args.log = 'reports/marco.log'

  def ooni_main(self, cmd):
    self.args.input.randomize = cmd.randomize
    self.args.input.ip = cmd.listfile
    main(self, self.args)

if __name__ == '__main__':
    if len(sys.argv) < 2:
        print >> sys.stderr, ("This script takes one or more networkstatus "
                              "files as an argument.")
    self = None
    main(self, sys.argv[1:])

# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are
# met:
#
#     * Redistributions of source code must retain the above copyright
# notice, this list of conditions and the following disclaimer.
#
#     * Redistributions in binary form must reproduce the above
# copyright notice, this list of conditions and the following disclaimer
# in the documentation and/or other materials provided with the
# distribution.
#
#     * Neither the names of the copyright owners nor the names of its
# contributors may be used to endorse or promote products derived from
# this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
# "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
# LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
# A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
# OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
# SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
# LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
# DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
# THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
##############################################################################
#
#    OpenERP, Open Source Management Solution
#    Copyright (C) 2004-2009 Tiny SPRL (<http://tiny.be>).
#
#    This program is free software: you can redistribute it and/or modify
#    it under the terms of the GNU Affero General Public License as
#    published by the Free Software Foundation, either version 3 of the
#    License, or (at your option) any later version.
#
#    This program is distributed in the hope that it will be useful,
#    but WITHOUT ANY WARRANTY; without even the implied warranty of
#    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#    GNU Affero General Public License for more details.
#
#    You should have received a copy of the GNU Affero General Public License
#    along with this program.  If not, see <http://www.gnu.org/licenses/>.
#
##############################################################################

import logging
import os
import signal
import sys
import threading
import time

import netrpc_server
import web_services

import openerp.cron
import openerp.modules
import openerp.netsvc
import openerp.osv
import openerp.tools
import openerp.service.wsgi_server

#.apidoc title: RPC Services

""" Classes of this module implement the network protocols that the
    OpenERP server uses to communicate with remote clients.

    Some classes are mostly utilities, whose API need not be visible to
    the average user/developer. Study them only if you are about to
    implement an extension to the network protocols, or need to debug some
    low-level behavior of the wire.
"""

_logger = logging.getLogger(__name__)

def load_server_wide_modules():
    for m in openerp.conf.server_wide_modules:
        try:
            openerp.modules.module.load_openerp_module(m)
        except Exception:
            msg = ''
            if m == 'web':
                msg = """
The `web` module is provided by the addons found in the `openerp-web` project.
Maybe you forgot to add those addons in your addons_path configuration."""
            _logger.exception('Failed to load server-wide module `%s`.%s', m, msg)

start_internal_done = False

def start_internal():
    global start_internal_done
    if start_internal_done:
        return
    openerp.netsvc.init_logger()
    openerp.modules.loading.open_openerp_namespace()
    # Instantiate local services (this is a legacy design).
    openerp.osv.osv.start_object_proxy()
    # Export (for RPC) services.
    web_services.start_web_services()
    load_server_wide_modules()
    start_internal_done = True

def start_services():
    """ Start all services including http, netrpc and cron """
    start_internal()

    # Initialize the HTTP stack.
    netrpc_server.init_servers()

    # Start the main cron thread.
    if openerp.conf.max_cron_threads:
        openerp.cron.start_master_thread()

    # Start the top-level servers threads (normally HTTP, HTTPS, and NETRPC).
    openerp.netsvc.Server.startAll()

    # Start the WSGI server.
    openerp.service.wsgi_server.start_server()

def stop_services():
    """ Stop all services. """
    # stop scheduling new jobs; we will have to wait for the jobs to complete below
    openerp.cron.cancel_all()

    openerp.netsvc.Server.quitAll()
    openerp.service.wsgi_server.stop_server()
    _logger.info("Initiating shutdown")
    _logger.info("Hit CTRL-C again or send a second signal to force the shutdown.")

    # Manually join() all threads before calling sys.exit() to allow a second signal
    # to trigger _force_quit() in case some non-daemon threads won't exit cleanly.
    # threading.Thread.join() should not mask signals (at least in python 2.5).
    me = threading.currentThread()
    _logger.debug('current thread: %r', me)
    for thread in threading.enumerate():
        _logger.debug('process %r (%r)', thread, thread.isDaemon())
        if thread != me and not thread.isDaemon():
            while thread.isAlive():
                _logger.debug('join and sleep')
                # Need a busyloop here as thread.join() masks signals
                # and would prevent the forced shutdown.
                thread.join(0.05)
                time.sleep(0.05)

    _logger.debug('--')
    openerp.modules.registry.RegistryManager.delete_all()
    logging.shutdown()

def start_services_workers():
    import openerp.service.workers
    openerp.multi_process = True
    openerp.service.workers.Multicorn(openerp.service.wsgi_server.application).run()

def _reexec():
    """reexecute openerp-server process with (nearly) the same arguments"""
    if openerp.tools.osutil.is_running_as_nt_service():
        os._exit(1)
    strip_args = ['-d', '-u']
    a = sys.argv[:]
    args = [x for i, x in enumerate(a) if x not in strip_args and a[max(i - 1, 0)] not in strip_args]
    os.execv(sys.executable, [sys.executable] + args)

def restart_server():
    if openerp.multi_process:
        raise NotImplementedError("Multicorn is not supported (but gunicorn was)")
        pid = openerp.wsgi.core.arbiter_pid
        os.kill(pid, signal.SIGHUP)
    else:
        if os.name == 'nt':
            def reborn():
                stop_services()
                _reexec()

            # run in a thread to let the current thread return response to the caller.
            threading.Thread(target=reborn).start()
        else:
            openerp.phoenix = True
            os.kill(os.getpid(), signal.SIGINT)

# vim:expandtab:smartindent:tabstop=4:softtabstop=4:shiftwidth=4:
<commit_msg>[FIX] services: call os.abort() for killing<commit_after># -*- coding: utf-8 -*-
##############################################################################
#
#    OpenERP, Open Source Management Solution
#    Copyright (C) 2004-2009 Tiny SPRL (<http://tiny.be>).
#
#    This program is free software: you can redistribute it and/or modify
#    it under the terms of the GNU Affero General Public License as
#    published by the Free Software Foundation, either version 3 of the
#    License, or (at your option) any later version.
#
#    This program is distributed in the hope that it will be useful,
#    but WITHOUT ANY WARRANTY; without even the implied warranty of
#    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#    GNU Affero General Public License for more details.
#
#    You should have received a copy of the GNU Affero General Public License
#    along with this program.  If not, see <http://www.gnu.org/licenses/>.
#
##############################################################################

import logging
import os
import signal
import sys
import threading
import time

import netrpc_server
import web_services

import openerp.cron
import openerp.modules
import openerp.netsvc
import openerp.osv
import openerp.tools
import openerp.service.wsgi_server

#.apidoc title: RPC Services

""" Classes of this module implement the network protocols that the
    OpenERP server uses to communicate with remote clients.

    Some classes are mostly utilities, whose API need not be visible to
    the average user/developer. Study them only if you are about to
    implement an extension to the network protocols, or need to debug some
    low-level behavior of the wire.
"""

_logger = logging.getLogger(__name__)

def load_server_wide_modules():
    for m in openerp.conf.server_wide_modules:
        try:
            openerp.modules.module.load_openerp_module(m)
        except Exception:
            msg = ''
            if m == 'web':
                msg = """
The `web` module is provided by the addons found in the `openerp-web` project.
Maybe you forgot to add those addons in your addons_path configuration."""
            _logger.exception('Failed to load server-wide module `%s`.%s', m, msg)

start_internal_done = False

def start_internal():
    global start_internal_done
    if start_internal_done:
        return
    openerp.netsvc.init_logger()
    openerp.modules.loading.open_openerp_namespace()
    # Instantiate local services (this is a legacy design).
    openerp.osv.osv.start_object_proxy()
    # Export (for RPC) services.
    web_services.start_web_services()
    load_server_wide_modules()
    start_internal_done = True

def start_services():
    """ Start all services including http, netrpc and cron """
    start_internal()

    # Initialize the HTTP stack.
    netrpc_server.init_servers()

    # Start the main cron thread.
    if openerp.conf.max_cron_threads:
        openerp.cron.start_master_thread()

    # Start the top-level servers threads (normally HTTP, HTTPS, and NETRPC).
    openerp.netsvc.Server.startAll()

    # Start the WSGI server.
    openerp.service.wsgi_server.start_server()

def stop_services():
    """ Stop all services. """
    # stop scheduling new jobs; we will have to wait for the jobs to complete below
    openerp.cron.cancel_all()

    openerp.netsvc.Server.quitAll()
    openerp.service.wsgi_server.stop_server()
    _logger.info("Initiating shutdown")
    _logger.info("Hit CTRL-C again or send a second signal to force the shutdown.")

    # Manually join() all threads before calling sys.exit() to allow a second signal
    # to trigger _force_quit() in case some non-daemon threads won't exit cleanly.
    # threading.Thread.join() should not mask signals (at least in python 2.5).
    me = threading.currentThread()
    _logger.debug('current thread: %r', me)
    for thread in threading.enumerate():
        _logger.debug('process %r (%r)', thread, thread.isDaemon())
        if thread != me and not thread.isDaemon():
            while thread.isAlive():
                _logger.debug('join and sleep')
                # Need a busyloop here as thread.join() masks signals
                # and would prevent the forced shutdown.
                thread.join(0.05)
                time.sleep(0.05)

    _logger.debug('--')
    openerp.modules.registry.RegistryManager.delete_all()
    logging.shutdown()

def start_services_workers():
    import openerp.service.workers
    openerp.multi_process = True
    openerp.service.workers.Multicorn(openerp.service.wsgi_server.application).run()

def _reexec():
    """reexecute openerp-server process with (nearly) the same arguments"""
    if openerp.tools.osutil.is_running_as_nt_service():
        os.abort()
    strip_args = ['-d', '-u']
    a = sys.argv[:]
    args = [x for i, x in enumerate(a) if x not in strip_args and a[max(i - 1, 0)] not in strip_args]
    os.execv(sys.executable, [sys.executable] + args)

def restart_server():
    if openerp.multi_process:
        raise NotImplementedError("Multicorn is not supported (but gunicorn was)")
        pid = openerp.wsgi.core.arbiter_pid
        os.kill(pid, signal.SIGHUP)
    else:
        if os.name == 'nt':
            def reborn():
                stop_services()
                _reexec()

            # run in a thread to let the current thread return response to the caller.
            threading.Thread(target=reborn).start()
        else:
            openerp.phoenix = True
            os.kill(os.getpid(), signal.SIGINT)

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


<commit_before># -*- coding: utf-8 -*-
# vim: tabstop=4 shiftwidth=4 softtabstop=4
#
# Copyright (C) 2017-2019 GEM Foundation
#
# OpenQuake is free software: you can redistribute it and/or modify it
# under the terms of the GNU Affero General Public License as published
# by the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# OpenQuake is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU Affero General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public License
# along with OpenQuake. If not, see <http://www.gnu.org/licenses/>.
import os
import sys
import time
import signal
import getpass
from openquake.baselib import sap, config
from openquake.commonlib import logs
from openquake.engine.utils import confirm
from openquake.server import dbserver
from openquake.commands.purge import purge_one, purge_all

@sap.script
def reset(yes):
    """
    Remove all the datastores and the database of the current user
    """
    ok = yes or confirm('Do you really want to destroy all your data? (y/n) ')
    if not ok:
        return

    status = dbserver.get_status()
    dbpath = os.path.realpath(os.path.expanduser(config.dbserver.file))
    if not os.path.isfile(dbpath):
        sys.exit('%s does not exist' % dbpath)
    elif status == 'running':
        user = getpass.getuser()
        for calc_id in logs.dbcmd('get_calc_ids', user):
            purge_one(calc_id, user, force=True)
        if os.access(dbpath, os.W_OK):   # single user mode
            purge_all(user)  # calculations in oqdata not in the db
            # stop the dbserver first
            pid = logs.dbcmd('getpid')
            os.kill(pid, signal.SIGTERM)
            time.sleep(.5)  # give time to stop
            assert dbserver.get_status() == 'not-running'
            print('dbserver stopped')
            # remove the database
            os.remove(dbpath)
            print('Removed %s' % dbpath)
    else:
        sys.exit('The dbserver is not running')

reset.flg('yes', 'confirmation')
<commit_msg>Make sure the DbServer is running when using 'oq reset'<commit_after># -*- coding: utf-8 -*-
# vim: tabstop=4 shiftwidth=4 softtabstop=4
#
# Copyright (C) 2017-2019 GEM Foundation
#
# OpenQuake is free software: you can redistribute it and/or modify it
# under the terms of the GNU Affero General Public License as published
# by the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# OpenQuake is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU Affero General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public License
# along with OpenQuake. If not, see <http://www.gnu.org/licenses/>.
import os
import sys
import time
import signal
import getpass
from openquake.baselib import sap, config
from openquake.commonlib import logs
from openquake.engine.utils import confirm
from openquake.server import dbserver
from openquake.commands.purge import purge_one, purge_all

@sap.script
def reset(yes):
    """
    Remove all the datastores and the database of the current user
    """
    ok = yes or confirm('Do you really want to destroy all your data? (y/n) ')
    if not ok:
        return

    dbpath = os.path.realpath(os.path.expanduser(config.dbserver.file))
    if not os.path.isfile(dbpath):
        sys.exit('%s does not exist' % dbpath)
    else:
        dbserver.ensure_on()  # start the dbserver in a subprocess
        user = getpass.getuser()
        for calc_id in logs.dbcmd('get_calc_ids', user):
            purge_one(calc_id, user, force=True)
        if os.access(dbpath, os.W_OK):   # single user mode
            purge_all(user)  # calculations in oqdata not in the db
            # stop the dbserver first
            pid = logs.dbcmd('getpid')
            os.kill(pid, signal.SIGTERM)
            time.sleep(.5)  # give time to stop
            assert dbserver.get_status() == 'not-running'
            print('dbserver stopped')
            # remove the database
            os.remove(dbpath)
            print('Removed %s' % dbpath)

reset.flg('yes', 'confirmation')
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
# vim: tabstop=4 shiftwidth=4 softtabstop=4
#
# Copyright (C) 2015-2020 GEM Foundation
#
# OpenQuake is free software: you can redistribute it and/or modify it
# under the terms of the GNU Affero General Public License as published
# by the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# OpenQuake is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU Affero General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public License
# along with OpenQuake. If not, see <http://www.gnu.org/licenses/>.
import numpy
from openquake.baselib import config, datastore
from openquake.commonlib import logs

F32 = numpy.float32

def read(calc_id, username=None):
    """
    :param calc_id: a calculation ID
    :param username: if given, restrict the search to the user's calculations
    :returns: the associated DataStore instance
    """
    if isinstance(calc_id, str) or calc_id < 0 and not username:
        # get the last calculation in the datastore of the current user
        return datastore.read(calc_id)
    job = logs.dbcmd('get_job', calc_id, username)
    if job:
        return datastore.read(job.ds_calc_dir + '.hdf5')
    else:
        # calc_id can be present in the datastore and not in the database:
        # this happens if the calculation was run with `oq run`
        return datastore.read(calc_id)

def max_rel_diff(curve_ref, curve, min_value=0.01):
    """
    Compute the maximum relative difference between two curves. Only values
    greather or equal than the min_value are considered.

    >>> curve_ref = [0.01, 0.02, 0.03, 0.05, 1.0]
    >>> curve = [0.011, 0.021, 0.031, 0.051, 1.0]
    >>> round(max_rel_diff(curve_ref, curve), 2)
    0.1
    """
    assert len(curve_ref) == len(curve), (len(curve_ref), len(curve))
    assert len(curve), 'The curves are empty!'
    max_diff = 0
    for c1, c2 in zip(curve_ref, curve):
        if c1 >= min_value:
            max_diff = max(max_diff, abs(c1 - c2) / c1)
    return max_diff

def max_rel_diff_index(curve_ref, curve, min_value=0.01):
    """
    Compute the maximum relative difference between two sets of curves.
    Only values greather or equal than the min_value are considered.
    Return both the maximum difference and its location (array index).

    >>> curve_refs = [[0.01, 0.02, 0.03, 0.05], [0.01, 0.02, 0.04, 0.06]]
    >>> curves = [[0.011, 0.021, 0.031, 0.051], [0.012, 0.022, 0.032, 0.051]]
    >>> max_rel_diff_index(curve_refs, curves)
    (0.2, 1)
    """
    assert len(curve_ref) == len(curve), (len(curve_ref), len(curve))
    assert len(curve), 'The curves are empty!'
    diffs = [max_rel_diff(c1, c2, min_value)
             for c1, c2 in zip(curve_ref, curve)]
    maxdiff = max(diffs)
    maxindex = diffs.index(maxdiff)
    return maxdiff, maxindex

def rmsep(array_ref, array, min_value=0):
    """
    Root Mean Square Error Percentage for two arrays.

    :param array_ref: reference array
    :param array: another array
    :param min_value: compare only the elements larger than min_value
    :returns: the relative distance between the arrays

    >>> curve_ref = numpy.array([[0.01, 0.02, 0.03, 0.05],
    ... [0.01, 0.02, 0.04, 0.06]])
    >>> curve = numpy.array([[0.011, 0.021, 0.031, 0.051],
    ... [0.012, 0.022, 0.032, 0.051]])
    >>> str(round(rmsep(curve_ref, curve, .01), 5))
    '0.11292'
    """
    bigvalues = array_ref > min_value
    reldiffsquare = (1. - array[bigvalues] / array_ref[bigvalues]) ** 2
    return numpy.sqrt(reldiffsquare.mean())

def log(array, cutoff):
    """
    Compute the logarithm of an array with a cutoff on the small values
    """
    arr = numpy.copy(array)
    arr[arr < cutoff] = cutoff
    return numpy.log(arr)

def closest_to_ref(arrays, ref, cutoff=1E-12):
    """
    :param arrays: a sequence of arrays
    :param ref: the reference array
    :returns: a list of indices ordered by closeness

    This function is used to extract the realization closest to the mean in
    disaggregation. For instance, if there are 2 realizations with indices
    0 and 1, the first hazard curve having values

    >>> c0 = numpy.array([.99, .97, .5, .1])

    and the second hazard curve having values

    >>> c1 = numpy.array([.98, .96, .45, .09])

    with weights 0.6 and 0.4 and mean

    >>> mean = numpy.average([c0, c1], axis=0, weights=[0.6, 0.4])

    then calling ``closest_to_ref`` will returns the indices 0 and 1
    respectively:

    >>> closest_to_ref([c0, c1], mean)
    [0, 1]

    This means that the realization 0 is the closest to the mean, as expected,
    since it has a larger weight. You can check that it is indeed true by
    computing the sum of the quadratic deviations:

    >>> ((c0 - mean)**2).sum()
    0.0004480000000000008
    >>> ((c1 - mean)**2).sum()
    0.0010079999999999985

    If the 2 realizations have equal weights the distance from the mean will be
    the same. In that case both the realizations will be okay; the one that
    will be chosen by closest_to_ref depends on the magic of floating point
    approximation (identical distances will likely be different as
    numpy.float64 numbers) or from the magic of Python list.sort.
    """
    dist = numpy.zeros(len(arrays))
    logref = log(ref, cutoff)
    pairs = []
    for idx, array in enumerate(arrays):
        diff = log(array, cutoff) - logref
        dist = numpy.sqrt((diff * diff).sum())
        pairs.append((dist, idx))
    pairs.sort()
    return [idx for dist, idx in pairs]

def compose_arrays(a1, a2, firstfield='etag'):
    """
    Compose composite arrays by generating an extended datatype containing
    all the fields. The two arrays must have the same length.
    """
    assert len(a1) == len(a2),  (len(a1), len(a2))
    if a1.dtype.names is None and len(a1.shape) == 1:
        # the first array is not composite, but it is one-dimensional
        a1 = numpy.array(a1, numpy.dtype([(firstfield, a1.dtype)]))

    fields1 = [(f, a1.dtype.fields[f][0]) for f in a1.dtype.names]
    if a2.dtype.names is None:  # the second array is not composite
        assert len(a2.shape) == 2, a2.shape
        width = a2.shape[1]
        fields2 = [('value%d' % i, a2.dtype) for i in range(width)]
        composite = numpy.zeros(a1.shape, numpy.dtype(fields1 + fields2))
        for f1 in dict(fields1):
            composite[f1] = a1[f1]
        for i in range(width):
            composite['value%d' % i] = a2[:, i]
        return composite

    fields2 = [(f, a2.dtype.fields[f][0]) for f in a2.dtype.names]
    composite = numpy.zeros(a1.shape, numpy.dtype(fields1 + fields2))
    for f1 in dict(fields1):
        composite[f1] = a1[f1]
    for f2 in dict(fields2):
        composite[f2] = a2[f2]
    return composite

def get_assets(dstore):
    """
    :param dstore: a datastore with keys 'assetcol'
    :returns: an array of records (id, tag1, ..., tagN, lon, lat)
    """
    assetcol = dstore['assetcol']
    tagnames = sorted(assetcol.tagnames)
    tag = {t: getattr(assetcol.tagcol, t) for t in tagnames}
    dtlist = [('id', '<S100')]
    for tagname in tagnames:
        dtlist.append((tagname, '<S100'))
    dtlist.extend([('lon', F32), ('lat', F32)])
    asset_data = []
    for a in assetcol.array:
        tup = tuple(b'%s' % tag[t][a[t]].encode('utf-8') for t in tagnames)
        asset_data.append((a['id'],) + tup + (a['lon'], a['lat']))
    return numpy.array(asset_data, dtlist)

def shared_dir_on():
    """
    :returns: True if a shared_dir has been set in openquake.cfg, else False
    """
    return config.directory.shared_dir
<commit_msg>Updated docstring [skip CI]<commit_after># -*- coding: utf-8 -*-
# vim: tabstop=4 shiftwidth=4 softtabstop=4
#
# Copyright (C) 2015-2020 GEM Foundation
#
# OpenQuake is free software: you can redistribute it and/or modify it
# under the terms of the GNU Affero General Public License as published
# by the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# OpenQuake is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU Affero General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public License
# along with OpenQuake. If not, see <http://www.gnu.org/licenses/>.
import numpy
from openquake.baselib import config, datastore
from openquake.commonlib import logs

F32 = numpy.float32

def read(calc_id, username=None):
    """
    :param calc_id: a calculation ID
    :param username: if given, restrict the search to the user's calculations
    :returns: the associated DataStore instance
    """
    if isinstance(calc_id, str) or calc_id < 0 and not username:
        # get the last calculation in the datastore of the current user
        return datastore.read(calc_id)
    job = logs.dbcmd('get_job', calc_id, username)
    if job:
        return datastore.read(job.ds_calc_dir + '.hdf5')
    else:
        # calc_id can be present in the datastore and not in the database:
        # this happens if the calculation was run with `oq run`
        return datastore.read(calc_id)

def max_rel_diff(curve_ref, curve, min_value=0.01):
    """
    Compute the maximum relative difference between two curves. Only values
    greather or equal than the min_value are considered.

    >>> curve_ref = [0.01, 0.02, 0.03, 0.05, 1.0]
    >>> curve = [0.011, 0.021, 0.031, 0.051, 1.0]
    >>> round(max_rel_diff(curve_ref, curve), 2)
    0.1
    """
    assert len(curve_ref) == len(curve), (len(curve_ref), len(curve))
    assert len(curve), 'The curves are empty!'
    max_diff = 0
    for c1, c2 in zip(curve_ref, curve):
        if c1 >= min_value:
            max_diff = max(max_diff, abs(c1 - c2) / c1)
    return max_diff

def max_rel_diff_index(curve_ref, curve, min_value=0.01):
    """
    Compute the maximum relative difference between two sets of curves.
    Only values greather or equal than the min_value are considered.
    Return both the maximum difference and its location (array index).

    >>> curve_refs = [[0.01, 0.02, 0.03, 0.05], [0.01, 0.02, 0.04, 0.06]]
    >>> curves = [[0.011, 0.021, 0.031, 0.051], [0.012, 0.022, 0.032, 0.051]]
    >>> max_rel_diff_index(curve_refs, curves)
    (0.2, 1)
    """
    assert len(curve_ref) == len(curve), (len(curve_ref), len(curve))
    assert len(curve), 'The curves are empty!'
    diffs = [max_rel_diff(c1, c2, min_value)
             for c1, c2 in zip(curve_ref, curve)]
    maxdiff = max(diffs)
    maxindex = diffs.index(maxdiff)
    return maxdiff, maxindex

def rmsep(array_ref, array, min_value=0):
    """
    Root Mean Square Error Percentage for two arrays.

    :param array_ref: reference array
    :param array: another array
    :param min_value: compare only the elements larger than min_value
    :returns: the relative distance between the arrays

    >>> curve_ref = numpy.array([[0.01, 0.02, 0.03, 0.05],
    ... [0.01, 0.02, 0.04, 0.06]])
    >>> curve = numpy.array([[0.011, 0.021, 0.031, 0.051],
    ... [0.012, 0.022, 0.032, 0.051]])
    >>> str(round(rmsep(curve_ref, curve, .01), 5))
    '0.11292'
    """
    bigvalues = array_ref > min_value
    reldiffsquare = (1. - array[bigvalues] / array_ref[bigvalues]) ** 2
    return numpy.sqrt(reldiffsquare.mean())

def log(array, cutoff):
    """
    Compute the logarithm of an array with a cutoff on the small values
    """
    arr = numpy.copy(array)
    arr[arr < cutoff] = cutoff
    return numpy.log(arr)

def closest_to_ref(arrays, ref, cutoff=1E-12):
    """
    :param arrays: a sequence of arrays
    :param ref: the reference array
    :returns: a list of indices ordered by closeness

    This function is used to extract the realization closest to the mean in
    disaggregation. For instance, if there are 2 realizations with indices
    0 and 1, the first hazard curve having values

    >>> c0 = numpy.array([.99, .97, .5, .1])

    and the second hazard curve having values

    >>> c1 = numpy.array([.98, .96, .45, .09])

    with weights 0.6 and 0.4 and mean

    >>> mean = numpy.average([c0, c1], axis=0, weights=[0.6, 0.4])

    then calling ``closest_to_ref`` will returns the indices 0 and 1
    respectively:

    >>> closest_to_ref([c0, c1], mean)
    [0, 1]

    This means that the realization 0 is the closest to the mean, as expected,
    since it has a larger weight. You can check that it is indeed true by
    computing the sum of the quadratic deviations:

    >>> ((c0 - mean)**2).sum()
    0.0004480000000000008
    >>> ((c1 - mean)**2).sum()
    0.0010079999999999985

    If the 2 realizations have equal weights the distance from the mean will be
    the same. In that case both the realizations will be okay; the one that
    will be chosen by ``closest_to_ref`` depends on the magic of floating point
    approximation (theoretically identical distances will likely be different
    as numpy.float64 numbers) or on the magic of Python ``list.sort``.
    """
    dist = numpy.zeros(len(arrays))
    logref = log(ref, cutoff)
    pairs = []
    for idx, array in enumerate(arrays):
        diff = log(array, cutoff) - logref
        dist = numpy.sqrt((diff * diff).sum())
        pairs.append((dist, idx))
    pairs.sort()
    return [idx for dist, idx in pairs]

def compose_arrays(a1, a2, firstfield='etag'):
    """
    Compose composite arrays by generating an extended datatype containing
    all the fields. The two arrays must have the same length.
    """
    assert len(a1) == len(a2),  (len(a1), len(a2))
    if a1.dtype.names is None and len(a1.shape) == 1:
        # the first array is not composite, but it is one-dimensional
        a1 = numpy.array(a1, numpy.dtype([(firstfield, a1.dtype)]))

    fields1 = [(f, a1.dtype.fields[f][0]) for f in a1.dtype.names]
    if a2.dtype.names is None:  # the second array is not composite
        assert len(a2.shape) == 2, a2.shape
        width = a2.shape[1]
        fields2 = [('value%d' % i, a2.dtype) for i in range(width)]
        composite = numpy.zeros(a1.shape, numpy.dtype(fields1 + fields2))
        for f1 in dict(fields1):
            composite[f1] = a1[f1]
        for i in range(width):
            composite['value%d' % i] = a2[:, i]
        return composite

    fields2 = [(f, a2.dtype.fields[f][0]) for f in a2.dtype.names]
    composite = numpy.zeros(a1.shape, numpy.dtype(fields1 + fields2))
    for f1 in dict(fields1):
        composite[f1] = a1[f1]
    for f2 in dict(fields2):
        composite[f2] = a2[f2]
    return composite

def get_assets(dstore):
    """
    :param dstore: a datastore with keys 'assetcol'
    :returns: an array of records (id, tag1, ..., tagN, lon, lat)
    """
    assetcol = dstore['assetcol']
    tagnames = sorted(assetcol.tagnames)
    tag = {t: getattr(assetcol.tagcol, t) for t in tagnames}
    dtlist = [('id', '<S100')]
    for tagname in tagnames:
        dtlist.append((tagname, '<S100'))
    dtlist.extend([('lon', F32), ('lat', F32)])
    asset_data = []
    for a in assetcol.array:
        tup = tuple(b'%s' % tag[t][a[t]].encode('utf-8') for t in tagnames)
        asset_data.append((a['id'],) + tup + (a['lon'], a['lat']))
    return numpy.array(asset_data, dtlist)

def shared_dir_on():
    """
    :returns: True if a shared_dir has been set in openquake.cfg, else False
    """
    return config.directory.shared_dir
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
# vim: tabstop=4 shiftwidth=4 softtabstop=4
#
# Copyright (C) 2015-2019 GEM Foundation
#
# OpenQuake is free software: you can redistribute it and/or modify it
# under the terms of the GNU Affero General Public License as published
# by the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# OpenQuake is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU Affero General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public License
# along with OpenQuake. If not, see <http://www.gnu.org/licenses/>.
import numpy
from openquake.baselib import config, datastore
from openquake.commonlib import logs

F32 = numpy.float32

def read(calc_id, username=None):
    """
    :param calc_id: a calculation ID
    :param username: if given, restrict the search to the user's calculations
    :returns: the associated DataStore instance
    """
    if isinstance(calc_id, str) or calc_id < 0 and not username:
        # get the last calculation in the datastore of the current user
        return datastore.read(calc_id)
    job = logs.dbcmd('get_job', calc_id, username)
    if job:
        return datastore.read(job.ds_calc_dir + '.hdf5')
    else:
        # calc_id can be present in the datastore and not in the database:
        # this happens if the calculation was run with `oq run`
        return datastore.read(calc_id)

def max_rel_diff(curve_ref, curve, min_value=0.01):
    """
    Compute the maximum relative difference between two curves. Only values
    greather or equal than the min_value are considered.

    >>> curve_ref = [0.01, 0.02, 0.03, 0.05, 1.0]
    >>> curve = [0.011, 0.021, 0.031, 0.051, 1.0]
    >>> round(max_rel_diff(curve_ref, curve), 2)
    0.1
    """
    assert len(curve_ref) == len(curve), (len(curve_ref), len(curve))
    assert len(curve), 'The curves are empty!'
    max_diff = 0
    for c1, c2 in zip(curve_ref, curve):
        if c1 >= min_value:
            max_diff = max(max_diff, abs(c1 - c2) / c1)
    return max_diff

def max_rel_diff_index(curve_ref, curve, min_value=0.01):
    """
    Compute the maximum relative difference between two sets of curves.
    Only values greather or equal than the min_value are considered.
    Return both the maximum difference and its location (array index).

    >>> curve_refs = [[0.01, 0.02, 0.03, 0.05], [0.01, 0.02, 0.04, 0.06]]
    >>> curves = [[0.011, 0.021, 0.031, 0.051], [0.012, 0.022, 0.032, 0.051]]
    >>> max_rel_diff_index(curve_refs, curves)
    (0.2, 1)
    """
    assert len(curve_ref) == len(curve), (len(curve_ref), len(curve))
    assert len(curve), 'The curves are empty!'
    diffs = [max_rel_diff(c1, c2, min_value)
             for c1, c2 in zip(curve_ref, curve)]
    maxdiff = max(diffs)
    maxindex = diffs.index(maxdiff)
    return maxdiff, maxindex

def rmsep(array_ref, array, min_value=0):
    """
    Root Mean Square Error Percentage for two arrays.

    :param array_ref: reference array
    :param array: another array
    :param min_value: compare only the elements larger than min_value
    :returns: the relative distance between the arrays

    >>> curve_ref = numpy.array([[0.01, 0.02, 0.03, 0.05],
    ... [0.01, 0.02, 0.04, 0.06]])
    >>> curve = numpy.array([[0.011, 0.021, 0.031, 0.051],
    ... [0.012, 0.022, 0.032, 0.051]])
    >>> str(round(rmsep(curve_ref, curve, .01), 5))
    '0.11292'
    """
    bigvalues = array_ref > min_value
    reldiffsquare = (1. - array[bigvalues] / array_ref[bigvalues]) ** 2
    return numpy.sqrt(reldiffsquare.mean())

def log(array, cutoff):
    """
    Compute the logarithm of an array with a cutoff on the small values
    """
    arr = numpy.copy(array)
    arr[arr < cutoff] = cutoff
    return numpy.log(arr)

def closest_to_ref(arrays, ref, nclosest=1, cutoff=1E-12):
    """
    :param arrays: a sequence of arrays
    :param ref: the reference array
    :param nclosest: the number of closest arrays to return (default 1)
    :returns: a list of nclosest elements [(dist, idx), ...]
    """
    dist = numpy.zeros(len(arrays))
    logref = log(ref, cutoff)
    pairs = []
    for idx, array in enumerate(arrays):
        diff = log(array, cutoff) - logref
        dist = numpy.sqrt((diff * diff).sum())
        pairs.append((dist, idx, array))
    return sorted(pairs)

def compose_arrays(a1, a2, firstfield='etag'):
    """
    Compose composite arrays by generating an extended datatype containing
    all the fields. The two arrays must have the same length.
    """
    assert len(a1) == len(a2),  (len(a1), len(a2))
    if a1.dtype.names is None and len(a1.shape) == 1:
        # the first array is not composite, but it is one-dimensional
        a1 = numpy.array(a1, numpy.dtype([(firstfield, a1.dtype)]))

    fields1 = [(f, a1.dtype.fields[f][0]) for f in a1.dtype.names]
    if a2.dtype.names is None:  # the second array is not composite
        assert len(a2.shape) == 2, a2.shape
        width = a2.shape[1]
        fields2 = [('value%d' % i, a2.dtype) for i in range(width)]
        composite = numpy.zeros(a1.shape, numpy.dtype(fields1 + fields2))
        for f1 in dict(fields1):
            composite[f1] = a1[f1]
        for i in range(width):
            composite['value%d' % i] = a2[:, i]
        return composite

    fields2 = [(f, a2.dtype.fields[f][0]) for f in a2.dtype.names]
    composite = numpy.zeros(a1.shape, numpy.dtype(fields1 + fields2))
    for f1 in dict(fields1):
        composite[f1] = a1[f1]
    for f2 in dict(fields2):
        composite[f2] = a2[f2]
    return composite

def get_assets(dstore):
    """
    :param dstore: a datastore with keys 'assetcol'
    :returns: an array of records (asset_ref, tag1, ..., tagN, lon, lat)
    """
    assetcol = dstore['assetcol']
    tagnames = sorted(assetcol.tagnames)
    tag = {t: getattr(assetcol.tagcol, t) for t in tagnames}
    dtlist = []
    for tagname in tagnames:
        dtlist.append((tagname, (bytes, 100)))
    dtlist.extend([('lon', F32), ('lat', F32)])
    asset_data = []
    for a in assetcol.array:
        tup = tuple(b'"%s"' % tag[t][a[t]].encode('utf-8') for t in tagnames)
        asset_data.append(tup + (a['lon'], a['lat']))
    return numpy.array(asset_data, dtlist)

def shared_dir_on():
    """
    :returns: True if a shared_dir has been set in openquake.cfg, else False
    """
    return config.directory.shared_dir
<commit_msg>[skip CI]<commit_after># -*- coding: utf-8 -*-
# vim: tabstop=4 shiftwidth=4 softtabstop=4
#
# Copyright (C) 2015-2019 GEM Foundation
#
# OpenQuake is free software: you can redistribute it and/or modify it
# under the terms of the GNU Affero General Public License as published
# by the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# OpenQuake is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU Affero General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public License
# along with OpenQuake. If not, see <http://www.gnu.org/licenses/>.
import numpy
from openquake.baselib import config, datastore
from openquake.commonlib import logs

F32 = numpy.float32

def read(calc_id, username=None):
    """
    :param calc_id: a calculation ID
    :param username: if given, restrict the search to the user's calculations
    :returns: the associated DataStore instance
    """
    if isinstance(calc_id, str) or calc_id < 0 and not username:
        # get the last calculation in the datastore of the current user
        return datastore.read(calc_id)
    job = logs.dbcmd('get_job', calc_id, username)
    if job:
        return datastore.read(job.ds_calc_dir + '.hdf5')
    else:
        # calc_id can be present in the datastore and not in the database:
        # this happens if the calculation was run with `oq run`
        return datastore.read(calc_id)

def max_rel_diff(curve_ref, curve, min_value=0.01):
    """
    Compute the maximum relative difference between two curves. Only values
    greather or equal than the min_value are considered.

    >>> curve_ref = [0.01, 0.02, 0.03, 0.05, 1.0]
    >>> curve = [0.011, 0.021, 0.031, 0.051, 1.0]
    >>> round(max_rel_diff(curve_ref, curve), 2)
    0.1
    """
    assert len(curve_ref) == len(curve), (len(curve_ref), len(curve))
    assert len(curve), 'The curves are empty!'
    max_diff = 0
    for c1, c2 in zip(curve_ref, curve):
        if c1 >= min_value:
            max_diff = max(max_diff, abs(c1 - c2) / c1)
    return max_diff

def max_rel_diff_index(curve_ref, curve, min_value=0.01):
    """
    Compute the maximum relative difference between two sets of curves.
    Only values greather or equal than the min_value are considered.
    Return both the maximum difference and its location (array index).

    >>> curve_refs = [[0.01, 0.02, 0.03, 0.05], [0.01, 0.02, 0.04, 0.06]]
    >>> curves = [[0.011, 0.021, 0.031, 0.051], [0.012, 0.022, 0.032, 0.051]]
    >>> max_rel_diff_index(curve_refs, curves)
    (0.2, 1)
    """
    assert len(curve_ref) == len(curve), (len(curve_ref), len(curve))
    assert len(curve), 'The curves are empty!'
    diffs = [max_rel_diff(c1, c2, min_value)
             for c1, c2 in zip(curve_ref, curve)]
    maxdiff = max(diffs)
    maxindex = diffs.index(maxdiff)
    return maxdiff, maxindex

def rmsep(array_ref, array, min_value=0):
    """
    Root Mean Square Error Percentage for two arrays.

    :param array_ref: reference array
    :param array: another array
    :param min_value: compare only the elements larger than min_value
    :returns: the relative distance between the arrays

    >>> curve_ref = numpy.array([[0.01, 0.02, 0.03, 0.05],
    ... [0.01, 0.02, 0.04, 0.06]])
    >>> curve = numpy.array([[0.011, 0.021, 0.031, 0.051],
    ... [0.012, 0.022, 0.032, 0.051]])
    >>> str(round(rmsep(curve_ref, curve, .01), 5))
    '0.11292'
    """
    bigvalues = array_ref > min_value
    reldiffsquare = (1. - array[bigvalues] / array_ref[bigvalues]) ** 2
    return numpy.sqrt(reldiffsquare.mean())

def log(array, cutoff):
    """
    Compute the logarithm of an array with a cutoff on the small values
    """
    arr = numpy.copy(array)
    arr[arr < cutoff] = cutoff
    return numpy.log(arr)

def closest_to_ref(arrays, ref, nclosest=1, cutoff=1E-12):
    """
    :param arrays: a sequence of arrays
    :param ref: the reference array
    :param nclosest: the number of closest arrays to return (default 1)
    :returns: a list of nclosest elements [(dist, idx), ...]
    """
    dist = numpy.zeros(len(arrays))
    logref = log(ref, cutoff)
    pairs = []
    for idx, array in enumerate(arrays):
        diff = log(array, cutoff) - logref
        dist = numpy.sqrt((diff * diff).sum())
        pairs.append((dist, idx))
    return sorted(pairs)

def compose_arrays(a1, a2, firstfield='etag'):
    """
    Compose composite arrays by generating an extended datatype containing
    all the fields. The two arrays must have the same length.
    """
    assert len(a1) == len(a2),  (len(a1), len(a2))
    if a1.dtype.names is None and len(a1.shape) == 1:
        # the first array is not composite, but it is one-dimensional
        a1 = numpy.array(a1, numpy.dtype([(firstfield, a1.dtype)]))

    fields1 = [(f, a1.dtype.fields[f][0]) for f in a1.dtype.names]
    if a2.dtype.names is None:  # the second array is not composite
        assert len(a2.shape) == 2, a2.shape
        width = a2.shape[1]
        fields2 = [('value%d' % i, a2.dtype) for i in range(width)]
        composite = numpy.zeros(a1.shape, numpy.dtype(fields1 + fields2))
        for f1 in dict(fields1):
            composite[f1] = a1[f1]
        for i in range(width):
            composite['value%d' % i] = a2[:, i]
        return composite

    fields2 = [(f, a2.dtype.fields[f][0]) for f in a2.dtype.names]
    composite = numpy.zeros(a1.shape, numpy.dtype(fields1 + fields2))
    for f1 in dict(fields1):
        composite[f1] = a1[f1]
    for f2 in dict(fields2):
        composite[f2] = a2[f2]
    return composite

def get_assets(dstore):
    """
    :param dstore: a datastore with keys 'assetcol'
    :returns: an array of records (asset_ref, tag1, ..., tagN, lon, lat)
    """
    assetcol = dstore['assetcol']
    tagnames = sorted(assetcol.tagnames)
    tag = {t: getattr(assetcol.tagcol, t) for t in tagnames}
    dtlist = []
    for tagname in tagnames:
        dtlist.append((tagname, (bytes, 100)))
    dtlist.extend([('lon', F32), ('lat', F32)])
    asset_data = []
    for a in assetcol.array:
        tup = tuple(b'"%s"' % tag[t][a[t]].encode('utf-8') for t in tagnames)
        asset_data.append(tup + (a['lon'], a['lat']))
    return numpy.array(asset_data, dtlist)

def shared_dir_on():
    """
    :returns: True if a shared_dir has been set in openquake.cfg, else False
    """
    return config.directory.shared_dir
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
"""
Wrapper around the OpenSHA-lite java library.

"""

import math
import os
import random
import numpy

from openquake import java
from openquake import kvs
from openquake import logs
from openquake import settings
from openquake import shapes

from openquake.hazard import job
from openquake.hazard import tasks
from openquake.job.mixins import Mixin
from openquake.output import geotiff

LOG = logs.LOG

def preload(fn): # pylint: disable=E0213
    """A decorator for preload steps that must run on the Jobber node"""
    def preloader(self, *args, **kwargs):
        """Validate job"""
        self.cache = java.jclass("KVS")(
                settings.KVS_HOST,
                settings.KVS_PORT)
        self.calc = java.jclass("CommandLineCalculator")(
                self.cache, self.key)
        return fn(self, *args, **kwargs) # pylint: disable=E1102
    return preloader

class BasePSHAMixin(Mixin):
    """Contains common functionality for PSHA Mixins."""

    def store_source_model(self, seed):
        """Generates an Earthquake Rupture Forecast, using the source zones and
        logic trees specified in the job config file. Note that this has to be
        done currently using the file itself, since it has nested references to
        other files."""

        LOG.info("Storing source model from job config")
        key = kvs.generate_product_key(self.id, kvs.tokens.SOURCE_MODEL_TOKEN)
        print "source model key is", key
        self.calc.sampleAndSaveERFTree(self.cache, key, seed)

    def store_gmpe_map(self, seed):
        """Generates a hash of tectonic regions and GMPEs, using the logic tree
        specified in the job config file."""
        key = kvs.generate_product_key(self.id, kvs.tokens.GMPE_TOKEN)
        print "GMPE map key is", key
        self.calc.sampleAndSaveGMPETree(self.cache, key, seed)

    def generate_erf(self):
        """Generate the Earthquake Rupture Forecast from the currently stored
        source model logic tree."""
        key = kvs.generate_product_key(self.id, kvs.tokens.SOURCE_MODEL_TOKEN)
        sources = java.jclass("JsonSerializer").getSourceListFromCache(
                    self.cache, key)
        erf = java.jclass("GEM1ERF")(sources)
        self.calc.setGEM1ERFParams(erf)
        return erf

    def set_gmpe_params(self, gmpe_map):
        """Push parameters from configuration file into the GMPE objects"""
        jpype = java.jvm()
        gmpe_lt_data = self.calc.createGmpeLogicTreeData()
        for tect_region in gmpe_map.keySet():
            gmpe = gmpe_map.get(tect_region)
            gmpe_lt_data.setGmpeParams(self.params['COMPONENT'],
                self.params['INTENSITY_MEASURE_TYPE'],
                jpype.JDouble(float(self.params['PERIOD'])),
                jpype.JDouble(float(self.params['DAMPING'])),
                self.params['GMPE_TRUNCATION_TYPE'],
                jpype.JDouble(float(self.params['TRUNCATION_LEVEL'])),
                self.params['STANDARD_DEVIATION_TYPE'],
                jpype.JDouble(float(self.params['REFERENCE_VS30_VALUE'])),
                jpype.JObject(gmpe, java.jclass("AttenuationRelationship")))
            gmpe_map.put(tect_region, gmpe)

    def generate_gmpe_map(self):
        """Generate the GMPE map from the stored GMPE logic tree."""
        key = kvs.generate_product_key(self.id, kvs.tokens.GMPE_TOKEN)
        gmpe_map = java.jclass("JsonSerializer").getGmpeMapFromCache(
                                                    self.cache,key)
        self.set_gmpe_params(gmpe_map)
        return gmpe_map

    def get_iml_list(self):
        """Build the appropriate Arbitrary Discretized Func from the IMLs,
        based on the IMT"""
        iml_vals = {'PGA' : numpy.log,  # pylint: disable=E1101
                    'MMI' : lambda iml: iml,
                    'PGV' : numpy.log, # pylint: disable=E1101
                    'PGD' : numpy.log, # pylint: disable=E1101
                    'SA' : numpy.log,  # pylint: disable=E1101
                     }

        iml_list = java.jclass("ArrayList")()
        for val in self.params['INTENSITY_MEASURE_LEVELS'].split(","):
            iml_list.add(
                iml_vals[self.params['INTENSITY_MEASURE_TYPE']](
                float(val)))
        return iml_list

    def site_list_generator(self):
        """Will subset and yield portions of the region, depending on the
        the computation mode."""
        verts = [float(x) for x in self.params['REGION_VERTEX'].split(",")]
        coords = zip(verts[1::2], verts[::2])
        region = shapes.Region.from_coordinates(coords)
        region.cell_size = float(self.params['REGION_GRID_SPACING'])
        yield [site for site in region]

    def parameterize_sites(self, site_list):
        """Convert python Sites to Java Sites, and add default parameters."""
        # TODO(JMC): There's Java code for this already, sets each site to have
        # the same default parameters

        jpype = java.jvm()
        jsite_list = java.jclass("ArrayList")()
        for x in site_list:
            site = x.to_java()

            vs30 = java.jclass("DoubleParameter")(jpype.JString("Vs30"))
            vs30.setValue(float(self.params['REFERENCE_VS30_VALUE']))
            depth25 = java.jclass("DoubleParameter")("Depth 2.5 km/sec")
            depth25.setValue(float(
                    self.params['REFERENCE_DEPTH_TO_2PT5KM_PER_SEC_PARAM']))
            sadigh = java.jclass("StringParameter")("Sadigh Site Type")
            sadigh.setValue(self.params['SADIGH_SITE_TYPE'])
            site.addParameter(vs30)
            site.addParameter(depth25)
            site.addParameter(sadigh)
            jsite_list.add(site)
        return jsite_list

class ClassicalMixin(BasePSHAMixin):
    """Classical PSHA method for performing Hazard calculations.

    Implements the JobMixin, which has a primary entry point of execute().
    Execute is responsible for dispatching celery tasks.

    Note that this Mixin, during execution, will always be an instance of the
    Job class, and thus has access to the self.params dict, full of config
    params loaded from the Job configuration file."""

    @preload
    def execute(self):

        results = []

        source_model_generator = random.Random()
        source_model_generator.seed(
                self.params.get('SOURCE_MODEL_LT_RANDOM_SEED', None))

        gmpe_generator = random.Random()
        gmpe_generator.seed(self.params.get('GMPE_LT_RANDOM_SEED', None))

        realizations = int(self.params['NUMBER_OF_LOGIC_TREE_SAMPLES'])

        LOG.info('Going to run classical PSHA hazard for %s realizations'
                % realizations)

        for realization in range(0, realizations):
            LOG.info('Calculatiing hazard curves for realization %s'
                     % realization)
            pending_tasks = []
            self.store_source_model(source_model_generator.getrandbits(32))
            self.store_gmpe_map(source_model_generator.getrandbits(32))
            for site_list in self.site_list_generator():
                pending_tasks.append(
                    tasks.compute_hazard_curve.delay(self.id,
                            site_list, realization))

            for task in pending_tasks:
                task.wait()
                if task.status != 'SUCCESS':
                    raise Exception(task.result)
                results.extend(task.result)
        return results

    @preload
    def compute_hazard_curve(self, site_list, realization):
        """ Compute hazard curves, write them to KVS as JSON,
        and return a list of the KVS keys for each curve. """
        jsite_list = self.parameterize_sites(site_list)
        hazard_curves = java.jclass("HazardCalculator").getHazardCurvesAsJson(
            jsite_list,
            self.generate_erf(),
            self.generate_gmpe_map(),
            self.get_iml_list(),
            float(self.params['MAXIMUM_DISTANCE']))

        # write the curves to the KVS and return a list of the keys
        kvs_client = kvs.get_client()
        curve_keys = []
        for i in xrange(0, len(hazard_curves)):
            curve = hazard_curves[i]
            site = site_list[i]
            lon = site.longitude
            lat = site.latitude
            curve_key = kvs.tokens.hazard_curve_key(self.id,
                                                    realization,
                                                    lon,
                                                    lat)
            kvs_client.set(curve_key, curve)
            curve_keys.append(curve_key)
        return curve_keys

class EventBasedMixin(BasePSHAMixin): # pylint: disable=W0232
    """Probabilistic Event Based method for performing Hazard calculations.

    Implements the JobMixin, which has a primary entry point of execute().
    Execute is responsible for dispatching celery tasks.

    Note that this Mixin, during execution, will always be an instance of the
    Job class, and thus has access to the self.params dict, full of config
    params loaded from the Job configuration file."""

    @preload
    def execute(self):
        """Main hazard processing block.

        Loops through various random realizations, spawning tasks to compute
        GMFs."""
        results = []

        source_model_generator = random.Random()
        source_model_generator.seed(
                self.params.get('SOURCE_MODEL_LT_RANDOM_SEED', None))

        gmpe_generator = random.Random()
        gmpe_generator.seed(self.params.get('GMPE_LT_RANDOM_SEED', None))

        gmf_generator = random.Random()
        gmf_generator.seed(self.params.get('GMF_RANDOM_SEED', None))

        histories = int(self.params['NUMBER_OF_SEISMICITY_HISTORIES'])
        realizations = int(self.params['NUMBER_OF_LOGIC_TREE_SAMPLES'])
        LOG.info("Going to run hazard for %s histories of %s realizations each."
                % (histories, realizations))

        for i in range(0, histories):
            pending_tasks = []
            for j in range(0, realizations):
                self.store_source_model(source_model_generator.getrandbits(32))
                self.store_gmpe_map(gmpe_generator.getrandbits(32))
                for site_list in self.site_list_generator():
                    stochastic_set_id = "%s!%s" % (i, j)
                    # pylint: disable=E1101
                    pending_tasks.append(
                        tasks.compute_ground_motion_fields.delay(
                            self.id,
                            site_list,
                            stochastic_set_id, gmf_generator.getrandbits(32)))

            for task in pending_tasks:
                task.wait()
                if task.status != 'SUCCESS':
                    raise Exception(task.result)

            # if self.params['OUTPUT_GMF_FILES']
            for j in range(0, realizations):
                stochastic_set_id = "%s!%s" % (i, j)
                stochastic_set_key = kvs.generate_product_key(
                    self.id, kvs.tokens.STOCHASTIC_SET_TOKEN, stochastic_set_id)
                print "Writing output for ses %s" % stochastic_set_key
                ses = kvs.get_value_json_decoded(stochastic_set_key)
                if ses:
                    results.extend(self.write_gmf_files(ses))
        return results

    def write_gmf_files(self, ses):
        """Generate a GeoTiff file for each GMF."""
        image_grid = self.region.grid
        iml_list = [float(param)
                    for param
                    in self.params['INTENSITY_MEASURE_LEVELS'].split(",")]

        LOG.debug("Generating GMF image, grid is %s col by %s rows" % (
                image_grid.columns, image_grid.rows))
        LOG.debug("IML: %s" % (iml_list))
        files = []
        for event_set in ses:
            for rupture in ses[event_set]:

                # NOTE(fab): we have to explicitly convert the JSON-decoded
                # tokens from Unicode to string, otherwise the path will not
                # be accepted by the GeoTiffFile constructor
                path = os.path.join(self.base_path, self['OUTPUT_DIR'],
                        "gmf-%s-%s.tiff" % (str(event_set.replace("!", "_")),
                                            str(rupture.replace("!", "_"))))
                gwriter = geotiff.GMFGeoTiffFile(path, image_grid,
                    init_value=0.0, normalize=True, iml_list=iml_list,
                    discrete=True)
                for site_key in ses[event_set][rupture]:
                    site = ses[event_set][rupture][site_key]
                    site_obj = shapes.Site(site['lon'], site['lat'])
                    point = image_grid.point_at(site_obj)
                    gwriter.write((point.row, point.column),
                        math.exp(float(site['mag'])))

                gwriter.close()
                files.append(path)
                files.append(gwriter.html_path)
        return files

    @preload
    def compute_ground_motion_fields(self, site_list, stochastic_set_id, seed):
        """Ground motion field calculation, runs on the workers."""
        jpype = java.jvm()

        jsite_list = self.parameterize_sites(site_list)
        key = kvs.generate_product_key(
                    self.id, kvs.tokens.STOCHASTIC_SET_TOKEN, stochastic_set_id)
        gmc = self.params['GROUND_MOTION_CORRELATION']
        correlate = (gmc == "true" and True or False)
        java.jclass("HazardCalculator").generateAndSaveGMFs(
                self.cache, key, stochastic_set_id, jsite_list,
                 self.generate_erf(),
                self.generate_gmpe_map(),
                java.jclass("Random")(seed),
                jpype.JBoolean(correlate))

def gmf_id(history_idx, realization_idx, rupture_idx):
    """ Return a GMF id suitable for use as a KVS key """
    return "%s!%s!%s" % (history_idx, realization_idx, rupture_idx)

job.HazJobMixin.register("Event Based", EventBasedMixin, order=0)
job.HazJobMixin.register("Classical", ClassicalMixin, order=1)
<commit_msg>Added some code to count the total number of sites<commit_after># -*- coding: utf-8 -*-
"""
Wrapper around the OpenSHA-lite java library.

"""

import math
import os
import random
import numpy

from openquake import java
from openquake import kvs
from openquake import logs
from openquake import settings
from openquake import shapes

from openquake.hazard import job
from openquake.hazard import tasks
from openquake.job.mixins import Mixin
from openquake.output import geotiff

LOG = logs.LOG

def preload(fn): # pylint: disable=E0213
    """A decorator for preload steps that must run on the Jobber node"""
    def preloader(self, *args, **kwargs):
        """Validate job"""
        self.cache = java.jclass("KVS")(
                settings.KVS_HOST,
                settings.KVS_PORT)
        self.calc = java.jclass("CommandLineCalculator")(
                self.cache, self.key)
        return fn(self, *args, **kwargs) # pylint: disable=E1102
    return preloader

class BasePSHAMixin(Mixin):
    """Contains common functionality for PSHA Mixins."""

    def store_source_model(self, seed):
        """Generates an Earthquake Rupture Forecast, using the source zones and
        logic trees specified in the job config file. Note that this has to be
        done currently using the file itself, since it has nested references to
        other files."""

        LOG.info("Storing source model from job config")
        key = kvs.generate_product_key(self.id, kvs.tokens.SOURCE_MODEL_TOKEN)
        print "source model key is", key
        self.calc.sampleAndSaveERFTree(self.cache, key, seed)

    def store_gmpe_map(self, seed):
        """Generates a hash of tectonic regions and GMPEs, using the logic tree
        specified in the job config file."""
        key = kvs.generate_product_key(self.id, kvs.tokens.GMPE_TOKEN)
        print "GMPE map key is", key
        self.calc.sampleAndSaveGMPETree(self.cache, key, seed)

    def generate_erf(self):
        """Generate the Earthquake Rupture Forecast from the currently stored
        source model logic tree."""
        key = kvs.generate_product_key(self.id, kvs.tokens.SOURCE_MODEL_TOKEN)
        sources = java.jclass("JsonSerializer").getSourceListFromCache(
                    self.cache, key)
        erf = java.jclass("GEM1ERF")(sources)
        self.calc.setGEM1ERFParams(erf)
        return erf

    def set_gmpe_params(self, gmpe_map):
        """Push parameters from configuration file into the GMPE objects"""
        jpype = java.jvm()
        gmpe_lt_data = self.calc.createGmpeLogicTreeData()
        for tect_region in gmpe_map.keySet():
            gmpe = gmpe_map.get(tect_region)
            gmpe_lt_data.setGmpeParams(self.params['COMPONENT'],
                self.params['INTENSITY_MEASURE_TYPE'],
                jpype.JDouble(float(self.params['PERIOD'])),
                jpype.JDouble(float(self.params['DAMPING'])),
                self.params['GMPE_TRUNCATION_TYPE'],
                jpype.JDouble(float(self.params['TRUNCATION_LEVEL'])),
                self.params['STANDARD_DEVIATION_TYPE'],
                jpype.JDouble(float(self.params['REFERENCE_VS30_VALUE'])),
                jpype.JObject(gmpe, java.jclass("AttenuationRelationship")))
            gmpe_map.put(tect_region, gmpe)

    def generate_gmpe_map(self):
        """Generate the GMPE map from the stored GMPE logic tree."""
        key = kvs.generate_product_key(self.id, kvs.tokens.GMPE_TOKEN)
        gmpe_map = java.jclass("JsonSerializer").getGmpeMapFromCache(
                                                    self.cache,key)
        self.set_gmpe_params(gmpe_map)
        return gmpe_map

    def get_iml_list(self):
        """Build the appropriate Arbitrary Discretized Func from the IMLs,
        based on the IMT"""
        iml_vals = {'PGA' : numpy.log,  # pylint: disable=E1101
                    'MMI' : lambda iml: iml,
                    'PGV' : numpy.log, # pylint: disable=E1101
                    'PGD' : numpy.log, # pylint: disable=E1101
                    'SA' : numpy.log,  # pylint: disable=E1101
                     }

        iml_list = java.jclass("ArrayList")()
        for val in self.params['INTENSITY_MEASURE_LEVELS'].split(","):
            iml_list.add(
                iml_vals[self.params['INTENSITY_MEASURE_TYPE']](
                float(val)))
        return iml_list

    def site_list_generator(self):
        """Will subset and yield portions of the region, depending on the
        the computation mode."""
        verts = [float(x) for x in self.params['REGION_VERTEX'].split(",")]
        coords = zip(verts[1::2], verts[::2])
        region = shapes.Region.from_coordinates(coords)
        region.cell_size = float(self.params['REGION_GRID_SPACING'])
        yield [site for site in region]

    def parameterize_sites(self, site_list):
        """Convert python Sites to Java Sites, and add default parameters."""
        # TODO(JMC): There's Java code for this already, sets each site to have
        # the same default parameters

        jpype = java.jvm()
        jsite_list = java.jclass("ArrayList")()
        for x in site_list:
            site = x.to_java()

            vs30 = java.jclass("DoubleParameter")(jpype.JString("Vs30"))
            vs30.setValue(float(self.params['REFERENCE_VS30_VALUE']))
            depth25 = java.jclass("DoubleParameter")("Depth 2.5 km/sec")
            depth25.setValue(float(
                    self.params['REFERENCE_DEPTH_TO_2PT5KM_PER_SEC_PARAM']))
            sadigh = java.jclass("StringParameter")("Sadigh Site Type")
            sadigh.setValue(self.params['SADIGH_SITE_TYPE'])
            site.addParameter(vs30)
            site.addParameter(depth25)
            site.addParameter(sadigh)
            jsite_list.add(site)
        return jsite_list

class ClassicalMixin(BasePSHAMixin):
    """Classical PSHA method for performing Hazard calculations.

    Implements the JobMixin, which has a primary entry point of execute().
    Execute is responsible for dispatching celery tasks.

    Note that this Mixin, during execution, will always be an instance of the
    Job class, and thus has access to the self.params dict, full of config
    params loaded from the Job configuration file."""

    @preload
    def execute(self):

        results = []

        source_model_generator = random.Random()
        source_model_generator.seed(
                self.params.get('SOURCE_MODEL_LT_RANDOM_SEED', None))

        gmpe_generator = random.Random()
        gmpe_generator.seed(self.params.get('GMPE_LT_RANDOM_SEED', None))

        realizations = int(self.params['NUMBER_OF_LOGIC_TREE_SAMPLES'])

        # tally and log the total number of sites
        total_sites = 0
        for site_list in self.site_list_generator():
            total_sites += len(site_list)

        LOG.info('Going to run classical PSHA hazard for \
%s realizations and %s sites'
                % (realizations, total_sites))

        for realization in range(0, realizations):
            LOG.info('Calculatiing hazard curves for realization %s'
                     % realization)
            pending_tasks = []
            self.store_source_model(source_model_generator.getrandbits(32))
            self.store_gmpe_map(source_model_generator.getrandbits(32))
            for site_list in self.site_list_generator():
                pending_tasks.append(
                    tasks.compute_hazard_curve.delay(self.id,
                            site_list, realization))

            for task in pending_tasks:
                task.wait()
                if task.status != 'SUCCESS':
                    raise Exception(task.result)
                results.extend(task.result)
        return results

    @preload
    def compute_hazard_curve(self, site_list, realization):
        """ Compute hazard curves, write them to KVS as JSON,
        and return a list of the KVS keys for each curve. """
        jsite_list = self.parameterize_sites(site_list)
        hazard_curves = java.jclass("HazardCalculator").getHazardCurvesAsJson(
            jsite_list,
            self.generate_erf(),
            self.generate_gmpe_map(),
            self.get_iml_list(),
            float(self.params['MAXIMUM_DISTANCE']))

        # write the curves to the KVS and return a list of the keys
        kvs_client = kvs.get_client()
        curve_keys = []
        for i in xrange(0, len(hazard_curves)):
            curve = hazard_curves[i]
            site = site_list[i]
            lon = site.longitude
            lat = site.latitude
            curve_key = kvs.tokens.hazard_curve_key(self.id,
                                                    realization,
                                                    lon,
                                                    lat)
            kvs_client.set(curve_key, curve)
            curve_keys.append(curve_key)
        return curve_keys

class EventBasedMixin(BasePSHAMixin): # pylint: disable=W0232
    """Probabilistic Event Based method for performing Hazard calculations.

    Implements the JobMixin, which has a primary entry point of execute().
    Execute is responsible for dispatching celery tasks.

    Note that this Mixin, during execution, will always be an instance of the
    Job class, and thus has access to the self.params dict, full of config
    params loaded from the Job configuration file."""

    @preload
    def execute(self):
        """Main hazard processing block.

        Loops through various random realizations, spawning tasks to compute
        GMFs."""
        results = []

        source_model_generator = random.Random()
        source_model_generator.seed(
                self.params.get('SOURCE_MODEL_LT_RANDOM_SEED', None))

        gmpe_generator = random.Random()
        gmpe_generator.seed(self.params.get('GMPE_LT_RANDOM_SEED', None))

        gmf_generator = random.Random()
        gmf_generator.seed(self.params.get('GMF_RANDOM_SEED', None))

        histories = int(self.params['NUMBER_OF_SEISMICITY_HISTORIES'])
        realizations = int(self.params['NUMBER_OF_LOGIC_TREE_SAMPLES'])
        LOG.info("Going to run hazard for %s histories of %s realizations each."
                % (histories, realizations))

        for i in range(0, histories):
            pending_tasks = []
            for j in range(0, realizations):
                self.store_source_model(source_model_generator.getrandbits(32))
                self.store_gmpe_map(gmpe_generator.getrandbits(32))
                for site_list in self.site_list_generator():
                    stochastic_set_id = "%s!%s" % (i, j)
                    # pylint: disable=E1101
                    pending_tasks.append(
                        tasks.compute_ground_motion_fields.delay(
                            self.id,
                            site_list,
                            stochastic_set_id, gmf_generator.getrandbits(32)))

            for task in pending_tasks:
                task.wait()
                if task.status != 'SUCCESS':
                    raise Exception(task.result)

            # if self.params['OUTPUT_GMF_FILES']
            for j in range(0, realizations):
                stochastic_set_id = "%s!%s" % (i, j)
                stochastic_set_key = kvs.generate_product_key(
                    self.id, kvs.tokens.STOCHASTIC_SET_TOKEN, stochastic_set_id)
                print "Writing output for ses %s" % stochastic_set_key
                ses = kvs.get_value_json_decoded(stochastic_set_key)
                if ses:
                    results.extend(self.write_gmf_files(ses))
        return results

    def write_gmf_files(self, ses):
        """Generate a GeoTiff file for each GMF."""
        image_grid = self.region.grid
        iml_list = [float(param)
                    for param
                    in self.params['INTENSITY_MEASURE_LEVELS'].split(",")]

        LOG.debug("Generating GMF image, grid is %s col by %s rows" % (
                image_grid.columns, image_grid.rows))
        LOG.debug("IML: %s" % (iml_list))
        files = []
        for event_set in ses:
            for rupture in ses[event_set]:

                # NOTE(fab): we have to explicitly convert the JSON-decoded
                # tokens from Unicode to string, otherwise the path will not
                # be accepted by the GeoTiffFile constructor
                path = os.path.join(self.base_path, self['OUTPUT_DIR'],
                        "gmf-%s-%s.tiff" % (str(event_set.replace("!", "_")),
                                            str(rupture.replace("!", "_"))))
                gwriter = geotiff.GMFGeoTiffFile(path, image_grid,
                    init_value=0.0, normalize=True, iml_list=iml_list,
                    discrete=True)
                for site_key in ses[event_set][rupture]:
                    site = ses[event_set][rupture][site_key]
                    site_obj = shapes.Site(site['lon'], site['lat'])
                    point = image_grid.point_at(site_obj)
                    gwriter.write((point.row, point.column),
                        math.exp(float(site['mag'])))

                gwriter.close()
                files.append(path)
                files.append(gwriter.html_path)
        return files

    @preload
    def compute_ground_motion_fields(self, site_list, stochastic_set_id, seed):
        """Ground motion field calculation, runs on the workers."""
        jpype = java.jvm()

        jsite_list = self.parameterize_sites(site_list)
        key = kvs.generate_product_key(
                    self.id, kvs.tokens.STOCHASTIC_SET_TOKEN, stochastic_set_id)
        gmc = self.params['GROUND_MOTION_CORRELATION']
        correlate = (gmc == "true" and True or False)
        java.jclass("HazardCalculator").generateAndSaveGMFs(
                self.cache, key, stochastic_set_id, jsite_list,
                 self.generate_erf(),
                self.generate_gmpe_map(),
                java.jclass("Random")(seed),
                jpype.JBoolean(correlate))

def gmf_id(history_idx, realization_idx, rupture_idx):
    """ Return a GMF id suitable for use as a KVS key """
    return "%s!%s!%s" % (history_idx, realization_idx, rupture_idx)

job.HazJobMixin.register("Event Based", EventBasedMixin, order=0)
job.HazJobMixin.register("Classical", ClassicalMixin, order=1)
<|endoftext|>