Skip to content

Commit f448338

Browse files
ryanmrichardryanmrichard
authored
Atom Blocking of QCUP (#158)
* non-symmetric case works * adds symmetric atom blocking
1 parent 9b4cd97 commit f448338

File tree

11 files changed

+1478
-17
lines changed

11 files changed

+1478
-17
lines changed

‎.gitignore‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,7 @@
2020
# These are directories used by IDEs for storing settings
2121
.idea/
2222
.vscode/
23+
.cache/
2324

2425
# These are common Python virtual enviornment directory names
2526
venv/

‎src/integrals/ao_integrals/ao_integrals.hpp‎

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -27,6 +27,8 @@ DECLARE_MODULE(KDensityFitted);
2727
DECLARE_MODULE(DFIntegral);
2828
DECLARE_MODULE(CoulombMetric);
2929
DECLARE_MODULE(UQDriver);
30+
DECLARE_MODULE(UQAtomBlockedDriver);
31+
DECLARE_MODULE(UQAtomSymmBlockedDriver);
3032

3133
inline void set_defaults(pluginplay::ModuleManager& mm) {
3234
mm.change_submod("AO integral driver", "Coulomb matrix",
@@ -41,6 +43,12 @@ inline void set_defaults(pluginplay::ModuleManager& mm) {
4143
"Coulomb Metric");
4244
mm.change_submod("UQ Driver", "ERIs", "ERI4");
4345
mm.change_submod("UQ Driver", "ERI Error", "Primitive Error Model");
46+
mm.change_submod("UQ Atom Blocked Driver", "ERIs", "ERI4");
47+
mm.change_submod("UQ Atom Blocked Driver", "ERI Error",
48+
"Primitive Error Model");
49+
mm.change_submod("UQ Atom Symm Blocked Driver", "ERIs", "ERI4");
50+
mm.change_submod("UQ Atom Symm Blocked Driver", "ERI Error",
51+
"Primitive Error Model");
4452
}
4553

4654
inline void load_modules(pluginplay::ModuleManager& mm) {
@@ -52,6 +60,8 @@ inline void load_modules(pluginplay::ModuleManager& mm) {
5260
mm.add_module<DFIntegral>("Density Fitting Integral");
5361
mm.add_module<CoulombMetric>("Coulomb Metric");
5462
mm.add_module<UQDriver>("UQ Driver");
63+
mm.add_module<UQAtomBlockedDriver>("UQ Atom Blocked Driver");
64+
mm.add_module<UQAtomSymmBlockedDriver>("UQ Atom Symm Blocked Driver");
5565
}
5666

5767
} // namespace integrals::ao_integrals

‎src/integrals/ao_integrals/uq_atom_blocked_driver.cpp‎

Lines changed: 213 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,213 @@
1+
/*
2+
* Copyright 2025 NWChemEx-Project
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* http://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
17+
#include "../utils/uncertainty_reductions.hpp"
18+
#include "ao_integrals.hpp"
19+
#include <integrals/integrals.hpp>
20+
#ifdef ENABLE_SIGMA
21+
#include <sigma/sigma.hpp>
22+
#endif
23+
24+
using namespace tensorwrapper;
25+
26+
namespace integrals::ao_integrals {
27+
namespace {
28+
29+
template<typename FloatType, typename T, typename Tensor>
30+
auto average_error(T&& strides, T&& nbf, T&& ao_i, Tensor&& error,
31+
utils::mean_type mean) {
32+
std::vector<FloatType> buffer;
33+
34+
for(std::size_t i = 0; i < nbf[0]; ++i) {
35+
auto ioffset = (ao_i[0] + i) * strides[0];
36+
37+
for(std::size_t j = 0; j < nbf[1]; ++j) {
38+
auto joffset = ioffset + (ao_i[1] + j) * strides[1];
39+
40+
for(std::size_t k = 0; k < nbf[2]; ++k) {
41+
auto koffset = joffset + (ao_i[2] + k) * strides[2];
42+
43+
for(std::size_t l = 0; l < nbf[3]; ++l) {
44+
auto loffset = koffset + (ao_i[3] + l) * strides[3];
45+
buffer.push_back(error[loffset]);
46+
}
47+
}
48+
}
49+
}
50+
51+
return utils::compute_mean(mean, buffer);
52+
}
53+
54+
template<typename FloatType, typename T, typename Tensor>
55+
void update_block(T&& strides, T&& nbf, T&& ao_i, std::vector<FloatType>& out,
56+
Tensor&& value, FloatType error) {
57+
for(std::size_t i = 0; i < nbf[0]; ++i) {
58+
auto ioffset = (ao_i[0] + i) * strides[0];
59+
60+
for(std::size_t j = 0; j < nbf[1]; ++j) {
61+
auto joffset = ioffset + (ao_i[1] + j) * strides[1];
62+
63+
for(std::size_t k = 0; k < nbf[2]; ++k) {
64+
auto koffset = joffset + (ao_i[2] + k) * strides[2];
65+
66+
for(std::size_t l = 0; l < nbf[3]; ++l) {
67+
auto loffset = koffset + (ao_i[3] + l) * strides[3];
68+
out[loffset] = value[loffset] + error;
69+
}
70+
}
71+
}
72+
}
73+
}
74+
75+
struct Kernel {
76+
using shape_type = buffer::Contiguous::shape_type;
77+
Kernel(shape_type shape, std::array<simde::type::ao_basis_set, 4> aos,
78+
utils::mean_type mean) :
79+
m_shape(std::move(shape)), m_aos(aos), m_mean(mean) {}
80+
81+
template<typename FloatType0, typename FloatType1>
82+
Tensor operator()(const std::span<FloatType0> t,
83+
const std::span<FloatType1> error) {
84+
throw std::runtime_error(
85+
"UQ Integrals Driver kernel only supports same float types");
86+
}
87+
88+
template<typename FloatType>
89+
auto operator()(const std::span<FloatType> t,
90+
const std::span<FloatType> error) {
91+
Tensor rv;
92+
93+
using float_type = std::decay_t<FloatType>;
94+
if constexpr(types::is_uncertain_v<float_type>) {
95+
throw std::runtime_error("Did not expect an uncertain type");
96+
} else {
97+
#ifdef ENABLE_SIGMA
98+
using tensorwrapper::buffer::make_contiguous;
99+
100+
std::array n_centers{m_aos[0].size(), m_aos[1].size(),
101+
m_aos[2].size(), m_aos[3].size()};
102+
103+
std::array<std::size_t, 4> centers{0, 0, 0, 0};
104+
std::array<std::size_t, 4> ao_i{0, 0, 0, 0};
105+
std::array<std::size_t, 4> nbf{0, 0, 0, 0};
106+
107+
using uq_type = sigma::Uncertain<float_type>;
108+
std::vector<uq_type> rv_data(m_shape.size());
109+
std::array<std::size_t, 4> strides{0, 0, 0, 1};
110+
strides[2] = strides[3] * m_aos[3].n_aos();
111+
strides[1] = strides[2] * m_aos[2].n_aos();
112+
strides[0] = strides[1] * m_aos[1].n_aos();
113+
114+
for(centers[0] = 0; centers[0] < n_centers[0]; ++centers[0]) {
115+
nbf[0] = m_aos[0][centers[0]].n_aos();
116+
117+
ao_i[1] = 0;
118+
for(centers[1] = 0; centers[1] < n_centers[1]; ++centers[1]) {
119+
nbf[1] = m_aos[1][centers[1]].n_aos();
120+
121+
ao_i[2] = 0;
122+
for(centers[2] = 0; centers[2] < n_centers[2];
123+
++centers[2]) {
124+
nbf[2] = m_aos[2][centers[2]].n_aos();
125+
126+
ao_i[3] = 0;
127+
for(centers[3] = 0; centers[3] < n_centers[3];
128+
++centers[3]) {
129+
nbf[3] = m_aos[3][centers[3]].n_aos();
130+
131+
auto block_error = average_error<float_type>(
132+
strides, nbf, ao_i, error, m_mean);
133+
uq_type max_uq{0.0, block_error};
134+
135+
update_block(strides, nbf, ao_i, rv_data, t,
136+
max_uq);
137+
138+
ao_i[3] += nbf[3];
139+
}
140+
ao_i[2] += nbf[2];
141+
}
142+
ao_i[1] += nbf[1];
143+
}
144+
ao_i[0] += nbf[0];
145+
}
146+
tensorwrapper::buffer::Contiguous t_w_contig(std::move(rv_data),
147+
m_shape);
148+
rv = tensorwrapper::Tensor(m_shape, std::move(t_w_contig));
149+
#else
150+
throw std::runtime_error("Sigma support not enabled!");
151+
#endif
152+
}
153+
154+
return rv;
155+
}
156+
shape_type m_shape;
157+
std::array<simde::type::ao_basis_set, 4> m_aos;
158+
utils::mean_type m_mean;
159+
};
160+
161+
const auto desc = R"(
162+
UQ Integrals Driver
163+
-------------------
164+
165+
)";
166+
167+
} // namespace
168+
169+
using eri_pt = simde::ERI4;
170+
using error_pt = integrals::property_types::Uncertainty<eri_pt>;
171+
172+
MODULE_CTOR(UQAtomBlockedDriver) {
173+
satisfies_property_type<eri_pt>();
174+
description(desc);
175+
add_submodule<eri_pt>("ERIs");
176+
add_submodule<error_pt>("ERI Error");
177+
add_input<std::string>("Mean Type").set_default("none");
178+
}
179+
180+
MODULE_RUN(UQAtomBlockedDriver) {
181+
const auto& [braket] = eri_pt::unwrap_inputs(inputs);
182+
auto mean_str = inputs.at("Mean Type").value<std::string>();
183+
auto mean = utils::mean_from_string(mean_str);
184+
185+
auto& eri_mod = submods.at("ERIs").value();
186+
auto tol = eri_mod.inputs().at("Threshold").value<double>();
187+
188+
const auto& t = eri_mod.run_as<eri_pt>(braket);
189+
const auto& error = submods.at("ERI Error").run_as<error_pt>(braket, tol);
190+
191+
using tensorwrapper::buffer::make_contiguous;
192+
const auto& t_buffer = make_contiguous(t.buffer());
193+
const auto& e_buffer = make_contiguous(error.buffer());
194+
195+
const auto& bra = braket.bra();
196+
const auto& ket = braket.ket();
197+
const auto& mu = bra.first.ao_basis_set();
198+
const auto& nu = bra.second.ao_basis_set();
199+
const auto& lam = ket.first.ao_basis_set();
200+
const auto& sig = ket.second.ao_basis_set();
201+
202+
std::array aos{mu, nu, lam, sig};
203+
204+
using buffer::visit_contiguous_buffer;
205+
shape::Smooth shape = t.buffer().layout().shape().as_smooth().make_smooth();
206+
207+
Kernel k(shape, aos, mean);
208+
auto t_w_error = visit_contiguous_buffer(k, t_buffer, e_buffer);
209+
210+
auto rv = results();
211+
return eri_pt::wrap_results(rv, t_w_error);
212+
}
213+
} // namespace integrals::ao_integrals

0 commit comments

Comments
 (0)