Add batching to constraint guesser

src/aggforce/constraints/constfinder.py

@@ -6,11 +6,13 @@
 """
 
 import numpy as np
-from ..util import distances
+from ..util import distances, chunker
 from .hints import Constraints
 
 
-def guess_pairwise_constraints(xyz: np.ndarray, threshold: float = 1e-3) -> Constraints:
+def guess_pairwise_constraints(
+    xyz: np.ndarray, threshold: float = 1e-3, n_batches: int = 1
+) -> Constraints:
     """Find pairs of sites which are likely constrained via fluctuations.
 
     Fluctuations are found by
@@ -27,14 +29,43 @@ def guess_pairwise_constraints(xyz: np.ndarray, threshold: float = 1e-3) -> Cons
     threshold (positive float):
         Distances with standard deviations lower than this value are considered
         to be constrainted. Has units of xyz.
+    n_batches (int):
+        number of baches over which to divide the number of sites. As the
+        constraint finder has O(n_sites^2) memory requirement, large systems
+        will require a lot of RAM, so batching is required for low memory systems
 
     Returns:
     -------
     A set of frozen sets, each of which contains a pair of indices of sites
     which are guessed to be pairwise constrained.
     """
-    dists = distances(xyz)
-    sds = np.sqrt(np.var(dists, axis=0))
-    np.fill_diagonal(sds, threshold * 2)
-    inds = np.nonzero(sds < threshold)
-    return {frozenset(v) for v in zip(*inds)}
+    if n_batches == 1:
+        dists = distances(xyz)
+        sds = np.sqrt(np.var(dists, axis=0))
+        np.fill_diagonal(sds, threshold * 2)
+        inds = np.nonzero(sds < threshold)
+        return {frozenset(v) for v in zip(*inds)}
+    else:
+        n_sites = xyz.shape[1]
+        elem_chunks = chunker(np.arange(n_sites), n_batches)
+        constraints = set()
+        for i, first_entry_chunk in enumerate(elem_chunks):
+            for j, second_entry_chunk in enumerate(elem_chunks):
+                displacement_matrix = (
+                    xyz[:, None, first_entry_chunk, :]
+                    - xyz[:, second_entry_chunk, None, :]
+                )
+                distance_matrix = np.linalg.norm(displacement_matrix, axis=-1)
+                sds = np.sqrt(np.var(distance_matrix, axis=0))
+                if i == j:
+                    np.fill_diagonal(sds, threshold * 2)
+                inds = np.nonzero(sds < threshold)
+                if len(inds[0]) > 0 and len(inds[1]) > 0:
+                    local_constraints = {
+                        frozenset(v)
+                        for v in zip(
+                            second_entry_chunk[inds[0]], first_entry_chunk[inds[1]]
+                        )
+                    }
+                    constraints = constraints.union(local_constraints)
+        return constraints

src/aggforce/util.py

@@ -6,6 +6,55 @@
 from typing import Union, Callable, TypeVar, Iterable, Any, List, Generic
 import numpy as np
 
+
+def chunker(array: np.ndarray, n_batches: int) -> List[np.ndarray]:
+    """
+    Chunks an input array into a specified number of batches.
+
+    This function divides the input array into approximately equal-sized chunks.
+    The last chunk may contain more elements if the array length is not perfectly
+    divisible by the number of batches.
+
+    Parameters:
+    -----------
+    array : np.ndarray or List
+        The input array to be chunked.
+    n_batches : int
+        The number of batches to divide the array into. Must be a positive
+        integer and less than or equal to the length of the array.
+
+    Returns:
+    --------
+    batched_array: List
+        A list of lists/arrays, where each inner list/array is a chunk of the original array.
+
+    Examples:
+    >>> chunker([1, 2, 3, 4, 5, 6, 7, 8, 9], 3)
+    [[1, 2, 3], [4, 5, 6], [7, 8, 9]]
+
+    >>> chunker([1, 2, 3, 4, 5], 2)
+    [[1, 2], [3, 4, 5]]
+
+    >>> chunker([1, 2, 3, 4, 5], 5)
+    [[1], [2], [3], [4], [5]]
+
+    >>> chunker([1, 2, 3, 4, 5], 1)
+    [[1, 2, 3, 4, 5]]
+    """
+    if n_batches == 1:
+        return [array]
+    assert n_batches <= len(
+        array
+    ), "n_batches needs to be smaller than the array to chunk"
+    batched_array = []
+    n_elts_per_batch = len(array) // n_batches
+    for i in range(n_batches - 1):
+        batched_array.append(array[i * n_elts_per_batch : (i + 1) * n_elts_per_batch])
+    # last batch might be larger, it contains the rest of the elements in the array
+    batched_array.append(array[(i + 1) * n_elts_per_batch :])
+    return batched_array
+
+
 T = TypeVar("T")
 
 
@@ -60,7 +109,6 @@ def distances(
         raise ValueError("Cross distances only supported when return_matrix is truthy.")
     if return_displacements and not return_matrix:
         raise ValueError("Displacements only supported when return_matrix is truthy.")
-
     if cross_xyz is None:
         displacement_matrix = xyz[:, None, :, :] - xyz[:, :, None, :]
     else: