Refactor integral code

docs/source/03_for_developers/errors.rst

@@ -50,3 +50,18 @@ the error is actually raised because the mixer in `xitorch` produces a
 Apparently, this occurs if the convergence criteria are too strict in single
 precision. The solution is to increase the convergence criteria to more than
 1e-6.
+
+
+RuntimeError: clone is not supported by NestedIntSymNode
+--------------------------------------------------------
+
+This is a bug in PyTorch 2.3.0 and 2.3.1 (see
+`PyTorch #128607 <<https://github.com/pytorch/pytorch/issues/128607>`__).
+To avoid this error, manually import `torch._dynamo` in the code. For example:
+
+.. code-block:: python
+
+    from tad_mctc._version import __tversion__
+
+    if __tversion__ in ((2, 3, 0), (2, 3, 1)):
+        import torch._dynamo

examples/limitation_xitorch.py

@@ -20,21 +20,30 @@
 from pathlib import Path
 
 import torch
+from tad_mctc._version import __tversion__
 from tad_mctc.io import read
 
 import dxtb
 from dxtb.typing import DD
 
+if __tversion__ in ((2, 3, 0), (2, 3, 1)):
+    import torch._dynamo
+
+
 dd: DD = {"device": torch.device("cpu"), "dtype": torch.double}
 
 f = Path(__file__).parent / "molecules" / "lih.xyz"
 numbers, positions = read.read_from_path(f, **dd)
 charge = read.read_chrg_from_path(f, **dd)
 
-opts = {"verbosity": 3, "scf_mode": "nonpure"}
+opts = {"verbosity": 0, "scf_mode": "nonpure"}
 
 ######################################################################
 
 calc = dxtb.Calculator(numbers, dxtb.GFN1_XTB, opts=opts, **dd)
 pos = positions.clone().requires_grad_(True)
-hess = calc.hessian(pos, chrg=charge, use_functorch=True)
+
+try:
+    calc.hessian(pos, chrg=charge, use_functorch=True)
+except RuntimeError as e:
+    print(f"RuntimeError:\n{str(e)}")

examples/profiling/batch-vs-seq-nicotine.py

@@ -29,7 +29,14 @@
 dxtb.timer.cuda_sync = False
 
 # read molecule from file
-f = Path(__file__).parent / "molecules" / "nicotine.xyz"
+p = Path(__file__).parent
+
+if "molecules" not in [x.name for x in p.iterdir()]:
+    p = p.parent
+
+f = p / "molecules" / "nicotine.xyz"
+
+
 numbers, positions = read.read_from_path(f, **dd)
 charge = read.read_chrg_from_path(f, **dd)
 

examples/profiling/importing.py

@@ -47,3 +47,5 @@
 print("dxtb", t2 - t1)
 print("Param", t3 - t2)
 print("scipy", t4 - t3)
+
+del scipy, torch, GFN1_XTB

examples/run-all.sh

@@ -0,0 +1,19 @@
+#!/bin/bash
+
+set -e
+
+# search recursively for all python files
+for example in $(find . -name "*.py"); do
+    title="Running $example"
+
+    # match the length of the title
+    line=$(printf '=%.0s' $(seq 1 ${#title}))
+
+    echo "$line"
+    echo "$title"
+    echo "$line"
+
+    python3 "$example"
+
+    printf "\n\n"
+done

setup.cfg

@@ -41,10 +41,10 @@ install_requires =
     numpy<2
     pydantic>=2.0.0
     scipy
-    tad-dftd3>=0.2.2
-    tad-dftd4>=0.1.0
-    tad-libcint>=0.0.1
-    tad-mctc>=0.1.5
+    tad-dftd3>=0.3.0
+    tad-dftd4>=0.2.0
+    tad-libcint>=0.1.0
+    tad-mctc>=0.2.0
     tad-multicharge
     tomli
     tomli-w

src/dxtb/_src/calculators/properties/vibration/analysis.py

@@ -112,8 +112,9 @@ def to_unit(self, value: Literal["freqs", "modes"], unit: str) -> Tensor:
         if value == "freqs":
             return self._convert(self.freqs, unit, self.converter_freqs)
 
+        # TODO: Implement conversion for normal modes (required?)
         # if value == "modes":
-        #   return self._convert(self.modes, unit, self.converter_modes)
+        #    return self._convert(self.modes, unit, self.converter_modes)
 
         raise ValueError(f"Unsupported value for conversion: {value}")
 
@@ -143,7 +144,7 @@ def _get_rotational_modes(mass: Tensor, positions: Tensor):
     _, paxes = storch.eighb(im)
 
     # make z-axis rotation vector with smallest moment of inertia
-    # w = torch.flip(w, [-1])
+    # _ = torch.flip(_, [-1])
     paxes = torch.flip(paxes, [-1])
     ex, ey, ez = paxes.mT
 
@@ -288,7 +289,7 @@ def vib_analysis(
     # by 2π (ω = √λ / 2π) in order to immediately get frequencies in
     # Hartree: E = hbar * ω with hbar = 1 in atomic units. Dividing by 2π
     # effectively converts from angular frequency (ω) to the frequency in
-    # cycles per second (ν, Hz), which would requires the following
+    # cycles per second (ν, Hz), which would require the following
     # conversion to cm^-1: 1e-2 / units.CODATA.c / units.AU2SECOND.
     sgn = torch.sign(force_const_au)
     freqs_au = torch.sqrt(torch.abs(force_const_au)) * sgn

src/dxtb/_src/calculators/types/analytical.py

@@ -168,7 +168,7 @@
         timer.stop("ograd")
 
         overlap = self.cache["overlap"]
-        assert isinstance(overlap, ints.types.Overlap)
+        assert isinstance(overlap, ints.types.OverlapIntegral)
         assert overlap.matrix is not None
 
         # density matrix
@@ -200,7 +200,7 @@
         assert self.integrals.hcore is not None
 
         cn = ncoord.cn_d3(self.numbers, positions)
-        dedcn, dedr = self.integrals.hcore.integral.get_gradient(
+        dedcn, dedr = self.integrals.hcore.get_gradient(
             positions,
             overlap.matrix,
             overlap_grad,
@@ -410,7 +410,7 @@
             self.ihelp,
             self.opts.scf,
             intmats,
-            self.integrals.hcore.integral.refocc,
+            self.integrals.hcore.refocc,
         )
 
         timer.stop("SCF")
@@ -462,7 +462,7 @@
         )
 
         cn = ncoord.cn_d3(self.numbers, positions)
-        dedcn, dedr = self.integrals.hcore.integral.get_gradient(
+        dedcn, dedr = self.integrals.hcore.get_gradient(
             positions,
             intmats.overlap,
             overlap_grad,

src/dxtb/_src/calculators/types/autograd.py

@@ -155,8 +155,18 @@
             # pylint: disable=import-outside-toplevel
             from tad_mctc.autograd import jacrev
 
-            # jacrev requires a scalar from `self.energy`!
-            deriv = jacrev(self.energy, argnums=0)(positions, chrg, spin, **kwargs)
+            try:
+                # jacrev requires a scalar from `self.energy`!
+                deriv = jacrev(self.energy, argnums=0)(positions, chrg, spin, **kwargs)
+            except RuntimeError as e:
+                if "clone is not supported by NestedIntSymNode" in str(e):
+                    raise RuntimeError(
+                        "This is a bug in PyTorch 2.3.0 and 2.3.1. "
+                        "Try manually importing `torch._dynamo` before running "
+                        "any calculation."
+                    ) from e
+                raise
+
             assert isinstance(deriv, Tensor)
 
         elif grad_mode == "row":

src/dxtb/_src/calculators/types/base.py

@@ -133,9 +133,9 @@ def __init__(
         raman_depol: Tensor | None = None,
         #
         hcore: Tensor | None = None,
-        overlap: ints.types.Overlap | None = None,
-        dipint: ints.types.Dipole | None = None,
-        quadint: ints.types.Quadrupole | None = None,
+        overlap: ints.types.OverlapIntegral | None = None,
+        dipint: ints.types.DipoleIntegral | None = None,
+        quadint: ints.types.QuadrupoleIntegral | None = None,
         #
         bond_orders: Tensor | None = None,
         coefficients: Tensor | None = None,
@@ -404,6 +404,12 @@ class BaseCalculator(GetPropertiesMixin, TensorLike):
     results: dict[str, Any]
     """Results container."""
 
+    _ncalcs: int
+    """
+    Number of calculations performed with the calculator. Helpful for keeping
+    track of cache hits and actual new calculations.
+    """
+
     def __init__(
         self,
         numbers: Tensor,
@@ -481,6 +487,7 @@ def __init__(
         if self.opts.batch_mode == 0 and numbers.ndim > 1:
             self.opts.batch_mode = 1
 
+        # TODO: Should the IndexHelper be a singleton?
         self.ihelp = IndexHelper.from_numbers(numbers, par, self.opts.batch_mode)
 
         ################
@@ -573,6 +580,7 @@ def __init__(
                     "interaction."
                 )
             self.opts.ints.level = max(labels.INTLEVEL_DIPOLE, self.opts.ints.level)
+
         if efield_grad.LABEL_EFIELD_GRAD in self.interactions.labels:
             if self.opts.ints.level < labels.INTLEVEL_DIPOLE:
                 OutputHandler.warn(
@@ -582,21 +590,29 @@ def __init__(
                 )
             self.opts.ints.level = max(labels.INTLEVEL_QUADRUPOLE, self.opts.ints.level)
 
-        # setup integral
-        driver = self.opts.ints.driver
-        self.integrals = ints.Integrals(
-            numbers, par, self.ihelp, driver=driver, intlevel=self.opts.ints.level, **dd
-        )
+        # setup integral driver and integral container
+        mgr = ints.DriverManager(self.opts.ints.driver, **dd)
+        mgr.create_driver(numbers, par, self.ihelp)
+
+        self.integrals = ints.Integrals(mgr, intlevel=self.opts.ints.level, **dd)
 
         if self.opts.ints.level >= labels.INTLEVEL_OVERLAP:
-            self.integrals.hcore = ints.types.HCore(numbers, par, self.ihelp, **dd)
-            self.integrals.overlap = ints.types.Overlap(driver=driver, **dd)
+            self.integrals.hcore = ints.factories.new_hcore(
+                numbers, par, self.ihelp, **dd
+            )
+            self.integrals.overlap = ints.factories.new_overlap(
+                driver=mgr.driver_type, **dd
+            )
 
         if self.opts.ints.level >= labels.INTLEVEL_DIPOLE:
-            self.integrals.dipole = ints.types.Dipole(driver=driver, **dd)
+            self.integrals.dipole = ints.factories.new_dipint(
+                driver=mgr.driver_type, **dd
+            )
 
         if self.opts.ints.level >= labels.INTLEVEL_QUADRUPOLE:
-            self.integrals.quadrupole = ints.types.Quadrupole(driver=driver, **dd)
+            self.integrals.quadrupole = ints.factories.new_quadint(
+                driver=mgr.driver_type, **dd
+            )
 
         OutputHandler.write_stdout("done\n", v=4)
 

src/dxtb/_src/calculators/types/energy.py

@@ -23,15 +23,11 @@
 
 from __future__ import annotations
 
-import logging
-
 import torch
 from tad_mctc.convert import any_to_tensor
 from tad_mctc.io.checks import content_checks, shape_checks
 
-from dxtb import OutputHandler
-from dxtb import integrals as ints
-from dxtb import labels
+from dxtb import OutputHandler, labels
 from dxtb._src import scf
 from dxtb._src.constants import defaults
 from dxtb._src.integral.container import IntegralMatrices
@@ -216,7 +212,7 @@
         # While one can theoretically skip the core Hamiltonian, the
         # current implementation does not account for this case because the
         # reference occupation is necessary for the SCF procedure.
-        if self.integrals.hcore is None or self.integrals.hcore.matrix is None:
+        if self.integrals.hcore is None:
             raise NotImplementedError(
                 "Core Hamiltonian missing. Skipping the Core Hamiltonian in "
                 "the SCF is currently not supported. Please increase the "
@@ -261,7 +257,7 @@
             self.ihelp,
             self.opts.scf,
             intmats,
-            self.integrals.hcore.integral.refocc,
+            self.integrals.hcore.refocc,
         )
 
         timer.stop("SCF")
@@ -326,12 +322,16 @@
 
         if kwargs.get("store_fock", copts.fock):
             self.cache["fock"] = scf_results["hamiltonian"]
+
         if kwargs.get("store_hcore", copts.hcore):
             self.cache["hcore"] = self.integrals.hcore
+
         if kwargs.get("store_overlap", copts.overlap):
             self.cache["overlap"] = self.integrals.overlap
+
         if kwargs.get("store_dipole", copts.dipole):
             self.cache["dipint"] = self.integrals.dipole
+
         if kwargs.get("store_quadrupole", copts.quadrupole):
             self.cache["quadint"] = self.integrals.quadrupole
 
@@ -403,18 +403,24 @@
         """
         self.singlepoint(positions, chrg, spin, **kwargs)
 
+        ovlp_msg = (
+            "Overlap matrix not found in cache. The overlap is not saved "
+            "per default. Enable saving either via the calculator options "
+            "(`calc.opts.cache.store.overlap = True`) or by passing the "
+            "`store_overlap=True` keyword argument to called method, e.g., "
+            "`calc.energy(positions, store_overlap=True)"
+        )
+
         overlap = self.cache["overlap"]
         if overlap is None:
-            raise RuntimeError(
-                "Overlap matrix not found in cache. The overlap is not saved "
-                "per default. Enable saving either via the calculator options "
-                "(`calc.opts.cache.store.overlap = True`) or by passing the "
-                "`store_overlap=True` keyword argument to called method, e.g., "
-                "`calc.energy(positions, store_overlap=True)"
-            )
+            raise RuntimeError(ovlp_msg)
+
+        # pylint: disable=import-outside-toplevel
+        from dxtb._src.integral.types import OverlapIntegral
 
-        assert isinstance(overlap, ints.types.Overlap)
-        assert overlap.matrix is not None
+        assert isinstance(overlap, OverlapIntegral)
+        if overlap.matrix is None:
+            raise RuntimeError(ovlp_msg)
 
         density = self.cache["density"]
         if density is None: