Merge pull request #3441 from scal444/nsgrid_debug

Add check for too many grids in NSGrid

Author: richardjgowers

.github/workflows/gh-ci.yaml

@@ -94,6 +94,7 @@ jobs:
         gfortran-9 -v
         echo "FC=gfortran-9" >> $GITHUB_ENV
         echo "numprocs=3" >> $GITHUB_ENV
+        echo "ENABLE_HIGH_MEM_UNIT_TESTS=true" >> $GITHUB_ENV
 
     - name: setup_linux
       if: startsWith(matrix.os, 'ubuntu')

package/AUTHORS

@@ -170,6 +170,7 @@ Chronological list of authors
   - Filip T. Szczypiński
   - Marcelo C. R. Melo
   - Mark D. Driver
+  - Kevin Boyd
 
 
 External code

package/CHANGELOG

@@ -13,11 +13,11 @@ The rules for this file:
   * release numbers follow "Semantic Versioning" http://semver.org
 
 ------------------------------------------------------------------------------
-??/??/?? IAlibay, melomcr, mdd31, ianmkenney, richardjgowers, hmacdope, orbeckst
-
+??/??/?? IAlibay, melomcr, mdd31, ianmkenney, richardjgowers, hmacdope, orbeckst, scal444
  * 2.1.0
 
 Fixes
+  * Avoid integer overflow in NSGrid with too many grids (Issue #3183)
   * Fix ITPParser charge reading (Issue #3419).
   * TRRWriter preserves timestep data (Issue #3350).
   * Fixed Universe creation from a custom object which only provided a topology,

package/MDAnalysis/lib/nsgrid.pyx

@@ -71,7 +71,9 @@ not reflect in the results.
 .. versionadded:: 0.19.0
 .. versionchanged:: 1.0.2
    Rewrote module
-
+.. versionchanged:: 2.1.0
+   Capped max grid dimension to 1290, which when cubed is the max value of
+   a 32 bit signed integer.
 
 Classes
 -------
@@ -90,6 +92,10 @@ DEF XZ = 6
 DEF YZ = 7
 DEF ZZ = 8
 
+# Cube root of the maximum size of a 32 bit signed integer. If the system is divided into more
+# grids than this, integer overflow will occur.
+DEF MAX_GRID_DIM = 1290
+
 ctypedef float coordinate[3];
 
 cdef extern from "calc_distances.h" nogil:
@@ -342,9 +348,11 @@ cdef class FastNS(object):
         # add 0.001 here to avoid floating point errors
         # will make cells slightly too large as a result, ah well
         min_cellsize += 0.001
-        self.ncells[0] = <int> math.floor(self.triclinic_dimensions[XX] / min_cellsize)
-        self.ncells[1] = <int> math.floor(self.triclinic_dimensions[YY] / min_cellsize)
-        self.ncells[2] = <int> math.floor(self.triclinic_dimensions[ZZ] / min_cellsize)
+        # If the cell size is too small, indexing overflow will occur. Limit the number
+        # of cells in any dimension to the cube root of the maximum of 32 bit integer values.
+        self.ncells[0] = <int> min(math.floor(self.triclinic_dimensions[XX] / min_cellsize), MAX_GRID_DIM)
+        self.ncells[1] = <int> min(math.floor(self.triclinic_dimensions[YY] / min_cellsize), MAX_GRID_DIM)
+        self.ncells[2] = <int> min(math.floor(self.triclinic_dimensions[ZZ] / min_cellsize), MAX_GRID_DIM)
 
         self.pbc = pbc
         # If there aren't enough cells in a given dimension it's equivalent to one

testsuite/MDAnalysisTests/lib/test_nsgrid.py

@@ -21,6 +21,9 @@
 # J. Comput. Chem. 32 (2011), 2319--2327, doi:10.1002/jcc.21787
 #
 
+from distutils.util import strtobool
+import os
+
 import pytest
 
 from collections import defaultdict, Counter
@@ -65,9 +68,11 @@ def test_pbc_box(box):
         nsgrid.FastNS(4.0, coords, box=box)
 
 
-@pytest.mark.parametrize('cutoff', [-4, 100000])
-def test_nsgrid_badcutoff(universe, cutoff):
-    with pytest.raises(ValueError):
+@pytest.mark.parametrize('cutoff, match', ((-4, "Cutoff must be positive"),
+                                           (100000,
+                                            "Cutoff 100000 too large for box")))
+def test_nsgrid_badcutoff(universe, cutoff, match):
+    with pytest.raises(ValueError, match=match):
         run_grid_search(universe, 0, cutoff)
 
 
@@ -371,3 +376,36 @@ def test_issue_2670():
 
     # should return the one atom overlap
     assert len(ag2.select_atoms('around 0.0 resid 3')) == 1
+
+
+def high_mem_tests_enabled():
+    """ Returns true if ENABLE_HIGH_MEM_UNIT_TESTS is set to true."""
+    env = os.getenv("ENABLE_HIGH_MEM_UNIT_TESTS", default="0")
+    try:
+        return strtobool(env)
+    except ValueError:
+        return False
+
+
+reason = ("Turned off by default. The test can be enabled by setting "
+          "the ENABLE_HIGH_MEM_UNIT_TESTS "
+          "environment variable. Make sure you have at least 10GB of RAM.")
+
+
+# Tests that with a tiny cutoff to box ratio, the number of grids is capped
+# to avoid indexing overflow. Expected results copied from test_nsgrid_search
+# with no box.
+@pytest.mark.skipif(not high_mem_tests_enabled(), reason=reason)
+def test_issue_3183():
+    np.random.seed(90003)
+    points = (np.random.uniform(low=0, high=1.0,
+                                size=(100, 3)) * (10.)).astype(np.float32)
+    cutoff = 2.0
+    query = np.array([1., 1., 1.], dtype=np.float32).reshape((1, 3))
+    box = np.array([10000., 10000., 10000., 90., 90., 90.])
+
+    searcher = nsgrid.FastNS(cutoff, points, box)
+    searchresults = searcher.search(query)
+    indices = searchresults.get_pairs()[:, 1]
+    want_results = [3, 13, 24]
+    assert_equal(np.sort(indices), want_results)