diff --git a/cxx/include/integrals/property_types.hpp b/cxx/include/integrals/property_types.hpp index f12e060e..b5710457 100644 --- a/cxx/include/integrals/property_types.hpp +++ b/cxx/include/integrals/property_types.hpp @@ -21,6 +21,7 @@ * types are added in the future. */ #pragma once +#include #include #include @@ -111,4 +112,18 @@ TEMPLATED_PROPERTY_TYPE_RESULTS(Normalize, T) { return rv; } +// PT used to produce a UQFactory (declared in uq_factory.hpp) configured by +// the module's own "UQ Type" and "Order" inputs. Has no property-type-level +// inputs: the only configuration is module-specific (see uq_initializer.cpp). +DECLARE_PROPERTY_TYPE(UQInitializer); +PROPERTY_TYPE_INPUTS(UQInitializer) { return pluginplay::declare_input(); } + +PROPERTY_TYPE_RESULTS(UQInitializer) { + auto rv = pluginplay::declare_result().add_field("UQ Factory"); + rv["UQ Factory"].set_description( + "A callable, comparable factory that constructs UQ values of the " + "module-configured kind ."); + return rv; +} + } // end namespace integrals::property_types diff --git a/cxx/include/integrals/uq_factory.hpp b/cxx/include/integrals/uq_factory.hpp new file mode 100644 index 00000000..a49eb5c9 --- /dev/null +++ b/cxx/include/integrals/uq_factory.hpp @@ -0,0 +1,275 @@ +/* + * 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. + */ + +/** @file uq_factory.hpp + * + * Defines integrals::property_types::UQFactory, the result type of the + * UQInitializer property type (declared in property_types.hpp). + */ +#pragma once +#include +#include +#include +#include +#include + +namespace integrals::property_types { + +/// The UQ representations UQInitializer knows how to build a factory for. +enum class UQKind { + uncertain, + interval, + affine, + thresholded_affine, + taylor_model +}; + +/** @brief Parses a "UQ Type" module-input string into a UQKind. + * + * This is the single place the accepted "UQ Type" string literals are + * spelled out; UQInitializer and UQAtomSymmBlockedDriver's tests both rely + * on this mapping staying in sync with UQFactory's construction logic. + * + * @throw std::runtime_error if @p uq_type does not match a known UQ kind. + */ +inline UQKind uq_kind_from_string(const std::string& uq_type) { + if(uq_type == "uncertain") return UQKind::uncertain; + if(uq_type == "interval") return UQKind::interval; + if(uq_type == "affine") return UQKind::affine; + if(uq_type == "thresholded affine") return UQKind::thresholded_affine; + if(uq_type == "taylor model") return UQKind::taylor_model; + throw std::runtime_error( + "integrals::property_types::uq_kind_from_string: Invalid UQ type name " + + uq_type); +} + +/// Inverse of uq_kind_from_string, kept alongside it for symmetry/debugging. +inline std::string to_string(UQKind kind) { + switch(kind) { + case UQKind::uncertain: return "uncertain"; + case UQKind::interval: return "interval"; + case UQKind::affine: return "affine"; + case UQKind::thresholded_affine: return "thresholded affine"; + case UQKind::taylor_model: return "taylor model"; + } + throw std::runtime_error( + "integrals::property_types::to_string(UQKind): Unrecognized UQKind"); +} + +/** @brief Polymorphic base class for constructing a UQ scalar from a center + * and a radius. + * + * The UQFactorBase class is used as a PIMPL for UQFactory (below) so that + * + * Each UQ representation (sigma::Uncertain, sigma::Interval, sigma::Affine, + * sigma::ThresholdedAffine, sigma::TaylorModel) is an independent template + * class with no shared base in `sigma`/`tensorwrapper` (they're unified only + * by compile-time traits). This class supplies the runtime-polymorphic base + * UQInitializer needs so it can hand back "a factory for whichever UQ kind + * the caller asked for" as a single result type, and so downstream code + * (UQAtomSymmBlockedDriver's Kernel) can construct UQ values without a + * compile-time dependence on which kind was selected. + * + * N.b., only `double` centers/radii are supported at present. If/when needed + * the API could be extended to other types. + */ +class UQFactoryBase { +public: + /// Virtual dtor + virtual ~UQFactoryBase() = default; + + /// Constructs a UQ scalar of *this's kind, centered at @p center with + /// radius @p radius, type-erased into a wtf::fp::Float. + virtual wtf::fp::Float operator()(double center, double radius) const = 0; + + /// Identifies which UQ representation *this constructs. + virtual UQKind kind() const noexcept = 0; + + /// Deep-copies *this. Kept for forward compatibility even though + /// UQFactory (below) uses shared, not owning, copies. + virtual std::unique_ptr clone() const = 0; + + /// True if *this and @p rhs would construct value-equal UQ scalars. + virtual bool equal(const UQFactoryBase& rhs) const noexcept { + return kind() == rhs.kind(); + } + + /// Imposes a strict weak ordering, for PluginPlay's AnyField. + virtual bool less(const UQFactoryBase& rhs) const noexcept { + return kind() < rhs.kind(); + } +}; + +namespace detail { + +/// Implements UQFactory for all UQ kinds except TaylorModel. +template typename UQType> +class SimpleFactoryImpl final : public UQFactoryBase { +public: + /// Dispatches to tensorwrapper::types::construct_uq_type + wtf::fp::Float operator()(double center, double radius) const override { + return wtf::fp::Float( + tensorwrapper::types::construct_uq_type>(center, + radius)); + } + + /// Returns the UQKind template parameter. + UQKind kind() const noexcept override { return Kind; } + + /// Returns a deep-copy of *this. + std::unique_ptr clone() const override { + return std::make_unique(*this); + } +}; + +/// Convenience aliases for the 4 SimpleFactoryImpl specializations. +using UncertainFactoryImpl = + SimpleFactoryImpl; +using IntervalFactoryImpl = + SimpleFactoryImpl; +using AffineFactoryImpl = + SimpleFactoryImpl; +using ThresholdedAffineFactoryImpl = + SimpleFactoryImpl; + +/** @brief Constructs a TaylorModel scalar spanning [center - radius, + * center + radius] with a maximum order of @p order. + * + * When Sigma is disabled tensorwrapper::types::taylor_model_type is just + * @p T, which has no order to set, so we defer to TensorWrapper's generic + * construction (which simply returns a scalar carrying no uncertainty). The + * dispatch lives in a function template because a discarded `if constexpr` + * branch is only left uninstantiated inside a template. + */ +template +T make_taylor_model(double center, double radius, std::size_t order) { + if constexpr(tensorwrapper::types::is_taylor_model_v) { + return T(center - radius, center + radius, typename T::Order(order)); + } else { + return tensorwrapper::types::construct_uq_type(center, radius); + } +} + +/** @brief Specializes UQFactoryBase to construct TaylorModel UQ scalars. + * + * TaylorModel objects need an order parameter. This class stores that + * parameter and passes it to each TaylorModel upon construction. + */ +class TaylorModelFactoryImpl final : public UQFactoryBase { +public: + /// Takes the TaylorModel order to use when constructing UQ scalars. + explicit TaylorModelFactoryImpl(std::size_t order = 2) : m_order_(order) {} + + /// Calls ctor for tensorwrapper::types::taylor_model_type + wtf::fp::Float operator()(double center, double radius) const override { + using uq_t = tensorwrapper::types::taylor_model_type; + return wtf::fp::Float( + make_taylor_model(center, radius, m_order_)); + } + + /// Returns UQKind::taylor_model. + UQKind kind() const noexcept override { return UQKind::taylor_model; } + + /// Returns the TaylorModel order stored in *this. + std::size_t order() const noexcept { return m_order_; } + + // Deep-copy *this + std::unique_ptr clone() const override { + return std::make_unique(*this); + } + + /// Additionally compares the order + bool equal(const UQFactoryBase& rhs) const noexcept override { + if(rhs.kind() != UQKind::taylor_model) return false; + return m_order_ == + static_cast(rhs).m_order_; + } + + bool less(const UQFactoryBase& rhs) const noexcept override { + if(kind() != rhs.kind()) return kind() < rhs.kind(); + return m_order_ < + static_cast(rhs).m_order_; + } + +private: + std::size_t m_order_; +}; + +} // namespace detail + +/** @brief A copyable, comparable factory for constructing UQ values of a + * caller-configured kind. + * + * This is the public API for creating UQ Forms without needing to know what + * kind of UQ Form is being created + * + * This class wraps one in a std::shared_ptr: copies are + * shallow (the underlying Impl is immutable after construction), giving + * cheap value semantics without deep-cloning on every copy. + */ +class UQFactory { +public: + /// Defaults to an "uncertain" factory so PluginPlay's requirement that + /// results be default constructible is met without a null state. + UQFactory() : m_impl_(std::make_shared()) {} + + /// Constructs a factory for the UQ kind specified by @p kind, with optional + /// TaylorModel order @p order (ignored for non-TaylorModel kinds). + explicit UQFactory(UQKind kind, std::size_t order = 2) : + m_impl_(make_impl_(kind, order)) {} + + /// Main API for building the UQ value + wtf::fp::Float operator()(double center, double radius) const { + return (*m_impl_)(center, radius); + } + + /// Returns the UQKind *this will build + UQKind kind() const noexcept { return m_impl_->kind(); } + + /// True if *this and @p rhs would build value-equal UQ scalars. + bool operator==(const UQFactory& rhs) const noexcept { + return m_impl_->equal(*rhs.m_impl_); + } + + /// Imposes a strict weak ordering + bool operator<(const UQFactory& rhs) const noexcept { + return m_impl_->less(*rhs.m_impl_); + } + +private: + static std::shared_ptr make_impl_(UQKind kind, + std::size_t order) { + switch(kind) { + case UQKind::uncertain: + return std::make_shared(); + case UQKind::interval: + return std::make_shared(); + case UQKind::affine: + return std::make_shared(); + case UQKind::thresholded_affine: + return std::make_shared(); + case UQKind::taylor_model: + return std::make_shared(order); + } + throw std::runtime_error( + "integrals::property_types::UQFactory: Unrecognized UQKind"); + } + + std::shared_ptr m_impl_; +}; + +} // namespace integrals::property_types diff --git a/cxx/src/integrals/ao_integrals/ao_integrals.hpp b/cxx/src/integrals/ao_integrals/ao_integrals.hpp index 7087cc9e..eb8ae3ff 100644 --- a/cxx/src/integrals/ao_integrals/ao_integrals.hpp +++ b/cxx/src/integrals/ao_integrals/ao_integrals.hpp @@ -28,6 +28,7 @@ DECLARE_MODULE(DFIntegral); DECLARE_MODULE(CoulombMetric); DECLARE_MODULE(UQDriver); DECLARE_MODULE(UQAtomSymmBlockedDriver); +DECLARE_MODULE(UQInitializer); inline void set_defaults(pluginplay::ModuleManager& mm) { mm.change_submod("AO integral driver", "Coulomb matrix", @@ -45,6 +46,8 @@ inline void set_defaults(pluginplay::ModuleManager& mm) { mm.change_submod("UQ Atom Symm Blocked Driver", "ERIs", "ERI4"); mm.change_submod("UQ Atom Symm Blocked Driver", "ERI Error", "Primitive Error Model"); + mm.change_submod("UQ Atom Symm Blocked Driver", "UQ Initializer", + "UQ Initializer"); } inline void load_modules(pluginplay::ModuleManager& mm) { @@ -57,6 +60,7 @@ inline void load_modules(pluginplay::ModuleManager& mm) { mm.add_module("Coulomb Metric"); mm.add_module("UQ Driver"); mm.add_module("UQ Atom Symm Blocked Driver"); + mm.add_module("UQ Initializer"); } } // namespace integrals::ao_integrals diff --git a/cxx/src/integrals/ao_integrals/uq_atom_symm_blocked_driver.cpp b/cxx/src/integrals/ao_integrals/uq_atom_symm_blocked_driver.cpp index b59a651f..ecf9fc1b 100644 --- a/cxx/src/integrals/ao_integrals/uq_atom_symm_blocked_driver.cpp +++ b/cxx/src/integrals/ao_integrals/uq_atom_symm_blocked_driver.cpp @@ -26,9 +26,12 @@ using namespace tensorwrapper; namespace integrals::ao_integrals { namespace { +using wtf::buffer::FloatBuffer; + template auto average_error(T&& strides, T&& nbf, T&& ao_i, Tensor&& error, - utils::mean_type mean) { + utils::mean_type mean, + const integrals::property_types::UQFactory& factory = {}) { std::string error_base = "integrals::ao_integrals::UQAtomSymmBlockedDriver: "; @@ -37,10 +40,18 @@ auto average_error(T&& strides, T&& nbf, T&& ao_i, Tensor&& error, auto n_elements = nbf[0] * nbf[1] * nbf[2] * nbf[3]; - using tensorwrapper::types::construct_uq_type; + // The factory always returns a type-erased wtf::fp::Float (since it can + // be configured, at runtime, to build any of the 5 UQ representations), + // so every element is unwrapped to the concrete UQType requested by the + // caller via float_cast. + auto make_elem = [&](float_type ei) -> UQType { + auto elem = factory(float_type(0.0), ei); + return wtf::fp::float_cast(elem); + }; + if(mean == utils::mean_type::none) { - std::vector result; - result.reserve(n_elements); + FloatBuffer result; + result.reserve(n_elements); for(std::size_t i = 0; i < nbf[0]; ++i) { auto ioffset = (ao_i[0] + i) * strides[0]; for(std::size_t j = 0; j < nbf[1]; ++j) { @@ -50,8 +61,7 @@ auto average_error(T&& strides, T&& nbf, T&& ao_i, Tensor&& error, for(std::size_t l = 0; l < nbf[3]; ++l) { auto loffset = koffset + (ao_i[3] + l) * strides[3]; auto ei = std::fabs(error[loffset]); - auto elem = construct_uq_type(0.0, ei); - result.push_back(elem); + result.push_back(make_elem(ei)); } } } @@ -75,20 +85,22 @@ auto average_error(T&& strides, T&& nbf, T&& ao_i, Tensor&& error, } } auto mean_value = utils::compute_mean(mean, buffer); - auto value = construct_uq_type(0.0, mean_value); - return std::vector(n_elements, value); + auto value = make_elem(mean_value); + return FloatBuffer(std::vector(n_elements, value)); #else throw std::runtime_error(error_base + "Sigma support not enabled!"); - return std::vector{}; + return FloatBuffer{}; #endif } #ifdef ENABLE_SIGMA template auto compute_block(T&& strides, T&& nbf, T&& ao_i, Tensor&& value, - const std::vector& errors) { - auto n_elements = nbf[0] * nbf[1] * nbf[2] * nbf[3]; - std::vector buffer(n_elements); + const FloatBuffer& errors) { + auto n_elements = nbf[0] * nbf[1] * nbf[2] * nbf[3]; + auto buffer = FloatBuffer(std::vector(n_elements)); + auto errors_span = errors.value(); + auto buffer_span = buffer.value(); for(std::size_t i = 0; i < nbf[0]; ++i) { auto ilocal = i * nbf[1] * nbf[2] * nbf[3]; @@ -103,9 +115,9 @@ auto compute_block(T&& strides, T&& nbf, T&& ao_i, Tensor&& value, auto koffset = joffset + (ao_i[2] + k) * strides[2]; for(std::size_t l = 0; l < nbf[3]; ++l) { - auto llocal = klocal + l; - auto loffset = koffset + (ao_i[3] + l) * strides[3]; - buffer[llocal] = errors[llocal] + value[loffset]; + auto llocal = klocal + l; + auto loffset = koffset + (ao_i[3] + l) * strides[3]; + buffer_span[llocal] = errors_span[llocal] + value[loffset]; } } } @@ -114,12 +126,11 @@ auto compute_block(T&& strides, T&& nbf, T&& ao_i, Tensor&& value, } #endif -template +template void set_block(T&& strides, T&& nbf, const std::array& permuted_ao_offsets, const std::array& sigma, - const std::vector& block, - std::vector& out) { + const FloatBuffer& block, FloatBuffer& out) { // sigma[d] = the original mode that for what is now mode d. Therefore, // sigma[d] maps us back to the original mode, e.g., if the permutation // took 0, 1, 2, 3 to 3, 2, 1, 0 then sigma[0] = 3, sigma[1] = 2, @@ -130,6 +141,8 @@ void set_block(T&& strides, T&& nbf, // Here we iterate in canonical (i,j,k,l) order — the same order the block // was filled — and then scatter to its permuted position in out. + auto block_span = block.value(); + auto out_span = out.value(); std::size_t block_idx = 0; for(std::size_t i = 0; i < nbf[0]; ++i) { for(std::size_t j = 0; j < nbf[1]; ++j) { @@ -142,21 +155,24 @@ void set_block(T&& strides, T&& nbf, (permuted_ao_offsets[1] + cidx[sigma[1]]) * strides[1] + (permuted_ao_offsets[2] + cidx[sigma[2]]) * strides[2] + (permuted_ao_offsets[3] + cidx[sigma[3]]) * strides[3]; - out[out_idx] = block[block_idx++]; + out_span[out_idx] = block_span[block_idx++]; } } } } } -template typename UQType> struct Kernel { using shape_type = buffer::Contiguous::shape_type; using demangler_type = ::utilities::printing::Demangler; Kernel(shape_type shape, std::array aos, - utils::mean_type mean) : - m_shape(std::move(shape)), m_aos(aos), m_mean(mean) {} + utils::mean_type mean, + integrals::property_types::UQFactory factory = {}) : + m_shape(std::move(shape)), + m_aos(aos), + m_mean(mean), + m_factory(std::move(factory)) {} template Tensor operator()(const std::span t, @@ -183,7 +199,6 @@ struct Kernel { type0); } else { #ifdef ENABLE_SIGMA - using tensorwrapper::buffer::make_contiguous; using utils::get_permutations_with_sigma; std::array n_centers{m_aos[0].size(), m_aos[1].size(), @@ -193,8 +208,6 @@ struct Kernel { std::array ao_offsets{0, 0, 0, 0}; std::array nbf{0, 0, 0, 0}; - using uq_type = UQType; - std::vector rv_data(m_shape.size()); std::array strides{0, 0, 0, 1}; strides[2] = strides[3] * m_aos[3].n_aos(); strides[1] = strides[2] * m_aos[2].n_aos(); @@ -208,60 +221,106 @@ struct Kernel { // centers[1] bool all_same = (m_aos[0] == m_aos[2]) && mu_is_nu && lam_is_sig; - for(centers[0] = 0; centers[0] < n_centers[0]; ++centers[0]) { - nbf[0] = m_aos[0][centers[0]].n_aos(); - - ao_offsets[1] = 0; - for(centers[1] = 0; centers[1] < n_centers[1]; ++centers[1]) { - // We restrict our bra pairs to centers[0] <= centers[1] - if(centers[1] > centers[0] && mu_is_nu) break; - nbf[1] = m_aos[1][centers[1]].n_aos(); - - ao_offsets[2] = 0; - for(centers[2] = 0; centers[2] < n_centers[2]; - ++centers[2]) { - // (c2, c3) <= (c0, c1) is impossible if c2 > c0 - if(centers[2] > centers[0] && all_same) break; - bool c2eqc0 = centers[2] == centers[0]; - nbf[2] = m_aos[2][centers[2]].n_aos(); - - ao_offsets[3] = 0; - for(centers[3] = 0; centers[3] < n_centers[3]; - ++centers[3]) { - // Restrict ket pairs to centers[2] <= centers[3] - if(centers[3] > centers[2] && lam_is_sig) break; - - nbf[3] = m_aos[3][centers[3]].n_aos(); - // Skip (c2,c3) > (c0,c1) lexicographically - bool pair_gt = (c2eqc0 && centers[3] > centers[1]); - if(pair_gt && all_same) break; - - auto block_errors = average_error( - strides, nbf, ao_offsets, error, m_mean); - - // Compute (ab|cd) - auto block = compute_block(strides, nbf, ao_offsets, - t, block_errors); - - // Set all symmetry equivalent blocks to `block` - auto perms = get_permutations_with_sigma( - ao_offsets, mu_is_nu, lam_is_sig, all_same); - for(auto& [perm, sigma] : perms) { - set_block(strides, nbf, perm, sigma, block, - rv_data); + // The UQ kind is fixed by m_factory for the whole run (it never + // varies block-to-block), so it's resolved to a concrete + // UQType exactly once here, via a single switch on + // UQFactory::kind(). UQFactoryBase::operator() always returns a + // wtf::fp::Float holding a *double*-valued UQ scalar (real ERI + // tensors are always double in practice), so uq_type is built + // from kind() + this Kernel's own float_type rather than by + // inspecting a seed value, which would incorrectly decouple the + // UQ scalar's width from float_type whenever this Kernel + // template is instantiated for float. + auto run_for_kind = [&] typename UQType> { + using uq_type = UQType; + + auto rv_data = + FloatBuffer(std::vector(m_shape.size())); + + for(centers[0] = 0; centers[0] < n_centers[0]; ++centers[0]) { + nbf[0] = m_aos[0][centers[0]].n_aos(); + + ao_offsets[1] = 0; + for(centers[1] = 0; centers[1] < n_centers[1]; + ++centers[1]) { + // We restrict our bra pairs to centers[0] <= centers[1] + if(centers[1] > centers[0] && mu_is_nu) break; + nbf[1] = m_aos[1][centers[1]].n_aos(); + + ao_offsets[2] = 0; + for(centers[2] = 0; centers[2] < n_centers[2]; + ++centers[2]) { + // (c2, c3) <= (c0, c1) is impossible if c2 > c0 + if(centers[2] > centers[0] && all_same) break; + bool c2eqc0 = centers[2] == centers[0]; + nbf[2] = m_aos[2][centers[2]].n_aos(); + + ao_offsets[3] = 0; + for(centers[3] = 0; centers[3] < n_centers[3]; + ++centers[3]) { + // Restrict ket pairs to centers[2] <= + // centers[3] + if(centers[3] > centers[2] && lam_is_sig) break; + + nbf[3] = m_aos[3][centers[3]].n_aos(); + // Skip (c2,c3) > (c0,c1) lexicographically + bool pair_gt = + (c2eqc0 && centers[3] > centers[1]); + if(pair_gt && all_same) break; + + auto block_errors = average_error( + strides, nbf, ao_offsets, error, m_mean, + m_factory); + + // Compute (ab|cd) + auto block = compute_block( + strides, nbf, ao_offsets, t, block_errors); + + // Set all symmetry equivalent blocks to + // `block` + auto perms = get_permutations_with_sigma( + ao_offsets, mu_is_nu, lam_is_sig, all_same); + for(auto& [perm, sigma] : perms) { + set_block(strides, nbf, perm, + sigma, block, rv_data); + } + + ao_offsets[3] += nbf[3]; } - - ao_offsets[3] += nbf[3]; + ao_offsets[2] += nbf[2]; } - ao_offsets[2] += nbf[2]; + ao_offsets[1] += nbf[1]; } - ao_offsets[1] += nbf[1]; + ao_offsets[0] += nbf[0]; } - ao_offsets[0] += nbf[0]; + tensorwrapper::buffer::Contiguous t_w_contig(std::move(rv_data), + m_shape); + rv = tensorwrapper::Tensor(m_shape, std::move(t_w_contig)); + }; + + using integrals::property_types::UQKind; + switch(m_factory.kind()) { + case UQKind::uncertain: + run_for_kind.template + operator()(); + break; + case UQKind::interval: + run_for_kind.template + operator()(); + break; + case UQKind::affine: + run_for_kind + .template operator()(); + break; + case UQKind::thresholded_affine: + run_for_kind.template + operator()(); + break; + case UQKind::taylor_model: + run_for_kind.template + operator()(); + break; } - tensorwrapper::buffer::Contiguous t_w_contig(std::move(rv_data), - m_shape); - rv = tensorwrapper::Tensor(m_shape, std::move(t_w_contig)); #else throw std::runtime_error(m_error_base + "Sigma support not enabled!"); @@ -273,6 +332,7 @@ struct Kernel { shape_type m_shape; std::array m_aos; utils::mean_type m_mean; + integrals::property_types::UQFactory m_factory; std::string m_error_base = "integrals::ao_integrals::UQAtomSymmBlockedDriver: "; }; @@ -293,7 +353,7 @@ MODULE_CTOR(UQAtomSymmBlockedDriver) { description(desc); add_submodule("ERIs"); add_submodule("ERI Error"); - add_input("UQ Type").set_default("uncertain"); + add_submodule("UQ Initializer"); add_input("Mean Type").set_default("none"); } @@ -301,7 +361,9 @@ MODULE_RUN(UQAtomSymmBlockedDriver) { const auto& [braket] = eri_pt::unwrap_inputs(inputs); auto mean_str = inputs.at("Mean Type").value(); auto mean = utils::mean_from_string(mean_str); - auto uq_type = inputs.at("UQ Type").value(); + + auto factory = submods.at("UQ Initializer") + .run_as(); auto& eri_mod = submods.at("ERIs").value(); auto tol = eri_mod.inputs().at("Threshold").value(); @@ -325,29 +387,10 @@ MODULE_RUN(UQAtomSymmBlockedDriver) { using buffer::visit_contiguous_buffer; shape::Smooth shape = t.buffer().layout().shape().as_smooth().make_smooth(); - simde::type::tensor t_w_error; - if(uq_type == "uncertain") { - Kernel k(shape, aos, mean); - t_w_error = visit_contiguous_buffer(k, t_buffer, e_buffer); - } else if(uq_type == "interval") { - Kernel k(shape, aos, mean); - t_w_error = visit_contiguous_buffer(k, t_buffer, e_buffer); - } else if(uq_type == "affine") { - Kernel k(shape, aos, mean); - t_w_error = visit_contiguous_buffer(k, t_buffer, e_buffer); - } else if(uq_type == "thresholded affine") { - Kernel k(shape, aos, - mean); - t_w_error = visit_contiguous_buffer(k, t_buffer, e_buffer); - } else if(uq_type == "taylor model") { - Kernel k(shape, aos, mean); - t_w_error = visit_contiguous_buffer(k, t_buffer, e_buffer); - } else { - throw std::runtime_error( - "integrals::ao_integrals::UQAtomSymmBlockedDriver: Invalid UQ type " - "name " + - uq_type); - } + Kernel k(shape, aos, mean, factory); + simde::type::tensor t_w_error = + visit_contiguous_buffer(k, t_buffer, e_buffer); + auto rv = results(); return eri_pt::wrap_results(rv, t_w_error); } diff --git a/cxx/src/integrals/ao_integrals/uq_initializer.cpp b/cxx/src/integrals/ao_integrals/uq_initializer.cpp new file mode 100644 index 00000000..584454cf --- /dev/null +++ b/cxx/src/integrals/ao_integrals/uq_initializer.cpp @@ -0,0 +1,52 @@ +/* + * 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 "ao_integrals.hpp" +#include + +namespace integrals::ao_integrals { +namespace { + +const auto desc = R"( +UQ Initializer +-------------- + +Produces a UQFactory configured by this module's "UQ Type" input (and, for +"taylor model", the "Order" input). The resulting factory can be used directly +in a module's hot loops, avoiding module-dispatch overhead. +)"; + +} // namespace + +using pt = integrals::property_types::UQInitializer; + +MODULE_CTOR(UQInitializer) { + satisfies_property_type(); + description(desc); + add_input("UQ Type").set_default("uncertain"); + add_input("Order").set_default(std::size_t(2)); +} + +MODULE_RUN(UQInitializer) { + auto uq_type = inputs.at("UQ Type").value(); + auto order = inputs.at("Order").value(); + auto kind = integrals::property_types::uq_kind_from_string(uq_type); + integrals::property_types::UQFactory factory(kind, order); + auto rv = results(); + return pt::wrap_results(rv, factory); +} + +} // namespace integrals::ao_integrals diff --git a/tests/cxx/unit/integrals/ao_integrals/test_uq_atom_symm_blocked_driver.cpp b/tests/cxx/unit/integrals/ao_integrals/test_uq_atom_symm_blocked_driver.cpp index e1637d2a..a83b041d 100644 --- a/tests/cxx/unit/integrals/ao_integrals/test_uq_atom_symm_blocked_driver.cpp +++ b/tests/cxx/unit/integrals/ao_integrals/test_uq_atom_symm_blocked_driver.cpp @@ -155,10 +155,15 @@ TEMPLATE_LIST_TEST_CASE("UQ Atom Symm Blocked Driver", "", uq_types) { }(); if constexpr(tensorwrapper::types::is_uq_type_v) { + // "UQ Type" now lives solely on the "UQ Initializer" submodule + // (UQAtomSymmBlockedDriver no longer has its own, redundant + // copy of this input), so it's set via the ModuleManager on + // that submodule rather than on `mod` directly. + mm.change_input("UQ Initializer", "UQ Type", uq_type); + // The errors for the integrals are on the order of 1e-5. // Subtracting results in differences between 1e-5 and 1e-4. SECTION("No Mean") { - mod.change_input("UQ Type", uq_type); auto T = mod.run_as(braket); auto tol = 1.0e-6; auto T_eri = mm.at("ERI4").run_as(braket); @@ -167,7 +172,6 @@ TEMPLATE_LIST_TEST_CASE("UQ Atom Symm Blocked Driver", "", uq_types) { REQUIRE(corr_answer_no_mean(T, T_eri, T_err)); } SECTION("Max Error") { - mod.change_input("UQ Type", uq_type); mod.change_input("Mean Type", "max"); auto T = mod.run_as(braket); @@ -175,13 +179,48 @@ TEMPLATE_LIST_TEST_CASE("UQ Atom Symm Blocked Driver", "", uq_types) { REQUIRE(approximately_equal(T_corr, T, 1E-4)); } SECTION("Geometric Mean") { - mod.change_input("UQ Type", uq_type); mod.change_input("Mean Type", "geometric"); auto T = mod.run_as(braket); auto T_corr = corr_answer(T); REQUIRE(approximately_equal(T_corr, T, 1E-4)); } + + if constexpr(tensorwrapper::types::is_taylor_model_v) { + SECTION("Taylor model order is configurable end-to-end") { + // Memoization must be off here: the two calls below share + // the same braket/"UQ Type" input and only differ by which + // module is bound to the "UQ Initializer" submodule slot, + // so a memoized run would (incorrectly, for this test) + // return the first call's cached result for the second. + mod.turn_off_memoization(); + auto T_order2 = mod.run_as(braket); // default + // Order=2 + + // Copying "UQ Initializer" after the "UQ Type" input was + // set on it above carries that setting into the copy, so + // only "Order" needs to be overridden here. + mm.copy_module("UQ Initializer", "UQ Initializer Order4"); + mm.change_input("UQ Initializer Order4", "Order", + std::size_t(4)); + mm.change_submod("UQ Atom Symm Blocked Driver", + "UQ Initializer", "UQ Initializer Order4"); + + auto mod4 = mm.at("UQ Atom Symm Blocked Driver"); + mod4.turn_off_memoization(); + auto T_order4 = mod4.run_as(braket); + + auto t2 = eigen_tensor<4, float_type>(T_order2.buffer()); + auto t4 = eigen_tensor<4, float_type>(T_order4.buffer()); + + // {0,0,0,1} carries a nonzero error estimate (see + // corr_answer above), so its Taylor-model max_order() + // must reflect the configured "UQ Initializer"/"Order" + // input rather than the compiled-in default of 2. + REQUIRE(t2(0, 0, 0, 1).max_order() == 2); + REQUIRE(t4(0, 0, 0, 1).max_order() == 4); + } + } } } } diff --git a/tests/cxx/unit/integrals/ao_integrals/test_uq_initializer.cpp b/tests/cxx/unit/integrals/ao_integrals/test_uq_initializer.cpp new file mode 100644 index 00000000..27b7867b --- /dev/null +++ b/tests/cxx/unit/integrals/ao_integrals/test_uq_initializer.cpp @@ -0,0 +1,89 @@ +/* + * 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 "../testing/testing.hpp" + +using namespace integrals::testing; + +using pt = integrals::property_types::UQInitializer; +using integrals::property_types::uq_kind_from_string; +using integrals::property_types::UQFactory; +using integrals::property_types::UQKind; + +TEST_CASE("UQInitializer") { + auto mm = initialize_integrals(); + auto& mod = mm.at("UQ Initializer"); + + SECTION("Default UQ Type/order matches compiled-in defaults") { + auto factory = mod.run_as(); + REQUIRE(factory.kind() == UQKind::uncertain); + } + + SECTION("Each UQ Type string produces a factory of the matching kind") { + std::vector> cases{ + {"uncertain", UQKind::uncertain}, + {"interval", UQKind::interval}, + {"affine", UQKind::affine}, + {"thresholded affine", UQKind::thresholded_affine}, + {"taylor model", UQKind::taylor_model}}; + + for(const auto& [uq_type, kind] : cases) { + auto copy = mod.unlocked_copy(); + copy.change_input("UQ Type", uq_type); + auto factory = copy.run_as(); + REQUIRE(factory.kind() == kind); + + // factory(center, radius) should not throw and (when Sigma + // supplies real UQ types) should reflect the requested + // center/radius via the underlying UQ type's bounds. + [[maybe_unused]] auto elem = factory(0.774606, 0.0000010000000000); +#ifdef ENABLE_SIGMA + wtf::fp::visit_float( + [](auto value) { + auto lo = tensorwrapper::types::uq_lower(value); + auto hi = tensorwrapper::types::uq_upper(value); + REQUIRE(lo <= 0.774606); + REQUIRE(hi >= 0.774606); + }, + elem); +#endif + } + } + + SECTION("Custom order is honored for taylor model") { + auto mod4 = mod.unlocked_copy(); + mod4.change_input("UQ Type", std::string("taylor model")); + mod4.change_input("Order", std::size_t(4)); + auto factory = mod4.run_as(); + REQUIRE(factory.kind() == UQKind::taylor_model); + + [[maybe_unused]] auto elem = factory(0.774606, 0.0000010000000000); + // Without Sigma the "taylor model" type is a plain double, which has + // no order to honor. +#ifdef ENABLE_SIGMA + auto tm = + wtf::fp::float_cast>( + elem); + REQUIRE(tm.max_order() == 4); +#endif + } + + SECTION("Invalid UQ Type throws") { + auto copy = mod.unlocked_copy(); + copy.change_input("UQ Type", std::string("not a real uq type")); + REQUIRE_THROWS_AS(copy.run_as(), std::runtime_error); + } +} diff --git a/tests/cxx/unit/integrals/test_uq_factory.cpp b/tests/cxx/unit/integrals/test_uq_factory.cpp new file mode 100644 index 00000000..e25f5ddf --- /dev/null +++ b/tests/cxx/unit/integrals/test_uq_factory.cpp @@ -0,0 +1,282 @@ +/* + * 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 "testing/testing.hpp" +#include + +using namespace integrals::property_types; + +namespace { + +/// The center/radius used throughout; the radius is on the order of the error +/// the UQ-aware integral drivers actually see. +constexpr double center = 0.774606; +constexpr double radius = 0.0000010000000000; + +/// All of the kinds, in the order they're declared (which is also the order +/// UQFactory::operator< sorts them in). +const std::vector all_kinds{UQKind::uncertain, UQKind::interval, + UQKind::affine, UQKind::thresholded_affine, + UQKind::taylor_model}; + +} // namespace + +TEST_CASE("uq_kind_from_string") { + SECTION("Valid names") { + REQUIRE(uq_kind_from_string("uncertain") == UQKind::uncertain); + REQUIRE(uq_kind_from_string("interval") == UQKind::interval); + REQUIRE(uq_kind_from_string("affine") == UQKind::affine); + REQUIRE(uq_kind_from_string("thresholded affine") == + UQKind::thresholded_affine); + REQUIRE(uq_kind_from_string("taylor model") == UQKind::taylor_model); + } + + SECTION("Invalid names throw") { + REQUIRE_THROWS_AS(uq_kind_from_string("not a uq type"), + std::runtime_error); + REQUIRE_THROWS_AS(uq_kind_from_string(""), std::runtime_error); + + // The mapping is case- and spelling-sensitive. + REQUIRE_THROWS_AS(uq_kind_from_string("Uncertain"), std::runtime_error); + REQUIRE_THROWS_AS(uq_kind_from_string("thresholded_affine"), + std::runtime_error); + REQUIRE_THROWS_AS(uq_kind_from_string("taylor_model"), + std::runtime_error); + } +} + +TEST_CASE("to_string(UQKind)") { + SECTION("Each kind") { + REQUIRE(to_string(UQKind::uncertain) == "uncertain"); + REQUIRE(to_string(UQKind::interval) == "interval"); + REQUIRE(to_string(UQKind::affine) == "affine"); + REQUIRE(to_string(UQKind::thresholded_affine) == "thresholded affine"); + REQUIRE(to_string(UQKind::taylor_model) == "taylor model"); + } + + SECTION("Is the inverse of uq_kind_from_string") { + for(auto kind : all_kinds) + REQUIRE(uq_kind_from_string(to_string(kind)) == kind); + } +} + +TEST_CASE("UQFactory") { + UQFactory defaulted; + UQFactory uncertain(UQKind::uncertain); + UQFactory interval(UQKind::interval); + UQFactory affine(UQKind::affine); + UQFactory taffine(UQKind::thresholded_affine); + UQFactory tm(UQKind::taylor_model); + UQFactory tm4(UQKind::taylor_model, 4); + + SECTION("Default ctor") { REQUIRE(defaulted.kind() == UQKind::uncertain); } + + SECTION("Value ctor") { + for(auto kind : all_kinds) REQUIRE(UQFactory(kind).kind() == kind); + } + + SECTION("Value ctor ignores order for non-TaylorModel kinds") { + for(auto kind : all_kinds) { + if(kind == UQKind::taylor_model) continue; + REQUIRE(UQFactory(kind, 7) == UQFactory(kind)); + } + } + + SECTION("Copy ctor") { + UQFactory copy(tm4); + REQUIRE(copy == tm4); + REQUIRE(copy.kind() == UQKind::taylor_model); + } + + SECTION("Copy assignment") { + UQFactory copy; + copy = interval; + REQUIRE(copy == interval); + } + +// Sigma provides the UQ representations; without it TensorWrapper's UQ types +// are plain doubles that carry neither bounds nor a TaylorModel order, leaving +// nothing here to check. +#ifdef ENABLE_SIGMA + SECTION("operator()") { + using tensorwrapper::types::uq_center; + using tensorwrapper::types::uq_lower; + using tensorwrapper::types::uq_upper; + + SECTION("uncertain") { + auto erased = uncertain(center, radius); + auto value = + wtf::fp::float_cast>( + erased); + REQUIRE(uq_center(value) == Catch::Approx(center)); + REQUIRE(value.sd() == Catch::Approx(radius)); + } + + SECTION("interval") { + auto erased = interval(center, radius); + auto value = + wtf::fp::float_cast>( + erased); + REQUIRE(uq_lower(value) == Catch::Approx(center - radius)); + REQUIRE(uq_upper(value) == Catch::Approx(center + radius)); + } + + SECTION("affine") { + auto erased = affine(center, radius); + auto value = + wtf::fp::float_cast>( + erased); + REQUIRE(uq_center(value) == Catch::Approx(center)); + REQUIRE(uq_lower(value) == Catch::Approx(center - radius)); + REQUIRE(uq_upper(value) == Catch::Approx(center + radius)); + } + + SECTION("thresholded affine") { + auto erased = taffine(center, radius); + auto value = wtf::fp::float_cast< + tensorwrapper::types::thresholded_affine_type>(erased); + REQUIRE(uq_center(value) == Catch::Approx(center)); + REQUIRE(uq_lower(value) == Catch::Approx(center - radius)); + REQUIRE(uq_upper(value) == Catch::Approx(center + radius)); + } + + SECTION("taylor model") { + using tm_type = tensorwrapper::types::taylor_model_type; + auto erased = tm(center, radius); + auto value = wtf::fp::float_cast(erased); + REQUIRE(uq_center(value) == Catch::Approx(center)); + REQUIRE(uq_lower(value) == Catch::Approx(center - radius)); + REQUIRE(uq_upper(value) == Catch::Approx(center + radius)); + + // The order defaults to 2 and is settable. + REQUIRE(value.max_order() == 2); + auto erased4 = tm4(center, radius); + auto value4 = wtf::fp::float_cast(erased4); + REQUIRE(value4.max_order() == 4); + } + } +#endif + + SECTION("kind") { + REQUIRE(defaulted.kind() == UQKind::uncertain); + REQUIRE(interval.kind() == UQKind::interval); + REQUIRE(tm4.kind() == UQKind::taylor_model); + } + + SECTION("operator==") { + SECTION("Same kind") { + REQUIRE(defaulted == uncertain); + REQUIRE(interval == UQFactory(UQKind::interval)); + REQUIRE(affine == UQFactory(UQKind::affine)); + REQUIRE(taffine == UQFactory(UQKind::thresholded_affine)); + } + + SECTION("Different kinds") { + for(auto lhs : all_kinds) { + for(auto rhs : all_kinds) { + if(lhs == rhs) continue; + REQUIRE_FALSE(UQFactory(lhs) == UQFactory(rhs)); + } + } + } + + SECTION("TaylorModel additionally compares order") { + REQUIRE(tm == UQFactory(UQKind::taylor_model, 2)); + REQUIRE_FALSE(tm == tm4); + REQUIRE(tm4 == UQFactory(UQKind::taylor_model, 4)); + } + + SECTION("Comparison is symmetric") { + REQUIRE_FALSE(uncertain == tm); + REQUIRE_FALSE(tm == uncertain); + REQUIRE_FALSE(tm4 == interval); + REQUIRE_FALSE(interval == tm4); + } + } + + SECTION("operator<") { + SECTION("Orders by kind") { + for(std::size_t i = 0; i < all_kinds.size(); ++i) { + for(std::size_t j = 0; j < all_kinds.size(); ++j) { + UQFactory lhs(all_kinds[i]), rhs(all_kinds[j]); + REQUIRE((lhs < rhs) == (i < j)); + } + } + } + + SECTION("Irreflexive") { + for(auto kind : all_kinds) { + UQFactory factory(kind); + REQUIRE_FALSE(factory < factory); + } + REQUIRE_FALSE(tm4 < tm4); + } + + SECTION("TaylorModel additionally orders by order") { + REQUIRE(tm < tm4); + REQUIRE_FALSE(tm4 < tm); + } + } +} + +TEST_CASE("UQFactoryBase") { + detail::UncertainFactoryImpl uncertain; + detail::IntervalFactoryImpl interval; + detail::TaylorModelFactoryImpl tm2; + detail::TaylorModelFactoryImpl tm4(4); + + SECTION("kind") { + REQUIRE(uncertain.kind() == UQKind::uncertain); + REQUIRE(interval.kind() == UQKind::interval); + REQUIRE(detail::AffineFactoryImpl{}.kind() == UQKind::affine); + REQUIRE(detail::ThresholdedAffineFactoryImpl{}.kind() == + UQKind::thresholded_affine); + REQUIRE(tm2.kind() == UQKind::taylor_model); + } + + SECTION("clone") { + auto pcopy = interval.clone(); + REQUIRE(pcopy.get() != &interval); // Deep, not shallow + REQUIRE(pcopy->kind() == UQKind::interval); + REQUIRE(pcopy->equal(interval)); + + auto ptm = tm4.clone(); + REQUIRE(ptm.get() != &tm4); + REQUIRE(ptm->equal(tm4)); + REQUIRE_FALSE(ptm->equal(tm2)); + } + + SECTION("equal") { + REQUIRE(uncertain.equal(detail::UncertainFactoryImpl{})); + REQUIRE_FALSE(uncertain.equal(interval)); + REQUIRE(tm2.equal(detail::TaylorModelFactoryImpl{})); + REQUIRE_FALSE(tm2.equal(tm4)); + REQUIRE_FALSE(tm2.equal(uncertain)); + } + + SECTION("less") { + REQUIRE(uncertain.less(interval)); + REQUIRE_FALSE(interval.less(uncertain)); + REQUIRE(tm2.less(tm4)); + REQUIRE_FALSE(tm4.less(tm2)); + REQUIRE_FALSE(tm2.less(tm2)); + } + + SECTION("order") { + REQUIRE(tm2.order() == 2); // Default + REQUIRE(tm4.order() == 4); + } +}