Fix tests

src/torchpme/_utils.py

@@ -1,4 +1,4 @@
-from typing import Optional, Union
+from typing import Union
 
 import torch
 
@@ -11,7 +11,6 @@ def _validate_parameters(
     neighbor_distances: torch.Tensor,
     smearing: Union[float, None],
 ) -> None:
-
     dtype = positions.dtype
     device = positions.device
 

src/torchpme/calculators/calculator.py

@@ -1,5 +1,3 @@
-from typing import Optional, Union
-
 import torch
 from torch import profiler
 

src/torchpme/calculators/ewald.py

@@ -1,5 +1,3 @@
-from typing import Optional, Union
-
 import torch
 
 from ..lib import generate_kvectors_for_ewald

src/torchpme/calculators/p3m.py

@@ -1,5 +1,3 @@
-from typing import Optional, Union
-
 import torch
 
 from ..lib.kspace_filter import P3MKSpaceFilter
@@ -66,7 +64,11 @@ def __init__(
             prefactor=prefactor,
         )
 
-        cell = torch.eye(3, device=self.potential.smearing.device, dtype=self.potential.smearing.dtype)
+        cell = torch.eye(
+            3,
+            device=self.potential.smearing.device,
+            dtype=self.potential.smearing.dtype,
+        )
         ns_mesh = torch.ones(3, dtype=int, device=cell.device)
 
         self.kspace_filter: P3MKSpaceFilter = P3MKSpaceFilter(

src/torchpme/calculators/pme.py

@@ -1,5 +1,3 @@
-from typing import Optional, Union
-
 import torch
 from torch import profiler
 
@@ -70,7 +68,11 @@ def __init__(
 
         self.mesh_spacing: float = mesh_spacing
 
-        cell = torch.eye(3, device=self.potential.smearing.device, dtype=self.potential.smearing.dtype)
+        cell = torch.eye(
+            3,
+            device=self.potential.smearing.device,
+            dtype=self.potential.smearing.dtype,
+        )
         ns_mesh = torch.ones(3, dtype=int, device=cell.device)
 
         self.kspace_filter: KSpaceFilter = KSpaceFilter(

src/torchpme/lib/splines.py

@@ -151,30 +151,25 @@ def _solve_tridiagonal(a, b, c, d):
 def compute_second_derivatives(
     x_points: torch.Tensor,
     y_points: torch.Tensor,
-    high_precision: Optional[bool] = True,
 ):
     """
     Computes second derivatives given the grid points of a cubic spline.
 
     :param x_points: Abscissas of the splining points for the real-space function
     :param y_points: Ordinates of the splining points for the real-space function
-    :param high_accuracy: bool, perform calculation in double precision
 
     :return: The second derivatives for the spline points
     """
     # Do the calculation in float64 if required
     x = x_points
     y = y_points
-    if high_precision:
-        x = x.to(dtype=torch.float64)
-        y = y.to(dtype=torch.float64)
 
     # Calculate intervals
     intervals = x[1:] - x[:-1]
     dy = (y[1:] - y[:-1]) / intervals
 
     # Create zero boundary conditions (natural spline)
-    d2y = torch.zeros_like(x)
+    torch.zeros_like(x)
 
     n = len(x)
     a = torch.zeros_like(x)  # Sub-diagonal (a[1..n-1])
@@ -195,18 +190,16 @@ def compute_second_derivatives(
         c[i] = intervals[i] / 6
         d[i] = dy[i] - dy[i - 1]
 
-    d2y = _solve_tridiagonal(a, b, c, d)
+    return _solve_tridiagonal(a, b, c, d)
 
     # Converts back to the original dtype
-    return d2y
 
 
 def compute_spline_ft(
     k_points: torch.Tensor,
     x_points: torch.Tensor,
     y_points: torch.Tensor,
     d2y_points: torch.Tensor,
-    high_precision: Optional[bool] = True,
 ):
     r"""
     Computes the Fourier transform of a splined radial function.
@@ -228,7 +221,6 @@ def compute_spline_ft(
     :param x_points: Abscissas of the splining points for the real-space function
     :param y_points: Ordinates of the splining points for the real-space function
     :param d2y_points:  Second derivatives for the spline points
-    :param high_accuracy: bool, perform calculation in double precision
 
     :return: The radial Fourier transform :math:`\hat{f}(k)` computed
         at the ``k_points`` provided.
@@ -244,8 +236,6 @@ def compute_spline_ft(
 
     # chooses precision for the FT evaluation
     dtype = x_points.dtype
-    if high_precision:
-        dtype = torch.float64
 
     # broadcast to compute at once on all k values.
     # all these are terms that enter the analytical integral.

src/torchpme/potentials/combined.py

@@ -1,4 +1,4 @@
-from typing import Optional, Union
+from typing import Optional
 
 import torch
 

src/torchpme/potentials/coulomb.py

@@ -1,4 +1,4 @@
-from typing import Optional, Union
+from typing import Optional
 
 import torch
 
@@ -59,7 +59,7 @@ def lr_from_dist(self, dist: torch.Tensor) -> torch.Tensor:
                 "Cannot compute long-range contribution without specifying `smearing`."
             )
 
-        return torch.erf(dist / self.smearing / 2.0 ** 0.5) / dist
+        return torch.erf(dist / self.smearing / 2.0**0.5) / dist
 
     def lr_from_k_sq(self, k_sq: torch.Tensor) -> torch.Tensor:
         r"""

src/torchpme/potentials/inversepowerlaw.py

@@ -1,4 +1,4 @@
-from typing import Optional, Union
+from typing import Optional
 
 import torch
 from torch.special import gammainc
@@ -43,9 +43,7 @@ def __init__(
 
         # function call to check the validity of the exponent
         gammaincc_over_powerlaw(exponent, torch.tensor(1.0))
-        self.register_buffer(
-            "exponent", torch.tensor(float(exponent))
-        )
+        self.register_buffer("exponent", torch.tensor(exponent, dtype=torch.float64))
 
     @torch.jit.export
     def from_dist(self, dist: torch.Tensor) -> torch.Tensor:
@@ -99,7 +97,9 @@ def lr_from_k_sq(self, k_sq: torch.Tensor) -> torch.Tensor:
             )
 
         peff = (3 - self.exponent) / 2
-        prefac = torch.pi**1.5 / gamma(self.exponent / 2) * (2 * self.smearing**2) ** peff
+        prefac = (
+            torch.pi**1.5 / gamma(self.exponent / 2) * (2 * self.smearing**2) ** peff
+        )
         x = 0.5 * self.smearing**2 * k_sq
 
         # The k=0 term often needs to be set separately since for exponents p<=3

src/torchpme/potentials/potential.py

@@ -1,7 +1,8 @@
-from typing import Optional, Union
+from typing import Optional
 
 import torch
 
+
 class Potential(torch.nn.Module):
     r"""
     Base class defining the interface for a pair potential energy function
@@ -39,7 +40,9 @@ def __init__(
         super().__init__()
 
         if smearing is not None:
-            self.register_buffer("smearing", torch.tensor(smearing))
+            self.register_buffer(
+                "smearing", torch.tensor(smearing, dtype=torch.float64)
+            )
         else:
             self.smearing = None
 

src/torchpme/potentials/spline.py

@@ -1,4 +1,4 @@
-from typing import Optional, Union
+from typing import Optional
 
 import torch
 
@@ -63,7 +63,7 @@ def __init__(
 
         if len(y_grid) != len(r_grid):
             raise ValueError("Length of radial grid and value array mismatch.")
-        
+
         self.register_buffer("r_grid", r_grid)
         self.register_buffer("y_grid", y_grid)
 
@@ -106,14 +106,14 @@ def __init__(
 
         if y_at_zero is None:
             self._y_at_zero = self._spline(
-                torch.zeros(1, dtype=self.dtype, device=self.device)
+                torch.zeros(1, dtype=self.r_grid.dtype, device=self.r_grid.device)
             )
         else:
             self._y_at_zero = y_at_zero
 
         if yhat_at_zero is None:
             self._yhat_at_zero = self._krn_spline(
-                torch.zeros(1, dtype=self.dtype, device=self.device)
+                torch.zeros(1, dtype=self.k_grid.dtype, device=self.k_grid.device)
             )
         else:
             self._yhat_at_zero = yhat_at_zero

src/torchpme/tuning/ewald.py

@@ -1,5 +1,5 @@
 import math
-from typing import Any, Optional, Union
+from typing import Any
 from warnings import warn
 
 import torch

src/torchpme/tuning/p3m.py

@@ -1,6 +1,6 @@
 import math
 from itertools import product
-from typing import Any, Optional, Union
+from typing import Any
 from warnings import warn
 
 import torch

src/torchpme/tuning/pme.py

@@ -1,6 +1,6 @@
 import math
 from itertools import product
-from typing import Any, Optional, Union
+from typing import Any
 from warnings import warn
 
 import torch

src/torchpme/tuning/tuner.py

@@ -1,6 +1,6 @@
 import math
 import time
-from typing import Optional, Union
+from typing import Optional
 
 import torch
 
@@ -234,7 +234,7 @@ def _timing(self, smearing: float, k_space_params: dict):
             ),
             **k_space_params,
         )
-
+        calculator.to(device=self.positions.device, dtype=self.positions.dtype)
         return self.time_func(calculator)
 
 

tests/calculators/test_calculator.py

@@ -78,9 +78,7 @@ def test_invalid_shape_cell():
 
 def test_invalid_dtype_cell():
     calculator = CalculatorTest()
-    match = (
-        r"type of `cell` \(torch.float64\) must be same as that of the `positions` class \(torch.float32\)"
-    )
+    match = r"type of `cell` \(torch.float64\) must be same as that of the `positions` class \(torch.float32\)"
     with pytest.raises(TypeError, match=match):
         calculator.forward(
             positions=POSITIONS_1,

tests/calculators/test_values_direct.py

@@ -18,7 +18,8 @@ class CalculatorTest(Calculator):
     def __init__(self, **kwargs):
         super().__init__(
             potential=CoulombPotential(
-                smearing=None, exclusion_radius=None,
+                smearing=None,
+                exclusion_radius=None,
             ),
             **kwargs,
         )