diff --git a/include/gauxc/cube_grid.hpp b/include/gauxc/cube_grid.hpp new file mode 100644 index 00000000..6b01f112 --- /dev/null +++ b/include/gauxc/cube_grid.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 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). + * + * @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..78d29611 --- /dev/null +++ b/include/gauxc/external/cube.hpp @@ -0,0 +1,76 @@ +/** + * 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 + +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..a510ee85 --- /dev/null +++ b/include/gauxc/orbital_evaluator.hpp @@ -0,0 +1,158 @@ +/** + * 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 size is derived in part from the OpenMP + * thread count, so results are bit-reproducible for a fixed thread count but + * may differ between thread counts by at most the basis-set shell tolerance. + * + * 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). + */ + static OrbitalEvaluator make_orbital_evaluator( ExecutionSpace ex, + BasisSet basis ); + +}; // class OrbitalEvaluatorFactory + +} // namespace GauXC diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index 1aed4b42..0ef34edb 100644 --- a/src/CMakeLists.txt +++ b/src/CMakeLists.txt @@ -41,6 +41,8 @@ add_library( gauxc molgrid_impl.cxx molgrid_defaults.cxx atomic_radii.cxx + cube_grid.cxx + orbital_evaluator.cxx ) target_include_directories( gauxc 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 46612c81..9df2df95 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) @@ -32,7 +35,7 @@ if( GAUXC_ENABLE_HDF5 ) FetchContent_MakeAvailable( HighFive ) endif() - target_sources( gauxc PRIVATE hdf5_write.cxx hdf5_read.cxx ) + target_sources( gauxc PRIVATE hdf5_write.cxx hdf5_read.cxx cube_hdf5.cxx ) target_link_libraries( gauxc PUBLIC HighFive ) else() message(WARNING "GAUXC_ENABLE_HDF5 was enabled, but HDF5 was not found, Disabling HDF5 Bindings") diff --git a/src/external/cube.cxx b/src/external/cube.cxx new file mode 100644 index 00000000..626169ee --- /dev/null +++ b/src/external/cube.cxx @@ -0,0 +1,214 @@ +/** + * 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 reproduces glibc + * "%13.5E" for every value except exact half-way ties in the 6th significant + * digit, where glibc rounds the exact binary value half-to-even while this + * routine rounds half-away-from-zero. Both agree to within one unit of the + * last printed digit. 3-digit exponents 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; + } + + // Exponent via floor(log10), with corrections for FP edge cases + // (e.g. 9.99999 rounding up across a power-of-ten boundary). + int exp10 = static_cast( std::floor( std::log10(absv) ) ); + double scale = std::pow( 10.0, -exp10 ); + double mant = absv * scale; + + long long mant_int = std::llround( mant * 1e5 ); + if( mant_int >= 1000000 ) { + mant_int = 100000; + ++exp10; + } else if( mant_int < 100000 ) { + --exp10; + scale = std::pow( 10.0, -exp10 ); + mant = absv * scale; + mant_int = std::llround( mant * 1e5 ); + if( mant_int >= 1000000 ) mant_int = 999999; + else if( mant_int < 100000 ) mant_int = 100000; + } + + // 3+ digit exponents have a different field layout; defer to snprintf. + if( exp10 > 99 or exp10 < -99 ) { + char tmp[32]; + const int n = std::snprintf( tmp, sizeof(tmp), "%13.5E", v ); + if( n >= 13 ) { + std::memcpy( out, tmp + (n - 13), 13 ); + } else { + const int pad = 13 - n; + for( int i = 0; i < pad; ++i ) out[i] = ' '; + std::memcpy( out + pad, tmp, static_cast(n) ); + } + 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..b3e68cc3 --- /dev/null +++ b/src/orbital_evaluator.cxx @@ -0,0 +1,538 @@ +/** + * 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 + +#ifdef _OPENMP +#include +#else +inline int omp_get_max_threads() { return 1; } +#endif + +#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 ) : + basis( std::move(bs) ) { + + 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. + * + * Two constraints, whichever is tighter: + * 1. Cache footprint: the AO scratch (nbf*batch doubles) should stay + * around 1 MiB per thread so it lives comfortably in L2/L3. + * 2. Load balance: enough batches that each thread gets several, i.e. + * batch <= npts / (4*nthreads). + */ +size_t choose_batch_size( int32_t nbf, size_t npts, int nthreads ) { + constexpr size_t kTargetAOBytesPerThread = 1024 * 1024; + constexpr size_t kMinBatch = 128; + constexpr size_t kMaxBatch = 8192; + + size_t batch = kMaxBatch; + if( nbf > 0 ) { + batch = kTargetAOBytesPerThread / ( sizeof(double) * static_cast(nbf) ); + } + if( nthreads > 0 ) { + batch = std::min( batch, npts / ( 4 * static_cast(nthreads) ) ); + } + 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; +} + +/** @brief Bbox of the contiguous CubeGrid index range [p0, p0+np). + * + * Exact (not merely conservative): with iz varying fastest, a range that + * crosses an ix boundary necessarily contains both iy=0 and iy=ny-1, and a + * range that crosses an iy boundary necessarily contains both iz=0 and + * iz=nz-1. Computing this analytically avoids rescanning the 3*np + * coordinates that were just generated. + */ +PointBbox bbox_for_range( const CubeGrid& g, size_t p0, size_t np ) { + const int64_t nz = g.nz, ny = g.ny; + const int64_t k0 = static_cast(p0); + const int64_t k1 = static_cast(p0 + np) - 1; + + int64_t lo_idx[3], hi_idx[3]; + lo_idx[0] = k0 / (ny * nz); + hi_idx[0] = k1 / (ny * nz); + if( hi_idx[0] > lo_idx[0] ) { + lo_idx[1] = 0; hi_idx[1] = ny - 1; + lo_idx[2] = 0; hi_idx[2] = nz - 1; + } else { + lo_idx[1] = (k0 / nz) % ny; + hi_idx[1] = (k1 / nz) % ny; + if( hi_idx[1] > lo_idx[1] ) { + lo_idx[2] = 0; hi_idx[2] = nz - 1; + } else { + lo_idx[2] = k0 % nz; + hi_idx[2] = k1 % nz; + } + } + + PointBbox b; + 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]); + b.lo[k] = std::min(a, c); + b.hi[k] = std::max(a, c); + } + 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; +} + +/// Point source backed by a caller-supplied AoS coordinate array. +struct RawPointSource { + static constexpr bool needs_scratch = false; + const double* points; + + const double* batch( size_t p0, size_t, double* ) const { + return points + 3 * p0; + } + PointBbox bbox( size_t, size_t np, const double* pts ) const { + return compute_bbox( pts, np ); + } +}; + +/// Point source that generates CubeGrid coordinates on the fly, avoiding the +/// 3*num_points()*8 byte temporary coordinate array. +struct GridPointSource { + static constexpr bool needs_scratch = true; + const CubeGrid* grid; + + const double* batch( size_t p0, size_t np, double* scr ) const { + // Incremental (ix,iy,iz) counters; coordinates are recomputed from the + // index (rather than accumulated) to stay bit-identical to + // CubeGrid::points_into. + const CubeGrid& g = *grid; + const int64_t nz = g.nz, ny = g.ny; + int64_t iz = static_cast(p0) % nz; + int64_t iy = (static_cast(p0) / nz) % ny; + int64_t ix = static_cast(p0) / (ny * nz); + double x = g.origin[0] + g.spacing[0] * static_cast(ix); + double y = g.origin[1] + g.spacing[1] * static_cast(iy); + + for( size_t i = 0; i < np; ++i ) { + scr[3 * i + 0] = x; + scr[3 * i + 1] = y; + scr[3 * i + 2] = g.origin[2] + g.spacing[2] * static_cast(iz); + if( ++iz == nz ) { + iz = 0; + if( ++iy == ny ) { + iy = 0; + ++ix; + x = g.origin[0] + g.spacing[0] * static_cast(ix); + } + y = g.origin[1] + g.spacing[1] * static_cast(iy); + } + } + return scr; + } + + PointBbox bbox( size_t p0, size_t np, const double* ) const { + return bbox_for_range( *grid, p0, np ); + } +}; + +/// 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: + + struct Scratch { std::vector C_compressed; }; + + 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 init( Scratch& scr, size_t ) const { + scr.C_compressed.resize( static_cast(impl_.nbf_) * nmo_ ); + } + + void zero( size_t p0, size_t np ) const { + for( int32_t j = 0; j < nmo_; ++j ) { + double* out_col = out_ + static_cast(j) * ldo_ + p0; + std::fill( out_col, out_col + np, 0.0 ); + } + } + + void apply( Scratch& scr, size_t p0, size_t np, int32_t nbe, + const std::vector& shells, const double* ao ) const { + + // Gather the rows of C for the surviving shells into a contiguous + // (nbe, nmo) col-major buffer so the contraction is a dense GEMM. + 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; + } + + blas::gemm( 'T', 'N', static_cast(np), nmo_, nbe, 1.0, + ao, nbe, scr.C_compressed.data(), nbe, 0.0, out_ + p0, + static_cast(ldo_) ); + } + +}; + +/// 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: + + struct Scratch { + std::vector dm_ao; + std::vector xmat_scr; + }; + + DensityContractor( const detail::OrbitalEvaluatorImpl& impl, const double* D, + size_t ldd, double* out ) : + impl_(impl), D_(D), ldd_(ldd), out_(out) {} + + void init( Scratch& scr, size_t batch_size ) const { + scr.dm_ao.resize( static_cast(impl_.nbf_) * batch_size ); + } + + void zero( size_t p0, size_t np ) const { + std::fill( out_ + p0, out_ + p0 + np, 0.0 ); + } + + void apply( Scratch& scr, size_t p0, size_t np, int32_t nbe, + const std::vector& shells, const double* ao ) const { + + // eval_xmat needs nbe*nbe scratch. Growing it 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. + const size_t scr_size = static_cast(nbe) * nbe; + if( scr.xmat_scr.size() < scr_size ) scr.xmat_scr.resize( scr_size ); + + 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 ); + + // 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) + impl_.host_driver->eval_uvvar_lda_rks( np, static_cast(nbe), ao, + scr.dm_ao.data(), static_cast(nbe), out_ + p0 ); + } + +}; + +/** @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, size_t npts, + const PointSource& src, const Contractor& contract ) { + + const BasisSet& basis = impl.basis; + const auto& shell_cutoff_r2 = impl.shell_cutoff_r2; + const int32_t nbf = impl.nbf_; + const int32_t nshells_total = basis.nshells(); + + const size_t batch_size = choose_batch_size( nbf, npts, omp_get_max_threads() ); + const int64_t n_batches = + static_cast( (npts + batch_size - 1) / batch_size ); + +#pragma omp parallel + { + std::vector pt_buf; + if constexpr( PointSource::needs_scratch ) pt_buf.resize( 3 * batch_size ); + std::vector ao_buf( static_cast(nbf) * batch_size ); + std::vector screened_shells; + screened_shells.reserve( nshells_total ); + typename Contractor::Scratch scr; + contract.init( scr, batch_size ); + +#pragma omp for schedule(dynamic, 1) + for( int64_t b = 0; b < n_batches; ++b ) { + + const size_t p0 = static_cast(b) * batch_size; + const size_t np = std::min( batch_size, npts - p0 ); + + const double* pts = src.batch( p0, np, pt_buf.data() ); + + // Per-batch shell screening: keep shells whose cutoff radius reaches + // any point of this batch's bounding box. + const PointBbox bbox = src.bbox( p0, np, pts ); + 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( p0, np ); continue; } + + // 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. This is a pre-existing inconsistency in the + // reference driver rather than one 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, p0, np, nbe, screened_shells, ao_buf.data() ); + + } + } + +} + +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()." ); +} + +} // 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 ) { + + switch(ex) { + + case ExecutionSpace::Host: + return OrbitalEvaluator( + std::make_unique( std::move(basis) ) + ); + + 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_, npts, RawPointSource{points}, + 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_, npts, RawPointSource{points}, + 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_, npts, GridPointSource{&grid}, + 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_, npts, GridPointSource{&grid}, + 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..fab041ea --- /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 ); + +}; // class OrbitalEvaluatorImpl + +} // namespace GauXC::detail diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index 46dbe487..3198d63c 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 ) target_link_libraries( gauxc_test PUBLIC gauxc gauxc_catch2 Eigen3::Eigen ) if(GAUXC_ENABLE_CUTLASS) @@ -78,6 +79,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..65c59a75 --- /dev/null +++ b/tests/orbital_evaluator_test.cxx @@ -0,0 +1,813 @@ +/** + * 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 +#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) { + return OrbitalEvaluatorFactory::make_orbital_evaluator(ExecutionSpace::Host, + basis); +} + +/// 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); + 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("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); + + // 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)); + for (auto& sh : basis) sh.set_shell_tolerance(1e-10); + const int32_t nbf = basis.nbf(); + auto eval = make_evaluator(basis); + + 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 screens against an analytically derived batch bbox + // while the pointer overload scans the coordinates; with screening active + // the two must still agree exactly. + 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] == orb_pts[p]); + CHECK(rho[p] == rho_pts[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; + // 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 is invariant to the OpenMP thread count", + "[orbital_evaluator]") { + // Bit-exact here because the grid encloses the molecule at a 1e-12 shell + // tolerance, so nothing screens and the batch decomposition cannot change + // the arithmetic. The screened case is covered separately below. + 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 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); + + 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 thread-count dependence is bounded by screening", + "[orbital_evaluator]") { + // Batch size is derived from the thread count, so when screening is active + // different thread counts screen against different batch bounding boxes and + // the results are not bit-identical. The discrepancy is bounded by the shell + // tolerance, which is what this pins down. + 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); + + 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(); + + 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 rho_serial(static_cast(npts)); + std::vector rho_par(static_cast(npts)); + +#ifdef _OPENMP + const int saved_threads = omp_get_max_threads(); + omp_set_num_threads(1); +#endif + eval.eval_density(grid, D.data(), nbf, rho_serial.data()); +#ifdef _OPENMP + omp_set_num_threads(saved_threads > 1 ? saved_threads : 2); +#endif + 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(rho_par[p] == Approx(rho_serial[p]).margin(shell_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("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(); + 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}; + + 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]; + 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 rounds half-way ties within one last digit", + "[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 + // The hand-rolled formatter rounds exact half-way ties in the 6th + // significant digit half-away-from-zero, whereas glibc rounds the exact + // binary value half-to-even. Both are within one unit of the last printed + // digit; 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@"