DijkstraFrom & EgoGraph for gonum

DijkstraFrom.go

// DijkstraFrom returns a shortest-path tree for a shortest path from u to all nodes in
// the graph g. If the graph does not implement Weighted, UniformCost is used.
// DijkstraFrom will panic if g has a u-reachable negative edge weight.
//
// If g is a graph.Graph, all nodes of the graph will be stored in the shortest-path
// tree, otherwise only nodes reachable from u will be stored.
//
// The time complexity of DijkstrFrom is O(|E|.log|V|).
func DijkstraFrom(u graph.Node, g traverse.Graph) Shortest {
	var path Shortest
	if h, ok := g.(graph.Graph); ok {
		if h.Node(u.ID()) == nil {
			return Shortest{from: u}
		}
		path = newShortestFrom(u, graph.NodesOf(h.Nodes()))
	} else {
		if g.From(u.ID()) == nil {
			return Shortest{from: u}
		}
		path = newShortestFrom(u, []graph.Node{u})
	}

	var weight Weighting
	if wg, ok := g.(Weighted); ok {
		weight = wg.Weight
	} else {
		weight = UniformCost(g)
	}

	// Dijkstra's algorithm here is implemented essentially as
	// described in Function B.2 in figure 6 of UTCS Technical
	// Report TR-07-54.
	//
	// This implementation deviates from the report as follows:
	// - the value of path.dist for the start vertex u is initialized to 0;
	// - outdated elements from the priority queue (i.e. with respect to the dist value)
	//   are skipped.
	//
	// http://www.cs.utexas.edu/ftp/techreports/tr07-54.pdf
	Q := priorityQueue{{node: u, dist: 0}}
	for Q.Len() != 0 {
		mid := heap.Pop(&Q).(distanceNode)
		k := path.indexOf[mid.node.ID()]
		if mid.dist > path.dist[k] {
			continue
		}
		mnid := mid.node.ID()
		for _, v := range graph.NodesOf(g.From(mnid)) {
			vid := v.ID()
			j, ok := path.indexOf[vid]
			if !ok {
				j = path.add(v)
			}
			w, ok := weight(mnid, vid)
			if !ok {
				panic("dijkstra: unexpected invalid weight")
			}
			if w < 0 {
				panic("dijkstra: negative edge weight")
			}
			joint := path.dist[k] + w
			if joint < path.dist[j] {
				heap.Push(&Q, distanceNode{node: v, dist: joint})
				path.set(j, joint, k)
			}
		}
	}

	return path
}

EgoFrom.go

func EgoFrom(u graph.Node, g graph.Graph, c float64) *simple.WeightedUndirectedGraph {
	path := newShortestFrom(u, graph.NodesOf(g.Nodes()))
	ego := simple.NewWeightedUndirectedGraph(0, math.Inf(1))

	if g.Node(u.ID()) == nil {
		panic("ego: from node doesn't have ID")
	}

	var weight Weighting
	if wg, ok := g.(Weighted); ok {
		weight = wg.Weight
	} else {
		weight = UniformCost(g)
	}
    
	Q := priorityQueue{{node: u, dist: 0}}
	for Q.Len() != 0 {
		mid := heap.Pop(&Q).(distanceNode)
        
		if mid.dist > c {
			break
		}
        
		k := path.indexOf[mid.node.ID()]
		if mid.dist > path.dist[k] {
			continue
		}
		mnid := mid.node.ID()
		for _, v := range graph.NodesOf(g.From(mnid)) {
			vid := v.ID()
			j, ok := path.indexOf[vid]
			if !ok {
				j = path.add(v)
			}
			w, ok := weight(mnid, vid)
			if !ok {
				panic("ego: unexpected invalid weight")
			}
			if w < 0 {
				panic("ego: negative edge weight")
			}
			joint := path.dist[k] + w
			if joint > c {
				continue
			}
			if joint < path.dist[j] {
				heap.Push(&Q, distanceNode{node: v, dist: joint})
				path.set(j, joint, k)
			}
		}
	}
	
	for id := range path.next {
		if path.next[id] != -1 {
			mnid := path.nodes[id]
			vid := path.nodes[path.next[id]]
			w, _ := weight(mnid.ID(), vid.ID())
			ego.SetWeightedEdge(simple.WeightedEdge{F: mnid, T: vid, W: w})
		}
	}

	return ego
}

egograph.go

package main

import (
    "fmt"
    "os"
    "io"
    "math"
    "strconv"
    "encoding/csv"
    "gonum.org/v1/gonum/graph/path"
    "gonum.org/v1/gonum/graph/simple"
)

func main() {

    // Define Graph
    var self float64 = 0
    var absent float64 = math.Inf(1)
    g := simple.NewWeightedUndirectedGraph(self, absent)

    // Open File
    fp, err := os.Open("./sample.weighted.edgelist")
    if err != nil {
       panic(err)
    }
    defer fp.Close()

    // Read rows and add edges
    reader := csv.NewReader(fp)
    reader.Comma = '\t'
    for {
        record, err := reader.Read()
        if err == io.EOF {
            break
        } else if err != nil {
            panic(err)
        }

        id_from, _ := strconv.ParseInt(record[0], 10, 64)
        id_to, _ := strconv.ParseInt(record[1], 10, 64)
        weight, _ := strconv.ParseFloat(record[2], 64)
        
        weightedEdge := simple.WeightedEdge{F: simple.Node(id_from), T: simple.Node(id_to), W: weight}
        g.SetWeightedEdge(weightedEdge)
    }

    // calculate ego graph
    ego := path.EgoFrom(g.Node(11), g, float64(100))
    edges := ego.WeightedEdges()

    // export to weighted edgelist
    fp2, err := os.Create("./ego-11.weighted.edgelist")
    if err != nil {
        panic(err)
    }
    defer fp2.Close()

    writer := csv.NewWriter(fp2)
    writer.Comma = '\t'

    for edges.Next() {
        var record []string
        record = append(record, fmt.Sprint(edges.WeightedEdge().From()))
        record = append(record, fmt.Sprint(edges.WeightedEdge().To()))
        record = append(record, fmt.Sprint(edges.WeightedEdge().Weight()))
        writer.Write(record)
    }
    writer.Flush()
    
}