Skip to content

Updates error model #164

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
ryanmrichard merged 13 commits into from
Apr 2, 2026
Merged

Updates error model #164

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
13 commits
Select commit Hold shift + click to select a range
cce3ab5
module for getting unnormalized decontracted ERIs
ryanmrichard Mar 31, 2026
73bec78
backup
ryanmrichard Mar 31, 2026
570bb9c
primitive eris work
ryanmrichard Mar 31, 2026
a62285b
backup
ryanmrichard Mar 31, 2026
e0c2150
can match libint
ryanmrichard Mar 31, 2026
450e0f4
Merge branch 'master' into test_prim_screening
ryanmrichard Mar 31, 2026
cbefd32
factor out primitive mapping
ryanmrichard Mar 31, 2026
079d7ff
coarse k refactored
ryanmrichard Apr 1, 2026
5723a76
screening works
ryanmrichard Apr 1, 2026
6c3d428
fix docs
ryanmrichard Apr 1, 2026
4ac15b4
update error model
ryanmrichard Apr 1, 2026
2ff8453
Merge branch 'master' into update_error_model
ryanmrichard Apr 1, 2026
fe7306d
adds missing fabs
ryanmrichard Apr 2, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
16 changes: 11 additions & 5 deletions examples/primitive_error_models/main.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -19,8 +19,11 @@

/* This example showcases how to:
*
* 1. Compute the analytic error in an ERI4 integral tensor owing to primitive
* pair screening.
* 1. Compute the analytic error in an ERI4 tensor from changing the Libint
* integral threshold (benchmark vs screened ERI4).
* 2. Compute a primitive-screening uncertainty model aligned with
* PrimitiveContractor (see module "Error estimate": Tolerance, Coarse,
* Fine).
*/

namespace {
Expand Down Expand Up @@ -78,13 +81,16 @@ int main(int argc, char* argv[]) {
// Make BraKet
chemist::braket::BraKet mnls(aos2, op, aos2);

// Compute the error by screening with tolerance "tol"
// Analytic: difference between unscreened and Libint-screened ERI4 at tol.
// Primitive error model: sum of per-skipped-quartet estimates (default
// "Tolerance" mode adds tol per skip; use module input "Error estimate").
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;
std::cout << "Analytic error (integral threshold): " << error << std::endl;
std::cout << "Primitive screening error model: " << approx_error
<< std::endl;

return 0;
}
3 changes: 0 additions & 3 deletions src/integrals/libint/libint.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -69,9 +69,6 @@ EXTERN_LIBINT(aos_squared, v_ee_type, aos_squared);
#undef EXTERN_LIBINT

void set_defaults(pluginplay::ModuleManager& mm) {
mm.change_submod("Primitive Error Model",
"Black Box Primitive Pair Estimator",
"Black Box Primitive Pair Estimator");
mm.change_submod("CauchySchwarz Estimator", "Decontract Basis Set",
"Decontract Basis Set");
mm.copy_module("ERI4", "Benchmark ERI4");
Expand Down
1 change: 1 addition & 0 deletions src/integrals/libint/primitive_contractor.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -161,6 +161,7 @@ MODULE_RUN(PrimitiveContractor) {
const auto gamma_ij = gamma_bra[pi][pj];
const auto Q_ij = Q_bra[pi][pj];

if(K_ij < thresh) continue;
for(std::size_t pshell_k = 0; pshell_k < nprims[2]; ++pshell_k) {
const auto lam = map2[pshell_k];
const auto pk = pmap2[pshell_k];
Expand Down
293 changes: 127 additions & 166 deletions src/integrals/libint/primitive_error_model.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -14,209 +14,170 @@
* limitations under the License.
*/

#include "../utils/primitive_index_helpers.hpp"
#include "detail_/primitive_pair_estimators.hpp"
#include "libint.hpp"
#include <array>
#include <cmath>
#include <integrals/integrals.hpp>
#include <stdexcept>
#include <string>

namespace integrals::libint {
namespace {

const auto desc = "Compute uncertainty estimates for ERI4 integrals";
const auto desc = R"(
Primitive Error Model
=====================

// Sums contributions Q_abQ_cd when Q_ab or Q_cd is below tol
template<typename ShellsType, typename OffsetTypes, typename TensorType>
auto shell_block_error(ShellsType&& shells, OffsetTypes&& prim_offsets,
TensorType&& K_ab, TensorType&& K_cd, double tol,
bool use_tol) {
// Get the number of primitives in each shell
const auto n_prim0 = shells[0].n_primitives();
const auto n_prim1 = shells[1].n_primitives();
const auto n_prim2 = shells[2].n_primitives();
const auto n_prim3 = shells[3].n_primitives();
Uncertainty tensor for ERI4 primitive screening aligned with
`PrimitiveContractor`: same decontracted primitive-AO quadruple loop and the
same coarse / fine gates as libint `ScreeningMethod::Original` (coarse
`coarse_k_ij`, fine `fine_k_ij` with `gamma_ij`).

using float_type = double; // TODO: get from Q_ab or Q_cd
using iter_type = std::decay_t<decltype(n_prim0)>; // Type of the loop index
For each decontracted quartet that would be skipped by the contractor, this
module accumulates one of three per-quartet estimates into the corresponding
contracted AO element: fixed `Tolerance`, coarse pair product
`K_ij * K_kl`, or fine metric `|Q_ij Q_kl| / sqrt(gamma_ij + gamma_kl)`.
)";

// Create an array to hold the primitive indices
std::array<iter_type, 4> prim{0, 0, 0, 0};

// This is the accumulated error for this shell block
double error = 0.0;

using wtf::fp::float_cast;
for(prim[0] = 0; prim[0] < n_prim0; ++prim[0]) {
const auto p0 = prim_offsets[0] + prim[0];

for(prim[1] = 0; prim[1] < n_prim1; ++prim[1]) {
const auto p1 = prim_offsets[1] + prim[1];

// Was the Q_ab element neglected for primitive pair (p0, p1)?
std::vector i01{p0, p1};

auto K_01 = float_cast<float_type>(K_ab.get_elem(i01));
auto K_01_mag = std::fabs(K_01);
bool K_01_neglected = K_01_mag < tol;

for(prim[2] = 0; prim[2] < n_prim2; ++prim[2]) {
const auto p2 = prim_offsets[2] + prim[2];

for(prim[3] = 0; prim[3] < n_prim3; ++prim[3]) {
const auto p3 = prim_offsets[3] + prim[3];

// Was the Q_cd element neglected for pair (p2, p3)?
std::vector i23{p2, p3};
auto K_23 = float_cast<float_type>(K_cd.get_elem(i23));
auto K_23_mag = std::fabs(K_23);
bool K_23_neglected = K_23_mag < tol;
/** @return True iff `PrimitiveContractor` would `continue` (skip) this quartet.
*/
inline bool primitive_quartet_skipped(double K_ij, double K_kl, double Q_ij,
double Q_kl, double gamma_ij,
double gamma_kl, double thresh) {
if(K_ij < thresh) return true;
if(K_kl < thresh) return true;
if(K_ij * K_kl <= thresh) return true;
const double gamma_ijkl = gamma_ij + gamma_kl;
const double pfac = std::abs(Q_ij * Q_kl / std::sqrt(gamma_ijkl));
if(pfac < thresh) return true;
return false;
}

auto prod_neglected = K_01_mag * K_23_mag < tol;
enum class ErrorEstimateKind { Tolerance, Coarse, Fine };

// If either was neglected we pick up an error Q_01 * Q_23
if(K_01_neglected || K_23_neglected || prod_neglected) {
if(use_tol)
error += tol;
else
error += K_01_mag * K_23_mag;
}
} // End loop over prim[3]
} // End loop over prim[2]
} // End loop over prim[1]
} // End loop over prim[0]
return error;
inline ErrorEstimateKind parse_error_estimate(const std::string& s) {
if(s == "Tolerance") return ErrorEstimateKind::Tolerance;
if(s == "Coarse") return ErrorEstimateKind::Coarse;
if(s == "Fine") return ErrorEstimateKind::Fine;
throw std::invalid_argument(
"Primitive Error Model: \"Error estimate\" must be \"Tolerance\", "
"\"Coarse\", or \"Fine\"");
}

// Sets all AO integrals in the shell block to the computed error
template<typename ShellsType, typename OffsetTypes, typename TensorType>
void fill_ao_block(ShellsType&& shells, OffsetTypes&& ao_offsets,
TensorType&& t, double error) {
// Get the number of AOs in each shell
const auto n_aos0 = shells[0].size();
const auto n_aos1 = shells[1].size();
const auto n_aos2 = shells[2].size();
const auto n_aos3 = shells[3].size();

// Make an array to hold the AO indices
using iter_type = std::decay_t<decltype(n_aos0)>; // Type of the loop index
std::array<iter_type, 4> ao{0, 0, 0, 0};

for(ao[0] = 0; ao[0] < n_aos0; ++ao[0]) {
const auto a0 = ao_offsets[0] + ao[0];

for(ao[1] = 0; ao[1] < n_aos1; ++ao[1]) {
const auto a1 = ao_offsets[1] + ao[1];

for(ao[2] = 0; ao[2] < n_aos2; ++ao[2]) {
const auto a2 = ao_offsets[2] + ao[2];

for(ao[3] = 0; ao[3] < n_aos3; ++ao[3]) {
const auto a3 = ao_offsets[3] + ao[3];

std::vector index{a0, a1, a2, a3};
t.set_elem(index, error);
} // End loop over ao[3]
} // End loop over ao[2]
} // End loop over ao[1]
} // End loop over ao[0]
inline double skip_increment(ErrorEstimateKind kind, double thresh, double K_ij,
double K_kl, double Q_ij, double Q_kl,
double gamma_ij, double gamma_kl) {
switch(kind) {
case ErrorEstimateKind::Tolerance: return thresh;
case ErrorEstimateKind::Coarse: return K_ij * K_kl;
case ErrorEstimateKind::Fine:
return std::abs(Q_ij * Q_kl / std::sqrt(gamma_ij + gamma_kl));
}
return 0.0;
}

} // namespace

using eri4_pt = simde::ERI4;
using pair_estimator_pt = integrals::property_types::PrimitivePairEstimator;
using pt = integrals::property_types::Uncertainty<eri4_pt>;
using eri4_pt = simde::ERI4;
using pt = integrals::property_types::Uncertainty<eri4_pt>;

MODULE_CTOR(PrimitiveErrorModel) {
satisfies_property_type<pt>();
description(desc);
// TODO citation for Chemist paper

add_submodule<pair_estimator_pt>("Black Box Primitive Pair Estimator")
.set_description("The module used to estimate the contributions of "
"primitive pairs to the overall integral values");
add_input<bool>("Use Tol").set_default(true);
add_input<std::string>("Error estimate")
.set_default("Tolerance")
.set_description(
"Per skipped primitive quartet: \"Tolerance\" adds the screening "
"threshold; \"Coarse\" adds K_ij*K_kl; \"Fine\" adds the fine-screen "
"metric |Q_ij Q_kl|/sqrt(gamma_ij+gamma_kl).");
}

MODULE_RUN(PrimitiveErrorModel) {
const auto& [braket, tol] = pt::unwrap_inputs(inputs);
const auto use_tol = inputs.at("Use Tol").value<bool>();
const auto bra0 = braket.bra().first.ao_basis_set();
const auto bra1 = braket.bra().second.ao_basis_set();
const auto ket0 = braket.ket().first.ao_basis_set();
const auto ket1 = braket.ket().second.ao_basis_set();

// Get the K_ab and K_cd tensors (used to determine which primitives are
// neglected)
auto& estimator = submods.at("Black Box Primitive Pair Estimator");
const auto& K_ab_tw = estimator.run_as<pair_estimator_pt>(bra0, bra1);
const auto& K_cd_tw = estimator.run_as<pair_estimator_pt>(ket0, ket1);

// XXX: Workaround for needing the contiguous buffers to access elements
using tensorwrapper::buffer::make_contiguous;
const auto& K_ab = make_contiguous(K_ab_tw.buffer());
const auto& K_cd = make_contiguous(K_cd_tw.buffer());

// Get the number of AOs in each basis set
const auto n_mu = bra0.n_aos();
const auto n_nu = bra1.n_aos();
const auto n_lambda = ket0.n_aos();
const auto n_sigma = ket1.n_aos();

// Make the return buffer
using float_type = double; // TODO: get from Q_ab or Q_cd
tensorwrapper::shape::Smooth shape({n_mu, n_nu, n_lambda, n_sigma});
const auto kind =
parse_error_estimate(inputs.at("Error estimate").value<std::string>());

auto bra = braket.bra();
auto ket = braket.ket();

const auto& bs0 = bra.first.ao_basis_set();
const auto& bs1 = bra.second.ao_basis_set();
const auto& bs2 = ket.first.ao_basis_set();
const auto& bs3 = ket.second.ao_basis_set();

const auto K_bra = detail_::coarse_k_ij(bs0, bs1);
const auto K_ket = detail_::coarse_k_ij(bs2, bs3);
const auto gamma_bra = detail_::gamma_ij(bs0, bs1);
const auto gamma_ket = detail_::gamma_ij(bs2, bs3);
const auto Q_bra = detail_::fine_k_ij(bs0, bs1);
const auto Q_ket = detail_::fine_k_ij(bs2, bs3);

auto map0 = utils::build_prim_ao_to_cgto_map(bs0);
auto map1 = utils::build_prim_ao_to_cgto_map(bs1);
auto map2 = utils::build_prim_ao_to_cgto_map(bs2);
auto map3 = utils::build_prim_ao_to_cgto_map(bs3);

auto pmap0 = utils::build_prim_ao_to_prim_shell_map(bs0);
auto pmap1 = utils::build_prim_ao_to_prim_shell_map(bs1);
auto pmap2 = utils::build_prim_ao_to_prim_shell_map(bs2);
auto pmap3 = utils::build_prim_ao_to_prim_shell_map(bs3);

std::array<std::size_t, 4> naos{bs0.n_aos(), bs1.n_aos(), bs2.n_aos(),
bs3.n_aos()};
std::array<std::size_t, 4> nprims{map0.size(), map1.size(), map2.size(),
map3.size()};

using float_type = double;
tensorwrapper::shape::Smooth shape({naos[0], naos[1], naos[2], naos[3]});
std::vector<float_type> raw_buffer(shape.size(), 0.0);
tensorwrapper::buffer::Contiguous buffer(std::move(raw_buffer), shape);

// Make arrays to hold the indices and offsets
using iter_type = std::decay_t<decltype(n_mu)>; // Type of the loop index
using shell_view = decltype(bra0.shell(0));
std::array<iter_type, 4> n_shells{bra0.n_shells(), bra1.n_shells(),
ket0.n_shells(), ket1.n_shells()};
std::array<iter_type, 4> shell_i{0, 0, 0, 0};
std::array<iter_type, 4> prim_offset{0, 0, 0, 0};
std::array<iter_type, 4> ao_offsets{0, 0, 0, 0};
auto ao_offset = [&](std::size_t i, std::size_t j, std::size_t k,
std::size_t l) {
return i * naos[1] * naos[2] * naos[3] + j * naos[2] * naos[3] +
k * naos[3] + l;
};

// We'll collect the shells into this array
std::array<shell_view, 4> shells;
for(std::size_t pshell_i = 0; pshell_i < nprims[0]; ++pshell_i) {
const auto mu = map0[pshell_i];
const auto pi = pmap0[pshell_i];

for(shell_i[0] = 0; shell_i[0] < n_shells[0]; ++shell_i[0]) {
shells[0] = bra0.shell(shell_i[0]);
for(std::size_t pshell_j = 0; pshell_j < nprims[1]; ++pshell_j) {
const auto nu = map1[pshell_j];
const auto pj = pmap1[pshell_j];

prim_offset[1] = 0;
ao_offsets[1] = 0;
for(shell_i[1] = 0; shell_i[1] < n_shells[1]; ++shell_i[1]) {
shells[1] = bra1.shell(shell_i[1]);
const double K_ij = K_bra[pi][pj];
const double gamma_ij = gamma_bra[pi][pj];
const double Q_ij = Q_bra[pi][pj];

prim_offset[2] = 0;
ao_offsets[2] = 0;
for(shell_i[2] = 0; shell_i[2] < n_shells[2]; ++shell_i[2]) {
shells[2] = ket0.shell(shell_i[2]);
for(std::size_t pshell_k = 0; pshell_k < nprims[2]; ++pshell_k) {
const auto lam = map2[pshell_k];
const auto pk = pmap2[pshell_k];

prim_offset[3] = 0;
ao_offsets[3] = 0;
for(shell_i[3] = 0; shell_i[3] < n_shells[3]; ++shell_i[3]) {
shells[3] = ket1.shell(shell_i[3]);
for(std::size_t pshell_l = 0; pshell_l < nprims[3];
++pshell_l) {
const auto sig = map3[pshell_l];
const auto pl = pmap3[pshell_l];

auto shell_error = shell_block_error(
shells, prim_offset, K_ab, K_cd, tol, use_tol);
fill_ao_block(shells, ao_offsets, buffer, shell_error);
const double K_kl = K_ket[pk][pl];
const double Q_kl = Q_ket[pk][pl];
const double gamma_kl = gamma_ket[pk][pl];

prim_offset[3] += shells[3].n_primitives();
ao_offsets[3] += shells[3].size();
} // End loop over shell_i[3]

prim_offset[2] += shells[2].n_primitives();
ao_offsets[2] += shells[2].size();
} // End loop over shell_i[2]

prim_offset[1] += shells[1].n_primitives();
ao_offsets[1] += shells[1].size();
} // End loop over shell_i[1]
if(!primitive_quartet_skipped(K_ij, K_kl, Q_ij, Q_kl,
gamma_ij, gamma_kl, tol)) {
continue;
}

prim_offset[0] += shells[0].n_primitives();
ao_offsets[0] += shells[0].size();
} // End loop over shell_i[0]
const double inc = skip_increment(
kind, tol, K_ij, K_kl, Q_ij, Q_kl, gamma_ij, gamma_kl);
raw_buffer[ao_offset(mu, nu, lam, sig)] += inc;
}
}
}
}

tensorwrapper::buffer::Contiguous buffer(std::move(raw_buffer), shape);
simde::type::tensor error(shape, std::move(buffer));
auto result = results();
return pt::wrap_results(result, error);
Expand Down
2 changes: 1 addition & 1 deletion src/integrals/utils/uncertainty_reductions.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -38,7 +38,7 @@ auto geometric_mean(const ContainerType& values) {
std::size_t non_zero_count = 0;
for(const auto& val : values) {
if(val == 0) continue;
log_sum += std::log(val);
log_sum += std::log(std::fabs(val));
++non_zero_count;
}
if(non_zero_count == 0 || log_sum == 0) return float_type{0};
Expand Down
Loading
Loading