Container for (multipolar) potentials

.pre-commit-config.yaml

@@ -12,13 +12,13 @@ repos:
         language_version: python3
 
   - repo: https://github.com/asottile/setup-cfg-fmt
-    rev: v2.2.0
+    rev: v2.4.0
     hooks:
       - id: setup-cfg-fmt
         args: [--include-version-classifiers, --max-py-version, "3.11"]
 
   - repo: https://github.com/asottile/pyupgrade
-    rev: v3.3.1
+    rev: v3.9.0
     hooks:
       - id: pyupgrade
         args: [--py37-plus, --keep-runtime-typing]
@@ -31,6 +31,6 @@ repos:
         args: ["--profile", "black", "--filter-files"]
 
   - repo: https://github.com/psf/black
-    rev: 23.1.0
+    rev: 23.7.0
     hooks:
       - id: black

setup.cfg

@@ -65,6 +65,8 @@ dev =
     pytest
     pytest-random-order
     tox
+pyscf =
+    pyscf
 
 [options.package_data]
 dxtb =

src/dxtb/_types.py

@@ -112,35 +112,20 @@ class Molecule(TypedDict):
     """Tensor of 3D coordinates of shape (n, 3)"""
 
 
-class SCFResult(TypedDict):
-    """Collection of SCF result variables."""
+class PotentialData(TypedDict):
+    """Shape and label information of Potentials."""
 
-    charges: Tensor
-    """Self-consistent orbital-resolved Mulliken partial charges."""
+    mono: torch.Size | None
+    """Shape of the monopolar potential."""
 
-    coefficients: Tensor
-    """LCAO-MO coefficients (eigenvectors of Fockian)."""
+    dipole: torch.Size | None
+    """Shape of the dipolar potential."""
 
-    density: Tensor
-    """Density matrix."""
+    quad: torch.Size | None
+    """Shape of the quadrupolar potential."""
 
-    emo: Tensor
-    """Energy of molecular orbitals (sorted by increasing energy)."""
-
-    energy: Tensor
-    """Energies of the self-consistent contributions (interactions)."""
-
-    fenergy: Tensor
-    """Atom-resolved electronic free energy from fractional occupation."""
-
-    hamiltonian: Tensor
-    """Full Hamiltonian matrix (H0 + H1)."""
-
-    occupation: Tensor
-    """Orbital occupations."""
-
-    potential: Tensor
-    """Self-consistent orbital-resolved potential."""
+    label: list[str] | str | None
+    """Labels for the interactions contributing to the potential."""
 
 
 class Slicers(TypedDict):

src/dxtb/constants/defaults.py

@@ -127,3 +127,9 @@
 
 TORCH_DEVICE_CHOICES = ["cpu", "cuda"]
 """List of possible choices for `TORCH_DEVICE`."""
+
+PAD = 0
+"""Defaults value to indicate padding."""
+
+PADNZ = -999999
+"""Default non-zero value to indicate padding."""

src/dxtb/integral/__init__.py

@@ -2,6 +2,7 @@
 Functions for calculation of overlap with McMurchie-Davidson algorithm.
 """
 
+from .container import Integrals
 from .mmd import overlap_gto
 from .mmd.explicit import mmd_explicit
 from .mmd.recursion import mmd_recursion

src/dxtb/integral/container.py

@@ -0,0 +1,141 @@
+"""
+Integral container
+==================
+
+A class that acts as a container for integrals.
+"""
+from __future__ import annotations
+
+import torch
+
+from .._types import Tensor, TensorLike
+
+__all__ = ["Integrals"]
+
+
+class Integrals(TensorLike):
+    """
+    Integral container.
+    """
+
+    __slots__ = ["_hcore", "_overlap", "_dipole", "_quad"]
+
+    def __init__(
+        self,
+        hcore: Tensor | None = None,
+        overlap: Tensor | None = None,
+        dipole: Tensor | None = None,
+        quad: Tensor | None = None,
+        device: torch.device | None = None,
+        dtype: torch.dtype | None = None,
+    ):
+        super().__init__(device, dtype)
+
+        self._hcore = hcore
+        self._overlap = overlap
+        self._dipole = dipole
+        self._quad = quad
+
+    @property
+    def hcore(self) -> Tensor | None:
+        return self._hcore
+
+    @hcore.setter
+    def hcore(self, hcore: Tensor) -> None:
+        self._hcore = hcore
+        self.checks()
+
+    @property
+    def overlap(self) -> Tensor | None:
+        return self._overlap
+
+    @overlap.setter
+    def overlap(self, overlap: Tensor) -> None:
+        self._overlap = overlap
+        self.checks()
+
+    @property
+    def dipole(self) -> Tensor | None:
+        return self._dipole
+
+    @dipole.setter
+    def dipole(self, dipole: Tensor) -> None:
+        self._dipole = dipole
+        self.checks()
+
+    @property
+    def quad(self) -> Tensor | None:
+        return self._quad
+
+    @quad.setter
+    def quad(self, quad: Tensor) -> None:
+        self._quad = quad
+        self.checks()
+
+    def checks(self) -> None:
+        """
+        Checks the shapes of the tensors.
+
+        Expected shapes:
+        - hcore and overlap: (batch_size, nao, nao) or (nao, nao)
+        - dipole: (batch_size, 3, nao, nao) or (3, nao, nao)
+        - quad: (batch_size, 6, nao, nao) or (6, nao, nao)
+
+        Raises
+        ------
+        ValueError:
+            If any of the tensors have incorrect shapes or inconsistent batch
+            sizes.
+        """
+        nao = None
+        batch_size = None
+
+        for name in ["hcore", "overlap", "dipole", "quad"]:
+            tensor = getattr(self, "_" + name)
+            if tensor is None:
+                continue
+
+            if name in ["hcore", "overlap"]:
+                if len(tensor.shape) not in [2, 3]:
+                    raise ValueError(
+                        f"Tensor '{name}' must have 2 or 3 dimensions. "
+                        f"Got {len(tensor.shape)}."
+                    )
+                if len(tensor.shape) == 3:
+                    if batch_size is not None and tensor.shape[0] != batch_size:
+                        raise ValueError(
+                            f"Tensor '{name}' has a different batch size. "
+                            f"Expected {batch_size}, got {tensor.shape[0]}."
+                        )
+                    batch_size = tensor.shape[0]
+                nao = tensor.shape[-1]
+            else:  # dipole or quad
+                if len(tensor.shape) not in [3, 4]:
+                    raise ValueError(
+                        f"Tensor '{name}' must have 3 or 4 dimensions. "
+                        f"Got {len(tensor.shape)}."
+                    )
+                if len(tensor.shape) == 4:
+                    if batch_size is not None and tensor.shape[0] != batch_size:
+                        raise ValueError(
+                            f"Tensor '{name}' has a different batch size. "
+                            f"Expected {batch_size}, got {tensor.shape[0]}."
+                        )
+                    batch_size = tensor.shape[0]
+                nao = tensor.shape[-2]
+
+            if tensor.shape[-2:] != (nao, nao):
+                raise ValueError(
+                    f"Tensor '{name}' last two dimensions should be "
+                    f"(nao, nao). Got {tensor.shape[-2:]}."
+                )
+            if name == "dipole" and tensor.shape[-3] != 3:
+                raise ValueError(
+                    f"Tensor '{name}' third to last dimension should be 3. "
+                    f"Got {tensor.shape[-3]}."
+                )
+            if name == "quad" and tensor.shape[-3] != 6:
+                raise ValueError(
+                    f"Tensor '{name}' third to last dimension should be 6. "
+                    f"Got {tensor.shape[-3]}."
+                )

src/dxtb/interaction/__init__.py

@@ -5,3 +5,4 @@
 from .base import Interaction
 from .external import ElectricField, new_efield
 from .list import InteractionList
+from .potential import Potential

src/dxtb/interaction/base.py

@@ -1,3 +1,4 @@
+# pylint: disable=assignment-from-none
 """
 Provides base class for interactions in the extended tight-binding Hamiltonian.
 The `Interaction` class is not purely abstract as its methods return zero.
@@ -8,6 +9,7 @@
 
 from .._types import Any, Tensor, TensorLike, TensorOrTensors
 from ..basis import IndexHelper
+from .potential import Potential
 
 __all__ = ["Interaction"]
 
@@ -98,8 +100,11 @@ def get_cache(
         return self.Cache()
 
     def get_potential(
-        self, charges: Tensor, cache: Interaction.Cache, ihelp: IndexHelper
-    ) -> Tensor:
+        self,
+        charges: Tensor,
+        cache: Interaction.Cache,
+        ihelp: IndexHelper,
+    ) -> Potential:
         """
         Compute the potential from the charges, all quantities are orbital-resolved.
 
@@ -117,14 +122,24 @@ def get_potential(
         Tensor
             Potential vector for each orbital partial charge.
         """
+
+        # monopole potential: shell-resolved
         qsh = ihelp.reduce_orbital_to_shell(charges)
         vsh = self.get_shell_potential(qsh, cache)
 
+        # monopole potential: atom-resolved
         qat = ihelp.reduce_shell_to_atom(qsh)
         vat = self.get_atom_potential(qat, cache)
 
+        # spread to orbital-resolution
         vsh += ihelp.spread_atom_to_shell(vat)
-        return ihelp.spread_shell_to_orbital(vsh)
+        vmono = ihelp.spread_shell_to_orbital(vsh)
+
+        # multipole potentials
+        vdipole = self.get_dipole_potential(qat, cache)
+        vquad = self.get_quadrupole_potential(qat, cache)
+
+        return Potential(vmono, vdipole=vdipole, vquad=vquad, label=self.label)
 
     def get_shell_potential(self, charges: Tensor, *_) -> Tensor:
         """
@@ -164,6 +179,44 @@ def get_atom_potential(self, charges: Tensor, *_) -> Tensor:
         """
         return torch.zeros_like(charges)
 
+    def get_dipole_potential(self, *_) -> Tensor | None:
+        """
+        Compute the dipole potential. All quantities are atom-resolved.
+
+        This method should be implemented by the subclass. Here, it serves
+        only to create an empty `Interaction` by returning None.
+
+        Parameters
+        ----------
+        charges : Tensor
+            Atom-resolved partial charges.
+
+        Returns
+        -------
+        Tensor | None
+            Atom-resolved potential vector for each atom or None if not needed.
+        """
+        return None
+
+    def get_quadrupole_potential(self, *_) -> Tensor | None:
+        """
+        Compute the quadrupole potential. All quantities are atom-resolved.
+
+        This method should be implemented by the subclass. Here, it serves
+        only to create an empty `Interaction` by returning None.
+
+        Parameters
+        ----------
+        charges : Tensor
+            Atom-resolved partial charges.
+
+        Returns
+        -------
+        Tensor | Nene
+            Atom-resolved potential vector for each atom or None if not needed.
+        """
+        return None
+
     def get_energy(
         self, charges: Tensor, cache: Interaction.Cache, ihelp: IndexHelper
     ) -> Tensor: