| // Copyright ©2019 The Gonum Authors. All rights reserved. |
| // Use of this source code is governed by a BSD-style |
| // license that can be found in the LICENSE file. |
| |
| package layout |
| |
| import ( |
| "math" |
| |
| "golang.org/x/exp/rand" |
| |
| "gonum.org/v1/gonum/graph" |
| "gonum.org/v1/gonum/spatial/barneshut" |
| "gonum.org/v1/gonum/spatial/r2" |
| ) |
| |
| // EadesR2 implements the graph layout algorithm essentially as |
| // described in "A heuristic for graph drawing", Congressus |
| // numerantium 42:149-160. |
| // The implementation here uses the Barnes-Hut approximation for |
| // global repulsion calculation, and edge weights are considered |
| // when calculating adjacent node attraction. |
| type EadesR2 struct { |
| // Updates is the number of updates to perform. |
| Updates int |
| |
| // Repulsion is the strength of the global |
| // repulsive force between nodes in the |
| // layout. It corresponds to C3 in the paper. |
| Repulsion float64 |
| |
| // Rate is the gradient descent rate. It |
| // corresponds to C4 in the paper. |
| Rate float64 |
| |
| // Theta is the Barnes-Hut theta constant. |
| Theta float64 |
| |
| // Src is the source of randomness used |
| // to initialize the nodes' locations. If |
| // Src is nil, the global random number |
| // generator is used. |
| Src rand.Source |
| |
| nodes graph.Nodes |
| indexOf map[int64]int |
| |
| particles []barneshut.Particle2 |
| forces []r2.Vec |
| } |
| |
| // Update is the EadesR2 spatial graph update function. |
| func (u *EadesR2) Update(g graph.Graph, layout LayoutR2) bool { |
| if u.Updates <= 0 { |
| return false |
| } |
| u.Updates-- |
| |
| if !layout.IsInitialized() { |
| var rnd func() float64 |
| if u.Src == nil { |
| rnd = rand.Float64 |
| } else { |
| rnd = rand.New(u.Src).Float64 |
| } |
| u.nodes = g.Nodes() |
| u.indexOf = make(map[int64]int, u.nodes.Len()) |
| if u.nodes.Len() >= 0 { |
| u.particles = make([]barneshut.Particle2, 0, u.nodes.Len()) |
| } |
| for u.nodes.Next() { |
| id := u.nodes.Node().ID() |
| u.indexOf[id] = len(u.particles) |
| u.particles = append(u.particles, eadesR2Node{id: id, pos: r2.Vec{X: rnd(), Y: rnd()}}) |
| } |
| u.forces = make([]r2.Vec, len(u.particles)) |
| } |
| u.nodes.Reset() |
| |
| // Apply global repulsion. |
| plane, err := barneshut.NewPlane(u.particles) |
| if err != nil { |
| return false |
| } |
| var updated bool |
| for i, p := range u.particles { |
| f := r2.Scale(-u.Repulsion, plane.ForceOn(p, u.Theta, barneshut.Gravity2)) |
| // Prevent marginal updates that can be caused by |
| // floating point error when nodes are very far apart. |
| if math.Hypot(f.X, f.Y) > 1e-12 { |
| updated = true |
| } |
| u.forces[i] = f |
| } |
| |
| // Handle edge weighting for attraction. |
| var weight func(uid, vid int64) float64 |
| if wg, ok := g.(graph.Weighted); ok { |
| if _, ok := g.(graph.Directed); ok { |
| weight = func(xid, yid int64) float64 { |
| var w float64 |
| f, ok := wg.Weight(xid, yid) |
| if ok { |
| w += f |
| } |
| r, ok := wg.Weight(yid, xid) |
| if ok { |
| w += r |
| } |
| return w |
| } |
| } else { |
| weight = func(xid, yid int64) float64 { |
| w, ok := wg.Weight(xid, yid) |
| if ok { |
| return w |
| } |
| return 0 |
| } |
| } |
| } else { |
| // This is only called when the adjacency is known so just return unit. |
| weight = func(_, _ int64) float64 { return 1 } |
| } |
| |
| seen := make(map[[2]int64]bool) |
| for u.nodes.Next() { |
| xid := u.nodes.Node().ID() |
| xidx := u.indexOf[xid] |
| to := g.From(xid) |
| for to.Next() { |
| yid := to.Node().ID() |
| if seen[[2]int64{xid, yid}] { |
| continue |
| } |
| seen[[2]int64{yid, xid}] = true |
| yidx := u.indexOf[yid] |
| |
| // Apply adjacent node attraction. |
| v := r2.Sub(u.particles[yidx].Coord2(), u.particles[xidx].Coord2()) |
| f := r2.Scale(weight(xid, yid)*math.Log(math.Hypot(v.X, v.Y)), v) |
| if math.IsInf(f.X, 0) || math.IsInf(f.Y, 0) { |
| return false |
| } |
| if math.Hypot(f.X, f.Y) > 1e-12 { |
| updated = true |
| } |
| u.forces[xidx] = r2.Add(u.forces[xidx], f) |
| u.forces[yidx] = r2.Sub(u.forces[yidx], f) |
| } |
| } |
| |
| if !updated { |
| return false |
| } |
| |
| rate := u.Rate |
| if rate == 0 { |
| rate = 0.1 |
| } |
| for i, f := range u.forces { |
| n := u.particles[i].(eadesR2Node) |
| n.pos = r2.Add(n.pos, r2.Scale(rate, f)) |
| u.particles[i] = n |
| layout.SetCoord2(n.id, n.pos) |
| } |
| return true |
| } |
| |
| type eadesR2Node struct { |
| id int64 |
| pos r2.Vec |
| } |
| |
| func (p eadesR2Node) Coord2() r2.Vec { return p.pos } |
| func (p eadesR2Node) Mass() float64 { return 1 } |