diff --git a/src/dxtb/basis/indexhelper.py b/src/dxtb/basis/indexhelper.py index 9796dad62..ade3d56db 100644 --- a/src/dxtb/basis/indexhelper.py +++ b/src/dxtb/basis/indexhelper.py @@ -129,6 +129,8 @@ def __init__( **_, ): super().__init__(device, dtype) + self.clear_cache() + self.unique_angular = unique_angular self.angular = angular self.atom_to_unique = atom_to_unique diff --git a/src/dxtb/basis/type.py b/src/dxtb/basis/type.py index 2b212f6e6..ba1c55244 100644 --- a/src/dxtb/basis/type.py +++ b/src/dxtb/basis/type.py @@ -72,14 +72,14 @@ def __init__( self, numbers: Tensor, par: Param, - angular: Tensor, + ihelp: IndexHelper, device: torch.device | None = None, dtype: torch.dtype | None = None, ) -> None: super().__init__(device, dtype) self.numbers = numbers self.meta = par.meta - self.angular = angular + self.ihelp = ihelp self.ngauss = get_elem_param( numbers, @@ -112,9 +112,12 @@ def create_cgtos(self) -> tuple[list[Tensor], list[Tensor]]: # maybe this can be batched too, but this loop is rather small # so it is probably not worth it - for i in range(self.angular.size(0)): + for i in range(self.ihelp.unique_angular.size(0)): alpha, coeff = slater.to_gauss( - self.ngauss[i], self.pqn[i], self.angular[i], self.slater[i] + self.ngauss[i], + self.pqn[i], + self.ihelp.unique_angular[i], + self.slater[i], ) # FIXME: @@ -138,7 +141,6 @@ def create_cgtos(self) -> tuple[list[Tensor], list[Tensor]]: def unique_shell_pairs( self, - ihelp: IndexHelper, mask: Tensor | None = None, uplo: Literal["n", "N", "u", "U", "l", "L"] = "l", ) -> tuple[Tensor, Tensor]: @@ -147,8 +149,6 @@ def unique_shell_pairs( Parameters ---------- - ihelp : IndexHelper - Helper class for indexing. mask : Tensor Mask for uplo : Literal["n", "N", "u", "U", "l", "L"] @@ -168,20 +168,20 @@ def unique_shell_pairs( """ # NOTE: Zeros correspond to padding only if batched. They have # meaning for single runs and cannot be ignored in this case. - sh2ush = ihelp.spread_shell_to_orbital(ihelp.shells_to_ushell) + sh2ush = self.ihelp.spread_shell_to_orbital(self.ihelp.shells_to_ushell) # FIXME: Maybe a bitwise operation is easier to understand? For now, # we convert unique shell indices to prime numbers to obtain unique # products upon multiplication (fundamental theorem of arithmetic). orbs = primes.to(self.device)[sh2ush] orbs = orbs.unsqueeze(-2) * orbs.unsqueeze(-1) - sh2orb = ihelp.spread_shell_to_orbital(ihelp.orbitals_per_shell) + sh2orb = self.ihelp.spread_shell_to_orbital(self.ihelp.orbitals_per_shell) # extra offset along only one dimension to distinguish (n, m) and # (m, n) of the same orbital block (e.g. 1x3 sp and 3x1 ps block) offset = 10000 - if ihelp.batched: + if self.ihelp.batched: orbs = torch.where( real_pairs(sh2ush), orbs + sh2orb.unsqueeze(-1) * offset, @@ -192,7 +192,7 @@ def unique_shell_pairs( # catch systems with one single orbital (e.g. Helium) if orbs.size(-1) == 1: - if ihelp.batched: + if self.ihelp.batched: raise NotImplementedError() return torch.tensor([[0]]), torch.tensor(1) @@ -223,7 +223,6 @@ def unique_shell_pairs( def to_bse( self, - ihelp: IndexHelper, qcformat: Literal["gaussian94", "nwchem"] = "nwchem", save: bool = False, overwrite: bool = False, @@ -273,10 +272,10 @@ def to_bse( if qcformat == "gaussian94": txt += f"{PSE[number]}\n" - shells = ihelp.shells_per_atom[i] + shells = self.ihelp.shells_per_atom[i] for _ in range(shells): alpha, coeff = slater.to_gauss( - self.ngauss[s], self.pqn[s], self.angular[s], self.slater[s] + self.ngauss[s], self.pqn[s], self.ihelp.angular[s], self.slater[s] ) if self.valence[s].item() is False: alpha, coeff = orthogonalize( @@ -286,7 +285,7 @@ def to_bse( alphas.append(alpha) coeffs.append(coeff) - l = angular2label[self.angular.tolist()[s]] + l = angular2label[self.ihelp.angular.tolist()[s]] if qcformat == "gaussian94": txt += f"{l} {len(alpha)} 1.00\n" elif qcformat == "nwchem": @@ -336,7 +335,7 @@ def to_bse( return fulltxt - def create_dqc(self, positions: Tensor, ihelp: IndexHelper) -> list[AtomCGTOBasis]: + def create_dqc(self, positions: Tensor) -> list[AtomCGTOBasis]: if self.numbers.ndim > 1: raise NotImplementedError("Batch mode not implemented.") @@ -350,11 +349,11 @@ def create_dqc(self, positions: Tensor, ihelp: IndexHelper) -> list[AtomCGTOBasi s = 0 for i, number in enumerate(self.numbers): bases: list[CGTOBasis] = [] - shells = ihelp.shells_per_atom[i] + shells = self.ihelp.shells_per_atom[i] for _ in range(shells): alpha, coeff = slater.to_gauss( - self.ngauss[s], self.pqn[s], ihelp.angular[s], self.slater[s] + self.ngauss[s], self.pqn[s], self.ihelp.angular[s], self.slater[s] ) # orthogonalize the H2s against the H1s @@ -367,7 +366,7 @@ def create_dqc(self, positions: Tensor, ihelp: IndexHelper) -> list[AtomCGTOBasi coeffs.append(coeff) cgto = CGTOBasis( - angmom=ihelp.angular.tolist()[s], # int! + angmom=self.ihelp.angular.tolist()[s], # int! alphas=alpha, coeffs=coeff, normalized=True, diff --git a/src/dxtb/constants/defaults.py b/src/dxtb/constants/defaults.py index 3e3defed6..d77a0161a 100644 --- a/src/dxtb/constants/defaults.py +++ b/src/dxtb/constants/defaults.py @@ -40,6 +40,9 @@ INTCUTOFF = 50.0 """Real-space cutoff (in Angstrom) for integral evaluation. (50.0)""" +INTDRIVER = "libcint" +"""Integral driver.""" + # SCF settings GUESS = "eeq" diff --git a/src/dxtb/integral/__init__.py b/src/dxtb/integral/__init__.py index ba877d072..3b1d9083a 100644 --- a/src/dxtb/integral/__init__.py +++ b/src/dxtb/integral/__init__.py @@ -5,4 +5,4 @@ from .mmd import overlap_gto from .mmd.explicit import mmd_explicit from .mmd.recursion import mmd_recursion -from .overlap import Overlap +from .overlap import * diff --git a/src/dxtb/integral/libcint/intor.py b/src/dxtb/integral/libcint/intor.py index e7b6d2f28..458b662ba 100644 --- a/src/dxtb/integral/libcint/intor.py +++ b/src/dxtb/integral/libcint/intor.py @@ -13,7 +13,11 @@ import numpy as np import torch -from dxtblibs import CGTO, CINT + +try: + from dxtblibs import CGTO, CINT +except ImportError: + pass from ..._types import Callable, Tensor from .namemanager import IntorNameManager diff --git a/src/dxtb/integral/libcint/namemanager.py b/src/dxtb/integral/libcint/namemanager.py index ac9371886..35ee4f37a 100644 --- a/src/dxtb/integral/libcint/namemanager.py +++ b/src/dxtb/integral/libcint/namemanager.py @@ -37,6 +37,9 @@ class IntorNameManager: "r_origj": (3,), "rr_origj": (9,), "rrr_origj": (27,), + "j": (3,), + "jj": (9,), + "jjj": (27,), }, ) op_comp = defaultdict( diff --git a/src/dxtb/integral/libcint/symmetry.py b/src/dxtb/integral/libcint/symmetry.py index 8f2d00ad7..3dd573616 100644 --- a/src/dxtb/integral/libcint/symmetry.py +++ b/src/dxtb/integral/libcint/symmetry.py @@ -5,6 +5,7 @@ In tight-binding, we do not require anything special here. Only S1 symmetry is of intereset. """ +from __future__ import annotations from abc import abstractmethod diff --git a/src/dxtb/integral/libcint/utils.py b/src/dxtb/integral/libcint/utils.py index b3f84989c..1059f480a 100644 --- a/src/dxtb/integral/libcint/utils.py +++ b/src/dxtb/integral/libcint/utils.py @@ -4,7 +4,6 @@ This module contains helpers required for calculating integrals with libcint. """ - from __future__ import annotations import ctypes diff --git a/src/dxtb/integral/libcint/wrapper.py b/src/dxtb/integral/libcint/wrapper.py index 39fde88cd..29f43c765 100644 --- a/src/dxtb/integral/libcint/wrapper.py +++ b/src/dxtb/integral/libcint/wrapper.py @@ -12,7 +12,11 @@ import numpy as np import torch -from dxtblibs import CINT + +try: + from dxtblibs import CINT # type: ignore +except ImportError: + pass from ..._types import Iterator, Tensor from ...basis import AtomCGTOBasis, CGTOBasis, IndexHelper diff --git a/src/dxtb/integral/overlap.py b/src/dxtb/integral/overlap.py index 9e7bee285..3b67deb90 100644 --- a/src/dxtb/integral/overlap.py +++ b/src/dxtb/integral/overlap.py @@ -17,6 +17,145 @@ from .mmd import overlap_gto, overlap_gto_grad from .utils import get_pairs, get_subblock_start +__all__ = ["Overlap", "OverlapLibcint"] + + +def create_libcint_wrapper( + bas: Basis, + positions: Tensor, + ihelp: IndexHelper, + spherical: bool = True, +) -> intor.LibcintWrapper: + atombases = bas.create_dqc(positions) + return intor.LibcintWrapper(atombases, ihelp, spherical=spherical) + + +class OverlapLibcint(TensorLike): + """ + Overlap integral from atomic orbitals. + """ + + def __init__( + self, + numbers: Tensor, + par: Param, + ihelp: IndexHelper, + driver: intor.LibcintWrapper | None = None, + device: torch.device | None = None, + dtype: torch.dtype | None = None, + ): + super().__init__(device, dtype) + self.numbers = numbers + self.unique = torch.unique(numbers) + self.par = par + self.ihelp = ihelp + self.driver = driver + + def build( + self, + positions: Tensor, + mask: Tensor | None = None, + ) -> Tensor: + """ + Overlap calculation of unique shells pairs, using the + McMurchie-Davidson algorithm. + + Parameters + ---------- + positions : Tensor + Cartesian coordinates of all atoms in the system. + mask : Tensor | None + Mask for positions to make batched computations easier. The overlap + does not work in a batched fashion. Hence, we loop over the batch + dimension and must remove the padding. Defaults to `None`, i.e., + `batch.deflate()` is used. + + Returns + ------- + Tensor + Overlap matrix. + """ + + # attempt to build the driver from scratch if none given yet + if self.driver is None: + bas = Basis( + self.numbers, + self.par, + self.ihelp, + dtype=self.dtype, + device=self.device, + ) + atombases = bas.create_dqc(positions) + self.driver = intor.LibcintWrapper(atombases, self.ihelp) + + if self.numbers.ndim > 1: + raise NotImplementedError() + # s = self._batch(wrapper, mask) + # s = self._single(wrapper) + else: + s = intor.overlap(self.driver) + norm = torch.pow(s.diagonal(dim1=-1, dim2=-2), -0.5) + s = torch.einsum("...ij,...i,...j->...ij", s, norm, norm) + + return s + + def get_gradient(self, positions: Tensor, mask: Tensor | None = None): + """ + Overlap gradient calculation of unique shells pairs, using the + McMurchie-Davidson algorithm. + + Parameters + ---------- + positions : Tensor + Cartesian coordinates of all atoms in the system. + mask : Tensor | None + Mask for positions to make batched computations easier. The overlap + does not work in a batched fashion. Hence, we loop over the batch + dimension and must remove the padding. Defaults to `None`, i.e., + `batch.deflate()` is used. + + Returns + ------- + Tensor + Overlap gradient of shape `(nb, norb, norb, 3)`. + """ + # attempt to build the driver from scratch if none given yet + if self.driver is None: + bas = Basis( + self.numbers, + self.par, + self.ihelp, + dtype=self.dtype, + device=self.device, + ) + atombases = bas.create_dqc(positions) + self.driver = intor.LibcintWrapper(atombases, self.ihelp) + + if self.numbers.ndim > 1: + raise NotImplementedError() + # grad = self._batch(wrapper, mask) + # grad = self._single(wrapper) + else: + s = intor.overlap(self.driver) + norm = torch.pow(s.diagonal(dim1=-1, dim2=-2), -0.5) + + # (3, norb, norb) + grad = intor.int1e("ipovlp", self.driver) + + # normalize and move xyz dimension to last, which is required for + # the reduction (only works with extra dimension in last) + grad = -torch.einsum("...xij,...i,...j->...ijx", grad, norm, norm) + + return grad + + def _single(self, wrapper: intor.LibcintWrapper) -> Tensor: + raise NotImplementedError() + + def _batch( + self, wrapper: intor.LibcintWrapper, mask: Tensor | None = None + ) -> Tensor: + raise NotImplementedError() + class OverlapFunction(Protocol): """ @@ -101,8 +240,6 @@ def __init__( ihelp: IndexHelper, uplo: Literal["n", "N", "u", "U", "l", "L"] = "l", cutoff: Tensor | float | int | None = defaults.INTCUTOFF, - driver: Literal["dxtb", "libcint"] = "libcint", - spherical: bool = True, device: torch.device | None = None, dtype: torch.dtype | None = None, ): @@ -112,14 +249,16 @@ def __init__( self.par = par self.ihelp = ihelp self.cutoff = cutoff if cutoff is None else cutoff * units.AA2AU - self.driver = driver - self.spherical = spherical if uplo not in ("n", "N", "u", "U", "l", "L"): raise ValueError(f"Unknown option for `uplo` chosen: '{uplo}'.") self.uplo = uplo.casefold() # type: ignore - def build(self, positions: Tensor, mask: Tensor | None = None) -> Tensor: + def build( + self, + positions: Tensor, + mask: Tensor | None = None, + ) -> Tensor: """ Overlap calculation of unique shells pairs, using the McMurchie-Davidson algorithm. @@ -177,42 +316,10 @@ def get_gradient(self, positions: Tensor, mask: Tensor | None = None): return grad def _single(self, func: OverlapFunction, positions: Tensor) -> Tensor: - if self.driver == "libcint": - import time - - start_time = time.perf_counter() - - bas = Basis( - self.numbers, - self.par, - self.ihelp.angular, - dtype=self.dtype, - device=self.device, - ) - - bas_time = time.perf_counter() - print(f"basis {bas_time - start_time:.3f}") - - atombases = bas.create_dqc(positions, self.ihelp) - wrapper = intor.LibcintWrapper(atombases, self.ihelp, spherical=self.spherical) # type: ignore - - wrap_time = time.perf_counter() - print(f"wrap {wrap_time - bas_time:.3f}") - - s = intor.overlap(wrapper) - norm = torch.pow(s.diagonal(dim1=-1, dim2=-2), -0.5) - s = torch.einsum("...ij,...i,...j->...ij", s, norm, norm) - - ovlp_time = time.perf_counter() - print(f"ovlp {ovlp_time - wrap_time:.3f}") - print(f"total {ovlp_time - start_time:.3f}") - - return s - bas = Basis( self.unique, self.par, - self.ihelp.unique_angular, + self.ihelp, dtype=self.dtype, device=self.device, ) @@ -240,7 +347,7 @@ def _batch( bas = Basis( torch.unique(nums), self.par, - ihelp.unique_angular, + ihelp, dtype=self.dtype, device=self.device, ) @@ -380,7 +487,7 @@ def overlap( dist = torch.cdist(pos, pos, compute_mode="donot_use_mm_for_euclid_dist") mask = (dist < cutoff) & (dist > 0.1) - umap, n_unique_pairs = bas.unique_shell_pairs(ihelp, mask=mask, uplo=uplo) + umap, n_unique_pairs = bas.unique_shell_pairs(mask=mask, uplo=uplo) # overlap calculation ovlp = torch.zeros(*umap.shape, dtype=positions.dtype, device=positions.device) @@ -486,7 +593,7 @@ def overlap_gradient( dist = torch.cdist(pos, pos, compute_mode="donot_use_mm_for_euclid_dist") mask = (dist < cutoff) & (dist > 0.1) - umap, n_unique_pairs = bas.unique_shell_pairs(ihelp, mask=mask, uplo=uplo) + umap, n_unique_pairs = bas.unique_shell_pairs(mask=mask, uplo=uplo) # overlap calculation ds = torch.zeros((3, *umap.shape), dtype=positions.dtype, device=positions.device) @@ -537,5 +644,4 @@ def overlap_gradient( ds = torch.triu(ds, diagonal=1) - torch.tril(ds.mT) # (3, norb, norb) -> (norb, norb, 3) - print("ds.shape", ds.shape) return torch.einsum("xij->ijx", ds) diff --git a/src/dxtb/mol/external/_pyscf.py b/src/dxtb/mol/external/_pyscf.py index 4538c301a..0d75b800e 100644 --- a/src/dxtb/mol/external/_pyscf.py +++ b/src/dxtb/mol/external/_pyscf.py @@ -27,7 +27,7 @@ import numpy as np try: - import pyscf # type: ignore + from pyscf import gto # type: ignore except ImportError as e: raise ImportError("PySCF is not installed") from e @@ -39,7 +39,7 @@ # Turn off PySCF's normalization since dxtb's normalization is different, # requiring a separate normalization anyway. But this way, the results are # equal to dxtb's libcint wrapper and we can immediately compare integrals. -pyscf.gto.mole.NORMALIZE_GTO = False +gto.mole.NORMALIZE_GTO = False __all__ = ["PyscfMol", "M"] @@ -69,7 +69,7 @@ def M( return mol -class PyscfMol(pyscf.gto.Mole): +class PyscfMol(gto.Mole): """ Pyscf's molecule representation that can be created upon passing only `numbers` and `positions`. Note that the basis set is created from a @@ -110,7 +110,7 @@ def __init__( p = Path(path, f"{number:02d}.nwchem") with open(p, encoding="utf8") as f: content = f.read() - basis[f"{PSE[number]}"] = pyscf.gto.parse(content) + basis[f"{PSE[number]}"] = gto.parse(content) self.atom = atom self.basis = basis diff --git a/src/dxtb/xtb/calculator.py b/src/dxtb/xtb/calculator.py index df07203f2..e5e06d2ad 100644 --- a/src/dxtb/xtb/calculator.py +++ b/src/dxtb/xtb/calculator.py @@ -4,12 +4,13 @@ from __future__ import annotations import warnings +from functools import wraps import torch from .. import ncoord, scf -from .._types import Any, Sequence, Tensor, TensorLike -from ..basis import IndexHelper +from .._types import Any, Callable, Sequence, Tensor, TensorLike +from ..basis import Basis, IndexHelper from ..classical import ( Classical, ClassicalList, @@ -22,7 +23,8 @@ from ..coulomb import new_es2, new_es3 from ..data import cov_rad_d3 from ..dispersion import Dispersion, new_dispersion -from ..integral import Overlap +from ..integral import Overlap, OverlapLibcint +from ..integral.libcint import LibcintWrapper from ..interaction import Interaction, InteractionList from ..param import Param, get_elem_angular from ..utils import Timers, ToleranceWarning @@ -30,6 +32,23 @@ from ..xtb.h0 import Hamiltonian from .h0 import Hamiltonian +__all__ = ["Calculator"] + + +def use_intdriver(driver_arg: int = 1) -> Callable: + def decorator(fcn: Callable) -> Callable: + @wraps(fcn) + def wrap(self: Calculator, *args: Any, **kwargs: Any) -> Any: + if self.opts["integral_driver"] == "libcint": + self.init_intdriver(args[driver_arg]) + + result = fcn(self, *args, **kwargs) + return result + + return wrap + + return decorator + class Result(TensorLike): """ @@ -258,6 +277,12 @@ def __init__( # setup calculator options opts = opts if opts is not None else {} + + # FIXME: Batch mode for libcint + drv = opts.get("integral_driver", defaults.INTDRIVER) + if numbers.ndim > 1: + drv = "dxtb" + self.opts = { "fwd_options": { "damp": opts.get("damp", defaults.DAMP), @@ -286,6 +311,7 @@ def __init__( "exclude": opts.get("exclude", defaults.EXCLUDE), "guess": opts.get("guess", defaults.GUESS), "spin": opts.get("spin", defaults.SPIN), + "integral_driver": drv, } # set tolerances separately to catch unreasonably small values @@ -294,7 +320,6 @@ def __init__( self.ihelp = IndexHelper.from_numbers(numbers, get_elem_angular(par.element)) self.hamiltonian = Hamiltonian(numbers, par, self.ihelp, **dd) - self.overlap = Overlap(numbers, par, self.ihelp, **dd) # setup self-consistent contributions es2 = ( @@ -309,7 +334,7 @@ def __init__( ) if interaction is None: - self.interactions = InteractionList(es2, es3, interaction) + self.interactions = InteractionList(es2, es3) else: self.interactions = InteractionList(es2, es3, *interaction) @@ -341,6 +366,15 @@ def __init__( self.timer.stop("setup calculator") + if self.opts["integral_driver"] == "libcint": + self.overlap = OverlapLibcint(numbers, par, self.ihelp, **dd) + self.basis = Basis(numbers, par, self.ihelp, **dd) + else: + self.overlap = Overlap(numbers, par, self.ihelp, **dd) + + self._intdriver = None + self._positions = None + def set_option(self, name: str, value: Any) -> None: if name not in self.opts: raise KeyError(f"Option '{name}' does not exist.") @@ -363,6 +397,33 @@ def set_tol(self, name: str, value: float) -> None: self.opts["fwd_options"][name] = value + def driver(self, positions: Tensor) -> LibcintWrapper: + # save current positions to check + self._positions = positions.detach().clone() + + atombases = self.basis.create_dqc(positions) + return LibcintWrapper(atombases, self.ihelp) + + def init_intdriver(self, positions: Tensor): + if self.opts["integral_driver"] != "libcint": + return + + diff = 0 + # create LibcintWrapper if it does not exist yet + if self._intdriver is None: + self._intdriver = self.driver(positions) + else: + assert self._positions is not None + + diff = (self._positions - positions).abs().sum() + if diff > 1e-10: + self._intdriver = self.driver(positions) + + if isinstance(self.overlap, OverlapLibcint): + if self.overlap.driver is None or diff > 1e-10: + self.overlap.driver = self._intdriver + + @use_intdriver() def singlepoint( self, numbers: Tensor, diff --git a/test/test_a_memory_leak/test_scf.py b/test/test_a_memory_leak/test_scf.py index f8c0363f1..adbf0cfb0 100644 --- a/test/test_a_memory_leak/test_scf.py +++ b/test/test_a_memory_leak/test_scf.py @@ -18,6 +18,7 @@ from .util import has_memleak_tensor opts = {"verbosity": 0, "maxiter": 50, "exclude": ["rep", "disp", "hal"]} +repeats = 5 device = None @@ -54,7 +55,12 @@ def fcn(): if create_graph is True: energy.backward() - assert not has_memleak_tensor(fcn, gccollect=run_gc), "Memory leak detected" + leak = has_memleak_tensor(fcn, repeats=repeats, gccollect=run_gc) + + # run garbage collector to avoid leaks across other tests + gc.collect() + + assert not leak, "Memory leak detected" @pytest.mark.filterwarnings("ignore") @@ -63,7 +69,9 @@ def fcn(): # FIXME: not calling the garbage collector also causes a memory leak # @pytest.mark.parametrize("run_gc", [False, True]) @pytest.mark.parametrize("create_graph", [False, True]) -def test_fulltracking(dtype: torch.dtype, run_gc: bool, create_graph: bool) -> None: +def skip_test_fulltracking( + dtype: torch.dtype, run_gc: bool, create_graph: bool +) -> None: dd: DD = {"device": device, "dtype": dtype} def fcn(): @@ -87,7 +95,7 @@ def fcn(): if create_graph is True: energy.backward() - leak = has_memleak_tensor(fcn, gccollect=run_gc) + leak = has_memleak_tensor(fcn, repeats=repeats, gccollect=run_gc) # run garbage collector to avoid leaks across other tests gc.collect() diff --git a/test/test_a_memory_leak/util.py b/test/test_a_memory_leak/util.py index fa8b2f482..22876207a 100644 --- a/test/test_a_memory_leak/util.py +++ b/test/test_a_memory_leak/util.py @@ -56,10 +56,12 @@ def _show_memsize(fcn, ntries: int = 10, gccollect: bool = False): gc.collect() size = _get_tensor_memory() - print(f"{i + 1:3d} iteration: {size - size0:.7f} MiB of tensors") + print(f"{i + 1:3d} iteration: {size - size0:.16f} MiB of tensors") -def has_memleak_tensor(fcn: Callable, gccollect: bool = False) -> bool: +def has_memleak_tensor( + fcn: Callable, repeats: int = 5, gccollect: bool = False +) -> bool: """ Assert no tensor memory leak when calling the function. @@ -82,9 +84,7 @@ def has_memleak_tensor(fcn: Callable, gccollect: bool = False) -> bool: gc.collect() size = _get_tensor_memory() - print(size, size0) if size0 != size: - ntries = 10 - _show_memsize(fcn, ntries, gccollect=gccollect) + _show_memsize(fcn, repeats, gccollect=gccollect) return size0 != size diff --git a/test/test_basis/test_export.py b/test/test_basis/test_export.py index 2902741cc..fde05acda 100644 --- a/test/test_basis/test_export.py +++ b/test/test_basis/test_export.py @@ -30,9 +30,9 @@ def test_export( assert False ihelp = IndexHelper.from_numbers(numbers, get_elem_angular(par.element)) - bas = Basis(numbers, par, ihelp.unique_angular, dtype=dtype) + bas = Basis(numbers, par, ihelp, dtype=dtype) - txt = bas.to_bse(ihelp, qcformat=qcformat) + txt = bas.to_bse(qcformat=qcformat) # check with saved basis files root = Path(__file__).parents[2] diff --git a/test/test_basis/test_general.py b/test/test_basis/test_general.py index c3df65cfb..068fd0522 100644 --- a/test/test_basis/test_general.py +++ b/test/test_basis/test_general.py @@ -69,7 +69,7 @@ def test_fail_higher_orbital_trafo(): number = torch.tensor([86]) ihelp = IndexHelper.from_numbers(number, get_elem_angular(par.element)) - bas = Basis(number, par, ihelp.unique_angular) + bas = Basis(number, par, ihelp) alpha, coeff = bas.create_cgtos() j = torch.tensor(5) diff --git a/test/test_dipole/__init__.py b/test/test_dipole/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/test/test_dipole/samples.py b/test/test_dipole/samples.py new file mode 100644 index 000000000..15e83ee8f --- /dev/null +++ b/test/test_dipole/samples.py @@ -0,0 +1,8 @@ +""" +Molecules for testing. +""" +from __future__ import annotations + +from ..molecules import mols + +samples = mols diff --git a/test/test_dipole/test_integral.py b/test/test_dipole/test_integral.py new file mode 100644 index 000000000..f630b6fd1 --- /dev/null +++ b/test/test_dipole/test_integral.py @@ -0,0 +1,73 @@ +""" +Test overlap from libcint. + +We cannot compare with our internal integrals or the reference integrals from +tblite, because the sorting of the p-orbitals are different. +""" +from __future__ import annotations + +from math import sqrt + +import pytest +import torch + +from dxtb._types import DD, Tensor +from dxtb.basis import Basis, IndexHelper +from dxtb.integral.libcint import LibcintWrapper, intor +from dxtb.param import GFN1_XTB as par +from dxtb.param import get_elem_angular +from dxtb.utils import numpy_to_tensor + +try: + from dxtb.mol.external._pyscf import M + + pyscf = True +except ImportError: + pyscf = False + +from .samples import samples + +sample_list = ["H2", "LiH", "Li2", "H2O", "S", "SiH4", "MB16_43_01", "C60"] + +device = None + + +def snorm(overlap: Tensor) -> Tensor: + return torch.pow(overlap.diagonal(dim1=-1, dim2=-2), -0.5) + + +@pytest.mark.skipif(pyscf is False, reason="PySCF not installed") +@pytest.mark.parametrize("dtype", [torch.float, torch.double]) +@pytest.mark.parametrize("name", sample_list) +def test_single(dtype: torch.dtype, name: str) -> None: + dd: DD = {"device": device, "dtype": dtype} + tol = sqrt(torch.finfo(dtype).eps) * 1e-2 + + sample = samples[name] + numbers = sample["numbers"].to(device) + positions = sample["positions"].to(**dd) + ihelp = IndexHelper.from_numbers(numbers, get_elem_angular(par.element)) + bas = Basis(numbers, par, ihelp, **dd) + atombases = bas.create_dqc(positions) + + # dxtb's libcint overlap + wrapper = LibcintWrapper(atombases, ihelp) + dxtb_dpint = intor.int1e("j", wrapper) + + # pyscf reference integral + mol = M(numbers, positions, parse_arg=False) + pyscf_dpint = numpy_to_tensor(mol.intor("int1e_r_origj"), **dd) + + assert dxtb_dpint.shape == pyscf_dpint.shape + assert pytest.approx(pyscf_dpint, abs=tol) == dxtb_dpint + + # normalize + norm = snorm(intor.overlap(wrapper)) + dxtb_dpint = torch.einsum("xij,i,j->xij", dxtb_dpint, norm, norm) + pyscf_dpint = torch.einsum("xij,i,j->xij", pyscf_dpint, norm, norm) + + print("") + print(dxtb_dpint) + + assert dxtb_dpint.shape == pyscf_dpint.shape + assert pytest.approx(pyscf_dpint, abs=tol) == dxtb_dpint diff --git a/test/test_libcint/test_grad_pos.py b/test/test_libcint/test_grad_pos.py index fd074d212..de5e5ae1e 100644 --- a/test/test_libcint/test_grad_pos.py +++ b/test/test_libcint/test_grad_pos.py @@ -35,13 +35,13 @@ def gradchecker( positions = sample["positions"].to(**dd) ihelp = IndexHelper.from_numbers(numbers, get_elem_angular(par.element)) - bas = Basis(numbers, par, ihelp.angular, **dd) + bas = Basis(numbers, par, ihelp, **dd) # variables to be differentiated positions.requires_grad_(True) def func(pos: Tensor) -> Tensor: - atombases = bas.create_dqc(pos, ihelp) + atombases = bas.create_dqc(pos) wrapper = intor.LibcintWrapper(atombases, ihelp, spherical=False) return intor.overlap(wrapper) diff --git a/test/test_libcint/test_overlap.py b/test/test_libcint/test_overlap.py index 8bd06f8bf..e5612e7d2 100644 --- a/test/test_libcint/test_overlap.py +++ b/test/test_libcint/test_overlap.py @@ -47,15 +47,15 @@ def test_single(dtype: torch.dtype, name: str) -> None: numbers = sample["numbers"].to(device) positions = sample["positions"].to(**dd) ihelp = IndexHelper.from_numbers(numbers, get_elem_angular(par.element)) - bas = Basis(numbers, par, ihelp.angular, **dd) - atombases = bas.create_dqc(positions, ihelp) + bas = Basis(numbers, par, ihelp, **dd) + atombases = bas.create_dqc(positions) # dxtb's libcint overlap wrapper = LibcintWrapper(atombases, ihelp) dxtb_overlap = intor.overlap(wrapper) # pyscf reference overlap ("sph" needed, implicit in dxtb) - mol = M(numbers, positions) + mol = M(numbers, positions, parse_arg=False) pyscf_overlap = numpy_to_tensor(mol.intor("int1e_ovlp"), **dd) assert dxtb_overlap.shape == pyscf_overlap.shape @@ -85,13 +85,13 @@ def test_grad(dtype: torch.dtype, name: str) -> None: positions = sample["positions"].to(**dd) ihelp = IndexHelper.from_numbers(numbers, get_elem_angular(par.element)) - bas = Basis(numbers, par, ihelp.angular, **dd) - atombases = bas.create_dqc(positions, ihelp) + bas = Basis(numbers, par, ihelp, **dd) + atombases = bas.create_dqc(positions) wrapper = LibcintWrapper(atombases, ihelp) int1 = intor.int1e("ipovlp", wrapper) - mol = M(numbers, positions) + mol = M(numbers, positions, parse_arg=False) int2 = numpy_to_tensor(mol.intor("int1e_ipovlp"), **dd) assert int1.shape == int2.shape diff --git a/test/test_libcint/test_overlap_grad.py b/test/test_libcint/test_overlap_grad.py index 5af65df2f..9abdf4186 100644 --- a/test/test_libcint/test_overlap_grad.py +++ b/test/test_libcint/test_overlap_grad.py @@ -32,8 +32,8 @@ def explicit(name: str, dd: DD, tol: float) -> None: ref = load_from_npz(ref_overlap, name, **dd) ihelp = IndexHelper.from_numbers(numbers, get_elem_angular(par.element)) - bas = Basis(numbers, par, ihelp.angular, **dd) - atombases = bas.create_dqc(positions, ihelp) + bas = Basis(numbers, par, ihelp, **dd) + atombases = bas.create_dqc(positions) wrapper = intor.LibcintWrapper(atombases, ihelp) s = intor.overlap(wrapper) @@ -42,7 +42,7 @@ def explicit(name: str, dd: DD, tol: float) -> None: # (3, norb, norb) grad = intor.int1e("ipovlp", wrapper) - # normalize and switch move xyz dimension to last, which is required for + # normalize and move xyz dimension to last, which is required for # the reduction (only works with extra dimension in last) grad = torch.einsum("...xij,...i,...j->...ijx", grad, norm, norm) @@ -82,8 +82,8 @@ def autograd(name: str, dd: DD, tol: float) -> None: positions.requires_grad_(True) ihelp = IndexHelper.from_numbers(numbers, get_elem_angular(par.element)) - bas = Basis(numbers, par, ihelp.angular, **dd) - atombases = bas.create_dqc(positions, ihelp) + bas = Basis(numbers, par, ihelp, **dd) + atombases = bas.create_dqc(positions) wrapper = intor.LibcintWrapper(atombases, ihelp) s = intor.overlap(wrapper) diff --git a/test/test_overlap/test_gradient_grad_pos.py b/test/test_overlap/test_gradient_grad_pos.py index 15252a67f..543bfec0c 100644 --- a/test/test_overlap/test_gradient_grad_pos.py +++ b/test/test_overlap/test_gradient_grad_pos.py @@ -33,7 +33,7 @@ def gradchecker( positions = sample["positions"].to(**dd) ihelp = IndexHelper.from_numbers(numbers, get_elem_angular(par.element)) - bas = Basis(torch.unique(numbers), par, ihelp.unique_angular, **dd) + bas = Basis(torch.unique(numbers), par, ihelp, **dd) # variables to be differentiated positions.requires_grad_(True) diff --git a/test/test_overlap/test_overlap_pairs.py b/test/test_overlap/test_overlap_pairs.py index 9fa33ca10..cc3f053d3 100644 --- a/test/test_overlap/test_overlap_pairs.py +++ b/test/test_overlap/test_overlap_pairs.py @@ -83,7 +83,7 @@ def test_overlap_higher_orbitals(dtype: torch.dtype): number = torch.tensor([86]) ihelp = IndexHelper.from_numbers(number, get_elem_angular(par.element)) - bas = Basis(number, par, ihelp.unique_angular) + bas = Basis(number, par, ihelp) alpha, coeff = bas.create_cgtos() ai = alpha[0] diff --git a/test/test_scf/test_hess.py b/test/test_scf/test_hess.py index 904793e28..ee78579ec 100644 --- a/test/test_scf/test_hess.py +++ b/test/test_scf/test_hess.py @@ -22,6 +22,7 @@ "xitorch_fatol": 1.0e-8, "xitorch_xatol": 1.0e-8, "verbosity": 0, + "integral_driver": "dxtb", } device = None