Replace pair_id in neighbors list samples with the actual cell shift

docs/extensions/rascaline_json_schema.py

@@ -3,9 +3,8 @@
 
 from docutils import nodes
 from docutils.parsers.rst import Directive
-from markdown_it import MarkdownIt
-
 from html_hidden import html_hidden
+from markdown_it import MarkdownIt
 from myst_parser.config.main import MdParserConfig
 from myst_parser.mdit_to_docutils.base import DocutilsRenderer
 
@@ -121,7 +120,7 @@ def _json_schema_to_nodes(self, schema, inline=False):
                         if isinstance(subfields, nodes.literal):
                             subfields = [subfields]
 
-                        for (i, sf) in enumerate(subfields):
+                        for i, sf in enumerate(subfields):
                             field += sf
 
                             if isinstance(sf, nodes.inline):
@@ -152,8 +151,11 @@ def _json_schema_to_nodes(self, schema, inline=False):
                     raise Exception(f"unknown integer format: {schema['format']}")
 
             elif schema["type"] == "string":
-                # TODO enums?
-                return nodes.literal(text="string")
+                if "enum" in schema:
+                    values = [f'"{v}"' for v in schema["enum"]]
+                    return nodes.literal(text=" | ".join(values))
+                else:
+                    return nodes.literal(text="string")
 
             elif schema["type"] == "boolean":
                 return nodes.literal(text="boolean")

python/rascaline/rascaline/_c_api.py

@@ -46,6 +46,7 @@ class rascal_pair_t(ctypes.Structure):
         ("second", c_uintptr_t),
         ("distance", ctypes.c_double),
         ("vector", ctypes.c_double * 3),
+        ("cell_shift_indices", ctypes.c_int32 * 3),
     ]
 
 

python/rascaline/rascaline/calculators.py

@@ -37,29 +37,48 @@ def __init__(self, cutoff, delta, name):
 
 class NeighborList(CalculatorBase):
     """
-    This calculator computes the neighbor list for a given spherical cutoff, and
-    returns the list of distance vectors between all pairs of atoms strictly
-    inside the cutoff.
-
-    Users can request either a "full" neighbor list (including an entry for both
-    ``i - j`` pairs and ``j - i`` pairs) or save memory/computational by only
-    working with "half" neighbor list (only including one entry for each ``i/j``
-    pair)
-
-    Self pairs (pairs between an atom and periodic copy itself) can appear when
-    the cutoff is larger than the cell under periodic boundary conditions. Self
-    pairs with a distance of 0 are only included when the user passes
-    ``self_pairs=True``, which is not the default behavior.
-
-    This sample produces a single property (``"distance"``) with three
-    components (``"pair_direction"``) containing the x, y, and z component of
-    the vector from the first atom in the pair to the second. In addition to the
-    atom indexes, the samples also contain a pair index, to be able to
-    distinguish between multiple pairs between the same atom (if the cutoff is
-    larger than the cell).
+    This calculator computes the neighbor list for a given spherical cutoff, and returns
+    the list of distance vectors between all pairs of atoms strictly inside the cutoff.
+
+    Users can request either a "full" neighbor list (including an entry for both ``i -
+    j`` pairs and ``j - i`` pairs) or save memory/computational by only working with
+    "half" neighbor list (only including one entry for each ``i/j`` pair)
+
+    Pairs between an atom and it's own periodic copy can appear when the cutoff is
+    larger than the cell under periodic boundary conditions. Self pairs with a distance
+    of 0 (i.e. self pairs inside the original unit cell) are only included when using
+    ``self_pairs=True``.
+
+    The ``quantity`` parameter determine what will be included in the output. It can
+    take one of three values:
+
+    - ``"Distance"``, to get the distance between the atoms, accounting for periodic
+      boundary conditions. This is the default.
+    - ``"CellShiftVector"``, to get the cell shift vector, which can then be used to
+      apply periodic boundary conditions and compute the distance vector.
+
+      If ``S`` is the cell shift vector, ``rij`` the pair distance vector, ``ri`` and
+      ``rj`` the position of the atoms, ``rij = rj - ri + S``.
+    - ``"CellShiftIndices"``, to get three integers indicating the number of cell
+      vectors (``A``, ``B``, and ``C``) entering the cell shift.
+
+      If the cell vectors are ``A``, ``B``, and ``C``, this returns three integers
+      ``sa``, ``sb``, and ``sc`` such that the cell shift ``S = sa * A + sb * B + sc *
+      C``.
+
+    This calculator produces a single property (``"distance"``, ``"cell_shift_vector"``,
+    or ``"cell_shift_indices"``) with three components (``"pair_direction"``) containing
+    the x, y, and z component of the requested vector. In addition to the atom indexes,
+    the samples also contain a pair index, to be able to distinguish between multiple
+    pairs between the same atom (if the cutoff is larger than the cell).
     """
 
-    def __init__(self, cutoff, full_neighbor_list, self_pairs=False):
+    def __init__(
+        self,
+        cutoff: float,
+        full_neighbor_list: bool,
+        self_pairs: bool = False,
+    ):
         parameters = {
             "cutoff": cutoff,
             "full_neighbor_list": full_neighbor_list,

python/rascaline/rascaline/systems/ase.py

@@ -89,21 +89,21 @@ def compute_neighbors(self, cutoff):
 
         self._pairs = []
 
-        nl_result = neighborlist.neighbor_list("ijdD", self._atoms, cutoff)
-        for i, j, d, D in zip(*nl_result):
+        nl_result = neighborlist.neighbor_list("ijdDS", self._atoms, cutoff)
+        for i, j, d, D, S in zip(*nl_result):
             if j < i:
                 # we want a half neighbor list, so drop all duplicated
                 # neighbors
                 continue
-            self._pairs.append((i, j, d, D))
+            self._pairs.append((i, j, d, D, S))
 
         self._pairs_by_center = []
         for _ in range(self.size()):
             self._pairs_by_center.append([])
 
-        for i, j, d, D in self._pairs:
-            self._pairs_by_center[i].append((i, j, d, D))
-            self._pairs_by_center[j].append((i, j, d, D))
+        for i, j, d, D, S in self._pairs:
+            self._pairs_by_center[i].append((i, j, d, D, S))
+            self._pairs_by_center[j].append((i, j, d, D, S))
 
     def pairs(self):
         return self._pairs

python/rascaline/rascaline/systems/base.py

@@ -314,17 +314,18 @@ def pairs(self):
         The pairs are those which were computed by the last call
         :py:func:`SystemBase.compute_neighbors`
 
-        Get all neighbor pairs in this system as a list of tuples ``(int, int,
-        float, (float, float, float))`` containing the indexes of the first and
-        second atom in the pair, the distance between the atoms, and the wrapped
-        between them. Alternatively, this function can return a 1D numpy array
-        with ``dtype=rascal_pair_t``.
-
-        The list of pair should only contain each pair once (and not twice as
-        ``i-j`` and ``j-i``), should not contain self pairs (``i-i``); and
-        should only contains pairs where the distance between atoms is actually
-        bellow the cutoff passed in the last call to
-        :py:func:`rascaline.SystemBase.compute_neighbors`.
+        Get all neighbor pairs in this system as a list of tuples ``(int, int, float,
+        (float, float, float), (int, int, int))`` containing the indexes of the first
+        and second atom in the pair, the distance between the atoms, the vector between
+        them, and the cell shift. The vector should be ``position[first] -
+        position[second] * + H * cell_shift`` where ``H`` is the cell matrix.
+        Alternatively, this function can return a 1D numpy array with
+        ``dtype=rascal_pair_t``.
+
+        The list of pair should only contain each pair once (and not twice as ``i-j``
+        and ``j-i``), should not contain self pairs (``i-i``); and should only contains
+        pairs where the distance between atoms is actually bellow the cutoff passed in
+        the last call to :py:func:`rascaline.SystemBase.compute_neighbors`.
 
         This function is only valid to call after a call to
         :py:func:`rascaline.SystemBase.compute_neighbors` to set the cutoff.

python/rascaline/tests/calculators/dummy_calculator.py

@@ -51,8 +51,8 @@ def test_compute():
 
     H_block = descriptor.block(species_center=1)
     assert H_block.values.shape == (2, 2)
-    assert np.all(H_block.values[0] == (2, 1))
-    assert np.all(H_block.values[1] == (3, 3))
+    assert np.all(H_block.values[0] == (2, 11))
+    assert np.all(H_block.values[1] == (3, 13))
 
     assert len(H_block.samples) == 2
     assert H_block.samples.names == ["structure", "center"]

python/rascaline/tests/calculators/sample_selection.py

@@ -49,8 +49,8 @@ def test_selection():
 
     H_block = descriptor.block(species_center=1)
     assert H_block.values.shape == (2, 2)
-    assert np.all(H_block.values[0] == (2, 1))
-    assert np.all(H_block.values[1] == (3, 3))
+    assert np.all(H_block.values[0] == (2, 11))
+    assert np.all(H_block.values[1] == (3, 13))
 
     O_block = descriptor.block(species_center=8)
     assert O_block.values.shape == (1, 2)
@@ -72,8 +72,8 @@ def test_subset_variables():
 
     H_block = descriptor.block(species_center=1)
     assert H_block.values.shape == (2, 2)
-    assert np.all(H_block.values[0] == (2, 1))
-    assert np.all(H_block.values[1] == (3, 3))
+    assert np.all(H_block.values[0] == (2, 11))
+    assert np.all(H_block.values[1] == (3, 13))
 
     O_block = descriptor.block(species_center=8)
     assert O_block.values.shape == (1, 2)
@@ -128,7 +128,7 @@ def test_predefined_selection():
 
     H_block = descriptor.block(species_center=1)
     assert H_block.values.shape == (1, 2)
-    assert np.all(H_block.values[0] == (3, 3))
+    assert np.all(H_block.values[0] == (3, 13))
 
     O_block = descriptor.block(species_center=8)
     assert O_block.values.shape == (1, 2)

python/rascaline/tests/test_systems.py

@@ -9,7 +9,7 @@ def species(self):
         return [1, 1, 8, 8]
 
     def positions(self):
-        return [[0, 0, 0], [0, 0, 1], [0, 0, 2], [0, 0, 3]]
+        return [[0, 0, 10], [0, 0, 1], [0, 0, 2], [0, 0, 3]]
 
     def cell(self):
         return [[10, 0, 0], [0, 10, 0], [0, 0, 10]]
@@ -19,29 +19,29 @@ def compute_neighbors(self, cutoff):
 
     def pairs(self):
         return [
-            (0, 1, 1.0, (0.0, 0.0, 1.0)),
-            (1, 2, 1.0, (0.0, 0.0, 1.0)),
-            (2, 3, 1.0, (0.0, 0.0, 1.0)),
+            (0, 1, 1.0, (0.0, 0.0, 1.0), (0, 0, 1)),
+            (1, 2, 1.0, (0.0, 0.0, 1.0), (0, 0, 0)),
+            (2, 3, 1.0, (0.0, 0.0, 1.0), (0, 0, 0)),
         ]
 
     def pairs_containing(self, center):
         if center == 0:
             return [
-                (0, 1, 1.0, (0.0, 0.0, 1.0)),
+                (0, 1, 1.0, (0.0, 0.0, 1.0), (0, 0, 1)),
             ]
         elif center == 1:
             return [
-                (0, 1, 1.0, (0.0, 0.0, 1.0)),
-                (1, 2, 1.0, (0.0, 0.0, 1.0)),
+                (0, 1, 1.0, (0.0, 0.0, 1.0), (0, 0, 1)),
+                (1, 2, 1.0, (0.0, 0.0, 1.0), (0, 0, 0)),
             ]
         elif center == 2:
             return [
-                (1, 2, 1.0, (0.0, 0.0, 1.0)),
-                (2, 3, 1.0, (0.0, 0.0, 1.0)),
+                (1, 2, 1.0, (0.0, 0.0, 1.0), (0, 0, 0)),
+                (2, 3, 1.0, (0.0, 0.0, 1.0), (0, 0, 0)),
             ]
         elif center == 3:
             return [
-                (2, 3, 1.0, (0.0, 0.0, 1.0)),
+                (2, 3, 1.0, (0.0, 0.0, 1.0), (0, 0, 0)),
             ]
         else:
             raise Exception("got invalid center")

rascaline-c-api/include/rascaline.h

@@ -131,10 +131,17 @@ typedef struct rascal_pair_t {
    */
   double distance;
   /**
-   * vector from the first atom to the second atom, wrapped inside the unit
-   * cell as required by periodic boundary conditions.
+   * vector from the first atom to the second atom, accounting for periodic
+   * boundary conditions. This should be
+   * `position[second] - position[first] + H * cell_shift`
+   * where `H` is the cell matrix.
    */
   double vector[3];
+  /**
+   * How many cell shift where applied to the `second` atom to create this
+   * pair.
+   */
+  int32_t cell_shift_indices[3];
 } rascal_pair_t;
 
 /**

rascaline-c-api/src/system.rs

@@ -19,9 +19,14 @@ pub struct rascal_pair_t {
     pub second: usize,
     /// distance between the two atoms
     pub distance: f64,
-    /// vector from the first atom to the second atom, wrapped inside the unit
-    /// cell as required by periodic boundary conditions.
+    /// vector from the first atom to the second atom, accounting for periodic
+    /// boundary conditions. This should be
+    /// `position[second] - position[first] + H * cell_shift`
+    /// where `H` is the cell matrix.
     pub vector: [f64; 3],
+    /// How many cell shift where applied to the `second` atom to create this
+    /// pair.
+    pub cell_shift_indices: [i32; 3],
 }
 
 /// A `rascal_system_t` deals with the storage of atoms and related information,

rascaline-c-api/tests/calculator.cpp

@@ -169,7 +169,7 @@ TEST_CASE("Compute descriptor") {
             1, 0, /**/ 0, 1,
         };
         auto values = std::vector<double>{
-            5, 9, /**/ 6, 18, /**/ 7, 15,
+            5, 39, /**/ 6, 18, /**/ 7, 15,
         };
         auto gradient_samples = std::vector<int32_t>{
             0, 0, 0, /**/ 0, 0, 1, /**/ 0, 0, 2,
@@ -194,7 +194,7 @@ TEST_CASE("Compute descriptor") {
             0, 0,
         };
         values = std::vector<double>{
-            4, 3,
+            4, 33,
         };
         gradient_samples = std::vector<int32_t>{
             0, 0, 0, /**/ 0, 0, 1,
@@ -260,7 +260,7 @@ TEST_CASE("Compute descriptor") {
             1, 0, /**/ 0, 1,
         };
         auto values = std::vector<double>{
-            5, 9, /**/ 7, 15,
+            5, 39, /**/ 7, 15,
         };
         auto gradient_samples = std::vector<int32_t>{
             0, 0, 0, /**/ 0, 0, 1, /**/ 0, 0, 2,
@@ -337,7 +337,7 @@ TEST_CASE("Compute descriptor") {
             0, 1,
         };
         auto values = std::vector<double>{
-            9, /**/ 18, /**/ 15,
+            39, /**/ 18, /**/ 15,
         };
         auto gradient_samples = std::vector<int32_t>{
             0, 0, 0, /**/ 0, 0, 1, /**/ 0, 0, 2,
@@ -362,7 +362,7 @@ TEST_CASE("Compute descriptor") {
             0, 0,
         };
         values = std::vector<double>{
-            3,
+            33,
         };
         gradient_samples = std::vector<int32_t>{
             0, 0, 0, /**/ 0, 0, 1,
@@ -575,7 +575,7 @@ TEST_CASE("Compute descriptor") {
             1, 0, 0, 1
         };
         values = std::vector<double>{
-            4, 3
+            4, 33
         };
         gradient_samples = std::vector<int32_t>{
             0, 0, 0, /**/ 0, 0, 1,