Adds a UQ model for primitive screening (#155)

* backup

* backukp

* backup

* refactor estimator

* uses normalized coeffs, refactor testing infrastructure

* cs screening seems to work

* backup

* module for screening pairs

* finally r2g

* run precommit/add missing header

* ifdef protect sigma code

* blank lines for precommit

* address comments

* call utils::set_defaults

Author: ryanmrichard

CMakeLists.txt

@@ -72,6 +72,12 @@ cmaize_add_library(
     DEPENDS "${project_depends}"
 )
 
+cmaize_add_executable(
+    primitive_error_models
+    SOURCE_DIR "examples/primitive_error_models"
+    DEPENDS "${PROJECT_NAME}"
+)
+
 include(nwx_pybind11)
 nwx_add_pybind11_module(
     ${PROJECT_NAME}

examples/primitive_error_models/main.cpp

@@ -0,0 +1,90 @@
+/*
+ * Copyright 2026 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 <integrals/integrals.hpp>
+#include <simde/simde.hpp>
+
+/* This example showcases how to:
+ *
+ *  1. Compute the analytic error in an ERI4 integral tensor owing to primitive
+ *     pair screening.
+ */
+
+namespace {
+
+// This makes a basis set for H2 (bond distance 1.40 a.u.) using STO-3G.
+inline simde::type::ao_basis_set h2_sto3g_basis_set() {
+    using ao_basis_t     = simde::type::ao_basis_set;
+    using atomic_basis_t = simde::type::atomic_basis_set;
+    using cg_t           = simde::type::contracted_gaussian;
+    using point_t        = simde::type::point;
+    using doubles_t      = std::vector<double>;
+
+    point_t r0{0.0, 0.0, 0.0};
+    point_t r1{0.0, 0.0, 1.40};
+
+    doubles_t cs{0.1543289673, 0.5353281423, 0.4446345422};
+    doubles_t es{3.425250914, 0.6239137298, 0.1688554040};
+    cg_t cg0(cs.begin(), cs.end(), es.begin(), es.end(), r0);
+    cg_t cg1(cs.begin(), cs.end(), es.begin(), es.end(), r1);
+    atomic_basis_t h0("sto-3g", 1, r0);
+    atomic_basis_t h1("sto-3g", 1, r1);
+    h0.add_shell(chemist::ShellType::cartesian, 0, cg0);
+    h1.add_shell(chemist::ShellType::cartesian, 0, cg1);
+
+    ao_basis_t bs;
+    bs.add_center(h0);
+    bs.add_center(h1);
+    return bs;
+}
+
+} // namespace
+
+// Property types for the ERI4 and the error in the ERI4
+using eri4_pt       = simde::ERI4;
+using eri4_error_pt = integrals::property_types::Uncertainty<eri4_pt>;
+
+int main(int argc, char* argv[]) {
+    // Makes sure the environment doesn't go out of scope before the end.
+    auto rt = std::make_unique<parallelzone::runtime::RuntimeView>();
+
+    // Initializes a ModuleManager object with the integrals plugin
+    pluginplay::ModuleManager mm(std::move(rt), nullptr);
+    integrals::load_modules(mm);
+    integrals::set_defaults(mm);
+
+    // Modules for computing analytic error and estimating error
+    auto& analytic_error_mod = mm.at("Analytic Error");
+    auto& error_model        = mm.at("Primitive Error Model");
+
+    // Makes: basis set, direct product of the basis set, and 1/r12 operator
+    simde::type::aos aos(h2_sto3g_basis_set());
+    simde::type::aos_squared aos2(aos, aos);
+    simde::type::v_ee_type op{};
+
+    // Make BraKet
+    chemist::braket::BraKet mnls(aos2, op, aos2);
+
+    // Compute the error by screening with tolerance "tol"
+    double tol        = 1E-10;
+    auto error        = analytic_error_mod.run_as<eri4_error_pt>(mnls, tol);
+    auto approx_error = error_model.run_as<eri4_error_pt>(mnls, tol);
+
+    std::cout << "Analytic error: " << error << std::endl;
+    std::cout << "Estimated error: " << approx_error << std::endl;
+
+    return 0;
+}

include/integrals/integrals.hpp

@@ -22,7 +22,8 @@
  *  your unit test also needs most of the headers included by it).
  */
 #pragma once
-#include "integrals/integrals_mm.hpp"
+#include <integrals/integrals_mm.hpp>
+#include <integrals/property_types.hpp>
 
 /** @namespace integrals
  *

include/integrals/integrals_mm.hpp

@@ -27,4 +27,6 @@ namespace integrals {
  */
 DECLARE_PLUGIN(integrals);
 
+void set_defaults(pluginplay::ModuleManager&);
+
 } // end namespace integrals

include/integrals/property_types.hpp

@@ -30,6 +30,64 @@
  */
 namespace integrals::property_types {
 
+// PT used to estimate the contribution of primitive pairs
+DECLARE_PROPERTY_TYPE(PrimitivePairEstimator);
+PROPERTY_TYPE_INPUTS(PrimitivePairEstimator) {
+    using ao_basis = const simde::type::ao_basis_set&;
+    auto rv        = pluginplay::declare_input()
+                .add_field<ao_basis>("Bra Basis Set")
+                .add_field<ao_basis>("Ket Basis Set");
+    rv["Bra Basis Set"].set_description(
+      "The atomic orbital basis set for the bra");
+    rv["Ket Basis Set"].set_description(
+      "The atomic orbital basis set for the ket");
+    return rv;
+}
+
+PROPERTY_TYPE_RESULTS(PrimitivePairEstimator) {
+    using tensor = simde::type::tensor;
+    auto rv      = pluginplay::declare_result().add_field<tensor>(
+      "Primitive Pair Estimates");
+    rv["Primitive Pair Estimates"].set_description(
+      "A tensor containing the estimated values for each primitive pair "
+      "integral");
+    return rv;
+}
+
+DECLARE_PROPERTY_TYPE(PairScreener);
+PROPERTY_TYPE_INPUTS(PairScreener) {
+    using ao_basis = const simde::type::ao_basis_set&;
+    auto rv        = pluginplay::declare_input()
+                .add_field<ao_basis>("Bra Basis Set")
+                .add_field<ao_basis>("Ket Basis Set")
+                .add_field<double>("Tolerance");
+    return rv;
+}
+
+PROPERTY_TYPE_RESULTS(PairScreener) {
+    using index_vector = std::vector<std::size_t>;
+    using pair_vector  = std::vector<index_vector>;
+    auto rv            = pluginplay::declare_result().add_field<pair_vector>(
+      "Primitive Pairs Passing Screening");
+    return rv;
+}
+
+template<typename BasePT>
+DECLARE_TEMPLATED_PROPERTY_TYPE(Uncertainty, BasePT);
+
+template<typename BasePT>
+TEMPLATED_PROPERTY_TYPE_INPUTS(Uncertainty, BasePT) {
+    auto rv  = BasePT::inputs();
+    auto rv0 = rv.template add_field<double>("Tolerance");
+    rv0["Tolerance"].set_description("The screening threshold");
+    return rv0;
+}
+
+template<typename BasePT>
+TEMPLATED_PROPERTY_TYPE_RESULTS(Uncertainty, BasePT) {
+    return BasePT::results();
+}
+
 using DecontractBasisSet =
   simde::Convert<simde::type::ao_basis_set, simde::type::ao_basis_set>;
 

src/integrals/ao_integrals/ao_integrals.hpp

@@ -40,6 +40,7 @@ inline void set_defaults(pluginplay::ModuleManager& mm) {
     mm.change_submod("Density Fitting Integral", "Coulomb Metric",
                      "Coulomb Metric");
     mm.change_submod("UQ Driver", "ERIs", "ERI4");
+    mm.change_submod("UQ Driver", "ERI Error", "Primitive Error Model");
 }
 
 inline void load_modules(pluginplay::ModuleManager& mm) {

src/integrals/ao_integrals/uq_driver.cpp

@@ -15,6 +15,10 @@
  */
 
 #include "ao_integrals.hpp"
+#include <integrals/integrals.hpp>
+#ifdef ENABLE_SIGMA
+#include <sigma/sigma.hpp>
+#endif
 
 using namespace tensorwrapper;
 
@@ -37,19 +41,25 @@ struct Kernel {
                     const std::span<FloatType> error) {
         Tensor rv;
 
-        if constexpr(types::is_uncertain_v<FloatType>) {
-            auto rv_buffer = buffer::make_contiguous<FloatType>(m_shape);
-            auto rv_data   = buffer::get_raw_data<FloatType>(rv_buffer);
+        using float_type = std::decay_t<FloatType>;
+        if constexpr(types::is_uncertain_v<float_type>) {
+            throw std::runtime_error("Did not expect an uncertain type");
+        } else {
+#ifdef ENABLE_SIGMA
+            using uq_type  = sigma::Uncertain<float_type>;
+            auto rv_buffer = buffer::make_contiguous<uq_type>(m_shape);
+            auto rv_data   = buffer::get_raw_data<uq_type>(rv_buffer);
             for(std::size_t i = 0; i < t.size(); ++i) {
-                const auto elem       = t[i].mean();
-                const auto elem_error = error[i].mean();
-                rv_data[i]            = FloatType(elem, elem_error);
+                const auto elem       = t[i];
+                const auto elem_error = error[i];
+                rv_data[i]            = uq_type(elem, elem_error);
             }
-
             rv = tensorwrapper::Tensor(m_shape, std::move(rv_buffer));
-        } else {
-            throw std::runtime_error("Expected an uncertain type");
+#else
+            throw std::runtime_error("Sigma support not enabled!");
+#endif
         }
+
         return rv;
     }
     shape_type m_shape;
@@ -63,35 +73,24 @@ UQ Integrals Driver
 
 } // namespace
 
-using eri_pt = simde::ERI4;
+using eri_pt   = simde::ERI4;
+using error_pt = integrals::property_types::Uncertainty<eri_pt>;
 
 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);
+    add_submodule<error_pt>("ERI Error");
 }
 
 MODULE_RUN(UQDriver) {
-    auto tau_0 = inputs.at("benchmark precision").value<double>();
-    auto tau   = inputs.at("precision").value<double>();
+    const auto& [braket] = eri_pt::unwrap_inputs(inputs);
 
     auto& eri_mod = submods.at("ERIs").value();
+    auto tol      = eri_mod.inputs().at("Threshold").value<double>();
 
-    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");
+    const auto& t     = eri_mod.run_as<eri_pt>(braket);
+    const auto& error = submods.at("ERI Error").run_as<error_pt>(braket, tol);
 
     using buffer::visit_contiguous_buffer;
     shape::Smooth shape = t.buffer().layout().shape().as_smooth().make_smooth();

src/integrals/integrals_mm.cpp

@@ -28,6 +28,9 @@ namespace integrals {
  *  @throw none No throw guarantee
  */
 void set_defaults(pluginplay::ModuleManager& mm) {
+    libint::set_defaults(mm);
+    ao_integrals::set_defaults(mm);
+    utils::set_defaults(mm);
     mm.change_submod("AO integral driver", "Kinetic", "Kinetic");
     mm.change_submod("AO integral driver", "Electron-Nuclear attraction",
                      "Nuclear");
@@ -41,8 +44,6 @@ void load_modules(pluginplay::ModuleManager& mm) {
     ao_integrals::load_modules(mm);
     libint::load_modules(mm);
     utils::load_modules(mm);
-    set_defaults(mm);
-    ao_integrals::set_defaults(mm);
 }
 
 } // namespace integrals

src/integrals/libint/analytic_error.cpp

@@ -0,0 +1,56 @@
+/*
+ * Copyright 2026 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 "libint.hpp"
+#include <integrals/integrals.hpp>
+
+namespace integrals::libint {
+namespace {
+
+const auto desc = "Uses the error in the ERI4 as the uncertainty.";
+}
+
+using eri4_pt = simde::ERI4;
+using pt      = integrals::property_types::Uncertainty<eri4_pt>;
+
+MODULE_CTOR(AnalyticError) {
+    satisfies_property_type<pt>();
+    description(desc);
+
+    add_submodule<eri4_pt>("ERI4s");
+}
+
+MODULE_RUN(AnalyticError) {
+    const auto& [braket, tol] = pt::unwrap_inputs(inputs);
+
+    auto& eri_mod = submods.at("ERI4s");
+
+    auto normal_mod = eri_mod.value().unlocked_copy();
+    normal_mod.change_input("Threshold", tol);
+
+    // N.b., t_0 is the benchmark value
+    const auto& t_0 = eri_mod.run_as<eri4_pt>(braket);
+    const auto& t   = normal_mod.run_as<eri4_pt>(braket);
+
+    simde::type::tensor error;
+    error("m,n,l,s") = t("m,n,l,s") - t_0("m,n,l,s");
+
+    // Wrap and return the results
+    auto rv = results();
+    return pt::wrap_results(rv, error);
+}
+
+} // namespace integrals::libint