diff --git a/docs/src/references/changelog.rst b/docs/src/references/changelog.rst index 5631f33c3..76d12dd07 100644 --- a/docs/src/references/changelog.rst +++ b/docs/src/references/changelog.rst @@ -27,6 +27,7 @@ changelog `_ 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. diff --git a/src/torchpme/calculators/calculator.py b/src/torchpme/calculators/calculator.py index 70ede3efd..42e5d720c 100644 --- a/src/torchpme/calculators/calculator.py +++ b/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( diff --git a/src/torchpme/lib/kspace_filter.py b/src/torchpme/lib/kspace_filter.py index 5beec40f6..94a9fa30e 100644 --- a/src/torchpme/lib/kspace_filter.py +++ b/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 """ diff --git a/src/torchpme/lib/kvectors.py b/src/torchpme/lib/kvectors.py index e77e5a5fa..44590da10 100644 --- a/src/torchpme/lib/kvectors.py +++ b/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. diff --git a/src/torchpme/lib/mesh_interpolator.py b/src/torchpme/lib/mesh_interpolator.py index 8124d85e0..82cbff8d4 100644 --- a/src/torchpme/lib/mesh_interpolator.py +++ b/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( diff --git a/src/torchpme/tuning/ewald.py b/src/torchpme/tuning/ewald.py index 6452fa7dd..adea3cca1 100644 --- a/src/torchpme/tuning/ewald.py +++ b/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 ------- diff --git a/src/torchpme/tuning/p3m.py b/src/torchpme/tuning/p3m.py index 5685ffaf4..a81d35605 100644 --- a/src/torchpme/tuning/p3m.py +++ b/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 ------- diff --git a/src/torchpme/tuning/pme.py b/src/torchpme/tuning/pme.py index 55a0556e0..b4237d8d7 100644 --- a/src/torchpme/tuning/pme.py +++ b/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 ------- diff --git a/src/torchpme/tuning/tuner.py b/src/torchpme/tuning/tuner.py index 8461fda5d..d25746cf7 100644 --- a/src/torchpme/tuning/tuner.py +++ b/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