Add OrbitalEvaluator and Gaussian cube writer utility

include/gauxc/external/cube.hpp

@@ -0,0 +1,133 @@
+/**
+ * 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-2026, Microsoft Corporation
+ *
+ * All rights reserved.
+ *
+ * See LICENSE.txt for details
+ */
+#pragma once
+
+#include <array>
+#include <cstdint>
+#include <string>
+#include <vector>
+
+#include <gauxc/molecule.hpp>
+
+namespace GauXC {
+
+/** @brief Specification of a Gaussian-cube-format 3D rectangular grid.
+ *
+ *  The grid is axis-aligned with the Cartesian frame (the standard
+ *  cube-format axis vectors are simply (spacing[0], 0, 0), (0, spacing[1], 0),
+ *  (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;
+};
+
+/** @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,126 @@
+/**
+ * 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-2026, Microsoft Corporation
+ *
+ * All rights reserved.
+ *
+ * See LICENSE.txt for details
+ */
+#pragma once
+
+#include <cstdint>
+#include <memory>
+
+#include <gauxc/basisset.hpp>
+#include <gauxc/enums.hpp>
+#include <gauxc/external/cube.hpp>
+
+namespace GauXC {
+
+/** @brief Evaluate molecular orbitals and densities on arbitrary point sets.
+ *
+ *  Wraps the host 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`.
+ *
+ *  Currently only `ExecutionSpace::Host` is supported; passing any other
+ *  value will throw on construction. Device support can be added later by
+ *  routing through a device-side collocation driver while keeping the same
+ *  signatures.
+ */
+class OrbitalEvaluator {
+ public:
+  /** @brief Construct an evaluator bound to a basis set.
+   *
+   *  @param basis Basis set (copied internally).
+   *  @param exec  Execution space; only `ExecutionSpace::Host` is supported.
+   */
+  explicit OrbitalEvaluator(BasisSet<double> basis,
+                            ExecutionSpace exec = ExecutionSpace::Host);
+
+  ~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 noexcept;
+
+  /// Underlying basis set.
+  const BasisSet<double>& basis() const noexcept;
+
+  /** @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(int64_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(int64_t npts, const double* points,
+                     int32_t nmo, const double* C, int64_t ldc,
+                     double* out, int64_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(int64_t npts, const double* points,
+                    const double* D, int64_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, int64_t ldc,
+                     double* out, int64_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, int64_t ldd,
+                    double* out) const;
+
+ private:
+  struct Impl;
+  std::unique_ptr<Impl> pimpl_;
+};
+
+}  // namespace GauXC

src/CMakeLists.txt

@@ -41,6 +41,7 @@ add_library( gauxc
   molgrid_impl.cxx 
   molgrid_defaults.cxx 
   atomic_radii.cxx 
+  orbital_evaluator.cxx
 )
 
 target_include_directories( gauxc

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")