From 6ac74984d648d930bc75bc4e6c0a4c515d1daaab Mon Sep 17 00:00:00 2001 From: Conrad Johnston Date: Wed, 29 Apr 2026 22:00:06 +0000 Subject: [PATCH] Add native cube generator backed by GauXC OrbitalEvaluator MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Replace PySCF cubegen with a native C++ implementation that uses GauXC's OrbitalEvaluator for MO and density evaluation on cube grids, and GauXC's write_cube for file output. This eliminates: - The vendored gauxc-private/ header (no longer needed — OrbitalEvaluator is the public API for point-set evaluation) - 600+ lines of hand-rolled collocation batching, GEMM, and text formatting code (all now in GauXC upstream) GauXC pin bumped to ConradJohnston/GauXC@10f59a8 which includes OrbitalEvaluator with shell screening, pipelined collocation+GEMM, CubeGrid-native overloads, and parallel cube writer. --- cpp/CMakeLists.txt | 1 + cpp/cmake/third_party.cmake | 4 +- .../qdk/chemistry/cube/cube_generator.hpp | 57 ++++++ cpp/src/qdk/chemistry/cube/CMakeLists.txt | 6 + cpp/src/qdk/chemistry/cube/cube_generator.cpp | 168 ++++++++++++++++++ 5 files changed, 234 insertions(+), 2 deletions(-) create mode 100644 cpp/include/qdk/chemistry/cube/cube_generator.hpp create mode 100644 cpp/src/qdk/chemistry/cube/CMakeLists.txt create mode 100644 cpp/src/qdk/chemistry/cube/cube_generator.cpp diff --git a/cpp/CMakeLists.txt b/cpp/CMakeLists.txt index d588883f4..b2bcc4727 100644 --- a/cpp/CMakeLists.txt +++ b/cpp/CMakeLists.txt @@ -147,6 +147,7 @@ endif() add_library(chemistry) add_subdirectory(src/qdk/chemistry/data) add_subdirectory(src/qdk/chemistry/algorithms) +add_subdirectory(src/qdk/chemistry/cube) add_subdirectory(src/qdk/chemistry/utils) target_include_directories(chemistry PUBLIC diff --git a/cpp/cmake/third_party.cmake b/cpp/cmake/third_party.cmake index 21e170a00..41f3eda01 100644 --- a/cpp/cmake/third_party.cmake +++ b/cpp/cmake/third_party.cmake @@ -72,8 +72,8 @@ set(GAUXC_ENABLE_MPI ${QDK_CHEMISTRY_ENABLE_MPI} CACHE BOOL "Enable gauxc MPI S set(GAUXC_ENABLE_OPENMP ${QDK_ENABLE_OPENMP} CACHE BOOL "Enable gauxc OpenMP Support" FORCE) handle_dependency(gauxc - GIT_REPOSITORY https://github.com/wavefunction91/gauxc.git - GIT_TAG 62fea07c9306dbd83dd18b6957358827ac9b3da0 + GIT_REPOSITORY https://github.com/ConradJohnston/gauxc.git + GIT_TAG 10f59a8f42ae5f66e72c287831980a6448e36aa2 BUILD_TARGET gauxc::gauxc INSTALL_TARGET gauxc::gauxc ${DEPENDENCY_BUILD_FLAGS} diff --git a/cpp/include/qdk/chemistry/cube/cube_generator.hpp b/cpp/include/qdk/chemistry/cube/cube_generator.hpp new file mode 100644 index 000000000..63d31a823 --- /dev/null +++ b/cpp/include/qdk/chemistry/cube/cube_generator.hpp @@ -0,0 +1,57 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. See LICENSE.txt in the project root for +// license information. + +#pragma once + +#include +#include +#include +#include + +namespace qdk::chemistry::data { +class BasisSet; +class Wavefunction; +} // namespace qdk::chemistry::data + +namespace qdk::chemistry::cube { + +using CubeField = std::vector; + +struct CubeGrid { + Eigen::Vector3d origin{0.0, 0.0, 0.0}; + Eigen::Vector3d spacing{0.2, 0.2, 0.2}; + std::size_t nx = 80, ny = 80, nz = 80; + + static CubeGrid from_basis_set(const data::BasisSet&, std::size_t nx = 80, + std::size_t ny = 80, std::size_t nz = 80, + double margin = 3.0); + std::size_t num_points() const { return nx * ny * nz; } +}; + +class CubeGenerator { + public: + explicit CubeGenerator(std::shared_ptr basis_set); + ~CubeGenerator() noexcept; + CubeGenerator(CubeGenerator&&) noexcept; + CubeGenerator& operator=(CubeGenerator&&) noexcept; + + CubeField orbital(const Eigen::VectorXd& mo_coeff, const std::string& outfile, + const CubeGrid& grid, + const std::string& comment = "") const; + + CubeField density(const Eigen::MatrixXd& density_matrix, + const std::string& outfile, const CubeGrid& grid, + const std::string& comment = "") const; + + private: + struct Impl; + std::unique_ptr _impl; +}; + +std::vector generate_orbital_cubes( + const data::Wavefunction&, const std::vector& indices, + const std::string& output_dir, const CubeGrid& grid, + const std::string& label_prefix = "orbital_"); + +} // namespace qdk::chemistry::cube diff --git a/cpp/src/qdk/chemistry/cube/CMakeLists.txt b/cpp/src/qdk/chemistry/cube/CMakeLists.txt new file mode 100644 index 000000000..3bf2e1bc6 --- /dev/null +++ b/cpp/src/qdk/chemistry/cube/CMakeLists.txt @@ -0,0 +1,6 @@ +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. + +target_sources(chemistry PRIVATE + cube_generator.cpp +) diff --git a/cpp/src/qdk/chemistry/cube/cube_generator.cpp b/cpp/src/qdk/chemistry/cube/cube_generator.cpp new file mode 100644 index 000000000..9b9d2b0d3 --- /dev/null +++ b/cpp/src/qdk/chemistry/cube/cube_generator.cpp @@ -0,0 +1,168 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. See LICENSE.txt in the project root for +// license information. + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace qdk::chemistry::cube { + +CubeGrid CubeGrid::from_basis_set(const data::BasisSet& basis_set, + std::size_t nx, std::size_t ny, + std::size_t nz, double margin) { + const auto structure = basis_set.get_structure(); + if (!structure) + throw std::runtime_error("CubeGrid: basis set has no structure."); + const Eigen::MatrixXd& coords = structure->get_coordinates(); + if (coords.rows() == 0) + throw std::runtime_error("CubeGrid: structure has no atoms."); + + Eigen::Vector3d lo = coords.colwise().minCoeff(); + Eigen::Vector3d extent = + (coords.colwise().maxCoeff() - lo).array() + 2.0 * margin; + + CubeGrid g; + g.origin = lo.array() - margin; + g.nx = nx; + g.ny = ny; + g.nz = nz; + g.spacing[0] = nx > 1 ? extent[0] / double(nx - 1) : 0.0; + g.spacing[1] = ny > 1 ? extent[1] / double(ny - 1) : 0.0; + g.spacing[2] = nz > 1 ? extent[2] / double(nz - 1) : 0.0; + return g; +} + +namespace { + +GauXC::BasisSet to_gauxc_basis(const data::BasisSet& qdk) { + using PA = GauXC::Shell::prim_array; + using CA = GauXC::Shell::cart_array; + + const auto st = qdk.get_structure(); + const Eigen::MatrixXd& coords = st->get_coordinates(); + const bool sph = qdk.get_atomic_orbital_type() == data::AOType::Spherical; + + GauXC::BasisSet basis; + for (std::size_t ia = 0; ia < qdk.get_num_atoms(); ++ia) { + CA center{coords(ia, 0), coords(ia, 1), coords(ia, 2)}; + for (const auto& sh : qdk.get_shells_for_atom(ia)) { + const int l = sh.get_angular_momentum(); + const auto np = static_cast(sh.exponents.size()); + PA alpha{}, coeff{}; + for (int i = 0; i < np; ++i) { + alpha[i] = sh.exponents[i]; + coeff[i] = sh.coefficients[i]; + } + basis.emplace_back(GauXC::PrimSize(np), GauXC::AngularMomentum(l), + GauXC::SphericalType(l > 1 && sph), alpha, coeff, + center, true); + } + } + return basis; +} + +GauXC::Molecule to_gauxc_mol(const data::Structure& st) { + const auto& coords = st.get_coordinates(); + const auto& elems = st.get_elements(); + GauXC::Molecule mol; + for (Eigen::Index i = 0; i < coords.rows(); ++i) + mol.push_back({GauXC::AtomicNumber(int64_t(elems[i])), coords(i, 0), + coords(i, 1), coords(i, 2)}); + return mol; +} + +GauXC::CubeGrid to_gauxc_grid(const CubeGrid& g) { + return {{g.origin[0], g.origin[1], g.origin[2]}, + {g.spacing[0], g.spacing[1], g.spacing[2]}, + int64_t(g.nx), + int64_t(g.ny), + int64_t(g.nz)}; +} + +} // namespace + +struct CubeGenerator::Impl { + std::shared_ptr basis_set; + GauXC::BasisSet gauxc_basis; + GauXC::Molecule gauxc_mol; + GauXC::OrbitalEvaluator evaluator; + int32_t nbf; + + explicit Impl(std::shared_ptr bs) + : basis_set(std::move(bs)), + gauxc_basis(to_gauxc_basis(*basis_set)), + gauxc_mol(to_gauxc_mol(*basis_set->get_structure())), + evaluator(gauxc_basis), + nbf(gauxc_basis.nbf()) { + for (auto& sh : gauxc_basis) sh.set_shell_tolerance(1e-12); + } +}; + +CubeGenerator::CubeGenerator(std::shared_ptr bs) + : _impl(std::make_unique(std::move(bs))) {} +CubeGenerator::~CubeGenerator() noexcept = default; +CubeGenerator::CubeGenerator(CubeGenerator&&) noexcept = default; +CubeGenerator& CubeGenerator::operator=(CubeGenerator&&) noexcept = default; + +CubeField CubeGenerator::orbital(const Eigen::VectorXd& C, + const std::string& outfile, + const CubeGrid& grid, + const std::string& comment) const { + if (C.size() != _impl->nbf) + throw std::invalid_argument("orbital: mo_coeff length mismatch."); + auto g = to_gauxc_grid(grid); + CubeField field(g.num_points()); + _impl->evaluator.eval_orbital(g, C.data(), field.data()); + if (!outfile.empty()) + GauXC::write_cube(outfile, _impl->gauxc_mol, g, field.data(), comment); + return field; +} + +CubeField CubeGenerator::density(const Eigen::MatrixXd& D, + const std::string& outfile, + const CubeGrid& grid, + const std::string& comment) const { + if (D.rows() != _impl->nbf || D.cols() != _impl->nbf) + throw std::invalid_argument("density: matrix shape mismatch."); + auto g = to_gauxc_grid(grid); + CubeField field(g.num_points()); + _impl->evaluator.eval_density(g, D.data(), _impl->nbf, field.data()); + if (!outfile.empty()) + GauXC::write_cube(outfile, _impl->gauxc_mol, g, field.data(), comment); + return field; +} + +std::vector generate_orbital_cubes( + const data::Wavefunction& wfn, const std::vector& indices, + const std::string& output_dir, const CubeGrid& grid, + const std::string& prefix) { + CubeGenerator gen(wfn.get_orbitals()->get_basis_set()); + auto [C_a, C_b] = wfn.get_orbitals()->get_coefficients(); + std::filesystem::create_directories(output_dir); + + std::vector paths; + paths.reserve(indices.size()); + for (auto p : indices) { + if (std::size_t(C_a.cols()) <= p) + throw std::out_of_range("generate_orbital_cubes: index OOB."); + char buf[64]; + std::snprintf(buf, sizeof(buf), "%s%04zu.cube", prefix.c_str(), p); + auto path = (std::filesystem::path(output_dir) / buf).string(); + gen.orbital(C_a.col(p), path, grid, + "Orbital " + std::to_string(p) + " (alpha)"); + paths.push_back(std::move(path)); + } + return paths; +} + +} // namespace qdk::chemistry::cube