Module for using analytic ERI error as UQ

src/integrals/ao_integrals/ao_integrals.hpp

@@ -26,6 +26,7 @@ DECLARE_MODULE(JDensityFitted);
 DECLARE_MODULE(KDensityFitted);
 DECLARE_MODULE(DFIntegral);
 DECLARE_MODULE(CoulombMetric);
+DECLARE_MODULE(UQDriver);
 
 inline void set_defaults(pluginplay::ModuleManager& mm) {
     mm.change_submod("AO integral driver", "Coulomb matrix",
@@ -38,6 +39,7 @@ inline void set_defaults(pluginplay::ModuleManager& mm) {
                      "Density Fitting Integral");
     mm.change_submod("Density Fitting Integral", "Coulomb Metric",
                      "Coulomb Metric");
+    mm.change_submod("UQ Driver", "ERIs", "ERI4");
 }
 
 inline void load_modules(pluginplay::ModuleManager& mm) {
@@ -48,7 +50,7 @@ inline void load_modules(pluginplay::ModuleManager& mm) {
     mm.add_module<KDensityFitted>("Density Fitted K builder");
     mm.add_module<DFIntegral>("Density Fitting Integral");
     mm.add_module<CoulombMetric>("Coulomb Metric");
-    set_defaults(mm);
+    mm.add_module<UQDriver>("UQ Driver");
 }
 
 } // namespace integrals::ao_integrals

src/integrals/ao_integrals/uq_driver.cpp

@@ -0,0 +1,96 @@
+/*
+ * 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 "ao_integrals.hpp"
+
+namespace integrals::ao_integrals {
+namespace {
+
+struct Kernel {
+    using buffer_base_type = tensorwrapper::buffer::BufferBase;
+    template<typename FloatType>
+    auto run(const buffer_base_type& t, const buffer_base_type& error) {
+        tensorwrapper::Tensor rv;
+
+        if constexpr(tensorwrapper::types::is_uncertain_v<FloatType>) {
+            using alloc_type = tensorwrapper::allocator::Eigen<FloatType>;
+            alloc_type alloc(t.allocator().runtime());
+
+            const auto& t_eigen     = alloc.rebind(t);
+            const auto& error_eigen = alloc.rebind(error);
+
+            auto rv_buffer = alloc.allocate(t_eigen.layout());
+            for(std::size_t i = 0; i < t_eigen.size(); ++i) {
+                const auto elem          = (t_eigen.data() + i)->mean();
+                const auto elem_error    = (error_eigen.data() + i)->mean();
+                *(rv_buffer->data() + i) = FloatType(elem, elem_error);
+            }
+
+            const auto& shape = t_eigen.layout().shape();
+            rv = tensorwrapper::Tensor(shape, std::move(rv_buffer));
+        } else {
+            throw std::runtime_error("Expected an uncertain type");
+        }
+        return rv;
+    }
+};
+
+const auto desc = R"(
+UQ Integrals Driver
+-------------------
+
+)";
+
+} // namespace
+
+using eri_pt = simde::ERI4;
+
+MODULE_CTOR(UQDriver) {
+    satisfies_property_type<eri_pt>();
+    description(desc);
+    add_submodule<eri_pt>("ERIs");
+    add_input<double>("benchmark precision").set_default(1.0e-16);
+    add_input<double>("precision").set_default(1.0e-16);
+}
+
+MODULE_RUN(UQDriver) {
+    auto tau_0 = inputs.at("benchmark precision").value<double>();
+    auto tau   = inputs.at("precision").value<double>();
+
+    auto& eri_mod = submods.at("ERIs").value();
+
+    auto benchmark_mod = eri_mod.unlocked_copy();
+    benchmark_mod.change_input("Threshold", tau_0);
+    benchmark_mod.change_input("With UQ?", true);
+
+    auto normal_mod = eri_mod.unlocked_copy();
+    normal_mod.change_input("Threshold", tau);
+    normal_mod.change_input("With UQ?", true);
+
+    const auto& [t_0] = eri_pt::unwrap_results(benchmark_mod.run(inputs));
+    const auto& [t]   = eri_pt::unwrap_results(normal_mod.run(inputs));
+
+    simde::type::tensor error;
+    error("m,n,l,s") = t("m,n,l,s") - t_0("m,n,l,s");
+
+    using tensorwrapper::utilities::floating_point_dispatch;
+    Kernel k;
+    auto t_w_error = floating_point_dispatch(k, t.buffer(), error.buffer());
+
+    auto rv = results();
+    return eri_pt::wrap_results(rv, t_w_error);
+}
+} // namespace integrals::ao_integrals

src/integrals/integrals_mm.cpp

@@ -40,6 +40,7 @@ void load_modules(pluginplay::ModuleManager& mm) {
     ao_integrals::load_modules(mm);
     libint::load_modules(mm);
     set_defaults(mm);
+    ao_integrals::set_defaults(mm);
 }
 
 } // namespace integrals

tests/cxx/unit/integrals/ao_integrals/test_uq_driver.cpp

@@ -0,0 +1,79 @@
+/*
+ * Copyright 2022 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.hpp"
+
+template<typename FloatType>
+auto corr_answer(const simde::type::tensor& T) {
+    if constexpr(std::is_same_v<FloatType, double>) {
+        return T;
+    } else {
+        simde::type::tensor T_corr(T);
+        using alloc_type           = tensorwrapper::allocator::Eigen<FloatType>;
+        auto& corr_buffer          = alloc_type::rebind(T_corr.buffer());
+        corr_buffer.at(0, 0, 0, 0) = FloatType{0.774606, 0};
+        corr_buffer.at(0, 0, 0, 1) = FloatType{0.265558, 2.49687e-06};
+        corr_buffer.at(0, 0, 1, 0) = FloatType{0.265558, 2.49687e-06};
+        corr_buffer.at(0, 0, 1, 1) = FloatType{0.446701, 0};
+        corr_buffer.at(0, 1, 0, 0) = FloatType{0.265558, 2.49687e-06};
+        corr_buffer.at(0, 1, 0, 1) = FloatType{0.120666, 1.10748e-05};
+        corr_buffer.at(0, 1, 1, 0) = FloatType{0.120666, 1.10748e-05};
+        corr_buffer.at(0, 1, 1, 1) = FloatType{0.265558, 2.49687e-06};
+        corr_buffer.at(1, 0, 0, 0) = FloatType{0.265558, 2.49687e-06};
+        corr_buffer.at(1, 0, 0, 1) = FloatType{0.120666, 1.10748e-05};
+        corr_buffer.at(1, 0, 1, 0) = FloatType{0.120666, 1.10748e-05};
+        corr_buffer.at(1, 0, 1, 1) = FloatType{0.265558, 2.49687e-06};
+        corr_buffer.at(1, 1, 0, 0) = FloatType{0.446701, 0};
+        corr_buffer.at(1, 1, 0, 1) = FloatType{0.265558, 2.49687e-06};
+        corr_buffer.at(1, 1, 1, 0) = FloatType{0.265558, 2.49687e-06};
+        corr_buffer.at(1, 1, 1, 1) = FloatType{0.774606, 0};
+        return T_corr;
+    }
+}
+
+TEST_CASE("UQ Driver") {
+    using float_type = tensorwrapper::types::udouble;
+    if constexpr(tensorwrapper::types::is_uncertain_v<float_type>) {
+        using test_pt = simde::ERI4;
+
+        auto rt = std::make_unique<parallelzone::runtime::RuntimeView>();
+        pluginplay::ModuleManager mm(std::move(rt), nullptr);
+        integrals::load_modules(mm);
+        REQUIRE(mm.count("UQ Driver"));
+
+        mm.change_input("UQ Driver", "precision", 1.0e-6);
+
+        // Get basis set
+        auto mol  = test::h2_molecule();
+        auto aobs = test::h2_sto3g_basis_set();
+
+        // Make AOS object
+        simde::type::aos aos(aobs);
+        simde::type::aos_squared aos_squared(aos, aos);
+
+        // Make Operator
+        simde::type::v_ee_type op{};
+
+        // Make BraKet Input
+        chemist::braket::BraKet braket(aos_squared, op, aos_squared);
+
+        // Call module
+        auto T      = mm.at("UQ Driver").run_as<test_pt>(braket);
+        auto T_corr = corr_answer<float_type>(T);
+        using tensorwrapper::operations::approximately_equal;
+        REQUIRE(approximately_equal(T_corr, T, 1E-6));
+    }
+}