concatenation and block-diagonal composition of dense tensors

src/tensor/dense_tensor.c

@@ -773,7 +773,7 @@ void transpose_dense_tensor(const int* restrict perm, const struct dense_tensor*
 	assert(nelem == dense_tensor_num_elements(t));
 
 	long* index_t = ct_calloc(dp_eff.ndim,  sizeof(long));
-	long* index_r =  ct_malloc(dp_eff.ndim * sizeof(long));
+	long* index_r = ct_malloc(dp_eff.ndim * sizeof(long));
 
 	for (long ot = 0; ot < nelem; ot += dp_eff.dim[dp_eff.ndim - 1])
 	{
@@ -1670,6 +1670,170 @@ void dense_tensor_kronecker_product(const struct dense_tensor* restrict s, const
 }
 
 
+//________________________________________________________________________________________________________________________
+///
+/// \brief Concatenate tensors along the specified axis. All other dimensions must respectively agree.
+///
+void dense_tensor_concatenate(const struct dense_tensor* restrict tlist, const int num_tensors, const int i_ax, struct dense_tensor* restrict r)
+{
+	assert(num_tensors >= 1);
+
+	const int ndim = tlist[0].ndim;
+	assert(0 <= i_ax && i_ax < ndim);
+
+	for (int j = 0; j < num_tensors - 1; j++)
+	{
+		// data types must match
+		assert(tlist[j].dtype == tlist[j + 1].dtype);
+		// degrees must match
+		assert(tlist[j].ndim == tlist[j + 1].ndim);
+		for (int i = 0; i < tlist[j].ndim; i++) {
+			if (i != i_ax) {
+				// other dimensions must match
+				assert(tlist[j].dim[i] == tlist[j + 1].dim[i]);
+			}
+		}
+	}
+
+	// dimensions of new tensor
+	long dim_concat = 0;
+	for (int j = 0; j < num_tensors; j++) {
+		dim_concat += tlist[j].dim[i_ax];
+	}
+	long* rdim = ct_malloc(ndim * sizeof(long));
+	for (int i = 0; i < ndim; i++)
+	{
+		if (i == i_ax) {
+			rdim[i] = dim_concat;
+		}
+		else {
+			rdim[i] = tlist[0].dim[i];
+		}
+	}
+	allocate_dense_tensor(tlist[0].dtype, ndim, rdim, r);
+	ct_free(rdim);
+
+	// leading dimensions
+	const long ld = integer_product(r->dim, i_ax);
+	// trailing dimensions times data type size
+	const long tdd = integer_product(&r->dim[i_ax + 1], ndim - (i_ax + 1)) * sizeof_numeric_type(r->dtype);
+
+	long offset = 0;
+	for (int j = 0; j < num_tensors; j++)
+	{
+		const long stride = tlist[j].dim[i_ax] * tdd;
+
+		for (long i = 0; i < ld; i++) {
+			memcpy((int8_t*)r->data + (i * r->dim[i_ax] + offset) * tdd,
+			       (int8_t*)tlist[j].data + i * stride, stride);
+		}
+
+		offset += tlist[j].dim[i_ax];
+	}
+}
+
+
+//________________________________________________________________________________________________________________________
+///
+/// \brief Compose tensors along the specified axes, creating a block-diagonal pattern. All other dimensions must respectively agree.
+///
+void dense_tensor_block_diag(const struct dense_tensor* restrict tlist, const int num_tensors, const int* i_ax, const int ndim_block, struct dense_tensor* restrict r)
+{
+	assert(num_tensors >= 1);
+
+	const int ndim = tlist[0].ndim;
+
+	assert(1 <= ndim_block && ndim_block <= ndim);
+	bool* i_ax_indicator = ct_calloc(ndim, sizeof(bool));
+	for (int i = 0; i < ndim_block; i++)
+	{
+		assert(0 <= i_ax[i] && i_ax[i] < ndim);
+		assert(!i_ax_indicator[i_ax[i]]);  // axis index can only appear once
+		i_ax_indicator[i_ax[i]] = true;
+	}
+
+	for (int j = 0; j < num_tensors - 1; j++)
+	{
+		// data types must match
+		assert(tlist[j].dtype == tlist[j + 1].dtype);
+		// degrees must match
+		assert(tlist[j].ndim == tlist[j + 1].ndim);
+		for (int i = 0; i < tlist[j].ndim; i++) {
+			if (!i_ax_indicator[i]) {
+				// other dimensions must match
+				assert(tlist[j].dim[i] == tlist[j + 1].dim[i]);
+			}
+		}
+	}
+
+	// dimensions of new tensor
+	long* rdim = ct_malloc(ndim * sizeof(long));
+	for (int i = 0; i < ndim; i++)
+	{
+		if (i_ax_indicator[i]) {
+			long dsum = 0;
+			for (int j = 0; j < num_tensors; j++) {
+				dsum += tlist[j].dim[i];
+			}
+			rdim[i] = dsum;
+		}
+		else {
+			rdim[i] = tlist[0].dim[i];
+		}
+	}
+	allocate_dense_tensor(tlist[0].dtype, ndim, rdim, r);
+	ct_free(rdim);
+
+	// effective dimensions used for indexing tensor slices
+	int ndim_eff = 0;
+	for (int i = 0; i < ndim; i++) {
+		if (i_ax_indicator[i]) {
+			ndim_eff = i + 1;
+		}
+	}
+	assert(ndim_eff >= 1);
+
+	// trailing dimensions times data type size
+	const long tdd = integer_product(&r->dim[ndim_eff], ndim - ndim_eff) * sizeof_numeric_type(r->dtype);
+
+	long* offset = ct_calloc(ndim_eff, sizeof(long));
+
+	long* index_t = ct_calloc(ndim_eff, sizeof(long));
+	long* index_r = ct_calloc(ndim_eff, sizeof(long));
+
+	for (int j = 0; j < num_tensors; j++)
+	{
+		// leading dimensions
+		const long ld = integer_product(tlist[j].dim, ndim_eff - 1);
+
+		const long stride = tlist[j].dim[ndim_eff - 1] * tdd;
+
+		memset(index_t, 0, ndim_eff * sizeof(long));
+		for (long ot = 0; ot < ld; ot++, next_tensor_index(ndim_eff - 1, tlist[j].dim, index_t))
+		{
+			for (int i = 0; i < ndim_eff; i++) {
+				index_r[i] = offset[i] + index_t[i];
+				assert(index_r[i] < r->dim[i]);
+			}
+			const long or = tensor_index_to_offset(ndim_eff, r->dim, index_r);
+			memcpy((int8_t*)r->data + or * tdd,
+			       (int8_t*)tlist[j].data + ot * stride, stride);
+		}
+
+		for (int i = 0; i < ndim_eff; i++) {
+			if (i_ax_indicator[i]) {
+				offset[i] += tlist[j].dim[i];
+			}
+		}
+	}
+
+	ct_free(index_r);
+	ct_free(index_t);
+	ct_free(offset);
+	ct_free(i_ax_indicator);
+}
+
+
 //________________________________________________________________________________________________________________________
 ///
 /// \brief Compute the QR decomposition of the matrix 'a', and store the result in 'q' and 'r' (will be allocated).

src/tensor/dense_tensor.h

@@ -174,6 +174,10 @@ void dense_tensor_dot_update(const void* alpha, const struct dense_tensor* restr
 
 void dense_tensor_kronecker_product(const struct dense_tensor* restrict s, const struct dense_tensor* restrict t, struct dense_tensor* restrict r);
 
+void dense_tensor_concatenate(const struct dense_tensor* restrict tlist, const int num_tensors, const int i_ax, struct dense_tensor* restrict r);
+
+void dense_tensor_block_diag(const struct dense_tensor* restrict tlist, const int num_tensors, const int* i_ax, const int ndim_block, struct dense_tensor* restrict r);
+
 
 //________________________________________________________________________________________________________________________
 //

test/run_tests.c

@@ -21,6 +21,8 @@ char* test_dense_tensor_dot();
 char* test_dense_tensor_dot_update();
 char* test_dense_tensor_kronecker_product();
 char* test_dense_tensor_kronecker_product_degree_zero();
+char* test_dense_tensor_concatenate();
+char* test_dense_tensor_block_diag();
 char* test_dense_tensor_qr();
 char* test_dense_tensor_rq();
 char* test_dense_tensor_svd();
@@ -99,6 +101,8 @@ int main()
 		TEST_FUNCTION_ENTRY(test_dense_tensor_dot_update),
 		TEST_FUNCTION_ENTRY(test_dense_tensor_kronecker_product),
 		TEST_FUNCTION_ENTRY(test_dense_tensor_kronecker_product_degree_zero),
+		TEST_FUNCTION_ENTRY(test_dense_tensor_concatenate),
+		TEST_FUNCTION_ENTRY(test_dense_tensor_block_diag),
 		TEST_FUNCTION_ENTRY(test_dense_tensor_qr),
 		TEST_FUNCTION_ENTRY(test_dense_tensor_rq),
 		TEST_FUNCTION_ENTRY(test_dense_tensor_svd),

test/tensor/test_dense_tensor.c

@@ -615,6 +615,114 @@ char* test_dense_tensor_kronecker_product_degree_zero()
 }
 
 
+char* test_dense_tensor_concatenate()
+{
+	hid_t file = H5Fopen("../test/tensor/data/test_dense_tensor_concatenate.hdf5", H5F_ACC_RDONLY, H5P_DEFAULT);
+	if (file < 0) {
+		return "'H5Fopen' in test_dense_tensor_concatenate failed";
+	}
+
+	struct dense_tensor tlist[3];
+	const long dims[3][4] = {
+		{ 5, 8, 7, 3 },
+		{ 5, 8, 9, 3 },
+		{ 5, 8, 2, 3 },
+	};
+	for (int i = 0; i < 3; i++)
+	{
+		allocate_dense_tensor(CT_SINGLE_REAL, 4, dims[i], &tlist[i]);
+		// read values from disk
+		char varname[1024];
+		sprintf(varname, "t%i", i);
+		if (read_hdf5_dataset(file, varname, H5T_NATIVE_FLOAT, tlist[i].data) < 0) {
+			return "reading tensor entries from disk failed";
+		}
+	}
+
+	struct dense_tensor r;
+	const int i_ax = 2;
+	dense_tensor_concatenate(tlist, 3, i_ax, &r);
+
+	// load reference values from disk
+	struct dense_tensor r_ref;
+	const long refdim[4] = { 5, 8, 18, 3 };
+	allocate_dense_tensor(CT_SINGLE_REAL, 4, refdim, &r_ref);
+	// read values from disk
+	if (read_hdf5_dataset(file, "r", H5T_NATIVE_FLOAT, r_ref.data) < 0) {
+		return "reading tensor entries from disk failed";
+	}
+
+	// compare
+	if (!dense_tensor_allclose(&r, &r_ref, 1e-13)) {
+		return "concatenated tensor does not match reference";
+	}
+
+	delete_dense_tensor(&r_ref);
+	delete_dense_tensor(&r);
+	for (int i = 0; i < 3; i++) {
+		delete_dense_tensor(&tlist[i]);
+	}
+
+	H5Fclose(file);
+
+	return 0;
+}
+
+
+char* test_dense_tensor_block_diag()
+{
+	hid_t file = H5Fopen("../test/tensor/data/test_dense_tensor_block_diag.hdf5", H5F_ACC_RDONLY, H5P_DEFAULT);
+	if (file < 0) {
+		return "'H5Fopen' in test_dense_tensor_block_diag failed";
+	}
+
+	struct dense_tensor tlist[3];
+	const long dims[3][5] = {
+		{ 5, 8, 7, 3, 4 },
+		{ 8, 2, 7, 4, 4 },
+		{ 6, 1, 7, 9, 4 },
+	};
+	for (int i = 0; i < 3; i++)
+	{
+		allocate_dense_tensor(CT_DOUBLE_REAL, 5, dims[i], &tlist[i]);
+		// read values from disk
+		char varname[1024];
+		sprintf(varname, "t%i", i);
+		if (read_hdf5_dataset(file, varname, H5T_NATIVE_DOUBLE, tlist[i].data) < 0) {
+			return "reading tensor entries from disk failed";
+		}
+	}
+
+	struct dense_tensor r;
+	const int i_ax[3] = { 0, 1, 3 };
+	dense_tensor_block_diag(tlist, 3, i_ax, 3, &r);
+
+	// load reference values from disk
+	struct dense_tensor r_ref;
+	const long refdim[5] = { 19, 11, 7, 16, 4 };
+	allocate_dense_tensor(CT_DOUBLE_REAL, 5, refdim, &r_ref);
+	// read values from disk
+	if (read_hdf5_dataset(file, "r", H5T_NATIVE_DOUBLE, r_ref.data) < 0) {
+		return "reading tensor entries from disk failed";
+	}
+
+	// compare
+	if (!dense_tensor_allclose(&r, &r_ref, 1e-13)) {
+		return "block-diagonal tensor does not match reference";
+	}
+
+	delete_dense_tensor(&r_ref);
+	delete_dense_tensor(&r);
+	for (int i = 0; i < 3; i++) {
+		delete_dense_tensor(&tlist[i]);
+	}
+
+	H5Fclose(file);
+
+	return 0;
+}
+
+
 char* test_dense_tensor_qr()
 {
 	hid_t file = H5Fopen("../test/tensor/data/test_dense_tensor_qr.hdf5", H5F_ACC_RDONLY, H5P_DEFAULT);

test/tensor/test_dense_tensor.py

@@ -179,6 +179,47 @@ def dense_tensor_kronecker_product_degree_zero_data():
         file["r"] = interleave_complex(r)
 
 
+def dense_tensor_concatenate_data():
+
+    # random number generator
+    rng = np.random.default_rng(201)
+
+    tlist = [
+        rng.standard_normal((5, 8, 7, 3)).astype(np.float32),
+        rng.standard_normal((5, 8, 9, 3)).astype(np.float32),
+        rng.standard_normal((5, 8, 2, 3)).astype(np.float32)]
+
+    r = np.concatenate(tlist, axis=2)
+
+    with h5py.File("data/test_dense_tensor_concatenate.hdf5", "w") as file:
+        file["t0"] = tlist[0]
+        file["t1"] = tlist[1]
+        file["t2"] = tlist[2]
+        file["r"] = r
+
+
+def dense_tensor_block_diag_data():
+
+    # random number generator
+    rng = np.random.default_rng(874)
+
+    tlist = [
+        rng.standard_normal((5, 8, 7, 3, 4)),
+        rng.standard_normal((8, 2, 7, 4, 4)),
+        rng.standard_normal((6, 1, 7, 9, 4))]
+
+    r = np.zeros((19, 11, 7, 16, 4))
+    r[  :5,    :8 , :,  :3, :] = tlist[0]
+    r[ 5:13,  8:10, :, 3:7, :] = tlist[1]
+    r[13:,   10:  , :, 7:,  :] = tlist[2]
+
+    with h5py.File("data/test_dense_tensor_block_diag.hdf5", "w") as file:
+        file["t0"] = tlist[0]
+        file["t1"] = tlist[1]
+        file["t2"] = tlist[2]
+        file["r"] = r
+
+
 def dense_tensor_qr_data():
 
     # random number generator
@@ -283,6 +324,8 @@ def main():
     dense_tensor_dot_update_data()
     dense_tensor_kronecker_product_data()
     dense_tensor_kronecker_product_degree_zero_data()
+    dense_tensor_concatenate_data()
+    dense_tensor_block_diag_data()
     dense_tensor_qr_data()
     dense_tensor_rq_data()
     dense_tensor_svd_data()