Add Explicit dtype and device Support for Calculators and Ensure Compatibility with Potentials (#143)

* Refactor parameter handling in calculators and potentials for improved dtype and device management
* Updated docstrings and changelog, added an assertion to check for an instance of the potential, and resolved the TorchScript Potential/Calculator incompatibility.
* Update changelog and add test for potential and calculator compatibility

Author: E-Rum

docs/extensions/versions_list.py

@@ -43,7 +43,6 @@ def run(self):
     :margin: 0 0 0 0\n"""
 
         for group_i, (version_short, group) in enumerate(grouped_versions.items()):
-
             if group_i < 3:
                 generated_content += f"""
     .. grid-item::

docs/src/references/changelog.rst

@@ -24,9 +24,16 @@ changelog <https://keepachangelog.com/en/1.1.0/>`_ format. This project follows
 `Unreleased <https://github.com/lab-cosmo/torch-pme/>`_
 -------------------------------------------------------
 
+Added
+#####
+
+* Added ``dtype`` and ``device`` for ``Calculator`` classses
+
 Fixed
 #####
 
+* Ensured consistency of ``dtype`` and ``device`` in the ``Potential`` and
+  ``Calculator`` classses
 * Fixed consistency of ``dtype`` and ``device`` in the ``SplinePotential`` class
 * Fix inconsistent ``cutoff`` in neighbor list example
 * All calculators now check if the cell is zero if the potential is range-separated

examples/5-autograd-demo.py

@@ -463,16 +463,16 @@ def forward(self, positions, cell, charges):
 
 print(
     f"""
-Delta-Value: {value-jit_value}
+Delta-Value: {value - jit_value}
 
 Delta-Position gradients:
-{positions.grad.T-jit_positions.grad.T}
+{positions.grad.T - jit_positions.grad.T}
 
 Delta-Cell gradients:
-{cell.grad-jit_cell.grad}
+{cell.grad - jit_cell.grad}
 
 Delta-Charges gradients:
-{charges.grad.T-jit_charges.grad.T}
+{charges.grad.T - jit_charges.grad.T}
 """
 )
 

src/torchpme/calculators/calculator.py

@@ -26,21 +26,37 @@ class Calculator(torch.nn.Module):
         will come from a full (True) or half (False, default) neighbor list.
     :param prefactor: electrostatics prefactor; see :ref:`prefactors` for details and
         common values.
+    :param dtype: type used for the internal buffers and parameters
+    :param device: device used for the internal buffers and parameters
     """
 
     def __init__(
         self,
         potential: Potential,
         full_neighbor_list: bool = False,
         prefactor: float = 1.0,
+        dtype: Optional[torch.dtype] = None,
+        device: Optional[torch.device] = None,
     ):
         super().__init__()
-        # TorchScript requires to initialize all attributes in __init__
-        self._device = torch.device("cpu")
-        self._dtype = torch.float32
 
+        assert isinstance(potential, Potential), (
+            f"Potential must be an instance of Potential, got {type(potential)}"
+        )
+
+        self.device = "cpu" if device is None else device
+        self.dtype = torch.get_default_dtype() if dtype is None else dtype
         self.potential = potential
 
+        assert self.dtype == self.potential.dtype, (
+            f"Potential and Calculator must have the same dtype, got {self.dtype} and "
+            f"{self.potential.dtype}"
+        )
+        assert self.device == self.potential.device, (
+            f"Potential and Calculator must have the same device, got {self.device} and "
+            f"{self.potential.device}"
+        )
+
         self.full_neighbor_list = full_neighbor_list
 
         self.prefactor = prefactor

src/torchpme/calculators/ewald.py

@@ -1,3 +1,5 @@
+from typing import Optional
+
 import torch
 
 from ..lib import generate_kvectors_for_ewald
@@ -53,6 +55,8 @@ class EwaldCalculator(Calculator):
         :obj:`False`, a "half" neighbor list is expected.
     :param prefactor: electrostatics prefactor; see :ref:`prefactors` for details and
         common values.
+    :param dtype: type used for the internal buffers and parameters
+    :param device: device used for the internal buffers and parameters
     """
 
     def __init__(
@@ -61,11 +65,15 @@ def __init__(
         lr_wavelength: float,
         full_neighbor_list: bool = False,
         prefactor: float = 1.0,
+        dtype: Optional[torch.dtype] = None,
+        device: Optional[torch.device] = None,
     ):
         super().__init__(
             potential=potential,
             full_neighbor_list=full_neighbor_list,
             prefactor=prefactor,
+            dtype=dtype,
+            device=device,
         )
         if potential.smearing is None:
             raise ValueError(

src/torchpme/calculators/p3m.py

@@ -1,3 +1,5 @@
+from typing import Optional
+
 import torch
 
 from ..lib.kspace_filter import P3MKSpaceFilter
@@ -40,6 +42,8 @@ class P3MCalculator(PMECalculator):
         set to :py:obj:`False`, a "half" neighbor list is expected.
     :param prefactor: electrostatics prefactor; see :ref:`prefactors` for details and
         common values.
+    :param dtype: type used for the internal buffers and parameters
+    :param device: device used for the internal buffers and parameters
 
     For an **example** on the usage for any calculator refer to :ref:`userdoc-how-to`.
     """
@@ -51,6 +55,8 @@ def __init__(
         interpolation_nodes: int = 4,
         full_neighbor_list: bool = False,
         prefactor: float = 1.0,
+        dtype: Optional[torch.dtype] = None,
+        device: Optional[torch.device] = None,
     ):
         self.mesh_spacing: float = mesh_spacing
 
@@ -62,6 +68,8 @@ def __init__(
             potential=potential,
             full_neighbor_list=full_neighbor_list,
             prefactor=prefactor,
+            dtype=dtype,
+            device=device,
         )
 
         self.kspace_filter: P3MKSpaceFilter = P3MKSpaceFilter(

src/torchpme/calculators/pme.py

@@ -1,3 +1,5 @@
+from typing import Optional
+
 import torch
 from torch import profiler
 
@@ -45,6 +47,8 @@ class PMECalculator(Calculator):
         set to :obj:`False`, a "half" neighbor list is expected.
     :param prefactor: electrostatics prefactor; see :ref:`prefactors` for details and
         common values.
+    :param dtype: type used for the internal buffers and parameters
+    :param device: device used for the internal buffers and parameters
     """
 
     def __init__(
@@ -54,11 +58,15 @@ def __init__(
         interpolation_nodes: int = 4,
         full_neighbor_list: bool = False,
         prefactor: float = 1.0,
+        dtype: Optional[torch.dtype] = None,
+        device: Optional[torch.device] = None,
     ):
         super().__init__(
             potential=potential,
             full_neighbor_list=full_neighbor_list,
             prefactor=prefactor,
+            dtype=dtype,
+            device=device,
         )
 
         if potential.smearing is None:
@@ -69,8 +77,8 @@ def __init__(
         self.mesh_spacing: float = mesh_spacing
 
         self.kspace_filter: KSpaceFilter = KSpaceFilter(
-            cell=torch.eye(3),
-            ns_mesh=torch.ones(3, dtype=int),
+            cell=torch.eye(3, dtype=self.dtype, device=self.device),
+            ns_mesh=torch.ones(3, dtype=int, device=self.device),
             kernel=self.potential,
             fft_norm="backward",
             ifft_norm="forward",
@@ -79,8 +87,8 @@ def __init__(
         self.interpolation_nodes: int = interpolation_nodes
 
         self.mesh_interpolator: MeshInterpolator = MeshInterpolator(
-            cell=torch.eye(3),
-            ns_mesh=torch.ones(3, dtype=int),
+            cell=torch.eye(3, dtype=self.dtype, device=self.device),
+            ns_mesh=torch.ones(3, dtype=int, device=self.device),
             interpolation_nodes=self.interpolation_nodes,
             method="Lagrange",  # convention for classic PME
         )

src/torchpme/lib/mesh_interpolator.py

@@ -432,8 +432,7 @@ def mesh_to_points(self, mesh_vals: torch.Tensor) -> torch.Tensor:
         """
         if mesh_vals.dim() != 4:
             raise ValueError(
-                f"`mesh_vals` of dimension {mesh_vals.dim()} has to be of "
-                "dimension 4"
+                f"`mesh_vals` of dimension {mesh_vals.dim()} has to be of dimension 4"
             )
 
         return (

src/torchpme/potentials/combined.py

@@ -47,10 +47,7 @@ def __init__(
             dtype=dtype,
             device=device,
         )
-        if dtype is None:
-            dtype = torch.get_default_dtype()
-        if device is None:
-            device = torch.device("cpu")
+
         smearings = [pot.smearing for pot in potentials]
         if not all(smearings) and any(smearings):
             raise ValueError(
@@ -76,7 +73,9 @@ def __init__(
                     "The number of initial weights must match the number of potentials being combined"
                 )
         else:
-            initial_weights = torch.ones(len(potentials), dtype=dtype, device=device)
+            initial_weights = torch.ones(
+                len(potentials), dtype=self.dtype, device=self.device
+            )
         # for torchscript
         self.potentials = torch.nn.ModuleList(potentials)
         if learnable_weights:

src/torchpme/potentials/coulomb.py

@@ -38,19 +38,17 @@ def __init__(
         device: Optional[torch.device] = None,
     ):
         super().__init__(smearing, exclusion_radius, dtype, device)
-        if dtype is None:
-            dtype = torch.get_default_dtype()
-        if device is None:
-            device = torch.device("cpu")
 
         # constants used in the forwward
         self.register_buffer(
             "_rsqrt2",
-            torch.rsqrt(torch.tensor(2.0, dtype=dtype, device=device)),
+            torch.rsqrt(torch.tensor(2.0, dtype=self.dtype, device=self.device)),
         )
         self.register_buffer(
             "_sqrt_2_on_pi",
-            torch.sqrt(torch.tensor(2.0 / torch.pi, dtype=dtype, device=device)),
+            torch.sqrt(
+                torch.tensor(2.0 / torch.pi, dtype=self.dtype, device=self.device)
+            ),
         )
 
     def from_dist(self, dist: torch.Tensor) -> torch.Tensor: