nodebpy 0.3.0__py3-none-any.whl → 0.4.0__py3-none-any.whl

nodebpy/lib/nodearrange/arrange/ranking.py

@@ -0,0 +1,256 @@
+# SPDX-License-Identifier: GPL-2.0-or-later
+
+# http://dx.doi.org/10.1109/32.221135
+
+from __future__ import annotations
+
+from functools import cache
+from math import sqrt
+from typing import TYPE_CHECKING
+
+import networkx as nx
+
+from ..utils import group_by
+from .graph import Kind, MultiEdge, Node, opposite
+
+if TYPE_CHECKING:
+    from .sugiyama import ClusterGraph
+
+
+# https://api.semanticscholar.org/CorpusID:14932050
+def get_nesting_graph(CG: ClusterGraph) -> nx.MultiDiGraph[Node]:
+    H = CG.G.copy()
+    for u, v in CG.T.edges:
+        if u.type == Kind.CLUSTER:
+            if v.type != Kind.CLUSTER:
+                H.add_edges_from(((u.left, v), (v, u.right)))
+            else:
+                H.add_edges_from(((u.left, v.left), (v.right, u.right)))
+
+    return H
+
+
+@cache
+def get_adj_edges_H(H: nx.MultiDiGraph[Node], v: Node) -> tuple[MultiEdge, ...]:
+    return (*H.in_edges(v, keys=True), *H.out_edges(v, keys=True))
+
+
+@cache
+def get_adj_edges_T(T: nx.MultiDiGraph[Node], v: Node) -> tuple[MultiEdge, ...]:
+    return (*T.in_edges(v, keys=True), *T.out_edges(v, keys=True))
+
+
+def get_slack(e: MultiEdge) -> int:
+    u, v, _ = e
+    min_length = 1
+    return v.rank - u.rank - min_length
+
+
+def tight_tree(
+  H: nx.MultiDiGraph[Node],
+  T: nx.MultiDiGraph[Node],
+  v: Node,
+  visited: set[MultiEdge] | None = None,
+) -> int:
+    if visited is None:
+        visited = set()
+
+    T.add_node(v)
+
+    for e in get_adj_edges_H(H, v):
+        if e in visited:
+            continue
+
+        visited.add(e)
+
+        u, w, _ = e
+        other = u if v != u else w
+        if e in T.edges:
+            tight_tree(H, T, other, visited)
+        elif other not in T and get_slack(e) == 0:
+            T.add_edge(*e)
+            tight_tree(H, T, other, visited)
+
+    return len(T)
+
+
+def set_post_order_numbers(v: Node, T: nx.MultiDiGraph[Node]) -> None:
+    visited = set()
+    num = 0
+
+    def recurse(w: Node) -> int:
+        nums = []
+        for e in get_adj_edges_T(T, w):
+            if e in visited:
+                continue
+
+            visited.add(e)
+            nums.append(recurse(opposite(w, e)))
+
+        nonlocal num
+        w.po_num = num
+        w.lowest_po_num = min(nums + [num])
+        num += 1
+        return w.lowest_po_num
+
+    recurse(v)
+
+
+def compute_cut_values(H: nx.MultiDiGraph[Node], T: nx.MultiDiGraph[Node]) -> None:
+    unknown_cut_values = {}
+    leaves = []
+    for v in H:
+        adj_edges = get_adj_edges_T(T, v)
+        unknown_cut_values[v] = list(adj_edges)
+        if len(adj_edges) == 1:
+            leaves.append(v)
+
+    for v in leaves:
+        while len(unknown_cut_values[v]) == 1:
+            to_determine = unknown_cut_values[v][0]
+            d = T.edges[to_determine]
+            d['cut_value'] = H.edges[to_determine]['weight']
+            u, w, _ = to_determine
+            for e in get_adj_edges_H(H, v):
+                if e == to_determine:
+                    continue
+
+                weight = H.edges[e]['weight']
+                if e in T.edges:
+                    if u == e[0] or w == e[1]:
+                        d['cut_value'] -= T.edges[e]['cut_value'] - weight
+                    else:
+                        d['cut_value'] += T.edges[e]['cut_value'] - weight
+                else:
+                    if (v == u and e[0] != v) or (v != u and e[0] == v):
+                        weight = -weight
+                    d['cut_value'] += weight
+
+            unknown_cut_values[u].remove(to_determine)
+            unknown_cut_values[w].remove(to_determine)
+            v = w if u == v else u
+
+
+def feasible_tree(H: nx.MultiDiGraph[Node]) -> nx.MultiDiGraph[Node]:
+    generations = nx.topological_generations(nx.reverse_view(H))  # type: ignore
+    for i, col in enumerate(reversed(tuple(generations))):
+        for v in col:
+            v.rank = i
+
+    T = nx.MultiDiGraph()
+    v_root = next(iter(H))
+
+    while tight_tree(H, T, v_root) < len(H):
+        incident_edges = [(u, v, k) for u, v, k in H.edges(keys=True) if (u in T) ^ (v in T)]
+        e = min(incident_edges, key=get_slack)
+        slack = -get_slack(e) if e[1] in T else get_slack(e)
+        for v in T:
+            v.rank += slack
+
+    set_post_order_numbers(v_root, T)
+    compute_cut_values(H, T)
+
+    return T
+
+
+def leave_edge(T: nx.MultiDiGraph[Node]) -> MultiEdge | None:
+    return next(((u, v, k) for u, v, k, c in T.edges.data('cut_value', keys=True) if c < 0), None)
+
+
+def is_in_head(v: Node, e: MultiEdge) -> bool:
+    u, w, _ = e
+
+    if u.lowest_po_num <= v.po_num and v.po_num <= u.po_num and w.lowest_po_num <= v.po_num and v.po_num <= w.po_num:
+        return u.po_num >= w.po_num
+
+    return u.po_num < w.po_num
+
+
+def enter_edge(H: nx.MultiDiGraph[Node], e: MultiEdge) -> MultiEdge:
+    edges = [f for f in H.edges(keys=True) if is_in_head(f[0], e) and not is_in_head(f[1], e)]
+    return min(edges, key=get_slack)
+
+
+def exchange(
+  H: nx.MultiDiGraph[Node],
+  T: nx.MultiDiGraph[Node],
+  leave: MultiEdge,
+  enter: MultiEdge,
+) -> None:
+    T.remove_edge(*leave)
+    T.add_edge(*enter)
+
+    slack = get_slack(enter)
+    if not is_in_head(enter[1], leave):
+        slack = -slack
+
+    for v in H:
+        if not is_in_head(v, leave):
+            v.rank += slack
+
+    get_adj_edges_T.cache_clear()
+
+    set_post_order_numbers(v, T)
+    compute_cut_values(H, T)
+
+
+def normalize_and_balance(CG: ClusterGraph, H: nx.DiGraph[Node]) -> None:
+    for cc in nx.weakly_connected_components(CG.G):
+        c = next(iter(cc)).cluster
+        assert c
+
+        if any(v.cluster != c for v in cc):
+            continue
+
+        ranked = group_by(cc, key=lambda v: v.rank, sort=True)
+
+        if c.node:
+            start = min(v.rank for v in CG.T[c] if v.type != Kind.CLUSTER)
+        else:
+            start = c.left.rank - (max(ranked.values()) - min(ranked.values()))
+
+        for i, col in enumerate(ranked, start):
+            for v in col:
+                v.rank = i
+
+    col_sizes = []
+    for i, col in enumerate(group_by(H, key=lambda v: v.rank, sort=True)):
+        col_sizes.append(len(col))
+        for v in col:
+            v.rank = i
+
+    for v in H:
+        if len(H.in_edges(v)) != len(H.out_edges(v)):
+            continue
+
+        start = v.rank - min([v.rank - u.rank for u in H.pred[v]], default=-1) + 1
+        stop = v.rank + min([w.rank - v.rank for w in H[v]], default=-1)
+        new_rank = max(range(start, stop), key=lambda i: col_sizes[i], default=v.rank)
+
+        if col_sizes[new_rank] < col_sizes[v.rank]:
+            col_sizes[v.rank] -= 1
+            col_sizes[new_rank] += 1
+            v.rank = new_rank
+
+
+_BASE_ITER_LIMIT = 50
+
+
+def compute_ranks(CG: ClusterGraph) -> None:
+    for i, layer in enumerate(nx.topological_generations(CG.T)):
+        for c in CG.S.intersection(layer):
+            c.nesting_level = i
+
+    H = get_nesting_graph(CG)
+    nx.set_edge_attributes(H, 1, 'weight')  # type: ignore
+
+    T = feasible_tree(H)
+    i = 0
+    iter_limit = _BASE_ITER_LIMIT * sqrt(len(H))
+    while (e := leave_edge(T)) and i < iter_limit:
+        exchange(H, T, e, enter_edge(H, e))
+        i += 1
+
+    root = next(c for c in CG.S if not CG.T.pred[c])
+    H.remove_nodes_from((root.left, root.right))
+    normalize_and_balance(CG, H)

nodebpy/lib/nodearrange/arrange/realize.py

@@ -0,0 +1,231 @@
+# SPDX-License-Identifier: GPL-2.0-or-later
+
+from __future__ import annotations
+
+from itertools import chain
+from math import isclose
+from statistics import fmean
+
+import networkx as nx
+from bpy.types import Node as BlenderNode
+from mathutils import Vector
+
+from .. import config
+from ..utils import get_ntree, move
+from .graph import (
+  FROM_SOCKET,
+  TO_SOCKET,
+  Cluster,
+  ClusterGraph,
+  Kind,
+  Node,
+  Socket,
+  add_dummy_edge,
+  get_reroute_paths,
+  is_real,
+  socket_graph,
+)
+
+
+def is_safe_to_remove(v: Node) -> bool:
+    if not is_real(v):
+        return True
+
+    if v.node.label:
+        return False
+
+    for val in config.multi_input_sort_ids.values():
+        if any(v == i[0].owner for i in val):
+            return False
+
+    return all(
+      s.node.select for s in chain(
+      config.linked_sockets[v.node.inputs[0]],
+      config.linked_sockets[v.node.outputs[0]],
+      ))
+
+
+def dissolve_reroute_edges(G: nx.DiGraph[Node], path: list[Node]) -> None:
+    if not G[path[-1]]:
+        return
+
+    try:
+        u, _, o = next(iter(G.in_edges(path[0], data=FROM_SOCKET)))
+    except StopIteration:
+        return
+
+    succ_inputs = [e[2] for e in G.out_edges(path[-1], data=TO_SOCKET)]
+
+    # Check if a reroute has been used to link the same output to the same multi-input multiple
+    # times
+    for *_, d in G.out_edges(u, data=True):
+        if d[FROM_SOCKET] == o and d[TO_SOCKET] in succ_inputs:
+            path.clear()
+            return
+
+    links = get_ntree().links
+    for i in succ_inputs:
+        G.add_edge(u, i.owner, from_socket=o, to_socket=i)
+        links.new(o.bpy, i.bpy)
+
+
+def remove_reroutes(CG: ClusterGraph) -> None:
+    reroute_clusters = {#
+      c for c in CG.S
+      if all(v.type != Kind.CLUSTER and v.is_reroute for v in CG.T[c])}
+    for path in get_reroute_paths(CG, is_safe_to_remove):
+        if path[0].cluster in reroute_clusters:
+            if len(path) > 2:
+                u, *between, v = path
+                add_dummy_edge(CG.G, u, v)
+                CG.remove_nodes_from(between)
+        else:
+            dissolve_reroute_edges(CG.G, path)
+            CG.remove_nodes_from(path)
+
+
+_Y_TOL = 5
+
+
+def simplify_path(CG: ClusterGraph, path: list[Node]) -> None:
+    G = CG.G
+    pred_output = lambda w: next(iter(G.in_edges(w, data=FROM_SOCKET)))[2]
+    succ_input = lambda w: next(iter(G.out_edges(w, data=TO_SOCKET)))[2]
+
+    if len(path) == 1:
+        v = path[0]
+
+        if not G.pred[v] or G.out_degree[v] != 1 or v.col is None or is_real(v):
+            return
+
+        p = pred_output(v)
+        q = succ_input(v)
+        if isclose(p.y, q.y, rel_tol=0, abs_tol=_Y_TOL):
+            G.add_edge(p.owner, q.owner, from_socket=p, to_socket=q)
+            CG.remove_nodes_from(path)
+            path.clear()
+
+        return
+
+    u, *between, v = path
+
+    if G.pred[u] and isclose((p := pred_output(u)).y, u.y, rel_tol=0, abs_tol=_Y_TOL):
+        between.append(u)
+    else:
+        p = Socket(u, 0, True)
+
+    if G.out_degree[v] == 1 and isclose(v.y, (q := succ_input(v)).y, rel_tol=0, abs_tol=_Y_TOL):
+        between.append(v)
+    else:
+        q = Socket(v, 0, False)
+
+    if p.owner != u or q.owner != v or between:
+        G.add_edge(p.owner, q.owner, from_socket=p, to_socket=q)
+
+    CG.remove_nodes_from(between)
+    for v in between:
+        path.remove(v)
+
+
+def add_reroute(v: Node) -> None:
+    reroute = get_ntree().nodes.new(type='NodeReroute')
+    assert v.cluster
+    reroute.parent = v.cluster.node
+    config.selected.append(reroute)
+    v.node = reroute
+    v.type = Kind.NODE
+
+
+def realize_edges(G: nx.DiGraph[Node]) -> None:
+    links = get_ntree().links
+    for u, v, d in G.edges.data():
+        if u.is_reroute or v.is_reroute:
+            links.new(d[FROM_SOCKET].bpy, d[TO_SOCKET].bpy)
+
+
+def realize_dummy_nodes(CG: ClusterGraph) -> None:
+    for path in get_reroute_paths(CG, is_safe_to_remove, aligned=True):
+        simplify_path(CG, path)
+
+        for v in path:
+            if not is_real(v):
+                add_reroute(v)
+
+    realize_edges(CG.G)
+
+
+def restore_multi_input_orders(G: nx.MultiDiGraph[Node]) -> None:
+    links = get_ntree().links
+    H = socket_graph(G)
+    for socket, sort_ids in config.multi_input_sort_ids.items():
+        multi_input = socket.bpy
+        assert multi_input
+
+        as_links = {l.from_socket: l for l in links if l.to_socket == multi_input}
+
+        for output in {s.bpy for s in H.pred[socket]} - as_links.keys():
+            assert output
+            as_links[output] = links.new(output, multi_input)
+
+        if len(as_links) != len({l.multi_input_sort_id for l in as_links.values()}):
+            for link in as_links.values():
+                links.remove(link)
+
+            for output in as_links:
+                as_links[output] = links.new(output, multi_input)
+
+        SH = H.subgraph({i[0] for i in sort_ids} | {socket} | {v for v in H if v.owner.is_reroute})
+        seen = set()
+        for base_from_socket, sort_id in sort_ids:
+            other = min(as_links.values(), key=lambda l: abs(l.multi_input_sort_id - sort_id))
+            from_socket = next(
+              s for s, t in nx.edge_dfs(SH, base_from_socket) if t == socket and s not in seen)
+            as_links[from_socket.bpy].swap_multi_input_sort_id(other)  # type: ignore
+            seen.add(from_socket)
+
+
+def realize_locations(G: nx.DiGraph[Node], old_center: Vector) -> None:
+    new_center = (fmean([v.x for v in G]), fmean([v.y for v in G]))
+    offset_x, offset_y = -Vector(new_center) + old_center
+
+    for v in G:
+        assert isinstance(v.node, BlenderNode)
+        assert v.cluster
+
+        # Optimization: avoid using bpy.ops for as many nodes as possible (see `utils.move()`)
+        v.node.parent = None
+
+        x, y = v.node.location
+        v.x += offset_x
+        v.y += offset_y
+        move(v.node, x=v.x - x, y=v.corrected_y() - y)
+
+        v.node.parent = v.cluster.node
+
+
+def resize_unshrunken_frame(CG: ClusterGraph, cluster: Cluster) -> None:
+    frame = cluster.node
+
+    if not frame or frame.shrink:
+        return
+
+    real_children = [v for v in CG.T[cluster] if is_real(v)]
+
+    for v in real_children:
+        v.node.parent = None
+
+    frame.shrink = False
+    frame.shrink = True
+
+    for v in real_children:
+        v.node.parent = frame
+
+
+def realize_layout(CG: ClusterGraph, old_center: Vector) -> None:
+    if config.SETTINGS.add_reroutes:
+        realize_dummy_nodes(CG)
+
+    restore_multi_input_orders(CG.G)
+    realize_locations(CG.G, old_center)
+    for c in CG.S:
+        resize_unshrunken_frame(CG, c)