From 9a6a617888cd859a44a265043dac1fee472708d6 Mon Sep 17 00:00:00 2001 From: sayeg84 Date: Wed, 28 May 2025 16:54:06 +0200 Subject: [PATCH 1/2] Add batching to constraint finder --- src/aggforce/constraints/constfinder.py | 45 ++++++++++++++++++---- src/aggforce/util.py | 50 ++++++++++++++++++++++++- 2 files changed, 87 insertions(+), 8 deletions(-) diff --git a/src/aggforce/constraints/constfinder.py b/src/aggforce/constraints/constfinder.py index 4cac0f1..e55b498 100644 --- a/src/aggforce/constraints/constfinder.py +++ b/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 diff --git a/src/aggforce/util.py b/src/aggforce/util.py index 36bc931..e874f99 100644 --- a/src/aggforce/util.py +++ b/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: From 0c25a7b1dbd41634b7a875abbd90e6444fa8625f Mon Sep 17 00:00:00 2001 From: sayeg84 Date: Wed, 28 May 2025 17:10:10 +0200 Subject: [PATCH 2/2] Black --- src/aggforce/constraints/constfinder.py | 2 +- src/aggforce/util.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/src/aggforce/constraints/constfinder.py b/src/aggforce/constraints/constfinder.py index e55b498..304b8f9 100644 --- a/src/aggforce/constraints/constfinder.py +++ b/src/aggforce/constraints/constfinder.py @@ -30,7 +30,7 @@ def guess_pairwise_constraints( 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 + 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 diff --git a/src/aggforce/util.py b/src/aggforce/util.py index e874f99..d339e11 100644 --- a/src/aggforce/util.py +++ b/src/aggforce/util.py @@ -6,6 +6,7 @@ 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. @@ -54,7 +55,6 @@ def chunker(array: np.ndarray, n_batches: int) -> List[np.ndarray]: return batched_array - T = TypeVar("T")