build.rna_tools.tools.PyMOL4RNA.external_flatten_object.ChainSet.map_chain() - Code Metrics - mmagnus/rna-tools - Measure and Improve Code Quality continuously with Scrutinizer

ChainSet.map_chain() A
last analyzed 2025-03-19 22:26 UTC

↳ Parent: build.rna_tools.tools.PyMOL4RNA.external_flatten_object

Complexity

Conditions

Size

Total Lines	8
Code Lines	2

Duplication

Lines	0
Ratio	0 %

Importance

Changes

Metric	Value
cc	1
eloc	2
nop	4
dl	0
loc	8
rs	10
c	0
b	0
f	0

"""
flatten_obj.py - Flatten multi-state pymol objects into a single state.

<https://pymolwiki.org/index.php/Flatten_obj>

This is particularly useful for dealing with biological assemblies, which are
loaded as multi-state objects when fetched using `fetch PDBID, type=pdb1`. It
can also be used as a quick way to combine multiple objects without causing
collisions between chain identifiers.

The command re-letters chains to avoid collisions. Older versions of PyMOL
restrict the chain id to a single character, so the script will fail for
assemblies with >62 chains. With more recent versions, this problem is solved
with multi-character chain IDs. Several options are available for how
re-lettering should occur.

Author: Spencer Bliven <[email protected]>
Date: October 30, 2015
Version: 1.0
License: Public Domain
"""
from pymol import cmd, stored
import re
try:
    from collections import OrderedDict
    _orderedDict = True
except ImportError:
    _orderedDict = False

# PyMOL 1.7.4 introduces support for multi-letter chains, so we can afford to
# use a smaller alphabet. In earlier versions, use lower-case letters if needed
# (requires running `set ignore_case, 0`)
_long_chains = cmd.get_version()[1] >= 1.74
_default_base = 36 if _long_chains else 62

class OutOfChainsError(Exception):
    def __init__(self,msg):
        self.msg=msg
    def __str__(self):
        return str(self.msg)

class ChainSet(object):
    """
    Base class for various methods to rename chains

    Contains _chains, which maps from the renamed chain to a tuple with the
    original (object,state,chain). All dict-like accessors work on ChainSets,
    e.g.
        chain_set["A"] -> ("obj",1,"A")

    """
    def __init__(self):
        # Use an OrderedDict in Python >= 1.7 for better printing
        if _orderedDict:
            self._chains = OrderedDict()
        else:
            self._chains = dict()

    def map_chain(self, obj, state, origChain ):
        """
        map_chain(string obj,int state, string chain]]) -> string

        Maps a chain letter to a unique chainID. Results are unique within each
        instance, and can be used as keys on this chain set.
        """
        raise NotImplementedError("Base class")

    # delegate most methods to _chains
    def __getattr__(self,at):
        if at in "pop popitem update setdefault".split():
            raise AttributeError("type object '%s' has no attribute '%s'"%(type(self),at))
        return getattr(self._chains,at)
    def __cmp__(self,other):    return self._chains.__cmp__(other)
    def __eq__(self,other):     return self._chains.__eq__(other)
    def __ge__(self,other):     return self._chains.__ge__(other)
    def __gt__(self,other):     return self._chains.__gt__(other)
    def __le__(self,other):     return self._chains.__le__(other)
    def __lt__(self,other):     return self._chains.__lt__(other)
    def __ne__(self,other):     return self._chains.__ne__(other)
    def __len__(self):          return self._chains.__len__()
    def __contains__(self,key): return self._chains.__contains__(key)
    def __getitem__(self,key):  return self._chains.__getitem__(key)
    def __iter__(self):         return self._chains.__iter__()
    def __str__(self):          return str(self._chains)

    @staticmethod
    def _int_to_chain(i,base=_default_base):
        """
        _int_to_chain(int,int) -> str

        Converts a positive integer to a chain ID. Chain IDs include uppercase
        characters, numbers, and optionally lowercase letters.

        i = a positive integer to convert
        base = the alphabet size to include. Typically 36 or 62.
        """
        if i < 0:
            raise ValueError("positive integers only")
        if base < 0 or 62 < base:
            raise ValueError("Invalid base")

        quot = int(i)//base
        rem = i%base
        if rem < 26:
            letter = chr( ord("A") + rem)
        elif rem < 36:
            letter = str( rem-26)
        else:
            letter = chr( ord("a") + rem - 36)
        if quot == 0:
            return letter
        else:
            return ChainSet._int_to_chain(quot-1,base) + letter


class DefaultChainSet(ChainSet):
    """
    Avoids relettering chains if possible. If a chain has been used, uses the
    next available chain letter. Note that this can potentially lead to
    cascading renames, e.g. if chains are sorted alphabetically rather than by
    object.

    Used for rename = 0.
    """
    def __init__(self):
        super(DefaultChainSet,self).__init__()
        self._next_chain = 0
    def map_chain(self, obj, state, origChain ):
        # Keep _next_chain up-to-date
        while ChainSet._int_to_chain(self._next_chain) in self:
            self._next_chain += 1
        # Map this chain
        if origChain in self:
            # Rename
            next_chain = ChainSet._int_to_chain(self._next_chain)
            self._next_chain += 1
        else:
            next_chain = origChain
        self._chains[next_chain] = (obj,state,origChain)
        return next_chain

class SequentialChainSet(ChainSet):
    """
    Renumbers all chains starting at A, continuing through the capital letters
    and numbers, and then adding additional letters through 9999 (the last
    valid chain for mmCIF) and beyond.

    Used for rename=1
    """
    def __init__(self):
        super(SequentialChainSet,self).__init__()
        self._next_chain = 0

    def map_chain(self, obj, state, origChain ):
        next_chain = ChainSet._int_to_chain(self._next_chain)
        self._chains[next_chain] = (obj,state,origChain)
        self._next_chain += 1
        return next_chain

class LongChainSet(ChainSet):
    """
    Uses long strings for the chain names. Chains are renamed like
    "%s_%s_%04d"%(original_chainid,objectname,state).

    Used for rename=2
    """
    def map_chain(self, obj, state, origChain ):
        ch = "%s_%s_%04d"%(origChain,obj,state)
        if ch in self:
            raise ValueError("Duplicate chain %s"%(ch))
        self._chains[ch] = (obj,state,origChain)
        return ch




def flatten_obj(name="",selection="",state=0,rename=0,quiet=1,chain_map=""):
    """
DESCRIPTION

    "flatten_obj" combines multiple objects or states into a single object,
    renaming chains where required

USAGE

    flatten_obj name, selection[, state[, rename[, quiet[, chain_map]]]]

ARGUMENTS

    name = a unique name for the flattened object {default: flat}

    selection = the set of objects to include in the flattening. The selection
        will be expanded to include all atoms of objects. {default: all}

    state = the source state to select. Use 0 or -1 to flatten all states {default: 0}

    rename = The scheme to use for renaming chains: {default: 0}
        (0) preserve chains IDs where possible, rename other chains
            alphabetically
        (1) rename all chains alphabetically
        (2) rename chains using the original chain letter, object name, and state

    quiet = If set to 0, print some additional information about progress and
        chain renaming {default: 1}

    chain_map = An attribute name for the 'stored' scratch object. If
        specified, `stored.<chain_map>` will be populated with a dictionary
        mapping the new chain names to a tuple giving the originated object,
        state, and chainID. {default: ""}

NOTES

    Like the select command, if name is omitted then the default object name
    ("flat") is used as the name argument.

    Chain renaming is tricky. PDB files originally limited chains to single
    letter identifiers containing [A-Za-z0-9]. When this was found to be
    limiting, multi-letter chains (ideally < 4 chars) were allowed. This is
    supported as of PyMOL 1.7. Earlier versions do not accept rename=2, and
    will raise an exception when flattening a structure with more than 62
    chains.

EXAMPLES

    flatten_obj flat, nmrObj
    flatten_obj ( obj1 or obj2 )

SEE ALSO

    split_states

    """

    # arguments

    # Single argument; treat as selection
    if name and not selection:
        selection = name
        name = ""
    # default name and selection
    if not name:
        name = "flat"
    if not selection:
        selection = "(all)"

    state = int(state)
    rename = int(rename)
    quiet = int(quiet)

    # Wrap in extra parantheses for get_object_list
    selection = "( %s )" % selection

    if rename == 0:
        chainSet = DefaultChainSet()
    elif rename == 1:
        chainSet = SequentialChainSet()
    elif rename == 2:
        chainSet = LongChainSet()
    else:
        raise ValueError("Unrecognized rename option (Valid: 0,1,2)")

    metaprefix = "temp" #TODO unique prefix

    # store original value of retain_order, which causes weird interleaving of
    # structures if enabled.
    retain_order = cmd.get("retain_order")
    try:
        cmd.set("retain_order",0)

        # create new object for each state
        for obj in cmd.get_object_list(selection):

            if state <= 0:
                # all states
                prefix = "%s_%s_"%(metaprefix,obj)
                cmd.split_states(obj,prefix=prefix)
            else:
                prefix = "%s_%s_%04d"%(metaprefix,obj,state)
                cmd.create(prefix, obj, state, 1)

        # renumber all states
        statere = re.compile("^%s_(.*)_(\d+)$" % metaprefix) # matches split object names

        warn_lowercase = False

        # Iterate over all objects with metaprefix
        try:
            for obj in cmd.get_object_list("(%s_*)"%(metaprefix) ):
                m = statere.match(obj)
                if m is None:
                    print(("Failed to match object %s" %obj))
                    continue
                origobj = m.group(1)
                statenum = int(m.group(2))

                chains = cmd.get_chains(obj)

                rev_chain_map = {} #old -> new, for this obj only
                for chain in sorted(chains,key=lambda x:(len(x),x)):
                    new_chain = chainSet.map_chain(origobj,statenum,chain)
                    rev_chain_map[chain] = new_chain
                    if not quiet:
                        print(("  %s state %d chain %s -> %s"%(origobj,statenum,chain, new_chain) ))
                    if not _long_chains:
                        if len(new_chain) > 1:
                            raise OutOfChainsError("No additional chains available (max 62).")

                space = {'rev_chain_map':rev_chain_map}
                cmd.alter(obj,"chain = rev_chain_map[chain]",space=space)

            print(("Creating object from %s_*"%metaprefix))
            # Recombine into a single object
            cmd.create(name,"%s_*"%metaprefix)

            # Set chain_map
            if chain_map:
                setattr(stored,chain_map,chainSet)

            # Warn if lowercase chains were generated
            if cmd.get("ignore_case") == "on" and any([c.upper() != c for c in list(chainSet.keys())]):
                print("Warning: using lower-case chain IDs. Consider running the "
                        "following command:\n  set ignore_case, 0" )

        finally:
            # Clean up
            print("Cleaning up intermediates")
            cmd.delete("%s_*"%metaprefix)
    finally:
        # restore original parameters
        print("Resetting variables")
        cmd.set("retain_order",retain_order)


cmd.extend('flatten_obj', flatten_obj)

# tab-completion of arguments
cmd.auto_arg[0]['flatten_obj'] = [ cmd.object_sc, 'name or selection', '']
cmd.auto_arg[1]['flatten_obj'] = [ cmd.object_sc, 'selection', '']


1			"""
2			flatten_obj.py - Flatten multi-state pymol objects into a single state.
3
4			<https://pymolwiki.org/index.php/Flatten_obj>
5
6			This is particularly useful for dealing with biological assemblies, which are
7			loaded as multi-state objects when fetched using `fetch PDBID, type=pdb1`. It
8			can also be used as a quick way to combine multiple objects without causing
9			collisions between chain identifiers.
10
11			The command re-letters chains to avoid collisions. Older versions of PyMOL
12			restrict the chain id to a single character, so the script will fail for
13			assemblies with >62 chains. With more recent versions, this problem is solved
14			with multi-character chain IDs. Several options are available for how
15			re-lettering should occur.
16
17			Author: Spencer Bliven <[email protected]>
18			Date: October 30, 2015
19			Version: 1.0
20			License: Public Domain
21			"""
22			from pymol import cmd, stored
23			import re
24			try:
25			from collections import OrderedDict
26			_orderedDict = True
27			except ImportError:
28			_orderedDict = False
29
30			# PyMOL 1.7.4 introduces support for multi-letter chains, so we can afford to
31			# use a smaller alphabet. In earlier versions, use lower-case letters if needed
32			# (requires running `set ignore_case, 0`)
33			_long_chains = cmd.get_version()[1] >= 1.74
34			_default_base = 36 if _long_chains else 62
35
36			class OutOfChainsError(Exception):
37			def __init__(self,msg):
38			self.msg=msg
39			def __str__(self):
40			return str(self.msg)
41
42			class ChainSet(object):
43			"""
44			Base class for various methods to rename chains
45
46			Contains _chains, which maps from the renamed chain to a tuple with the
47			original (object,state,chain). All dict-like accessors work on ChainSets,
48			e.g.
49			chain_set["A"] -> ("obj",1,"A")
50
51			"""
52			def __init__(self):
53			# Use an OrderedDict in Python >= 1.7 for better printing
54			if _orderedDict:
55			self._chains = OrderedDict()
56			else:
57			self._chains = dict()
58
59			def map_chain(self, obj, state, origChain ):
60			"""
61			map_chain(string obj,int state, string chain]]) -> string
62
63			Maps a chain letter to a unique chainID. Results are unique within each
64			instance, and can be used as keys on this chain set.
65			"""
66			raise NotImplementedError("Base class")
67
68			# delegate most methods to _chains
69			def __getattr__(self,at):
70			if at in "pop popitem update setdefault".split():
71			raise AttributeError("type object '%s' has no attribute '%s'"%(type(self),at))
72			return getattr(self._chains,at)
73			def __cmp__(self,other): return self._chains.__cmp__(other)
74			def __eq__(self,other): return self._chains.__eq__(other)
75			def __ge__(self,other): return self._chains.__ge__(other)
76			def __gt__(self,other): return self._chains.__gt__(other)
77			def __le__(self,other): return self._chains.__le__(other)
78			def __lt__(self,other): return self._chains.__lt__(other)
79			def __ne__(self,other): return self._chains.__ne__(other)
80			def __len__(self): return self._chains.__len__()
81			def __contains__(self,key): return self._chains.__contains__(key)
82			def __getitem__(self,key): return self._chains.__getitem__(key)
83			def __iter__(self): return self._chains.__iter__()
84			def __str__(self): return str(self._chains)
85
86			@staticmethod
87			def _int_to_chain(i,base=_default_base):
88			"""
89			_int_to_chain(int,int) -> str
90
91			Converts a positive integer to a chain ID. Chain IDs include uppercase
92			characters, numbers, and optionally lowercase letters.
93
94			i = a positive integer to convert
95			base = the alphabet size to include. Typically 36 or 62.
96			"""
97			if i < 0:
98			raise ValueError("positive integers only")
99			if base < 0 or 62 < base:
100			raise ValueError("Invalid base")
101
102			quot = int(i)//base
103			rem = i%base
104			if rem < 26:
105			letter = chr( ord("A") + rem)
106			elif rem < 36:
107			letter = str( rem-26)
108			else:
109			letter = chr( ord("a") + rem - 36)
110			if quot == 0:
111			return letter
112			else:
113			return ChainSet._int_to_chain(quot-1,base) + letter
114
115
116			class DefaultChainSet(ChainSet):
117			"""
118			Avoids relettering chains if possible. If a chain has been used, uses the
119			next available chain letter. Note that this can potentially lead to
120			cascading renames, e.g. if chains are sorted alphabetically rather than by
121			object.
122
123			Used for rename = 0.
124			"""
125			def __init__(self):
126			super(DefaultChainSet,self).__init__()
127			self._next_chain = 0
128			def map_chain(self, obj, state, origChain ):
129			# Keep _next_chain up-to-date
130			while ChainSet._int_to_chain(self._next_chain) in self:
131			self._next_chain += 1
132			# Map this chain
133			if origChain in self:
134			# Rename
135			next_chain = ChainSet._int_to_chain(self._next_chain)
136			self._next_chain += 1
137			else:
138			next_chain = origChain
139			self._chains[next_chain] = (obj,state,origChain)
140			return next_chain
141
142			class SequentialChainSet(ChainSet):
143			"""
144			Renumbers all chains starting at A, continuing through the capital letters
145			and numbers, and then adding additional letters through 9999 (the last
146			valid chain for mmCIF) and beyond.
147
148			Used for rename=1
149			"""
150			def __init__(self):
151			super(SequentialChainSet,self).__init__()
152			self._next_chain = 0
153
154			def map_chain(self, obj, state, origChain ):
155			next_chain = ChainSet._int_to_chain(self._next_chain)
156			self._chains[next_chain] = (obj,state,origChain)
157			self._next_chain += 1
158			return next_chain
159
160			class LongChainSet(ChainSet):
161			"""
162			Uses long strings for the chain names. Chains are renamed like
163			"%s_%s_%04d"%(original_chainid,objectname,state).
164
165			Used for rename=2
166			"""
167			def map_chain(self, obj, state, origChain ):
168			ch = "%s_%s_%04d"%(origChain,obj,state)
169			if ch in self:
170			raise ValueError("Duplicate chain %s"%(ch))
171			self._chains[ch] = (obj,state,origChain)
172			return ch
173
174
175
176
177			def flatten_obj(name="",selection="",state=0,rename=0,quiet=1,chain_map=""):
178			"""
179			DESCRIPTION
180
181			"flatten_obj" combines multiple objects or states into a single object,
182			renaming chains where required
183
184			USAGE
185
186			flatten_obj name, selection[, state[, rename[, quiet[, chain_map]]]]
187
188			ARGUMENTS
189
190			name = a unique name for the flattened object {default: flat}
191
192			selection = the set of objects to include in the flattening. The selection
193			will be expanded to include all atoms of objects. {default: all}
194
195			state = the source state to select. Use 0 or -1 to flatten all states {default: 0}
196
197			rename = The scheme to use for renaming chains: {default: 0}
198			(0) preserve chains IDs where possible, rename other chains
199			alphabetically
200			(1) rename all chains alphabetically
201			(2) rename chains using the original chain letter, object name, and state
202
203			quiet = If set to 0, print some additional information about progress and
204			chain renaming {default: 1}
205
206			chain_map = An attribute name for the 'stored' scratch object. If
207			specified, `stored.<chain_map>` will be populated with a dictionary
208			mapping the new chain names to a tuple giving the originated object,
209			state, and chainID. {default: ""}
210
211			NOTES
212
213			Like the select command, if name is omitted then the default object name
214			("flat") is used as the name argument.
215
216			Chain renaming is tricky. PDB files originally limited chains to single
217			letter identifiers containing [A-Za-z0-9]. When this was found to be
218			limiting, multi-letter chains (ideally < 4 chars) were allowed. This is
219			supported as of PyMOL 1.7. Earlier versions do not accept rename=2, and
220			will raise an exception when flattening a structure with more than 62
221			chains.
222
223			EXAMPLES
224
225			flatten_obj flat, nmrObj
226			flatten_obj ( obj1 or obj2 )
227
228			SEE ALSO
229
230			split_states
231
232			"""
233
234			# arguments
235
236			# Single argument; treat as selection
237			if name and not selection:
238			selection = name
239			name = ""
240			# default name and selection
241			if not name:
242			name = "flat"
243			if not selection:
244			selection = "(all)"
245
246			state = int(state)
247			rename = int(rename)
248			quiet = int(quiet)
249
250			# Wrap in extra parantheses for get_object_list
251			selection = "( %s )" % selection
252
253			if rename == 0:
254			chainSet = DefaultChainSet()
255			elif rename == 1:
256			chainSet = SequentialChainSet()
257			elif rename == 2:
258			chainSet = LongChainSet()
259			else:
260			raise ValueError("Unrecognized rename option (Valid: 0,1,2)")
261
262			metaprefix = "temp" #TODO unique prefix
263
264			# store original value of retain_order, which causes weird interleaving of
265			# structures if enabled.
266			retain_order = cmd.get("retain_order")
267			try:
268			cmd.set("retain_order",0)
269
270			# create new object for each state
271			for obj in cmd.get_object_list(selection):
272
273			if state <= 0:
274			# all states
275			prefix = "%s_%s_"%(metaprefix,obj)
276			cmd.split_states(obj,prefix=prefix)
277			else:
278			prefix = "%s_%s_%04d"%(metaprefix,obj,state)
279			cmd.create(prefix, obj, state, 1)
280
281			# renumber all states
282			statere = re.compile("^%s_(.*)_(\d+)$" % metaprefix) # matches split object names
283
284			warn_lowercase = False
285
286			# Iterate over all objects with metaprefix
287			try:
288			for obj in cmd.get_object_list("(%s_*)"%(metaprefix) ):
289			m = statere.match(obj)
290			if m is None:
291			print(("Failed to match object %s" %obj))
292			continue
293			origobj = m.group(1)
294			statenum = int(m.group(2))
295
296			chains = cmd.get_chains(obj)
297
298			rev_chain_map = {} #old -> new, for this obj only
299			for chain in sorted(chains,key=lambda x:(len(x),x)):
300			new_chain = chainSet.map_chain(origobj,statenum,chain)
301			rev_chain_map[chain] = new_chain
302			if not quiet:
303			print((" %s state %d chain %s -> %s"%(origobj,statenum,chain, new_chain) ))
304			if not _long_chains:
305			if len(new_chain) > 1:
306			raise OutOfChainsError("No additional chains available (max 62).")
307
308			space = {'rev_chain_map':rev_chain_map}
309			cmd.alter(obj,"chain = rev_chain_map[chain]",space=space)
310
311			print(("Creating object from %s_*"%metaprefix))
312			# Recombine into a single object
313			cmd.create(name,"%s_*"%metaprefix)
314
315			# Set chain_map
316			if chain_map:
317			setattr(stored,chain_map,chainSet)
318
319			# Warn if lowercase chains were generated
320			if cmd.get("ignore_case") == "on" and any([c.upper() != c for c in list(chainSet.keys())]):
321			print("Warning: using lower-case chain IDs. Consider running the "
322			"following command:\n set ignore_case, 0" )
323
324			finally:
325			# Clean up
326			print("Cleaning up intermediates")
327			cmd.delete("%s_*"%metaprefix)
328			finally:
329			# restore original parameters
330			print("Resetting variables")
331			cmd.set("retain_order",retain_order)
332
333
334			cmd.extend('flatten_obj', flatten_obj)
335
336			# tab-completion of arguments
337			cmd.auto_arg[0]['flatten_obj'] = [ cmd.object_sc, 'name or selection', '']
338			cmd.auto_arg[1]['flatten_obj'] = [ cmd.object_sc, 'selection', '']
339

mmagnus / rna-tools

ChainSet.map_chain() A last analyzed 2025-03-19 22:26 UTC

Complexity

Size

Duplication

Importance

Duplication Side-by-Side

Filter issues like

ChainSet.map_chain() A
last analyzed 2025-03-19 22:26 UTC