Minor refactor

README.rst

@@ -159,7 +159,7 @@ The following example shows how to calculate the EEQ partial charges and the cor
     cn = torch.tensor([3.0, 3.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0])
 
     eeq_model = eeq.EEQModel.param2019()
-    energy, qat = eeq.solve(numbers, positions, total_charge, eeq_model, cn)
+    energy, qat = eeq_model.solve(numbers, positions, total_charge, cn)
 
     print(torch.sum(energy, -1))
     # tensor(-0.1750)

docs/index.rst

@@ -80,7 +80,7 @@ The following example shows how to calculate the EEQ partial charges and the cor
     cn = torch.tensor([3.0, 3.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0])
 
     eeq_model = eeq.EEQModel.param2019()
-    energy, qat = eeq.solve(numbers, positions, total_charge, eeq_model, cn)
+    energy, qat = eeq_model.solve(numbers, positions, total_charge, cn)
 
     print(torch.sum(energy, -1))
     # tensor(-0.1750)

docs/modules/index.rst

@@ -9,6 +9,5 @@ The following modules are contained with `tad_multicharge`.
 
    param/index
    defaults
-   eeq
-   model
+   model/index
    typing/index

docs/modules/model/base.rst

@@ -0,0 +1,2 @@
+.. automodule:: tad_multicharge.model.base
+   :members:

docs/modules/model/eeq.rst

@@ -0,0 +1,2 @@
+.. automodule:: tad_multicharge.model.eeq
+   :members:

docs/modules/model/index.rst

@@ -0,0 +1,8 @@
+.. _model:
+
+.. automodule:: tad_multicharge.model
+
+.. toctree::
+
+   base
+   eeq

examples/eeq-batch.py

@@ -3,7 +3,7 @@
 from tad_mctc.batch import pack
 from tad_mctc.convert import symbol_to_number
 
-from tad_multicharge import eeq
+from tad_multicharge.model import eeq
 
 # S22 system 4: formamide dimer
 numbers = pack(

examples/eeq-single.py

@@ -1,7 +1,7 @@
 # SPDX-Identifier: CC0-1.0
 import torch
 
-from tad_multicharge import eeq
+from tad_multicharge.model import eeq
 
 numbers = torch.tensor([7, 7, 1, 1, 1, 1, 1, 1])
 
@@ -23,7 +23,7 @@
 cn = torch.tensor([3.0, 3.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0])
 
 eeq_model = eeq.EEQModel.param2019()
-energy, qat = eeq.solve(numbers, positions, total_charge, eeq_model, cn)
+energy, qat = eeq_model.solve(numbers, positions, total_charge, cn)
 
 print(torch.sum(energy, -1))
 # tensor(-0.1750)

src/tad_multicharge/__init__.py

@@ -70,18 +70,16 @@
 >>> # total charge of both systems
 >>> charge = torch.tensor([0.0, 0.0])
 >>>
->>> # calculate dispersion energy in Hartree
->>> energy = torch.sum(d4.dftd4(numbers, positions, charge, param), -1)
+>>> # calculate electrostatic energy in Hartree
+>>> energy = torch.sum(eeq.get_energy(numbers, positions, charge), -1)
 >>>
 >>> torch.set_printoptions(precision=10)
 >>> print(energy)
-tensor([-0.0088341432, -0.0027013607])
->>> print(energy[0] - 2*energy[1])
-tensor(-0.0034314217)
+>>> # tensor([-0.2086755037, -0.0972094536])
+>>> print(energy[0] - 2 * energy[1])
+>>> # tensor(-0.0142565966)
 """
 import torch
 
-from . import eeq, model
 from .__version__ import __version__
-from .eeq import get_charges as get_eeq_charges
-from .eeq import get_eeq
+from .model.eeq import get_charges as get_eeq_charges

src/tad_multicharge/model/__init__.py

@@ -0,0 +1,26 @@
+# This file is part of tad-multicharge.
+#
+# SPDX-Identifier: LGPL-3.0
+# Copyright (C) 2023 Marvin Friede
+#
+# tad-multicharge is free software: you can redistribute it and/or modify
+# it under the terms of the GNU Lesser General Public License as published by
+# the Free Software Foundation, either version 3 of the License, or
+# (at your option) any later version.
+#
+# tad-multicharge is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# GNU Lesser General Public License for more details.
+#
+# You should have received a copy of the GNU Lesser General Public License
+# along with tad-multicharge. If not, see <https://www.gnu.org/licenses/>.
+"""
+Model
+=====
+
+This module contains all available charge models. Currently, only the
+electronegativity equilibration model (EEQ) is implemented.
+"""
+from .base import *
+from .eeq import *

src/tad_multicharge/model.py → src/tad_multicharge/model/base.py

@@ -16,17 +16,18 @@
 # You should have received a copy of the GNU Lesser General Public License
 # along with tad-multicharge. If not, see <https://www.gnu.org/licenses/>.
 """
-Charge Model
-============
+Model: Base Charge Model
+========================
 
-Implementation of the electronegativity equlibration model for obtaining
-atomic partial charges as well as atom-resolved electrostatic energies.
+Implementation of a base class for charge models.
 """
 from __future__ import annotations
 
+from abc import abstractmethod
+
 import torch
 
-from .typing import Tensor, TensorLike
+from ..typing import Tensor, TensorLike
 
 __all__ = ["ChargeModel"]
 
@@ -76,3 +77,35 @@
             for tensor in (self.chi, self.kcn, self.eta, self.rad)
         ):
             raise RuntimeError("All tensors must have the same dtype!")
+
+    @abstractmethod
+    def solve(
+        self,
+        numbers: Tensor,
+        positions: Tensor,
+        total_charge: Tensor,
+        cn: Tensor,
+    ) -> tuple[Tensor, Tensor]:
+        """
+        Solve the electronegativity equilibration for the partial charges
+        minimizing the electrostatic energy.
+
+        Parameters
+        ----------
+        numbers : Tensor
+            Atomic numbers of all atoms in the system.
+        positions : Tensor
+            Cartesian coordinates of the atoms in the system (batch, natoms, 3).
+        total_charge : Tensor
+            Total charge of the system.
+        model : ChargeModel
+            Charge model to use.
+        cn : Tensor
+            Coordination numbers for all atoms in the system.
+
+        Returns
+        -------
+        (Tensor, Tensor)
+            Tuple of electrostatic energies and partial charges.
+        """
+        ...