<commit_before>'''
Channels is where we store information for mapping virtual (qubit) channel to real channels.

Created on Jan 19, 2012

Original Author: Colm Ryan

Copyright 2013 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

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.
'''

from . import PulseShapes
import numpy as np

from math import tan,cos,pi
# Python 2/3 compatibility: use the py3 meaning of 'str'
from builtins import str

from atom.api import Atom, Str, Float, Instance, \
    Dict, Enum, Bool, Typed, Int

from copy import deepcopy

class Channel(Atom):
    '''
    Every channel has a label and some printers.
    '''
    label = Str()
    enabled = Bool(True)

    def __repr__(self):
        return str(self)

    def __str__(self):
        return "{0}: {1}".format(self.__class__.__name__, self.label)

    def json_encode(self):
        jsonDict = self.__getstate__()

        #Turn objects into labels
        for member in ["physChan", "gateChan", "trigChan", "source", "target"]:
            if member in jsonDict:
                obj = jsonDict.pop(member)
                if obj:
                    jsonDict[member] = obj.label

        #We want the name of shape functions
        if "pulseParams" in jsonDict:
            pulseParams = deepcopy(jsonDict.pop("pulseParams"))
            if "shapeFun" in pulseParams:
                pulseParams["shapeFun"] = pulseParams["shapeFun"].__name__
            jsonDict["pulseParams"] = pulseParams

        return jsonDict

class PhysicalChannel(Channel):
    '''
    The main class for actual AWG channels.
    '''
    AWG = Str()
    translator = Str()
    generator = Str()
    samplingRate = Float(default=1.2e9)
    delay = Float()

class LogicalChannel(Channel):
    '''
    The main class from which we will generate sequences.
    At some point it needs to be assigned to a physical channel.
    '''
    #During initilization we may just have a string reference to the channel
    physChan = Instance((str,PhysicalChannel))

    def __init__(self, **kwargs):
        super(LogicalChannel, self).__init__(**kwargs)
        if self.physChan is None:
            self.physChan = PhysicalChannel(label=kwargs['label']+'-phys')

class PhysicalMarkerChannel(PhysicalChannel):
    '''
    An digital output channel on an AWG.
    '''
    gateBuffer = Float(0.0).tag(desc="How much extra time should be added onto the beginning of a gating pulse")
    gateMinWidth = Float(0.0).tag(desc="The minimum marker pulse width")

class PhysicalQuadratureChannel(PhysicalChannel):
    '''
    Something used to implement a standard qubit channel with two analog channels and a microwave gating channel.
    '''
    IChannel = Str()
    QChannel = Str()
    #During initilization we may just have a string reference to the channel
    ampFactor = Float(1.0)
    phaseSkew = Float(0.0)

    @property
    def correctionT(self):
        return np.array([[self.ampFactor, self.ampFactor*tan(self.phaseSkew*pi/180)], [0, 1/cos(self.phaseSkew*pi/180)]])

class LogicalMarkerChannel(LogicalChannel):
    '''
    A class for digital channels for gating sources or triggering other things.
    '''
    pulseParams = Dict(default={'shapeFun': PulseShapes.square, 'length':100e-9})

class Qubit(LogicalChannel):
    '''
    The main class for generating qubit pulses.  Effectively a logical "QuadratureChannel".
    '''
    pulseParams = Dict(default={'length':20e-9, 'piAmp':1.0, 'pi2Amp':0.5, 'shapeFun':PulseShapes.gaussian, 'cutoff':2, 'dragScaling':0, 'sigma':5e-9})
    gateChan = Instance((str, LogicalMarkerChannel))
    frequency = Float(0.0).tag(desc='modulation frequency of the channel (can be positive or negative)')

    def __init__(self, **kwargs):
        super(Qubit, self).__init__(**kwargs)
        if self.gateChan is None:
            self.gateChan = LogicalMarkerChannel(label=kwargs['label']+'-gate')

class Measurement(LogicalChannel):
    '''
    A class for measurement channels.
    Measurements are special because they can be different types:
    autodyne which needs an IQ pair or hetero/homodyne which needs just a marker channel.
    '''
    measType = Enum('autodyne','homodyne').tag(desc='Type of measurment (autodyne, homodyne)')
    autodyneFreq = Float(0.0).tag(desc='use to bake the modulation into the pulse, so that it has constant phase')
    frequency = Float(0.0).tag(desc='use frequency to asssociate modulation with the channel')
    pulseParams = Dict(default={'length':100e-9, 'amp':1.0, 'shapeFun':PulseShapes.tanh, 'cutoff':2, 'sigma':1e-9})
    gateChan = Instance((str, LogicalMarkerChannel))
    trigChan = Instance((str, LogicalMarkerChannel))

    def __init__(self, **kwargs):
        super(Measurement, self).__init__(**kwargs)
        if self.gateChan is None:
            self.gateChan = LogicalMarkerChannel(label=kwargs['label']+'-gate')
        if self.trigChan is None:
            self.trigChan = LogicalMarkerChannel(label='digitizerTrig')

class Edge(LogicalChannel):
    '''
    Defines an arc/directed edge between qubit vertices. If a device supports bi-directional
    connectivity, that is represented with two independent Edges.

    An Edge is also effectively an abstract channel, so it carries the same properties as a
    Qubit channel.
    '''
    # allow string in source and target so that we can store a label or an object
    source = Instance((str, Qubit))
    target = Instance((str, Qubit))
    pulseParams = Dict(default={'length':20e-9, 'amp':1.0, 'phase':0.0, 'shapeFun':PulseShapes.gaussian, 'cutoff':2, 'dragScaling':0, 'sigma':5e-9, 'riseFall': 20e-9})
    gateChan = Instance((str, LogicalMarkerChannel))
    frequency = Float(0.0).tag(desc='modulation frequency of the channel (can be positive or negative)')

    def __init__(self, **kwargs):
        super(Edge, self).__init__(**kwargs)
        if self.gateChan is None:
            self.gateChan = LogicalMarkerChannel(label=kwargs['label']+'-gate')

    def isforward(self, source, target):
        ''' Test whether (source, target) matches the directionality of the edge. '''
        nodes = (self.source, self.target)
        if (source not in nodes) or (target not in nodes):
            raise ValueError('One of {0} is not a node in the edge'.format((source, target)))
        if (self.source, self.target) == (source, target):
            return True
        else:
            return False

NewLogicalChannelList = [Qubit, Edge, LogicalMarkerChannel, Measurement]
NewPhysicalChannelList = [PhysicalMarkerChannel, PhysicalQuadratureChannel]
<commit_msg>Check if channel property is already string<commit_after>'''
Channels is where we store information for mapping virtual (qubit) channel to real channels.

Created on Jan 19, 2012

Original Author: Colm Ryan

Copyright 2013 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

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.
'''

from . import PulseShapes
import numpy as np

from math import tan,cos,pi
# Python 2/3 compatibility: use the py3 meaning of 'str'
from builtins import str

from atom.api import Atom, Str, Float, Instance, \
    Dict, Enum, Bool, Typed, Int

from copy import deepcopy

class Channel(Atom):
    '''
    Every channel has a label and some printers.
    '''
    label = Str()
    enabled = Bool(True)

    def __repr__(self):
        return str(self)

    def __str__(self):
        return "{0}: {1}".format(self.__class__.__name__, self.label)

    def json_encode(self):
        jsonDict = self.__getstate__()

        #Turn objects into labels
        for member in ["physChan", "gateChan", "trigChan", "source", "target"]:
            if member in jsonDict and not isinstance(jsonDict[member], str):
                obj = jsonDict.pop(member)
                if obj:
                    jsonDict[member] = obj.label

        #We want the name of shape functions
        if "pulseParams" in jsonDict:
            pulseParams = deepcopy(jsonDict.pop("pulseParams"))
            if "shapeFun" in pulseParams:
                pulseParams["shapeFun"] = pulseParams["shapeFun"].__name__
            jsonDict["pulseParams"] = pulseParams

        return jsonDict

class PhysicalChannel(Channel):
    '''
    The main class for actual AWG channels.
    '''
    AWG = Str()
    translator = Str()
    generator = Str()
    samplingRate = Float(default=1.2e9)
    delay = Float()

class LogicalChannel(Channel):
    '''
    The main class from which we will generate sequences.
    At some point it needs to be assigned to a physical channel.
    '''
    #During initilization we may just have a string reference to the channel
    physChan = Instance((str,PhysicalChannel))

    def __init__(self, **kwargs):
        super(LogicalChannel, self).__init__(**kwargs)
        if self.physChan is None:
            self.physChan = PhysicalChannel(label=kwargs['label']+'-phys')

class PhysicalMarkerChannel(PhysicalChannel):
    '''
    An digital output channel on an AWG.
    '''
    gateBuffer = Float(0.0).tag(desc="How much extra time should be added onto the beginning of a gating pulse")
    gateMinWidth = Float(0.0).tag(desc="The minimum marker pulse width")

class PhysicalQuadratureChannel(PhysicalChannel):
    '''
    Something used to implement a standard qubit channel with two analog channels and a microwave gating channel.
    '''
    IChannel = Str()
    QChannel = Str()
    #During initilization we may just have a string reference to the channel
    ampFactor = Float(1.0)
    phaseSkew = Float(0.0)

    @property
    def correctionT(self):
        return np.array([[self.ampFactor, self.ampFactor*tan(self.phaseSkew*pi/180)], [0, 1/cos(self.phaseSkew*pi/180)]])

class LogicalMarkerChannel(LogicalChannel):
    '''
    A class for digital channels for gating sources or triggering other things.
    '''
    pulseParams = Dict(default={'shapeFun': PulseShapes.square, 'length':100e-9})

class Qubit(LogicalChannel):
    '''
    The main class for generating qubit pulses.  Effectively a logical "QuadratureChannel".
    '''
    pulseParams = Dict(default={'length':20e-9, 'piAmp':1.0, 'pi2Amp':0.5, 'shapeFun':PulseShapes.gaussian, 'cutoff':2, 'dragScaling':0, 'sigma':5e-9})
    gateChan = Instance((str, LogicalMarkerChannel))
    frequency = Float(0.0).tag(desc='modulation frequency of the channel (can be positive or negative)')

    def __init__(self, **kwargs):
        super(Qubit, self).__init__(**kwargs)
        if self.gateChan is None:
            self.gateChan = LogicalMarkerChannel(label=kwargs['label']+'-gate')

class Measurement(LogicalChannel):
    '''
    A class for measurement channels.
    Measurements are special because they can be different types:
    autodyne which needs an IQ pair or hetero/homodyne which needs just a marker channel.
    '''
    measType = Enum('autodyne','homodyne').tag(desc='Type of measurment (autodyne, homodyne)')
    autodyneFreq = Float(0.0).tag(desc='use to bake the modulation into the pulse, so that it has constant phase')
    frequency = Float(0.0).tag(desc='use frequency to asssociate modulation with the channel')
    pulseParams = Dict(default={'length':100e-9, 'amp':1.0, 'shapeFun':PulseShapes.tanh, 'cutoff':2, 'sigma':1e-9})
    gateChan = Instance((str, LogicalMarkerChannel))
    trigChan = Instance((str, LogicalMarkerChannel))

    def __init__(self, **kwargs):
        super(Measurement, self).__init__(**kwargs)
        if self.gateChan is None:
            self.gateChan = LogicalMarkerChannel(label=kwargs['label']+'-gate')
        if self.trigChan is None:
            self.trigChan = LogicalMarkerChannel(label='digitizerTrig')

class Edge(LogicalChannel):
    '''
    Defines an arc/directed edge between qubit vertices. If a device supports bi-directional
    connectivity, that is represented with two independent Edges.

    An Edge is also effectively an abstract channel, so it carries the same properties as a
    Qubit channel.
    '''
    # allow string in source and target so that we can store a label or an object
    source = Instance((str, Qubit))
    target = Instance((str, Qubit))
    pulseParams = Dict(default={'length':20e-9, 'amp':1.0, 'phase':0.0, 'shapeFun':PulseShapes.gaussian, 'cutoff':2, 'dragScaling':0, 'sigma':5e-9, 'riseFall': 20e-9})
    gateChan = Instance((str, LogicalMarkerChannel))
    frequency = Float(0.0).tag(desc='modulation frequency of the channel (can be positive or negative)')

    def __init__(self, **kwargs):
        super(Edge, self).__init__(**kwargs)
        if self.gateChan is None:
            self.gateChan = LogicalMarkerChannel(label=kwargs['label']+'-gate')

    def isforward(self, source, target):
        ''' Test whether (source, target) matches the directionality of the edge. '''
        nodes = (self.source, self.target)
        if (source not in nodes) or (target not in nodes):
            raise ValueError('One of {0} is not a node in the edge'.format((source, target)))
        if (self.source, self.target) == (source, target):
            return True
        else:
            return False

NewLogicalChannelList = [Qubit, Edge, LogicalMarkerChannel, Measurement]
NewPhysicalChannelList = [PhysicalMarkerChannel, PhysicalQuadratureChannel]
<|endoftext|>


<commit_before># -*- coding: UTF-8 -*-
'''
Created on Feb 18, 2015

@author: cryan@bbn.com and bjohnson@bbn.com

Common plotting tools.

Copyright 2012-2015 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

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.
'''

import os.path, uuid, tempfile
import bokeh.plotting as bk
from bokeh.layouts import column
import numpy as np

from . import config

def output_notebook():
    bk.output_notebook()

def output_file():
    bk.output_file(os.path.join(tempfile.gettempdir(), str(uuid.uuid4()) +
                                ".html"))

# Default to output_file
# Commented out because recent versions of Bokeh remember the state of
# output_notebook/output_file on a per module basis.

# output_file()

def build_waveforms(seq):
    # import here to avoid circular imports
    from . import Compiler, PulseSequencer, BlockLabel, ControlFlow
    wires = Compiler.compile_sequence(seq)

    # build a concatenated waveform for each channel
    channels = wires.keys()
    concatShapes = {q: np.array([0], dtype=np.complex128) for q in channels}
    for q in channels:
        frame = 0
        repeat = 0
        label_idx = {}
        for entry in wires[q]:
            if isinstance(entry, PulseSequencer.Pulse):
                shape = np.exp(1j *
                               (frame + entry.phase)) * entry.amp * entry.shape
                frame += entry.frameChange
                concatShapes[q] = np.append(concatShapes[q], shape)
            elif isinstance(entry, BlockLabel.BlockLabel):
                label_idx[entry] = len(concatShapes[q]) + 1
            elif isinstance(entry, ControlFlow.LoadRepeat):
                repeat = entry.value - 1
            elif isinstance(entry, ControlFlow.Repeat):
                concatShapes[q] = np.append(concatShapes[q], np.tile(
                    concatShapes[q][label_idx[entry.target]:], repeat))

    # add an extra zero to make things look more normal
    for q in channels:
        concatShapes[q] = np.append(concatShapes[q], 0)
    return concatShapes

def plot_waveforms(waveforms, figTitle=''):
    channels = waveforms.keys()
    # plot
    plots = []
    for (ct, chan) in enumerate(channels):
        fig = bk.figure(title=figTitle + repr(chan),
                        plot_width=800,
                        plot_height=350,
                        y_range=[-1.05, 1.05],
                        x_axis_label=u'Time (μs)')
        fig.background_fill_color = config.plotBackground
        if config.gridColor:
            fig.xgrid.grid_line_color = config.gridColor
            fig.ygrid.grid_line_color = config.gridColor
        waveformToPlot = waveforms[chan]
        xpts = np.linspace(0, len(waveformToPlot) / chan.physChan.samplingRate
                           / 1e-6, len(waveformToPlot))
        fig.line(xpts, np.real(waveformToPlot), color='red')
        fig.line(xpts, np.imag(waveformToPlot), color='blue')
        plots.append(fig)
    bk.show(column(*plots))

def show(seq):
    waveforms = build_waveforms(seq)
    plot_waveforms(waveforms)
<commit_msg>Load bokeh js from local copies, suppress bokeh specific warnings if desired.<commit_after># -*- coding: UTF-8 -*-
'''
Created on Feb 18, 2015

@author: cryan@bbn.com and bjohnson@bbn.com

Common plotting tools.

Copyright 2012-2015 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

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.
'''

import os.path, uuid, tempfile
import bokeh.plotting as bk
from bokeh.layouts import column
from bokeh.util.warnings import BokehUserWarning
from bokeh.resources import INLINE
import numpy as np
import warnings
from . import config

def output_notebook(local=True, suppress_warnings=True):
    if suppress_warnings:
        warnings.simplefilter("ignore", BokehUserWarning)
    if local:
        bk.output_notebook(resources=INLINE)
    else:
        bk.output_notebook()

def output_file():
    bk.output_file(os.path.join(tempfile.gettempdir(), str(uuid.uuid4()) +
                                ".html"))

# Default to output_file
# Commented out because recent versions of Bokeh remember the state of
# output_notebook/output_file on a per module basis.

# output_file()

def build_waveforms(seq):
    # import here to avoid circular imports
    from . import Compiler, PulseSequencer, BlockLabel, ControlFlow
    wires = Compiler.compile_sequence(seq)

    # build a concatenated waveform for each channel
    channels = wires.keys()
    concatShapes = {q: np.array([0], dtype=np.complex128) for q in channels}
    for q in channels:
        frame = 0
        repeat = 0
        label_idx = {}
        for entry in wires[q]:
            if isinstance(entry, PulseSequencer.Pulse):
                shape = np.exp(1j *
                               (frame + entry.phase)) * entry.amp * entry.shape
                frame += entry.frameChange
                concatShapes[q] = np.append(concatShapes[q], shape)
            elif isinstance(entry, BlockLabel.BlockLabel):
                label_idx[entry] = len(concatShapes[q]) + 1
            elif isinstance(entry, ControlFlow.LoadRepeat):
                repeat = entry.value - 1
            elif isinstance(entry, ControlFlow.Repeat):
                concatShapes[q] = np.append(concatShapes[q], np.tile(
                    concatShapes[q][label_idx[entry.target]:], repeat))

    # add an extra zero to make things look more normal
    for q in channels:
        concatShapes[q] = np.append(concatShapes[q], 0)
    return concatShapes

def plot_waveforms(waveforms, figTitle=''):
    channels = waveforms.keys()
    # plot
    plots = []
    for (ct, chan) in enumerate(channels):
        fig = bk.figure(title=figTitle + repr(chan),
                        plot_width=800,
                        plot_height=350,
                        y_range=[-1.05, 1.05],
                        x_axis_label=u'Time (μs)')
        fig.background_fill_color = config.plotBackground
        if config.gridColor:
            fig.xgrid.grid_line_color = config.gridColor
            fig.ygrid.grid_line_color = config.gridColor
        waveformToPlot = waveforms[chan]
        xpts = np.linspace(0, len(waveformToPlot) / chan.physChan.samplingRate
                           / 1e-6, len(waveformToPlot))
        fig.line(xpts, np.real(waveformToPlot), color='red')
        fig.line(xpts, np.imag(waveformToPlot), color='blue')
        plots.append(fig)
    bk.show(column(*plots))

def show(seq):
    waveforms = build_waveforms(seq)
    plot_waveforms(waveforms)
<|endoftext|>


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

###
# Copyright (c) 2002-2007 Systems in Motion
#
# Permission to use, copy, modify, and distribute this software for any
# purpose with or without fee is hereby granted, provided that the above
# copyright notice and this permission notice appear in all copies.
#
# THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES
# WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF
# MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR
# ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES
# WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN
# ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF
# OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.
#

###
# This is an example from the Inventor Mentor,
# chapter 14, example 3.
#
# This example illustrates the creation of motion hierarchies
# using nodekits by creating a model of a balance-style scale.
#
# It adds an SoEventCallback to the "callback" list in the
#     nodekit called 'support.'
# The callback will have the following response to events:
# Pressing right arrow key == lower the right pan
# Pressing left arrow key  == lower the left pan
# The pans are lowered by animating three rotations in the
#     motion hierarchy.
# Use an SoText2Kit to print instructions to the user as part
#     of the scene.
#

import sys

from pivy.coin import *
from pivy.sogui import *

# Callback Function for Animating the Balance Scale.
# -- used to make the balance tip back and forth
# -- Note: this routine is only called in response to KeyPress
#    events since the call 'setEventInterest(KeyPressMask)' is
#    made on the SoEventCallback node that uses it.
# -- The routine checks if the key pressed was left arrow (which
#    is XK_Left in X-windows talk), or right arrow (which is
#    XK_Right)
# -- The balance is made to tip by rotating the beam part of the
#    scale (to tip it) and then compensating (making the strings
#    vertical again) by rotating the string parts in the opposite
#    direction.
def tipTheBalance(support, # The nodekit representing 'support', the
                  # fulcrum of the balance. Passed in during
                  # main routine, below.
                  eventCB):

    ev = eventCB.getEvent()

    # Which Key was pressed?
    # If Right or Left Arrow key, then continue...
    if SoKeyboardEvent.isKeyPressEvent(ev, SoKeyboardEvent.RIGHT_ARROW) or \
       SoKeyboardEvent.isKeyPressEvent(ev, SoKeyboardEvent.LEFT_ARROW):
        startRot, beamIncrement, stringIncrement = SbRotation(), SbRotation(), SbRotation()

        # These three parts are extracted based on knowledge of the
        # motion hierarchy (see the diagram in the main routine.
        beam1   = support.getPart("childList[0]",TRUE)
        string1 = beam1.getPart("childList[0]",TRUE)
        string2 = beam1.getPart("childList[1]",TRUE)

        # Set angular increments to be .1 Radians about the Z-Axis
        # The strings rotate opposite the beam, and the two types
        # of key press produce opposite effects.
        if SoKeyboardEvent.isKeyPressEvent(ev, SoKeyboardEvent.RIGHT_ARROW):
            beamIncrement.setValue(SbVec3f(0, 0, 1), -.1)
            stringIncrement.setValue(SbVec3f(0, 0, 1), .1)
        else:
            beamIncrement.setValue(SbVec3f(0, 0, 1), .1)
            stringIncrement.setValue(SbVec3f(0, 0, 1), -.1)

        # Use SO_GET_PART to find the transform for each of the
        # rotating parts and modify their rotations.

        xf = beam1.getPart("transform", TRUE)
        startRot = xf.rotation.getValue()
        startRot *= beamIncrement
        xf.rotation = startRot

        xf = string1.getPart("transform", TRUE)
        startRot = xf.rotation.getValue()
        startRot *= stringIncrement
        xf.rotation = startRot

        xf = string2.getPart("transform", TRUE)
        startRot = xf.rotation.getValue()
        startRot *= stringIncrement
        xf.rotation = startRot

        # eventCB.setHandled()

def main():
    # Initialize Inventor and Qt
    myWindow = SoGui.init(sys.argv[0])
    if myWindow == None: sys.exit(1)

    myScene = SoSceneKit()

    myScene.setPart("lightList[0]", SoLightKit())
    myScene.setPart("cameraList[0]", SoCameraKit())
    myScene.setCameraNumber(0)

    # Create the Balance Scale -- put each part in the
    # childList of its parent, to build up this hierarchy:
    #
    #                    myScene
    #                       |
    #                     support
    #                       |
    #                     beam
    #                       |
    #                   --------
    #                   |       |
    #                string1  string2
    #                   |       |
    #                tray1     tray2

    support = SoShapeKit()
    support.setPart("shape", SoCone())
    support.set("shape { height 3 bottomRadius .3 }")
    myScene.setPart("childList[0]", support)

    beam = SoShapeKit()
    beam.setPart("shape", SoCube())
    beam.set("shape { width 3 height .2 depth .2 }")
    beam.set("transform { translation 0 1.5 0 } ")
    support.setPart("childList[0]", beam)

    string1 = SoShapeKit()
    string1.setPart("shape", SoCylinder())
    string1.set("shape { radius .05 height 2}")
    string1.set("transform { translation -1.5 -1 0 }")
    string1.set("transform { center 0 1 0 }")
    beam.setPart("childList[0]", string1)

    string2 = SoShapeKit()
    string2.setPart("shape", SoCylinder())
    string2.set("shape { radius .05 height 2}")
    string2.set("transform { translation 1.5 -1 0 } ")
    string2.set("transform { center 0 1 0 } ")
    beam.setPart("childList[1]", string2)

    tray1 = SoShapeKit()
    tray1.setPart("shape", SoCylinder())
    tray1.set("shape { radius .75 height .1 }")
    tray1.set("transform { translation 0 -1 0 } ")
    string1.setPart("childList[0]", tray1)

    tray2 = SoShapeKit()
    tray2.setPart("shape", SoCylinder())
    tray2.set("shape { radius .75 height .1 }")
    tray2.set("transform { translation 0 -1 0 } ")
    string2.setPart("childList[0]", tray2)

    # Add EventCallback so Balance Responds to Events
    myCallbackNode = SoEventCallback()
    myCallbackNode.addEventCallback(SoKeyboardEvent.getClassTypeId(),
                                    tipTheBalance, support)
    support.setPart("callbackList[0]", myCallbackNode)

    # Add Instructions as Text in the Scene...
    myText = SoShapeKit()
    myText.setPart("shape", SoText2())
    myText.set("shape { string \"Press Left or Right Arrow Key\" }")
    myText.set("shape { justification CENTER }")
    myText.set("font { name \"Helvetica-Bold\" }")
    myText.set("font { size 16.0 }")
    myText.set("transform { translation 0 -2 0 }")
    myScene.setPart("childList[1]", myText)

    myRenderArea = SoGuiRenderArea(myWindow)

    # Get camera from scene and tell it to viewAll...
    myCamera = myScene.getPart("cameraList[0].camera", TRUE)
    myCamera.viewAll(myScene, myRenderArea.getViewportRegion())

    myRenderArea.setSceneGraph(myScene)
    myRenderArea.setTitle("Balance Scale Made of Nodekits")
    myRenderArea.show()

    SoGui.show(myWindow)
    SoGui.mainLoop()

if __name__ == "__main__":
    main()
<commit_msg>reset example file<commit_after>#!/usr/bin/env python

###
# Copyright (c) 2002-2007 Systems in Motion
#
# Permission to use, copy, modify, and distribute this software for any
# purpose with or without fee is hereby granted, provided that the above
# copyright notice and this permission notice appear in all copies.
#
# THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES
# WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF
# MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR
# ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES
# WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN
# ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF
# OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.
#

###
# This is an example from the Inventor Mentor,
# chapter 14, example 3.
#
# This example illustrates the creation of motion hierarchies
# using nodekits by creating a model of a balance-style scale.
#
# It adds an SoEventCallback to the "callback" list in the
#     nodekit called 'support.'
# The callback will have the following response to events:
# Pressing right arrow key == lower the right pan
# Pressing left arrow key  == lower the left pan
# The pans are lowered by animating three rotations in the
#     motion hierarchy.
# Use an SoText2Kit to print instructions to the user as part
#     of the scene.
#

import sys

from pivy.coin import *
from pivy.sogui import *

# Callback Function for Animating the Balance Scale.
# -- used to make the balance tip back and forth
# -- Note: this routine is only called in response to KeyPress
#    events since the call 'setEventInterest(KeyPressMask)' is
#    made on the SoEventCallback node that uses it.
# -- The routine checks if the key pressed was left arrow (which
#    is XK_Left in X-windows talk), or right arrow (which is
#    XK_Right)
# -- The balance is made to tip by rotating the beam part of the
#    scale (to tip it) and then compensating (making the strings
#    vertical again) by rotating the string parts in the opposite
#    direction.
def tipTheBalance(support, # The nodekit representing 'support', the
                  # fulcrum of the balance. Passed in during
                  # main routine, below.
                  eventCB):

    ev = eventCB.getEvent()

    # Which Key was pressed?
    # If Right or Left Arrow key, then continue...
    if SoKeyboardEvent.isKeyPressEvent(ev, SoKeyboardEvent.RIGHT_ARROW) or \
       SoKeyboardEvent.isKeyPressEvent(ev, SoKeyboardEvent.LEFT_ARROW):
        startRot, beamIncrement, stringIncrement = SbRotation(), SbRotation(), SbRotation()

        # These three parts are extracted based on knowledge of the
        # motion hierarchy (see the diagram in the main routine.
        beam1   = support.getPart("childList[0]",TRUE)
        string1 = beam1.getPart("childList[0]",TRUE)
        string2 = beam1.getPart("childList[1]",TRUE)

        # Set angular increments to be .1 Radians about the Z-Axis
        # The strings rotate opposite the beam, and the two types
        # of key press produce opposite effects.
        if SoKeyboardEvent.isKeyPressEvent(ev, SoKeyboardEvent.RIGHT_ARROW):
            beamIncrement.setValue(SbVec3f(0, 0, 1), -.1)
            stringIncrement.setValue(SbVec3f(0, 0, 1), .1)
        else:
            beamIncrement.setValue(SbVec3f(0, 0, 1), .1)
            stringIncrement.setValue(SbVec3f(0, 0, 1), -.1)

        # Use SO_GET_PART to find the transform for each of the
        # rotating parts and modify their rotations.

        xf = beam1.getPart("transform", TRUE)
        startRot = xf.rotation.getValue()
        startRot *= beamIncrement
        xf.rotation = startRot

        xf = string1.getPart("transform", TRUE)
        startRot = xf.rotation.getValue()
        startRot *= stringIncrement
        xf.rotation = startRot

        xf = string2.getPart("transform", TRUE)
        startRot = xf.rotation.getValue()
        startRot *= stringIncrement
        xf.rotation = startRot

        eventCB.setHandled()

def main():
    # Initialize Inventor and Qt
    myWindow = SoGui.init(sys.argv[0])
    if myWindow == None: sys.exit(1)

    myScene = SoSceneKit()

    myScene.setPart("lightList[0]", SoLightKit())
    myScene.setPart("cameraList[0]", SoCameraKit())
    myScene.setCameraNumber(0)

    # Create the Balance Scale -- put each part in the
    # childList of its parent, to build up this hierarchy:
    #
    #                    myScene
    #                       |
    #                     support
    #                       |
    #                     beam
    #                       |
    #                   --------
    #                   |       |
    #                string1  string2
    #                   |       |
    #                tray1     tray2

    support = SoShapeKit()
    support.setPart("shape", SoCone())
    support.set("shape { height 3 bottomRadius .3 }")
    myScene.setPart("childList[0]", support)

    beam = SoShapeKit()
    beam.setPart("shape", SoCube())
    beam.set("shape { width 3 height .2 depth .2 }")
    beam.set("transform { translation 0 1.5 0 } ")
    support.setPart("childList[0]", beam)

    string1 = SoShapeKit()
    string1.setPart("shape", SoCylinder())
    string1.set("shape { radius .05 height 2}")
    string1.set("transform { translation -1.5 -1 0 }")
    string1.set("transform { center 0 1 0 }")
    beam.setPart("childList[0]", string1)

    string2 = SoShapeKit()
    string2.setPart("shape", SoCylinder())
    string2.set("shape { radius .05 height 2}")
    string2.set("transform { translation 1.5 -1 0 } ")
    string2.set("transform { center 0 1 0 } ")
    beam.setPart("childList[1]", string2)

    tray1 = SoShapeKit()
    tray1.setPart("shape", SoCylinder())
    tray1.set("shape { radius .75 height .1 }")
    tray1.set("transform { translation 0 -1 0 } ")
    string1.setPart("childList[0]", tray1)

    tray2 = SoShapeKit()
    tray2.setPart("shape", SoCylinder())
    tray2.set("shape { radius .75 height .1 }")
    tray2.set("transform { translation 0 -1 0 } ")
    string2.setPart("childList[0]", tray2)

    # Add EventCallback so Balance Responds to Events
    myCallbackNode = SoEventCallback()
    myCallbackNode.addEventCallback(SoKeyboardEvent.getClassTypeId(),
                                    tipTheBalance, support)
    support.setPart("callbackList[0]", myCallbackNode)

    # Add Instructions as Text in the Scene...
    myText = SoShapeKit()
    myText.setPart("shape", SoText2())
    myText.set("shape { string \"Press Left or Right Arrow Key\" }")
    myText.set("shape { justification CENTER }")
    myText.set("font { name \"Helvetica-Bold\" }")
    myText.set("font { size 16.0 }")
    myText.set("transform { translation 0 -2 0 }")
    myScene.setPart("childList[1]", myText)

    myRenderArea = SoGuiRenderArea(myWindow)

    # Get camera from scene and tell it to viewAll...
    myCamera = myScene.getPart("cameraList[0].camera", TRUE)
    myCamera.viewAll(myScene, myRenderArea.getViewportRegion())

    myRenderArea.setSceneGraph(myScene)
    myRenderArea.setTitle("Balance Scale Made of Nodekits")
    myRenderArea.show()

    SoGui.show(myWindow)
    SoGui.mainLoop()

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


<commit_before>#!/usr/bin/env python
#
# BSD 3-Clause License
#
# Copyright (c) 2016-21, University of Liverpool
# All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
#
# * Redistributions of source code must retain the above copyright notice, this
#   list of conditions and the following disclaimer.
#
# * Redistributions in binary form must reproduce the above copyright notice,
#   this list of conditions and the following disclaimer in the documentation
#   and/or other materials provided with the distribution.
#
# * Neither the name of the copyright holder nor the names of its
#   contributors may be used to endorse or promote products derived from
#   this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
"""This script provides a model validation pipeline. It will
take a sequence, a predicted distogram and a PDB model, then
compare the distances observed in the model and those predicted.
It will then report regions in the model where an outlier was
detected, and use map_align to provide a solution for potential
register errors.

It uses one external programs to perform this task:

   map_align for contact map alignment

*** This program needs to be installed separately from https://github.com/sokrypton/map_align***

"""

import argparse
import os
import shutil
from prettytable import PrettyTable

import conkit.applications
import conkit.command_line
import conkit.io
import conkit.plot

logger = None

def create_argument_parser():
    """Create a parser for the command line arguments used in conkit-validate"""

    parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.ArgumentDefaultsHelpFormatter)
    parser.add_argument("seqfile", type=check_file_exists, help="Path to sequence file")
    parser.add_argument("seqformat", type=str, help="Sequence format")
    parser.add_argument("distfile", type=check_file_exists, help="Path to distance prediction file")
    parser.add_argument("distformat", type=str, help="Format of distance prediction file")
    parser.add_argument("pdbfile", type=check_file_exists, help="Path to structure file")
    parser.add_argument("pdbformat", type=str, help="Format of structure file")
    parser.add_argument("--outdir", dest="outdir", type=str, help="Output directory",
                        default=os.path.join(os.getcwd(), 'conkit-validate'))
    parser.add_argument("--overwrite", dest="overwrite", default=False, action="store_true",
                        help="overwrite output directory if exists")
    parser.add_argument("--skip_alignment", dest="skip_alignment", default=False, action="store_true",
                        help="skip contact map alignment step")
    parser.add_argument("--map_align_exe", dest="map_align_exe", default="map_align",
                        type=check_file_exists, help="Path to the map_align executable")
    parser.add_argument("--gap_opening_penalty", dest="gap_opening_penalty", default=-1, type=float,
                        help="Gap opening penalty")
    parser.add_argument("--gap_extension_penalty", dest="gap_extension_penalty", default=-0.01, type=float,
                        help="Gap extension penalty")
    parser.add_argument("--seq_separation_cutoff", dest="seq_separation_cutoff", default=3, type=int,
                        help="Sequence separation cutoff"),
    parser.add_argument("--n_iterations", dest="n_iterations", default=20, type=int,
                        help="Number of iterations")

    return parser

def check_file_exists(input_path):
    """Check if a given path exists

    Parameters
    ----------
    input_path : str, None
       Location of the file to be tested

    Returns
    -------
    abspath : str, None
       The absolute path of the file if it exists, None if the input is None

    Raises
    ------
    :exc:`FileNotFoundError`
        The file doesn't exist
    """

    if input_path is None:
        return None
    if os.path.isfile(os.path.abspath(input_path)):
        return os.path.abspath(input_path)
    else:
        raise FileNotFoundError("{} cannot be found".format(input_path))

def prepare_output_directory(outdir, overwrite=False):
    """Prepare the output directory for conkit-validate.

    Parameters
    ----------
    outdir : str
       Path to the output directory
    overwrite : bool
       Whether the output directory should be overwritten or not [default: False]

    Raises
    ------
    :exc:`ValueError`
        The output directory already exists and overwrite is False
    """

    if os.path.isdir(outdir):
        if not overwrite:
            raise ValueError('Output directory already exists')
        else:
            shutil.rmtree(outdir)
    os.mkdir(outdir)

def parse_map_align_stdout(stdout):
    """Parse the stdout of map_align and extract the alignment of residues.

    Parameters
    ----------
    stdout : str
       Starndard output created with map_align

    Returns
    ------
    alignment: dict
        A dictionary where the aligned residue numbers of map_b are the keys and the residue numbers of map_a the values
    """

    alignment = {}
    for line in stdout.split('\n'):
        if line and line.split()[0] == "MAX":
            line = line.rstrip().lstrip().split()
            for residue in line[8:]:
                resnum = residue.split(":")
                alignment[int(resnum[1])] = int(resnum[0])

    return alignment

def main():
    """The main routine for conkit-validate functionality"""
    parser = create_argument_parser()
    args = parser.parse_args()

    global logger
    logger = conkit.command_line.setup_logging(level="info")

    logger.info("Output directory:                           %s", args.outdir)

    prepare_output_directory(args.outdir, args.overwrite)

    sequence = conkit.io.read(args.seqfile, args.seqformat).top
    prediction = conkit.io.read(args.distfile, args.distformat).top
    model = conkit.io.read(args.pdbfile, args.pdbformat).top
    logger.info("Input PDB model:                            %s", args.pdbfile)
    logger.info("Input distance prediction:                  %s", args.distfile)
    logger.info("Length of the sequence:                     %d", len(sequence))

    fig_fname = os.path.join(args.outdir, 'conkit.png')
    figure = conkit.plot.ModelValidationFigure(model, prediction, sequence, use_weights=True)
    figure.savefig(fig_fname, overwrite=args.overwrite)
    logger.info("Validation plot written to %s", fig_fname)

    if args.skip_alignment:
        logger.info("Skipping contact map alignment, no changes suggested.")

    elif any(figure.outliers):
        table = PrettyTable()
        table.field_names = ["Outlier no.", "Residue no.", "wRMSD", "FN Count"]
        for idx, outlier in enumerate(figure.outliers, 1):
            table.add_row([idx, str(outlier), '{0:.2f}'.format(figure.rmsd_profile[outlier]),
                           '{0:.2f}'.format(figure.fn_profile[outlier])])
        logger.info("\nList of detected outliers:")
        logger.info(table)

        contact_map_a = os.path.join(args.outdir, 'contact_map_a.mapalign')
        contact_map_b = os.path.join(args.outdir, 'contact_map_b.mapalign')
        conkit.io.write(contact_map_a, 'mapalign', prediction)
        conkit.io.write(contact_map_b, 'mapalign', model)

        map_align_cline = conkit.applications.MapAlignCommandline(
            cmd=args.map_align_exe,
            contact_map_a=contact_map_a,
            contact_map_b=contact_map_b,
            gap_opening_penalty=args.gap_opening_penalty,
            gap_extension_penalty=args.gap_extension_penalty,
            seq_separation_cutoff=args.seq_separation_cutoff,
            n_iterations=args.n_iterations
        )

        logger.info("\nExecuting: %s", map_align_cline)
        stdout, stderr = map_align_cline()
        map_align_log = os.path.join(args.outdir, 'map_align.log')
        with open(map_align_log, 'w') as fhandle:
            fhandle.write(stdout)
        alignment = parse_map_align_stdout(stdout)
        for idx, outlier in enumerate(figure.outliers, 1):
            table = PrettyTable()
            table.field_names = ["Current Residue", "New Residue"]
            start_outlier = outlier - 20 if outlier > 20 else 0
            stop_outlier = outlier + 20 if outlier + 20 < len(sequence) else len(sequence)
            for resnum in range(start_outlier, stop_outlier + 1):
                if resnum in alignment.keys() and alignment[resnum] != resnum:
                    table.add_row(['{} ({})'.format(sequence.seq[resnum - 1], resnum),
                                   '{} ({})'.format(sequence.seq[alignment[resnum] - 1], alignment[resnum])])
            if table._rows:
                logger.info("\nList of proposed changes to fix outlier no. {}:".format(idx))
                logger.info(table)
            else:
                logger.info("\nCannot find optimal re-alignment for outlier no. {}:".format(idx))

    else:
        logger.info("No outliers were detected, finishing now.")

if __name__ == "__main__":
    import sys
    import traceback

    try:
        main()
        sys.exit(0)
    except Exception as e:
        if not isinstance(e, SystemExit):
            msg = "".join(traceback.format_exception(*sys.exc_info()))
            logger.critical(msg)
        sys.exit(1)
<commit_msg>use os.linesep<commit_after>#!/usr/bin/env python
#
# BSD 3-Clause License
#
# Copyright (c) 2016-21, University of Liverpool
# All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
#
# * Redistributions of source code must retain the above copyright notice, this
#   list of conditions and the following disclaimer.
#
# * Redistributions in binary form must reproduce the above copyright notice,
#   this list of conditions and the following disclaimer in the documentation
#   and/or other materials provided with the distribution.
#
# * Neither the name of the copyright holder nor the names of its
#   contributors may be used to endorse or promote products derived from
#   this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
"""This script provides a model validation pipeline. It will
take a sequence, a predicted distogram and a PDB model, then
compare the distances observed in the model and those predicted.
It will then report regions in the model where an outlier was
detected, and use map_align to provide a solution for potential
register errors.

It uses one external programs to perform this task:

   map_align for contact map alignment

*** This program needs to be installed separately from https://github.com/sokrypton/map_align***

"""

import argparse
import os
import shutil
from prettytable import PrettyTable

import conkit.applications
import conkit.command_line
import conkit.io
import conkit.plot

logger = None

def create_argument_parser():
    """Create a parser for the command line arguments used in conkit-validate"""

    parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.ArgumentDefaultsHelpFormatter)
    parser.add_argument("seqfile", type=check_file_exists, help="Path to sequence file")
    parser.add_argument("seqformat", type=str, help="Sequence format")
    parser.add_argument("distfile", type=check_file_exists, help="Path to distance prediction file")
    parser.add_argument("distformat", type=str, help="Format of distance prediction file")
    parser.add_argument("pdbfile", type=check_file_exists, help="Path to structure file")
    parser.add_argument("pdbformat", type=str, help="Format of structure file")
    parser.add_argument("--outdir", dest="outdir", type=str, help="Output directory",
                        default=os.path.join(os.getcwd(), 'conkit-validate'))
    parser.add_argument("--overwrite", dest="overwrite", default=False, action="store_true",
                        help="overwrite output directory if exists")
    parser.add_argument("--skip_alignment", dest="skip_alignment", default=False, action="store_true",
                        help="skip contact map alignment step")
    parser.add_argument("--map_align_exe", dest="map_align_exe", default="map_align",
                        type=check_file_exists, help="Path to the map_align executable")
    parser.add_argument("--gap_opening_penalty", dest="gap_opening_penalty", default=-1, type=float,
                        help="Gap opening penalty")
    parser.add_argument("--gap_extension_penalty", dest="gap_extension_penalty", default=-0.01, type=float,
                        help="Gap extension penalty")
    parser.add_argument("--seq_separation_cutoff", dest="seq_separation_cutoff", default=3, type=int,
                        help="Sequence separation cutoff"),
    parser.add_argument("--n_iterations", dest="n_iterations", default=20, type=int,
                        help="Number of iterations")

    return parser

def check_file_exists(input_path):
    """Check if a given path exists

    Parameters
    ----------
    input_path : str, None
       Location of the file to be tested

    Returns
    -------
    abspath : str, None
       The absolute path of the file if it exists, None if the input is None

    Raises
    ------
    :exc:`FileNotFoundError`
        The file doesn't exist
    """

    if input_path is None:
        return None
    if os.path.isfile(os.path.abspath(input_path)):
        return os.path.abspath(input_path)
    else:
        raise FileNotFoundError("{} cannot be found".format(input_path))

def prepare_output_directory(outdir, overwrite=False):
    """Prepare the output directory for conkit-validate.

    Parameters
    ----------
    outdir : str
       Path to the output directory
    overwrite : bool
       Whether the output directory should be overwritten or not [default: False]

    Raises
    ------
    :exc:`ValueError`
        The output directory already exists and overwrite is False
    """

    if os.path.isdir(outdir):
        if not overwrite:
            raise ValueError('Output directory already exists')
        else:
            shutil.rmtree(outdir)
    os.mkdir(outdir)

def parse_map_align_stdout(stdout):
    """Parse the stdout of map_align and extract the alignment of residues.

    Parameters
    ----------
    stdout : str
       Starndard output created with map_align

    Returns
    ------
    alignment: dict
        A dictionary where the aligned residue numbers of map_b are the keys and the residue numbers of map_a the values
    """

    alignment = {}
    for line in stdout.split('\n'):
        if line and line.split()[0] == "MAX":
            line = line.rstrip().lstrip().split()
            for residue in line[8:]:
                resnum = residue.split(":")
                alignment[int(resnum[1])] = int(resnum[0])

    return alignment

def main():
    """The main routine for conkit-validate functionality"""
    parser = create_argument_parser()
    args = parser.parse_args()

    global logger
    logger = conkit.command_line.setup_logging(level="info")

    logger.info("Output directory:                           %s", args.outdir)

    prepare_output_directory(args.outdir, args.overwrite)

    sequence = conkit.io.read(args.seqfile, args.seqformat).top
    prediction = conkit.io.read(args.distfile, args.distformat).top
    model = conkit.io.read(args.pdbfile, args.pdbformat).top
    logger.info("Input PDB model:                            %s", args.pdbfile)
    logger.info("Input distance prediction:                  %s", args.distfile)
    logger.info("Length of the sequence:                     %d", len(sequence))

    fig_fname = os.path.join(args.outdir, 'conkit.png')
    figure = conkit.plot.ModelValidationFigure(model, prediction, sequence, use_weights=True)
    figure.savefig(fig_fname, overwrite=args.overwrite)
    logger.info("Validation plot written to %s", fig_fname)

    if args.skip_alignment:
        logger.info("Skipping contact map alignment, no changes suggested.")

    elif any(figure.outliers):
        table = PrettyTable()
        table.field_names = ["Outlier no.", "Residue no.", "wRMSD", "FN Count"]
        for idx, outlier in enumerate(figure.outliers, 1):
            table.add_row([idx, str(outlier), '{0:.2f}'.format(figure.rmsd_profile[outlier]),
                           '{0:.2f}'.format(figure.fn_profile[outlier])])
        logger.info(os.linesep + "List of detected outliers:")
        logger.info(table)

        contact_map_a = os.path.join(args.outdir, 'contact_map_a.mapalign')
        contact_map_b = os.path.join(args.outdir, 'contact_map_b.mapalign')
        conkit.io.write(contact_map_a, 'mapalign', prediction)
        conkit.io.write(contact_map_b, 'mapalign', model)

        map_align_cline = conkit.applications.MapAlignCommandline(
            cmd=args.map_align_exe,
            contact_map_a=contact_map_a,
            contact_map_b=contact_map_b,
            gap_opening_penalty=args.gap_opening_penalty,
            gap_extension_penalty=args.gap_extension_penalty,
            seq_separation_cutoff=args.seq_separation_cutoff,
            n_iterations=args.n_iterations
        )

        logger.info(os.linesep + "Executing: %s", map_align_cline)
        stdout, stderr = map_align_cline()
        map_align_log = os.path.join(args.outdir, 'map_align.log')
        with open(map_align_log, 'w') as fhandle:
            fhandle.write(stdout)
        alignment = parse_map_align_stdout(stdout)
        for idx, outlier in enumerate(figure.outliers, 1):
            table = PrettyTable()
            table.field_names = ["Current Residue", "New Residue"]
            start_outlier = outlier - 20 if outlier > 20 else 0
            stop_outlier = outlier + 20 if outlier + 20 < len(sequence) else len(sequence)
            for resnum in range(start_outlier, stop_outlier + 1):
                if resnum in alignment.keys() and alignment[resnum] != resnum:
                    table.add_row(['{} ({})'.format(sequence.seq[resnum - 1], resnum),
                                   '{} ({})'.format(sequence.seq[alignment[resnum] - 1], alignment[resnum])])
            if table._rows:
                logger.info(os.linesep + "List of proposed changes to fix outlier no. {}:".format(idx))
                logger.info(table)
            else:
                logger.info(os.linesep + "Cannot find optimal re-alignment for outlier no. {}:".format(idx))

    else:
        logger.info("No outliers were detected, finishing now.")

if __name__ == "__main__":
    import sys
    import traceback

    try:
        main()
        sys.exit(0)
    except Exception as e:
        if not isinstance(e, SystemExit):
            msg = "".join(traceback.format_exception(*sys.exc_info()))
            logger.critical(msg)
        sys.exit(1)
<|endoftext|>


<commit_before>import boto3
from flask import Blueprint
from flask import jsonify
from flask import current_app

blueprint = Blueprint('content', __name__)

@blueprint.route('/content/<section>/<sub_section>/<fragment>')
def content(section, sub_section, fragment):
    """ Piece of Content
    ---
    tags:
      - content
    parameters:
      - name: section
        in: path
        description: The section that the content is for
        type: string
        required: true
      - name: sub_section
        in: path
        description: The sub_section that the content is for
        type: string
        required: true
      - name: fragment
        in: path
        description: The actual fragment of content
        type: string
        required: true
    responses:
      200:
        description: content data
        schema:
          $ref: '#/definitions/ContentData'
      404:
        description: Content not found
      default:
        description: Unexpected error
        schema:
          $ref: '#/definitions/ApiError'
    """
    config = current_app.config
    client = _getS3Client(config)
    key = "{sec}/{sub_sec}/{frag}{fext}".format(sec=section, sub_sec=sub_section, frag=fragment, fext=config['CONTENT_FILE_EXTENSTION'])
    kwargs = {"Bucket": config['BUCKET_NAME'], "Key": key}
    content_obj = client.get_object(**kwargs)
    content_length = content_obj['ContentLength']
    content_body = content_obj['Body'].read().decode('utf8')
    return jsonify({"bucket": config['BUCKET_NAME'], "key": key, "content_length": content_length, "body": content_body})

def _getS3Client(config):
    kwargs = {'service_name': 's3', 'aws_access_key_id': config['AWS_KEY'], 'aws_secret_access_key': config['AWS_SECRET']}
    if 'S3_URL' in config and config['S3_URL']:
        kwargs['endpoint_url'] = config['S3_URL']
    session = boto3.session.Session()
    return session.client(**kwargs)
<commit_msg>Code Style update as requested by @romabysen - nice catch, thanks!<commit_after>import boto3
from flask import Blueprint
from flask import jsonify
from flask import current_app

blueprint = Blueprint('content', __name__)

@blueprint.route('/content/<section>/<sub_section>/<fragment>')
def content(section, sub_section, fragment):
    """ Piece of Content
    ---
    tags:
      - content
    parameters:
      - name: section
        in: path
        description: The section that the content is for
        type: string
        required: true
      - name: sub_section
        in: path
        description: The sub_section that the content is for
        type: string
        required: true
      - name: fragment
        in: path
        description: The actual fragment of content
        type: string
        required: true
    responses:
      200:
        description: content data
        schema:
          $ref: '#/definitions/ContentData'
      404:
        description: Content not found
      default:
        description: Unexpected error
        schema:
          $ref: '#/definitions/ApiError'
    """
    config = current_app.config
    client = _get_s3_client(config)
    key = "{sec}/{sub_sec}/{frag}{fext}".format(sec=section, sub_sec=sub_section, frag=fragment, fext=config['CONTENT_FILE_EXTENSTION'])
    kwargs = {"Bucket": config['BUCKET_NAME'], "Key": key}
    content_obj = client.get_object(**kwargs)
    content_length = content_obj['ContentLength']
    content_body = content_obj['Body'].read().decode('utf8')
    return jsonify({"bucket": config['BUCKET_NAME'], "key": key, "content_length": content_length, "body": content_body})

def _get_s3_client(config):
    kwargs = {'service_name': 's3', 'aws_access_key_id': config['AWS_KEY'], 'aws_secret_access_key': config['AWS_SECRET']}
    if 'S3_URL' in config and config['S3_URL']:
        kwargs['endpoint_url'] = config['S3_URL']
    session = boto3.session.Session()
    return session.client(**kwargs)
<|endoftext|>


<commit_before>#!/usr/bin/env python
"""
                        __                             __ __
.----.-----.---.-.-----|  |_.-----.----.--------.-----|  |  |_
|  __|  _  |  _  |__ --|   _|  -__|   _|        |  -__|  |   _|
|____|_____|___._|_____|____|_____|__| |__|__|__|_____|__|____|
--IPython Shell for Interactive Exploration--------------------

Read, write, or fill ARM memory. Numbers are hex. Trailing _ is
short for 0000, leading _ adds 'pad' scratchpad RAM offset.
Internal _ are ignored so you can use them as separators.

    rd 1ff_ 100
    wr _ fe00
    fill _10 55aa_55aa 4
    ALSO: peek, poke, read_block, watch

Assemble and disassemble ARM instructions:

    dis 3100
    asm _4 mov r3, #0x14
    dis _4 10
    ALSO: assemble, disassemble, blx

Or compile and invoke C++ code:

    ec 0x42
    ec ((uint16_t*)pad)[40]++
    ALSO: compile, evalc

The 'defines' and 'includes' dicts keep things you can define
in Python but use when compiling C++ and ASM snippets:

    defines['buffer'] = pad + 0x10000
    includes += ['int slide(int x) { return x << 8; }']
    ec slide(buffer)
    asm _ ldr r0, =buffer; bx lr

You can script the device's SCSI interface too:

    sc c ac              # Backdoor signature
    sc 8 ff 00 ff        # Undocumented firmware version
    ALSO: sc_eject, sc_sense, sc_read, scsi_in, scsi_out

Happy hacking!
~MeS`14
"""

from IPython import embed
from IPython.core import magic
from IPython.core.magic_arguments import magic_arguments, argument, parse_argstring
from IPython.terminal.embed import InteractiveShellEmbed

import remote, struct, sys
from hilbert import hilbert
from dump import *
from code import *
from watch import *

from binascii import a2b_hex, b2a_hex

__all__ = ['hexint', 'ShellMagics', 'peek', 'poke', 'setup_hexoutput']

# Placeholder for global Device in the IPython shell
d = None

def setup_hexoutput(ipy):
    """Configure an IPython shell to display integers in hex by default.
       You can convert to decimal with str() still.
       """
    def handler(n, p, cycle):
        return p.text("0x%x" % n)

    formatter = ipy.display_formatter.formatters['text/plain']
    formatter.for_type(int, handler)

def hexint(s):
    """This takes a bunch of weird number formats, as explained in the module docs"""
    if s.startswith('_'):
        base = pad
    else:
        base = 0
    if s.endswith('_'):
        s += '0000'
    return base + int(s.replace('_',''), 16)

def hexint_tuple(s):
    """A tuple of colon-separated hexints"""
    return tuple(hexint(i) for i in s.split(':'))

def get_signature(d):
    """Return a 12-byte signature identifying the firmware patch."""
    return d.get_signature()

def scsi_out(d, cdb, data):
    """Send a low-level SCSI packet with outgoing data."""
    return d.scsi_out(cdb, data)

def scsi_in(d, cdb, size):
    """Send a low-level SCSI packet that expects incoming data."""
    return d.scsi_in(cdb, size)

def peek(d, address):
    """Read one 32-bit word from ARM memory. Return it as an unsigned integer."""
    return d.peek(address)

def poke(d, address, word):
    """Write one 32-bit word to ARM memory."""
    return d.poke(address, word);

def peek_byte(d, address):
    """Read one 8-bit byte from ARM memory. Return it as an unsigned integer."""
    return d.peek(address)

def poke_byte(d, address, byte):
    """Write one 8-bit byte to ARM memory."""
    return d.poke(address, byte);

def blx(d, address, r0=0):
    """Invoke a function with one argument word and two return words."""
    return d.blx(address, r0)

@magic.magics_class
class ShellMagics(magic.Magics):

    @magic.line_magic
    @magic_arguments()
    @argument('address', type=hexint)
    @argument('size', type=hexint, nargs='?', default=0x100, help='Address to read from. Hexadecimal. Not necessarily aligned.')

    def rd(self, line):
        """Read ARM memory block"""
        args = parse_argstring(self.rd, line)
        dump(d, args.address, args.size)

    @magic.line_cell_magic
    @magic_arguments()
    @argument('address', type=hexint, help='Hex address')
    @argument('word', type=hexint, nargs='*', help='Hex words')

    def wr(self, line, cell=''):
        """Write hex words into ARM memory"""
        args = parse_argstring(self.wr, line)
        args.word.extend(map(hexint, cell.split()))
        for i, w in enumerate(args.word):
            d.poke(args.address + i*4, w)

    @magic.line_magic
    @magic_arguments()
    @argument('address', type=hexint, help='Hex address, word aligned')
    @argument('word', type=hexint, help='Hex word')
    @argument('count', type=hexint, help='Hex wordcount')

    def fill(self, line):
        """Fill contiguous words in ARM memory with the same value"""
        args = parse_argstring(self.fill, line)
        for i in range(args.count):
            d.poke(args.address + i*4, args.word)

    @magic.line_magic
    @magic_arguments()
    @argument('address', type=hexint_tuple, nargs='*', help='Single hex address, or a range start:end including both endpoints')

    def watch(self, line):
        """Watch memory for changes"""
        args = parse_argstring(self.watch, line)
        changes = watch_scanner(d, args.address)
        try:
            for line in watch_tabulator(changes):
                print line
        except KeyboardInterrupt:
            pass

    @magic.line_magic
    @magic_arguments()
    @argument('address', type=hexint, help='Hex address')
    @argument('size', type=hexint, nargs='?', default=0x40, help='Hex byte count')

    def dis(self, line):
        """Disassemble ARM instructions"""
        args = parse_argstring(self.dis, line)
        print disassemble(d, args.address, args.size)

    @magic.line_cell_magic
    @magic_arguments()
    @argument('-b', '--base', type=int, default=0, help='First address in map')
    @argument('-s', '--scale', type=int, default=256, help='Scale in bytes per pixel')
    @argument('-w', '--width', type=int, default=4096, help='Size of square hilbert map, in pixels')
    @argument('x', type=int)
    @argument('y', type=int)

    def msl(self, line, cell=''):
        """Memsquare lookup"""
        args = parse_argstring(self.msl, line)
        print args.base, args.scale, args.width
        addr = args.base + args.scale * hilbert(args.x, args.y, args.width)
        print "%08x" % addr

    @magic.line_cell_magic
    @magic_arguments()
    @argument('address', type=hexint)
    @argument('code', nargs='*')

    def asm(self, line, cell=''):
        """Assemble one or more ARM instructions

        NOTE that the assembled instructions will be padded to the
        nearest 32-bit word.

        Use with line or cell mode:

            %asm address op; op; op

            %%asm address
            op
            op
            op
        """
        args = parse_argstring(self.asm, line)
        assert 0 == (args.address & 3)
        code = ' '.join(args.code) + '\n' + cell
        assemble(d, args.address, code)

    @magic.line_magic
    def ec(self, line):
        """Evaluate a 32-bit C++ expression on the target"""
        print "0x%08x" % evalc(d, line)

    @magic.line_magic
    @magic_arguments()
    @argument('len', type=hexint, help='Length of input transfer')
    @argument('cdb', type=hexint, nargs='*', help='Up to 12 SCSI CDB bytes')

    def sc(self, line, cell=''):
        """Send a low-level SCSI command with a 12-byte CDB"""
        args = parse_argstring(self.sc, line)
        cdb = struct.pack('12B', *(args.cdb + [0]*12)[:12])
        data = d.scsi_in(cdb, args.len)
        sys.stdout.write(hexdump(data))

    @magic.line_magic
    def sc_eject(self, line):
        """Ask the drive to eject its disc."""
        self.sc('0 1b 0 0 0 2')

    @magic.line_magic
    def sc_sense(self, line):
        """Send a Request Sense command."""
        self.sc('20 3 0 0 0 20')

    @magic.line_magic
    @magic_arguments()
    @argument('lba', type=hexint, help='Logical Block Address')
    @argument('length', type=hexint, nargs='?', default=1, help='Transfer length, in 2kb blocks')
    @argument('-f', type=str, default=None, metavar='FILE', help='Log binary data to a file also')

    def sc_read(self, line, cell=''):
        """Read blocks from the SCSI device."""
        args = parse_argstring(self.sc_read, line)
        cdb = struct.pack('>BBIIBB', 0xA8, 0, args.lba, args.length, 0, 0)
        data = d.scsi_in(cdb, args.length * 2048)
        sys.stdout.write(hexdump(data, log_file=args.f))

if __name__ == '__main__':
    d = remote.Device()
    shell = InteractiveShellEmbed()
    setup_hexoutput(shell)
    shell.register_magics(ShellMagics)
    shell.mainloop(display_banner = __doc__)
<commit_msg>Watch shell UI tweaks<commit_after>#!/usr/bin/env python
"""
                        __                             __ __
.----.-----.---.-.-----|  |_.-----.----.--------.-----|  |  |_
|  __|  _  |  _  |__ --|   _|  -__|   _|        |  -__|  |   _|
|____|_____|___._|_____|____|_____|__| |__|__|__|_____|__|____|
--IPython Shell for Interactive Exploration--------------------

Read, write, or fill ARM memory. Numbers are hex. Trailing _ is
short for 0000, leading _ adds 'pad' scratchpad RAM offset.
Internal _ are ignored so you can use them as separators.

    rd 1ff_ 100
    wr _ fe00
    fill _10 55aa_55aa 4
    ALSO: peek, poke, read_block, watch

Assemble and disassemble ARM instructions:

    dis 3100
    asm _4 mov r3, #0x14
    dis _4 10
    ALSO: assemble, disassemble, blx

Or compile and invoke C++ code:

    ec 0x42
    ec ((uint16_t*)pad)[40]++
    ALSO: compile, evalc

The 'defines' and 'includes' dicts keep things you can define
in Python but use when compiling C++ and ASM snippets:

    defines['buffer'] = pad + 0x10000
    includes += ['int slide(int x) { return x << 8; }']
    ec slide(buffer)
    asm _ ldr r0, =buffer; bx lr

You can script the device's SCSI interface too:

    sc c ac              # Backdoor signature
    sc 8 ff 00 ff        # Undocumented firmware version
    ALSO: sc_eject, sc_sense, sc_read, scsi_in, scsi_out

Happy hacking!
~MeS`14
"""

from IPython import embed
from IPython.core import magic
from IPython.core.magic_arguments import magic_arguments, argument, parse_argstring
from IPython.terminal.embed import InteractiveShellEmbed

import remote, struct, sys
from hilbert import hilbert
from dump import *
from code import *
from watch import *

from binascii import a2b_hex, b2a_hex

__all__ = ['hexint', 'ShellMagics', 'peek', 'poke', 'setup_hexoutput']

# Placeholder for global Device in the IPython shell
d = None

def setup_hexoutput(ipy):
    """Configure an IPython shell to display integers in hex by default.
       You can convert to decimal with str() still.
       """
    def handler(n, p, cycle):
        return p.text("0x%x" % n)

    formatter = ipy.display_formatter.formatters['text/plain']
    formatter.for_type(int, handler)

def hexint(s):
    """This takes a bunch of weird number formats, as explained in the module docs"""
    if s.startswith('_'):
        base = pad
    else:
        base = 0
    if s.endswith('_'):
        s += '0000'
    return base + int(s.replace('_',''), 16)

def hexint_tuple(s):
    """A tuple of colon-separated hexints"""
    return tuple(hexint(i) for i in s.split(':'))

def get_signature(d):
    """Return a 12-byte signature identifying the firmware patch."""
    return d.get_signature()

def scsi_out(d, cdb, data):
    """Send a low-level SCSI packet with outgoing data."""
    return d.scsi_out(cdb, data)

def scsi_in(d, cdb, size):
    """Send a low-level SCSI packet that expects incoming data."""
    return d.scsi_in(cdb, size)

def peek(d, address):
    """Read one 32-bit word from ARM memory. Return it as an unsigned integer."""
    return d.peek(address)

def poke(d, address, word):
    """Write one 32-bit word to ARM memory."""
    return d.poke(address, word);

def peek_byte(d, address):
    """Read one 8-bit byte from ARM memory. Return it as an unsigned integer."""
    return d.peek(address)

def poke_byte(d, address, byte):
    """Write one 8-bit byte to ARM memory."""
    return d.poke(address, byte);

def blx(d, address, r0=0):
    """Invoke a function with one argument word and two return words."""
    return d.blx(address, r0)

@magic.magics_class
class ShellMagics(magic.Magics):

    @magic.line_magic
    @magic_arguments()
    @argument('address', type=hexint)
    @argument('size', type=hexint, nargs='?', default=0x100, help='Address to read from. Hexadecimal. Not necessarily aligned.')

    def rd(self, line):
        """Read ARM memory block"""
        args = parse_argstring(self.rd, line)
        dump(d, args.address, args.size)

    @magic.line_cell_magic
    @magic_arguments()
    @argument('address', type=hexint, help='Hex address')
    @argument('word', type=hexint, nargs='*', help='Hex words')

    def wr(self, line, cell=''):
        """Write hex words into ARM memory"""
        args = parse_argstring(self.wr, line)
        args.word.extend(map(hexint, cell.split()))
        for i, w in enumerate(args.word):
            d.poke(args.address + i*4, w)

    @magic.line_magic
    @magic_arguments()
    @argument('address', type=hexint, help='Hex address, word aligned')
    @argument('word', type=hexint, help='Hex word')
    @argument('count', type=hexint, help='Hex wordcount')

    def fill(self, line):
        """Fill contiguous words in ARM memory with the same value"""
        args = parse_argstring(self.fill, line)
        for i in range(args.count):
            d.poke(args.address + i*4, args.word)

    @magic.line_magic
    @magic_arguments()
    @argument('address', type=hexint_tuple, nargs='+', help='Single hex address, or a range start:end including both endpoints')

    def watch(self, line):
        """Watch memory for changes, shows the results in an ASCII data table.

        To use the results programmatically, see the watch_scanner() and
        watch_tabulator() functions.

        Keeps running until you kill it with a KeyboardInterrupt.
        """
        args = parse_argstring(self.watch, line)
        changes = watch_scanner(d, args.address)
        try:
            for line in watch_tabulator(changes):
                print line
        except KeyboardInterrupt:
            pass

    @magic.line_magic
    @magic_arguments()
    @argument('address', type=hexint, help='Hex address')
    @argument('size', type=hexint, nargs='?', default=0x40, help='Hex byte count')

    def dis(self, line):
        """Disassemble ARM instructions"""
        args = parse_argstring(self.dis, line)
        print disassemble(d, args.address, args.size)

    @magic.line_cell_magic
    @magic_arguments()
    @argument('-b', '--base', type=int, default=0, help='First address in map')
    @argument('-s', '--scale', type=int, default=256, help='Scale in bytes per pixel')
    @argument('-w', '--width', type=int, default=4096, help='Size of square hilbert map, in pixels')
    @argument('x', type=int)
    @argument('y', type=int)

    def msl(self, line, cell=''):
        """Memsquare lookup"""
        args = parse_argstring(self.msl, line)
        print args.base, args.scale, args.width
        addr = args.base + args.scale * hilbert(args.x, args.y, args.width)
        print "%08x" % addr

    @magic.line_cell_magic
    @magic_arguments()
    @argument('address', type=hexint)
    @argument('code', nargs='*')

    def asm(self, line, cell=''):
        """Assemble one or more ARM instructions

        NOTE that the assembled instructions will be padded to the
        nearest 32-bit word.

        Use with line or cell mode:

            %asm address op; op; op

            %%asm address
            op
            op
            op
        """
        args = parse_argstring(self.asm, line)
        assert 0 == (args.address & 3)
        code = ' '.join(args.code) + '\n' + cell
        assemble(d, args.address, code)

    @magic.line_magic
    def ec(self, line):
        """Evaluate a 32-bit C++ expression on the target"""
        print "0x%08x" % evalc(d, line)

    @magic.line_magic
    @magic_arguments()
    @argument('len', type=hexint, help='Length of input transfer')
    @argument('cdb', type=hexint, nargs='*', help='Up to 12 SCSI CDB bytes')

    def sc(self, line, cell=''):
        """Send a low-level SCSI command with a 12-byte CDB"""
        args = parse_argstring(self.sc, line)
        cdb = struct.pack('12B', *(args.cdb + [0]*12)[:12])
        data = d.scsi_in(cdb, args.len)
        sys.stdout.write(hexdump(data))

    @magic.line_magic
    def sc_eject(self, line):
        """Ask the drive to eject its disc."""
        self.sc('0 1b 0 0 0 2')

    @magic.line_magic
    def sc_sense(self, line):
        """Send a Request Sense command."""
        self.sc('20 3 0 0 0 20')

    @magic.line_magic
    @magic_arguments()
    @argument('lba', type=hexint, help='Logical Block Address')
    @argument('length', type=hexint, nargs='?', default=1, help='Transfer length, in 2kb blocks')
    @argument('-f', type=str, default=None, metavar='FILE', help='Log binary data to a file also')

    def sc_read(self, line, cell=''):
        """Read blocks from the SCSI device."""
        args = parse_argstring(self.sc_read, line)
        cdb = struct.pack('>BBIIBB', 0xA8, 0, args.lba, args.length, 0, 0)
        data = d.scsi_in(cdb, args.length * 2048)
        sys.stdout.write(hexdump(data, log_file=args.f))

if __name__ == '__main__':
    d = remote.Device()
    shell = InteractiveShellEmbed()
    setup_hexoutput(shell)
    shell.register_magics(ShellMagics)
    shell.mainloop(display_banner = __doc__)
<|endoftext|>


<commit_before>HEADER_INDEX = 0
BEGIN_INDEX = 1
STOP = "stop"
ERROR = "error"

class RelationalData (object):
    """This class represents a relational data set and offers comparison to another.

    This class expects to be handed a list-like object of same-length list-
    like objects. The first entry represents names that correspond to the items
    of the same index. One such name should correspond to a unique entry per each
    row (as compared to that entry of each other row), in other words, the primary
    key.

    For example:

    [
    ('id', 'type', 'value'),
    (1, 'cat', 5),
    (2, 'dog', 3),
    ]

    One means of achieving such a data set is by querying a SQL database with
    SQLAlchemy and prepending the column names. We can then use the comparison
    methods to determine the differences between the data sets (each a
    RelationalData object).
    """

    def __init__(self, data, pkey):
        """Initialization expects data of lists of lists and the key, a string.

        The formats expected are described more in-depth in the class docstring."""
        self.data = data
        # Compute length once so it does not have to be recomputed every time
        # matching_row is called.
        self.length = len(data)
        self.pkey = pkey
        # TODO will this be slow? Would it be (generally) faster to check if it's sorted first?
        self.sort_by_pkey()

        self.headers = data[HEADER_INDEX]

    def sort_by_pkey(self):
        """Sort the data by primary key.

        We make a later assumption that our data is in sorted by primary key order.
        Later methods will fail if sorting has not been completed."""
        self.data = sorted(self.data, key=lambda row: self.pkey_val(row))

    def val(self, row, attr_name):
        """Return the value of the row associated with the given header name."""
        i = self.headers.index(attr_name)
        return row[i]

    def pkey_val(self, row):
        """Return the primary key value of the given row."""
        return self.val(row, self.pkey)

    def compare(self, comparand):
        """Compare this relational data set to the comparand relational data set.

        This comparison is currently limited by the following assumptions:

        1) Only columns with identical names will be compared. At some point
           in the future, we may add the ability to specify column name
           equivalences.
        2) A given row will only be compared to a row that has the same primary
           key value. A row that does not have a corresponding row in the
           compared data set will be entered entirely into the errors."""
        # This is the beginning index of the comparand data, starting after the
        # header index.
        comparand_index = BEGIN_INDEX
        self.comparand = comparand
        self.shared_headers = set(self.headers) & set(comparand.headers)

        self.errors = []
        comparand.errors = []

        for row in self.data[BEGIN_INDEX:]:
            # I don't love that we return a tuple here, but I prefer it to setting
            # a variable on the comparand. Can we do better?
            comparand_row, comparand_index = \
                comparand.matching_row(self.pkey_val(row), comparand_index)

            if comparand_row == STOP:
                break
            elif comparand_row == ERROR:
                self.errors.append({'missing_row': row})
            else:
                self.compare_row(row, comparand_row)
                comparand_index += 1

        for row in comparand.data[comparand_index:]:
            self.errors.append({'comparand_missing_row': row})

    def matching_row(self, pkey_to_match, i):
        """Return the row that has the given primary key value.

        This method relies on the sort by primary key we did at initialization.
        If we reach the end the data before finding a match, return instructions
        to cease comparisons. Add each row that has a key less than the given key
        to our errors. If we reach a row that has a key greater than the given key,
        return instructions to add the originating row to the errors."""
        # We assume here that the primary key columns are integers.
        if i == self.length:
            return STOP, i

        row = self.data[i]
        row_pkey = self.val(row, self.pkey)

        if row_pkey < pkey_to_match:
            # As this method is typically called from the comparand, how can we set
            # errors on the original item, and not the comparand?
            self.errors.append({'missing_row': row})
            i += 1
            return self.matching_row(pkey_to_match, i)
        elif row_pkey > pkey_to_match:
            return ERROR, i
        else:
            return row, i

    def compare_row(self, row, comparand_row):
        """Compare the values of each row for each shared header.

        Add both values to our errors if there is a mismatch."""
        for header in self.shared_headers:
            val = self.val(row, header)
            comparand_val = self.comparand.val(comparand_row, header)

            if val != comparand_val:
                self.errors.append({
                    '{}_{}_original'.format(header, self.pkey_val(row)): val,
                    '{}_{}_comparand'.format(header,
                        self.pkey_val(comparand_row)): comparand_val,
                    })<commit_msg>REF: move sort_by_pkey after headers instantiation and make sure the headers are not included in the sorting.<commit_after>HEADER_INDEX = 0
BEGIN_INDEX = 1
STOP = "stop"
ERROR = "error"

class RelationalData (object):
    """This class represents a relational data set and offers comparison to another.

    This class expects to be handed a list-like object of same-length list-
    like objects. The first entry represents names that correspond to the items
    of the same index. One such name should correspond to a unique entry per each
    row (as compared to that entry of each other row), in other words, the primary
    key.

    For example:

    [
    ('id', 'type', 'value'),
    (1, 'cat', 5),
    (2, 'dog', 3),
    ]

    One means of achieving such a data set is by querying a SQL database with
    SQLAlchemy and prepending the column names. We can then use the comparison
    methods to determine the differences between the data sets (each a
    RelationalData object).
    """

    def __init__(self, data, pkey):
        """Initialization expects data of lists of lists and the key, a string.

        The formats expected are described more in-depth in the class docstring."""
        self.data = data
        # Compute length once so it does not have to be recomputed every time
        # matching_row is called.
        self.length = len(data)
        self.pkey = pkey

        self.headers = data[HEADER_INDEX]

        # TODO will this be slow? Would it be (generally) faster to check if it's sorted first?
        self.sort_by_pkey()

    def sort_by_pkey(self):
        """Sort the data by primary key.

        We make a later assumption that our data is in sorted by primary key order.
        Later methods will fail if sorting has not been completed."""
        sorted_no_header = sorted(self.data[BEGIN_INDEX:], key=lambda row: self.pkey_val(row))
        self.data = [self.data[HEADER_INDEX]] + sorted_no_header

    def val(self, row, attr_name):
        """Return the value of the row associated with the given header name."""
        i = self.headers.index(attr_name)
        return row[i]

    def pkey_val(self, row):
        """Return the primary key value of the given row."""
        return self.val(row, self.pkey)

    def compare(self, comparand):
        """Compare this relational data set to the comparand relational data set.

        This comparison is currently limited by the following assumptions:

        1) Only columns with identical names will be compared. At some point
           in the future, we may add the ability to specify column name
           equivalences.
        2) A given row will only be compared to a row that has the same primary
           key value. A row that does not have a corresponding row in the
           compared data set will be entered entirely into the errors."""
        # This is the beginning index of the comparand data, starting after the
        # header index.
        comparand_index = BEGIN_INDEX
        self.comparand = comparand
        self.shared_headers = set(self.headers) & set(comparand.headers)

        self.errors = []
        comparand.errors = []

        for row in self.data[BEGIN_INDEX:]:
            # I don't love that we return a tuple here, but I prefer it to setting
            # a variable on the comparand. Can we do better?
            comparand_row, comparand_index = \
                comparand.matching_row(self.pkey_val(row), comparand_index)

            if comparand_row == STOP:
                break
            elif comparand_row == ERROR:
                self.errors.append({'missing_row': row})
            else:
                self.compare_row(row, comparand_row)
                comparand_index += 1

        for row in comparand.data[comparand_index:]:
            self.errors.append({'comparand_missing_row': row})

    def matching_row(self, pkey_to_match, i):
        """Return the row that has the given primary key value.

        This method relies on the sort by primary key we did at initialization.
        If we reach the end the data before finding a match, return instructions
        to cease comparisons. Add each row that has a key less than the given key
        to our errors. If we reach a row that has a key greater than the given key,
        return instructions to add the originating row to the errors."""
        # We assume here that the primary key columns are integers.
        if i == self.length:
            return STOP, i

        row = self.data[i]
        row_pkey = self.val(row, self.pkey)

        if row_pkey < pkey_to_match:
            # As this method is typically called from the comparand, how can we set
            # errors on the original item, and not the comparand?
            self.errors.append({'missing_row': row})
            i += 1
            return self.matching_row(pkey_to_match, i)
        elif row_pkey > pkey_to_match:
            return ERROR, i
        else:
            return row, i

    def compare_row(self, row, comparand_row):
        """Compare the values of each row for each shared header.

        Add both values to our errors if there is a mismatch."""
        for header in self.shared_headers:
            val = self.val(row, header)
            comparand_val = self.comparand.val(comparand_row, header)

            if val != comparand_val:
                self.errors.append({
                    '{}_{}_original'.format(header, self.pkey_val(row)): val,
                    '{}_{}_comparand'.format(header,
                        self.pkey_val(comparand_row)): comparand_val,
                    })<|endoftext|>