Customize layer assignment problem

src/layering.jl

@@ -1,9 +1,40 @@
 # This contains algorithms for breaking up a DAG into layers
-
-function layer_by_longest_path_to_source(graph)
+"Calculate the layer of each node"
+function layer_by_longest_path_to_source(graph, opt_layer_assign)
     dists = longest_paths(graph, sources(graph))
-    layer_groups = IterTools.groupby(i->dists[i], sort(vertices(graph), by=i->dists[i]))
-    return collect.(layer_groups)
+    layer_groups = collect.(IterTools.groupby(i->dists[i], sort(vertices(graph), by=i->dists[i])))
+    agree_with_opt_layer_assign!(layer_groups, graph, opt_layer_assign)
+    return layer_groups
+end
+
+"Correct the layer of each node, according to the optional parameter by user"
+function agree_with_opt_layer_assign!(layer_groups, graph, opt_layer_assign)
+    if length(opt_layer_assign) > 0
+        key_opt = collect(keys(opt_layer_assign))
+        values_opt = collect(values(opt_layer_assign))
+
+        opt_ids_by_layers = sortperm(collect(values(opt_layer_assign)), rev=true)
+        key_opt = key_opt[opt_ids_by_layers]
+        values_opt = values_opt[opt_ids_by_layers]
+
+        # change layers if needed
+        for (k, l) in zip(key_opt, values_opt)
+            #current layer of node
+            curr_layer = findfirst(k .∈ layer_groups)
+            if !isnothing(curr_layer)
+                if curr_layer > l
+                    @warn "Ignored opt_layer_assign for node $k; curr layer ($curr_layer) > desired layer ($l)"
+                elseif any(has_edge(graph, k, v) for v in vcat(layer_groups[curr_layer:min(l, end)]...))
+                    @warn "opt_layer_assign node $k incompatible with edge order"
+                elseif l > length(layer_groups)
+                    @warn "opt_layer_assign node $k optional layer larger than the available layers"
+                elseif curr_layer != l
+                    filter!(x->x != k, layer_groups[curr_layer])
+                    push!(layer_groups[l], k)
+                end
+            end
+        end
+    end
 end
 
 
@@ -23,12 +54,12 @@ function add_dummy_nodes!(graph, layer2nodes)
     dag_or_error(graph)
     nondummy_nodes = vertices(graph)
     # mapping from edges in the original graph to paths in the graph with dummy nodes
-    edge_to_paths = Dict(e => eltype(graph)[] for e in edges(graph)) 
+    edge_to_paths = Dict(e => eltype(graph)[] for e in edges(graph))
     node2layer = node2layer_lookup(layer2nodes)  # doesn't have dummy nodes, doesn't need them
     for cur_node in vertices(graph)
         cur_layer = node2layer[cur_node]
         for out_edge in filter(e -> src(e) == cur_node, collect(edges(graph)))  # need to copy as outwise will mutate when the graph is mutated
-            out_node = out_edge.dst 
+            out_node = out_edge.dst
             out_layer = node2layer[out_node]
             cur_layer < out_layer || throw(DomainError(node2layer, "Layer assigmenment must be strictly monotonic"))
             path = get!(edge_to_paths, out_edge, eltype(graph)[])
@@ -54,4 +85,4 @@ function add_dummy_nodes!(graph, layer2nodes)
     is_dummy_mask = trues(nv(graph))
     is_dummy_mask[nondummy_nodes] .= false
     return is_dummy_mask, edge_to_paths
-end
+end

src/zarate.jl

@@ -32,6 +32,7 @@ Returns:
  - `xs`: the xs coordinates of vertices in the layout
  - `ys`: the ys coordinates of vertices in the layout
  - `paths`: a Dict which, for each edge in `graph`, contains a Tuple of coordinate vectors (xs, ys).
+ - opt_layer_assign: dictionary specifying the requested custom layer for a given node.
 
 The layout is calculated on a graph where dummy nodes can be added to the different layers.
 As a result, plotting edges as straight lines between two nodes can result in
@@ -52,11 +53,12 @@ for e in edges(g)
 end
 ```
 """
-function solve_positions(layout::Zarate, original_graph)
+function solve_positions(layout::Zarate, original_graph;
+        opt_layer_assign::Dict{Int, Int} = Dict{Int, Int}())
     graph = copy(original_graph)
 
     # 1. Layer Assigment
-    layer2nodes = layer_by_longest_path_to_source(graph)
+    layer2nodes = layer_by_longest_path_to_source(graph, opt_layer_assign)
     is_dummy_mask, edge_to_path = add_dummy_nodes!(graph, layer2nodes)
 
     # 2. Layer Ordering