better docs for cell, charges, positions

docs/src/references/changelog.rst

@@ -27,6 +27,7 @@ changelog <https://keepachangelog.com/en/1.1.0/>`_ format. This project follows
 Added
 #####
 
+* Better documentation for for ``cell``, ``charges`` and ``positions`` parameters
 * Require consistent ``dtype`` between ``positions`` and ``neighbor_distances`` in
   ``Calculator`` classes and tuning functions.
 

src/torchpme/calculators/calculator.py

@@ -69,30 +69,6 @@ def _compute_rspace(
         neighbor_indices: torch.Tensor,
         neighbor_distances: torch.Tensor,
     ) -> torch.Tensor:
-        """
-        Computes the "real space" part of the potential. Depending on the
-        ``smearing`` in the potential class, it will either evaluate the
-        full potential, or only the short-range, "local" part.
-
-        :param charges: 2D tensor or list of 2D tensor of shape
-            ``(n_channels,len(positions))``.
-            ``n_channels`` is the number of charge channels the
-            potential should be calculated for a standard potential ``n_channels=1``.
-            If more than one "channel" is provided multiple potentials for the same
-            position are computed depending on the charges and the potentials.
-        :param neighbor_indices: Single or list of 2D tensors of shape ``(n, 2)``, where
-            ``n`` is the number of neighbors. The two columns correspond to the indices
-            of a **half neighbor list** for the two atoms which are considered neighbors
-            (e.g. within a cutoff distance) if ``full_neighbor_list=False`` (default).
-            Otherwise, a full neighbor list is expected.
-        :param neighbor_distances: single or list of 1D tensors containing the distance
-            between the ``n`` pairs corresponding to a **half (or full) neighbor list**
-            (see ``neighbor_indices``).
-        :return: Torch tensor containing the potential(s) for all
-            positions. Each tensor in the list is of shape ``(len(positions),
-            len(charges))``, where If the inputs are only single tensors only a single
-            torch tensor with the potentials is returned.
-        """
         # Compute the pair potential terms V(r_ij) for each pair of atoms (i,j)
         # contained in the neighbor list
         with profiler.record_function("compute bare potential"):
@@ -163,13 +139,19 @@ def forward(
         calculator implements a ``_compute_kspace`` method, it will also evaluate the
         long-range part using a Fourier-domain method.
 
-        :param charges: torch.Tensor, atomic (pseudo-)charges
-        :param cell: torch.Tensor, periodic supercell for the system
-        :param positions: torch.Tensor, Cartesian coordinates of the particles within
-            the supercell.
-        :param neighbor_indices: torch.Tensor with the ``i,j`` indices of neighbors for
+        :param charges: torch.Tensor of shape ``(n_channels, len(positions))``
+            containaing the atomic (pseudo-)charges. ``n_channels`` is the number of
+            charge channels the potential should be calculated. For a standard potential
+            ``n_channels = 1``. If more than one "channel" is provided multiple
+            potentials for the same position are computed depending on the charges and
+            the potentials.
+        :param cell: torch.tensor of shape ``(3, 3)``, where ``cell[i]`` is the i-th basis
+            vector of the unit cell
+        :param positions: torch.tensor of shape ``(N, 3)`` containing the Cartesian
+            coordinates of the ``N`` particles within the supercell.
+        :param neighbor_indices: torch.tensor with the ``i,j`` indices of neighbors for
             which the potential should be computed in real space.
-        :param neighbor_distances: torch.Tensor with the pair distances of the neighbors
+        :param neighbor_distances: torch.tensor with the pair distances of the neighbors
             for which the potential should be computed in real space.
         """
         _validate_parameters(

src/torchpme/lib/kspace_filter.py

@@ -110,8 +110,8 @@ def update(
         ``ns_mesh`` are given, the instance's attributes required by these will also be
         updated accordingly.
 
-        :param cell: torch.tensor of shape ``(3, 3)``, where `
-            `cell[i]`` is the i-th basis vector of the unit cell
+        :param cell: torch.tensor of shape ``(3, 3)``, where ``cell[i]`` is the i-th
+            basis vector of the unit cell
         :param ns_mesh: toch.tensor of shape ``(3,)``
             Number of mesh points to use along each of the three axes
         """

src/torchpme/lib/kvectors.py

@@ -6,9 +6,8 @@ def get_ns_mesh(cell: torch.Tensor, mesh_spacing: float):
     Computes the mesh size given a target mesh spacing and cell
     getting the closest powers of 2 to help with FFT.
 
-    :param cell: torch.tensor of shape ``(3, 3)``
-        Tensor specifying the real space unit cell of a structure, where ``cell[i]`` is
-        the i-th basis vector
+    :param cell: torch.tensor of shape ``(3, 3)``, where ``cell[i]`` is the i-th basis
+        vector of the unit cell
     :param mesh_spacing: float
     :param differentiable: boll
 
@@ -81,9 +80,8 @@ def generate_kvectors_for_mesh(cell: torch.Tensor, ns: torch.Tensor) -> torch.Te
     This variant is used in combination with **mesh based calculators** using the fast
     fourier transform (FFT) algorithm.
 
-    :param cell: torch.tensor of shape ``(3, 3)``
-        Tensor specifying the real space unit cell of a structure, where ``cell[i]`` is
-        the i-th basis vector
+    :param cell: torch.tensor of shape ``(3, 3)``, where ``cell[i]`` is the i-th basis
+        vector of the unit cell
     :param ns: torch.tensor of shape ``(3,)`` and dtype int
         ``ns = [nx, ny, nz]`` contains the number of mesh points in the x-, y- and
         z-direction, respectively. For faster performance during the Fast Fourier
@@ -119,9 +117,8 @@ def generate_kvectors_for_ewald(
     reciprocal space vectors is returned, rather than the FFT-optimized set that roughly
     contains only half of the vectors.
 
-    :param cell: torch.tensor of shape ``(3, 3)``
-        Tensor specifying the real space unit cell of a structure, where ``cell[i]`` is
-        the i-th basis vector
+    :param cell: torch.tensor of shape ``(3, 3)``, where ``cell[i]`` is the i-th basis
+        vector of the unit cell
     :param ns: torch.tensor of shape ``(3,)`` and dtype int
         ``ns = [nx, ny, nz]`` contains the number of mesh points in the x-, y- and
         z-direction, respectively.

src/torchpme/lib/mesh_interpolator.py

@@ -80,8 +80,8 @@ def update(
         the mesh resolution changes. If neither ``cell`` nor ``ns_mesh`` are passed
         there is nothing to be done.
 
-        :param cell: torch.tensor of shape ``(3, 3)``, where ``cell[i]`` is the i-th
-            basis vector of the unit cell
+        :param cell: torch.tensor of shape ``(3, 3)``, where ``cell[i]`` is the i-th basis
+            vector of the unit cell
         :param ns_mesh: toch.tensor of shape ``(3,)`` Number of mesh points to use along
             each of the three axes
         """
@@ -297,8 +297,8 @@ def compute_weights(self, positions: torch.Tensor):
         when calling the forward (:func:`points_to_mesh`) and backward
         (:func:`mesh_to_points`) interpolation functions.
 
-        :param positions: torch.tensor of shape ``(N, 3)``
-            Absolute positions of atoms in Cartesian coordinates
+        :param positions: torch.tensor of shape ``(N, 3)`` containing the Cartesian
+            coordinates of the ``N`` particles within the supercell.
         """
         if positions.device != self._device:
             raise ValueError(

src/torchpme/tuning/ewald.py

@@ -36,10 +36,12 @@ def tune_ewald(
         :func:`torchpme.tuning.ewald.EwaldErrorBounds.err_kspace`, which takes
         ``torch.Tensor`` as parameters.
 
-    :param charges: torch.Tensor, atomic (pseudo-)charges
-    :param cell: torch.Tensor, periodic supercell for the system
-    :param positions: torch.Tensor, Cartesian coordinates of the particles within the
-        supercell.
+    :param charges: torch.tensor of shape ``(len(positions, 1))`` containing the atomic
+        (pseudo-)charges
+    :param cell: torch.tensor of shape ``(3, 3)``, where ``cell[i]`` is the i-th basis
+        vector of the unit cell
+    :param positions: torch.tensor of shape ``(N, 3)`` containing the Cartesian
+        coordinates of the ``N`` particles within the supercell.
     :param cutoff: float, cutoff distance for the neighborlist
     :param exponent: :math:`p` in :math:`1/r^p` potentials, currently only :math:`p=1`
         is supported
@@ -122,10 +124,12 @@ class EwaldErrorBounds(TuningErrorBounds):
             \text{Error}_{\text{total}} = \sqrt{\text{Error}_{\text{real space}}^2 +
             \text{Error}_{\text{Fourier space}}^2
 
-    :param charges: atomic charges
-    :param cell: single tensor of shape (3, 3), describing the bounding
-    :param positions: single tensor of shape (``len(charges), 3``) containing the
-        Cartesian positions of all point charges in the system.
+    :param charges: torch.tensor of shape ``(len(positions, 1))`` containing the atomic
+        (pseudo-)charges
+    :param cell: torch.tensor of shape ``(3, 3)``, where ``cell[i]`` is the i-th basis
+        vector of the unit cell
+    :param positions: torch.tensor of shape ``(N, 3)`` containing the Cartesian
+        coordinates of the ``N`` particles within the supercell.
 
     Example
     -------

src/torchpme/tuning/p3m.py

@@ -92,10 +92,12 @@ def tune_p3m(
 
         \alpha = \left(\sqrt{2}\,\mathrm{smearing} \right)^{-1}
 
-    :param charges: torch.Tensor, atomic (pseudo-)charges
-    :param cell: torch.Tensor, periodic supercell for the system
-    :param positions: torch.Tensor, Cartesian coordinates of the particles within the
-        supercell.
+    :param charges: torch.tensor of shape ``(len(positions, 1))`` containing the atomic
+        (pseudo-)charges
+    :param cell: torch.tensor of shape ``(3, 3)``, where ``cell[i]`` is the i-th basis
+        vector of the unit cell
+    :param positions: torch.tensor of shape ``(N, 3)`` containing the Cartesian
+        coordinates of the ``N`` particles within the supercell.
     :param neighbor_indices: torch.Tensor with the ``i,j`` indices of neighbors for
         which the potential should be computed in real space.
     :param neighbor_distances: torch.Tensor with the pair distances of the neighbors for
@@ -203,10 +205,12 @@ class P3MErrorBounds(TuningErrorBounds):
         :func:`torchpme.tuning.p3m.P3MErrorBounds.err_kspace`, which takes
         ``torch.Tensor`` as parameters.
 
-    :param charges: atomic charges
-    :param cell: single tensor of shape (3, 3), describing the bounding
-    :param positions: single tensor of shape (``len(charges), 3``) containing the
-        Cartesian positions of all point charges in the system.
+    :param charges: torch.tensor of shape ``(len(positions, 1))`` containing the atomic
+        (pseudo-)charges
+    :param cell: torch.tensor of shape ``(3, 3)``, where ``cell[i]`` is the i-th basis
+        vector of the unit cell
+    :param positions: torch.tensor of shape ``(N, 3)`` containing the Cartesian
+        coordinates of the ``N`` particles within the supercell.
 
     Example
     -------

src/torchpme/tuning/pme.py

@@ -35,10 +35,12 @@ def tune_pme(
 
         \alpha = \left(\sqrt{2}\,\mathrm{smearing} \right)^{-1}
 
-    :param charges: torch.Tensor, atomic (pseudo-)charges
-    :param cell: torch.Tensor, periodic supercell for the system
-    :param positions: torch.Tensor, Cartesian coordinates of the particles within the
-        supercell.
+    :param charges: torch.tensor of shape ``(len(positions, 1))`` containing the atomic
+        (pseudo-)charges
+    :param cell: torch.tensor of shape ``(3, 3)``, where ``cell[i]`` is the i-th basis
+        vector of the unit cell
+    :param positions: torch.tensor of shape ``(N, 3)`` containing the Cartesian
+        coordinates of the ``N`` particles within the supercell.
     :param cutoff: float, cutoff distance for the neighborlist
     :param neighbor_indices: torch.Tensor with the ``i,j`` indices of neighbors for
         which the potential should be computed in real space.
@@ -146,10 +148,12 @@ class PMEErrorBounds(TuningErrorBounds):
         :func:`torchpme.tuning.pme.PMEErrorBounds.err_kspace`, which takes
         ``torch.Tensor`` as parameters.
 
-    :param charges: atomic charges
-    :param cell: single tensor of shape (3, 3), describing the bounding
-    :param positions: single tensor of shape (``len(charges), 3``) containing the
-        Cartesian positions of all point charges in the system.
+    :param charges: torch.tensor of shape ``(len(positions, 1))`` containing the atomic
+        (pseudo-)charges
+    :param cell: torch.tensor of shape ``(3, 3)``, where ``cell[i]`` is the i-th basis
+        vector of the unit cell
+    :param positions: torch.tensor of shape ``(N, 3)`` containing the Cartesian
+        coordinates of the ``N`` particles within the supercell.
 
     Example
     -------

src/torchpme/tuning/tuner.py

@@ -16,10 +16,12 @@ class TuningErrorBounds(torch.nn.Module):
     process. It can also be used with the :class:`torchpme.tuning.tuner.TunerBase` to
     build up a custom parameter tuner.
 
-    :param charges: atomic charges
-    :param cell: single tensor of shape (3, 3), describing the bounding
-    :param positions: single tensor of shape (``len(charges), 3``) containing the
-        Cartesian positions of all point charges in the system.
+    :param charges: torch.tensor of shape ``(len(positions, 1))`` containing the atomic
+        (pseudo-)charges
+    :param cell: torch.tensor of shape ``(3, 3)``, where ``cell[i]`` is the i-th basis
+        vector of the unit cell
+    :param positions: torch.tensor of shape ``(N, 3)`` containing the Cartesian
+        coordinates of the ``N`` particles within the supercell.
     """
 
     def __init__(
@@ -49,10 +51,12 @@ class TunerBase:
     real space error formula. The :func:`TunerBase.tune` defines the interface for a
     sophisticated tuning process, which takes a value of the desired accuracy.
 
-    :param charges: atomic charges
-    :param cell: single tensor of shape (3, 3), describing the bounding
-    :param positions: single tensor of shape (``len(charges), 3``) containing the
-        Cartesian positions of all point charges in the system.
+    :param charges: torch.tensor of shape ``(len(positions, 1))`` containing the atomic
+        (pseudo-)charges
+    :param cell: torch.tensor of shape ``(3, 3)``, where ``cell[i]`` is the i-th basis
+        vector of the unit cell
+    :param positions: torch.tensor of shape ``(N, 3)`` containing the Cartesian
+        coordinates of the ``N`` particles within the supercell.
     :param cutoff: real space cutoff, serves as a hyperparameter here.
     :param calculator: the calculator to be tuned
     :param exponent: exponent of the potential, only exponent = 1 is supported
@@ -154,17 +158,19 @@ class GridSearchTuner(TunerBase):
         cutoff, one could instantiate the tuner with different cutoff values and
         manually pick the best from the tuning results.
 
-    :param charges: atomic charges
-    :param cell: single tensor of shape (3, 3), describing the bounding
-    :param positions: single tensor of shape (``len(charges), 3``) containing the
-        Cartesian positions of all point charges in the system.
+    :param charges: torch.tensor of shape ``(len(positions, 1))`` containing the atomic
+        (pseudo-)charges
+    :param cell: torch.tensor of shape ``(3, 3)``, where ``cell[i]`` is the i-th basis
+        vector of the unit cell
+    :param positions: torch.tensor of shape ``(N, 3)`` containing the Cartesian
+        coordinates of the ``N`` particles within the supercell.
     :param cutoff: real space cutoff, serves as a hyperparameter here.
     :param calculator: the calculator to be tuned
     :param error_bounds: error bounds for the calculator
     :param params: list of Fourier space parameter sets for which the error is estimated
-    :param neighbor_indices: torch.Tensor with the ``i,j`` indices of neighbors for
+    :param neighbor_indices: torch.tensor with the ``i,j`` indices of neighbors for
         which the potential should be computed in real space.
-    :param neighbor_distances: torch.Tensor with the pair distances of the neighbors for
+    :param neighbor_distances: torch.tensor with the pair distances of the neighbors for
         which the potential should be computed in real space.
     :param exponent: exponent of the potential, only exponent = 1 is supported
     """
@@ -255,14 +261,16 @@ class TuningTimings(torch.nn.Module):
     warmup runs. The class takes the information of the structure that one wants to
     benchmark on, and the configuration of the timing process as inputs.
 
-    :param charges: atomic charges
-    :param cell: single tensor of shape (3, 3), describing the bounding
-    :param positions: single tensor of shape (``len(charges), 3``) containing the
-        Cartesian positions of all point charges in the system.
+    :param charges: torch.tensor of shape ``(len(positions, 1))`` containing the atomic
+        (pseudo-)charges
+    :param cell: torch.tensor of shape ``(3, 3)``, where ``cell[i]`` is the i-th basis
+        vector of the unit cell
+    :param positions: torch.tensor of shape ``(N, 3)`` containing the Cartesian
+        coordinates of the ``N`` particles within the supercell.
     :param cutoff: real space cutoff, serves as a hyperparameter here.
-    :param neighbor_indices: torch.Tensor with the ``i,j`` indices of neighbors for
+    :param neighbor_indices: torch.tensor with the ``i,j`` indices of neighbors for
         which the potential should be computed in real space.
-    :param neighbor_distances: torch.Tensor with the pair distances of the neighbors for
+    :param neighbor_distances: torch.tensor with the pair distances of the neighbors for
         which the potential should be computed in real space.
     :param n_repeat: number of times to repeat to estimate the average timing
     :param n_warmup: number of warmup runs, recommended to be at least 4