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 get_energy(atomz, with_ii=True):
atomzs = [atomz, atomz]
poss = torch.tensor([[-0.5, 0.0, 0.0], [0.5, 0.0, 0.0]], dtype=dtype)
mol = Mol((atomzs, poss), basis=basis, spin=0, dtype=dtype)
qc = HF(mol, restricted=True).run()
ene = qc.energy()
if with_ii:
ene = ene - mol.get_nuclei_energy()
return ene
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 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 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 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_grad_vxc(xccls, xcparams, atomz):
# check if the gradients w.r.t. vxc parameters are obtained correctly
poss = torch.tensor([[0.0, 0.0, 0.0]], dtype=dtype)
mol = Mol(([atomz], poss), basis="3-21G", dtype=dtype, grid=3)
mol.setup_grid()
def get_energy(*params):
xc = xccls(*params)
qc = KS(mol, xc=xc, restricted=False).run()
ene = qc.energy()
return ene
params = tuple(
torch.tensor(p, dtype=dtype).requires_grad_() for p in xcparams)
torch.autograd.gradcheck(get_energy, params)
def test_uhf_energy_same_as_rhf(atomzs, dist, energy_true, variational):
# test to see if uhf energy gets the same energy as rhf 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=basis, dtype=dtype)
qc = HF(mol, restricted=False, variational=variational).run()
ene = qc.energy()
assert torch.allclose(ene, ene * 0 + energy_true, rtol=1e-8)
def test_uhf_energy_atoms(atomz, spin, energy_true, variational):
# check the energy of atoms with non-0 spins
torch.manual_seed(123)
poss = torch.tensor([[0.0, 0.0, 0.0]], dtype=dtype)
mol = Mol(([atomz], poss), basis=basis, dtype=dtype, spin=spin)
qc = HF(mol, restricted=False, variational=variational).run()
ene = qc.energy()
assert torch.allclose(ene, ene * 0 + energy_true, atol=0.0, rtol=1e-7)
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_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_uhf_energy_mols(atomzs, dist, spin, energy_true, variational):
# check the energy of molecules with non-0 spins
# NOTE: O2 iteration gets into excited state (probably)
torch.manual_seed(123)
poss = torch.tensor([[-0.5, 0.0, 0.0], [0.5, 0.0, 0.0]], dtype=dtype) * dist
mol = Mol((atomzs, poss), basis=basis, dtype=dtype, spin=spin)
qc = HF(mol, restricted=False, variational=variational).run()
ene = qc.energy()
assert torch.allclose(ene, ene * 0 + energy_true, rtol=1e-8, atol=0.0)
def test_rks_grad_vext(xc, atomzs, dist, vext_p):
# check if the gradient w.r.t. vext is obtained correctly (only check 1st
# grad because we don't need 2nd grad at the moment)
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, grid=3)
mol.setup_grid()
rgrid = mol.get_grid().get_rgrid() # (ngrid, ndim)
rgrid_norm = torch.norm(rgrid, dim=-1) # (ngrid,)
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
vext_params = torch.tensor(vext_p, dtype=dtype).requires_grad_()
torch.autograd.gradcheck(get_energy, (vext_params, ))
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_hf_mols():
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**", dtype=dtype, spin=spin)
qc = HF(mol).run()
ene = qc.energy()
# see if grad and backward create memleak
grads = torch.autograd.grad(ene, (poss, ), create_graph=True)
ene.backward() # no create_graph here because of known pytorch's leak
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_mol_cache():
# test if cache is stored correctly
cache_fname = "_temp_cache.h5"
# remove the cache if exists
if os.path.exists(cache_fname):
os.remove(cache_fname)
moldesc = "H 0 0 0"
mol = Mol(moldesc, basis="3-21G").set_cache(cache_fname)
h = mol.get_hamiltonian()
h.build()
# read the stored file
with h5py.File(cache_fname, "r") as f:
olp_cache = torch.as_tensor(f["hamilton/overlap"])
olp = h.get_overlap().fullmatrix()
assert torch.allclose(olp, olp_cache)
# try again with different atom, if cache is set, then it should be same
# as previous (although it is a wrong result)
# TODO: raise a warning if mol properties do not match
moldesc1 = "Li 0 0 0"
mol1 = Mol(moldesc1, basis="3-21G").set_cache(cache_fname)
h1 = mol1.get_hamiltonian()
h1.build()
# remove the cache
if os.path.exists(cache_fname):
os.remove(cache_fname)
olp1 = h1.get_overlap().fullmatrix()
assert torch.allclose(olp, olp1)
def test_rhf_basis_inputs():
# test to see if the various basis inputs produce the same results
atomzs = [1, 1]
poss = torch.tensor([[-0.5, 0.0, 0.0], [0.5, 0.0, 0.0]], dtype=dtype) * 1.0
mol0 = Mol((atomzs, poss), basis=basis, dtype=dtype)
ene0 = HF(mol0, restricted=True).run().energy()
mol1 = Mol((atomzs, poss), basis=[basis, basis], dtype=dtype)
ene1 = HF(mol1, restricted=True).run().energy()
assert torch.allclose(ene0, ene1)
mol1 = Mol((atomzs, poss), basis=loadbasis("1:" + basis), dtype=dtype)
ene1 = HF(mol1, restricted=True).run().energy()
assert torch.allclose(ene0, ene1)
mol2 = Mol((atomzs, poss), basis={"H": basis}, dtype=dtype)
ene2 = HF(mol2, restricted=True).run().energy()
assert torch.allclose(ene0, ene2)
mol2 = Mol((atomzs, poss), basis={1: basis}, dtype=dtype)
ene2 = HF(mol2, restricted=True).run().energy()
assert torch.allclose(ene0, ene2)
mol2 = Mol((atomzs, poss), basis={1: loadbasis("1:3-21G")}, dtype=dtype)
ene2 = HF(mol2, restricted=True).run().energy()
assert torch.allclose(ene0, ene2)
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_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 test_rhf_energy_overcomplete(atomzs, dist, energy_true, variational):
# test to see if the energy calculated by DQC agrees with PySCF with
# overcomplete basis
torch.manual_seed(123)
# double the basis to make it overcomplete
basis0 = loadbasis(f"{atomzs[0]}:{basis}") * 2
basis1 = loadbasis(f"{atomzs[1]}:{basis}") * 2
poss = torch.tensor([[-0.5, 0.0, 0.0], [0.5, 0.0, 0.0]], dtype=dtype) * dist
# orthogonalize_basis = True to handle overcomplete
mol = Mol((atomzs, poss), basis=[basis0, basis1], dtype=dtype, orthogonalize_basis=True)
qc = HF(mol, restricted=True, variational=variational).run()
ene = qc.energy()
assert torch.allclose(ene, ene * 0 + energy_true, rtol=1e-7)
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 test_mol_grid(moldesc):
# default: level 4
m = Mol(moldesc, basis="6-311++G**", dtype=dtype)
m.setup_grid()
rgrid = m.get_grid().get_rgrid()
# only check the dimension and the type because the number of grid points
# can be changed
assert rgrid.shape[1] == 3
assert m.get_grid().coord_type == "cart"
def h2o_qc():
# run the self-consistent HF iteration for h2o
atomzs = torch.tensor([8, 1, 1], dtype=torch.int64)
# from CCCBDB (calculated geometry for H2O)
atomposs = torch.tensor([
[0.0, 0.0, 0.2156],
[0.0, 1.4749, -0.8625],
[0.0, -1.4749, -0.8625],
], dtype=dtype).requires_grad_()
efield = torch.zeros(3, dtype=dtype).requires_grad_()
grad_efield = torch.zeros((3, 3), dtype=dtype).requires_grad_()
efields = (efield, grad_efield)
mol = Mol(moldesc=(atomzs, atomposs), basis="3-21G", dtype=dtype, efield=efields)
qc = HF(mol).run()
return qc
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 test_rhf_energy(atomzs, dist, energy_true, variational):
# test to see if the energy calculated by DQC agrees with PySCF
torch.manual_seed(123)
# only set debugging mode only in one case to save time
if atomzs == [1, 1]:
xt.set_debug_mode(True)
poss = torch.tensor([[-0.5, 0.0, 0.0], [0.5, 0.0, 0.0]], dtype=dtype) * dist
mol = Mol((atomzs, poss), basis=basis, dtype=dtype)
qc = HF(mol, restricted=True, variational=variational).run()
ene = qc.energy()
assert torch.allclose(ene, ene * 0 + energy_true, rtol=1e-7)
if atomzs == [1, 1]:
xt.set_debug_mode(False)
def system1(request):
poss = torch.tensor([[0.0, 0.0, 0.8], [0.0, 0.0, -0.8]], dtype=dtype)
moldesc = ([1, 1], poss)
# untruncated grid is required to pass the hamiltonian tests
m = Mol(moldesc, basis="3-21G", dtype=dtype, grid=3,
orthogonalize_basis=request.param["basis_ortho"])
m.setup_grid()
hamilton = m.get_hamiltonian()
hamilton.build()
hamilton.setup_grid(m.get_grid())
return m