Skip to content

Commit

Permalink
MolecularData: factor out independent run methods (#99)
Browse files Browse the repository at this point in the history
  • Loading branch information
@kevinsung
kevinsung authored Feb 13, 2024
1 parent 11e1027 commit eeb80ff
Show file tree
Hide file tree
Showing 4 changed files with 111 additions and 101 deletions.
69 changes: 35 additions & 34 deletions docs/tutorials/04-lucj.ipynb
View file
Original file line number Diff line number Diff line change
Expand Up @@ -19,7 +19,7 @@
"output_type": "stream",
"text": [
"converged SCF energy = -77.4456267643962\n",
"CASCI E = -77.6290254326717 E(CI) = -3.57322412553862 S^2 = 0.0000000\n"
"CASCI E = -77.6290254326717 E(CI) = -3.57322412553863 S^2 = 0.0000000\n"
]
}
],
Expand Down Expand Up @@ -53,12 +53,13 @@
"active_space = range(mol.nelectron // 2 - 2, mol.nelectron // 2 + 2)\n",
"\n",
"# Get molecular data and molecular Hamiltonian (one- and two-body tensors)\n",
"mol_data = ffsim.MolecularData.from_scf(\n",
" hartree_fock, active_space=active_space, fci=True\n",
")\n",
"mol_data = ffsim.MolecularData.from_scf(hartree_fock, active_space=active_space)\n",
"norb = mol_data.norb\n",
"nelec = mol_data.nelec\n",
"mol_hamiltonian = mol_data.hamiltonian"
"mol_hamiltonian = mol_data.hamiltonian\n",
"\n",
"# Compute FCI energy\n",
"mol_data.run_fci()"
]
},
{
Expand Down Expand Up @@ -96,8 +97,8 @@
"name": "stdout",
"output_type": "stream",
"text": [
"E(CCSD) = -77.49387212754468 E_corr = -0.04824536314851346\n",
"Energy at initialization: -77.46975600021715\n"
"E(CCSD) = -77.49387212754473 E_corr = -0.04824536314851485\n",
"Energy at initialization: -77.46975600021644\n"
]
}
],
Expand Down Expand Up @@ -150,10 +151,10 @@
" message: STOP: TOTAL NO. of ITERATIONS REACHED LIMIT\n",
" success: False\n",
" status: 1\n",
" fun: -77.5096393678589\n",
" x: [-1.486e-01 1.914e-01 ... 6.534e-03 -1.987e+00]\n",
" fun: -77.50964466939116\n",
" x: [-1.479e-01 1.876e-01 ... 3.824e-03 6.462e-01]\n",
" nit: 5\n",
" jac: [-1.194e-03 -1.067e-03 ... -5.684e-04 -2.913e-04]\n",
" jac: [ 1.333e-03 -2.134e-03 ... -5.372e-04 -3.055e-04]\n",
" nfev: 584\n",
" njev: 8\n",
" hess_inv: <72x72 LbfgsInvHessProduct with dtype=float64>\n"
Expand Down Expand Up @@ -216,10 +217,10 @@
" message: STOP: TOTAL NO. of ITERATIONS REACHED LIMIT\n",
" success: False\n",
" status: 1\n",
" fun: -77.45741314558897\n",
" x: [-4.128e-01 8.249e-02 ... -3.067e-02 -1.935e+00]\n",
" fun: -77.45739668895403\n",
" x: [-4.078e-01 8.271e-02 ... -1.694e-02 7.034e-01]\n",
" nit: 5\n",
" jac: [ 8.285e-04 -1.137e-05 ... 3.837e-05 -2.345e-04]\n",
" jac: [ 1.815e-03 -2.700e-05 ... 3.553e-05 -8.384e-05]\n",
" nfev: 423\n",
" njev: 9\n",
" hess_inv: <46x46 LbfgsInvHessProduct with dtype=float64>\n"
Expand Down Expand Up @@ -276,34 +277,34 @@
"Number of parameters: 46\n",
" message: Stop: Total number of iterations reached limit.\n",
" success: False\n",
" fun: -77.49749625463323\n",
" x: [ 1.063e+00 -8.659e-02 ... -1.679e+01 2.643e+01]\n",
" fun: -77.47158320347539\n",
" x: [-6.631e-01 8.976e-02 ... -2.173e-01 9.995e-01]\n",
" nit: 5\n",
" jac: [ 9.021e-03 1.711e-03 ... 6.221e-03 -5.828e-03]\n",
" nfev: 856\n",
" jac: [ 7.259e-03 -1.181e-03 ... 3.433e-03 -5.383e-03]\n",
" nfev: 826\n",
" njev: 5\n",
" nlinop: 626\n",
" nlinop: 596\n",
"\n",
"Iteration 1\n",
" Energy: -77.41517984348127\n",
" Norm of gradient: 0.10920578601840765\n",
" Regularization hyperparameter: 6.687981176805703e-07\n",
" Variation hyperparameter: 0.696485198650709\n",
" Energy: -77.45724086956001\n",
" Norm of gradient: 0.010199505442545696\n",
" Regularization hyperparameter: 0.026890960056319083\n",
" Variation hyperparameter: 0.9931578556353272\n",
"Iteration 2\n",
" Energy: -77.4387582021023\n",
" Norm of gradient: 0.10615951394818811\n",
" Regularization hyperparameter: 0.037348668091859744\n",
" Variation hyperparameter: 0.6955252427321973\n",
" Energy: -77.4580971045074\n",
" Norm of gradient: 0.0076454977282466455\n",
" Regularization hyperparameter: 0.0006243190909239268\n",
" Variation hyperparameter: 0.9931620654125893\n",
"Iteration 3\n",
" Energy: -77.46917421145808\n",
" Norm of gradient: 0.058657242826842995\n",
" Regularization hyperparameter: 0.005576706874662407\n",
" Variation hyperparameter: 0.9943325006472397\n",
" Energy: -77.45812347703834\n",
" Norm of gradient: 0.007443152735598327\n",
" Regularization hyperparameter: 1.0505969544705465\n",
" Variation hyperparameter: 0.9931574003267095\n",
"Iteration 4\n",
" Energy: -77.48627050747123\n",
" Norm of gradient: 0.05746841139489317\n",
" Regularization hyperparameter: 0.015337740819403546\n",
" Variation hyperparameter: 0.9941142453075373\n"
" Energy: -77.46243342363793\n",
" Norm of gradient: 0.02558853103196335\n",
" Regularization hyperparameter: 0.0033642512286005625\n",
" Variation hyperparameter: 0.9991474891911944\n"
]
}
],
Expand Down
15 changes: 9 additions & 6 deletions docs/tutorials/05-entanglement-forging.ipynb
View file
Original file line number Diff line number Diff line change
Expand Up @@ -18,7 +18,7 @@
"name": "stdout",
"output_type": "stream",
"text": [
"converged SCF energy = -75.6787887956297\n",
"converged SCF energy = -75.6787887956296\n",
"CASCI E = -75.7288249991515 E(CI) = -23.6332495815006 S^2 = 0.0000000\n"
]
}
Expand Down Expand Up @@ -53,10 +53,13 @@
"active_space = range(1, mol.nao_nr())\n",
"\n",
"# Get molecular data and molecular Hamiltonian (one- and two-body tensors)\n",
"mol_data = ffsim.MolecularData.from_mole(mol, active_space=active_space, fci=True)\n",
"mol_data = ffsim.MolecularData.from_mole(mol, active_space=active_space)\n",
"norb = mol_data.norb\n",
"nelec = mol_data.nelec\n",
"mol_hamiltonian = mol_data.hamiltonian"
"mol_hamiltonian = mol_data.hamiltonian\n",
"\n",
"# Compute FCI energy\n",
"mol_data.run_fci()"
]
},
{
Expand All @@ -68,7 +71,7 @@
"name": "stdout",
"output_type": "stream",
"text": [
"Energy at initialialization: -74.20656273321596\n"
"Energy at initialialization: -74.20656273321593\n"
]
}
],
Expand Down Expand Up @@ -114,10 +117,10 @@
" message: STOP: TOTAL NO. of f AND g EVALUATIONS EXCEEDS LIMIT\n",
" success: False\n",
" status: 1\n",
" fun: -75.68085225170408\n",
" fun: -75.68085348348217\n",
" x: [ 2.996e+00 -7.549e-01 ... 2.650e+00 8.012e-01]\n",
" nit: 6\n",
" jac: [ 1.756e-03 9.118e-03 ... -1.192e-02 9.521e-04]\n",
" jac: [ 1.771e-03 9.121e-03 ... -1.192e-02 9.578e-04]\n",
" nfev: 112\n",
" njev: 7\n",
" hess_inv: <15x15 LbfgsInvHessProduct with dtype=float64>\n"
Expand Down
125 changes: 65 additions & 60 deletions python/ffsim/molecular_data.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -14,7 +14,8 @@
from collections.abc import Iterable, Sequence

import numpy as np
from pyscf import ao2mo, cc, gto, mcscf, mp, scf, symm
import pyscf.scf
from pyscf import ao2mo, cc, gto, mcscf, mp, symm
from pyscf.scf.hf import SCF

from ffsim.hamiltonians import MolecularHamiltonian
Expand Down Expand Up @@ -53,7 +54,7 @@ def orbital_symmetries(hartree_fock: SCF, orbitals: Sequence[int]) -> list[int]
return [MOLPRO_ID[hartree_fock.mol.groupname][i] for i in idx]


@dataclasses.dataclass(frozen=True)
@dataclasses.dataclass
class MolecularData:
"""Class for storing molecular data.
Expand All @@ -62,10 +63,10 @@ class MolecularData:
basis: The basis set, e.g. "sto-6g".
spin: The spin of the molecule.
symmetry: The symmetry of the molecule.
norb: The number of spatial orbitals.
nelec: The number of alpha and beta electrons.
mo_coeff: Hartree-Fock canonical orbital coefficients in the AO basis.
mo_occ: Hartree-Fock canonical orbital occupancies.
norb: The number of spatial orbitals.
nelec: The number of alpha and beta electrons.
active_space: The orbitals included in the active space.
core_energy: The core energy.
one_body_tensor: The one-body tensor.
Expand All @@ -87,18 +88,18 @@ class MolecularData:
basis: str
spin: int
symmetry: str | None
# Hartree-Fock data
mo_coeff: np.ndarray
mo_occ: np.ndarray
hf_energy: float
# active space information
norb: int
nelec: tuple[int, int]
active_space: list[int]
# Hamiltonian coefficients
# molecular integrals
core_energy: float
one_body_integrals: np.ndarray
two_body_integrals: np.ndarray
# Hartree-Fock data
hf_energy: float
mo_coeff: np.ndarray
mo_occ: np.ndarray
# MP2 data
mp2_energy: float | None = None
mp2_t2: np.ndarray | None = None
Expand All @@ -122,23 +123,29 @@ def hamiltonian(self) -> MolecularHamiltonian:
constant=self.core_energy,
)

@property
def mole(self) -> gto.Mole:
"""The pySCF Mole class for this molecular data."""
mol = gto.Mole()
return mol.build(atom=self.atom, basis=self.basis, symmetry=self.symmetry)

@property
def scf(self) -> gto.Mole:
"""The pySCF SCF class for this molecular data."""
hartree_fock = pyscf.scf.RHF(self.mole)
hartree_fock.mo_occ = self.mo_occ
hartree_fock.mo_coeff = self.mo_coeff
return hartree_fock

@staticmethod
def from_scf(
hartree_fock: SCF,
active_space: Iterable[int] | None = None,
*,
mp2: bool = False,
ccsd: bool = False,
fci: bool = False,
hartree_fock: SCF, active_space: Iterable[int] | None = None
) -> "MolecularData":
"""Initialize a MolecularData object from a Hartree-Fock calculation.
Args:
hartree_fock: The Hartree-Fock object.
active_space: An optional list of orbitals to use for the active space.
mp2: Whether to calculate and store the MP2 energy and t2 amplitudes.
ccsd: Whether to calculate and store the CCSD energy, t1, and t2 amplitudes.
fci: Whether to calculate and store the FCI energy and state vector.
"""
if not hartree_fock.e_tot:
raise ValueError(
Expand All @@ -161,31 +168,6 @@ def from_scf(
one_body_tensor, core_energy = cas.get_h1cas(mo)
two_body_integrals = cas.get_h2cas(mo)

# run MP2 if requested
frozen = [i for i in range(hartree_fock.mol.nao_nr()) if i not in active_space]
mp2_energy = None
mp2_t2 = None
if mp2:
mp2_solver = mp.MP2(hartree_fock, frozen=frozen)
mp2_energy, mp2_t2 = mp2_solver.kernel()

# run CCSD if requested
ccsd_energy = None
ccsd_t1 = None
ccsd_t2 = None
if ccsd:
ccsd_solver = cc.CCSD(
hartree_fock,
frozen=frozen,
)
ccsd_energy, ccsd_t1, ccsd_t2 = ccsd_solver.kernel()

# run FCI if requested
fci_energy = None
fci_vec = None
if fci:
fci_energy, _, fci_vec, _, _ = cas.kernel()

# compute dipole integrals
charges = hartree_fock.mol.atom_charges()
coords = hartree_fock.mol.atom_coords()
Expand Down Expand Up @@ -213,13 +195,6 @@ def from_scf(
one_body_integrals=one_body_tensor,
two_body_integrals=two_body_integrals,
hf_energy=hf_energy,
mp2_energy=mp2_energy,
mp2_t2=mp2_t2,
ccsd_energy=ccsd_energy,
ccsd_t1=ccsd_t1,
ccsd_t2=ccsd_t2,
fci_energy=fci_energy,
fci_vec=fci_vec,
dipole_integrals=dipole_integrals,
orbital_symmetries=orbsym,
)
Expand All @@ -228,23 +203,53 @@ def from_scf(
def from_mole(
molecule: gto.Mole,
active_space: Iterable[int] | None = None,
mp2: bool = False,
ccsd: bool = False,
fci: bool = False,
scf_func=scf.RHF,
scf_func=pyscf.scf.RHF,
) -> "MolecularData":
"""Initialize a MolecularData object from a pySCF molecule.
Args:
molecule: The molecule.
active_space: An optional list of orbitals to use for the active space.
mp2: Whether to calculate and store the MP2 energy and t2 amplitudes.
ccsd: Whether to calculate and store the CCSD energy, t1, and t2 amplitudes.
fci: Whether to calculate and store the FCI energy and state vector.
scf_func: The pySCF SCF function to use for the Hartree-Fock calculation.
"""
hartree_fock = scf_func(molecule)
hartree_fock.run()
return MolecularData.from_scf(
hartree_fock, active_space=active_space, mp2=mp2, ccsd=ccsd, fci=fci
)
return MolecularData.from_scf(hartree_fock, active_space=active_space)

def run_mp2(self, *, store_t2: bool = False):
"""Run MP2 and store results."""
cas = mcscf.CASCI(self.scf, ncas=self.norb, nelecas=self.nelec)
mo = cas.sort_mo(self.active_space, mo_coeff=self.mo_coeff, base=0)
frozen = [i for i in range(self.norb) if i not in self.active_space]
mp2_solver = mp.MP2(self.scf, frozen=frozen)
mp2_energy, mp2_t2 = mp2_solver.kernel(mo_coeff=mo)
self.mp2_energy = mp2_energy
if store_t2:
self.mp2_t2 = mp2_t2

def run_fci(self, *, store_fci_vec: bool = False) -> None:
"""Run FCI and store results."""
cas = mcscf.CASCI(self.scf, ncas=self.norb, nelecas=self.nelec)
mo = cas.sort_mo(self.active_space, mo_coeff=self.mo_coeff, base=0)
fci_energy, _, fci_vec, _, _ = cas.kernel(mo_coeff=mo)
self.fci_energy = fci_energy
if store_fci_vec:
self.fci_vec = fci_vec

def run_ccsd(
self,
t1: np.ndarray | None = None,
t2: np.ndarray | None = None,
*,
store_t1: bool = False,
store_t2: bool = False,
) -> None:
"""Run CCSD and store results."""
frozen = [i for i in range(self.norb) if i not in self.active_space]
ccsd_solver = cc.CCSD(self.scf, frozen=frozen)
ccsd_energy, ccsd_t1, ccsd_t2 = ccsd_solver.kernel(t1=t1, t2=t2)
self.ccsd_energy = ccsd_energy + self.hf_energy
if store_t1:
self.ccsd_t1 = ccsd_t1
if store_t2:
self.ccsd_t2 = ccsd_t2
3 changes: 2 additions & 1 deletion tests/gates/orbital_rotation_test.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -274,7 +274,8 @@ def test_apply_orbital_rotation_nitrogen():
symmetry="d2h",
)
active_space = range(2, mol.nao_nr())
mol_data = ffsim.MolecularData.from_mole(mol, active_space=active_space, ccsd=True)
mol_data = ffsim.MolecularData.from_mole(mol, active_space=active_space)
mol_data.run_ccsd(store_t1=True, store_t2=True)
norb = mol_data.norb
nelec = mol_data.nelec
t1 = mol_data.ccsd_t1
Expand Down

0 comments on commit eeb80ff

Please sign in to comment.