Add labels as a keyword argument into plot functions (#32)

* Add interaction: display labels on edge hover * Add labels as argument into plot functions * Fix from last commit * Apply @mofeing suggestions from code review Co-authored-by: Sergio Sánchez Ramírez <[email protected]> --------- Co-authored-by: Sergio Sánchez Ramírez <[email protected]>
bsc-quantic · Feb 20, 2023 · f55e951 · f55e951
1 parent 818f7db
commit f55e951
Showing 1 changed file with 9 additions and 3 deletions.
diff --git a/src/Visualization.jl b/src/Visualization.jl
@@ -9,17 +9,17 @@ const MAX_EDGE_WIDTH = 10.
 const MAX_ARROW_SIZE = 25.
 const MAX_NODE_SIZE = 40.
 
-function Makie.plot(path::ContractionPath; colormap = to_colormap(:viridis)[begin:end-10], kwargs...)
+function Makie.plot(path::ContractionPath; kwargs...)
     f = Figure()
 
-    p, ax = plot!(f[1,1], path; colormap, kwargs...)
+    p, ax = plot!(f[1,1], path; kwargs...)
     display(f)
 
     return f, ax, p
 end
 
 # TODO replace `to_colormap(:viridis)[begin:end-10]` with a custom colormap
-function Makie.plot!(f::GridPosition, path::ContractionPath; colormap = to_colormap(:viridis)[begin:end-10], kwargs...)
+function Makie.plot!(f::GridPosition, path::ContractionPath; colormap = to_colormap(:viridis)[begin:end-10], labels = false, kwargs...)
     scene = Scene()
     default_attrs = default_theme(scene, GraphPlot)
 
@@ -39,6 +39,8 @@ function Makie.plot!(f::GridPosition, path::ContractionPath; colormap = to_color
     kwargs[:node_size] = (log_flop/max_flop) * MAX_NODE_SIZE
     kwargs[:node_color] = log_flop
 
+    elabels = [join(OptimizedEinsum.labels(path, i)) for i in 1:ne(graph)]
+
     if haskey(kwargs, :layout) && kwargs[:layout] isa GraphMakie.NetworkLayout.IterativeLayout{3}
         ax = LScene(f[1,1])
     else
@@ -55,6 +57,10 @@ function Makie.plot!(f::GridPosition, path::ContractionPath; colormap = to_color
         node_attr = (colorrange=(min_flop, max_flop),
         # TODO replace `to_colormap(:plasma)[begin:end-50]), kwargs...)` with a custom colormap
         colormap = to_colormap(:plasma)[begin:end-50]),
+        elabels = labels ? elabels : nothing,
+        elabels_color = [:black for i in 1:ne(graph)],
+        # TODO configurable `elabels_textsize`
+        elabels_textsize = [(log_size[i]/max_size) * 5 + 12 for i in 1:ne(graph)],
         kwargs...)
 
     # TODO configurable `labelsize`