From 78d05d5863c11d90712586fddb77dc3f648db1ce Mon Sep 17 00:00:00 2001 From: "Jonathan M. Waldrop" Date: Thu, 17 Jul 2025 15:32:38 -0500 Subject: [PATCH] block diagonal matrix function --- .../utilities/block_diagonal_matrix.hpp | 32 +++++++ include/tensorwrapper/utilities/utilities.hpp | 1 + .../utilities/block_diagonal_matrix.cpp | 82 ++++++++++++++++++ .../tensorwrapper/testing/inputs.hpp | 6 +- .../utilities/block_diagonal_matrix.cpp | 85 +++++++++++++++++++ 5 files changed, 203 insertions(+), 3 deletions(-) create mode 100644 include/tensorwrapper/utilities/block_diagonal_matrix.hpp create mode 100644 src/tensorwrapper/utilities/block_diagonal_matrix.cpp create mode 100644 tests/cxx/unit_tests/tensorwrapper/utilities/block_diagonal_matrix.cpp diff --git a/include/tensorwrapper/utilities/block_diagonal_matrix.hpp b/include/tensorwrapper/utilities/block_diagonal_matrix.hpp new file mode 100644 index 00000000..8b8895f6 --- /dev/null +++ b/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 + +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 matrices); + +} // namespace tensorwrapper::utilities \ No newline at end of file diff --git a/include/tensorwrapper/utilities/utilities.hpp b/include/tensorwrapper/utilities/utilities.hpp index 415185ee..1c56040a 100644 --- a/include/tensorwrapper/utilities/utilities.hpp +++ b/include/tensorwrapper/utilities/utilities.hpp @@ -15,6 +15,7 @@ */ #pragma once +#include #include #include diff --git a/src/tensorwrapper/utilities/block_diagonal_matrix.cpp b/src/tensorwrapper/utilities/block_diagonal_matrix.cpp new file mode 100644 index 00000000..b3c922e4 --- /dev/null +++ b/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 +#include +#include +#include +#include +#include + +namespace tensorwrapper::utilities { + +namespace { + +struct BlockDiagonalMatrixKernel { + template + auto run(const buffer::BufferBase& b, const std::vector& matrices) { + using allocator_type = tensorwrapper::allocator::Eigen; + + // 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 matrices) { + const auto& buffer0 = matrices[0].buffer(); + BlockDiagonalMatrixKernel kernel; + return floating_point_dispatch(kernel, buffer0, matrices); +} + +} // namespace tensorwrapper::utilities \ No newline at end of file diff --git a/tests/cxx/unit_tests/tensorwrapper/testing/inputs.hpp b/tests/cxx/unit_tests/tensorwrapper/testing/inputs.hpp index 71b9dab1..d3991f77 100644 --- a/tests/cxx/unit_tests/tensorwrapper/testing/inputs.hpp +++ b/tests/cxx/unit_tests/tensorwrapper/testing/inputs.hpp @@ -53,9 +53,9 @@ inline auto smooth_vector_alt() { } template -inline auto smooth_matrix_() { - auto buffer = eigen_matrix(); - shape::Smooth shape{2, 2}; +inline auto smooth_matrix_(std::size_t n = 2, std::size_t m = 2) { + auto buffer = eigen_matrix(n, m); + shape::Smooth shape{n, m}; return detail_::TensorInput(shape, std::move(buffer)); } diff --git a/tests/cxx/unit_tests/tensorwrapper/utilities/block_diagonal_matrix.cpp b/tests/cxx/unit_tests/tensorwrapper/utilities/block_diagonal_matrix.cpp new file mode 100644 index 00000000..075b9f74 --- /dev/null +++ b/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 + +using namespace tensorwrapper; +using namespace testing; + +using tensorwrapper::utilities::block_diagonal_matrix; + +namespace { + +template +struct TestTraits { + using other_float = float; +}; + +template<> +struct TestTraits { + using other_float = double; +}; + +} // namespace + +TEMPLATE_LIST_TEST_CASE("block_diagonal_matrix", "", + types::floating_point_types) { + using other_float = typename TestTraits::other_float; + Tensor square_matrix1(smooth_matrix_()); + Tensor square_matrix2(smooth_matrix_(3, 3)); + Tensor vector1(smooth_vector_()); + Tensor vector2(smooth_vector_()); + Tensor rectangular_matrix1(smooth_matrix_(2, 3)); + std::vector inputs1{square_matrix1, square_matrix2}; + std::vector inputs2{square_matrix1, vector1}; + std::vector inputs3{square_matrix1, vector2}; + std::vector 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(); + 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)); + } +} \ No newline at end of file