def test_pbc_becke_grid_dvol(rgrid_integrator, rgrid_transform):
dtype = torch.float64
nr = 40
prec = 7
radgrid = RadialGrid(nr, rgrid_integrator, rgrid_transform, dtype=dtype)
sphgrid = LebedevGrid(radgrid, prec=prec)
atompos = torch.tensor([[0.0, 0.0, 0.0], [1.0, 0.0, 0.0]], dtype=dtype)
natoms = atompos.shape[0]
lattice = Lattice(torch.eye(3, dtype=dtype) * 3)
grid = PBCBeckeGrid([sphgrid] * natoms, atompos, lattice=lattice)
dvol = grid.get_dvolume() # (ngrid,)
rgrid = grid.get_rgrid() # (ngrid, ndim)
ls = lattice.get_lattice_ls(rcut=5) # (nls, ndim)
atomposs = (atompos.unsqueeze(1) + ls).reshape(-1, 1,
3) # (natoms * nls, ndim)
# test gaussian integration
fcn = torch.exp(-((rgrid - atomposs)**2).sum(dim=-1)).sum(dim=0) # (ngrid)
# fcn = rgrid[:, 0] * 0 + 1
int1 = (fcn * dvol).sum()
val1 = int1 * 0 + 2 * np.pi**1.5 # analytical function
# TODO: rtol is relatively large, maybe inspect the Becke integration grid?
assert torch.allclose(int1, int1 * 0 + val1, rtol=1e-2)
def pbc_h1():
# get the hamiltonian for pbc system
# setup the environment
pos = torch.tensor([0.0, 0.0, 0.0], dtype=dtype)
atomz = 2
atom = "He"
bases = loadbasis("%d:3-21G" % atomz, dtype=dtype)
atombases = [AtomCGTOBasis(atomz=atomz, bases=bases, pos=pos)]
kpts = torch.tensor([[0.0, 0.0, 0.0], [0.1, 0.2, 0.15]], dtype=dtype)
# setup the lattice
a = torch.tensor([[1.0, 0.0, 0.0], [0.0, 1.0, 0.0], [0.0, 0.0, 1.0]], dtype=dtype) * 3
latt = Lattice(a)
# setup the auxbasis
auxbases = loadbasis("%d:def2-sv(p)-jkfit" % atomz, dtype=dtype)
atomauxbases = [AtomCGTOBasis(atomz=atomz, bases=auxbases, pos=pos)]
df = DensityFitInfo(method="gdf", auxbases=atomauxbases)
# build the hamiltonian
h = HamiltonCGTO_PBC(atombases, latt=latt, df=df, kpts=kpts,
lattsum_opt={"precision": 1e-8})
h.build()
# pyscf system build
mol = pyscf.pbc.gto.C(atom="%s 0 0 0" % atom, a=a.numpy(), basis="3-21G", unit="Bohr", spin=0)
df = pyscf.pbc.df.GDF(mol, kpts=kpts.detach().numpy())
df.auxbasis = "def2-svp-jkfit"
df.build()
return h, df
def __init__(
self,
soldesc: Union[str, Tuple[AtomZsType, AtomPosType]],
alattice: torch.Tensor,
basis: Union[str, List[CGTOBasis], List[str],
List[List[CGTOBasis]]],
*,
grid: Union[int, str] = "sg3",
spin: Optional[ZType] = None,
lattsum_opt: Optional[Union[PBCIntOption, Dict]] = None,
dtype: torch.dtype = torch.float64,
device: torch.device = torch.device('cpu'),
):
self._dtype = dtype
self._device = device
self._grid_inp = grid
self._grid: Optional[BaseGrid] = None
charge = 0 # we can't have charged solids for now
# get the AtomCGTOBasis & the hamiltonian
# atomzs: (natoms,) dtype: torch.int or dtype for floating point
# atompos: (natoms, ndim)
atomzs, atompos = parse_moldesc(soldesc, dtype, device)
allbases = _parse_basis(atomzs, basis) # list of list of CGTOBasis
atombases = [
AtomCGTOBasis(atomz=atz, bases=bas, pos=atpos)
for (atz, bas, atpos) in zip(atomzs, allbases, atompos)
]
self._atombases = atombases
self._atompos = atompos # (natoms, ndim)
self._atomzs = atomzs # (natoms,) int-type
nelecs_tot: torch.Tensor = torch.sum(atomzs)
# get the number of electrons and spin and orbital weights
nelecs, spin, frac_mode = _get_nelecs_spin(nelecs_tot, spin, charge)
assert not frac_mode, "Fractional Z mode for pbc is not supported"
_orb_weights, _orb_weights_u, _orb_weights_d = _get_orb_weights(
nelecs, spin, frac_mode, dtype, device)
# initialize cache
self._cache = Cache()
# save the system's properties
self._spin = spin
self._charge = charge
self._numel = nelecs
self._orb_weights = _orb_weights
self._orb_weights_u = _orb_weights_u
self._orb_weights_d = _orb_weights_d
self._lattice = Lattice(alattice)
self._lattsum_opt = PBCIntOption.get_default(lattsum_opt)
def test_lattice():
# testing various properties of the lattice object
a = torch.tensor([[1.0, 0.0, 0.0], [1.0, 1.0, 0.0], [2.0, 1.0, 1.0]], dtype=dtype)
b = torch.inverse(a.T) * (2 * np.pi)
latt = Lattice(a)
assert torch.allclose(latt.lattice_vectors(), a)
assert torch.allclose(latt.recip_vectors(), b)
assert torch.allclose(latt.volume(), torch.det(a))
# check the lattice_ls function returns the correct shape
ls0 = latt.get_lattice_ls(rcut=1.5) # (nb, ndim)
assert ls0.ndim == 2
assert ls0.shape[1] == 3
# check the ls has no repeated coordinates
ls0_dist = torch.norm(ls0[:, None, :] - ls0[None, :, :], dim=-1) # (nb, nb)
ls0_dist = ls0_dist + torch.eye(ls0_dist.shape[0], dtype=dtype)
assert torch.all(ls0_dist.abs() > 1e-9)