Add ROHF GDM

cpp/src/qdk/chemistry/algorithms/microsoft/scf/include/qdk/chemistry/scf/core/scf_algorithm.h

@@ -98,10 +98,12 @@ class SCFAlgorithm {
       int num_molecular_orbitals, int idx_spin);
 
   /**
-   * @brief Update the density matrix for restricted or unrestricted
-   * calculations. For ASAHF and ROHF calculations, this method will be
-   * overridden to implement the specific density matrix construction for those
-   * methods.
+   * @brief Update the density matrix for RHF, UHF, or ROHF calculations.
+   *
+   * The base implementation handles standard RHF/UHF construction and also
+   * reconstructs spin-blocked alpha/beta densities for ROHF from a shared
+   * molecular-orbital coefficient matrix. Algorithms with custom density
+   * construction (for example ASAHF) can still override this method.
    *
    * @param[in,out] P Reference to the density matrix to be updated
    * @param[in] C Reference to the molecular orbital coefficients
@@ -113,6 +115,20 @@ class SCFAlgorithm {
                                      bool unrestricted, int nelec_alpha,
                                      int nelec_beta);
 
+  /**
+   * @brief Provide ROHF convergence matrices for OG evaluation
+   *
+   * For ROHF, some algorithms evaluate convergence using an effective Fock
+   * matrix and total density matrix rather than the spin-blocked SCFImpl
+   * matrices. This helper returns references to those matrices.
+   *
+   * @param[in] scf_impl SCF implementation object
+   * @return Pair of references to convergence Fock and density matrices
+   * @throws std::logic_error If ROHF convergence matrices are unavailable
+   */
+  std::pair<const RowMajorMatrix&, const RowMajorMatrix&>
+  try_get_rohf_convergence_matrices(const SCFImpl& scf_impl);
+
   /**
    * @brief Calculate orbital gradient (OG) error for convergence checking
    *
@@ -137,6 +153,45 @@ class SCFAlgorithm {
                                     RowMajorMatrix& error_matrix,
                                     int num_orbital_sets);
 
+  /**
+   * @brief Build ROHF convergence matrices from spin-blocked SCF matrices
+   *
+   * Converts spin-blocked Fock/density matrices into the effective ROHF Fock
+   * and total-density representation used for OG evaluation.
+   *
+   * @param[in] F Spin-blocked Fock matrix in AO basis with alpha and beta
+   * blocks stacked by row
+   * @param[in] C Molecular-orbital coefficient matrix used for AO<->MO
+   * transformations
+   * @param[in] P Spin-blocked density matrix in AO basis with alpha and beta
+   * blocks stacked by row
+   * @param[in] nelec_alpha Number of alpha electrons
+   * @param[in] nelec_beta Number of beta electrons
+   * @param[out] effective_fock Effective ROHF Fock matrix in AO basis
+   * @param[out] total_density Total AO density matrix (P_alpha + P_beta)
+   */
+  static void build_rohf_f_p_matrix(const RowMajorMatrix& F,
+                                    const RowMajorMatrix& C,
+                                    const RowMajorMatrix& P, int nelec_alpha,
+                                    int nelec_beta,
+                                    RowMajorMatrix& effective_fock,
+                                    RowMajorMatrix& total_density);
+
+  /**
+   * @brief Access cached ROHF effective Fock matrix
+   */
+  const RowMajorMatrix& get_rohf_convergence_fock_matrix() const;
+
+  /**
+   * @brief Access cached ROHF total density matrix
+   */
+  const RowMajorMatrix& get_rohf_convergence_density_matrix() const;
+
+  /**
+   * @brief Mutable access to cached ROHF total density matrix
+   */
+  RowMajorMatrix& rohf_convergence_density_matrix();
+
  protected:
   const SCFContext& ctx_;  ///< Reference to SCF context
   double og_error_ = 0.0;  ///< Current orbital gradient error
@@ -149,5 +204,7 @@ class SCFAlgorithm {
   double delta_energy_ =
       std::numeric_limits<double>::infinity();  ///< Energy change
   double density_rms_ = 0.0;                    ///< Last calculated density RMS
+  RowMajorMatrix rohf_effective_fock_;  ///< Cached ROHF effective Fock (AO)
+  RowMajorMatrix rohf_total_density_;   ///< Cached ROHF total density (P_a+P_b)
 };
 }  // namespace qdk::chemistry::scf

cpp/src/qdk/chemistry/algorithms/microsoft/scf/src/scf/ks_impl.cpp

@@ -36,6 +36,11 @@ KSImpl::KSImpl(std::shared_ptr<Molecule> mol, const SCFConfig& cfg,
                std::shared_ptr<BasisSet> raw_basis_set)
     : SCFImpl(mol, cfg, basis_set, raw_basis_set, true) {
   QDK_LOG_TRACE_ENTERING();
+  if (cfg.scf_orbital_type == SCFOrbitalType::RestrictedOpenShell) {
+    throw std::invalid_argument(
+        "ROKS (Restricted Open-Shell Kohn-Sham) is not supported. "
+        "Use ROHF (scf_type=\"restricted\") without DFT functionals.");
+  }
 #ifdef ENABLE_NVTX3
   NVTX3_FUNC_RANGE();
 #endif
@@ -128,12 +133,14 @@ double KSImpl::total_energy_() {
 
 std::pair<double, RowMajorMatrix>
 KSImpl::evaluate_trial_density_energy_and_fock(
-    const RowMajorMatrix& P_matrix, const std::source_location& loc) const {
+    const RowMajorMatrix& P_matrix, double* J_out, double* K_out,
+    const std::source_location& loc) const {
   QDK_LOG_TRACE_ENTERING();
   // Fock matrix from base class does not include XC contributions; XC terms are
   // added below
   auto [total_energy, F_matrix] =
-      SCFImpl::evaluate_trial_density_energy_and_fock(P_matrix, loc);
+      SCFImpl::evaluate_trial_density_energy_and_fock(P_matrix, J_out, K_out,
+                                                      loc);
   // Do not update XC_ here: XC_ is a member variable and must not be modified
   // in this const trial evaluation.
   double scf_xc_energy = 0.0;

cpp/src/qdk/chemistry/algorithms/microsoft/scf/src/scf/ks_impl.h

@@ -86,13 +86,16 @@ class KSImpl : public SCFImpl {
    * must be fully initialized before calling it.
    *
    * @param P_matrix Trial density matrix
+   * @param J_out Optional output buffer for Coulomb matrix J
+   * @param K_out Optional output buffer for exchange matrix K
    * @param loc Source location of the caller (automatically captured)
    * @return std::pair containing:
    *   - first: total energy including XC in Hartree
    *   - second: Fock matrix in AO basis
    */
   std::pair<double, RowMajorMatrix> evaluate_trial_density_energy_and_fock(
-      const RowMajorMatrix& P_matrix,
+      const RowMajorMatrix& P_matrix, double* J_out = nullptr,
+      double* K_out = nullptr,
       const std::source_location& loc =
           std::source_location::current()) const override;
 

cpp/src/qdk/chemistry/algorithms/microsoft/scf/src/scf/scf_impl.cpp

@@ -16,6 +16,7 @@
 #endif
 #include <qdk/chemistry/scf/util/env_helper.h>
 
+#include <cstring>
 #include <filesystem>
 #include <fstream>
 #include <iostream>
@@ -618,10 +619,20 @@ void SCFImpl::iterate_() {
     auto [alpha, beta, omega] = get_hyb_coeff_();
     eri_->build_JK(P_.data(), J_.data(), K_.data(), alpha, beta, omega);
     update_fock_();
-    for (int i = 0; i < num_orbital_spin_blocks_; ++i) {
-      scf_algorithm_->solve_fock_eigenproblem(F_, S_, X_, C_, eigenvalues_, P_,
-                                              nelec_, num_atomic_orbitals_,
-                                              num_molecular_orbitals_, i);
+
+    if (ctx_.cfg->scf_orbital_type == SCFOrbitalType::RestrictedOpenShell) {
+      const auto rohf_convergence_matrices =
+          scf_algorithm_->try_get_rohf_convergence_matrices(*this);
+      const auto& effective_fock = std::get<0>(rohf_convergence_matrices);
+      scf_algorithm_->solve_fock_eigenproblem(
+          effective_fock, S_, X_, C_, eigenvalues_, P_, nelec_,
+          num_atomic_orbitals_, num_molecular_orbitals_, 0);
+    } else {
+      for (int i = 0; i < num_orbital_spin_blocks_; ++i) {
+        scf_algorithm_->solve_fock_eigenproblem(
+            F_, S_, X_, C_, eigenvalues_, P_, nelec_, num_atomic_orbitals_,
+            num_molecular_orbitals_, i);
+      }
     }
     scf_algorithm_->update_density_matrix(
         P_, C_, ctx_.cfg->scf_orbital_type == SCFOrbitalType::Unrestricted,
@@ -1035,9 +1046,18 @@ void SCFImpl::write_gradients_(const std::vector<double>& gradients,
 
 std::pair<double, RowMajorMatrix>
 SCFImpl::evaluate_trial_density_energy_and_fock(
-    const RowMajorMatrix& P_matrix, const std::source_location& loc) const {
+    const RowMajorMatrix& P_matrix, double* J_out, double* K_out,
+    const std::source_location& loc) const {
   QDK_LOG_TRACE_ENTERING();
 
+  // TODO: This function is called within iterate_ which runs only on rank 0,
+  // but build_jk_matrices requires participation from all MPI ranks.
+  if (ctx_.cfg->mpi.world_size > 1) {
+    throw std::runtime_error(
+        "Temporary limitation: evaluate_trial_density_energy_and_fock is not "
+        "supported with MPI world_size > 1.");
+  }
+
   QDK_LOGGER().debug(
       "Computing energy and Fock matrix by trial density matrix (called from "
       "function '{}' at line {})",
@@ -1056,12 +1076,22 @@ SCFImpl::evaluate_trial_density_energy_and_fock(
       num_density_matrices_ * num_atomic_orbitals_, num_atomic_orbitals_);
   eri_->build_JK(P_matrix.data(), J_matrix.data(), K_matrix.data(), alpha, beta,
                  omega);
+
+  if (J_out != nullptr) {
+    std::memcpy(J_out, J_matrix.data(),
+                sizeof(double) * static_cast<size_t>(J_matrix.size()));
+  }
+  if (K_out != nullptr) {
+    std::memcpy(K_out, K_matrix.data(),
+                sizeof(double) * static_cast<size_t>(K_matrix.size()));
+  }
 #ifdef QDK_CHEMISTRY_ENABLE_PCM
   throw std::runtime_error("PCM is not supported in trial density evaluation.");
 #endif
 
   if (ctx_.cfg->mpi.world_rank == 0) {
-    if (ctx_.cfg->scf_orbital_type == SCFOrbitalType::Unrestricted) {
+    if (ctx_.cfg->scf_orbital_type == SCFOrbitalType::Unrestricted ||
+        ctx_.cfg->scf_orbital_type == SCFOrbitalType::RestrictedOpenShell) {
       F_matrix +=
           (J_matrix.block(0, 0, num_atomic_orbitals_, num_atomic_orbitals_) +
            J_matrix.block(num_atomic_orbitals_, 0, num_atomic_orbitals_,
@@ -1097,4 +1127,23 @@ SCFImpl::evaluate_trial_density_energy_and_fock(
   return {total_energy, F_matrix};
 }
 
+void SCFImpl::build_jk_matrices(const RowMajorMatrix& density_matrix,
+                                RowMajorMatrix& J, RowMajorMatrix& K) const {
+  QDK_LOG_TRACE_ENTERING();
+
+  const int expected_rows = num_density_matrices_ * num_atomic_orbitals_;
+  const int expected_cols = num_atomic_orbitals_;
+  VERIFY_INPUT(density_matrix.rows() == expected_rows &&
+                   density_matrix.cols() == expected_cols,
+               "density_matrix shape should be (num_density_matrices_ * "
+               "num_atomic_orbitals_, num_atomic_orbitals_)");
+  VERIFY_INPUT(J.rows() == expected_rows && J.cols() == expected_cols,
+               "J shape should match density_matrix shape");
+  VERIFY_INPUT(K.rows() == expected_rows && K.cols() == expected_cols,
+               "K shape should match density_matrix shape");
+
+  auto [alpha, beta, omega] = get_hyb_coeff_();
+  eri_->build_JK(density_matrix.data(), J.data(), K.data(), alpha, beta, omega);
+}
+
 }  // namespace qdk::chemistry::scf

cpp/src/qdk/chemistry/algorithms/microsoft/scf/src/scf/scf_impl.h

@@ -148,6 +148,29 @@ class SCFImpl {
    */
   const RowMajorMatrix& get_fock_matrix() const { return F_; }
 
+  /**
+   * @brief Get the core Hamiltonian matrix
+   * @return Reference to core Hamiltonian H
+   */
+  const RowMajorMatrix& get_core_hamiltonian() const { return H_; }
+
+  /**
+   * @brief Build Coulomb (J) and exchange (K) matrices for a given density
+   *
+   * Utility JK build for a caller-supplied density matrix.
+   * Currently used by the GDM trial-rotation gradient path
+   * (GDMLineFunctor / ROHF generalized-gradient evaluation).
+   * It is not part of the main SCFImpl::iterate_ update loop.
+   * Keeps this trial-path two-electron build logic in one place.
+   *
+   * @param[in] density_matrix Density matrix
+   * (num_density_matrices x NAO x NAO)
+   * @param[out] J Output Coulomb matrix (same size as density_matrix)
+   * @param[out] K Output exchange matrix (same size as density_matrix)
+   */
+  void build_jk_matrices(const RowMajorMatrix& density_matrix,
+                         RowMajorMatrix& J, RowMajorMatrix& K) const;
+
   /**
    * @brief Get the molecular orbital coefficient matrix
    * @see SCF::get_orbitals_matrix() for API details
@@ -198,14 +221,19 @@ class SCFImpl {
    * must be fully initialized before calling it.
    *
    * @param P_matrix Trial density matrix
+   * @param J_out Optional output buffer for Coulomb matrix J (same shape as
+   *   P_matrix)
+   * @param K_out Optional output buffer for exchange matrix K (same shape as
+   *   P_matrix)
    * @param loc Source location of the caller (automatically captured)
    * @return std::pair containing:
    *   - first: total energy in Hartree
    *   - second: Fock matrix in AO basis
    */
   virtual std::pair<double, RowMajorMatrix>
   evaluate_trial_density_energy_and_fock(
-      const RowMajorMatrix& P_matrix,
+      const RowMajorMatrix& P_matrix, double* J_out = nullptr,
+      double* K_out = nullptr,
       const std::source_location& loc = std::source_location::current()) const;
 
  protected: