diff --git a/cpp/src/qdk/chemistry/algorithms/microsoft/scf/include/qdk/chemistry/scf/core/scf_algorithm.h b/cpp/src/qdk/chemistry/algorithms/microsoft/scf/include/qdk/chemistry/scf/core/scf_algorithm.h index 96f7d7884..ec305947e 100644 --- a/cpp/src/qdk/chemistry/algorithms/microsoft/scf/include/qdk/chemistry/scf/core/scf_algorithm.h +++ b/cpp/src/qdk/chemistry/algorithms/microsoft/scf/include/qdk/chemistry/scf/core/scf_algorithm.h @@ -7,11 +7,42 @@ #include #include +#include #include #include +#include namespace qdk::chemistry::scf { +/** + * @brief Compute C = A^T * B * A using BLAS GEMM only + * + * Matrix shapes: + * A: m x n + * B: m x m + * C: n x n + * + * @param[in] m Row count of A and dimension of square B block. + * @param[in] n Column count of A and dimension of square C block. + * + * Assumed leading dimensions for contiguous storage: + * RowMajor: lda = n, ldb = m, ldc = n + * ColMajor: lda = m, ldb = m, ldc = n + * + * Pointers may reference sub-block starts (not necessarily matrix origin). + * The intermediate workspace is provided by caller to avoid repeated + * allocations across consecutive calls. + * + * @param[in] A Pointer to the A matrix block. + * @param[in] B Pointer to the B matrix block. + * @param[out] C Pointer to output C matrix block. + * @param[in,out] workspace Temporary buffer of at least m*n doubles. + * @param[in] layout BLAS matrix layout (RowMajor by default). + */ +void compute_atba_gemm(const double* A, const double* B, double* C, int m, + int n, std::vector& workspace, + blas::Layout layout = blas::Layout::RowMajor); + // Forward declaration class SCFImpl; /** @@ -98,10 +129,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 +146,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 + build_rohf_convergence_matrices(const SCFImpl& scf_impl); + /** * @brief Calculate orbital gradient (OG) error for convergence checking * @@ -137,6 +184,55 @@ 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(); + + /** + * @brief Copy precomputed ROHF convergence matrices into this object's cache + * + * Used by composite algorithms (e.g. DIIS_GDM) to propagate the + * already-computed ROHF matrices from the outer algorithm to inner + * sub-algorithm objects without expensive recomputation. + */ + void set_rohf_convergence_cache(const RowMajorMatrix& fock, + const RowMajorMatrix& density); + protected: const SCFContext& ctx_; ///< Reference to SCF context double og_error_ = 0.0; ///< Current orbital gradient error @@ -149,5 +245,7 @@ class SCFAlgorithm { double delta_energy_ = std::numeric_limits::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 diff --git a/cpp/src/qdk/chemistry/algorithms/microsoft/scf/src/scf/ks_impl.cpp b/cpp/src/qdk/chemistry/algorithms/microsoft/scf/src/scf/ks_impl.cpp index f1d43a6d5..2e902faff 100644 --- a/cpp/src/qdk/chemistry/algorithms/microsoft/scf/src/scf/ks_impl.cpp +++ b/cpp/src/qdk/chemistry/algorithms/microsoft/scf/src/scf/ks_impl.cpp @@ -36,6 +36,11 @@ KSImpl::KSImpl(std::shared_ptr mol, const SCFConfig& cfg, std::shared_ptr 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 KSImpl::evaluate_trial_density_energy_and_fock( - const RowMajorMatrix& P_matrix, const std::source_location& loc) const { + const RowMajorMatrix& P_matrix, RowMajorMatrix& J_out, + RowMajorMatrix& 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; diff --git a/cpp/src/qdk/chemistry/algorithms/microsoft/scf/src/scf/ks_impl.h b/cpp/src/qdk/chemistry/algorithms/microsoft/scf/src/scf/ks_impl.h index 85264864a..b0bccb38b 100644 --- a/cpp/src/qdk/chemistry/algorithms/microsoft/scf/src/scf/ks_impl.h +++ b/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 Output Coulomb matrix J + * @param K_out Output 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 evaluate_trial_density_energy_and_fock( - const RowMajorMatrix& P_matrix, + const RowMajorMatrix& P_matrix, RowMajorMatrix& J_out, + RowMajorMatrix& K_out, const std::source_location& loc = std::source_location::current()) const override; diff --git a/cpp/src/qdk/chemistry/algorithms/microsoft/scf/src/scf/scf_impl.cpp b/cpp/src/qdk/chemistry/algorithms/microsoft/scf/src/scf/scf_impl.cpp index be74058ed..90d22a7e2 100644 --- a/cpp/src/qdk/chemistry/algorithms/microsoft/scf/src/scf/scf_impl.cpp +++ b/cpp/src/qdk/chemistry/algorithms/microsoft/scf/src/scf/scf_impl.cpp @@ -16,6 +16,7 @@ #endif #include +#include #include #include #include @@ -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_->build_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, @@ -1036,8 +1047,27 @@ void SCFImpl::write_gradients_(const std::vector& gradients, std::pair SCFImpl::evaluate_trial_density_energy_and_fock( const RowMajorMatrix& P_matrix, const std::source_location& loc) const { + RowMajorMatrix J_out = RowMajorMatrix::Zero( + num_density_matrices_ * num_atomic_orbitals_, num_atomic_orbitals_); + RowMajorMatrix K_out = RowMajorMatrix::Zero( + num_density_matrices_ * num_atomic_orbitals_, num_atomic_orbitals_); + return evaluate_trial_density_energy_and_fock(P_matrix, J_out, K_out, loc); +} + +std::pair +SCFImpl::evaluate_trial_density_energy_and_fock( + const RowMajorMatrix& P_matrix, RowMajorMatrix& J_out, + RowMajorMatrix& K_out, const std::source_location& loc) const { QDK_LOG_TRACE_ENTERING(); + 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. This function is called within " + "iterate_ which runs only on rank 0, but build_jk_matrices requires " + "all MPI ranks."); + } + QDK_LOGGER().debug( "Computing energy and Fock matrix by trial density matrix (called from " "function '{}' at line {})", @@ -1049,27 +1079,23 @@ SCFImpl::evaluate_trial_density_energy_and_fock( RowMajorMatrix H_matrix = H_.eval(); RowMajorMatrix F_matrix = H_matrix; - auto [alpha, beta, omega] = get_hyb_coeff_(); - RowMajorMatrix J_matrix = RowMajorMatrix::Zero( - num_density_matrices_ * num_atomic_orbitals_, num_atomic_orbitals_); - RowMajorMatrix K_matrix = RowMajorMatrix::Zero( - num_density_matrices_ * num_atomic_orbitals_, num_atomic_orbitals_); - eri_->build_JK(P_matrix.data(), J_matrix.data(), K_matrix.data(), alpha, beta, - omega); + build_jk_matrices(P_matrix, J_out, K_out); + #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_, - num_atomic_orbitals_)) + (J_out.block(0, 0, num_atomic_orbitals_, num_atomic_orbitals_) + + J_out.block(num_atomic_orbitals_, 0, num_atomic_orbitals_, + num_atomic_orbitals_)) .replicate(2, 1) - - K_matrix; + K_out; } else { - F_matrix += J_matrix - 0.5 * K_matrix; + F_matrix += J_out - 0.5 * K_out; } } @@ -1097,4 +1123,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 diff --git a/cpp/src/qdk/chemistry/algorithms/microsoft/scf/src/scf/scf_impl.h b/cpp/src/qdk/chemistry/algorithms/microsoft/scf/src/scf/scf_impl.h index 67f5ade5d..3b61a946c 100644 --- a/cpp/src/qdk/chemistry/algorithms/microsoft/scf/src/scf/scf_impl.h +++ b/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,6 +221,8 @@ class SCFImpl { * must be fully initialized before calling it. * * @param P_matrix Trial density matrix + * @param J_out Output Coulomb matrix J (same shape as P_matrix) + * @param K_out Output 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 @@ -205,6 +230,22 @@ class SCFImpl { */ virtual std::pair evaluate_trial_density_energy_and_fock( + const RowMajorMatrix& P_matrix, RowMajorMatrix& J_out, + RowMajorMatrix& K_out, + const std::source_location& loc = std::source_location::current()) const; + + /** + * @brief Convenience overload that allocates J and K internally. + * + * Use this when the caller does not need the J and K matrices. + * + * @param P_matrix Trial density 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 + */ + std::pair evaluate_trial_density_energy_and_fock( const RowMajorMatrix& P_matrix, const std::source_location& loc = std::source_location::current()) const; diff --git a/cpp/src/qdk/chemistry/algorithms/microsoft/scf/src/scf_algorithm/diis.cpp b/cpp/src/qdk/chemistry/algorithms/microsoft/scf/src/scf_algorithm/diis.cpp index 67cf66f85..a07040ac8 100644 --- a/cpp/src/qdk/chemistry/algorithms/microsoft/scf/src/scf_algorithm/diis.cpp +++ b/cpp/src/qdk/chemistry/algorithms/microsoft/scf/src/scf_algorithm/diis.cpp @@ -6,13 +6,10 @@ #include -#include #include -#include #include #include #include -#include #include "../scf/scf_impl.h" #include "util/macros.h" @@ -121,164 +118,6 @@ class DIIS { std::numeric_limits::infinity(); ///< Current DIIS error }; -/** - * @brief Rebuild the effective ROHF Fock and density matrices - * - * Implements the averaged-block construction (Guest & Saunders 1974, - * Plakhutin & Davidson 2014) to obtain the single-density ROHF view and - * caches the total density $P_\alpha + P_\beta$ alongside it. - * - * @param F Spin-blocked Fock matrix with alpha and beta blocks stacked - * @param C Molecular-orbital coefficients used for block transformations - * @param P Spin-blocked density matrix - * @param nelec_alpha Number of alpha electrons - * @param nelec_beta Number of beta electrons - * @param effective_fock Output effective ROHF Fock matrix (AO basis) - * @param total_density Output total density matrix (P_alpha + P_beta) - */ -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) { - QDK_LOG_TRACE_ENTERING(); - const int num_atomic_orbitals = static_cast(C.rows()); - const int num_molecular_orbitals = static_cast(C.cols()); - if (num_atomic_orbitals != num_molecular_orbitals) { - throw std::runtime_error( - "ROHF build requires number of atomic orbitals to equal number of " - "molecular orbitals!"); - } - - total_density = - P.block(0, 0, num_atomic_orbitals, num_atomic_orbitals) + - P.block(num_atomic_orbitals, 0, num_atomic_orbitals, num_atomic_orbitals); - - if (effective_fock.rows() != num_atomic_orbitals || - effective_fock.cols() != num_atomic_orbitals) { - effective_fock = - RowMajorMatrix::Zero(num_atomic_orbitals, num_atomic_orbitals); - } - - if (C.isZero()) { - effective_fock.noalias() = - F.block(0, 0, num_atomic_orbitals, num_atomic_orbitals); - return; - } - - RowMajorMatrix F_up_mo = - RowMajorMatrix::Zero(num_molecular_orbitals, num_molecular_orbitals); - RowMajorMatrix F_dn_mo = F_up_mo; - RowMajorMatrix effective_F_mo = F_up_mo; - - F_up_mo.noalias() = C.transpose() * - F.block(0, 0, num_atomic_orbitals, num_atomic_orbitals) * - C; - F_dn_mo.noalias() = C.transpose() * - F.block(num_atomic_orbitals, 0, num_atomic_orbitals, - num_atomic_orbitals) * - C; - - auto average_block = [&effective_F_mo, &F_up_mo, &F_dn_mo]( - int row, int col, int rows, int cols) { - if (rows <= 0 || cols <= 0) return; - effective_F_mo.block(row, col, rows, cols).noalias() = - 0.5 * (F_up_mo.block(row, col, rows, cols) + - F_dn_mo.block(row, col, rows, cols)); - }; - auto copy_block = [&effective_F_mo](const RowMajorMatrix& src, int row, - int col, int rows, int cols) { - if (rows <= 0 || cols <= 0) return; - effective_F_mo.block(row, col, rows, cols) = - src.block(row, col, rows, cols); - }; - - const int nd = nelec_beta; - const int ns = nelec_alpha - nelec_beta; - const int nv = num_molecular_orbitals - nelec_alpha; - - average_block(0, 0, nd, nd); - average_block(0, nd + ns, nd, nv); - average_block(nd + ns, 0, nv, nd); - average_block(nd + ns, nd + ns, nv, nv); - average_block(nd, nd, ns, ns); - copy_block(F_dn_mo, 0, nd, nd, ns); - copy_block(F_dn_mo, nd, 0, ns, nd); - copy_block(F_up_mo, nd, nd + ns, ns, nv); - copy_block(F_up_mo, nd + ns, nd, nv, ns); - - // Transform the effective Fock matrix back to AO basis by solving - // C^{-T} * F_MO * C^{-1} = F_AO - // We use LAPACK's getrf/getrs to solve the linear systems involving C^T and - // C without explicitly inverting C - const int matrix_dim = num_molecular_orbitals; - using ColMajorMatrix = - Eigen::Matrix; - // LAPACK expects column-major layout, so we copy the row-major data into a - // column-major matrix without transposing the logical layout - ColMajorMatrix Ct = - Eigen::Map(C.data(), matrix_dim, C.rows()); - // F_MO is symmetric, so we can use it directly as the right-hand side - // without transposing - ColMajorMatrix temp_rhs = effective_F_mo; - std::vector ipiv(matrix_dim); - - auto info = - lapack::getrf(matrix_dim, matrix_dim, Ct.data(), matrix_dim, ipiv.data()); - if (info != 0) { - throw std::runtime_error("getrf failed while factorizing C^T"); - } - - info = lapack::getrs(lapack::Op::NoTrans, matrix_dim, matrix_dim, Ct.data(), - matrix_dim, ipiv.data(), temp_rhs.data(), matrix_dim); - if (info != 0) { - throw std::runtime_error("getrs failed while solving C^T X = F_mo"); - } - - temp_rhs.transposeInPlace(); - info = lapack::getrs(lapack::Op::NoTrans, matrix_dim, matrix_dim, Ct.data(), - matrix_dim, ipiv.data(), temp_rhs.data(), matrix_dim); - if (info != 0) { - throw std::runtime_error("getrs failed while solving C^T X = M^T"); - } - - effective_fock = temp_rhs.transpose(); - if (!effective_fock.isApprox(effective_fock.transpose())) { - effective_fock = 0.5 * (effective_fock + effective_fock.transpose().eval()); - } -} - -/** - * @brief Reconstruct spin-blocked densities from the ROHF MO matrix - * - * Generates $P_\alpha$ and $P_\beta$ blocks so we can hand the updated - * density back to SCFImpl after diagonalization. - * - * @param P Spin-blocked density matrix to overwrite - * @param C Molecular-orbital coefficients from latest diagonalization - * @param nelec_alpha Number of alpha electrons - * @param nelec_beta Number of beta electrons - */ -void update_rohf_density_matrix(RowMajorMatrix& P, const RowMajorMatrix& C, - int nelec_alpha, int nelec_beta) { - QDK_LOG_TRACE_ENTERING(); - const int num_atomic_orbitals = static_cast(C.rows()); - - auto build_density = [&](auto&& target, int n_occ) { - if (n_occ <= 0) { - target.setZero(); - return; - } - target.noalias() = C.block(0, 0, num_atomic_orbitals, n_occ) * - C.block(0, 0, num_atomic_orbitals, n_occ).transpose(); - }; - - auto P_alpha = P.block(0, 0, num_atomic_orbitals, num_atomic_orbitals); - auto P_beta = - P.block(num_atomic_orbitals, 0, num_atomic_orbitals, num_atomic_orbitals); - build_density(P_alpha, nelec_alpha); - build_density(P_beta, nelec_beta); -} - DIIS::DIIS(const SCFContext& ctx, const size_t subspace_size) : ctx_(ctx), subspace_size_(subspace_size) { QDK_LOG_TRACE_ENTERING(); @@ -430,14 +269,6 @@ DIIS::DIIS(const SCFContext& ctx, const size_t subspace_size) : SCFAlgorithm(ctx), diis_impl_(std::make_unique(ctx, subspace_size)) { QDK_LOG_TRACE_ENTERING(); - if (ctx.cfg->scf_orbital_type == SCFOrbitalType::RestrictedOpenShell) { - const int num_atomic_orbitals = - static_cast(ctx.basis_set->num_atomic_orbitals); - rohf_effective_fock_ = - RowMajorMatrix::Zero(num_atomic_orbitals, num_atomic_orbitals); - rohf_total_density_ = - RowMajorMatrix::Zero(num_atomic_orbitals, num_atomic_orbitals); - } } DIIS::~DIIS() noexcept = default; @@ -447,8 +278,7 @@ void DIIS::iterate(SCFImpl& scf_impl) { // Decide which Fock/density view to feed into Pulay: RHF/UHF use the direct // SCFImpl matrices, while ROHF works on the cached total-density view. - RowMajorMatrix& working_density = select_working_density(scf_impl); - const RowMajorMatrix& working_fock = select_working_fock(scf_impl); + auto [working_density, working_fock] = select_working_matrices(scf_impl); auto& C = scf_impl.orbitals_matrix(); const auto& S = scf_impl.overlap(); @@ -490,77 +320,21 @@ void DIIS::iterate(SCFImpl& scf_impl) { } } -void DIIS::update_density_matrix(RowMajorMatrix& P, const RowMajorMatrix& C, - bool unrestricted, int nelec_alpha, - int nelec_beta) { - QDK_LOG_TRACE_ENTERING(); - if (ctx_.cfg->scf_orbital_type == SCFOrbitalType::RestrictedOpenShell) { - impl::update_rohf_density_matrix(P, C, nelec_alpha, nelec_beta); - return; - } - SCFAlgorithm::update_density_matrix(P, C, unrestricted, nelec_alpha, - nelec_beta); -} - -void DIIS::build_rohf_f_p_matrix(const RowMajorMatrix& F, - const RowMajorMatrix& C, - const RowMajorMatrix& P, int nelec_alpha, - int nelec_beta) { - QDK_LOG_TRACE_ENTERING(); - if (ctx_.cfg->scf_orbital_type != SCFOrbitalType::RestrictedOpenShell) { - throw std::logic_error("ROHF matrix build requested for non-ROHF run"); - } - impl::build_rohf_f_p_matrix(F, C, P, nelec_alpha, nelec_beta, - rohf_effective_fock_, rohf_total_density_); -} - -const RowMajorMatrix& DIIS::get_rohf_fock_matrix() const { - QDK_LOG_TRACE_ENTERING(); - if (rohf_effective_fock_.size() == 0) { - throw std::logic_error("ROHF cache not initialized"); - } - return rohf_effective_fock_; -} - -const RowMajorMatrix& DIIS::get_rohf_density_matrix() const { - QDK_LOG_TRACE_ENTERING(); - if (rohf_total_density_.size() == 0) { - throw std::logic_error("ROHF cache not initialized"); - } - return rohf_total_density_; -} - -RowMajorMatrix& DIIS::rohf_density_matrix() { - QDK_LOG_TRACE_ENTERING(); - if (rohf_total_density_.size() == 0) { - throw std::logic_error("ROHF cache not initialized"); - } - return rohf_total_density_; -} - double DIIS::current_diis_error() const { QDK_LOG_TRACE_ENTERING(); return diis_impl_->get_diis_error(); } -RowMajorMatrix& DIIS::select_working_density(SCFImpl& scf_impl) { - QDK_LOG_TRACE_ENTERING(); - if (ctx_.cfg->scf_orbital_type == SCFOrbitalType::RestrictedOpenShell) { - // The ROHF helper is refreshed inside SCFAlgorithm::check_convergence(), - // so by the time iterate() is invoked the total density view is already - // up to date. - (void)scf_impl; // kept for symmetry with the unrestricted branch - return rohf_density_matrix(); - } - return scf_impl.density_matrix(); -} - -const RowMajorMatrix& DIIS::select_working_fock(const SCFImpl& scf_impl) { +std::pair DIIS::select_working_matrices( + SCFImpl& scf_impl) { QDK_LOG_TRACE_ENTERING(); if (ctx_.cfg->scf_orbital_type == SCFOrbitalType::RestrictedOpenShell) { - return get_rohf_fock_matrix(); + // Cache is populated by check_convergence() for pure DIIS, or by + // DIIS_GDM::iterate via set_rohf_convergence_cache() for composite use. + return {rohf_convergence_density_matrix(), + get_rohf_convergence_fock_matrix()}; } - return scf_impl.get_fock_matrix(); + return {scf_impl.density_matrix(), scf_impl.get_fock_matrix()}; } } // namespace qdk::chemistry::scf diff --git a/cpp/src/qdk/chemistry/algorithms/microsoft/scf/src/scf_algorithm/diis.h b/cpp/src/qdk/chemistry/algorithms/microsoft/scf/src/scf_algorithm/diis.h index 6b1628d7f..3cbf4781f 100644 --- a/cpp/src/qdk/chemistry/algorithms/microsoft/scf/src/scf_algorithm/diis.h +++ b/cpp/src/qdk/chemistry/algorithms/microsoft/scf/src/scf_algorithm/diis.h @@ -8,6 +8,7 @@ #include #include +#include namespace qdk::chemistry::scf { @@ -53,60 +54,6 @@ class DIIS : public SCFAlgorithm { */ void iterate(SCFImpl& scf_impl) override; - /** - * @brief Update the density matrix - * - * Default implementation handles restricted and unrestricted cases, while - * subclasses such as ASAHF may override it - * - * @param[in,out] P Density matrix to overwrite - * @param[in] C Molecular orbital coefficient matrix - * @param[in] unrestricted True if two spin blocks are present - * @param[in] nelec_alpha Number of alpha electrons - * @param[in] nelec_beta Number of beta electrons - */ - void update_density_matrix(RowMajorMatrix& P, const RowMajorMatrix& C, - bool unrestricted, int nelec_alpha, - int nelec_beta) override; - - /** - * @brief Build cached ROHF Fock/density matrices for convergence checks - * - * Converts spin-blocked Fock/density matrices into the total-density / - * effective-Fock representation, then saved internally for use in DIIS - * iterations and convergence checks - * - * @param[in] F Spin-blocked Fock matrix from `SCFImpl` - * @param[in] C Molecular orbital coefficients - * @param[in] P Spin-blocked density matrix - * @param[in] nelec_alpha Number of alpha electrons - * @param[in] nelec_beta Number of beta electrons - */ - void build_rohf_f_p_matrix(const RowMajorMatrix& F, const RowMajorMatrix& C, - const RowMajorMatrix& P, int nelec_alpha, - int nelec_beta); - - /** - * @brief Access the cached ROHF-effective Fock matrix - * - * @return Const reference to the total-density ROHF Fock matrix - */ - const RowMajorMatrix& get_rohf_fock_matrix() const; - - /** - * @brief Access the cached total (alpha+beta) density matrix - * - * @return Const reference to the cached ROHF density - */ - const RowMajorMatrix& get_rohf_density_matrix() const; - - /** - * @brief Mutable access to the total density matrix (used inside iterate) - * - * @return Reference to the cached ROHF density - */ - RowMajorMatrix& rohf_density_matrix(); - private: /** * @brief Return most recent Pulay error metric for damping logic @@ -114,17 +61,12 @@ class DIIS : public SCFAlgorithm { double current_diis_error() const; /** - * @brief Select the density view used for the current iteration - */ - RowMajorMatrix& select_working_density(SCFImpl& scf_impl); - /** - * @brief Select the Fock matrix view used for the current iteration + * @brief Select the density/fock views used for the current iteration */ - const RowMajorMatrix& select_working_fock(const SCFImpl& scf_impl); + std::pair select_working_matrices( + SCFImpl& scf_impl); std::unique_ptr diis_impl_; ///< Pulay DIIS core implementation - 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 diff --git a/cpp/src/qdk/chemistry/algorithms/microsoft/scf/src/scf_algorithm/diis_gdm.cpp b/cpp/src/qdk/chemistry/algorithms/microsoft/scf/src/scf_algorithm/diis_gdm.cpp index c80cefdda..8f6ec9d78 100644 --- a/cpp/src/qdk/chemistry/algorithms/microsoft/scf/src/scf_algorithm/diis_gdm.cpp +++ b/cpp/src/qdk/chemistry/algorithms/microsoft/scf/src/scf_algorithm/diis_gdm.cpp @@ -68,6 +68,14 @@ void DIIS_GDM::iterate(SCFImpl& scf_impl) { if (use_gdm_) { gdm_algorithm_->iterate(scf_impl); } else { + // For ROHF, propagate the ROHF convergence cache computed by our + // check_convergence() to the inner DIIS object so it can read from + // its own cache without expensive recomputation. + if (ctx_.cfg->scf_orbital_type == SCFOrbitalType::RestrictedOpenShell) { + diis_algorithm_->set_rohf_convergence_cache( + get_rohf_convergence_fock_matrix(), + get_rohf_convergence_density_matrix()); + } diis_algorithm_->iterate(scf_impl); } } diff --git a/cpp/src/qdk/chemistry/algorithms/microsoft/scf/src/scf_algorithm/gdm.cpp b/cpp/src/qdk/chemistry/algorithms/microsoft/scf/src/scf_algorithm/gdm.cpp index 74a6ac52b..29d2fdf9c 100644 --- a/cpp/src/qdk/chemistry/algorithms/microsoft/scf/src/scf_algorithm/gdm.cpp +++ b/cpp/src/qdk/chemistry/algorithms/microsoft/scf/src/scf_algorithm/gdm.cpp @@ -17,6 +17,7 @@ #include "../scf/scf_impl.h" #include "line_search.h" #include "qdk/chemistry/scf/core/scf.h" +#include "qdk/chemistry/scf/core/scf_algorithm.h" #include "qdk/chemistry/scf/core/types.h" #include "util/matrix_exp.h" #ifdef QDK_CHEMISTRY_ENABLE_GPU @@ -34,13 +35,22 @@ namespace impl { * @param[in] spin_index Spin index (0 for alpha, 1 for beta) * @param[in] kappa_vector The kappa vector to apply for rotation * @param[in] num_occupied_orbitals Number of occupied orbitals for this spin - * @param[in] num_molecular_orbitals Number of molecular orbitals + * @param[in] num_orbital_spin_blocks Number of spin blocks (1 for RHF, 2 for + * UHF) */ -static void apply_orbital_rotation(RowMajorMatrix& C, const int spin_index, - const Eigen::VectorXd& kappa_vector, - const int num_occupied_orbitals, - const int num_molecular_orbitals) { +static void apply_restricted_unrestricted_orbital_rotation( + RowMajorMatrix& C, const int spin_index, + const Eigen::VectorXd& kappa_vector, const int num_occupied_orbitals, + const int num_orbital_spin_blocks) { QDK_LOG_TRACE_ENTERING(); + const int num_molecular_orbitals = C.cols(); + const int num_atomic_orbitals = + (num_orbital_spin_blocks == 2) ? C.rows() / 2 : C.rows(); + + if (num_orbital_spin_blocks == 2 && (C.rows() % 2 != 0)) { + throw std::invalid_argument("In UHF or UKS, C row count must be even!"); + } + const int num_virtual_orbitals = num_molecular_orbitals - num_occupied_orbitals; @@ -60,18 +70,296 @@ static void apply_orbital_rotation(RowMajorMatrix& C, const int spin_index, matrix_exp(kappa_complete.data(), exp_kappa.data(), num_molecular_orbitals); // Rotate C: C' = C * exp(kappa) + const int C_block_offset = + num_atomic_orbitals * num_molecular_orbitals * spin_index; RowMajorMatrix C_before_rotate = - C.block(num_molecular_orbitals * spin_index, 0, num_molecular_orbitals, + C.block(spin_index * num_atomic_orbitals, 0, num_atomic_orbitals, num_molecular_orbitals); + const double* C_before_rotate_ptr = C_before_rotate.data(); blas::gemm(blas::Layout::RowMajor, blas::Op::NoTrans, blas::Op::NoTrans, - num_molecular_orbitals, num_molecular_orbitals, + num_atomic_orbitals, num_molecular_orbitals, + num_molecular_orbitals, 1.0, C_before_rotate_ptr, + num_molecular_orbitals, exp_kappa.data(), num_molecular_orbitals, + 0.0, C.data() + C_block_offset, num_molecular_orbitals); +} + +/** + * @brief Construct the ROHF kappa matrix and apply rotation C * exp(kappa) + * + * @param[in,out] C Molecular orbital coefficient matrix + * @param[in] num_electrons Occupied orbital counts (alpha, beta) + * @param[in] kappa_vector Concatenated ROHF rotation parameters + */ +static void apply_restricted_open_shell_orbital_rotation( + RowMajorMatrix& C, const std::vector& num_electrons, + const Eigen::VectorXd& kappa_vector) { + QDK_LOG_TRACE_ENTERING(); + const int num_atomic_orbitals = C.rows(); + const int num_molecular_orbitals = C.cols(); + const int num_closed_orbitals = num_electrons[1]; + const int num_open_orbitals = num_electrons[0] - num_closed_orbitals; + const int num_virtual_orbitals = num_molecular_orbitals - num_electrons[0]; + + int offset = 0; + const int iw_size = num_closed_orbitals * num_open_orbitals; + const int wa_size = num_open_orbitals * num_virtual_orbitals; + const int ia_size = num_closed_orbitals * num_virtual_orbitals; + + const auto kappa_iw = Eigen::Map( + kappa_vector.data() + offset, num_closed_orbitals, num_open_orbitals); + offset += iw_size; + + const auto kappa_wa = Eigen::Map( + kappa_vector.data() + offset, num_open_orbitals, num_virtual_orbitals); + offset += wa_size; + + const auto kappa_ia = Eigen::Map( + kappa_vector.data() + offset, num_closed_orbitals, num_virtual_orbitals); + offset += ia_size; + + RowMajorMatrix kappa_complete = + RowMajorMatrix::Zero(num_molecular_orbitals, num_molecular_orbitals); + + kappa_complete.block(0, num_closed_orbitals, num_closed_orbitals, + num_open_orbitals) = kappa_iw; + kappa_complete.block(num_closed_orbitals, 0, num_open_orbitals, + num_closed_orbitals) = -kappa_iw.transpose(); + + const int open_start = num_closed_orbitals; + const int virtual_start = num_closed_orbitals + num_open_orbitals; + + kappa_complete.block(open_start, virtual_start, num_open_orbitals, + num_virtual_orbitals) = kappa_wa; + kappa_complete.block(virtual_start, open_start, num_virtual_orbitals, + num_open_orbitals) = -kappa_wa.transpose(); + + kappa_complete.block(0, virtual_start, num_closed_orbitals, + num_virtual_orbitals) = kappa_ia; + kappa_complete.block(virtual_start, 0, num_virtual_orbitals, + num_closed_orbitals) = -kappa_ia.transpose(); + + // kappa_vector stores only unique (upper-triangular) rotation parameters. + // kappa_complete is the full antisymmetric matrix (K = -K^T), so each unique + // parameter appears twice. Scale by 0.5 to match the gradient convention. + kappa_complete *= 0.5; + + RowMajorMatrix exp_kappa = + RowMajorMatrix::Zero(num_molecular_orbitals, num_molecular_orbitals); + matrix_exp(kappa_complete.data(), exp_kappa.data(), num_molecular_orbitals); + + RowMajorMatrix C_before_rotate = C; + blas::gemm(blas::Layout::RowMajor, blas::Op::NoTrans, blas::Op::NoTrans, + num_atomic_orbitals, num_molecular_orbitals, num_molecular_orbitals, 1.0, C_before_rotate.data(), num_molecular_orbitals, exp_kappa.data(), num_molecular_orbitals, - 0.0, - C.block(num_molecular_orbitals * spin_index, 0, - num_molecular_orbitals, num_molecular_orbitals) - .data(), - num_molecular_orbitals); + 0.0, C.data(), num_molecular_orbitals); +} + +/** + * @brief Compute restricted/unrestricted orbital gradients for all spins + * @param[in] F Fock matrix in AO basis + * @param[in] C Molecular orbital coefficient matrix + * @param[in] num_electrons Occupied orbital counts per spin component + * @param[in] rotation_offset Starting index for each spin's rotation slice + * @param[in] rotation_size Number of rotation parameters per spin + * (n_occ*n_virt) + * @param[in] num_orbital_spin_blocks Number of spin blocks to iterate + * @param[out] gradient Output gradient vector (concatenated across spins) + */ +static void compute_restricted_unrestricted_gradient( + const RowMajorMatrix& F, const RowMajorMatrix& C, + const std::vector& num_electrons, + const std::vector& rotation_offset, + const std::vector& rotation_size, int num_orbital_spin_blocks, + Eigen::VectorXd& gradient) { + const int num_molecular_orbitals = C.cols(); + const int num_atomic_orbitals = + (num_orbital_spin_blocks == 2) ? C.rows() / 2 : C.rows(); + + if (num_orbital_spin_blocks == 2 && (C.rows() % 2 != 0)) { + throw std::invalid_argument("In UHF or UKS, C row count must be even!"); + } + + int total_rotation_size = 0; + for (int i = 0; i < num_orbital_spin_blocks; ++i) { + total_rotation_size += rotation_size[i]; + } + gradient.setZero(total_rotation_size); + + std::vector atba_workspace(static_cast(num_atomic_orbitals) * + num_molecular_orbitals); + + for (int spin_index = 0; spin_index < num_orbital_spin_blocks; ++spin_index) { + const int num_occupied_orbitals = num_electrons[spin_index]; + const int num_virtual_orbitals = + num_molecular_orbitals - num_occupied_orbitals; + const int spin_rotation_size = rotation_size[spin_index]; + + if (spin_rotation_size == 0) { + continue; + } + + RowMajorMatrix F_MO = + RowMajorMatrix::Zero(num_molecular_orbitals, num_molecular_orbitals); + const int C_block_offset = + num_atomic_orbitals * num_molecular_orbitals * spin_index; + const int F_block_offset = + num_atomic_orbitals * num_atomic_orbitals * spin_index; + const double* C_block_ptr = C.data() + C_block_offset; + const double* F_block_ptr = F.data() + F_block_offset; + compute_atba_gemm(C_block_ptr, F_block_ptr, F_MO.data(), + num_atomic_orbitals, num_molecular_orbitals, + atba_workspace); + + // Extract occupied-virtual block and compute gradient + // The -4.0 before F_{ia} comes from derivative of energy w.r.t. kappa + // Reference: Helgaker, T., Jørgensen, P., & Olsen, J. (2000). Molecular + // electronic-structure theory, Eq. 10.8.34 (2013 reprint edition) + // -4.0 is for restricted closed-shell system. For unrestricted systems, the + // gradient is computed separately for each spin component, in that case the + // coefficient before F_{ia, spin} is -2.0 + RowMajorMatrix gradient_matrix = + -((num_orbital_spin_blocks == 2) ? 2.0 : 4.0) * + F_MO.block(0, num_occupied_orbitals, num_occupied_orbitals, + num_virtual_orbitals); + + gradient.segment(rotation_offset[spin_index], spin_rotation_size) = + Eigen::Map(gradient_matrix.data(), + spin_rotation_size) + .eval(); + } +} + +/** + * @brief Compute ROHF orbital gradients using the generalized Fock matrix + * + * The gradient is packed as (iw, wa, ia) blocks following the same + * segmentation used by the ROHF kappa vector. + * + * @param[in] scf_impl SCF implementation for core Hamiltonian access + * @param[in] C Molecular orbital coefficient matrix + * @param[in] J_ao Coulomb matrix in AO basis for alpha/beta blocks + * @param[in] K_ao Exchange matrix in AO basis for alpha/beta blocks + * @param[in] num_electrons Occupied orbital counts (alpha, beta) + * @param[in] rotation_size Rotation size for the ROHF kappa vector + * @param[out] gradient Output gradient vector + */ +static void compute_restricted_open_shell_gradient( + const SCFImpl& scf_impl, const RowMajorMatrix& C, + const RowMajorMatrix& J_ao, const RowMajorMatrix& K_ao, + const std::vector& num_electrons, + const std::vector& rotation_size, Eigen::VectorXd& gradient) { + const int total_rotation_size = rotation_size[0]; + gradient.setZero(total_rotation_size); + + const int num_molecular_orbitals = static_cast(C.cols()); + const int num_closed_orbitals = num_electrons[1]; + const int num_open_orbitals = num_electrons[0] - num_closed_orbitals; + const int num_virtual_orbitals = num_molecular_orbitals - num_electrons[0]; + RowMajorMatrix generalized_fock_mo = + RowMajorMatrix::Zero(num_molecular_orbitals, num_molecular_orbitals); + + const int num_atomic_orbitals = scf_impl.get_num_atomic_orbitals(); + + if (J_ao.size() < 2 * num_atomic_orbitals * num_atomic_orbitals || + K_ao.size() < 2 * num_atomic_orbitals * num_atomic_orbitals) { + throw std::invalid_argument( + "J_ao and K_ao must each contain alpha and beta blocks " + "(size >= 2 * num_atomic_orbitals^2)."); + } + + // H_ is stored as (num_density_matrices * num_atomic_orbitals) x + // num_atomic_orbitals, but the core Hamiltonian is spin-independent so both + // blocks are identical. We only use the first block via .data(). + const auto& H_ao_full = scf_impl.get_core_hamiltonian(); + + const auto J_alpha_ao = Eigen::Map( + J_ao.data(), num_atomic_orbitals, num_atomic_orbitals); + const auto J_beta_ao = Eigen::Map( + J_ao.data() + num_atomic_orbitals * num_atomic_orbitals, + num_atomic_orbitals, num_atomic_orbitals); + const auto K_alpha_ao = Eigen::Map( + K_ao.data(), num_atomic_orbitals, num_atomic_orbitals); + const auto K_beta_ao = Eigen::Map( + K_ao.data() + num_atomic_orbitals * num_atomic_orbitals, + num_atomic_orbitals, num_atomic_orbitals); + + const auto C_ao_mo = + C.block(0, 0, num_atomic_orbitals, num_molecular_orbitals); + const double* C_ao_mo_ptr = C_ao_mo.data(); + std::vector atba_workspace(static_cast(num_atomic_orbitals) * + num_molecular_orbitals); + + // Calculate Generalized Fock matrix in MO basis + RowMajorMatrix H_mo = + RowMajorMatrix::Zero(num_molecular_orbitals, num_molecular_orbitals); + compute_atba_gemm(C_ao_mo_ptr, H_ao_full.data(), H_mo.data(), + num_atomic_orbitals, num_molecular_orbitals, + atba_workspace); + + RowMajorMatrix J_alpha_mo = + RowMajorMatrix::Zero(num_molecular_orbitals, num_molecular_orbitals); + compute_atba_gemm(C_ao_mo_ptr, J_alpha_ao.data(), J_alpha_mo.data(), + num_atomic_orbitals, num_molecular_orbitals, + atba_workspace); + + RowMajorMatrix J_beta_mo = + RowMajorMatrix::Zero(num_molecular_orbitals, num_molecular_orbitals); + compute_atba_gemm(C_ao_mo_ptr, J_beta_ao.data(), J_beta_mo.data(), + num_atomic_orbitals, num_molecular_orbitals, + atba_workspace); + + RowMajorMatrix K_alpha_mo = + RowMajorMatrix::Zero(num_molecular_orbitals, num_molecular_orbitals); + compute_atba_gemm(C_ao_mo_ptr, K_alpha_ao.data(), K_alpha_mo.data(), + num_atomic_orbitals, num_molecular_orbitals, + atba_workspace); + + RowMajorMatrix K_beta_mo = + RowMajorMatrix::Zero(num_molecular_orbitals, num_molecular_orbitals); + compute_atba_gemm(C_ao_mo_ptr, K_beta_ao.data(), K_beta_mo.data(), + num_atomic_orbitals, num_molecular_orbitals, + atba_workspace); + + RowMajorMatrix F_I = H_mo + 2.0 * J_beta_mo - K_beta_mo; + RowMajorMatrix F_A = J_alpha_mo - J_beta_mo - 0.5 * (K_alpha_mo - K_beta_mo); + RowMajorMatrix Q = J_alpha_mo - J_beta_mo - (K_alpha_mo - K_beta_mo); + + RowMajorMatrix F_sum = F_I + F_A; + generalized_fock_mo.block(0, 0, num_closed_orbitals, num_molecular_orbitals) = + 2.0 * F_sum.block(0, 0, num_closed_orbitals, num_molecular_orbitals); + generalized_fock_mo.block(num_closed_orbitals, 0, num_open_orbitals, + num_molecular_orbitals) = + (F_I + Q).block(num_closed_orbitals, 0, num_open_orbitals, + num_molecular_orbitals); + + int offset = 0; + RowMajorMatrix grad_iw = + -2.0 * + (generalized_fock_mo.block(0, num_closed_orbitals, num_closed_orbitals, + num_open_orbitals) - + generalized_fock_mo + .block(num_closed_orbitals, 0, num_open_orbitals, + num_closed_orbitals) + .transpose()); + gradient.segment(offset, num_closed_orbitals * num_open_orbitals) = + Eigen::Map(grad_iw.data(), grad_iw.size()).eval(); + offset += num_closed_orbitals * num_open_orbitals; + + RowMajorMatrix grad_wa = + -2.0 * generalized_fock_mo.block(num_closed_orbitals, + num_closed_orbitals + num_open_orbitals, + num_open_orbitals, num_virtual_orbitals); + gradient.segment(offset, num_open_orbitals * num_virtual_orbitals) = + Eigen::Map(grad_wa.data(), grad_wa.size()).eval(); + offset += num_open_orbitals * num_virtual_orbitals; + + RowMajorMatrix grad_ia = + -2.0 * + generalized_fock_mo.block(0, num_closed_orbitals + num_open_orbitals, + num_closed_orbitals, num_virtual_orbitals); + gradient.segment(offset, num_closed_orbitals * num_virtual_orbitals) = + Eigen::Map(grad_ia.data(), grad_ia.size()).eval(); } /** @@ -85,29 +373,44 @@ class GDMLineFunctor { /** * @brief Bind functor to a specific SCF state for line search evaluations. * @param scf_impl Reference to `SCFImpl` used to evaluate trial densities - * @param C_pseudo_canonical Molecular orbitals in pseudo-canonical basis + * @param C_pseudo_canonical Pseudo-canonical molecular orbital coefficients * @param num_electrons Occupied orbital counts per spin component * @param rotation_offset Starting index for each spin's rotation slice * @param rotation_size Number of rotation parameters per spin (n_occ*n_virt) - * @param num_molecular_orbitals Total molecular orbitals in the system - * @param unrestricted Whether alpha/beta densities are treated separately + * @param scf_orbital_type Spin symmetry used across SCF algorithms */ GDMLineFunctor(const SCFImpl& scf_impl, const RowMajorMatrix& C_pseudo_canonical, const std::vector& num_electrons, const std::vector& rotation_offset, const std::vector& rotation_size, - int num_molecular_orbitals, bool unrestricted) + SCFOrbitalType scf_orbital_type) : scf_impl_(scf_impl), C_pseudo_canonical_(C_pseudo_canonical), num_electrons_(num_electrons), rotation_offset_(rotation_offset), rotation_size_(rotation_size), - num_density_matrices_(unrestricted ? 2 : 1), - num_molecular_orbitals_(num_molecular_orbitals), - unrestricted_(unrestricted), + num_atomic_orbitals_(scf_orbital_type == SCFOrbitalType::Unrestricted + ? static_cast(C_pseudo_canonical.rows()) / + 2 + : static_cast(C_pseudo_canonical.rows())), + num_molecular_orbitals_(static_cast(C_pseudo_canonical.cols())), + scf_orbital_type_(scf_orbital_type), cached_kappa_(Eigen::VectorXd()), - cached_energy_(std::numeric_limits::infinity()) {} + cached_J_(scf_orbital_type == SCFOrbitalType::RestrictedOpenShell + ? RowMajorMatrix::Zero(2 * num_atomic_orbitals_, + num_atomic_orbitals_) + : RowMajorMatrix()), + cached_K_(scf_orbital_type == SCFOrbitalType::RestrictedOpenShell + ? RowMajorMatrix::Zero(2 * num_atomic_orbitals_, + num_atomic_orbitals_) + : RowMajorMatrix()), + cached_energy_(std::numeric_limits::infinity()) { + if (scf_orbital_type == SCFOrbitalType::Unrestricted && + (C_pseudo_canonical.rows() % 2 != 0)) { + throw std::invalid_argument("In UHF or UKS, C row count must be even!"); + } + } /** * @brief Evaluate energy at given kappa vector x @@ -157,14 +460,23 @@ class GDMLineFunctor { const std::vector& rotation_size_; // Value parameters - const int num_density_matrices_; + const int num_atomic_orbitals_; const int num_molecular_orbitals_; - const bool unrestricted_; + const SCFOrbitalType scf_orbital_type_; + + int num_orbital_spin_blocks() const { + return scf_orbital_type_ == SCFOrbitalType::Unrestricted ? 2 : 1; + } + int num_density_matrices() const { + return scf_orbital_type_ == SCFOrbitalType::Restricted ? 1 : 2; + } // Cache for avoiding redundant Fock matrix computation Eigen::VectorXd cached_kappa_; // Cached kappa vector double cached_energy_; RowMajorMatrix cached_F_; // Needed for gradient computation + RowMajorMatrix cached_J_; // Needed for ROHF gradient computation + RowMajorMatrix cached_K_; // Needed for ROHF gradient computation RowMajorMatrix cached_C_; // For writing back to scf_impl RowMajorMatrix cached_P_; // For writing back to scf_impl }; @@ -181,35 +493,56 @@ double GDMLineFunctor::eval(const Eigen::VectorXd& x) { cached_C_ = C_pseudo_canonical_; // Apply rotation for all spins with kappa_trial - for (int i = 0; i < num_density_matrices_; i++) { - auto kappa_spin = - kappa_trial.segment(rotation_offset_[i], rotation_size_[i]); - apply_orbital_rotation(cached_C_, i, kappa_spin, num_electrons_[i], - num_molecular_orbitals_); + if (scf_orbital_type_ == SCFOrbitalType::RestrictedOpenShell) { + apply_restricted_open_shell_orbital_rotation(cached_C_, num_electrons_, + kappa_trial); + } else { + for (int i = 0; i < num_orbital_spin_blocks(); i++) { + auto kappa_spin = + kappa_trial.segment(rotation_offset_[i], rotation_size_[i]); + apply_restricted_unrestricted_orbital_rotation(cached_C_, i, kappa_spin, + num_electrons_[i], + num_orbital_spin_blocks()); + } } // Compute P_trial from rotated C (for all spins) cached_P_ = RowMajorMatrix::Zero( - num_density_matrices_ * num_molecular_orbitals_, num_molecular_orbitals_); + num_density_matrices() * num_atomic_orbitals_, num_atomic_orbitals_); - for (int i = 0; i < num_density_matrices_; i++) { + for (int i = 0; i < num_density_matrices(); i++) { const int num_occupied_orbitals = num_electrons_[i]; - const double occupation_factor = unrestricted_ ? 1.0 : 2.0; - - cached_P_.block(num_molecular_orbitals_ * i, 0, num_molecular_orbitals_, - num_molecular_orbitals_) = - occupation_factor * - cached_C_.block(num_molecular_orbitals_ * i, 0, num_molecular_orbitals_, - num_occupied_orbitals) * - cached_C_ - .block(num_molecular_orbitals_ * i, 0, num_molecular_orbitals_, - num_occupied_orbitals) - .transpose(); + const double occupation_factor = + (scf_orbital_type_ == SCFOrbitalType::Restricted) ? 2.0 : 1.0; + auto P_block = cached_P_.block(num_atomic_orbitals_ * i, 0, + num_atomic_orbitals_, num_atomic_orbitals_); + + // For ROHF, both alpha and beta densities share the same C block (row 0). + // For UHF, each spin has its own C block offset by num_atomic_orbitals_. + const int C_row_offset = + (scf_orbital_type_ == SCFOrbitalType::RestrictedOpenShell) + ? 0 + : num_atomic_orbitals_ * i; + const double* C_block_data = + cached_C_.data() + C_row_offset * num_molecular_orbitals_; + blas::gemm(blas::Layout::RowMajor, blas::Op::NoTrans, blas::Op::Trans, + num_atomic_orbitals_, num_atomic_orbitals_, + num_occupied_orbitals, occupation_factor, C_block_data, + num_molecular_orbitals_, C_block_data, num_molecular_orbitals_, + 0.0, P_block.data(), num_atomic_orbitals_); } // Evaluate energy and Fock matrix using trial density matrix - auto [energy, F_trial] = - scf_impl_.evaluate_trial_density_energy_and_fock(cached_P_); + double energy; + RowMajorMatrix F_trial; + if (scf_orbital_type_ == SCFOrbitalType::RestrictedOpenShell) { + std::tie(energy, F_trial) = + scf_impl_.evaluate_trial_density_energy_and_fock(cached_P_, cached_J_, + cached_K_); + } else { + std::tie(energy, F_trial) = + scf_impl_.evaluate_trial_density_energy_and_fock(cached_P_); + } // Cache all results for potential grad() call at same kappa cached_energy_ = energy; @@ -226,46 +559,15 @@ Eigen::VectorXd GDMLineFunctor::grad(const Eigen::VectorXd& x) { eval(x); } - // Initialize the full gradient vector (concatenated for all spins) - int total_rotation_size = 0; - for (int i = 0; i < num_density_matrices_; i++) { - total_rotation_size += rotation_size_[i]; - } - Eigen::VectorXd gradient = Eigen::VectorXd::Zero(total_rotation_size); - - // Compute gradient for each spin component - for (int i = 0; i < num_density_matrices_; i++) { - const int num_occupied_orbitals = num_electrons_[i]; - const int num_virtual_orbitals = - num_molecular_orbitals_ - num_occupied_orbitals; - - // Transform Fock matrix to MO basis: F_MO = C^T * F * C - RowMajorMatrix F_MO = - cached_C_ - .block(num_molecular_orbitals_ * i, 0, num_molecular_orbitals_, - num_molecular_orbitals_) - .transpose() * - cached_F_.block(num_molecular_orbitals_ * i, 0, num_molecular_orbitals_, - num_molecular_orbitals_) * - cached_C_.block(num_molecular_orbitals_ * i, 0, num_molecular_orbitals_, - num_molecular_orbitals_); - - // Extract occupied-virtual block and compute gradient - // The -4.0 before F_{ia} comes from derivative of energy w.r.t. kappa - // Reference: Helgaker, T., Jørgensen, P., & Olsen, J. (2000). Molecular - // electronic-structure theory, Eq. 10.8.34 (2013 reprint edition) - // -4.0 is for restricted closed-shell system. For unrestricted systems, the - // gradient is computed separately for each spin component, in that case the - // coefficient before F_{ia, spin} is -2.0 - RowMajorMatrix gradient_matrix = - -(unrestricted_ ? 2.0 : 4.0) * F_MO.block(0, num_occupied_orbitals, - num_occupied_orbitals, - num_virtual_orbitals); - - // Flatten matrix to vector and store in appropriate segment - gradient.segment(rotation_offset_[i], rotation_size_[i]) = - Eigen::Map(gradient_matrix.data(), - rotation_size_[i]); + Eigen::VectorXd gradient; + if (scf_orbital_type_ == SCFOrbitalType::RestrictedOpenShell) { + compute_restricted_open_shell_gradient(scf_impl_, cached_C_, cached_J_, + cached_K_, num_electrons_, + rotation_size_, gradient); + } else { + compute_restricted_unrestricted_gradient( + cached_F_, cached_C_, num_electrons_, rotation_offset_, rotation_size_, + num_orbital_spin_blocks(), gradient); } return gradient; @@ -311,25 +613,29 @@ class GDM { private: /** * @brief Transform history matrices (either history_dgrad or history_kappa) - * using current rotation matrices Uoo and Uvv to transform into the - * pseudo-canonical orbital basis, K_new = Uoo^T * K_old * Uvv + * using current rotation matrices to transform into the + * pseudo-canonical orbital basis, K_new = U_left^T * K_old * U_right * @param[in,out] history History matrix block to be transformed (either * history_dgrad or history_kappa) + * @param[in] u_left Left rotation matrix (e.g., Uii or Uaa) + * @param[in] u_right Right rotation matrix (e.g., Uaa or Uww) * @param[in] history_size Number of history entries - * @param[in] num_occupied_orbitals Number of electrons for current spin - * @param[in] num_molecular_orbitals Number of molecular orbitals + * @param[in] num_rows Number of rows in each unpacked history block + * @param[in] num_cols Number of cols in each unpacked history block * */ void transform_history_(Eigen::Block& history, - const int history_size, - const int num_occupied_orbitals, - const int num_molecular_orbitals); + const RowMajorMatrix& u_left, + const RowMajorMatrix& u_right, const int history_size, + const int num_rows, const int num_cols); /** * @brief Generate pseudo-canonical orbitals and apply transformations * @param[in] F Fock matrix in AO basis * @param[in,out] C Molecular orbital coefficient matrix * @param[in] spin_index Spin index (0 for alpha, 1 for beta) + * @param[in] num_orbital_spin_blocks Number of spin blocks (1 for RHF, 2 + * for UHF) * @param[in,out] history_kappa_spin Block reference to history kappa for this * spin * @param[in,out] history_dgrad_spin Block reference to history dgrad for this @@ -338,12 +644,57 @@ class GDM { * for this spin * */ - void generate_pseudo_canonical_orbital_( + void generate_restricted_unrestricted_pseudo_canonical_orbital_( const RowMajorMatrix& F, RowMajorMatrix& C, const int spin_index, + const int num_orbital_spin_blocks, Eigen::Block history_kappa_spin, Eigen::Block history_dgrad_spin, Eigen::VectorBlock current_gradient_spin); + /** + * @brief Build ROHF pseudo-canonical orbitals and transform ROHF history + * + * Diagonalizes closed/open/virtual MO blocks, rotates orbital coefficients + * to the pseudo-canonical basis, and transforms packed ROHF BFGS history and + * current gradient blocks (iw, wa, ia) into the updated basis. + * + * @param[in] F Spin-blocked Fock matrix in AO basis + * @param[in,out] C ROHF orbital coefficient matrix to rotate + * @param[in,out] history_kappa Packed ROHF kappa-history matrix + * @param[in,out] history_dgrad Packed ROHF gradient-difference history + * @param[in,out] current_gradient Packed ROHF current gradient vector + */ + void generate_restricted_open_shell_pseudo_canonical_orbital_( + const RowMajorMatrix& F, RowMajorMatrix& C, RowMajorMatrix& history_kappa, + RowMajorMatrix& history_dgrad, Eigen::VectorXd& current_gradient); + + /** + * @brief Build pseudo-canonical orbitals and initial Hessian across spins + * @param[in] F Fock matrix in AO basis + * @param[in,out] C Molecular orbital coefficient matrix + * @param[in] num_molecular_orbitals Total number of molecular orbitals + * @param[out] initial_hessian Output concatenated initial Hessian + */ + void build_restricted_unrestricted_pseudo_canonical_orbitals_hessian_( + const RowMajorMatrix& F, RowMajorMatrix& C, int num_molecular_orbitals, + Eigen::VectorXd& initial_hessian); + + /** + * @brief Build ROHF initial Hessian in pseudo-canonical basis + * + * Calls ROHF pseudo-canonical transformation, then fills the diagonal + * preconditioner for packed ROHF rotations (iw, wa, ia) using orbital-energy + * differences and the current absolute energy change. + * + * @param[in] F Spin-blocked Fock matrix in AO basis + * @param[in,out] C ROHF orbital coefficient matrix (updated in place) + * @param[in] num_molecular_orbitals Total number of molecular orbitals + * @param[out] initial_hessian Output ROHF initial Hessian vector + */ + void build_restricted_open_shell_pseudo_canonical_orbitals_hessian_( + const RowMajorMatrix& F, RowMajorMatrix& C, int num_molecular_orbitals, + Eigen::VectorXd& initial_hessian); + /// Reference to SCFContext const SCFContext& ctx_; ///< Reference to SCFContext /// Energy change from the last step @@ -378,19 +729,17 @@ class GDM { /// Eigenvalues of pseudo-canonical orbitals, used for building Hessian Eigen::VectorXd pseudo_canonical_eigenvalues_; - /// Horizontal rotation matrix of occupied orbitals - RowMajorMatrix Uoo_; - /// Horizontal rotation matrix of virtual orbitals - RowMajorMatrix Uvv_; - Eigen::VectorXd kappa_; // vertical rotation matrix of this step /// Energy of the last accepted step, used to decide if we rescale the kappa /// vector in this step double last_accepted_energy_; - int gdm_step_count_; // GDM iteration counter - int num_density_matrices_; // Number of density matrices (1 for restricted, 2 - // for unrestricted) + int gdm_step_count_; // GDM iteration counter + + /// Derive number of orbital spin blocks from context + int num_orbital_spin_blocks() const { + return ctx_.cfg->scf_orbital_type == SCFOrbitalType::Unrestricted ? 2 : 1; + } }; GDM::GDM(const SCFContext& ctx, int history_size_limit) @@ -424,36 +773,57 @@ GDM::GDM(const SCFContext& ctx, int history_size_limit) QDK_LOGGER().debug("GDM initialized with history_size_limit = {}", history_size_limit_); - num_density_matrices_ = - (cfg.scf_orbital_type == SCFOrbitalType::Unrestricted) ? 2 : 1; // Calculate rotation sizes for each spin - rotation_size_.resize(num_density_matrices_); - rotation_offset_.resize(num_density_matrices_); - - total_rotation_size_ = 0; - for (int spin_index = 0; spin_index < num_density_matrices_; spin_index++) { - const int num_occupied_orbitals = num_electrons_[spin_index]; - const int num_virtual_orbitals = - num_molecular_orbitals - num_occupied_orbitals; - // Validate dimensions (negative values indicate invalid input) - // Zero occupied or virtual orbitals is valid for unrestricted calculations - // (e.g., H atom has 0 beta electrons) - if (num_occupied_orbitals < 0) { + rotation_size_.resize(num_orbital_spin_blocks()); + rotation_offset_.resize(num_orbital_spin_blocks()); + + if (cfg.scf_orbital_type == SCFOrbitalType::RestrictedOpenShell) { + int num_closed_orbitals = num_electrons_[1]; + int num_open_orbitals = num_electrons_[0] - num_closed_orbitals; + int num_virtual_orbitals = num_molecular_orbitals - num_electrons_[0]; + if (num_closed_orbitals < 0 || num_open_orbitals < 0 || + num_virtual_orbitals < 0) { throw std::invalid_argument( - std::string("GDM: num_occupied_orbitals must be non-negative, got ") + - std::to_string(num_occupied_orbitals) + " for spin " + - std::to_string(spin_index)); + "Invalid ROHF system: " + "num_closed_orbitals=" + + std::to_string(num_closed_orbitals) + + ", num_open_orbitals=" + std::to_string(num_open_orbitals) + + ", num_virtual_orbitals=" + std::to_string(num_virtual_orbitals)); } - if (num_virtual_orbitals < 0) { - throw std::invalid_argument( - std::string("GDM: num_virtual_orbitals must be non-negative, got ") + - std::to_string(num_virtual_orbitals) + " for spin " + - std::to_string(spin_index)); + rotation_size_[0] = num_closed_orbitals * num_open_orbitals + + num_open_orbitals * num_virtual_orbitals + + num_closed_orbitals * num_virtual_orbitals; + rotation_offset_[0] = 0; + total_rotation_size_ = rotation_size_[0]; + } else { + total_rotation_size_ = 0; + for (int spin_index = 0; spin_index < num_orbital_spin_blocks(); + spin_index++) { + const int num_occupied_orbitals = num_electrons_[spin_index]; + const int num_virtual_orbitals = + num_molecular_orbitals - num_occupied_orbitals; + // Validate dimensions (negative values indicate invalid input) + // Zero occupied or virtual orbitals is valid for unrestricted + // calculations (e.g., H atom has 0 beta electrons) + if (num_occupied_orbitals < 0) { + throw std::invalid_argument( + std::string( + "GDM: num_occupied_orbitals must be non-negative, got ") + + std::to_string(num_occupied_orbitals) + " for spin " + + std::to_string(spin_index)); + } + if (num_virtual_orbitals < 0) { + throw std::invalid_argument( + std::string( + "GDM: num_virtual_orbitals must be non-negative, got ") + + std::to_string(num_virtual_orbitals) + " for spin " + + std::to_string(spin_index)); + } + rotation_size_[spin_index] = num_occupied_orbitals * num_virtual_orbitals; + rotation_offset_[spin_index] = total_rotation_size_; + total_rotation_size_ += rotation_size_[spin_index]; } - rotation_size_[spin_index] = num_occupied_orbitals * num_virtual_orbitals; - rotation_offset_[spin_index] = total_rotation_size_; - total_rotation_size_ += rotation_size_[spin_index]; } // Initialize concatenated matrices and vectors @@ -467,48 +837,126 @@ GDM::GDM(const SCFContext& ctx, int history_size_limit) } void GDM::transform_history_(Eigen::Block& history, - const int history_size, - const int num_occupied_orbitals, - const int num_molecular_orbitals) { + const RowMajorMatrix& u_left, + const RowMajorMatrix& u_right, + const int history_size, const int num_rows, + const int num_cols) { QDK_LOG_TRACE_ENTERING(); - const int num_virtual_orbitals = - num_molecular_orbitals - num_occupied_orbitals; // Validate dimensions (negative values indicate invalid input) - if (num_occupied_orbitals < 0 || num_virtual_orbitals < 0) { + if (num_rows < 0 || num_cols < 0) { throw std::invalid_argument( - std::string( - "transform_history_: invalid dimensions (num_occupied_orbitals=") + - std::to_string(num_occupied_orbitals) + - ", num_virtual_orbitals=" + std::to_string(num_virtual_orbitals) + ")"); + std::string("transform_history_: invalid dimensions (num_rows=") + + std::to_string(num_rows) + ", num_cols=" + std::to_string(num_cols) + + ")"); } - // Skip transformation if either dimension is zero (no rotations for this - // spin) - if (num_occupied_orbitals == 0 || num_virtual_orbitals == 0) { + // Skip transformation if either dimension is zero + if (num_rows == 0 || num_cols == 0) { return; } - RowMajorMatrix temp = - RowMajorMatrix::Zero(num_occupied_orbitals, num_virtual_orbitals); + RowMajorMatrix temp = RowMajorMatrix::Zero(num_rows, num_cols); for (int line = 0; line < history_size; line++) { double* history_line_ptr = history.row(line).data(); - // K_ov (new) = Uoo^T * K_ov * Uvv + // K_new = U_left^T * K_old * U_right blas::gemm(blas::Layout::RowMajor, blas::Op::NoTrans, blas::Op::NoTrans, - num_occupied_orbitals, num_virtual_orbitals, - num_virtual_orbitals, 1.0, history_line_ptr, - num_virtual_orbitals, Uvv_.data(), num_virtual_orbitals, 0.0, - temp.data(), num_virtual_orbitals); + num_rows, num_cols, num_cols, 1.0, history_line_ptr, num_cols, + u_right.data(), num_cols, 0.0, temp.data(), num_cols); blas::gemm(blas::Layout::RowMajor, blas::Op::Trans, blas::Op::NoTrans, - num_occupied_orbitals, num_virtual_orbitals, - num_occupied_orbitals, 1.0, Uoo_.data(), num_occupied_orbitals, - temp.data(), num_virtual_orbitals, 0.0, history_line_ptr, - num_virtual_orbitals); + num_rows, num_cols, num_rows, 1.0, u_left.data(), num_rows, + temp.data(), num_cols, 0.0, history_line_ptr, num_cols); } } -void GDM::generate_pseudo_canonical_orbital_( +/** + * @brief Diagonalize an MO sub-block and rotate the corresponding orbital block + * + * This helper computes eigenpairs of a symmetric MO-space sub-block using + * LAPACK `syev`, stores the resulting eigenvalues into a target slice, and + * applies the eigenvector rotation to the corresponding orbital columns: + * C_block <- C_block * U. + * + * @param[in,out] input_block_output_eigenvectors On input: symmetric block to + * diagonalize. On output: eigenvectors used for rotation. + * @param[in] block_size Dimension of the square sub-block to diagonalize. + * @param[in,out] eigenvalues Vector receiving eigenvalues for this block. + * @param[in] eigenvalue_start_index Start index in `eigenvalues` where this + * block's eigenvalues are written. + * @param[in,out] transformed_orbitals Orbital coefficient sub-block rotated in + * place. + * @param[in] num_atomic_orbitals Row count used for GEMM in the orbital + * rotation. + * @param[in] num_molecular_orbitals Leading dimension of the parent orbital + * coefficient matrix. + */ +static void calculate_pseudo_canonical_orbital_block( + RowMajorMatrix& input_block_output_eigenvectors, const int block_size, + Eigen::VectorXd& eigenvalues, const int eigenvalue_start_index, + Eigen::Block transformed_orbitals, + const int num_atomic_orbitals, const int num_molecular_orbitals) { + lapack::syev(lapack::Job::Vec, lapack::Uplo::Lower, block_size, + input_block_output_eigenvectors.data(), block_size, + eigenvalues.data() + eigenvalue_start_index); + input_block_output_eigenvectors.transposeInPlace(); + RowMajorMatrix copied_orbitals = transformed_orbitals; + blas::gemm(blas::Layout::RowMajor, blas::Op::NoTrans, blas::Op::NoTrans, + num_atomic_orbitals, block_size, block_size, 1.0, + copied_orbitals.data(), block_size, + input_block_output_eigenvectors.data(), block_size, 0.0, + transformed_orbitals.data(), num_molecular_orbitals); +} + +/** + * @brief Rotate a packed gradient sub-block to the pseudo-canonical basis + * + * Interprets a contiguous slice of current_vector (starting at start_index) + * as a num_rows x num_cols matrix G, applies G' = u_left^T * G * u_right, + * and writes the flattened result back into the corresponding slice of + * transformed_vector. + * + * This keeps the gradient representation aligned with the updated + * pseudo-canonical orbitals after block diagonalization and orbital rotation. + * + * @param[in] current_vector Source packed gradient vector. + * @param[in] start_index Start index of the sub-block inside the packed vector. + * @param[in] num_rows Row count of the unpacked gradient block. + * @param[in] num_cols Column count of the unpacked gradient block. + * @param[in] u_left Left rotation matrix. + * @param[in] u_right Right rotation matrix. + * @param[out] transformed_vector Destination packed vector receiving the + * transformed sub-block. + */ +static void rotate_gradient_to_pseudo_canonical_basis( + const Eigen::Ref& current_vector, + const int start_index, const int num_rows, const int num_cols, + const RowMajorMatrix& u_left, const RowMajorMatrix& u_right, + Eigen::Ref transformed_vector) { + RowMajorMatrix current_matrix = Eigen::Map( + current_vector.data() + start_index, num_rows, num_cols); + RowMajorMatrix temp_matrix = RowMajorMatrix::Zero(num_rows, num_cols); + RowMajorMatrix transformed_matrix = RowMajorMatrix::Zero(num_rows, num_cols); + + // temp_matrix = current_matrix * u_right + blas::gemm(blas::Layout::RowMajor, blas::Op::NoTrans, blas::Op::NoTrans, + num_rows, num_cols, num_cols, 1.0, current_matrix.data(), num_cols, + u_right.data(), num_cols, 0.0, temp_matrix.data(), num_cols); + + // transformed_matrix = u_left^T * temp_matrix + blas::gemm(blas::Layout::RowMajor, blas::Op::Trans, blas::Op::NoTrans, + num_rows, num_cols, num_rows, 1.0, u_left.data(), num_rows, + temp_matrix.data(), num_cols, 0.0, transformed_matrix.data(), + num_cols); + transformed_vector.segment(start_index, num_rows * num_cols) = + Eigen::Map(transformed_matrix.data(), + num_rows * num_cols); +} + +void GDM::generate_restricted_unrestricted_pseudo_canonical_orbital_( const RowMajorMatrix& F, RowMajorMatrix& C, const int spin_index, + const int num_orbital_spin_blocks, Eigen::Block history_kappa_spin, Eigen::Block history_dgrad_spin, Eigen::VectorBlock current_gradient_spin) { + const int num_atomic_orbitals = + (num_orbital_spin_blocks == 2) ? C.rows() / 2 : C.rows(); const int num_molecular_orbitals = C.cols(); const int num_occupied_orbitals = num_electrons_[spin_index]; const int num_virtual_orbitals = @@ -525,74 +973,272 @@ void GDM::generate_pseudo_canonical_orbital_( if (num_occupied_orbitals == 0 || num_virtual_orbitals == 0) { return; } - const int rotation_size = num_occupied_orbitals * num_virtual_orbitals; RowMajorMatrix F_MO = - C.block(num_molecular_orbitals * spin_index, 0, num_molecular_orbitals, - num_molecular_orbitals) - .transpose() * - F.block(num_molecular_orbitals * spin_index, 0, num_molecular_orbitals, - num_molecular_orbitals) * - C.block(num_molecular_orbitals * spin_index, 0, num_molecular_orbitals, - num_molecular_orbitals); + RowMajorMatrix::Zero(num_molecular_orbitals, num_molecular_orbitals); + const double* C_block_ptr = + C.data() + num_atomic_orbitals * num_molecular_orbitals * spin_index; + const double* F_block_ptr = + F.data() + num_atomic_orbitals * num_atomic_orbitals * spin_index; + std::vector atba_workspace(static_cast(num_atomic_orbitals) * + num_molecular_orbitals); + compute_atba_gemm(C_block_ptr, F_block_ptr, F_MO.data(), num_atomic_orbitals, + num_molecular_orbitals, atba_workspace); + + // Perform pseudo-canonical transformation + // Diagonalize occupied/virtual blocks and rotate orbitals to the + // pseudo-canonical basis. + RowMajorMatrix Uii = + F_MO.block(0, 0, num_occupied_orbitals, num_occupied_orbitals); + auto C_occ_view = C.block(num_atomic_orbitals * spin_index, 0, + num_atomic_orbitals, num_occupied_orbitals); + calculate_pseudo_canonical_orbital_block( + Uii, num_occupied_orbitals, pseudo_canonical_eigenvalues_, 0, C_occ_view, + num_atomic_orbitals, num_molecular_orbitals); + + RowMajorMatrix Uaa = F_MO.block(num_occupied_orbitals, num_occupied_orbitals, + num_virtual_orbitals, num_virtual_orbitals); + auto C_virt_view = + C.block(num_atomic_orbitals * spin_index, num_occupied_orbitals, + num_atomic_orbitals, num_virtual_orbitals); + calculate_pseudo_canonical_orbital_block( + Uaa, num_virtual_orbitals, pseudo_canonical_eigenvalues_, + num_occupied_orbitals, C_virt_view, num_atomic_orbitals, + num_molecular_orbitals); + + // Transform the vectors in history_kappa, history_dgrad, and + // current_gradient_spin to accommodate current pseudo-canonical orbitals + transform_history_(history_kappa_spin, Uii, Uaa, history_size_, + num_occupied_orbitals, num_virtual_orbitals); + transform_history_(history_dgrad_spin, Uii, Uaa, history_size_, + num_occupied_orbitals, num_virtual_orbitals); + + rotate_gradient_to_pseudo_canonical_basis( + current_gradient_spin, 0, num_occupied_orbitals, num_virtual_orbitals, + Uii, Uaa, current_gradient_spin); +} - // Perform pseudo-canonical transformation and BFGS - // Obtain pseudo-canonical orbitals. Foo and Fvv are symmetric matrices, but - // the output eigenvectors are column-major - Uoo_ = F_MO.block(0, 0, num_occupied_orbitals, num_occupied_orbitals); - Uvv_ = F_MO.block(num_occupied_orbitals, num_occupied_orbitals, - num_virtual_orbitals, num_virtual_orbitals); - - // Compute eigenvalues/eigenvectors of occupied-occupied and virtual-virtual - // blocks for pseudo-canonical orbital transformation - lapack::syev(lapack::Job::Vec, lapack::Uplo::Lower, num_occupied_orbitals, - Uoo_.data(), num_occupied_orbitals, - pseudo_canonical_eigenvalues_.data()); - lapack::syev(lapack::Job::Vec, lapack::Uplo::Lower, num_virtual_orbitals, - Uvv_.data(), num_virtual_orbitals, - pseudo_canonical_eigenvalues_.data() + num_occupied_orbitals); - - // Transpose to convert column-major eigenvectors to row-major format - Uoo_.transposeInPlace(); - Uvv_.transposeInPlace(); - - // Transform occupied orbitals - auto C_occ_view = C.block(num_molecular_orbitals * spin_index, 0, - num_molecular_orbitals, num_occupied_orbitals); - RowMajorMatrix C_occ = C_occ_view.eval(); - blas::gemm(blas::Layout::RowMajor, blas::Op::NoTrans, blas::Op::NoTrans, - num_molecular_orbitals, num_occupied_orbitals, - num_occupied_orbitals, 1.0, C_occ.data(), num_occupied_orbitals, - Uoo_.data(), num_occupied_orbitals, 0.0, C_occ_view.data(), - num_molecular_orbitals); +void GDM::build_restricted_unrestricted_pseudo_canonical_orbitals_hessian_( + const RowMajorMatrix& F, RowMajorMatrix& C, int num_molecular_orbitals, + Eigen::VectorXd& initial_hessian) { + initial_hessian.setZero(total_rotation_size_); - // Transform virtual orbitals - auto C_virt_view = - C.block(num_molecular_orbitals * spin_index, num_occupied_orbitals, - num_molecular_orbitals, num_virtual_orbitals); - RowMajorMatrix C_virt = C_virt_view.eval(); - blas::gemm(blas::Layout::RowMajor, blas::Op::NoTrans, blas::Op::NoTrans, - num_molecular_orbitals, num_virtual_orbitals, num_virtual_orbitals, - 1.0, C_virt.data(), num_virtual_orbitals, Uvv_.data(), - num_virtual_orbitals, 0.0, C_virt_view.data(), - num_molecular_orbitals); - - // Transform the vectors in history_kappa and history_dgrad to - // accommodate current pseudo-canonical orbitals - transform_history_(history_kappa_spin, history_size_, num_occupied_orbitals, - num_molecular_orbitals); - transform_history_(history_dgrad_spin, history_size_, num_occupied_orbitals, - num_molecular_orbitals); - - // Transform the gradient to accommodate current pseudo-canonical orbitals - RowMajorMatrix current_gradient_matrix = - Eigen::Map(current_gradient_spin.data(), - num_occupied_orbitals, num_virtual_orbitals); - RowMajorMatrix current_gradient_transformed_matrix = - Uoo_.transpose() * current_gradient_matrix * Uvv_; - Eigen::VectorXd current_gradient_transformed = Eigen::Map( - current_gradient_transformed_matrix.data(), rotation_size); - current_gradient_spin = current_gradient_transformed; + for (int i = 0; i < num_orbital_spin_blocks(); ++i) { + const int num_occupied_orbitals = num_electrons_[i]; + const int num_virtual_orbitals = + num_molecular_orbitals - num_occupied_orbitals; + + auto history_kappa_spin = history_kappa_.block( + 0, rotation_offset_[i], history_size_limit_, rotation_size_[i]); + auto history_dgrad_spin = history_dgrad_.block( + 0, rotation_offset_[i], history_size_limit_, rotation_size_[i]); + auto current_gradient_spin = + current_gradient_.segment(rotation_offset_[i], rotation_size_[i]); + + // Generate pseudo-canonical orbitals and transform gradient and history + generate_restricted_unrestricted_pseudo_canonical_orbital_( + F, C, i, num_orbital_spin_blocks(), history_kappa_spin, + history_dgrad_spin, current_gradient_spin); + + // Build this spin's segment of initial Hessian + // Reference: Helgaker, T., Jorgensen, P., & Olsen, J. (2000). Molecular + // electronic-structure theory, Eq. 10.8.56 (2013 reprint edition) + // 4.0 is for restricted closed-shell system. For unrestricted systems, the + // gradient is computed separately for each spin component, in that case the + // coefficient should be 2.0 + double initial_hessian_coeff = (num_orbital_spin_blocks() == 2) ? 2.0 : 4.0; + for (int j = 0; j < num_occupied_orbitals; j++) { + for (int v = 0; v < num_virtual_orbitals; v++) { + int index = rotation_offset_[i] + j * num_virtual_orbitals + v; + double pseudo_canonical_energy_diff = + std::abs(pseudo_canonical_eigenvalues_(num_occupied_orbitals + v) - + pseudo_canonical_eigenvalues_(j)); + initial_hessian(index) = + std::max(initial_hessian_coeff * (std::abs(delta_energy_) + + pseudo_canonical_energy_diff), + nonpositive_threshold_); + } + } + } +} + +void GDM::generate_restricted_open_shell_pseudo_canonical_orbital_( + const RowMajorMatrix& F, RowMajorMatrix& C, RowMajorMatrix& history_kappa, + RowMajorMatrix& history_dgrad, Eigen::VectorXd& current_gradient) { + const int num_molecular_orbitals = C.cols(); + const int num_atomic_orbitals = C.rows(); + const int num_closed_orbitals = num_electrons_[1]; + const int num_open_orbitals = num_electrons_[0] - num_closed_orbitals; + const int num_occupied_orbitals = num_electrons_[0]; + const int num_virtual_orbitals = num_molecular_orbitals - num_electrons_[0]; + const int size_closed_open = num_closed_orbitals * num_open_orbitals; + const int size_open_virtual = num_open_orbitals * num_virtual_orbitals; + const int size_closed_virtual = num_closed_orbitals * num_virtual_orbitals; + + RowMajorMatrix F_up_mo = + RowMajorMatrix::Zero(num_molecular_orbitals, num_molecular_orbitals); + RowMajorMatrix F_dn_mo = F_up_mo; + + const double* C_block_ptr = C.data(); + const double* F_up_block_ptr = F.data(); + const double* F_dn_block_ptr = + F.data() + num_atomic_orbitals * num_atomic_orbitals; + std::vector atba_workspace(static_cast(num_atomic_orbitals) * + num_molecular_orbitals); + + compute_atba_gemm(C_block_ptr, F_up_block_ptr, F_up_mo.data(), + num_atomic_orbitals, num_molecular_orbitals, + atba_workspace); + compute_atba_gemm(C_block_ptr, F_dn_block_ptr, F_dn_mo.data(), + num_atomic_orbitals, num_molecular_orbitals, + atba_workspace); + + // Perform pseudo-canonical transformation + // Diagonalize occupied/virtual blocks and rotate orbitals to the + // pseudo-canonical basis. + RowMajorMatrix Uii = + 0.5 * (F_up_mo.block(0, 0, num_closed_orbitals, num_closed_orbitals) + + F_dn_mo.block(0, 0, num_closed_orbitals, num_closed_orbitals)); + auto C_closed_view = C.block(0, 0, num_atomic_orbitals, num_closed_orbitals); + calculate_pseudo_canonical_orbital_block( + Uii, num_closed_orbitals, pseudo_canonical_eigenvalues_, 0, C_closed_view, + num_atomic_orbitals, num_molecular_orbitals); + + RowMajorMatrix Uww = + 0.5 * (F_up_mo.block(num_closed_orbitals, num_closed_orbitals, + num_open_orbitals, num_open_orbitals) + + F_dn_mo.block(num_closed_orbitals, num_closed_orbitals, + num_open_orbitals, num_open_orbitals)); + auto C_open_view = + C.block(0, num_closed_orbitals, num_atomic_orbitals, num_open_orbitals); + calculate_pseudo_canonical_orbital_block( + Uww, num_open_orbitals, pseudo_canonical_eigenvalues_, + num_closed_orbitals, C_open_view, num_atomic_orbitals, + num_molecular_orbitals); + + RowMajorMatrix Uaa = + 0.5 * (F_up_mo.block(num_occupied_orbitals, num_occupied_orbitals, + num_virtual_orbitals, num_virtual_orbitals) + + F_dn_mo.block(num_occupied_orbitals, num_occupied_orbitals, + num_virtual_orbitals, num_virtual_orbitals)); + auto C_virt_view = C.block(0, num_occupied_orbitals, num_atomic_orbitals, + num_virtual_orbitals); + calculate_pseudo_canonical_orbital_block( + Uaa, num_virtual_orbitals, pseudo_canonical_eigenvalues_, + num_occupied_orbitals, C_virt_view, num_atomic_orbitals, + num_molecular_orbitals); + + // Transform the vectors in history_kappa, history_dgrad, and current_gradient + // to accommodate current pseudo-canonical orbitals + int offset = 0; + auto history_kappa_iw = + history_kappa.block(0, 0, history_size_, size_closed_open); + transform_history_(history_kappa_iw, Uii, Uww, history_size_, + num_closed_orbitals, num_open_orbitals); + auto history_dgrad_iw = + history_dgrad.block(0, 0, history_size_, size_closed_open); + transform_history_(history_dgrad_iw, Uii, Uww, history_size_, + num_closed_orbitals, num_open_orbitals); + offset += size_closed_open; + + auto history_kappa_wa = + history_kappa.block(0, offset, history_size_, size_open_virtual); + transform_history_(history_kappa_wa, Uww, Uaa, history_size_, + num_open_orbitals, num_virtual_orbitals); + auto history_dgrad_wa = + history_dgrad.block(0, offset, history_size_, size_open_virtual); + transform_history_(history_dgrad_wa, Uww, Uaa, history_size_, + num_open_orbitals, num_virtual_orbitals); + offset += size_open_virtual; + + auto history_kappa_ia = + history_kappa.block(0, offset, history_size_, size_closed_virtual); + transform_history_(history_kappa_ia, Uii, Uaa, history_size_, + num_closed_orbitals, num_virtual_orbitals); + auto history_dgrad_ia = + history_dgrad.block(0, offset, history_size_, size_closed_virtual); + transform_history_(history_dgrad_ia, Uii, Uaa, history_size_, + num_closed_orbitals, num_virtual_orbitals); + + Eigen::VectorXd current_gradient_transformed = + Eigen::VectorXd(current_gradient.size()); + + offset = 0; + rotate_gradient_to_pseudo_canonical_basis( + current_gradient, offset, num_closed_orbitals, num_open_orbitals, Uii, + Uww, current_gradient_transformed); + offset += size_closed_open; + + rotate_gradient_to_pseudo_canonical_basis( + current_gradient, offset, num_open_orbitals, num_virtual_orbitals, Uww, + Uaa, current_gradient_transformed); + offset += size_open_virtual; + + rotate_gradient_to_pseudo_canonical_basis( + current_gradient, offset, num_closed_orbitals, num_virtual_orbitals, Uii, + Uaa, current_gradient_transformed); + + current_gradient = current_gradient_transformed; +} + +void GDM::build_restricted_open_shell_pseudo_canonical_orbitals_hessian_( + const RowMajorMatrix& F, RowMajorMatrix& C, int num_molecular_orbitals, + Eigen::VectorXd& initial_hessian) { + initial_hessian.setZero(total_rotation_size_); + + generate_restricted_open_shell_pseudo_canonical_orbital_( + F, C, history_kappa_, history_dgrad_, current_gradient_); + + const int num_closed_orbitals = num_electrons_[1]; + const int num_open_orbitals = num_electrons_[0] - num_closed_orbitals; + const int num_virtual_orbitals = num_molecular_orbitals - num_electrons_[0]; + + // between closed and open shells, or between open and virtual orbitals, the + // coefficient should be 2.0 + double initial_hessian_coeff = 2.0; + int offset = 0; + for (int j = 0; j < num_closed_orbitals; j++) { + for (int v = 0; v < num_open_orbitals; v++) { + int index = j * num_open_orbitals + v; + double pseudo_canonical_energy_diff = + std::abs(pseudo_canonical_eigenvalues_(num_closed_orbitals + v) - + pseudo_canonical_eigenvalues_(j)); + initial_hessian(index) = + std::max(initial_hessian_coeff * + (std::abs(delta_energy_) + pseudo_canonical_energy_diff), + nonpositive_threshold_); + } + } + offset += num_closed_orbitals * num_open_orbitals; + + for (int v = 0; v < num_open_orbitals; v++) { + for (int a = 0; a < num_virtual_orbitals; a++) { + int index = offset + v * num_virtual_orbitals + a; + double pseudo_canonical_energy_diff = + std::abs(pseudo_canonical_eigenvalues_(num_electrons_[0] + a) - + pseudo_canonical_eigenvalues_(num_closed_orbitals + v)); + initial_hessian(index) = + std::max(initial_hessian_coeff * + (std::abs(delta_energy_) + pseudo_canonical_energy_diff), + nonpositive_threshold_); + } + } + offset += num_open_orbitals * num_virtual_orbitals; + + // between closed and virtual orbitals, the coefficient should be 4.0 + initial_hessian_coeff = 4.0; + for (int j = 0; j < num_closed_orbitals; j++) { + for (int a = 0; a < num_virtual_orbitals; a++) { + int index = offset + j * num_virtual_orbitals + a; + double pseudo_canonical_energy_diff = + std::abs(pseudo_canonical_eigenvalues_(num_electrons_[0] + a) - + pseudo_canonical_eigenvalues_(j)); + initial_hessian(index) = + std::max(initial_hessian_coeff * + (std::abs(delta_energy_) + pseudo_canonical_energy_diff), + nonpositive_threshold_); + } + } } void GDM::iterate(SCFImpl& scf_impl) { @@ -601,10 +1247,15 @@ void GDM::iterate(SCFImpl& scf_impl) { const auto& F = scf_impl.get_fock_matrix(); const auto* cfg = ctx_.cfg; + // TODO: This function is called within iterate_ which runs only on rank 0, + // but build_jk_matrices requires participation from all MPI ranks. + if (cfg->mpi.world_size > 1) { + throw std::runtime_error( + "Temporary limitation: build_jk_matrices is not " + "supported with MPI world_size > 1."); + } const int num_molecular_orbitals = static_cast(ctx_.num_molecular_orbitals); - const int num_density_matrices = - (cfg->scf_orbital_type == SCFOrbitalType::Unrestricted) ? 2 : 1; // Check if there are any virtual orbitals for any spin component // If not, orbital rotation is not possible and we should skip GDM iteration @@ -615,44 +1266,25 @@ void GDM::iterate(SCFImpl& scf_impl) { return; } - // Compute current gradient and dgrad for each spin - for (int i = 0; i < num_density_matrices; ++i) { - const int num_occupied_orbitals = num_electrons_[i]; - const int num_virtual_orbitals = - num_molecular_orbitals - num_occupied_orbitals; - const int rotation_size = num_occupied_orbitals * num_virtual_orbitals; - RowMajorMatrix F_MO = - C.block(num_molecular_orbitals * i, 0, num_molecular_orbitals, - num_molecular_orbitals) - .transpose() * - F.block(num_molecular_orbitals * i, 0, num_molecular_orbitals, - num_molecular_orbitals) * - C.block(num_molecular_orbitals * i, 0, num_molecular_orbitals, - num_molecular_orbitals); + if (cfg->scf_orbital_type == SCFOrbitalType::RestrictedOpenShell) { + const int num_atomic_orbitals = scf_impl.get_num_atomic_orbitals(); + RowMajorMatrix J_ao = + RowMajorMatrix::Zero(2 * num_atomic_orbitals, num_atomic_orbitals); + RowMajorMatrix K_ao = + RowMajorMatrix::Zero(2 * num_atomic_orbitals, num_atomic_orbitals); + scf_impl.build_jk_matrices(scf_impl.get_density_matrix(), J_ao, K_ao); + compute_restricted_open_shell_gradient(scf_impl, C, J_ao, K_ao, + num_electrons_, rotation_size_, + current_gradient_); + } else { + compute_restricted_unrestricted_gradient( + F, C, num_electrons_, rotation_offset_, rotation_size_, + num_orbital_spin_blocks(), current_gradient_); + } - // Extract occupied-virtual block and compute gradient - // The -4.0 before F_{ia} comes from derivative of energy w.r.t. kappa - // Reference: Helgaker, T., Jørgensen, P., & Olsen, J. (2000). Molecular - // electronic-structure theory, Eq. 10.8.34 (2013 reprint edition) - // -4.0 is for restricted closed-shell system. For unrestricted systems, the - // gradient is computed separately for each spin component, in that case the - // coefficient before F_{ia, spin} is -2.0 - RowMajorMatrix current_gradient_matrix = - -((cfg->scf_orbital_type == SCFOrbitalType::Unrestricted) ? 2.0 : 4.0) * - F_MO.block(0, num_occupied_orbitals, num_occupied_orbitals, - num_virtual_orbitals); - current_gradient_.segment(rotation_offset_[i], rotation_size_[i]) = - Eigen::Map(current_gradient_matrix.data(), - rotation_size); - - if (gdm_step_count_ != 0) { - // Add new gradient difference to history for this spin - history_dgrad_ - .block(0, rotation_offset_[i], history_size_limit_, rotation_size_[i]) - .row(history_size_) = - current_gradient_.segment(rotation_offset_[i], rotation_size_[i]) - - previous_gradient_.segment(rotation_offset_[i], rotation_size_[i]); - } + if (gdm_step_count_ != 0) { + // Add new gradient difference to history for all spins + history_dgrad_.row(history_size_) = current_gradient_ - previous_gradient_; } // Update history size and manage history overflow. History for both spins are @@ -674,45 +1306,13 @@ void GDM::iterate(SCFImpl& scf_impl) { } } - // Build concatenated initial Hessian for all spins - Eigen::VectorXd initial_hessian = Eigen::VectorXd::Zero(total_rotation_size_); - - for (int i = 0; i < num_density_matrices; ++i) { - const int num_occupied_orbitals = num_electrons_[i]; - const int num_virtual_orbitals = - num_molecular_orbitals - num_occupied_orbitals; - - auto history_kappa_spin = history_kappa_.block( - 0, rotation_offset_[i], history_size_limit_, rotation_size_[i]); - auto history_dgrad_spin = history_dgrad_.block( - 0, rotation_offset_[i], history_size_limit_, rotation_size_[i]); - auto current_gradient_spin = - current_gradient_.segment(rotation_offset_[i], rotation_size_[i]); - - // Generate pseudo-canonical orbitals and transform gradient and history - generate_pseudo_canonical_orbital_( - F, C, i, history_kappa_spin, history_dgrad_spin, current_gradient_spin); - - // Build this spin's segment of initial Hessian - // Reference: Helgaker, T., Jørgensen, P., & Olsen, J. (2000). Molecular - // electronic-structure theory, Eq. 10.8.56 (2013 reprint edition) - // 4.0 is for restricted closed-shell system. For unrestricted systems, the - // gradient is computed separately for each spin component, in that case the - // coefficient should be 2.0 - double initial_hessian_coeff = - (cfg->scf_orbital_type == SCFOrbitalType::Unrestricted) ? 2.0 : 4.0; - for (int j = 0; j < num_occupied_orbitals; j++) { - for (int v = 0; v < num_virtual_orbitals; v++) { - int index = rotation_offset_[i] + j * num_virtual_orbitals + v; - double pseudo_canonical_energy_diff = - std::abs(pseudo_canonical_eigenvalues_(num_occupied_orbitals + v) - - pseudo_canonical_eigenvalues_(j)); - initial_hessian(index) = - std::max(initial_hessian_coeff * (std::abs(delta_energy_) + - pseudo_canonical_energy_diff), - nonpositive_threshold_); - } - } + Eigen::VectorXd initial_hessian; + if (cfg->scf_orbital_type == SCFOrbitalType::RestrictedOpenShell) { + build_restricted_open_shell_pseudo_canonical_orbitals_hessian_( + F, C, num_molecular_orbitals, initial_hessian); + } else { + build_restricted_unrestricted_pseudo_canonical_orbitals_hessian_( + F, C, num_molecular_orbitals, initial_hessian); } double latest_inverse_rho = 1.0; @@ -828,10 +1428,9 @@ void GDM::iterate(SCFImpl& scf_impl) { RowMajorMatrix C_pseudo_canonical = C.eval(); // Create line search functor for energy evaluation - GDMLineFunctor line_functor( - scf_impl, C_pseudo_canonical, num_electrons_, rotation_offset_, - rotation_size_, num_molecular_orbitals, - cfg->scf_orbital_type == SCFOrbitalType::Unrestricted); + GDMLineFunctor line_functor(scf_impl, C_pseudo_canonical, num_electrons_, + rotation_offset_, rotation_size_, + cfg->scf_orbital_type); Eigen::VectorXd start_kappa = Eigen::VectorXd::Zero(kappa_.size()); Eigen::VectorXd kappa_dir = kappa_; // Search direction diff --git a/cpp/src/qdk/chemistry/algorithms/microsoft/scf/src/scf_algorithm/line_search.h b/cpp/src/qdk/chemistry/algorithms/microsoft/scf/src/scf_algorithm/line_search.h index c3bd239fa..0dad7721e 100644 --- a/cpp/src/qdk/chemistry/algorithms/microsoft/scf/src/scf_algorithm/line_search.h +++ b/cpp/src/qdk/chemistry/algorithms/microsoft/scf/src/scf_algorithm/line_search.h @@ -9,7 +9,9 @@ #include #include -namespace qdk::chemistry::scf::impl { +namespace qdk::chemistry::scf { + +namespace impl { /** * @brief Nocedal-Wright line search with strong Wolfe conditions. @@ -154,4 +156,6 @@ void nocedal_wright_line_search(Functor& op, if (!converged) throw std::runtime_error("Line Search Failed"); } -} // namespace qdk::chemistry::scf::impl +} // namespace impl + +} // namespace qdk::chemistry::scf diff --git a/cpp/src/qdk/chemistry/algorithms/microsoft/scf/src/scf_algorithm/scf_algorithm.cpp b/cpp/src/qdk/chemistry/algorithms/microsoft/scf/src/scf_algorithm/scf_algorithm.cpp index f1f4250ff..2613a9d15 100644 --- a/cpp/src/qdk/chemistry/algorithms/microsoft/scf/src/scf_algorithm/scf_algorithm.cpp +++ b/cpp/src/qdk/chemistry/algorithms/microsoft/scf/src/scf_algorithm/scf_algorithm.cpp @@ -7,9 +7,11 @@ #include #include +#include #include #include #include +#include #include "../scf/scf_impl.h" #ifdef QDK_CHEMISTRY_ENABLE_GPU @@ -30,8 +32,44 @@ #include #endif +#include + namespace qdk::chemistry::scf { +void compute_atba_gemm(const double* A, const double* B, double* C, int m, + int n, std::vector& workspace, + blas::Layout layout) { + if (A == nullptr || B == nullptr || C == nullptr) { + throw std::invalid_argument("compute_atba_gemm: null matrix pointer."); + } + if (m < 0 || n < 0) { + throw std::invalid_argument("compute_atba_gemm: negative dimensions."); + } + if (m == 0 || n == 0) { + return; + } + + const size_t required_workspace_size = static_cast(m) * n; + if (workspace.size() < required_workspace_size) { + throw std::invalid_argument( + "compute_atba_gemm: workspace is smaller than m * n."); + } + + const int lda = (layout == blas::Layout::RowMajor) ? n : m; + const int ldb = m; + const int ldc = n; + + const int ld_workspace = (layout == blas::Layout::RowMajor) ? n : m; + + // workspace = B * A + blas::gemm(layout, blas::Op::NoTrans, blas::Op::NoTrans, m, n, m, 1.0, B, ldb, + A, lda, 0.0, workspace.data(), ld_workspace); + + // C = A^T * workspace + blas::gemm(layout, blas::Op::Trans, blas::Op::NoTrans, n, n, m, 1.0, A, lda, + workspace.data(), ld_workspace, 0.0, C, ldc); +} + SCFAlgorithm::SCFAlgorithm(const SCFContext& ctx) : ctx_(ctx), step_count_(0), @@ -47,6 +85,13 @@ SCFAlgorithm::SCFAlgorithm(const SCFContext& ctx) : 1; P_last_ = RowMajorMatrix::Zero(num_density_matrices * num_atomic_orbitals, num_atomic_orbitals); + + if (ctx.cfg->scf_orbital_type == SCFOrbitalType::RestrictedOpenShell) { + rohf_effective_fock_ = + RowMajorMatrix::Zero(num_atomic_orbitals, num_atomic_orbitals); + rohf_total_density_ = + RowMajorMatrix::Zero(num_atomic_orbitals, num_atomic_orbitals); + } } SCFAlgorithm::~SCFAlgorithm() noexcept = default; @@ -60,7 +105,7 @@ std::shared_ptr SCFAlgorithm::create(const SCFContext& ctx) { switch (cfg.scf_algorithm.method) { case SCFAlgorithmName::ASAHF: if (rohf_enabled) { - throw std::runtime_error("ROHF-enabled ASAHF is not supported!"); + throw std::invalid_argument("ROHF-enabled ASAHF is not supported!"); } return std::make_shared( ctx, cfg.scf_algorithm.diis_subspace_size); @@ -69,15 +114,9 @@ std::shared_ptr SCFAlgorithm::create(const SCFContext& ctx) { return std::make_shared(ctx, cfg.scf_algorithm.diis_subspace_size); case SCFAlgorithmName::GDM: - if (rohf_enabled) { - throw std::runtime_error("ROHF-enabled GDM is not supported!"); - } return std::make_shared(ctx, cfg.scf_algorithm.gdm_config); case SCFAlgorithmName::DIIS_GDM: - if (rohf_enabled) { - throw std::runtime_error("ROHF-enabled DIIS_GDM is not supported!"); - } return std::make_shared(ctx, cfg.scf_algorithm.diis_subspace_size, cfg.scf_algorithm.gdm_config); @@ -167,6 +206,36 @@ void SCFAlgorithm::update_density_matrix(RowMajorMatrix& P, bool unrestricted, int nelec_alpha, int nelec_beta) { QDK_LOG_TRACE_ENTERING(); + if (ctx_.cfg->scf_orbital_type == SCFOrbitalType::RestrictedOpenShell) { + const int num_atomic_orbitals = + static_cast(ctx_.basis_set->num_atomic_orbitals); + if (C.rows() != num_atomic_orbitals) { + throw std::invalid_argument( + "ROHF coefficient matrix row count does not match AO dimension"); + } + if (P.rows() != 2 * num_atomic_orbitals || + P.cols() != num_atomic_orbitals) { + throw std::invalid_argument( + "ROHF density matrix must contain alpha and beta AO blocks"); + } + + auto build_density_block = [&](auto&& target, int n_occ) { + if (n_occ <= 0) { + target.setZero(); + return; + } + target.noalias() = C.block(0, 0, num_atomic_orbitals, n_occ) * + C.block(0, 0, num_atomic_orbitals, n_occ).transpose(); + }; + + auto P_alpha = P.block(0, 0, num_atomic_orbitals, num_atomic_orbitals); + auto P_beta = P.block(num_atomic_orbitals, 0, num_atomic_orbitals, + num_atomic_orbitals); + build_density_block(P_alpha, nelec_alpha); + build_density_block(P_beta, nelec_beta); + return; + } + const int num_orbital_sets = unrestricted ? 2 : 1; const int num_atomic_orbitals = static_cast(ctx_.basis_set->num_atomic_orbitals); @@ -176,8 +245,6 @@ void SCFAlgorithm::update_density_matrix(RowMajorMatrix& P, "Coefficient matrix rows do not match orbital set count"); } - // For ASAHF and ROHF, the density matrix construction is different and - // will be handled in the overridden methods const double occupancy_factor = unrestricted ? 1.0 : 2.0; for (int i = 0; i < num_orbital_sets; ++i) { const int n_occ = (i == 0) ? nelec_alpha : nelec_beta; @@ -194,6 +261,176 @@ void SCFAlgorithm::update_density_matrix(RowMajorMatrix& P, } } +std::pair +SCFAlgorithm::build_rohf_convergence_matrices(const SCFImpl& scf_impl) { + QDK_LOG_TRACE_ENTERING(); + if (ctx_.cfg->scf_orbital_type != SCFOrbitalType::RestrictedOpenShell) { + throw std::logic_error( + "ROHF convergence matrices are only available for " + "RestrictedOpenShell calculations"); + } + + const auto nelec_vec = scf_impl.get_num_electrons(); + build_rohf_f_p_matrix( + scf_impl.get_fock_matrix(), scf_impl.get_orbitals_matrix(), + scf_impl.get_density_matrix(), nelec_vec[0], nelec_vec[1], + rohf_effective_fock_, rohf_total_density_); + + return {get_rohf_convergence_fock_matrix(), + get_rohf_convergence_density_matrix()}; +} + +void SCFAlgorithm::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) { + QDK_LOG_TRACE_ENTERING(); + const int num_atomic_orbitals = static_cast(C.rows()); + const int num_molecular_orbitals = static_cast(C.cols()); + if (num_atomic_orbitals != num_molecular_orbitals) { + throw std::invalid_argument( + "ROHF build requires number of atomic orbitals to equal number of " + "molecular orbitals!"); + } + + total_density = + P.block(0, 0, num_atomic_orbitals, num_atomic_orbitals) + + P.block(num_atomic_orbitals, 0, num_atomic_orbitals, num_atomic_orbitals); + + if (effective_fock.rows() != num_atomic_orbitals || + effective_fock.cols() != num_atomic_orbitals) { + effective_fock = + RowMajorMatrix::Zero(num_atomic_orbitals, num_atomic_orbitals); + } + + if (C.isZero()) { + effective_fock.noalias() = + F.block(0, 0, num_atomic_orbitals, num_atomic_orbitals); + return; + } + + RowMajorMatrix F_up_mo = + RowMajorMatrix::Zero(num_molecular_orbitals, num_molecular_orbitals); + RowMajorMatrix F_dn_mo = F_up_mo; + RowMajorMatrix effective_F_mo = F_up_mo; + + const double* C_block_ptr = C.data(); + const double* F_up_block_ptr = F.data(); + const double* F_dn_block_ptr = + F.data() + num_atomic_orbitals * num_atomic_orbitals; + std::vector atba_workspace(static_cast(num_atomic_orbitals) * + num_molecular_orbitals); + compute_atba_gemm(C_block_ptr, F_up_block_ptr, F_up_mo.data(), + num_atomic_orbitals, num_molecular_orbitals, atba_workspace, + blas::Layout::RowMajor); + compute_atba_gemm(C_block_ptr, F_dn_block_ptr, F_dn_mo.data(), + num_atomic_orbitals, num_molecular_orbitals, atba_workspace, + blas::Layout::RowMajor); + + auto average_block = [&effective_F_mo, &F_up_mo, &F_dn_mo]( + int row, int col, int rows, int cols) { + if (rows <= 0 || cols <= 0) return; + effective_F_mo.block(row, col, rows, cols).noalias() = + 0.5 * (F_up_mo.block(row, col, rows, cols) + + F_dn_mo.block(row, col, rows, cols)); + }; + auto copy_block = [&effective_F_mo](const RowMajorMatrix& src, int row, + int col, int rows, int cols) { + if (rows <= 0 || cols <= 0) return; + effective_F_mo.block(row, col, rows, cols) = + src.block(row, col, rows, cols); + }; + + const int nd = nelec_beta; + const int ns = nelec_alpha - nelec_beta; + const int nv = num_molecular_orbitals - nelec_alpha; + + average_block(0, 0, nd, nd); + average_block(0, nd + ns, nd, nv); + average_block(nd + ns, 0, nv, nd); + average_block(nd + ns, nd + ns, nv, nv); + average_block(nd, nd, ns, ns); + copy_block(F_dn_mo, 0, nd, nd, ns); + copy_block(F_dn_mo, nd, 0, ns, nd); + copy_block(F_up_mo, nd, nd + ns, ns, nv); + copy_block(F_up_mo, nd + ns, nd, nv, ns); + + // Transform the effective Fock matrix back to AO basis by solving + // C^{-T} * F_MO * C^{-1} = F_AO + // We use LAPACK's getrf/getrs to solve the linear systems involving C^T and + // C without explicitly inverting C + const int matrix_dim = num_molecular_orbitals; + using ColMajorMatrix = + Eigen::Matrix; + // LAPACK expects column-major layout, so we copy the row-major data into a + // column-major matrix without transposing the logical layout + ColMajorMatrix Ct = + Eigen::Map(C.data(), matrix_dim, C.rows()); + // F_MO is symmetric, so we can use it directly as the right-hand side + // without transposing + ColMajorMatrix temp_rhs = effective_F_mo; + std::vector ipiv(matrix_dim); + + auto info = + lapack::getrf(matrix_dim, matrix_dim, Ct.data(), matrix_dim, ipiv.data()); + if (info != 0) { + throw std::runtime_error("getrf failed while factorizing C^T"); + } + + info = lapack::getrs(lapack::Op::NoTrans, matrix_dim, matrix_dim, Ct.data(), + matrix_dim, ipiv.data(), temp_rhs.data(), matrix_dim); + if (info != 0) { + throw std::runtime_error("getrs failed while solving C^T X = F_mo"); + } + + temp_rhs.transposeInPlace(); + info = lapack::getrs(lapack::Op::NoTrans, matrix_dim, matrix_dim, Ct.data(), + matrix_dim, ipiv.data(), temp_rhs.data(), matrix_dim); + if (info != 0) { + throw std::runtime_error("getrs failed while solving C^T X = M^T"); + } + + effective_fock = temp_rhs.transpose(); + if (!effective_fock.isApprox(effective_fock.transpose())) { + QDK_LOGGER().warn( + "Effective Fock matrix in AO is far from symmetric. Symmetrizing..."); + effective_fock = 0.5 * (effective_fock + effective_fock.transpose().eval()); + } +} + +const RowMajorMatrix& SCFAlgorithm::get_rohf_convergence_fock_matrix() const { + QDK_LOG_TRACE_ENTERING(); + if (rohf_effective_fock_.size() == 0) { + throw std::logic_error("ROHF convergence cache not initialized"); + } + return rohf_effective_fock_; +} + +const RowMajorMatrix& SCFAlgorithm::get_rohf_convergence_density_matrix() + const { + QDK_LOG_TRACE_ENTERING(); + if (rohf_total_density_.size() == 0) { + throw std::logic_error("ROHF convergence cache not initialized"); + } + return rohf_total_density_; +} + +RowMajorMatrix& SCFAlgorithm::rohf_convergence_density_matrix() { + QDK_LOG_TRACE_ENTERING(); + if (rohf_total_density_.size() == 0) { + throw std::logic_error("ROHF convergence cache not initialized"); + } + return rohf_total_density_; +} + +void SCFAlgorithm::set_rohf_convergence_cache(const RowMajorMatrix& fock, + const RowMajorMatrix& density) { + rohf_effective_fock_ = fock; + rohf_total_density_ = density; +} + double SCFAlgorithm::calculate_og_error_(const RowMajorMatrix& F, const RowMajorMatrix& P, const RowMajorMatrix& S, @@ -250,24 +487,13 @@ bool SCFAlgorithm::check_convergence(const SCFImpl& scf_impl) { RowMajorMatrix error_matrix; int num_orbital_sets = scf_impl.get_num_orbital_spin_blocks(); - const RowMajorMatrix* F_ptr; - const RowMajorMatrix* P_ptr; - std::vector nelec_vec = scf_impl.get_num_electrons(); - const int nelec[2] = {nelec_vec[0], nelec_vec[1]}; + const RowMajorMatrix* F_ptr = nullptr; + const RowMajorMatrix* P_ptr = nullptr; if (ctx_.cfg->scf_orbital_type == SCFOrbitalType::RestrictedOpenShell) { - // To be modified when ROHFGDM is implemented: in that case, the pointer - // will come from the DIIS instance saved in DIIS_GDM, like the current - // DIIS_GDM implementation - DIIS* rohf_diis = dynamic_cast(this); - if (rohf_diis == nullptr) { - throw std::logic_error("ROHF convergence requires DIIS implementation"); - } - rohf_diis->build_rohf_f_p_matrix( - scf_impl.get_fock_matrix(), scf_impl.get_orbitals_matrix(), - scf_impl.get_density_matrix(), nelec[0], nelec[1]); - F_ptr = &rohf_diis->get_rohf_fock_matrix(); - P_ptr = &rohf_diis->get_rohf_density_matrix(); + const auto rohf_matrices = build_rohf_convergence_matrices(scf_impl); + F_ptr = &rohf_matrices.first; + P_ptr = &rohf_matrices.second; } else { F_ptr = &scf_impl.get_fock_matrix(); P_ptr = &scf_impl.get_density_matrix(); diff --git a/cpp/tests/test_scf.cpp b/cpp/tests/test_scf.cpp index ad2de41e0..4adf59e47 100644 --- a/cpp/tests/test_scf.cpp +++ b/cpp/tests/test_scf.cpp @@ -151,14 +151,15 @@ TEST_F(ScfTest, OH_ROHF_INCORE_DIIS) { EXPECT_TRUE(wfn_doublet->get_orbitals()->is_restricted()); } -TEST_F(ScfTest, OH_ROHF_Invalid_GDM) { +TEST_F(ScfTest, OH_ROKS_invalid) { auto oh = testing::create_oh_structure(); auto scf_solver = ScfSolverFactory::create(); scf_solver->settings().set("enable_gdm", true); - scf_solver->settings().set("method", "hf"); + scf_solver->settings().set("method", "pbe"); scf_solver->settings().set("scf_type", "restricted"); - EXPECT_THROW(scf_solver->run(oh, 0, 2, "sto-3g"), std::runtime_error); + // Restricted ROKS should reject this open-shell doublet case. + EXPECT_THROW(scf_solver->run(oh, 0, 2, "sto-3g"), std::invalid_argument); } TEST_F(ScfTest, Oxygen_atom_gdm) { @@ -228,6 +229,34 @@ TEST_F(ScfTest, Oxygen_atom_charged_doublet_gdm) { EXPECT_FALSE(wfn_doublet->get_orbitals()->is_restricted()); } +TEST_F(ScfTest, OH_ROHF_GDM) { + auto oh = testing::create_oh_structure(); + auto scf_solver = ScfSolverFactory::create(); + scf_solver->settings().set("enable_gdm", true); + scf_solver->settings().set("method", "hf"); + scf_solver->settings().set("scf_type", "restricted"); + auto [E_doublet, wfn_doublet] = scf_solver->run(oh, 0, 2, "sto-3g"); + + EXPECT_NEAR(E_doublet, -74.361530753176, testing::scf_energy_tolerance); + + // Check doublet orbitals + EXPECT_TRUE(wfn_doublet->get_orbitals()->is_restricted()); +} + +TEST_F(ScfTest, Oxygen_atom_ROHF_GDM) { + auto oxygen = testing::create_oxygen_structure(); + auto scf_solver = ScfSolverFactory::create(); + scf_solver->settings().set("enable_gdm", true); + scf_solver->settings().set("method", "hf"); + scf_solver->settings().set("scf_type", "restricted"); + auto [E_triplet, wfn_triplet] = scf_solver->run(oxygen, 0, 3, "cc-pvdz"); + + EXPECT_NEAR(E_triplet, -74.787513074624, testing::scf_energy_tolerance); + + // Check triplet orbitals are restricted (ROHF) + EXPECT_TRUE(wfn_triplet->get_orbitals()->is_restricted()); +} + TEST_F(ScfTest, Oxygen_atom_invalid_energy_thresh_diis_switch_gdm) { auto oxygen = testing::create_oxygen_structure(); auto scf_solver = ScfSolverFactory::create(); @@ -235,9 +264,8 @@ TEST_F(ScfTest, Oxygen_atom_invalid_energy_thresh_diis_switch_gdm) { scf_solver->settings().set("method", "pbe"); scf_solver->settings().set("enable_gdm", true); scf_solver->settings().set("energy_thresh_diis_switch", -2e-4); - // Default should be a singlet - EXPECT_THROW(scf_solver->run(oxygen, 0, 1, "cc-pvdz"), - std::invalid_argument); // open-shell dublet + + EXPECT_THROW(scf_solver->run(oxygen, 0, 1, "cc-pvdz"), std::invalid_argument); } TEST_F(ScfTest, Oxygen_atom_invalid_bfgs_history_size_limit_gdm) { @@ -247,7 +275,7 @@ TEST_F(ScfTest, Oxygen_atom_invalid_bfgs_history_size_limit_gdm) { scf_solver->settings().set("method", "pbe"); scf_solver->settings().set("enable_gdm", true); scf_solver->settings().set("gdm_bfgs_history_size_limit", 0); - // Default should be a singlet + EXPECT_THROW(scf_solver->run(oxygen, 0, 1, "cc-pvdz"), std::invalid_argument); } diff --git a/python/tests/test_scf.py b/python/tests/test_scf.py index 077cbc7f5..3882ef5bb 100644 --- a/python/tests/test_scf.py +++ b/python/tests/test_scf.py @@ -273,6 +273,64 @@ def test_scf_solver_oh_rohf_diis(self): assert abs(energy - (-74.361530753176)) < scf_energy_tolerance assert orbitals.is_restricted() + def test_scf_solver_oh_rohf_incore_diis(self): + """Test SCF solver on OH system with ROHF/sto-3g using incore ERI.""" + oh_structure = create_oh_structure() + scf_solver = algorithms.create("scf_solver") + + scf_solver.settings().set("enable_gdm", False) + scf_solver.settings().set("method", "hf") + scf_solver.settings().set("scf_type", "restricted") + scf_solver.settings().set("eri_method", "incore") + + energy, wavefunction = scf_solver.run(oh_structure, 0, 2, "sto-3g") + orbitals = wavefunction.get_orbitals() + + assert abs(energy - (-74.361530753176)) < scf_energy_tolerance + assert orbitals.is_restricted() + + def test_scf_solver_oh_rohf_gdm(self): + """Test SCF solver on OH system with ROHF/sto-3g and GDM enabled.""" + oh_structure = create_oh_structure() + scf_solver = algorithms.create("scf_solver") + + scf_solver.settings().set("enable_gdm", True) + scf_solver.settings().set("method", "hf") + scf_solver.settings().set("scf_type", "restricted") + + energy, wavefunction = scf_solver.run(oh_structure, 0, 2, "sto-3g") + orbitals = wavefunction.get_orbitals() + + assert abs(energy - (-74.361530753176)) < scf_energy_tolerance + assert orbitals.is_restricted() + + def test_scf_solver_oxygen_atom_rohf_gdm(self): + """Test SCF solver on oxygen atom triplet with ROHF/cc-pvdz and GDM.""" + oxygen = create_oxygen_structure() + scf_solver = algorithms.create("scf_solver") + + scf_solver.settings().set("enable_gdm", True) + scf_solver.settings().set("method", "hf") + scf_solver.settings().set("scf_type", "restricted") + + energy, wavefunction = scf_solver.run(oxygen, 0, 3, "cc-pvdz") + orbitals = wavefunction.get_orbitals() + + assert abs(energy - (-74.787513074624)) < scf_energy_tolerance + assert orbitals.is_restricted() + + def test_scf_solver_oh_roks_invalid(self): + """Test restricted open-shell KS request on OH doublet raises error.""" + oh_structure = create_oh_structure() + scf_solver = algorithms.create("scf_solver") + + scf_solver.settings().set("enable_gdm", True) + scf_solver.settings().set("method", "pbe") + scf_solver.settings().set("scf_type", "restricted") + + with pytest.raises(ValueError, match="Restricted open-shell calculations are only supported"): + scf_solver.run(oh_structure, 0, 2, "sto-3g") + def test_scf_solver_oxygen_atom_gdm(self): """Test SCF solver on oxygen atom with PBE/cc-pvdz.""" oxygen = create_oxygen_structure()