Block Diagonal Matrix Function

include/tensorwrapper/utilities/block_diagonal_matrix.hpp

@@ -0,0 +1,32 @@
+/*
+ * Copyright 2025 NWChemEx-Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#pragma once
+#include <tensorwrapper/tensor/tensor.hpp>
+
+namespace tensorwrapper::utilities {
+
+/** @brief Produce a block diagonal matrix from a set of square matrices.
+ *
+ *  @param[in] matrices The vector of matrices to be placed along the output
+ *                      diagonal.
+ *
+ *  @return A block diagonal matrix whose block values are equal to the input
+ *          matrices.
+ */
+Tensor block_diagonal_matrix(std::vector<Tensor> matrices);
+
+} // namespace tensorwrapper::utilities

include/tensorwrapper/utilities/utilities.hpp

@@ -15,6 +15,7 @@
  */
 
 #pragma once
+#include <tensorwrapper/utilities/block_diagonal_matrix.hpp>
 #include <tensorwrapper/utilities/floating_point_dispatch.hpp>
 #include <tensorwrapper/utilities/to_json.hpp>
 

src/tensorwrapper/utilities/block_diagonal_matrix.cpp

@@ -0,0 +1,82 @@
+/*
+ * Copyright 2025 NWChemEx-Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+#include <tensorwrapper/allocator/allocator.hpp>
+#include <tensorwrapper/buffer/buffer.hpp>
+#include <tensorwrapper/layout/layout.hpp>
+#include <tensorwrapper/shape/shape.hpp>
+#include <tensorwrapper/utilities/block_diagonal_matrix.hpp>
+#include <tensorwrapper/utilities/floating_point_dispatch.hpp>
+
+namespace tensorwrapper::utilities {
+
+namespace {
+
+struct BlockDiagonalMatrixKernel {
+    template<typename FloatType>
+    auto run(const buffer::BufferBase& b, const std::vector<Tensor>& matrices) {
+        using allocator_type = tensorwrapper::allocator::Eigen<FloatType>;
+
+        // All inputs must be Rank 2, square, and the same floating point type.
+        // If so, sum their extent sizes.
+        std::size_t size = 0;
+        for(const auto& matrix : matrices) {
+            if(!allocator_type::can_rebind(matrix.buffer()))
+                throw std::runtime_error(
+                  "All inputs must have the same floating point type");
+
+            if(matrix.rank() != 2)
+                throw std::runtime_error(
+                  "All inputs must be matrices (Rank == 2)");
+
+            const auto& mshape = matrix.buffer().layout().shape().as_smooth();
+            if(mshape.extent(0) != mshape.extent(1))
+                throw std::runtime_error("All inputs must be square matrices");
+
+            size += mshape.extent(0);
+        }
+
+        // Allocate new buffer
+        allocator_type allocator(b.allocator().runtime());
+        shape::Smooth oshape{size, size};
+        layout::Physical olayout(oshape);
+        auto obuffer = allocator.construct(olayout, 0.0);
+
+        // Copy values from input into corresponding blocks
+        std::size_t offset = 0;
+        for(const auto& matrix : matrices) {
+            const auto& mbuffer = allocator.rebind(matrix.buffer());
+            auto extent = mbuffer.layout().shape().as_smooth().extent(0);
+            for(std::size_t i = 0; i < extent; ++i) {
+                for(std::size_t j = 0; j < extent; ++j) {
+                    obuffer->set_elem({offset + i, offset + j},
+                                      mbuffer.get_elem({i, j}));
+                }
+            }
+            offset += extent;
+        }
+        return Tensor(oshape, std::move(obuffer));
+    }
+};
+
+} // namespace
+
+Tensor block_diagonal_matrix(std::vector<Tensor> matrices) {
+    const auto& buffer0 = matrices[0].buffer();
+    BlockDiagonalMatrixKernel kernel;
+    return floating_point_dispatch(kernel, buffer0, matrices);
+}
+
+} // namespace tensorwrapper::utilities

tests/cxx/unit_tests/tensorwrapper/testing/inputs.hpp

@@ -53,9 +53,9 @@ inline auto smooth_vector_alt() {
 }
 
 template<typename FloatType>
-inline auto smooth_matrix_() {
-    auto buffer = eigen_matrix<FloatType>();
-    shape::Smooth shape{2, 2};
+inline auto smooth_matrix_(std::size_t n = 2, std::size_t m = 2) {
+    auto buffer = eigen_matrix<FloatType>(n, m);
+    shape::Smooth shape{n, m};
     return detail_::TensorInput(shape, std::move(buffer));
 }
 

tests/cxx/unit_tests/tensorwrapper/utilities/block_diagonal_matrix.cpp

@@ -0,0 +1,85 @@
+/*
+ * Copyright 2025 NWChemEx-Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include "../testing/testing.hpp"
+#include <tensorwrapper/utilities/block_diagonal_matrix.hpp>
+
+using namespace tensorwrapper;
+using namespace testing;
+
+using tensorwrapper::utilities::block_diagonal_matrix;
+
+namespace {
+
+template<typename FloatType>
+struct TestTraits {
+    using other_float = float;
+};
+
+template<>
+struct TestTraits<float> {
+    using other_float = double;
+};
+
+} // namespace
+
+TEMPLATE_LIST_TEST_CASE("block_diagonal_matrix", "",
+                        types::floating_point_types) {
+    using other_float = typename TestTraits<TestType>::other_float;
+    Tensor square_matrix1(smooth_matrix_<TestType>());
+    Tensor square_matrix2(smooth_matrix_<TestType>(3, 3));
+    Tensor vector1(smooth_vector_<other_float>());
+    Tensor vector2(smooth_vector_<TestType>());
+    Tensor rectangular_matrix1(smooth_matrix_<TestType>(2, 3));
+    std::vector<Tensor> inputs1{square_matrix1, square_matrix2};
+    std::vector<Tensor> inputs2{square_matrix1, vector1};
+    std::vector<Tensor> inputs3{square_matrix1, vector2};
+    std::vector<Tensor> inputs4{square_matrix1, rectangular_matrix1};
+
+    SECTION("All matrices are square") {
+        shape::Smooth corr_shape{5, 5};
+        layout::Physical corr_layout(corr_shape);
+        auto allocator   = make_allocator<TestType>();
+        auto corr_buffer = allocator.allocate(corr_layout);
+        double counter1 = 1.0, counter2 = 1.0;
+        for(std::size_t i = 0; i < 5; ++i) {
+            for(std::size_t j = 0; j < 5; ++j) {
+                if(i >= 2 and j >= 2)
+                    corr_buffer->set_elem({i, j}, counter1++);
+                else if(i < 2 and j < 2)
+                    corr_buffer->set_elem({i, j}, counter2++);
+                else
+                    corr_buffer->set_elem({i, j}, 0.0);
+            }
+        }
+        Tensor corr(corr_shape, std::move(corr_buffer));
+
+        auto result = block_diagonal_matrix(inputs1);
+        REQUIRE(result == corr);
+    }
+
+    SECTION("Input has different floating point types") {
+        REQUIRE_THROWS(block_diagonal_matrix(inputs2));
+    }
+
+    SECTION("Input has non-matrix Tensor") {
+        REQUIRE_THROWS(block_diagonal_matrix(inputs3));
+    }
+
+    SECTION("Input has reactangular matrix") {
+        REQUIRE_THROWS(block_diagonal_matrix(inputs4));
+    }
+}