From 350c5a50216ac3c8387086b13b091e3fdabf3b4f Mon Sep 17 00:00:00 2001 From: Jia-Xin Zhu Date: Wed, 18 Dec 2024 13:42:31 +0800 Subject: [PATCH 1/5] add `relprop` for atom selection and corresponding UT --- package/MDAnalysis/core/selection.py | 45 +++++++++++++++++++ .../core/test_atomselections.py | 11 +++++ 2 files changed, 56 insertions(+) diff --git a/package/MDAnalysis/core/selection.py b/package/MDAnalysis/core/selection.py index 591c074030..950fe01739 100644 --- a/package/MDAnalysis/core/selection.py +++ b/package/MDAnalysis/core/selection.py @@ -1353,6 +1353,51 @@ def _apply(self, group): return group[mask] +class RelPropertySelection(PropertySelection): + """Some of the possible properties: + x, y, z, + + .. versionadded:: 2.9.0 + """ + + token = "relprop" + precedence = 1 + + def __init__(self, parser, tokens): + super().__init__(parser, tokens) + self.sel = parser.parse_expression(self.precedence) + # self.ori_value = self.value + + def _apply(self, group): + try: + values = getattr(group, self.props[self.prop]) + except KeyError: + errmsg = f"Expected one of {list(self.props.keys())}" + raise SelectionError(errmsg) from None + except NoDataError: + attr = self.props[self.prop] + errmsg = f"This Universe does not contain {attr} information" + raise SelectionError(errmsg) from None + + try: + col = {"x": 0, "y": 1, "z": 2}[self.prop] + except KeyError: + pass + else: + values = values[:, col] + sel = self.sel.apply(group) + rel_value = ( + sel.center_of_geometry().reshape(3).astype(np.float32)[col] + ) + values -= rel_value + + if self.absolute: + values = np.abs(values) + mask = self.operator(values, self.value) + + return group[mask] + + class SameSelection(Selection): """ Selects all atoms that have the same subkeyword value as any atom in selection diff --git a/testsuite/MDAnalysisTests/core/test_atomselections.py b/testsuite/MDAnalysisTests/core/test_atomselections.py index bced4c43bd..879297309e 100644 --- a/testsuite/MDAnalysisTests/core/test_atomselections.py +++ b/testsuite/MDAnalysisTests/core/test_atomselections.py @@ -253,6 +253,17 @@ def test_prop(self, universe): assert_equal(len(sel), 3194) assert_equal(len(sel2), 2001) + def test_relprop(self, universe): + sel1 = universe.select_atoms("relprop z <= 1 index 0") + sel2 = universe.select_atoms("relprop abs z <= 1 index 0") + + positions = universe.trajectory[0].positions + ref = positions[0, 2] + mask_1 = (positions[:, 2] - ref) <= 1 + assert_equal(len(sel1), np.count_nonzero(mask_1)) + mask_2 = np.abs(positions[:, 2] - ref) <= 1 + assert_equal(len(sel2), np.count_nonzero(mask_2)) + def test_bynum(self, universe): "Tests the bynum selection, also from AtomGroup instances (Issue 275)" sel = universe.select_atoms('bynum 5') From 08f8b98fe5941ce4fc26e8279bce6f51558e27a3 Mon Sep 17 00:00:00 2001 From: Jia-Xin Zhu Date: Wed, 18 Dec 2024 15:55:28 +0800 Subject: [PATCH 2/5] add UT for exception --- testsuite/MDAnalysisTests/core/test_atomselections.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/testsuite/MDAnalysisTests/core/test_atomselections.py b/testsuite/MDAnalysisTests/core/test_atomselections.py index 879297309e..3aea7befe7 100644 --- a/testsuite/MDAnalysisTests/core/test_atomselections.py +++ b/testsuite/MDAnalysisTests/core/test_atomselections.py @@ -1087,6 +1087,12 @@ def test_invalid_prop_selection(self, universe): with pytest.raises(SelectionError, match="Expected one of"): universe.select_atoms("prop parsnip < 2") + def test_invalid_relprop_selection(self, universe): + with pytest.raises(SelectionError, match="Expected one of"): + universe.select_atoms("relprop parsnip < 2 index 0") + with pytest.raises(SelectionError, match="Unknown selection token"): + universe.select_atoms("relprop z < 2") + def test_segid_and_resid(): u = make_Universe(('segids', 'resids')) From a5b7bf9be1dada29cba2feea051c2dcceecf6499 Mon Sep 17 00:00:00 2001 From: Jia-Xin Zhu Date: Sun, 29 Dec 2024 22:06:19 +0800 Subject: [PATCH 3/5] resolve comments --- package/MDAnalysis/core/groups.py | 10 ++++++++++ package/MDAnalysis/core/selection.py | 18 +++++++++++------- .../core/test_atomselections.py | 9 +++++++++ 3 files changed, 30 insertions(+), 7 deletions(-) diff --git a/package/MDAnalysis/core/groups.py b/package/MDAnalysis/core/groups.py index e8bf30ba11..8fd74be121 100644 --- a/package/MDAnalysis/core/groups.py +++ b/package/MDAnalysis/core/groups.py @@ -3254,6 +3254,16 @@ def select_atoms(self, sel, *othersel, periodic=True, rtol=1e-05, For example, ``prop z >= 5.0`` selects all atoms with z coordinate greater than 5.0; ``prop abs z <= 5.0`` selects all atoms within -5.0 <= z <= 5.0. + relprop [abs] *property* *operator* *value* *selection* + selects atoms based on position relative to the center of + geometry (COG) of a given selection, using *property* + **x**, **y**, or **z** coordinate. Supports the **abs** + keyword (for absolute value) and the following + *operators*: **<, >, <=, >=, ==, !=**. + For example, ``relprop z >= 5.0 protein`` selects all atoms + with z coordinate greater than 5.0 relative to the COG + of protein; ``relprop abs z <= 5.0 protein`` selects all + atoms within -5.0 <= z <= 5.0 relative to the COG of protein. sphzone *radius* *selection* Selects all atoms that are within *radius* of the center of geometry of *selection* diff --git a/package/MDAnalysis/core/selection.py b/package/MDAnalysis/core/selection.py index 950fe01739..1ff3bab16e 100644 --- a/package/MDAnalysis/core/selection.py +++ b/package/MDAnalysis/core/selection.py @@ -1366,7 +1366,7 @@ class RelPropertySelection(PropertySelection): def __init__(self, parser, tokens): super().__init__(parser, tokens) self.sel = parser.parse_expression(self.precedence) - # self.ori_value = self.value + self.periodic = parser.periodic def _apply(self, group): try: @@ -1382,14 +1382,18 @@ def _apply(self, group): try: col = {"x": 0, "y": 1, "z": 2}[self.prop] except KeyError: - pass + errmsg = f"Expected one of x y z for property, got {self.prop}" + raise SelectionError(errmsg) from None else: - values = values[:, col] sel = self.sel.apply(group) - rel_value = ( - sel.center_of_geometry().reshape(3).astype(np.float32)[col] - ) - values -= rel_value + ref_value = sel.center_of_geometry().reshape(3).astype(np.float32) + box = group.dimensions if self.periodic else None + if box is not None: + values = distances.minimize_vectors( + values - ref_value[None, :], box=box + )[:, col] + else: + values = values[:, col] - ref_value[col] if self.absolute: values = np.abs(values) diff --git a/testsuite/MDAnalysisTests/core/test_atomselections.py b/testsuite/MDAnalysisTests/core/test_atomselections.py index 3aea7befe7..be1ad48f0e 100644 --- a/testsuite/MDAnalysisTests/core/test_atomselections.py +++ b/testsuite/MDAnalysisTests/core/test_atomselections.py @@ -1092,6 +1092,15 @@ def test_invalid_relprop_selection(self, universe): universe.select_atoms("relprop parsnip < 2 index 0") with pytest.raises(SelectionError, match="Unknown selection token"): universe.select_atoms("relprop z < 2") + with pytest.raises(SelectionError, match="Expected one of x y z for property"): + universe.select_atoms("relprop id < 2 index 0") + empty_universe = mda.Universe.empty( + 6, 2, atom_resindex=[0, 0, 0, 1, 1, 1] + ) + with pytest.raises( + SelectionError, match="This Universe does not contain" + ): + empty_universe.select_atoms("relprop z <= 1 index 0") def test_segid_and_resid(): From cda4dcac3500641becaeaba5600720a59e052061 Mon Sep 17 00:00:00 2001 From: Jia-Xin Zhu Date: Sun, 29 Dec 2024 22:09:31 +0800 Subject: [PATCH 4/5] formatting --- package/MDAnalysis/core/groups.py | 14 +++++++------- .../MDAnalysisTests/core/test_atomselections.py | 2 +- 2 files changed, 8 insertions(+), 8 deletions(-) diff --git a/package/MDAnalysis/core/groups.py b/package/MDAnalysis/core/groups.py index 8fd74be121..a74d955f5d 100644 --- a/package/MDAnalysis/core/groups.py +++ b/package/MDAnalysis/core/groups.py @@ -3255,14 +3255,14 @@ def select_atoms(self, sel, *othersel, periodic=True, rtol=1e-05, coordinate greater than 5.0; ``prop abs z <= 5.0`` selects all atoms within -5.0 <= z <= 5.0. relprop [abs] *property* *operator* *value* *selection* - selects atoms based on position relative to the center of - geometry (COG) of a given selection, using *property* - **x**, **y**, or **z** coordinate. Supports the **abs** - keyword (for absolute value) and the following + selects atoms based on position relative to the center of + geometry (COG) of a given selection, using *property* + **x**, **y**, or **z** coordinate. Supports the **abs** + keyword (for absolute value) and the following *operators*: **<, >, <=, >=, ==, !=**. - For example, ``relprop z >= 5.0 protein`` selects all atoms - with z coordinate greater than 5.0 relative to the COG - of protein; ``relprop abs z <= 5.0 protein`` selects all + For example, ``relprop z >= 5.0 protein`` selects all atoms + with z coordinate greater than 5.0 relative to the COG + of protein; ``relprop abs z <= 5.0 protein`` selects all atoms within -5.0 <= z <= 5.0 relative to the COG of protein. sphzone *radius* *selection* Selects all atoms that are within *radius* of the center of diff --git a/testsuite/MDAnalysisTests/core/test_atomselections.py b/testsuite/MDAnalysisTests/core/test_atomselections.py index be1ad48f0e..45f15823f2 100644 --- a/testsuite/MDAnalysisTests/core/test_atomselections.py +++ b/testsuite/MDAnalysisTests/core/test_atomselections.py @@ -1092,7 +1092,7 @@ def test_invalid_relprop_selection(self, universe): universe.select_atoms("relprop parsnip < 2 index 0") with pytest.raises(SelectionError, match="Unknown selection token"): universe.select_atoms("relprop z < 2") - with pytest.raises(SelectionError, match="Expected one of x y z for property"): + with pytest.raises(SelectionError, match="Expected one of"): universe.select_atoms("relprop id < 2 index 0") empty_universe = mda.Universe.empty( 6, 2, atom_resindex=[0, 0, 0, 1, 1, 1] From 2b4ab7ad1a9910823597aa5765be00cb21b03129 Mon Sep 17 00:00:00 2001 From: Jia-Xin Zhu Date: Sun, 29 Dec 2024 23:30:04 +0800 Subject: [PATCH 5/5] fix bug in UT --- testsuite/MDAnalysisTests/core/test_atomselections.py | 7 ++----- 1 file changed, 2 insertions(+), 5 deletions(-) diff --git a/testsuite/MDAnalysisTests/core/test_atomselections.py b/testsuite/MDAnalysisTests/core/test_atomselections.py index 45f15823f2..f8d5f98b97 100644 --- a/testsuite/MDAnalysisTests/core/test_atomselections.py +++ b/testsuite/MDAnalysisTests/core/test_atomselections.py @@ -1093,14 +1093,11 @@ def test_invalid_relprop_selection(self, universe): with pytest.raises(SelectionError, match="Unknown selection token"): universe.select_atoms("relprop z < 2") with pytest.raises(SelectionError, match="Expected one of"): - universe.select_atoms("relprop id < 2 index 0") - empty_universe = mda.Universe.empty( - 6, 2, atom_resindex=[0, 0, 0, 1, 1, 1] - ) + universe.select_atoms("relprop resid < 2 index 0") with pytest.raises( SelectionError, match="This Universe does not contain" ): - empty_universe.select_atoms("relprop z <= 1 index 0") + universe.select_atoms("relprop z <= 1 index 0") def test_segid_and_resid():