gԚdZddlmZddlmZmZddlmZddlZ ddl Z dgZ dd Z dd Z ddZ d dZ d!dZ d"dZd#dZdS)$z'Draw a rustworkx graph with matplotlib.)Iterable)islicecycle)NumberNmpl_drawTFc  ddlm}n"#t$r}td|d}~wwxYw||}n|}|d|1|jr|}n|d}t|f||||d|| | | r||SdS)aDraw a graph with Matplotlib. .. note:: Matplotlib is an optional dependency and will not be installed with rustworkx by default. If you intend to use this function make sure that you install matplotlib with either ``pip install matplotlib`` or ``pip install 'rustworkx[mpl]'`` :param graph: A rustworkx graph, either a :class:`~rustworkx.PyGraph` or a :class:`~rustworkx.PyDiGraph`. :param dict pos: An optional dictionary (or a :class:`~rustworkx.Pos2DMapping` object) with nodes as keys and positions as values. If not specified a spring layout positioning will be computed. See `layout_functions` for functions that compute node positions. :param matplotlib.Axes ax: An optional Matplotlib Axes object to draw the graph in. :param bool arrows: For :class:`~rustworkx.PyDiGraph` objects if ``True`` draw arrowheads. (defaults to ``True``) Note, that the Arrows will be the same color as edges. :param str arrowstyle: An optional string for directed graphs to choose the style of the arrowsheads. See :class:`matplotlib.patches.ArrowStyle` for more options. By default the value is set to ``'-\|>'``. :param int arrow_size: For directed graphs, choose the size of the arrow head's length and width. See :class:`matplotlib.patches.FancyArrowPatch` attribute and constructor kwarg ``mutation_scale`` for more info. Defaults to 10. :param bool with_labels: Set to ``True`` to draw labels on the nodes. Edge labels will only be drawn if the ``edge_labels`` parameter is set to a function. Defaults to ``False``. :param list node_list: An optional list of node indices in the graph to draw. If not specified all nodes will be drawn. :param list edge_list: An option list of edges in the graph to draw. If not specified all edges will be drawn :param int|list node_size: Optional size of nodes. If an array is specified it must be the same length as node_list. Defaults to 300 :param node_color: Optional node color. Can be a single color or a sequence of colors with the same length as node_list. Color can be string or rgb (or rgba) tuple of floats from 0-1. If numeric values are specified they will be mapped to colors using the ``cmap`` and ``vmin``,``vmax`` parameters. See :func:`matplotlib.scatter` for more details. Defaults to ``'#1f78b4'``) :param str node_shape: The optional shape node. The specification is the same as the :func:`matplotlib.pyplot.scatter` function's ``marker`` kwarg, valid options are one of ``['s', 'o', '^', '>', 'v', '<', 'd', 'p', 'h', '8']``. Defaults to ``'o'`` :param float alpha: Optional value for node and edge transparency :param matplotlib.colors.Colormap cmap: An optional Matplotlib colormap object for mapping intensities of nodes :param float vmin: Optional minimum value for node colormap scaling :param float vmax: Optional minimum value for node colormap scaling :param float|sequence linewidths: An optional line width for symbol borders. If a sequence is specified it must be the same length as node_list. Defaults to 1.0 :param float|sequence width: An optional width to use for edges. Can either be a float or sequence of floats. If a sequence is specified it must be the same length as node_list. Defaults to 1.0 :param str|sequence edge_color: color or array of colors (default='k') Edge color. Can be a single color or a sequence of colors with the same length as edge_list. Color can be string or rgb (or rgba) tuple of floats from 0-1. If numeric values are specified they will be mapped to colors using the ``edge_cmap`` and ``edge_vmin``, ``edge_vmax`` parameters. :param matplotlib.colors.Colormap edge_cmap: An optional Matplotlib colormap for mapping intensities of edges. :param float edge_vmin: Optional minimum value for edge colormap scaling :param float edge_vmax: Optional maximum value for node colormap scaling :param str style: An optional string to specify the edge line style. For example, ``'-'``, ``'--'``, ``'-.'``, ``':'`` or words like ``'solid'`` or ``'dashed'``. See the :class:`matplotlib.patches.FancyArrowPatch` attribute and kwarg ``linestyle`` for more details. Defaults to ``'solid'``. :param func labels: An optional callback function that will be passed a node payload and return a string label for the node. For example:: labels=str could be used to just return a string cast of the node's data payload. Or something like:: labels=lambda node: node['label'] could be used if the node payloads are dictionaries. :param func edge_labels: An optional callback function that will be passed an edge payload and return a string label for the edge. For example:: edge_labels=str could be used to just return a string cast of the edge's data payload. Or something like:: edge_labels=lambda edge: edge['label'] could be used if the edge payloads are dictionaries. If this is set edge labels will be drawn in the visualization. :param int font_size: An optional fontsize to use for text labels, By default a value of 12 is used for nodes and 10 for edges. :param str font_color: An optional font color for strings. By default ``'k'`` (ie black) is set. :param str font_weight: An optional string used to specify the font weight. By default a value of ``'normal'`` is used. :param str font_family: An optional font family to use for strings. By default ``'sans-serif'`` is used. :param str label: An optional string label to use for the graph legend. :param str connectionstyle: An optional value used to create a curved arc of rounding radius rad. For example, ``connectionstyle='arc3,rad=0.2'``. See :class:`matplotlib.patches.ConnectionStyle` and :class:`matplotlib.patches.FancyArrowPatch` for more info. By default this is set to ``"arc3"``. :returns: A matplotlib figure for the visualization if not running with an interactive backend (like in jupyter) or if ``ax`` is not set. :rtype: matplotlib.figure.Figure For Example: .. jupyter-execute:: import matplotlib.pyplot as plt import rustworkx as rx from rustworkx.visualization import mpl_draw G = rx.generators.directed_path_graph(25) mpl_draw(G) plt.draw() rNmatplotlib needs to be installed prior to running rustworkx.visualization.mpl_draw(). You can install matplotlib with: 'pip install matplotlib'w)rrr )posaxarrows with_labels) matplotlib.pyplotpyplot ImportErrorgcf get_figure set_facecoloraxesgcaadd_axes draw_graph set_axis_offdraw_if_interactive isinteractive) graphr r rrkwdspltecfs X/granary/edev/quantum/lib/python3.11/site-packages/rustworkx/visualization/matplotlib.pyrrBs&H'''''''  9      z WWYY ]]__S z 7 +BB\**BuU#"VUUPTUUUOO     "* #-* (#(c P  ddlm}n"#t$r}td|d}~wwxYwhdhdhdhdzzztfd|Dr3d fd |D}t d ||d d r$ fd D|d <|dd r# fdD|d<fd| D}fd| D} t| dtr| d=fd| D} fd| D} |tj}t|fi|t!|fd|i| |rt#|fi|  rt%|fi| |dS)aDraw the graph using Matplotlib. Draw the graph with Matplotlib with options for node positions, labeling, titles, and many other drawing features. See draw() for simple drawing without labels or axes. Parameters ---------- graph: A rustworkx :class:`~rustworkx.PyDiGraph` or :class:`~rustworkx.PyGraph` pos : dictionary, optional A dictionary with nodes as keys and positions as values. If not specified a spring layout positioning will be computed. See :mod:`rustworkx.drawing.layout` for functions that compute node positions. Notes ----- For directed graphs, arrows are drawn at the head end. Arrows can be turned off with keyword arrows=False. rNr > r cmapvmaxvminalphalabel node_list node_size edgecolors linewidths node_color node_shape>r r(r)stylewidth edge_cmap edge_list edge_vmax edge_vminr*r+ arrow_size arrowstyle edge_colorr/connectionstylemin_source_marginmin_target_margin> r bboxr(labels font_size font_color font_family font_weightverticalalignmenthorizontalalignment> r r<r(rotater>r? edge_labelsr@rArBrCcg|]}|vSrG.0k valid_kwdss r" zdraw_graph..Cs . . .AAZ  . . .z, cg|]}|v| SrGrGrHs r"rLzdraw_graph..Ds#!I!I!IQj5H5H!5H5H5HrMzReceived invalid argument(s): r=c4i|]}||SrGrG)rIxrlabel_fns r" zdraw_graph..Is)NNNA!XXeAh//NNNrMrEcPi|]"}|d|df|d#S)rr rG)rIrP edge_label_fns r"rRzdraw_graph..LsB   23QqT1Q4L--!--   rMc$i|] \}}|v || SrGrG)rIrJvvalid_node_kwdss r"rRzdraw_graph..P)GGG$!Q!2F2FA2F2F2FrMc$i|] \}}|v || SrGrG)rIrJrWvalid_edge_kwdss r"rRzdraw_graph..QrYrMr(c$i|] \}}|v || SrGrG)rIrJrWvalid_label_kwdss r"rRzdraw_graph..Ts*III41a18H3H3H!Q3H3H3HrMc$i|] \}}|v || SrGrG)rIrJrWvalid_edge_label_kwdss r"rRzdraw_graph..Us*SSS1=R8R8Rq!8R8R8RrMr)rrranyjoin ValueErrorpop node_indicesweighted_edge_listitems isinstancegetlist rustworkx spring_layout draw_nodes draw_edges draw_labelsdraw_edge_labelsr)rr rrrrr invalid_args node_kwds edge_kwds label_kwdsedge_label_kwdsrUrQr[r_rKr]rXs` @@@@@@@r"rrs 2'''''''  9       OO*      !?25EEH]]J . . . . . . .//Jyy!I!I!I!IT!I!I!IJJ H,HHIIIxx$''HONNNNN9K9K9M9MNNNXHH]D11M     7<7O7O7Q7Q   ]HGGG$**,,GGGIGGGG$**,,GGGI)--(($// g IIII4::<<IIIJSSSS SSSO {%e,,uc''Y'''uc66&6I666.E3--*---877777s 0+0,#1f78b4oc ddl}ddl}ddlm}n"#t$r}t d|d}~wwxYw| |} ||}t|dkr|j dS tj fd|D}n&#t$r}td|d|d}~wwxYwt|trt!|||||| }d}| |dddf|dddf|||||| || | |  }| d d d d d d |d |S) aDraw the nodes of the graph. This draws only the nodes of the graph. :param graph: A rustworkx graph, either a :class:`~rustworkx.PyGraph` or a :class:`~rustworkx.PyDiGraph`. :param dict pos: A dictionary with nodes as keys and positions as values. Positions should be sequences of length 2. :param Axes ax: An optional Matplotlib Axes object, if specified it will draw the graph in the specified Matplotlib axes. :param list node_list: If specified only draw the specified node indices. If not specified all nodes in the graph will be drawn. :param float|array node_size: Size of nodes. If an array it must be the same length as node_list. Defaults to 300 node_color : color or array of colors (default='#1f78b4') Node color. Can be a single color or a sequence of colors with the same length as node_list. Color can be string or rgb (or rgba) tuple of floats from 0-1. If numeric values are specified they will be mapped to colors using the cmap and vmin,vmax parameters. See matplotlib.scatter for more details. node_shape : string (default='o') The shape of the node. Specification is as matplotlib.scatter marker, one of 'so^>v.s333AQ333rMzNode z has no position.r ) scmarkerr%r'r&r(r-r,r)bothFaxiswhichbottomleft labelbottom labelleftrT) matplotlibmatplotlib.collectionsrrrrrdlen collectionsPathCollectionnpasarrayKeyError IndexErrorrgr apply_alphascatter tick_params set_zorder)rr r*r+r.r/r(r%r'r&r r-r,r)mplrrr xynode_collections ` r"rlrlcsR     %%%%'''''''  9     z WWYY&&((  9~~--d333> Z3333333 4 4 >>>555566A=>%"" UItT4PP jj 111a4 111a4     !  ONN   q!!! s( 1,1 B// C9C  C?rJsolid arc3c 789:; ddl;ddl}ddl}ddl}ddlm}n"#t $r}t d|d}~wwxYw|!t|tj r| rd}nd} | || }t|dkrgS|!t|}|d}t}|D];}d|t!||dt!||df<<|}t%j|rt|t|krt%jd|Dr t ;jjsJn| | t1|} | t3|} ;j| | 88 fd |D}d }g}|};j|7t%j|:;jj||}tA|D]\}}|dd|dd}!} |dd|dd}#}"d}$d}%t%j|rb||dd \}&}'||!|&}(||!|'})||(|}$||)|}%n|||x}$}%|$|kr|}$|%|kr|}%t|t|kr ||}*n4t|dkr |d}*n||t|z}*t%j|rBt|t|kr ||}+n||t|z}+n|}+t!tE||vrd },nd},;j#| |!f|"|#f||$|%||*|+|d|,z|d }-|$|- %|- t%j&t!|}t%j't%j(|dddddf}.t%j)t%j(|dddddf}/t%j't%j(|dddddf}0t%j)t%j(|dddddf}1|/|.z }2|1|0z 9 79:;fd}3d|2zd9z}5}4|.|4z |0|5z f|/|4z|1|5zff}6 *|6 + ,dddddd|S)a! Draw the edges of the graph. This draws only the edges of the graph. Parameters ---------- graph: A rustworkx graph pos : dictionary A dictionary with nodes as keys and positions as values. Positions should be sequences of length 2. edge_list : collection of edge tuples (default=graph.edge_list()) Draw only specified edges width : float or array of floats (default=1.0) Line width of edges edge_color : color or array of colors (default='k') Edge color. Can be a single color or a sequence of colors with the same length as edge_list. Color can be string or rgb (or rgba) tuple of floats from 0-1. If numeric values are specified they will be mapped to colors using the edge_cmap and edge_vmin,edge_vmax parameters. style : string (default=solid line) Edge line style e.g.: '-', '--', '-.', ':' or words like 'solid' or 'dashed'. (See `matplotlib.patches.FancyArrowPatch`: `linestyle`) alpha : float or None (default=None) The edge transparency edge_cmap : Matplotlib colormap, optional Colormap for mapping intensities of edges edge_vmin,edge_vmax : floats, optional Minimum and maximum for edge colormap scaling ax : Matplotlib Axes object, optional Draw the graph in the specified Matplotlib axes. arrows : bool, optional (default=True) For directed graphs, if True set default to drawing arrowheads. Otherwise set default to no arrowheads. Ignored if `arrowstyle` is set. Note: Arrows will be the same color as edges. arrowstyle : str (default='-\|>' if directed else '-') For directed graphs and `arrows==True` defaults to '-\|>', otherwise defaults to '-'. See `matplotlib.patches.ArrowStyle` for more options. arrow_size : int (default=10) For directed graphs, choose the size of the arrow head's length and width. See `matplotlib.patches.FancyArrowPatch` for attribute ``mutation_scale`` for more info. node_size : scalar or array (default=300) Size of nodes. Though the nodes are not drawn with this function, the node size is used in determining edge positioning. node_list : list, optional (default=graph.node_indices()) This provides the node order for the `node_size` array (if it is an array). node_shape : string (default='o') The marker used for nodes, used in determining edge positioning. Specification is as a `matplotlib.markers` marker, e.g. one of 'so^>v-rJr c8g|]}t|tSrG)rgr)rIr{s r"rLzdraw_edges..s">>>aJq&))>>>rM)r'r&c8g|]}|SrGrG)rIr color_normalr2s r"rLzdraw_edges..s+EEEQii Q00EEErMcl|dvrtjd|zdz Stj|dz S)Nzs^>v.to_marker_edges; X  71{?++a/ /7;''!+ +rMrT?z, rad = ) r7shrinkAshrinkBmutation_scalecolor linewidthr9 linestylezorderc ptj||kr dkrd zn } j|}d|z}|dz}|tjd|gz|tj||gz|tj|dgz||tj| dgz|tj| |gz|tjd|gzg}j j|gd} n  ||g|Ri|} | S)Nrg{Gzt?g??)r rrrrr)rall transDatainverted transformrpathPath)posAposBargskwargs selfloop_htdata_locv_shifth_shiftrretr base_connectionstyleh max_nodesizers r"_connectionstylez$draw_edges.._connectionstylesX 6$$,   D34q&&%,..aK|,,..88>>HK'GmG2:q'l3332:w&89992:wl3332:xm4442:x&9:::2:q'l333 D(-- 6 6t < <>S>S>STTCC'&tTCDCCCFCCC rMg?r}Fr~)-rmatplotlib.colorsmatplotlib.patchesmatplotlib.pathrrrrgrj PyDiGraphrr3rrirddicttuplekeysriterablercolorsColormapget_cmapminmax NormalizepatchesConnectionStylearraycolorConverter to_rgba_array enumerateindexreversedFancyArrowPatchappend add_patchraminravelamaxupdate_datalimautoscale_viewr)B uS1Y''s1Q4y)9)9:;;!!##H JF __H - - F>>:>>> ? ?.  i)<== = = = = I  JI  JIz+++KK EEEEE*EEE ,,,N;66GG8I&&**,,L:,:::uMMLX&&884aT!WQZBaT!WQZB  ;y ! ! R&q\"1"-NFF()@)@A ()@)@A *N+;ZHHM*N+;ZHHMM,:N9j,Q,Q QMM , , ,-M , , ,-M |  H - -&q/KK   ! # #&q/KK&q3|+<+<'<=K ;u   5zzS]]**"1X "1s5zz>2 J $ H , ,CCC ++ H H!!!) +.>.>.>>,    &&& Uz%//**H'"(8AAAqqq!G#45566K 728HQQQ1W-.. / /D 728HQQQ1W-.. / /D 728HQQQ1W-.. / /D {A t AD4!8$DT!4$;/$+td{1KKGgNN    s ?:? sans-serifnormalcenterc  ddlm} n"#t$r}td|d}~wwxYw| | } |d|D}i}|D]`\}}||\}}t |tst|}| ||||||||| | | j ||  }|||<a| dddddd|S) aDraw node labels on the graph. Parameters ---------- graph: A rustworkx graph pos : dictionary A dictionary with nodes as keys and positions as values. Positions should be sequences of length 2. labels : dictionary (default={n: n for n in graph}) Node labels in a dictionary of text labels keyed by node. Node-keys in labels should appear as keys in `pos`. If needed use: `{n:lab for n,lab in labels.items() if n in pos}` font_size : int (default=12) Font size for text labels font_color : string (default='k' black) Font color string font_weight : string (default='normal') Font weight font_family : string (default='sans-serif') Font family alpha : float or None (default=None) The text transparency bbox : Matplotlib bbox, (default is Matplotlib's ax.text default) Specify text box properties (e.g. shape, color etc.) for node labels. horizontalalignment : string (default='center') Horizontal alignment {'center', 'right', 'left'} verticalalignment : string (default='center') Vertical alignment {'center', 'top', 'bottom', 'baseline', 'center_baseline'} ax : Matplotlib Axes object, optional Draw the graph in the specified Matplotlib axes. clip_on : bool (default=True) Turn on clipping of node labels at axis boundaries Returns ------- dict `dict` of labels keyed on the nodes rNr ci|]}||SrGrG)rIns r"rRzdraw_labels..rs5551!Q555rM) sizerfamilyweightr(rCrBrr<clip_onr}Fr~) rrrrrdrfrgstrtextrr)rr r=r>r?r@rAr(r<rCrBr rrr text_itemsrr)rPyts r"rnrn#s]D'''''''  9     z WWYY ~55 2 2 4 4555JLLNN5QA%%% JJE GG  3/l    1 NN    r#rcP ddlm}n"#t$r}td|d}~wwxYw| |} |d|D}n|}i}|D]\\}}}||\}}||\}}||z|d|z zz||z|d|z zz}}||f|vr.tj||z }||kr |d|zz}n|d|zz }| rtj ||z ||z dtj zz dz}|d kr|d z}|d kr|d z }tj ||f}| j tj |f|d d}nd }| tddd} t!|t"st#|}| ||||||||| | || j | d|} | |||f<| dddddd|S)aDraw edge labels. Parameters ---------- graph: A rustworkx graph pos : dictionary A dictionary with nodes as keys and positions as values. Positions should be sequences of length 2. edge_labels : dictionary (default={}) Edge labels in a dictionary of labels keyed by edge two-tuple. Only labels for the keys in the dictionary are drawn. label_pos : float (default=0.5) Position of edge label along edge (0=head, 0.5=center, 1=tail) font_size : int (default=10) Font size for text labels font_color : string (default='k' black) Font color string font_weight : string (default='normal') Font weight font_family : string (default='sans-serif') Font family alpha : float or None (default=None) The text transparency bbox : Matplotlib bbox, optional Specify text box properties (e.g. shape, color etc.) for edge labels. Default is {boxstyle='round', ec=(1.0, 1.0, 1.0), fc=(1.0, 1.0, 1.0)}. horizontalalignment : string (default='center') Horizontal alignment {'center', 'right', 'left'} verticalalignment : string (default='center') Vertical alignment {'center', 'top', 'bottom', 'baseline', 'center_baseline'} ax : Matplotlib Axes object, optional Draw the graph in the specified Matplotlib axes. rotate : bool (default=True) Rotate edge labels to lie parallel to edges clip_on : bool (default=True) Turn on clipping of edge labels at axis boundaries Returns ------- dict `dict` of labels keyed by edge rNr c i|] \}}}||f| SrGrG)rIurWds r"rRz$draw_edge_labels..s$FFF1a1a&!FFFrMrrg@ihZi)r rTrround)rrr)boxstyleecfcr ) rrrrr(rCrBrotationrr<rrr}Fr~)rrrrrerfrrabsarctan2pirrtransform_anglesreshaperrgrrr)!rr rE label_posr>r?r@rAr(r<rCrBr rDrrr r=rn1n2r)rrrrrPrdyangler trans_angler s! r"rorosT'''''''  9     z WWYYFF5+C+C+E+EFFFJ!<<>>2!2!R%r7Rr7R NR3?3 3 NR3?3 3  8v{{}} $ $RBBwwTBYTBY  JrBwR00C"%K@3FErzz s{{ 1a&!!B,77%8J8JBJJW]L^L^__`abKKK <_QQQD%%% JJE GG  3/ l   " ! B8NN    r#c ddl ddl}ddl}n"#t$r}td|d}~wwxYwt |t |krbt |dt rG j|}| ||| |} nd tj j j |g} n0#t$r#tj fd|D} YnwxYw t |t | ks| jt |krQtj| t |df} | d| dddf<| d | dddf<| d | ddd f<t%t't)|t | | ddd f<n#t*$r || ddd f<YnwxYw| S)a!Apply an alpha (or list of alphas) to the colors provided. Parameters ---------- colors : color string or array of floats (default='r') Color of element. Can be a single color format string, or a sequence of colors with the same length as node_list. If numeric values are specified they will be mapped to colors using the cmap and vmin,vmax parameters. See matplotlib.scatter for more details. alpha : float or array of floats Alpha values for elements. This can be a single alpha value, in which case it will be applied to all the elements of color. Otherwise, if it is an array, the elements of alpha will be applied to the colors in order (cycling through alpha multiple times if necessary). elem_list : array of rustworkx objects The list of elements which are being colored. These could be nodes, edges or labels. cmap : matplotlib colormap Color map for use if colors is a list of floats corresponding to points on a color mapping. vmin, vmax : float Minimum and maximum values for normalizing colors if a colormap is used Returns ------- rgba_colors : numpy ndarray Array containing RGBA format values for each of the node colours. rNr )r%cNg|]!}jj|"SrG)rrto_rgba)rIrrs r"rLzapply_alpha..ms-#a#a#aQVCJ$=$E$Ee$L$L#a#a#arMr)rrr )rr )rrTrT)rr matplotlib.cmrrrgrcmScalarMappableset_climr"rrrrrbrresizerirr TypeError) rr( elem_listr%r'r&rr mapper rgba_colorsrs @r"rr0sMJ           9     6{{c)nn$$F1Iv)F)F$&&D&11d###nnV,,  c(CJ$=$E$Ef$M$M#NOOKK c c c(#a#a#a#aZ`#a#a#abbKKK c # u::K(( ( (K,<I,N,N)K#i..!1DEEK!,T!2KA !,T!2KA !,T!2KA  e c+6F6F!G!GHH AAAqD ###" AAArE# s5  /*/52C((*DDCG!!G76G7)NNTF)NTF) NrurvrwNNNNNNNN)NrrJrNNrNNNNTNruNrwrrr) NrrJrrNNrrNT) NrrrJrrNNrrNTT)NNN)__doc__collections.abcr itertoolsrrnumbersrnumpyrrj__all__rrrlrmrnrorrGrMr"r4sd.-$$$$$$########  [[[[|@@@@L    J     +{{{{B     ppppl     WWWWtLLLLLLrM