CollapseSel

Overview

CollapseSel is a small utility function to compress strings for selections. So, if you have a selection with residues 1+2+3+4+20+21+22+100-120 this will return, 1-4+20-22+100-120.

CollapseIDs is a small utility function to compress strings for an array of IDs. This does NOT have the logic for detecting duplicates, CollapseSel does.

Example

run /dir/to/collapseSel.py
fetch 1cll
select EE, resn GLU
print collapseSel("EE")
#
# and PyMOL should output:
#
# 6-7+11+14+31+45+47+54+67+82-84+87+104+114+119-120+123+127+139-140
#

The Code

import pymol
from pymol import stored

def collapseIDs(ids):
	"""
	Helper function to make a smart list of IDs: eg turns 1+2+3+4+5 into 1-5.
	"""
	rVal = []
	if len(ids)==0:
		return ""

	scanning=False
	anchor = 0
	start = 0
	# 1-5 7-10 12 21-23
	for cur in range(0,len(ids)-1):
		if ids[cur]+1 == ids[cur+1]:
			if scanning:
				scanning=True
				continue
			else:
				scanning=True
				start = cur
		else:
			if scanning:
				scanning=False
				rVal.append(str(ids[start]) + "-" + str(ids[cur]))
				start = cur
			else:
				scanning=False
				rVal.append(str(ids[cur]))
	if scanning:
		rVal.append( str(ids[start]) + "-" + str(ids[cur+1]))
	else:
		rVal.append(str(ids[-1]))
	return rVal

def collapseSel(sel=None,lType="resi"):
        """
        collapseSel -- given a valid PyMOL selection and list type, return a collapsed
                list of numbers corresponding to the lType.  For example, to compactly
                print the residue identifiers for all the waters, try:
                        select theWaters, resn HOH
                        print collapseSel( "theWaters" )

                This will convert: 1+2+3+4+10+11+12 into 1-4+10-12

        PARAMS
                sel
                        The selection to collapse
                lType
                        The identifier type: 'resi', 'ID', 'id', or any numerical property.

        RETURNS
                a string of collapsed IDs.
        """
        if sel==None:
                return ""

        stored.s = set()
        cmd.iterate( sel, "stored.s.add(int(float(%s)))" % lType)
        l = list(stored.s)
        l.sort()
        return "+".join(collapseIDs(list(l)))

cmd.extend("collapseSel", collapseSel)