Fix psi4#1989

psi4/src/psi4/libmints/wavefunction.cc

@@ -945,7 +945,7 @@ SharedVector Wavefunction::epsilon_subset_helper(SharedVector epsilon, const Dim
 }
 
 SharedMatrix Wavefunction::matrix_subset_helper(SharedMatrix M, SharedMatrix C, const std::string &basis,
-                                                const std::string matrix_basename) const {
+                                                const std::string matrix_basename, bool MO_as_overlap) const {
     if (basis == "AO") {
         auto temp = std::vector<double>(AO2SO_->max_ncol() * AO2SO_->max_nrow());
         std::string m2_name = matrix_basename + " (AO basis)";
@@ -998,30 +998,37 @@ SharedMatrix Wavefunction::matrix_subset_helper(SharedMatrix M, SharedMatrix C,
         return M2;
     } else if (basis == "MO") {
         std::string m2_name = matrix_basename + " (MO basis)";
-        auto M2 = std::make_shared<Matrix>(m2_name, C->colspi(), C->colspi());
-
-        int symm = M->symmetry();
-        int nirrep = M->nirrep();
-
-        auto SC = std::vector<double>(C->max_ncol() * C->max_nrow());
-        auto temp = std::vector<double>(C->max_ncol() * C->max_nrow());
-        for (int h = 0; h < nirrep; h++) {
-            int nmol = C->colspi()[h];
-            int nmor = C->colspi()[h ^ symm];
-            int nsol = C->rowspi()[h];
-            int nsor = C->rowspi()[h ^ symm];
-            if (!nmol || !nmor || !nsol || !nsor) continue;
-            double **Slp = S_->pointer(h);
-            double **Srp = S_->pointer(h ^ symm);
-            double **Clp = C->pointer(h);
-            double **Crp = C->pointer(h ^ symm);
-            double **Mmop = M2->pointer(h);
-            double **Msop = M->pointer(h);
-
-            C_DGEMM('N', 'N', nsor, nmor, nsor, 1.0, Srp[0], nsor, Crp[0], nmor, 0.0, SC.data(), nmor);
-            C_DGEMM('N', 'N', nsol, nmor, nsor, 1.0, Msop[0], nsor, SC.data(), nmor, 0.0, temp.data(), nmor);
-            C_DGEMM('N', 'N', nsol, nmol, nsol, 1.0, Slp[0], nsol, Clp[0], nmol, 0.0, SC.data(), nmol);
-            C_DGEMM('T', 'N', nmol, nmor, nsol, 1.0, SC.data(), nmol, temp.data(), nmor, 0.0, Mmop[0], nmor);
+        SharedMatrix M2;
+        if (MO_as_overlap) {
+            M2 = std::make_shared<Matrix>(m2_name, C->colspi(), C->colspi());
+
+            int symm = M->symmetry();
+            int nirrep = M->nirrep();
+
+            auto SC = std::vector<double>(C->max_ncol() * C->max_nrow());
+            auto temp = std::vector<double>(C->max_ncol() * C->max_nrow());
+            for (int h = 0; h < nirrep; h++) {
+                int nmol = C->colspi()[h];
+                int nmor = C->colspi()[h ^ symm];
+                int nsol = C->rowspi()[h];
+                int nsor = C->rowspi()[h ^ symm];
+                if (!nmol || !nmor || !nsol || !nsor) continue;
+                double **Slp = S_->pointer(h);
+                double **Srp = S_->pointer(h ^ symm);
+                double **Clp = C->pointer(h);
+                double **Crp = C->pointer(h ^ symm);
+                double **Mmop = M2->pointer(h);
+                double **Msop = M->pointer(h);
+
+                C_DGEMM('N', 'N', nsor, nmor, nsor, 1.0, Srp[0], nsor, Crp[0], nmor, 0.0, SC.data(), nmor);
+                C_DGEMM('N', 'N', nsol, nmor, nsor, 1.0, Msop[0], nsor, SC.data(), nmor, 0.0, temp.data(), nmor);
+                C_DGEMM('N', 'N', nsol, nmol, nsol, 1.0, Slp[0], nsol, Clp[0], nmol, 0.0, SC.data(), nmol);
+                C_DGEMM('T', 'N', nmol, nmor, nsol, 1.0, SC.data(), nmol, temp.data(), nmor, 0.0, Mmop[0], nmor);
+            }
+        } else {
+            M2 = M->clone();
+            M2->transform(C);
+            M2->set_name(m2_name);
         }
         return M2;
     } else {
@@ -1188,11 +1195,11 @@ SharedMatrix Wavefunction::Db_subset(const std::string &basis) const {
 }
 
 SharedMatrix Wavefunction::Fa_subset(const std::string &basis) const {
-    return matrix_subset_helper(Fa_, Ca_, basis, "Fock");
+    return matrix_subset_helper(Fa_, Ca_, basis, "Fock", false);
 }
 
 SharedMatrix Wavefunction::Fb_subset(const std::string &basis) const {
-    return matrix_subset_helper(Fb_, Cb_, basis, "Fock");
+    return matrix_subset_helper(Fb_, Cb_, basis, "Fock", false);
 }
 
 SharedVector Wavefunction::epsilon_a_subset(const std::string &basis, const std::string &subset) const {

psi4/src/psi4/libmints/wavefunction.h

@@ -544,10 +544,12 @@ class PSI_API Wavefunction : public std::enable_shared_from_this<Wavefunction> {
      * @param C matrix in the SO basis to use for transforms to MO basis
      * @param basis the symmetry basis to use
      *  AO, SO, MO, CartAO
+     * @param MO_as_overlap whether the AO-to-MO transformation requires the overlap.
+     *   Only used when MO basis requested.
      * @return the matrix M in the desired basis
      **/
     SharedMatrix matrix_subset_helper(SharedMatrix M, SharedMatrix C, const std::string& basis,
-                                      const std::string matrix_basename) const;
+                                      const std::string matrix_basename, bool MO_as_overlap = true) const;
 
     /**
      * Return the alpha orbital eigenvalues in the desired basis

tests/pytests/test_wfn.py

@@ -0,0 +1,24 @@
+import pytest
+
+import numpy as np
+import psi4
+
+from utils import compare_arrays
+
+pytestmark = [pytest.mark.psi, pytest.mark.api]
+
+def test_export_ao_elec_dip_deriv():
+    h2o = psi4.geometry("""
+        O
+        H 1 1.0
+        H 1 1.0 2 101.5
+    """)
+
+    rhf_e, wfn = psi4.energy('SCF/cc-pVDZ', molecule=h2o, return_wfn=True)
+
+    F_diagonals = []
+    for h in wfn.epsilon_a().nph:
+        F_diagonals.append(np.diag(h))
+    F_expected = psi4.core.Matrix.from_array(F_diagonals)
+    assert psi4.compare_matrices(wfn.Fa_subset("MO"), F_expected, 8, "Alpha Fock Matrix")  # TEST
+    assert psi4.compare_matrices(wfn.Fb_subset("MO"), F_expected, 8, "Beta Fock Matrix")  # TEST