NOs

main.py

@@ -1,6 +1,7 @@
 #!/usr/bin/env python3
 from qe.drivers import *
 from qe.io import *
+from qe.no import *
 import sys
 import argparse
 
@@ -33,6 +34,14 @@
         help="Way in which Hamiltonian is generated. Integral driven: local set of integrals, determinants are found on each node. Determinant driven: local set of determinants, all integrals are on each node.",
     )
 
+    parser.add_argument(
+        "-NOs",
+        type=bool,
+        default=False,
+        required=False,
+        help="Construct Natural Orbitals?"
+    )
+
     args = parser.parse_args()
     # Load integrals
     n_ord, E0, d_one_e_integral, d_two_e_integral = load_integrals(args.fcidump_path)
@@ -56,3 +65,6 @@
             driven_by=args.driven_by,
         )
         print(f"N_det: {len(psi_det)}, E {E}")
+
+    if do_NOs:
+        no_occ, no_coeff = natural_orbitals(comm, lewis, n_orb, psi_coef, psi_det, n)

qe/drivers.py

@@ -121,9 +121,33 @@ def gen_all_connected_spindet(self, spindet: Spin_determinant, ed: int) -> Itera
         apply_excitation_to_spindet = partial(Excitation.apply_excitation, spindet)
         return map(apply_excitation_to_spindet, l_exc)
 
+    def gen_all_connected_single_det_from_det(self, det: Determinant) -> Iterator[Determinant]:
+        """
+        Generate all the determinant who are single excitation from the input determinant
+
+        >>> sorted(Excitation(3).gen_all_connected_single_det_from_det( Determinant((0, 1), (0,))))
+        [Determinant(alpha=(0, 1), beta=(1,)),
+         Determinant(alpha=(0, 1), beta=(2,)),
+         Determinant(alpha=(0, 2), beta=(0,)),
+         Determinant(alpha=(1, 2), beta=(0,))]
+        """
+
+        # All single excitation from alpha or for beta determinant
+
+        # We use l_single_a, and l_single_b twice. So we store them.
+        l_single_a = set(self.gen_all_connected_spindet(det.alpha, 1))
+
+        s_a = (Determinant(det_alpha, det.beta) for det_alpha in l_single_a)
+
+        l_single_b = set(self.gen_all_connected_spindet(det.beta, 1))
+
+        s_b = (Determinant(det.alpha, det_beta) for det_beta in l_single_b)
+
+        return chain(s_a, s_b)
+
     def gen_all_connected_det_from_det(self, det: Determinant) -> Iterator[Determinant]:
         """
-        Generate all the determinant who are single or double exictation (aka connected) from the input determinant
+        Generate all the determinant who are single or double excitation (aka connected) from the input determinant
 
         >>> sorted(Excitation(3).gen_all_connected_det_from_det( Determinant((0, 1), (0,))))
         [Determinant(alpha=(0, 1), beta=(1,)),
@@ -1970,6 +1994,56 @@ def psi_external_pt2(self, psi_coef: Psi_coef) -> Tuple[Psi_det, List[Energy]]:
             "i,i,i -> i", nominator, nominator, denominator
         )  # vector * vector * vector -> scalar
 
+    def MO_rdm1(self, psi_coef: Psi_coef, n, n_orb) -> MO_1rdm:
+        """
+        Returns the 1-rdm in the MO basis,
+        Based off of implementation found here: 10.48550/ARXIV.1311.6244
+
+         psi_coef: Psi_coef
+            wave function coefficient
+         n: int
+            total number of determinants
+         n_orb: int
+            number of molecular orbitals
+        """
+        c = np.array(psi_coef, dtype="float")  # Coef. vector as np.array
+        # Coeffs. of local determinants
+        c_i = c[self.offsets[self.rank] : (self.offsets[self.rank] + self.distribution[self.rank])]
+        # Pre-allocate MO 1RDM
+        local_mo_rdm = np.zeros(n_orb, n_orb)
+
+        # Naive
+        # for I in psi_local
+        #   do diagonal contribution
+        #   for single excitations from I
+        #       do off-diagonal contributions
+        # reduce on local_mo_rdms
+
+        # Loop over local dets
+        for I, det_I in enumerate(self.psi_local):
+            # do diagonal contribution
+            coeff_sq = c_i[I] * c_i[I]
+            for orb_i in det_I.alpha:
+                local_mo_rdm[orb_i, orb_i] += coeff_sq
+            for orb_i in det_I.beta:
+                local_mo_rdm[orb_i, orb_i] += coeff_sq
+            #  for single excitations from current determinants, this returns iterator
+            single_excitations = Excitation(n_orb).gen_all_connected_single_det_from_det(det_I)
+            for J, det_j in enumerate(single_excitations):
+                coeff_ij = c_i[I] * c_i[J]
+                # calculate phase, index of hole, and index of particle between two determinants
+                if det_I.alpha == det_J.alpha:
+                    # compute phase and hole particle index, returned as Tuple[phase,hole,particle].
+                    php_tuple = PhaseIdx.single_exc(det_I.alpha, det_J.alpha)
+                else:  # not alpha_excitation:
+                    php_tuple = PhaseIdx.single_exc(det_I.beta, det_J.beta)
+                # account for spin average
+                local_mo_rdm[php_tuple[1], php_tuple[2]] += php_tuple[0] * coeff_ij * 2
+                local_mo_rdm[php_tuple[2], php_tuple[1]] += php_tuple[0] * coeff_ij * 2
+
+        # reduce on local_rdm?
+        return False
+
     def E_pt2(self, psi_coef: Psi_coef) -> Energy:
         # The sum of the pt2 contribution of each external determinant
         psi_external_energy = self.psi_external_pt2(psi_coef)

qe/fundamental_types.py

@@ -1,4 +1,5 @@
 from typing import Tuple, Dict, NamedTuple, List, NewType
+import numpy as np
 
 # Orbital index (0,1,2,...,n_orb-1)
 OrbitalIdx = NewType("OrbitalIdx", int)
@@ -25,6 +26,11 @@ class Determinant(NamedTuple):
 
 Psi_det = List[Determinant]
 Psi_coef = List[float]
+
+MO_1rdm = np.ndarray
+no_occ = np.ndarray
+no_coeff = np.ndarray
+
 # We have two type of energy.
 # The varitional Energy who correpond Psi_det
 # The pt2 Energy who correnpond to the pertubative energy induce by each determinant connected to Psi_det

qe/no.py

@@ -0,0 +1,20 @@
+from drivers import Powerplant_manager
+from fundamental_types import *
+import numpy as np
+
+
+def get_MO_1rdm(comm, lewis, n_orb, psi_coef: Psi_coef, psi_det: Psi_det, n) -> MO_1rdm:
+    return Powerplant_manager(comm, lewis).MO_1rdm(psi_coef, n, n_orb)
+
+
+def natural_orbitals(
+    comm, lewis: Hamiltonian_generator, n_orb, psi_coef: Psi_coef, psi_det: Psi_det, n
+) -> Tuple[no_occ, no_coeff]:
+    # create MO 1 rdm
+    mo_1rdm = get_MO_1rdm(comm, lewis, n_orb, psi_coef, psi_det, n)
+    # diagonalize MO 1 rdm
+    no_occ, no_coeff_MO_basis = np.linalg.eigh(mo_1rdm)
+    no_occ = no_occ[::-1]
+    no_coeff = np.fliplr(no_coeff)
+    # transform MO 1 rdm to ao basis
+    # NEED AO BASIS FOR THIS