Fix examples and tests

Fix again

examples/01-charges-example.py

@@ -42,12 +42,12 @@
 # %%
 #
 # Create the properties CsCl unit cell
-
+dtype = torch.float64
 symbols = ("Cs", "Cl")
 types = torch.tensor([55, 17])
-charges = torch.tensor([[1.0], [-1.0]], dtype=torch.float64)
-positions = torch.tensor([(0, 0, 0), (0.5, 0.5, 0.5)], dtype=torch.float64)
-cell = torch.eye(3, dtype=torch.float64)
+charges = torch.tensor([[1.0], [-1.0]], dtype=dtype)
+positions = torch.tensor([(0, 0, 0), (0.5, 0.5, 0.5)], dtype=dtype)
+cell = torch.eye(3, dtype=dtype)
 pbc = torch.tensor([True, True, True])
 
 
@@ -72,6 +72,7 @@
     cutoff=cutoff,
     neighbor_indices=neighbor_indices,
     neighbor_distances=neighbor_distances,
+    dtype=dtype,
 )
 
 # %%
@@ -101,7 +102,7 @@
 # will be used to *compute* the potential energy of the system.
 
 calculator = torchpme.PMECalculator(
-    torchpme.CoulombPotential(smearing=smearing), **pme_params
+    torchpme.CoulombPotential(smearing=smearing, dtype=dtype), dtype=dtype, **pme_params
 )
 
 # %%
@@ -112,7 +113,7 @@
 # As a first application of multiple charge channels, we start simply by using the
 # classic definition of one charge channel per atom.
 
-charges = torch.tensor([[1.0], [-1.0]], dtype=torch.float64)
+charges = torch.tensor([[1.0], [-1.0]], dtype=dtype)
 
 # %%
 #
@@ -160,7 +161,7 @@
 # species-specific potentials and facilitating the learning process for machine learning
 # algorithms.
 
-charges_one_hot = torch.tensor([[1.0, 0.0], [0.0, 1.0]], dtype=torch.float64)
+charges_one_hot = torch.tensor([[1.0, 0.0], [0.0, 1.0]], dtype=dtype)
 
 # %%
 #
@@ -205,7 +206,7 @@
 # creating a new calculator with the metatensor interface.
 
 calculator_metatensor = torchpme.metatensor.PMECalculator(
-    torchpme.CoulombPotential(smearing=smearing), **pme_params
+    torchpme.CoulombPotential(smearing=smearing, dtype=dtype), dtype=dtype, **pme_params
 )
 
 # %%

examples/02-neighbor-lists-usage.py

@@ -55,6 +55,7 @@
 #
 # As a test system, we use a 2x2x2 supercell of an CsCl crystal in a cubic cell.
 
+dtype = torch.float64
 atoms_unitcell = ase.Atoms(
     symbols=["Cs", "Cl"],
     positions=np.array([(0, 0, 0), (0.5, 0.5, 0.5)]),
@@ -97,7 +98,7 @@
 nl = vesin.torch.NeighborList(cutoff=cutoff, full_list=False)
 neighbor_indices, neighbor_distances = nl.compute(
     points=positions.to(dtype=torch.float64, device="cpu"),
-    box=cell.to(dtype=torch.float64, device="cpu"),
+    box=cell.to(dtype=dtype, device="cpu"),
     periodic=True,
     quantities="Pd",
 )
@@ -109,6 +110,7 @@
     cutoff=cutoff,
     neighbor_indices=neighbor_indices,
     neighbor_distances=neighbor_distances,
+    dtype=dtype,
 )
 
 # %%
@@ -193,7 +195,9 @@ def distances(
 # compute the potential.
 
 pme = torchpme.PMECalculator(
-    potential=torchpme.CoulombPotential(smearing=smearing), **pme_params
+    potential=torchpme.CoulombPotential(smearing=smearing, dtype=dtype),
+    dtype=dtype,
+    **pme_params,
 )
 potential = pme(
     charges=charges,

examples/08-combined-potential.py

@@ -28,6 +28,8 @@
 from torchpme import CombinedPotential, EwaldCalculator, InversePowerLawPotential
 from torchpme.prefactors import eV_A
 
+dtype = torch.float64
+
 # %%
 # Combined potentials
 # -------------------
@@ -65,10 +67,10 @@
 # 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, smearing=smearing)
-pot_2 = InversePowerLawPotential(exponent=2, smearing=smearing)
+pot_1 = InversePowerLawPotential(exponent=1, smearing=smearing, dtype=dtype)
+pot_2 = InversePowerLawPotential(exponent=2, smearing=smearing, dtype=dtype)
 
-potential = CombinedPotential(potentials=[pot_1, pot_2], smearing=smearing)
+potential = CombinedPotential(potentials=[pot_1, pot_2], smearing=smearing, dtype=dtype)
 
 # Note also that :class:`CombinedPotential` can be used with any combination of
 # potentials, as long they are all either direct or range separated. For instance, one
@@ -80,7 +82,7 @@
 # We now plot of the individual and combined ``potential`` functions together with an
 # explicit sum of the two potentials.
 
-dist = torch.logspace(-3, 2, 1000)
+dist = torch.logspace(-3, 2, 1000, dtype=dtype)
 
 fig, ax = plt.subplots()
 
@@ -115,7 +117,7 @@
 # combines all terms in a range-separated potential, including the k-space
 # kernel.
 
-k = torch.logspace(-2, 2, 1000)
+k = torch.logspace(-2, 2, 1000, dtype=dtype)
 
 fig, ax = plt.subplots()
 
@@ -154,9 +156,8 @@
 # much bigger system.
 
 calculator = EwaldCalculator(
-    potential=potential, lr_wavelength=lr_wavelength, prefactor=eV_A
+    potential=potential, lr_wavelength=lr_wavelength, prefactor=eV_A, dtype=dtype
 )
-calculator.to(dtype=torch.float64)
 
 
 # %%

examples/10-tuning.py

@@ -120,7 +120,9 @@
 pme_params = {"mesh_spacing": 1.0, "interpolation_nodes": 4}
 
 pme = torchpme.PMECalculator(
-    potential=torchpme.CoulombPotential(smearing=smearing),
+    potential=torchpme.CoulombPotential(smearing=smearing, device=device, dtype=dtype),
+    device=device,
+    dtype=dtype,
     **pme_params,  # type: ignore[arg-type]
 )
 
@@ -168,6 +170,8 @@
     neighbor_indices=neighbor_indices,
     neighbor_distances=neighbor_distances,
     run_backward=True,
+    device=device,
+    dtype=dtype,
 )
 estimated_timing = timings(pme)
 
@@ -210,15 +214,19 @@ def filter_neighbors(cutoff, neighbor_indices, neighbor_distances):
     return neighbor_indices[filter_idx], neighbor_distances[filter_idx]
 
 
-def timed_madelung(cutoff, smearing, mesh_spacing, interpolation_nodes):
+def timed_madelung(cutoff, smearing, mesh_spacing, interpolation_nodes, device, dtype):
     filter_indices, filter_distances = filter_neighbors(
         cutoff, neighbor_indices, neighbor_distances
     )
 
     pme = torchpme.PMECalculator(
-        potential=torchpme.CoulombPotential(smearing=smearing),
+        potential=torchpme.CoulombPotential(
+            smearing=smearing, device=device, dtype=dtype
+        ),
         mesh_spacing=mesh_spacing,
         interpolation_nodes=interpolation_nodes,
+        device=device,
+        dtype=dtype,
     )
     potential = pme(
         charges=charges,
@@ -239,6 +247,8 @@ def timed_madelung(cutoff, smearing, mesh_spacing, interpolation_nodes):
         run_backward=True,
         n_warmup=1,
         n_repeat=4,
+        device=device,
+        dtype=dtype,
     )
     estimated_timing = timings(pme)
     return madelung, estimated_timing
@@ -251,7 +261,9 @@ def timed_madelung(cutoff, smearing, mesh_spacing, interpolation_nodes):
 bounds = np.zeros((len(smearing_grid), len(spacing_grid)))
 for ism, smearing in enumerate(smearing_grid):
     for isp, spacing in enumerate(spacing_grid):
-        results[ism, isp], timings[ism, isp] = timed_madelung(8.0, smearing, spacing, 4)
+        results[ism, isp], timings[ism, isp] = timed_madelung(
+            8.0, smearing, spacing, 4, device, dtype
+        )
         bounds[ism, isp] = error_bounds(8.0, smearing, spacing, 4)
 
 # %%
@@ -374,7 +386,7 @@ def timed_madelung(cutoff, smearing, mesh_spacing, interpolation_nodes):
 for inint, nint in enumerate(nint_grid):
     for isp, spacing in enumerate(spacing_grid):
         results[inint, isp], timings[inint, isp] = timed_madelung(
-            5.0, 1.0, spacing, nint
+            5.0, 1.0, spacing, nint, device=device, dtype=dtype
         )
         bounds[inint, isp] = error_bounds(5.0, 1.0, spacing, nint)
 
@@ -445,15 +457,19 @@ def timed_madelung(cutoff, smearing, mesh_spacing, interpolation_nodes):
     cutoff=5.0,
     neighbor_indices=neighbor_indices,
     neighbor_distances=neighbor_distances,
+    device=device,
+    dtype=dtype,
 )
 
-print(f"""
+print(
+    f"""
 Estimated PME parameters (cutoff={5.0} Å):
 Smearing: {smearing} Å
 Mesh spacing: {parameters["mesh_spacing"]} Å
 Interpolation order: {parameters["interpolation_nodes"]}
 Estimated time per step: {timing} s
-""")
+"""
+)
 
 # %%
 # What is the best cutoff?
@@ -476,6 +492,8 @@ def timed_madelung(cutoff, smearing, mesh_spacing, interpolation_nodes):
         cutoff=cutoff,
         neighbor_indices=filter_indices,
         neighbor_distances=filter_distances,
+        device=device,
+        dtype=dtype,
     )
     timings_grid.append(timing)
 

examples/basic-usage.py

@@ -146,7 +146,7 @@
 # contains all the necessary functions (such as those defining the short-range and
 # long-range splits) for this potential and makes them useable in the rest of the code.
 
-potential = CoulombPotential(smearing=smearing)
+potential = CoulombPotential(smearing=smearing, device=device, dtype=dtype)
 
 # %%
 #
@@ -193,7 +193,9 @@
 # Since our structure is relatively small, we use the :class:`EwaldCalculator`.
 # We start by the initialization of the class.
 
-calculator = EwaldCalculator(potential=potential, lr_wavelength=lr_wavelength)
+calculator = EwaldCalculator(
+    potential=potential, lr_wavelength=lr_wavelength, device=device, dtype=dtype
+)
 
 # %%
 #

tests/helpers.py

@@ -256,7 +256,10 @@ 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.to(dtype=torch.float64, device="cpu"),
+        box=box.to(dtype=torch.float64, device="cpu"),
+        periodic=periodic,
+        quantities="PdS",
     )
 
     neighbor_indices = torch.from_numpy(neighbor_indices.astype(int)).to(