def propagate(self, start_state, control, num_steps, integration_step):
'''
Integrate system dynamics
:param start_state: numpy array with the start state for the integration
:param control: numpy array with constant controls to be applied during integration
:param num_steps: number of steps to integrate
:param integration_step: dt of integration
:return: new state of the system
'''
action = control
delta_state = torch.zeros_like(start_state)
last_state = start_state.clone()
for i in range(num_steps):
init_graph_features(self.G1,
self.graph_feat_size,
self.node_feat_size,
self.edge_feat_size,
cuda=True,
bs=1)
load_graph_features(self.G1,
action,
last_state.clone(),
None,
bs=1,
noise=0,
std=None)
self.gn.eval()
# with torch.no_grad():
G_out = self.gn(self.in_normalizer.normalize(self.G1))
for node in G_out.nodes():
delta_state[0, [node, node +
2]] = G_out.nodes[node]['feat'][0].cpu()
last_state += delta_state
return last_state
in_normalizer = normalizers['in_normalizer']
out_normalizer = normalizers['out_normalizer']
std = in_normalizer.get_std()
for epoch in range(300):
for i, data in tqdm(enumerate(dl), total=len(dset) / 200 + 1):
optimizer.zero_grad()
action, delta_state, last_state = data
action, delta_state, last_state = action.float(),\
delta_state.float(), last_state.float()
if use_cuda:
action, delta_state, last_state = action.cuda(),\
delta_state.cuda(), last_state.cuda()
init_graph_features(G1,
graph_feat_size,
node_feat_size,
edge_feat_size,
cuda=True,
bs=200)
load_graph_features(G1,
action,
last_state,
delta_state,
bs=200,
noise=0.02,
std=std)
G_out = gn(in_normalizer.normalize(G1))
init_graph_features(G_target,
graph_feat_size,
node_feat_size,
edge_feat_size,
if use_cuda:
action, delta_state, last_state = action.cuda(), delta_state.cuda(
), last_state.cuda()
state = last_state[1].unsqueeze(0)
state_gt = last_state[1].unsqueeze(0).clone()
for i in range(1, 50):
print(i)
action_i = action[i].unsqueeze(0)
delta_state_i = delta_state[i].unsqueeze(0)
last_state_i = last_state[i].unsqueeze(0)
init_graph_features(G1,
graph_feat_size,
node_feat_size,
edge_feat_size,
cuda=True,
bs=1)
load_graph_features(G1, action_i, state, None, bs=1, noise=0)
G_out = gn(in_normalizer.normalize(G1))
G_out = out_normalizer.inormalize(G_out)
delta_state_pred = get_graph_features(G_out)
state_gt += delta_state_i
state += delta_state_pred
img = draw_state(state_gt)
img_pred = draw_state(state)
writer.append_data(np.concatenate([img, img_pred], axis=1))