def test_vext(atomzs, dist):
# check if the external potential gives the correct density profile
# (energy is different because the energy of xc is not integral of potentials
# times density)
def get_dens(qc):
dm = qc.aodm()
rgrid = mol.get_grid().get_rgrid()
dens = mol.get_hamiltonian().aodm2dens(dm, rgrid)
return dens
with xt.enable_debug():
poss = torch.tensor([[-0.5, 0.0, 0.0], [0.5, 0.0, 0.0]],
dtype=dtype) * dist
mol = Mol((atomzs, poss), basis="3-21G", dtype=dtype)
xc = get_xc("lda_x")
qc = KS(mol, xc=xc).run()
# get the density profile
dens = get_dens(qc)
ldax_pot = xc.get_vxc(ValGrad(value=dens)).value
# ldax_pot2 = -0.7385587663820223 * (4.0 / 3) * dens ** (1.0 / 3)
# assert torch.allclose(ldax_pot, ldax_pot2)
# calculate the energy with the external potential
mol2 = Mol((atomzs, poss), basis="3-21G", dtype=dtype, vext=ldax_pot)
qc2 = KS(mol2, xc=None).run()
dens2 = get_dens(qc2)
# make sure the densities agree
assert torch.allclose(dens, dens2)
def test_no_xc(atomzs, dist, restricted):
# check if xc == None produces the same result as "0*lda"
with xt.enable_debug():
poss = torch.tensor([[-0.5, 0.0, 0.0], [0.5, 0.0, 0.0]],
dtype=dtype) * dist
mol1 = Mol((atomzs, poss), basis="3-21G", dtype=dtype)
qc1 = KS(mol1, xc="0*lda_x", restricted=restricted).run()
ene1 = qc1.energy()
mol2 = Mol((atomzs, poss), basis="3-21G", dtype=dtype)
qc2 = KS(mol2, xc=None, restricted=restricted).run()
ene2 = qc2.energy()
assert torch.allclose(ene1, ene2)
def test_uks_energy_same_as_rks(xc, atomzs, dist, energy_true):
# test to see if uks energy gets the same energy as rks for non-polarized systems
poss = torch.tensor([[-0.5, 0.0, 0.0], [0.5, 0.0, 0.0]],
dtype=dtype) * dist
mol = Mol((atomzs, poss), basis="6-311++G**", dtype=dtype)
qc = KS(mol, xc=xc, restricted=False).run()
ene = qc.energy()
assert torch.allclose(ene, ene * 0 + energy_true)
def test_rks_energy(xc, atomzs, dist, energy_true, grid):
# test to see if the energy calculated by DQC agrees with PySCF
# for this test only we test for different types of grids to see if any error is raised
poss = torch.tensor([[-0.5, 0.0, 0.0], [0.5, 0.0, 0.0]],
dtype=dtype) * dist
mol = Mol((atomzs, poss), basis="6-311++G**", dtype=dtype, grid=grid)
qc = KS(mol, xc=xc, restricted=True).run()
ene = qc.energy()
assert torch.allclose(ene, ene * 0 + energy_true)
def test_rks_energy_df(xc, atomzs, dist, energy_true, grid):
# test to see if the energy calculated by DQC agrees with PySCF using
# density fitting
poss = torch.tensor([[-0.5, 0.0, 0.0], [0.5, 0.0, 0.0]],
dtype=dtype) * dist
mol = Mol((atomzs, poss), basis="6-311++G**", dtype=dtype, grid=grid)
mol.densityfit(method="coulomb", auxbasis="def2-sv(p)-jkfit")
qc = KS(mol, xc=xc, restricted=True).run()
ene = qc.energy()
assert torch.allclose(ene, ene * 0 + energy_true)
def get_energy(atomzs):
poss = torch.tensor([[-0.5, 0.0, 0.0], [0.5, 0.0, 0.0]], dtype=dtype)
mol = Mol((atomzs, poss),
basis="6-311++G**",
spin=0,
dtype=dtype,
grid="sg3")
qc = KS(mol, xc="lda_x", restricted=True).run()
ene = qc.energy() - mol.get_nuclei_energy()
return ene
def get_energy(vext_params):
vext = rgrid_norm * rgrid_norm * vext_params # (ngrid,)
mol = Mol((atomzs, poss),
basis="3-21G",
dtype=dtype,
grid=3,
vext=vext)
qc = KS(mol, xc=xc, restricted=True).run()
ene = qc.energy()
return ene
def test_uks_energy_atoms(xc, atomz, spin, energy_true):
# check the energy of atoms with non-0 spins
poss = torch.tensor([[0.0, 0.0, 0.0]], dtype=dtype)
mol = Mol(([atomz], poss),
basis="6-311++G**",
grid=4,
dtype=dtype,
spin=spin)
qc = KS(mol, xc=xc, restricted=False).run()
ene = qc.energy()
assert torch.allclose(ene, ene * 0 + energy_true, atol=0.0, rtol=1e-6)
def test_uks_energy_mols(xc, atomzs, dist, spin, energy_true):
# check the energy of molecules with non-0 spins
poss = torch.tensor([[-0.5, 0.0, 0.0], [0.5, 0.0, 0.0]],
dtype=dtype) * dist
mol = Mol((atomzs, poss),
basis="6-311++G**",
grid=3,
dtype=dtype,
spin=spin)
qc = KS(mol, xc=xc, restricted=False).run()
ene = qc.energy()
assert torch.allclose(ene, ene * 0 + energy_true, rtol=1e-6, atol=0.0)
def test_uks_energy_mols_df(xc, atomzs, dist, spin, energy_true):
# check the energy of molecules with non-0 spins with density fitting
poss = torch.tensor([[-0.5, 0.0, 0.0], [0.5, 0.0, 0.0]],
dtype=dtype) * dist
mol = Mol((atomzs, poss),
basis="6-311++G**",
grid=3,
dtype=dtype,
spin=spin)
mol.densityfit(method="coulomb", auxbasis="def2-sv(p)-jkfit")
qc = KS(mol, xc=xc, restricted=False).run()
ene = qc.energy()
assert torch.allclose(ene, ene * 0 + energy_true, rtol=1e-6, atol=0.0)
def get_energy(alphas, coeffs):
basis = {
'H': [
CGTOBasis(angmom=0,
alphas=alphas,
coeffs=coeffs,
normalized=True)
]
}
moldesc = "H 0.0000 0.0000 0.0000; H 0.0000 0.0000 1.1163"
mol = Mol(moldesc, basis=basis, dtype=dtype, grid=3)
# qc = HF(mol, restricted=True).run()
qc = KS(mol, xc="lda_x", restricted=True).run()
return qc.energy()
def test_rks_energy(xc, atomzs, dist, energy_true, grid, variational):
# test to see if the energy calculated by DQC agrees with PySCF
# for this test only we test for different types of grids to see if any error is raised
if xc == "mgga_x_scan":
if atomzs == [1, 1]:
pytest.xfail("Psi4 and PySCF don't converge")
poss = torch.tensor([[-0.5, 0.0, 0.0], [0.5, 0.0, 0.0]],
dtype=dtype) * dist
mol = Mol((atomzs, poss), basis="6-311++G**", dtype=dtype, grid=grid)
qc = KS(mol, xc=xc, restricted=True,
variational=variational).run(fwd_options={"verbose": True})
ene = qc.energy()
# < 1 kcal/mol
assert torch.allclose(ene, ene * 0 + energy_true, atol=1.3e-3, rtol=0)
def _test_ks_mols():
# setting up xc
a = torch.tensor(-0.7385587663820223, dtype=dtype, requires_grad=True)
p = torch.tensor(1.3333333333333333, dtype=dtype, requires_grad=True)
xc = PseudoLDA(a=a, p=p)
poss = torch.tensor([[-0.5, 0.0, 0.0], [0.5, 0.0, 0.0]], dtype=dtype) * dist
poss = poss.requires_grad_()
mol = Mol((atomzs, poss), basis="6-311++G**", grid=3, dtype=dtype, spin=spin)
qc = KS(mol, xc=xc).run()
ene = qc.energy()
# see if grad and backward create memleak
grads = torch.autograd.grad(ene, (poss, a, p), create_graph=True)
ene.backward() # no create_graph here because of known pytorch's leak
def atest_pbc_rks_energy(xc, atomzs, spin, alattice, energy_true, grid):
# test to see if the energy calculated by DQC agrees with PySCF
# for this test only we test for different types of grids to see if any error is raised
alattice = torch.as_tensor(alattice, dtype=dtype)
poss = torch.tensor([[0.0, 0.0, 0.0]], dtype=dtype)
mol = Sol((atomzs, poss),
basis="3-21G",
spin=spin,
alattice=alattice,
dtype=dtype,
grid=grid)
mol.densityfit(method="gdf", auxbasis="def2-sv(p)-jkfit")
qc = KS(mol, xc=xc, restricted=False).run()
ene = qc.energy()
assert torch.allclose(ene, ene * 0 + energy_true)
def test_uks_energy_atoms(xc, atomz, spin, energy_true):
# check the energy of atoms with non-0 spins
if xc == "mgga_x_scan":
if atomz == 3:
pytest.xfail("No benchmark converges")
poss = torch.tensor([[0.0, 0.0, 0.0]], dtype=dtype)
mol = Mol(([atomz], poss),
basis="6-311++G**",
grid="sg3",
dtype=dtype,
spin=spin)
qc = KS(mol, xc=xc, restricted=False).run()
ene = qc.energy()
# < 1 kcal/mol
assert torch.allclose(ene, ene * 0 + energy_true, atol=1e-3, rtol=0)
def dqcelrepxc(atom: str, spin: int = 0, xc: str = "lda_x", basis: str = BASIS):
# returns the electron repulsion and xc operator using DQC
mol = Mol(atom, spin=spin, basis=basis, dtype=DTYPE, grid=4)
qc = KS(mol, xc=xc)
hamilt = mol.get_hamiltonian()
if spin == 0:
# set dm to be an identity matrix
dm = torch.eye(hamilt.nao, dtype=DTYPE)
velrepxc = hamilt.get_vxc(dm) + hamilt.get_elrep(dm)
return velrepxc.fullmatrix()
else:
dmu = torch.eye(hamilt.nao, dtype=DTYPE)
dm = SpinParam(u=dmu, d=dmu)
vxc = hamilt.get_vxc(dm)
elrep = hamilt.get_elrep(dm.u + dm.d)
return torch.cat(((vxc.u + elrep).fullmatrix().unsqueeze(0),
(vxc.d + elrep).fullmatrix().unsqueeze(0)), dim=0)
def test_pbc_rks_energy(xc, atomzs, spin, alattice, energy_true, grid):
# test to see if the energy calculated by DQC agrees with PySCF
# for this test only we test for different types of grids to see if any error is raised
alattice = torch.as_tensor(alattice, dtype=dtype)
poss = torch.tensor([[0.0, 0.0, 0.0]], dtype=dtype)
mol = Sol((atomzs, poss),
basis="3-21G",
spin=spin,
alattice=alattice,
dtype=dtype,
grid=grid)
mol.densityfit(method="gdf", auxbasis="def2-sv(p)-jkfit")
qc = KS(mol, xc=xc, restricted=False).run()
ene = qc.energy()
energy_true = ene * 0 + energy_true
# rtol is a bit too high. Investigation indicates that this might be due to
# the difference in grid (i.e. changing from sg2 to sg3 changes the values),
# not from the eta of the compensating charge (i.e. changing it does not
# change the value)
# TODO: make a better grid
assert torch.allclose(ene, energy_true, rtol=1e-3)
def test_rks_energy_df(xc, atomzs, dist, energy_true, grid):
# test to see if the energy calculated by DQC agrees with PySCF using
# density fitting
if xc == "mgga_x_scan":
if atomzs == [1, 1]:
pytest.xfail("Psi4 and PySCF don't converge")
for lowmem in [False, True]: # simulating low memory condition
if lowmem:
init_value = config.THRESHOLD_MEMORY
config.THRESHOLD_MEMORY = 1000000 # 1 MB
poss = torch.tensor([[-0.5, 0.0, 0.0], [0.5, 0.0, 0.0]],
dtype=dtype) * dist
mol = Mol((atomzs, poss), basis="6-311++G**", dtype=dtype, grid=grid)
mol.densityfit(method="coulomb", auxbasis="def2-sv(p)-jkfit")
qc = KS(mol, xc=xc, restricted=True).run()
ene = qc.energy()
# error to be < 1 kcal/mol
assert torch.allclose(ene, ene * 0 + energy_true, atol=1.1e-3, rtol=0)
if lowmem:
# restore the value
config.THRESHOLD_MEMORY = init_value
def get_energy(vext_params):
vext = rgrid_norm * rgrid_norm * vext_params # (ngrid,)
qc = KS(mol, xc=xc, vext=vext, restricted=True).run()
ene = qc.energy()
return ene
def run_ks_forward(moldesc, basis="6-311++G**", xc="lda_x", grid="sg3"):
# run a simple KS energy calculation
mol = Mol(moldesc, basis=basis, grid=grid)
qc = KS(mol, xc=xc).run()
ene = qc.energy()
return ene
mol.densityfit(method="gdf", auxbasis="def2-sv(p)-jkfit")
qc = KS(mol, xc=xc, restricted=False).run()
ene = qc.energy()
assert torch.allclose(ene, ene * 0 + energy_true)
if __name__ == "__main__":
import time
xc = "lda_x"
basis = "3-21G"
atomzs = [3]
spin = 1
alattice = np.array([[1., 1., -1.], [-1., 1., 1.], [1., -1., 1.]
]) * 0.5 * 6.6329387300636
grid = "sg2"
# test to see if the energy calculated by DQC agrees with PySCF
# for this test only we test for different types of grids to see if any error is raised
alattice = torch.as_tensor(alattice, dtype=dtype)
poss = torch.tensor([[0.0, 0.0, 0.0]], dtype=dtype)
mol = Sol((atomzs, poss),
basis="3-21G",
spin=spin,
alattice=alattice,
dtype=dtype,
grid=grid)
mol.densityfit(method="gdf", auxbasis="def2-sv(p)-jkfit")
qc = KS(mol, xc=xc, restricted=False).run()
ene = qc.energy()
assert torch.allclose(ene, ene * 0 + energy_true)
def get_energy(dist_tensor):
poss_tensor = torch.tensor([[-0.5, 0.0, 0.0], [0.5, 0.0, 0.0]],
dtype=dtype) * dist_tensor
mol = Mol((atomzs, poss_tensor), basis="3-21G", dtype=dtype, grid=3)
qc = KS(mol, xc=xc, restricted=True).run(bck_options=bck_options)
return qc.energy()
def get_energy(*params):
xc = xccls(*params)
qc = KS(mol, xc=xc, restricted=False).run()
ene = qc.energy()
return ene
def get_energy(efield):
mol = Mol(moldesc, basis="3-21G", dtype=dtype, efield=efield)
qc = KS(mol, xc="lda_x").run()
ene = qc.energy()
return ene