def seq_rnn_embed(vxs, exs, birnn, return_seqs=False):
"""Embed given sequences using rnn."""
# vxs.shape == (..., S)
# exs.shape == (..., S, E)
assert vxs.shape == exs.shape[:
-1], "Sequence embedding dimensions do not match."
lengths = np.sum(vxs != 0, -1).flatten() # (X,)
seqs = F.reshape(exs, (-1, ) + exs.shape[-2:]) # (X, S, E)
toembed = [
s[..., :l, :] for s, l in zip(F.separate(seqs, 0), lengths) if l != 0
] # Y x [(S1, E), (S2, E), ...]
hs, ys = birnn(None, toembed) # (2, Y, E), Y x [(S1, 2*E), (S2, 2*E), ...]
if return_seqs:
ys = F.pad_sequence(ys) # (Y, S, 2*E)
ys = F.reshape(ys, ys.shape[:-1] + (2, EMBED)) # (Y, S, 2, E)
ys = F.mean(ys, -2) # (Y, S, E)
if ys.shape[0] == lengths.size:
ys = F.reshape(ys, exs.shape) # (..., S, E)
return ys
embeds = np.zeros((lengths.size, vxs.shape[-1], EMBED),
dtype=np.float32) # (X, S, E)
idxs = np.nonzero(lengths) # (Y,)
embeds = F.scatter_add(embeds, idxs, ys) # (X, S, E)
embeds = F.reshape(embeds, exs.shape) # (..., S, E)
return embeds # (..., S, E)
hs = F.mean(hs, 0) # (Y, E)
if hs.shape[0] == lengths.size:
hs = F.reshape(hs, vxs.shape[:-1] + (EMBED, )) # (..., E)
return hs
# Add zero values back to match original shape
embeds = np.zeros((lengths.size, EMBED), dtype=np.float32) # (X, E)
idxs = np.nonzero(lengths) # (Y,)
embeds = F.scatter_add(embeds, idxs, hs) # (X, E)
embeds = F.reshape(embeds, vxs.shape[:-1] + (EMBED, )) # (..., E)
return embeds # (..., E)
def __call__(self, atoms_feat, pair_feat, global_feat,
atom_idx, pair_idx, start_idx, end_idx):
# 1) Pass the Dense layer
a_f_d = self.dense_for_atom(atoms_feat)
p_f_d = self.dense_for_pair(pair_feat)
g_f_d = self.dense_for_global(global_feat)
# 2) Update the edge vector
start_node = a_f_d[start_idx]
end_node = a_f_d[end_idx]
g_f_extend_with_pair_idx = g_f_d[pair_idx]
concat_p_v = functions.concat((p_f_d, start_node, end_node,
g_f_extend_with_pair_idx))
update_p = self.update_for_pair(concat_p_v)
# 3) Update the node vector
# 1. get sum edge feature of all nodes using scatter_add method
zero = self.xp.zeros(a_f_d.shape, dtype=self.xp.float32)
sum_edeg_vec = functions.scatter_add(zero, start_idx, update_p) + \
functions.scatter_add(zero, end_idx, update_p)
# 2. get degree of all nodes using scatter_add method
one = self.xp.ones(p_f_d.shape, dtype=self.xp.float32)
degree = functions.scatter_add(zero, start_idx, one) + \
functions.scatter_add(zero, end_idx, one)
# 3. get mean edge feature of all nodes
mean_edge_vec = sum_edeg_vec / degree
# 4. concating
g_f_extend_with_atom_idx = g_f_d[atom_idx]
concat_a_v = functions.concat((a_f_d, mean_edge_vec,
g_f_extend_with_atom_idx))
update_a = self.update_for_atom(concat_a_v)
# 4) Update the global vector
out_shape = g_f_d.shape
ave_p = get_mean_feat(update_p, pair_idx, out_shape, self.xp)
ave_a = get_mean_feat(update_a, atom_idx, out_shape, self.xp)
concat_g_v = functions.concat((ave_a, ave_p, g_f_d), axis=1)
update_g = self.update_for_global(concat_g_v)
# 5) Skip connection
if self.skip_intermediate:
# Skip intermediate feature, used for first layer.
new_a_f = update_a + a_f_d
new_p_f = update_p + p_f_d
new_g_f = update_g + g_f_d
else:
# Skip input feature, used all layer except first layer.
# input feature must be same dimension with updated feature.
new_a_f = update_a + atoms_feat
new_p_f = update_p + pair_feat
new_g_f = update_g + global_feat
# 6) dropout
if self.dropout_ratio > 0.0:
new_a_f = functions.dropout(new_a_f, ratio=self.dropout_ratio)
new_p_f = functions.dropout(new_p_f, ratio=self.dropout_ratio)
new_g_f = functions.dropout(new_g_f, ratio=self.dropout_ratio)
return new_a_f, new_p_f, new_g_f
def get_mean_feat(feat, idx, out_shape, xp):
"""Return mean node or edge feature in each graph.
This method is the same as average pooling
about node or edge feature in each graph.
"""
zero = xp.zeros(out_shape, dtype=xp.float32)
sum_vec = functions.scatter_add(zero, idx, feat)
one = xp.ones(feat.shape, dtype=xp.float32)
degree = functions.scatter_add(zero, idx, one)
return sum_vec / degree
def forward(self, xinit, env, train_warm_start_idxs):
""" forward function
:param xinit:
:param env:Pendulum_DX
:param train_warm_start_idxs:
:return:
"""
Q = self.xp.zeros((self.n_sc, self.n_sc))
index_diag = self.xp.zeros((self.n_sc, self.n_sc), dtype=bool)
self.xp.fill_diagonal(index_diag, True)
index_not_diag = self.xp.zeros((self.n_sc, self.n_sc), dtype=bool)
index_not_diag[self.xp.tril_indices(self.n_sc, -1)] = True
Q = F.scatter_add(Q, self.xp.tril_indices(self.n_sc, -1), self.lower_without_diag)
diag_q = F.sigmoid(self.learn_q_logit)
Q = F.where(index_diag, diag_q, Q)
Q = Q @ Q.T
p = self.learn_p
# print("p.shape", p.shape)
# print(self.n_sc)
# p = F.concat((p, self.xp.array(0.0).reshape(1,)), axis=0)
# print("Q ", Q)
# print("p", p)
x_mpc, u_mpc = env.mpc_Q(env.true_dx, xinit, Q, p,
u_init=self.xp.transpose(train_warm_start_idxs, axes=(1, 0, 2)))
return x_mpc, u_mpc
def seq_rnn_embed(vxs, exs, rnn_layer, initial_state=None, reverse=False):
"""Embed given sequences using rnn."""
# vxs.shape == (..., S)
# exs.shape == (..., S, E)
# initial_state == (..., E)
assert vxs.shape == exs.shape[:
-1], "Sequence embedding dimensions do not match."
lengths = np.sum(vxs != 0, -1).flatten() # (X,)
seqs = F.reshape(exs, (-1, ) + exs.shape[-2:]) # (X, S, E)
if reverse:
toembed = [
F.flip(s[..., :l, :], -2)
for s, l in zip(F.separate(seqs, 0), lengths) if l != 0
] # Y x [(S1, E), (S2, E), ...]
else:
toembed = [
s[..., :l, :] for s, l in zip(F.separate(seqs, 0), lengths)
if l != 0
] # Y x [(S1, E), (S2, E), ...]
if initial_state is not None:
initial_state = F.reshape(initial_state, (-1, EMBED)) # (X, E)
initial_state = initial_state[None,
np.flatnonzero(lengths)] # (1, Y, E)
hs, ys = rnn_layer(initial_state,
toembed) # (1, Y, E), Y x [(S1, 2*E), (S2, 2*E), ...]
hs = hs[0] # (Y, E)
if hs.shape[0] == lengths.size:
hs = F.reshape(hs, vxs.shape[:-1] + (EMBED, )) # (..., E)
return hs
# Add zero values back to match original shape
embeds = np.zeros((lengths.size, EMBED), dtype=np.float32) # (X, E)
idxs = np.nonzero(lengths) # (Y,)
embeds = F.scatter_add(embeds, idxs, hs) # (X, E)
embeds = F.reshape(embeds, vxs.shape[:-1] + (EMBED, )) # (..., E)
return embeds # (..., E)
def __call__(self, h, edge_index):
# add self node feature
new_h = h
messages = h[edge_index[0]]
new_h = functions.scatter_add(new_h, edge_index[1], messages)
# apply MLP
new_h = self.mlp(new_h)
if self.dropout_ratio > 0.0:
new_h = functions.dropout(new_h, ratio=self.dropout_ratio)
return new_h
def compute_cos_angle(vert,face,theta,xp):
# Obsolite: replaced with compute_cos_angle_sub which is more efficient
cos_angle = Variable(xp.cos(theta))
sin_angle = Variable(xp.sin(theta))
for f in face:
n = len(f)
id_p = xp.array([f[(i-1)%n] for i in range(n+1)])
id = xp.array([f[i%n] for i in range(n+1)])
id_n = xp.array([f[(i+1)%n] for i in range(n)])
L = F.sum((vert[id_p] - vert[id])**2, axis=1)
D = F.sum((vert[id_n] - vert[ id_p[:-1] ])**2, axis=1)
c1 = (L[:n]+L[1:]-D)/(2*F.sqrt(L[:n]*L[1:]))
s1 = F.sqrt(1-c1**2)
# trigonometric addition formula
c0 = cos_angle[f]
s0 = sin_angle[f]
cos_angle = F.scatter_add(cos_angle,f,-c0 + c0*c1 - s0*s1)
sin_angle = F.scatter_add(sin_angle,f,-s0 + c0*s1 + s0*c1)
return cos_angle
def check_forward(self, a_data, b_data):
a = chainer.Variable(a_data)
b = chainer.Variable(b_data)
y = functions.scatter_add(a, self.slices, b)
self.assertEqual(y.data.dtype, numpy.float32)
# Test to make sure that the input values are not changed
numpy.testing.assert_equal(cuda.to_cpu(a.data), self.a_data_original)
a_data_copy = cuda.to_cpu(a_data).copy()
numpy.add.at(a_data_copy, self.slices, cuda.to_cpu(b_data))
numpy.testing.assert_equal(a_data_copy, cuda.to_cpu(y.data))
def check_forward(self, a_data, b_data):
a = chainer.Variable(a_data)
b = chainer.Variable(b_data)
y = functions.scatter_add(a, self.slices, b)
self.assertEqual(y.data.dtype, numpy.float32)
# Test to make sure that the input values are not changed
numpy.testing.assert_equal(cuda.to_cpu(a.data), self.a_data_original)
a_data_copy = cuda.to_cpu(a_data).copy()
numpy.add.at(a_data_copy, self.slices, cuda.to_cpu(b_data))
numpy.testing.assert_equal(a_data_copy, cuda.to_cpu(y.data))
def __call__(self, h, batch, h0=None, is_real_node=None):
# --- Readout part ---
h1 = functions.concat((h, h0), axis=1) if h0 is not None else h
g1 = functions.sigmoid(self.i_layer(h1))
g2 = self.activation(self.j_layer(h1))
g = g1 * g2
# sum along node axis
y = self.xp.zeros((int(batch[-1]) + 1, self.out_dim),
dtype=numpy.float32)
y = functions.scatter_add(y, batch, g)
y = self.activation_agg(y)
if self.concat_n_info:
n_nodes = self.xp.zeros(y.shape[0], dtype=self.xp.float32)
n_nodes = functions.scatter_add(n_nodes, batch,
self.xp.ones(batch.shape[0]))
y = functions.concat((y, n_nodes.reshape((-1, 1))))
return y
def compute_curvature(vert,face,xp):
# Obsolite: replaced with compute_curvature_sub which is more efficient
K = Variable(xp.full((len(vert),), 2*xp.pi))
for f in face:
n = len(f)
id_p = xp.array([f[(i-1)%n] for i in range(n+1)])
id = xp.array([f[i%n] for i in range(n+1)])
id_n = xp.array([f[(i+1)%n] for i in range(n)])
L = F.sum((vert[id_p] - vert[id])**2, axis=1)
D = F.sum((vert[id_n] - vert[ id_p[:-1] ])**2, axis=1)
c1 = (L[:n]+L[1:]-D)/(2*F.sqrt(L[:n]*L[1:]))
K = F.scatter_add(K,f,-F.arccos(c1))
return K
def __call__(self, h, batch, axis=1, **kwargs):
if self.activation is not None:
h = self.activation(h)
else:
h = h
if self.mode == 'sum':
y = self.xp.zeros((batch[-1] + 1, h.shape[1]),
dtype=self.xp.float32)
y = functions.scatter_add(y, batch, h)
else:
raise ValueError('mode {} is not supported'.format(self.mode))
return y
def interpolate_trilinear(grid, lin_ind_frustrum, voxel_coords, img_shape,
frustrum_depth):
xp = chainer.cuda.get_array_module(voxel_coords)
batch, num_feats, height, width, depth = grid.shape
lin_ind_frustrum = lin_ind_frustrum.astype(
"int32") # indexを指定するだけなので勾配を流す必要はない
x_indices = voxel_coords[2, :]
y_indices = voxel_coords[1, :]
z_indices = voxel_coords[0, :]
x0 = x_indices.astype("int32")
y0 = y_indices.astype("int32")
z0 = z_indices.astype("int32")
x1 = x0 + 1
y1 = y0 + 1
z1 = z0 + 1
x1 = xp.clip(x1, 0, width - 1)
y1 = xp.clip(y1, 0, height - 1)
z1 = xp.clip(z1, 0, depth - 1)
x = x_indices - x0
y = y_indices - y0
z = z_indices - z0
# output = torch.zeros(batch, num_feats, img_shape[0]*img_shape[1]*depth).cuda()
output = xp.zeros(
(batch, num_feats, img_shape[0] * img_shape[1] * frustrum_depth),
dtype="float32")
added = grid[:, :, x0, y0, z0] * (1 - x) * (1 - y) * (1 - z) + \
grid[:, :, x1, y0, z0] * x * (1 - y) * (1 - z) + \
grid[:, :, x0, y1, z0] * (1 - x) * y * (1 - z) + \
grid[:, :, x0, y0, z1] * (1 - x) * (1 - y) * z + \
grid[:, :, x1, y0, z1] * x * (1 - y) * z + \
grid[:, :, x0, y1, z1] * (1 - x) * y * z + \
grid[:, :, x1, y1, z0] * x * y * (1 - z) + \
grid[:, :, x1, y1, z1] * x * y * z
make_slice = MakeSlice()
output = F.scatter_add(output, make_slice[:, :, lin_ind_frustrum], added)
output = output.reshape(batch, num_feats, frustrum_depth, img_shape[0],
img_shape[1])
return output
def forward(self, xinit, env, train_warm_start_idxs):
""" forward function
:param xinit:
:param env:Pendulum_DX
:param train_warm_start_idxs:
:return:
"""
Q = self.xp.zeros((self.n_sc, self.n_sc))
index_diag = self.xp.zeros((self.n_sc, self.n_sc), dtype=bool)
self.xp.fill_diagonal(index_diag, True)
index_not_diag = self.xp.zeros((self.n_sc, self.n_sc), dtype=bool)
index_not_diag[self.xp.tril_indices(self.n_sc, -1)] = True
Q = F.scatter_add(Q, self.xp.tril_indices(self.n_sc, -1), self.lower_without_diag)
diag_q = F.sigmoid(self.learn_q_logit)
Q = F.where(index_diag, diag_q, Q)
Q = Q @ Q.T
p = self.learn_p
Q = OBSERVATION_MATRIX.T @ Q @ OBSERVATION_MATRIX
p = p @ OBSERVATION_MATRIX
x_mpc, u_mpc = env.mpc_Q(env.true_dx, xinit, Q, p,
u_init=self.xp.transpose(train_warm_start_idxs, axes=(1, 0, 2)))
return x_mpc, u_mpc
def test_too_many_indices(self):
with self.assertRaises(type_check.InvalidType):
functions.scatter_add(self.a_data, (0, 0, 0, 0), self.b_data)
def test_multiple_ellipsis(self):
with self.assertRaises(ValueError):
functions.scatter_add(
self.a_data, (Ellipsis, Ellipsis), self.b_data)
def test_multiple_ellipsis(self):
with self.assertRaises(ValueError):
functions.scatter_add(self.a_data, (Ellipsis, Ellipsis),
self.b_data)
def f(a, b):
y = functions.scatter_add(a, self.slices, b)
return y * y
def forward(self, inputs, device):
a, b = inputs
y = functions.scatter_add(a, self.slices, b)
return y,
def f(a, b):
return functions.scatter_add(a, self.slices, b)
def test_requires_broadcasting(self):
with self.assertRaises(ValueError):
functions.scatter_add(self.a_data, slice(0, 2), self.b_data)
def test_too_many_indices(self):
with self.assertRaises(type_check.InvalidType):
functions.scatter_add(self.a_data, (0, 0, 0, 0), self.b_data)
def test_requires_broadcasting(self):
with self.assertRaises(ValueError):
functions.scatter_add(self.a_data, slice(0, 2), self.b_data)
def __call__(self, h, edge_index, edge_attr):
next_h = self.root_weight(h)
features = self.edge_weight(h).reshape(-1, self.n_edge_types,
self.out_channels)
messages = features[edge_index[0], edge_attr, :]
return functions.scatter_add(next_h, edge_index[1], messages)
def dataset_loss(self, data_iter, warmstart=None):
""" calculate loss and set warm start
this loss is
:param loader:
:param warmstart:
:return:
"""
true_q, true_p = self.env.true_dx.get_true_obj()
losses = []
iter_before = data_iter.epoch
while data_iter.epoch < iter_before + 1:
next_batch = data_iter.next()
next_batch = dataset.concat_examples(next_batch)
x_inits = next_batch[0]
xs = next_batch[1]
us = next_batch[2]
idxs = next_batch[3]
n_batch = x_inits.shape[0]
dx = self.env.true_dx
if not self.is_lower_triangle and not self.is_strange_observation:
print("not lower triangle and not strange observation")
assert type(self.net) == Pendulum_Net_cost_logit
q = F.sigmoid(self.net.learn_q_logit)
p = F.sqrt(q) * self.net.learn_p
_, pred_u = self.env.mpc(self.env.true_dx,
x_inits,
q,
p,
u_init=xp.transpose(warmstart[idxs],
axes=(1, 0, 2)))
pred_u = F.transpose(pred_u, axes=(1, 0, 2)).array
warmstart[idxs] = pred_u
elif not self.is_strange_observation:
print("lower triangle and not strange observation")
assert type(self.net) == Pendulum_Net_cost_lower_triangle
Q = xp.zeros((self.n_sc, self.n_sc))
index_diag = xp.zeros((self.n_sc, self.n_sc), dtype=bool)
xp.fill_diagonal(index_diag, True)
index_not_diag = xp.zeros((self.n_sc, self.n_sc), dtype=bool)
index_not_diag[xp.tril_indices(self.n_sc, -1)] = True
Q = F.scatter_add(Q, xp.tril_indices(self.n_sc, -1),
self.net.lower_without_diag)
# index_not_diag = xp.zeros((self.n_state, self.n_state), dtype=bool)
# index_not_diag[xp.tril_indices(self.n_state, -1)] = True
# Q = F.scatter_add(Q, xp.tril_indices(self.n_state, -1), self.net.lower_without_diag)
diag_q = F.sigmoid(self.net.learn_q_logit)
Q = F.where(index_diag, diag_q, Q)
Q = Q @ Q.T
p = self.net.learn_p
p = F.concat((p, xp.array(0.0).reshape(1, )), axis=0)
_, pred_u = self.env.mpc_Q(self.env.true_dx,
x_inits,
Q,
p,
u_init=xp.transpose(warmstart[idxs],
axes=(1, 0, 2)))
pred_u = F.transpose(pred_u, axes=(1, 0, 2)).array
warmstart[idxs] = pred_u
elif not self.is_lower_triangle:
print("not lower triangle and strange observation")
assert type(
self.net) == Pendulum_Net_cost_logit_strange_obervation
q = F.sigmoid(self.net.learn_q_logit)
p = F.sqrt(q) * self.net.learn_p
# p = self.learn_p
Q = xp.zeros((self.n_sc, self.n_sc))
index_diag = xp.zeros((self.n_sc, self.n_sc), dtype=bool)
xp.fill_diagonal(index_diag, True)
Q = F.where(index_diag, q, Q)
Q = OBSERVATION_MATRIX.T @ Q @ OBSERVATION_MATRIX
p = p @ OBSERVATION_MATRIX
_, pred_u = self.env.mpc_Q(self.env.true_dx,
x_inits,
Q,
p,
u_init=xp.transpose(warmstart[idxs],
axes=(1, 0, 2)))
pred_u = F.transpose(pred_u, axes=(1, 0, 2)).array
warmstart[idxs] = pred_u
else:
print("lower triangle and strange observation")
assert type(
self.net
) == Pendulum_Net_cost_lower_triangle_strange_obervation
Q = xp.zeros((self.n_sc, self.n_sc))
index_diag = xp.zeros((self.n_sc, self.n_sc), dtype=bool)
xp.fill_diagonal(index_diag, True)
index_not_diag = xp.zeros((self.n_sc, self.n_sc), dtype=bool)
index_not_diag[xp.tril_indices(self.n_sc, -1)] = True
Q = F.scatter_add(Q, xp.tril_indices(self.n_sc, -1),
self.net.lower_without_diag)
diag_q = F.sigmoid(self.net.learn_q_logit)
Q = F.where(index_diag, diag_q, Q)
Q = Q @ Q.T
Q = OBSERVATION_MATRIX.T @ Q @ OBSERVATION_MATRIX
p = self.net.learn_p
p = p @ OBSERVATION_MATRIX
_, pred_u = self.env.mpc_Q(self.env.true_dx,
x_inits,
Q,
p,
u_init=xp.transpose(warmstart[idxs],
axes=(1, 0, 2)))
pred_u = F.transpose(pred_u, axes=(1, 0, 2)).array
warmstart[idxs] = pred_u
assert pred_u.shape == us.shape
squared_loss = (us - pred_u) * (us - pred_u)
# print(squared_loss.shape)
loss = xp.mean(squared_loss)
losses.append(loss)
loss = xp.stack(losses).mean()
return loss
def f(a, b):
y = functions.scatter_add(a, self.slices, b)
return y * y
def f(a, b):
return functions.scatter_add(a, self.slices, b)
