Remove own System class

docs/src/references/lib/index.rst

@@ -9,4 +9,3 @@ are used for the meshLODE calculators.
 
    fourier_convolution
    mesh_interpolator
-   system

examples/library-tutorial.py

@@ -18,6 +18,7 @@
 import matplotlib.pyplot as plt
 import numpy as np
 import torch
+from metatensor.torch.atomistic import System
 
 import meshlode
 
@@ -50,20 +51,18 @@ def sliceplot(mesh, sz=12, cmap="viridis", vmin=None, vmax=None):
 # --------------------
 #
 # Builds a CsCl structure by replicating the primitive cell using ase and convert it to
-# a :py:class:`list` of :py:class:`meshlode.System`. We add a bit of noise to make
-# it less boring!
+# a :py:class:`list` of :py:class:`metatensor.torch.atomistic.System`. We add a bit of
+# noise to make it less boring!
 #
 
 positions = torch.tensor([[0, 0, 0], [0.5, 0.5, 0.5]]) * 4
-atomic_types = torch.tensor([55, 17])  # Cs and Cl
+types = torch.tensor([55, 17])  # Cs and Cl
 cell = torch.eye(3) * 4
-ase_frame = ase.Atoms(positions=positions, cell=cell, numbers=atomic_types).repeat(
-    [2, 2, 2]
-)
+ase_frame = ase.Atoms(positions=positions, cell=cell, numbers=types).repeat([2, 2, 2])
 ase_frame.positions[:] += np.random.normal(size=ase_frame.positions.shape) * 0.1
 charges = torch.tensor([1.0, -1.0] * 8)
-frame = meshlode.System(
-    species=torch.tensor(ase_frame.numbers),
+system = System(
+    types=torch.tensor(ase_frame.numbers),
     positions=torch.tensor(np.array(ase_frame.positions)),
     cell=torch.tensor(ase_frame.cell),
 )
@@ -95,15 +94,15 @@ def sliceplot(mesh, sz=12, cmap="viridis", vmin=None, vmax=None):
 #
 
 interpol = meshlode.lib.mesh_interpolator.MeshInterpolator(
-    frame.cell, torch.tensor([16, 16, 16]), interpolation_order=3
+    system.cell, torch.tensor([16, 16, 16]), interpolation_order=3
 )
 
-interpol.compute_interpolation_weights(frame.positions)
+interpol.compute_interpolation_weights(system.positions)
 
 
 # %%
 # We use two sets of weights: ones (giving the atom density irrespective
-# of the species) and charges (giving a smooth representation of the point
+# of the types) and charges (giving a smooth representation of the point
 # charges).
 #
 
@@ -137,15 +136,15 @@ def sliceplot(mesh, sz=12, cmap="viridis", vmin=None, vmax=None):
 # %%
 # plain atomic_smearing
 rho_mesh = fsc.compute(
-    mesh_values=mesh, cell=frame.cell, potential_exponent=0, atomic_smearing=1
+    mesh_values=mesh, cell=system.cell, potential_exponent=0, atomic_smearing=1
 )
 
 sliceplot(rho_mesh[0, :, :, :5])
 
 # %%
 # coulomb-like potential, no atomic_smearing
 coulomb_mesh = fsc.compute(
-    mesh_values=mesh, cell=frame.cell, potential_exponent=1, atomic_smearing=0
+    mesh_values=mesh, cell=system.cell, potential_exponent=1, atomic_smearing=0
 )
 
 sliceplot(coulomb_mesh[1, :, :, :5], cmap="seismic")
@@ -175,13 +174,13 @@ def sliceplot(mesh, sz=12, cmap="viridis", vmin=None, vmax=None):
 #
 
 interpol_slice = meshlode.lib.mesh_interpolator.MeshInterpolator(
-    frame.cell, torch.tensor([16, 16, 16]), interpolation_order=4
+    system.cell, torch.tensor([16, 16, 16]), interpolation_order=4
 )
 
 # Compute a denser grid on a 2D slice
 n_points = 50
-x = torch.linspace(0, frame.cell[0, 0], n_points + 1)[:n_points]
-y = torch.linspace(0, frame.cell[1, 1], n_points + 1)[:n_points]
+x = torch.linspace(0, system.cell[0, 0], n_points + 1)[:n_points]
+y = torch.linspace(0, system.cell[1, 1], n_points + 1)[:n_points]
 xx, yy = torch.meshgrid(x, y, indexing="ij")
 
 # Flatten xx and yy, and concatenate with a zero column for the z-coordinate

examples/madelung.py

@@ -11,24 +11,24 @@
 import math
 
 import torch
+from metatensor.torch.atomistic import System
 
 import meshlode
 
 
 # %%
 # Define simple example structure having the CsCl structure and compute the reference
-# values. MeshPotential by default outputs the species sorted according to the atomic
+# values. MeshPotential by default outputs the types sorted according to the atomic
 # number. Thus, we input the compound "CsCl" and "ClCs" since Cl and Cs have atomic
 # numbers 17 and 55, respectively.
-atomic_types = torch.tensor([17, 55])  # Cl and Cs
+types = torch.tensor([17, 55])  # Cl and Cs
+positions = torch.tensor([[0, 0, 0], [0.5, 0.5, 0.5]])
 charges = torch.tensor([-1.0, 1.0])
 cell = torch.eye(3)
-positions = torch.tensor([[0, 0, 0], [0.5, 0.5, 0.5]])
-frame = meshlode.System(species=atomic_types, positions=positions, cell=torch.eye(3))
 
 # %%
 # Define the expected values of the energy
-n_atoms = len(positions)
+n_atoms = len(types)
 madelung = 2 * 1.7626 / math.sqrt(3)
 energies_ref = -madelung * torch.ones((n_atoms, 1))
 
@@ -52,7 +52,7 @@
     interpolation_order=interpolation_order,
     subtract_self=True,
 )
-potentials_torch = MP.compute(frame)
+potentials_torch = MP.compute(types=types, positions=positions, cell=cell)
 
 # %%
 # The "potentials" that have been computed so far are not the actual electrostatic
@@ -66,7 +66,7 @@
     for idx_n in range(n_atoms):
         # The coulomb potential between atoms i and j is charge_i * charge_j / d_ij
         # The features are simply computing a pure 1/r potential with no prefactors.
-        # Thus, to compute the energy between atoms of species i and j, we need to
+        # Thus, to compute the energy between atoms of types i and j, we need to
         # multiply by the charges of i and j.
         print(charges[idx_c] * charges[idx_n], potentials_torch[idx_n, idx_c])
         atomic_energies_torch[idx_c] += (
@@ -89,32 +89,36 @@
 # Computation using ``meshlode.metatensor``
 # -----------------------------------------
 #
-# We now compute the same constants using the metatensor based calculator
+# We now compute the same constants using the metatensor based calculator. To achieve
+# this we first store our system parameters like the ``types``, ``positions`` and the
+# ``cell`` defined above into a :py:class`from metatensor.torch.atomistic.System: class.
+
+system = System(types=types, positions=positions, cell=cell)
 
 MP = meshlode.metatensor.MeshPotential(
     atomic_smearing=atomic_smearing,
     mesh_spacing=mesh_spacing,
     interpolation_order=interpolation_order,
     subtract_self=True,
 )
-potential_metatensor = MP.compute(frame)
+potential_metatensor = MP.compute(system)
 
 
 # %%
 # To get the Madelung constant, we again need to take a linear combination
 # of the "potentials" weighted by the charges of the atoms.
 
 atomic_energies_metatensor = torch.zeros((n_atoms, 1))
-for idx_c, c in enumerate(atomic_types):
-    for idx_n, n in enumerate(atomic_types):
-        # Take the coefficients with the correct center atom and neighbor atom species
+for idx_c, c in enumerate(types):
+    for idx_n, n in enumerate(types):
+        # Take the coefficients with the correct center atom and neighbor atom types
         block = potential_metatensor.block(
             {"center_type": int(c), "neighbor_type": int(n)}
         )
 
         # The coulomb potential between atoms i and j is charge_i * charge_j / d_ij
         # The features are simply computing a pure 1/r potential with no prefactors.
-        # Thus, to compute the energy between atoms of species i and j, we need to
+        # Thus, to compute the energy between atoms of types i and j, we need to
         # multiply by the charges of i and j.
         print(c, n, charges[idx_c] * charges[idx_n], block.values[0, 0])
         atomic_energies_metatensor[idx_c] += (

pyproject.toml

@@ -44,7 +44,7 @@ examples = [
     "matplotlib",
 ]
 metatensor = [
-    "metatensor[torch]",
+    "metatensor[torch] > 0.3",
 ]
 
 [project.urls]

src/meshlode/__init__.py

@@ -1,11 +1,10 @@
 from .calculators.meshpotential import MeshPotential
-from .lib.system import System
 
 try:
     from . import metatensor  # noqa
 except ImportError:
     pass
 
 
-__all__ = ["MeshPotential", "System"]
+__all__ = ["MeshPotential"]
 __version__ = "0.0.0-dev"

src/meshlode/calculators/meshpotential.py

@@ -4,7 +4,6 @@
 
 from meshlode.lib.fourier_convolution import FourierSpaceConvolution
 from meshlode.lib.mesh_interpolator import MeshInterpolator
-from meshlode.lib.system import System
 
 
 def _1d_tolist(x: torch.Tensor) -> List[int]:
@@ -47,14 +46,14 @@ class MeshPotential(torch.nn.Module):
 
     Define simple example structure having the CsCl (Cesium Chloride) structure
 
-    >>> atomic_types = torch.tensor([55, 17])  # Cs and Cl
+    >>> types = torch.tensor([55, 17])  # Cs and Cl
     >>> positions = torch.tensor([[0, 0, 0], [0.5, 0.5, 0.5]])
     >>> cell = torch.eye(3)
 
     Compute features
 
     >>> MP = MeshPotential(atomic_smearing=0.2, mesh_spacing=0.1, interpolation_order=4)
-    >>> MP.compute(atomic_types=atomic_types, positions=positions, cell=cell)
+    >>> MP.compute(types=types, positions=positions, cell=cell)
     tensor([[-0.5467,  1.3755],
             [ 1.3755, -0.5467]])
     """
@@ -91,9 +90,7 @@ def __init__(
         if all_types is None:
             self.all_types = None
         else:
-            self.all_types = _1d_tolist(
-                torch.unique(torch.tensor(all_types))
-            )
+            self.all_types = _1d_tolist(torch.unique(torch.tensor(all_types)))
 
         # Initilize auxiliary objects
         self.fourier_space_convolution = FourierSpaceConvolution()
@@ -116,60 +113,103 @@ def compute(
         positions: Union[List[torch.Tensor], torch.Tensor],
         cell: Union[List[torch.Tensor], torch.Tensor],
     ) -> Union[torch.Tensor, List[torch.Tensor]]:
-        """Compute the potential at the position of each atom for all provided systems.
+        """compute potential for all provided "systems" stacked inside list.
 
-        :param types: TODO
-        :param positions: TODO
-        :param cell: TODO
+        :param types: single or list of 1D tensor of integer representing the
+            particles identity. For atoms, this is typically their atomic numbers.
+        :param positions: single or 2D tensor of shape (len(types), 3) containing the
+            Cartesian positions of all particles in the system.
+        :param cell: single or 2D tensor of shape (3, 3), describing the bounding
+            box/unit cell of the system. Each row should be one of the bounding box
+            vector; and columns should contain the x, y, and z components of these
+            vectors (i.e. the cell should be given in row-major order).
 
         :return: List of torch Tensors containing the potentials for all frames and all
-            atoms. Each tensor in the list is of shape (n_atoms,n_species), where
-            n_species is the number of species in all systems combined. If the input was
+            atoms. Each tensor in the list is of shape (n_atoms, n_types), where
+            n_types is the number of types in all systems combined. If the input was
             a single system only a single torch tensor with the potentials is returned.
 
-            IMPORTANT: If multiple species are present, the different "species-channels"
+            IMPORTANT: If multiple types are present, the different "types-channels"
             are ordered according to atomic number. For example, if a structure contains
-            a water molecule with atoms 0, 1, 2 being of species O, H, H, then for this
+            a water molecule with atoms 0, 1, 2 being of types O, H, H, then for this
             system, the feature tensor will be of shape (3, 2) = (``n_atoms``,
-            ``n_species``), where ``features[0, 0]`` is the potential at the position of
+            ``n_types``), where ``features[0, 0]`` is the potential at the position of
             the Oxygen atom (atom 0, first index) generated by the HYDROGEN(!) atoms,
             while ``features[0,1]`` is the potential at the position of the Oxygen atom
             generated by the Oxygen atom(s).
         """
-        # Make sure that the compute function also works if only a single frame is
-        # provided as input (for convenience of users testing out the code)
+        # make sure compute function works if only a single tensor are provided as input
         if not isinstance(types, list):
             types = [types]
-        if not isinstance(species, list):
+        if not isinstance(positions, list):
             positions = [positions]
-        if not isinstance(species, list):
+        if not isinstance(cell, list):
             cell = [cell]
 
-        # TODO check that all inputs are on the same device and that positions and cell
-        # have the same dtype!
+        for types_single, positions_single, cell_single in zip(types, positions, cell):
+            if len(types_single.shape) != 1:
+                raise ValueError(
+                    "each `type` must be a 1 dimensional tensor, got at least "
+                    f"one tensor with {len(types_single.shape)} dimensions"
+                )
+
+            if positions_single.shape != (len(types_single), 3):
+                raise ValueError(
+                    "each `positions` must be a (n_types x 3) tensor, got at least "
+                    f"one tensor with shape {list(positions_single.shape)}"
+                )
+
+            if cell_single.shape != (3, 3):
+                raise ValueError(
+                    "each `cell` must be a (3 x 3) tensor, got at least "
+                    f"one tensor with shape {list(cell_single.shape)}"
+                )
+
+            if cell_single.dtype != positions_single.dtype:
+                raise ValueError(
+                    "`cell` must be have the same dtype as `positions`, got "
+                    f"{cell_single.dtype} and {positions_single.dtype}"
+                )
+
+            if (
+                positions_single.device != types_single.device
+                or cell_single.device != types_single.device
+            ):
+                raise ValueError(
+                    "`types`, `positions`, and `cell` must be on the same device, got "
+                    f"{types_single.device}, {positions_single.device} and "
+                    f"{cell_single.device}."
+                )
+
+        # We don't require and test that all dtypes and devices are consistent if a list
+        # of inputs. Each "frame" is processed independently.
 
         requested_types = self._get_requested_types(types)
-        n_types = len(atomic_types)
+        n_types = len(requested_types)
 
         potentials = []
         for types_single, positions_single, cell_single in zip(types, positions, cell):
             # One-hot encoding of charge information
-            charges = torch.zeros((len(types_single), n_types), dtype=positions_single.dtype)
+            charges = torch.zeros(
+                (len(types_single), n_types), dtype=positions_single.dtype
+            )
             for i_type, atomic_type in enumerate(requested_types):
                 charges[types_single == atomic_type, i_type] = 1.0
 
             # Compute the potentials
             potentials.append(
-                self._compute_single_frame(positions=positions_single, charges=charges, cell=cell_single)
+                self._compute_single_system(
+                    positions=positions_single, charges=charges, cell=cell_single
+                )
             )
 
-        if len(species) == 1:
+        if len(types) == 1:
             return potentials[0]
         else:
             return potentials
 
     def _get_requested_types(self, types: List[torch.Tensor]) -> List[int]:
-        """Extract a list of all present types from the list of types."""
+        """Extract a list of all unique and present types from the list of types."""
         all_types = torch.hstack(types)
         types_requested = _1d_tolist(torch.unique(all_types))
 
@@ -183,7 +223,7 @@ def _get_requested_types(self, types: List[torch.Tensor]) -> List[int]:
         else:
             return types_requested
 
-    def _compute_single_frame(
+    def _compute_single_system(
         self,
         positions: torch.Tensor,
         charges: torch.Tensor,
@@ -201,8 +241,8 @@ def _compute_single_frame(
             charge of atom i. More generally, the potential for the same atom positions
             is computed for n_channels independent meshes, and one can specify the
             "charge" of each atom on each of the meshes independently. For standard LODE
-            that treats all atomic species separately, one example could be: If n_atoms
-            = 4 and the species are [Na, Cl, Cl, Na], one could set n_channels=2 and use
+            that treats all (atomic) types separately, one example could be: If n_atoms
+            = 4 and the types are [Na, Cl, Cl, Na], one could set n_channels=2 and use
             the one-hot encoding charges = torch.tensor([[1,0],[0,1],[0,1],[1,0]]) for
             the charges. This would then separately compute the "Na" potential and "Cl"
             potential. Subtracting these from each other, one could recover the more