text:
# -*- coding: utf-8 -*-
r"""
Graph Plotting

*(For LaTeX drawings of graphs, see the* :mod:`~sage.graphs.graph_latex` *module.)*

All graphs have an associated Sage graphics object, which you can display::

    sage: G = graphs.WheelGraph(15)
    sage: P = G.plot()
    sage: P.show() # long time

.. PLOT::

    sphinx_plot(graphs.WheelGraph(15))

If you create a graph in Sage using the ``Graph`` command, then plot that graph,
the positioning of nodes is determined using the spring-layout algorithm. For
the special graph constructors, which you get using ``graphs.[tab]``, the
positions are preset. For example, consider the Petersen graph with default node
positioning vs. the Petersen graph constructed by this database::

    sage: petersen_spring = Graph(':I`ES@obGkqegW~')
    sage: petersen_spring.show() # long time

.. PLOT::

    petersen_spring = Graph(':I`ES@obGkqegW~')
    sphinx_plot(petersen_spring)

::

    sage: petersen_database = graphs.PetersenGraph()
    sage: petersen_database.show() # long time

.. PLOT::

    petersen_database = graphs.PetersenGraph()
    sphinx_plot(petersen_database)

For all the constructors in this database (except some random graphs), the
position dictionary is filled in, instead of using the spring-layout algorithm.

**Plot options**

Here is the list of options accepted by
:meth:`~sage.graphs.generic_graph.GenericGraph.plot` and the constructor of
:class:`GraphPlot`. Those two functions also accept all options of
:meth:`sage.plot.graphics.Graphics.show`.

.. csv-table::
    :class: contentstable
    :widths: 30, 70
    :delim: |

{PLOT_OPTIONS_TABLE}

**Default options**

This module defines two dictionaries containing default options for the
:meth:`~sage.graphs.generic_graph.GenericGraph.plot` and
:meth:`~sage.graphs.generic_graph.GenericGraph.show` methods. These two
dictionaries are ``sage.graphs.graph_plot.DEFAULT_PLOT_OPTIONS`` and
``sage.graphs.graph_plot.DEFAULT_SHOW_OPTIONS``, respectively.

Obviously, these values are overruled when arguments are given explicitly.

Here is how to define the default size of a graph drawing to be ``[6,6]``. The
first two calls to :meth:`~sage.graphs.generic_graph.GenericGraph.show` use this
option, while the third does not (a value for ``figsize`` is explicitly given)::

    sage: import sage.graphs.graph_plot
    sage: sage.graphs.graph_plot.DEFAULT_SHOW_OPTIONS['figsize'] = [6,6]
    sage: graphs.PetersenGraph().show() # long time
    sage: graphs.ChvatalGraph().show()  # long time
    sage: graphs.PetersenGraph().show(figsize=[4,4]) # long time

We can now reset the default to its initial value, and now display graphs as
previously::

    sage: sage.graphs.graph_plot.DEFAULT_SHOW_OPTIONS['figsize'] = [4,4]
    sage: graphs.PetersenGraph().show() # long time
    sage: graphs.ChvatalGraph().show()  # long time

.. NOTE::

    * While ``DEFAULT_PLOT_OPTIONS`` affects both ``G.show()`` and ``G.plot()``,
      settings from ``DEFAULT_SHOW_OPTIONS`` only affects ``G.show()``.

    * In order to define a default value permanently, you can add a couple of
      lines to `Sage's startup scripts <../../../repl/startup.html>`_. Example::

       sage: import sage.graphs.graph_plot
       sage: sage.graphs.graph_plot.DEFAULT_SHOW_OPTIONS['figsize'] = [4,4]

**Index of methods and functions**

.. csv-table::
    :class: contentstable
    :widths: 30, 70
    :delim: |

    :meth:`GraphPlot.set_pos` | Set the position plotting parameters for this GraphPlot.
    :meth:`GraphPlot.set_vertices` | Set the vertex plotting parameters for this GraphPlot.
    :meth:`GraphPlot.set_edges` | Set the edge (or arrow) plotting parameters for the GraphPlot object.
    :meth:`GraphPlot.show` | Show the (Di)Graph associated with this GraphPlot object.
    :meth:`GraphPlot.plot` | Return a graphics object representing the (di)graph.
    :meth:`GraphPlot.layout_tree` | Compute a nice layout of a tree.
"""

layout_options =   {
    'layout': 'A layout algorithm -- one of : "acyclic", "circular" (plots the '
        'graph with vertices evenly distributed on a circle), "ranked", '
        '"graphviz", "planar", "spring" (traditional spring layout, using the '
        'graph\'s current positions as initial positions), or "tree" (the tree '
        'will be plotted in levels, depending on minimum distance for the root).',
    'iterations': 'The number of times to execute the spring layout algorithm.',
    'heights': 'A dictionary mapping heights to the list of vertices at this height.',
    'spring': 'Use spring layout to finalize the current layout.',
    'tree_root': 'A vertex designation for drawing trees. A vertex of the tree '
        'to be used as the root for the ``layout=\'tree\'`` option. If no root '
        'is specified, then one is chosen close to the center of the tree. '
        'Ignored unless ``layout=\'tree\'``.',
    'forest_roots': 'An iterable specifying which vertices to use as roots for '
        'the ``layout=\'forest\'`` option. If no root is specified for a tree, '
        'then one is chosen close to the center of the tree. '
        'Ignored unless ``layout=\'forest\'``.',
    'tree_orientation': 'The direction of tree branches -- \'up\', \'down\', '
        '\'left\' or \'right\'.',
    'save_pos': 'Whether or not to save the computed position for the graph.',
    'dim': 'The dimension of the layout -- 2 or 3.',
    'prog': 'Which graphviz layout program to use -- one of "circo", "dot", '
        '"fdp", "neato", or "twopi".',
    'by_component': 'Whether to do the spring layout by connected component '
        '-- a boolean.',
    }

graphplot_options = layout_options.copy()

graphplot_options.update(
                   {'pos': 'The position dictionary of vertices',
                    'vertex_labels': 'Whether or not to draw vertex labels.',
                    'vertex_color': 'Default color for vertices not listed '
                        'in vertex_colors dictionary.',
                    'vertex_colors': 'Dictionary of vertex coloring : each '
                        'key is a color recognizable by matplotlib, and each '
                        'corresponding entry is a list of vertices. ',
                    'vertex_size': 'The size to draw the vertices.',
                    'vertex_shape': 'The shape to draw the vertices. '
                        'Currently unavailable for Multi-edged DiGraphs.',
                    'edge_labels': 'Whether or not to draw edge labels.',
                    'edge_style': 'The linestyle of the edges. It should be '
                        'one of "solid", "dashed", "dotted", dashdot", or '
                        '"-", "--", ":", "-.", respectively. ',
                    'edge_thickness': 'The thickness of the edges.',
                    'edge_color': 'The default color for edges not listed in edge_colors.',
                    'edge_colors': 'a dictionary specifying edge colors: each '
                        'key is a color recognized by matplotlib, and each '
                        'entry is a list of edges.',
                    'color_by_label': 'Whether to color the edges according '
                        'to their labels. This also accepts a function or '
                        'dictionary mapping labels to colors.',
                    'partition': 'A partition of the vertex set. If specified, '
                        'plot will show each cell in a different color. '
                        'vertex_colors takes precedence.',
                    'loop_size': 'The radius of the smallest loop.',
                    'dist': 'The distance between multiedges.',
                    'max_dist': 'The max distance range to allow multiedges.',
                    'talk': 'Whether to display the vertices in talk mode '
                        '(larger and white).',
                    'graph_border': 'Whether or not to draw a frame around the graph.',
                    'edge_labels_background' : 'The color of the background of the edge labels'})

_PLOT_OPTIONS_TABLE = ""
for key, value in graphplot_options.items():
    _PLOT_OPTIONS_TABLE += "    ``"+str(key)+"`` | "+str(value)+"\n"
__doc__ = __doc__.format(PLOT_OPTIONS_TABLE=_PLOT_OPTIONS_TABLE)

# ****************************************************************************
#      Copyright (C) 2009   Emily Kirkman
#                    2009   Robert L. Miller <rlmillster@gmail.com>
#
#  Distributed under the terms of the GNU General Public License (GPL)
#
#    This code is distributed in the hope that it will be useful,
#    but WITHOUT ANY WARRANTY; without even the implied warranty of
#    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
#    General Public License for more details.
#
#  The full text of the GPL is available at:
#
#                  https://www.gnu.org/licenses/
# ****************************************************************************
from sage.structure.sage_object import SageObject
from sage.plot.all import Graphics, scatter_plot, bezier_path, line, arrow, text, circle
from math import sqrt, cos, sin, atan, pi

DEFAULT_SHOW_OPTIONS = {
    "figsize"             : [4,4]
    }

DEFAULT_PLOT_OPTIONS = {
    "vertex_size"         : 200,
    "vertex_labels"       : True,
    "layout"              : None,
    "edge_style"          : 'solid',
    "edge_thickness"      : 1,
    "edge_color"          : 'black',
    "edge_colors"         : None,
    "edge_labels"         : False,
    "iterations"          : 50,
    "tree_orientation"    : 'down',
    "heights"             : None,
    "graph_border"        : False,
    "talk"                : False,
    "color_by_label"      : False,
    "partition"           : None,
    "dist"                : .075,
    "max_dist"            : 1.5,
    "loop_size"           : .075,
    "edge_labels_background" : "white"
    }

class GraphPlot(SageObject):
    def __init__(self, graph, options):
        """
        Return a ``GraphPlot`` object, which stores all the parameters needed
        for plotting (Di)Graphs.

        A ``GraphPlot`` has a plot and show function, as well as some functions
        to set parameters for vertices and edges.  This constructor assumes
        default options are set.  Defaults are shown in the example below.

        EXAMPLES::

            sage: from sage.graphs.graph_plot import GraphPlot
            sage: options = {
            ....:   'vertex_size': 200,
            ....:   'vertex_labels': True,
            ....:   'layout': None,
            ....:   'edge_style': 'solid',
            ....:   'edge_color': 'black',
            ....:   'edge_colors': None,
            ....:   'edge_labels': False,
            ....:   'iterations': 50,
            ....:   'tree_orientation': 'down',
            ....:   'heights': None,
            ....:   'graph_border': False,
            ....:   'talk': False,
            ....:   'color_by_label': False,
            ....:   'partition': None,
            ....:   'dist': .075,
            ....:   'max_dist': 1.5,
            ....:   'loop_size': .075,
            ....:   'edge_labels_background': 'transparent'}
            sage: g = Graph({0:[1, 2], 2:[3], 4:[0, 1]})
            sage: GP = GraphPlot(g, options)

        """
        # Setting the default values if needed
        for k, value in DEFAULT_PLOT_OPTIONS.items():
            if k not in options:
                options[k] = value
        self._plot_components = {}
        self._nodelist = list(graph)
        self._graph = graph
        self._options = options # contains both plot and show options
        self.set_pos()
        self._arcs = self._graph.has_multiple_edges(to_undirected=True)
        self._loops = self._graph.has_loops()
        self._arcdigraph = self._graph.is_directed() and self._arcs

        self.set_vertices()
        self.set_edges()

    def _repr_(self):
        """
        Return a string representation of a ``GraphPlot`` object.

        EXAMPLES:

        This function is called implicitly by the code below::

            sage: g = Graph({0:[1,2], 2:[3], 4:[0,1]})
            sage: g.graphplot() # indirect doctest
            GraphPlot object for Graph on 5 vertices
        """
        return "GraphPlot object for %s"%self._graph

    def set_pos(self):
        """
        Set the position plotting parameters for this GraphPlot.

        EXAMPLES:

        This function is called implicitly by the code below::

            sage: g = Graph({0:[1,2], 2:[3], 4:[0,1]})
            sage: g.graphplot(save_pos=True, layout='circular') # indirect doctest
            GraphPlot object for Graph on 5 vertices

        The following illustrates the format of a position dictionary, but due
        to numerical noise we do not check the values themselves::

            sage: g.get_pos()
            {0: (0.0, 1.0),
             1: (-0.951..., 0.309...),
             2: (-0.587..., -0.809...),
             3: (0.587..., -0.809...),
             4: (0.951..., 0.309...)}

        ::

            sage: T = list(graphs.trees(7))
            sage: t = T[3]
            sage: t.plot(heights={0:[0], 1:[4,5,1], 2:[2], 3:[3,6]})
            Graphics object consisting of 14 graphics primitives

        .. PLOT::

            g = Graph({0:[1,2], 2:[3], 4:[0,1]})
            g.graphplot(save_pos=True, layout='circular') # indirect doctest
            T = list(graphs.trees(7))
            t = T[3]
            P = t.plot(heights={0:[0], 1:[4,5,1], 2:[2], 3:[3,6]})
            sphinx_plot(P)

        TESTS:

        Make sure that vertex locations are floats.  Not being floats isn't a
        bug in itself but makes it too easy to accidentally introduce a bug
        elsewhere, such as in :meth:`set_edges` (:trac:`10124`), via silent
        truncating division of integers::

            sage: g = graphs.FruchtGraph()
            sage: gp = g.graphplot()
            sage: set(map(type, flatten(gp._pos.values())))
            {<... 'float'>}
            sage: g = graphs.BullGraph()
            sage: gp = g.graphplot(save_pos=True)
            sage: set(map(type, flatten(gp._pos.values())))
            {<... 'float'>}

        Non-ascii labels are also possible using unicode (:trac:`21008`)::

            sage: Graph({u'où': [u'là', u'ici']}).plot()
            Graphics object consisting of 6 graphics primitives
        """
        self._pos = self._graph.layout(**self._options)
        # make sure the positions are floats (trac #10124)
        self._pos = {k: (float(v[0]), float(v[1]))
                         for k, v in self._pos.items()}

    def set_vertices(self, **vertex_options):
        """
        Set the vertex plotting parameters for this ``GraphPlot``.

        This function is called by the constructor but can also be called to
        make updates to the vertex options of an existing ``GraphPlot`` object.
        Note that the changes are cumulative.

        EXAMPLES::

            sage: g = Graph({}, loops=True, multiedges=True, sparse=True)
            sage: g.add_edges([(0,0,'a'),(0,0,'b'),(0,1,'c'),(0,1,'d'),
            ....:              (0,1,'e'),(0,1,'f'),(0,1,'f'),(2,1,'g'),(2,2,'h')])
            sage: GP = g.graphplot(vertex_size=100, edge_labels=True, color_by_label=True,
            ....:                  edge_style='dashed')
            sage: GP.set_vertices(talk=True)
            sage: GP.plot()
            Graphics object consisting of 26 graphics primitives
            sage: GP.set_vertices(vertex_color='green', vertex_shape='^')
            sage: GP.plot()
            Graphics object consisting of 26 graphics primitives

        .. PLOT::

                g = Graph({}, loops=True, multiedges=True, sparse=True)
                g.add_edges([(0,0,'a'),(0,0,'b'),(0,1,'c'),(0,1,'d'),(0,1,'e'),(0,1,'f'),
                             (0,1,'f'),(2,1,'g'),(2,2,'h')])
                GP = g.graphplot(vertex_size=100, edge_labels=True, color_by_label=True,
                                 edge_style='dashed')
                GP.set_vertices(talk=True)
                sphinx_plot(GP)

        .. PLOT::

                g = Graph({}, loops=True, multiedges=True, sparse=True)
                g.add_edges([(0,0,'a'),(0,0,'b'),(0,1,'c'),(0,1,'d'),(0,1,'e'),(0,1,'f'),
                            (0,1,'f'),(2,1,'g'),(2,2,'h')])
                GP = g.graphplot(vertex_size=100, edge_labels=True, color_by_label=True,
                                 edge_style='dashed')
                GP.set_vertices(talk=True)
                GP.set_vertices(vertex_color='green', vertex_shape='^')
                sphinx_plot(GP)

        """
        # Handle base vertex options
        voptions = {}

        for arg in vertex_options:
            self._options[arg] = vertex_options[arg]

        # First set defaults for styles
        vertex_colors = None
        if self._options['talk']:
            voptions['markersize'] = 500
            if self._options['partition'] is None:
                vertex_colors = '#ffffff'
        else:
            voptions['markersize'] = self._options['vertex_size']

        if 'vertex_color' not in self._options or self._options['vertex_color'] is None:
            vertex_color = '#fec7b8'
        else:
            vertex_color = self._options['vertex_color']

        if 'vertex_colors' not in self._options or self._options['vertex_colors'] is None:
            if self._options['partition'] is not None:
                from sage.plot.colors import rainbow
                partition = self._options['partition']
                l = len(partition)
                R = rainbow(l)
                vertex_colors = {R[i]: partition[i] for i in range(l)}
            elif not vertex_colors:
                vertex_colors = vertex_color
        else:
            vertex_colors = self._options['vertex_colors']

        if 'vertex_shape' in self._options:
            voptions['marker'] = self._options['vertex_shape']

        if self._graph.is_directed():
            self._vertex_radius = sqrt(voptions['markersize'] / pi)
            self._arrowshorten = 2 * self._vertex_radius
            if self._arcdigraph:
                self._vertex_radius = sqrt(voptions['markersize'] / (20500 * pi))

        voptions['zorder'] = 7

        if not isinstance(vertex_colors, dict):
            voptions['facecolor'] = vertex_colors
            if self._arcdigraph:
                self._plot_components['vertices'] = [circle(center,
                                                            self._vertex_radius,
                                                            fill=True,
                                                            facecolor=vertex_colors,
                                                            edgecolor='black',
                                                            clip=False)
                                                     for center in self._pos.values()]
            else:
                self._plot_components['vertices'] = scatter_plot(list(self._pos.values()),
                                                                 clip=False, **voptions)
        else:
            # Color list must be ordered:
            pos = []
            colors = []
            for i in vertex_colors:
                pos += [self._pos[j] for j in vertex_colors[i]]
                colors += [i] * len(vertex_colors[i])

            # If all the vertices have not been assigned a color
            if len(self._pos) != len(pos):
                leftovers = [j for j in self._pos.values() if j not in pos]
                pos += leftovers
                colors += [vertex_color] * len(leftovers)

            if self._arcdigraph:
                self._plot_components['vertices'] = [circle(pos[i],
                                                            self._vertex_radius,
                                                            fill=True,
                                                            facecolor=colors[i],
                                                            edgecolor='black',
                                                            clip=False)
                                                     for i in range(len(pos))]
            else:
                self._plot_components['vertices'] = scatter_plot(pos,
                                                                 facecolor=colors,
                                                                 clip=False, **voptions)

        if self._options['vertex_labels']:
            self._plot_components['vertex_labels'] = []
            # TODO: allow text options
            for v in self._nodelist:
                self._plot_components['vertex_labels'].append(text(str(v),
                    self._pos[v], rgbcolor=(0,0,0), zorder=8))

    def set_edges(self, **edge_options):
        """
        Set the edge (or arrow) plotting parameters for the ``GraphPlot``
        object.

        This function is called by the constructor but can also be called to
        make updates to the vertex options of an existing ``GraphPlot`` object.
        Note that the changes are cumulative.

        EXAMPLES::

            sage: g = Graph(loops=True, multiedges=True, sparse=True)
            sage: g.add_edges([(0,0,'a'),(0,0,'b'),(0,1,'c'),(0,1,'d'),
            ....:  (0,1,'e'),(0,1,'f'),(0,1,'f'),(2,1,'g'),(2,2,'h')])
            sage: GP = g.graphplot(vertex_size=100, edge_labels=True, color_by_label=True,
            ....:  edge_style='dashed')
            sage: GP.set_edges(edge_style='solid')
            sage: GP.plot()
            Graphics object consisting of 26 graphics primitives

        .. PLOT::

            g = Graph(loops=True, multiedges=True, sparse=True)
            g.add_edges([(0,0,'a'),(0,0,'b'),(0,1,'c'),(0,1,'d'),
                (0,1,'e'),(0,1,'f'),(0,1,'f'),(2,1,'g'),(2,2,'h')])
            GP = g.graphplot(vertex_size=100, edge_labels=True,
                color_by_label=True, edge_style='dashed')
            GP.set_edges(edge_style='solid')
            sphinx_plot(GP)

        ::

            sage: GP.set_edges(edge_color='black')
            sage: GP.plot()
            Graphics object consisting of 26 graphics primitives

        .. PLOT::

            g = Graph(loops=True, multiedges=True, sparse=True)
            g.add_edges([(0,0,'a'),(0,0,'b'),(0,1,'c'),(0,1,'d'),
                (0,1,'e'),(0,1,'f'),(0,1,'f'),(2,1,'g'),(2,2,'h')])
            GP = g.graphplot(vertex_size=100, edge_labels=True,
                color_by_label=True, edge_style='dashed')
            GP.set_edges(edge_style='solid')
            GP.set_edges(edge_color='black')
            sphinx_plot(GP)

        ::

            sage: d = DiGraph(loops=True, multiedges=True, sparse=True)
            sage: d.add_edges([(0,0,'a'),(0,0,'b'),(0,1,'c'),(0,1,'d'),
            ....:   (0,1,'e'),(0,1,'f'),(0,1,'f'),(2,1,'g'),(2,2,'h')])
            sage: GP = d.graphplot(vertex_size=100, edge_labels=True, color_by_label=True,
            ....:  edge_style='dashed')
            sage: GP.set_edges(edge_style='solid')
            sage: GP.plot()
            Graphics object consisting of 28 graphics primitives

        .. PLOT::

            d = DiGraph(loops=True, multiedges=True, sparse=True)
            d.add_edges([(0,0,'a'),(0,0,'b'),(0,1,'c'),(0,1,'d'),
                (0,1,'e'),(0,1,'f'),(0,1,'f'),(2,1,'g'),(2,2,'h')])
            GP = d.graphplot(vertex_size=100, edge_labels=True, color_by_label=True,
                edge_style='dashed')
            GP.set_edges(edge_style='solid')
            sphinx_plot(GP)

        ::

            sage: GP.set_edges(edge_color='black')
            sage: GP.plot()
            Graphics object consisting of 28 graphics primitives

        .. PLOT::

            d = DiGraph(loops=True, multiedges=True, sparse=True)
            d.add_edges([(0,0,'a'),(0,0,'b'),(0,1,'c'),(0,1,'d'),
                (0,1,'e'),(0,1,'f'),(0,1,'f'),(2,1,'g'),(2,2,'h')])
            GP = d.graphplot(vertex_size=100, edge_labels=True, color_by_label=True,
                edge_style='dashed')
            GP.set_edges(edge_style='solid')
            GP.set_edges(edge_color='black')
            sphinx_plot(GP)

        TESTS::

            sage: G = Graph("Fooba")
            sage: G.show(edge_colors={'red':[(3,6),(2,5)]})

        Verify that default edge labels are pretty close to being between the vertices
        in some cases where they weren't due to truncating division (:trac:`10124`)::

            sage: test_graphs = graphs.FruchtGraph(), graphs.BullGraph()
            sage: tol = 0.001
            sage: for G in test_graphs:
            ....:     E = G.edges()
            ....:     for e0, e1, elab in E:
            ....:         G.set_edge_label(e0, e1, '%d %d' % (e0, e1))
            ....:     gp = G.graphplot(save_pos=True, edge_labels=True)
            ....:     vx = gp._plot_components['vertices'][0].xdata
            ....:     vy = gp._plot_components['vertices'][0].ydata
            ....:     for elab in gp._plot_components['edge_labels']:
            ....:         textobj = elab[0]
            ....:         x, y, s = textobj.x, textobj.y, textobj.string
            ....:         v0, v1 = map(int, s.split())
            ....:         vn = vector(((x-(vx[v0]+vx[v1])/2.), y-(vy[v0]+vy[v1])/2.)).norm()
            ....:         assert vn < tol

        Ticket :trac:`24051` is fixed::

            sage: G = Graph([(0, 1), (0, 1)], multiedges=True)
            sage: G.plot(edge_colors={"red": [(1, 0)]})
            Graphics object consisting of 5 graphics primitives
        """
        for arg in edge_options:
            self._options[arg] = edge_options[arg]
        if 'edge_colors' in edge_options:
            self._options['color_by_label'] = False
        if self._options['edge_labels_background'] == "transparent":
            self._options['edge_labels_background'] = "None"

        # Handle base edge options: thickness, linestyle
        eoptions = {}
        if 'edge_style' in self._options:
            from sage.plot.misc import get_matplotlib_linestyle
            eoptions['linestyle'] = get_matplotlib_linestyle(
                                        self._options['edge_style'],
                                        return_type='long')
        if 'edge_thickness' in self._options:
            eoptions['thickness'] = self._options['edge_thickness']

        # Set labels param to add labels on the fly
        labels = False
        if self._options['edge_labels']:
            labels = True
            self._plot_components['edge_labels'] = []

        # Make dict collection of all edges (keep label and edge color)
        edges_to_draw = {}

        def append_or_set(key, label, color, head):
            if key in edges_to_draw:
                edges_to_draw[key].append((label, color, head))
            else:
                edges_to_draw[key] = [(label, color, head)]

        v_to_int = {v: i for i,v in enumerate(self._graph)}

        if self._options['color_by_label'] or isinstance(self._options['edge_colors'], dict):
            if self._options['color_by_label']:
                edge_colors = self._graph._color_by_label(format=self._options['color_by_label'])
            else:
                edge_colors = self._options['edge_colors']
            edges_drawn = []
            for color in edge_colors:
                for edge in edge_colors[color]:
                    if v_to_int[edge[0]] < v_to_int[edge[1]]:
                        key = (edge[0], edge[1])
                        head = 1
                    else:
                        key = (edge[1], edge[0])
                        head = 0
                    if len(edge) < 3:
                        label = self._graph.edge_label(edge[0], edge[1])
                        if isinstance(label, list):
                            append_or_set(key, label[-1], color, head)
                            edges_drawn.append((edge[0], edge[1], label[-1]))
                            for i in range(len(label) - 1):
                                edges_to_draw[key].append((label[i], color, head))
                                edges_drawn.append((edge[0], edge[1], label[i]))
                        else:
                            append_or_set(key, label, color, head)
                            edges_drawn.append((edge[0], edge[1], label))
                    else:
                        label = edge[2]
                        labelList = self._graph.edge_label(edge[0], edge[1])
                        if isinstance(labelList, list):
                            for l in labelList:
                                if l == label:
                                    append_or_set(key, label, color, head)
                                    edges_drawn.append((edge[0], edge[1], label))
                        else:
                            if labelList == label:
                                append_or_set(key, label, color, head)
                                edges_drawn.append((edge[0], edge[1], label))

            # Add unspecified edges (default color black set in DEFAULT_PLOT_OPTIONS)
            for edge in self._graph.edge_iterator():
                if ((edge[0], edge[1], edge[2]) not in edges_drawn and
                    (self._graph.is_directed() or
                     (edge[1], edge[0], edge[2]) not in edges_drawn
                    )):
                    if v_to_int[edge[0]] < v_to_int[edge[1]]:
                        key = (edge[0], edge[1])
                        head = 1
                    else:
                        key = (edge[1], edge[0])
                        head = 0
                    append_or_set(key, edge[2], self._options['edge_color'], head)

        else:
            for edge in self._graph.edge_iterator():
                if v_to_int[edge[0]] < v_to_int[edge[1]]:
                    key = (edge[0], edge[1])
                    head = 1
                else:
                    key = (edge[1], edge[0])
                    head = 0
                append_or_set(key, edge[2], self._options['edge_color'], head)

        if edges_to_draw:
            self._plot_components['edges'] = []
        else:
            return

        # Check for multi-edges or loops
        if self._arcs or self._loops:
            tmp = edges_to_draw.copy()
            dist = self._options['dist'] * 2.
            loop_size = self._options['loop_size']
            max_dist = self._options['max_dist']
            from sage.functions.all import sqrt
            for a, b in tmp:
                if a == b:
                    # Loops
                    distance = dist
                    local_labels = edges_to_draw.pop((a, b))
                    len_local_labels = len(local_labels)
                    if len_local_labels * dist > max_dist:
                        distance = float(max_dist) / len_local_labels
                    curr_loop_size = loop_size
                    for i in range(len_local_labels):
                        self._plot_components['edges'].append(circle((self._pos[a][0],
                            self._pos[a][1]-curr_loop_size), curr_loop_size,
                            rgbcolor=local_labels[i][1], **eoptions))
                        if labels:
                            self._plot_components['edge_labels'].append(text(local_labels[i][0],
                                (self._pos[a][0], self._pos[a][1] - 2 * curr_loop_size),
                                background_color=self._options['edge_labels_background']))
                        curr_loop_size += distance / 4
                elif len(edges_to_draw[a,b]) > 1:
                    # Multi-edge
                    local_labels = edges_to_draw.pop((a,b))

                    # Compute perpendicular bisector
                    p1 = self._pos[a]
                    p2 = self._pos[b]
                    M = ((p1[0]+p2[0])/2., (p1[1]+p2[1])/2.) # midpoint
                    if not p1[1] == p2[1]:
                        S = float(p1[0]-p2[0]) / (p2[1]-p1[1]) # perp slope
                        y = lambda x: S*x-S*M[0]+M[1] # perp bisector line

                        # f,g are functions of distance d to determine x values
                        # on line y at d from point M
                        f = lambda d: sqrt(d**2/(1.+S**2)) + M[0]
                        g = lambda d: -sqrt(d**2/(1.+S**2)) + M[0]

                        odd_x = f
                        even_x = g
                        if p1[0] == p2[0]:
                            odd_y = lambda d: M[1]
                            even_y = odd_y
                        else:
                            odd_y = lambda x: y(f(x))
                            even_y = lambda x: y(g(x))
                    else:
                        odd_x = lambda d: M[0]
                        even_x = odd_x
                        odd_y = lambda d: M[1] + d
                        even_y = lambda d: M[1] - d

                    # We now have the control points for each bezier curve
                    # in terms of distance parameter d.
                    # Also note that the label for each edge should be drawn at d/2.
                    # (This is because we're using the perp bisectors).
                    distance = dist
                    len_local_labels = len(local_labels)
                    if len_local_labels * dist > max_dist:
                        distance = float(max_dist) / len_local_labels
                    for i in range(len_local_labels // 2):
                        k = (i + 1.0) * distance
                        if self._arcdigraph:
                            odd_start = self._polar_hack_for_multidigraph(p1,
                                [odd_x(k), odd_y(k)], self._vertex_radius)[0]
                            odd_end = self._polar_hack_for_multidigraph([odd_x(k), odd_y(k)],
                                p2, self._vertex_radius)[1]
                            even_start = self._polar_hack_for_multidigraph(p1,
                                [even_x(k), even_y(k)], self._vertex_radius)[0]
                            even_end = self._polar_hack_for_multidigraph([even_x(k), even_y(k)],
                                p2, self._vertex_radius)[1]
                            self._plot_components['edges'].append(arrow(path=[[odd_start,
                                [odd_x(k), odd_y(k)], odd_end]], head=local_labels[2*i][2],
                                zorder=1, rgbcolor=local_labels[2*i][1], **eoptions))
                            self._plot_components['edges'].append(arrow(path=[[even_start,
                                [even_x(k), even_y(k)], even_end]], head=local_labels[2*i+1][2],
                                zorder=1, rgbcolor=local_labels[2*i+1][1], **eoptions))
                        else:
                            self._plot_components['edges'].append(bezier_path([[p1,
                                [odd_x(k), odd_y(k)], p2]], zorder=1,
                                rgbcolor=local_labels[2*i][1], **eoptions))
                            self._plot_components['edges'].append(bezier_path([[p1,
                                [even_x(k), even_y(k)], p2]], zorder=1,
                                rgbcolor=local_labels[2*i+1][1], **eoptions))
                        if labels:
                            j = k / 2.0
                            self._plot_components['edge_labels'].append(text(local_labels[2*i][0],
                                [odd_x(j), odd_y(j)], background_color=self._options['edge_labels_background']))
                            self._plot_components['edge_labels'].append(text(local_labels[2*i+1][0],
                                [even_x(j), even_y(j)],
                                background_color=self._options['edge_labels_background']))
                    if len_local_labels % 2:
                        edges_to_draw[a,b] = [local_labels[-1]] # draw line for last odd

        is_directed = self._graph.is_directed()
        for a,b in edges_to_draw:
            if self._arcdigraph:
                C, D = self._polar_hack_for_multidigraph(self._pos[a], self._pos[b], self._vertex_radius)
                self._plot_components['edges'].append(arrow(C, D,
                    rgbcolor=edges_to_draw[a,b][0][1], head=edges_to_draw[a,b][0][2],
                    **eoptions))
                if labels:
                    self._plot_components['edge_labels'].append(text(str(edges_to_draw[a,b][0][0]),
                        [(C[0]+D[0])/2., (C[1]+D[1])/2.],
                        background_color=self._options['edge_labels_background']))
            elif is_directed:
                self._plot_components['edges'].append(arrow(self._pos[a], self._pos[b],
                    rgbcolor=edges_to_draw[a,b][0][1], arrowshorten=self._arrowshorten,
                    head=edges_to_draw[a,b][0][2], **eoptions))
            else:
                self._plot_components['edges'].append(line([self._pos[a], self._pos[b]],
                    rgbcolor=edges_to_draw[a,b][0][1], **eoptions))
            if labels and not self._arcdigraph:
                self._plot_components['edge_labels'].append(text(str(edges_to_draw[a,b][0][0]),
                    [(self._pos[a][0] + self._pos[b][0])/2.,
                    (self._pos[a][1] + self._pos[b][1])/2.],
                    background_color=self._options['edge_labels_background']))

    def _polar_hack_for_multidigraph(self, A, B, VR):
        """
        Helper function to quickly compute the two points of intersection of a
        line segment from A to B (xy tuples) and circles centered at A and B,
        both with radius VR.  Returns a tuple of xy tuples representing the two
        points.

        EXAMPLES::

            sage: d = DiGraph(loops=True, multiedges=True, sparse=True)
            sage: d.add_edges([(0,0,'a'),(0,0,'b'),(0,1,'c'),(0,1,'d'),
            ....:   (0,1,'e'),(0,1,'f'),(0,1,'f'),(2,1,'g'),(2,2,'h')])
            sage: GP = d.graphplot(vertex_size=100, edge_labels=True, color_by_label=True,
            ....:  edge_style='dashed')
            sage: GP._polar_hack_for_multidigraph((0, 1), (1, 1), .1)
            ([0.10..., 1.00...], [0.90..., 1.00...])

        TESTS:

        Make sure that Python ints are acceptable arguments (:trac:`10124`)::

            sage: GP = DiGraph().graphplot()
            sage: GP._polar_hack_for_multidigraph((0, 1), (2, 2), .1)
            ([0.08..., 1.04...], [1.91..., 1.95...])
            sage: GP._polar_hack_for_multidigraph((int(0), int(1)), (int(2), int(2)), .1)
            ([0.08..., 1.04...], [1.91..., 1.95...])

        """
        D = [float(B[i] - A[i]) for i in range(2)]
        R = sqrt(D[0]**2 + D[1]**2)
        theta = 3 * pi / 2
        if D[0] > 0:
            theta = atan(D[1] / D[0])
            if D[1] < 0:
                theta += 2 * pi
        elif D[0] < 0:
            theta = atan(D[1] / D[0]) + pi
        elif D[1] > 0:
            theta = pi / 2
        cos_theta = cos(theta)
        sin_theta = sin(theta)
        return ([VR * cos_theta + A[0], VR * sin_theta + A[1]],
                [(R - VR) * cos_theta + A[0], (R - VR) * sin_theta + A[1]])

    def show(self, **kwds):
        """
        Show the (Di)Graph associated with this ``GraphPlot`` object.

        INPUT:

        This method accepts all parameters of
        :meth:`sage.plot.graphics.Graphics.show`.

        .. NOTE::

            - See :mod:`the module's documentation <sage.graphs.graph_plot>` for
              information on default values of this method.

            - Any options not used by plot will be passed on to the
              :meth:`~sage.plot.graphics.Graphics.show` method.

        EXAMPLES::

            sage: C = graphs.CubeGraph(8)
            sage: P = C.graphplot(vertex_labels=False, vertex_size=0, graph_border=True)
            sage: P.show()

        .. PLOT::

            C = graphs.CubeGraph(8)
            P = C.graphplot(vertex_labels=False, vertex_size=0, graph_border=True)
            sphinx_plot(P)

        """
        # Setting the default values if needed
        for k, value in DEFAULT_SHOW_OPTIONS.items():
            if k not in kwds:
                kwds[k] = value

        self.plot().show(**kwds)

    def plot(self, **kwds):
        """
        Return a graphics object representing the (di)graph.

        INPUT:

        The options accepted by this method are to be found in the documentation
        of the :mod:`sage.graphs.graph_plot` module, and the
        :meth:`~sage.plot.graphics.Graphics.show` method.

        .. NOTE::

            See :mod:`the module's documentation <sage.graphs.graph_plot>` for
            information on default values of this method.

        We can specify some pretty precise plotting of familiar graphs::

            sage: from math import sin, cos, pi
            sage: P = graphs.PetersenGraph()
            sage: d = {'#FF0000':[0,5], '#FF9900':[1,6], '#FFFF00':[2,7], '#00FF00':[3,8],
            ....:  '#0000FF':[4,9]}
            sage: pos_dict = {}
            sage: for i in range(5):
            ....:  x = float(cos(pi/2 + ((2*pi)/5)*i))
            ....:  y = float(sin(pi/2 + ((2*pi)/5)*i))
            ....:  pos_dict[i] = [x,y]
            ...
            sage: for i in range(5,10):
            ....:  x = float(0.5*cos(pi/2 + ((2*pi)/5)*i))
            ....:  y = float(0.5*sin(pi/2 + ((2*pi)/5)*i))
            ....:  pos_dict[i] = [x,y]
            ...
            sage: pl = P.graphplot(pos=pos_dict, vertex_colors=d)
            sage: pl.show()

        .. PLOT::

            from math import sin, cos, pi
            P = graphs.PetersenGraph()
            d = {'#FF0000':[0,5], '#FF9900':[1,6], '#FFFF00':[2,7], '#00FF00':[3,8],
                '#0000FF':[4,9]}
            pos_dict = {}
            for i in range(5):
                x = float(cos(pi/2 + ((2*pi)/5)*i))
                y = float(sin(pi/2 + ((2*pi)/5)*i))
                pos_dict[i] = [x,y]

            for i in range(5,10):
                x = float(0.5*cos(pi/2 + ((2*pi)/5)*i))
                y = float(0.5*sin(pi/2 + ((2*pi)/5)*i))
                pos_dict[i] = [x,y]

            pl = P.graphplot(pos=pos_dict, vertex_colors=d)
            sphinx_plot(pl)

        Here are some more common graphs with typical options::

            sage: C = graphs.CubeGraph(8)
            sage: P = C.graphplot(vertex_labels=False, vertex_size=0, graph_border=True)
            sage: P.show()

        .. PLOT::

            C = graphs.CubeGraph(8)
            P = C.graphplot(vertex_labels=False, vertex_size=0, graph_border=True)
            sphinx_plot(P)

        ::

            sage: G = graphs.HeawoodGraph().copy(sparse=True)
            sage: for u,v,l in G.edges():
            ....:  G.set_edge_label(u,v,'(' + str(u) + ',' + str(v) + ')')
            sage: G.graphplot(edge_labels=True).show()

        .. PLOT::

            G = graphs.HeawoodGraph().copy(sparse=True)
            for u,v,l in G.edges():
                G.set_edge_label(u,v,'(' + str(u) + ',' + str(v) + ')')
            sphinx_plot(G.graphplot(edge_labels=True))

        The options for plotting also work with directed graphs::

            sage: D = DiGraph( { 0: [1, 10, 19], 1: [8, 2], 2: [3, 6], 3: [19, 4],
            ....:  4: [17, 5], 5: [6, 15], 6: [7], 7: [8, 14], 8: [9], 9: [10, 13],
            ....:  10: [11], 11: [12, 18], 12: [16, 13], 13: [14], 14: [15], 15: [16],
            ....:  16: [17], 17: [18], 18: [19], 19: []})
            sage: for u,v,l in D.edges():
            ....:  D.set_edge_label(u,v,'(' + str(u) + ',' + str(v) + ')')
            sage: D.graphplot(edge_labels=True, layout='circular').show()

        .. PLOT::

            D = DiGraph( { 0: [1, 10, 19], 1: [8, 2], 2: [3, 6], 3: [19, 4], 4: [17, 5],
                5: [6, 15], 6: [7], 7: [8, 14], 8: [9], 9: [10, 13], 10: [11],
                11: [12, 18],12: [16, 13], 13: [14], 14: [15], 15: [16], 16: [17],
                17: [18], 18: [19], 19: []})
            for u,v,l in D.edges():
                D.set_edge_label(u,v,'(' + str(u) + ',' + str(v) + ')')
            sphinx_plot(D.graphplot(edge_labels=True, layout='circular'))

        This example shows off the coloring of edges::

            sage: from sage.plot.colors import rainbow
            sage: C = graphs.CubeGraph(5)
            sage: R = rainbow(5)
            sage: edge_colors = {}
            sage: for i in range(5):
            ....:  edge_colors[R[i]] = []
            sage: for u,v,l in C.edges():
            ....:  for i in range(5):
            ....:      if u[i] != v[i]:
            ....:          edge_colors[R[i]].append((u,v,l))
            sage: C.graphplot(vertex_labels=False, vertex_size=0, edge_colors=edge_colors).show()

        .. PLOT::

            from sage.plot.colors import rainbow
            C = graphs.CubeGraph(5)
            R = rainbow(5)
            edge_colors = {}
            for i in range(5):
                edge_colors[R[i]] = []
            for u,v,l in C.edges():
                for i in range(5):
                    if u[i] != v[i]:
                        edge_colors[R[i]].append((u,v,l))
            sphinx_plot(C.graphplot(vertex_labels=False, vertex_size=0,
                edge_colors=edge_colors))

        With the ``partition`` option, we can separate out same-color groups
        of vertices::

            sage: D = graphs.DodecahedralGraph()
            sage: Pi = [[6,5,15,14,7],[16,13,8,2,4],[12,17,9,3,1],[0,19,18,10,11]]
            sage: D.show(partition=Pi)

        .. PLOT::

            D = graphs.DodecahedralGraph()
            Pi = [[6,5,15,14,7],[16,13,8,2,4],[12,17,9,3,1],[0,19,18,10,11]]
            sphinx_plot(D.plot(partition=Pi))

        Loops are also plotted correctly::

            sage: G = graphs.PetersenGraph()
            sage: G.allow_loops(True)
            sage: G.add_edge(0,0)
            sage: G.show()

        .. PLOT::

            G = graphs.PetersenGraph()
            G.allow_loops(True)
            G.add_edge(0,0)
            sphinx_plot(G)

        ::

            sage: D = DiGraph({0:[0,1], 1:[2], 2:[3]}, loops=True)
            sage: D.show()
            sage: D.show(edge_colors={(0,1,0):[(0,1,None),(1,2,None)],(0,0,0):[(2,3,None)]})

        .. PLOT::

            D = DiGraph({0:[0,1], 1:[2], 2:[3]}, loops=True)
            P = D.plot(edge_colors={(0,1,0):[(0,1,None),(1,2,None)],(0,0,0):[(2,3,None)]})
            sphinx_plot(P)

        More options::

            sage: pos = {0:[0.0, 1.5], 1:[-0.8, 0.3], 2:[-0.6, -0.8],
            ....:  3:[0.6, -0.8], 4:[0.8, 0.3]}
            sage: g = Graph({0:[1], 1:[2], 2:[3], 3:[4], 4:[0]})
            sage: g.graphplot(pos=pos, layout='spring', iterations=0).plot()
            Graphics object consisting of 11 graphics primitives

        .. PLOT::

            pos = {0:[0.0, 1.5], 1:[-0.8, 0.3], 2:[-0.6, -0.8],
                3:[0.6, -0.8], 4:[0.8, 0.3]}
            g = Graph({0:[1], 1:[2], 2:[3], 3:[4], 4:[0]})
            P = g.graphplot(pos=pos, layout='spring', iterations=0).plot()
            sphinx_plot(P)

        ::

            sage: G = Graph()
            sage: P = G.graphplot().plot()
            sage: P.axes()
            False
            sage: G = DiGraph()
            sage: P = G.graphplot().plot()
            sage: P.axes()
            False

        We can plot multiple graphs::

            sage: T = list(graphs.trees(7))
            sage: t = T[3]
            sage: t.graphplot(heights={0:[0], 1:[4,5,1], 2:[2], 3:[3,6]}).plot()
            Graphics object consisting of 14 graphics primitives

        .. PLOT::

            T = list(graphs.trees(7))
            t = T[3]
            sphinx_plot(t.graphplot(heights={0:[0], 1:[4,5,1], 2:[2], 3:[3,6]}))

        ::

            sage: T = list(graphs.trees(7))
            sage: t = T[3]
            sage: t.graphplot(heights={0:[0], 1:[4,5,1], 2:[2], 3:[3,6]}).plot()
            Graphics object consisting of 14 graphics primitives

        .. PLOT::

            T = list(graphs.trees(7))
            t = T[3]
            sphinx_plot(t.graphplot(heights={0:[0], 1:[4,5,1], 2:[2], 3:[3,6]}))

        ::

            sage: t.set_edge_label(0,1,-7)
            sage: t.set_edge_label(0,5,3)
            sage: t.set_edge_label(0,5,99)
            sage: t.set_edge_label(1,2,1000)
            sage: t.set_edge_label(3,2,'spam')
            sage: t.set_edge_label(2,6,3/2)
            sage: t.set_edge_label(0,4,66)
            sage: t.graphplot(heights={0:[0], 1:[4,5,1], 2:[2], 3:[3,6]}, edge_labels=True).plot()
            Graphics object consisting of 20 graphics primitives

        .. PLOT::

            T = list(graphs.trees(7))
            t = T[3]
            t.set_edge_label(0,1,-7)
            t.set_edge_label(0,5,3)
            t.set_edge_label(0,5,99)
            t.set_edge_label(1,2,1000)
            t.set_edge_label(3,2,'spam')
            t.set_edge_label(2,6,3/2)
            t.set_edge_label(0,4,66)
            sphinx_plot(t.graphplot(heights={0:[0], 1:[4,5,1], 2:[2], 3:[3,6]}, edge_labels=True))

        ::

            sage: T = list(graphs.trees(7))
            sage: t = T[3]
            sage: t.graphplot(layout='tree').show()

        .. PLOT::

            T = list(graphs.trees(7))
            t = T[3]
            sphinx_plot(t.graphplot(layout='tree'))

        The tree layout is also useful::

            sage: t = DiGraph('JCC???@A??GO??CO??GO??')
            sage: t.graphplot(layout='tree', tree_root=0, tree_orientation="up").show()

        .. PLOT::

            t = DiGraph('JCC???@A??GO??CO??GO??')
            sphinx_plot(t.graphplot(layout='tree', tree_root=0, tree_orientation="up"))

        More examples::

            sage: D = DiGraph({0:[1,2,3], 2:[1,4], 3:[0]})
            sage: D.graphplot().show()

        .. PLOT::

            D = DiGraph({0:[1,2,3], 2:[1,4], 3:[0]})
            sphinx_plot(D.graphplot())

        ::

            sage: D = DiGraph(multiedges=True, sparse=True)
            sage: for i in range(5):
            ....:   D.add_edge((i,i+1,'a'))
            ....:   D.add_edge((i,i-1,'b'))
            sage: D.graphplot(edge_labels=True,edge_colors=D._color_by_label()).plot()
            Graphics object consisting of 34 graphics primitives

        .. PLOT::

            D = DiGraph(multiedges=True, sparse=True)
            for i in range(5):
                 D.add_edge((i,i+1,'a'))
                 D.add_edge((i,i-1,'b'))
            sphinx_plot(D.graphplot(edge_labels=True,edge_colors=D._color_by_label()))

        ::

            sage: g = Graph({}, loops=True, multiedges=True, sparse=True)
            sage: g.add_edges([(0,0,'a'),(0,0,'b'),(0,1,'c'),(0,1,'d'),
            ....:   (0,1,'e'),(0,1,'f'),(0,1,'f'),(2,1,'g'),(2,2,'h')])
            sage: g.graphplot(edge_labels=True, color_by_label=True, edge_style='dashed').plot()
            Graphics object consisting of 26 graphics primitives

        .. PLOT::

            g = Graph({}, loops=True, multiedges=True, sparse=True)
            g.add_edges([(0,0,'a'),(0,0,'b'),(0,1,'c'),(0,1,'d'),
                 (0,1,'e'),(0,1,'f'),(0,1,'f'),(2,1,'g'),(2,2,'h')])
            sphinx_plot(g.graphplot(edge_labels=True, color_by_label=True, edge_style='dashed'))

        The ``edge_style`` option may be provided in the short format too::

            sage: g.graphplot(edge_labels=True, color_by_label=True, edge_style='--').plot()
            Graphics object consisting of 26 graphics primitives

        TESTS:

        Make sure that show options work with plot also::

            sage: g = Graph({})
            sage: g.plot(title='empty graph', axes=True)
            Graphics object consisting of 0 graphics primitives

        Check for invalid inputs::

            sage: p = graphs.PetersenGraph().plot(egabrag='garbage')
            Traceback (most recent call last):
            ...
            ValueError: Invalid input 'egabrag=garbage'

        Make sure that no graphics primitive is clipped::

            sage: tadpole = Graph({0:[0,1]}).plot()
            sage: bbox = tadpole.get_minmax_data()
            sage: for part in tadpole:
            ....:      part_bbox = part.get_minmax_data()
            ....:      assert bbox['xmin'] <= part_bbox['xmin'] <= part_bbox['xmax'] <= bbox['xmax']
            ....:      assert bbox['ymin'] <= part_bbox['ymin'] <= part_bbox['ymax'] <= bbox['ymax']

        Check that one can plot immutable graphs (:trac:`17340`)::

            sage: Graph({0:[0]},immutable=True).plot()
            Graphics object consisting of 3 graphics primitives
        """
        G = Graphics()
        options = self._options.copy()
        options.update(kwds)
        G._set_extra_kwds(Graphics._extract_kwds_for_show(options))

        # Check the arguments
        for o in options:
            if o not in graphplot_options and o not in G._extra_kwds:
                raise ValueError("Invalid input '{}={}'".format(o, options[o]))

        for comp in self._plot_components.values():
            if not isinstance(comp, list):
                G += comp
            else:
                for item in comp:
                    G += item

        if self._options['graph_border']:
            xmin = G.xmin()
            xmax = G.xmax()
            ymin = G.ymin()
            ymax = G.ymax()
            dx = (xmax - xmin) / 10.0
            dy = (ymax - ymin) / 10.0
            border = (line([(xmin - dx, ymin - dy), (xmin - dx, ymax + dy),
                            (xmax + dx, ymax + dy), (xmax + dx, ymin - dy),
                            (xmin - dx, ymin - dy)], thickness=1.3))
            border.axes_range(xmin=(xmin - dx), xmax=(xmax + dx),
                              ymin=(ymin - dy), ymax=(ymax + dy))
            G += border
        G.set_aspect_ratio(1)
        G.axes(False)
        return G

    def layout_tree(self, root, orientation):
        """
        Compute a nice layout of a tree.

        INPUT:

        - ``root`` -- the root vertex.

        - ``orientation`` -- whether to place the root at the top or at the
          bottom:

          * ``orientation="down"`` -- children are placed below their parent
          * ``orientation="top"`` -- children are placed above their parent

        EXAMPLES::

            sage: from sage.graphs.graph_plot import GraphPlot
            sage: G = graphs.HoffmanSingletonGraph()
            sage: T = Graph()
            sage: T.add_edges(G.min_spanning_tree(starting_vertex=0))
            sage: T.show(layout='tree', tree_root=0) # indirect doctest

        """
        T = self._graph

        if not self._graph.is_tree():
            raise RuntimeError("Cannot use tree layout on this graph: self.is_tree() returns False.")

        children = {root: T.neighbors(root)}

        #always make a copy of the children because they get eaten
        stack = [[u for u in children[root]]]
        stick = [root]
        parent = {u: root for u in children[root]}
        pos = {}
        obstruction = [0.0] * T.num_verts()
        if orientation == 'down':
            o = -1
        else:
            o = 1

        def slide(v, dx):
            """
            Shift the vertex v and its descendants to the right by dx.

            Precondition: v and its descendents have already had their
            positions computed.

            """
            level = [v]
            while level:
                nextlevel = []
                for u in level:
                    x, y = pos[u]
                    x += dx
                    obstruction[y] = max(x + 1, obstruction[y])
                    pos[u] = x, y
                    nextlevel += children[u]

                level = nextlevel

        while stack:
            C = stack[-1]
            if not C:
                p = stick.pop()
                stack.pop()
                cp = children[p]
                y = o * len(stack)
                if not cp:
                    x = obstruction[y]
                    pos[p] = x, y
                else:
                    x = sum([pos[c][0] for c in cp]) / float(len(cp))
                    pos[p] = x,y
                    ox = obstruction[y]
                    if x < ox:
                        slide(p, ox - x)
                        x = ox
                obstruction[y] = x+1
                continue

            t = C.pop()
            pt = parent[t]

            ct = [u for u in T.neighbors(t) if u != pt]
            for c in ct:
                parent[c] = t
            children[t] = ct

            stack.append(ct)
            stick.append(t)

        return pos

max_issues_repo_name: Ivo-Maffei/sage

max_forks_repo_name: dimpase/sage

reason: import sage,from sage


text:
#!/usr/bin/env python3
# Copyright (c) Facebook, Inc. and its affiliates.
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.
import sys

import mock
import numpy as np

from habitat.config.default import get_config
from habitat.sims import make_sim

class CameraMock:
    def get_rgb(self):
        return np.zeros((256, 256, 3))

    def get_depth(self):
        return np.zeros((256, 256, 1))

    def reset(self):
        pass

    def step(self, *args, **kwargs):
        pass

class RobotMock:
    def __init__(self, *args, **kwargs):
        self.camera = CameraMock()
        self.base = BaseMock()

class BaseMock:
    def __init__(self, *args, **kwargs):
        self.base_state = mock.MagicMock()
        self.base_state.bumper = False

    def go_to_relative(self, *args, **kwargs):
        pass

def test_pyrobot(mocker):
    if "pyrobot" not in sys.modules:
        # Mock pyrobot package if it is not installed
        mock_pyrobot = mocker.MagicMock()
        mock_pyrobot.Robot = RobotMock
        sys.modules["pyrobot"] = mock_pyrobot

        # Re-register pyrobot with mock
        from habitat.sims.registration import _try_register_pyrobot

        _try_register_pyrobot()

    config = get_config()
    with make_sim("PyRobot-v0", config=config.PYROBOT) as reality:

        observations = reality.reset()
        observations = reality.step(
            "go_to_relative",
            {
                "xyt_position": [0, 0, (10 / 180) * np.pi],
                "use_map": False,
                "close_loop": True,
                "smooth": False,
            },
        )

ext: py

max_stars_repo_name: ruffsl/habitat-api

max_issues_repo_name: ruffsl/habitat-api

max_forks_repo_name: ruffsl/habitat-api

reason: import numpy


text:
from numbers import Number

import numpy as np
import numpy.testing as npt

from quaternion import Quaternion

class DualQuaternion(object):
  """ Clifford dual quaternion denoted as dq = q_rot + epsilon * q_dual.

  Can be instantiated by:
  >>> dq = DualQuaternion(q_rot, q_dual)
  >>> dq = DualQuaternion.from_vector([qrx, qry, qrz, qrw, qtx, qty, qtz, qtw])
  >>> dq = DualQuaternion.from_pose(x, y, z, q_x, q_y, q_z, q_w)
  >>> dq = DualQuaternion.from_pose_vector([x, y, z, q_x, q_y, q_z, q_w])
  >>> dq = DualQuaternion(q_rot, q_dual)
  >>> # Given a [4x4] transformation matrix T.
  >>> dq = DualQuaternion.from_transformation_matrix(T)
  """
  dq = np.array([0., 0., 0., 1.0, 0., 0., 0., 0.]).T

  def __init__(self, q_rot, q_dual):
    for i in [q_rot, q_dual]:
      assert isinstance(
          i, Quaternion), "q_rot and q_dual should be quaternions."

    self.dq = np.array([0., 0., 0., 1.0, 0., 0., 0., 0.]).T

    # Assign real part (rotation).
    self.dq[0:4] = q_rot.q.copy()

    # Assign dual part (translation).
    self.dq[4:8] = q_dual.q.copy()

    self.assert_normalization()

  def __str__(self):
    return "[q_rot: {}, q_dual: {}]".format(np.str(self.q_rot), np.str(self.q_dual))

  def __repr__(self):
    return ("<Dual quaternion q_rot {} q_dual {}>").format(self.q_rot, self.q_dual)

  def __add__(self, other):
    """ Dual quaternion addition. """
    dq_added = self.dq + other.dq
    return DualQuaternion.from_vector(dq_added)

  def __sub__(self, other):
    """ Dual quaternion subtraction. """
    dq_sub = self.dq - other.dq
    return DualQuaternion.from_vector(dq_sub)

  def __mul__(self, other):
    """ Dual quaternion multiplication.

    The multiplication with a scalar returns the dual quaternion with all
    elements multiplied by the scalar.

    The multiplication of two dual quaternions dq1 and dq2 as:
    q1_rot * q2_rot + epsilon * (q1_rot * q2_trans + q1_trans * q2_rot),
    where dq1 and dq2 are defined as:
    dq1 = q1_rot + epsilon * q1_trans,
    dq2 = q2_rot + epsilon * q2_trans.
    """
    if isinstance(other, DualQuaternion):
      rotational_part = self.q_rot * other.q_rot
      translational_part = (self.q_rot * other.q_dual +
                            self.q_dual * other.q_rot)
      return DualQuaternion(rotational_part.copy(), translational_part.copy())
    elif isinstance(other, Number):
      dq = self.dq.copy()
      dq_out = dq * np.float64(other)
      return DualQuaternion.from_vector(dq_out)
    else:
      assert False, ("Multiplication is only defined for scalars or dual " "quaternions.")

  def __rmul__(self, other):
    """ Scalar dual quaternion multiplication.

    The multiplication with a scalar returns the dual quaternion with all
    elements multiplied by the scalar.
    """
    if isinstance(other, Number):
      dq = self.dq.copy()
      dq_out = np.float64(other) * dq
      return DualQuaternion.from_vector(dq_out)
    else:
      assert False, ("Multiplication is only defined for scalars or dual " "quaternions.")

  def __truediv__(self, other):
    """ Quaternion division with either scalars or quaternions.

    The division with a scalar returns the dual quaternion with all
    translational elements divided by the scalar.

    The division with a dual quaternion returns dq = dq1/dq2 = dq1 * dq2^-1,
    hence other divides on the right.
    """
    # TODO(ff): Check if this is correct.
    print("WARNING: This might not be properly implemented.")
    if isinstance(other, DualQuaternion):
      return self * other.inverse()
    elif isinstance(other, Number):
      dq = self.dq.copy()
      dq_out = dq / np.float64(other)
      return DualQuaternion.from_vector(dq_out)
    else:
      assert False, "Division is only defined for scalars or dual quaternions."

  def __div__(self, other):
    """ Quaternion division with either scalars or quaternions.

    The division with a scalar returns the dual quaternion with all
    translational elements divided by the scalar.

    The division with a dual quaternion returns dq = dq1 / dq2 = dq1 * dq2^-1.
    """
    return self.__truediv__(other)

  def __eq__(self, other):
    """ Check equality. """
    if isinstance(other, DualQuaternion):
      return np.allclose(self.dq, other.dq)
    else:
      return False

  @classmethod
  def from_vector(cls, dq):
    dual_quaternion_vector = None
    if isinstance(dq, np.ndarray):
      dual_quaternion_vector = dq.copy()
    else:
      dual_quaternion_vector = np.array(dq)
    return cls(Quaternion(q=dual_quaternion_vector[0:4]),
               Quaternion(q=dual_quaternion_vector[4:8]))

  @classmethod
  def from_pose(cls, x, y, z, rx, ry, rz, rw):
    """ Create a normalized dual quaternion from a pose. """
    qr = Quaternion(rx, ry, rz, rw)
    qr.normalize()
    qt = (Quaternion(x, y, z, 0) * qr) * 0.5
    return cls(qr, qt)

  @classmethod
  def from_pose_vector(cls, pose):
    """ Create a normalized dual quaternion from a pose vector. """
    return cls.from_pose(pose[0], pose[1], pose[2], pose[3], pose[4], pose[5], pose[6])

  @classmethod
  def from_transformation_matrix(cls, transformation_matrix):
    q_rot = Quaternion.from_rotation_matrix(transformation_matrix[0:3, 0:3])

    pose_vec = np.zeros(7)
    pose_vec[3:7] = q_rot.q
    pose_vec[0:3] = transformation_matrix[0:3, 3]

    return cls.from_pose_vector(pose_vec)

  @classmethod
  def identity(cls):
    identity_q_rot = Quaternion(0., 0., 0., 1.)
    identity_q_dual = Quaternion(0., 0., 0., 0.)
    return cls(identity_q_rot, identity_q_dual)

  # def conjugate_translation(self):
  #   """ Dual quaternion translation conjugate. """
  #   return DualQuaternion(self.q_rot.conjugate(), -self.q_dual.conjugate())

  def conjugate(self):
    """ Dual quaternion multiplication conjugate. """
    return DualQuaternion(self.q_rot.conjugate(), self.q_dual.conjugate())

  def inverse(self):
    """ Dual quaternion inverse. """
    assert self.norm()[0] > 1e-8
    return DualQuaternion(self.q_rot.inverse(),
                          -self.q_rot.inverse() * self.q_dual * self.q_rot.inverse())

  def enforce_positive_q_rot_w(self):
    """ Enforce a positive real part of the rotation quaternion. """
    assert self.norm()[0] > 1e-8
    if self.q_rot.w < 0.0:
      self.dq = -self.dq

  def norm(self):
    """ The norm of a dual quaternion. """
    assert self.q_rot.norm() > 1e-8, (
        "Dual quaternion has rotational part equal to zero, hence the norm is"
        "not defined.")
    real_norm = self.q_rot.norm()
    dual_norm = np.dot(self.q_rot.q, self.q_dual.q) / real_norm
    return (real_norm, dual_norm)

  def is_normalized(self):
    real_part = np.absolute(self.norm()[0] - 1.0) < 1e-8
    dual_part = np.absolute(self.norm()[1]) < 1e-8
    return real_part and dual_part

  def assert_normalization(self):
    assert self.is_normalized, "Something went wrong, the dual quaternion is not normalized!"

  def normalize(self):
    """ Normalize the dual quaternion. """
    real_norm = self.q_rot.norm()
    self.dq[0:4] = self.q_rot.q / real_norm
    self.dq[4:8] = self.q_dual.q / real_norm
    self.assert_normalization()

  def scalar(self):
    """ The scalar part of the dual quaternion.

    Defined as: scalar(dq) := 0.5*(dq+dq.conjugate())
    """
    scalar_part = 0.5 * (self + self.conjugate())
    npt.assert_allclose(
        scalar_part.dq[[0, 1, 2, 4, 5, 6]], np.zeros(6), atol=1e-6)
    return scalar_part.copy()

  def screw_axis(self):
    """ The rotation, translation and screw axis from the dual quaternion. """
    rotation = 2. * np.degrees(np.arccos(self.q_rot.w))
    rotation = np.mod(rotation, 360.)

    if (rotation > 1.e-12):
      translation = -2. * self.q_dual.w / np.sin(rotation / 2. * np.pi / 180.)
      screw_axis = self.q_rot.q[0:3] / np.sin(rotation / 2. * np.pi / 180.)
    else:
      translation = 2. * np.sqrt(np.sum(np.power(self.q_dual.q[0:3], 2.)))
      if (translation > 1.e-12):
        screw_axis = 2. * self.q_dual.q[0:3] / translation
      else:
        screw_axis = np.zeros(3)

    # TODO(ntonci): Add axis point for completeness

    return screw_axis, rotation, translation

  def passive_transform_point(self, point):
    """ Applies the passive transformation of the dual quaternion to a point.
    """
    # TODO(ff): Check if the rotation is in the right direction.
    point_dq = DualQuaternion.from_pose(
        point[0], point[1], point[2], 0, 0, 0, 1)
    dq_in_new_frame = self * point_dq
    return dq_in_new_frame.to_pose()[0:3]

  def active_transform_point(self, point):
    """ Applies the active transformation of the dual quaternion to a point.
    """
    return self.inverse().passive_transform_point(point)

  # TODO(ff): Implement translational velocity, rotational velocity.
  # See for instance:
  # https://www.ncbi.nlm.nih.gov/pmc/articles/PMC3576712/pdf/fnbeh-07-00007.pdf

  def to_matrix(self):
    """ Returns a [4x4] transformation matrix. """
    self.normalize()
    matrix_out = np.identity(4)
    matrix_out[0:3, 0:3] = self.q_rot.to_rotation_matrix()
    matrix_out[0:3, 3] = self.to_pose()[0:3]
    return matrix_out.copy()

  def to_pose(self):
    """ Returns a [7x1] pose vector.

    In the form: pose = [x, y, z, qx, qy, qz, qw].T.
    """
    self.normalize()

    pose = np.zeros(7)
    q_rot = self.q_rot
    if (q_rot.w < 0.):
      q_rot = -q_rot
    translation = (2.0 * self.q_dual) * self.q_rot.conjugate()

    pose[0:3] = translation.q[0:3].copy()
    pose[3:7] = q_rot.q.copy()
    return pose.copy()

  def copy(self):
    """ Copy dual quaternion. """
    return DualQuaternion.from_vector(self.dq)

  @property
  def q_rot(self):
    return Quaternion(q=self.dq[0:4])

  @property
  def q_dual(self):
    return Quaternion(q=self.dq[4:8])

  @property
  def r_x(self):
    return self.dq[0]

  @property
  def r_y(self):
    return self.dq[1]

  @property
  def r_z(self):
    return self.dq[2]

  @property
  def r_w(self):
    return self.dq[3]

  @property
  def t_x(self):
    return self.dq[4]

  @property
  def t_y(self):
    return self.dq[5]

  @property
  def t_z(self):
    return self.dq[6]

  @property
  def t_w(self):
    return self.dq[7]


text:
"""Calculate the distance between sequences between seasons.
"""
import argparse
from augur.frequency_estimators import float_to_datestring, timestamp_to_float
from augur.utils import annotate_parents_for_tree, read_node_data, write_json
import Bio.Phylo
import json
import numpy as np
import pandas as pd

def get_titer_distance_between_nodes(tree, past_node, current_node, titer_attr="dTiter"):
    # Find MRCA of tips from one tip up. Sum the titer attribute of interest
    # while walking up to the MRCA, to avoid an additional pass later. The loop
    # below stops when the past node is found in the list of the candidate
    # MRCA's terminals. This test should always evaluate to true when the MRCA
    # is the root node, so we should not have to worry about trying to find the
    # parent of the root.
    current_node_branch_sum = 0.0
    mrca = current_node
    while past_node.name not in mrca.terminals:
        current_node_branch_sum += mrca.attr[titer_attr]
        mrca = mrca.parent

    # Sum the node weights for the other tip from the bottom up until we reach
    # the MRCA. The value of the MRCA is intentionally excluded here, as it
    # would represent the branch leading to the MRCA and would be outside the
    # path between the two tips.
    past_node_branch_sum = 0.0
    current_node = past_node
    while current_node != mrca:
        past_node_branch_sum += current_node.attr[titer_attr]
        current_node = current_node.parent

    final_sum = past_node_branch_sum + current_node_branch_sum
    return final_sum

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--tree", help="Newick tree", required=True)
    parser.add_argument("--frequencies", help="frequencies JSON", required=True)
    parser.add_argument("--model-attribute-name", help="name of attribute to use from titer model file", default="dTiter")
    parser.add_argument("--attribute-name", help="name to store distances", required=True)
    parser.add_argument("--model", help="JSON providing the titer tree model", required=True)
    parser.add_argument("--date-annotations", help="JSON of branch lengths and date annotations from augur refine for samples in the given tree; required for comparisons to earliest or latest date", required=True)
    parser.add_argument("--months-back-for-current-samples", type=int, help="number of months prior to the last date with estimated frequencies to include samples as current", required=True)
    parser.add_argument("--years-back-to-compare", type=int, help="number of years prior to the current season to search for samples to calculate pairwise comparisons with", required=True)
    parser.add_argument("--max-past-samples", type=int, default=200, help="maximum number of past samples to randomly select for comparison to current samples")
    parser.add_argument("--min-frequency", type=float, default=0.0, help="minimum frequency to consider a sample alive")
    parser.add_argument("--output", help="JSON file with calculated distances stored by node name and attribute name", required=True)

    args = parser.parse_args()

    # Load tree and annotate parents.
    tree = Bio.Phylo.read(args.tree, "newick")
    tree = annotate_parents_for_tree(tree)

    # Make a single pass through the tree in postorder to store a set of all
    # terminals descending from each node. This uses more memory, but it allows
    # faster identification of MRCAs between any pair of tips in the tree and
    # speeds up pairwise distance calculations by orders of magnitude.
    for node in tree.find_clades(order="postorder"):
        node.terminals = set()
        for child in node.clades:
            if child.is_terminal():
                node.terminals.add(child.name)
            else:
                node.terminals.update(child.terminals)

    # Load frequencies.
    with open(args.frequencies, "r") as fh:
        frequencies = json.load(fh)

    pivots = np.array(frequencies.pop("pivots"))

    # Identify pivots that belong within our search window for past samples.
    # First, calculate dates associated with the interval for current samples
    # based on the number of months back requested. Then, calculate interval for
    # past samples with an upper bound based on the earliest current samples and
    # a lower bound based on the years back requested.
    last_pivot_datetime = pd.to_datetime(float_to_datestring(pivots[-1]))
    last_current_datetime = last_pivot_datetime - pd.DateOffset(months=args.months_back_for_current_samples)
    last_past_datetime = last_pivot_datetime - pd.DateOffset(years=args.years_back_to_compare)

    # Find the pivot indices that correspond to the current and past pivots.
    current_pivot_indices = np.array([
        pd.to_datetime(float_to_datestring(pivot)) > last_current_datetime
        for pivot in pivots
    ])
    past_pivot_indices = np.array([
        ((pd.to_datetime(float_to_datestring(pivot)) >= last_past_datetime) &
         (pd.to_datetime(float_to_datestring(pivot)) <= last_current_datetime))
        for pivot in pivots
    ])

    # Load date and titer model annotations and annotate tree with them.
    annotations = read_node_data([args.date_annotations, args.model])
    for node in tree.find_clades():
        node.attr = annotations["nodes"][node.name]
        node.attr["num_date"] = node.attr["numdate"]

    # Identify samples to compare including those in the current timepoint
    # (pivot) and those in previous timepoints.
    current_samples = []
    past_samples = []
    date_by_sample = {}
    tips_by_sample = {}
    for tip in tree.find_clades(terminal=True):
        # Samples with nonzero frequencies in the last timepoint are current
        # samples. Those with one or more nonzero frequencies in the search
        # window of the past timepoints are past samples.
        frequencies[tip.name]["frequencies"] = np.array(frequencies[tip.name]["frequencies"])
        if (frequencies[tip.name]["frequencies"][current_pivot_indices] > args.min_frequency).sum() > 0:
            current_samples.append(tip.name)
            tips_by_sample[tip.name] = tip
        elif (frequencies[tip.name]["frequencies"][past_pivot_indices] > args.min_frequency).sum() > 0:
            past_samples.append(tip.name)
            tips_by_sample[tip.name] = tip

        date_by_sample[tip.name] = tip.attr["numdate"]

    print("Expecting %i comparisons for %i current and %i past samples" % (len(current_samples) * len(past_samples), len(current_samples), len(past_samples)))
    distances_by_node = {}
    comparisons = 0

    for current_sample in current_samples:
        if not current_sample in distances_by_node:
            distances_by_node[current_sample] = {}

        if not args.attribute_name in distances_by_node[current_sample]:
            distances_by_node[current_sample][args.attribute_name] = {}

        for past_sample in past_samples:
            # The past is in the past.
            if date_by_sample[past_sample] < date_by_sample[current_sample]:
                distances_by_node[current_sample][args.attribute_name][past_sample] = np.around(get_titer_distance_between_nodes(
                    tree,
                    tips_by_sample[past_sample],
                    tips_by_sample[current_sample],
                    args.model_attribute_name
                ), 4)

                comparisons += 1
                if comparisons % 10000 == 0:
                    print("Completed", comparisons, "comparisons, with last distance of", distances_by_node[current_sample][args.attribute_name][past_sample], flush=True)

    print("Calculated %i comparisons" % comparisons)
    # Prepare params for export.
    params = {
        "attribute": args.attribute_name,
        "years_back_to_compare": args.years_back_to_compare
    }

    # Export distances to JSON.
    write_json({"params": params, "nodes": distances_by_node}, args.output, indent=None)

max_stars_repo_name: blab/flu-forecasting-all-lineages


text:
from torch.utils.data import DataLoader
from sentence_transformers import losses, util, models
from sentence_transformers import SentencesDataset, LoggingHandler, SentenceTransformer, evaluation
from sentence_transformers.readers import InputExample
import logging
from datetime import datetime
import os
from shutil import copyfile
import sys
import math
import gzip
import random
import tqdm
from transformers import AutoTokenizer, AutoModel, BertModel
import transformers
import torch
from SPARTA import SPARTA
import json
import numpy as np
from torch.cuda.amp import autocast
import os
from shutil import copyfile
import datetime
from collections import defaultdict
from scipy.sparse import csc_matrix, csr_matrix

random.seed(42)

scaler = torch.cuda.amp.GradScaler()

#### Just some code to print debug information to stdout
logging.basicConfig(format='%(asctime)s - %(message)s',
                    datefmt='%Y-%m-%d %H:%M:%S',
                    level=logging.INFO,
                    handlers=[LoggingHandler()])
#### /print debug information to stdout

# Fill GPU

device = 'cuda' if torch.cuda.is_available() else 'cpu'

model_name = sys.argv[1]
model = SPARTA(model_name, device)

model_save_path = "output/msmarco-{}-{}".format(model_name.rstrip("/").split("/")[-1], datetime.datetime.now().strftime("%Y-%m-%d_%H-%M-%S"))
model.tokenizer.save_pretrained(model_save_path)

##Distil setting
if 'distil' in model_name:
    batch_size, num_negatives = 4, 35
else:
    batch_size, num_negatives = 3, 20

logging.info(f"batch_size: {batch_size}")
logging.info(f"num_neg: {num_negatives}")

# Write self to path
os.makedirs(model_save_path, exist_ok=True)

train_script_path = os.path.join(model_save_path, 'train_script.py')
copyfile(__file__, train_script_path)
with open(train_script_path, 'a') as fOut:
    fOut.write("\n\n# Script was called via:\n#python " + " ".join(sys.argv))

########################
corpus = {}
train_queries = {}

#### Read dev file
logging.info("Create dev dataset")
dev_corpus_max_size = 100*1000

dev_queries_file = 'data/queries.dev.small.tsv'
needed_pids = set()
needed_qids = set()
dev_qids = set()

dev_queries = {}
dev_corpus = {}
dev_rel_docs = {}

with open(dev_queries_file) as fIn:
    for line in fIn:
        qid, query = line.strip().split("\t")
        dev_qids.add(qid)

with open('data/qrels.dev.tsv') as fIn:
    for line in fIn:
        qid, _, pid, _ = line.strip().split('\t')

        if qid not in dev_qids:
            continue

        if qid not in dev_rel_docs:
            dev_rel_docs[qid] = set()
        dev_rel_docs[qid].add(pid)

        needed_pids.add(pid)
        needed_qids.add(qid)

with open(dev_queries_file) as fIn:
    for line in fIn:
        qid, query = line.strip().split("\t")
        if qid in needed_qids:
            dev_queries[qid] = query

with gzip.open('data/collection-rnd.tsv.gz', 'rt') as fIn:
    for line in fIn:
        pid, passage = line.strip().split("\t")
        if pid in needed_pids or dev_corpus_max_size <= 0 or len(dev_corpus) <= dev_corpus_max_size:
            dev_corpus[pid] = passage

dev_corpus_pids = list(dev_corpus.keys())
dev_corpus = [dev_corpus[pid] for pid in dev_corpus_pids]

########### Eval functions

def compute_passage_emb(passages):
    sparse_embeddings = []
    bert_input_emb = model.bert_model.embeddings.word_embeddings(torch.tensor(list(range(0, len(model.tokenizer))), device=device))
    sparse_vec_size = 2000

    # Set Special tokens [CLS] [MASK] etc. to zero
    for special_id in model.tokenizer.all_special_ids:
        bert_input_emb[special_id] = 0 * bert_input_emb[special_id]

    with torch.no_grad():
        tokens = model.tokenizer(passages, padding=True, truncation=True, return_tensors='pt', max_length=500).to(device)
        passage_embeddings = model.bert_model(**tokens).last_hidden_state
        for passage_emb in passage_embeddings:
            scores = torch.matmul(bert_input_emb, passage_emb.transpose(0, 1))
            max_scores = torch.max(scores, dim=-1).values
            relu_scores = torch.relu(max_scores) #Eq. 5
            final_scores = torch.log(relu_scores + 1)  # Eq. 6, final score

            top_results = torch.topk(final_scores, k=sparse_vec_size, sorted=True)
            passage_emb = defaultdict(float)
            for score, idx in zip(top_results[0].cpu().tolist(), top_results[1].cpu().tolist()):
                if score > 0:
                    passage_emb[idx] = score
                else:
                    break

            sparse_embeddings.append(passage_emb)

    return sparse_embeddings

def evaluate_msmarco():
    passage_embs_sorted = []
    batch_size = 32

    length_sorted_idx = np.argsort([-len(pas) for pas in dev_corpus])
    dev_corpus_sorted = [dev_corpus[idx] for idx in length_sorted_idx]

    for start_idx in tqdm.trange(0, len(dev_corpus_sorted), batch_size, desc='encode corpus'):
        passage_embs_sorted.extend(compute_passage_emb(dev_corpus_sorted[start_idx:start_idx + batch_size]))

    passage_embs = [passage_embs_sorted[idx] for idx in np.argsort(length_sorted_idx)]

    logging.info("Create sparse matrix")
    row = []
    col = []
    values = []
    for pid, emb in enumerate(passage_embs):
        for tid, score in emb.items():
            row.append(tid)
            col.append(pid)
            values.append(score)

    sparse = csr_matrix((values, (row, col)), shape=(len(model.tokenizer), len(passage_embs)), dtype=np.float)
    logging.info("Scores: {}".format(sparse.shape))

    mrr = []
    k = 10
    for qid, question in tqdm.tqdm(dev_queries.items(), desc="score"):
        token_ids = model.tokenizer(question, add_special_tokens=False)['input_ids']

        # Get the candidate passages
        scores = np.asarray(sparse[token_ids, :].sum(axis=0)).squeeze(0)
        top_k_ind = np.argpartition(scores, -k)[-k:]
        hits = sorted([(dev_corpus_pids[pid], scores[pid]) for pid in top_k_ind], key=lambda x: x[1], reverse=True)

        mrr_score = 0
        for rank, hit in enumerate(hits[0:10]):
            pid = hit[0]
            if pid in dev_rel_docs[qid]:
                mrr_score = 1 / (rank + 1)
                break
        mrr.append(mrr_score)

    assert len(mrr) == len(dev_queries)
    mrr = np.mean(mrr)
    logging.info("MRR@10: {:.4f}".format(mrr))
    return mrr

best_score = 0 #evaluate_msmarco()

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

#### Read train file

with gzip.open('data/collection.tsv.gz', 'rt') as fIn:
    for line in fIn:
        pid, passage = line.strip().split("\t")
        corpus[pid] = passage

with open('data/queries.train.tsv', 'r') as fIn:
    for line in fIn:
        qid, query = line.strip().split("\t")
        train_queries[qid] = {'query': query,
                              'pos': set(),
                              'soft-pos': set(),
                              'neg': set()}

#Read qrels file for relevant positives per query
with open('data/qrels.train.tsv') as fIn:
    for line in fIn:
        qid, _, pid, _ = line.strip().split()
        train_queries[qid]['pos'].add(pid)

logging.info("Clean train queries")
deleted_queries = 0
for qid in list(train_queries.keys()):
    if len(train_queries[qid]['pos']) == 0:
        deleted_queries += 1
        del train_queries[qid]
        continue

logging.info("Deleted queries pos-empty: {}".format(deleted_queries))

for hard_neg_file in ['data/hard-negatives-all.jsonl.gz']: #'data/hard-negatives-ann-roberta.jsonl.gz']: #['data/hard-negatives-ann-msmarco-distilbert-base-v2.jsonl.gz', 'data/hard-negatives-ann.jsonl.gz', 'data/hard-negatives-ann-no_idnt.jsonl.gz', 'data/hard-negatives-all.jsonl.gz']:
    logging.info("Read hard negatives: "+hard_neg_file)
    with gzip.open(hard_neg_file, 'rt') as fIn:
        try:
            for line in fIn:
                try:
                    data = json.loads(line)
                except:
                    continue
                qid = data['qid']

                if qid in train_queries:
                    neg_added = 0
                    max_neg_added = 100

                    hits = sorted(data['hits'], key=lambda x: x['score'] if 'score' in x else x['bm25-score'], reverse=True)
                    for hit in hits:
                        pid = hit['corpus_id'] if 'corpus_id' in hit else hit['pid']

                        if pid in train_queries[qid]['pos']:    #Skip entries we have as positives
                            continue

                        if hit['bert-score'] < 0.1 and neg_added < max_neg_added:
                            train_queries[qid]['neg'].add(pid)
                            neg_added += 1
                        elif hit['bert-score'] > 0.9:
                            train_queries[qid]['soft-pos'].add(pid)
        except:
            pass

logging.info("Clean train queries with empty neg set")
deleted_queries = 0
for qid in list(train_queries.keys()):
    if len(train_queries[qid]['neg']) == 0:
        deleted_queries += 1
        del train_queries[qid]
        continue

logging.info("Deleted queries neg empty: {}".format(deleted_queries))

train_queries = list(train_queries.values())
for idx in range(len(train_queries)):
    train_queries[idx]['pos'] = list(train_queries[idx]['pos'])
    train_queries[idx]['neg'] = list(train_queries[idx]['neg'])
    train_queries[idx]['soft-pos'] = list(train_queries[idx]['soft-pos'])

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

####
# Prepare optimizers
param_optimizer = list(model.named_parameters())
no_decay = ['bias', 'LayerNorm.bias', 'LayerNorm.weight']
optimizer_grouped_parameters = [
    {'params': [p for n, p in param_optimizer if not any(nd in n for nd in no_decay)], 'weight_decay': 0.01},
    {'params': [p for n, p in param_optimizer if any(nd in n for nd in no_decay)], 'weight_decay': 0.0}
]

grad_acc_steps, lr = 1, 2e-5
#grad_acc_steps, lr = 16, 2e-5

num_epochs = 1
optimizer = transformers.AdamW(model.parameters(), lr=lr, eps=1e-6)   #optimizer_grouped_parameters
t_total = math.ceil(len(train_queries)/batch_size*num_epochs)
num_warmup_steps = int(t_total/grad_acc_steps * 0.1)    #10% for warm up
scheduler = transformers.get_linear_schedule_with_warmup(optimizer, num_warmup_steps=num_warmup_steps, num_training_steps=t_total)
loss_fct = torch.nn.CrossEntropyLoss()
max_grad_norm = 1

for epoch in tqdm.trange(num_epochs, desc='Epochs'):
    random.shuffle(train_queries)
    idx = 0
    for start_idx in tqdm.trange(0, len(train_queries), batch_size):
        idx += 1
        if (idx) % 5000 == 0:
            score = evaluate_msmarco()
            if score > best_score:
                best_score = score
                model.bert_model.save_pretrained(model_save_path)
                logging.info(f"Save to {model_save_path}")

        batch = train_queries[start_idx:start_idx+batch_size]
        queries = [b['query'] for b in batch]

        #First the positives
        passages = [corpus[random.choice(b['pos'])] for b in batch]

        #Then the negatives
        for b in batch:
            for pid in random.sample(b['neg'], k=min(len(b['neg']), num_negatives)):
                passages.append(corpus[pid])

        label = torch.tensor(list(range(len(batch))), device=device)

        ##FP16
        with autocast():
            final_scores = model(queries, passages)
            final_scores = 10*final_scores
            loss_value = loss_fct(final_scores, label) / grad_acc_steps

        scaler.scale(loss_value).backward()
        if (idx + 1) % grad_acc_steps == 0:
            scaler.unscale_(optimizer)
            torch.nn.utils.clip_grad_norm_(model.parameters(), max_grad_norm)
            scaler.step(optimizer)
            scaler.update()
            model.zero_grad()
            scheduler.step()

        """
        #Normal FP32 with grad acc
        final_scores = model(query, passages)
        #Compute loss
        loss_value = loss_fct(final_scores, label)
        if grad_acc_steps > 1:
            loss_value /= grad_acc_steps
        loss_value.backward()

        if (idx+1) % grad_acc_steps == 0:
            torch.nn.utils.clip_grad_norm_(model.parameters(), max_grad_norm)
            optimizer.step()
            model.zero_grad()
            scheduler.step()
        """

logging.info("Final eval:")
evaluate_msmarco()

max_stars_repo_name: nreimers/beir-sparta

max_issues_repo_name: tonellotto/beir-sparta

max_forks_repo_name: tonellotto/beir-sparta

reason: import numpy,from scipy


text:
#!/usr/bin/env python

# call as: python cov2ensemble.py
# include plot code: plot_cov2ensemble.py

# =======================================
# Version 0.16
# 28 July, 2019
# https://patternizer.github.io/
# michael.taylor AT reading DOT ac DOT uk
# =======================================

import os
import os.path
from os import fsync, remove
import sys
import glob
import optparse
from  optparse import OptionParser
from netCDF4 import Dataset
import numpy as np
from numpy import array_equal, savetxt, loadtxt, frombuffer, save as np_save, load as np_load, savez_compressed, array
from pandas import Series, DataFrame, Panel
import xarray
import scipy.linalg as la
from scipy.special import erfinv
import sklearn
from sklearn.decomposition import PCA
from sklearn.preprocessing import StandardScaler
from sklearn import datasets
from sklearn.datasets import make_classification
import seaborn as sns; sns.set(style='darkgrid')
import matplotlib.pyplot as plt; plt.close('all')
import matplotlib.dates as mdates
import matplotlib.colors as mcolors
import matplotlib.ticker as mticker
from matplotlib.collections import PolyCollection
from mpl_toolkits import mplot3d
from mpl_toolkits.mplot3d import Axes3D
from mpl_toolkits.mplot3d import proj3d

#----------------------------------------------------------------------------------
import ensemble_func as ensemble_func # functions for ensemble generation
import convert_func as convert_func # functions for L<-->BT conversion & HAR + CCI Meas Eqn
#----------------------------------------------------------------------------------

# =======================================
# INCLUDE PLOT CODE:
exec(open('plot_ensemble.py').read())
exec(open('plot_cov2ensemble.py').read())
# =======================================

if __name__ == "__main__":

    #------------------------------------------------------------------------------
    # parser = OptionParser("usage: %prog ch npop nens har_file")
    # (options, args) = parser.parse_args()
    #------------------------------------------------------------------------------
    # ch = args[0]
    # npop = args[1]
    # nens_file = args[2]
    # har_file = args[3]
    #------------------------------------------------------------------------------

    #------------------------------------------------------------------------------
    # RUN PARAMETERS
    #------------------------------------------------------------------------------
    ch = 37
    npop = 1000000
    nens = 10

    # /gws/nopw/j04/fiduceo/Users/jmittaz/FCDR/Mike/FCDR_AVHRR/GBCS/dat_cci/
    har_file = 'FIDUCEO_Harmonisation_Data_' + str(ch) + '.nc'

    n = int(nens/2) # --> (2*n) = nens = number of ensemble members
    c = 0.99 # variance_threshold

    FLAG_load_mc = 0          # 0-->run, 1-->load
    FLAG_mns = 1              # 0-->CRS, 1-->MNS
    FLAG_load_draws = 1
    FLAG_load_ensemble = 1
    FLAG_write_netcdf = 0
    FLAG_plot = 1

    # software_tag = '3e8c463' # job dir=job_avhxx_v6_EIV_10x_11 (old runs)
    # software_tag = '4d111a1' # job_dir=job_avhxx_v6_EIV_10x_11 (new runs)
    software_tag = 'v0.3Bet' # job_dir=job_avhxx_v6_EIV_10x_11 (new runs)
    #--------------------------------------------------------------------------

    senspr = [b'N12',b'N14',b'N15',b'N16',b'N17',b'N18',b'N19',b'MTA',b'MTB']

    if FLAG_mns:
        plotstem = '_'+str(ch)+'_'+'MNS'+'_'+software_tag+'.png'
        mc_file = 'MC_Harmonisation_MNS.nc'
    else:
        plotstem = '_'+str(ch)+'_'+'CRS'+'_'+software_tag+'.png'
        mc_file = 'MC_Harmonisation_CRS.nc'

    #
    # Load harmonisation parameters
    #

    har = xarray.open_dataset(har_file, decode_cf=True)
    a_ave = np.array(har.parameter)
    a_cov = np.array(har.parameter_covariance_matrix)
    a_u = np.array(har.parameter_uncertainty)

    if FLAG_load_mc:

        # Load ensemble

        ncout = Dataset(mc_file,'r')
        if ch == 37:
            da = ncout.variables['delta_params3'][:]
        elif ch == 11:
            da = ncout.variables['delta_params4'][:]
        else:
            da = ncout.variables['delta_params5'][:]
        ncout.close()

    if FLAG_mns:

        filestr_draws = "draws_" + str(ch) + "_" + str(npop) + ".npy"
        filestr_ensemble = "ensemble_" + str(ch) + "_" + str(npop) + ".npy"

        # Load / Generate draws

        if FLAG_load_draws:
            draws = np_load(filestr_draws)
        else:
            draws = ensemble_func.cov2draws(har,npop)
            np_save(filestr_draws, draws, allow_pickle=False)

        # Load / Generate ensemble

        if FLAG_load_ensemble:
            ens = np_load(filestr_ensemble, allow_pickle=True).item()
        else:
            ens = ensemble_func.draws2ensemble(draws,nens)
            np_save(filestr_ensemble, ens, allow_pickle=True)

        ensemble = ens['ensemble']
        ensemble_idx = ens['ensemble_idx']
        decile = ens['decile']
        Z = ens['Z']
        Z_norm = ens['Z_norm']
        da = ensemble - a_ave # [nens,npar]

    else:

        ev = ensemble_func.cov2ev(a_cov,c)
        da = ensemble_func.ev2da(ev,n) # sum of nPC (99% variance)

        # Projection onto first 2 PCs:

        da_pc12 = ensemble_func.da2pc12(ev,n)
        plot_ensemble_deltas_pc12(da_pc12,a_u)

    # Ensemble closure

    da_cov = np.cov(da.T) # [npar,npar]
    da_u = np.sqrt(np.diag(da_cov)) # [npar]
    norm_u = np.linalg.norm(a_u - da_u)
    norm_cov = np.linalg.norm(a_cov - da_cov)
    print('|a_u - da_u|=',norm_u)
    print('|a_cov - da_cov|=',norm_cov)

    # Open ensemble file and overwrite channel deltas and uuid:

    if FLAG_write_netcdf:

        ncout = Dataset(mc_file,'r+')
        if ch == 37:
            ncout.variables['delta_params3'][:] = da
            ncout.HARM_UUID3 = har.uuid
        elif ch == 11:
            ncout.variables['delta_params4'][:] = da
            ncout.HARM_UUID4 = har.uuid
        else:
            ncout.variables['delta_params5'][:] = da
            ncout.HARM_UUID5 = har.uuid
        ncout.close()

    if FLAG_plot:

        if FLAG_mns:
            plot_ensemble_closure(da,draws,har)
            plot_ensemble_decile_selection(Z_norm, ensemble_idx, nens)
            plot_ensemble_decile_distribution(Z, decile, npop, nens)
        else:
            plot_eigenspectrum(ev)
        plot_ensemble_deltas(da)
        plot_ensemble_deltas_an(da,a_u)
        plot_ensemble_deltas_normalised(da,a_u)
        plot_crs()

        plot_bestcase_covariance(har)
        plot_bestcase_parameters(har, sensor)
        plot_ensemble_deltas(har, ensemble, sensor, nens)
        plot_population_cdf(har, draws, sensor, nens, npop)
        plot_population_coefficients(har, draws, sensor, npop)
        plot_population_histograms(har, draws, sensor, nens)

    #------------------------------------------------
    print('** end')

reason: import numpy,from numpy,import scipy,from scipy