Checkpoint with working tests and linting

examples/8-combined-potential.py

@@ -65,8 +65,8 @@
 # evaluation, and so one has to set it also for the combined potential, even if it is
 # not used explicitly in the evaluation of the combination.
 
-pot_1 = InversePowerLawPotential(exponent=1.0, smearing=smearing)
-pot_2 = InversePowerLawPotential(exponent=2.0, smearing=smearing)
+pot_1 = InversePowerLawPotential(exponent=1, smearing=smearing)
+pot_2 = InversePowerLawPotential(exponent=2, smearing=smearing)
 
 potential = CombinedPotential(potentials=[pot_1, pot_2], smearing=smearing)
 

src/torchpme/calculators/ewald.py

@@ -129,6 +129,6 @@ def _compute_kspace(
         ivolume = torch.abs(cell.det()).pow(-1)
         charge_tot = torch.sum(charges, dim=0)
         prefac = self.potential.background_correction()
-        energy -= 2 * prefac * charge_tot * ivolume if charge_tot != 0 else 0
+        energy -= 2 * prefac * charge_tot * ivolume
         # Compensate for double counting of pairs (i,j) and (j,i)
         return energy / 2

src/torchpme/potentials/integerspline.py

@@ -59,7 +59,7 @@ def __init__(
             dtype = torch.get_default_dtype()
         if device is None:
             device = torch.device("cpu")
-        self.r_grid = r_grid
+        self.r_grid = r_grid.to(dtype=dtype, device=device)
         self.register_buffer(
             "exponent", torch.tensor(exponent, dtype=dtype, device=device)
         )
@@ -108,7 +108,12 @@ def lr_from_dist(self, dist: torch.Tensor) -> torch.Tensor:
     @torch.jit.export
     def lr_from_k_sq(self, k_sq: torch.Tensor) -> torch.Tensor:
         r"""TODO: Fourier transform of the LR part potential in terms of :math:`\mathbf{k^2}`."""
-        spline = SplinePotential(self.r_grid, self.lr_from_dist(self.r_grid))
+        spline = SplinePotential(
+            self.r_grid,
+            self.lr_from_dist(self.r_grid),
+            device=self.r_grid.device,
+            dtype=self.r_grid.dtype,
+        )
         return spline.lr_from_k_sq(k_sq)
 
     def self_contribution(self) -> torch.Tensor:

src/torchpme/potentials/inversepowerlaw.py

@@ -25,7 +25,12 @@ class CustomE1(torch.autograd.Function):
     @staticmethod
     def forward(ctx, input):
         ctx.save_for_backward(input)
-        return exp1(input)
+        if input.is_cuda:
+            input_cpu = input.cpu()
+            result = torch.tensor(exp1(input_cpu.numpy()), device=input.device)
+        else:
+            result = torch.tensor(exp1(input.numpy()))
+        return result
 
     @staticmethod
     def backward(ctx, grad_output):
@@ -36,9 +41,6 @@ def backward(ctx, grad_output):
 # Auxilary function for stable Fourier transform implementation
 def gammaincc_over_powerlaw(exponent: torch.Tensor, z: torch.Tensor) -> torch.Tensor:
     """Function to compute the regularized incomplete gamma function complement for integer exponents."""
-    if exponent not in [1, 2, 3, 4, 5, 6]:
-        raise ValueError(f"Unsupported exponent: {exponent}")
-
     if exponent == 1:
         return torch.exp(-z) / z
     if exponent == 2:
@@ -56,7 +58,7 @@ def gammaincc_over_powerlaw(exponent: torch.Tensor, z: torch.Tensor) -> torch.Te
             (2 - 4 * z) * torch.exp(-z)
             + 4 * torch.sqrt(torch.pi) * z**1.5 * torch.erfc(torch.sqrt(z))
         ) / 3
-    return None
+    raise ValueError(f"Unsupported exponent: {exponent}")
 
 
 class InversePowerLawPotential(Potential):

tests/helpers.py

@@ -257,7 +257,7 @@ def neighbor_list(
 
     nl = NeighborList(cutoff=cutoff, full_list=full_neighbor_list)
     neighbor_indices, d, S = nl.compute(
-        points=positions, box=box, periodic=periodic, quantities="pdS"
+        points=positions, box=box, periodic=periodic, quantities="PdS"
     )
 
     neighbor_indices = torch.from_numpy(neighbor_indices.astype(int)).to(

tests/test_potentials.py

@@ -30,7 +30,7 @@ def gamma(x):
 # Electron mass m_e = 9.1094 * 1e-31 kg
 # TODO: for the moment, InversePowerLawPotential only works for exponent 0<p<3
 # ps = [1.0, 2.0, 3.0, 6.0] + [0.12345, 0.54321, 2.581304, 4.835909, 6.674311, 9.109431]
-ps = [1.0, 2.0, 3.0] + [0.12345, 0.54321, 2.581304]
+ps = [1, 2, 3]
 
 # Define range of smearing parameters covering relevant values
 smearinges = [0.1, 0.5, 1.0, 1.56]
@@ -83,7 +83,7 @@ def test_sr_lr_split(exponent, smearing):
     assert_close(potential_from_dist, potential_from_sum, rtol=rtol, atol=atol)
 
 
-@pytest.mark.parametrize("exponent", [1.0, 2.0, 3.0])
+@pytest.mark.parametrize("exponent", [1, 2, 3])
 @pytest.mark.parametrize("smearing", smearinges)
 def test_exact_sr(exponent, smearing):
     """
@@ -103,11 +103,11 @@ def test_exact_sr(exponent, smearing):
     # Compute exact analytical expression obtained for relevant exponents
     potential_1 = erfc(dists / SQRT2 / smearing) / dists
     potential_2 = torch.exp(-0.5 * dists_sq / smearing**2) / dists_sq
-    if exponent == 1.0:
+    if exponent == 1:
         potential_exact = potential_1
-    elif exponent == 2.0:
+    elif exponent == 2:
         potential_exact = potential_2
-    elif exponent == 3.0:
+    elif exponent == 3:
         prefac = SQRT2 / torch.sqrt(PI) / smearing
         potential_exact = potential_1 / dists_sq + prefac * potential_2
 
@@ -117,7 +117,7 @@ def test_exact_sr(exponent, smearing):
     assert_close(potential_sr_from_dist, potential_exact, rtol=rtol, atol=atol)
 
 
-@pytest.mark.parametrize("exponent", [1.0, 2.0, 3.0])
+@pytest.mark.parametrize("exponent", [1, 2, 3])
 @pytest.mark.parametrize("smearing", smearinges)
 def test_exact_lr(exponent, smearing):
     """
@@ -137,11 +137,11 @@ def test_exact_lr(exponent, smearing):
     # Compute exact analytical expression obtained for relevant exponents
     potential_1 = erf(dists / SQRT2 / smearing) / dists
     potential_2 = torch.exp(-0.5 * dists_sq / smearing**2) / dists_sq
-    if exponent == 1.0:
+    if exponent == 1:
         potential_exact = potential_1
-    elif exponent == 2.0:
+    elif exponent == 2:
         potential_exact = 1 / dists_sq - potential_2
-    elif exponent == 3.0:
+    elif exponent == 3:
         prefac = SQRT2 / torch.sqrt(PI) / smearing
         potential_exact = potential_1 / dists_sq - prefac * potential_2
 
@@ -151,7 +151,7 @@ def test_exact_lr(exponent, smearing):
     assert_close(potential_lr_from_dist, potential_exact, rtol=rtol, atol=atol)
 
 
-@pytest.mark.parametrize("exponent", [1.0, 2.0])
+@pytest.mark.parametrize("exponent", [1, 2])
 @pytest.mark.parametrize("smearing", smearinges)
 def test_exact_fourier(exponent, smearing):
     """
@@ -169,11 +169,11 @@ def test_exact_fourier(exponent, smearing):
     fourier_from_class = ipl.lr_from_k_sq(ks_sq)
 
     # Compute exact analytical expression obtained for relevant exponents
-    if exponent == 1.0:
+    if exponent == 1:
         fourier_exact = 4 * PI / ks_sq * torch.exp(-0.5 * smearing**2 * ks_sq)
-    elif exponent == 2.0:
+    elif exponent == 2:
         fourier_exact = 2 * PI**2 / ks * erfc(smearing * ks / SQRT2)
-    elif exponent == 3.0:
+    elif exponent == 3:
         fourier_exact = -2 * PI * expi(-0.5 * smearing**2 * ks_sq)
 
     # Compare results. Large tolerance due to singular division
@@ -183,7 +183,7 @@ def test_exact_fourier(exponent, smearing):
 
 
 @pytest.mark.parametrize("smearing", smearinges)
-@pytest.mark.parametrize("exponent", ps[:-1])  # for p=9.11, the results are unstable
+@pytest.mark.parametrize("exponent", ps[:-1])
 def test_lr_value_at_zero(exponent, smearing):
     """
     The LR part of the potential should no longer have a singularity as r-->0. Instead,
@@ -213,18 +213,18 @@ def test_lr_value_at_zero(exponent, smearing):
 
 
 def test_exponent_out_of_range():
-    match = r"`exponent` p=.* has to satisfy 0 < p <= 3"
+    match = r"Unsupported exponent: .*"
     with pytest.raises(ValueError, match=match):
         InversePowerLawPotential(exponent=-1.0, smearing=0.0)
 
     with pytest.raises(ValueError, match=match):
-        InversePowerLawPotential(exponent=4, smearing=0.0)
+        InversePowerLawPotential(exponent=7, smearing=0.0)
 
 
 @pytest.mark.parametrize("potential", [CoulombPotential, InversePowerLawPotential])
 def test_range_none(potential):
     if potential is InversePowerLawPotential:
-        pot = potential(exponent=2.0)
+        pot = potential(exponent=2)
     else:
         pot = potential()
 
@@ -250,16 +250,16 @@ class NoImplPotential(Potential):
     with pytest.raises(
         NotImplementedError, match="from_dist is not implemented for NoImplPotential"
     ):
-        mypot.from_dist(torch.tensor([1, 2.0, 3.0]))
+        mypot.from_dist(torch.tensor([1, 2, 3]))
     with pytest.raises(
         NotImplementedError, match="lr_from_dist is not implemented for NoImplPotential"
     ):
-        mypot.lr_from_dist(torch.tensor([1, 2.0, 3.0]))
+        mypot.lr_from_dist(torch.tensor([1, 2, 3]))
     with pytest.raises(
         NotImplementedError,
         match="lr_from_k_sq is not implemented for NoImplPotential",
     ):
-        mypot.lr_from_k_sq(torch.tensor([1, 2.0, 3.0]))
+        mypot.lr_from_k_sq(torch.tensor([1, 2, 3]))
     with pytest.raises(
         NotImplementedError,
         match="self_contribution is not implemented for NoImplPotential",
@@ -274,7 +274,7 @@ class NoImplPotential(Potential):
         ValueError,
         match="Cannot compute cutoff function when `exclusion_radius` is not set",
     ):
-        mypot.f_cutoff(torch.tensor([1, 2.0, 3.0]))
+        mypot.f_cutoff(torch.tensor([1, 2, 3]))
 
 
 @pytest.mark.parametrize("exclusion_radius", [0.5, 1.0, 2.0])
@@ -294,7 +294,7 @@ def test_inverserp_coulomb(smearing):
     """
     # Compute LR part of Coulomb potential using the potentials class working for any
     # exponent
-    ipl = InversePowerLawPotential(exponent=1.0, smearing=smearing, dtype=dtype)
+    ipl = InversePowerLawPotential(exponent=1, smearing=smearing, dtype=dtype)
     coul = CoulombPotential(smearing=smearing, dtype=dtype)
 
     ipl_from_dist = ipl.from_dist(dists)
@@ -439,8 +439,8 @@ def forward(self, x: torch.Tensor):
 
 @pytest.mark.parametrize("smearing", smearinges)
 def test_combined_potential(smearing):
-    ipl_1 = InversePowerLawPotential(exponent=1.0, smearing=smearing, dtype=dtype)
-    ipl_2 = InversePowerLawPotential(exponent=2.0, smearing=smearing, dtype=dtype)
+    ipl_1 = InversePowerLawPotential(exponent=1, smearing=smearing, dtype=dtype)
+    ipl_2 = InversePowerLawPotential(exponent=2, smearing=smearing, dtype=dtype)
 
     ipl_1_from_dist = ipl_1.from_dist(dists)
     ipl_1_sr_from_dist = ipl_1.sr_from_dist(dists)
@@ -585,7 +585,7 @@ def test_potential_device_dtype(potential_class, device, dtype):
         pytest.skip("CUDA is not available")
 
     smearing = 1.0
-    exponent = 1.0
+    exponent = 2
 
     if potential_class is InversePowerLawPotential:
         potential = potential_class(