End-to-end UQ

src/scf/driver/scf_loop.cpp

@@ -107,8 +107,8 @@ MODULE_RUN(SCFLoop) {
     }
     auto e_nuclear = V_nn_mod.run_as<v_nn_pt>(qs_lhs, qs_rhs);
 
-    wf_type psi_old             = psi0;
-    double e_old                = 0;
+    wf_type psi_old = psi0;
+    simde::type::tensor e_old;
     const unsigned int max_iter = 3;
     unsigned int iter           = 0;
 
@@ -154,8 +154,23 @@ MODULE_RUN(SCFLoop) {
         psi_old = new_psi;
         ++iter;
     }
-    auto rv = results();
-    return pt<wf_type>::wrap_results(rv, e_old + e_nuclear, psi_old);
+    simde::type::tensor e_total;
+
+    // e_nuclear is a double. This hack converts it to udouble (if needed)
+    tensorwrapper::allocator::Eigen<double> dalloc(get_runtime());
+    using tensorwrapper::types::udouble;
+    tensorwrapper::allocator::Eigen<udouble> ualloc(get_runtime());
+
+    if(ualloc.can_rebind(e_old.buffer())) {
+        simde::type::tensor temp(e_old);
+        auto val = dalloc.rebind(e_nuclear.buffer()).at();
+        ualloc.rebind(temp.buffer()).at() = val;
+        e_nuclear                         = temp;
+    }
+
+    e_total("") = e_old("") + e_nuclear("");
+    auto rv     = results();
+    return pt<wf_type>::wrap_results(rv, e_total, psi_old);
 }
 
 } // namespace scf::driver

src/scf/eigen_solver/eigen_generalized.cpp

@@ -20,6 +20,60 @@
 
 namespace scf::eigen_solver {
 
+namespace {
+struct Kernel {
+    template<typename FloatType>
+    auto run(const tensorwrapper::buffer::BufferBase& A,
+             const tensorwrapper::buffer::BufferBase& B,
+             parallelzone::runtime::RuntimeView& rv) {
+        // Convert to Eigen buffers
+        tensorwrapper::allocator::Eigen<FloatType> allocator(rv);
+        const auto& eigen_A = allocator.rebind(A);
+        const auto& eigen_B = allocator.rebind(B);
+
+        // Wrap the tensors in Eigen::Map objects to avoid copy
+        const auto* pA      = eigen_A.data();
+        const auto* pB      = eigen_B.data();
+        const auto& shape_A = eigen_A.layout().shape().as_smooth();
+        auto rows           = shape_A.extent(0);
+        auto cols           = shape_A.extent(1);
+
+        constexpr auto rmajor = Eigen::RowMajor;
+        constexpr auto edynam = Eigen::Dynamic;
+        using matrix_type = Eigen::Matrix<FloatType, edynam, edynam, rmajor>;
+        using map_type    = Eigen::Map<const matrix_type>;
+
+        map_type A_map(pA, rows, cols);
+        map_type B_map(pB, rows, cols);
+
+        // Compute
+        Eigen::GeneralizedSelfAdjointEigenSolver<matrix_type> ges(A_map, B_map);
+        auto eigen_values  = ges.eigenvalues();
+        auto eigen_vectors = ges.eigenvectors();
+
+        // Wrap in TensorWrapper Tensor
+        tensorwrapper::shape::Smooth vector_shape{rows};
+        tensorwrapper::shape::Smooth matrix_shape{rows, cols};
+        tensorwrapper::layout::Physical vector_layout(vector_shape);
+        tensorwrapper::layout::Physical matrix_layout(matrix_shape);
+
+        auto pvalues_buffer  = allocator.allocate(vector_layout);
+        auto pvectors_buffer = allocator.allocate(matrix_layout);
+
+        for(auto i = 0; i < rows; ++i) {
+            pvalues_buffer->at(i) = eigen_values(i);
+            for(auto j = 0; j < cols; ++j) {
+                pvectors_buffer->at(i, j) = eigen_vectors(i, j);
+            }
+        }
+
+        simde::type::tensor values(vector_shape, std::move(pvalues_buffer));
+        simde::type::tensor vectors(matrix_shape, std::move(pvectors_buffer));
+        return std::make_pair(values, vectors);
+    }
+};
+} // namespace
+
 using pt = simde::GeneralizedEigenSolve;
 
 const auto desc = R"(
@@ -37,43 +91,13 @@ MODULE_CTOR(EigenGeneralized) {
 MODULE_RUN(EigenGeneralized) {
     auto&& [A, B] = pt::unwrap_inputs(inputs);
 
-    // Convert to Eigen buffers
-    tensorwrapper::allocator::Eigen<double> allocator(get_runtime());
-    const auto& eigen_A = allocator.rebind(A.buffer());
-    const auto& eigen_B = allocator.rebind(B.buffer());
-
-    // Wrap the tensors in Eigen::Map objects to avoid copy
-    const auto* pA      = eigen_A.data();
-    const auto* pB      = eigen_B.data();
-    const auto& shape_A = eigen_A.layout().shape().as_smooth();
-    auto rows           = shape_A.extent(0);
-    auto cols           = shape_A.extent(1);
-    Eigen::Map<const Eigen::MatrixXd> A_map(pA, rows, cols);
-    Eigen::Map<const Eigen::MatrixXd> B_map(pB, rows, cols);
-
-    // Compute
-    Eigen::GeneralizedSelfAdjointEigenSolver<Eigen::MatrixXd> ges(A_map, B_map);
-    auto eigen_values  = ges.eigenvalues();
-    auto eigen_vectors = ges.eigenvectors();
-
-    // Wrap in TensorWrapper Tensor
-    tensorwrapper::shape::Smooth vector_shape{rows};
-    tensorwrapper::shape::Smooth matrix_shape{rows, cols};
-    tensorwrapper::layout::Physical vector_layout(vector_shape);
-    tensorwrapper::layout::Physical matrix_layout(matrix_shape);
-
-    auto pvalues_buffer  = allocator.allocate(vector_layout);
-    auto pvectors_buffer = allocator.allocate(matrix_layout);
-
-    for(auto i = 0; i < rows; ++i) {
-        pvalues_buffer->at(i) = eigen_values(i);
-        for(auto j = 0; j < cols; ++j) {
-            pvectors_buffer->at(i, j) = eigen_vectors(i, j);
-        }
-    }
+    using tensorwrapper::utilities::floating_point_dispatch;
 
-    simde::type::tensor values(vector_shape, std::move(pvalues_buffer));
-    simde::type::tensor vectors(matrix_shape, std::move(pvectors_buffer));
+    auto r = get_runtime();
+    Kernel k;
+    const auto& A_buffer   = A.buffer();
+    const auto& B_buffer   = B.buffer();
+    auto [values, vectors] = floating_point_dispatch(k, A_buffer, B_buffer, r);
 
     auto rv = results();
     return pt::wrap_results(rv, values, vectors);

src/scf/matrix_builder/density_matrix.cpp

@@ -23,7 +23,40 @@ namespace {
 const auto desc = R"(
 )";
 
-}
+struct Kernel {
+    template<typename FloatType>
+    auto run(const tensorwrapper::buffer::BufferBase& c, std::size_t n_aos,
+             std::size_t n_occ) {
+        constexpr auto rmajor = Eigen::RowMajor;
+        constexpr auto edynam = Eigen::Dynamic;
+        using allocator_type  = tensorwrapper::allocator::Eigen<FloatType>;
+        using tensor_type    = Eigen::Matrix<FloatType, edynam, edynam, rmajor>;
+        using map_type       = Eigen::Map<tensor_type>;
+        using const_map_type = Eigen::Map<const tensor_type>;
+        auto rv              = c.allocator().runtime();
+        allocator_type alloc(rv);
+
+        tensorwrapper::shape::Smooth p_shape(n_aos, n_aos);
+        tensorwrapper::layout::Physical l(p_shape);
+        auto pp_buffer = alloc.allocate(l);
+
+        // Step 2: Grab the orbitals in the ensemble
+        auto& c_buffer = alloc.rebind(c);
+
+        const_map_type c_eigen(c_buffer.data(), n_aos, n_aos);
+        map_type p_eigen(pp_buffer->data(), n_aos, n_aos);
+        auto slice = c_eigen.block(0, 0, n_aos, n_occ);
+
+        // Step 3: CC_dagger
+        using index_pair_t = Eigen::IndexPair<int>;
+        Eigen::array<index_pair_t, 1> modes{index_pair_t(1, 1)};
+        p_eigen = slice * slice.transpose();
+
+        return simde::type::tensor(p_shape, std::move(pp_buffer));
+    }
+};
+
+} // namespace
 
 using pt = simde::aos_rho_e_aos<simde::type::cmos>;
 
@@ -63,35 +96,10 @@ MODULE_RUN(DensityMatrix) {
               "Please shuffle your orbitals so that the ensemble is a "
               "contiguous slice of the orbitals.");
 
-    using allocator_type    = tensorwrapper::allocator::Eigen<double>;
-    using tensor_type       = Eigen::Tensor<double, 2, Eigen::RowMajor>;
-    using const_tensor_type = Eigen::Tensor<const double, 2, Eigen::RowMajor>;
-    using map_type          = Eigen::TensorMap<tensor_type>;
-    using const_map_type    = Eigen::TensorMap<const_tensor_type>;
-
-    allocator_type alloc(get_runtime());
-
-    tensorwrapper::shape::Smooth p_shape(n_aos, n_aos);
-    tensorwrapper::layout::Physical l(p_shape);
-    auto pp_buffer = alloc.allocate(l);
-
-    // Step 2: Grab the orbitals in the ensemble
-    const auto& c_buffer = alloc.rebind(c.buffer());
-    const_map_type c_eigen(c_buffer.data(), n_aos, n_aos);
-    map_type p_eigen(pp_buffer->data(), n_aos, n_aos);
-
-    using slice_t = Eigen::array<Eigen::Index, 2>;
-    slice_t offsets{0, 0};
-    slice_t extents{n_aos, Eigen::Index(participants.size())};
-    auto slice = c_eigen.slice(offsets, extents);
-
-    // Step 3: CC_dagger
-    using index_pair_t = Eigen::IndexPair<int>;
-    Eigen::array<index_pair_t, 1> modes{index_pair_t(1, 1)};
-    p_eigen = slice.contract(slice, modes);
-
-    simde::type::tensor p(p_shape, std::move(pp_buffer));
-
+    // TODO: The need to dispatch like this goes away when TW supports slicing
+    using tensorwrapper::utilities::floating_point_dispatch;
+    Kernel k;
+    auto p = floating_point_dispatch(k, c.buffer(), n_aos, participants.size());
     auto rv = results();
     return pt::wrap_results(rv, p);
 }

src/scf/matrix_builder/determinant_driver.cpp

@@ -103,7 +103,6 @@ MODULE_CTOR(DeterminantDriver) {
 }
 
 MODULE_RUN(DeterminantDriver) {
-    using float_type      = double;
     using wf_type         = simde::type::rscf_wf;
     const auto&& [braket] = pt<wf_type>::unwrap_inputs(inputs);
     const auto& bra       = braket.bra();
@@ -130,11 +129,8 @@ MODULE_RUN(DeterminantDriver) {
     tensorwrapper::Tensor x;
     x("") = rho("i,j") * t("i,j");
 
-    using allocator_type = tensorwrapper::allocator::Eigen<float_type>;
-    auto& x_buffer       = allocator_type::rebind(x.buffer());
-
     auto rv = results();
-    return pt<wf_type>::wrap_results(rv, x_buffer.at());
+    return pt<wf_type>::wrap_results(rv, x);
 }
 
 } // namespace scf::matrix_builder

src/scf/matrix_builder/electronic_energy.cpp

@@ -48,7 +48,7 @@ MODULE_RUN(ElectronicEnergy) {
     const auto& ket     = H_ij.ket();
 
     auto& mod = submods.at("determinant driver");
-    double e  = 0.0;
+    simde::type::tensor e;
 
     auto n_ops = H.size();
     for(decltype(n_ops) i = 0; i < n_ops; ++i) {
@@ -57,7 +57,11 @@ MODULE_RUN(ElectronicEnergy) {
 
         chemist::braket::BraKet O_ij(bra, O_i, ket);
         const auto o = mod.run_as<det_pt<wf_type>>(O_ij);
-        e += ci * o;
+        simde::type::tensor temp;
+        temp("") = o("") * ci;
+        if(i == 0)
+            e = temp;
+        else { e("") = e("") + temp(""); }
     }
 
     auto rv = results();

tests/cxx/integration_tests/coulombs_law.cpp

@@ -19,9 +19,15 @@
 using pt = simde::charge_charge_interaction;
 using Catch::Matchers::WithinAbs;
 
-TEST_CASE("CoulombsLaw") {
-    auto mm   = test_scf::load_modules();
-    auto& mod = mm.at("Coulomb's Law");
+TEMPLATE_LIST_TEST_CASE("CoulombsLaw", "", test_scf::float_types) {
+    using float_type = double;
+    auto mm          = test_scf::load_modules<float_type>();
+    auto& mod        = mm.at("Coulomb's Law");
+
+    tensorwrapper::allocator::Eigen<float_type> alloc(mm.get_runtime());
+    tensorwrapper::shape::Smooth shape_corr{};
+    auto pcorr = alloc.allocate(tensorwrapper::layout::Physical(shape_corr));
+    using tensorwrapper::operations::approximately_equal;
 
     auto h2_nuclei = test_scf::make_h2<simde::type::nuclei>();
     // TODO: Conversions are missing in Chemist. Use those when they're in place
@@ -30,5 +36,8 @@ TEST_CASE("CoulombsLaw") {
         qs.push_back(nucleus.as_point_charge());
 
     auto e_nuclear = mod.run_as<pt>(qs, qs);
-    REQUIRE_THAT(e_nuclear, WithinAbs(0.71510297482837526, 1E-6));
+
+    pcorr->at() = 0.71510297482837526;
+    tensorwrapper::Tensor corr(shape_corr, std::move(pcorr));
+    REQUIRE(approximately_equal(corr, e_nuclear, 1E-6));
 }

tests/cxx/integration_tests/driver/scf_driver.cpp

@@ -16,15 +16,21 @@
 
 #include "../integration_tests.hpp"
 
-using Catch::Matchers::WithinAbs;
-
 using pt = simde::AOEnergy;
 
-TEST_CASE("SCFDriver") {
-    auto mm  = test_scf::load_modules();
-    auto h2  = test_scf::make_h2<simde::type::chemical_system>();
-    auto aos = test_scf::h2_aos().ao_basis_set();
+TEMPLATE_LIST_TEST_CASE("SCFDriver", "", test_scf::float_types) {
+    using float_type = TestType;
+    auto mm          = test_scf::load_modules<float_type>();
+    auto h2          = test_scf::make_h2<simde::type::chemical_system>();
+    auto aos         = test_scf::h2_aos().ao_basis_set();
+
+    tensorwrapper::allocator::Eigen<float_type> alloc(mm.get_runtime());
+    tensorwrapper::shape::Smooth shape_corr{};
+    auto pcorr = alloc.allocate(tensorwrapper::layout::Physical(shape_corr));
+    using tensorwrapper::operations::approximately_equal;
+    pcorr->at() = -1.1167592336;
+    tensorwrapper::Tensor corr(shape_corr, std::move(pcorr));
 
-    const auto e = mm.run_as<pt>("SCF Driver", aos, h2);
-    REQUIRE_THAT(e, WithinAbs(-1.1167592336, 1E-6));
+    const auto e = mm.template run_as<pt>("SCF Driver", aos, h2);
+    REQUIRE(approximately_equal(corr, e, 1E-6));
 }

tests/cxx/integration_tests/driver/scf_loop.cpp

@@ -24,17 +24,26 @@ using egy_pt = simde::eval_braket<WFType, simde::type::hamiltonian, WFType>;
 template<typename WFType>
 using pt = simde::Optimize<egy_pt<WFType>, WFType>;
 
-TEST_CASE("SCFLoop") {
-    using wf_type = simde::type::rscf_wf;
-    auto mm       = test_scf::load_modules();
-    auto& mod     = mm.at("Loop");
+TEMPLATE_LIST_TEST_CASE("SCFLoop", "", test_scf::float_types) {
+    using float_type = TestType;
+    using wf_type    = simde::type::rscf_wf;
+    auto mm          = test_scf::load_modules<float_type>();
+    auto& mod        = mm.at("Loop");
+
+    tensorwrapper::allocator::Eigen<float_type> alloc(mm.get_runtime());
+    tensorwrapper::shape::Smooth shape_corr{};
+    auto pcorr = alloc.allocate(tensorwrapper::layout::Physical(shape_corr));
+    using tensorwrapper::operations::approximately_equal;
 
     using index_set = typename wf_type::orbital_index_set_type;
-    wf_type psi0(index_set{0}, test_scf::h2_cmos());
+    wf_type psi0(index_set{0}, test_scf::h2_cmos<float_type>());
 
     auto H = test_scf::h2_hamiltonian();
 
     chemist::braket::BraKet H_00(psi0, H, psi0);
-    const auto& [e, psi] = mod.run_as<pt<wf_type>>(H_00, psi0);
-    REQUIRE_THAT(e, WithinAbs(-1.1167592336, 1E-6));
+    const auto& [e, psi] = mod.template run_as<pt<wf_type>>(H_00, psi0);
+
+    pcorr->at() = -1.1167592336;
+    tensorwrapper::Tensor corr(shape_corr, std::move(pcorr));
+    REQUIRE(approximately_equal(corr, e, 1E-6));
 }