Rework JK builder

forte/CMakeLists.txt

@@ -169,6 +169,9 @@ integrals/make_integrals.cc
 integrals/one_body_integrals.cc
 integrals/parallel_ccvv_algorithms.cc
 integrals/paralleldfmo.cc
+integrals/jk_builder.cc
+integrals/customjk_builder.cc
+integrals/psi4jk_builder.cc
 mrdsrg-helper/dsrg_mem.cc
 mrdsrg-helper/dsrg_source.cc
 mrdsrg-helper/dsrg_time.cc

forte/api/forte_python_module.cc

@@ -44,8 +44,10 @@
 #include "base_classes/dynamic_correlation_solver.h"
 #include "base_classes/state_info.h"
 #include "base_classes/scf_info.h"
+#include "base_classes/orbitals.h"
 
 #include "integrals/make_integrals.h"
+#include "integrals/jk_builder.h"
 
 #include "orbital-helpers/aosubspace.h"
 #include "orbital-helpers/mp2_nos.h"
@@ -84,6 +86,7 @@ void export_StateInfo(py::module& m);
 void export_SigmaVector(py::module& m);
 void export_SparseCISolver(py::module& m);
 void export_ForteCubeFile(py::module& m);
+void export_Orbitals(py::module& m);
 void export_OrbitalTransform(py::module& m);
 void export_Localize(py::module& m);
 void export_SemiCanonical(py::module& m);
@@ -188,10 +191,12 @@ PYBIND11_MODULE(_forte, m) {
     m.def("make_fragment_projector", &make_fragment_projector,
           "Make a fragment(embedding) projector");
     m.def("make_embedding", &make_embedding, "Apply fragment projector to embed");
+    m.def("make_psi4jk", &make_psi4jk, "Make a Forte JK object based on psi4");
     m.def("make_custom_ints", &make_custom_forte_integrals,
           "Make a custom Forte integral object from arrays");
     m.def("make_ints_from_psi4", &make_forte_integrals_from_psi4, "ref_wfn"_a, "options"_a,
-          "mo_space_info"_a, "int_type"_a = "", "Make a Forte integral object from psi4");
+          "mo_space_info"_a, "orbitals"_a, "int_type"_a = "",
+          "Make a Forte integral object from psi4");
     m.def("make_active_space_method", &make_active_space_method, "Make an active space method");
     m.def("make_active_space_solver", &make_active_space_solver, "Make an active space solver",
           "method"_a, "state_nroots_map"_a, "scf_info"_a, "mo_space_info"_a, "options"_a,
@@ -290,6 +295,8 @@ PYBIND11_MODULE(_forte, m) {
     export_ForteIntegrals(m);
 
     export_Symmetry(m);
+    export_Orbitals(m);
+    m.def("make_orbitals_from_psi", &make_orbitals_from_psi);
     export_OrbitalTransform(m);
     export_Localize(m);
     export_SemiCanonical(m);

forte/api/integrals_api.cc

@@ -31,6 +31,7 @@
 
 #include "helpers/helpers.h"
 #include "integrals/integrals.h"
+#include "base_classes/orbitals.h"
 
 namespace py = pybind11;
 
@@ -40,8 +41,7 @@ namespace forte {
 void export_ForteIntegrals(py::module& m) {
     py::class_<ForteIntegrals, std::shared_ptr<ForteIntegrals>>(m, "ForteIntegrals")
         .def("rotate_orbitals", &ForteIntegrals::rotate_orbitals, "Rotate MOs during contructor")
-        .def("Ca", &ForteIntegrals::Ca, "Return the alpha MO coefficients")
-        .def("Cb", &ForteIntegrals::Cb, "Return the beta MO coefficients")
+        // .def("orbitals", &ForteIntegrals::orbitals, "Return the orbitals object")
         .def("nmo", &ForteIntegrals::nmo, "Return the total number of moleuclar orbitals")
         .def("ncmo", &ForteIntegrals::ncmo, "Return the number of correlated orbitals")
         .def("frozen_core_energy", &ForteIntegrals::frozen_core_energy,

forte/api/orbital_api.cc

@@ -34,6 +34,7 @@
 #include "base_classes/orbital_transform.h"
 #include "base_classes/forte_options.h"
 #include "base_classes/mo_space_info.h"
+#include "base_classes/orbitals.h"
 #include "integrals/integrals.h"
 
 #include "orbital-helpers/localize.h"
@@ -45,6 +46,19 @@ using namespace pybind11::literals;
 
 namespace forte {
 
+void export_Orbitals(py::module& m) {
+    py::class_<Orbitals, std::shared_ptr<Orbitals>>(m, "Orbitals")
+        .def(py::init<const std::shared_ptr<psi::Matrix>&, const std::shared_ptr<psi::Matrix>&>())
+        .def("Ca", &Orbitals::Ca, "Return the alpha orbital coefficient matrix")
+        .def("Cb", &Orbitals::Cb, "Return the beta orbital coefficient matrix")
+        .def("set", &Orbitals::set, "Set the alpha and beta orbital coefficient matrices")
+        .def("rotate", &Orbitals::rotate,
+             "Rotate the orbitals using the given transformation matrices")
+        .def("copy", &Orbitals::copy, "Copy the orbitals")
+        .def("are_spin_restricted", &Orbitals::are_spin_restricted,
+             "Check if the orbitals are spin restricted");
+}
+
 /// Export the OrbitalTransform class
 void export_OrbitalTransform(py::module& m) {
     py::class_<OrbitalTransform>(m, "OrbitalTransform")
@@ -56,8 +70,8 @@ void export_OrbitalTransform(py::module& m) {
 /// Export the ForteOptions class
 void export_Localize(py::module& m) {
     py::class_<Localize>(m, "Localize")
-        .def(py::init<std::shared_ptr<ForteOptions>, std::shared_ptr<ForteIntegrals>,
-                      std::shared_ptr<MOSpaceInfo>>())
+        .def(py::init<std::shared_ptr<ForteOptions>, std::shared_ptr<MOSpaceInfo>,
+                      std::shared_ptr<Orbitals>, std::shared_ptr<ForteIntegrals>>())
         .def("compute_transformation", &Localize::compute_transformation,
              "Compute the transformation")
         .def("set_orbital_space",

forte/base_classes/README.md

@@ -5,6 +5,7 @@ Canonical order of base classes
     - StateInfo 
     - SCFInfo 
     - MOSpaceInfo 
+    - Orbitals
     - Integrals/ActiveSpaceIntegrals 
     - ForteOptions
 

forte/base_classes/active_space_solver.cc

@@ -112,7 +112,7 @@ const std::map<StateInfo, std::vector<double>>& ActiveSpaceSolver::compute_energ
     }
 
     // initialize multipole integrals
-    if (as_ints_->ints()->integral_type() != Custom) {
+    if (as_ints_->ints()->integral_type() != IntegralType::Custom) {
         if (not as_mp_ints_) {
             auto mp_ints = std::make_shared<MultipoleIntegrals>(as_ints_->ints(), mo_space_info_);
             as_mp_ints_ = std::make_shared<ActiveMultipoleIntegrals>(mp_ints);
@@ -160,7 +160,7 @@ const std::map<StateInfo, std::vector<double>>& ActiveSpaceSolver::compute_energ
     }
     print_energies();
 
-    if (as_ints_->ints()->integral_type() != Custom and
+    if (as_ints_->ints()->integral_type() != IntegralType::Custom and
         options_->get_str("ACTIVE_SPACE_SOLVER") != "EXTERNAL") {
         compute_dipole_moment(as_mp_ints_);
         compute_quadrupole_moment(as_mp_ints_);

forte/base_classes/orbital_transform.cc

@@ -42,26 +42,26 @@
 
 namespace forte {
 
-OrbitalTransform::OrbitalTransform(std::shared_ptr<ForteIntegrals> ints,
-                                   std::shared_ptr<MOSpaceInfo> mo_space_info)
-    : ints_(ints), mo_space_info_(mo_space_info) {}
+OrbitalTransform::OrbitalTransform(std::shared_ptr<MOSpaceInfo> mo_space_info,
+                                   std::shared_ptr<Orbitals> orbitals,
+                                   std::shared_ptr<ForteIntegrals> ints)
+    : mo_space_info_(mo_space_info), orbitals_(orbitals), ints_(ints) {}
 
-std::unique_ptr<OrbitalTransform>
-make_orbital_transformation(const std::string& type, std::shared_ptr<SCFInfo> scf_info,
-                            std::shared_ptr<ForteOptions> options,
-                            std::shared_ptr<ForteIntegrals> ints,
-                            std::shared_ptr<MOSpaceInfo> mo_space_info) {
+std::unique_ptr<OrbitalTransform> make_orbital_transformation(
+    const std::string& type, std::shared_ptr<SCFInfo> scf_info,
+    std::shared_ptr<ForteOptions> options, std::shared_ptr<ForteIntegrals> ints,
+    std::shared_ptr<MOSpaceInfo> mo_space_info, std::shared_ptr<Orbitals> orbitals) {
 
     std::unique_ptr<OrbitalTransform> orb_t;
 
     if (type == "LOCAL") {
-        orb_t = std::make_unique<Localize>(options, ints, mo_space_info);
+        orb_t = std::make_unique<Localize>(options, mo_space_info, orbitals, ints);
     } else if (type == "MP2NO") {
-        orb_t = std::make_unique<MP2_NOS>(scf_info, options, ints, mo_space_info);
+        orb_t = std::make_unique<MP2_NOS>(scf_info, options, mo_space_info, orbitals, ints);
     } else if (type == "CINO") {
-        orb_t = std::make_unique<CINO>(options, ints, mo_space_info);
+        orb_t = std::make_unique<CINO>(options, mo_space_info, orbitals, ints);
     } else if (type == "MRCINO") {
-        orb_t = std::make_unique<MRCINO>(scf_info, options, ints, mo_space_info);
+        orb_t = std::make_unique<MRCINO>(scf_info, options, mo_space_info, orbitals, ints);
     } else if (type == "MRPT2NO") {
         // perform reference CI calculation
         auto as_type = options->get_str("ACTIVE_SPACE_SOLVER");
@@ -78,7 +78,7 @@ make_orbital_transformation(const std::string& type, std::shared_ptr<SCFInfo> sc
             as_solver->compute_average_rdms(state_weights_map, rdm_level, RDMsType::spin_free);
 
         // initialize
-        orb_t = std::make_unique<MRPT2_NOS>(rdms, scf_info, options, ints, mo_space_info);
+        orb_t = std::make_unique<MRPT2_NOS>(scf_info, options, mo_space_info, orbitals, ints, rdms);
     } else {
         throw std::runtime_error("Orbital type " + type + " is not supported!");
     }

forte/base_classes/orbital_transform.h

@@ -1,4 +1,4 @@
-#pragma once 
+#pragma once
 
 namespace psi {
 class Matrix;
@@ -7,16 +7,17 @@ class Matrix;
 namespace forte {
 
 class ForteIntegrals;
+class ForteOptions;
 class MOSpaceInfo;
 class SCFInfo;
-class ForteOptions;
+class Orbitals;
 
 class OrbitalTransform {
 
   public:
     /// Constructor
-    OrbitalTransform(std::shared_ptr<ForteIntegrals> ints,
-                     std::shared_ptr<MOSpaceInfo> mo_space_info);
+    OrbitalTransform(std::shared_ptr<MOSpaceInfo> mo_space_info, std::shared_ptr<Orbitals> orbitals,
+                     std::shared_ptr<ForteIntegrals> ints = nullptr);
 
     /// Default constructor
     OrbitalTransform() = default;
@@ -31,21 +32,22 @@ class OrbitalTransform {
     std::shared_ptr<psi::Matrix> get_Ub() { return Ub_; };
 
   protected:
-    // The integrals
-    std::shared_ptr<ForteIntegrals> ints_;
     /// The MOSpace info
     std::shared_ptr<MOSpaceInfo> mo_space_info_;
+    /// The orbitals
+    std::shared_ptr<Orbitals> orbitals_;
+    // The integrals
+    std::shared_ptr<ForteIntegrals> ints_;
 
     /// @brief Unitary matrix for alpha orbital rotations
     std::shared_ptr<psi::Matrix> Ua_;
     /// @brief Unitary matrix for beta orbital rotations
     std::shared_ptr<psi::Matrix> Ub_;
 };
 
-std::unique_ptr<OrbitalTransform>
-make_orbital_transformation(const std::string& type, std::shared_ptr<SCFInfo> scf_info,
-                            std::shared_ptr<ForteOptions> options,
-                            std::shared_ptr<ForteIntegrals> ints,
-                            std::shared_ptr<MOSpaceInfo> mo_space_info);
+std::unique_ptr<OrbitalTransform> make_orbital_transformation(
+    const std::string& type, std::shared_ptr<SCFInfo> scf_info,
+    std::shared_ptr<ForteOptions> options, std::shared_ptr<ForteIntegrals> ints,
+    std::shared_ptr<MOSpaceInfo> mo_space_info, std::shared_ptr<Orbitals> orbitals);
 
 } // namespace forte

forte/base_classes/orbitals.cc

@@ -27,19 +27,51 @@
  */
 
 #include "psi4/libmints/wavefunction.h"
+#include "helpers/helpers.h"
 
 #include "orbitals.h"
 
 namespace forte {
-Orbitals::Orbitals(const std::shared_ptr<const psi::Wavefunction>& wfn, bool restricted) {
-    // Grab the MO coefficients from psi and enforce spin restriction if necessary
-    Ca_ = wfn->Ca()->clone();
-    Cb_ = restricted ? wfn->Ca()->clone() : wfn->Cb()->clone();
+
+Orbitals::Orbitals(const std::shared_ptr<psi::Matrix>& Ca, const std::shared_ptr<psi::Matrix>& Cb) {
+    Ca_ = Ca->clone();
+    Cb_ = Cb->clone();
+}
+
+std::shared_ptr<psi::Matrix> Orbitals::Ca() { return Ca_; }
+
+std::shared_ptr<psi::Matrix> Orbitals::Cb() { return Cb_; }
+
+void Orbitals::set(const std::shared_ptr<psi::Matrix>& Ca, const std::shared_ptr<psi::Matrix>& Cb) {
+    if (not(elementwise_compatible_matrices(Ca_, Ca) and
+            elementwise_compatible_matrices(Cb_, Cb))) {
+        throw std::runtime_error(
+            "Orbital::set: orbital coefficient matrices have different dimensions!");
+    }
+    Ca_ = Ca->clone();
+    Cb_ = Cb->clone();
+}
+
+void Orbitals::rotate(std::shared_ptr<psi::Matrix> Ua, std::shared_ptr<psi::Matrix> Ub) {
+    // 1. Rotate the orbital coefficients and store them in the ForteIntegral object
+    auto Ca_rotated = psi::linalg::doublet(Ca_, Ua);
+    auto Cb_rotated = psi::linalg::doublet(Cb_, Ub);
+    Ca_->copy(Ca_rotated);
+    Cb_->copy(Cb_rotated);
+}
+
+void Orbitals::copy(const Orbitals& other) {
+    Ca_->copy(*other.Ca_);
+    Cb_->copy(*other.Cb_);
+}
+
+bool Orbitals::are_spin_restricted(double threshold) const {
+    return matrix_distance(Ca_, Cb_) < threshold;
 }
 
-std::unique_ptr<Orbitals> make_orbitals(const std::shared_ptr<const psi::Wavefunction>& wfn,
-                                        bool restricted) {
-    return std::make_unique<Orbitals>(wfn, restricted);
+std::shared_ptr<Orbitals> make_orbitals_from_psi(const std::shared_ptr<psi::Wavefunction>& wfn,
+                                                 bool restricted) {
+    return std::make_shared<Orbitals>(wfn->Ca(), restricted ? wfn->Ca() : wfn->Cb());
 }
 
 } // namespace forte

forte/base_classes/orbitals.h

@@ -42,13 +42,26 @@ namespace forte {
 class Orbitals {
   public:
     // ==> Class Constructor <==
-    Orbitals(const std::shared_ptr<const psi::Wavefunction>& wfn, bool restricted);
+    Orbitals(const std::shared_ptr<psi::Matrix>& Ca, const std::shared_ptr<psi::Matrix>& Cb);
 
     // ==> Class Methods <==
     /// @return The alpha orbital coefficient matrix
-    const std::shared_ptr<const psi::Matrix> Ca() const { return Ca_; }
+    std::shared_ptr<psi::Matrix> Ca();
     /// @return The beta orbital coefficient matrix
-    const std::shared_ptr<const psi::Matrix> Cb() const { return Cb_; }
+    std::shared_ptr<psi::Matrix> Cb();
+
+    void set(const std::shared_ptr<psi::Matrix>& Ca, const std::shared_ptr<psi::Matrix>& Cb);
+
+    /// @brief Rotate the orbitals using the given transformation matrices
+    /// The orbital coefficients are rotated according to the following formula:
+    /// C(new) = C(old) U
+    /// @param Ua Alpha orbital transformation matrix
+    /// @param Ub Beta orbital transformation matrix
+    void rotate(std::shared_ptr<psi::Matrix> Ua, std::shared_ptr<psi::Matrix> Ub);
+
+    void copy(const Orbitals& other);
+
+    bool are_spin_restricted(double threshold = 1.0e-8) const;
 
   private:
     // ==> Class Data <==
@@ -58,7 +71,7 @@ class Orbitals {
     std::shared_ptr<psi::Matrix> Cb_;
 };
 
-std::unique_ptr<Orbitals> make_orbitals(const std::shared_ptr<const psi::Wavefunction>& wfn,
-                                        bool restricted);
+std::shared_ptr<Orbitals> make_orbitals_from_psi(const std::shared_ptr<psi::Wavefunction>& wfn,
+                                                 bool restricted);
 
 } // namespace forte

forte/data.py

@@ -3,14 +3,15 @@
 from psi4.core import Wavefunction, Molecule
 
 from ._forte import (
-    ForteOptions,
-    SCFInfo,
-    MOSpaceInfo,
-    ForteIntegrals,
     ActiveSpaceIntegrals,
-    StateInfo,
     ActiveSpaceSolver,
+    ForteOptions,
+    ForteIntegrals,
+    MOSpaceInfo,
+    Orbitals,
     RDMs,
+    SCFInfo,
+    StateInfo,
     Symmetry,
 )
 from forte import Model, Results
@@ -34,6 +35,8 @@ class ForteData:
         The model object.
     scf_info: SCFInfo
         The SCF information.
+    orbitals: Orbitals
+        The molecular orbitals.
     mo_space_info: MOSpaceInfo
         The molecular orbital space information.
     ints: ForteIntegrals
@@ -54,6 +57,7 @@ class ForteData:
     molecule: Molecule = None
     model: Model = None
     scf_info: SCFInfo = None
+    orbitals: Orbitals = None
     mo_space_info: MOSpaceInfo = None
     ints: ForteIntegrals = None
     as_ints: ActiveSpaceIntegrals = None

forte/helpers/determinant_helpers.cc

@@ -69,7 +69,8 @@ make_hamiltonian_matrix(const std::vector<Determinant>& dets,
     auto H = std::make_shared<psi::Matrix>("H", n, n);
 
     // If we are running DiskDF then we need to revert to a single thread loop
-    auto threads = (as_ints->get_integral_type() == DiskDF) ? 1 : omp_get_max_threads();
+    auto threads =
+        (as_ints->get_integral_type() == IntegralType::DiskDF) ? 1 : omp_get_max_threads();
 
 #pragma omp parallel for schedule(dynamic) num_threads(threads)
     for (size_t I = 0; I < n; I++) {