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