wehub-resource-sync
180 lines
5.5 KiB
Python
180 lines
5.5 KiB
Python
from __future__ import absolute_import, division, print_function
|
|
|
|
import argparse
|
|
import os
|
|
from math import cos, pi, radians, sin
|
|
|
|
import numpy as np
|
|
|
|
"""
|
|
This adapted from comes from https://github.com/jsikyoon/Interaction-networks_tensorflow
|
|
which generates multi-body dynamic simulation data for Interaction network
|
|
"""
|
|
|
|
# 5 features on the state [mass,x,y,x_vel,y_vel]
|
|
fea_num = 5
|
|
# G stand for Gravity constant 10**5 can help numerical stability
|
|
G = 10**5
|
|
# time step
|
|
diff_t = 0.001
|
|
|
|
|
|
def init(total_state, n_body, fea_num, orbit):
|
|
data = np.zeros((total_state, n_body, fea_num), dtype=float)
|
|
if orbit:
|
|
data[0][0][0] = 100
|
|
data[0][0][1:5] = 0.0
|
|
# The position are initialized randomly.
|
|
for i in range(1, n_body):
|
|
data[0][i][0] = np.random.rand() * 8.98 + 0.02
|
|
distance = np.random.rand() * 90.0 + 10.0
|
|
theta = np.random.rand() * 360
|
|
theta_rad = pi / 2 - radians(theta)
|
|
data[0][i][1] = distance * cos(theta_rad)
|
|
data[0][i][2] = distance * sin(theta_rad)
|
|
data[0][i][3] = (
|
|
-1
|
|
* data[0][i][2]
|
|
/ norm(data[0][i][1:3])
|
|
* (G * data[0][0][0] / norm(data[0][i][1:3]) ** 2)
|
|
* distance
|
|
/ 1000
|
|
)
|
|
data[0][i][4] = (
|
|
data[0][i][1]
|
|
/ norm(data[0][i][1:3])
|
|
* (G * data[0][0][0] / norm(data[0][i][1:3]) ** 2)
|
|
* distance
|
|
/ 1000
|
|
)
|
|
else:
|
|
for i in range(n_body):
|
|
data[0][i][0] = np.random.rand() * 8.98 + 0.02
|
|
distance = np.random.rand() * 90.0 + 10.0
|
|
theta = np.random.rand() * 360
|
|
theta_rad = pi / 2 - radians(theta)
|
|
data[0][i][1] = distance * cos(theta_rad)
|
|
data[0][i][2] = distance * sin(theta_rad)
|
|
data[0][i][3] = np.random.rand() * 6.0 - 3.0
|
|
data[0][i][4] = np.random.rand() * 6.0 - 3.0
|
|
return data
|
|
|
|
|
|
def norm(x):
|
|
return np.sqrt(np.sum(x**2))
|
|
|
|
|
|
def get_f(reciever, sender):
|
|
diff = sender[1:3] - reciever[1:3]
|
|
distance = norm(diff)
|
|
if distance < 1:
|
|
distance = 1
|
|
return G * reciever[0] * sender[0] / (distance**3) * diff
|
|
|
|
|
|
# Compute stat according to the paper for normalization
|
|
def compute_stats(train_curr):
|
|
data = np.vstack(train_curr).reshape(-1, fea_num)
|
|
stat_median = np.median(data, axis=0)
|
|
stat_max = np.quantile(data, 0.95, axis=0)
|
|
stat_min = np.quantile(data, 0.05, axis=0)
|
|
return stat_median, stat_max, stat_min
|
|
|
|
|
|
def calc(cur_state, n_body):
|
|
next_state = np.zeros((n_body, fea_num), dtype=float)
|
|
f_mat = np.zeros((n_body, n_body, 2), dtype=float)
|
|
f_sum = np.zeros((n_body, 2), dtype=float)
|
|
acc = np.zeros((n_body, 2), dtype=float)
|
|
for i in range(n_body):
|
|
for j in range(i + 1, n_body):
|
|
if j != i:
|
|
f = get_f(cur_state[i][:3], cur_state[j][:3])
|
|
f_mat[i, j] += f
|
|
f_mat[j, i] -= f
|
|
f_sum[i] = np.sum(f_mat[i], axis=0)
|
|
acc[i] = f_sum[i] / cur_state[i][0]
|
|
next_state[i][0] = cur_state[i][0]
|
|
next_state[i][3:5] = cur_state[i][3:5] + acc[i] * diff_t
|
|
next_state[i][1:3] = cur_state[i][1:3] + next_state[i][3:5] * diff_t
|
|
return next_state
|
|
|
|
|
|
# The state is [mass,pos_x,pos_y,vel_x,vel_y]* n_body
|
|
def gen(n_body, num_steps, orbit):
|
|
# initialization on just first state
|
|
data = init(num_steps, n_body, fea_num, orbit)
|
|
for i in range(1, num_steps):
|
|
data[i] = calc(data[i - 1], n_body)
|
|
return data
|
|
|
|
|
|
if __name__ == "__main__":
|
|
argparser = argparse.ArgumentParser()
|
|
argparser.add_argument("--num_bodies", type=int, default=6)
|
|
argparser.add_argument("--num_traj", type=int, default=10)
|
|
argparser.add_argument("--steps", type=int, default=1000)
|
|
argparser.add_argument("--data_path", type=str, default="data")
|
|
|
|
args = argparser.parse_args()
|
|
if not os.path.exists(args.data_path):
|
|
os.mkdir(args.data_path)
|
|
|
|
# Generate data
|
|
data_curr = []
|
|
data_next = []
|
|
|
|
for i in range(args.num_traj):
|
|
raw_traj = gen(args.num_bodies, args.steps, True)
|
|
data_curr.append(raw_traj[:-1])
|
|
data_next.append(raw_traj[1:])
|
|
print("Train Traj: ", i)
|
|
|
|
# Compute normalization statistic from data
|
|
stat_median, stat_max, stat_min = compute_stats(data_curr)
|
|
data = np.vstack(data_curr)
|
|
label = np.vstack(data_next)[:, :, 3:5]
|
|
shuffle_idx = np.arange(data.shape[0])
|
|
np.random.shuffle(shuffle_idx)
|
|
train_split = int(0.9 * data.shape[0])
|
|
valid_split = train_split + 300
|
|
data = data[shuffle_idx]
|
|
label = label[shuffle_idx]
|
|
|
|
train_data = data[:train_split]
|
|
train_label = label[:train_split]
|
|
|
|
valid_data = data[train_split:valid_split]
|
|
valid_label = label[train_split:valid_split]
|
|
|
|
test_data = data[valid_split:]
|
|
test_label = label[valid_split:]
|
|
|
|
np.savez(
|
|
args.data_path + "/n_body_train.npz",
|
|
data=train_data,
|
|
label=train_label,
|
|
n_particles=args.num_bodies,
|
|
median=stat_median,
|
|
max=stat_max,
|
|
min=stat_min,
|
|
)
|
|
|
|
np.savez(
|
|
args.data_path + "/n_body_valid.npz",
|
|
data=valid_data,
|
|
label=valid_label,
|
|
n_particles=args.num_bodies,
|
|
)
|
|
|
|
test_traj = gen(args.num_bodies, args.steps, True)
|
|
|
|
np.savez(
|
|
args.data_path + "/n_body_test.npz",
|
|
data=test_data,
|
|
label=test_label,
|
|
n_particles=args.num_bodies,
|
|
first_frame=test_traj[0],
|
|
test_traj=test_traj,
|
|
)
|