Add OrbitalEvaluator and Gaussian cube writer utility

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 <array>
+#include <cstdint>
+#include <vector>
+
+#include <gauxc/molecule.hpp>
+
+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<double, 3> origin{0.0, 0.0, 0.0};   ///< (0,0,0) corner, Bohr
+  std::array<double, 3> 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<double> 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

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 <string>
+
+#include <gauxc/cube_grid.hpp>
+#include <gauxc/molecule.hpp>
+
+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

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 <cstddef>
+#include <cstdint>
+#include <memory>
+
+#include <gauxc/basisset.hpp>
+#include <gauxc/cube_grid.hpp>
+#include <gauxc/enums.hpp>
+
+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_type>;
+  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<double>& 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<double> basis );
+
+}; // class OrbitalEvaluatorFactory
+
+}  // namespace GauXC

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