diff --git a/psico/editing.py b/psico/editing.py index 4c951bc..94e919b 100644 --- a/psico/editing.py +++ b/psico/editing.py @@ -254,7 +254,7 @@ def stub2ala(selection='all', quiet=1, *, _self=cmd): _self.sort() -def remove_alt(selection='all', keep='A', quiet=1, *, _self=cmd): +def remove_alt(selection='all', keep='first', quiet=1, *, _self=cmd): ''' DESCRIPTION @@ -268,13 +268,39 @@ def remove_alt(selection='all', keep='A', quiet=1, *, _self=cmd): selection = string: atom selection - keep = string: AltLoc to keep {default: A} + keep = string: AltLoc to keep, or 'first' to keep the first observed AltLoc {default: first} ''' + if keep == "first": + return remove_alt_keep_first(selection, quiet=quiet, _self=_self) + + if len(keep) != 1: + raise CmdException( + f"keep must be 'first' or a single letter, got {keep!r}") + _self.remove('(%s) and not alt +%s' % (selection, keep), quiet=int(quiet)) _self.alter(selection, '(alt,q)=("",1.0)') _self.sort() +def remove_alt_keep_first(selection='*', *, quiet=1, _self=cmd): + ''' + Remove alternative location atoms, keep the first observed. + ''' + alts = {} + expr = "(alt, q) = callback((model, segi, chain, resi, resn, name), alt)" + + def callback(namekey, alt): + return ("#", 0.) if alts.setdefault(namekey, alt) != alt else ("", 1.) + + tmpsele = _self.get_unused_name("_sele") + _self.select(tmpsele, selection) + try: + _self.alter(tmpsele, expr, space={"callback": callback}) + _self.remove(f"{tmpsele} & not alt ''", quiet=quiet) + finally: + _self.delete(tmpsele) + + def _common_ss_alter(selection, ss_dict, ss_map, raw='', *, _self=cmd): ''' DESCRIPTION diff --git a/tests/test_editing.py b/tests/test_editing.py index c04ee10..7979003 100644 --- a/tests/test_editing.py +++ b/tests/test_editing.py @@ -2,7 +2,7 @@ import psico.querying import os import pytest -from pymol import cmd +from pymol import cmd, CmdException from pytest import approx DATA_DIR = os.path.join(os.path.dirname(__file__), 'data') @@ -147,10 +147,20 @@ def test_remove_alt(): cmd.alter('resn ALA', 'alt="A"') cmd.alter('resn CYS', 'alt="B"') cmd.create('m2', 'm1') - psico.editing.remove_alt('m1') + cmd.create('m3', 'm1') + cmd.alter('m3', 'alt = {"A": "X", "B": "Y"}[alt]') + cmd.create('m4', 'm1 m3') + assert cmd.count_atoms('m4') == 42 + psico.editing.remove_alt('m1', keep='A') psico.editing.remove_alt('m2', keep='B') assert cmd.count_atoms('m1') == 10 assert cmd.count_atoms('m2') == 11 + psico.editing.remove_alt('m3', keep='first') + assert cmd.count_atoms('m3') == 21 + psico.editing.remove_alt('m4', keep='first') + assert cmd.count_atoms('m4') == 21 + with pytest.raises(CmdException, match="single letter"): + psico.editing.remove_alt('m1', keep='last') @pytest.mark.exe