Skip to content

Add backends support for minimize_vectors #5235

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
PardhavMaradani wants to merge 12 commits into
base: develop
Choose a base branch
from
Open

Add backends support for minimize_vectors #5235

Show file tree
Hide file tree
Changes from 5 commits
Commits
Show all changes
12 commits
Select commit Hold shift + click to select a range
da2595a
Add backends support for minimize_vectors
PardhavMaradani Feb 18, 2026
5b63e60
* Explicitly set language to c++ to avoid darwin build error
PardhavMaradani Feb 18, 2026
d7d8745
* Add cpp options to openmp distances as well
PardhavMaradani Feb 18, 2026
b24d0ce
* Fix source suffix
PardhavMaradani Feb 18, 2026
b0bb72e
* Include .hpp files for sdist
PardhavMaradani Feb 18, 2026
de5dfa0
* Actually pass the backend param to the tests
PardhavMaradani Feb 20, 2026
ef8cb4f
Merge branch 'develop' into minimize-vectors-backend-support
PardhavMaradani Feb 20, 2026
cc20dd8
* Remove distutils designator as it is not needed (after changing lan…
PardhavMaradani Feb 20, 2026
81b90e3
Merge branch 'minimize-vectors-backend-support'
PardhavMaradani Feb 20, 2026
3454f6d
* Add asv benchmarks
PardhavMaradani Feb 21, 2026
be2b4db
* Skip ValueError as well for missing runtime backends
PardhavMaradani Feb 21, 2026
3baccdb
Merge branch 'develop' into minimize-vectors-backend-support
PardhavMaradani Feb 26, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions package/AUTHORS
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -271,6 +271,7 @@ Chronological list of authors
- Mohammad Ayaan
- Khushi Phougat
- Kushagar Garg
- Pardhav Maradani

External code
-------------
Expand Down
3 changes: 2 additions & 1 deletion package/CHANGELOG
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -15,7 +15,7 @@ The rules for this file:

-------------------------------------------------------------------------------
??/??/?? IAlibay, orbeckst, marinegor, tylerjereddy, ljwoods2, marinegor,
spyke7, talagayev, tanii1125, BradyAJohnston
spyke7, talagayev, tanii1125, BradyAJohnston, PardhavMaradani

* 2.11.0

Expand All @@ -37,6 +37,7 @@ Fixes
DSSP by porting upstream PyDSSP 0.9.1 fix (Issue #4913)

Enhancements
* Adds support for backend options in `minimize_vectors` (Issue #5234, PR #5235)
* Adds support for parsing `.tpr` files produced by GROMACS 2026.0
* Enables parallelization for analysis.diffusionmap.DistanceMatrix
(Issue #4679, PR #4745)
Expand Down
1 change: 1 addition & 0 deletions package/MANIFEST.in
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@ global-include *.pyx
global-include *.pxd
global-include *.c
global-include *.h
global-include *.hpp
global-include *.cpp
include MDAnalysis/lib/src/transformations/*
recursive-include doc *
Expand Down
12 changes: 12 additions & 0 deletions package/MDAnalysis/lib/_distopia.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -190,3 +190,15 @@ def calc_self_distance_array_triclinic(
coords: np.ndarray, box: np.ndarray, results: np.ndarray
) -> None:
distopia.self_distance_array_triclinic(coords, box, results=results)


def calc_minimize_vectors_ortho(
vectors: np.ndarray, box: np.ndarray, output: np.ndarray
) -> None:
distopia.minimize_vectors_ortho(vectors, box[:3], output)


def calc_minimize_vectors_triclinic(
vectors: np.ndarray, box: np.ndarray, output: np.ndarray
) -> None:
distopia.minimize_vectors_triclinic(vectors, box, output)
7 changes: 6 additions & 1 deletion package/MDAnalysis/lib/c_distances.pxd
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,3 +1,4 @@
# distutils: language = c++
from libc.stdint cimport uint64_t, UINT64_MAX


Expand Down Expand Up @@ -26,4 +27,8 @@ cdef extern from "calc_distances.h":
void _ortho_pbc(coordinate* coords, uint64_t numcoords, float* box)
void _triclinic_pbc(coordinate* coords, uint64_t numcoords, float* box)
void minimum_image(double* x, float* box, float* inverse_box)
void minimum_image_triclinic(float* x, float* box, float* inverse_box)
void minimum_image_triclinic(float* x, float* box, float* inverse_box)

cdef extern from "calc_distances.hpp":
void _calc_minimize_vectors_ortho[T](T* vectors, uint64_t numvectors, T* box, T* output)
void _calc_minimize_vectors_triclinic[T](T* vectors, uint64_t numvectors, T* box, T* output)
149 changes: 13 additions & 136 deletions package/MDAnalysis/lib/c_distances.pyx
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -33,6 +33,7 @@ cimport cython
from libc.stdint cimport uint64_t, UINT64_MAX
import numpy
cimport numpy
from cython cimport floating
numpy.import_array()

from libc.math cimport round as cround
Expand Down Expand Up @@ -272,143 +273,19 @@ def contact_matrix_pbc(coord, sparse_contacts, box, cutoff):

@cython.boundscheck(False)
@cython.wraparound(False)
cdef inline void _minimum_image_orthogonal(cython.floating[:] dx,
cython.floating[:] box,
cython.floating[:] inverse_box) nogil:
"""Minimize dx to be the shortest vector

Parameters
----------
dx : numpy.array, shape (3,)
vector to minimize
box : numpy.array, shape (3,)
box length in each dimension
inverse_box : numpy.array, shape (3,)
inverse of box

Operates in-place on dx!
"""
cdef int i
cdef cython.floating s

for i in range(3):
if box[i] > 0:
s = inverse_box[i] * dx[i]
dx[i] = box[i] * (s - cround(s))
def calc_minimize_vectors_ortho(
floating[:, ::1] vectors,
floating[::1] box,
floating[:, ::1] output) -> None:
cdef uint64_t numvectors = vectors.shape[0]
_calc_minimize_vectors_ortho(&vectors[0, 0], numvectors, &box[0], &output[0, 0])


@cython.boundscheck(False)
@cython.wraparound(False)
cdef inline void _minimum_image_triclinic(cython.floating[:] dx,
cython.floating[:] box,
cython.floating[:] inverse_box) nogil:
"""Minimise dx to be the shortest vector

Parameters
----------
dx : numpy.array, shape (3,)
the vector to apply minimum image convention to
box : numpy.array, shape (9,)
flattened 3x3 representation of the unit cell
inverse_box : numpy.array, shape (3,)
inverse of the **diagonal** of the 3x3 representation

This version is near identical to the version in calc_distances.h, with the
difference being that this version does not assume that coordinates are at
most a single box length apart.

Operates in-place on dx!
"""
cdef cython.floating dx_min[3]
cdef cython.floating s, dsq, dsq_min, rx
cdef cython.floating ry[2]
cdef cython.floating rz[3]
cdef int ix, iy, iz

# first make shift only 1 cell in any direction
s = cround(inverse_box[2] * dx[2])
dx[0] -= s * box[6]
dx[1] -= s * box[7]
dx[2] -= s * box[8]
s = cround(inverse_box[1] * dx[1])
dx[0] -= s * box[3]
dx[1] -= s * box[4]
s = cround(inverse_box[0] * dx[0])
dx[0] -= s * box[0]

if cython.floating is float:
dsq_min = FLT_MAX
else:
dsq_min = DBL_MAX
dx_min[0] = 0.0
dx_min[1] = 0.0
dx_min[2] = 0.0

# then check all images to see which combination of 1 cell shifts gives the best shift
for ix in range(-1, 2):
rx = dx[0] + box[0] * ix
for iy in range(-1, 2):
ry[0] = rx + box[3] * iy
ry[1] = dx[1] + box[4] * iy
for iz in range(-1, 2):
rz[0] = ry[0] + box[6] * iz
rz[1] = ry[1] + box[7] * iz
rz[2] = dx[2] + box[8] * iz
dsq = rz[0] * rz[0] + rz[1] * rz[1] + rz[2] * rz[2]
if (dsq < dsq_min):
dsq_min = dsq
dx_min[0] = rz[0]
dx_min[1] = rz[1]
dx_min[2] = rz[2]

dx[0] = dx_min[0]
dx[1] = dx_min[1]
dx[2] = dx_min[2]


@cython.boundscheck(False)
@cython.wraparound(False)
@cython.cdivision(True)
def _minimize_vectors_ortho(cython.floating[:, :] vectors not None, cython.floating[:] box not None,
cython.floating[:, :] output not None):
cdef int i, n
cdef cython.floating box_inverse[3]
cdef cython.floating[:] box_inverse_view

box_inverse[0] = 1.0 / box[0]
box_inverse[1] = 1.0 / box[1]
box_inverse[2] = 1.0 / box[2]

box_inverse_view = box_inverse

n = len(vectors)
with nogil:
for i in range(n):
output[i, 0] = vectors[i, 0]
output[i, 1] = vectors[i, 1]
output[i, 2] = vectors[i, 2]
_minimum_image_orthogonal(output[i, :], box, box_inverse_view)


@cython.boundscheck(False)
@cython.wraparound(False)
def _minimize_vectors_triclinic(cython.floating[:, :] vectors not None, cython.floating[:] box not None,
cython.floating[:, :] output not None):
cdef int i, n
cdef cython.floating box_inverse[3]
cdef cython.floating[:] box_inverse_view

# this is only inverse of diagonal, used for initial shift to ensure vector is within a single shift away
box_inverse[0] = 1.0 / box[0]
box_inverse[1] = 1.0 / box[4]
box_inverse[2] = 1.0 / box[8]

box_inverse_view = box_inverse

n = len(vectors)
with nogil:
for i in range(n):
output[i, 0] = vectors[i, 0]
output[i, 1] = vectors[i, 1]
output[i, 2] = vectors[i, 2]
_minimum_image_triclinic(output[i, :], box, box_inverse_view)
def calc_minimize_vectors_triclinic(
floating[:, ::1] vectors,
floating[:, ::1] box,
floating[:, ::1] output) -> None:
cdef uint64_t numvectors = vectors.shape[0]
_calc_minimize_vectors_triclinic(&vectors[0, 0], numvectors, &box[0][0], &output[0, 0])
27 changes: 27 additions & 0 deletions package/MDAnalysis/lib/c_distances_openmp.pyx
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -21,6 +21,7 @@
# J. Comput. Chem. 32 (2011), 2319--2327, doi:10.1002/jcc.21787
#
#
# distutils: language = c++
Comment thread
PardhavMaradani marked this conversation as resolved.
Outdated

"""
Parallel distance calculation library --- :mod:`MDAnalysis.lib.c_distances_openmp`
Expand All @@ -30,9 +31,11 @@ Parallel distance calculation library --- :mod:`MDAnalysis.lib.c_distances_openm
Contains OpenMP versions of the contents of "calc_distances.h"
"""

cimport cython
from libc.stdint cimport uint64_t
import numpy
cimport numpy
from cython cimport floating
numpy.import_array()

cdef extern from "string.h":
Expand Down Expand Up @@ -60,6 +63,10 @@ cdef extern from "calc_distances.h":
void _ortho_pbc(coordinate* coords, uint64_t numcoords, float* box)
void _triclinic_pbc(coordinate* coords, uint64_t numcoords, float* box)

cdef extern from "calc_distances.hpp":
void _calc_minimize_vectors_ortho[T](T* vectors, uint64_t numvectors, T* box, T* output)
void _calc_minimize_vectors_triclinic[T](T* vectors, uint64_t numvectors, T* box, T* output)


OPENMP_ENABLED = True if USED_OPENMP else False

Expand Down Expand Up @@ -241,3 +248,23 @@ def triclinic_pbc(numpy.ndarray coords, numpy.ndarray box):
numcoords = coords.shape[0]

_triclinic_pbc(<coordinate*> coords.data, numcoords, <float*> box.data)


@cython.boundscheck(False)
@cython.wraparound(False)
def calc_minimize_vectors_ortho(
floating[:, ::1] vectors,
floating[::1] box,
floating[:, ::1] output):
cdef uint64_t numvectors = vectors.shape[0]
_calc_minimize_vectors_ortho(&vectors[0, 0], numvectors, &box[0], &output[0, 0])


@cython.boundscheck(False)
@cython.wraparound(False)
def calc_minimize_vectors_triclinic(
floating[:, ::1] vectors,
floating[:, ::1] box,
floating[:, ::1] output):
cdef uint64_t numvectors = vectors.shape[0]
_calc_minimize_vectors_triclinic(&vectors[0, 0], numvectors, &box[0][0], &output[0, 0])
27 changes: 23 additions & 4 deletions package/MDAnalysis/lib/distances.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -151,7 +151,10 @@
from .mdamath import triclinic_vectors
from ._augment import augment_coordinates, undo_augment
from .nsgrid import FastNS
from .c_distances import _minimize_vectors_ortho, _minimize_vectors_triclinic
from .c_distances import (
calc_minimize_vectors_ortho,
calc_minimize_vectors_triclinic,
)
from ._distopia import HAS_DISTOPIA


Expand Down Expand Up @@ -2103,7 +2106,11 @@ def apply_PBC(


@check_coords("vectors", enforce_copy=False, enforce_dtype=False)
def minimize_vectors(vectors: npt.NDArray, box: npt.NDArray) -> npt.NDArray:
def minimize_vectors(
vectors: npt.NDArray,
box: npt.NDArray,
backend: str = "serial",
) -> npt.NDArray:
"""Apply minimum image convention to an array of vectors

This function is required for calculating the correct vectors between two
Expand All @@ -2121,6 +2128,8 @@ def minimize_vectors(vectors: npt.NDArray, box: npt.NDArray) -> npt.NDArray:
triclinic and must be provided in the same format as returned by
:attr:`MDAnalysis.coordinates.timestep.Timestep.dimensions`:
``[lx, ly, lz, alpha, beta, gamma]``.
backend : {'serial', 'OpenMP'}, optional
Copy link
Copy Markdown
Member

@tylerjereddy tylerjereddy Feb 20, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Maybe I missed it, but how is the user to control the number of threads? OMP_NUM_THREADS? joblib context?

Copy link
Copy Markdown
Member

@tylerjereddy tylerjereddy Feb 20, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

how many threads were used in your benchmarks?

Copy link
Copy Markdown
Contributor Author

@PardhavMaradani PardhavMaradani Feb 20, 2026 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I will have to check this. I followed what was being done for other methods, but didn't quite look into the details yet - I will now. A very crude way I checked that that it was using multiple cores was by checking top and unlike the serial case (one core), I saw it using 4 8 cores in the openmp case - though I'll now have to check why only 4, because the machine has 8 cores. I learnt about the asv benchmarks for a different PR, but have been struggling to get it running so far, but will add those benchmarks after I can verify. Thanks

Copy link
Copy Markdown
Contributor Author

@PardhavMaradani PardhavMaradani Feb 20, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

but how is the user to control the number of threads? OMP_NUM_THREADS? joblib context?

I am not quite sure how to answer this. The only references I found in MDA documentation to control the number of threads is for the transfomation classes citing the use of threadpoolctl internally and also suggesting the use of the OMP_NUM_THREADS env variable to control the thread limit. I was however able to use both of them (the OMP_NUM_THREADS env var (at runtime) and threadpoolctl.threadpool_limits (in code) and verify that this controls the number of threads for these functions that support openmp backend.

Here is a sample top output for OMP_NUM_THREADS=4 when running the above compact tests:

Details 
top - 11:29:01 up 1 day, 16:19,  1 user,  load average: 2.70, 2.03, 1.49
Tasks: 356 total,   2 running, 353 sleeping,   0 stopped,   1 zombie
%Cpu0  : 57.6 us,  0.5 sy,  0.0 ni, 41.8 id,  0.0 wa,  0.0 hi,  0.0 si,  0.0 st
%Cpu1  :  7.1 us,  0.0 sy,  0.0 ni, 92.9 id,  0.0 wa,  0.0 hi,  0.0 si,  0.0 st
%Cpu2  : 60.3 us,  0.0 sy,  0.0 ni, 39.7 id,  0.0 wa,  0.0 hi,  0.0 si,  0.0 st
%Cpu3  :  1.6 us,  0.5 sy,  0.0 ni, 97.8 id,  0.0 wa,  0.0 hi,  0.0 si,  0.0 st
%Cpu4  :  1.6 us,  0.5 sy,  0.0 ni, 95.7 id,  2.2 wa,  0.0 hi,  0.0 si,  0.0 st
%Cpu5  : 62.0 us,  0.0 sy,  0.0 ni, 38.0 id,  0.0 wa,  0.0 hi,  0.0 si,  0.0 st
%Cpu6  :  0.5 us,  0.5 sy,  0.0 ni, 98.9 id,  0.0 wa,  0.0 hi,  0.0 si,  0.0 st
%Cpu7  : 87.0 us, 13.0 sy,  0.0 ni,  0.0 id,  0.0 wa,  0.0 hi,  0.0 si,  0.0 st
MiB Mem :   7770.6 total,   1025.2 free,   3988.1 used,   2757.3 buff/cache
MiB Swap:   2048.0 total,      0.2 free,   2047.8 used.   2909.3 avail Mem 

    PID USER      PR  NI    VIRT    RES    SHR S  %CPU  %MEM     TIME+ COMMAND                                           
  83009 pardhav   20   0 1136224 287580  78976 R 277.3   3.6   0:28.82 python                                            
.....

Here is a similar one for OMP_NUM_THREADS=8:

Details 
top - 11:30:53 up 1 day, 16:20,  1 user,  load average: 2.92, 2.10, 1.56
Tasks: 356 total,   2 running, 353 sleeping,   0 stopped,   1 zombie
%Cpu0  : 53.8 us,  1.9 sy,  0.0 ni, 43.3 id,  0.0 wa,  0.0 hi,  1.0 si,  0.0 st
%Cpu1  : 89.7 us, 10.3 sy,  0.0 ni,  0.0 id,  0.0 wa,  0.0 hi,  0.0 si,  0.0 st
%Cpu2  : 45.3 us,  0.5 sy,  0.0 ni, 54.2 id,  0.0 wa,  0.0 hi,  0.0 si,  0.0 st
%Cpu3  : 50.7 us,  0.5 sy,  0.0 ni, 48.8 id,  0.0 wa,  0.0 hi,  0.0 si,  0.0 st
%Cpu4  : 47.6 us,  3.8 sy,  0.0 ni, 48.6 id,  0.0 wa,  0.0 hi,  0.0 si,  0.0 st
%Cpu5  : 48.6 us,  0.5 sy,  0.0 ni, 50.9 id,  0.0 wa,  0.0 hi,  0.0 si,  0.0 st
%Cpu6  : 48.4 us,  0.0 sy,  0.0 ni, 51.6 id,  0.0 wa,  0.0 hi,  0.0 si,  0.0 st
%Cpu7  : 51.7 us,  0.0 sy,  0.0 ni, 48.3 id,  0.0 wa,  0.0 hi,  0.0 si,  0.0 st
MiB Mem :   7770.6 total,    967.5 free,   4045.2 used,   2758.0 buff/cache
MiB Swap:   2048.0 total,      0.2 free,   2047.8 used.   2852.2 avail Mem 

    PID USER      PR  NI    VIRT    RES    SHR S  %CPU  %MEM     TIME+ COMMAND                                           
  83031 pardhav   20   0 1823332 383052  78464 R 436.0   4.8   1:47.77 python                                            
.....

I believe this should be the same with all other distance methods that support accelerated backends. Thanks

Keyword selecting the type of acceleration.

Returns
-------
Expand All @@ -2130,6 +2139,8 @@ def minimize_vectors(vectors: npt.NDArray, box: npt.NDArray) -> npt.NDArray:


.. versionadded:: 2.1.0
.. versionchanged:: 2.11.0
Added *backend* keyword.
"""
boxtype, box = check_box(box)
output = np.empty_like(vectors)
Expand All @@ -2138,8 +2149,16 @@ def minimize_vectors(vectors: npt.NDArray, box: npt.NDArray) -> npt.NDArray:
box = box.astype(vectors.dtype)

if boxtype == "ortho":
_minimize_vectors_ortho(vectors, box, output)
_run(
"calc_minimize_vectors_ortho",
args=(vectors, box, output),
backend=backend,
)
else:
_minimize_vectors_triclinic(vectors, box.ravel(), output)
_run(
"calc_minimize_vectors_triclinic",
args=(vectors, box, output),
backend=backend,
)

return output
Loading
Loading