def test_modes_ops():
X = np.array([[1, 0], [0, 1], [1, 1], [0, 0], [1, 2]])
I = np.array([0, 0, 0, 1, 1])
A_data = [1, 1, 1, 1, 1]
A_row = [0, 1, 2, 3, 4]
A_col = [1, 0, 1, 4, 3]
A_sparse = csr_matrix((A_data, (A_row, A_col)), shape=(5, 5))
A_sparse_tensor = ops.sp_matrix_to_sp_tensor(A_sparse)
# Disjoint signal to batch
expected_result = np.array(
[[[1.0, 0.0], [0.0, 1.0], [1.0, 1.0]], [[0.0, 0.0], [1.0, 2.0], [0.0, 0.0]]]
)
result = ops.disjoint_signal_to_batch(X, I).numpy()
assert expected_result.shape == result.shape
assert np.allclose(expected_result, result)
# Disjoint adjacency to batch
expected_result = np.array(
[
[[0.0, 1.0, 0.0], [1.0, 0.0, 0.0], [0.0, 1.0, 0.0]],
[[0.0, 1.0, 0.0], [1.0, 0.0, 0.0], [0.0, 0.0, 0.0]],
]
)
result = ops.disjoint_adjacency_to_batch(A_sparse_tensor, I).numpy()
assert expected_result.shape == result.shape
assert np.allclose(expected_result, result)
def call(self, inputs):
if self.data_mode == 'disjoint':
X, I = inputs
X = ops.disjoint_signal_to_batch(X, I)
else:
X = inputs
if self.data_mode == 'single':
X = tf.expand_dims(X, 0)
N = tf.shape(X)[-2]
sort_perm = tf.argsort(X[..., -1], direction='DESCENDING')
X_sorted = tf.gather(X, sort_perm, axis=-2, batch_dims=1)
def truncate():
_X_out = X_sorted[..., :self.k, :]
return _X_out
def pad():
padding = [[0, 0], [0, self.k - N], [0, 0]]
_X_out = tf.pad(X_sorted, padding)
return _X_out
X_out = tf.cond(tf.less_equal(self.k, N), truncate, pad)
if self.data_mode == 'single':
X_out = tf.squeeze(X_out, [0])
X_out.set_shape((self.k, self.F))
elif self.data_mode == 'batch' or self.data_mode == 'disjoint':
X_out.set_shape((None, self.k, self.F))
return X_out
def test_modes_ops():
ns = [10, 5, 8]
x_list = [np.random.randn(n, 3) for n in ns]
a_list = [sp.csr_matrix(np.ones((n, n))) for n in ns]
x, a, i = to_disjoint(x_list, a_list)
a = sp_matrix_to_sp_tensor(a)
expected_x = np.zeros((len(ns), max(ns), 3))
for i_, n in enumerate(ns):
expected_x[i_, :n] = x_list[i_]
expected_a = np.zeros((len(ns), max(ns), max(ns)))
for i_, n in enumerate(ns):
expected_a[i_, :n, :n] = a_list[i_].A
# Disjoint signal to batch
result = ops.disjoint_signal_to_batch(x, i).numpy()
assert expected_x.shape == result.shape
assert np.allclose(expected_x, result, atol=tol)
# Disjoint adjacency to batch
result = ops.disjoint_adjacency_to_batch(a, i).numpy()
assert expected_a.shape == result.shape
assert np.allclose(expected_a, result, atol=tol)
def call(self, inputs, **kwargs):
X, A, I = inputs
batch_X = ops.disjoint_signal_to_batch(X, I)
batch_A = ops.disjoint_adjacency_to_batch(A, I)
# Ensure that the channel dimension is known
batch_X.set_shape((None, None, X.shape[-1]))
batch_A.set_shape((None, None, None))
return batch_X, batch_A
