diff --git a/testsuite/MDAnalysisTests/core/test_atomselections.py b/testsuite/MDAnalysisTests/core/test_atomselections.py index 214b97e1ce..c00282fdcf 100644 --- a/testsuite/MDAnalysisTests/core/test_atomselections.py +++ b/testsuite/MDAnalysisTests/core/test_atomselections.py @@ -33,6 +33,7 @@ from MDAnalysis import SelectionError, SelectionWarning from MDAnalysis.core.selection import Parser from MDAnalysis.lib.distances import distance_array +from MDAnalysis.lib.pkdtree import PeriodicKDTree from MDAnalysis.tests.datafiles import ( DCD, GRO, @@ -57,6 +58,7 @@ ) from numpy.lib import NumpyVersion from numpy.testing import assert_equal +from numpy.testing import assert_array_equal from MDAnalysisTests import make_Universe @@ -1774,3 +1776,32 @@ def test_formal_charge_selection(sel, size, name): assert len(ag) == size assert ag.atoms[0].name == name + + +def test_cylayer_selection_matches_periodic_kdtree(): + u = mda.Universe(PSF, DCD) + u.dimensions = np.array([100.0, 100.0, 100.0, 90.0, 90.0, 90.0]) + + sel = u.select_atoms("cylayer 5 10 15 -15 name CA") + + ref = u.select_atoms("name CA") + center = ref.center_of_geometry() + + inner, outer, zmax, zmin = 5.0, 10.0, 15.0, -15.0 + search_radius = np.sqrt(outer**2 + zmax**2) + + kdtree = PeriodicKDTree(u.dimensions) + kdtree.set_coords(u.atoms.positions, cutoff=search_radius) + + idx = kdtree.search(center, search_radius) + + pos = u.atoms.positions[idx] + dxy = np.sqrt((pos[:, 0] - center[0]) ** 2 + (pos[:, 1] - center[1]) ** 2) + dz = pos[:, 2] - center[2] + + mask = (inner <= dxy) & (dxy <= outer) & (zmin <= dz) & (dz <= zmax) + + expected = np.sort(idx[mask]) + result = np.sort(sel.indices) + + assert_array_equal(result, expected)