diff --git a/package/AUTHORS b/package/AUTHORS index ad93e39ca9c..b8b42b9299f 100644 --- a/package/AUTHORS +++ b/package/AUTHORS @@ -163,7 +163,8 @@ Chronological list of authors - Estefania Barreto-Ojeda - Paarth Thadani - Henry Kobin - + - Kosuke Kudo + External code ------------- diff --git a/package/CHANGELOG b/package/CHANGELOG index a8948666c37..ea3db368ad8 100644 --- a/package/CHANGELOG +++ b/package/CHANGELOG @@ -17,7 +17,8 @@ The rules for this file: lilyminium, daveminh, jbarnoud, yuxuanzhuang, VOD555, ianmkenney, calcraven，xiki-tempula, mieczyslaw, manuel.nuno.melo, PicoCentauri, hanatok, rmeli, aditya-kamath, tirkarthi, LeonardoBarneschi, hejamu, - biogen98, orioncohen, z3y50n, hp115, ojeda-e, thadanipaarth, HenryKobin + biogen98, orioncohen, z3y50n, hp115, ojeda-e, thadanipaarth, HenryKobin, + 1ut * 2.0.0 @@ -98,6 +99,7 @@ Fixes * Fix syntax warning over comparison of literals using is (Issue #3066) Enhancements + * Added sort method to the atomgroup (Issue #2976, PR #3188) * ITPParser now reads [ atomtypes ] sections in ITP files, used for charges and masses not defined in the [ atoms ] sections * Add `set_dimensions` transformation class for setting constant diff --git a/package/MDAnalysis/core/groups.py b/package/MDAnalysis/core/groups.py index 0671bcef3ad..53e9e8529b6 100644 --- a/package/MDAnalysis/core/groups.py +++ b/package/MDAnalysis/core/groups.py @@ -3324,6 +3324,83 @@ def write(self, filename=None, file_format=None, raise ValueError("No writer found for format: {}".format(filename)) + def sort(self, key='ix', keyfunc=None): + """ + Returns a sorted ``AtomGroup`` using a specified attribute as the key. + + Parameters + ---------- + key: str, optional + The name of the ``AtomGroup`` attribute to sort by (e.g. ``ids``, + ``ix``. default= ``ix`` ). + keyfunc: callable, optional + A function to convert multidimensional arrays to a single + dimension. This 1D array will be used as the sort key and + is required when sorting with an ``AtomGroup`` attribute + key which has multiple dimensions. Note: this argument + is ignored when the attribute is one dimensional. + + Returns + ---------- + :class:`AtomGroup` + Sorted ``AtomGroup``. + + Example + ---------- + + .. code-block:: python + + >>> import MDAnalysis as mda + >>> from MDAnalysisTests.datafiles import PDB_small + >>> u = mda.Universe(PDB_small) + >>> ag = u.atoms[[3, 2, 1, 0]] + >>> ag.ix + array([3 2 1 0]) + >>> ag = ag.sort() + >>> ag.ix + array([0 1 2 3]) + >>> ag.positions + array([[-11.921, 26.307, 10.41 ], + [-11.447, 26.741, 9.595], + [-12.44 , 27.042, 10.926], + [-12.632, 25.619, 10.046]], dtype=float32) + >>> ag = ag.sort("positions", lambda x: x[:, 1]) + >>> ag.positions + array([[-12.632, 25.619, 10.046], + [-11.921, 26.307, 10.41 ], + [-11.447, 26.741, 9.595], + [-12.44 , 27.042, 10.926]], dtype=float32) + + Note + ---------- + This uses a stable sort as implemented by + `numpy.argsort(kind='stable')`. + + + .. versionadded:: 2.0.0 + """ + idx = getattr(self.atoms, key) + if len(idx) != len(self.atoms): + raise ValueError("The array returned by the attribute '{}' " + "must have the same length as the number of " + "atoms in the input AtomGroup".format(key)) + if idx.ndim == 1: + order = np.argsort(idx, kind='stable') + elif idx.ndim > 1: + if keyfunc is None: + raise NameError("The {} attribute returns a multidimensional " + "array. In order to sort it, a function " + "returning a 1D array (to be used as the sort " + "key) must be passed to the keyfunc argument" + .format(key)) + sortkeys = keyfunc(idx) + if sortkeys.ndim != 1: + raise ValueError("The function assigned to the argument " + "'keyfunc':{} doesn't return a 1D array." + .format(keyfunc)) + order = np.argsort(sortkeys, kind='stable') + return self.atoms[order] + class ResidueGroup(GroupBase): """ResidueGroup base class. diff --git a/testsuite/MDAnalysisTests/core/test_atomgroup.py b/testsuite/MDAnalysisTests/core/test_atomgroup.py index 1a3d401f0ee..a719ccdfbb5 100644 --- a/testsuite/MDAnalysisTests/core/test_atomgroup.py +++ b/testsuite/MDAnalysisTests/core/test_atomgroup.py @@ -1674,3 +1674,80 @@ def test_partial_timestep(self, universe): ag.ts.velocities, self.prec, err_msg="Partial timestep coordinates wrong") + + +class TestAtomGroupSort(object): + """Tests the AtomGroup.sort attribute""" + + @pytest.fixture() + def universe(self): + u = mda.Universe.empty( + n_atoms=7, + n_residues=3, + n_segments=2, + atom_resindex=np.array([0, 0, 0, 1, 1, 1, 2]), + residue_segindex=np.array([0, 0, 1]), + trajectory=True, + velocities=True, + forces=True + ) + attributes = ["id", "charge", "mass", "tempfactor"] + + for i in (attributes): + u.add_TopologyAttr(i, [6, 5, 4, 3, 2, 1, 0]) + + u.add_TopologyAttr('resid', [2, 1, 0]) + u.add_TopologyAttr('segid', [1, 0]) + u.add_TopologyAttr('bonds', [(0, 1)]) + + return u + + @pytest.fixture() + def ag(self, universe): + ag = universe.atoms + ag.positions = (-np.arange(21)).reshape(7, 3) + return ag + + test_ids = [ + "ix", + "ids", + "resids", + "segids", + "charges", + "masses", + "tempfactors" + ] + + test_data = [ + ("ix", np.array([0, 1, 2, 3, 4, 5, 6])), + ("ids", np.array([6, 5, 4, 3, 2, 1, 0])), + ("resids", np.array([6, 3, 4, 5, 0, 1, 2])), + ("segids", np.array([6, 0, 1, 2, 3, 4, 5])), + ("charges", np.array([6, 5, 4, 3, 2, 1, 0])), + ("masses", np.array([6, 5, 4, 3, 2, 1, 0])), + ("tempfactors", np.array([6, 5, 4, 3, 2, 1, 0])), + ] + + @pytest.mark.parametrize("inputs, expected", test_data, ids=test_ids) + def test_sort(self, ag, inputs, expected): + agsort = ag.sort(inputs) + assert np.array_equal(expected, agsort.ix) + + def test_sort_bonds(self, ag): + with pytest.raises(ValueError, match=r"The array returned by the " + "attribute"): + ag.sort("bonds") + + def test_sort_positions_2D(self, ag): + with pytest.raises(ValueError, match=r"The function assigned to"): + ag.sort("positions", keyfunc=lambda x: x) + + def test_sort_position_no_keyfunc(self, ag): + with pytest.raises(NameError, match=r"The .* attribute returns a " + "multidimensional array. In order to sort it, "): + ag.sort("positions") + + def test_sort_position(self, ag): + ref = [6, 5, 4, 3, 2, 1, 0] + agsort = ag.sort("positions", keyfunc=lambda x: x[:, 1]) + assert np.array_equal(ref, agsort.ix)