-
Notifications
You must be signed in to change notification settings - Fork 0
/
Copy pathgather_collision_data.py
97 lines (81 loc) · 3.16 KB
/
gather_collision_data.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
from matplotlib import pyplot as plt
import numpy as np
from collison_detection import move, confine_particles, handle_collisions_given_indices, \
get_collision_indices_q_tree, reformat_indices, get_adjacency_list
x_max = y_max = 20
n_particles = 300
radius = 0.5
positions = np.random.random((n_particles, 2))
positions[:, 0] *= x_max
positions[:, 1] *= y_max
angles = np.random.random(n_particles) * 2*np.pi
speed = 0.1
v = np.zeros((n_particles, 2))
v[:, 0] = np.cos(angles) * speed
v[:, 1] = np.sin(angles) * speed
colormap = plt.cm.get_cmap('plasma')
colors = [colormap(0.2 + 0.6 * i/(n_particles-1)) for i in range(n_particles)]
fig, ax = plt.subplots()
ax.set_xlim(0, x_max)
ax.set_ylim(0, y_max)
ax.set_aspect('equal', adjustable='box')
scatter_plot = ax.scatter([None]*n_particles, [None]*n_particles, s=750*radius**2, color=colors)
fig.canvas.draw()
plt.show(block=False)
n_train = 5000
x_train = [] # Features: (x, y, v_x, v_y)
edge_index = []
y_train = [] # Output: (v_x, v_y)
x_train_detection = []
y_train_detection = []
while len(y_train) < n_train:
confine_particles(positions, v, x_max, y_max, radius)
features = np.append(positions, v, axis=1)
features[:, 0] /= x_max
features[:, 1] /= y_max
features[:, 2:] /= speed
collision_indices = get_collision_indices_q_tree(positions, radius, [x_max, y_max])
if collision_indices.size > 0:
collision_indices = collision_indices.astype('intc')
handle_collisions_given_indices(positions, v, radius, collision_indices)
collision_indices = collision_indices.T
unique_indices = np.unique(collision_indices)
features = features[unique_indices]
outputs = v[unique_indices] / speed
edge_indices = reformat_indices(get_adjacency_list(collision_indices))
edge_index.append(edge_indices)
x_train.append(features)
y_train.append(outputs)
x_train_detection.append(positions)
y_detect = np.zeros(len(positions))
y_detect[unique_indices] = 1
y_train_detection.append(y_detect)
print(len(y_train))
move(positions, v)
scatter_plot.set_offsets(positions)
fig.canvas.draw()
fig.canvas.flush_events()
np.savez('edge_index.npz', *edge_index)
np.savez('x_train.npz', *x_train)
np.savez('y_train.npz', *y_train)
np.savez('x_train_detection.npz', *x_train_detection)
np.savez('y_train_detection.npz', *y_train_detection)
'''while len(y_train) < n_train:
confine_particles(positions, v, x_max, y_max, radius)
features = np.append(positions, v, axis=1)
features[:, 0] /= x_max
features[:, 1] /= y_max
features[:, 2:] /= speed
collision_indices = handle_collisions(positions, v, radius)
if collision_indices.size > 0:
collision_indices = collision_indices.astype('intc')
features = features[collision_indices[0, :], :]
outputs = v[collision_indices[0, :], :] / speed
edge_index.append(np.argsort(np.argsort(collision_indices, axis=1), axis=1))
x_train.append(features)
y_train.append(outputs)
print(len(y_train))
move(positions, v)
scatter_plot.set_offsets(positions)
fig.canvas.draw()
fig.canvas.flush_events()'''