__init__.py 28.5 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-2011 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

Summary
+++++++

.. autosummary::
   :nosignatures:

   graph_draw
32
   fruchterman_reingold_layout
33
34
35
36
37
   arf_layout
   random_layout

Contents
++++++++
38
39
"""

40
41
42
43
44
45
46
import sys
import os
import os.path
import time
import warnings
import ctypes
import ctypes.util
47
import tempfile
48
from .. import _degree, _prop, PropertyMap, _check_prop_vector,\
49
     _check_prop_scalar, _check_prop_writable, group_vector_property,\
50
51
     ungroup_vector_property, GraphView
from .. topology import label_components
Tiago Peixoto's avatar
Tiago Peixoto committed
52
from .. decorators import _limit_args
53
import numpy.random
54
from numpy import *
55
import copy
56
57
58

from .. dl_import import dl_import
dl_import("import libgraph_tool_layout")
59

60
61
62
63
64
65
try:
    import matplotlib.cm
    import matplotlib.colors
except ImportError:
    warnings.warn("error importing matplotlib module... " + \
                  "graph_draw() will not work.", ImportWarning)
Tiago Peixoto's avatar
Tiago Peixoto committed
66

67
68
69
try:
    libname = ctypes.util.find_library("c")
    libc = ctypes.CDLL(libname)
70
71
    if hasattr(libc, "open_memstream"):
        libc.open_memstream.restype = ctypes.POINTER(ctypes.c_char)
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
except OSError:
    pass

try:
    libname = ctypes.util.find_library("gvc")
    if libname is None:
        raise OSError()
    libgv = ctypes.CDLL(libname)
    # properly set the return types of certain functions
    ptype = ctypes.POINTER(ctypes.c_char)
    libgv.gvContext.restype = ptype
    libgv.agopen.restype = ptype
    libgv.agnode.restype = ptype
    libgv.agedge.restype = ptype
    libgv.agget.restype = ptype
    # create a context to use the whole time (if we keep freeing and recreating
    # it, we will hit a memory leak in graphviz)
    gvc = libgv.gvContext()
except OSError:
    warnings.warn("error importing graphviz C library (libgvc)... " + \
                  "graph_draw() will not work.", ImportWarning)


95
96
__all__ = ["graph_draw", "fruchterman_reingold_layout", "arf_layout",
           "random_layout"]
97

Tiago Peixoto's avatar
Tiago Peixoto committed
98

99
100
101
102
103
104
105
106
107
def aset(elem, attr, value):
    v = str(value)
    libgv.agsafeset(elem, str(attr), v, v)


def aget(elem, attr):
    return ctypes.string_at(libgv.agget(elem, str(attr)))


108
109
110
def graph_draw(g, pos=None, size=(15, 15), pin=False, layout=None, maxiter=None,
               ratio="fill", overlap="prism", sep=None, splines=False,
               vsize=0.105, penwidth=1.0, elen=None, gprops={}, vprops={},
111
112
113
114
               eprops={}, vcolor="#a40000", ecolor="#2e3436", vcmap=None,
               vnorm=True, ecmap=None, enorm=True, vorder=None, eorder=None,
               output="", output_format="auto", fork=False,
               return_string=False, seed=0):
115
116
117
118
119
120
    r"""Draw a graph using graphviz.

    Parameters
    ----------
    g : Graph
        Graph to be used.
121
    pos : PropertyMap or tuple of PropertyMaps (optional, default: None)
122
123
124
125
126
        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.
127
    layout : string (default: "neato" if g.num_vertices() <= 1000 else "sfdp")
128
        Layout engine to be used. Possible values are "neato", "fdp", "dot",
129
        "circo", "twopi" and "arf".
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
    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
159
        drawn on a single page, then size is set to an "ideal" value. In
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
187
188
189
190
        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.
191
192
193
194
    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)
195
196
        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
197
198
        on text. If a tuple is specified, the first value should be a property
        map, and the second is a scale factor.
199
200
201
202
203
204
205
206
207
208
209
210
211
    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.
212
    vcolor : string or PropertyMap (default: "#a40000")
213
214
215
        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.
216
    ecolor : string or PropertyMap (default: "#2e3436")
217
218
219
220
221
222
223
224
225
226
227
        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.
228
229
230
231
232
233
    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.
234
235
236
237
238
239
    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'
240
241
        parameter, or 'xlib' if it is empty. If the value is None, no output is
        produced.
242
    fork : bool (default: False)
243
        If True, the program is forked before drawing. This is used as a
244
245
246
        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'.
247
248
249
    return_string : bool (default: False)
        If True, a string containing the rendered graph as binary data is
        returned (defaults to png format).
250
251
252

    Returns
    -------
253
254
    pos : PropertyMap
        Vector vertex property map with the x and y coordinates of the vertices.
255
256
257
258
259
260
    gv : gv.digraph or gv.graph (optional, only if returngv == True)
        Internally used graphviz graph.


    Notes
    -----
261
262
263
    This function is a wrapper for the [graphviz] routines. Extensive additional
    documentation for the graph, vertex and edge properties is available at:
    http://www.graphviz.org/doc/info/attrs.html.
264
265
266
267


    Examples
    --------
268
    >>> from numpy import *
269
270
271
    >>> from numpy.random import seed, zipf
    >>> seed(42)
    >>> g = gt.random_graph(1000, lambda: min(zipf(2.4), 40),
272
    ...                     lambda i, j: exp(abs(i - j)), directed=False)
273
    >>> # extract largest component
274
    >>> g = gt.GraphView(g, vfilt=gt.label_largest_component(g))
275
    >>> deg = g.degree_property_map("out")
276
    >>> deg.a = 2 * (sqrt(deg.a) * 0.5 + 0.4)
277
    >>> ebet = gt.betweenness(g)[1]
278
279
280
281
282
    >>> ebet.a *= 4000
    >>> ebet.a += 10
    >>> gt.graph_draw(g, vsize=deg, vcolor=deg, vorder=deg, elen=10,
    ...               ecolor=ebet, eorder=ebet, penwidth=ebet,
    ...               overlap="prism", output="graph-draw.png")
283
    <...>
284
285
286
287
288
289
290
291
292
293
294

    .. 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
    ----------
295
    .. [graphviz] http://www.graphviz.org
296
297

    """
Tiago Peixoto's avatar
Tiago Peixoto committed
298

299
    if output != "" and output != None:
300
        output = os.path.expanduser(output)
301
        # check opening file for writing, since graphviz will bork if it is not
302
303
304
305
306
        # 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))

307
308
309
    has_layout = False
    try:
        gvg = libgv.agopen("G", 1 if g.is_directed() else 0)
310

311
312
        if layout is None:
            layout = "neato" if g.num_vertices() <= 1000 else "sfdp"
313

314
315
316
317
        if layout == "arf":
            layout = "neato"
            pos = arf_layout(g, pos=pos)
            pin = True
Tiago Peixoto's avatar
Tiago Peixoto committed
318

319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
        if pos != None:
            # copy user-supplied property
            if isinstance(pos, PropertyMap):
                pos = ungroup_vector_property(pos, [0, 1])
            else:
                pos = (g.copy_property(pos[0]), g.copy_property(pos[1]))

        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:
            s = g.new_edge_property("double")
            g.copy_property(penwidth[0], s)
            s.a *= penwidth[1]
            penwidth = s

        # main graph properties
        aset(gvg, "outputorder", "edgesfirst")
        aset(gvg, "mode", "major")
        if overlap == False:
            overlap = "false"
343
        else:
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
            overlap = "true"
        if isinstance(overlap, str):
            aset(gvg, "overlap", overlap)
        if sep != None:
            aset(gvg, "sep", sep)
        if splines:
            aset(gvg, "splines", "true")
        aset(gvg, "ratio", ratio)
        # size is in centimeters... convert to inches
        aset(gvg, "size", "%f,%f" % (size[0] / 2.54, size[1] / 2.54))
        if maxiter != None:
            aset(gvg, "maxiter", maxiter)

        seed = numpy.random.randint(sys.maxint)
        aset(gvg, "start", "%d" % seed)

        # apply all user supplied graph properties
        for k, val in gprops.iteritems():
            if isinstance(val, PropertyMap):
                aset(gvg, k, val[g])
            else:
                aset(gvg, k, 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]
                minmax[0] = min(c, minmax[0])
                minmax[1] = max(c, minmax[1])
            if minmax[0] == minmax[1]:
                minmax[1] += 1
            if vnorm:
                vnorm = matplotlib.colors.normalize(vmin=minmax[0], vmax=minmax[1])
            else:
                vnorm = lambda x: x

        if ecolor != None and not isinstance(ecolor, str):
            minmax = [float("inf"), -float("inf")]
            for e in g.edges():
                c = ecolor[e]
                minmax[0] = min(c, minmax[0])
                minmax[1] = max(c, minmax[1])
            if minmax[0] == minmax[1]:
                minmax[1] += 1
            if enorm:
                enorm = matplotlib.colors.normalize(vmin=minmax[0],
                                                    vmax=minmax[1])
            else:
                enorm = lambda x: x
394

395
396
        if vcmap is None:
            vcmap = matplotlib.cm.jet
Tiago Peixoto's avatar
Tiago Peixoto committed
397

398
399
        if ecmap is None:
            ecmap = matplotlib.cm.jet
400

401
402
403
        # add nodes
        if vorder != None:
            vertices = sorted(g.vertices(), lambda a, b: cmp(vorder[a], vorder[b]))
404
        else:
405
406
407
            vertices = g.vertices()
        for v in vertices:
            n = libgv.agnode(gvg, str(int(v)))
Tiago Peixoto's avatar
Tiago Peixoto committed
408

409
410
            if type(vsize) == PropertyMap:
                vw = vh = vsize[v]
Tiago Peixoto's avatar
Tiago Peixoto committed
411
            else:
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
                vw = vh = vsize

            aset(n, "shape", "circle")
            aset(n, "width", "%g" % vw)
            aset(n, "height", "%g" % vh)
            aset(n, "style", "filled")
            aset(n, "color", ecolor if isinstance(ecolor, str) else "#2e3436")
            # apply color
            if isinstance(vcolor, str):
                aset(n, "fillcolor", vcolor)
            else:
                color = tuple([int(c * 255.0) for c in vcmap(vnorm(vcolor[v]))])
                aset(n, "fillcolor", "#%.2x%.2x%.2x%.2x" % color)
            aset(n, "label", "")

            # user supplied position
428
            if pos is not None:
429
430
431
432
433
434
435
436
437
438
439
440
441
                aset(n, "pos", "%f,%f" % (pos[0][v], pos[1][v]))
                aset(n, "pin", pin)

            # apply all user supplied properties
            for k, val in vprops.iteritems():
                if isinstance(val, PropertyMap):
                    aset(n, k, val[v])
                else:
                    aset(n, k, val)

        # add edges
        if eorder != None:
            edges = sorted(g.edges(), lambda a, b: cmp(eorder[a], eorder[b]))
442
        else:
443
444
445
446
447
448
449
450
451
452
453
454
            edges = g.edges()
        for e in edges:
            ge = libgv.agedge(gvg,
                              libgv.agnode(gvg, str(int(e.source()))),
                              libgv.agnode(gvg, str(int(e.target()))))
            aset(ge, "arrowsize", "0.3")
            if g.is_directed():
                aset(ge, "arrowhead", "vee")

            # apply color
            if isinstance(ecolor, str):
                aset(ge, "color", ecolor)
Tiago Peixoto's avatar
Tiago Peixoto committed
455
            else:
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
                color = tuple([int(c * 255.0) for c in ecmap(enorm(ecolor[e]))])
                aset(ge, "color", "#%.2x%.2x%.2x%.2x" % color)

            # apply edge length
            if elen != None:
                if isinstance(elen, PropertyMap):
                    aset(ge, "len", elen[e])
                else:
                    aset(ge, "len", elen)

            # apply width
            if penwidth != None:
                if isinstance(penwidth, PropertyMap):
                    aset(ge, "penwidth", penwidth[e])
                else:
                    aset(ge, "penwidth", penwidth)

            # apply all user supplied properties
            for k, v in eprops.iteritems():
                if isinstance(v, PropertyMap):
                    aset(ge, k, v[e])
                else:
                    aset(ge, k, v)

        libgv.gvLayout(gvc, gvg, layout)
        has_layout = True
        retv = libgv.gvRender(gvc, gvg, "dot", None)  # retrieve positions only

        if pos == None:
            pos = (g.new_vertex_property("double"),
                   g.new_vertex_property("double"))
        for v in g.vertices():
            n = libgv.agnode(gvg, str(int(v)))
            p = aget(n, "pos")
            p = p.split(",")
            pos[0][v] = float(p[0])
            pos[1][v] = float(p[1])

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

        pos = group_vector_property(pos)

        if return_string:
            if output_format == "auto":
                output_format = "png"
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
            if hasattr(libc, "open_memstream"):
                buf = ctypes.c_char_p()
                buf_len = ctypes.c_size_t()
                fstream = libc.open_memstream(ctypes.byref(buf),
                                              ctypes.byref(buf_len))
                libgv.gvRender(gvc, gvg, output_format, fstream)
                libc.fclose(fstream)
                data = copy.copy(ctypes.string_at(buf, buf_len.value))
                libc.free(buf)
            else:
                # write to temporary file, if open_memstream is not available
                output = tempfile.mkstemp()[1]
                libgv.gvRenderFilename(gvc, gvg, output_format, output)
                data = open(output).read()
                os.remove(output)
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
        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:
                    libgv.gvRenderFilename(gvc, gvg, output_format, output)
                    os._exit(0)  # since we forked, it's good to be sure
                if output_format != "xlib":
                    os.wait()
            elif output != None:
                libgv.gvRenderFilename(gvc, gvg, output_format, output)
Tiago Peixoto's avatar
Tiago Peixoto committed
536

537
538
539
        ret = [pos]
        if return_string:
            ret.append(data)
Tiago Peixoto's avatar
Tiago Peixoto committed
540

541
542
543
544
    finally:
        if has_layout:
            libgv.gvFreeLayout(gvc, gvg)
        libgv.agclose(gvg)
545
546
547
548
549

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

Tiago Peixoto's avatar
Tiago Peixoto committed
551

552
def random_layout(g, shape=None, pos=None, dim=2):
553
554
555
556
557
558
    r"""Performs a random layout of the graph.

    Parameters
    ----------
    g : Graph
        Graph to be used.
Tiago Peixoto's avatar
Tiago Peixoto committed
559
560
561
562
563
    shape : tuple or list (optional, default: None)
        Rectangular shape of the bounding area. The size of this parameter must
        match `dim`, and each element can be either a pair specifying a range,
        or a single value specifying a range starting from zero. If None is
        passed, a square of linear size :math:`\sqrt{N}` is used.
564
565
566
567
568
569
570
571
572
573
574
575
576
    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)`.
Tiago Peixoto's avatar
Tiago Peixoto committed
577
578
579
580
581
582
583
584
585
586
587

    Examples
    --------
    >>> from numpy.random import seed
    >>> seed(42)
    >>> g = gt.random_graph(100, lambda: (3, 3))
    >>> shape = [[50, 100], [1, 2], 4]
    >>> pos = gt.random_layout(g, shape=shape, dim=3)
    >>> pos[g.vertex(0)].a
    array([ 86.59969709,   1.31435598,   0.64651486])

588
589
    """

590
    if pos == None:
Tiago Peixoto's avatar
Tiago Peixoto committed
591
592
        pos = g.new_vertex_property("vector<double>")
    _check_prop_vector(pos, name="pos")
593

Tiago Peixoto's avatar
Tiago Peixoto committed
594
    pos = ungroup_vector_property(pos, range(0, dim))
595
596

    if shape == None:
Tiago Peixoto's avatar
Tiago Peixoto committed
597
        shape = [sqrt(g.num_vertices())] * dim
598
599

    for i in xrange(dim):
Tiago Peixoto's avatar
Tiago Peixoto committed
600
601
602
603
604
605
606
607
        if hasattr(shape[i], "__len__"):
            if len(shape[i]) != 2:
                raise ValueError("The elements of 'shape' must have size 2.")
            r = [min(shape[i]), max(shape[i])]
        else:
            r = [min(shape[i], 0), max(shape[i], 0)]
        d = r[1] - r[0]
        pos[i].a = numpy.random.random(g.num_vertices()) * d + r[0]
608

Tiago Peixoto's avatar
Tiago Peixoto committed
609
    pos = group_vector_property(pos)
610
611
    return pos

Tiago Peixoto's avatar
Tiago Peixoto committed
612

613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
def fruchterman_reingold_layout(g, weight=None, a=None, r=1., scale=None,
                                circular=False, grid=True, t_range=None,
                                n_iter=100, pos=None):
    r"""Calculate the Fruchterman-Reingold 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.
    a : float (optional, default: :math:`V`)
        Attracting force between adjacent vertices.
    r : float (optional, default: 1.0)
        Repulsive force between vertices.
    scale : float (optional, default: :math:`\sqrt{V}`)
        Total scale of the layout (either square side or radius).
    circular : bool (optional, default: False)
        If `True`, the layout will have a circular shape. Otherwise the shape
        will be a square.
    grid : bool (optional, default: True)
        If `True`, the repulsive forces will only act on vertices which are on
        the same site on a grid. Otherwise they will act on all vertex pairs.
    t_range : tuple of floats (optional, default: (scale / 10, scale / 1000))
        Temperature range used in annealing. The temperature limits the
        displacement at each iteration.
    n_iter : int (optional, default: 100)
        Total number of iterations.
    pos : PropertyMap (optional, default: None)
        Vector vertex property maps where the coordinates should be stored. If
        provided, this will also be used as the initial position of the
        vertices.

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

    Notes
    -----
    This algorithm is defined in [fruchterman-reingold]_, and has
Tiago Peixoto's avatar
Tiago Peixoto committed
654
655
    complexity :math:`O(\text{n-iter}\times V^2)` if `grid=False` or
    :math:`O(\text{n-iter}\times (V + E))` otherwise.
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701

    Examples
    --------
    >>> from numpy.random import seed, zipf
    >>> seed(42)
    >>> g = gt.price_network(300)
    >>> pos = gt.fruchterman_reingold_layout(g, n_iter=1000)
    >>> gt.graph_draw(g, pos=pos, pin=True, output="graph-draw-fr.png")
    <...>

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

        Fruchterman-Reingold layout of a Price network.

    References
    ----------
    .. [fruchterman-reingold] Fruchterman, Thomas M. J.; Reingold, Edward M.
       "Graph Drawing by Force-Directed Placement". Software – Practice & Experience
       (Wiley) 21 (11): 1129–1164. (1991) :doi:`10.1002/spe.4380211102`
    """

    if pos == None:
        pos = random_layout(g, dim=2)
    _check_prop_vector(pos, name="pos", floating=True)

    if a is None:
        a = float(g.num_vertices())

    if scale is None:
        scale = sqrt(g.num_vertices())

    if t_range is None:
        t_range = (scale / 10, scale / 1000)

    ug = GraphView(g, directed=False)
    libgraph_tool_layout.fruchterman_reingold_layout(ug._Graph__graph,
                                                     _prop("v", g, pos),
                                                     _prop("e", g, weight),
                                                     a, r, not circular, scale,
                                                     grid, t_range[0],
                                                     t_range[1], n_iter)
    return pos


def arf_layout(g, weight=None, d=0.5, a=10, dt=0.001, epsilon=1e-6,
702
               max_iter=1000, pos=None, dim=2):
703
704
705
706
707
708
709
710
    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.
711
    d : float (optional, default: 0.5)
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
        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
    -----
734
    This algorithm is defined in [geipel-self-organization-2007]_, and has
735
736
737
738
739
740
    complexity :math:`O(V^2)`.

    Examples
    --------
    >>> from numpy.random import seed, zipf
    >>> seed(42)
741
742
    >>> g = gt.price_network(300)
    >>> pos = gt.arf_layout(g, max_iter=0)
743
744
745
746
747
748
    >>> gt.graph_draw(g, pos=pos, pin=True, output="graph-draw-arf.png")
    <...>

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

749
        ARF layout of a Price network.
750
751
752

    References
    ----------
753
    .. [geipel-self-organization-2007] Markus M. Geipel, "Self-Organization
754
755
756
       applied to Dynamic Network Layout", International Journal of Modern
       Physics C vol. 18, no. 10 (2007), pp. 1537-1549,
       :doi:`10.1142/S0129183107011558`, :arxiv:`0704.1748v5`
757
758
759
    .. _arf: http://www.sg.ethz.ch/research/graphlayout
    """

760
    if pos is None:
761
762
763
764
        if dim != 2:
            pos = random_layout(g, dim=dim)
        else:
            pos = graph_draw(g, output=None)
765
766
    _check_prop_vector(pos, name="pos", floating=True)

767
768
769
770
    ug = GraphView(g, directed=False)
    libgraph_tool_layout.arf_layout(ug._Graph__graph, _prop("v", g, pos),
                                    _prop("e", g, weight), d, a, dt, max_iter,
                                    epsilon, dim)
771
    return pos