Merge pull request #225 from nlesc-nano/adf_periodic

ENH: Add support for periodic ADF calculations with `r_max != inf`

FOX/classes/multi_mol.py

@@ -1630,11 +1630,6 @@ def init_adf(
 
         .. _DASK: https://dask.org/
 
-        Note
-        ----
-        Periodic ADF calculations (see the **periodic** parameter) are currently
-        only supported for systems with cuboid latices and ``r_max != inf``.
-
         Parameters
         ----------
         mol_subset : :class:`slice`, optional
@@ -1713,37 +1708,22 @@ def init_adf(
             else:
                 lattice = cast("np.ndarray[Any, np.dtype[np.float64]]", lattice[m_subset])
 
-            # Scipy's `cKDTree` only supports cuboid lattices
-            if r_max_:
-                is0 = np.abs(lattice - 0) < 1e-8
-                if not (np.count_nonzero(is0, axis=-1) == 2).all():
-                    raise NotImplementedError("non-cuboid lattices are not supported")
-
             # Set the vector-length of all absent axes to `inf`
             slc = [i for i in range(3) if i not in periodic_ar]
             lattice[..., slc, :] = np.inf
 
-            # Perform a translation to remove negative elements, as `cKDTree` cannot
-            # handle the latter if `boxsize` is specified
-            mol -= mol.min(axis=1)[..., None, :]
-            if lattice.ndim == 2:
-                boxsize_iter = repeat(np.linalg.norm(lattice, axis=-1))
-                lattice_iter = repeat(lattice)
-            else:
-                boxsize_iter = iter(np.linalg.norm(lattice, axis=-1))
-                lattice_iter = iter(lattice)
+            lattice_iter = repeat(lattice) if lattice.ndim == 2 else iter(lattice)
             periodic_iter = repeat(periodic_ar)
         else:
             lattice_iter = repeat(None)
             periodic_iter = repeat(range(3))
-            boxsize_iter = repeat(None)
 
         # Construct the angular distribution function
         # Perform the task in parallel (with dask) if possible
         if DASK_EX is None and r_max_:
             func = dask.delayed(_adf_inner_cdktree)
-            jobs = [func(m, n, r_max_, atom_pairs.values(), b, weight) for m, b in
-                    zip(mol, boxsize_iter)]
+            jobs = [func(m, n, r_max_, atom_pairs.values(), l, p, weight) for m, l, p in
+                    zip(mol, lattice_iter, periodic_iter)]
             results = dask.compute(*jobs)
         elif DASK_EX is None and not r_max_:
             func = dask.delayed(_adf_inner)
@@ -1752,8 +1732,8 @@ def init_adf(
             results = dask.compute(*jobs)
         elif DASK_EX is not None and r_max_:
             func = _adf_inner_cdktree
-            results = [func(m, n, r_max_, atom_pairs.values(), b, weight) for m, b in
-                       zip(mol, boxsize_iter)]
+            results = [func(m, n, r_max_, atom_pairs.values(), l, p, weight) for m, l, p in
+                       zip(mol, lattice_iter, periodic_iter)]
         elif DASK_EX is not None and not r_max_:
             func = _adf_inner
             results = [func(m, atom_pairs.values(), l, p, weight) for m, l, p in

FOX/functions/adf.py

@@ -55,25 +55,32 @@ def _adf_inner_cdktree(
     n: int,
     r_max: float,
     idx_list: Iterable[_3Tuple[NDArray[np.bool_]]],
-    boxsize: None | NDArray[f8],
+    lattice: None | NDArray[f8],
+    periodicity: Iterable[Literal[0, 1, 2]] = range(3),
     weight: None | Callable[[NDArray[f8]], NDArray[f8]] = None,
 ) -> List[NDArray[f8]]:
     """Perform the loop of :meth:`.init_adf` with a distance cutoff."""
     # Construct slices and a distance matrix
-    tree = cKDTree(m, boxsize=boxsize)
-    dist, idx = tree.query(m, n, distance_upper_bound=r_max, p=2)  # type: NDArray[f8], NDArray[i8]  # noqa: E501
-    dist[dist == np.inf] = 0.0
-    idx[idx == m.shape[0]] = 0
-
-    # Slice the Cartesian coordinates
-    coords13: NDArray[f8] = m[idx]
-    coords2: NDArray[f8] = m[..., None, :]
-
-    # Construct (3D) angle- and distance-matrices
-    with np.errstate(divide='ignore', invalid='ignore'):
-        vec: NDArray[f8] = ((coords13 - coords2) / dist[..., None])
-        ang: NDArray[f8] = np.arccos(np.einsum('jkl,jml->jkm', vec, vec))
-        dist = np.maximum(dist[..., None], dist[..., None, :])
+    if lattice is not None:
+        with np.errstate(divide='ignore', invalid='ignore'):
+            dist, vec, idx = _adf_inner_cdktree_periodic(m, n, r_max, lattice, periodicity)
+            ang: NDArray[f8] = np.arccos(np.einsum('jkl,jml->jkm', vec, vec))
+            dist = np.maximum(dist[..., None], dist[..., None, :])
+    else:
+        tree = cKDTree(m)
+        dist, idx = tree.query(m, n, distance_upper_bound=r_max, p=2)
+        dist[dist == np.inf] = 0.0
+        idx[idx == len(m)] = 0
+
+        # Slice the Cartesian coordinates
+        coords13: NDArray[f8] = m[idx]
+        coords2: NDArray[f8] = m[..., None, :]
+
+        # Construct (3D) angle- and distance-matrices
+        with np.errstate(divide='ignore', invalid='ignore'):
+            vec = ((coords13 - coords2) / dist[..., None])
+            ang = np.arccos(np.einsum('jkl,jml->jkm', vec, vec))
+            dist = np.maximum(dist[..., None], dist[..., None, :])
     ang[np.isnan(ang)] = 0.0
 
     # Radian (float) to degrees (int)
@@ -88,6 +95,34 @@ def _adf_inner_cdktree(
     return ret
 
 
+def _adf_inner_cdktree_periodic(
+    m: NDArray[f8],
+    n: int,
+    r_max: float,
+    lattice: NDArray[f8],
+    periodicity: Iterable[Literal[0, 1, 2]],
+) -> Tuple[NDArray[f8], NDArray[f8], NDArray[np.intp]]:
+    # Construct the (full) distance matrix and vectors
+    dist, vec = _adf_inner_periodic(m, lattice, periodicity)
+
+    # Apply `n` and `r_max`: truncate the number of distances/vectors
+    idx1 = np.argsort(dist, axis=1)
+    if n < idx1.shape[1]:
+        idx1 = idx1[:, :n]
+    dist = np.take_along_axis(dist, idx1, axis=1)
+    mask = dist > r_max
+    idx1[mask] = 0
+    dist[mask] = 0.0
+
+    # Return the subsets
+    idx0 = np.empty_like(idx1)
+    idx0[:] = np.arange(len(idx0))[..., None]
+    i = idx0.ravel()
+    j = idx1.ravel()
+    vec_ret = vec[i, j].reshape(*dist.shape, 3)
+    return dist, vec_ret, idx1
+
+
 def _adf_inner(
     m: NDArray[f8],
     idx_list: Iterable[_3Tuple[NDArray[np.bool_]]],
@@ -128,7 +163,7 @@ def _adf_inner(
 def _adf_inner_periodic(
     m: NDArray[f8],
     lattice: NDArray[f8],
-    periodicity: Iterable[Literal[0, 1, 2]] = range(3),
+    periodicity: Iterable[Literal[0, 1, 2]],
 ) -> Tuple[NDArray[f8], NDArray[f8]]:
     """Construct the distance matrix and angle-defining vectors for periodic systems."""
     vec = m - m[..., None, :]

tests/test_multi_mol.py

@@ -34,14 +34,6 @@
 MOL_LATTICE_2D = MOL_LATTICE_3D.copy()
 MOL_LATTICE_2D.lattice = MOL_LATTICE_2D.lattice[0]
 
-MOL_LATTICE_3D_ORTH = MOL_LATTICE_3D.copy()
-MOL_LATTICE_3D_ORTH.lattice = np.zeros_like(MOL_LATTICE_3D_ORTH.lattice)
-MOL_LATTICE_3D_ORTH.lattice[..., 0] = np.linalg.norm(MOL_LATTICE_3D.lattice, axis=-1)
-
-MOL_LATTICE_2D_ORTH = MOL_LATTICE_3D.copy()
-MOL_LATTICE_2D_ORTH.lattice = np.zeros_like(MOL_LATTICE_2D_ORTH.lattice)
-MOL_LATTICE_2D_ORTH.lattice[..., 0] = np.linalg.norm(MOL_LATTICE_2D.lattice, axis=-1)
-
 
 @delete_finally(join(PATH, '.tmp.xyz'))
 def test_mol_to_file():
@@ -289,8 +281,8 @@ class TestADF:
         [
             ("adf_weighted", MOL, {"atom_subset": ("Cd", "Se")}),
             ("adf", MOL, {"atom_subset": ("Cd", "Se"), "weight": None}),
-            ("adf_periodic_2d", MOL_LATTICE_2D_ORTH, {"atom_subset": ("Pb",), "periodic": "xyz"}),
-            ("adf_periodic_3d", MOL_LATTICE_3D_ORTH, {"atom_subset": ("Pb",), "periodic": "xy"}),
+            ("adf_periodic_2d", MOL_LATTICE_2D, {"atom_subset": ("Pb",), "periodic": "xyz"}),
+            ("adf_periodic_3d", MOL_LATTICE_3D, {"atom_subset": ("Pb",), "periodic": "xy"}),
             ("adf_2d_inf", MOL_LATTICE_2D, {"atom_subset": ("Pb",), "r_max": np.inf}),
             ("adf_3d_inf", MOL_LATTICE_3D, {"atom_subset": ("Pb",), "r_max": np.inf}),
             ("adf_periodic_2d_inf", MOL_LATTICE_2D,
@@ -310,7 +302,6 @@ def test_passes(self, name: str, mol: MultiMolecule, kwargs: Mapping[str, Any])
         "kwargs,exc",
         [
             ({"periodic": "bob"}, ValueError),
-            ({"periodic": "xyz"}, NotImplementedError),
         ]
     )
     @pytest.mark.parametrize("mol", [MOL_LATTICE_2D, MOL_LATTICE_3D])