Changed types of sr_cutoff, smearing, and exponent from `torch.…

…tensor` to `float` (#19)

- Deleted redundant transpose.

src/meshlode/calculators/base.py

@@ -23,23 +23,23 @@ def _compute_sr(
         charges: torch.Tensor,
         cell: torch.Tensor,
         smearing: float,
-        sr_cutoff: torch.Tensor,
+        sr_cutoff: float,
         neighbor_indices: Optional[torch.Tensor] = None,
         neighbor_shifts: Optional[torch.Tensor] = None,
     ) -> torch.Tensor:
         if neighbor_indices is None or neighbor_shifts is None:
             # Get list of neighbors
             struc = Atoms(positions=positions.detach().numpy(), cell=cell, pbc=True)
             atom_is, atom_js, neighbor_shifts = neighbor_list(
-                "ijS", struc, sr_cutoff.item(), self_interaction=False
+                "ijS", struc, sr_cutoff, self_interaction=False
             )
             atom_is = torch.tensor(atom_is)
             atom_js = torch.tensor(atom_js)
             shifts = torch.tensor(neighbor_shifts, dtype=cell.dtype)  # N x 3
         else:
             atom_is = neighbor_indices[0]
             atom_js = neighbor_indices[1]
-            shifts = neighbor_shifts.type(cell.dtype).T
+            shifts = neighbor_shifts.type(cell.dtype)
 
         # Compute energy
         potential = torch.zeros_like(charges)

src/meshlode/calculators/ewaldpotential.py

@@ -10,7 +10,7 @@ class _EwaldPotentialImpl(_ShortRange):
     def __init__(
         self,
         exponent: float,
-        sr_cutoff: Union[None, torch.Tensor],
+        sr_cutoff: Union[None, float],
         atomic_smearing: Union[None, float],
         lr_wavelength: Union[None, float],
         subtract_self: bool,
@@ -52,7 +52,8 @@ def _compute_single_system(
         # structures.
         if self.sr_cutoff is None:
             cell_dimensions = torch.linalg.norm(cell, dim=1)
-            sr_cutoff = torch.min(cell_dimensions) / 2 - 1e-6
+            cutoff_max = torch.min(cell_dimensions) / 2 - 1e-6
+            sr_cutoff = cutoff_max.item()
         else:
             sr_cutoff = self.sr_cutoff
 

src/meshlode/calculators/pmepotential.py

@@ -11,7 +11,7 @@ class _PMEPotentialImpl(_ShortRange):
     def __init__(
         self,
         exponent: float,
-        sr_cutoff: Union[None, torch.Tensor],
+        sr_cutoff: Union[None, float],
         atomic_smearing: Union[None, float],
         mesh_spacing: Union[None, float],
         interpolation_order: int,
@@ -29,15 +29,6 @@ def __init__(
         _ShortRange.__init__(
             self, exponent=exponent, subtract_interior=subtract_interior
         )
-        self.atomic_smearing = atomic_smearing
-        self.mesh_spacing = mesh_spacing
-        self.interpolation_order = interpolation_order
-        self.sr_cutoff = sr_cutoff
-
-        # If interior contributions are to be subtracted, also do so for self term
-        if self.subtract_interior:
-            subtract_self = True
-        self.subtract_self = subtract_self
 
         self.atomic_smearing = atomic_smearing
         self.mesh_spacing = mesh_spacing
@@ -70,7 +61,7 @@ def _compute_single_system(
         if self.sr_cutoff is None:
             cell_dimensions = torch.linalg.norm(cell, dim=1)
             cutoff_max = torch.min(cell_dimensions) / 2 - 1e-6
-            sr_cutoff = cutoff_max
+            sr_cutoff = cutoff_max.item()
         else:
             sr_cutoff = self.sr_cutoff
 
@@ -113,8 +104,8 @@ def _compute_lr(
         positions: torch.Tensor,
         charges: torch.Tensor,
         cell: torch.Tensor,
-        smearing: torch.Tensor,
-        lr_wavelength: torch.Tensor,
+        smearing: float,
+        lr_wavelength: float,
         subtract_self=True,
     ) -> torch.Tensor:
         # Step 0 (Preparation): Compute number of times each basis vector of the

src/meshlode/lib/potentials.py

@@ -1,5 +1,4 @@
 import math
-from typing import Union
 
 import torch
 from torch.special import gammainc, gammaincc, gammaln
@@ -26,11 +25,8 @@ class InversePowerLawPotential:
     :param exponent: the exponent :math:`p` in :math:`1/r^p` potentials
     """
 
-    def __init__(self, exponent: Union[float, torch.Tensor]):
-        if type(exponent) is float:
-            self.exponent = torch.tensor(exponent)
-        else:
-            self.exponent = exponent
+    def __init__(self, exponent: float):
+        self.exponent = exponent
 
     def potential_from_dist(self, dist: torch.Tensor) -> torch.Tensor:
         """
@@ -52,7 +48,7 @@ def potential_from_dist_sq(self, dist_sq: torch.Tensor) -> torch.Tensor:
         return torch.pow(dist_sq, -self.exponent / 2.0)
 
     def potential_sr_from_dist(
-        self, dist: torch.Tensor, smearing: torch.Tensor
+        self, dist: torch.Tensor, smearing: float
     ) -> torch.Tensor:
         """
         Short-range (SR) part of the range-separated 1/r^p potential as a function of r.
@@ -64,23 +60,24 @@ def potential_sr_from_dist(
 
         :param dist: torch.tensor containing the distances at which the potential is to
             be evaluated.
-        :param smearing: torch.tensor containing the parameter often called "sigma" in
+        :param smearing: float containing the parameter often called "sigma" in
             publications, which determines the length-scale at which the short-range and
             long-range parts of the naive 1/r^p potential are separated. For the Coulomb
             potential (p=1), this potential can be interpreted as the effective
             potential generated by a Gaussian charge density, in which case this
             smearing parameter corresponds to the "width" of the Gaussian.
         """
+        exponent = torch.tensor(self.exponent, device=dist.device, dtype=dist.dtype)
         x = 0.5 * dist**2 / smearing**2
-        peff = self.exponent / 2
+        peff = exponent / 2
         prefac = 1.0 / (2 * smearing**2) ** peff
         potential = prefac * gammaincc(peff, x) / x**peff
 
         # potential = erfc(dist / torch.sqrt(torch.tensor(2.)) / smearing) / dist
         return potential
 
     def potential_lr_from_dist(
-        self, dist: torch.Tensor, smearing: torch.Tensor
+        self, dist: torch.Tensor, smearing: float
     ) -> torch.Tensor:
         """
         Long-range (LR) part of the range-separated 1/r^p potential as a function of r.
@@ -93,21 +90,22 @@ def potential_lr_from_dist(
 
         :param dist: torch.tensor containing the distances at which the potential is to
             be evaluated.
-        :param smearing: torch.tensor containing the parameter often called "sigma" in
+        :param smearing: float containing the parameter often called "sigma" in
             publications, which determines the length-scale at which the short-range and
             long-range parts of the naive 1/r^p potential are separated. For the Coulomb
             potential (p=1), this potential can be interpreted as the effective
             potential generated by a Gaussian charge density, in which case this
             smearing parameter corresponds to the "width" of the Gaussian.
         """
+        exponent = torch.tensor(self.exponent, device=dist.device, dtype=dist.dtype)
         x = 0.5 * dist**2 / smearing**2
-        peff = self.exponent / 2
+        peff = exponent / 2
         prefac = 1.0 / (2 * smearing**2) ** peff
         potential = prefac * gammainc(peff, x) / x**peff
         return potential
 
     def potential_fourier_from_k_sq(
-        self, k_sq: torch.Tensor, smearing: torch.Tensor
+        self, k_sq: torch.Tensor, smearing: float
     ) -> torch.Tensor:
         """
         Fourier transform of the long-range (LR) part potential parametrized in terms of
@@ -117,21 +115,22 @@ def potential_fourier_from_k_sq(
 
         :param k_sq: torch.tensor containing the squared lengths (2-norms) of the wave
             vectors k at which the Fourier-transformed potential is to be evaluated
-        :param smearing: torch.tensor containing the parameter often called "sigma" in
+        :param smearing: float containing the parameter often called "sigma" in
             publications, which determines the length-scale at which the short-range and
             long-range parts of the naive 1/r^p potential are separated. For the Coulomb
             potential (p=1), this potential can be interpreted as the effective
             potential generated by a Gaussian charge density, in which case this
             smearing parameter corresponds to the "width" of the Gaussian.
         """
-        peff = (3 - self.exponent) / 2
-        prefac = (math.pi) ** 1.5 / gamma(self.exponent / 2) * (2 * smearing**2) ** peff
+        exponent = torch.tensor(self.exponent, device=k_sq.device, dtype=k_sq.dtype)
+        peff = (3 - exponent) / 2
+        prefac = (math.pi) ** 1.5 / gamma(exponent / 2) * (2 * smearing**2) ** peff
         x = 0.5 * smearing**2 * k_sq
         fourier = prefac * gammaincc(peff, x) / x**peff * gamma(peff)
 
         return fourier
 
-    def potential_fourier_at_zero(self, smearing: torch.Tensor) -> torch.Tensor:
+    def potential_fourier_at_zero(self, smearing: float) -> torch.Tensor:
         """
         The value of the Fourier-transformed potential (LR part implemented above) as
         k --> 0 often needs to be set separately since for exponents p <= 3 = dimension,
@@ -142,7 +141,7 @@ def potential_fourier_at_zero(self, smearing: torch.Tensor) -> torch.Tensor:
         diverge as k --> 0, and one could instead assign the correct limit.
         This is not implemented for now for consistency reasons.
 
-        :param smearing: torch.tensor containing the parameter often called "sigma" in
+        :param smearing: float containing the parameter often called "sigma" in
             publications, which determines the length-scale at which the short-range and
             long-range parts of the naive 1/r^p potential are separated. For the Coulomb
             potential (p=1), this potential can be interpreted as the effective

src/meshlode/metatensor/base.py

@@ -102,7 +102,7 @@ def compute(self, systems: Union[List[System], System]) -> TensorMap:
                     neighbor_list = system.get_neighbor_list(neighbor_list_options)
 
                     neighbor_indices = neighbor_list.samples.values[:, :2].T
-                    neighbor_shifts = neighbor_list.samples.values[:, 2:].T
+                    neighbor_shifts = neighbor_list.samples.values[:, 2:]
 
                     break
 

tests/calculators/test_values_periodic.py

@@ -280,7 +280,7 @@ def define_crystal(crystal_name="CsCl"):
     return types, positions, charges, cell, madelung_ref, num_formula_units
 
 
-scaling_factors = torch.tensor([1 / 2.0353610, 1.0, 3.4951291], dtype=torch.float64)
+scaling_factors = [1 / 2.0353610, 1.0, 3.4951291]
 neutral_crystals = ["CsCl", "NaCl_primitive", "NaCl_cubic", "zincblende", "wurtzite"]
 neutral_crystals += ["cu2o", "fluorite"]
 
@@ -307,11 +307,11 @@ def test_madelung(crystal_name, scaling_factor, calc_name):
 
     # Define calculator and tolerances
     if calc_name == "ewald":
-        sr_cutoff = scaling_factor * torch.tensor(1.0, dtype=dtype)
+        sr_cutoff = scaling_factor
         calc = EwaldPotential(sr_cutoff=sr_cutoff)
         rtol = 4e-6
     elif calc_name == "pme":
-        sr_cutoff = scaling_factor * torch.tensor(2.0, dtype=dtype)
+        sr_cutoff = scaling_factor * 2
         calc = PMEPotential(sr_cutoff=sr_cutoff)
         rtol = 9e-4
 
@@ -333,7 +333,7 @@ def test_madelung(crystal_name, scaling_factor, calc_name):
     "wigner_bcc_cubiccell",
 ]
 
-scaling_factors = torch.tensor([0.4325, 1.0, 2.0353610], dtype=torch.float64)
+scaling_factors = [0.4325, 1.0, 2.0353610]
 
 
 @pytest.mark.parametrize("crystal_name", wigner_crystals)
@@ -372,15 +372,15 @@ def test_wigner(crystal_name, scaling_factor):
         smeareff *= scaling_factor
 
         # Compute potential and compare against reference
-        calc = EwaldPotential(atomic_smearing=smeareff)
+        calc = EwaldPotential(atomic_smearing=smeareff.item())
         potentials = calc.compute(positions=positions, charges=charges, cell=cell)
         energies = potentials * charges
         energies_ref = -torch.ones_like(energies) * madelung_ref
         torch.testing.assert_close(energies, energies_ref, atol=0.0, rtol=rtol)
 
 
 orthogonal_transformations = generate_orthogonal_transformations()
-scaling_factors = torch.tensor([0.4325, 2.0353610], dtype=dtype)
+scaling_factors = [0.4325, 2.0353610]
 
 
 @pytest.mark.parametrize("sr_cutoff", [2.01, 5.5])
@@ -420,7 +420,7 @@ def test_random_structure(sr_cutoff, frame_index, scaling_factor, ortho, calc_na
     charges = torch.tensor([1, 1, 1, 1, -1, -1, -1, -1], dtype=dtype).reshape((-1, 1))
 
     # Compute potential using MeshLODE and compare against reference values
-    sr_cutoff = scaling_factor * torch.tensor(sr_cutoff, dtype=dtype)
+    sr_cutoff = scaling_factor * sr_cutoff
     if calc_name == "ewald":
         calc = EwaldPotential(sr_cutoff=sr_cutoff)
         rtol_e = 2e-5

tests/lib/test_potentials.py

@@ -8,6 +8,7 @@
 
 
 def gamma(x):
+    x = torch.tensor(x, dtype=dtype)
     return torch.exp(torch.special.gammaln(x))
 
 
@@ -54,7 +55,6 @@ def test_potential_from_squared_argument(exponent):
     The potentials class can either compute the potential by taking the distance r or
     its square r^2 as an argument. This test makes sure that both implementations agree.
     """
-    exponent = torch.tensor(exponent, dtype=dtype)
 
     # Compute diverse potentials for this inverse power law
     ipl = InversePowerLawPotential(exponent=exponent)
@@ -75,8 +75,6 @@ def test_sr_lr_split(exponent, smearing):
     whether the sum of the SR and LR parts combine to the standard inverse power-law
     potential.
     """
-    exponent = torch.tensor(exponent, dtype=dtype)
-    smearing = torch.tensor(smearing, dtype=dtype)
 
     # Compute diverse potentials for this inverse power law
     ipl = InversePowerLawPotential(exponent=exponent)
@@ -105,8 +103,6 @@ def test_exact_sr(exponent, smearing):
     distance range (the variable dist_min) is increased, since the potential has a
     (removable) singularity at r=0.
     """
-    exponent = torch.tensor(exponent, dtype=dtype)
-    smearing = torch.tensor(smearing, dtype=dtype)
 
     # Compute SR part of Coulomb potential using the potentials class working for any
     # exponent
@@ -141,8 +137,6 @@ def test_exact_lr(exponent, smearing):
     distance range (the variable dist_min) is increased, since the potential has a
     (removable) singularity at r=0.
     """
-    exponent = torch.tensor(exponent, dtype=dtype)
-    smearing = torch.tensor(smearing, dtype=dtype)
 
     # Compute LR part of Coulomb potential using the potentials class working for any
     # exponent
@@ -161,8 +155,8 @@ def test_exact_lr(exponent, smearing):
         potential_exact = potential_1 / dists_sq - prefac * potential_2
 
     # Compare results. Large tolerance due to singular division
-    rtol = 8e-12
-    atol = 3e-16
+    rtol = 1e-10
+    atol = 1e-12
     assert_close(potential_lr_from_dist, potential_exact, rtol=rtol, atol=atol)
 
 
@@ -177,8 +171,6 @@ def test_exact_fourier(exponent, smearing):
     distance range (the variable dist_min) is increased, since the potential has a
     (removable) singularity at r=0.
     """
-    exponent = torch.tensor(exponent, dtype=dtype)
-    smearing = torch.tensor(smearing, dtype=dtype)
 
     # Compute LR part of Coulomb potential using the potentials class working for any
     # exponent
@@ -194,8 +186,8 @@ def test_exact_fourier(exponent, smearing):
         fourier_exact = -2 * PI * expi(-0.5 * smearing**2 * ks_sq)
 
     # Compare results. Large tolerance due to singular division
-    rtol = 10 * machine_epsilon
-    atol = 7e-16
+    rtol = 1e-14
+    atol = 1e-14
     assert_close(fourier_from_class, fourier_exact, rtol=rtol, atol=atol)
 
 
@@ -217,11 +209,8 @@ def test_lr_value_at_zero(exponent, smearing):
     In practice, this should not be such an issue since no two atoms should approach
     each other until their distance is 1e-5 (the value used here).
     """
-    exponent = torch.tensor(exponent, dtype=dtype)
-    smearing = torch.tensor(smearing, dtype=dtype)
-
     # Get atomic density at tiny distance
-    dist_small = torch.tensor(1e-8)
+    dist_small = torch.tensor(1e-8, dtype=dtype)
     ipl = InversePowerLawPotential(exponent=exponent)
     potential_close_to_zero = ipl.potential_lr_from_dist(dist_small, smearing=smearing)