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:
#
# 123+67+114+6-7+104+11+140+45+14+47+82-84+54+87+120+127+139+119+31
#

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(%s))" % lType)
	return "+".join(collapseIDs(list(stored.s)))

cmd.extend("collapseSel", collapseSel)
cmd.extend("collapsIDs", collapseIDs)