__init__.py 22.7 KB
Newer Older
Tiago Peixoto's avatar
Tiago Peixoto committed
1
#! /usr/bin/env python
2
# -*- coding: utf-8 -*-
Tiago Peixoto's avatar
Tiago Peixoto committed
3
#
4
5
# graph_tool -- a general graph manipulation python module
#
Tiago Peixoto's avatar
Tiago Peixoto committed
6
# Copyright (C) 2007-2010 Tiago de Paula Peixoto <tiago@skewed.de>
Tiago Peixoto's avatar
Tiago Peixoto committed
7
8
9
10
11
12
13
14
15
16
17
18
19
20
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program 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.
#
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.

21
"""
22
``graph_tool.draw`` - Graph drawing
23
-----------------------------------
24
25
26
27
28
29
30
31
32
33
34
35
36

Summary
+++++++

.. autosummary::
   :nosignatures:

   graph_draw
   arf_layout
   random_layout

Contents
++++++++
37
38
"""

39
import sys, os, os.path, time, warnings, tempfile
40
41
42
from .. core import _degree, _prop, PropertyMap, _check_prop_vector,\
     _check_prop_scalar, _check_prop_writable, group_vector_property,\
     ungroup_vector_property
Tiago Peixoto's avatar
Tiago Peixoto committed
43
from .. decorators import _limit_args
44
import numpy.random
45
46
47
48
from numpy import *

from .. dl_import import dl_import
dl_import("import libgraph_tool_layout")
49
50
51
52
53
54

try:
    import gv
except ImportError:
    warnings.warn("error importing gv module... graph_draw() will not work.",
                  ImportWarning)
55
56
57
try:
    import matplotlib.cm
    import matplotlib.colors
58
    from pylab import imread
59
60
61
except ImportError:
    warnings.warn("error importing matplotlib module... " + \
                  "graph_draw() will not work.", ImportWarning)
Tiago Peixoto's avatar
Tiago Peixoto committed
62

63
64
__all__ = ["graph_draw", "arf_layout", "random_layout"]

Tiago Peixoto's avatar
Tiago Peixoto committed
65

66
67
def graph_draw(g, pos=None, size=(15, 15), pin=False, layout=None,
               maxiter=None, ratio="fill", overlap="prism", sep=None,
68
69
               splines=False, vsize=0.1, penwidth=1.0, elen=None, gprops={},
               vprops={}, eprops={}, vcolor=None, ecolor=None,
Tiago Peixoto's avatar
Tiago Peixoto committed
70
               vcmap=matplotlib.cm.jet, vnorm=True, ecmap=matplotlib.cm.jet,
71
72
               enorm=True, vorder=None, eorder=None, output="",
               output_format="auto", returngv=False, fork=False,
73
               return_bitmap=False, seed=0):
74
75
76
77
78
79
    r"""Draw a graph using graphviz.

    Parameters
    ----------
    g : Graph
        Graph to be used.
80
    pos : PropertyMap or tuple of PropertyMaps (optional, default: None)
81
82
83
84
85
        Vertex property maps containing the x and y coordinates of the vertices.
    size : tuple of scalars (optional, default: (15,15))
        Size (in centimeters) of the canvas.
    pin : bool (default: False)
        If True, the vertices are not moved from their initial position.
86
    layout : string (default: "neato" if g.num_vertices() <= 1000 else "sfdp")
87
        Layout engine to be used. Possible values are "neato", "fdp", "dot",
88
        "circo", "twopi" and "arf".
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
    maxiter : int (default: None)
        If specified, limits the maximum number of iterations.
    ratio : string or float (default: "fill")
        Sets the aspect ratio (drawing height/drawing width) for the
        drawing. Note that this is adjusted before the 'size' attribute
        constraints are enforced.

        If ratio is numeric, it is taken as the desired aspect ratio. Then, if
        the actual aspect ratio is less than the desired ratio, the drawing
        height is scaled up to achieve the desired ratio; if the actual ratio is
        greater than that desired ratio, the drawing width is scaled up.

        If ratio = "fill" and the size attribute is set, node positions are
        scaled, separately in both x and y, so that the final drawing exactly
        fills the specified size.

        If ratio = "compress" and the size attribute is set, dot attempts to
        compress the initial layout to fit in the given size. This achieves a
        tighter packing of nodes but reduces the balance and symmetry.
        This feature only works in dot.

        If ratio = "expand", the size attribute is set, and both the width and
        the height of the graph are less than the value in size, node positions
        are scaled uniformly until at least one dimension fits size exactly.
        Note that this is distinct from using size as the desired size, as here
        the drawing is expanded before edges are generated and all node and text
        sizes remain unchanged.

        If ratio = "auto", the page attribute is set and the graph cannot be
118
        drawn on a single page, then size is set to an "ideal" value. In
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
        particular, the size in a given dimension will be the smallest integral
        multiple of the page size in that dimension which is at least half the
        current size. The two dimensions are then scaled independently to the
        new size. This feature only works in dot.
    overlap : bool or string (default: "prism")
        Determines if and how node overlaps should be removed. Nodes are first
        enlarged using the sep attribute. If True, overlaps are retained. If
        the value is "scale", overlaps are removed by uniformly scaling in x and
        y. If the value is False, node overlaps are removed by a Voronoi-based
        technique. If the value is "scalexy", x and y are separately scaled to
        remove overlaps.

        If sfdp is available, one can set overlap to "prism" to use a proximity
        graph-based algorithm for overlap removal. This is the preferred
        technique, though "scale" and False can work well with small graphs.
        This technique starts with a small scaling up, controlled by the
        overlap_scaling attribute, which can remove a significant portion of the
        overlap. The prism option also accepts an optional non-negative integer
        suffix. This can be used to control the number of attempts made at
        overlap removal. By default, overlap="prism" is equivalent to
        overlap="prism1000". Setting overlap="prism0" causes only the scaling
        phase to be run.

        If the value is "compress", the layout will be scaled down as much as
        possible without introducing any overlaps, obviously assuming there are
        none to begin with.
    sep : float (default: None)
        Specifies margin to leave around nodes when removing node overlap. This
        guarantees a minimal non-zero distance between nodes.
    splines : bool (default: False)
        If True, the edges are drawn as splines and routed around the vertices.
150
151
152
153
    vsize : float, PropertyMap, or tuple (default: 0.1)
        Default vertex size (width and height). If a tuple is specified, the
        first value should be a property map, and the second is a scale factor.
    penwidth : float, PropertyMap or tuple (default: 1.0)
154
155
        Specifies the width of the pen, in points, used to draw lines and
        curves, including the boundaries of edges and clusters. It has no effect
Tiago Peixoto's avatar
Tiago Peixoto committed
156
157
        on text. If a tuple is specified, the first value should be a property
        map, and the second is a scale factor.
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
    elen : float or PropertyMap (default: None)
        Preferred edge length, in inches.
    gprops : dict (default: {})
        Additional graph properties, as a dictionary. The keys are the property
        names, and the values must be convertible to string.
    vprops : dict (default: {})
        Additional vertex properties, as a dictionary. The keys are the property
        names, and the values must be convertible to string, or vertex property
        maps, with values convertible to strings.
    eprops : dict (default: {})
        Additional edge properties, as a dictionary. The keys are the property
        names, and the values must be convertible to string, or edge property
        maps, with values convertible to strings.
    vcolor : string or PropertyMap (default: None)
        Drawing color for vertices. If the valued supplied is a property map,
        the values must be scalar types, whose color values are obtained from
        the 'vcmap' argument.
    ecolor : string or PropertyMap (default: None)
        Drawing color for edges. If the valued supplied is a property map,
        the values must be scalar types, whose color values are obtained from
        the 'ecmap' argument.
    vcmap : matplotlib.colors.Colormap (default: matplotlib.cm.jet)
        Vertex color map.
    vnorm : bool (default: True)
        Normalize vertex color values to the [0,1] range.
    ecmap : matplotlib.colors.Colormap (default: matplotlib.cm.jet)
        Edge color map.
    enorm : bool (default: True)
        Normalize edge color values to the [0,1] range.
187
188
189
190
191
192
    vorder : PropertyMap (default: None)
        Scalar vertex property map which specifies the order with which vertices
        are drawn.
    eorder : PropertyMap (default: None)
        Scalar edge property map which specifies the order with which edges
        are drawn.
193
194
195
196
197
198
    output : string (default: "")
        Output file name.
    output_format : string (default: "auto")
        Output file format. Possible values are "auto", "xlib", "ps", "svg",
        "svgz", "fig", "mif", "hpgl", "pcl", "png", "gif", "dia", "imap",
        "cmapx". If the value is "auto", the format is guessed from the 'output'
199
200
        parameter, or 'xlib' if it is empty. If the value is None, no output is
        produced.
201
202
203
    returngv : bool (default: False)
        Return the graph object used internally with the gv module.
    fork : bool (default: False)
204
        If True, the program is forked before drawing. This is used as a
205
206
207
        work-around for a bug in graphviz, where the exit() function is called,
        which would cause the calling program to end. This is always assumed
        'True', if output_format = 'xlib'.
208
209
210
    return_bitmap : bool (default: False)
        If True, a bitmap (:class:`~numpy.ndarray`) of the rendered graph is
        returned.
211
212
213

    Returns
    -------
214
215
    pos : PropertyMap
        Vector vertex property map with the x and y coordinates of the vertices.
216
217
218
219
220
221
    gv : gv.digraph or gv.graph (optional, only if returngv == True)
        Internally used graphviz graph.


    Notes
    -----
222
    This function is a wrapper for the [graphviz] python
223
224
225
226
227
228
    routines. Extensive additional documentation for the graph, vertex and edge
    properties is available at: http://www.graphviz.org/doc/info/attrs.html.


    Examples
    --------
229
    >>> from numpy import *
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
    >>> from numpy.random import seed, zipf
    >>> seed(42)
    >>> g = gt.random_graph(1000, lambda: min(zipf(2.4), 40),
    ...                     lambda i,j: exp(abs(i-j)), directed=False)
    >>> # extract largest component
    >>> comp = gt.label_components(g)
    >>> h = gt.vertex_hist(g, comp)
    >>> max_comp = h[1][list(h[0]).index(max(h[0]))]
    >>> g.remove_vertex_if(lambda v: comp[v] != max_comp)
    >>>
    >>> deg = g.degree_property_map("out")
    >>> deg.get_array()[:] = 2*(sqrt(deg.get_array()[:])*0.5 + 0.4)
    >>> ebet = gt.betweenness(g)[1]
    >>> ebet.get_array()[:] *= 4000
    >>> ebet.get_array()[:] += 10
    >>> gt.graph_draw(g, vsize=deg, vcolor=deg, elen=10, ecolor=ebet,
    ...               penwidth=ebet, overlap="prism", output="graph-draw.png")
247
    <...>
248
249
250
251
252
253
254
255
256
257
258

    .. figure:: graph-draw.png
        :align: center

        Kamada-Kawai force-directed layout of a graph with a power-law degree
        distribution, and dissortative degree correlation. The vertex size and
        color indicate the degree, and the edge color and width the edge
        betweeness centrality.

    References
    ----------
259
    .. [graphviz] http://www.graphviz.org
260
261

    """
Tiago Peixoto's avatar
Tiago Peixoto committed
262

263
    if output != "" and output != None:
264
        output = os.path.expanduser(output)
265
        # check opening file for writing, since graphviz will bork if it is not
266
267
268
269
270
        # possible to open file
        if os.path.dirname(output) != "" and \
               not os.access(os.path.dirname(output), os.W_OK):
            raise IOError("cannot write to " + os.path.dirname(output))

Tiago Peixoto's avatar
Tiago Peixoto committed
271
272
273
274
275
    if g.is_directed():
        gvg = gv.digraph("G")
    else:
        gvg = gv.graph("G")

276
277
278
    if layout is None:
        layout = "neato" if g.num_vertices() <= 1000 else "sfdp"

279
280
281
282
283
    if layout == "arf":
        layout = "neato"
        pos = arf_layout(g, pos=pos)
        pin = True

284
285
    if pos != None:
        # copy user-supplied property
286
        if isinstance(pos, PropertyMap):
287
            pos = ungroup_vector_property(pos, [0,1])
288
289
        else:
            pos = (g.copy_property(pos[0]), g.copy_property(pos[1]))
290

291
292
293
294
295
296
297
    if type(vsize) == tuple:
        s = g.new_vertex_property("double")
        g.copy_property(vsize[0], s)
        s.a *= vsize[1]
        vsize = s

    if type(penwidth) == tuple:
298
        s = g.new_edge_property("double")
299
300
301
302
        g.copy_property(penwidth[0], s)
        s.a *= penwidth[1]
        penwidth = s

Tiago Peixoto's avatar
Tiago Peixoto committed
303
    # main graph properties
Tiago Peixoto's avatar
Tiago Peixoto committed
304
305
    gv.setv(gvg, "outputorder", "edgesfirst")
    gv.setv(gvg, "mode", "major")
306
    if overlap == False:
307
        overlap = "false"
308
309
    else:
        overlap = "true"
Tiago Peixoto's avatar
Tiago Peixoto committed
310
311
    if isinstance(overlap, str):
        gv.setv(gvg, "overlap", overlap)
312
    if sep != None:
Tiago Peixoto's avatar
Tiago Peixoto committed
313
        gv.setv(gvg, "sep", str(sep))
Tiago Peixoto's avatar
Tiago Peixoto committed
314
    if splines:
Tiago Peixoto's avatar
Tiago Peixoto committed
315
316
317
318
        gv.setv(gvg, "splines", "true")
    gv.setv(gvg, "ratio", str(ratio))
    # size is in centimeters... convert to inches
    gv.setv(gvg, "size", "%f,%f" % (size[0] / 2.54, size[1] / 2.54))
Tiago Peixoto's avatar
Tiago Peixoto committed
319
    if maxiter != None:
Tiago Peixoto's avatar
Tiago Peixoto committed
320
        gv.setv(gvg, "maxiter", str(maxiter))
321

322
323
    seed = numpy.random.randint(sys.maxint)
    gv.setv(gvg, "start", "%d" % seed)
Tiago Peixoto's avatar
Tiago Peixoto committed
324
325

    # apply all user supplied properties
Tiago Peixoto's avatar
Tiago Peixoto committed
326
    for k, val in gprops.iteritems():
Tiago Peixoto's avatar
Tiago Peixoto committed
327
328
329
330
331
332
333
334
335
336
        if isinstance(val, PropertyMap):
            gv.setv(gvg, k, str(val[g]))
        else:
            gv.setv(gvg, k, str(val))

    # normalize color properties
    if vcolor != None and not isinstance(vcolor, str):
        minmax = [float("inf"), -float("inf")]
        for v in g.vertices():
            c = vcolor[v]
Tiago Peixoto's avatar
Tiago Peixoto committed
337
338
            minmax[0] = min(c, minmax[0])
            minmax[1] = max(c, minmax[1])
339
340
        if minmax[0] == minmax[1]:
            minmax[1] += 1
Tiago Peixoto's avatar
Tiago Peixoto committed
341
342
        if vnorm:
            vnorm = matplotlib.colors.normalize(vmin=minmax[0], vmax=minmax[1])
343
344
        else:
            vnorm = lambda x: x
Tiago Peixoto's avatar
Tiago Peixoto committed
345
346
347
348

    if ecolor != None and not isinstance(ecolor, str):
        minmax = [float("inf"), -float("inf")]
        for e in g.edges():
349
            c = ecolor[e]
Tiago Peixoto's avatar
Tiago Peixoto committed
350
351
            minmax[0] = min(c, minmax[0])
            minmax[1] = max(c, minmax[1])
352
353
        if minmax[0] == minmax[1]:
            minmax[1] += 1
Tiago Peixoto's avatar
Tiago Peixoto committed
354
355
        if enorm:
            enorm = matplotlib.colors.normalize(vmin=minmax[0], vmax=minmax[1])
356
357
        else:
            enorm = lambda x: x
Tiago Peixoto's avatar
Tiago Peixoto committed
358

359
    nodes = {}
Tiago Peixoto's avatar
Tiago Peixoto committed
360
361

    # add nodes
362
363
364
365
366
    if vorder != None:
        vertices = sorted(g.vertices(), lambda a, b: cmp(vorder[a], vorder[b]))
    else:
        vertices = g.vertices()
    for v in vertices:
Tiago Peixoto's avatar
Tiago Peixoto committed
367
        n = gv.node(gvg, str(g.vertex_index[v]))
368
369
370

        if type(vsize) == PropertyMap:
            vw = vh = vsize[v]
371
        else:
372
            vw = vh = vsize
373
374
375

        gv.setv(n, "width", "%g" % vw)
        gv.setv(n, "height", "%g" % vh)
Tiago Peixoto's avatar
Tiago Peixoto committed
376
377
378
379
        gv.setv(n, "style", "filled")
        gv.setv(n, "color", "black")
        # apply color
        if vcolor != None:
Tiago Peixoto's avatar
Tiago Peixoto committed
380
            if isinstance(vcolor, str):
Tiago Peixoto's avatar
Tiago Peixoto committed
381
382
                gv.setv(n, "fillcolor", vcolor)
            else:
Tiago Peixoto's avatar
Tiago Peixoto committed
383
                color = tuple([int(c * 255.0) for c in vcmap(vnorm(vcolor[v]))])
Tiago Peixoto's avatar
Tiago Peixoto committed
384
385
386
387
388
389
390
                gv.setv(n, "fillcolor", "#%.2x%.2x%.2x%.2x" % color)
        else:
            gv.setv(n, "fillcolor", "red")
        gv.setv(n, "label", "")

        # user supplied position
        if pos != None:
Tiago Peixoto's avatar
Tiago Peixoto committed
391
            gv.setv(n, "pos", "%f,%f" % (pos[0][v], pos[1][v]))
Tiago Peixoto's avatar
Tiago Peixoto committed
392
393
394
            gv.setv(n, "pin", str(pin))

        # apply all user supplied properties
Tiago Peixoto's avatar
Tiago Peixoto committed
395
        for k, val in vprops.iteritems():
Tiago Peixoto's avatar
Tiago Peixoto committed
396
397
398
399
            if isinstance(val, PropertyMap):
                gv.setv(n, k, str(val[v]))
            else:
                gv.setv(n, k, str(val))
400
        nodes[v] = n
401

402
403
404
405
406
407
    # add edges
    if eorder != None:
        edges = sorted(g.edges(), lambda a, b: cmp(eorder[a], eorder[b]))
    else:
        edges = g.edges()
    for e in edges:
408
409
        ge = gv.edge(nodes[e.source()],
                     nodes[e.target()])
Tiago Peixoto's avatar
Tiago Peixoto committed
410
        gv.setv(ge, "arrowsize", "0.3")
411
412
        if g.is_directed():
            gv.setv(ge, "arrowhead", "vee")
413

Tiago Peixoto's avatar
Tiago Peixoto committed
414
415
        # apply color
        if ecolor != None:
Tiago Peixoto's avatar
Tiago Peixoto committed
416
            if isinstance(ecolor, str):
Tiago Peixoto's avatar
Tiago Peixoto committed
417
418
                gv.setv(ge, "color", ecolor)
            else:
Tiago Peixoto's avatar
Tiago Peixoto committed
419
                color = tuple([int(c * 255.0) for c in ecmap(enorm(ecolor[e]))])
Tiago Peixoto's avatar
Tiago Peixoto committed
420
421
                gv.setv(ge, "color", "#%.2x%.2x%.2x%.2x" % color)

422
423
424
425
        # apply edge length
        if elen != None:
            if isinstance(elen, PropertyMap):
                gv.setv(ge, "len", str(elen[e]))
Tiago Peixoto's avatar
Tiago Peixoto committed
426
            else:
427
                gv.setv(ge, "len", str(elen))
Tiago Peixoto's avatar
Tiago Peixoto committed
428
429

        # apply width
430
431
432
        if penwidth != None:
            if isinstance(penwidth, PropertyMap):
                gv.setv(ge, "penwidth", str(penwidth[e]))
Tiago Peixoto's avatar
Tiago Peixoto committed
433
            else:
434
                gv.setv(ge, "penwidth", str(penwidth))
Tiago Peixoto's avatar
Tiago Peixoto committed
435
436

        # apply all user supplied properties
Tiago Peixoto's avatar
Tiago Peixoto committed
437
        for k, v in eprops.iteritems():
Tiago Peixoto's avatar
Tiago Peixoto committed
438
439
440
441
            if isinstance(v, PropertyMap):
                gv.setv(ge, k, str(v[e]))
            else:
                gv.setv(ge, k, str(v))
442

Tiago Peixoto's avatar
Tiago Peixoto committed
443
    gv.layout(gvg, layout)
Tiago Peixoto's avatar
Tiago Peixoto committed
444
    gv.render(gvg, "dot", "/dev/null")  # retrieve positions
Tiago Peixoto's avatar
Tiago Peixoto committed
445
446
447

    if pos == None:
        pos = (g.new_vertex_property("double"), g.new_vertex_property("double"))
448
449
    for n, n_gv in nodes.iteritems():
        p = gv.getv(n_gv, "pos")
Tiago Peixoto's avatar
Tiago Peixoto committed
450
        p = p.split(",")
451
452
        pos[0][n] = float(p[0])
        pos[1][n] = float(p[1])
453

454
    # I don't get this, but it seems necessary
455
456
    pos[0].a /= 100
    pos[1].a /= 100
457

458
    pos = group_vector_property(pos)
459

460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
    if return_bitmap:
        # This is a not-so-nice hack which obtains an image buffer from a png
        # file. It is a pity that graphviz does not give access to its internal
        # buffers.
        tmp = tempfile.mkstemp(suffix=".png")[1]
        gv.render(gvg, "png", tmp)
        img = imread(tmp)
        os.remove(tmp)
    else:
        if output_format == "auto":
            if output == "":
                output_format = "xlib"
            elif output != None:
                output_format = output.split(".")[-1]

        # if using xlib we need to fork the process, otherwise good ol' graphviz
        # will call exit() when the window is closed
        if output_format == "xlib" or fork:
            pid = os.fork()
            if pid == 0:
                gv.render(gvg, output_format, output)
Tiago Peixoto's avatar
Tiago Peixoto committed
481
                os._exit(0)  # since we forked, it's good to be sure
482
483
484
485
486
487
488
489
490
            if output_format != "xlib":
                os.wait()
        elif output != None:
            gv.render(gvg, output_format, output)

    ret = [pos]
    if return_bitmap:
        ret.append(img)

491
    if returngv:
492
        ret.append(gv)
493
494
    else:
        gv.rm(gvg)
495
        del gvg
496
497
498
499
500

    if len(ret) > 1:
        return tuple(ret)
    else:
        return ret[0]
501

Tiago Peixoto's avatar
Tiago Peixoto committed
502

503
def random_layout(g, shape=None, pos=None, dim=2):
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
    r"""Performs a random layout of the graph.

    Parameters
    ----------
    g : Graph
        Graph to be used.
    shape : tuple (optional, default: None)
        Rectangular shape of the bounding area. If None, a square of linear size
        :math:`\sqrt{N}` is used.
    pos : PropertyMap (optional, default: None)
        Vector vertex property maps where the coordinates should be stored.
    dim : int (optional, default: 2)
        Number of coordinates per vertex.

    Returns
    -------
    pos : A vector vertex property map
        Vertex property map with the coordinates of the vertices.

    Notes
    -----
    This algorithm has complexity :math:`O(V)`.
    """

528
529
530
531
532
533
534
    if pos == None:
        pos = [g.new_vertex_property("double") for i in xrange(dim)]

    if isinstance(pos, PropertyMap) and "vector" in pos.value_type():
        pos = ungroup_vector_property(pos)

    if shape == None:
Tiago Peixoto's avatar
Tiago Peixoto committed
535
        shape = [sqrt(g.num_vertices())] * dim
536
537
538
539
540

    for i in xrange(dim):
        _check_prop_scalar(pos[i], name="pos[%d]" % i)
        _check_prop_writable(pos[i], name="pos[%d]" % i)
        a = pos[i].get_array()
Tiago Peixoto's avatar
Tiago Peixoto committed
541
        a[:] = numpy.random.random(len(a)) * shape[i]
542
543
544
545

    pos = group_vector_property(g, pos)
    return pos

Tiago Peixoto's avatar
Tiago Peixoto committed
546

547
548
def arf_layout(g, weight=None, d=0.1, a=10, dt=0.001, epsilon=1e-6,
               max_iter=1000, pos=None, dim=2):
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
    r"""Calculate the ARF spring-block layout of the graph.

    Parameters
    ----------
    g : Graph
        Graph to be used.
    weight : PropertyMap (optional, default: None)
        An edge property map with the respective weights.
    d : float (optional, default: 0.1)
        Opposing force between vertices.
    a : float (optional, default: 10)
        Attracting force between adjacent vertices.
    dt : float (optional, default: 0.001)
        Iteration step size.
    epsilon : float (optional, default: 1e-6)
        Convergence criterion.
    max_iter : int (optional, default: 1000)
        Maximum number of iterations. If this value is 0, it runs until
        convergence.
    pos : PropertyMap (optional, default: None)
        Vector vertex property maps where the coordinates should be stored.
    dim : int (optional, default: 2)
        Number of coordinates per vertex.

    Returns
    -------
    pos : A vector vertex property map
        Vertex property map with the coordinates of the vertices.

    Notes
    -----
580
    This algorithm is defined in [geipel-self-organization-2007]_, and has
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
    complexity :math:`O(V^2)`.

    Examples
    --------
    >>> from numpy.random import seed, zipf
    >>> seed(42)
    >>> g = gt.random_graph(100, lambda: 3, directed=False)
    >>> t = gt.min_spanning_tree(g)
    >>> g.set_edge_filter(t)
    >>> pos = gt.graph_draw(g, output=None) # initial configuration
    >>> pos = gt.arf_layout(g, pos=pos, max_iter=0)
    >>> gt.graph_draw(g, pos=pos, pin=True, output="graph-draw-arf.png")
    <...>

    .. figure:: graph-draw-arf.png
        :align: center

        ARF layout of a minimum spanning tree of a random graph.

    References
    ----------
602
    .. [geipel-self-organization-2007] Markus M. Geipel, "Self-Organization
603
604
605
606
607
       applied to Dynamic Network Layout" , International Journal of Modern
       Physics C vol. 18, no. 10 (2007), pp. 1537-1549, arXiv:0704.1748v5
    .. _arf: http://www.sg.ethz.ch/research/graphlayout
    """

608
    if pos == None:
609
610
611
612
        if dim != 2:
            pos = random_layout(g, dim=dim)
        else:
            pos = graph_draw(g, output=None)
613
614
615
    _check_prop_vector(pos, name="pos", floating=True)

    g.stash_filter(directed=True)
616
617
618
619
620
621
622
    try:
        g.set_directed(False)
        libgraph_tool_layout.arf_layout(g._Graph__graph, _prop("v", g, pos),
                                        _prop("e", g, weight), d, a, dt,
                                        max_iter, epsilon, dim)
    finally:
        g.pop_filter(directed=True)
623
    return pos