Merge pull request #15 from davide-f/change-from-SimpleWeightedGraph-…

…to-MetaGraph

Change from simple weighted graph to meta graph

Project.toml

@@ -4,20 +4,16 @@ authors = ["Daniel Schwabeneder <[email protected]> and contributors"]
 version = "0.1.3"
 
 [deps]
+Graphs = "86223c79-3864-5bf0-83f7-82e725a168b6"
 LayeredLayouts = "f4a74d36-062a-4d48-97cd-1356bad1de4e"
-LightGraphs = "093fc24a-ae57-5d10-9952-331d41423f4d"
+MetaGraphs = "626554b9-1ddb-594c-aa3c-2596fe9399a5"
 Plots = "91a5bcdd-55d7-5caf-9e0b-520d859cae80"
 RecipesBase = "3cdcf5f2-1ef4-517c-9805-6587b60abb01"
-SimpleWeightedGraphs = "47aef6b3-ad0c-573a-a1e2-d07658019622"
 SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 
 [compat]
-LayeredLayouts = "0.1.1"
-LightGraphs = "1"
 Plots = "1"
 RecipesBase = "1"
-SimpleWeightedGraphs = "1"
-julia = "1.5"
 
 [extras]
 FileIO = "5789e2e9-d7fb-5bc7-8068-2c6fae9b9549"

src/SankeyPlots.jl

@@ -1,17 +1,16 @@
 module SankeyPlots
 
 using LayeredLayouts
-using LightGraphs
 using Plots
 using RecipesBase
-using SimpleWeightedGraphs
+using Graphs, MetaGraphs
 using SparseArrays
 
 export sankey, sankey!
 
 """
     sankey(src, dst, weights; kwargs..., plotattributes...)
-    sankey(g::SimpleWeightedDiGraph; kwargs..., plotattributes...)
+    sankey(g::AbstractMetaGraph; kwargs..., plotattributes...)
 
 Plot a sankey diagram.
 
@@ -79,8 +78,8 @@ In addition to [Plots.jl attributes](http://docs.juliaplots.org/latest/attribute
             for (j, w) in enumerate(vertices(g))
                 if has_edge(g, v, w)
                     y_src = y[i] + h - src_offsets[i, j]
-                    h_edge = g.weights[j, i] / (2m)
-
+                    edge_it = Edge(v, w)
+                    h_edge = get_prop(g, edge_it, :weight) / (2m)
 
                     sankey_y = Float64[]
                     x_start = x[i] + 0.1
@@ -190,8 +189,41 @@ In addition to [Plots.jl attributes](http://docs.juliaplots.org/latest/attribute
     ()
 end
 
-sankey_graph(src, dst, w) = SimpleWeightedDiGraph(src, dst, w)
-sankey_graph(g::SimpleWeightedDiGraph) = copy(g)
+"Function to add a weighted edge"
+function add_weighted_edge!(g, src, dst, weight)
+    new_edge = Edge(src, dst)
+    add_edge!(g, new_edge)
+    set_prop!(g, new_edge, :weight, weight)
+    return g
+end
+
+"""
+Function to create a MetaGraphs from source nodes (src), destination nodes (dst) and weights (w)
+"""
+function sankey_graph(src::Vector{T}, dst::Vector{T}, w) where {T <: Integer}
+    # get list of unique nodes
+    unique_nodes = sort(unique([src; dst]))
+
+    n_nodes = length(unique_nodes)
+    n_edges = length(src)
+    list_nodes = 1:n_nodes
+
+    # verify that the node numbers is appropriate
+    @assert(unique_nodes==unique_nodes)
+    # verify length of vectors
+    @assert(length(src)==length(dst)==length(w))
+
+    # initialize graph
+    g = MetaDiGraph(n_nodes)
+
+    # add edges iteratively
+    for i = 1:n_edges
+        add_weighted_edge!(g, src[i], dst[i], w[i])
+    end
+
+    return g
+end
+sankey_graph(g::AbstractMetaGraph) = copy(g)
 sankey_graph(args...) = error("Check `?sankey` for supported signatures.")
 
 sankey_names(g, names) = names
@@ -203,17 +235,18 @@ function sankey_layout!(g)
     for (edge, path) in paths
         s = edge.src
         px, py = path
+        weight_path = get_prop(g, edge, :weight)
         if length(px) > 2
             for i in 2:length(px)-1
                 add_vertex!(g)
                 v = last(vertices(g))
-                add_edge!(g, s, v, edge.weight)
+                add_weighted_edge!(g, s, v, weight_path)
                 push!(xs, px[i])
                 push!(ys, py[i])
                 push!(mask, true)
                 s = v
             end
-            add_edge!(g, s, edge.dst, edge.weight)
+            add_weighted_edge!(g, s, edge.dst, weight_path)
             rem_edge!(g, edge)
         end
     end
@@ -222,8 +255,8 @@ end
 
 function vertex_weight(g, v)
     max(
-        sum0(weight, Iterators.filter(e -> src(e) == v, edges(g))),
-        sum0(weight, Iterators.filter(e -> dst(e) == v, edges(g))),
+        sum0(x->get_prop(g, x, :weight), Iterators.filter(e -> src(e) == v, edges(g))),
+        sum0(x->get_prop(g, x, :weight), Iterators.filter(e -> dst(e) == v, edges(g))),
     )
 end
 sum0(f, x) = isempty(x) ? 0.0 : sum(f, x)
@@ -242,7 +275,9 @@ function get_src_offsets(g, perm)
                 if offset > 0
                     p[i, j] = offset
                 end
-                offset += g.weights[j, i]
+                edge_it = Edge(v, verts[j])
+                # add to offset if edge is available
+                offset += get_prop(g, edge_it, :weight)
             end
         end
     end
@@ -260,7 +295,7 @@ function get_dst_offsets(g, perm)
                 if offset > 0
                     p[i, j] = offset
                 end
-                offset += g.weights[i, j]
+                offset += get_prop(g, Edge(verts[j], i), :weight)
             end
         end
     end