Expose CI RDMs (#440)

* Update code and tutorial * Fix two interface errors * Added test of new function accessors * Update gitignore
evangelistalab · Jan 23, 2025 · 259d7b8 · 259d7b8
1 parent ad08d94
commit 259d7b8
Show file tree

Hide file tree

Showing 20 changed files with 231 additions and 40 deletions.
diff --git a/.gitignore b/.gitignore
@@ -14,6 +14,12 @@ forte/version.h
 forte.egg-info
 docs/build
 build/**
+*.cube
+*.npy
+*.dat
+*.bad
+*.old
+*.tga
 
 # Qt Creator project files
 *.config

diff --git a/forte/CMakeLists.txt b/forte/CMakeLists.txt
@@ -83,6 +83,7 @@ pybind11_add_module(_forte
 api/active_space_integrals_api.cc
 api/active_space_method_api.cc
 api/active_space_solver_api.cc
+api/ci_rdms_api.cc
 api/configuration_api.cc
 api/cube_file_api.cc
 api/determinant_api.cc

diff --git a/forte/api/ci_rdms_api.cc b/forte/api/ci_rdms_api.cc
@@ -0,0 +1,61 @@
+/*
+ * @BEGIN LICENSE
+ *
+ * Forte: an open-source plugin to Psi4 (https://github.com/psi4/psi4)
+ * that implements a variety of quantum chemistry methods for strongly
+ * correlated electrons.
+ *
+ * Copyright (c) 2012-2025 by its authors (see LICENSE, AUTHORS).
+ *
+ * The copyrights for code used from other parties are included in
+ * the corresponding files.
+ *
+ * This program is free software: you can redistribute it and/or modify
+ * it under the terms of the GNU General Public License as published by
+ * the Free Software Foundation, either version 3 of the License, or
+ * (at your option) any later version.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ * GNU General Public License for more details.
+ *
+ * You should have received a copy of the GNU General Public License
+ * along with this program.  If not, see http://www.gnu.org/licenses/.
+ *
+ * @END LICENSE
+ */
+
+#include <pybind11/pybind11.h>
+#include <pybind11/numpy.h>
+#include <pybind11/stl.h>
+
+#include "ci_rdm/ci_rdms.h"
+
+namespace py = pybind11;
+using namespace pybind11::literals;
+
+namespace forte {
+void export_CI_RDMS(py::module& m) {
+    py::class_<CI_RDMS, std::shared_ptr<CI_RDMS>>(m, "CI_RDMS")
+        .def(py::init<const std::vector<int>&, const std::vector<Determinant>&,
+                      std::shared_ptr<psi::Matrix>, int, int>(),
+             "mo_symmetry"_a, "det_space"_a, "evecs"_a, "root1"_a, "root2"_a)
+        .def(
+            "compute_1rdm",
+            [](CI_RDMS& rdms) {
+                std::vector<double> opdm_a, opdm_b;
+                rdms.compute_1rdm(opdm_a, opdm_b);
+                return std::make_tuple(opdm_a, opdm_b);
+            },
+            "Compute the 1RDM")
+        .def(
+            "compute_2rdm",
+            [](CI_RDMS& rdms) {
+                std::vector<double> tpdm_aa, tpdm_ab, tpdm_bb;
+                rdms.compute_2rdm(tpdm_aa, tpdm_ab, tpdm_bb);
+                return std::make_tuple(tpdm_aa, tpdm_ab, tpdm_bb);
+            },
+            "Compute the 2RDM");
+}
+} // namespace forte
diff --git a/forte/api/determinant_api.cc b/forte/api/determinant_api.cc
@@ -87,6 +87,7 @@ void export_Determinant(py::module& m) {
         .def("destroy_beta_bit", &Determinant::destroy_beta_bit, "n"_a, "Destroy a beta bit")
         .def("count_alfa", &Determinant::count_alfa, "Count the number of set alpha bits")
         .def("count_beta", &Determinant::count_beta, "Count the number of set beta bits")
+        .def("count_pairs", &Determinant::npair, "Count the number of set pairs of bits")
         .def("symmetry", &Determinant::symmetry, "Get the symmetry")
         .def(
             "gen_excitation",

diff --git a/forte/api/forte_python_module.cc b/forte/api/forte_python_module.cc
@@ -209,6 +209,7 @@ PYBIND11_MODULE(_forte, m) {
     export_Determinant(m);
     export_Configuration(m);
     export_String(m);
+    export_CI_RDMS(m);
 
     export_SQOperatorString(m);
     export_SparseExp(m);

diff --git a/forte/api/forte_python_module.h b/forte/api/forte_python_module.h
@@ -58,6 +58,7 @@ void export_MCSCF(py::module& m);
 void export_Determinant(py::module& m);
 void export_String(py::module& m);
 void export_Configuration(py::module& m);
+void export_CI_RDMS(py::module& m);
 
 // Sparse Operators, States, and related classes
 void export_SQOperatorString(py::module& m);

diff --git a/forte/api/sparse_operator_api.cc b/forte/api/sparse_operator_api.cc
@@ -62,12 +62,15 @@ struct ExpOperator {
 
 void export_SparseOperator(py::module& m) {
     py::class_<SparseOperator>(m, "SparseOperator", "A class to represent a sparse operator")
+        // Constructors
         .def(py::init<>(), "Default constructor")
         .def(py::init<SparseOperator>(), "Copy constructor")
         .def(py::init<const SparseOperator::container&>(),
              "Create a SparseOperator from a container of terms")
         .def(py::init<const SQOperatorString&, sparse_scalar_t>(), "sqop"_a,
              "coefficient"_a = sparse_scalar_t(1), "Create a SparseOperator with a single term")
+
+        // Add/Remove terms
         .def("add",
              py::overload_cast<const SQOperatorString&, sparse_scalar_t>(&SparseOperator::add),
              "sqop"_a, "coefficient"_a = sparse_scalar_t(1), "Add a term to the operator")
@@ -76,21 +79,43 @@ void export_SparseOperator(py::module& m) {
                  &SparseOperator::add_term_from_str),
              "str"_a, "coefficient"_a = sparse_scalar_t(1), "allow_reordering"_a = false,
              "Add a term to the operator from a string representation")
+        .def(
+            "add",
+            [](SparseOperator& op, const std::vector<size_t>& acre, const std::vector<size_t>& bcre,
+               const std::vector<size_t>& aann, const std::vector<size_t>& bann,
+               sparse_scalar_t coeff) {
+                op.add(SQOperatorString({acre.begin(), acre.end()}, {bcre.begin(), bcre.end()},
+                                        {aann.begin(), aann.end()}, {bann.begin(), bann.end()}),
+                       coeff);
+            },
+            "acre"_a, "bcre"_a, "aann"_a, "bann"_a, "coeff"_a = sparse_scalar_t(1),
+            "Add a term to the operator by passing lists of creation and annihilation indices. "
+            "This version is faster than the string version and does not check for reordering")
+        .def(
+            "remove",
+            [](SparseOperator& op, const std::string& s) {
+                const auto [sqop, _] = make_sq_operator_string(s, false);
+                op.remove(sqop);
+            },
+            "Remove a term")
+
+        // Accessors
         .def(
             "__iter__",
             [](const SparseOperator& v) {
                 return py::make_iterator(v.elements().begin(), v.elements().end());
             },
             py::keep_alive<0, 1>()) // Essential: keep object alive while iterator exists
         .def(
-            "coefficient",
+            "__getitem__",
             [](const SparseOperator& op, const std::string& s) {
                 const auto [sqop, factor] = make_sq_operator_string(s, false);
                 return factor * op[sqop];
             },
             "Get the coefficient of a term")
+        .def("__len__", &SparseOperator::size, "Get the number of terms in the operator")
         .def(
-            "__getitem__",
+            "coefficient",
             [](const SparseOperator& op, const std::string& s) {
                 const auto [sqop, factor] = make_sq_operator_string(s, false);
                 return factor * op[sqop];
@@ -103,23 +128,9 @@ void export_SparseOperator(py::module& m) {
                 op[sqop] = factor * value;
             },
             "Set the coefficient of a term")
-        .def(
-            "remove",
-            [](SparseOperator& op, const std::string& s) {
-                const auto [sqop, _] = make_sq_operator_string(s, false);
-                op.remove(sqop);
-            },
-            "Remove a term")
-        .def(
-            "__matmul__",
-            [](const SparseOperator& lhs, const SparseOperator& rhs) { return lhs * rhs; },
-            "Multiply two SparseOperator objects")
-        .def(
-            "commutator",
-            [](const SparseOperator& lhs, const SparseOperator& rhs) {
-                return commutator(lhs, rhs);
-            },
-            "Compute the commutator of two SparseOperator objects")
+
+        // Arithmetic Operations
+        .def("__add__", &SparseOperator::operator+, "Add two SparseOperators")
         .def("__iadd__", &SparseOperator::operator+=, "Add a SparseOperator to this SparseOperator")
         .def("__isub__", &SparseOperator::operator-=,
              "Subtract a SparseOperator from this SparseOperator")
@@ -142,6 +153,16 @@ void export_SparseOperator(py::module& m) {
                 return self;
             },
             "Multiply this SparseOperator by another SparseOperator")
+        .def(
+            "__matmul__",
+            [](const SparseOperator& lhs, const SparseOperator& rhs) { return lhs * rhs; },
+            "Multiply two SparseOperator objects")
+        .def(
+            "commutator",
+            [](const SparseOperator& lhs, const SparseOperator& rhs) {
+                return commutator(lhs, rhs);
+            },
+            "Compute the commutator of two SparseOperator objects")
         .def(
             "__itruediv__",
             [](SparseOperator& self, sparse_scalar_t scalar) {
@@ -167,6 +188,12 @@ void export_SparseOperator(py::module& m) {
                 return self * scalar; // Reuse the __mul__ logic
             },
             "Multiply a scalar by a SparseOperator")
+        .def(
+            "__rdiv__",
+            [](const SparseOperator& self, sparse_scalar_t scalar) {
+                return self * (1.0 / scalar); // This uses the operator* we defined
+            },
+            "Divide a scalar by a SparseOperator")
         .def(
             "__mul__",
             [](const SparseOperator& self, const SparseOperator& other) {
@@ -179,18 +206,10 @@ void export_SparseOperator(py::module& m) {
                 return C;
             },
             "Multiply two SparseOperators")
-        .def(
-            "__rdiv__",
-            [](const SparseOperator& self, sparse_scalar_t scalar) {
-                return self * (1.0 / scalar); // This uses the operator* we defined
-            },
-            "Divide a scalar by a SparseOperator")
-        .def("__add__", &SparseOperator::operator+, "Add two SparseOperators")
         .def(py::self - py::self, "Subtract two SparseOperators")
         .def(
             "__neg__", [](const SparseOperator& self) { return -self; }, "Negate the operator")
         .def("copy", &SparseOperator::copy, "Create a copy of this SparseOperator")
-        .def("__len__", &SparseOperator::size, "Get the number of terms in the operator")
         .def(
             "norm", [](const SparseOperator& op) { return op.norm(); },
             "Compute the norm of the operator")

diff --git a/forte/api/sparse_operator_list_api.cc b/forte/api/sparse_operator_list_api.cc
@@ -77,6 +77,26 @@ void export_SparseOperatorList(py::module& m) {
         .def("add", &SparseOperatorList::add)
         .def("add", &SparseOperatorList::add_term_from_str, "str"_a,
              "coefficient"_a = sparse_scalar_t(1), "allow_reordering"_a = false)
+        // .def("add",
+        //      [](SparseOperatorList& op, const, sparse_scalar_t value, bool allow_reordering) {
+        //          make_sq_operator_string_from_list op.add(sqop, value);
+        //      })
+        .def("add_term",
+             py::overload_cast<const std::vector<std::tuple<bool, bool, int>>&, double, bool>(
+                 &SparseOperatorList::add_term),
+             "op_list"_a, "value"_a = 0.0, "allow_reordering"_a = false)
+        .def(
+            "add",
+            [](SparseOperatorList& op, const std::vector<size_t>& acre,
+               const std::vector<size_t>& bcre, const std::vector<size_t>& aann,
+               const std::vector<size_t>& bann, sparse_scalar_t coeff) {
+                op.add(SQOperatorString({acre.begin(), acre.end()}, {bcre.begin(), bcre.end()},
+                                        {aann.begin(), aann.end()}, {bann.begin(), bann.end()}),
+                       coeff);
+            },
+            "acre"_a, "bcre"_a, "aann"_a, "bann"_a, "coeff"_a = sparse_scalar_t(1),
+            "Add a term to the operator by passing lists of creation and annihilation indices. "
+            "This version is faster than the string version and does not check for reordering")
         .def("to_operator", &SparseOperatorList::to_operator)
         .def(
             "remove",
@@ -97,6 +117,13 @@ void export_SparseOperatorList(py::module& m) {
         .def(
             "__getitem__", [](const SparseOperatorList& op, const size_t n) { return op[n]; },
             "Get the coefficient of a term")
+        .def(
+            "__getitem__",
+            [](const SparseOperatorList& op, const std::string& s) {
+                const auto [sqop, factor] = make_sq_operator_string(s, false);
+                return factor * op[sqop];
+            },
+            "Get the coefficient of a term")
         .def(
             "__setitem__",
             [](SparseOperatorList& op, const size_t n, sparse_scalar_t value) { op[n] = value; },

diff --git a/forte/api/sq_operator_string_api.cc b/forte/api/sq_operator_string_api.cc
@@ -75,15 +75,15 @@ void export_SQOperatorString(py::module& m) {
              "Get the string representation of the operator string")
         .def(
             "__mul__",
-            [](const SQOperatorString& sqop, const std::complex<double>& scalar) {
+            [](const SQOperatorString& sqop, const sparse_scalar_t& scalar) {
                 SparseOperator sop;
                 sop.add(sqop, scalar);
                 return sop;
             },
             py::is_operator(), "Multiply an operator string by a scalar")
         .def(
             "__rmul__",
-            [](const SQOperatorString& sqop, const std::complex<double>& scalar) {
+            [](const SQOperatorString& sqop, const sparse_scalar_t& scalar) {
                 SparseOperator sop;
                 sop.add(sqop, scalar);
                 return sop;

diff --git a/forte/ci_rdm/ci_rdms.cc b/forte/ci_rdm/ci_rdms.cc
@@ -62,6 +62,8 @@ CI_RDMS::CI_RDMS(const std::vector<int>& mo_symmetry, DeterminantHashVec& wfn,
 
 CI_RDMS::~CI_RDMS() {}
 
+size_t CI_RDMS::norb() const { return norb_; }
+
 void CI_RDMS::startup() {
     // The number of correlated molecular orbitals
     norb_ = mo_symmetry_.size();

diff --git a/forte/ci_rdm/ci_rdms.h b/forte/ci_rdm/ci_rdms.h
@@ -115,6 +115,8 @@ class CI_RDMS {
                   std::vector<double>& tprdm_aab, std::vector<double>& tprdm_abb,
                   std::vector<double>& tprdm_bbb);
 
+    size_t norb() const;
+
     void set_print(bool print) { print_ = print; }
 
     void set_max_rdm(int rdm);

diff --git a/forte/helpers/determinant_helpers.cc b/forte/helpers/determinant_helpers.cc
@@ -120,12 +120,12 @@ std::vector<Determinant> make_hilbert_space(size_t nmo, size_t na, size_t nb, si
     }
     // find the maximum value in mo_symmetry and check that it is less than nirrep
     int max_sym = *std::max_element(mo_symmetry.begin(), mo_symmetry.end());
-    if (max_sym >= nirrep) {
+    if (max_sym >= static_cast<int>(nirrep)) {
         throw std::runtime_error("The symmetry of the MOs is greater than the number of irreps.");
     }
     // implement other sensible checks, like making sure that symmetry is less than nirrep and na <=
     // nmo, nb <= nmo
-    if (symmetry >= nirrep) {
+    if (symmetry >= static_cast<int>(nirrep)) {
         throw std::runtime_error(
             "The symmetry of the determinants is greater than the number of irreps.");
     }

diff --git a/forte/helpers/math_structures.h b/forte/helpers/math_structures.h
@@ -363,6 +363,17 @@ template <typename Derived, typename T, typename F> class VectorSpaceList {
     /// @return an element of the vector
     const auto& operator()(size_t n) const { return elements_[n]; }
 
+    /// @return the value of the element corresponding to the key e. Use the search function to
+    /// check if the element is present
+    F operator[](const T& e) const {
+        auto it = std::find_if(elements_.begin(), elements_.end(),
+                               [&e](const auto& p) { return p.first == e; });
+        if (it == elements_.end()) {
+            return zero_;
+        }
+        return it->second;
+    }
+
     /// @return the norm of the vector space
     /// @param p the norm to calculate (default is 2, -1 is infinity norm)
     double norm(int p = 2) const {

diff --git a/forte/sparse_ci/determinant.hpp b/forte/sparse_ci/determinant.hpp
@@ -122,6 +122,15 @@ template <size_t N> class DeterminantImpl : public BitArray<N> {
         }
     }
 
+    explicit DeterminantImpl(std::vector<size_t> alfa_list, std::vector<size_t> beta_list) {
+        for (auto i : alfa_list) {
+            set_alfa_bit(i, true);
+        }
+        for (auto i : beta_list) {
+            set_beta_bit(i, true);
+        }
+    }
+
     /// Construct the determinant from an occupation vector that
     /// specifies the alpha and beta strings.  occupation = [Ia,Ib]
     explicit DeterminantImpl(const BitArray<nbits_half>& Ia, const BitArray<nbits_half>& Ib) {

diff --git a/forte/sparse_ci/sparse_operator.cc b/forte/sparse_ci/sparse_operator.cc
@@ -70,6 +70,8 @@ std::string format_term_in_sum(sparse_scalar_t coefficient, const std::string& t
     return std::format("({} + {}i) * {}", std::real(coefficient), std::imag(coefficient), term);
 }
 
+// void SparseOperator::add_term(const SQOperator& sqop) { op_list_.push_back(sqop); }
+
 std::vector<std::string> SparseOperator::str() const {
     std::vector<std::string> v;
     for (const auto& [sqop, c] : this->elements()) {
@@ -165,4 +167,10 @@ std::vector<std::string> SparseOperatorList::str() const {
     return v;
 }
 
+void SparseOperatorList::add_term(const std::vector<std::tuple<bool, bool, int>>& op_list,
+                                  double coefficient, bool allow_reordering) {
+    auto [sqop, sign] = make_sq_operator_string_from_list(op_list, allow_reordering);
+    add(sqop, sign * coefficient);
+}
+
 } // namespace forte