<commit_before>import numpy as np
import libpinocchio_pywrap as se3
import utils
from explog import exp

class RobotWrapper:
    def __init__(self,filename):
        self.modelFileName = filename
        self.model = se3.buildModelFromUrdf(filename,True)
        self.data = self.model.createData()
        self.v0 = utils.zero(self.nv)
        self.q0 = utils.zero(self.nq)

    def increment(self,q,dq):
        M = se3.SE3( se3.Quaternion(q[6,0],q[3,0],q[4,0],q[5,0]).matrix(), q[:3])
        dM = exp(dq[:6])
        M = M*dM
        q[:3] = M.translation
        q[3:7] = se3.Quaternion(M.rotation).coeffs()
        q[7:] += dq[6:]

    @property
    def nq(self):
        return self.model.nq
    @property
    def nv(self):
        return self.model.nv

    def com(self,q):
        return se3.centerOfMass(self.model,self.data,q)
    def Jcom(self,q):
        return se3.jacobianCenterOfMass(self.model,self.data,q)

    def mass(self,q):
        return se3.crba(self.model,self.data,q)
    def biais(self,q,v):
        return se3.rnea(self.model,self.data,q,v,self.v0)
    def gravity(self,q):
        return se3.rnea(self.model,self.data,q,self.v0,self.v0)

    def position(self,q,index):
        se3.kinematics(self.model,self.data,q,self.v0)
        return self.data.oMi[index]
    def velocity(self,q,v,index):
        se3.kinematics(self.model,self.data,q,v)
        return self.data.v[index]
    def jacobian(self,q,index):
        return se3.jacobian(self.model,self.data,index,q,True)

    # --- ACCESS TO NAMES ----
    # Return the index of the joint whose name is given in argument.
    def index(self,name):
        return [ i for i,n in enumerate(self.model.names) if n == name ][0]

    # --- VIEWER ---
    # For each joint/body index, returns the corresponding name of the node in Gepetto-viewer.
    def viewerNodeNames(self,index):
        assert( self.model.hasVisual[index] )
        return self.viewerRootNodeName+'/'+self.model.bodyNames[index]

    def initDisplay(self,viewerRootNodeName = "world/pinocchio", loadModel = False):
        import gepetto.corbaserver
        try:
            self.viewer=gepetto.corbaserver.Client()
            self.viewerRootNodeName = viewerRootNodeName
            if loadModel:
                self.loadDisplayModel(viewerRootNodeName)
        except:
            if 'viewer' in self.__dict__: del self.viewer
            print "Error while starting the viewer client. "
            print "Check wheter gepetto-viewer is properly started"

    # Create the scene displaying the robot meshes in Gepetto-viewer.
    def loadDisplayModel(self, nodeName, windowName = "pinocchio"):
        import os
        if not self.viewer.gui.createWindow (windowName):
            print "Warning: window '"+windowName+"' already created. Cannot (re-)load the model."
            return
        self.viewer.gui.createSceneWithFloor("world")
        self.viewer.gui.addSceneToWindow("world",windowName)
        self.viewer.gui.addURDF(nodeName,
                                self.modelFileName,
                                os.path.dirname(self.modelFileName)+"/")

    # Display in gepetto-view the robot at configuration q, by placing all the bodies.
    def display(self,q):
        if 'viewer' not in self.__dict__: return
        # Update the robot geometry.
        se3.kinematics(self.model,self.data,q,self.v0)
        # Iteratively place the robot bodies.
        for i in range(1,self.model.nbody):
            if self.model.hasVisual[i]:
                M = self.data.oMi[i]
                self.viewer.gui.applyConfiguration(self.viewerNodeNames(i),
                                                   utils.se3ToXYZQUAT(M))
        self.viewer.gui.refresh()

__all__ = [ 'RobotWrapper' ]
<commit_msg>[Python][Major] Modified the viewer init to read the meshes in a different root dir.<commit_after>import numpy as np
import libpinocchio_pywrap as se3
import utils
from explog import exp

class RobotWrapper:
    def __init__(self,filename):
        self.modelFileName = filename
        self.model = se3.buildModelFromUrdf(filename,True)
        self.data = self.model.createData()
        self.v0 = utils.zero(self.nv)
        self.q0 = utils.zero(self.nq)

    def increment(self,q,dq):
        M = se3.SE3( se3.Quaternion(q[6,0],q[3,0],q[4,0],q[5,0]).matrix(), q[:3])
        dM = exp(dq[:6])
        M = M*dM
        q[:3] = M.translation
        q[3:7] = se3.Quaternion(M.rotation).coeffs()
        q[7:] += dq[6:]

    @property
    def nq(self):
        return self.model.nq
    @property
    def nv(self):
        return self.model.nv

    def com(self,q):
        return se3.centerOfMass(self.model,self.data,q)
    def Jcom(self,q):
        return se3.jacobianCenterOfMass(self.model,self.data,q)

    def mass(self,q):
        return se3.crba(self.model,self.data,q)
    def biais(self,q,v):
        return se3.rnea(self.model,self.data,q,v,self.v0)
    def gravity(self,q):
        return se3.rnea(self.model,self.data,q,self.v0,self.v0)

    def position(self,q,index):
        se3.kinematics(self.model,self.data,q,self.v0)
        return self.data.oMi[index]
    def velocity(self,q,v,index):
        se3.kinematics(self.model,self.data,q,v)
        return self.data.v[index]
    def jacobian(self,q,index):
        return se3.jacobian(self.model,self.data,index,q,True)

    # --- ACCESS TO NAMES ----
    # Return the index of the joint whose name is given in argument.
    def index(self,name):
        return [ i for i,n in enumerate(self.model.names) if n == name ][0]

    # --- VIEWER ---
    # For each joint/body index, returns the corresponding name of the node in Gepetto-viewer.
    def viewerNodeNames(self,index):
        assert( self.model.hasVisual[index] )
        return self.viewerRootNodeName+'/'+self.model.bodyNames[index]

    def initDisplay(self,viewerRootNodeName = "world/pinocchio", loadModel = False):
        import gepetto.corbaserver
        try:
            self.viewer=gepetto.corbaserver.Client()
            self.viewerRootNodeName = viewerRootNodeName
            if loadModel:
                self.loadDisplayModel(viewerRootNodeName)
        except:
            if 'viewer' in self.__dict__: del self.viewer
            print "Error while starting the viewer client. "
            print "Check wheter gepetto-viewer is properly started"

    # Create the scene displaying the robot meshes in Gepetto-viewer.
    def loadDisplayModel(self, nodeName, windowName = "pinocchio", meshDir = None):
        import os
        if not self.viewer.gui.createWindow (windowName):
            print "Warning: window '"+windowName+"' already created. Cannot (re-)load the model."
            return
        if not meshDir: meshDir = os.path.dirname(self.modelFileName)+"/"
        self.viewer.gui.createSceneWithFloor("world")
        self.viewer.gui.addSceneToWindow("world",windowName)
        self.viewer.gui.addURDF(nodeName,
                                self.modelFileName,
                                meshDir)

    # Display in gepetto-view the robot at configuration q, by placing all the bodies.
    def display(self,q):
        if 'viewer' not in self.__dict__: return
        # Update the robot geometry.
        se3.kinematics(self.model,self.data,q,self.v0)
        # Iteratively place the robot bodies.
        for i in range(1,self.model.nbody):
            if self.model.hasVisual[i]:
                M = self.data.oMi[i]
                self.viewer.gui.applyConfiguration(self.viewerNodeNames(i),
                                                   utils.se3ToXYZQUAT(M))
        self.viewer.gui.refresh()

__all__ = [ 'RobotWrapper' ]
<|endoftext|>


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

# <https://bbs.archlinux.org/viewtopic.php?id=75365#p579926>,
# <https://bbs.archlinux.org/viewtopic.php?pid=783507#p783507>.
import pygtk
pygtk.require('2.0')
import gtk
import os

class DoTheShutdown:

    # Cancel/exit
    def delete_event(self, widget, event, data=None):
        gtk.main_quit()
        return False

    # Logout
    def logout(self, widget):
        os.system("skill -TERM $DESKTOP_SESSION")
        gtk.main_quit()

    # Reboot
    def reboot(self, widget):
        # os.system("dbus-send --system --print-reply --dest=org.freedesktop.ConsoleKit /org/freedesktop/ConsoleKit/Manager org.freedesktop.ConsoleKit.Manager.Restart")
        os.system("sudo shutdown -r now")
        gtk.main_quit()

    # Shutdown
    def shutdown(self, widget):
        # os.system("dbus-send --system --print-reply --dest=org.freedesktop.ConsoleKit /org/freedesktop/ConsoleKit/Manager org.freedesktop.ConsoleKit.Manager.Stop")
        os.system("sudo shutdown -h now")
        gtk.main_quit()

    # S3 aka Suspend to RAM aka Sleep
    def suspend(self, widget):
        os.system("xscreensaver-command -lock && sudo hibernate-ram")
        gtk.main_quit()

    # Hibernate
    def hibernate(self, widget):
        os.system("xscreensaver-command -lock && sudo hibernate")
        gtk.main_quit()

    def __init__(self):
        width = 100

        # Create a new window
        self.window = gtk.Window(gtk.WINDOW_TOPLEVEL)
        self.window.set_title("Shut down")
        self.window.set_resizable(False)
        self.window.set_position(1)
        self.window.connect("delete_event", self.delete_event)
        # self.window.set_border_width(20)

        # Create an accelgroup and add it to the window
        accel_group = gtk.AccelGroup()
        self.window.add_accel_group(accel_group)

        # Create a box to pack widgets into
        self.box1 = gtk.HBox(False, 0)
        self.window.add(self.box1)

        # Create cancel button
        self.button1 = gtk.Button("Cancel")
        self.button1.set_border_width(10)
        self.button1.set_size_request(width,-1)
        self.button1.connect("clicked", self.delete_event, "closed")
        self.box1.pack_start(self.button1, True, True, 0)
        self.button1.add_accelerator("clicked", accel_group,
            gtk.gdk.keyval_from_name('Escape'), 0, 0)
        self.button1.add_accelerator("clicked", accel_group,
            gtk.gdk.keyval_from_name('c'), 0, 0)
        self.button1.show()

        # Create logout button
        self.button2 = gtk.Button("Logout")
        self.button2.set_border_width(10)
        self.button2.set_size_request(width,-1)
        self.button2.connect("clicked", self.logout)
        self.box1.pack_start(self.button2, True, True, 0)
        self.button2.add_accelerator("clicked", accel_group,
            gtk.gdk.keyval_from_name('l'), 0, 0)
        self.button2.show()

        # Create reboot button
        self.button3 = gtk.Button("Reboot")
        self.button3.set_border_width(10)
        self.button3.set_size_request(width,-1)
        self.button3.connect("clicked", self.reboot)
        self.box1.pack_start(self.button3, True, True, 0)
        self.button3.add_accelerator("clicked", accel_group,
            gtk.gdk.keyval_from_name('r'), 0, 0)
        self.button3.show()

        # Create shutdown button
        self.button4 = gtk.Button("Shutdown")
        self.button4.set_border_width(10)
        self.button4.set_size_request(width,-1)
        self.button4.connect("clicked", self.shutdown)
        self.box1.pack_start(self.button4, True, True, 0)
        self.button4.add_accelerator("clicked", accel_group,
            gtk.gdk.keyval_from_name('s'), 0, 0)
        self.button4.show()

        # Create suspend button
        self.button5 = gtk.Button("Suspend")
        self.button5.set_border_width(10)
        self.button5.set_size_request(width,-1)
        self.button5.connect("clicked", self.suspend)
        self.box1.pack_start(self.button5, True, True, 0)
        self.button5.add_accelerator("clicked", accel_group,
            gtk.gdk.keyval_from_name('h'), 0, 0)
        self.button5.show()

        # Create hibernate button
        self.button6 = gtk.Button("Hibernate")
        self.button6.set_border_width(10)
        self.button6.set_size_request(width,-1)
        self.button6.connect("clicked", self.hibernate)
        self.box1.pack_start(self.button6, True, True, 0)
        self.button6.add_accelerator("clicked", accel_group,
            gtk.gdk.keyval_from_name('h'), 0, 0)
        self.button6.show()

        self.box1.show()
        self.window.show()

def main():
    gtk.main()

if __name__ == "__main__":
    run_it = DoTheShutdown()
    main()
<commit_msg>shutdown_menu.py<commit_after>#!/usr/bin/env python
# -*- coding: utf-8 -*-

# <https://bbs.archlinux.org/viewtopic.php?id=75365#p579926>,
# <https://bbs.archlinux.org/viewtopic.php?pid=783507#p783507>.
import pygtk
pygtk.require('2.0')
import gtk
import os

class DoTheShutdown:

    # Cancel/exit
    def delete_event(self, widget, event, data=None):
        gtk.main_quit()
        return False

    # Logout
    def logout(self, widget):
        os.system("skill -TERM $DESKTOP_SESSION")
        gtk.main_quit()

    # Reboot
    def reboot(self, widget):
        # os.system("dbus-send --system --print-reply --dest=org.freedesktop.ConsoleKit /org/freedesktop/ConsoleKit/Manager org.freedesktop.ConsoleKit.Manager.Restart")
        os.system("sudo shutdown -r now")
        gtk.main_quit()

    # Shutdown
    def shutdown(self, widget):
        # os.system("dbus-send --system --print-reply --dest=org.freedesktop.ConsoleKit /org/freedesktop/ConsoleKit/Manager org.freedesktop.ConsoleKit.Manager.Stop")
        os.system("sudo shutdown -h now")
        gtk.main_quit()

    # S3 aka Suspend to RAM aka Sleep
    def suspend(self, widget):
        os.system("xscreensaver-command -lock && sudo hibernate-ram")
        gtk.main_quit()

    # Hibernate
    def hibernate(self, widget):
        os.system("xscreensaver-command -lock && sudo hibernate")
        gtk.main_quit()

    # Lock
    def lock(self, widget):
        os.system("xscreensaver-command -lock && sleep 1 && xset dpms force off")
        gtk.main_quit()

    def __init__(self):
        border_width = 10
        button_width = 100

        # Create a new window
        self.window = gtk.Window(gtk.WINDOW_TOPLEVEL)
        self.window.set_title("Shut down")
        self.window.set_resizable(False)
        self.window.set_position(1)
        self.window.connect("delete_event", self.delete_event)
        # self.window.set_border_width(20)

        # Create an accelgroup and add it to the window
        accel_group = gtk.AccelGroup()
        self.window.add_accel_group(accel_group)

        # Create a box to pack widgets into
        self.box1 = gtk.HBox(False, 0)
        self.window.add(self.box1)

        # Create cancel button
        self.button1 = gtk.Button("Cancel")
        self.button1.set_border_width(border_width)
        self.button1.set_size_request(button_width,-1)
        self.button1.connect("clicked", self.delete_event, "closed")
        self.box1.pack_start(self.button1, True, True, 0)
        self.button1.add_accelerator("clicked", accel_group,
            gtk.gdk.keyval_from_name('Escape'), 0, 0)
        self.button1.add_accelerator("clicked", accel_group,
            gtk.gdk.keyval_from_name('c'), 0, 0)
        self.button1.show()

        # Create logout button
        self.button2 = gtk.Button("Logout")
        self.button2.set_border_width(border_width)
        self.button2.set_size_request(button_width,-1)
        self.button2.connect("clicked", self.logout)
        self.box1.pack_start(self.button2, True, True, 0)
        self.button2.add_accelerator("clicked", accel_group,
            gtk.gdk.keyval_from_name('o'), 0, 0)
        self.button2.show()

        # Create reboot button
        self.button3 = gtk.Button("Reboot")
        self.button3.set_border_width(border_width)
        self.button3.set_size_request(button_width,-1)
        self.button3.connect("clicked", self.reboot)
        self.box1.pack_start(self.button3, True, True, 0)
        self.button3.add_accelerator("clicked", accel_group,
            gtk.gdk.keyval_from_name('r'), 0, 0)
        self.button3.show()

        # Create shutdown button
        self.button4 = gtk.Button("Shutdown")
        self.button4.set_border_width(border_width)
        self.button4.set_size_request(button_width,-1)
        self.button4.connect("clicked", self.shutdown)
        self.box1.pack_start(self.button4, True, True, 0)
        self.button4.add_accelerator("clicked", accel_group,
            gtk.gdk.keyval_from_name('h'), 0, 0)
        self.button4.show()

        # Create suspend button
        self.button5 = gtk.Button("Suspend")
        self.button5.set_border_width(border_width)
        self.button5.set_size_request(button_width,-1)
        self.button5.connect("clicked", self.suspend)
        self.box1.pack_start(self.button5, True, True, 0)
        self.button5.add_accelerator("clicked", accel_group,
            gtk.gdk.keyval_from_name('s'), 0, 0)
        self.button5.show()

        # Create hibernate button
        self.button6 = gtk.Button("Hibernate")
        self.button6.set_border_width(border_width)
        self.button6.set_size_request(button_width,-1)
        self.button6.connect("clicked", self.hibernate)
        self.box1.pack_start(self.button6, True, True, 0)
        self.button6.add_accelerator("clicked", accel_group,
            gtk.gdk.keyval_from_name('i'), 0, 0)
        self.button6.show()

        # Create lock button
        self.button7 = gtk.Button("Lock")
        self.button7.set_border_width(border_width)
        self.button7.set_size_request(button_width,-1)
        self.button7.connect("clicked", self.lock)
        self.box1.pack_start(self.button7, True, True, 0)
        self.button7.add_accelerator("clicked", accel_group,
            gtk.gdk.keyval_from_name('l'), 0, 0)
        self.button7.show()

        self.box1.show()
        self.window.show()

def main():
    gtk.main()

if __name__ == "__main__":
    run_it = DoTheShutdown()
    main()
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-

import sys
import zmq
import time
import contextlib
try:
    import simplejson as json
except ImportError:
    import json

DEBUG = False

def debug(prefix, ss):
    if (DEBUG):
        for s in ss.splitlines():
            sys.stderr.write("[DEBUG] %s %s\n" % (prefix, s))

def escape(s):
    return(s.replace("\\", "\\\\").replace(")", "\\)").replace("]", "\\]"))

def context():
    ctx = zmq.Context()
    return(ctx)

@contextlib.contextmanager
def connection(ctx, endpoint):
    conn = Connection(ctx, endpoint)
    try:
        yield(conn)
    finally:
        conn.close()

def pp(data):
    if ("done" in data):
        return(".done")
    elif ("fail" in data):
        return("fail: %d" % data["fail"]["code"])
    else:
        return("chunk: ...")

class Cursor(object):

    def __init__(self, conn, chan, when):
        self.conn   = conn
        self.chan   = chan
        self.when0  = when
        self.read   = 0
        self.when   = when
        self.closed = False

    def close(self):
        self.closed = True
        self.conn.send("close %d;" % self.chan)
        t = time.time() - self.when0
        debug("timing", "total: %.4f / %d / %d/s" % (t, self.read, self.read/t))
        return(self.conn.recv())

    def _eof(self, data):
        return("done" in data or "fail" in data)

    def fetchall(self):
        while (True):
            for row in self.fetch():
                yield(row)
            if (self.closed):
                break

    def fetch(self, limit=64):
        if (self.closed):
            raise(RuntimeError("closed"))
        self.when = time.time()
        self.conn.send("fetch %d %d;" % (self.chan, limit))
        msg = self.conn.recv()
        self.read += len(msg)
        debug("timing", "fetch: %.4f / items: %d" % (time.time() - self.when, len(msg)))
        if (self._eof(msg[-1])):
            self.close()
        return(msg)

    def fetchone(self):
        return(self.fetch(1))

class Stmt(object):

    def __init__(self, conn):
        self.conn  = conn

    def cursor(self):
        when = time.time()
        self.conn.send(";")
        data = self.conn.recv()[0]
        if ("done" not in data):
            raise(RuntimeError(pp(data)))
        return(Cursor(self.conn, data["done"]["channel"], when))

    def done(self):
        for _ in self.cursor().fetchall():
            pass

    def send(self, line, nl=True):
        if (nl):
            nl = "\n"
        else:
            nl = ""
        self.conn.send("%s%s" % (nl, line), zmq.SNDMORE)

    def create(self):
        return(CreateStmt(self))

    def match(self):
        return(MatchStmt(self))

    def deref(self, key=None):
        stmt = DerefStmt(self)
        if (key is not None):
            stmt.deref(key)
        return(stmt)

class DerefStmt(object):

    def __init__(self, stmt, *args, **kwargs):
        self.stmt = stmt

    def deref(self, k):
        self.stmt.send("deref %s" % (k,))
        return(self)

    def cursor(self):
        return(self.stmt.cursor())

    def done(self):
        return(self.stmt.done())

class CreateStmt(object):

    def __init__(self, stmt, *args, **kwargs):
        self.stmt = stmt

    def node(self, a):
        self.stmt.send("create (%s)" % escape(a))
        return(self)

    def link(self, a, l, b):
        self.stmt.send("create (%s) -[%s]> (%s)" % (escape(a), escape(l), escape(b)))
        return(self)

    def cursor(self):
        return(self.stmt.cursor())

    def done(self):
        return(self.stmt.done())

class MatchStmt(object):

    def __init__(self, stmt):
        self.stmt = stmt
        self.stmt.send("match ")

    def select(self, node):
        self.stmt.send("(%s)" % escape(node), nl=False)
        return(self)

    def out(self, label, node=""):
        self.stmt.send(" -[%s]> (%s)" % (escape(label), escape(node)), nl=False)
        return(self)

    def or_(self):
        self.stmt.send("match ")
        return(self)

    def cursor(self):
        return(self.stmt.cursor())

    def done(self):
        return(self.stmt.done())

class Connection(object):

    def __init__(self, context, endpoint):
        self.socket = context.socket(zmq.REQ)
        self.socket.connect(endpoint)
        self.debug = []

    def send(self, s, flags=0):
        self.debug.append(s)
        if (flags == 0):
           debug(">", "".join(self.debug))
           self.debug = []
        self.socket.send_string(s, flags)

    def _recv(self):
        msg = self.socket.recv()
        debug("<", msg)
        return(json.loads(msg))

    def recv(self):
        tmp = [self._recv()]
        while (self.socket.getsockopt(zmq.RCVMORE)):
            tmp.append(self._recv())
        return(tmp)

    def using(self, username, namespace):
        self.send("begin using (%s) (%s)" % (escape(username), namespace), zmq.SNDMORE)
        return(Stmt(self))

    def close(self):
        self.socket.close()

<commit_msg>Implemented the Leela Network Iterator user-friendly access<commit_after># -*- coding: utf-8 -*-

# TODO:
#   - Return a list of lists to the user. Why?
#   - Write Functional Tests
#       - Populate the test database
#       - Authentication tests
#       - Thread tests
#   - Handle the Authentication process. How?
#   - Handle the Zookeeper nodes as a Cursor
#   - Write the Man documentation
#   - Write the Pretty Printer (WarpDrive Protocol)
#   - Provide Threads
#   - Use Timeout
#   - Handle Deref

import sys
import zmq
import time
import contextlib
import unittest

DEBUG = True
#DEBUG = False

def debug(prefix, ss):
    if (DEBUG):
        for s in ss.splitlines():
            sys.stderr.write("[DEBUG] %s %s\n" % (prefix, s))

def escape(s):
    return(s.replace("\\", "\\\\").replace(")", "\\)").replace("]", "\\]"))

def context():
    ctx = zmq.Context()
    return(ctx)

def pp(data):
    if ("fail" in data):
        return("fail: %s" % data[2])

class Connection(object):

    def __init__(self, context, endpoint):
        self.context = context
        self.endpoint = endpoint
        self.debug = []

    def connect(self):
        self.socket = self.context.socket(zmq.REQ)
        self.socket.connect(self.endpoint[0])

    def channel(self):
        when = time.time()
        self.send("")
        data = self.recv()
        if ("done" not in data):
            raise(RuntimeError(pp(data)))
        return(Channel(self, int(data[1]), when, self.auth()))

    def execute(self, query):
        self.send(self.auth(), zmq.SNDMORE)
        self.send("begin", zmq.SNDMORE)
        self.send("%s" % (query), zmq.SNDMORE)
        return (self.channel())

    def send(self, s, flags=0):
        self.debug.append(s)
        if (flags == 0):
           debug(">", "".join(self.debug))
           self.debug = []
        self.socket.send_string(s, flags)

    def _recv(self):
        msg = self.socket.recv()
        debug("<", msg)
        return(msg)

    def recv(self):
        tmp = [self._recv()]
        while (self.socket.getsockopt(zmq.RCVMORE)):
            tmp.append(self._recv())
        return(tmp)

    def close(self):
        self.socket.close()

    def auth(self):
        return("usertest:0:0 0")

    @contextlib.contextmanager
    def cursor(self, query, timeout=-1, nowait=False):
        self.connect()
        try:
            channel = self.execute(query)
            yield(channel.fetchall())
            #channel.close(nowait)
        finally:
            channel.close(nowait)
            self.close()

class Channel(object):

    def __init__(self, conn, chan, when, auth):
        self.conn   = conn
        self.chan   = chan
        self.when0  = when
        self.read   = 0
        self.when   = when
        self.closed = False
        self.auth   = auth

    def close(self, nowait=False):
        self.closed = True
        self.conn.send(self.auth, zmq.SNDMORE)
        self.conn.send("close", zmq.SNDMORE)
        if nowait:
            self.conn.send("%d" % self.chan, zmq.SNDMORE)
            self.conn.send("nowait")
        else:
            self.conn.send("%d" % self.chan)

        t = time.time() - self.when0
        debug("timing", "total: %.4f / %d / %d/s" % (t, self.read, self.read/t))
        return(self.conn.recv())

    def _eof(self, data):
        return("done" in data or "fail" in data)

    def fetchall(self):
        while (True):
            if (self.closed):
                break
            for el in self.fetch():
                yield el

    def fetch(self):
        if (self.closed):
            raise(RuntimeError("closed"))
        self.when = time.time()
        self.conn.send(self.auth, zmq.SNDMORE)
        self.conn.send("fetch", zmq.SNDMORE)
        self.conn.send("%d" % self.chan)
        msg = self.conn.recv()
        self.read += len(msg)
        debug("timing", "fetch: %.4f / items: %d" % (time.time() - self.when, len(msg)))
        if (self._eof(msg)):
            if "done" in msg:
                msg.remove('done')
            else:
                msg.remove('fail')
            self.close()
        else:
            msg.remove('item')

        lt = []
        more = 1
        while len(msg) > 0 and more > 0:
            if 'list' == msg[0]:
                more = int(msg[1])
                msg  = msg[2:]
            elif 'name' == msg[0]:
                lt.append(msg)
                more = 0
            else:
                lt.append((msg[0], msg[2:2+int(msg[1])]))
                msg   = msg[2+int(msg[1]):]
                more -= 1
        return lt

    def fetchone(self):
        return(self.fetch())

#Library Tests

def setUpModule ():

    def gen_mc_name(num):
        for x in range(num):
            yield "hm" + "{0:04d}".format(x)

    def gen_ni_name(num):
        for k in range(8):
            x = num * 8 + k
            lt = []
            while x > 0:
                lt.append("%02x" % (x%256))
                x = x // 256
            while (len(lt) < 6):
                lt.append("00")
            yield ":".join(lt)

    def gen_sw_name(num):
        num *= 8
        return ("sw%04d" % (num//64))

    def create(orig, label, dest):
        return (["make (%s)" % orig, "make (%s)" % dest, "make (%s) -[%s]> (%s)" % (orig, label, dest)])

    def create_mc(x):
        machine = "hm%04d" % x
        switch = gen_sw_name(x)

        lt = create("ITA", "machine", machine)
        lt.extend(create("ITA", "switch", switch))

        for mac in gen_ni_name(x):
            lt.extend(create(machine, "nic", mac))
            lt.extend(create(mac, "link", switch))

        return("using (test_namespace) %s;" % "\n".join(lt))

    if len(sys.argv) > 1:
        conn = Connection(context(), ["tcp://warp0017.locaweb.com.br:4080"])
        conn.connect()
        for x in range(10):
            query = create_mc(x)
            chann = conn.execute(query)
            chann.close(nowait=True)

def tearDownModule ():
    pass

class UserLibTest (unittest.TestCase):

    def setUp(self):
        self.conn = Connection(context(), ["tcp://warp0017.locaweb.com.br:4080"])

    def test_database_is_empty (self):
        with self.conn.cursor("using (test_database) path (test_ITA);") as cursor:
            for row in cursor:
                self.assertEqual(row, [])

    def test_make_node(self):
        with self.conn.cursor("using (test_database) make (test_ITA);") as cursor:
            for row in cursor:
                pass
        with self.conn.cursor("using (test_database) path (test_ITA);") as cursor:
            for row in cursor:
                print row

if __name__ == '__main__':
    unittest.main(argv=[sys.argv[0]])
<|endoftext|>


<commit_before>#!/usr/bin/python
# -*- coding: utf-8 -*-
##
# distributions.py: module for probability distributions
##
# © 2012 Chris Ferrie (csferrie@gmail.com) and
#        Christopher E. Granade (cgranade@gmail.com)
#
# This file is a part of the Qinfer project.
# Licensed under the AGPL version 3.
##
# 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/>.
##

## IMPORTS ####################################################################

import numpy as np
import scipy.stats as st
import scipy.linalg as la
import abc

import warnings

## CLASSES ####################################################################

class Distribution(object):
    """
    Abstract base class for probability distributions on one or more random
    variables.
    """
    __metaclass__ = abc.ABCMeta

    @abc.abstractproperty
    def n_rvs(self):
        """
        The number of random variables that this distribution is over.

        :rtype: `int`
        """
        pass

    @abc.abstractmethod
    def sample(self, n=1):
        """
        Returns one or more samples from this probability distribution.

        :param int n: Number of samples to return.
        :return numpy.ndarray: An array containing samples from the
            distribution of shape ``(n, d)``, where ``d`` is the number of
            random variables.
        """
        pass

class ProductDistribution(Distribution):
    r"""
    Returns the Cartesian product of two distributions :math:`A` and
    :math:`B`, :math:`\Pr(A, B) = \Pr(A) \Pr(B)`.

    :param Distribution A: Distribution object representing :math:`A`.
    :param Distribution B: Distribution object representing :math:`B`.
    """
    def __init__(self, A, B):
        self.A = A
        self.B = B

    @property
    def n_rvs(self):
        return self.A.n_rvs + self.B.n_rvs

    def sample(self, n=1):
        A_sample = self.A.sample(n)
        B_sample = self.B.sample(n)
        return np.hstack((A_sample, B_sample))

_DEFAULT_RANGES = np.array([[0, 1]])
_DEFAULT_RANGES.flags.writeable = False # Prevent anyone from modifying the
                                        # default ranges.
class UniformDistribution(Distribution):
    """
    Uniform distribution on a given rectangular region.

    :param numpy.ndarray ranges: Array of shape ``(n_rvs, 2)``, where ``n_rvs``
        is the number of random variables, specifying the upper and lower limits
        for each variable.
    """

    def __init__(self, ranges=_DEFAULT_RANGES):
        if not isinstance(ranges, np.ndarray):
            ranges = np.array(ranges)

        if len(ranges.shape) == 1:
            ranges = ranges[np.newaxis, ...]

        self._ranges = ranges
        self._n_rvs = ranges.shape[0]
        self._delta = ranges[:, 1] - ranges[:, 0]

    @property
    def n_rvs(self):
        return self._n_rvs

    def sample(self, n=1):
        shape = (n, self._n_rvs)# if n == 1 else (self._n_rvs, n)
        z = np.random.random(shape)
        return self._ranges[:, 0] + z * self._delta

class ConstantDistribution(Distribution):
    """
    Represents a determinstic variable; useful for combining with other
    distributions, marginalizing, etc.

    :param values: Shape ``(n,)`` array or list of values :math:`X_0` such that
        :math:`\Pr(X) = \delta(X - X_0)`.
    """

    def __init__(self, values):
        self._values = np.array(values)[np.newaxis, :]

    @property
    def n_rvs(self):
        return self._values.shape[1]

    def sample(self, n=1):
        return np.repeat(self._values, n, axis=0)

class UniformDistributionWith0(Distribution):
    """
    Uniform distribution on a given rectangular region with padded zeros.

    :param numpy.ndarray ranges: Array of shape ``(n_rvs, 2)``, where ``n_rvs``
        is the number of random variables, specifying the upper and lower limits
        for each variable.
    """

    def __init__(self, ranges=_DEFAULT_RANGES, zeros = 0):
        warnings.warn(
            "This class has been superceded by ProductDistribution and ConstantDistribution.",
            DeprecationWarning
        )

        if not isinstance(ranges, np.ndarray):
            ranges = np.array(ranges)

        if len(ranges.shape) == 1:
            ranges = ranges[np.newaxis, ...]

        self._ranges = ranges
        self._n_rvs = ranges.shape[0]
        self._delta = ranges[:, 1] - ranges[:, 0]

        self.zeros = zeros

    @property
    def n_rvs(self):
        return self._n_rvs

    def sample(self, n=1):
        shape = (n, self._n_rvs)# if n == 1 else (self._n_rvs, n)
        z = np.random.random(shape)
        foo =  self._ranges[:, 0] + z * self._delta
        return np.pad(foo,((0,0), (0, self.zeros)), mode = 'constant')

    def grad_log_pdf(self, var):
        # THIS IS NOT TECHNICALLY LEGIT; BCRB doesn't technically work with a
        # prior that doesn't go to 0 at its end points.  But we do it anyway.
        if var.shape[0] == 1:
            return 12/(self._delta)**2
        else:
            return np.zeros(var.shape)

class NormalDistribution(Distribution):
    def __init__(self, mean, var):
        self.mean = mean
        self.var = var

        self.dist = st.norm(mean, np.sqrt(var))

    @property
    def n_rvs(self):
        return 1

    def sample(self, n=1):
        return self.dist.rvs(size=n)[:, np.newaxis]

    def grad_log_pdf(self, x):
        return -(x - self.mean) / self.var

class SlantedNormalDistribution(Distribution):
    """
    Uniform distribution on a given rectangular region.

    :param numpy.ndarray ranges: Array of shape ``(n_rvs, 2)``, where ``n_rvs``
        is the number of random variables, specifying the upper and lower limits
        for each variable.
    """

    def __init__(self, ranges=_DEFAULT_RANGES, weight=0.01):
        if not isinstance(ranges, np.ndarray):
            ranges = np.array(ranges)

        if len(ranges.shape) == 1:
            ranges = ranges[np.newaxis, ...]

        self._ranges = ranges
        self._n_rvs = ranges.shape[0]
        #self._delta = ranges[:, 1] - ranges[:, 0]
        self._weight = weight

    @property
    def n_rvs(self):
        return self._n_rvs

    def sample(self, n=1):
        shape = (n, self._n_rvs)# if n == 1 else (self._n_rvs, n)
        z = np.random.randn(n,self._n_rvs)
        return self._ranges[:, 0] +self._weight*z+np.random.rand(n)*self._ranges[:, 1];

class MVUniformDistribution(object):

    def __init__(self, dim = 6):
        self.dim = dim

    def sample(self, n = 1):
        return np.random.mtrand.dirichlet(np.ones(self.dim),n)

class DiscreteUniformDistribution(Distribution):
    def __init__(self, num_bits):
        self._num_bits = num_bits

    def sample(self, n=1):
        z = np.random.randint(2**self._num_bits,n)
        return z

# TODO: make the following into Distributions.
class HilbertSchmidtUniform(object):
    """
    Creates a new Hilber-Schmidt uniform prior on state space of dimension ``dim``.
    See e.g. [Mez06]_ and [Mis12]_.

    :param int dim: Dimension of the state space.
    """
    def __init__(self, dim=2):
        self.dim = dim
        self.paulis1Q = np.array([[[1,0],[0,1]],[[1,0],[0,-1]],[[0,-1j],[1j,0]],[[0,1],[1,0]]])

        self.paulis = self.make_Paulis(self.paulis1Q, 4)

    def sample(self):
        #Generate random unitary (see e.g. http://arxiv.org/abs/math-ph/0609050v2)
        g = (np.random.randn(self.dim,self.dim) + 1j*np.random.randn(self.dim,self.dim))/np.sqrt(2.0)
        q,r = la.qr(g)
        d = np.diag(r)

        ph = d/np.abs(d)
        ph = np.diag(ph)

        U = np.dot(q,ph)

        #Generate random matrix
        z = np.random.randn(self.dim,self.dim) + 1j*np.random.randn(self.dim,self.dim)

        rho = np.dot(np.dot(np.identity(self.dim)+U,np.dot(z,z.conj().transpose())),np.identity(self.dim)+U.conj().transpose())
        rho = rho/np.trace(rho)

        x = np.zeros([self.dim**2-1])
        for idx in xrange(self.dim**2-1):
            x[idx] = np.real(np.trace(np.dot(rho,self.paulis[idx+1])))

        return x

    def make_Paulis(self,paulis,d):
        if d == self.dim*2:
            return paulis
        else:
            temp = np.zeros([d**2,d,d],dtype='complex128')
            for idx in xrange(temp.shape[0]):
                temp[idx,:] = np.kron(paulis[np.trunc(idx/d)], self.paulis1Q[idx % 4])
            return self.make_Paulis(temp,d*2)

class HaarUniform(object):
    """
    Creates a new Haar uniform prior on state space of dimension ``dim``.

    :param int dim: Dimension of the state space.
    """
    def __init__(self,dim = 2):
        self.dim = dim

    def sample(self):
        #Generate random unitary (see e.g. http://arxiv.org/abs/math-ph/0609050v2)
        z = (np.random.randn(self.dim,self.dim) + 1j*np.random.randn(self.dim,self.dim))/np.sqrt(2.0)
        q,r = la.qr(z)
        d = np.diag(r)

        ph = d/np.abs(d)
        ph = np.diag(ph)

        U = np.dot(q,ph)

        #TODO: generalize this to general dimensions
        #Apply Haar random unitary to |0> state to get random pure state
        psi = np.dot(U,np.array([1,0]))
        z = np.real(np.dot(psi.conj(),np.dot(np.array([[1,0],[0,-1]]),psi)))
        y = np.real(np.dot(psi.conj(),np.dot(np.array([[0,-1j],[1j,0]]),psi)))
        x = np.real(np.dot(psi.conj(),np.dot(np.array([[0,1],[1,0]]),psi)))

        return np.array([x,y,z])

class GinibreUniform(object):
    """
    Creates a prior on state space of dimension dim according to the Ginibre
    ensemble with parameter ``k``.
    See e.g. [Mis12]_.

    :param int dim: Dimension of the state space.
    """
    def __init__(self,dim = 2, k = 2):
        self.dim = dim
        self.k = k

    def sample(self):
        #Generate random matrix
        z = np.random.randn(self.dim,self.k) + 1j*np.random.randn(self.dim,self.k)

        rho = np.dot(z,z.conj().transpose())
        rho = rho/np.trace(rho)

        z = np.real(np.trace(np.dot(rho,np.array([[1,0],[0,-1]]))))
        y = np.real(np.trace(np.dot(rho,np.array([[0,-1j],[1j,0]]))))
        x = np.real(np.trace(np.dot(rho,np.array([[0,1],[1,0]]))))

        return np.array([x,y,z])

class PostselectedDistribution(Distribution):
    """
    Postselects a distribution based on validity within a given model.
    """
    # TODO: rewrite LiuWestResampler in terms of this and a
    #       new MixtureDistribution.

    def __init__(self, distribution, model, maxiters=100):
        self._dist = distribution
        self._model = model

        self._maxiters = 100

    @property
    def n_rvs(self):
        return self._dist.n_rvs

    def sample(self, n=1):
        """
        Returns one or more samples from this probability distribution.

        :param int n: Number of samples to return.
        :return numpy.ndarray: An array containing samples from the
            distribution of shape ``(n, d)``, where ``d`` is the number of
            random variables.
        """
        samples = np.empty((n, self.n_rvs))
        idxs_to_sample = np.arange(n)

        iters = 0

        while idxs_to_sample.size and iters < self._maxiters:
            samples[idxs_to_sample] = self._dist.sample(len(idxs_to_sample))

            idxs_to_sample = idxs_to_sample[np.nonzero(np.logical_not(
                self._model.are_models_valid(samples[idxs_to_sample, :])
            ))[0]]

            iters += 1

        if idxs_to_sample.size:
            raise RuntimeError("Did not successfully postselect within {} iterations.".format(self._maxiters))

        return samples
<commit_msg>Added multivariate normal distribution<commit_after>#!/usr/bin/python
# -*- coding: utf-8 -*-
##
# distributions.py: module for probability distributions
##
# © 2012 Chris Ferrie (csferrie@gmail.com) and
#        Christopher E. Granade (cgranade@gmail.com)
#
# This file is a part of the Qinfer project.
# Licensed under the AGPL version 3.
##
# 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/>.
##

## IMPORTS ####################################################################

import numpy as np
import scipy.stats as st
import scipy.linalg as la
import abc

import warnings

## CLASSES ####################################################################

class Distribution(object):
    """
    Abstract base class for probability distributions on one or more random
    variables.
    """
    __metaclass__ = abc.ABCMeta

    @abc.abstractproperty
    def n_rvs(self):
        """
        The number of random variables that this distribution is over.

        :rtype: `int`
        """
        pass

    @abc.abstractmethod
    def sample(self, n=1):
        """
        Returns one or more samples from this probability distribution.

        :param int n: Number of samples to return.
        :return numpy.ndarray: An array containing samples from the
            distribution of shape ``(n, d)``, where ``d`` is the number of
            random variables.
        """
        pass

class ProductDistribution(Distribution):
    r"""
    Returns the Cartesian product of two distributions :math:`A` and
    :math:`B`, :math:`\Pr(A, B) = \Pr(A) \Pr(B)`.

    :param Distribution A: Distribution object representing :math:`A`.
    :param Distribution B: Distribution object representing :math:`B`.
    """
    def __init__(self, A, B):
        self.A = A
        self.B = B

    @property
    def n_rvs(self):
        return self.A.n_rvs + self.B.n_rvs

    def sample(self, n=1):
        A_sample = self.A.sample(n)
        B_sample = self.B.sample(n)
        return np.hstack((A_sample, B_sample))

_DEFAULT_RANGES = np.array([[0, 1]])
_DEFAULT_RANGES.flags.writeable = False # Prevent anyone from modifying the
                                        # default ranges.
class UniformDistribution(Distribution):
    """
    Uniform distribution on a given rectangular region.

    :param numpy.ndarray ranges: Array of shape ``(n_rvs, 2)``, where ``n_rvs``
        is the number of random variables, specifying the upper and lower limits
        for each variable.
    """

    def __init__(self, ranges=_DEFAULT_RANGES):
        if not isinstance(ranges, np.ndarray):
            ranges = np.array(ranges)

        if len(ranges.shape) == 1:
            ranges = ranges[np.newaxis, ...]

        self._ranges = ranges
        self._n_rvs = ranges.shape[0]
        self._delta = ranges[:, 1] - ranges[:, 0]

    @property
    def n_rvs(self):
        return self._n_rvs

    def sample(self, n=1):
        shape = (n, self._n_rvs)# if n == 1 else (self._n_rvs, n)
        z = np.random.random(shape)
        return self._ranges[:, 0] + z * self._delta

    def grad_log_pdf(self, var):
        # THIS IS NOT TECHNICALLY LEGIT; BCRB doesn't technically work with a
        # prior that doesn't go to 0 at its end points.  But we do it anyway.
        if var.shape[0] == 1:
            return 12/(self._delta)**2
        else:
            return np.zeros(var.shape)

class ConstantDistribution(Distribution):
    """
    Represents a determinstic variable; useful for combining with other
    distributions, marginalizing, etc.

    :param values: Shape ``(n,)`` array or list of values :math:`X_0` such that
        :math:`\Pr(X) = \delta(X - X_0)`.
    """

    def __init__(self, values):
        self._values = np.array(values)[np.newaxis, :]

    @property
    def n_rvs(self):
        return self._values.shape[1]

    def sample(self, n=1):
        return np.repeat(self._values, n, axis=0)

class UniformDistributionWith0(Distribution):
    """
    Uniform distribution on a given rectangular region with padded zeros.

    :param numpy.ndarray ranges: Array of shape ``(n_rvs, 2)``, where ``n_rvs``
        is the number of random variables, specifying the upper and lower limits
        for each variable.
    """

    def __init__(self, ranges=_DEFAULT_RANGES, zeros = 0):
        warnings.warn(
            "This class has been superceded by ProductDistribution and ConstantDistribution.",
            DeprecationWarning
        )

        if not isinstance(ranges, np.ndarray):
            ranges = np.array(ranges)

        if len(ranges.shape) == 1:
            ranges = ranges[np.newaxis, ...]

        self._ranges = ranges
        self._n_rvs = ranges.shape[0]
        self._delta = ranges[:, 1] - ranges[:, 0]

        self.zeros = zeros

    @property
    def n_rvs(self):
        return self._n_rvs

    def sample(self, n=1):
        shape = (n, self._n_rvs)# if n == 1 else (self._n_rvs, n)
        z = np.random.random(shape)
        foo =  self._ranges[:, 0] + z * self._delta
        return np.pad(foo,((0,0), (0, self.zeros)), mode = 'constant')

    def grad_log_pdf(self, var):
        # THIS IS NOT TECHNICALLY LEGIT; BCRB doesn't technically work with a
        # prior that doesn't go to 0 at its end points.  But we do it anyway.
        if var.shape[0] == 1:
            return 12/(self._delta)**2
        else:
            return np.zeros(var.shape)

class NormalDistribution(Distribution):
    def __init__(self, mean, var):
        self.mean = mean
        self.var = var

        self.dist = st.norm(mean, np.sqrt(var))

    @property
    def n_rvs(self):
        return 1

    def sample(self, n=1):
        return self.dist.rvs(size=n)[:, np.newaxis]

    def grad_log_pdf(self, x):
        return -(x - self.mean) / self.var

class MultivariateNormalDistribution(Distribution):
    def __init__(self, mean, cov):

        self.mean = mean
        self.cov = cov
        self.invcov = la.inv(cov)

    @property
    def n_rvs(self):
        return self.mean.shape[1]

    def sample(self):

        return np.dot(la.sqrtm(self.cov), np.random.randn(self.n_rvs)) + self.mean

    def grad_log_pdf(self, x):
        return -np.dot(self.invcov,(x - self.mean[0]))

class SlantedNormalDistribution(Distribution):
    """
    Uniform distribution on a given rectangular region.

    :param numpy.ndarray ranges: Array of shape ``(n_rvs, 2)``, where ``n_rvs``
        is the number of random variables, specifying the upper and lower limits
        for each variable.
    """

    def __init__(self, ranges=_DEFAULT_RANGES, weight=0.01):
        if not isinstance(ranges, np.ndarray):
            ranges = np.array(ranges)

        if len(ranges.shape) == 1:
            ranges = ranges[np.newaxis, ...]

        self._ranges = ranges
        self._n_rvs = ranges.shape[0]
        #self._delta = ranges[:, 1] - ranges[:, 0]
        self._weight = weight

    @property
    def n_rvs(self):
        return self._n_rvs

    def sample(self, n=1):
        shape = (n, self._n_rvs)# if n == 1 else (self._n_rvs, n)
        z = np.random.randn(n,self._n_rvs)
        return self._ranges[:, 0] +self._weight*z+np.random.rand(n)*self._ranges[:, 1];

class MVUniformDistribution(object):

    def __init__(self, dim = 6):
        self.dim = dim

    def sample(self, n = 1):
        return np.random.mtrand.dirichlet(np.ones(self.dim),n)

class DiscreteUniformDistribution(Distribution):
    def __init__(self, num_bits):
        self._num_bits = num_bits

    def sample(self, n=1):
        z = np.random.randint(2**self._num_bits,n)
        return z

# TODO: make the following into Distributions.
class HilbertSchmidtUniform(object):
    """
    Creates a new Hilber-Schmidt uniform prior on state space of dimension ``dim``.
    See e.g. [Mez06]_ and [Mis12]_.

    :param int dim: Dimension of the state space.
    """
    def __init__(self, dim=2):
        self.dim = dim
        self.paulis1Q = np.array([[[1,0],[0,1]],[[1,0],[0,-1]],[[0,-1j],[1j,0]],[[0,1],[1,0]]])

        self.paulis = self.make_Paulis(self.paulis1Q, 4)

    def sample(self):
        #Generate random unitary (see e.g. http://arxiv.org/abs/math-ph/0609050v2)
        g = (np.random.randn(self.dim,self.dim) + 1j*np.random.randn(self.dim,self.dim))/np.sqrt(2.0)
        q,r = la.qr(g)
        d = np.diag(r)

        ph = d/np.abs(d)
        ph = np.diag(ph)

        U = np.dot(q,ph)

        #Generate random matrix
        z = np.random.randn(self.dim,self.dim) + 1j*np.random.randn(self.dim,self.dim)

        rho = np.dot(np.dot(np.identity(self.dim)+U,np.dot(z,z.conj().transpose())),np.identity(self.dim)+U.conj().transpose())
        rho = rho/np.trace(rho)

        x = np.zeros([self.dim**2-1])
        for idx in xrange(self.dim**2-1):
            x[idx] = np.real(np.trace(np.dot(rho,self.paulis[idx+1])))

        return x

    def make_Paulis(self,paulis,d):
        if d == self.dim*2:
            return paulis
        else:
            temp = np.zeros([d**2,d,d],dtype='complex128')
            for idx in xrange(temp.shape[0]):
                temp[idx,:] = np.kron(paulis[np.trunc(idx/d)], self.paulis1Q[idx % 4])
            return self.make_Paulis(temp,d*2)

class HaarUniform(object):
    """
    Creates a new Haar uniform prior on state space of dimension ``dim``.

    :param int dim: Dimension of the state space.
    """
    def __init__(self,dim = 2):
        self.dim = dim

    def sample(self):
        #Generate random unitary (see e.g. http://arxiv.org/abs/math-ph/0609050v2)
        z = (np.random.randn(self.dim,self.dim) + 1j*np.random.randn(self.dim,self.dim))/np.sqrt(2.0)
        q,r = la.qr(z)
        d = np.diag(r)

        ph = d/np.abs(d)
        ph = np.diag(ph)

        U = np.dot(q,ph)

        #TODO: generalize this to general dimensions
        #Apply Haar random unitary to |0> state to get random pure state
        psi = np.dot(U,np.array([1,0]))
        z = np.real(np.dot(psi.conj(),np.dot(np.array([[1,0],[0,-1]]),psi)))
        y = np.real(np.dot(psi.conj(),np.dot(np.array([[0,-1j],[1j,0]]),psi)))
        x = np.real(np.dot(psi.conj(),np.dot(np.array([[0,1],[1,0]]),psi)))

        return np.array([x,y,z])

class GinibreUniform(object):
    """
    Creates a prior on state space of dimension dim according to the Ginibre
    ensemble with parameter ``k``.
    See e.g. [Mis12]_.

    :param int dim: Dimension of the state space.
    """
    def __init__(self,dim = 2, k = 2):
        self.dim = dim
        self.k = k

    def sample(self):
        #Generate random matrix
        z = np.random.randn(self.dim,self.k) + 1j*np.random.randn(self.dim,self.k)

        rho = np.dot(z,z.conj().transpose())
        rho = rho/np.trace(rho)

        z = np.real(np.trace(np.dot(rho,np.array([[1,0],[0,-1]]))))
        y = np.real(np.trace(np.dot(rho,np.array([[0,-1j],[1j,0]]))))
        x = np.real(np.trace(np.dot(rho,np.array([[0,1],[1,0]]))))

        return np.array([x,y,z])

class PostselectedDistribution(Distribution):
    """
    Postselects a distribution based on validity within a given model.
    """
    # TODO: rewrite LiuWestResampler in terms of this and a
    #       new MixtureDistribution.

    def __init__(self, distribution, model, maxiters=100):
        self._dist = distribution
        self._model = model

        self._maxiters = 100

    @property
    def n_rvs(self):
        return self._dist.n_rvs

    def sample(self, n=1):
        """
        Returns one or more samples from this probability distribution.

        :param int n: Number of samples to return.
        :return numpy.ndarray: An array containing samples from the
            distribution of shape ``(n, d)``, where ``d`` is the number of
            random variables.
        """
        samples = np.empty((n, self.n_rvs))
        idxs_to_sample = np.arange(n)

        iters = 0

        while idxs_to_sample.size and iters < self._maxiters:
            samples[idxs_to_sample] = self._dist.sample(len(idxs_to_sample))

            idxs_to_sample = idxs_to_sample[np.nonzero(np.logical_not(
                self._model.are_models_valid(samples[idxs_to_sample, :])
            ))[0]]

            iters += 1

        if idxs_to_sample.size:
            raise RuntimeError("Did not successfully postselect within {} iterations.".format(self._maxiters))

        return samples
<|endoftext|>


<commit_before># Copyright 2013-2014 MongoDB, Inc.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Receives documents from the oplog worker threads and indexes them
into the backend.

This file is a document manager for the Solr search engine, but the intent
is that this file can be used as an example to add on different backends.
To extend this to other systems, simply implement the exact same class and
replace the method definitions with API calls for the desired backend.
"""
import re
import json

from pysolr import Solr, SolrError

from mongo_connector import errors
from mongo_connector.compat import u
from mongo_connector.constants import (DEFAULT_COMMIT_INTERVAL,
                                       DEFAULT_MAX_BULK)
from mongo_connector.util import retry_until_ok
from mongo_connector.doc_managers import DocManagerBase, exception_wrapper
from mongo_connector.doc_managers.formatters import DocumentFlattener

# pysolr only has 1 exception: SolrError
wrap_exceptions = exception_wrapper({
    SolrError: errors.OperationFailed})

ADMIN_URL = 'admin/luke?show=schema&wt=json'

decoder = json.JSONDecoder()

class DocManager(DocManagerBase):
    """The DocManager class creates a connection to the backend engine and
    adds/removes documents, and in the case of rollback, searches for them.

    The reason for storing id/doc pairs as opposed to doc's is so that multiple
    updates to the same doc reflect the most up to date version as opposed to
    multiple, slightly different versions of a doc.
    """

    def __init__(self, url, auto_commit_interval=DEFAULT_COMMIT_INTERVAL,
                 unique_key='_id', chunk_size=DEFAULT_MAX_BULK, **kwargs):
        """Verify Solr URL and establish a connection.
        """
        self.solr = Solr(url)
        self.unique_key = unique_key
        # pysolr does things in milliseconds
        if auto_commit_interval is not None:
            self.auto_commit_interval = auto_commit_interval * 1000
        else:
            self.auto_commit_interval = None
        self.chunk_size = chunk_size
        self.field_list = []
        self._build_fields()
        self._formatter = DocumentFlattener()

    def _parse_fields(self, result, field_name):
        """ If Schema access, parse fields and build respective lists
        """
        field_list = []
        for key, value in result.get('schema', {}).get(field_name, {}).items():
            if key not in field_list:
                field_list.append(key)
        return field_list

    @wrap_exceptions
    def _build_fields(self):
        """ Builds a list of valid fields
        """
        declared_fields = self.solr._send_request('get', ADMIN_URL)
        result = decoder.decode(declared_fields)
        self.field_list = self._parse_fields(result, 'fields')

        # Build regular expressions to match dynamic fields.
        # dynamic field names may have exactly one wildcard, either at
        # the beginning or the end of the name
        self._dynamic_field_regexes = []
        for wc_pattern in self._parse_fields(result, 'dynamicFields'):
            if wc_pattern[0] == "*":
                self._dynamic_field_regexes.append(
                    re.compile(".*%s\Z" % wc_pattern[1:]))
            elif wc_pattern[-1] == "*":
                self._dynamic_field_regexes.append(
                    re.compile("\A%s.*" % wc_pattern[:-1]))

    def _clean_doc(self, doc):
        """Reformats the given document before insertion into Solr.

        This method reformats the document in the following ways:
          - removes extraneous fields that aren't defined in schema.xml
          - unwinds arrays in order to find and later flatten sub-documents
          - flattens the document so that there are no sub-documents, and every
            value is associated with its dot-separated path of keys

        An example:
          {"a": 2,
           "b": {
             "c": {
               "d": 5
             }
           },
           "e": [6, 7, 8]
          }

        becomes:
          {"a": 2, "b.c.d": 5, "e.0": 6, "e.1": 7, "e.2": 8}

        """

        # Translate the _id field to whatever unique key we're using.
        # _id may not exist in the doc, if we retrieved it from Solr
        # as part of update.
        if '_id' in doc:
            doc[self.unique_key] = str(doc.pop("_id"))

        # SOLR cannot index fields within sub-documents, so flatten documents
        # with the dot-separated path to each value as the respective key
        flat_doc = self._formatter.format_document(doc)

        # Only include fields that are explicitly provided in the
        # schema or match one of the dynamic field patterns, if
        # we were able to retrieve the schema
        if len(self.field_list) + len(self._dynamic_field_regexes) > 0:
            def include_field(field):
                return field in self.field_list or any(
                    regex.match(field) for regex in self._dynamic_field_regexes
                )
            return dict((k, v) for k, v in flat_doc.items() if include_field(k))
        return flat_doc

    def stop(self):
        """ Stops the instance
        """
        pass

    def apply_update(self, doc, update_spec):
        """Override DocManagerBase.apply_update to have flat documents."""
        # Replace a whole document
        if not '$set' in update_spec and not '$unset' in update_spec:
            # update spec contains the new document
            update_spec['_ts'] = doc['_ts']
            update_spec['ns'] = doc['ns']
            update_spec['_id'] = doc['_id']
            return update_spec
        for to_set in update_spec.get("$set", []):
            value = update_spec['$set'][to_set]
            # Find dotted-path to the value, remove that key from doc, then
            # put value at key:
            keys_to_pop = []
            for key in doc:
                if key.startswith(to_set):
                    if key == to_set or key[len(to_set)] == '.':
                        keys_to_pop.append(key)
            for key in keys_to_pop:
                doc.pop(key)
            doc[to_set] = value
        for to_unset in update_spec.get("$unset", []):
            # MongoDB < 2.5.2 reports $unset for fields that don't exist within
            # the document being updated.
            keys_to_pop = []
            for key in doc:
                if key.startswith(to_unset):
                    if key == to_unset or key[len(to_unset)] == '.':
                        keys_to_pop.append(key)
            for key in keys_to_pop:
                doc.pop(key)
        return doc

    @wrap_exceptions
    def update(self, doc, update_spec):
        """Apply updates given in update_spec to the document whose id
        matches that of doc.

        """
        # Commit outstanding changes so that the document to be updated is the
        # same version to which the changes apply.
        self.commit()
        query = "%s:%s" % (self.unique_key, u(doc['_id']))
        results = self.solr.search(query)
        if not len(results):
            # Document may not be retrievable yet
            self.commit()
            results = self.solr.search(query)
        # Results is an iterable containing only 1 result
        for doc in results:
            updated = self.apply_update(doc, update_spec)
            # A _version_ of 0 will always apply the update
            updated['_version_'] = 0
            self.upsert(updated)
            return updated

    @wrap_exceptions
    def upsert(self, doc):
        """Update or insert a document into Solr

        This method should call whatever add/insert/update method exists for
        the backend engine and add the document in there. The input will
        always be one mongo document, represented as a Python dictionary.
        """
        if self.auto_commit_interval is not None:
            self.solr.add([self._clean_doc(doc)],
                          commit=(self.auto_commit_interval == 0),
                          commitWithin=u(self.auto_commit_interval))
        else:
            self.solr.add([self._clean_doc(doc)], commit=False)

    @wrap_exceptions
    def bulk_upsert(self, docs):
        """Update or insert multiple documents into Solr

        docs may be any iterable
        """
        if self.auto_commit_interval is not None:
            add_kwargs = {
                "commit": (self.auto_commit_interval == 0),
                "commitWithin": str(self.auto_commit_interval)
            }
        else:
            add_kwargs = {"commit": False}

        cleaned = (self._clean_doc(d) for d in docs)
        if self.chunk_size > 0:
            batch = list(next(cleaned) for i in range(self.chunk_size))
            while batch:
                self.solr.add(batch, **add_kwargs)
                batch = list(next(cleaned)
                             for i in range(self.chunk_size))
        else:
            self.solr.add(cleaned, **add_kwargs)

    @wrap_exceptions
    def remove(self, doc):
        """Removes documents from Solr

        The input is a python dictionary that represents a mongo document.
        """
        self.solr.delete(id=u(doc["_id"]),
                         commit=(self.auto_commit_interval == 0))

    @wrap_exceptions
    def _remove(self):
        """Removes everything
        """
        self.solr.delete(q='*:*', commit=(self.auto_commit_interval == 0))

    @wrap_exceptions
    def _stream_search(self, query):
        """Helper method for iterating over Solr search results."""
        for doc in self.solr.search(query, rows=100000000):
            if self.unique_key != "_id":
                doc["_id"] = doc.pop(self.unique_key)
            yield doc

    @wrap_exceptions
    def search(self, start_ts, end_ts):
        """Called to query Solr for documents in a time range."""
        query = '_ts: [%s TO %s]' % (start_ts, end_ts)
        return self._stream_search(query)

    @wrap_exceptions
    def _search(self, query):
        """For test purposes only. Performs search on Solr with given query
            Does not have to be implemented.
        """
        return self._stream_search(query)

    def commit(self):
        """This function is used to force a commit.
        """
        retry_until_ok(self.solr.commit)

    @wrap_exceptions
    def get_last_doc(self):
        """Returns the last document stored in the Solr engine.
        """
        #search everything, sort by descending timestamp, return 1 row
        try:
            result = self.solr.search('*:*', sort='_ts desc', rows=1)
        except ValueError:
            return None

        for r in result:
            r['_id'] = r.pop(self.unique_key)
            return r
<commit_msg>Use u function to make _id unicode in Solr DocManager.<commit_after># Copyright 2013-2014 MongoDB, Inc.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Receives documents from the oplog worker threads and indexes them
into the backend.

This file is a document manager for the Solr search engine, but the intent
is that this file can be used as an example to add on different backends.
To extend this to other systems, simply implement the exact same class and
replace the method definitions with API calls for the desired backend.
"""
import re
import json

from pysolr import Solr, SolrError

from mongo_connector import errors
from mongo_connector.compat import u
from mongo_connector.constants import (DEFAULT_COMMIT_INTERVAL,
                                       DEFAULT_MAX_BULK)
from mongo_connector.util import retry_until_ok
from mongo_connector.doc_managers import DocManagerBase, exception_wrapper
from mongo_connector.doc_managers.formatters import DocumentFlattener

# pysolr only has 1 exception: SolrError
wrap_exceptions = exception_wrapper({
    SolrError: errors.OperationFailed})

ADMIN_URL = 'admin/luke?show=schema&wt=json'

decoder = json.JSONDecoder()

class DocManager(DocManagerBase):
    """The DocManager class creates a connection to the backend engine and
    adds/removes documents, and in the case of rollback, searches for them.

    The reason for storing id/doc pairs as opposed to doc's is so that multiple
    updates to the same doc reflect the most up to date version as opposed to
    multiple, slightly different versions of a doc.
    """

    def __init__(self, url, auto_commit_interval=DEFAULT_COMMIT_INTERVAL,
                 unique_key='_id', chunk_size=DEFAULT_MAX_BULK, **kwargs):
        """Verify Solr URL and establish a connection.
        """
        self.solr = Solr(url)
        self.unique_key = unique_key
        # pysolr does things in milliseconds
        if auto_commit_interval is not None:
            self.auto_commit_interval = auto_commit_interval * 1000
        else:
            self.auto_commit_interval = None
        self.chunk_size = chunk_size
        self.field_list = []
        self._build_fields()
        self._formatter = DocumentFlattener()

    def _parse_fields(self, result, field_name):
        """ If Schema access, parse fields and build respective lists
        """
        field_list = []
        for key, value in result.get('schema', {}).get(field_name, {}).items():
            if key not in field_list:
                field_list.append(key)
        return field_list

    @wrap_exceptions
    def _build_fields(self):
        """ Builds a list of valid fields
        """
        declared_fields = self.solr._send_request('get', ADMIN_URL)
        result = decoder.decode(declared_fields)
        self.field_list = self._parse_fields(result, 'fields')

        # Build regular expressions to match dynamic fields.
        # dynamic field names may have exactly one wildcard, either at
        # the beginning or the end of the name
        self._dynamic_field_regexes = []
        for wc_pattern in self._parse_fields(result, 'dynamicFields'):
            if wc_pattern[0] == "*":
                self._dynamic_field_regexes.append(
                    re.compile(".*%s\Z" % wc_pattern[1:]))
            elif wc_pattern[-1] == "*":
                self._dynamic_field_regexes.append(
                    re.compile("\A%s.*" % wc_pattern[:-1]))

    def _clean_doc(self, doc):
        """Reformats the given document before insertion into Solr.

        This method reformats the document in the following ways:
          - removes extraneous fields that aren't defined in schema.xml
          - unwinds arrays in order to find and later flatten sub-documents
          - flattens the document so that there are no sub-documents, and every
            value is associated with its dot-separated path of keys

        An example:
          {"a": 2,
           "b": {
             "c": {
               "d": 5
             }
           },
           "e": [6, 7, 8]
          }

        becomes:
          {"a": 2, "b.c.d": 5, "e.0": 6, "e.1": 7, "e.2": 8}

        """

        # Translate the _id field to whatever unique key we're using.
        # _id may not exist in the doc, if we retrieved it from Solr
        # as part of update.
        if '_id' in doc:
            doc[self.unique_key] = u(doc.pop("_id"))

        # SOLR cannot index fields within sub-documents, so flatten documents
        # with the dot-separated path to each value as the respective key
        flat_doc = self._formatter.format_document(doc)

        # Only include fields that are explicitly provided in the
        # schema or match one of the dynamic field patterns, if
        # we were able to retrieve the schema
        if len(self.field_list) + len(self._dynamic_field_regexes) > 0:
            def include_field(field):
                return field in self.field_list or any(
                    regex.match(field) for regex in self._dynamic_field_regexes
                )
            return dict((k, v) for k, v in flat_doc.items() if include_field(k))
        return flat_doc

    def stop(self):
        """ Stops the instance
        """
        pass

    def apply_update(self, doc, update_spec):
        """Override DocManagerBase.apply_update to have flat documents."""
        # Replace a whole document
        if not '$set' in update_spec and not '$unset' in update_spec:
            # update spec contains the new document
            update_spec['_ts'] = doc['_ts']
            update_spec['ns'] = doc['ns']
            update_spec['_id'] = doc['_id']
            return update_spec
        for to_set in update_spec.get("$set", []):
            value = update_spec['$set'][to_set]
            # Find dotted-path to the value, remove that key from doc, then
            # put value at key:
            keys_to_pop = []
            for key in doc:
                if key.startswith(to_set):
                    if key == to_set or key[len(to_set)] == '.':
                        keys_to_pop.append(key)
            for key in keys_to_pop:
                doc.pop(key)
            doc[to_set] = value
        for to_unset in update_spec.get("$unset", []):
            # MongoDB < 2.5.2 reports $unset for fields that don't exist within
            # the document being updated.
            keys_to_pop = []
            for key in doc:
                if key.startswith(to_unset):
                    if key == to_unset or key[len(to_unset)] == '.':
                        keys_to_pop.append(key)
            for key in keys_to_pop:
                doc.pop(key)
        return doc

    @wrap_exceptions
    def update(self, doc, update_spec):
        """Apply updates given in update_spec to the document whose id
        matches that of doc.

        """
        # Commit outstanding changes so that the document to be updated is the
        # same version to which the changes apply.
        self.commit()
        query = "%s:%s" % (self.unique_key, u(doc['_id']))
        results = self.solr.search(query)
        if not len(results):
            # Document may not be retrievable yet
            self.commit()
            results = self.solr.search(query)
        # Results is an iterable containing only 1 result
        for doc in results:
            updated = self.apply_update(doc, update_spec)
            # A _version_ of 0 will always apply the update
            updated['_version_'] = 0
            self.upsert(updated)
            return updated

    @wrap_exceptions
    def upsert(self, doc):
        """Update or insert a document into Solr

        This method should call whatever add/insert/update method exists for
        the backend engine and add the document in there. The input will
        always be one mongo document, represented as a Python dictionary.
        """
        if self.auto_commit_interval is not None:
            self.solr.add([self._clean_doc(doc)],
                          commit=(self.auto_commit_interval == 0),
                          commitWithin=u(self.auto_commit_interval))
        else:
            self.solr.add([self._clean_doc(doc)], commit=False)

    @wrap_exceptions
    def bulk_upsert(self, docs):
        """Update or insert multiple documents into Solr

        docs may be any iterable
        """
        if self.auto_commit_interval is not None:
            add_kwargs = {
                "commit": (self.auto_commit_interval == 0),
                "commitWithin": str(self.auto_commit_interval)
            }
        else:
            add_kwargs = {"commit": False}

        cleaned = (self._clean_doc(d) for d in docs)
        if self.chunk_size > 0:
            batch = list(next(cleaned) for i in range(self.chunk_size))
            while batch:
                self.solr.add(batch, **add_kwargs)
                batch = list(next(cleaned)
                             for i in range(self.chunk_size))
        else:
            self.solr.add(cleaned, **add_kwargs)

    @wrap_exceptions
    def remove(self, doc):
        """Removes documents from Solr

        The input is a python dictionary that represents a mongo document.
        """
        self.solr.delete(id=u(doc["_id"]),
                         commit=(self.auto_commit_interval == 0))

    @wrap_exceptions
    def _remove(self):
        """Removes everything
        """
        self.solr.delete(q='*:*', commit=(self.auto_commit_interval == 0))

    @wrap_exceptions
    def _stream_search(self, query):
        """Helper method for iterating over Solr search results."""
        for doc in self.solr.search(query, rows=100000000):
            if self.unique_key != "_id":
                doc["_id"] = doc.pop(self.unique_key)
            yield doc

    @wrap_exceptions
    def search(self, start_ts, end_ts):
        """Called to query Solr for documents in a time range."""
        query = '_ts: [%s TO %s]' % (start_ts, end_ts)
        return self._stream_search(query)

    @wrap_exceptions
    def _search(self, query):
        """For test purposes only. Performs search on Solr with given query
            Does not have to be implemented.
        """
        return self._stream_search(query)

    def commit(self):
        """This function is used to force a commit.
        """
        retry_until_ok(self.solr.commit)

    @wrap_exceptions
    def get_last_doc(self):
        """Returns the last document stored in the Solr engine.
        """
        #search everything, sort by descending timestamp, return 1 row
        try:
            result = self.solr.search('*:*', sort='_ts desc', rows=1)
        except ValueError:
            return None

        for r in result:
            r['_id'] = r.pop(self.unique_key)
            return r
<|endoftext|>


<commit_before># $Id$

"""User Datagram Protocol."""

import dpkt

UDP_PORT_MAX	= 65535

class UDP(dpkt.Packet):
    __hdr__ = (
        ('sport', 'H', 0xdead),
        ('dport', 'H', 0),
        ('ulen', 'H', 8),
        ('sum', 'H', 0)
        )
<commit_msg>Added support for DNS in udp layer<commit_after># $Id: udp.py 23 2006-11-08 15:45:33Z dugsong $

"""User Datagram Protocol."""

import dpkt
import dns

UDP_PORT_MAX	= 65535

class UDP(dpkt.Packet):
    __hdr__ = (
        ('sport', 'H', 0xdead),
        ('dport', 'H', 0),
        ('ulen', 'H', 8),
        ('sum', 'H', 0)
        )

    def unpack(self, buf):
        dpkt.Packet.unpack(self, buf)
        if self.sport == 53 or self.dport == 53:
            # Assume it's DNS
            self.data = dns.DNS(self.data)
<|endoftext|>


<commit_before>#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""Utility functions for data releases."""
from __future__ import division, print_function, unicode_literals
import numpy as np
import socket
import os

def sphere_dist(lon1, lat1, lon2, lat2):
    """
    Haversine formula for angular distance on a sphere: more stable at poles.

    Inputs must be in DEGREES.

    Credit: https://github.com/astropy/astropy/pull/881
    (pull request wasn't merged at the time of writing this code,
    hence the function was copied here)
    """
    sdlat = np.sin(np.radians(lat2 - lat1) * 0.5)
    sdlon = np.sin(np.radians(lon2 - lon1) * 0.5)
    coslats = np.cos(np.radians(lat1)) * np.cos(np.radians(lat2))

    return np.degrees(2 * np.arcsin((sdlat**2 + coslats * sdlon**2) ** 0.5))

def sphere_dist_fast(lon1, lat1, lon2, lat2):
    """
    Euclidean angular distance "on a sphere" - only valid on sphere in the
    small-angle approximation.
    """
    # Bugfix: crossing meridian
    if isinstance(lon1, np.ndarray) and len(lon1) > 1:
        lon1[(lon1 - lon2) > 180] -= 360
    elif isinstance(lon2, np.ndarray) and len(lon2) > 1:
        lon2[lon1 - lon2 > 180] -= 360
    elif lon1 - lon2 > 180:
        lon1 -= 360

    dlat = lat2 - lat1
    dlon = (lon2 - lon1) * np.cos(np.radians(lat1 + lat2) * 0.5)

    return (dlat ** 2 + dlon ** 2) ** 0.5

def crossmatch(ra, dec, ra_array, dec_array, matchdist=0.5):
    """Returns the index of the matched source.

    ra/dec in degrees.
    matching distance in arcseconds.
    """
    precut = np.abs(dec - dec_array) < (matchdist/3600.)  # Optimized
    if not precut.any():
        return None
    dist = sphere_dist(ra, dec, ra_array[precut], dec_array[precut])
    idx_closest = dist.argmin()
    if dist[idx_closest] < (matchdist / 3600.):
        return np.argwhere(precut)[idx_closest]
    else:
        return None

def get_pid():
    """Returns the hostname and process identifier.

    Returns
    -------
    pid : string
        A string of the form "hostname/process_id".
    """
    pid = '{0}/{1}'.format(socket.gethostname(),
                           os.getpid())
<commit_msg>bugfix: get_pid did not return<commit_after>#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""Utility functions for data releases."""
from __future__ import division, print_function, unicode_literals
import numpy as np
import socket
import os

def sphere_dist(lon1, lat1, lon2, lat2):
    """
    Haversine formula for angular distance on a sphere: more stable at poles.

    Inputs must be in DEGREES.

    Credit: https://github.com/astropy/astropy/pull/881
    (pull request wasn't merged at the time of writing this code,
    hence the function was copied here)
    """
    sdlat = np.sin(np.radians(lat2 - lat1) * 0.5)
    sdlon = np.sin(np.radians(lon2 - lon1) * 0.5)
    coslats = np.cos(np.radians(lat1)) * np.cos(np.radians(lat2))

    return np.degrees(2 * np.arcsin((sdlat**2 + coslats * sdlon**2) ** 0.5))

def sphere_dist_fast(lon1, lat1, lon2, lat2):
    """
    Euclidean angular distance "on a sphere" - only valid on sphere in the
    small-angle approximation.
    """
    # Bugfix: crossing meridian
    if isinstance(lon1, np.ndarray) and len(lon1) > 1:
        lon1[(lon1 - lon2) > 180] -= 360
    elif isinstance(lon2, np.ndarray) and len(lon2) > 1:
        lon2[lon1 - lon2 > 180] -= 360
    elif lon1 - lon2 > 180:
        lon1 -= 360

    dlat = lat2 - lat1
    dlon = (lon2 - lon1) * np.cos(np.radians(lat1 + lat2) * 0.5)

    return (dlat ** 2 + dlon ** 2) ** 0.5

def crossmatch(ra, dec, ra_array, dec_array, matchdist=0.5):
    """Returns the index of the matched source.

    ra/dec in degrees.
    matching distance in arcseconds.
    """
    precut = np.abs(dec - dec_array) < (matchdist/3600.)  # Optimized
    if not precut.any():
        return None
    dist = sphere_dist(ra, dec, ra_array[precut], dec_array[precut])
    idx_closest = dist.argmin()
    if dist[idx_closest] < (matchdist / 3600.):
        return np.argwhere(precut)[idx_closest]
    else:
        return None

def get_pid():
    """Returns the hostname and process identifier.

    Returns
    -------
    pid : string
        A string of the form "hostname/process_id".
    """
    pid = '{0}/{1}'.format(socket.gethostname(),
                           os.getpid())
    return pid
<|endoftext|>