-
Notifications
You must be signed in to change notification settings - Fork 0
/
mynode.py
185 lines (146 loc) · 7.12 KB
/
mynode.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
import bpy
import sys
import math
def set_socket_value_range(socket: bpy.types.NodeSocket,
default_value: float = 0.0,
min_value: float = 0.0,
max_value: float = 1.0) -> None:
assert socket.type == "VALUE"
socket.default_value = default_value
socket.min_value = min_value
socket.max_value = max_value
def clean_nodes(nodes: bpy.types.Nodes) -> None:
for node in nodes:
nodes.remove(node)
def arrange_nodes(node_tree: bpy.types.NodeTree, verbose: bool = False) -> None:
max_num_iters = 2000
epsilon = 1e-05
target_space = 50.0
second_stage = False
fix_horizontal_location = True
fix_vertical_location = True
fix_overlaps = True
if verbose:
print("-----------------")
print("Target nodes:")
for node in node_tree.nodes:
print("- " + node.name)
# In the first stage, expand nodes overly
target_space *= 2.0
# Gauss-Seidel-style iterations
previous_squared_deltas_sum = sys.float_info.max
for i in range(max_num_iters):
squared_deltas_sum = 0.0
if fix_horizontal_location:
for link in node_tree.links:
k = 0.9 if not second_stage else 0.5
threshold_factor = 2.0
x_from = link.from_node.location[0]
x_to = link.to_node.location[0]
w_from = link.from_node.width
signed_space = x_to - x_from - w_from
C = signed_space - target_space
grad_C_x_from = -1.0
grad_C_x_to = 1.0
# Skip if the distance is sufficiently large
if C >= target_space * threshold_factor:
continue
lagrange = C / (grad_C_x_from * grad_C_x_from + grad_C_x_to * grad_C_x_to)
delta_x_from = -lagrange * grad_C_x_from
delta_x_to = -lagrange * grad_C_x_to
link.from_node.location[0] += k * delta_x_from
link.to_node.location[0] += k * delta_x_to
squared_deltas_sum += k * k * (delta_x_from * delta_x_from + delta_x_to * delta_x_to)
if fix_vertical_location:
k = 0.5 if not second_stage else 0.05
socket_offset = 20.0
def get_from_socket_index(node: bpy.types.Node, node_socket: bpy.types.NodeSocket) -> int:
for i in range(len(node.outputs)):
if node.outputs[i] == node_socket:
return i
assert False
def get_to_socket_index(node: bpy.types.Node, node_socket: bpy.types.NodeSocket) -> int:
for i in range(len(node.inputs)):
if node.inputs[i] == node_socket:
return i
assert False
for link in node_tree.links:
from_socket_index = get_from_socket_index(link.from_node, link.from_socket)
to_socket_index = get_to_socket_index(link.to_node, link.to_socket)
y_from = link.from_node.location[1] - socket_offset * from_socket_index
y_to = link.to_node.location[1] - socket_offset * to_socket_index
C = y_from - y_to
grad_C_y_from = 1.0
grad_C_y_to = -1.0
lagrange = C / (grad_C_y_from * grad_C_y_from + grad_C_y_to * grad_C_y_to)
delta_y_from = -lagrange * grad_C_y_from
delta_y_to = -lagrange * grad_C_y_to
link.from_node.location[1] += k * delta_y_from
link.to_node.location[1] += k * delta_y_to
squared_deltas_sum += k * k * (delta_y_from * delta_y_from + delta_y_to * delta_y_to)
if fix_overlaps and second_stage:
k = 0.9
margin = 0.5 * target_space
# Examine all node pairs
for node_1 in node_tree.nodes:
for node_2 in node_tree.nodes:
if node_1 == node_2:
continue
x_1 = node_1.location[0]
x_2 = node_2.location[0]
w_1 = node_1.width
w_2 = node_2.width
cx_1 = x_1 + 0.5 * w_1
cx_2 = x_2 + 0.5 * w_2
rx_1 = 0.5 * w_1 + margin
rx_2 = 0.5 * w_2 + margin
# Note: "dimensions" and "height" may not be correct depending on the situation
def get_height(node: bpy.types.Node) -> float:
if node.dimensions.y > epsilon:
return node.dimensions.y
elif math.fabs(node.height - 100.0) > epsilon:
return node.height
else:
return 200.0
y_1 = node_1.location[1]
y_2 = node_2.location[1]
h_1 = get_height(node_1)
h_2 = get_height(node_2)
cy_1 = y_1 - 0.5 * h_1
cy_2 = y_2 - 0.5 * h_2
ry_1 = 0.5 * h_1 + margin
ry_2 = 0.5 * h_2 + margin
C_x = math.fabs(cx_1 - cx_2) - (rx_1 + rx_2)
C_y = math.fabs(cy_1 - cy_2) - (ry_1 + ry_2)
# If no collision, just skip
if C_x >= 0.0 or C_y >= 0.0:
continue
# Solve collision for the "easier" direction
if C_x > C_y:
grad_C_x_1 = 1.0 if cx_1 - cx_2 >= 0.0 else -1.0
grad_C_x_2 = -1.0 if cx_1 - cx_2 >= 0.0 else 1.0
lagrange = C_x / (grad_C_x_1 * grad_C_x_1 + grad_C_x_2 * grad_C_x_2)
delta_x_1 = -lagrange * grad_C_x_1
delta_x_2 = -lagrange * grad_C_x_2
node_1.location[0] += k * delta_x_1
node_2.location[0] += k * delta_x_2
squared_deltas_sum += k * k * (delta_x_1 * delta_x_1 + delta_x_2 * delta_x_2)
else:
grad_C_y_1 = 1.0 if cy_1 - cy_2 >= 0.0 else -1.0
grad_C_y_2 = -1.0 if cy_1 - cy_2 >= 0.0 else 1.0
lagrange = C_y / (grad_C_y_1 * grad_C_y_1 + grad_C_y_2 * grad_C_y_2)
delta_y_1 = -lagrange * grad_C_y_1
delta_y_2 = -lagrange * grad_C_y_2
node_1.location[1] += k * delta_y_1
node_2.location[1] += k * delta_y_2
squared_deltas_sum += k * k * (delta_y_1 * delta_y_1 + delta_y_2 * delta_y_2)
if verbose:
print("Iteration #" + str(i) + ": " + str(previous_squared_deltas_sum - squared_deltas_sum))
# Check the termination conditiion
if math.fabs(previous_squared_deltas_sum - squared_deltas_sum) < epsilon:
if second_stage:
break
else:
target_space = 0.5 * target_space
second_stage = True
previous_squared_deltas_sum = squared_deltas_sum