<commit_before>import shutil
import json
import os
import subprocess
import sys
#comment these two lines if you don't want to queue.
from rq import Connection, Queue
from redis import Redis
#Hello, World
#E.D.C.O.M EDCOM rocks and we love them!
preset="slow"
ac3=0
aac=0
vid=0

aacstr=0
ac3str=0
vidstr=0

filesto=[]
fil=sys.argv[1]

print(fil)

# Run FFProbe to get all of the available streams for any given file.
out=json.loads(subprocess.Popen(["/usr/bin/ffprobe","-v", "quiet", "-print_format", "json", "-show_format", "-show_streams",fil], stdout=subprocess.PIPE, stderr=subprocess.STDOUT).communicate()[0])["streams"]
print(out) # Print that
streams_audio=[]

#A counter to see how "good" the file is. If it's 3, then we don't need to do anything to the file.
good=0
streams_final=[]
subs_streams=[]
for stream in out:
	curstream={}
	if stream['codec_type']=='audio':
		if stream["codec_name"]=="ac3" or stream["codec_name"]=="dts" or stream["codec_name"]=="dca": # check for high quality/special codecs.
#																										This check is more complicated than it needs to be because some DVDs are really, really strange.
			curstream["type"]="audio"
			curstream["codec"]=stream["codec_name"]
			curstream["index"]=stream["index"]
			if "channel_layout" not in list(stream.keys()):
				stream["channel_layout"]="stereo"
			curstream["channel_layout"]=stream["channel_layout"]
			if stream["channels"] >= 5:
				curstream["channel_layout"]="surround"
				curstream["newcodec"]="copy"
				good=good+1
			else:
				curstream["channel_layout"]="stereo"
				curstream["newcodec"]="copy"
				good=good+1
		elif stream["codec_name"]=="aac":
			curstream["type"]="audio"
			curstream["codec"]=stream["codec_name"]
			curstream["index"]=stream["index"]
			curstream["channel_layout"]=stream["channel_layout"]
			if stream["channels"] == 6: # Surround AAC is strange, but oddly common in the warez scene.
				curstream["channel_layout"]="surround"
				curstream["newcodec"]="aac"
				curstream["channels"]="2"
			else:
				curstream["newcodec"]="copy"
				good=good+1
		else:
			curstream["type"]="audio"
			curstream["index"]=stream["index"]
			curstream["newcodec"]="aac"
			curstream["codec"]=stream["codec_name"]
			curstream["channel_layout"]=stream["channels"]
		streams_audio.append(curstream)
	elif stream["codec_type"]=="video": # Don't care unless it's already h.264.
		curstream["type"]="video"
		curstream["index"]=stream['index']
		if stream['codec_name']=="h264":
			curstream["newcodec"]= "copy"
			good=good+1
		else:
			curstream["newcodec"]="h264"
		streams_final.append(curstream)
	elif stream["codec_type"]=="subtitle":
		subs_streams.append(stream["index"])
if good==3:
	sys.exit(0)

# Potential for increasing if multi-streamed audio becomes prevalent. (DVDs/Blu-Ray)
aacStreams = [x for x in streams_audio if x["codec"] == "aac"]
ac3Streams = [x for x in streams_audio if x["codec"] == "ac3"]

if len(ac3Streams) == 0:
	ac3Streams = [x for x in streams_audio if x["codec"] == "dts"] # prefer AC3, otherwise do DTS/DCA. Most files will only have one.
if len(ac3Streams) == 0:
	ac3Streams = [x for x in streams_audio if x["codec"] == "dca"]

aac=0
ac3=0

if len(aacStreams)>0:
	aac=aacStreams[0]
if len(ac3Streams)>0:
	ac3=ac3Streams[0]
if len(ac3Streams)==0 or (ac3["channel_layout"]=="stereo" and aac!=0):
	#no AC3 stream. See if we can convert that strange surround AAC to AC3.
	try:
		if aac["channel_layout"] == "surround":
			newstr={}
			newstr["newcodec"]="ac3"
			newstr["channels"]="6"
			newstr["index"]=aac["index"]
			newstr["type"]="audio"
			newstr["codec"]="NewStream"
			ac3=newstr
	except:
		pass

if aac==0 and ac3==0:
	newstr={}
	newstr["newcodec"]="aac"
	newstr["channels"]="2"
	newstr["index"]=streams_audio[0]["index"]
	newstr["type"]="audio"
	newstr["codec"]="NewStream"
	aac=newstr

if aac==0:
	#no AAC stream. There must be an AC3 stream.
	newstr={}
	newstr["newcodec"]="aac"
	newstr["channels"]="2"
	newstr["index"]=ac3["index"]
	newstr["type"]="audio"
	newstr["codec"]="NewStream"
	aac=newstr

if aac != 0:
	streams_final.append(aac)
if ac3 !=0:
	streams_final.append(ac3)

#build an ffmpeg command.

numaudio=0
video=1
audio=0
ffmpeg=[]
for stream in streams_final:
	ffmpeg.append("-map")
	ffmpeg.append("0:"+str(stream["index"]))
	if stream["type"]=="video":
		if stream["newcodec"]=="copy":
			ffmpeg.append("-c:v")
			ffmpeg.append("copy")
			video=0
		elif stream["newcodec"]=="h264":
			ffmpeg.append("-c:v")
			ffmpeg.append("libx264")
			ffmpeg.append("-crf")
			ffmpeg.append("20") #adjust this up/down depending on your percieved quality.
			ffmpeg.append("-level:v")
			ffmpeg.append("4.1")
			ffmpeg.append("-preset")
			ffmpeg.append(preset)
			ffmpeg.append("-bf") # these slow down encoding, but reduce file size pretty significantly.
			ffmpeg.append("16")
			ffmpeg.append("-b_strategy")
			ffmpeg.append("2")
			ffmpeg.append("-subq")
			ffmpeg.append("10")
		else:
			print("Unknown codec:"+stream["newcodec"])
			sys.exit(1)
	elif stream["type"]=="audio":
		if stream["newcodec"]=="aac":
			ffmpeg.append("-c:a:"+str(numaudio))
			ffmpeg.append("libfdk_aac")
			ffmpeg.append("-ac:a:"+str(numaudio))
			ffmpeg.append("2")
			ffmpeg.append("-vbr")
			ffmpeg.append("5")
			#ffmpeg.append("-b:a:"+str(numaudio))
			#ffmpeg.append("320k")
			audio=1
			numaudio=numaudio+1
		elif stream["newcodec"]=="ac3":
			ffmpeg.append("-c:a:"+str(numaudio))
			ffmpeg.append("ac3")
			ffmpeg.append("-ac:a:"+str(numaudio))
			ffmpeg.append("6")
			ffmpeg.append("-b:a:"+str(numaudio))
			ffmpeg.append("640k")
			audio=1
			numaudio=numaudio+1
		elif stream["newcodec"]=="copy":
			ffmpeg.append("-c:a:"+str(numaudio))
			ffmpeg.append("copy")
			numaudio=numaudio+1
		else:
			print("Unknown codec:"+stream["newcodec"])
ffmpeg.append("-movflags")
ffmpeg.append("faststart")
job={}
job["path"]=fil
job["opts"]=ffmpeg

head,tail=os.path.split(fil)

if fil[-4:]==".mpg" or fil[-4:]==".vob": # These are usually OTA  recordings or DVD rips, which are in 1080i. There isn't a good way to test this.
	print("Deinterlacing!")
	ffmpeg.append("-vf")
	ffmpeg.append("yadif=0:-1:0")

if len(subs_streams) > 0:
	ffmpeg.append("-scodec")
	ffmpeg.append("mov_text")
	for stream in subs_streams:
		ffmpeg.append("-map")
		ffmpeg.append("0:"+str(stream))

#uncomment these next few lines if you want to just run ffmpeg.

#res=subprocess.call(["ffmpeg","-i",job["path"]]+job["opts"]+["out.mp4"])
#if res != 0:
#	print "FFMPEG WENT BAD!"
#else:
#	shutil.move("out.mp4",job["path"]+".mp4")
#	shutil.remove(job["path"])
# exit()

#Delete the rest of the file if you don't want to enqueue.

# enqueue the file.
redis_conn = Redis()
if video==0 and audio==0:
	if fil[-4:]!=".mp4":
		q = Queue('mux-core',connection=redis_conn)
		job["opts"]=["-acodec","copy","-vcodec","copy"]
		q.enqueue_call('tasks.ffmpeg',args=(job,),timeout=360000)
		print('Enqueued: '+ fil+" for Remux"+str(job))

if video==1:
	q = Queue('video-core',connection=redis_conn)
	q.enqueue_call('tasks.ffmpeg',args=(job,),timeout=360000)
	print('Enqueued: '+ fil+" for Video"+str(job))
elif audio!=0:
	q = Queue('audio-core',connection=redis_conn)
	q.enqueue_call('tasks.ffmpeg',args=(job,),timeout=360000)
	print('Enqueued: '+ fil+" for Audio"+str(job))
<commit_msg>Fixed issue #1<commit_after>import shutil
import json
import os
import subprocess
import sys
#comment these two lines if you don't want to queue.
from rq import Connection, Queue
from redis import Redis
#Hello, World
#E.D.C.O.M EDCOM rocks and we love them!
preset="slow"
ac3=0
aac=0
vid=0

aacstr=0
ac3str=0
vidstr=0

filesto=[]
fil=sys.argv[1]

print(fil)

# Run FFProbe to get all of the available streams for any given file.
out=json.loads(subprocess.Popen(["/usr/bin/ffprobe","-v", "quiet", "-print_format", "json", "-show_format", "-show_streams",fil], stdout=subprocess.PIPE, stderr=subprocess.STDOUT).communicate()[0])["streams"]
print(out) # Print that
streams_audio=[]

#A counter to see how "good" the file is. If it's 3, then we don't need to do anything to the file.
good=0
streams_final=[]
subs_streams=[]
for stream in out:
	curstream={}
	if stream['codec_type']=='audio':
		if stream["codec_name"]=="ac3" or stream["codec_name"]=="dts" or stream["codec_name"]=="dca": # check for high quality/special codecs.
#																										This check is more complicated than it needs to be because some DVDs are really, really strange.
			curstream["type"]="audio"
			curstream["codec"]=stream["codec_name"]
			curstream["index"]=stream["index"]
			if "channel_layout" not in list(stream.keys()):
				stream["channel_layout"]="stereo"
			curstream["channel_layout"]=stream["channel_layout"]
			if stream["channels"] >= 5:
				curstream["channel_layout"]="surround"
				curstream["newcodec"]="copy"
				good=good+1
			else:
				curstream["channel_layout"]="stereo"
				curstream["newcodec"]="copy"
				good=good+1
		elif stream["codec_name"]=="aac":
			curstream["type"]="audio"
			curstream["codec"]=stream["codec_name"]
			curstream["index"]=stream["index"]
			curstream["channel_layout"]=stream["channel_layout"]
			if stream["channels"] == 6: # Surround AAC is strange, but oddly common in the warez scene.
				curstream["channel_layout"]="surround"
				curstream["newcodec"]="aac"
				curstream["channels"]="2"
			else:
				curstream["newcodec"]="copy"
				good=good+1
		else:
			curstream["type"]="audio"
			curstream["index"]=stream["index"]
			curstream["newcodec"]="aac"
			curstream["codec"]=stream["codec_name"]
			curstream["channel_layout"]=stream["channels"]
		streams_audio.append(curstream)
	elif stream["codec_type"]=="video": # Don't care unless it's already h.264.
		curstream["type"]="video"
		curstream["index"]=stream['index']
		if stream['codec_name']=="h264":
			curstream["newcodec"]= "copy"
			good=good+1
		else:
			curstream["newcodec"]="h264"
		streams_final.append(curstream)
	elif stream["codec_type"]=="subtitle":
		subs_streams.append(stream["index"])
if good==3:
	sys.exit(0)

# Potential for increasing if multi-streamed audio becomes prevalent. (DVDs/Blu-Ray)
aacStreams = [x for x in streams_audio if x["codec"] == "aac"]
ac3Streams = [x for x in streams_audio if x["codec"] == "ac3"]

if len(ac3Streams) == 0:
	ac3Streams = [x for x in streams_audio if x["codec"] == "dts"] # prefer AC3, otherwise do DTS/DCA. Most files will only have one.
if len(ac3Streams) == 0:
	ac3Streams = [x for x in streams_audio if x["codec"] == "dca"]

aac=0
ac3=0

if len(aacStreams)>0:
	aac=aacStreams[0]
if len(ac3Streams)>0:
	ac3=ac3Streams[0]
if len(ac3Streams)==0 or (ac3["channel_layout"]=="stereo" and aac!=0):
	#no AC3 stream. See if we can convert that strange surround AAC to AC3.
	try:
		if aac["channel_layout"] == "surround":
			newstr={}
			newstr["newcodec"]="ac3"
			newstr["channels"]="6"
			newstr["index"]=aac["index"]
			newstr["type"]="audio"
			newstr["codec"]="NewStream"
			ac3=newstr
	except:
		pass

if aac==0 and ac3==0:
	newstr={}
	newstr["newcodec"]="aac"
	newstr["channels"]="2"
	newstr["index"]=streams_audio[0]["index"]
	newstr["type"]="audio"
	newstr["codec"]="NewStream"
	aac=newstr

if aac==0:
	#no AAC stream. There must be an AC3 stream.
	newstr={}
	newstr["newcodec"]="aac"
	newstr["channels"]="2"
	newstr["index"]=ac3["index"]
	newstr["type"]="audio"
	newstr["codec"]="NewStream"
	aac=newstr

if aac != 0:
	streams_final.append(aac)
if ac3 !=0:
	streams_final.append(ac3)

#build an ffmpeg command.

numaudio=0
video=1
audio=0
ffmpeg=[]
for stream in streams_final:
	ffmpeg.append("-map")
	ffmpeg.append("0:"+str(stream["index"]))
	if stream["type"]=="video":
		if stream["newcodec"]=="copy":
			ffmpeg.append("-c:v")
			ffmpeg.append("copy")
			video=0
		elif stream["newcodec"]=="h264":
			ffmpeg.append("-c:v")
			ffmpeg.append("libx264")
			ffmpeg.append("-crf")
			ffmpeg.append("20") #adjust this up/down depending on your percieved quality.
			ffmpeg.append("-level:v")
			ffmpeg.append("4.1")
			ffmpeg.append("-preset")
			ffmpeg.append(preset)
			ffmpeg.append("-bf") # these slow down encoding, but reduce file size pretty significantly.
			ffmpeg.append("16")
			ffmpeg.append("-b_strategy")
			ffmpeg.append("2")
			ffmpeg.append("-subq")
			ffmpeg.append("10")
		else:
			print("Unknown codec:"+stream["newcodec"])
			sys.exit(1)
	elif stream["type"]=="audio":
		if stream["newcodec"]=="aac":
			ffmpeg.append("-c:a:"+str(numaudio))
			ffmpeg.append("libfdk_aac")
			ffmpeg.append("-ac:a:"+str(numaudio))
			ffmpeg.append("2")
			ffmpeg.append("-vbr")
			ffmpeg.append("5")
			#ffmpeg.append("-b:a:"+str(numaudio))
			#ffmpeg.append("320k")
			audio=1
			numaudio=numaudio+1
		elif stream["newcodec"]=="ac3":
			ffmpeg.append("-c:a:"+str(numaudio))
			ffmpeg.append("ac3")
			ffmpeg.append("-ac:a:"+str(numaudio))
			ffmpeg.append("6")
			ffmpeg.append("-b:a:"+str(numaudio))
			ffmpeg.append("640k")
			audio=1
			numaudio=numaudio+1
		elif stream["newcodec"]=="copy":
			ffmpeg.append("-c:a:"+str(numaudio))
			ffmpeg.append("copy")
			numaudio=numaudio+1
		else:
			print("Unknown codec:"+stream["newcodec"])
ffmpeg.append("-movflags")
ffmpeg.append("faststart")
job={}
job["path"]=fil
job["opts"]=ffmpeg

head,tail=os.path.split(fil)

if fil[-4:]==".mpg" or fil[-4:]==".vob": # These are usually OTA  recordings or DVD rips, which are in 1080i. There isn't a good way to test this.
	print("Deinterlacing!")
	ffmpeg.append("-vf")
	ffmpeg.append("yadif=0:-1:0")

if len(subs_streams) > 0:
	ffmpeg.append("-scodec")
	ffmpeg.append("mov_text")
	for stream in subs_streams:
		ffmpeg.append("-map")
		ffmpeg.append("0:"+str(stream))

#uncomment these next few lines if you want to just run ffmpeg.

#res=subprocess.call(["ffmpeg","-i",job["path"]]+job["opts"]+["out.mp4"])
#if res != 0:
#	print "FFMPEG FAILURE!"
#else:
#	shutil.move("out.mp4",job["path"]+".mp4")
#	os.remove(job["path"])
# exit()

#Delete the rest of the file if you don't want to enqueue.

# enqueue the file.
redis_conn = Redis()
if video==0 and audio==0:
	if fil[-4:]!=".mp4":
		q = Queue('mux-core',connection=redis_conn)
		job["opts"]=["-acodec","copy","-vcodec","copy"]
		q.enqueue_call('tasks.ffmpeg',args=(job,),timeout=360000)
		print('Enqueued: '+ fil+" for Remux"+str(job))

if video==1:
	q = Queue('video-core',connection=redis_conn)
	q.enqueue_call('tasks.ffmpeg',args=(job,),timeout=360000)
	print('Enqueued: '+ fil+" for Video"+str(job))
elif audio!=0:
	q = Queue('audio-core',connection=redis_conn)
	q.enqueue_call('tasks.ffmpeg',args=(job,),timeout=360000)
	print('Enqueued: '+ fil+" for Audio"+str(job))
<|endoftext|>


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

# (C) 2004 British Broadcasting Corporation and Kamaelia Contributors(1)
#     All Rights Reserved.
#
# You may only modify and redistribute this under the terms of any of the
# following licenses(2): Mozilla Public License, V1.1, GNU General
# Public License, V2.0, GNU Lesser General Public License, V2.1
#
# (1) Kamaelia Contributors are listed in the AUTHORS file and at
#     http://kamaelia.sourceforge.net/AUTHORS - please extend this file,
#     not this notice.
# (2) Reproduced in the COPYING file, and at:
#     http://kamaelia.sourceforge.net/COPYING
# Under section 3.5 of the MPL, we are using this text since we deem the MPL
# notice inappropriate for this file. As per MPL/GPL/LGPL removal of this
# notice is prohibited.
#
# Please contact us via: kamaelia-list-owner@lists.sourceforge.net
# to discuss alternative licensing.
# -------------------------------------------------------------------------

# physics code for forces between particles
#
# unbonded force acts between all non bonded particles
# bonded force acts between bonded particles

from SpatialIndexer import SpatialIndexer

from operator import sub as _sub
from operator import add as _add
from operator import mul as _mul

class ParticleSystem(object):
    """System of particles.

    Maintains the set of particles and runs the physics simulation over them
    the specified laws.
    """

    def __init__(self, laws, initialParticles = [], initialTick = 0):
        """Initialise the particle system.
           laws             = laws object
           initialParticles = list of particles
           initialTick      = start value for tick counter
        """
        self.indexer = SpatialIndexer(laws.maxInteractRadius)

        self.laws         = laws
        self.particles    = []
        self.tick         = initialTick
        self.particleDict = {}
        self.add(*initialParticles)

    def add(self, *newParticles):
        """Add the specified particle(s) into the system"""
        self.particles.extend(newParticles)
        for p in newParticles:
           self.particleDict[p.ID] = p
        self.indexer.updateLoc(*newParticles)

    def remove(self, *oldParticles):
        """Remove the specified particle(s) from the system.
           Note that this method does not destroy bonds from other particles to these ones.
        """
        for particle in oldParticles:
            self.particles.remove(particle)
            del self.particleDict[particle.ID]
        self.indexer.remove(*oldParticles)

    def removeByID(self, *ids):
        """Remove particle(s) as specified by id(s) from the system.
           Note that this method does not destroy bonds from other particles to these ones.
        """
        particles = [self.particleDict[id] for id in ids]
        self.remove( *particles )

    def updateLoc(self, *particles):
        """Notify this physics system that the specified particle(s)
           have changed position.

           Must be called if you change a particle's position,
           before calling run().
        """
        self.indexer.updateLoc(*particles)

    def withinRadius(self, centre, radius, filter=(lambda particle:True)):
        """Returns a list of zero or more (particle, distSquared) tuples,
           representing those particles within radius distance of the
           specified centre coords.

           distance-squared from the centre coords is returned too to negate
           any need you may have to calculate it again yourself.

           You can specify a filter function that takes a candidate particle
           as an argument and should return True if it is to be included
           (if it is within the radius, of course). This is to allow efficient
           pre-filtering of the particles before the distance test is done.
        """
        return self.indexer.withinRadius(centre, radius, filter)

    def run(self, cycles = 1):
        """Run the simulation for a given number of cycles"""
        _indexer = self.indexer
        _laws    = self.laws
        while cycles > 0:
            cycles -= 1
            self.tick += 1
            _tick = self.tick
            for p in self.particles:
                p.doInteractions(_indexer, _laws, _tick)
            for p in self.particles:
                p.update(_laws)
        _indexer.updateAll()

class Particle(object):
    """Particle within a physics system with an arbitrary number of dimensions.

    Represents a particle that interacts with other particles. One set of forces are applied for
    those particles that are unbonded. Interactions between bonded particles are controlled by another
    set of forces.

    Bonds are bi-directional. Establishing a bond from A to B, will also establish it back from B to A.
    Similarly, breaking the bond will do so in both directions too.
    """

    def __init__(self, position, initialTick = 0, ptype = None, velocity = None, ID = None):
        """Initialise particle.
           position    = tuple of coordinates
           initialTick = initial value for tick counter (should be same as for particle system)
           ptype       = particle type identifier
           velocity    = tuple of initial velocity vectors
           ID          = unique identifier
        """
        self.pos    = position
        self.tick   = initialTick
        self.static = False
        self.ptype  = ptype
        self.bondedTo   = []
        self.bondedFrom = []
        self.latestInteractWith = (None,None)

        if velocity != None:
            self.velocity = list(velocity)
        else:
            self.velocity = [ 0.0 for xx in self.pos ]
        if ID is None:
           self.ID = str(id(self))
        else:
           self.ID = ID

    def getBonded(self):
        """Return list of particles this one is bonded to (outgoing bonds)."""
        return self.bondedTo

    getBondedTo = getBonded

    def getBondedFrom(self):
        """Return list of particles that bond to this one (incoming bonds)."""
        return self.bondedFrom

    def makeBond(self,particles, index):
       """Make bond between this particle and another
       Specified by its index into the 'particles' entity you provide.

       If the bond already exists, then this method does nothing.
       """

       target = particles[index]
       if not target in self.bondedTo:
           self.bondedTo += [particles[index]]
           particles[index].bondedFrom += [self]

    def breakBond(self, particles, index):
        """Break bond between this particle and another
        Specified by its index into the 'particles' entity you provide.

        If the bond doesnt already exist, this method will fail.
        """
        self.bondedTo.remove( particles[index] )
        particles[index].bondedFrom.remove(self)

    def breakAllBonds(self, outgoing=True, incoming=True):
        """Breaks all bonds between this particle to others.
          Default behaviour is to break all bonds both from this node
          to others (outgoing) and from others to this (incoming).

          If outgoing or incoming are set to false, then bonds in those
          directions will not be broken
        """
        if outgoing:
            for bondTo in self.bondedTo:
                bondTo.bondedFrom.remove(self)
            self.bondedTo = []

        if incoming:
            for bondFrom in self.bondedFrom:
                bondFrom.bondedTo.remove(self)
            self.bondedFrom = []

    def getLoc(self):
        """Return current possition"""
        return self.pos

    def freeze(self):
        """Lock the particle in place"""
        self.static = True

    def unFreeze(self):
        """Allow the particle to move freely"""
        self.static = False

    def distSquared(self, altpos):
        """Returns the distance squared of this particle from the specified position"""
        return sum(map(lambda x1,x2 : (x1-x2)*(x1-x2), self.pos, altpos))

    def doInteractions(self, particleIndex, laws, tick):
        """Apply laws in relation to this particle with respect to other particles,
        to update the velocity of this particle.

        particleIndex is an object with a withinRadius(centre, radius, filter) method that
        returns a list of (particles, distSquared) tuples, listing particles within that distance
        of the specified coordinates. The filter argument is a function that returns true if a
        given particle is to be included in the list.

        laws.maxInteractRadius is the max distance at which unbonded interactions are considered
        laws.unbonded(ptype, ptype, dist, distanceSquared) is the velocity change applied to both particles.
        +ve = attraction -ve = repulsion
        laws.bonded(dist, distanceSquared) is the same but for bonded particles

        Tick is the current tick counter value. Any particles this one encounters that already
        have reached 'tick' will not be interacted with since it will be assumed that
        that particle has already performed the interaction math.

        This code relies on any bonds being registered in both directions (a->b and b->a).
        If you override the bonding code in this class, then make sure you maintain this property.
        """
        self.tick = tick
        _bonded   = laws.bonded
        _unbonded = laws.unbonded
        __add, __sub, __mul = _add, _sub, _mul

        # bonded interactions with bonded particles
        bonded = self.getBondedTo() + self.getBondedFrom()
        for particle in bonded:
            if particle.tick != self.tick and particle.latestInteractWith != (self.tick, self):
                particle.latestInteractWith != (self.tick, self)
                ds = self.distSquared(particle.pos)
                if ds > 0.0:
                    dist = ds ** 0.5
                    dvelocity = _bonded(self.ptype, particle.ptype, dist, ds)
                    deltas = map(__sub, particle.pos, self.pos)
                    dv_d = dvelocity / dist
                    scaleddeltas = map(__mul, deltas, [dv_d]*len(deltas))
                    self.velocity     = map(__add, self.velocity,     scaleddeltas)
                    particle.velocity = map(__sub, particle.velocity, scaleddeltas)
#                    self.velocity     = map(lambda delta,v : v + (delta*dv_d), deltas, self.velocity)
#                    particle.velocity = map(lambda delta,v : v - (delta*dv_d), deltas, particle.velocity)
                else:
                    pass # dunno, ought to have an error i guess

        # repulsion of other particles (not self, or those bonded to)
        filter = lambda particle : (particle.tick != self.tick) and not (particle in (bonded + [self]))

        particles = particleIndex.withinRadius(self.pos, laws.particleMaxInteractRadius(self.ptype), filter)
        for (particle, ds) in particles:
            if ds > 0.0:
                dist = ds ** 0.5
                dvelocity   = _unbonded(self.ptype, particle.ptype, dist, ds)
                deltas = map(__sub, particle.pos, self.pos)
                dv_d = dvelocity / dist
                scaleddeltas = map(__mul, deltas, [dv_d]*len(deltas))
                self.velocity     = map(__add, self.velocity,     scaleddeltas)
                particle.velocity = map(__sub, particle.velocity, scaleddeltas)
#                self.velocity     = map(lambda delta,v : v+(+delta*dv_d), deltas, self.velocity)
#                particle.velocity = map(lambda delta,v : v+(-delta*dv_d), deltas, particle.velocity)
            else:
                pass # dunno, ought to have an error i guess

    def update(self, laws):
        """Update this particle's position, also apply dampening to velocity

        laws.dampening( ptype, velocity ) should return the new velocity, that is then applied.
        """
        if self.static:
            self.velocity = [0 for x in self.velocity]
        else:
            self.velocity = laws.dampening(self.ptype, self.velocity)
            self.pos      = map(_add, self.pos, self.velocity)

class SimpleLaws(object):
    """Implements a simple set of physics laws for the particle system.
       repulsion force : force proportional to 1/distance_squared
       bonds : force proportional to extension (Hook's law)

       All force strengths etc. are set up to vaguely sensible values on the basis of
       the specified bond length
    """
    def __init__(self, bondLength        = 100,
                       maxRepelRadius    = None,
                       repulsionStrength = None,
                       maxBondForce      = None,
                       damp              = None,
                       dampcutoff        = None,
                       maxVelocity       = None
                 ):
        """bondlength argument is the 'master control'
              maxInteractRadius is the maximum distance repulsions are considered over
              repulsionStrength = strength of repulsion (when at bondLength distance)
              maxBondForce      = force at distance 0 or 2*bondLength (ie. bond spring constant)
              damp              = dampening (drag) 0.0 = none, 1.0 = no movement at all!
              dampcutoff        = minimum velocity below which friction stops movement all together
              maxVelocity       = maximum velocity limit

           maxVelocity is a fiddle factor to help you stop the system spiralling out of control

           damp and dampcutoff are effectively wind resistance and friction

           bondLengthsInteractRadius gives the physics system a way to reduce the amount of

           computation needed because particles are greater distances are deemed to not interact
        """
        self.bondLength = bondLength

        scale = 100.0 / self.bondLength

        def defaultIfNone(value, default):
            if value == None:
                return default
            else:
                return value

        self.maxInteractRadius          = defaultIfNone(maxRepelRadius,    200  / scale)
        self.repulsionForceAtBondLength = defaultIfNone(repulsionStrength, 3.2  / scale)
        self.maxBondForce               = defaultIfNone(maxBondForce,      20.0 / scale)
        self.damp                       = 1.0 - defaultIfNone(damp, 0.2)
        self.dampcutoff                 = defaultIfNone(dampcutoff,        0.4  / scale)
        self.maxVelocity                = defaultIfNone(maxVelocity,       32   / scale)

        self.bondForce         = self.maxBondForce / self.bondLength
        self.maxRepulsionForce = self.repulsionForceAtBondLength * self.bondLength**2

    def particleMaxInteractRadius(self, ptype):
        return self.maxInteractRadius

    def unbonded(self, ptype1, ptype2, dist, distSquared):
        """1/distance_squared unbonded repulsion force"""
        if distSquared < 1.0:
            return -self.maxRepulsionForce
        else:
            return -self.maxRepulsionForce / distSquared #* (self.particleRadius*self.particleRadius)

    def bonded(self, ptype1, ptype2, dist, distSquared):
        """proportional to extension bond force"""
        # note, its import that this retains its sign, so the direction of the force is determined
        f = self.bondForce * (dist - self.bondLength)
        if f < -self.maxBondForce:
            return -self.maxBondForce
        elif f > self.maxBondForce:
            return self.maxBondForce
        else:
            return f

    def dampening(self, ptype, velocity):
        """velocity dampening and minimal velocity (friction-like) cutoff"""
        vmag = reduce(lambda a,b:abs(a)+abs(b), velocity)
        if vmag < self.dampcutoff:
            return [0] * len(velocity) #[0 for a in velocity]
        else:
            damp = self.damp
            if damp * vmag > self.maxVelocity:
                damp = self.maxVelocity / vmag

            return map(_mul, velocity, [damp] * len(velocity) )
#            return [ damp * v for v in velocity ]

class MultipleLaws(object):
    """Laws framework for systems containing multiple particle types
    """
    def __init__(self, typesToLaws, defaultLaw = None):
        """Initialisation.
        typesToLaws = dictionary mapping particle type name pairs (type1, type2) to laws

        if you supply a pairing (t1, t2) it will also be applied to the case (t2,t1) without
        you needing to explicitly specify it. If you do, then your choice takes precedence.

        If you do not provide enough mappings to build a complete mapping from all types to all types,
        then the gaps will be automatically filled with mappings to defaultLaw.
        """

        self.laws = {}

        types = []

        # build specified mapping, reversing the pairings if not specified
        for ((type1, type2), law) in typesToLaws.items():
            self.laws[(type1,type2)] = law
            if not (type2,type1) in typesToLaws.keys():
                self.laws[(type2,type1)] = law

            # build a list of the different types of particles
            if not type1 in types:
                types.append(type1)
            if not type2 in types:
                types.append(type2)

        # go through the built links and check all combinations exist
        for type1 in types:
            for type2 in types:
                if not self.laws.has_key((type1, type2)):
                    self.laws[(type1,type2)] = defaultLaw

        # determine the maxInteractRadius
        self.maxInteractRadius = max( [law.maxInteractRadius for law in self.laws.values()] )

    def particleMaxInteractRadius(self, ptype):
        return self.laws[(ptype,ptype)].maxInteractRadius

    def unbonded(self, ptype1, ptype2, dist, distSquared):
        law = self.laws[(ptype1, ptype2)]
        return law.unbonded(ptype1, ptype2, dist, distSquared)

    def bonded(self, ptype1, ptype2, dist, distSquared):
        law = self.laws[(ptype1, ptype2)]
        return law.bonded(ptype1, ptype2, dist, distSquared)

    def dampening(self, ptype, velocity):
        law = self.laws[(ptype, ptype)]
        return law.dampening(ptype, velocity)
<commit_msg>Split into constiuent files into a subdirectory. (Makes manually browsing/viewing the hierarchy simpler) Will CVS add the change shortly.<commit_after><|endoftext|>


<commit_before>import time
import os

from celery import shared_task

from django.conf import settings
from django.core.files.storage import FileSystemStorage
from django.core.files.base import ContentFile
from django.utils.translation import ugettext_lazy as _
from django.contrib import messages

from excerptconverter.baseexcerptconverter import BaseExcerptConverter

from osmaxx.excerptexport import models

private_storage = FileSystemStorage(location=settings.PRIVATE_MEDIA_ROOT)

class DummyExcerptConverter(BaseExcerptConverter):
    @staticmethod
    def name():
        return 'Dummy'

    @staticmethod
    def export_formats():
        return {
            'txt': {
                'name': 'Text (.txt)',
                'file_extension': 'txt',
                'mime_type': 'text/plain'
            },
            'markdown': {
                'name': 'Markdown (.md)',
                'file_extension': 'md',
                'mime_type': 'text/markdown'
            }
        }

    @staticmethod
    def export_options():
        return {
            'detail_level': {
                'label': 'Detail level',
                'type': 'choice',
                'default': 'verbatim',
                'values': [
                    {'name': 'verbatim', 'label': 'Verbatim'},
                    {'name': 'simplified', 'label': 'Simplified'}
                ]
            }
        }

    @staticmethod
    def create_output_files(execution_configuration, extraction_order, supported_export_formats):
        for format_key in execution_configuration['formats']:
            output_file = models.OutputFile.objects.create(
                mime_type=supported_export_formats[format_key]['mime_type'],
                extraction_order=extraction_order
            )

            if not os.path.exists(private_storage.location):
                os.makedirs(private_storage.location)

            file_name = str(output_file.public_identifier) + '.' + \
                supported_export_formats[format_key]['file_extension']
            file_content = ContentFile(str('detail level: ' + execution_configuration['options']['detail_level']))
            fs_file = private_storage.save(file_name, file_content)

            # file must be committed, so reopen to attach to model
            output_file.file = fs_file
            output_file.save()

            if private_storage.exists(file_name):
                message_text = _('"%s" created successful' % file_name)
                BaseExcerptConverter.inform_user(extraction_order.orderer, messages.SUCCESS, message_text, False)
            else:
                message_text = _('Creation of "%s" failed!' % file_name)
                BaseExcerptConverter.inform_user(extraction_order.orderer, messages.ERROR, message_text, False)

    @staticmethod
    @shared_task
    def execute_task(extraction_order_id, supported_export_formats, execution_configuration):
        wait_time = 0
        # wait for the db to be updated!
        extraction_order = None
        while extraction_order is None:
            try:
                extraction_order = models.ExtractionOrder.objects.get(pk=extraction_order_id)
            except models.ExtractionOrder.DoesNotExist:
                time.sleep(5)
                wait_time += 5
                if wait_time > 30:
                    raise

        fake_work_waiting_time_in_seconds = 5

        # now set the new state
        extraction_order.state = models.ExtractionOrderState.WAITING
        extraction_order.save()

        time.sleep(fake_work_waiting_time_in_seconds)

        # now set the new state
        extraction_order.state = models.ExtractionOrderState.PROCESSING
        extraction_order.save()

        message_text = _('Your extraction order "%s" has been started' % extraction_order)
        BaseExcerptConverter.inform_user(extraction_order.orderer, messages.INFO, message_text, False)

        DummyExcerptConverter.create_output_files(execution_configuration, extraction_order, supported_export_formats)

        time.sleep(fake_work_waiting_time_in_seconds)

        # now set the new state
        extraction_order.state = models.ExtractionOrderState.FINISHED
        extraction_order.save()

        # inform the user of the status change.
        message_text = _('Your extraction order %s has been processed' % extraction_order)
        BaseExcerptConverter.inform_user(extraction_order.orderer, messages.SUCCESS, message_text)
<commit_msg>Bugfix: prevent DummyExcerptConverter setting status and creating notifications without files to convert<commit_after>import time
import os

from celery import shared_task

from django.conf import settings
from django.core.files.storage import FileSystemStorage
from django.core.files.base import ContentFile
from django.utils.translation import ugettext_lazy as _
from django.contrib import messages

from excerptconverter.baseexcerptconverter import BaseExcerptConverter

from osmaxx.excerptexport import models

private_storage = FileSystemStorage(location=settings.PRIVATE_MEDIA_ROOT)

class DummyExcerptConverter(BaseExcerptConverter):
    @staticmethod
    def name():
        return 'Dummy'

    @staticmethod
    def export_formats():
        return {
            'txt': {
                'name': 'Text (.txt)',
                'file_extension': 'txt',
                'mime_type': 'text/plain'
            },
            'markdown': {
                'name': 'Markdown (.md)',
                'file_extension': 'md',
                'mime_type': 'text/markdown'
            }
        }

    @staticmethod
    def export_options():
        return {
            'detail_level': {
                'label': 'Detail level',
                'type': 'choice',
                'default': 'verbatim',
                'values': [
                    {'name': 'verbatim', 'label': 'Verbatim'},
                    {'name': 'simplified', 'label': 'Simplified'}
                ]
            }
        }

    @staticmethod
    def create_output_files(execution_configuration, extraction_order, supported_export_formats):
        for format_key in execution_configuration['formats']:
            output_file = models.OutputFile.objects.create(
                mime_type=supported_export_formats[format_key]['mime_type'],
                extraction_order=extraction_order
            )

            if not os.path.exists(private_storage.location):
                os.makedirs(private_storage.location)

            file_name = str(output_file.public_identifier) + '.' + \
                supported_export_formats[format_key]['file_extension']
            file_content = ContentFile(str('detail level: ' + execution_configuration['options']['detail_level']))
            fs_file = private_storage.save(file_name, file_content)

            # file must be committed, so reopen to attach to model
            output_file.file = fs_file
            output_file.save()

            if private_storage.exists(file_name):
                message_text = _('"%s" created successful' % file_name)
                BaseExcerptConverter.inform_user(extraction_order.orderer, messages.SUCCESS, message_text, False)
            else:
                message_text = _('Creation of "%s" failed!' % file_name)
                BaseExcerptConverter.inform_user(extraction_order.orderer, messages.ERROR, message_text, False)

    @staticmethod
    @shared_task
    def execute_task(extraction_order_id, supported_export_formats, execution_configuration):
        if len(execution_configuration['formats']) > 0:
            wait_time = 0
            # wait for the db to be updated!
            extraction_order = None
            while extraction_order is None:
                try:
                    extraction_order = models.ExtractionOrder.objects.get(pk=extraction_order_id)
                except models.ExtractionOrder.DoesNotExist:
                    time.sleep(5)
                    wait_time += 5
                    if wait_time > 30:
                        raise

            fake_work_waiting_time_in_seconds = 5

            # now set the new state
            extraction_order.state = models.ExtractionOrderState.WAITING
            extraction_order.save()

            time.sleep(fake_work_waiting_time_in_seconds)

            # now set the new state
            extraction_order.state = models.ExtractionOrderState.PROCESSING
            extraction_order.save()

            message_text = _('Your extraction order "%s" has been started' % extraction_order)
            BaseExcerptConverter.inform_user(extraction_order.orderer, messages.INFO, message_text, False)

            DummyExcerptConverter.create_output_files(
                execution_configuration,
                extraction_order,
                supported_export_formats
            )

            time.sleep(fake_work_waiting_time_in_seconds)

            # now set the new state
            extraction_order.state = models.ExtractionOrderState.FINISHED
            extraction_order.save()

            # inform the user of the status change.
            message_text = _('Your extraction order %s has been processed' % extraction_order)
            BaseExcerptConverter.inform_user(extraction_order.orderer, messages.SUCCESS, message_text)
<|endoftext|>


<commit_before>"""engine.SCons.Tool.f90

Tool-specific initialization for the generic Posix f90 Fortran compiler.

There normally shouldn't be any need to import this module directly.
It will usually be imported through the generic SCons.Tool.Tool()
selection method.

"""

#
# __COPYRIGHT__
#
# Permission is hereby granted, free of charge, to any person obtaining
# a copy of this software and associated documentation files (the
# "Software"), to deal in the Software without restriction, including
# without limitation the rights to use, copy, modify, merge, publish,
# distribute, sublicense, and/or sell copies of the Software, and to
# permit persons to whom the Software is furnished to do so, subject to
# the following conditions:
#
# The above copyright notice and this permission notice shall be included
# in all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY
# KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE
# WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
# NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE
# LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION
# OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION
# WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
#

__revision__ = "__FILE__ __REVISION__ __DATE__ __DEVELOPER__"

import SCons.Defaults
import SCons.Scanner.Fortran
import SCons.Tool
import SCons.Util
from SCons.Tool.FortranCommon import add_all_to_env, add_f90_to_env

compilers = ['f90']

def generate(env):
    add_all_to_env(env)
    add_f90_to_env(env)

    fc = env.Detect(compilers) or 'f90'
    env['F90']  = fcomp
    env['SHF90']  = fcomp

    env['FORTRAN']  = fcomp
    env['SHFORTRAN']  = fcomp

def exists(env):
    return env.Detect(compilers)
<commit_msg>Fix typo in f90.py tool.<commit_after>"""engine.SCons.Tool.f90

Tool-specific initialization for the generic Posix f90 Fortran compiler.

There normally shouldn't be any need to import this module directly.
It will usually be imported through the generic SCons.Tool.Tool()
selection method.

"""

#
# __COPYRIGHT__
#
# Permission is hereby granted, free of charge, to any person obtaining
# a copy of this software and associated documentation files (the
# "Software"), to deal in the Software without restriction, including
# without limitation the rights to use, copy, modify, merge, publish,
# distribute, sublicense, and/or sell copies of the Software, and to
# permit persons to whom the Software is furnished to do so, subject to
# the following conditions:
#
# The above copyright notice and this permission notice shall be included
# in all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY
# KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE
# WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
# NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE
# LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION
# OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION
# WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
#

__revision__ = "__FILE__ __REVISION__ __DATE__ __DEVELOPER__"

import SCons.Defaults
import SCons.Scanner.Fortran
import SCons.Tool
import SCons.Util
from SCons.Tool.FortranCommon import add_all_to_env, add_f90_to_env

compilers = ['f90']

def generate(env):
    add_all_to_env(env)
    add_f90_to_env(env)

    fc = env.Detect(compilers) or 'f90'
    env['F90']  = fc
    env['SHF90']  = fc

    env['FORTRAN']  = fc
    env['SHFORTRAN']  = fc

def exists(env):
    return env.Detect(compilers)
<|endoftext|>


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

import numpy as np
import pandas as pd

from utils.preprocessing import one_hot_encoding

def get_weekday(date):
    return date.weekday()

path = '../data/raw/'
train_users = pd.read_csv(path + 'train_users.csv', nrows=1000)
test_users = pd.read_csv(path + 'test_users.csv', nrows=1000)
sessions = pd.read_csv(path + 'sessions.csv', nrows=400)

# Join users
users = pd.concat((train_users, test_users), axis=0, ignore_index=True)

# Drop useless column
users = users.drop('date_first_booking', axis=1)

# Replace NaNs
users['gender'].replace('-unknown-', np.nan, inplace=True)
users['language'].replace('-unknown-', np.nan, inplace=True)
sessions.replace('-unknown-', np.nan, inplace=True)

# Remove weird age values
users.loc[users['age'] > 100, 'age'] = np.nan
users.loc[users['age'] < 14, 'age'] = np.nan

# List categorical features
categorical_features = [
    'affiliate_channel',
    'affiliate_provider',
    'country_destination',
    'first_affiliate_tracked',
    'first_browser',
    'first_device_type',
    'gender',
    'language',
    'signup_app',
    'signup_method'
]

# Change categorical features
for categorical_feature in categorical_features:
    users[categorical_feature] = users[categorical_feature].astype('category')

# Change type to date
users['date_account_created'] = pd.to_datetime(users['date_account_created'])
users['date_first_active'] = pd.to_datetime(users['timestamp_first_active'],
                                            format='%Y%m%d%H%M%S')

users['weekday_account_created'] = users[
    'date_account_created'].apply(get_weekday)
users['weekday_first_active'] = users['date_first_active'].apply(get_weekday)

# Split dates into day, month, year
year_account_created = pd.DatetimeIndex(users['date_account_created']).year
users['year_account_created'] = year_account_created
month_account_created = pd.DatetimeIndex(users['date_account_created']).month
users['month_account_created'] = month_account_created
day_account_created = pd.DatetimeIndex(users['date_account_created']).day
users['day_account_created'] = day_account_created
year_first_active = pd.DatetimeIndex(users['date_first_active']).year
users['year_first_active'] = year_first_active
month_first_active = pd.DatetimeIndex(users['date_first_active']).month
users['month_first_active'] = month_first_active
day_first_active = pd.DatetimeIndex(users['date_first_active']).day
users['day_first_active'] = day_first_active

user_sessions = pd.DataFrame()

for user in sessions['user_id'].unique():

    # Get the user session
    user_session = sessions.loc[sessions['user_id'] == user]

    # Length of the session
    users.loc[users['id'] == user, 'session_length'] = int(len(user_session))

    action_type = user_session.groupby('action')
    action_type_secs_elapsed = action_type.secs_elapsed.sum()
    action_type_secs_elapsed.rename(
        lambda x: x + '_secs_elapsed', inplace=True)
    action_type_secs_elapsed['id'] = user

    user_sessions = user_sessions.append(action_type_secs_elapsed)

    action_type = user_session.groupby('action_type')
    action_type_secs_elapsed = action_type.secs_elapsed.sum()
    action_type_secs_elapsed.rename(
        lambda x: x + '_secs_elapsed', inplace=True)
    action_type_secs_elapsed['id'] = user

    user_sessions = user_sessions.append(action_type_secs_elapsed)

    action_type = user_session.groupby('action_detail')
    action_type_secs_elapsed = action_type.secs_elapsed.sum()
    action_type_secs_elapsed.rename(
        lambda x: x + '_secs_elapsed', inplace=True)
    action_type_secs_elapsed['id'] = user

    user_sessions = user_sessions.append(action_type_secs_elapsed)

    action_type = user_session.groupby('device_type')
    action_type_secs_elapsed = action_type.secs_elapsed.sum()
    action_type_secs_elapsed.rename(
        lambda x: x + '_secs_elapsed', inplace=True)
    action_type_secs_elapsed['id'] = user

    user_sessions = user_sessions.append(action_type_secs_elapsed)

    # Get the most used device
    if user_session['device_type'].value_counts().sum() is not 0:
        most_used_device = user_session['device_type'].value_counts().index[0]
        users.loc[users['id'] == user, 'most_used_device'] = most_used_device

user_sessions = user_sessions.groupby('id').sum()

users = users.set_index('id')

users = pd.concat([users, user_sessions], axis=1)

# phones = [
#     'Opera Phone', 'Blackberry', 'Windows Phone',
#     'Android App Unknown Phone/Tablet'
# ]
# desktop = ['Mac Desktop', 'Windows Desktop', 'Linux Desktop', 'Chromebook']
# sessions.device_type.replace(phones, 'Mobile Phone').value_counts()

# Get the count of general session information
user_sessions = sessions.groupby('user_id')
general_session_info = user_sessions.count()
general_session_info.rename(columns=lambda x: x + '_count', inplace=True)
users = pd.concat([users, general_session_info], axis=1)

secs_elapsed_by_user = sessions.groupby('user_id')['secs_elapsed']

# Elapsed seconds sum
secs_elapsed_sum = secs_elapsed_by_user.sum()
secs_elapsed_sum.name = 'secs_elapsed_sum'
users = pd.concat([users, secs_elapsed_sum], axis=1)

# Elapsed seconds mean
secs_elapsed_average = secs_elapsed_by_user.mean()
secs_elapsed_average.name = 'secs_elapsed_average'
users = pd.concat([users, secs_elapsed_average], axis=1)

# Elapsed seconds min
min_secs_elapsed = secs_elapsed_by_user.min()
min_secs_elapsed.name = 'secs_elapsed_min'
users = pd.concat([users, min_secs_elapsed], axis=1)

# Elapsed seconds max
max_secs_elapsed = secs_elapsed_by_user.max()
max_secs_elapsed.name = 'secs_elapsed_max'
users = pd.concat([users, max_secs_elapsed], axis=1)

# Elapsed seconds first_quantile
first_quantile = secs_elapsed_by_user.quantile(0.25)
first_quantile.name = 'secs_elapsed_first_quantile'
users = pd.concat([users, first_quantile], axis=1)

# Elapsed seconds second_quantile
second_quantile = secs_elapsed_by_user.quantile(0.5)
second_quantile.name = 'secs_elapsed_second_quantile'
users = pd.concat([users, second_quantile], axis=1)

# Elapsed seconds third_quantile
third_quantile = secs_elapsed_by_user.quantile(0.75)
third_quantile.name = 'secs_elapsed_third_quantile'
users = pd.concat([users, third_quantile], axis=1)

# Elapsed seconds fourth_quantile
fourth_quantile = secs_elapsed_by_user.quantile(0.9)
fourth_quantile.name = 'secs_elapsed_fourth_quantile'
users = pd.concat([users, fourth_quantile], axis=1)

# Elapsed seconds median
median = secs_elapsed_by_user.median()
median.name = 'secs_elapsed_median'
users = pd.concat([users, median], axis=1)

# Elapsed seconds std
std = secs_elapsed_by_user.std()
std.name = 'secs_elapsed_std'
users = pd.concat([users, std], axis=1)

# Elapsed seconds var
var = secs_elapsed_by_user.var()
var.name = 'secs_elapsed_var'
users = pd.concat([users, var], axis=1)

# Elapsed seconds skew
skew = secs_elapsed_by_user.skew()
skew.name = 'secs_elapsed_skew'
users = pd.concat([users, skew], axis=1)

# First value
first_secs_elapsed = secs_elapsed_by_user.first()
first_secs_elapsed.name = 'first_secs_elapsed'
users = pd.concat([users, first_secs_elapsed], axis=1)

# Last value
last_secs_elapsed = secs_elapsed_by_user.last()
last_secs_elapsed.name = 'last_secs_elapsed'
users = pd.concat([users, last_secs_elapsed], axis=1)

# Number of elapsed seconds greater than 1 day
query = sessions['secs_elapsed'] > 86400
day_pauses = sessions.loc[query].groupby('user_id').count()['secs_elapsed']
day_pauses.name = 'day_pauses'
users = pd.concat([users, day_pauses], axis=1)

# Number of elapsed seconds lesser than 1 hour
query = sessions['secs_elapsed'] < 3600
short_sessions = sessions.loc[query].groupby('user_id').count()['secs_elapsed']
short_sessions.name = 'short_sessions'
users = pd.concat([users, short_sessions], axis=1)

# Users not returning in a big time
query = sessions['secs_elapsed'] > 300000
long_sessions = sessions.loc[query].groupby('user_id').count()['secs_elapsed']
long_sessions.name = 'long_sessions'
users = pd.concat([users, long_sessions], axis=1)

train_users = train_users.set_index('id')
test_users = test_users.set_index('id')

processed_train_users = users.loc[train_users.index]
processed_test_users = users.loc[test_users.index]
processed_test_users.drop('country_destination', inplace=True, axis=1)

path = '../data/processed/'
processed_train_users.to_csv(path + 'train_users_without_encoding.csv')
processed_test_users.to_csv(path + 'test_users_without_encoding.csv')

drop_list = [
    'date_account_created',
    'date_first_active',
    'timestamp_first_active'
]

# Drop columns
users = users.drop(drop_list, axis=1)

# Encode categorical features
categorical_features = [
    'gender', 'signup_method', 'signup_flow', 'language', 'affiliate_channel',
    'affiliate_provider', 'first_affiliate_tracked', 'signup_app',
    'first_device_type', 'first_browser', 'most_used_device'
]

users = one_hot_encoding(users, categorical_features)

# from sklearn.preprocessing import StandardScaler
# scaler = StandardScaler()
# scaler.fit_transform(users)

users.index.name = 'id'
processed_train_users = users.loc[train_users.index]
processed_test_users = users.loc[test_users.index]
processed_test_users.drop('country_destination', inplace=True, axis=1)

processed_train_users.to_csv(path + 'processed_train_users.csv')
processed_test_users.to_csv(path + 'processed_test_users.csv')
<commit_msg>Process full data<commit_after>#!/usr/bin/env python

import numpy as np
import pandas as pd

from utils.preprocessing import one_hot_encoding

def get_weekday(date):
    return date.weekday()

path = '../data/raw/'
train_users = pd.read_csv(path + 'train_users.csv')
test_users = pd.read_csv(path + 'test_users.csv')
sessions = pd.read_csv(path + 'sessions.csv')

# Join users
users = pd.concat((train_users, test_users), axis=0, ignore_index=True)

# Drop useless column
users = users.drop('date_first_booking', axis=1)

# Replace NaNs
users['gender'].replace('-unknown-', np.nan, inplace=True)
users['language'].replace('-unknown-', np.nan, inplace=True)
sessions.replace('-unknown-', np.nan, inplace=True)

# Remove weird age values
users.loc[users['age'] > 100, 'age'] = np.nan
users.loc[users['age'] < 14, 'age'] = np.nan

# List categorical features
categorical_features = [
    'affiliate_channel',
    'affiliate_provider',
    'country_destination',
    'first_affiliate_tracked',
    'first_browser',
    'first_device_type',
    'gender',
    'language',
    'signup_app',
    'signup_method'
]

# Change categorical features
for categorical_feature in categorical_features:
    users[categorical_feature] = users[categorical_feature].astype('category')

# Change type to date
users['date_account_created'] = pd.to_datetime(users['date_account_created'])
users['date_first_active'] = pd.to_datetime(users['timestamp_first_active'],
                                            format='%Y%m%d%H%M%S')

users['weekday_account_created'] = users[
    'date_account_created'].apply(get_weekday)
users['weekday_first_active'] = users['date_first_active'].apply(get_weekday)

# Split dates into day, month, year
year_account_created = pd.DatetimeIndex(users['date_account_created']).year
users['year_account_created'] = year_account_created
month_account_created = pd.DatetimeIndex(users['date_account_created']).month
users['month_account_created'] = month_account_created
day_account_created = pd.DatetimeIndex(users['date_account_created']).day
users['day_account_created'] = day_account_created
year_first_active = pd.DatetimeIndex(users['date_first_active']).year
users['year_first_active'] = year_first_active
month_first_active = pd.DatetimeIndex(users['date_first_active']).month
users['month_first_active'] = month_first_active
day_first_active = pd.DatetimeIndex(users['date_first_active']).day
users['day_first_active'] = day_first_active

user_sessions = pd.DataFrame()

for user in sessions['user_id'].unique():

    # Get the user session
    user_session = sessions.loc[sessions['user_id'] == user]

    # Length of the session
    users.loc[users['id'] == user, 'session_length'] = int(len(user_session))

    action_type = user_session.groupby('action')
    action_type_secs_elapsed = action_type.secs_elapsed.sum()
    action_type_secs_elapsed.rename(
        lambda x: x + '_secs_elapsed', inplace=True)
    action_type_secs_elapsed['id'] = user

    user_sessions = user_sessions.append(action_type_secs_elapsed)

    action_type = user_session.groupby('action_type')
    action_type_secs_elapsed = action_type.secs_elapsed.sum()
    action_type_secs_elapsed.rename(
        lambda x: x + '_secs_elapsed', inplace=True)
    action_type_secs_elapsed['id'] = user

    user_sessions = user_sessions.append(action_type_secs_elapsed)

    action_type = user_session.groupby('action_detail')
    action_type_secs_elapsed = action_type.secs_elapsed.sum()
    action_type_secs_elapsed.rename(
        lambda x: x + '_secs_elapsed', inplace=True)
    action_type_secs_elapsed['id'] = user

    user_sessions = user_sessions.append(action_type_secs_elapsed)

    action_type = user_session.groupby('device_type')
    action_type_secs_elapsed = action_type.secs_elapsed.sum()
    action_type_secs_elapsed.rename(
        lambda x: x + '_secs_elapsed', inplace=True)
    action_type_secs_elapsed['id'] = user

    user_sessions = user_sessions.append(action_type_secs_elapsed)

    # Get the most used device
    if user_session['device_type'].value_counts().sum() is not 0:
        most_used_device = user_session['device_type'].value_counts().index[0]
        users.loc[users['id'] == user, 'most_used_device'] = most_used_device

user_sessions = user_sessions.groupby('id').sum()

users = users.set_index('id')

users = pd.concat([users, user_sessions], axis=1)

# phones = [
#     'Opera Phone', 'Blackberry', 'Windows Phone',
#     'Android App Unknown Phone/Tablet'
# ]
# desktop = ['Mac Desktop', 'Windows Desktop', 'Linux Desktop', 'Chromebook']
# sessions.device_type.replace(phones, 'Mobile Phone').value_counts()

# Get the count of general session information
user_sessions = sessions.groupby('user_id')
general_session_info = user_sessions.count()
general_session_info.rename(columns=lambda x: x + '_count', inplace=True)
users = pd.concat([users, general_session_info], axis=1)

secs_elapsed_by_user = sessions.groupby('user_id')['secs_elapsed']

# Elapsed seconds sum
secs_elapsed_sum = secs_elapsed_by_user.sum()
secs_elapsed_sum.name = 'secs_elapsed_sum'
users = pd.concat([users, secs_elapsed_sum], axis=1)

# Elapsed seconds mean
secs_elapsed_average = secs_elapsed_by_user.mean()
secs_elapsed_average.name = 'secs_elapsed_average'
users = pd.concat([users, secs_elapsed_average], axis=1)

# Elapsed seconds min
min_secs_elapsed = secs_elapsed_by_user.min()
min_secs_elapsed.name = 'secs_elapsed_min'
users = pd.concat([users, min_secs_elapsed], axis=1)

# Elapsed seconds max
max_secs_elapsed = secs_elapsed_by_user.max()
max_secs_elapsed.name = 'secs_elapsed_max'
users = pd.concat([users, max_secs_elapsed], axis=1)

# Elapsed seconds first_quantile
first_quantile = secs_elapsed_by_user.quantile(0.25)
first_quantile.name = 'secs_elapsed_first_quantile'
users = pd.concat([users, first_quantile], axis=1)

# Elapsed seconds second_quantile
second_quantile = secs_elapsed_by_user.quantile(0.5)
second_quantile.name = 'secs_elapsed_second_quantile'
users = pd.concat([users, second_quantile], axis=1)

# Elapsed seconds third_quantile
third_quantile = secs_elapsed_by_user.quantile(0.75)
third_quantile.name = 'secs_elapsed_third_quantile'
users = pd.concat([users, third_quantile], axis=1)

# Elapsed seconds fourth_quantile
fourth_quantile = secs_elapsed_by_user.quantile(0.9)
fourth_quantile.name = 'secs_elapsed_fourth_quantile'
users = pd.concat([users, fourth_quantile], axis=1)

# Elapsed seconds median
median = secs_elapsed_by_user.median()
median.name = 'secs_elapsed_median'
users = pd.concat([users, median], axis=1)

# Elapsed seconds std
std = secs_elapsed_by_user.std()
std.name = 'secs_elapsed_std'
users = pd.concat([users, std], axis=1)

# Elapsed seconds var
var = secs_elapsed_by_user.var()
var.name = 'secs_elapsed_var'
users = pd.concat([users, var], axis=1)

# Elapsed seconds skew
skew = secs_elapsed_by_user.skew()
skew.name = 'secs_elapsed_skew'
users = pd.concat([users, skew], axis=1)

# First value
first_secs_elapsed = secs_elapsed_by_user.first()
first_secs_elapsed.name = 'first_secs_elapsed'
users = pd.concat([users, first_secs_elapsed], axis=1)

# Last value
last_secs_elapsed = secs_elapsed_by_user.last()
last_secs_elapsed.name = 'last_secs_elapsed'
users = pd.concat([users, last_secs_elapsed], axis=1)

# Number of elapsed seconds greater than 1 day
query = sessions['secs_elapsed'] > 86400
day_pauses = sessions.loc[query].groupby('user_id').count()['secs_elapsed']
day_pauses.name = 'day_pauses'
users = pd.concat([users, day_pauses], axis=1)

# Number of elapsed seconds lesser than 1 hour
query = sessions['secs_elapsed'] < 3600
short_sessions = sessions.loc[query].groupby('user_id').count()['secs_elapsed']
short_sessions.name = 'short_sessions'
users = pd.concat([users, short_sessions], axis=1)

# Users not returning in a big time
query = sessions['secs_elapsed'] > 300000
long_sessions = sessions.loc[query].groupby('user_id').count()['secs_elapsed']
long_sessions.name = 'long_sessions'
users = pd.concat([users, long_sessions], axis=1)

train_users = train_users.set_index('id')
test_users = test_users.set_index('id')

processed_train_users = users.loc[train_users.index]
processed_test_users = users.loc[test_users.index]
processed_test_users.drop('country_destination', inplace=True, axis=1)

path = '../data/processed/'
processed_train_users.to_csv(path + 'train_users_without_encoding.csv')
processed_test_users.to_csv(path + 'test_users_without_encoding.csv')

drop_list = [
    'date_account_created',
    'date_first_active',
    'timestamp_first_active'
]

# Drop columns
users = users.drop(drop_list, axis=1)

# Encode categorical features
categorical_features = [
    'gender', 'signup_method', 'signup_flow', 'language', 'affiliate_channel',
    'affiliate_provider', 'first_affiliate_tracked', 'signup_app',
    'first_device_type', 'first_browser', 'most_used_device'
]

users = one_hot_encoding(users, categorical_features)

# from sklearn.preprocessing import StandardScaler
# scaler = StandardScaler()
# scaler.fit_transform(users)

users.index.name = 'id'
processed_train_users = users.loc[train_users.index]
processed_test_users = users.loc[test_users.index]
processed_test_users.drop('country_destination', inplace=True, axis=1)

processed_train_users.to_csv(path + 'processed_train_users.csv')
processed_test_users.to_csv(path + 'processed_test_users.csv')
<|endoftext|>


<commit_before># Copyright 2016 Raytheon BBN Technologies
#
# 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

from QGL import *

from auspex.exp_factory import QubitExpFactory
from auspex.analysis.io import load_from_HDF5

def calibrate(qubit, calibrations):
    """Takes in a qubit (as a string) and list of calibrations (as instantiated classes).
    e.g. calibrate_pulses("q1", [RabiAmp, PhaseEstimation])"""
    for calibration in calibrations:
        if calibration not isinstance(PulseCalibration):
            raise TypeError("calibrate_pulses was passed a calibration that is not actually a calibration.")
        calibration.calibrate()

class PulseCalibration(object):
    """Base class for calibration of qubit control pulses."""
    def __init__(self, qubit_name):
        super(PulseCalibration, self).__init__()
        self.qubit_name = qubit
        self.qubit      = QubitFactory(qubit)
        self.filename   = 'None'

    def sequence(self):
        """Returns the sequence for the given calibration, must be overridden"""
        return [[Id(self.qubit), MEAS(self.qubit)]]

    def run(self):
        meta_file = compile_to_hardware(self.sequence(), fileName=self.filename)
        exp = QubitExpFactory.create(meta_file=meta_file)
        exp.run_sweeps()

        filename  = exp.qubit_to_writer[self.qubit_name].filename.value
        groupname = exp.qubit_to_writer[self.qubit_name].groupname.value
        dataset, descriptor = load_from_HDF5(filename, groupname=groupname)
        # TODO: get the name of the relevant data from the graph
        data, var = dataset['data']['M1'], dataset['data']['M1-var']

        # Return data and variance of the mean
        return data, var/descriptor.metadata["num_averages"]

    def calibrate(self):
        """Runs the actual calibration routine, must be overridden"""
        pass

class RabiAmpCalibration(PulseCalibration):
    def __init__(self, qubit_name, amps=np.linspace(0.0, 1.0, 51)):
        super(RabiAmpCalibration, self).__init__(qubit_name)
        self.amps = amps

    def sequence(self):
        return [[Xtheta(self.qubit, amp=a), MEAS(self.qubit)] for a in self.amps]

class RamseyCalibration(PulseCalibration):
    def __init__(self, qubit_name, delays=np.linspace(0.0, 1.0, 51)):
        super(RamseyCalibration, self).__init__(qubit_name)
        self.delays = delays

    def sequence(self):
        return [[X90(self.qubit), Id(self.qubit, delay), MEAS(self.qubit)] for delay in self.delays]

class PhaseEstimation(PulseCalibration):
    """Estimates pulse rotation angle from a sequence of P^k experiments, where
    k is of the form 2^n. Uses the modified phase estimation algorithm from
    Kimmel et al, quant-ph/1502.02677 (2015). Every experiment i doubled.
    vardata should be the variance of the mean"""

    def __init__(self, qubit_name, num_pulses='1', amplitude='0.1'):
        """Phase estimation calibration. Direction is either 'X' or 'Y',
        num_pulses is log2(n) of the longest sequence n,
        and amplitude is self-exaplanatory."""

        super(PhaseEstimation, self).__init__(qubit_name)
        self.filename        = 'RepeatCal/RepeatCal'
        self.direction       = direction
        self.amplitude       = amplitude
        self.num_pulses      = num_pulses
        self.target          = np.pi/2.0
        self.iteration_limit = 5

    def sequence(self):
        # Exponentially growing repetitions of the target pulse, e.g.
        # (1, 2, 4, 8, 16, 32, 64, 128, ...) x X90
        seqs = [[Xtheta(self.qubit, amp=self.amplitude)]*n for n in 2**np.arange(self.num_pulses+1)]
        # measure each along Z or Y
        seqs = [s + m for s in seqs for m in [ [MEAS(q)], [X90m(self.qubit), MEAS(q)] ]]
        # tack on calibrations to the beginning
        seqs = [[Id(q), MEAS(q)], [X(q), MEAS(q)]] + seqs
        # repeat each
        return [copy(s) for s in seqs for _ in range(2)]

    def optimize(update_data):
        """Attempts to optimize the pulse amplitude for a pi/2 or pi pulse about X or Y. """

        ct = 1
        amp = self.amplitude
        while True:
            [phase, sigma] = phase_estimation(*self.run())
            print("Phase: %.4f Sigma: %.4f"%(phase,sigma))
            # correct for some errors related to 2pi uncertainties
            if np.sign(phase) != np.sign(amp):
                phase += np.sign(amp)*2*np.pi
            angle_error = phase - self.target;
            print('Angle error: %.4f'%angle_error);

            amp_target = self.target/phase * amp
            amp_error = amp - amp_target
            print('Amplitude error: %.4f\n'%amp_error)

            amp = amp_target
            ct += 1

            # check for stopping condition
            phase_error = phase - self.target
            if np.abs(phase_error) < 1e-2 or np.abs(phase_error/sigma) < 1 or ct > self.iteration_limit:
                if np.abs(phase_error) < 1e-2:
                    print('Reached target rotation angle accuracy');
                elif abs(phase_error/sigma) < 1:
                    print('Reached phase uncertainty limit');
                else:
                    print('Hit max iteration count');
                break
        print('Amp',amp)

        return amp

class Pi2Calibration(PhaseEstimation):
    def __init__(self, qubit_name, num_pulses='1'):
        super(Pi2Calibration, self).__init__(qubit_name)
        self.amplitude = self.qubit.pulseParams['pi2Amp']
        self.target    = np.pi/2.0

class PiCalibration(PulseCalibration):
    def __init__(self, qubit_name, num_pulses='1'):
        super(PiCalibration, self).__init__(qubit_name)
        self.amplitude = self.qubit.pulseParams['piAmp']
        self.target    = np.pi

class DRAGCalibration(PulseCalibration):
    pass

def restrict(phase):
    out = np.mod( phase + np.pi, 2*np.pi, ) - np.pi
    return out

def phase_estimation( data_in, vardata_in, verbose=False):
    """Estimates pulse rotation angle from a sequence of P^k experiments, where
    k is of the form 2^n. Uses the modified phase estimation algorithm from
    Kimmel et al, quant-ph/1502.02677 (2015). Every experiment i doubled.
    vardata should be the variance of the mean"""

    #average together pairs of data points
    avgdata = (data_in[0::2] + data_in[1::2])/2

    # normalize data using the first two pulses to calibrate the "meter"
    data = 1 + 2*(avgdata[2:] - avgdata[0]) / (avgdata[0] - avgdata[1])
    zdata = data[0::2]
    xdata = data[1::2]

    # similar scaling with variances
    vardata = (vardata_in[0::2] + vardata_in[1::2])/2
    vardata = vardata[2:] * 2 / abs(avgdata[0] - avgdata[1])**2
    zvar = vardata[0::2]
    xvar = vardata[1::2]

    phases = np.arctan2(xdata, zdata)
    distances = np.sqrt(xdata**2 + zdata**2)

    curGuess = phases[0]
    phase = curGuess
    sigma = np.pi

    if verbose == True:
        print('Current Guess: %f'%(curGuess))

    for k in range(1,len(phases)):

        if verbose == True:
            print('k: %d'%(k))

        # Each step of phase estimation needs to assign the measured phase to
        # the correct half circle. We will conservatively require that the
        # (x,z) tuple is long enough that we can assign it to the correct
        # quadrant of the circle with 2σ confidence

        if distances[k] < 2*np.sqrt(xvar[k] + zvar[k]):
            print('Phase estimation terminated at %dth pulse because the (x,z) vector is too short'%(k))
            break

        lowerBound = restrict(curGuess - np.pi/2**(k))
        upperBound = restrict(curGuess + np.pi/2**(k))
        possiblesTest = [ restrict((phases[k] + 2*n*np.pi)/2**(k)) for n in range(0,2**(k)+1)]

        if verbose == True:
            print('Lower Bound: %f'%lowerBound)
            print('Upper Bound: %f'%upperBound)

        possibles=[]
        for p in possiblesTest:
            # NOTE: previous code did not handle upperbound == lowerBound
            if lowerBound >= upperBound:
                satisfiesLB = p > lowerBound or p < 0.
                satisfiesUP = p < upperBound or p > 0.
            else:
                satisfiesLB = p > lowerBound
                satisfiesUP = p < upperBound

            if satisfiesLB == True and satisfiesUP == True:
                possibles.append(p)

        curGuess = possibles[0]
        if verbose == True:
            print('Current Guess: %f'%(curGuess))
    phase = curGuess
    sigma = np.maximum(np.abs(restrict(curGuess - lowerBound)), np.abs(restrict(curGuess - upperBound)))

    return phase, sigma

calibrate("q1", [Pi2Cal(direction="Y"),
                 PiCal()])
<commit_msg>Some initial debugging<commit_after># Copyright 2016 Raytheon BBN Technologies
#
# 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

from QGL import *
from QGL import config as QGLconfig
import auspex.config as config
from copy import copy
import os

from auspex.exp_factory import QubitExpFactory
from auspex.analysis.io import load_from_HDF5

def calibrate(qubit, calibrations):
    """Takes in a qubit (as a string) and list of calibrations (as instantiated classes).
    e.g. calibrate_pulses("q1", [RabiAmp, PhaseEstimation])"""
    for calibration in calibrations:
        if not isinstance(calibration, PulseCalibration):
            raise TypeError("calibrate_pulses was passed a calibration that is not actually a calibration.")
        calibration.calibrate()

class PulseCalibration(object):
    """Base class for calibration of qubit control pulses."""
    def __init__(self, qubit_name):
        super(PulseCalibration, self).__init__()
        self.qubit_name = qubit
        self.qubit      = QubitFactory(qubit)
        self.filename   = 'None'

    def sequence(self):
        """Returns the sequence for the given calibration, must be overridden"""
        return [[Id(self.qubit), MEAS(self.qubit)]]

    def run(self):
        seq_files = compile_to_hardware(self.sequence(), fileName=self.filename)
        metafileName = os.path.join(QGLconfig.AWGDir, self.filename + '-meta.json')
        exp = QubitExpFactory.create(meta_file=metafileName)
        exp.run_sweeps()
        # TODO: there should be no need for saving the calibration data
        wrs = [w for w in exp.writers if w.name == exp.qubit_to_writer[self.qubit_name]]
        filename = wrs[0].filename.value
        groupname = wrs[0].groupname.value

        dataset, descriptor = load_from_HDF5(filename, groupname=groupname)
        # TODO: get the name of the relevant data from the graph
        data, var = dataset['data']['M1'], dataset['data']['M1-var']

        # Return data and variance of the mean
        return data, var/descriptor.metadata["num_averages"]

    def calibrate(self):
        """Runs the actual calibration routine, must be overridden"""
        pass

class RabiAmpCalibration(PulseCalibration):
    def __init__(self, qubit_name, amps=np.linspace(0.0, 1.0, 51)):
        super(RabiAmpCalibration, self).__init__(qubit_name)
        self.amps = amps

    def sequence(self):
        return [[Xtheta(self.qubit, amp=a), MEAS(self.qubit)] for a in self.amps]

class RamseyCalibration(PulseCalibration):
    def __init__(self, qubit_name, delays=np.linspace(0.0, 1.0, 51)):
        super(RamseyCalibration, self).__init__(qubit_name)
        self.delays = delays

    def sequence(self):
        return [[X90(self.qubit), Id(self.qubit, delay), MEAS(self.qubit)] for delay in self.delays]

class PhaseEstimation(PulseCalibration):
    """Estimates pulse rotation angle from a sequence of P^k experiments, where
    k is of the form 2^n. Uses the modified phase estimation algorithm from
    Kimmel et al, quant-ph/1502.02677 (2015). Every experiment i doubled.
    vardata should be the variance of the mean"""

    def __init__(self, qubit_name, num_pulses= 1, amplitude= 0.1, direction = 'X'):
        """Phase estimation calibration. Direction is either 'X' or 'Y',
        num_pulses is log2(n) of the longest sequence n,
        and amplitude is self-exaplanatory."""

        super(PhaseEstimation, self).__init__(qubit_name)
        self.filename        = 'RepeatCal/RepeatCal'
        self.direction       = direction
        self.amplitude       = amplitude
        self.num_pulses      = num_pulses
        self.target          = np.pi/2.0
        self.iteration_limit = 5

    def sequence(self):
        # Exponentially growing repetitions of the target pulse, e.g.
        # (1, 2, 4, 8, 16, 32, 64, 128, ...) x X90
        seqs = [[Xtheta(self.qubit, amp=self.amplitude)]*n for n in 2**np.arange(self.num_pulses+1)]
        # measure each along Z or Y
        seqs = [s + m for s in seqs for m in [ [MEAS(self.qubit)], [X90m(self.qubit), MEAS(self.qubit)] ]]
        # tack on calibrations to the beginning
        seqs = [[Id(self.qubit), MEAS(self.qubit)], [X(self.qubit), MEAS(self.qubit)]] + seqs
        # repeat each
        return [copy(s) for s in seqs for _ in range(2)]

    def calibrate(self):
        """Attempts to optimize the pulse amplitude for a pi/2 or pi pulse about X or Y. """

        ct = 1
        amp = self.amplitude
        while True:
            [phase, sigma] = phase_estimation(*self.run())
            print("Phase: %.4f Sigma: %.4f"%(phase,sigma))
            # correct for some errors related to 2pi uncertainties
            if np.sign(phase) != np.sign(amp):
                phase += np.sign(amp)*2*np.pi
            angle_error = phase - self.target;
            print('Angle error: %.4f'%angle_error);

            amp_target = self.target/phase * amp
            amp_error = amp - amp_target
            print('Amplitude error: %.4f\n'%amp_error)

            amp = amp_target
            ct += 1

            # check for stopping condition
            phase_error = phase - self.target
            if np.abs(phase_error) < 1e-2 or np.abs(phase_error/sigma) < 1 or ct > self.iteration_limit:
                if np.abs(phase_error) < 1e-2:
                    print('Reached target rotation angle accuracy');
                elif abs(phase_error/sigma) < 1:
                    print('Reached phase uncertainty limit');
                else:
                    print('Hit max iteration count');
                break
        print('Amp',amp)

        return amp

class Pi2Calibration(PhaseEstimation):
    def __init__(self, qubit_name, num_pulses= 9):
        super(Pi2Calibration, self).__init__(qubit_name)
        self.amplitude = self.qubit.pulseParams['pi2Amp']
        self.target    = np.pi/2.0

class PiCalibration(PulseCalibration):
    def __init__(self, qubit_name, num_pulses= 9):
        super(PiCalibration, self).__init__(qubit_name)
        self.amplitude = self.qubit.pulseParams['piAmp']
        self.target    = np.pi

class DRAGCalibration(PulseCalibration):
    pass

def restrict(phase):
    out = np.mod( phase + np.pi, 2*np.pi, ) - np.pi
    return out

def phase_estimation( data_in, vardata_in, verbose=False):
    """Estimates pulse rotation angle from a sequence of P^k experiments, where
    k is of the form 2^n. Uses the modified phase estimation algorithm from
    Kimmel et al, quant-ph/1502.02677 (2015). Every experiment i doubled.
    vardata should be the variance of the mean"""

    #average together pairs of data points
    avgdata = (data_in[0::2] + data_in[1::2])/2

    # normalize data using the first two pulses to calibrate the "meter"
    data = 1 + 2*(avgdata[2:] - avgdata[0]) / (avgdata[0] - avgdata[1])
    zdata = data[0::2]
    xdata = data[1::2]

    # similar scaling with variances
    vardata = (vardata_in[0::2] + vardata_in[1::2])/2
    vardata = vardata[2:] * 2 / abs(avgdata[0] - avgdata[1])**2
    zvar = vardata[0::2]
    xvar = vardata[1::2]

    phases = np.arctan2(xdata, zdata)
    distances = np.sqrt(xdata**2 + zdata**2)

    curGuess = phases[0]
    phase = curGuess
    sigma = np.pi

    if verbose == True:
        print('Current Guess: %f'%(curGuess))

    for k in range(1,len(phases)):

        if verbose == True:
            print('k: %d'%(k))

        # Each step of phase estimation needs to assign the measured phase to
        # the correct half circle. We will conservatively require that the
        # (x,z) tuple is long enough that we can assign it to the correct
        # quadrant of the circle with 2σ confidence

        if distances[k] < 2*np.sqrt(xvar[k] + zvar[k]):
            print('Phase estimation terminated at %dth pulse because the (x,z) vector is too short'%(k))
            break

        lowerBound = restrict(curGuess - np.pi/2**(k))
        upperBound = restrict(curGuess + np.pi/2**(k))
        possiblesTest = [ restrict((phases[k] + 2*n*np.pi)/2**(k)) for n in range(0,2**(k)+1)]

        if verbose == True:
            print('Lower Bound: %f'%lowerBound)
            print('Upper Bound: %f'%upperBound)

        possibles=[]
        for p in possiblesTest:
            # NOTE: previous code did not handle upperbound == lowerBound
            if lowerBound >= upperBound:
                satisfiesLB = p > lowerBound or p < 0.
                satisfiesUP = p < upperBound or p > 0.
            else:
                satisfiesLB = p > lowerBound
                satisfiesUP = p < upperBound

            if satisfiesLB == True and satisfiesUP == True:
                possibles.append(p)

        curGuess = possibles[0]
        if verbose == True:
            print('Current Guess: %f'%(curGuess))
    phase = curGuess
    sigma = np.maximum(np.abs(restrict(curGuess - lowerBound)), np.abs(restrict(curGuess - upperBound)))

    return phase, sigma

<|endoftext|>