Generate truncated hilbert spaces (#448)

* add truncation to hilbert space

* ordering

* add sanity checks, large test

* Missed static cast

* review response

* review response

* add truncated test with sym

Author: brianz98

forte/api/determinant_api.cc

@@ -166,9 +166,26 @@ void export_Determinant(py::module& m) {
         .def("get_idx", &DeterminantHashVec::get_idx, " Return the index of a determinant");
 
     m.def(
-        "hilbert_space", &make_hilbert_space, "nmo"_a, "na"_a, "nb"_a, "nirrep"_a = 1,
+        "hilbert_space",
+        [](size_t nmo, size_t na, size_t nb, size_t nirrep, std::vector<int> mo_symmetry,
+           int symmetry) { return make_hilbert_space(nmo, na, nb, nirrep, mo_symmetry, symmetry); },
+        "nmo"_a, "na"_a, "nb"_a, "nirrep"_a = 1, "mo_symmetry"_a = std::vector<int>(),
+        "symmetry"_a = 0,
+        "Generate the Hilbert space for a given number of electrons and orbitals."
+        "If information about the symmetry of the MOs is not provided, it assumes that all MOs "
+        "have symmetry 0.");
+    m.def(
+        "hilbert_space",
+        [](size_t nmo, size_t na, size_t nb, Determinant ref, int truncation, size_t nirrep,
+           std::vector<int> mo_symmetry, int symmetry) {
+            return make_hilbert_space(nmo, na, nb, ref, truncation, nirrep, mo_symmetry, symmetry);
+        },
+        "nmo"_a, "na"_a, "nb"_a, "ref"_a, "truncation"_a, "nirrep"_a = 1,
         "mo_symmetry"_a = std::vector<int>(), "symmetry"_a = 0,
-        "Generate the Hilbert space for a given number of electrons and orbitals. If information "
-        "about the symmetry of the MOs is not provided, it assumes that all MOs have symmetry 0.");
+        "Generate the Hilbert space for a given number of electrons, orbitals, and the truncation "
+        "level."
+        "If information about the symmetry of the MOs is not provided, it assumes that all MOs "
+        "have symmetry 0."
+        "A reference determinant must be provided to establish the excitation rank.");
 }
 } // namespace forte

forte/helpers/determinant_helpers.cc

@@ -112,6 +112,52 @@ std::vector<std::vector<String>> make_strings(int n, int k, size_t nirrep,
     return strings;
 }
 
+std::vector<Determinant> make_hilbert_space(size_t nmo, size_t na, size_t nb, Determinant ref, int truncation,
+                                            size_t nirrep, std::vector<int> mo_symmetry, int symmetry) {
+    std::vector<Determinant> dets;
+    if (mo_symmetry.size() != nmo) {
+        mo_symmetry = std::vector<int>(nmo, 0);
+    }
+    // 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 >= 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 >= static_cast<int>(nirrep)) {
+        throw std::runtime_error(
+            "The symmetry of the determinants is greater than the number of irreps.");
+    }
+    if (na > nmo) {
+        throw std::runtime_error(
+            "The number of alpha electrons is greater than the number of MOs.");
+    }
+    if (nb > nmo) {
+        throw std::runtime_error("The number of beta electrons is greater than the number of MOs.");
+    }
+    if (truncation < 0 || truncation > static_cast<int>(na + nb)) {
+        throw std::runtime_error("The truncation level must an integer between 0 and na + nb.");
+    }
+
+    auto strings_a = make_strings(nmo, na, nirrep, mo_symmetry);
+    auto strings_b = make_strings(nmo, nb, nirrep, mo_symmetry);
+    for (size_t ha = 0; ha < nirrep; ha++) {
+        int hb = symmetry ^ ha;
+        for (const auto& Ia : strings_a[ha]) {
+            Determinant det;
+            det.set_alfa_str(Ia);
+            for (const auto& Ib : strings_b[hb]) {
+                det.set_beta_str(Ib);
+                if (det.fast_a_xor_b_count(ref) / 2 <= truncation) {
+                    dets.push_back(det);
+                } 
+            }
+        }
+    }
+    return dets;
+}
+
 std::vector<Determinant> make_hilbert_space(size_t nmo, size_t na, size_t nb, size_t nirrep,
                                             std::vector<int> mo_symmetry, int symmetry) {
     std::vector<Determinant> dets;

forte/helpers/determinant_helpers.h

@@ -65,12 +65,26 @@ std::vector<std::vector<String>> make_strings(int n, int k, size_t nirrep,
 /// @param nmo The number of orbitals
 /// @param na The number of alpha electrons
 /// @param nb The number of beta electrons
-/// @param nirrep The number of irreps
-/// @param mo_symmetry The symmetry of the MOs
-/// @param symmetry The symmetry of the determinants
+/// @param nirrep The number of irreps (optional)
+/// @param mo_symmetry The symmetry of the MOs (optional)
+/// @param symmetry The symmetry of the determinants (optional)
 /// @return A vector of determinants
 std::vector<Determinant> make_hilbert_space(size_t nmo, size_t na, size_t nb, size_t nirrep = 1,
                                             std::vector<int> mo_symmetry = std::vector<int>(),
                                             int symmetry = 0);
 
+/// @brief Generate the Hilbert space for a given number of electrons and orbitals
+/// @param nmo The number of orbitals
+/// @param na The number of alpha electrons
+/// @param nb The number of beta electrons
+/// @param ref The reference determinant 
+/// @param truncation The excitation level truncation
+/// @param nirrep The number of irreps (optional)
+/// @param mo_symmetry The symmetry of the MOs (optional)
+/// @param symmetry The symmetry of the determinants (optional)
+/// @return A vector of determinants
+std::vector<Determinant> make_hilbert_space(size_t nmo, size_t na, size_t nb, Determinant ref, int truncation,
+                                            size_t nirrep = 1, std::vector<int> mo_symmetry = std::vector<int>(),
+                                            int symmetry = 0);
+
 } // namespace forte

forte/sparse_ci/determinant.hpp

@@ -177,6 +177,18 @@ template <size_t N> class DeterminantImpl : public BitArray<N> {
         }
     }
 
+    void set_alfa_str(const BitArray<nbits_half>& sa) {
+        for (size_t n = 0; n < nwords_half; n++) {
+            words_[n] = sa.get_word(n);
+        }
+    }
+
+    void set_beta_str(const BitArray<nbits_half>& sb) {
+        for (size_t n = 0; n < nwords_half; n++) {
+            words_[n + nwords_half] = sb.get_word(n);
+        }
+    }
+
     void set(std::initializer_list<size_t> alfa_list, std::initializer_list<size_t> beta_list) {
         zero();
         for (auto i : alfa_list) {

tests/pytest/determinant/test_hilbert_space.py

@@ -3,13 +3,43 @@
 from forte import det
 
 
-def test_determinant_hilber_space():
+def test_determinant_hilbert_space():
     dets = forte.hilbert_space(2, 1, 1)
     # compare with the expected result of the determinant
     expected = [det("20"), det("+-"), det("-+"), det("02")]
     assert len(dets) == 4
     assert sorted(dets) == sorted(expected)
 
+    # non-aufbau case
+    ref = det("002")
+    dets = forte.hilbert_space(3, 1, 1, ref, truncation=2)
+    expected = [
+        det("002"),
+        det("0+-"),
+        det("0-+"),
+        det("+0-"),
+        det("-0+"),
+        det("020"),
+        det("200"),
+        det("+-0"),
+        det("-+0"),
+    ]
+    assert len(dets) == 9
+    assert sorted(dets) == sorted(expected)
+
+    ref = det("0+-")
+    dets = forte.hilbert_space(3, 1, 1, ref, truncation=1)
+    expected = [det("0+-"), det("-+0"), det("+0-"), det("002"), det("020")]
+    assert len(dets) == 5
+    assert sorted(dets) == sorted(expected)
+
+    ref = det("20")
+    dets = forte.hilbert_space(2, 1, 1, ref, truncation=1)
+    # compare with the expected result of the determinant
+    expected = [det("20"), det("+-"), det("-+")]
+    assert len(dets) == 3, dets
+    assert sorted(dets) == sorted(expected)
+
     dets = forte.hilbert_space(2, 1, 1, nirrep=2, mo_symmetry=[0, 1], symmetry=0)
     # compare with the expected result of the determinant
     expected = [det("20"), det("02")]
@@ -22,12 +52,29 @@ def test_determinant_hilber_space():
     assert len(dets) == 2
     assert sorted(dets) == sorted(expected)
 
+    ref = det("22000")
+    dets = forte.hilbert_space(5, 2, 2, ref, truncation=1, nirrep=4, mo_symmetry=[0, 1, 2, 3, 0], symmetry=2)
+    # compare with the expected result of the determinant
+    expected = [det("2+0-0"), det("2-0+0"), det("+2-00"), det("-2+00")]
+    assert len(dets) == 4
+    assert sorted(dets) == sorted(expected)
+
     # test with the case of 6 electrons and 6 orbitals
     dets = forte.hilbert_space(6, 3, 3, 8, [0, 2, 3, 5, 6, 7], 0)
     assert len(dets) == 56
 
+    # test with the case of 6 electrons and 7 orbitals, truncated to 2 excitations
+    ref = det("2220000")
+    dets = forte.hilbert_space(7, 3, 3, ref, truncation=2)
+    assert len(dets) == 205
+
+    # test with the case of 16 electrons and 18 orbitals, truncated to 3 excitations
+    ref = det("2" * 8 + "0" * 10)
+    dets = forte.hilbert_space(18, 8, 8, ref, truncation=3)
+    assert len(dets) == 224121
 
-def test_determinant_hilber_space_edge_cases():
+
+def test_determinant_hilbert_space_edge_cases():
     dets = forte.hilbert_space(1, 1, 1)
     # compare with the expected result of the determinant
     expected = [det("2")]
@@ -57,7 +104,14 @@ def test_determinant_hilber_space_edge_cases():
     assert len(dets) == 6
     assert sorted(dets) == sorted(expected)
 
+    ref = det("++00")
+    dets = forte.hilbert_space(4, 2, 0, ref, truncation=1)
+    # compare with the expected result of the determinant
+    expected = [det("++00"), det("+0+0"), det("+00+"), det("0++0"), det("0+0+")]
+    assert len(dets) == 5
+    assert sorted(dets) == sorted(expected)
+
 
 if __name__ == "__main__":
-    test_determinant_hilber_space()
-    test_determinant_hilber_space_edge_cases()
+    test_determinant_hilbert_space()
+    test_determinant_hilbert_space_edge_cases()