<commit_before>import matplotlib.pyplot as plt
import copy
import sys
import numpy as np

def set_node_position(G, **kwargs):
    """
    Moves node to position along with its edges.

    Parameters
    ----------
    G        : Graph whose nodes are to be moved

    kwargs
    ------
    node     : The node to move if not specified moves all the nodes to the
               corresponding position in position
    position : The new position of the node. If node specified should be in
               form (x,y) or [x,y]. If node not specified should be in form
               [(x1,y1), (x2,y2), (x3,y3)]. It must have new position for
               each node in the graph.

    Example
    -------
    set_node_position(G, node=1, position=[0.5, 0.5])
    set_node_position(G, node=1, position=(0.5, 0.5))
    #Considering graph G has 3 nodes
    set_node_position(G, position=[(0.1,0.1), (0.2, 0.2), (0.3, 0.3)]
    """
    try:
        node = kwargs['node']
    except KeyError:
        node = None
    try:
        position = kwargs['position']
    except KeyError:
        print("postion argument required")
        sys.exit(0)

    if node:
        _set_single_node_position(G, node=node, position=position)
    else:
        for node_no in range(G.number_of_nodes()):
            _set_single_node_position(G, node=G.nodes()[node_no],
                                      position=position[node_no])

def _set_single_node_position(G, **kwargs):
    """
    Moves a single node specified in node to the specified position.
    """
    try:
        node = kwargs['node']
    except KeyError:
        node = None
    try:
        position = kwargs['position']
    except KeyError:
        print("position argument required")
        sys.exit(0)

    axes = plt.gca()
    node_index = G.nodes().index(node)
    no_of_nodes = G.number_of_nodes()

    def compare(li1, li2):
        for i, j in zip(li1, li2):
            if i != j:
                return False
        return True

    if node:
        # move node
        x, y = position
        nodes_collection = axes.get_children()[no_of_nodes + 2]
        prev_positions = nodes_collection.get_offsets()
        prev_position = copy.deepcopy(prev_positions[node_index])
        prev_positions[node_index] = [x, y]

        # move edge
        edges_collection = axes.get_children()[no_of_nodes + 3]
        for edge in edges_collection.get_paths():
            if compare(edge.vertices[0], prev_position):
                edge.vertices[0] = position
            elif compare(edge.vertices[1], prev_position):
                edge.vertices[1] = position

        # move label
        label = axes.get_children()[node_index + 2]
        label.set_position([x, y])

    plt.draw()

def set_node_facecolor(G, **kwargs):
    """
    Changes node's facecolor

    Parameters
    ----------
    G     : Graph whose node's facecolor needs to be changed

    kwargs
    ------
    node  : The node whose facecolor needs to be changed. If node not specified
            and only one color is specified changes all the nodes to that
            facecolor. If node not specified and color is a list or tuple of
            colors, each node will be changed to its corresponding color.

    color : Can be a single value of a list of colors.

    The following color abbreviations are supported:
    ==========  ========
    character   color
    ==========  ========
    'b'         blue
    'g'         green
    'r'         red
    'c'         cyan
    'm'         magenta
    'y'         yellow
    'k'         black
    'w'         white
    ==========  ========

    Examples:
    set_node_facecolor(G, node=1, color='b')
    set_node_facecolor(G, node=1, color=['b'])
    set_node_facecolor(G, color=['r', 'b', 'g', 'y', 'm', 'k', 'c'])
    set_node_facecolor(G, color='b')
    set_node_facecolor(G, color=('b'))
    """
    try:
        node = kwargs['node']
    except KeyError:
        node = None
    try:
        color = kwargs['color']
    except KeyError:
        print("color argument required")
        sys.exit(0)

    if (isinstance(color, list) or isinstance(color, tuple)) and \
            len(color) == 1:
        color = color[0]

    same_colors = {'blue': 'b',
                   'green': 'g',
                   'red': 'r',
                   'cyan': 'c',
                   'magenta': 'm',
                   'yellow': 'y',
                   'black': 'k',
                   'white': 'w'}

    if isinstance(color, str) and len(color) != 1:
        color = same_colors[color]

    colors_dict = {'b': (0.,    0.,    1.,  1.),
                   'g': (0.,    0.5,   0.,  1.),
                   'r': (1.,    0.,    0.,  1.),
                   'c': (0.,    0.75,  0.75,  1.),
                   'm': (0.75,  0.,    0.75,  1.),
                   'y': (0.75,  0.75,  0.,  1.),
                   'k': (0.,    0.,    0.,  1.),
                   'w': (1.,    1.,    1.,  1.)}

    fig = plt.gcf()
    axes = plt.gca()
    no_of_nodes = G.number_of_nodes()
    nodes_collection = axes.get_children()[no_of_nodes + 2]

    # if node is specified manually changing the facecolor array
    # so that the colors for other nodes are retained
    if node:
        node_index = G.nodes().index(node)
        facecolor_array = nodes_collection.get_facecolor().tolist()
        facecolor_array = [tuple(x) for x in facecolor_array]
        if len(facecolor_array) == 1:
            facecolor_array = [copy.deepcopy(facecolor_array[0])
                               for i in range(no_of_nodes)]
        facecolor_array[node_index] = colors_dict[color]
        nodes_collection.set_facecolor(facecolor_array)

    # if node not specified call the matplotlib's set_facecolor function
    else:
        nodes_collection.set_facecolor(color)

    plt.draw()

def set_node_edgecolor(G, **kwargs):
    """
    Changes node edgecolors

    Parameters
    ----------
    G     : Graph whose node's edgecolor needs to be changed

    kwargs
    ------
    node  : The node whose edgecolor needs to be changed. If node not specified
            and only one color is specified changes all the nodes to that
            color. If node not specified and color is a list or tuple of
            colors, each node will be changed to its corresponding color.

    color : Can be a single value of a list of colors.

    The following color abbreviations are supported:
    ==========  ========
    character   color
    ==========  ========
    'b'         blue
    'g'         green
    'r'         red
    'c'         cyan
    'm'         magenta
    'y'         yellow
    'k'         black
    'w'         white
    ==========  ========

    Examples:
    set_node_edgecolor(G, node=1, color='b')
    set_node_edgecolor(G, node=1, color=['b'])
    set_node_edgecolor(G, color=['r', 'b', 'g', 'y', 'm', 'k', 'c'])
    set_node_edgecolor(G, color='b')
    set_node_edgecolor(G, color=('b'))
    """
    try:
        node = kwargs['node']
    except KeyError:
        node = None
    try:
        color = kwargs['color']
    except KeyError:
        print("color argument required")
        sys.exit(0)

    if (isinstance(color, list) or isinstance(color, tuple)) and \
            len(color) == 1:
        color = color[0]

    same_colors = {'blue': 'b',
                   'green': 'g',
                   'red': 'r',
                   'cyan': 'c',
                   'magenta': 'm',
                   'yellow': 'y',
                   'black': 'k',
                   'white': 'w'}

    if isinstance(color, str) and len(color) != 1:
        color = same_colors[color]

    colors_dict = {'b': (0.,    0.,    1.,  1.),
                   'g': (0.,    0.5,   0.,  1.),
                   'r': (1.,    0.,    0.,  1.),
                   'c': (0.,    0.75,  0.75,  1.),
                   'm': (0.75,  0.,    0.75,  1.),
                   'y': (0.75,  0.75,  0.,  1.),
                   'k': (0.,    0.,    0.,  1.),
                   'w': (1.,    1.,    1.,  1.)}

    fig = plt.gcf()
    axes = plt.gca()
    no_of_nodes = G.number_of_nodes()
    nodes_collection = axes.get_children()[no_of_nodes + 2]

    # if node is specified manually changing the edgecolor array
    # so that the colors for other nodes are retained
    if node:
        node_index = G.nodes().index(node)
        edgecolor_array = nodes_collection.get_edgecolor().tolist()
        edgecolor_array = [tuple(x) for x in edgecolor_array]
        if len(edgecolor_array) == 1:
            edgecolor_array = [copy.deepcopy(edgecolor_array[0])
                               for i in range(no_of_nodes)]
        edgecolor_array[node_index] = colors_dict[color]
        nodes_collection.set_edgecolor(edgecolor_array)

    # if node not specified call the matplotlib's set_edgecolor function
    else:
        nodes_collection.set_edgecolor(color)

    plt.draw()

def set_node_color(G, **kwargs):
    """
    Changes node colors

    Parameters
    ----------
    G     : Graph whose node's color needs to be changed

    kwargs
    ------
    node  : The node whose color needs to be changed. If node not specified
            and only one color is specified changes all the nodes to that
            color. If node not specified and color is a list or tuple of
            colors, each node will be changed to its corresponding color.

    color : Can be a single value of a list of colors.

    The following color abbreviations are supported:
    ==========  ========
    character   color
    ==========  ========
    'b'         blue
    'g'         green
    'r'         red
    'c'         cyan
    'm'         magenta
    'y'         yellow
    'k'         black
    'w'         white
    ==========  ========

    Examples:
    set_node_color(G, node=1, color='b')
    set_node_color(G, node=1, color=['b'])
    set_node_color(G, color=['r', 'b', 'g', 'y', 'm', 'k', 'c'])
    set_node_color(G, color='b')
    set_node_color(G, color=('b'))
    """
    try:
        node = kwargs['node']
    except KeyError:
        node = None
    try:
        color = kwargs['color']
    except KeyError:
        print("color argument required")
        sys.exit(0)

    if (isinstance(color, list) or isinstance(color, tuple)) and \
            len(color) == 1:
        color = color[0]

    same_colors = {'blue': 'b',
                   'green': 'g',
                   'red': 'r',
                   'cyan': 'c',
                   'magenta': 'm',
                   'yellow': 'y',
                   'black': 'k',
                   'white': 'w'}

    if isinstance(color, str) and len(color) != 1:
        color = same_colors[color]

    fig = plt.gcf()
    axes = plt.gca()
    no_of_nodes = G.number_of_nodes()
    nodes_collection = axes.get_children()[no_of_nodes + 2]
    if node:
        set_node_facecolor(G, **kwargs)
        set_node_edgecolor(G, **kwargs)
    else:
        nodes_collection.set_color(color)

    plt.draw()

def set_node_size(G, **kwargs):
    """
    Change node size

    Parameters
    ----------
    G    : The graph whose nodes are to be resized

    kwargs
    ------
    node : The node whose size is to be changed. If node not specified size
           of all the nodes will be changed.
    size : If size is an int and node specified, that node's size would be
           changed. If size is a list of tuple and no node specified each
           node will be changed to it's corresponding size. If node not
           specified and size is int all the nodes will be changed to that
           value.

    Examples
    --------
    set_node_size(G, node=1, size=10)
    set_node_size(G, size=5)
    set_node_size(G, size=[1,2,3,4,5,6,7])
    """
    try:
        node = kwargs['node']
    except KeyError:
        node = None
    try:
        size = kwargs['size']
    except KeyError:
        print("size argument required")
        sys.exit(0)

    fig = plt.gcf()
    axes = plt.gca()
    no_of_nodes = G.number_of_nodes()
    nodes_collection = axes.get_children()[no_of_nodes + 2]

    if node:
        node_index = G.nodes().index(node)
        line_width_arr = nodes_collection.get_linewidth()
        if len(line_width_arr) == 1:
            element = line_width_arr[0]
            line_width_arr = [copy.deepcopy(element)
                              for i in range(no_of_nodes)]
        line_width_arr[node_index] = size
        nodes_collection.set_linewidth(line_width_arr)

    else:
        nodes_collection.set_linewidth(size)

    plt.draw()

def set_node_style(G, **kwargs):
    try:
        node = kwargs['node']
    except KeyError:
        node = None
    try:
        style = kwargs['position']
    except KeyError:
        print("style argument required")
        sys.exit(0)

    fig = plt.gcf()
    axes = plt.gca()
    no_of_nodes = G.number_of_nodes()
    nodes_collection = axes.get_children()[no_of_nodes + 2]

    linestyle_dict = {'solid':   (None, None),
                      'dotted':  (0, (1.0, 3.0)),
                      'dashdot': (0, (3.0, 5.0, 1.0, 5.0))}
    # TODO: Still incomplete Check what does the tuple values do

def set_node_alpha(G, alpha):
    fig = plt.gcf()
    axes = plt.gca()
    no_of_nodes = G.number_of_nodes()
    nodes_collection = axes.get_children()[no_of_nodes + 2]

    nodes_collection.set_alpha(alpha)
    plt.draw()

# TODO set_edge_style

def set_edge_color(G, **kwargs):
    """
    Change the color of the edges

    Parameters
    ----------
    G    : The graph whose nodes are to be resized

    kwargs
    ------
    edge : The edge whose color is to be changed. If edge
           not specified and a single color value is given
           all the edges are changed to the specified color.
           If edge not given and color list is given the
           color of the corresponding edge is changed. The
           order of the nodes can be seen from G.edges()

    color: Color value. Either a list of a single value
    """
    try:
        edge = kwargs['edge']
    except:
        edge = None
    try:
        color = kwargs['color']
    except:
        print("color argument required")
        sys.exit(0)

    fig = plt.gcf()
    axes = plt.gca()
    no_of_edges = G.number_of_edges()
    edges_collection = axes.get_children()[no_of_edges + 3]

    if edge:
#        TODO: complete this
    else:
        edges_collection.set_color(color)

# TODO set_edge_linewidth

# TODO set_edge_alpha

# problem with mapping edges to the one's that matplotlib selects by default
<commit_msg>add _get_node_collection_object<commit_after>import matplotlib.pyplot as plt
import copy
import sys
import numpy as np

def _get_node_collection_object(G):
    fig = plt.gcf()
    axes = plt.gca()
    no_of_nodes = G.number_of_nodes()
    nodes_collection = axes.get_children()[no_of_nodes + 2]
    return (fig, axes, no_of_nodes, nodes_collection)

def set_node_position(G, **kwargs):
    """
    Moves node to position along with its edges.

    Parameters
    ----------
    G        : Graph whose nodes are to be moved

    kwargs
    ------
    node     : The node to move if not specified moves all the nodes to the
               corresponding position in position
    position : The new position of the node. If node specified should be in
               form (x,y) or [x,y]. If node not specified should be in form
               [(x1,y1), (x2,y2), (x3,y3)]. It must have new position for
               each node in the graph.

    Example
    -------
    set_node_position(G, node=1, position=[0.5, 0.5])
    set_node_position(G, node=1, position=(0.5, 0.5))
    #Considering graph G has 3 nodes
    set_node_position(G, position=[(0.1,0.1), (0.2, 0.2), (0.3, 0.3)]
    """
    try:
        node = kwargs['node']
    except KeyError:
        node = None
    try:
        position = kwargs['position']
    except KeyError:
        print("postion argument required")
        sys.exit(0)

    if node:
        _set_single_node_position(G, node=node, position=position)
    else:
        for node_no in range(G.number_of_nodes()):
            _set_single_node_position(G, node=G.nodes()[node_no],
                                      position=position[node_no])

def _set_single_node_position(G, **kwargs):
    """
    Moves a single node specified in node to the specified position.
    """
    try:
        node = kwargs['node']
    except KeyError:
        node = None
    try:
        position = kwargs['position']
    except KeyError:
        print("position argument required")
        sys.exit(0)

    fig = plt.gcf()
    axes = plt.gca()
    node_index = G.nodes().index(node)
    no_of_nodes = G.number_of_nodes()

    def compare(li1, li2):
        for i, j in zip(li1, li2):
            if i != j:
                return False
        return True

    if node:
        # move node
        x, y = position
        nodes_collection = axes.get_children()[no_of_nodes + 2]
        prev_positions = nodes_collection.get_offsets()
        prev_position = copy.deepcopy(prev_positions[node_index])
        prev_positions[node_index] = [x, y]

        # move edge
        edges_collection = axes.get_children()[no_of_nodes + 3]
        for edge in edges_collection.get_paths():
            if compare(edge.vertices[0], prev_position):
                edge.vertices[0] = position
            elif compare(edge.vertices[1], prev_position):
                edge.vertices[1] = position

        # move label
        label = axes.get_children()[node_index + 2]
        label.set_position([x, y])

    plt.draw()

def set_node_facecolor(G, **kwargs):
    """
    Changes node's facecolor

    Parameters
    ----------
    G     : Graph whose node's facecolor needs to be changed

    kwargs
    ------
    node  : The node whose facecolor needs to be changed. If node not specified
            and only one color is specified changes all the nodes to that
            facecolor. If node not specified and color is a list or tuple of
            colors, each node will be changed to its corresponding color.

    color : Can be a single value of a list of colors.

    The following color abbreviations are supported:
    ==========  ========
    character   color
    ==========  ========
    'b'         blue
    'g'         green
    'r'         red
    'c'         cyan
    'm'         magenta
    'y'         yellow
    'k'         black
    'w'         white
    ==========  ========

    Examples:
    set_node_facecolor(G, node=1, color='b')
    set_node_facecolor(G, node=1, color=['b'])
    set_node_facecolor(G, color=['r', 'b', 'g', 'y', 'm', 'k', 'c'])
    set_node_facecolor(G, color='b')
    set_node_facecolor(G, color=('b'))
    """
    try:
        node = kwargs['node']
    except KeyError:
        node = None
    try:
        color = kwargs['color']
    except KeyError:
        print("color argument required")
        sys.exit(0)

    if (isinstance(color, list) or isinstance(color, tuple)) and \
            len(color) == 1:
        color = color[0]

    same_colors = {'blue': 'b',
                   'green': 'g',
                   'red': 'r',
                   'cyan': 'c',
                   'magenta': 'm',
                   'yellow': 'y',
                   'black': 'k',
                   'white': 'w'}

    if isinstance(color, str) and len(color) != 1:
        color = same_colors[color]

    colors_dict = {'b': (0.,    0.,    1.,  1.),
                   'g': (0.,    0.5,   0.,  1.),
                   'r': (1.,    0.,    0.,  1.),
                   'c': (0.,    0.75,  0.75,  1.),
                   'm': (0.75,  0.,    0.75,  1.),
                   'y': (0.75,  0.75,  0.,  1.),
                   'k': (0.,    0.,    0.,  1.),
                   'w': (1.,    1.,    1.,  1.)}

    fig = plt.gcf()
    axes = plt.gca()
    no_of_nodes = G.number_of_nodes()
    nodes_collection = axes.get_children()[no_of_nodes + 2]

    # if node is specified manually changing the facecolor array
    # so that the colors for other nodes are retained
    if node:
        node_index = G.nodes().index(node)
        facecolor_array = nodes_collection.get_facecolor().tolist()
        facecolor_array = [tuple(x) for x in facecolor_array]
        if len(facecolor_array) == 1:
            facecolor_array = [copy.deepcopy(facecolor_array[0])
                               for i in range(no_of_nodes)]
        facecolor_array[node_index] = colors_dict[color]
        nodes_collection.set_facecolor(facecolor_array)

    # if node not specified call the matplotlib's set_facecolor function
    else:
        nodes_collection.set_facecolor(color)

    plt.draw()

def set_node_edgecolor(G, **kwargs):
    """
    Changes node edgecolors

    Parameters
    ----------
    G     : Graph whose node's edgecolor needs to be changed

    kwargs
    ------
    node  : The node whose edgecolor needs to be changed. If node not specified
            and only one color is specified changes all the nodes to that
            color. If node not specified and color is a list or tuple of
            colors, each node will be changed to its corresponding color.

    color : Can be a single value of a list of colors.

    The following color abbreviations are supported:
    ==========  ========
    character   color
    ==========  ========
    'b'         blue
    'g'         green
    'r'         red
    'c'         cyan
    'm'         magenta
    'y'         yellow
    'k'         black
    'w'         white
    ==========  ========

    Examples:
    set_node_edgecolor(G, node=1, color='b')
    set_node_edgecolor(G, node=1, color=['b'])
    set_node_edgecolor(G, color=['r', 'b', 'g', 'y', 'm', 'k', 'c'])
    set_node_edgecolor(G, color='b')
    set_node_edgecolor(G, color=('b'))
    """
    try:
        node = kwargs['node']
    except KeyError:
        node = None
    try:
        color = kwargs['color']
    except KeyError:
        print("color argument required")
        sys.exit(0)

    if (isinstance(color, list) or isinstance(color, tuple)) and \
            len(color) == 1:
        color = color[0]

    same_colors = {'blue': 'b',
                   'green': 'g',
                   'red': 'r',
                   'cyan': 'c',
                   'magenta': 'm',
                   'yellow': 'y',
                   'black': 'k',
                   'white': 'w'}

    if isinstance(color, str) and len(color) != 1:
        color = same_colors[color]

    colors_dict = {'b': (0.,    0.,    1.,  1.),
                   'g': (0.,    0.5,   0.,  1.),
                   'r': (1.,    0.,    0.,  1.),
                   'c': (0.,    0.75,  0.75,  1.),
                   'm': (0.75,  0.,    0.75,  1.),
                   'y': (0.75,  0.75,  0.,  1.),
                   'k': (0.,    0.,    0.,  1.),
                   'w': (1.,    1.,    1.,  1.)}

    fig = plt.gcf()
    axes = plt.gca()
    no_of_nodes = G.number_of_nodes()
    nodes_collection = axes.get_children()[no_of_nodes + 2]

    # if node is specified manually changing the edgecolor array
    # so that the colors for other nodes are retained
    if node:
        node_index = G.nodes().index(node)
        edgecolor_array = nodes_collection.get_edgecolor().tolist()
        edgecolor_array = [tuple(x) for x in edgecolor_array]
        if len(edgecolor_array) == 1:
            edgecolor_array = [copy.deepcopy(edgecolor_array[0])
                               for i in range(no_of_nodes)]
        edgecolor_array[node_index] = colors_dict[color]
        nodes_collection.set_edgecolor(edgecolor_array)

    # if node not specified call the matplotlib's set_edgecolor function
    else:
        nodes_collection.set_edgecolor(color)

    plt.draw()

def set_node_color(G, **kwargs):
    """
    Changes node colors

    Parameters
    ----------
    G     : Graph whose node's color needs to be changed

    kwargs
    ------
    node  : The node whose color needs to be changed. If node not specified
            and only one color is specified changes all the nodes to that
            color. If node not specified and color is a list or tuple of
            colors, each node will be changed to its corresponding color.

    color : Can be a single value of a list of colors.

    The following color abbreviations are supported:
    ==========  ========
    character   color
    ==========  ========
    'b'         blue
    'g'         green
    'r'         red
    'c'         cyan
    'm'         magenta
    'y'         yellow
    'k'         black
    'w'         white
    ==========  ========

    Examples:
    set_node_color(G, node=1, color='b')
    set_node_color(G, node=1, color=['b'])
    set_node_color(G, color=['r', 'b', 'g', 'y', 'm', 'k', 'c'])
    set_node_color(G, color='b')
    set_node_color(G, color=('b'))
    """
    try:
        node = kwargs['node']
    except KeyError:
        node = None
    try:
        color = kwargs['color']
    except KeyError:
        print("color argument required")
        sys.exit(0)

    if (isinstance(color, list) or isinstance(color, tuple)) and \
            len(color) == 1:
        color = color[0]

    same_colors = {'blue': 'b',
                   'green': 'g',
                   'red': 'r',
                   'cyan': 'c',
                   'magenta': 'm',
                   'yellow': 'y',
                   'black': 'k',
                   'white': 'w'}

    if isinstance(color, str) and len(color) != 1:
        color = same_colors[color]

    fig = plt.gcf()
    axes = plt.gca()
    no_of_nodes = G.number_of_nodes()
    nodes_collection = axes.get_children()[no_of_nodes + 2]
    if node:
        set_node_facecolor(G, **kwargs)
        set_node_edgecolor(G, **kwargs)
    else:
        nodes_collection.set_color(color)

    plt.draw()

def set_node_size(G, **kwargs):
    """
    Change node size

    Parameters
    ----------
    G    : The graph whose nodes are to be resized

    kwargs
    ------
    node : The node whose size is to be changed. If node not specified size
           of all the nodes will be changed.
    size : If size is an int and node specified, that node's size would be
           changed. If size is a list of tuple and no node specified each
           node will be changed to it's corresponding size. If node not
           specified and size is int all the nodes will be changed to that
           value.

    Examples
    --------
    set_node_size(G, node=1, size=10)
    set_node_size(G, size=5)
    set_node_size(G, size=[1,2,3,4,5,6,7])
    """
    try:
        node = kwargs['node']
    except KeyError:
        node = None
    try:
        size = kwargs['size']
    except KeyError:
        print("size argument required")
        sys.exit(0)

    fig = plt.gcf()
    axes = plt.gca()
    no_of_nodes = G.number_of_nodes()
    nodes_collection = axes.get_children()[no_of_nodes + 2]

    if node:
        node_index = G.nodes().index(node)
        line_width_arr = nodes_collection.get_linewidth()
        if len(line_width_arr) == 1:
            element = line_width_arr[0]
            line_width_arr = [copy.deepcopy(element)
                              for i in range(no_of_nodes)]
        line_width_arr[node_index] = size
        nodes_collection.set_linewidth(line_width_arr)

    else:
        nodes_collection.set_linewidth(size)

    plt.draw()

def set_node_style(G, **kwargs):
    try:
        node = kwargs['node']
    except KeyError:
        node = None
    try:
        style = kwargs['position']
    except KeyError:
        print("style argument required")
        sys.exit(0)

    fig = plt.gcf()
    axes = plt.gca()
    no_of_nodes = G.number_of_nodes()
    nodes_collection = axes.get_children()[no_of_nodes + 2]

    linestyle_dict = {'solid':   (None, None),
                      'dotted':  (0, (1.0, 3.0)),
                      'dashdot': (0, (3.0, 5.0, 1.0, 5.0))}
    # TODO: Still incomplete Check what does the tuple values do

def set_node_alpha(G, alpha):
    fig = plt.gcf()
    axes = plt.gca()
    no_of_nodes = G.number_of_nodes()
    nodes_collection = axes.get_children()[no_of_nodes + 2]

    nodes_collection.set_alpha(alpha)
    plt.draw()

# TODO set_edge_style

def set_edge_color(G, **kwargs):
    """
    Change the color of the edges

    Parameters
    ----------
    G    : The graph whose nodes are to be resized

    kwargs
    ------
    edge : The edge whose color is to be changed. If edge
           not specified and a single color value is given
           all the edges are changed to the specified color.
           If edge not given and color list is given the
           color of the corresponding edge is changed. The
           order of the nodes can be seen from G.edges()

    color: Color value. Either a list of a single value
    """
    try:
        edge = kwargs['edge']
    except:
        edge = None
    try:
        color = kwargs['color']
    except:
        print("color argument required")
        sys.exit(0)

    fig = plt.gcf()
    axes = plt.gca()
    no_of_edges = G.number_of_edges()
    edges_collection = axes.get_children()[no_of_edges + 3]

    if edge:
        pass
#        TODO: complete this
    else:
        edges_collection.set_color(color)

# TODO set_edge_linewidth

# TODO set_edge_alpha

# problem with mapping edges to the one's that matplotlib selects by default
<|endoftext|>


<commit_before>
def MOS_ASYMPTOTE(states, LB="17.2450", VB="13.0207", title="Title", draw_band=False, draw_occupation=False, draw_energy=False):
    """

    :states: TODO
    :returns: TODO

    """
    energies    = "{"
    spins       = "{"
    occupations = "{"
    bands       = "{"

    STATES_STRING=""
    for state in states:
        energy     = str(state["energy"])
        spin       = str(state["spin"])
        occupation = str(state["occupation"])
        band       = str(state["number"])
        STATES_STRING+="""\
//energy, spin, occupation, band
state(%s, %s, %s, %s)
.setAutoPosition()
.draw( draw_band       = DRAW_BAND,
    draw_occupation = DRAW_OCCUPATION,
    draw_energy     = DRAW_ENERGY
);\n\n"""%(energy, spin, occupation, band)
    return """

string LUMO_TITLE="%s";

real ENERGIE_LB_PRISTINE   = %s ;
real ENERGIE_VB_PRISTINE   = %s ;

real OBERKANTE     = 100;
real UNTERKANTE    = 0;
real IMG_WIDTH     = 50;
real KANTEN_HEIGHT = 20;

bool DRAW_ENERGY     = %s;
bool DRAW_BAND       = %s;
bool DRAW_OCCUPATION = %s;

unitsize(.2cm);

struct state {
  static int state_count;
  real energy;
  real occupation;
  real band;
  real value;
  pen color;
  pen spin_color;
  string title           = "";
  real spin              = 0;
  real VB                = ENERGIE_VB_PRISTINE;
  real LB                = ENERGIE_LB_PRISTINE;
  real DASH_WIDTH        = 25;
  real DASH_HEIGHT       = 1.8;
  real X_COORD           = 0;
  real Y_OFFSET          = 0;
  real OCCUPATION_CUTOFF = 0.1;
  real getPlottingValue (){
    real val = 100*(energy - VB)/(LB-VB);
    return val + Y_OFFSET;
  };
  bool isOccupied(){
    if ( occupation >= OCCUPATION_CUTOFF ) {
      return true;
    } else {
      return false;
    }
  };
  state setStyle(){
    if ( isOccupied() ) {
      color = red;
    } else {
      color = gray;
    }
    pen unoccupied_style = 0.7*white+dashed;
    pen occupied_style   = black;
    if ( isOccupied() ) {
      spin_color = occupied_style+linewidth(3);
    } else {
      spin_color = unoccupied_style+linewidth(3);
    }
    return this;
  };
  void operator init(real energy, real spin, real occupation, real band){
    this.energy     = energy;
    this.occupation = occupation;
    this.band       = band;
    this.spin       = spin;
    this.value      = getPlottingValue();
    state_count    += 1;
    setStyle();
  };
  pair getMiddlePoint (  ){
    real x,y;
    x = X_COORD+(DASH_WIDTH)/2;
    y = value + (DASH_HEIGHT)/2;
    return (x,y);
  };
  state setColors (pen color, pen spin_color=this.spin_color){
    this.color = color;
    this.spin_color = spin_color;
    return this;
  };
  state setAutoPosition (){
    int controller = state_count%%2;
    X_COORD=0+controller*(DASH_WIDTH);
    return this;
  };
  bool  isLeft (){
    if ( getMiddlePoint().x >= IMG_WIDTH/2 ) {
      return false;
    }
    else{
      return true;
    }
  };
  bool isUp (){   return spin == 1?true:false; };
  bool isDown (){ return spin == 2?true:false; };
  pair getSpinPosition (bool up=false){
    real x_deviation = 0.25*DASH_WIDTH;
    pair middle      = getMiddlePoint();
    if (up) {
      return (middle - (-x_deviation,0));
    } else {
      return (middle + (-x_deviation,0));
    }
  };
  path getSpinArrow (){
    bool up = isUp();
    pair position = getSpinPosition(up);
    real height = 3*DASH_HEIGHT;
    if (isUp()) {
      return position - (0,height/2) -- position + (0,height/2);
    } else {
      return position + (0,height/2) -- position - (0,height/2);
    }
  };
  state draw_energy (){
    if ( isLeft() ) {
      label((string)energy, (X_COORD,value), W, red);
    } else {
      label((string)energy, (X_COORD+DASH_WIDTH, value), E, red);
    }
    return this;
  };
  state draw_spin(){
    path spinArrow = getSpinArrow();
    draw(spinArrow, spin_color,Arrow(15));
    return this;
  };
  state draw (
      bool draw_band       = false,
      bool draw_occupation = true,
      bool draw_energy     = true
      ){
    filldraw(
        box(
          (X_COORD,value)
          ,(X_COORD+DASH_WIDTH,value+DASH_HEIGHT)
          ),
        color,color*0
        );
    if ( draw_band )
      label(scale(1)*(string)band       , getMiddlePoint() - (DASH_WIDTH/4 , 0) , black);
    if ( draw_occupation && occupation != 0)
      label(scale(1)*(string)occupation , getSpinPosition(!isUp()) , black);
    if ( draw_energy ) draw_energy();
    if ( spin != 0 ) draw_spin();
    return this;
  };
};

//----------------------------
//-  Valence and Cond bands  -
//----------------------------

label(LUMO_TITLE, (25, 100+KANTEN_HEIGHT/1.1), 0.8*blue);

path UNTERKANTE_BOX = box((0 , UNTERKANTE) , (IMG_WIDTH , UNTERKANTE - KANTEN_HEIGHT));
path OBERKANTE_BOX  = box((0 , OBERKANTE)  , (IMG_WIDTH , OBERKANTE + KANTEN_HEIGHT));

pen bandStyle = .8*white;
filldraw(OBERKANTE_BOX  , bandStyle, bandStyle);
filldraw(UNTERKANTE_BOX , bandStyle, bandStyle);

/* DRAW STATES {{{1 */
/********************/

%s

//---------------
//-  SCALE  {{{1-
//---------------

real pointsToEnergy ( real point ){
  return (ENERGIE_LB_PRISTINE-ENERGIE_VB_PRISTINE)*point/100 + ENERGIE_VB_PRISTINE;
};
int steps = 2;
real width = 100/steps;

// Bandgap

draw((50,0)--(50,100),dashed+linewidth(.5), Bars(4mm));
label((string)(ENERGIE_LB_PRISTINE-ENERGIE_VB_PRISTINE)+" eV", (50,50), Fill(white));
label("VB", (IMG_WIDTH,0)+UNTERKANTE, N, Fill(white));
label("CB", (IMG_WIDTH,100)+UNTERKANTE, S, Fill(white));

// SCALE
draw((0,0)--(0,100), linewidth(1));
for ( int i = 0; i <= steps; i+=1 ) {
  // SCALE TICKS
  draw((0,width*i)--(2,width*i));
  // SCALE LABELS
  label(scale(0.7)*(string)pointsToEnergy(width*i), (1,width*i), E, Fill(white));
}

// vim: nospell
//vim-run: asy -f pdf %% && mupdf $(basename %% .asy).pdf &
//vim-run: asy -batchView

"""%(title,
        LB,
        VB,
        str(draw_energy).lower(),
        str(draw_band).lower(),
        str(draw_occupation).lower(),
        STATES_STRING
        )

<commit_msg>colors varying<commit_after>
def MOS_ASYMPTOTE(states, LB="17.2450", VB="13.0207", title="Title", draw_band=False, draw_occupation=False, draw_energy=False):
    """

    :states: TODO
    :returns: TODO

    """
    energies    = "{"
    spins       = "{"
    occupations = "{"
    bands       = "{"

    STATES_STRING=""
    for state in states:
        energy     = str(state["energy"])
        spin       = str(state["spin"])
        occupation = str(state["occupation"])
        band       = str(state["number"])
        STATES_STRING+="""\
//energy, spin, occupation, band
state(%s, %s, %s, %s)
.setAutoPosition()
.draw( draw_band       = DRAW_BAND,
    draw_occupation = DRAW_OCCUPATION,
    draw_energy     = DRAW_ENERGY
);\n\n"""%(energy, spin, occupation, band)
    return """

string LUMO_TITLE="%s";

real ENERGIE_LB_PRISTINE   = %s ;
real ENERGIE_VB_PRISTINE   = %s ;

real OBERKANTE     = 100;
real UNTERKANTE    = 0;
real IMG_WIDTH     = 50;
real KANTEN_HEIGHT = 20;

bool DRAW_ENERGY     = %s;
bool DRAW_BAND       = %s;
bool DRAW_OCCUPATION = %s;

unitsize(.2cm);

struct state {
  static int state_count;
  real energy;
  real occupation;
  real band;
  real value;
  pen color;
  pen spin_color;
  pen spin_occupied_color   = black;
  pen spin_unoccupied_color = 0.7*white+dashed;
  pen occupied_color   = red;
  pen unoccupied_color = gray;
  string title         = "";
  real spin            = 0;
  real VB              = ENERGIE_VB_PRISTINE;
  real LB              = ENERGIE_LB_PRISTINE;
  real DASH_WIDTH      = 25;
  real DASH_HEIGHT     = 1.8;
  real X_COORD         = 0;
  real Y_OFFSET        = 0;
  real OCCUPATION_CUTOFF = 0.1;
  real getPlottingValue (){
    real val = 100*(energy - VB)/(LB-VB);
    return val + Y_OFFSET;
  };
  bool isOccupied(){
    if ( occupation >= OCCUPATION_CUTOFF ) {
      return true;
    } else {
      return false;
    }
  };
  state setStyle(){
    real parameter;
    if ( spin != 0 ) {
      parameter = occupation;
    } else {
      parameter = occupation/2;
    }
    color = parameter*occupied_color+(1-parameter)*unoccupied_color;
    spin_color = parameter*spin_occupied_color+(1-parameter)*spin_unoccupied_color+linewidth(3);
    return this;
  };
  void operator init(real energy, real spin, real occupation, real band){
    this.energy     = energy;
    this.occupation = occupation;
    this.band       = band;
    this.spin       = spin;
    this.value      = getPlottingValue();
    state_count    += 1;
    setStyle();
  };
  pair getMiddlePoint (  ){
    real x,y;
    x = X_COORD+(DASH_WIDTH)/2;
    y = value + (DASH_HEIGHT)/2;
    return (x,y);
  };
  state setColors (pen color, pen spin_color=this.spin_color){
    this.color = color;
    this.spin_color = spin_color;
    return this;
  };
  state setAutoPosition (){
    int controller = state_count%2;
    X_COORD=0+controller*(DASH_WIDTH);
    return this;
  };
  bool  isLeft (){
    if ( getMiddlePoint().x >= IMG_WIDTH/2 ) {
      return false;
    }
    else{
      return true;
    }
  };
  bool isUp (){   return spin == 1?true:false; };
  bool isDown (){ return spin == 2?true:false; };
  pair getSpinPosition (bool up=false){
    real x_deviation = 0.25*DASH_WIDTH;
    pair middle      = getMiddlePoint();
    if (up) {
      return (middle - (-x_deviation,0));
    } else {
      return (middle + (-x_deviation,0));
    }
  };
  path getSpinArrow (){
    bool up = isUp();
    pair position = getSpinPosition(up);
    real height = 3*DASH_HEIGHT;
    if (isUp()) {
      return position - (0,height/2) -- position + (0,height/2);
    } else {
      return position + (0,height/2) -- position - (0,height/2);
    }
  };
  state draw_energy (){
    if ( isLeft() ) {
      label((string)energy, (X_COORD,value), W, red);
    } else {
      label((string)energy, (X_COORD+DASH_WIDTH, value), E, red);
    }
    return this;
  };
  state draw_spin(){
    path spinArrow = getSpinArrow();
    draw(spinArrow, spin_color,Arrow(15));
    return this;
  };
  state draw (
      bool draw_band       = false,
      bool draw_occupation = true,
      bool draw_energy     = true
      ){
    filldraw(
        box(
          (X_COORD,value)
          ,(X_COORD+DASH_WIDTH,value+DASH_HEIGHT)
          ),
        color,color*0
        );
    if ( draw_band )
      label(scale(1)*(string)band       , getMiddlePoint() - (DASH_WIDTH/4 , 0) , black);
    if ( draw_occupation && occupation != 0)
      label(scale(1)*(string)occupation , getSpinPosition(!isUp()) , black);
    if ( draw_energy ) draw_energy();
    if ( spin != 0 ) draw_spin();
    return this;
  };
};

//----------------------------
//-  Valence and Cond bands  -
//----------------------------

label(LUMO_TITLE, (25, 100+KANTEN_HEIGHT/1.1), 0.8*blue);

path UNTERKANTE_BOX = box((0 , UNTERKANTE) , (IMG_WIDTH , UNTERKANTE - KANTEN_HEIGHT));
path OBERKANTE_BOX  = box((0 , OBERKANTE)  , (IMG_WIDTH , OBERKANTE + KANTEN_HEIGHT));

pen bandStyle = .8*white;
filldraw(OBERKANTE_BOX  , bandStyle, bandStyle);
filldraw(UNTERKANTE_BOX , bandStyle, bandStyle);

/* DRAW STATES {{{1 */
/********************/

%s

//---------------
//-  SCALE  {{{1-
//---------------

real pointsToEnergy ( real point ){
  return (ENERGIE_LB_PRISTINE-ENERGIE_VB_PRISTINE)*point/100 + ENERGIE_VB_PRISTINE;
};
int steps = 2;
real width = 100/steps;

// Bandgap

draw((50,0)--(50,100),dashed+linewidth(.5), Bars(4mm));
label((string)(ENERGIE_LB_PRISTINE-ENERGIE_VB_PRISTINE)+" eV", (50,50), Fill(white));
label("VB", (IMG_WIDTH,0)+UNTERKANTE, N, Fill(white));
label("CB", (IMG_WIDTH,100)+UNTERKANTE, S, Fill(white));

// SCALE
draw((0,0)--(0,100), linewidth(1));
for ( int i = 0; i <= steps; i+=1 ) {
  // SCALE TICKS
  draw((0,width*i)--(2,width*i));
  // SCALE LABELS
  label(scale(0.7)*(string)pointsToEnergy(width*i), (1,width*i), E, Fill(white));
}

// vim: nospell
//vim-run: asy -f pdf %% && mupdf $(basename %% .asy).pdf &
//vim-run: asy -batchView

"""%(title,
        LB,
        VB,
        str(draw_energy).lower(),
        str(draw_band).lower(),
        str(draw_occupation).lower(),
        STATES_STRING
        )

<|endoftext|>


<commit_before>"""
Django-social-auth application, allows OpenId or OAuth user
registration/authentication just adding a few configurations.
"""
version = (0, 1, 6)
__version__ = '.'.join(map(str, version))
<commit_msg>v0.1.7<commit_after>"""
Django-social-auth application, allows OpenId or OAuth user
registration/authentication just adding a few configurations.
"""
version = (0, 1, 7)
__version__ = '.'.join(map(str, version))
<|endoftext|>


<commit_before>import json
import six
import functools

import django

from django.core.exceptions import ValidationError
from django.conf import settings
from django.db import models

try:
    from django.utils.encoding import smart_unicode as smart_text
    smart_text  # placate pyflakes
except ImportError:
    from django.utils.encoding import smart_text

# SubfieldBase causes RemovedInDjango110Warning in 1.8 and 1.9, and
# will not work in 1.10 or later
if django.VERSION[:2] >= (1, 8):
    field_metaclass = type
else:
    from django.db.models import SubfieldBase
    field_metaclass = SubfieldBase

field_class = functools.partial(six.with_metaclass, field_metaclass)

if getattr(settings, 'SOCIAL_DJANGO_POSTGRES_JSONFIELD', False):
    from django.contrib.postgres.fields import JSONField as JSONFieldBase
else:
    JSONFieldBase = field_class(models.TextField)

class JSONField(JSONFieldBase):
    """Simple JSON field that stores python structures as JSON strings
    on database.
    """

    def __init__(self, *args, **kwargs):
        kwargs.setdefault('default', {})
        super(JSONField, self).__init__(*args, **kwargs)

    def from_db_value(self, value, expression, connection, context):
        return self.to_python(value)

    def to_python(self, value):
        """
        Convert the input JSON value into python structures, raises
        django.core.exceptions.ValidationError if the data can't be converted.
        """
        if self.blank and not value:
            return {}
        value = value or '{}'
        if isinstance(value, six.binary_type):
            value = six.text_type(value, 'utf-8')
        if isinstance(value, six.string_types):
            try:
                # with django 1.6 i have '"{}"' as default value here
                if value[0] == value[-1] == '"':
                    value = value[1:-1]

                return json.loads(value)
            except Exception as err:
                raise ValidationError(str(err))
        else:
            return value

    def validate(self, value, model_instance):
        """Check value is a valid JSON string, raise ValidationError on
        error."""
        if isinstance(value, six.string_types):
            super(JSONField, self).validate(value, model_instance)
            try:
                json.loads(value)
            except Exception as err:
                raise ValidationError(str(err))

    def get_prep_value(self, value):
        """Convert value to JSON string before save"""
        try:
            return json.dumps(value)
        except Exception as err:
            raise ValidationError(str(err))

    def value_to_string(self, obj):
        """Return value from object converted to string properly"""
        return smart_text(self.get_prep_value(self._get_val_from_obj(obj)))

    def value_from_object(self, obj):
        """Return value dumped to string."""
        return self.get_prep_value(self._get_val_from_obj(obj))
<commit_msg>settings variable name changed to social_core.utils.setting_name mechanism<commit_after>import json
import six
import functools

import django

from django.core.exceptions import ValidationError
from django.conf import settings
from django.db import models

from social_core.utils import setting_name

try:
    from django.utils.encoding import smart_unicode as smart_text
    smart_text  # placate pyflakes
except ImportError:
    from django.utils.encoding import smart_text

# SubfieldBase causes RemovedInDjango110Warning in 1.8 and 1.9, and
# will not work in 1.10 or later
if django.VERSION[:2] >= (1, 8):
    field_metaclass = type
else:
    from django.db.models import SubfieldBase
    field_metaclass = SubfieldBase

field_class = functools.partial(six.with_metaclass, field_metaclass)

if getattr(settings, setting_name('POSTGRES_JSONFIELD'), False):
    from django.contrib.postgres.fields import JSONField as JSONFieldBase
else:
    JSONFieldBase = field_class(models.TextField)

class JSONField(JSONFieldBase):
    """Simple JSON field that stores python structures as JSON strings
    on database.
    """

    def __init__(self, *args, **kwargs):
        kwargs.setdefault('default', {})
        super(JSONField, self).__init__(*args, **kwargs)

    def from_db_value(self, value, expression, connection, context):
        return self.to_python(value)

    def to_python(self, value):
        """
        Convert the input JSON value into python structures, raises
        django.core.exceptions.ValidationError if the data can't be converted.
        """
        if self.blank and not value:
            return {}
        value = value or '{}'
        if isinstance(value, six.binary_type):
            value = six.text_type(value, 'utf-8')
        if isinstance(value, six.string_types):
            try:
                # with django 1.6 i have '"{}"' as default value here
                if value[0] == value[-1] == '"':
                    value = value[1:-1]

                return json.loads(value)
            except Exception as err:
                raise ValidationError(str(err))
        else:
            return value

    def validate(self, value, model_instance):
        """Check value is a valid JSON string, raise ValidationError on
        error."""
        if isinstance(value, six.string_types):
            super(JSONField, self).validate(value, model_instance)
            try:
                json.loads(value)
            except Exception as err:
                raise ValidationError(str(err))

    def get_prep_value(self, value):
        """Convert value to JSON string before save"""
        try:
            return json.dumps(value)
        except Exception as err:
            raise ValidationError(str(err))

    def value_to_string(self, obj):
        """Return value from object converted to string properly"""
        return smart_text(self.get_prep_value(self._get_val_from_obj(obj)))

    def value_from_object(self, obj):
        """Return value dumped to string."""
        return self.get_prep_value(self._get_val_from_obj(obj))
<|endoftext|>


<commit_before>"""Spatial functions included in ODDT
Mainly used by other modules, but can be accessed directly.
"""

from math import sin, cos
import numpy as np
from scipy.spatial.distance import cdist as distance

__all__ = ['angle', 'angle_2v', 'dihedral', 'distance']

# angle functions
def angle(p1,p2,p3):
    """Returns an angle from a series of 3 points (point #2 is centroid).Angle is returned in degrees.

    Parameters
    ----------
    p1,p2,p3 : numpy arrays, shape = [n_points, n_dimensions]
        Triplets of points in n-dimensional space, aligned in rows.

    Returns
    -------
    angles : numpy array, shape = [n_points]
        Series of angles in degrees
    """
    v1 = p1-p2
    v2 = p3-p2
    return angle_2v(v1,v2)

def angle_2v(v1, v2):
    """Returns an angle between two vecors.Angle is returned in degrees.

    Parameters
    ----------
    v1,v2 : numpy arrays, shape = [n_vectors, n_dimensions]
        Pairs of vectors in n-dimensional space, aligned in rows.

    Returns
    -------
    angles : numpy array, shape = [n_vectors]
        Series of angles in degrees
    """
    dot = (v1*v2).sum(axis=-1) # better than np.dot(v1, v2), multiple vectors can be applied
    norm = np.linalg.norm(v1, axis=-1)* np.linalg.norm(v2, axis=-1)
    return np.degrees(np.arccos(np.clip(dot/norm, -1, 1)))

def dihedral(p1,p2,p3,p4):
    """Returns an dihedral angle from a series of 4 points. Dihedral is returned in degrees.
    Function distingishes clockwise and antyclockwise dihedrals.

    Parameters
    ----------
    p1,p2,p3,p4 : numpy arrays, shape = [n_points, n_dimensions]
        Quadruplets of points in n-dimensional space, aligned in rows.

    Returns
    -------
    angles : numpy array, shape = [n_points]
        Series of angles in degrees
    """
    v12 = (p1-p2)/np.linalg.norm(p1-p2)
    v23 = (p2-p3)/np.linalg.norm(p2-p3)
    v34 = (p3-p4)/np.linalg.norm(p3-p4)
    c1 = np.cross(v12, v23)
    c2 = np.cross(v23, v34)
    out = angle_2v(c1, c2)
    # check clockwise and anticlockwise
    n1 = c1/np.linalg.norm(c1)
    mask = (n1*v34).sum(axis=-1) > 0
    if len(mask.shape) == 0:
        if mask:
            out = -out
    else:
        out[mask] = -out[mask]
    return out

def rmsd(ref, mol, ignore_h = True, canonize = False, normalize = False):
    """Computes root mean square deviation (RMSD) between two molecules (including or excluding Hydrogens). No symmetry checks are performed.

    Parameters
    ----------
    ref : oddt.toolkit.Molecule object
        Reference molecule for the RMSD calculation

    mol : oddt.toolkit.Molecule object
        Query molecule for RMSD calculation

    ignore_h : bool (default=False)
        Flag indicating to ignore Hydrogen atoms while performing RMSD calculation

    canonize : bool (default=False)
        Match heavy atoms using OB canonical ordering

    normalize : bool (default=False)
        Normalize RMSD by square root of rot. bonds

    Returns
    -------
    rmsd : float
        RMSD between two molecules
    """
    if ignore_h:
        if canonize:
            ref_hvy = ref.coords[ref.canonic_order]
            mol_hvy = mol.coords[mol.canonic_order]
        else:
            hvy_map = np.array([atom.idx-1 for atom in mol if atom.atomicnum != 1])
            mol_hvy = mol.coords[hvy_map]
            ref_hvy = ref.coords[hvy_map]
        if mol_hvy.shape == ref_hvy.shape:
            rmsd = np.sqrt(((mol_hvy - ref_hvy)**2).sum(axis=-1).mean())
            if normalize:
                rmsd /= np.sqrt(mol.num_rotors)
            return rmsd
    else:
        if mol.coords.shape == ref.coords.shape:
            rmsd = np.sqrt(((mol.coords - ref.coords)**2).sum(axis=-1).mean())
            if normalize:
                    rmsd /= np.sqrt(mol.num_rotors)
            return rmsd
    # at this point raise an exception
    raise Exception('Unequal number of atoms in molecules')

def distance_complex(x, y):
    """ Computes distance between points, similar to distance(cdist), with major difference - allows higher dimmentions of input (cdist supports 2). But it's 2-6 times slower, so use distance unless you have to nest it wit a for loop."""
    return np.sqrt(((x[...,np.newaxis,:]-y)**2).sum(axis=-1))

def rotate(coords, alpha, beta, gamma):
    """Returns an angle from a series of 3 points (point #2 is centroid).Angle is returned in degrees.

    Parameters
    ----------
    coords : numpy arrays, shape = [n_points, 3]
        Coordinates in 3-dimensional space.

    alpha, beta, gamma: float
        Angles to rotate the coordinates along X,Y and Z axis. Angles are specified in radians.

    Returns
    -------
    new_coords : numpy arrays, shape = [n_points, 3]
        Rorated coordinates in 3-dimensional space.
    """
    centroid = coords.mean(axis=0)
    coords = coords - centroid

    sin_alpha = sin(alpha)
    cos_alpha = cos(alpha)
    sin_beta = sin(beta)
    cos_beta = cos(beta)
    sin_gamma = sin(gamma)
    cos_gamma = cos(gamma)

    rot_matrix = np.array([[cos_beta*cos_gamma, sin_alpha*sin_beta*cos_gamma-cos_alpha*sin_gamma, cos_alpha*sin_beta*cos_gamma+sin_alpha*sin_gamma],
                           [cos_beta*sin_gamma, sin_alpha*sin_beta*sin_gamma+cos_alpha*cos_gamma, cos_alpha*sin_beta*sin_gamma-sin_alpha*cos_gamma],
                           [-sin_beta,          sin_alpha*cos_beta,                               cos_alpha*cos_beta]])

    return (coords[:,np.newaxis,:] * rot_matrix).sum(axis=2) + centroid
<commit_msg>Initial hRMSD<commit_after>"""Spatial functions included in ODDT
Mainly used by other modules, but can be accessed directly.
"""

from math import sin, cos
import numpy as np
from scipy.spatial.distance import cdist as distance
# for Hungarian algorithm, in future use scipy.optimize.linear_sum_assignment (in scipy 0.17+)
try:
    from scipy.optimize import linear_sum_assignment
except ImportError:
    from sklearn.utils.linear_assignment_ import linear_assignment
    def linear_sum_assignment(M):
        o = linear_assignment(M)
        return o[:,0], o[:,1]

__all__ = ['angle', 'angle_2v', 'dihedral', 'distance', 'rmsd', 'rotate']

# angle functions
def angle(p1,p2,p3):
    """Returns an angle from a series of 3 points (point #2 is centroid).Angle is returned in degrees.

    Parameters
    ----------
    p1,p2,p3 : numpy arrays, shape = [n_points, n_dimensions]
        Triplets of points in n-dimensional space, aligned in rows.

    Returns
    -------
    angles : numpy array, shape = [n_points]
        Series of angles in degrees
    """
    v1 = p1-p2
    v2 = p3-p2
    return angle_2v(v1,v2)

def angle_2v(v1, v2):
    """Returns an angle between two vecors.Angle is returned in degrees.

    Parameters
    ----------
    v1,v2 : numpy arrays, shape = [n_vectors, n_dimensions]
        Pairs of vectors in n-dimensional space, aligned in rows.

    Returns
    -------
    angles : numpy array, shape = [n_vectors]
        Series of angles in degrees
    """
    dot = (v1*v2).sum(axis=-1) # better than np.dot(v1, v2), multiple vectors can be applied
    norm = np.linalg.norm(v1, axis=-1)* np.linalg.norm(v2, axis=-1)
    return np.degrees(np.arccos(np.clip(dot/norm, -1, 1)))

def dihedral(p1,p2,p3,p4):
    """Returns an dihedral angle from a series of 4 points. Dihedral is returned in degrees.
    Function distingishes clockwise and antyclockwise dihedrals.

    Parameters
    ----------
    p1,p2,p3,p4 : numpy arrays, shape = [n_points, n_dimensions]
        Quadruplets of points in n-dimensional space, aligned in rows.

    Returns
    -------
    angles : numpy array, shape = [n_points]
        Series of angles in degrees
    """
    v12 = (p1-p2)/np.linalg.norm(p1-p2)
    v23 = (p2-p3)/np.linalg.norm(p2-p3)
    v34 = (p3-p4)/np.linalg.norm(p3-p4)
    c1 = np.cross(v12, v23)
    c2 = np.cross(v23, v34)
    out = angle_2v(c1, c2)
    # check clockwise and anticlockwise
    n1 = c1/np.linalg.norm(c1)
    mask = (n1*v34).sum(axis=-1) > 0
    if len(mask.shape) == 0:
        if mask:
            out = -out
    else:
        out[mask] = -out[mask]
    return out

def rmsd(ref, mol, ignore_h = True, method=None, normalize=False):
    """Computes root mean square deviation (RMSD) between two molecules (including or excluding Hydrogens). No symmetry checks are performed.

    Parameters
    ----------
    ref : oddt.toolkit.Molecule object
        Reference molecule for the RMSD calculation

    mol : oddt.toolkit.Molecule object
        Query molecule for RMSD calculation

    ignore_h : bool (default=False)
        Flag indicating to ignore Hydrogen atoms while performing RMSD calculation

    method : str (default=None)
        The method to be used for atom asignment between ref and mol.
        None means that direct matching is applied, which is the default behavior.
        Available methods:
            - canonize - match heavy atoms using OB canonical ordering (it forces ignoring H's)
            - hungarian - minimize RMSD using Hungarian algorithm

    normalize : bool (default=False)
        Normalize RMSD by square root of rot. bonds

    Returns
    -------
    rmsd : float
        RMSD between two molecules
    """

    if method == 'canonize':
        ref_atoms = ref.coords[ref.canonic_order]
        mol_atoms = mol.coords[mol.canonic_order]
    elif method == 'hungarian':
        mol_map = []
        ref_map = []
        for a_type in np.unique(mol.atom_dict['atomtype']):
            if a_type != 'H' or not ignore_h:
                mol_idx = np.argwhere(mol.atom_dict['atomtype'] == a_type).flatten()
                ref_idx = np.argwhere(ref.atom_dict['atomtype'] == a_type).flatten()
                if len(mol_idx) == 1:
                    mol_map.append(mol_idx)
                    ref_map.append(ref_idx)
                    continue
                M = distance(mol.atom_dict['coords'][ref_idx], ref.atom_dict['coords'][ref_idx])
                M = M - M.min(axis=0) - M.min(axis=1).reshape(-1,1)
                tmp_mol, tmp_ref = linear_sum_assignment(M)
                mol_map.append(mol_idx[tmp_mol])
                ref_map.append(ref_idx[tmp_ref])
        mol_atoms = mol.atom_dict['coords'][np.hstack(mol_map)]
        ref_atoms = ref.atom_dict['coords'][np.hstack(ref_map)]
    elif ignore_h:
        hvy_map = np.argwhere(mol.atom_dict['atomicnum'] != 1).flatten()
        mol_atoms = mol.coords[hvy_map]
        ref_atoms = ref.coords[hvy_map]
    else:
        mol_atoms = mol.coords
        ref_atoms = ref.coords
    if mol_atoms.shape == ref_atoms.shape:
        rmsd = np.sqrt(((mol_atoms - ref_atoms)**2).sum(axis=-1).mean())
        if normalize:
            rmsd /= np.sqrt(mol.num_rotors)
        return rmsd
    # at this point raise an exception
    raise Exception('Unequal number of atoms in molecules')

def distance_complex(x, y):
    """ Computes distance between points, similar to distance(cdist), with major difference - allows higher dimmentions of input (cdist supports 2). But it's 2-6 times slower, so use distance unless you have to nest it wit a for loop."""
    return np.sqrt(((x[...,np.newaxis,:]-y)**2).sum(axis=-1))

def rotate(coords, alpha, beta, gamma):
    """Returns an angle from a series of 3 points (point #2 is centroid).Angle is returned in degrees.

    Parameters
    ----------
    coords : numpy arrays, shape = [n_points, 3]
        Coordinates in 3-dimensional space.

    alpha, beta, gamma: float
        Angles to rotate the coordinates along X,Y and Z axis. Angles are specified in radians.

    Returns
    -------
    new_coords : numpy arrays, shape = [n_points, 3]
        Rorated coordinates in 3-dimensional space.
    """
    centroid = coords.mean(axis=0)
    coords = coords - centroid

    sin_alpha = sin(alpha)
    cos_alpha = cos(alpha)
    sin_beta = sin(beta)
    cos_beta = cos(beta)
    sin_gamma = sin(gamma)
    cos_gamma = cos(gamma)

    rot_matrix = np.array([[cos_beta*cos_gamma, sin_alpha*sin_beta*cos_gamma-cos_alpha*sin_gamma, cos_alpha*sin_beta*cos_gamma+sin_alpha*sin_gamma],
                           [cos_beta*sin_gamma, sin_alpha*sin_beta*sin_gamma+cos_alpha*cos_gamma, cos_alpha*sin_beta*sin_gamma-sin_alpha*cos_gamma],
                           [-sin_beta,          sin_alpha*cos_beta,                               cos_alpha*cos_beta]])

    return (coords[:,np.newaxis,:] * rot_matrix).sum(axis=2) + centroid
<|endoftext|>


<commit_before>import asyncio
from aiohttp import ClientSession
from arq import Actor, BaseWorker, concurrent

class Downloader(Actor):
    async def startup(self):
        self.session = ClientSession(loop=self.loop)

    @concurrent
    async def download_content(self, url):
        async with self.session.get(url) as response:
            content = await response.read()
            print('{}: {:.80}...'.format(url, content.decode()))
        return len(content)

    async def shutdown(self):
        self.session.close()

class Worker(BaseWorker):
    shadows = [Downloader]

async def download_lots():
    d = Downloader()
    for url in ('https://facebook.com', 'https://microsoft.com', 'https://github.com'):
        await d.download_content(url)
    await d.close()

if __name__ == '__main__':
    loop = asyncio.get_event_loop()
    loop.run_until_complete(download_lots())
<commit_msg>Use f-string in main_demo code example (#106)<commit_after>import asyncio
from aiohttp import ClientSession
from arq import Actor, BaseWorker, concurrent

class Downloader(Actor):
    async def startup(self):
        self.session = ClientSession(loop=self.loop)

    @concurrent
    async def download_content(self, url):
        async with self.session.get(url) as response:
            content = await response.read()
            print(f'{url}: {content.decode():.80}...')
        return len(content)

    async def shutdown(self):
        self.session.close()

class Worker(BaseWorker):
    shadows = [Downloader]

async def download_lots():
    d = Downloader()
    for url in ('https://facebook.com', 'https://microsoft.com', 'https://github.com'):
        await d.download_content(url)
    await d.close()

if __name__ == '__main__':
    loop = asyncio.get_event_loop()
    loop.run_until_complete(download_lots())
<|endoftext|>


<commit_before>"""A MayaVi example of how to generate a 3D structured points dataset
using numpy arrays.  Also shown is a way to visualize this data with
mayavi2.  The script can be run like so:

  $ mayavi2 -x structured_points3d.py

Alternatively, it can be run as:

  $ python structured_points3d.py

Author: Prabhu Ramachandran <prabhu at aero dot iitb dot ac dot in>

Copyright (c) 2007, Enthought, Inc.
License: BSD style.
"""

# The following *optional* two lines allow a user to call this script
# as either `python script.py` or `mayavi2 script.py`.  These two
# lines must be placed before any other mayavi imports.
from enthought.mayavi.scripts import mayavi2
mayavi2.standalone(globals())

from enthought.tvtk.api import tvtk
from enthought.tvtk.array_handler import get_vtk_array_type
from numpy import array, ogrid, sin, ravel

# Make the data.
dims = array((128, 128, 128))
vol = array((-5., 5, -5, 5, -5, 5))
origin = vol[::2]
spacing = (vol[1::2] - origin)/(dims -1)
xmin, xmax, ymin, ymax, zmin, zmax = vol
x, y, z = ogrid[xmin:xmax:dims[0]*1j,
                ymin:ymax:dims[1]*1j,
                zmin:zmax:dims[2]*1j]
x, y, z = [t.astype('f') for t in (x, y, z)]
scalars = sin(x*y*z)/(x*y*z)

# Make the tvtk dataset.
spoints = tvtk.StructuredPoints(origin=origin, spacing=spacing,
                                dimensions=dims)
# The copy makes the data contiguous and the transpose makes it
# suitable for display via tvtk.  Note that it is not necessary to
# make the data contiguous since in that case the array is copied
# internally.
s = scalars.transpose().copy()
spoints.point_data.scalars = ravel(s)
spoints.point_data.scalars.name = 'scalars'

# This is needed in slightly older versions of VTK (like the 5.0.2
# release) to prevent a segfault.  VTK does not detect the correct
# data type.
spoints.scalar_type = get_vtk_array_type(s.dtype.char)

# Uncomment the next two lines to save the dataset to a VTK XML file.
#w = tvtk.XMLImageDataWriter(input=spoints, file_name='spoints3d.vti')
#w.write()

# Now view the data.
def view():
    from enthought.mayavi.sources.vtk_data_source import VTKDataSource
    from enthought.mayavi.modules.outline import Outline
    from enthought.mayavi.modules.image_plane_widget import ImagePlaneWidget

    mayavi.new_scene()
    src = VTKDataSource(data = spoints)
    mayavi.add_source(src)
    mayavi.add_module(Outline())
    mayavi.add_module(ImagePlaneWidget())

view()
<commit_msg>BUG: The array type is determined now by the dtype and not the deprecated typecode character.  Thanks to Ilan Schnell for reporting the bug and demonstrating the problem.<commit_after>"""A MayaVi example of how to generate a 3D structured points dataset
using numpy arrays.  Also shown is a way to visualize this data with
mayavi2.  The script can be run like so:

  $ mayavi2 -x structured_points3d.py

Alternatively, it can be run as:

  $ python structured_points3d.py

Author: Prabhu Ramachandran <prabhu at aero dot iitb dot ac dot in>

Copyright (c) 2007, Enthought, Inc.
License: BSD style.
"""

# The following *optional* two lines allow a user to call this script
# as either `python script.py` or `mayavi2 script.py`.  These two
# lines must be placed before any other mayavi imports.
from enthought.mayavi.scripts import mayavi2
mayavi2.standalone(globals())

from enthought.tvtk.api import tvtk
from enthought.tvtk.array_handler import get_vtk_array_type
from numpy import array, ogrid, sin, ravel

# Make the data.
dims = array((128, 128, 128))
vol = array((-5., 5, -5, 5, -5, 5))
origin = vol[::2]
spacing = (vol[1::2] - origin)/(dims -1)
xmin, xmax, ymin, ymax, zmin, zmax = vol
x, y, z = ogrid[xmin:xmax:dims[0]*1j,
                ymin:ymax:dims[1]*1j,
                zmin:zmax:dims[2]*1j]
x, y, z = [t.astype('f') for t in (x, y, z)]
scalars = sin(x*y*z)/(x*y*z)

# Make the tvtk dataset.
spoints = tvtk.StructuredPoints(origin=origin, spacing=spacing,
                                dimensions=dims)
# The copy makes the data contiguous and the transpose makes it
# suitable for display via tvtk.  Note that it is not necessary to
# make the data contiguous since in that case the array is copied
# internally.
s = scalars.transpose().copy()
spoints.point_data.scalars = ravel(s)
spoints.point_data.scalars.name = 'scalars'

# This is needed in slightly older versions of VTK (like the 5.0.2
# release) to prevent a segfault.  VTK does not detect the correct
# data type.
spoints.scalar_type = get_vtk_array_type(s.dtype)

# Uncomment the next two lines to save the dataset to a VTK XML file.
#w = tvtk.XMLImageDataWriter(input=spoints, file_name='spoints3d.vti')
#w.write()

# Now view the data.
def view():
    from enthought.mayavi.sources.vtk_data_source import VTKDataSource
    from enthought.mayavi.modules.outline import Outline
    from enthought.mayavi.modules.image_plane_widget import ImagePlaneWidget

    mayavi.new_scene()
    src = VTKDataSource(data = spoints)
    mayavi.add_source(src)
    mayavi.add_module(Outline())
    mayavi.add_module(ImagePlaneWidget())

view()
<|endoftext|>


<commit_before>#!/usr/bin/env python
from __future__ import print_function
from selenium import webdriver
from selenium.webdriver.common.keys import Keys
from time import sleep
from subprocess import call
import webbrowser
"""
test the libraries' Google Forms which are hooked up to Google Apps Script
need to run this periodically since the forms tend to lose their ability to
send emails without warning, this checks that scripts still have permissions
"""

form_root = 'https://docs.google.com/a/cca.edu/forms/d/{0}/viewform'
forms = {
    'acquisitions': '17pKoZ_U8AHbFjEg4blXiJod8uZwKI0s1DFJBSSSB7rE',
    'ill': '1pyhf7krGP3PQX1McJIt1fwLUYAwUd88vGG1brC33qe8',
    'instruction': '1i6Xn0q3MI7EXM6VpxUr64oywifGDcU-4NZoitp2y7-g'
}
# values submitted through form
vals = {
    'first title': 'TEST Sweet Citation to First Title\n\ndoobie doowop bedoo',
    'name': 'Testy Testerson',
    'email': 'ephetteplace@cca.edu'
}

def test_recommendations():
    """
    test the recommended acquisitions form
    """
    browser.get(form_root.format(forms['acquisitions']))

    # first citation
    browser.find_element_by_id('entry_420214461').send_keys(vals['first title'])
    # leave blank a bunch of optional fields...
    # name
    browser.find_element_by_id('entry_967350359').send_keys(vals['name'])
    # CCA affilliation => Grad Student
    browser.find_element_by_id('group_942677086_1').click()
    # email
    browser.find_element_by_id('entry_1030731656').send_keys(vals['email'])
    # program => Writing/Literature/Comics
    browser.find_element_by_id('group_1569646515_7').click()

    # submit form
    browser.find_element_by_id('ss-submit').click()

def test_ill():
    """
    test the inter-library loan request form
    """
    browser.get(form_root.format(forms['ill']))
    # citation
    browser.find_element_by_id('entry_17105932').send_keys(vals['first title'])
    # name
    browser.find_element_by_id('entry_1904144554').send_keys(vals['name'])
    # email
    browser.find_element_by_id('entry_1309367559').send_keys(vals['email'])
    # CCA affilliation => Grad Student
    browser.find_element_by_id('group_2072708086_2').click()
    # department => Architecture
    # <select> drop-down, get appropriate <option>
    browser.find_element_by_id('entry_443514869').find_elements_by_tag_name(
        'option')[1].click()
    # submit form
    browser.find_element_by_id('ss-submit').click()

def test_instruction():
    """
    test the instruction session request form
    """
    browser.get(form_root.format(forms['instruction']))
    # wait for user to manually sign into Google
    # @TODO should automate this, probably by passing credentials on the CLI
    print('Waiting to sign into Google, hit Return when done...')
    call('read')  # gotta be a native Python way to do this w/o shelling out
    # course name
    browser.find_element_by_id('entry_93795666').send_keys('TESTS-101')
    # date/time, string sent looks weird here but it's just the date
    # 01/01/2000 10:00 AM with the date/time autofill behavior
    browser.find_element_by_id('entry_608498536').send_keys('01012000' + Keys.TAB + '1000A')
    browser.find_element_by_id('ss-submit').click()

def open_sheets():
    """
    Open the back-end spreadsheets for each form
    to see if entries were copied over properly
    """
    sheet_root = 'https://docs.google.com/spreadsheets/d/'
    sheets = {
        'acquisitions': '1A0bBB4_mYWtcWoqfB5HP6H7WbIWjTknzHnVr-5EPw8s/edit#gid=1311627534',
        'ill': '1V4yaJajQ37oGEFZXYxGcTtlF6Jq_XhQvZVGF4xBWO1k/edit#gid=275875544',
        'instruction': '1WpdyyvrVjCrWzHprUOZ1IlDmU2pQEb2qrzNML9G2JRw/edit#gid=1190654636'
    }
    for sheet in sheets:
        webbrowser.open(sheet_root + sheets[sheet])

if __name__ == '__main__':
    # PhantomJS is a preferable, faster option but it seems to choke
    # whereas an actual browser window does not
    # browser = webdriver.PhantomJS()
    browser = webdriver.Chrome()
    browser.implicitly_wait(30)

    for test in (test_recommendations, test_ill, test_instruction):
        test()
        sleep(1)

    open_sheets()
    browser.quit()
<commit_msg>gforms - add helper functions<commit_after>#!/usr/bin/env python
from __future__ import print_function
from selenium import webdriver
from selenium.webdriver.common.keys import Keys
from time import sleep
from subprocess import call
import webbrowser
"""
test the libraries' Google Forms which are hooked up to Google Apps Script
need to run this periodically since the forms tend to lose their ability to
send emails without warning, this checks that scripts still have permissions
"""

form_root = 'https://docs.google.com/a/cca.edu/forms/d/{0}/viewform'
forms = {
    'acquisitions': '17pKoZ_U8AHbFjEg4blXiJod8uZwKI0s1DFJBSSSB7rE',
    'ill': '1pyhf7krGP3PQX1McJIt1fwLUYAwUd88vGG1brC33qe8',
    'instruction': '1i6Xn0q3MI7EXM6VpxUr64oywifGDcU-4NZoitp2y7-g'
}
# values submitted through form
vals = {
    'first title': 'TEST Sweet Citation to First Title\n\ndoobie doowop bedoo',
    'name': 'Testy Testerson',
    'email': 'ephetteplace@cca.edu'
}

def fill_in(label, value):
    """
    Given ARIA label, fill in all matching form inputs with provided value
    """
    # find_elements type method return lists
    for el in browser.find_elements_by_css_selector('[aria-label="' + label + '"]'):
        el.send_keys(value)

def click(label):
    """
    Given ARIA label, click the DOM element
    """
    browser.find_elements_by_css_selector('[aria-label="' + label + '"]')[0].click()

def test_recommendations():
    """
    test the recommended acquisitions form
    """
    browser.get(form_root.format(forms['acquisitions']))

    # first citation
    browser.find_element_by_id('entry_420214461').send_keys(vals['first title'])
    # leave blank a bunch of optional fields...
    # name
    browser.find_element_by_id('entry_967350359').send_keys(vals['name'])
    # CCA affilliation => Grad Student
    browser.find_element_by_id('group_942677086_1').click()
    # email
    browser.find_element_by_id('entry_1030731656').send_keys(vals['email'])
    # program => Writing/Literature/Comics
    browser.find_element_by_id('group_1569646515_7').click()

    # submit form
    browser.find_element_by_id('ss-submit').click()

def test_ill():
    """
    test the inter-library loan request form
    """
    browser.get(form_root.format(forms['ill']))
    # citation
    browser.find_element_by_id('entry_17105932').send_keys(vals['first title'])
    # name
    browser.find_element_by_id('entry_1904144554').send_keys(vals['name'])
    # email
    browser.find_element_by_id('entry_1309367559').send_keys(vals['email'])
    # CCA affilliation => Grad Student
    browser.find_element_by_id('group_2072708086_2').click()
    # department => Architecture
    # <select> drop-down, get appropriate <option>
    browser.find_element_by_id('entry_443514869').find_elements_by_tag_name(
        'option')[1].click()
    # submit form
    browser.find_element_by_id('ss-submit').click()

def test_instruction():
    """
    test the instruction session request form
    """
    browser.get(form_root.format(forms['instruction']))
    # wait for user to manually sign into Google
    # @TODO should automate this, probably by passing credentials on the CLI
    print('Waiting to sign into Google, hit Return when done...')
    call('read')  # gotta be a native Python way to do this w/o shelling out
    # course name
    browser.find_element_by_id('entry_93795666').send_keys('TESTS-101')
    # date/time, string sent looks weird here but it's just the date
    # 01/01/2000 10:00 AM with the date/time autofill behavior
    browser.find_element_by_id('entry_608498536').send_keys('01012000' + Keys.TAB + '1000A')
    browser.find_element_by_id('ss-submit').click()

def open_sheets():
    """
    Open the back-end spreadsheets for each form
    to see if entries were copied over properly
    """
    sheet_root = 'https://docs.google.com/spreadsheets/d/'
    sheets = {
        'acquisitions': '1A0bBB4_mYWtcWoqfB5HP6H7WbIWjTknzHnVr-5EPw8s/edit#gid=1311627534',
        'ill': '1V4yaJajQ37oGEFZXYxGcTtlF6Jq_XhQvZVGF4xBWO1k/edit#gid=275875544',
        'instruction': '1WpdyyvrVjCrWzHprUOZ1IlDmU2pQEb2qrzNML9G2JRw/edit#gid=1190654636'
    }
    for sheet in sheets:
        webbrowser.open(sheet_root + sheets[sheet])

if __name__ == '__main__':
    # PhantomJS is a preferable, faster option but it seems to choke
    # whereas an actual browser window does not
    # browser = webdriver.PhantomJS()
    browser = webdriver.Chrome()
    browser.implicitly_wait(30)

    for test in (test_recommendations, test_ill, test_instruction):
        test()
        sleep(1)

    open_sheets()
    browser.quit()
<|endoftext|>