diff --git a/include/gauxc/cube_grid.hpp b/include/gauxc/cube_grid.hpp new file mode 100644 index 00000000..58511876 --- /dev/null +++ b/include/gauxc/cube_grid.hpp @@ -0,0 +1,83 @@ +/** + * GauXC Copyright (c) 2020-2024, The Regents of the University of California, + * through Lawrence Berkeley National Laboratory (subject to receipt of + * any required approvals from the U.S. Dept. of Energy). + * + * (c) 2024-2025, Microsoft Corporation + * + * All rights reserved. + * + * See LICENSE.txt for details + */ +#pragma once + +#include +#include +#include + +#include + +namespace GauXC { + +/** @brief Specification of an axis-aligned 3D rectangular point grid. + * + * The grid axes are parallel to the Cartesian frame, i.e. the axis vectors + * are (spacing[0], 0, 0), (0, spacing[1], 0) and (0, 0, spacing[2]). All + * quantities are in atomic units (Bohr). + * + * Total number of points: nx*ny*nz. Storage cost per scalar field: + * 8*nx*ny*nz bytes (double precision). + */ +struct CubeGrid { + std::array origin{0.0, 0.0, 0.0}; ///< (0,0,0) corner, Bohr + std::array spacing{0.2, 0.2, 0.2}; ///< Step on each axis, Bohr + int64_t nx = 80; + int64_t ny = 80; + int64_t nz = 80; + + /// Total number of grid points. + int64_t num_points() const noexcept { return nx * ny * nz; } + + /** @brief Build a default grid that tightly encloses a molecule. + * + * The bounding box of the atomic centres is extended by `margin` Bohr on + * each side and discretised with the requested number of points along + * each axis. Spacing is chosen so that the first and last grid points + * coincide with the extended bounding-box corners (PySCF cubegen + * convention). + * + * The margin is the same on all three axes, so a flat molecule gets a box + * that is thin in the direction it is flat in. An orbital that reaches out + * that way, such as the pi system of an aromatic ring, can still be large + * where the box ends and will look cut off when plotted. Use a bigger + * margin, or set origin, spacing and point counts yourself, if that matters. + * + * @param mol Molecule whose atomic centres define the bounding box. + * @param nx,ny,nz Number of grid points along each axis. + * @param margin Margin (Bohr) added on each side. Default 3.0 matches + * the PySCF cubegen default. + */ + static CubeGrid from_molecule( const Molecule& mol, + int64_t nx = 80, int64_t ny = 80, + int64_t nz = 80, double margin = 3.0 ); + + /** @brief Materialise the grid points as an AoS coordinate array. + * + * Returns a vector of length 3*num_points() laid out in row-major + * (ix, iy, iz) order with iz varying fastest, matching the field-element + * ordering expected by `write_cube`. Suitable to be passed directly as + * the `points` argument to `OrbitalEvaluator::eval_orbital` / + * `eval_density`. + */ + std::vector points() const; + + /** @brief Write grid-point coordinates into a caller-supplied buffer. + * + * Same layout as `points()` but writes into `out`, which must have room + * for at least 3*num_points() doubles. Use this to avoid a heap + * allocation when the caller already owns a suitably sized buffer. + */ + void points_into( double* out ) const; +}; + +} // namespace GauXC diff --git a/include/gauxc/external/cube.hpp b/include/gauxc/external/cube.hpp new file mode 100644 index 00000000..dd11fc4f --- /dev/null +++ b/include/gauxc/external/cube.hpp @@ -0,0 +1,77 @@ +/** + * GauXC Copyright (c) 2020-2024, The Regents of the University of California, + * through Lawrence Berkeley National Laboratory (subject to receipt of + * any required approvals from the U.S. Dept. of Energy). + * + * (c) 2024-2025, Microsoft Corporation + * + * All rights reserved. + * + * See LICENSE.txt for details + */ +#pragma once + +#include + +#include +#include +#include + +namespace GauXC { + +/** @brief Write a Gaussian cube file. + * + * Field layout: row-major (ix, iy, iz) with iz varying fastest. Length must + * equal `grid.num_points()`. Values are written in fixed `%13.5E` format + * with six values per line, matching the standard cube-file convention + * produced by PySCF / Gaussian / NWChem. + * + * This routine performs only file I/O; it has no dependency on the GauXC + * numerical kernels and may be used independently of `OrbitalEvaluator`. + * + * @param path Output path. Parent directory must exist. + * @param mol Molecule (atomic numbers + Cartesian coordinates in Bohr) + * written into the cube-file header. + * @param grid Grid specification. + * @param field Length-`grid.num_points()` scalar field. + * @param comment Optional first comment line (the second comment line is + * always "Generated by GauXC"). If empty, a default first + * line is used. + */ +void write_cube( const std::string& path, + const Molecule& mol, + const CubeGrid& grid, + const double* field, + const std::string& comment = "" ); + +#ifdef GAUXC_HAS_HDF5 +/** @brief Write a cube field to an HDF5 file. + * + * Stores the grid specification, molecular geometry, and scalar field in a + * single HDF5 file for downstream analysis. The field is stored as a 3D + * dataset of shape (nx, ny, nz) in row-major order with iz varying fastest + * (matching the cube-file convention). + * + * HDF5 layout: + * /field (nx, ny, nz) float64 — the scalar field + * /grid/origin (3,) float64 + * /grid/spacing (3,) float64 + * /grid/shape (3,) int64 — {nx, ny, nz} + * /atoms/Z (natom,) int64 + * /atoms/coords (natom, 3) float64 + * /comment scalar string attribute on /field + * + * @param path Output path. Parent directory must exist. + * @param mol Molecule (atomic numbers + Cartesian coordinates in Bohr). + * @param grid Grid specification. + * @param field Length-`grid.num_points()` scalar field. + * @param comment Optional comment stored as an attribute on /field. + */ +void write_cube_hdf5( const std::string& path, + const Molecule& mol, + const CubeGrid& grid, + const double* field, + const std::string& comment = "" ); +#endif + +} // namespace GauXC diff --git a/include/gauxc/orbital_evaluator.hpp b/include/gauxc/orbital_evaluator.hpp new file mode 100644 index 00000000..ebecb2a8 --- /dev/null +++ b/include/gauxc/orbital_evaluator.hpp @@ -0,0 +1,180 @@ +/** + * GauXC Copyright (c) 2020-2024, The Regents of the University of California, + * through Lawrence Berkeley National Laboratory (subject to receipt of + * any required approvals from the U.S. Dept. of Energy). + * + * (c) 2024-2025, Microsoft Corporation + * + * All rights reserved. + * + * See LICENSE.txt for details + */ +#pragma once + +#include +#include +#include + +#include +#include +#include + +namespace GauXC { + +namespace detail { + /// OrbitalEvaluator Implementation class + class OrbitalEvaluatorImpl; +} + +/** @brief Evaluate molecular orbitals and densities on arbitrary point sets. + * + * Wraps the collocation kernel exposed by the local work driver with a + * thread-parallel batched evaluation loop, hiding the driver factory and + * the per-thread AO scratch from callers. The class is designed to be + * constructed once per molecule and reused across many evaluations (e.g. + * one per active-space orbital) without re-initialising the driver. + * + * The evaluator is intentionally decoupled from any particular file format: + * it returns plain numerical arrays. For Gaussian cube file I/O, see + * `gauxc/external/cube.hpp`. + * + * Points are evaluated in batches and each batch is screened against the + * per-shell cutoff radii. The batch decomposition is derived from the point + * count and the basis alone, never from the thread count, so results are + * bitwise reproducible from run to run and across thread counts. + * + * The CubeGrid overloads are not bitwise equal to the pointer overloads given + * the same points, however. A CubeGrid is walked in spatially compact tiles, + * which screen better, whereas a caller-supplied point array is batched in + * the order it arrives, so the two drop different shells. Each neglected + * shell contributes at most the shell tolerance, which makes the difference + * of the order of the tolerance rather than bounded by it: it grows with the + * number of shells dropped, and reaches roughly twice the tolerance on + * benzene in cc-pVDZ. + * + * Screening cost is governed by the shell tolerance, which the evaluator + * applies to its own private copy of the basis, so the caller's basis (and + * any SCF setup sharing it) is left untouched. Orbital error tracks the + * tolerance directly, density error goes as its square, so density tolerates + * a looser setting for the same accuracy: on a 110-atom system 1e-6 instead + * of the 1e-10 default runs ~2.5x faster on density. + * + * Instances are produced by OrbitalEvaluatorFactory, which selects the + * implementation for a given ExecutionSpace. + */ +class OrbitalEvaluator { + + using pimpl_type = detail::OrbitalEvaluatorImpl; + using pimpl_ptr_type = std::unique_ptr; + pimpl_ptr_type pimpl_; ///< Pointer to implementation instance + +public: + + // Delete default ctor + OrbitalEvaluator() = delete; + + /// Construct an OrbitalEvaluator instance from a preconstructed implementation + OrbitalEvaluator( pimpl_ptr_type&& pimpl ); + + ~OrbitalEvaluator() noexcept; + + // Non-copyable, movable + OrbitalEvaluator( const OrbitalEvaluator& ) = delete; + OrbitalEvaluator& operator=( const OrbitalEvaluator& ) = delete; + OrbitalEvaluator( OrbitalEvaluator&& ) noexcept; + OrbitalEvaluator& operator=( OrbitalEvaluator&& ) noexcept; + + /// Number of basis functions (rows of the AO matrix). + int32_t nbf() const; + + /// Underlying basis set. + const BasisSet& basis() const; + + /** @brief Evaluate a single MO chi(r) = sum_mu C[mu] * phi_mu(r). + * + * @param[in] npts Number of evaluation points. + * @param[in] points AoS array of length 3*npts: (x0,y0,z0,x1,y1,z1,...) + * in atomic units (Bohr). + * @param[in] C MO coefficient vector, length nbf(). + * @param[out] out Length-npts array of MO values. + */ + void eval_orbital( size_t npts, const double* points, + const double* C, double* out ) const; + + /** @brief Evaluate `nmo` MOs simultaneously. + * + * Storage layout: + * - `C` : (nbf, nmo) column-major, leading dimension `ldc` (>= nbf). + * - `out` : (npts, nmo) column-major, leading dimension `ldo` (>= npts). + * + * Equivalent to calling `eval_orbital` `nmo` times but amortises the AO + * collocation evaluation across all MOs (single AO buffer, GEMM contraction). + */ + void eval_orbitals( size_t npts, const double* points, + int32_t nmo, const double* C, size_t ldc, + double* out, size_t ldo ) const; + + /** @brief Evaluate the electron density + * rho(r) = sum_{mu,nu} D[mu,nu] * phi_mu(r) * phi_nu(r). + * + * @param[in] npts Number of evaluation points. + * @param[in] points AoS array of length 3*npts in Bohr. + * @param[in] D (nbf, nbf) symmetric density matrix, column-major, + * leading dimension `ldd` (>= nbf). + * @param[out] out Length-npts array of density values. + */ + void eval_density( size_t npts, const double* points, + const double* D, size_t ldd, + double* out ) const; + + /** @brief Evaluate a single MO on a CubeGrid without materialising all 3*N + * grid-point coordinates. + */ + void eval_orbital( const CubeGrid& grid, + const double* C, double* out ) const; + + /** @brief Evaluate `nmo` MOs on a CubeGrid without materialising all 3*N + * grid-point coordinates. + */ + void eval_orbitals( const CubeGrid& grid, + int32_t nmo, const double* C, size_t ldc, + double* out, size_t ldo ) const; + + /** @brief Evaluate the electron density on a CubeGrid without materialising + * all 3*N grid-point coordinates. + */ + void eval_density( const CubeGrid& grid, + const double* D, size_t ldd, + double* out ) const; + +}; // class OrbitalEvaluator + + +/// A factory to generate OrbitalEvaluator instances +class OrbitalEvaluatorFactory { + +public: + + // Delete default ctor + OrbitalEvaluatorFactory() = delete; + + /** + * @brief Construct an OrbitalEvaluator for a given execution space + * + * @param[in] ex Execution space in which to evaluate orbitals/densities. + * Currently only ExecutionSpace::Host is implemented; + * anything else throws. + * @param[in] basis Basis set (copied into the evaluator). + * @param[in] screening_tolerance Shell tolerance applied to the evaluator's + * own copy of `basis`, which sets the cutoff radii used + * for per-batch screening. A basis carried in from an SCF + * setup may be far tighter than cube-file precision needs + * and costs several-fold here for no benefit. + */ + static OrbitalEvaluator make_orbital_evaluator( ExecutionSpace ex, + BasisSet basis, + double screening_tolerance = detail::default_shell_tolerance ); + +}; // class OrbitalEvaluatorFactory + +} // namespace GauXC diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index f5b0a8b6..61968923 100644 --- a/src/CMakeLists.txt +++ b/src/CMakeLists.txt @@ -40,6 +40,8 @@ add_library( gauxc molgrid_impl.cxx molgrid_defaults.cxx atomic_radii.cxx + cube_grid.cxx + orbital_evaluator.cxx ) if( GAUXC_ENABLE_C ) diff --git a/src/cube_grid.cxx b/src/cube_grid.cxx new file mode 100644 index 00000000..bd59c1a3 --- /dev/null +++ b/src/cube_grid.cxx @@ -0,0 +1,77 @@ +/** + * GauXC Copyright (c) 2020-2024, The Regents of the University of California, + * through Lawrence Berkeley National Laboratory (subject to receipt of + * any required approvals from the U.S. Dept. of Energy). + * + * (c) 2024-2025, Microsoft Corporation + * + * All rights reserved. + * + * See LICENSE.txt for details + */ +#include + +#include + +#include + +namespace GauXC { + +CubeGrid CubeGrid::from_molecule( const Molecule& mol, int64_t nx, int64_t ny, + int64_t nz, double margin ) { + if( mol.empty() ) { + GAUXC_GENERIC_EXCEPTION("CubeGrid::from_molecule: molecule has no atoms."); + } + if( nx < 1 or ny < 1 or nz < 1 ) { + GAUXC_GENERIC_EXCEPTION("CubeGrid::from_molecule: nx, ny, nz must be >= 1."); + } + + double xmin = mol[0].x, xmax = mol[0].x; + double ymin = mol[0].y, ymax = mol[0].y; + double zmin = mol[0].z, zmax = mol[0].z; + for( const auto& a : mol ) { + xmin = std::min(xmin, a.x); + xmax = std::max(xmax, a.x); + ymin = std::min(ymin, a.y); + ymax = std::max(ymax, a.y); + zmin = std::min(zmin, a.z); + zmax = std::max(zmax, a.z); + } + + CubeGrid grid; + grid.origin = {xmin - margin, ymin - margin, zmin - margin}; + grid.nx = nx; + grid.ny = ny; + grid.nz = nz; + const double ex = (xmax - xmin) + 2.0 * margin; + const double ey = (ymax - ymin) + 2.0 * margin; + const double ez = (zmax - zmin) + 2.0 * margin; + grid.spacing[0] = nx > 1 ? ex / static_cast(nx - 1) : 0.0; + grid.spacing[1] = ny > 1 ? ey / static_cast(ny - 1) : 0.0; + grid.spacing[2] = nz > 1 ? ez / static_cast(nz - 1) : 0.0; + return grid; +} + +std::vector CubeGrid::points() const { + std::vector pts(static_cast(num_points()) * 3); + points_into(pts.data()); + return pts; +} + +void CubeGrid::points_into( double* out ) const { + size_t k = 0; + for( int64_t ix = 0; ix < nx; ++ix ) { + const double x = origin[0] + spacing[0] * static_cast(ix); + for( int64_t iy = 0; iy < ny; ++iy ) { + const double y = origin[1] + spacing[1] * static_cast(iy); + for( int64_t iz = 0; iz < nz; ++iz ) { + out[3 * k + 0] = x; + out[3 * k + 1] = y; + out[3 * k + 2] = origin[2] + spacing[2] * static_cast(iz); + ++k; + } + } + } +} + +} // namespace GauXC diff --git a/src/external/CMakeLists.txt b/src/external/CMakeLists.txt index 65fbff72..afd76887 100644 --- a/src/external/CMakeLists.txt +++ b/src/external/CMakeLists.txt @@ -9,6 +9,9 @@ # # See LICENSE.txt for details # +# Cube file writer is a self-contained utility (no third-party deps); always build it. +target_sources( gauxc PRIVATE cube.cxx ) + if( GAUXC_ENABLE_HDF5 ) include(FetchContent) find_package(HDF5) @@ -45,7 +48,7 @@ if( GAUXC_ENABLE_HDF5 ) endif() endif() - target_sources( gauxc PRIVATE hdf5_write.cxx hdf5_read.cxx ) + target_sources( gauxc PRIVATE hdf5_write.cxx hdf5_read.cxx cube_hdf5.cxx ) if( GAUXC_ENABLE_C ) target_sources( gauxc PRIVATE c_hdf5_read.cxx c_hdf5_write.cxx ) endif() diff --git a/src/external/cube.cxx b/src/external/cube.cxx new file mode 100644 index 00000000..f6a47562 --- /dev/null +++ b/src/external/cube.cxx @@ -0,0 +1,207 @@ +/** + * GauXC Copyright (c) 2020-2024, The Regents of the University of California, + * through Lawrence Berkeley National Laboratory (subject to receipt of + * any required approvals from the U.S. Dept. of Energy). + * + * (c) 2024-2025, Microsoft Corporation + * + * All rights reserved. + * + * See LICENSE.txt for details + */ +#include + +#include +#include +#include +#include +#include +#include +#include + +#include + +namespace GauXC { + +namespace { + +/** @brief Format `v` as a fixed 13-character "%13.5E" field. + * + * Layout: sign + D + '.' + 5 digits + 'E' + sign + 2 digits. snprintf in the + * inner loop dominates write time on large grids. This matches glibc + * "%13.5E" except for values lying within a few ulp of a rounding boundary + * in the 6th significant digit, where scaling the mantissa by pow(10,-exp10) + * can tip it across the boundary; the two then differ by one unit in the last + * printed digit. Values needing a 3-digit exponent defer to snprintf. + * + * @param[in] v Value to format. + * @param[out] out Exactly 13 bytes are written (no trailing NUL). + */ +inline void format_e13_5( double v, char* out ) { + + if( std::isnan(v) ) { + std::memcpy( out, std::signbit(v) ? " -NAN" : " NAN", 13 ); + return; + } + if( std::isinf(v) ) { + std::memcpy( out, v < 0 ? " -INF" : " INF", 13 ); + return; + } + + const bool negative = std::signbit(v); + const double absv = std::fabs(v); + + if( absv == 0.0 ) { + // glibc "%13.5E": 0.0 -> " 0.00000E+00"; -0.0 -> " -0.00000E+00". + std::memcpy( out, negative ? " -0.00000E+00" : " 0.00000E+00", 13 ); + return; + } + + int exp10 = static_cast( std::floor( std::log10(absv) ) ); + + // Tested before scaling: pow(10,-exp10) overflows for subnormal inputs. + bool wide_exponent = ( exp10 > 99 or exp10 < -99 ); + + long long mant_int = 0; + if( not wide_exponent ) { + const double mant = absv * std::pow( 10.0, -exp10 ); + // Over this range pow() is accurate to a few ulp, so mant cannot fall + // below 1 and only the carry out of 9.999995 needs correcting. + mant_int = std::llround( mant * 1e5 ); + if( mant_int >= 1000000 ) { + mant_int = 100000; + wide_exponent = ( ++exp10 > 99 ); + } + } + + // 3-digit exponents have a different field layout; defer to snprintf. + if( wide_exponent ) { + char tmp[32]; + // "%13.5E" sets a minimum field width of 13, so n is never below 13. + const int n = std::snprintf( tmp, sizeof(tmp), "%13.5E", v ); + std::memcpy( out, tmp + (n - 13), 13 ); + return; + } + + // Fast path. Layout: [sp][sign][D][.][d4 d3 d2 d1 d0][E][esign][e1 e0] + out[0] = ' '; + out[1] = negative ? '-' : ' '; + + char digits[6]; + for( int i = 5; i >= 0; --i ) { + digits[i] = static_cast( '0' + (mant_int % 10) ); + mant_int /= 10; + } + out[2] = digits[0]; + out[3] = '.'; + out[4] = digits[1]; + out[5] = digits[2]; + out[6] = digits[3]; + out[7] = digits[4]; + out[8] = digits[5]; + out[9] = 'E'; + out[10] = exp10 < 0 ? '-' : '+'; + + const int aexp = exp10 < 0 ? -exp10 : exp10; + out[11] = static_cast( '0' + (aexp / 10) ); + out[12] = static_cast( '0' + (aexp % 10) ); + +} + +} // namespace + + +void write_cube( const std::string& path, const Molecule& mol, + const CubeGrid& grid, const double* field, + const std::string& comment ) { + + if( not field ) { + GAUXC_GENERIC_EXCEPTION("write_cube: field pointer is null."); + } + if( grid.num_points() <= 0 ) { + GAUXC_GENERIC_EXCEPTION("write_cube: grid has zero points."); + } + + std::FILE* f = std::fopen( path.c_str(), "w" ); + if( not f ) { + GAUXC_GENERIC_EXCEPTION("write_cube: failed to open output file: " + path); + } + std::unique_ptr fh( f, &std::fclose ); + + // --- Header --- + std::fprintf( f, "%s\n", + comment.empty() ? "GauXC cube file" : comment.c_str() ); + std::fprintf( f, "Generated by GauXC\n" ); + + // natoms + origin (Bohr). + std::fprintf( f, "%5lld %12.6f %12.6f %12.6f\n", + static_cast(mol.size()), grid.origin[0], grid.origin[1], + grid.origin[2] ); + + // Three voxel-axis lines (axis-aligned grid). + std::fprintf( f, "%5lld %12.6f %12.6f %12.6f\n", + static_cast(grid.nx), grid.spacing[0], 0.0, 0.0 ); + std::fprintf( f, "%5lld %12.6f %12.6f %12.6f\n", + static_cast(grid.ny), 0.0, grid.spacing[1], 0.0 ); + std::fprintf( f, "%5lld %12.6f %12.6f %12.6f\n", + static_cast(grid.nz), 0.0, 0.0, grid.spacing[2] ); + + // One line per atom: Z, partial charge (0.0), x, y, z (Bohr). + for( const auto& atom : mol ) { + std::fprintf( f, "%5lld %12.6f %12.6f %12.6f %12.6f\n", + static_cast(atom.Z.get()), 0.0, atom.x, atom.y, atom.z ); + } + + // --- Data block --- + // Each (ix, iy) row is grid.nz values, six per line, %13.5E. The cube format + // requires a newline at the end of every row regardless of how many values + // land on the last line, so a row occupies exactly nz*13 + ceil(nz/6) bytes. + // Rows are therefore equally sized and independent: a chunk of rows is + // formatted in parallel at known offsets and committed with a single fwrite, + // with no compaction pass and a staging buffer that stays bounded regardless + // of grid size. + const int64_t nz = grid.nz; + const int64_t bytes_per_row = nz * 13 + (nz + 5) / 6; + const int64_t n_rows = grid.nx * grid.ny; + + // 8 MiB is well past the point where sequential fwrite stops caring about + // block size, and keeps the staging buffer small enough to stay cache- and + // NUMA-friendly on large grids. + constexpr int64_t target_chunk_bytes = 8ll * 1024 * 1024; + const int64_t rows_per_chunk = + std::clamp( target_chunk_bytes / bytes_per_row, 1, n_rows ); + std::vector buf( static_cast(rows_per_chunk * bytes_per_row) ); + + for( int64_t r0 = 0; r0 < n_rows; r0 += rows_per_chunk ) { + + const int64_t nr = std::min( rows_per_chunk, n_rows - r0 ); + +#pragma omp parallel for schedule(static) + for( int64_t r = 0; r < nr; ++r ) { + const double* row_data = + field + static_cast(r0 + r) * static_cast(nz); + char* dst = + buf.data() + static_cast(r) * static_cast(bytes_per_row); + int64_t off = 0; + for( int64_t iz = 0; iz < nz; ++iz ) { + format_e13_5( row_data[iz], dst + off ); + off += 13; + // Every 6 values OR at the end of the row -> newline. + if( (iz + 1) % 6 == 0 or iz + 1 == nz ) dst[off++] = '\n'; + } + } + + const size_t nbytes = static_cast( nr * bytes_per_row ); + if( std::fwrite( buf.data(), 1, nbytes, f ) != nbytes ) { + GAUXC_GENERIC_EXCEPTION("write_cube: short write to " + path); + } + + } + + if( std::fclose( fh.release() ) != 0 ) { + GAUXC_GENERIC_EXCEPTION("write_cube: failed to close " + path); + } + +} + +} // namespace GauXC diff --git a/src/external/cube_hdf5.cxx b/src/external/cube_hdf5.cxx new file mode 100644 index 00000000..db909d57 --- /dev/null +++ b/src/external/cube_hdf5.cxx @@ -0,0 +1,81 @@ +/** + * GauXC Copyright (c) 2020-2024, The Regents of the University of California, + * through Lawrence Berkeley National Laboratory (subject to receipt of + * any required approvals from the U.S. Dept. of Energy). + * + * (c) 2024-2025, Microsoft Corporation + * + * All rights reserved. + * + * See LICENSE.txt for details + */ +#include +#ifdef GAUXC_HAS_HDF5 + +#include +#include + +#include + +#include + +namespace GauXC { + +void write_cube_hdf5( const std::string& path, + const Molecule& mol, + const CubeGrid& grid, + const double* field, + const std::string& comment ) { + if( grid.num_points() <= 0 ) { + GAUXC_GENERIC_EXCEPTION("write_cube_hdf5: grid has zero points."); + } + if( not field ) { + GAUXC_GENERIC_EXCEPTION("write_cube_hdf5: null field pointer."); + } + + HighFive::File file(path, HighFive::File::Overwrite); + + // --- Field: 3D dataset (nx, ny, nz) --- + const auto nx = static_cast(grid.nx); + const auto ny = static_cast(grid.ny); + const auto nz = static_cast(grid.nz); + + auto ds_field = file.createDataSet( + "field", HighFive::DataSpace({nx, ny, nz})); + ds_field.write_raw(field); + + // Comment attribute on /field. + const std::string cmt = + comment.empty() ? "Cube field generated by GauXC" : comment; + ds_field.createAttribute("comment", cmt); + + // --- Grid metadata --- + auto grp_grid = file.createGroup("grid"); + grp_grid.createDataSet("origin", + std::vector(grid.origin.begin(), grid.origin.end())); + grp_grid.createDataSet("spacing", + std::vector(grid.spacing.begin(), grid.spacing.end())); + grp_grid.createDataSet("shape", + std::vector{grid.nx, grid.ny, grid.nz}); + + // --- Atoms --- + const size_t natom = mol.size(); + auto grp_atoms = file.createGroup("atoms"); + + std::vector atomic_numbers(natom); + std::vector coords(natom * 3); + for( size_t i = 0; i < natom; ++i ) { + atomic_numbers[i] = static_cast(mol[i].Z.get()); + coords[3 * i + 0] = mol[i].x; + coords[3 * i + 1] = mol[i].y; + coords[3 * i + 2] = mol[i].z; + } + + grp_atoms.createDataSet("Z", atomic_numbers); + grp_atoms.createDataSet("coords", + HighFive::DataSpace({natom, 3ul})).write_raw(coords.data()); +} + +} // namespace GauXC + +#endif // GAUXC_HAS_HDF5 diff --git a/src/orbital_evaluator.cxx b/src/orbital_evaluator.cxx new file mode 100644 index 00000000..c67ff5c2 --- /dev/null +++ b/src/orbital_evaluator.cxx @@ -0,0 +1,707 @@ +/** + * GauXC Copyright (c) 2020-2024, The Regents of the University of California, + * through Lawrence Berkeley National Laboratory (subject to receipt of + * any required approvals from the U.S. Dept. of Energy). + * + * (c) 2024-2025, Microsoft Corporation + * + * All rights reserved. + * + * See LICENSE.txt for details + */ +#include "orbital_evaluator_impl.hpp" + +#include +#include +#include +#include +#include +#include +#include + +#include +#include + +#include "xc_integrator/integrator_util/integrator_common.hpp" +#include "xc_integrator/local_work_driver/host/blas.hpp" + +namespace GauXC { + +namespace detail { + +OrbitalEvaluatorImpl::OrbitalEvaluatorImpl( BasisSet bs, + double screening_tolerance ) : + basis( std::move(bs) ) { + + // Retune our own copy: the caller's basis, which an SCF setup may share, + // is never touched. + for( auto& sh : basis ) sh.set_shell_tolerance( screening_tolerance ); + + driver_owner = LocalWorkDriverFactory::make_local_work_driver( + ExecutionSpace::Host, "Reference" ); + host_driver = dynamic_cast( driver_owner.get() ); + if( not host_driver ) { + GAUXC_GENERIC_EXCEPTION("OrbitalEvaluator: LocalWorkDriverFactory did not " + "return a LocalHostWorkDriver."); + } + + nbf_ = basis.nbf(); + + // BasisSetMap powers shell -> AO range lookups for the screened submat_map. + // No Molecule is available here, so an empty one is passed; the only field + // that needs it (shell_to_center) is unused by this class. + basis_map = std::make_unique( basis, Molecule{} ); + + // Cache squared cutoff radii so the screening loop is a comparison rather + // than a square root per shell per batch. + shell_cutoff_r2.resize( basis.size() ); + for( size_t s = 0; s < basis.size(); ++s ) { + const double r = basis[s].cutoff_radius(); + shell_cutoff_r2[s] = r * r; + } + +} + +} // namespace detail + + +namespace { + +/** @brief Choose the number of points evaluated per batch. + * + * Three bounds, listed by which one binds as the basis grows. + * + * 1. Cache footprint, on small and medium bases. Each thread holds `nscratch` + * live blocks of nbf*batch doubles -- the AO block, the collocation + * kernel's own transpose staging, and, for density, the D*AO product -- + * sized to share ~1 MiB, a typical per-core L2, rather than allowing each + * block that much on its own. + * + * 2. A minimum batch count, so every thread of a wide machine gets several. + * Fixed rather than a multiple of the thread count: batch shape decides + * which shells screen out, so taking the thread count here would make the + * numbers depend on it. 1024 covers 256 threads at four batches each and + * costs a few percent either way on a narrower one. + * + * 3. A floor, which is what binds on large bases. The cache bound keeps + * shrinking with nbf -- six points at ubiquitin/cc-pVDZ -- while the + * measured optimum does not: a batch that thin cannot amortise the + * per-batch nbe*nbe gather and leaves the GEMM too skinny to run well. + * Serial CPU time over batches of 32/64/128/256/512 puts the optimum at + * 32-64 for both taxol (nbf 1099) and ubiquitin (nbf 11577), with 128 + * costing ~10% and 512 costing 2.5x. Above nbf ~340 this floor, not the + * cache bound, sets the batch, so the ~1 MiB target is not met there. + * It does not defeat bound 2: at 64 points a batch, any grid of 65536 + * points or more already yields the 1024 batches that bound asks for. + * + * Only the batch-proportional blocks are budgeted here. The gathers that + * depend on nbe alone -- D compressed to nbe*nbe, C compressed to nbe*nmo -- + * are unaffected by the batch size, so no choice made here bounds them; on a + * large basis with weak screening the nbe*nbe gather dominates this budget. + */ +size_t choose_batch_size( int32_t nbf, size_t npts, int nscratch ) { + constexpr size_t kTargetScratchBytesPerThread = 1024 * 1024; + constexpr size_t kMinBatch = 64; + constexpr size_t kMaxBatch = 8192; + constexpr size_t kMinBatches = 1024; + + size_t batch = kMaxBatch; + if( nbf > 0 ) { + batch = kTargetScratchBytesPerThread / + ( sizeof(double) * static_cast(nbf) * + static_cast(nscratch) ); + } + batch = std::min( batch, npts / kMinBatches ); + return std::clamp( batch, kMinBatch, kMaxBatch ); +} + +/// Axis-aligned bounding box of a point set. +struct PointBbox { + std::array lo; + std::array hi; +}; + +/// Bbox of `npts` AoS points (array of length 3*npts). +PointBbox compute_bbox( const double* points, size_t npts ) { + PointBbox b{ {points[0], points[1], points[2]}, + {points[0], points[1], points[2]} }; + for( size_t p = 1; p < npts; ++p ) { + const double* xyz = points + 3 * p; + for( int k = 0; k < 3; ++k ) { + if( xyz[k] < b.lo[k] ) b.lo[k] = xyz[k]; + if( xyz[k] > b.hi[k] ) b.hi[k] = xyz[k]; + } + } + return b; +} + +/// Squared distance from `center` to the nearest point of the bbox. Zero if +/// the center lies inside. +double dist2_center_to_bbox( const double* center, const PointBbox& bbox ) { + double d2 = 0.0; + for( int k = 0; k < 3; ++k ) { + const double c = center[k]; + if( c < bbox.lo[k] ) { + const double dx = bbox.lo[k] - c; + d2 += dx * dx; + } else if( c > bbox.hi[k] ) { + const double dx = c - bbox.hi[k]; + d2 += dx * dx; + } + } + return d2; +} + +/** @brief The output indices covered by one batch. + * + * `nruns` runs of `run_len` consecutive indices starting at run_off[i]. A + * contiguous batch is the single-run case; a grid tile contributes one run + * per (ix,iy) pair. + */ +struct BatchSpan { + const size_t* run_off = nullptr; + size_t nruns = 0; + size_t run_len = 0; + size_t npts = 0; +}; + +/// Batch source over a caller-supplied AoS coordinate array. +struct RawPointSource { + static constexpr bool needs_scratch = false; + const double* points; + size_t npts_total; + size_t batch_size; + + size_t num_batches() const { + return ( npts_total + batch_size - 1 ) / batch_size; + } + size_t max_points() const { return batch_size; } + + const double* batch( size_t b, double*, std::vector& runs, + BatchSpan& span, PointBbox& bbox ) const { + const size_t p0 = b * batch_size; + const size_t np = std::min( batch_size, npts_total - p0 ); + runs.assign( 1, p0 ); + span = BatchSpan{ runs.data(), 1, np, np }; + const double* pts = points + 3 * p0; + bbox = compute_bbox( pts, np ); + return pts; + } +}; + +/** @brief Batch source that walks a CubeGrid in spatially compact tiles. + * + * A contiguous index range is a needle along z, and as soon as it crosses a + * row boundary its bbox spans the whole z extent, so distant shells survive + * screening. A tile of the same point count draws a far tighter box, which + * cuts nbe and therefore both the collocation (~nbe) and the density + * contraction (~nbe^2). The cost is that a tile's points are no longer + * contiguous in the output and have to be scattered. + * + * Tiling is the only grid traversal because a contiguous one was never + * measured to win. Counting nbe*np and nbe^2*np over every batch (a + * deterministic proxy for collocation and GEMM cost) across benzene and + * taxol, grids of 64^3 to 200^3, margins of 3 and 6 Bohr, and a single-plane + * grid where the scatter degenerates to runs of one point, tiling never did + * more work: from a 2-3% saving on benzene at 128^3 to a 53% saving on taxol + * at 200^3. Paired single-threaded timings agreed in direction throughout. + */ +struct GridTileSource { + static constexpr bool needs_scratch = true; + const CubeGrid* grid; + int64_t tx, ty, tz; ///< tile extent in grid points + int64_t ntx, nty, ntz; ///< tile counts per axis + + size_t num_batches() const { + return static_cast(ntx) * static_cast(nty) * + static_cast(ntz); + } + size_t max_points() const { + return static_cast(tx) * static_cast(ty) * + static_cast(tz); + } + + const double* batch( size_t b, double* scr, std::vector& runs, + BatchSpan& span, PointBbox& bbox ) const { + const CubeGrid& g = *grid; + const int64_t bi = static_cast(b); + const int64_t ix0 = ( bi / (nty * ntz) ) * tx; + const int64_t iy0 = ( ( bi / ntz ) % nty ) * ty; + const int64_t iz0 = ( bi % ntz ) * tz; + const int64_t ix1 = std::min( ix0 + tx, g.nx ); + const int64_t iy1 = std::min( iy0 + ty, g.ny ); + const int64_t iz1 = std::min( iz0 + tz, g.nz ); + const int64_t nzr = iz1 - iz0; + + runs.clear(); + size_t n = 0; + for( int64_t ix = ix0; ix < ix1; ++ix ) { + const double x = g.origin[0] + g.spacing[0] * static_cast(ix); + for( int64_t iy = iy0; iy < iy1; ++iy ) { + const double y = g.origin[1] + g.spacing[1] * static_cast(iy); + runs.push_back( + static_cast( (ix * g.ny + iy) * g.nz + iz0 ) ); + for( int64_t iz = iz0; iz < iz1; ++iz ) { + scr[3 * n + 0] = x; + scr[3 * n + 1] = y; + scr[3 * n + 2] = + g.origin[2] + g.spacing[2] * static_cast(iz); + ++n; + } + } + } + span = BatchSpan{ runs.data(), runs.size(), + static_cast(nzr), n }; + + const int64_t lo_idx[3] = { ix0, iy0, iz0 }; + const int64_t hi_idx[3] = { ix1 - 1, iy1 - 1, iz1 - 1 }; + for( int k = 0; k < 3; ++k ) { + const double a = + g.origin[k] + g.spacing[k] * static_cast(lo_idx[k]); + const double c = + g.origin[k] + g.spacing[k] * static_cast(hi_idx[k]); + bbox.lo[k] = std::min(a, c); + bbox.hi[k] = std::max(a, c); + } + return scr; + } +}; + +/// Tile extents holding roughly `target` points and as close to cubic in +/// physical space as the (generally anisotropic) grid spacing allows. +GridTileSource make_tile_source( const CubeGrid& g, size_t target ) { + const int64_t n[3] = { g.nx, g.ny, g.nz }; + std::array s{}; + for( int k = 0; k < 3; ++k ) { + s[k] = std::fabs( g.spacing[k] ); + if( !(s[k] > 0.0) ) s[k] = 1.0; // degenerate axis: treat as unconstrained + } + + std::array t{ 1, 1, 1 }; + double scale = std::cbrt( static_cast(target) * s[0] * s[1] * s[2] ); + int64_t prod = 1; + for( int attempt = 0; attempt < 4; ++attempt ) { + prod = 1; + for( int k = 0; k < 3; ++k ) { + t[k] = std::clamp( + static_cast( std::llround( scale / s[k] ) ), 1, n[k] ); + prod *= t[k]; + } + if( static_cast(prod) <= target ) break; + scale *= std::cbrt( static_cast(target) / + static_cast(prod) ); + } + + // Each rescale shrinks the overshoot as r -> r^(2/3), so a grid anisotropic + // enough to start far above target is still above it after four. Halving the + // longest axis converges regardless, at the cost of a less cubic tile. + while( static_cast(prod) > target ) { + const int kmax = static_cast( + std::max_element( t.begin(), t.end() ) - t.begin() ); + if( t[kmax] == 1 ) break; // nothing left to halve + prod /= t[kmax]; + t[kmax] = ( t[kmax] + 1 ) / 2; + prod *= t[kmax]; + } + + GridTileSource src; + src.grid = &g; + src.tx = t[0]; src.ty = t[1]; src.tz = t[2]; + src.ntx = (n[0] + t[0] - 1) / t[0]; + src.nty = (n[1] + t[1] - 1) / t[1]; + src.ntz = (n[2] + t[2] - 1) / t[2]; + return src; +} + +/// Contraction of the AO batch against MO coefficients: +/// out(np,nmo) = ao^T(np,nbe) @ C(nbe,nmo). +class OrbitalContractor { + + const detail::OrbitalEvaluatorImpl& impl_; + int32_t nmo_; + const double* C_; + size_t ldc_; + double* out_; + size_t ldo_; + +public: + + /// nbf*batch blocks this contraction keeps live, beyond the AO block. + static constexpr int scratch_blocks = 0; + + struct Scratch { + std::vector C_compressed; + std::vector staging; + }; + + OrbitalContractor( const detail::OrbitalEvaluatorImpl& impl, int32_t nmo, + const double* C, size_t ldc, double* out, size_t ldo ) : + impl_(impl), nmo_(nmo), C_(C), ldc_(ldc), out_(out), ldo_(ldo) {} + + void zero( const BatchSpan& span ) const { + for( int32_t j = 0; j < nmo_; ++j ) { + double* out_col = out_ + static_cast(j) * ldo_; + for( size_t r = 0; r < span.nruns; ++r ) + std::fill( out_col + span.run_off[r], + out_col + span.run_off[r] + span.run_len, 0.0 ); + } + } + + void apply( Scratch& scr, const BatchSpan& span, int32_t nbe, + const std::vector& shells, const double* ao ) const { + + const size_t np = span.npts; + + // Gather the rows of C for the surviving shells into a contiguous + // (nbe, nmo) col-major buffer so the contraction is a dense GEMM. + // Skipping this when no shell is screened out was measured and is not + // worth it: that case is 0-5% of batches on anything larger than water, + // and it is precisely the case where there is nothing to gather. + const size_t c_size = static_cast(nbe) * nmo_; + if( scr.C_compressed.size() < c_size ) scr.C_compressed.resize( c_size ); + + const BasisSetMap& basis_map = *impl_.basis_map; + int32_t row = 0; + for( int32_t ish : shells ) { + const auto rng = basis_map.shell_to_ao_range(ish); + for( int32_t j = 0; j < nmo_; ++j ) { + const double* C_col = C_ + static_cast(j) * ldc_; + double* dst = scr.C_compressed.data() + + static_cast(j) * nbe + row; + std::copy( C_col + rng.first, C_col + rng.second, dst ); + } + row += rng.second - rng.first; + } + + if( span.nruns == 1 ) { + blas::gemm( 'T', 'N', static_cast(np), nmo_, nbe, 1.0, + ao, nbe, scr.C_compressed.data(), nbe, 0.0, + out_ + span.run_off[0], static_cast(ldo_) ); + return; + } + + const size_t need = np * static_cast(nmo_); + if( scr.staging.size() < need ) scr.staging.resize( need ); + blas::gemm( 'T', 'N', static_cast(np), nmo_, nbe, 1.0, + ao, nbe, scr.C_compressed.data(), nbe, 0.0, + scr.staging.data(), static_cast(np) ); + + for( int32_t j = 0; j < nmo_; ++j ) { + const double* src = scr.staging.data() + static_cast(j) * np; + double* out_col = out_ + static_cast(j) * ldo_; + for( size_t r = 0; r < span.nruns; ++r ) + std::copy( src + r * span.run_len, src + (r + 1) * span.run_len, + out_col + span.run_off[r] ); + } + } + +}; + +/// Contraction of the AO batch against the density matrix: +/// rho(r) = sum_{mu,nu} D[mu,nu] phi_mu(r) phi_nu(r). +class DensityContractor { + + const detail::OrbitalEvaluatorImpl& impl_; + const double* D_; + size_t ldd_; + double* out_; + +public: + + /// dm_ao is a second nbf*batch block alongside the AO block. + static constexpr int scratch_blocks = 1; + + struct Scratch { + std::vector dm_ao; + std::vector xmat_scr; + std::vector staging; + }; + + DensityContractor( const detail::OrbitalEvaluatorImpl& impl, const double* D, + size_t ldd, double* out ) : + impl_(impl), D_(D), ldd_(ldd), out_(out) {} + + void zero( const BatchSpan& span ) const { + for( size_t r = 0; r < span.nruns; ++r ) + std::fill( out_ + span.run_off[r], + out_ + span.run_off[r] + span.run_len, 0.0 ); + } + + void apply( Scratch& scr, const BatchSpan& span, int32_t nbe, + const std::vector& shells, const double* ao ) const { + + const size_t np = span.npts; + + const int32_t nbf = impl_.nbf_; + LocalHostWorkDriver::submat_map_t submat_map; + std::tie( submat_map, std::ignore ) = + gen_compressed_submat_map( *impl_.basis_map, shells, nbf, nbf ); + + // eval_xmat gathers D into scratch only when the surviving shells span + // more than one contiguous AO block; a single block is read in place from + // D itself, so nbe*nbe is never touched. Growing on demand inside the + // parallel region keeps the high-water mark at the largest nbe this thread + // actually saw (not nbf) and first-touches the pages on the owning thread. + if( submat_map.size() > 1 ) { + const size_t scr_size = static_cast(nbe) * nbe; + if( scr.xmat_scr.size() < scr_size ) scr.xmat_scr.resize( scr_size ); + } + const size_t dm_ao_size = static_cast(nbe) * np; + if( scr.dm_ao.size() < dm_ao_size ) scr.dm_ao.resize( dm_ao_size ); + + // dm_ao = D_compressed @ ao + impl_.host_driver->eval_xmat( np, static_cast(nbf), + static_cast(nbe), submat_map, 1.0, D_, ldd_, + ao, static_cast(nbe), + scr.dm_ao.data(), static_cast(nbe), scr.xmat_scr.data() ); + + // rho[p] = sum_mu ao(mu, p) * dm_ao(mu, p) + double* den = out_ + span.run_off[0]; + if( span.nruns > 1 ) { + if( scr.staging.size() < np ) scr.staging.resize( np ); + den = scr.staging.data(); + } + impl_.host_driver->eval_uvvar_lda_rks( np, static_cast(nbe), ao, + scr.dm_ao.data(), static_cast(nbe), den ); + + if( span.nruns > 1 ) { + for( size_t r = 0; r < span.nruns; ++r ) + std::copy( den + r * span.run_len, den + (r + 1) * span.run_len, + out_ + span.run_off[r] ); + } + } + +}; + +/** @brief The single batched evaluation loop shared by all entry points. + * + * Batches of points are screened against the per-shell cutoff radii, the + * surviving shells are collocated into a per-thread AO buffer, and the + * result is handed to `contract` for the orbital- or density-specific + * reduction. All scratch is per-thread and per-call. + */ +template +void batched_eval( const detail::OrbitalEvaluatorImpl& impl, + const BatchSource& src, const Contractor& contract ) { + + const BasisSet& basis = impl.basis; + const auto& shell_cutoff_r2 = impl.shell_cutoff_r2; + const int32_t nshells_total = basis.nshells(); + + const size_t max_pts = src.max_points(); + const int64_t n_batches = static_cast( src.num_batches() ); + +#pragma omp parallel + { + std::vector pt_buf; + if constexpr( BatchSource::needs_scratch ) pt_buf.resize( 3 * max_pts ); + std::vector run_storage; + std::vector ao_buf; + std::vector screened_shells; + screened_shells.reserve( nshells_total ); + typename Contractor::Scratch scr; + +#pragma omp for schedule(dynamic, 1) + for( int64_t b = 0; b < n_batches; ++b ) { + + BatchSpan span; + PointBbox bbox; + const double* pts = + src.batch( static_cast(b), pt_buf.data(), run_storage, span, + bbox ); + const size_t np = span.npts; + + // Per-batch shell screening: keep shells whose cutoff radius reaches + // any point of this batch's bounding box. + screened_shells.clear(); + int32_t nbe = 0; + for( int32_t s = 0; s < nshells_total; ++s ) + if( dist2_center_to_bbox( basis[s].O_data(), bbox ) < shell_cutoff_r2[s] ) { + screened_shells.push_back(s); + nbe += basis[s].size(); + } + + // All shells screen out -> the result is identically zero here. + if( not nbe ) { contract.zero( span ); continue; } + + // Only the surviving nbe rows are written, so the AO block tracks the + // largest nbe seen rather than the nbf worst case. + const size_t ao_size = static_cast(nbe) * np; + if( ao_buf.size() < ao_size ) ao_buf.resize( ao_size ); + + // Both contractions assume eval_collocation packs AO rows by cumulative + // shell size in shell_list order. That holds for the gau2grid path; the + // non-gau2grid fallback in gau2grid_collocation.cxx instead uses the + // global shell_to_first_ao offset, which is only equivalent when no + // shell is screened out. That fallback is unreachable in any supported + // build (the top-level CMakeLists makes gau2grid a hard dependency), and + // the discrepancy is pre-existing rather than introduced here. + impl.host_driver->eval_collocation( np, screened_shells.size(), + static_cast(nbe), pts, basis, screened_shells.data(), + ao_buf.data() ); + + contract.apply( scr, span, nbe, screened_shells, ao_buf.data() ); + + } + } + +} + +/// Target points per batch for a given contraction, accounting for every +/// nbf*batch block it keeps live alongside the AO block. +template +size_t batch_target( int32_t nbf, size_t npts ) { + constexpr int collocation_scratch_blocks = 1; + const int nscratch = + 1 + collocation_scratch_blocks + Contractor::scratch_blocks; + return choose_batch_size( nbf, npts, nscratch ); +} + +void check_orbital_args( const std::string& ctx, size_t npts, int32_t nbf, + const double* C, size_t ldc, const double* out, + size_t ldo ) { + if( not C or not out ) + GAUXC_GENERIC_EXCEPTION( ctx + ": null pointer argument." ); + if( ldc < static_cast(nbf) ) + GAUXC_GENERIC_EXCEPTION( ctx + ": ldc must be >= nbf()." ); + if( ldo < npts ) + GAUXC_GENERIC_EXCEPTION( ctx + ": ldo must be >= npts." ); + // ldo is passed to blas::gemm, whose LDC parameter is int. + if( ldo > static_cast(std::numeric_limits::max()) ) + GAUXC_GENERIC_EXCEPTION( ctx + ": ldo exceeds BLAS int range." ); +} + +void check_density_args( const std::string& ctx, int32_t nbf, const double* D, + size_t ldd, const double* out ) { + if( not D or not out ) + GAUXC_GENERIC_EXCEPTION( ctx + ": null pointer argument." ); + if( ldd < static_cast(nbf) ) + GAUXC_GENERIC_EXCEPTION( ctx + ": ldd must be >= nbf()." ); + // ldd reaches blas::gemm through eval_xmat, which narrows it to int32_t. + if( ldd > static_cast(std::numeric_limits::max()) ) + GAUXC_GENERIC_EXCEPTION( ctx + ": ldd exceeds BLAS int range." ); +} + +} // namespace + + +OrbitalEvaluator::OrbitalEvaluator( pimpl_ptr_type&& pimpl ) : + pimpl_( std::move(pimpl) ) { + if( not pimpl_ ) GAUXC_PIMPL_NOT_INITIALIZED(); +} + +OrbitalEvaluator::~OrbitalEvaluator() noexcept = default; +OrbitalEvaluator::OrbitalEvaluator( OrbitalEvaluator&& ) noexcept = default; +OrbitalEvaluator& OrbitalEvaluator::operator=( OrbitalEvaluator&& ) noexcept = + default; + +int32_t OrbitalEvaluator::nbf() const { return pimpl_->nbf_; } + +const BasisSet& OrbitalEvaluator::basis() const { + return pimpl_->basis; +} + + +OrbitalEvaluator OrbitalEvaluatorFactory::make_orbital_evaluator( + ExecutionSpace ex, BasisSet basis, double screening_tolerance ) { + + switch(ex) { + + case ExecutionSpace::Host: + return OrbitalEvaluator( + std::make_unique( std::move(basis), + screening_tolerance ) + ); + + default: + GAUXC_GENERIC_EXCEPTION("OrbitalEvaluator: only ExecutionSpace::Host is " + "currently supported."); + + } + +} + + +// --------------------------------------------------------------------------- +// Caller-supplied point sets +// --------------------------------------------------------------------------- + +void OrbitalEvaluator::eval_orbital( size_t npts, const double* points, + const double* C, double* out ) const { + eval_orbitals( npts, points, /*nmo=*/1, C, + /*ldc=*/static_cast(pimpl_->nbf_), out, /*ldo=*/npts ); +} + +void OrbitalEvaluator::eval_orbitals( size_t npts, const double* points, + int32_t nmo, const double* C, size_t ldc, + double* out, size_t ldo ) const { + if( not npts or nmo < 1 ) return; + if( not points ) + GAUXC_GENERIC_EXCEPTION( + "OrbitalEvaluator::eval_orbitals: null pointer argument."); + check_orbital_args( "OrbitalEvaluator::eval_orbitals", npts, pimpl_->nbf_, C, + ldc, out, ldo ); + + batched_eval( *pimpl_, + RawPointSource{ points, npts, + batch_target( pimpl_->nbf_, npts ) }, + OrbitalContractor( *pimpl_, nmo, C, ldc, out, ldo ) ); +} + +void OrbitalEvaluator::eval_density( size_t npts, const double* points, + const double* D, size_t ldd, + double* out ) const { + if( not npts ) return; + if( not points ) + GAUXC_GENERIC_EXCEPTION( + "OrbitalEvaluator::eval_density: null pointer argument."); + check_density_args( "OrbitalEvaluator::eval_density", pimpl_->nbf_, D, ldd, + out ); + + batched_eval( *pimpl_, + RawPointSource{ points, npts, + batch_target( pimpl_->nbf_, npts ) }, + DensityContractor( *pimpl_, D, ldd, out ) ); +} + + +// --------------------------------------------------------------------------- +// CubeGrid overloads: generate per-batch point coordinates on-the-fly, +// avoiding the 3*num_points()*8 byte temporary coordinate array. +// --------------------------------------------------------------------------- + +void OrbitalEvaluator::eval_orbital( const CubeGrid& grid, const double* C, + double* out ) const { + eval_orbitals( grid, /*nmo=*/1, C, + /*ldc=*/static_cast(pimpl_->nbf_), out, + /*ldo=*/static_cast(grid.num_points()) ); +} + +void OrbitalEvaluator::eval_orbitals( const CubeGrid& grid, int32_t nmo, + const double* C, size_t ldc, double* out, + size_t ldo ) const { + if( grid.num_points() <= 0 or nmo < 1 ) return; + const size_t npts = static_cast( grid.num_points() ); + check_orbital_args( "OrbitalEvaluator::eval_orbitals(grid)", npts, + pimpl_->nbf_, C, ldc, out, ldo ); + + batched_eval( *pimpl_, + make_tile_source( grid, + batch_target( pimpl_->nbf_, npts ) ), + OrbitalContractor( *pimpl_, nmo, C, ldc, out, ldo ) ); +} + +void OrbitalEvaluator::eval_density( const CubeGrid& grid, const double* D, + size_t ldd, double* out ) const { + if( grid.num_points() <= 0 ) return; + const size_t npts = static_cast( grid.num_points() ); + check_density_args( "OrbitalEvaluator::eval_density(grid)", pimpl_->nbf_, D, + ldd, out ); + + batched_eval( *pimpl_, + make_tile_source( grid, + batch_target( pimpl_->nbf_, npts ) ), + DensityContractor( *pimpl_, D, ldd, out ) ); +} + +} // namespace GauXC diff --git a/src/orbital_evaluator_impl.hpp b/src/orbital_evaluator_impl.hpp new file mode 100644 index 00000000..60fe6b38 --- /dev/null +++ b/src/orbital_evaluator_impl.hpp @@ -0,0 +1,42 @@ +/** + * GauXC Copyright (c) 2020-2024, The Regents of the University of California, + * through Lawrence Berkeley National Laboratory (subject to receipt of + * any required approvals from the U.S. Dept. of Energy). + * + * (c) 2024-2025, Microsoft Corporation + * + * All rights reserved. + * + * See LICENSE.txt for details + */ +#pragma once + +#include +#include + +#include +#include +#include +#include + +#include "xc_integrator/local_work_driver/host/local_host_work_driver.hpp" + +namespace GauXC::detail { + +/// Host implementation state for OrbitalEvaluator +class OrbitalEvaluatorImpl { + +public: + + BasisSet basis; ///< Basis set to evaluate + std::unique_ptr driver_owner; + LocalHostWorkDriver* host_driver = nullptr; ///< Non-owning view of the driver + std::unique_ptr basis_map; ///< shell -> AO range lookups + std::vector shell_cutoff_r2; ///< Per-shell squared cutoff radius + int32_t nbf_ = 0; + + OrbitalEvaluatorImpl( BasisSet bs, double screening_tolerance ); + +}; // class OrbitalEvaluatorImpl + +} // namespace GauXC::detail diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index 817b1043..49e78829 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -69,6 +69,7 @@ add_executable( gauxc_test basis/parse_basis.cxx dd_psi_potential_test.cxx 2nd_derivative_test.cxx + orbital_evaluator_test.cxx c_api_test.cxx ) target_link_libraries( gauxc_test PUBLIC gauxc gauxc_catch2 Eigen3::Eigen ) @@ -82,6 +83,9 @@ endif() set( GAUXC_REF_DATA_PATH "${PROJECT_SOURCE_DIR}/tests/ref_data" ) +# Scratch directory for tests which write files (cube I/O) +set( GAUXC_TEST_TMP_PATH "${PROJECT_BINARY_DIR}/tests/tmp" ) +file( MAKE_DIRECTORY "${GAUXC_TEST_TMP_PATH}" ) configure_file( ut_common.hpp.in ${PROJECT_BINARY_DIR}/tests/ut_common.hpp ) target_include_directories( gauxc_test PRIVATE ${PROJECT_BINARY_DIR}/tests ) target_include_directories( gauxc_test PRIVATE ${PROJECT_SOURCE_DIR}/tests ) diff --git a/tests/orbital_evaluator_test.cxx b/tests/orbital_evaluator_test.cxx new file mode 100644 index 00000000..e0a09383 --- /dev/null +++ b/tests/orbital_evaluator_test.cxx @@ -0,0 +1,1532 @@ +/** + * GauXC Copyright (c) 2020-2024, The Regents of the University of California, + * through Lawrence Berkeley National Laboratory (subject to receipt of + * any required approvals from the U.S. Dept. of Energy). + * + * (c) 2024-2025, Microsoft Corporation + * + * All rights reserved. + * + * See LICENSE.txt for details + */ +#include "ut_common.hpp" +#include "catch2/catch.hpp" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include +#include +#include +#include +#ifdef _OPENMP +#include +#endif +#ifdef GAUXC_HAS_HDF5 +#include +#endif +#ifdef GAUXC_HAS_MPI +#include +#endif + +#include "standards.hpp" + +// Reference implementation lives in src/, behind the public header surface. +#include "xc_integrator/local_work_driver/host/local_host_work_driver.hpp" + +using namespace GauXC; + +namespace { + +/// Unique path for a test output file, inside the build-tree scratch +/// directory configured by CMake. +std::string make_temp_path(const char* suffix) { + static int counter = 0; + return std::string(GAUXC_TEST_TMP_PATH) + "/gauxc_test_" + + std::to_string(++counter) + suffix; +} + +OrbitalEvaluator make_evaluator(const BasisSet& basis, double tol) { + return OrbitalEvaluatorFactory::make_orbital_evaluator(ExecutionSpace::Host, + basis, tol); +} + +/// AO collocation over every shell, i.e. the unscreened reference. +std::vector reference_collocation(const BasisSet& basis, + int64_t npts, const double* pts) { + const int32_t nbf = basis.nbf(); + std::vector ao(static_cast(nbf) * npts); + auto drv = LocalWorkDriverFactory::make_local_work_driver( + ExecutionSpace::Host, "Reference"); + auto* host_drv = dynamic_cast(drv.get()); + REQUIRE(host_drv != nullptr); + std::vector shell_list(basis.size()); + for (size_t i = 0; i < shell_list.size(); ++i) + shell_list[i] = static_cast(i); + host_drv->eval_collocation(static_cast(npts), + static_cast(basis.nshells()), + static_cast(nbf), pts, basis, + shell_list.data(), ao.data()); + return ao; +} + +/// Build a deterministic PRNG-based set of points within a small bounding box +/// around the molecule. Avoids putting samples too close to nuclei to keep +/// values numerically well-behaved. +std::vector make_random_points(int64_t npts, unsigned seed = 1234u) { + std::mt19937 gen(seed); + std::uniform_real_distribution u(-2.5, 2.5); + std::vector pts(static_cast(npts) * 3); + for (int64_t p = 0; p < npts; ++p) { + pts[3 * p + 0] = u(gen); + pts[3 * p + 1] = u(gen); + pts[3 * p + 2] = u(gen); + } + return pts; +} + +std::vector make_random_vector(size_t n, unsigned seed) { + std::mt19937 gen(seed); + std::uniform_real_distribution u(-1.0, 1.0); + std::vector v(n); + for (auto& x : v) x = u(gen); + return v; +} + +} // namespace + +TEST_CASE("OrbitalEvaluator / Water cc-pVDZ matches eval_collocation", + "[orbital_evaluator]") { + auto mol = make_water(); + auto basis = make_ccpvdz(mol, SphericalType(true)); + for (auto& sh : basis) sh.set_shell_tolerance(1e-12); + + const int32_t nbf = basis.nbf(); + REQUIRE(nbf > 0); + + const int64_t npts = 137; // not a multiple of any batch size + const auto pts = make_random_points(npts); + + // Reference: AO collocation directly via the LocalHostWorkDriver. + const auto ao_ref = reference_collocation(basis, npts, pts.data()); + + auto eval = make_evaluator(basis, 1e-12); + REQUIRE(eval.nbf() == nbf); + + SECTION("eval_orbital with one-hot coefficient reproduces single AO column") { + std::vector C(nbf, 0.0); + std::vector out(static_cast(npts), 0.0); + for (int32_t mu : {0, nbf / 3, nbf / 2, nbf - 1}) { + std::fill(C.begin(), C.end(), 0.0); + C[static_cast(mu)] = 1.0; + std::fill(out.begin(), out.end(), 0.0); + eval.eval_orbital(npts, pts.data(), C.data(), out.data()); + for (int64_t p = 0; p < npts; ++p) { + const double ref = + ao_ref[static_cast(p) * nbf + static_cast(mu)]; + CHECK(out[static_cast(p)] == Approx(ref).margin(1e-12)); + } + } + } + + SECTION("eval_orbitals with random C matches AO ^T @ C") { + const int32_t nmo = 4; + const auto C = make_random_vector(static_cast(nbf) * nmo, 99u); + + std::vector out(static_cast(npts) * nmo, 0.0); + eval.eval_orbitals(npts, pts.data(), nmo, C.data(), nbf, out.data(), npts); + + for (int32_t j = 0; j < nmo; ++j) { + for (int64_t p = 0; p < npts; ++p) { + double ref = 0.0; + for (int32_t mu = 0; mu < nbf; ++mu) { + ref += C[static_cast(j) * nbf + mu] * + ao_ref[static_cast(p) * nbf + mu]; + } + const double got = + out[static_cast(j) * npts + static_cast(p)]; + CHECK(got == Approx(ref).margin(1e-10)); + } + } + } + + SECTION("eval_density with identity D equals sum of squared AO values") { + std::vector D(static_cast(nbf) * nbf, 0.0); + for (int32_t mu = 0; mu < nbf; ++mu) { + D[static_cast(mu) * nbf + mu] = 1.0; + } + std::vector out(static_cast(npts), 0.0); + eval.eval_density(npts, pts.data(), D.data(), nbf, out.data()); + + for (int64_t p = 0; p < npts; ++p) { + double ref = 0.0; + for (int32_t mu = 0; mu < nbf; ++mu) { + const double a = ao_ref[static_cast(p) * nbf + mu]; + ref += a * a; + } + CHECK(out[static_cast(p)] == Approx(ref).margin(1e-10)); + } + } + + SECTION("eval_density with rank-1 D = c c^T equals (c.AO)^2") { + const auto c = make_random_vector(static_cast(nbf), 7u); + + std::vector D(static_cast(nbf) * nbf, 0.0); + for (int32_t mu = 0; mu < nbf; ++mu) { + for (int32_t nu = 0; nu < nbf; ++nu) { + D[static_cast(mu) * nbf + nu] = c[mu] * c[nu]; + } + } + std::vector out(static_cast(npts), 0.0); + eval.eval_density(npts, pts.data(), D.data(), nbf, out.data()); + + std::vector orb(static_cast(npts), 0.0); + eval.eval_orbital(npts, pts.data(), c.data(), orb.data()); + + for (int64_t p = 0; p < npts; ++p) { + const double ref = orb[static_cast(p)] * orb[static_cast(p)]; + CHECK(out[static_cast(p)] == Approx(ref).margin(1e-10)); + } + } + + SECTION("invalid leading dimensions throw") { + std::vector C(nbf, 0.0); + std::vector out(static_cast(npts), 0.0); + CHECK_THROWS(eval.eval_orbitals(npts, pts.data(), 1, C.data(), nbf - 1, + out.data(), npts)); + CHECK_THROWS(eval.eval_orbitals(npts, pts.data(), 1, C.data(), nbf, + out.data(), npts - 1)); + // ldo is narrowed to int for the BLAS call; oversized values must throw + // rather than silently truncate. + const size_t huge_ldo = + static_cast(std::numeric_limits::max()) + 1; + CHECK_THROWS(eval.eval_orbitals(npts, pts.data(), 1, C.data(), nbf, + out.data(), huge_ldo)); + } +} + +TEST_CASE("OrbitalEvaluator public API surface", "[orbital_evaluator]") { + auto mol = make_water(); + constexpr double shell_tol = 1e-10; + auto basis = make_ccpvdz(mol, SphericalType(true)); + for (auto& sh : basis) sh.set_shell_tolerance(shell_tol); + const int32_t nbf = basis.nbf(); + + SECTION("basis() exposes the evaluator's own retuned copy") { + auto tight = make_ccpvdz(mol, SphericalType(true)); + for (auto& sh : tight) sh.set_shell_tolerance(1e-12); + auto eval = make_evaluator(tight, 1e-4); + + REQUIRE(eval.basis().nbf() == nbf); + REQUIRE(eval.basis().size() == tight.size()); + + // The caller's basis keeps the tolerance it was given, so the evaluator's + // looser one must show up as shorter cutoff radii on its copy alone. + bool any_shorter = false; + for (size_t i = 0; i < tight.size(); ++i) { + CHECK(eval.basis()[i].cutoff_radius() <= tight[i].cutoff_radius()); + if (eval.basis()[i].cutoff_radius() < tight[i].cutoff_radius()) + any_shorter = true; + } + CHECK(any_shorter); + } + + SECTION("move construction and assignment carry the implementation") { + const int64_t npts = 64; + const auto pts = make_random_points(npts, 7u); + const auto C = make_random_vector(static_cast(nbf), 11u); + + auto eval = make_evaluator(basis, shell_tol); + std::vector ref(static_cast(npts)); + eval.eval_orbital(npts, pts.data(), C.data(), ref.data()); + + OrbitalEvaluator moved(std::move(eval)); + REQUIRE(moved.nbf() == nbf); + std::vector out(static_cast(npts)); + moved.eval_orbital(npts, pts.data(), C.data(), out.data()); + for (int64_t p = 0; p < npts; ++p) CHECK(out[p] == ref[p]); + + // Assigned over an evaluator built with a looser tolerance, so the result + // would shift if the target's original implementation survived. + auto target = make_evaluator(basis, 1e-3); + target = std::move(moved); + REQUIRE(target.nbf() == nbf); + std::vector out2(static_cast(npts)); + target.eval_orbital(npts, pts.data(), C.data(), out2.data()); + for (int64_t p = 0; p < npts; ++p) CHECK(out2[p] == ref[p]); + } + + SECTION("a non-Host execution space is rejected") { + CHECK_THROWS(OrbitalEvaluatorFactory::make_orbital_evaluator( + ExecutionSpace::Device, basis)); + } + + SECTION("null pointers and undersized leading dimensions throw") { + auto eval = make_evaluator(basis, shell_tol); + const int64_t npts = 8; + const auto pts = make_random_points(npts, 3u); + std::vector C(static_cast(nbf), 1.0); + std::vector out(static_cast(npts), 0.0); + std::vector D(static_cast(nbf) * nbf, 0.0); + + CHECK_THROWS(eval.eval_orbital(npts, nullptr, C.data(), out.data())); + CHECK_THROWS(eval.eval_orbital(npts, pts.data(), nullptr, out.data())); + CHECK_THROWS(eval.eval_orbital(npts, pts.data(), C.data(), nullptr)); + CHECK_THROWS(eval.eval_density(npts, nullptr, D.data(), nbf, out.data())); + CHECK_THROWS(eval.eval_density(npts, pts.data(), nullptr, nbf, out.data())); + CHECK_THROWS(eval.eval_density(npts, pts.data(), D.data(), nbf, nullptr)); + CHECK_THROWS( + eval.eval_density(npts, pts.data(), D.data(), nbf - 1, out.data())); + + const auto grid = CubeGrid::from_molecule(mol, 3, 3, 3); + std::vector gout(static_cast(grid.num_points()), 0.0); + CHECK_THROWS(eval.eval_orbital(grid, nullptr, gout.data())); + CHECK_THROWS(eval.eval_density(grid, D.data(), nbf - 1, gout.data())); + } + + SECTION("empty work is a no-op rather than an error") { + auto eval = make_evaluator(basis, shell_tol); + std::vector C(static_cast(nbf), 1.0); + std::vector D(static_cast(nbf) * nbf, 0.0); + constexpr double untouched = -7.0; + std::vector out(4, untouched); + const auto pts = make_random_points(4, 5u); + + CHECK_NOTHROW( + eval.eval_orbitals(0, nullptr, 1, C.data(), nbf, out.data(), 0)); + CHECK_NOTHROW(eval.eval_density(0, nullptr, D.data(), nbf, out.data())); + CHECK_NOTHROW( + eval.eval_orbitals(4, pts.data(), 0, C.data(), nbf, out.data(), 4)); + + CubeGrid empty; + empty.nx = 0; + CHECK_NOTHROW(eval.eval_orbital(empty, C.data(), out.data())); + CHECK_NOTHROW(eval.eval_density(empty, D.data(), nbf, out.data())); + + for (double v : out) CHECK(v == untouched); + } +} + +TEST_CASE("CubeGrid eval overloads match pointer-based eval", + "[orbital_evaluator]") { + auto mol = make_water(); + auto basis = make_ccpvdz(mol, SphericalType(true)); + for (auto& sh : basis) sh.set_shell_tolerance(1e-12); + const int32_t nbf = basis.nbf(); + auto eval = make_evaluator(basis, 1e-12); + + // 6x7x8 fits in a single batch; 20x20x20 = 8000 points spreads over many + // batches at any thread count, exercising the trailing partial batch and + // the dynamic schedule. + const std::vector grids = {CubeGrid::from_molecule(mol, 6, 7, 8), + CubeGrid::from_molecule(mol, 20, 20, 20)}; + + const auto C = make_random_vector(static_cast(nbf), 42u); + + SECTION("eval_orbital(grid) matches eval_orbital(npts, points)") { + for (const auto& grid : grids) { + const int64_t npts = grid.num_points(); + const auto pts = grid.points(); + + std::vector ref(static_cast(npts)); + eval.eval_orbital(npts, pts.data(), C.data(), ref.data()); + + std::vector out(static_cast(npts)); + eval.eval_orbital(grid, C.data(), out.data()); + + for (int64_t p = 0; p < npts; ++p) { + CHECK(out[p] == Approx(ref[p]).margin(1e-12)); + } + } + } + + SECTION("eval_density(grid) matches eval_density(npts, points)") { + // Identity density. + std::vector D(static_cast(nbf) * nbf, 0.0); + for (int32_t i = 0; i < nbf; ++i) D[i * nbf + i] = 1.0; + + for (const auto& grid : grids) { + const int64_t npts = grid.num_points(); + const auto pts = grid.points(); + + std::vector ref(static_cast(npts)); + eval.eval_density(npts, pts.data(), D.data(), nbf, ref.data()); + + std::vector out(static_cast(npts)); + eval.eval_density(grid, D.data(), nbf, out.data()); + + for (int64_t p = 0; p < npts; ++p) { + CHECK(out[p] == Approx(ref[p]).margin(1e-12)); + } + } + } +} + +TEST_CASE("OrbitalEvaluator shell screening", "[orbital_evaluator]") { + // Two well-separated centres of different elements: batches near one centre + // screen out every shell of the other, so 0 < nbe < nbf is reached, and + // batches in the gap screen out everything (nbe == 0). + Molecule mol; + mol.emplace_back(AtomicNumber(8), 0.0, 0.0, 0.0); + mol.emplace_back(AtomicNumber(1), 20.0, 0.0, 0.0); + + auto basis = make_ccpvdz(mol, SphericalType(true)); + constexpr double shell_tol = 1e-10; + for (auto& sh : basis) sh.set_shell_tolerance(shell_tol); + const int32_t nbf = basis.nbf(); + auto eval = make_evaluator(basis, shell_tol); + + const auto C = make_random_vector(static_cast(nbf), 2024u); + std::vector D(static_cast(nbf) * nbf, 0.0); + for (int32_t i = 0; i < nbf; ++i) D[i * nbf + i] = 1.0; + + SECTION("a grid far from every atom evaluates to exactly zero") { + CubeGrid far_grid; + far_grid.origin = {100.0, 100.0, 100.0}; + far_grid.spacing = {0.5, 0.5, 0.5}; + far_grid.nx = far_grid.ny = far_grid.nz = 12; + const int64_t npts = far_grid.num_points(); + + std::vector orb(static_cast(npts), 1.0); + eval.eval_orbital(far_grid, C.data(), orb.data()); + for (int64_t p = 0; p < npts; ++p) CHECK(orb[p] == 0.0); + + std::vector rho(static_cast(npts), 1.0); + eval.eval_density(far_grid, D.data(), nbf, rho.data()); + for (int64_t p = 0; p < npts; ++p) CHECK(rho[p] == 0.0); + } + + SECTION("partial screening agrees with the unscreened reference") { + // Grid elongated along x (the slow axis), so a batch of consecutive + // points is a narrow x-slab: near one centre, near the other, or in the + // empty gap between them. + CubeGrid grid; + grid.origin = {-4.0, -4.0, -4.0}; + grid.spacing = {0.8, 2.0, 2.0}; + grid.nx = 40; + grid.ny = 4; + grid.nz = 4; + const int64_t npts = grid.num_points(); + const auto pts = grid.points(); + const auto ao_ref = reference_collocation(basis, npts, pts.data()); + + std::vector orb(static_cast(npts)); + eval.eval_orbital(grid, C.data(), orb.data()); + + std::vector rho(static_cast(npts)); + eval.eval_density(grid, D.data(), nbf, rho.data()); + + // The grid overload walks spatial tiles while the pointer overload takes + // contiguous index ranges, so the two screen against different bounding + // boxes. They agree to within the shell tolerance, not bitwise. + std::vector orb_pts(static_cast(npts)); + eval.eval_orbital(npts, pts.data(), C.data(), orb_pts.data()); + std::vector rho_pts(static_cast(npts)); + eval.eval_density(npts, pts.data(), D.data(), nbf, rho_pts.data()); + + bool any_nonzero = false; + for (int64_t p = 0; p < npts; ++p) { + CHECK(orb[p] == Approx(orb_pts[p]).margin(shell_tol)); + CHECK(rho[p] == Approx(rho_pts[p]).margin(shell_tol)); + + double orb_ref = 0.0, rho_ref = 0.0; + for (int32_t mu = 0; mu < nbf; ++mu) { + const double a = ao_ref[static_cast(p) * nbf + mu]; + orb_ref += C[static_cast(mu)] * a; + rho_ref += a * a; + } + if (std::fabs(orb_ref) > 1e-3) any_nonzero = true; + // Screening discards shells whose magnitude is below the shell + // tolerance across the batch bbox, so agreement is at that level. + CHECK(orb[p] == Approx(orb_ref).margin(1e-9)); + CHECK(rho[p] == Approx(rho_ref).margin(1e-9)); + } + CHECK(any_nonzero); + } +} + +TEST_CASE("OrbitalEvaluator tiled grid traversal matches the reference", + "[orbital_evaluator]") { + // Grids are walked in spatial tiles rather than contiguous index ranges, so + // a tile's points are scattered into the output. The two centres here are + // far enough apart that screening is active, and the 32 Bohr z extent is + // several times the largest cc-pVDZ cutoff radius (13.2 Bohr), so tiles at + // opposite ends of the grid see different shell sets. + Molecule mol; + mol.emplace_back(AtomicNumber(8), 0.0, 0.0, 0.0); + mol.emplace_back(AtomicNumber(1), 0.0, 0.0, 20.0); + + constexpr double shell_tol = 1e-10; + auto basis = make_ccpvdz(mol, SphericalType(true)); + for (auto& sh : basis) sh.set_shell_tolerance(shell_tol); + const int32_t nbf = basis.nbf(); + auto eval = make_evaluator(basis, shell_tol); + + CubeGrid grid; + grid.origin = {-4.0, -4.0, -6.0}; + grid.spacing = {2.0, 2.0, 0.8}; + grid.nx = 4; + grid.ny = 4; + grid.nz = 40; + const int64_t npts = grid.num_points(); + const auto pts = grid.points(); + const auto ao_ref = reference_collocation(basis, npts, pts.data()); + + const auto C = make_random_vector(static_cast(nbf), 8675u); + std::vector D(static_cast(nbf) * nbf, 0.0); + for (int32_t i = 0; i < nbf; ++i) D[i * nbf + i] = 1.0; + + std::vector orb(static_cast(npts)); + std::vector rho(static_cast(npts)); + eval.eval_orbital(grid, C.data(), orb.data()); + eval.eval_density(grid, D.data(), nbf, rho.data()); + + // The pointer overload always uses contiguous batches, so it cross-checks + // the tiled result against a different decomposition. + std::vector orb_pts(static_cast(npts)); + std::vector rho_pts(static_cast(npts)); + eval.eval_orbital(npts, pts.data(), C.data(), orb_pts.data()); + eval.eval_density(npts, pts.data(), D.data(), nbf, rho_pts.data()); + + bool any_nonzero = false; + for (int64_t p = 0; p < npts; ++p) { + double orb_ref = 0.0, rho_ref = 0.0; + for (int32_t mu = 0; mu < nbf; ++mu) { + const double a = ao_ref[static_cast(p) * nbf + mu]; + orb_ref += C[static_cast(mu)] * a; + rho_ref += a * a; + } + if (std::fabs(orb_ref) > 1e-3) any_nonzero = true; + CHECK(orb[p] == Approx(orb_ref).margin(1e-9)); + CHECK(rho[p] == Approx(rho_ref).margin(1e-9)); + CHECK(orb[p] == Approx(orb_pts[p]).margin(shell_tol)); + CHECK(rho[p] == Approx(rho_pts[p]).margin(shell_tol)); + } + CHECK(any_nonzero); +} + +TEST_CASE("OrbitalEvaluator scatters multiple orbitals into a padded output", + "[orbital_evaluator]") { + // A tiled batch is not contiguous in the output, so its result is staged and + // scattered column by column. nmo > 1 with a padded ldo is the only + // combination that exercises both strides of that scatter. Padding entries + // hold sentinels: C must never be read past nbf, out never written past npts. + Molecule mol; + mol.emplace_back(AtomicNumber(8), 0.0, 0.0, 0.0); + mol.emplace_back(AtomicNumber(1), 0.0, 0.0, 20.0); + + constexpr double shell_tol = 1e-10; + auto basis = make_ccpvdz(mol, SphericalType(true)); + for (auto& sh : basis) sh.set_shell_tolerance(shell_tol); + const int32_t nbf = basis.nbf(); + auto eval = make_evaluator(basis, shell_tol); + + CubeGrid grid; + grid.origin = {-4.0, -4.0, -6.0}; + grid.spacing = {2.0, 2.0, 0.8}; + grid.nx = 4; + grid.ny = 4; + grid.nz = 40; + const int64_t npts = grid.num_points(); + const auto pts = grid.points(); + const auto ao_ref = reference_collocation(basis, npts, pts.data()); + + constexpr int32_t nmo = 3; + constexpr double sentinel = -12345.0; + const size_t ldc = static_cast(nbf) + 7; + const size_t ldo = static_cast(npts) + 5; + + const auto Craw = make_random_vector(static_cast(nbf) * nmo, 99u); + std::vector C(ldc * nmo, sentinel); + for (int32_t j = 0; j < nmo; ++j) + for (int32_t mu = 0; mu < nbf; ++mu) + C[j * ldc + mu] = Craw[static_cast(j) * nbf + mu]; + + std::vector out(ldo * nmo, sentinel); + eval.eval_orbitals(grid, nmo, C.data(), ldc, out.data(), ldo); + + // The pointer overload always batches contiguously, so it checks the tiled + // scatter against a different decomposition of the same grid. + std::vector out_pts(ldo * nmo, sentinel); + eval.eval_orbitals(npts, pts.data(), nmo, C.data(), ldc, out_pts.data(), ldo); + + for (int32_t j = 0; j < nmo; ++j) { + for (int64_t p = 0; p < npts; ++p) { + double ref = 0.0; + for (int32_t mu = 0; mu < nbf; ++mu) + ref += Craw[static_cast(j) * nbf + mu] * + ao_ref[static_cast(p) * nbf + mu]; + const size_t k = static_cast(j) * ldo + static_cast(p); + CHECK(out[k] == Approx(ref).margin(1e-9)); + CHECK(out[k] == Approx(out_pts[k]).margin(shell_tol)); + } + for (size_t p = static_cast(npts); p < ldo; ++p) { + CHECK(out[static_cast(j) * ldo + p] == sentinel); + CHECK(out_pts[static_cast(j) * ldo + p] == sentinel); + } + } +} + +TEST_CASE("OrbitalEvaluator handles single-plane grids", "[orbital_evaluator]") { + // An axis of one point gets zero spacing from from_molecule, which the + // tile-shaping code must treat as unconstrained rather than divide by. + auto mol = make_water(); + constexpr double shell_tol = 1e-10; + auto basis = make_ccpvdz(mol, SphericalType(true)); + for (auto& sh : basis) sh.set_shell_tolerance(shell_tol); + const int32_t nbf = basis.nbf(); + auto eval = make_evaluator(basis, shell_tol); + + const auto C = make_random_vector(static_cast(nbf), 4242u); + std::vector D(static_cast(nbf) * nbf, 0.0); + for (int32_t i = 0; i < nbf; ++i) D[static_cast(i) * nbf + i] = 1.0; + + const std::vector grids = {CubeGrid::from_molecule(mol, 1, 9, 9), + CubeGrid::from_molecule(mol, 9, 1, 9), + CubeGrid::from_molecule(mol, 9, 9, 1)}; + + for (const auto& grid : grids) { + const int64_t npts = grid.num_points(); + const auto pts = grid.points(); + const auto ao = reference_collocation(basis, npts, pts.data()); + + std::vector orb(static_cast(npts)); + std::vector rho(static_cast(npts)); + eval.eval_orbital(grid, C.data(), orb.data()); + eval.eval_density(grid, D.data(), nbf, rho.data()); + + for (int64_t p = 0; p < npts; ++p) { + double orb_ref = 0.0, rho_ref = 0.0; + for (int32_t mu = 0; mu < nbf; ++mu) { + const double a = ao[static_cast(p) * nbf + mu]; + orb_ref += C[static_cast(mu)] * a; + rho_ref += a * a; + } + CHECK(orb[p] == Approx(orb_ref).margin(1e-9)); + CHECK(rho[p] == Approx(rho_ref).margin(1e-9)); + } + } +} + +TEST_CASE("OrbitalEvaluator density with a general D and padded ldd", + "[orbital_evaluator]") { + // Every other density test uses a diagonal D and ldd == nbf, and neither + // can see a whole class of defect. Permuting an identity on both sides + // leaves it unchanged, so a scrambled compressed-submatrix map is invisible + // to it; and ldd == nbf makes a leading dimension dropped in favour of nbf + // a no-op. Screening has to be active for the compressed submatrix to + // differ from D at all, hence the two well-separated centres. + Molecule mol; + mol.emplace_back(AtomicNumber(8), 0.0, 0.0, 0.0); + mol.emplace_back(AtomicNumber(8), 0.0, 0.0, 20.0); + + constexpr double shell_tol = 1e-10; + auto basis = make_ccpvdz(mol, SphericalType(true)); + for (auto& sh : basis) sh.set_shell_tolerance(shell_tol); + const int32_t nbf = basis.nbf(); + auto eval = make_evaluator(basis, shell_tol); + + CubeGrid grid; + grid.origin = {-3.0, -3.0, -4.0}; + grid.spacing = {2.0, 2.0, 1.0}; + grid.nx = 4; + grid.ny = 4; + grid.nz = 28; + const int64_t npts = grid.num_points(); + const auto pts = grid.points(); + const auto ao = reference_collocation(basis, npts, pts.data()); + + constexpr double sentinel = -4321.0; + const size_t ldd = static_cast(nbf) + 6; + + const auto rnd = make_random_vector(static_cast(nbf) * nbf, 31337u); + std::vector Dsq(static_cast(nbf) * nbf); + for (int32_t i = 0; i < nbf; ++i) + for (int32_t j = 0; j < nbf; ++j) + Dsq[static_cast(j) * nbf + i] = + rnd[static_cast(j) * nbf + i] + + rnd[static_cast(i) * nbf + j]; + + // Rows past nbf are poisoned, so an ldd mistaken for nbf reads them. + std::vector D(ldd * nbf, sentinel); + for (int32_t j = 0; j < nbf; ++j) + for (int32_t i = 0; i < nbf; ++i) + D[static_cast(j) * ldd + i] = Dsq[static_cast(j) * nbf + i]; + + std::vector rho_grid(static_cast(npts)); + std::vector rho_pts(static_cast(npts)); + eval.eval_density(grid, D.data(), ldd, rho_grid.data()); + eval.eval_density(npts, pts.data(), D.data(), ldd, rho_pts.data()); + + double max_ref = 0.0; + for (int64_t p = 0; p < npts; ++p) { + const double* a = ao.data() + static_cast(p) * nbf; + double ref = 0.0; + for (int32_t mu = 0; mu < nbf; ++mu) { + double acc = 0.0; + for (int32_t nu = 0; nu < nbf; ++nu) + acc += Dsq[static_cast(nu) * nbf + mu] * a[nu]; + ref += acc * a[mu]; + } + max_ref = std::max(max_ref, std::fabs(ref)); + CHECK(rho_grid[p] == Approx(ref).margin(1e-8)); + CHECK(rho_pts[p] == Approx(ref).margin(1e-8)); + } + CHECK(max_ref > 1e-2); +} + +TEST_CASE("OrbitalEvaluator boxes a batch out to its far face", + "[orbital_evaluator]") { + // A batch is screened against the box drawn round its own points. Building + // that box one grid step short on any axis would drop a shell sitting on the + // far face, and the points there would lose it entirely. + // + // One step has to be decisive for this to be visible at all. A shell only + // dropped by a one-step shrink lies between cutoff-step and cutoff of the + // box, where by construction it contributes about the shell tolerance, so a + // fine grid hides the error inside the tolerance it is measured against. + // The spacing is therefore set to 1.5x the largest cutoff radius, measured + // from the basis rather than assumed, which puts a shell on a tile face + // either fully in or fully out. + // + // Which grid point lands on a tile face depends on how the grid is + // subdivided, which is not observable from here and moves with the batch + // size, so the atom is swept over every grid point instead. Six points an + // axis puts it on a low face, a far face and an interior point of every + // axis under any tiling the batch size can produce. + constexpr double shell_tol = 1e-10; + + double radius = 0.0; + { + Molecule probe; + probe.emplace_back(AtomicNumber(8), 0.0, 0.0, 0.0); + auto pb = make_ccpvdz(probe, SphericalType(true)); + for (auto& sh : pb) sh.set_shell_tolerance(shell_tol); + for (const auto& sh : pb) radius = std::max(radius, sh.cutoff_radius()); + } + REQUIRE(radius > 1.0); + const double step = 1.5 * radius; + + CubeGrid grid; + grid.origin = {0.0, 0.0, 0.0}; + grid.spacing = {step, step, step}; + grid.nx = 6; + grid.ny = 6; + grid.nz = 6; + const int64_t npts = grid.num_points(); + const auto pts = grid.points(); + + double worst_orb = 0.0, worst_rho = 0.0, max_rho = 0.0; + for (int64_t site = 0; site < npts; ++site) { + Molecule mol; + mol.emplace_back(AtomicNumber(8), pts[3 * site + 0], pts[3 * site + 1], + pts[3 * site + 2]); + + auto basis = make_ccpvdz(mol, SphericalType(true)); + for (auto& sh : basis) sh.set_shell_tolerance(shell_tol); + const int32_t nbf = basis.nbf(); + auto eval = make_evaluator(basis, shell_tol); + const auto ao = reference_collocation(basis, npts, pts.data()); + + const auto C = make_random_vector(static_cast(nbf), 606u); + std::vector D(static_cast(nbf) * nbf, 0.0); + for (int32_t i = 0; i < nbf; ++i) D[static_cast(i) * nbf + i] = 1.0; + + std::vector orb(static_cast(npts)); + std::vector rho(static_cast(npts)); + eval.eval_orbital(grid, C.data(), orb.data()); + eval.eval_density(grid, D.data(), nbf, rho.data()); + + for (int64_t p = 0; p < npts; ++p) { + const double* a = ao.data() + static_cast(p) * nbf; + double orb_ref = 0.0, rho_ref = 0.0; + for (int32_t mu = 0; mu < nbf; ++mu) { + orb_ref += C[static_cast(mu)] * a[mu]; + rho_ref += a[mu] * a[mu]; + } + max_rho = std::max(max_rho, rho_ref); + worst_orb = std::max(worst_orb, std::fabs(orb[p] - orb_ref)); + worst_rho = std::max(worst_rho, std::fabs(rho[p] - rho_ref)); + } + } + + CHECK(worst_orb < 1e-9); + CHECK(worst_rho < 1e-9); + // The atom has to register on the grid, or nothing is being proved about + // whether its shells survived screening. + CHECK(max_rho > 1e-2); +} + +TEST_CASE("OrbitalEvaluator screens a tile spanning part of a long axis", + "[orbital_evaluator]") { + // A tile covers a sub-range of each axis, so on a grid far longer in z than + // the tile is, the box has to close around that sub-range rather than the + // whole axis. Three things have to line up for that to matter. The axis must + // outlast a tile, so nz is far above any tile extent. The grid must span + // well past a cutoff radius, taken here from the basis rather than + // hard-coded. And the molecule must sit at the far end of z, so that a box + // reaching further than its own points changes what survives screening. + auto mol = make_water(); + + constexpr double shell_tol = 1e-10; + auto basis = make_ccpvdz(mol, SphericalType(true)); + for (auto& sh : basis) sh.set_shell_tolerance(shell_tol); + const int32_t nbf = basis.nbf(); + + std::vector radii; + radii.reserve(basis.size()); + for (const auto& sh : basis) radii.push_back(sh.cutoff_radius()); + const size_t mid = radii.size() / 2; + std::nth_element(radii.begin(), radii.begin() + mid, radii.end()); + const double median_radius = radii[mid]; + + const double z_extent = 4.0 * median_radius; + + CubeGrid grid; + grid.nx = 1; + grid.ny = 3; + grid.nz = 3000; + grid.spacing = {1.0, 1.0, z_extent / static_cast(grid.nz)}; + // Offset in x so no point lands on a nucleus; the row runs from z_extent + // below the molecule up to 1 Bohr short of it. + grid.origin = {0.3, 0.0, -(z_extent + 1.0)}; + + const int64_t npts = grid.num_points(); + const auto pts = grid.points(); + const auto ao_ref = reference_collocation(basis, npts, pts.data()); + + const auto C = make_random_vector(static_cast(nbf), 4099u); + std::vector D(static_cast(nbf) * nbf, 0.0); + for (int32_t i = 0; i < nbf; ++i) D[i * nbf + i] = 1.0; + + std::vector orb(static_cast(npts)); + std::vector rho(static_cast(npts)); + std::vector orb_pts(static_cast(npts)); + std::vector rho_pts(static_cast(npts)); + + auto eval = make_evaluator(basis, shell_tol); + eval.eval_orbital(grid, C.data(), orb.data()); + eval.eval_density(grid, D.data(), nbf, rho.data()); + // The pointer overload boxes the points it is handed instead of deriving a + // box from tile indices, so it checks the index arithmetic independently. + eval.eval_orbital(npts, pts.data(), C.data(), orb_pts.data()); + eval.eval_density(npts, pts.data(), D.data(), nbf, rho_pts.data()); + + for (int64_t p = 0; p < npts; ++p) { + double orb_ref = 0.0, rho_ref = 0.0; + for (int32_t mu = 0; mu < nbf; ++mu) { + const double a = ao_ref[static_cast(p) * nbf + mu]; + orb_ref += C[static_cast(mu)] * a; + rho_ref += a * a; + } + CHECK(orb[p] == Approx(orb_ref).margin(1e-9)); + CHECK(rho[p] == Approx(rho_ref).margin(1e-9)); + CHECK(orb[p] == Approx(orb_pts[p]).margin(shell_tol)); + CHECK(rho[p] == Approx(rho_pts[p]).margin(shell_tol)); + } + + // Pins the span the tile box is being asked to resolve: the row ends on + // the molecule and starts far enough away for shells to have died off. + const size_t row_end = static_cast(grid.nz) - 1; + CHECK(rho[row_end] > 1e-2); + CHECK(rho[0] < 1e-4 * rho[row_end]); +} + +TEST_CASE("OrbitalEvaluator survives a thread-count change after construction", + "[orbital_evaluator]") { + // Regression guard: scratch must follow the thread count in force at + // evaluation time, not at construction. Nothing about the batch + // decomposition depends on the thread count, so the two runs screen + // identically and have to agree bitwise. + auto mol = make_water(); + auto basis = make_ccpvdz(mol, SphericalType(true)); + constexpr double shell_tol = 1e-12; + for (auto& sh : basis) sh.set_shell_tolerance(shell_tol); + const int32_t nbf = basis.nbf(); + + auto grid = CubeGrid::from_molecule(mol, 16, 16, 16); + const int64_t npts = grid.num_points(); + + const auto C = make_random_vector(static_cast(nbf), 5150u); + std::vector D(static_cast(nbf) * nbf, 0.0); + for (int32_t i = 0; i < nbf; ++i) D[i * nbf + i] = 1.0; + +#ifdef _OPENMP + const int saved_threads = omp_get_max_threads(); + // Construct while the thread count is 1, then raise it. Per-call scratch + // must follow the thread count in force at evaluation time. + omp_set_num_threads(1); +#endif + auto eval = make_evaluator(basis, shell_tol); + + std::vector orb_serial(static_cast(npts)); + std::vector rho_serial(static_cast(npts)); + eval.eval_orbital(grid, C.data(), orb_serial.data()); + eval.eval_density(grid, D.data(), nbf, rho_serial.data()); + +#ifdef _OPENMP + omp_set_num_threads(saved_threads > 1 ? saved_threads : 2); +#endif + + std::vector orb_par(static_cast(npts)); + std::vector rho_par(static_cast(npts)); + eval.eval_orbital(grid, C.data(), orb_par.data()); + eval.eval_density(grid, D.data(), nbf, rho_par.data()); + +#ifdef _OPENMP + omp_set_num_threads(saved_threads); +#endif + + for (int64_t p = 0; p < npts; ++p) { + CHECK(orb_par[p] == orb_serial[p]); + CHECK(rho_par[p] == rho_serial[p]); + } +} + +TEST_CASE("OrbitalEvaluator is safe to invoke concurrently through const", + "[orbital_evaluator]") { + // The evaluator holds no mutable state, so a const invocation has to be safe + // both from several threads at once and from inside an existing parallel + // region. Neither mode crashed when scratch was shared between them; they + // silently returned wrong numbers, which is why they are pinned here rather + // than left to the thread-count test to catch. + // + // Nothing here changes the thread count, so every run derives the same batch + // shape and screens identically. Agreement is therefore bitwise: any + // difference at all means state leaked between calls. + auto mol = make_water(); + constexpr double shell_tol = 1e-10; + auto basis = make_ccpvdz(mol, SphericalType(true)); + for (auto& sh : basis) sh.set_shell_tolerance(shell_tol); + const int32_t nbf = basis.nbf(); + auto eval = make_evaluator(basis, shell_tol); + const OrbitalEvaluator& ceval = eval; + + auto grid = CubeGrid::from_molecule(mol, 12, 12, 12); + const auto npts = static_cast(grid.num_points()); + + const auto C = make_random_vector(static_cast(nbf), 271828u); + std::vector D(static_cast(nbf) * nbf, 0.0); + for (int32_t i = 0; i < nbf; ++i) D[static_cast(i) * nbf + i] = 1.0; + + auto run = [&](double* orb, double* rho) { + ceval.eval_orbital(grid, C.data(), orb); + ceval.eval_density(grid, D.data(), nbf, rho); + }; + + std::vector orb_ref(npts), rho_ref(npts); + run(orb_ref.data(), rho_ref.data()); + double max_ref = 0.0; + for (size_t p = 0; p < npts; ++p) max_ref = std::max(max_ref, rho_ref[p]); + REQUIRE(max_ref > 1e-2); + + constexpr int nrunners = 3; + std::vector orb(npts * nrunners), rho(npts * nrunners); + + SECTION("concurrent std::threads on the same evaluator") { + std::vector pool; + for (int t = 0; t < nrunners; ++t) + pool.emplace_back([&, t] { + run(orb.data() + t * npts, rho.data() + t * npts); + }); + for (auto& th : pool) th.join(); + } + +#ifdef _OPENMP + SECTION("evaluation from inside an enclosing parallel region") { +#pragma omp parallel for num_threads(nrunners) schedule(static, 1) + for (int t = 0; t < nrunners; ++t) + run(orb.data() + static_cast(t) * npts, + rho.data() + static_cast(t) * npts); + } +#endif + + size_t orb_diff = 0, rho_diff = 0; + for (size_t t = 0; t < nrunners; ++t) + for (size_t p = 0; p < npts; ++p) { + if (orb[t * npts + p] != orb_ref[p]) ++orb_diff; + if (rho[t * npts + p] != rho_ref[p]) ++rho_diff; + } + CHECK(orb_diff == 0); + CHECK(rho_diff == 0); +} + +TEST_CASE("OrbitalEvaluator results do not depend on the thread count", + "[orbital_evaluator]") { + // The batch decomposition is derived from the point count and the basis + // alone, so every thread count screens against the same set of bounding + // boxes and has to return the same bits. Deriving it from the thread count + // instead is the easy mistake, and two things have to hold for that mistake + // to be visible. Screening must be active, hence two centres 20 Bohr apart + // with plenty of empty grid around them. And the batch size must be free to + // move: below kMinBatches * kMinBatch points the minimum-batch-count bound + // rounds away and the batch floor pins the batch whatever the thread count, + // so the grid is deliberately larger than that. + Molecule mol; + mol.emplace_back(AtomicNumber(8), 0.0, 0.0, 0.0); + mol.emplace_back(AtomicNumber(1), 20.0, 0.0, 0.0); + auto basis = make_ccpvdz(mol, SphericalType(true)); + constexpr double shell_tol = 1e-10; + for (auto& sh : basis) sh.set_shell_tolerance(shell_tol); + const int32_t nbf = basis.nbf(); + auto eval = make_evaluator(basis, shell_tol); + + CubeGrid grid; + grid.origin = {-4.0, -4.0, -4.0}; + grid.spacing = {0.2, 0.5, 0.5}; + grid.nx = 160; + grid.ny = 32; + grid.nz = 32; + const auto npts = static_cast(grid.num_points()); + REQUIRE(npts > 65536); + + const auto C = make_random_vector(static_cast(nbf), 1618u); + std::vector D(static_cast(nbf) * nbf, 0.0); + for (int32_t i = 0; i < nbf; ++i) D[i * nbf + i] = 1.0; + + std::vector orb_ref(npts), rho_ref(npts); + std::vector orb(npts), rho(npts); + +#ifdef _OPENMP + const int saved_threads = omp_get_max_threads(); + omp_set_num_threads(1); +#endif + eval.eval_orbital(grid, C.data(), orb_ref.data()); + eval.eval_density(grid, D.data(), nbf, rho_ref.data()); + + // Screening has to be doing something, or the invariance is vacuous. + double max_rho = 0.0; + for (size_t p = 0; p < npts; ++p) max_rho = std::max(max_rho, rho_ref[p]); + REQUIRE(max_rho > 1e-2); + + size_t orb_diff = 0, rho_diff = 0; + for (int nthreads : {2, 3, 5, 8, 13, 16, 32}) { +#ifdef _OPENMP + omp_set_num_threads(nthreads); +#else + (void)nthreads; +#endif + eval.eval_orbital(grid, C.data(), orb.data()); + eval.eval_density(grid, D.data(), nbf, rho.data()); + for (size_t p = 0; p < npts; ++p) { + if (orb[p] != orb_ref[p]) ++orb_diff; + if (rho[p] != rho_ref[p]) ++rho_diff; + } + } +#ifdef _OPENMP + omp_set_num_threads(saved_threads); +#endif + + CHECK(orb_diff == 0); + CHECK(rho_diff == 0); +} + +TEST_CASE("OrbitalEvaluator screening error scales with the shell tolerance", + "[orbital_evaluator]") { + // Pins the guidance on the class: orbital error is bounded by the shell + // tolerance, density error by its square, since dropping a shell with + // |phi| < t perturbs sum_uv D phi phi by ~t^2. + Molecule mol; + mol.emplace_back(AtomicNumber(8), 0.0, 0.0, 0.0); + mol.emplace_back(AtomicNumber(1), 20.0, 0.0, 0.0); + + constexpr double loose_tol = 1e-6; + auto basis = make_ccpvdz(mol, SphericalType(true)); + const int32_t nbf = basis.nbf(); + + std::vector radii_before; + for (const auto& sh : basis) radii_before.push_back(sh.cutoff_radius()); + + auto eval_tight = make_evaluator(basis, 1e-14); + auto eval_loose = make_evaluator(basis, loose_tol); + + // Retuning happens on each evaluator's own copy; a basis shared with an SCF + // setup must come back unchanged. + for (size_t i = 0; i < basis.size(); ++i) + CHECK(basis[i].cutoff_radius() == radii_before[i]); + + CubeGrid grid; + grid.origin = {-4.0, -4.0, -4.0}; + grid.spacing = {0.8, 2.0, 2.0}; + grid.nx = 40; + grid.ny = 4; + grid.nz = 4; + const int64_t npts = grid.num_points(); + + const auto C = make_random_vector(static_cast(nbf), 31337u); + std::vector D(static_cast(nbf) * nbf, 0.0); + for (int32_t i = 0; i < nbf; ++i) D[i * nbf + i] = 1.0; + + std::vector orb_tight(npts), orb_loose(npts); + std::vector rho_tight(npts), rho_loose(npts); + eval_tight.eval_orbital(grid, C.data(), orb_tight.data()); + eval_loose.eval_orbital(grid, C.data(), orb_loose.data()); + eval_tight.eval_density(grid, D.data(), nbf, rho_tight.data()); + eval_loose.eval_density(grid, D.data(), nbf, rho_loose.data()); + + double orb_err = 0.0, rho_err = 0.0; + for (int64_t p = 0; p < npts; ++p) { + orb_err = std::max(orb_err, std::fabs(orb_loose[p] - orb_tight[p])); + rho_err = std::max(rho_err, std::fabs(rho_loose[p] - rho_tight[p])); + } + + INFO("orbital error " << orb_err << ", density error " << rho_err); + CHECK(orb_err > 0.0); // screening is genuinely active at the loose tolerance + CHECK(orb_err <= 10.0 * loose_tol); + CHECK(rho_err <= 100.0 * loose_tol * loose_tol); +} + +TEST_CASE("CubeGrid construction and grid-points layout", "[cube]") { + auto mol = make_water(); + CubeGrid g = CubeGrid::from_molecule(mol, /*nx=*/16, /*ny=*/12, /*nz=*/8, + /*margin=*/2.0); + REQUIRE(g.num_points() == 16 * 12 * 8); + + auto pts = g.points(); + REQUIRE(pts.size() == static_cast(g.num_points()) * 3); + + // Check first and last grid points. + CHECK(pts[0] == Approx(g.origin[0])); + CHECK(pts[1] == Approx(g.origin[1])); + CHECK(pts[2] == Approx(g.origin[2])); + + const size_t last = static_cast(g.num_points()) - 1; + CHECK(pts[3 * last + 0] == + Approx(g.origin[0] + g.spacing[0] * (g.nx - 1))); + CHECK(pts[3 * last + 1] == + Approx(g.origin[1] + g.spacing[1] * (g.ny - 1))); + CHECK(pts[3 * last + 2] == + Approx(g.origin[2] + g.spacing[2] * (g.nz - 1))); + + // Check ordering: iz varies fastest. Point (1, 0, 0) should be at offset + // ny*nz, point (0, 1, 0) at nz, point (0, 0, 1) at 1. + const int64_t off_x = g.ny * g.nz; + const int64_t off_y = g.nz; + CHECK(pts[3 * static_cast(off_x) + 0] == + Approx(g.origin[0] + g.spacing[0])); + CHECK(pts[3 * static_cast(off_y) + 1] == + Approx(g.origin[1] + g.spacing[1])); + CHECK(pts[3 * 1 + 2] == Approx(g.origin[2] + g.spacing[2])); +} + +TEST_CASE("CubeGrid rejects degenerate specifications", "[cube]") { + auto mol = make_water(); + CHECK_THROWS(CubeGrid::from_molecule(Molecule{}, 4, 4, 4)); + CHECK_THROWS(CubeGrid::from_molecule(mol, 0, 4, 4)); + CHECK_THROWS(CubeGrid::from_molecule(mol, 4, 0, 4)); + CHECK_THROWS(CubeGrid::from_molecule(mol, 4, 4, 0)); + + const auto grid = CubeGrid::from_molecule(mol, 3, 4, 5); + const auto pts = grid.points(); + REQUIRE(pts.size() == static_cast(grid.num_points()) * 3); + + std::vector buf(pts.size(), -1.0); + grid.points_into(buf.data()); + for (size_t i = 0; i < pts.size(); ++i) CHECK(buf[i] == pts[i]); +} + +TEST_CASE("write_cube rejects bad input and defaults its comment", "[cube]") { +#ifdef GAUXC_HAS_MPI + int world_rank; + MPI_Comm_rank(MPI_COMM_WORLD, &world_rank); + if (world_rank) return; // File I/O; only run on root rank +#endif + auto mol = make_water(); + const auto grid = CubeGrid::from_molecule(mol, 2, 2, 2); + std::vector field(static_cast(grid.num_points()), 1.0); + const auto path = make_temp_path(".cube"); + + CHECK_THROWS(write_cube(path, mol, grid, nullptr)); + + CubeGrid empty; + empty.nx = 0; + CHECK_THROWS(write_cube(path, mol, empty, field.data())); + CHECK_THROWS( + write_cube("/nonexistent-gauxc-dir/out.cube", mol, grid, field.data())); + +#ifdef GAUXC_HAS_HDF5 + const auto h5 = make_temp_path(".h5"); + CHECK_THROWS(write_cube_hdf5(h5, mol, grid, nullptr)); + CHECK_THROWS(write_cube_hdf5(h5, mol, empty, field.data())); +#endif + + // An omitted comment falls back to a fixed first line. + write_cube(path, mol, grid, field.data()); + std::ifstream in(path); + REQUIRE(in.good()); + std::string line; + std::getline(in, line); + CHECK(line == "GauXC cube file"); + std::getline(in, line); + CHECK(line == "Generated by GauXC"); +} + +TEST_CASE("write_cube round-trips header and field data", "[cube]") { +#ifdef GAUXC_HAS_MPI + int world_rank; + MPI_Comm_rank(MPI_COMM_WORLD, &world_rank); + if (world_rank) return; // File I/O; only run on root rank +#endif + auto mol = make_water(); + CubeGrid grid = CubeGrid::from_molecule(mol, /*nx=*/5, /*ny=*/4, /*nz=*/7, + /*margin=*/2.0); + std::vector field(static_cast(grid.num_points())); + for (int64_t i = 0; i < grid.num_points(); ++i) { + // Mix of magnitudes to exercise the formatter (negative, near-zero, etc.) + field[static_cast(i)] = std::sin(0.13 * static_cast(i)) * + std::pow(10.0, (i % 5) - 2); + } + + const std::string path = make_temp_path(".cube"); + write_cube(path, mol, grid, field.data(), "Test cube"); + + // Parse back. + std::ifstream in(path); + REQUIRE(in.is_open()); + std::string l1, l2; + std::getline(in, l1); + std::getline(in, l2); + CHECK(l1 == "Test cube"); + CHECK(l2 == "Generated by GauXC"); + + long long natoms_read; + double ox, oy, oz; + in >> natoms_read >> ox >> oy >> oz; + CHECK(natoms_read == static_cast(mol.size())); + CHECK(ox == Approx(grid.origin[0])); + CHECK(oy == Approx(grid.origin[1])); + CHECK(oz == Approx(grid.origin[2])); + + // 3 axis lines. + for (int axis = 0; axis < 3; ++axis) { + long long n; + double a, b, c; + in >> n >> a >> b >> c; + if (axis == 0) { + CHECK(n == grid.nx); + CHECK(a == Approx(grid.spacing[0])); + CHECK(b == Approx(0.0)); + CHECK(c == Approx(0.0)); + } else if (axis == 1) { + CHECK(n == grid.ny); + CHECK(a == Approx(0.0)); + CHECK(b == Approx(grid.spacing[1])); + CHECK(c == Approx(0.0)); + } else { + CHECK(n == grid.nz); + CHECK(a == Approx(0.0)); + CHECK(b == Approx(0.0)); + CHECK(c == Approx(grid.spacing[2])); + } + } + + // Atoms. + for (size_t i = 0; i < mol.size(); ++i) { + long long Z; + double q, x, y, z; + in >> Z >> q >> x >> y >> z; + CHECK(Z == mol[i].Z.get()); + CHECK(q == Approx(0.0)); + CHECK(x == Approx(mol[i].x)); + CHECK(y == Approx(mol[i].y)); + CHECK(z == Approx(mol[i].z)); + } + + // Field. Read all remaining whitespace-separated doubles and compare. + std::vector field_read; + field_read.reserve(static_cast(grid.num_points())); + double v; + while (in >> v) field_read.push_back(v); + + REQUIRE(field_read.size() == field.size()); + // %13.5E gives 5 significant digits → relative tolerance ~1e-5 for the + // round-trip. + for (size_t i = 0; i < field.size(); ++i) { + CHECK(field_read[i] == Approx(field[i]).epsilon(1e-4).margin(1e-30)); + } + in.close(); + + // Line structure: each (ix,iy) row spans ceil(nz/6) lines, full lines carry + // six 13-char fields and the row's last line carries the remainder. This + // pins the exact per-row byte count that write_cube relies on to place rows + // without a compaction pass. + { + std::ifstream lin(path); + REQUIRE(lin.is_open()); + std::string line; + for (size_t i = 0; i < 6 + mol.size(); ++i) std::getline(lin, line); + + const int64_t lines_per_row = (grid.nz + 5) / 6; + for (int64_t row = 0; row < grid.nx * grid.ny; ++row) { + for (int64_t l = 0; l < lines_per_row; ++l) { + REQUIRE(static_cast(std::getline(lin, line))); + const int64_t nvals = (l + 1 == lines_per_row) ? (grid.nz - l * 6) : 6; + CHECK(line.size() == static_cast(13 * nvals)); + } + } + CHECK_FALSE(static_cast(std::getline(lin, line))); + } + + std::remove(path.c_str()); +} + +namespace { + +/// Write `field` as a single-row cube file and return the raw data block. +std::string cube_data_block(const Molecule& mol, + const std::vector& field) { + CubeGrid grid; + grid.origin = {0.0, 0.0, 0.0}; + grid.spacing = {0.1, 0.1, 0.1}; + grid.nx = 1; + grid.ny = 1; + grid.nz = static_cast(field.size()); + + const std::string path = make_temp_path(".cube"); + write_cube(path, mol, grid, field.data(), "fmt"); + + std::ifstream in(path); + REQUIRE(in.is_open()); + // Skip header (2 comments + 1 natoms line + 3 axis lines + natoms atoms). + std::string skip; + for (int i = 0; i < 6; ++i) std::getline(in, skip); + for (size_t i = 0; i < mol.size(); ++i) std::getline(in, skip); + + std::string block((std::istreambuf_iterator(in)), + std::istreambuf_iterator()); + in.close(); + std::remove(path.c_str()); + return block; +} + +} // namespace + +TEST_CASE("write_cube agrees with snprintf %13.5E formatting", "[cube]") { +#ifdef GAUXC_HAS_MPI + int world_rank; + MPI_Comm_rank(MPI_COMM_WORLD, &world_rank); + if (world_rank) return; // File I/O; only run on root rank +#endif + // Spot-check the custom formatter against snprintf for a battery of values + // by writing a tiny cube file and parsing it back. This is a stronger check + // than the round-trip above since we compare the byte stream. + auto mol = make_water(); + + const double qnan = std::numeric_limits::quiet_NaN(); + const double inf = std::numeric_limits::infinity(); + // The last five exercise the snprintf deferral: subnormals (where scaling + // the mantissa would overflow), a 3-digit exponent either side of zero, and + // a mantissa that carries from E+99 up into a 3-digit exponent. + std::vector field = {0.0, + -0.0, + 1.23456e-10, + -9.99995e-1, + 1.0e+99, + -1.0e+99, + 3.14159265358979, + -2.71828, + 1.0, + -1.0, + 1e-300, + 1.234e+05, + qnan, + -qnan, + inf, + -inf, + std::numeric_limits::denorm_min(), + -1e-310, + 1.0e+300, + -1.0e-305, + 9.9999999e+99}; + + const std::string data_block = cube_data_block(mol, field); + + std::ostringstream expected; + for (size_t i = 0; i < field.size(); ++i) { + char buf[32]; + // C99 leaves the sign of a printed NaN implementation-defined: glibc emits + // it, BSD libc does not, so snprintf is not a portable reference for that + // one value. The formatter deliberately emits it, which is pinned here. + if (std::isnan(field[i])) + std::snprintf(buf, sizeof(buf), "%13s", + std::signbit(field[i]) ? "-NAN" : "NAN"); + else + std::snprintf(buf, sizeof(buf), "%13.5E", field[i]); + expected << buf; + if ((i + 1) % 6 == 0 || (i + 1) == field.size()) expected << '\n'; + } + + CHECK(data_block == expected.str()); +} + +TEST_CASE("write_cube spans multiple output chunks", "[cube]") { +#ifdef GAUXC_HAS_MPI + int world_rank; + MPI_Comm_rank(MPI_COMM_WORLD, &world_rank); + if (world_rank) return; // File I/O; only run on root rank +#endif + // Sized to exceed write_cube's internal staging-buffer target so the chunk + // loop runs more than once, with a partial final chunk. + auto mol = make_water(); + CubeGrid grid; + grid.origin = {0.0, 0.0, 0.0}; + grid.spacing = {0.1, 0.1, 0.1}; + grid.nx = 2; + grid.ny = 100; + grid.nz = 4096; + const int64_t npts = grid.num_points(); + + std::vector field(static_cast(npts)); + for (int64_t i = 0; i < npts; ++i) + field[static_cast(i)] = std::sin(1e-3 * static_cast(i)); + + const std::string path = make_temp_path(".cube"); + write_cube(path, mol, grid, field.data(), "chunked"); + + std::ifstream in(path, std::ios::binary); + REQUIRE(in.is_open()); + std::string skip; + for (size_t i = 0; i < 6 + mol.size(); ++i) std::getline(in, skip); + + std::string data_block((std::istreambuf_iterator(in)), + std::istreambuf_iterator()); + in.close(); + + const int64_t bytes_per_row = grid.nz * 13 + (grid.nz + 5) / 6; + const int64_t n_rows = grid.nx * grid.ny; + REQUIRE(data_block.size() == static_cast(bytes_per_row * n_rows)); + + // Full byte comparison catches any misplacement at a chunk seam. + std::string expected; + expected.reserve(data_block.size()); + char buf[32]; + for (int64_t row = 0; row < n_rows; ++row) { + for (int64_t iz = 0; iz < grid.nz; ++iz) { + std::snprintf(buf, sizeof(buf), "%13.5E", + field[static_cast(row * grid.nz + iz)]); + expected += buf; + if ((iz + 1) % 6 == 0 || iz + 1 == grid.nz) expected += '\n'; + } + } + CHECK(data_block == expected); + + std::remove(path.c_str()); +} + +TEST_CASE("write_cube formatter stays within one last digit at a rounding " + "boundary", + "[cube]") { +#ifdef GAUXC_HAS_MPI + int world_rank; + MPI_Comm_rank(MPI_COMM_WORLD, &world_rank); + if (world_rank) return; // File I/O; only run on root rank +#endif + // For values sitting within a few ulp of a rounding boundary in the 6th + // significant digit, scaling the mantissa can tip it across the boundary, + // so the formatter and glibc may pick different last digits. Both stay + // within one unit of it; this pins that bound rather than byte equality. + auto mol = make_water(); + const std::vector field = {123456.5, 1.234575, -123456.5, -1.234575, + 9.9999949999999998642e-98, 0.0}; + + const std::string data_block = cube_data_block(mol, field); + REQUIRE(data_block.size() >= field.size() * 13); + + for (size_t i = 0; i < field.size(); ++i) { + const std::string tok = data_block.substr(i * 13, 13); + const double got = std::stod(tok); + const double v = field[i]; + const double last_digit = + v == 0.0 ? 1.0 + : std::pow(10.0, std::floor(std::log10(std::fabs(v))) - 5.0); + INFO("value " << v << " formatted as '" << tok << "'"); + CHECK(std::fabs(got - v) <= 0.5000001 * last_digit); + } +} + +#ifdef GAUXC_HAS_HDF5 +TEST_CASE("write_cube_hdf5 round-trip", "[cube]") { +#ifdef GAUXC_HAS_MPI + int world_rank; + MPI_Comm_rank(MPI_COMM_WORLD, &world_rank); + if (world_rank) return; // File I/O; only run on root rank +#endif + auto mol = make_water(); + auto grid = CubeGrid::from_molecule(mol, 4, 5, 6); + + // Fill a small field with known values. + const int64_t npts = grid.num_points(); + std::vector field(npts); + for (int64_t i = 0; i < npts; ++i) field[i] = 0.01 * i - 0.5; + + const std::string path = make_temp_path(".h5"); + write_cube_hdf5(path, mol, grid, field.data(), "test cube"); + + // Read back and verify. + HighFive::File file(path, HighFive::File::ReadOnly); + + // Field shape and values. + auto ds = file.getDataSet("field"); + auto dims = ds.getDimensions(); + REQUIRE(dims.size() == 3); + CHECK(dims[0] == static_cast(grid.nx)); + CHECK(dims[1] == static_cast(grid.ny)); + CHECK(dims[2] == static_cast(grid.nz)); + + std::vector read_field(npts); + ds.read(read_field.data()); + for (int64_t i = 0; i < npts; ++i) { + CHECK(read_field[i] == Approx(field[i]).epsilon(1e-14)); + } + + // Comment attribute. + std::string cmt; + ds.getAttribute("comment").read(cmt); + CHECK(cmt == "test cube"); + + // Grid metadata. + auto grp_grid = file.getGroup("grid"); + std::vector origin, spacing; + std::vector shape; + grp_grid.getDataSet("origin").read(origin); + grp_grid.getDataSet("spacing").read(spacing); + grp_grid.getDataSet("shape").read(shape); + REQUIRE(origin.size() == 3); + REQUIRE(spacing.size() == 3); + REQUIRE(shape.size() == 3); + for (int k = 0; k < 3; ++k) { + CHECK(origin[k] == Approx(grid.origin[k]).epsilon(1e-14)); + CHECK(spacing[k] == Approx(grid.spacing[k]).epsilon(1e-14)); + } + CHECK(shape[0] == grid.nx); + CHECK(shape[1] == grid.ny); + CHECK(shape[2] == grid.nz); + + // Atoms. + auto grp_atoms = file.getGroup("atoms"); + std::vector Z; + grp_atoms.getDataSet("Z").read(Z); + REQUIRE(Z.size() == mol.size()); + for (size_t i = 0; i < mol.size(); ++i) { + CHECK(Z[i] == static_cast(mol[i].Z.get())); + } + + std::vector coords(mol.size() * 3); + grp_atoms.getDataSet("coords").read(coords.data()); + for (size_t i = 0; i < mol.size(); ++i) { + CHECK(coords[3 * i + 0] == Approx(mol[i].x).epsilon(1e-14)); + CHECK(coords[3 * i + 1] == Approx(mol[i].y).epsilon(1e-14)); + CHECK(coords[3 * i + 2] == Approx(mol[i].z).epsilon(1e-14)); + } + + std::remove(path.c_str()); +} +#endif diff --git a/tests/ut_common.hpp.in b/tests/ut_common.hpp.in index f899aeef..b18e67ce 100644 --- a/tests/ut_common.hpp.in +++ b/tests/ut_common.hpp.in @@ -18,3 +18,4 @@ #include #cmakedefine GAUXC_REF_DATA_PATH "@GAUXC_REF_DATA_PATH@" +#cmakedefine GAUXC_TEST_TMP_PATH "@GAUXC_TEST_TMP_PATH@"