diff --git a/main.py b/main.py index b3520cf..3b276a4 100755 --- a/main.py +++ b/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) diff --git a/qe/drivers.py b/qe/drivers.py index d9f114a..1eef96a 100644 --- a/qe/drivers.py +++ b/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) diff --git a/qe/fundamental_types.py b/qe/fundamental_types.py index f1e1f9b..ea19789 100644 --- a/qe/fundamental_types.py +++ b/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 diff --git a/qe/no.py b/qe/no.py new file mode 100644 index 0000000..218b924 --- /dev/null +++ b/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