#!/usr/bin/env python """Provides the entities, the building blocks of the SMRCA hierachy representation of a macromolecular structure. The MultiEntity class is a special Entity class to hold multiple instances of other entities. All Entities apart from the Atom can hold others and inherit from the MultiEntity. The Entity is the most basic class to deal with structural and molecular data. Do not use it directly since some functions depend on methods provided by sub-classes. Classes inheriting from MultiEntity have to provide some attributes during init e.g: self.level = a valid string inside the SMCRA hierarchy). Holders of entities are like normal MultiEntities, but are temporary and are outside the parent-children axes. """ from numpy import (sqrt, arctan2, power, array, mean, sum) from cogent.data.protein_properties import AA_NAMES, AA_ATOM_BACKBONE_ORDER, \ AA_ATOM_REMOTE_ORDER, AREAIMOL_VDW_RADII, \ DEFAULT_AREAIMOL_VDW_RADIUS from cogent.data.ligand_properties import HOH_NAMES, LIGAND_AREAIMOL_VDW_RADII from operator import itemgetter, gt, ge, lt, le, eq, ne, or_, and_, contains, \ is_, is_not from collections import defaultdict from itertools import izip from copy import copy, deepcopy __author__ = "Marcin Cieslik" __copyright__ = "Copyright 2009, The Cogent Project" __credits__ = ["Marcin Cieslik"] __license__ = "GPL" __version__ = "1.4.1" __maintainer__ = "Marcin Cieslik" __email__ = "mpc4p@virginia.edu" __status__ = "Development" ALPHABET = 'ABCDEFGHIJKLMNOPQRSTUVWXYZ_ ' HIERARCHY = ['H', 'S', 'M', 'C', 'R', 'A'] AREAIMOL_VDW_RADII.update(LIGAND_AREAIMOL_VDW_RADII) # error while creating a structure (non-recoverable error) class ConstructionError(Exception): """Cannot unambiguously create a structure.""" pass # warning while creating a structure # (something wrong with the input, but recoverable) class ConstructionWarning(Exception): """Input violates some construction rules (contiguity).""" pass def sort_id_list(id_list, sort_tuple): """Sorts lists of id tuples. The order is defined by the PDB file specification.""" (hol_loc, str_loc, mod_loc, chn_loc, res_loc, at_loc) = sort_tuple # even a simple id is a tuple, this makes sorting general def space_last(ch_id1, ch_id2): # this is for chain sorting if ch_id1 == ' ' and ch_id2 != ' ': return 1 if ch_id2 == ' ' and ch_id1 != ' ': return - 1 if ch_id1 == ' ' and ch_id2 == ' ': return 0 return cmp(ch_id1, ch_id2) def atom(at_id1, at_id2): # hydrogen atoms come last is_hydrogen1 = (at_id1[0] == 'H') is_hydrogen2 = (at_id2[0] == 'H') diff = cmp(is_hydrogen1, is_hydrogen2) # back bone come first if not diff: order1 = AA_ATOM_BACKBONE_ORDER.get(at_id1) order2 = AA_ATOM_BACKBONE_ORDER.get(at_id2) diff = cmp(order2, order1) # (B)eta, (D)elta, (G)amma, .... o(X)t if not diff: remote1 = AA_ATOM_REMOTE_ORDER.get(at_id1[1:2]) remote2 = AA_ATOM_REMOTE_ORDER.get(at_id2[1:2]) diff = cmp(remote1, remote2) # branching comes last if not diff: diff = cmp(at_id1[2:4], at_id2[2:4]) return diff # SE vs CE - selenium first if not diff: alpha1 = ALPHABET.index(at_id1[0:1]) alpha2 = ALPHABET.index(at_id2[0:1]) diff = cmp(alpha2, alpha1) def residue(res_id1, res_id2): r1, r2 = 1, 1 if res_id1 in AA_NAMES: r1 = 2 if res_id1 in HOH_NAMES: r1 = 0 if res_id2 in AA_NAMES: r2 = 2 if res_id2 in HOH_NAMES: r2 = 0 if r1 is r2: return cmp(res_id1, res_id2) else: return cmp(r2, r1) # this assumes that the implementation of sorting is stable. # does it work for others then cPython. if res_loc or res_loc is 0: id_list.sort(key=itemgetter(res_loc), cmp=lambda x, y: residue(x[0], y[0])) # by res_name if at_loc or at_loc is 0: id_list.sort(key=itemgetter(at_loc), cmp=lambda x, y: space_last(x[1], y[1])) # by alt_loc if at_loc or at_loc is 0: id_list.sort(key=itemgetter(at_loc), cmp=lambda x, y: atom(x[0], y[0])) # by at_id if res_loc or res_loc is 0: id_list.sort(key=itemgetter(res_loc), cmp=lambda x, y: cmp(x[2], y[2])) # by res_ic if res_loc or res_loc is 0: id_list.sort(key=itemgetter(res_loc), cmp=lambda x, y: cmp(x[1], y[1])) # by res_id if chn_loc or chn_loc is 0: id_list.sort(key=itemgetter(chn_loc), cmp=space_last) # by chain if mod_loc or mod_loc is 0: id_list.sort(key=itemgetter(mod_loc)) # by model if str_loc or str_loc is 0: id_list.sort(key=itemgetter(str_loc)) # by structure return id_list def merge(dicts): """Merges multiple dictionaries into a new one.""" master_dict = {} for dict_ in dicts: master_dict.update(dict_) return master_dict def unique(lists): """Merges multiple iterables into a unique sorted tuple (sorted set).""" master_set = set() for set_ in lists: master_set.update(set_) return tuple(sorted(master_set)) class Entity(dict): """Container object all entities inherit from it. Inherits from dict.""" def __init__(self, id, name=None, *args): # This class has to be sub-classed! # the masked attribute has to be set before the __init__ of an Entity # because during __setstate__, __getstate__ sub-entities are iterated # by .values(), which relies on the attribute masked. to decide which # children should be omitted. self.masked = False self.parent = None # mandatory parent attribute self.modified = True # modified on creation self.id = (id,) # ids are non-zero lenght tuples self.name = (name or id) # prefer name over duplicate id self.xtra = {} # mandatory xtra dict attribute # Dictionary that keeps additional properties dict.__init__(self, *args) # finish init as dictionary def __copy__(self): return deepcopy(self) def __deepcopy__(self, memo): new_state = self.__getstate__() new_instance = self.__new__(type(self)) new_instance.__setstate__(new_state) return new_instance def __getstate__(self): new_state = copy(self.__dict__) # shallow new_state['parent'] = None return new_state def __setstate__(self, new_state): self.__dict__.update(new_state) def __repr__(self): """Default representation.""" # mandatory getLevel from sub-class return "" % (self.getId(), self.getLevel()) def __sub__(self, entity): """Override "-" as Euclidean distance between coordinates.""" return sqrt(sum(pow(self.coords - entity.coords, 2))) def _setId(self, id): self.name = id[0] def _getId(self): return (self.name,) def getId(self): """Return the id.""" return self._getId() def getFull_id(self): """Return the full id.""" parent = self.getParent() if parent: full_id = parent.getFull_id() else: full_id = () # we create a tuple on the top full_id = full_id + self.getId() # merge tuples from the left return full_id def setId(self, id_=None): """Set the id. Calls the ``_setId`` method.""" if (id_ and id_ != self.id) or (not id_ and (self.getId() != self.id)): self.id = (id_ or self.getId()) self.setModified(True, True) self._setId(self.id) if self.parent: self.parent.updateIds() def _setMasked(self, masked, force=False): if masked != self.masked or force: self.masked = masked # mask or unmask self.setModified(True, False) # set parents as modified def setMasked(self, *args, **kwargs): """Set masked flag (``masked``) ``True``.""" self._setMasked(True, *args, **kwargs) def setUnmasked(self, *args, **kwargs): """Set masked flag (``masked``) ``False``.""" self._setMasked(False, *args, **kwargs) def setModified(self, up=True, down=False): """Set modified flag (``modified``) ``True``.""" self.modified = True if up and self.parent: self.parent.setModified(True, False) def setUnmodified(self, up=False, down=False): """Set modified flag (``modified``) ``False``.""" self.modified = False if up and self.parent: self.parent.setUnmodified(True, False) def setParent(self, entity): """Set the parent ``Entity`` and adds oneself as the child.""" if self.parent != entity: # delete old parent self.delParent() # add new parent self.parent = entity self.parent.addChild(self) self.setModified(False, True) def delParent(self): """Detach mutually from the parent. Sets both child and parent modified flags (``modified``) as ``True``.""" if self.parent: self.parent.pop(self.getId()) self.parent.setModified(True, False) self.parent = None self.setModified(False, True) def getModified(self): """Return value of the modified flag (``modified``).""" return self.modified def getMasked(self): """Return value of the masked flag (``masked``).""" return self.masked def setLevel(self, level): """Set level (``level``).""" self.level = level def getLevel(self): """Return level (``level``)in the hierarchy.""" return self.level def setName(self, name): """Set name.""" self.name = name self.setId() def getName(self): """Return name.""" return self.name def getParent(self, level=None): """Return the parent ``Entity`` instance.""" if not level: return self.parent elif level == self.level: return self return self.parent.getParent(level) def move(self, origin): """Subtract the origin coordinates from the coordintats (``coords``).""" self.coords = self.coords - origin def setCoords(self, coords): """Set the entity coordinates. Coordinates should be a ``numpy.array``.""" self.coords = coords def getCoords(self): """Get the entity coordinates.""" return self.coords def getScoords(self): """Return spherical (r, theta, phi) coordinates.""" x, y, z = self.coords x2, y2, z2 = power(self.coords, 2) scoords = array([sqrt(x2 + y2 + z2), \ arctan2(sqrt(x2 + y2), z), \ arctan2(y, x)]) return scoords def getCcoords(self): """Return redundant, polar, clustering-coordinates on the unit-sphere. This is only useful for clustering.""" x, y, z = self.coords x2, y2, z2 = power(self.coords, 2) ccoords = array([arctan2(sqrt(y2 + z2), x), \ arctan2(sqrt(x2 + z2), y), \ arctan2(sqrt(x2 + y2), z) ]) return ccoords def setScoords(self): """Set ``entity.scoords``, see: getScoords.""" self.scoords = self.getScoords() def setCcoords(self): """Set ``entity.ccoords``, see: getCcoords.""" self.ccoords = self.getCcoords() class MultiEntity(Entity): """The ``MultiEntity`` contains other ``Entity`` or ``MultiEntity`` instances.""" def __init__(self, long_id, short_id=None, *args): self.index = HIERARCHY.index(self.level) # index corresponding to the hierarchy level self.table = dict([(level, {}) for level in HIERARCHY[self.index + 1:]]) # empty table Entity.__init__(self, long_id, short_id, *args) def __repr__(self): id_ = self.getId() return "" % (id_, len(self)) def _link(self): """Recursively adds a parent pointer to children.""" for child in self.itervalues(unmask=True): child.parent = self try: child._link() except AttributeError: pass def _unlink(self): """Recursively deletes the parent pointer from children.""" for child in self.itervalues(unmask=True): child.parent = None try: child._unlink() except AttributeError: pass def __getstate__(self): new_dict = copy(self.__dict__) # shallow copy new_dict['parent'] = None # remove recursion new_children = [] for child in self.itervalues(unmask=True): new_child_instance = deepcopy(child) new_children.append(new_child_instance) return (new_children, new_dict) def __setstate__(self, new_state): new_children, new_dict = new_state self.__dict__.update(new_dict) for child in new_children: self.addChild(child) def __copy__(self): return deepcopy(self) def __deepcopy__(self, memo): new_state = self.__getstate__() new_instance = self.__new__(type(self)) new_instance.__setstate__(new_state) return new_instance def __iter__(self): return self.itervalues() def setSort_tuple(self, sort_tuple=None): """Set the ``sort_tuple attribute``. The ``sort_tuple`` is a tuple needed by the ``sort_id_list`` function to correctly sort items within entities.""" if sort_tuple: self.sort_tuple = sort_tuple else: # making the sort tuple, ugly, uggly, uaughhlly ble sort_tuple = [None, None, None, None, None, None] key_lenght = len(self.keys()[0]) stop_i = self.index + 2 # next level, open right [) start_i = stop_i - key_lenght # before all nones indexes = range(start_i, stop_i) # Nones to change for value, index in enumerate(indexes): sort_tuple[index] = value self.sort_tuple = sort_tuple def getSort_tuple(self): """Return the ``sort_tuple`` attribute. If not set calls the ``setSort_tuple`` method first. See: ``setSort_tuple``.""" if not hasattr(self, 'sort_tuple'): self.setSort_tuple() return self.sort_tuple def itervalues(self, unmask=False): return (v for v in super(MultiEntity, self).itervalues() if not v.masked or unmask) def iteritems(self, unmask=False): return ((k, v) for k, v in super(MultiEntity, self).iteritems() if not v.masked or unmask) def iterkeys(self, unmask=False): return (k for k, v in super(MultiEntity, self).iteritems() if not v.masked or unmask) def values(self, *args, **kwargs): return list(self.itervalues(*args, **kwargs)) def items(self, *args, **kwargs): return list(self.iteritems(*args, **kwargs)) def keys(self, *args, **kwargs): return list(self.iterkeys(*args, **kwargs)) def __contains__(self, key, *args, **kwargs): return key in self.keys(*args, **kwargs) def sortedkeys(self, *args, **kwargs): list_ = sort_id_list(self.keys(*args, **kwargs), self.getSort_tuple()) return list_ def sortedvalues(self, *args, **kwargs): values = [self[i] for i in self.sortedkeys(*args, **kwargs)] return values def sorteditems(self, *args, **kwargs): items = [(i, self[i]) for i in self.sortedkeys()] return items def _setMasked(self, masked, force=False): """Set the masked flag (``masked``) recursively. If forced proceed even if the flag is already set correctly.""" if masked != self.masked or force: # the second condition is when if masked: # an entity has all children masked # we have to mask children # but is not masked itself for child in self.itervalues(): # only unmasked children child.setMasked() child.setModified(False, False) else: # we have to unmask children for child in self.itervalues(unmask=True): if child.masked or force: # only masked children child.setUnmasked(force=force) child.setModified(False, False) self.masked = masked self.setModified(True, False) # set parents as modified def setModified(self, up=True, down=True): """Set the modified flag (``modified``) ``True``. If down proceeds recursively for all children. If up proceeds recursively for all parents.""" self.modified = True if up and self.parent: self.parent.setModified(True, False) if down: for child in self.itervalues(unmask=True): child.setModified(False, True) def setUnmodified(self, up=False, down=False): """Set the modified (``modified``) flag ``False``. If down proceeds recursively for all children. If up proceeds recursively for all parents.""" self.modified = False if up and self.parent: self.parent.setUnmodified(True, False) if down: for child in self.itervalues(unmask=True): child.setUnmodified(False, True) def _initChild(self, child): """Initialize a child (during construction).""" child.parent = self self[child.getId()] = child def addChild(self, child): """Add a child.""" child.setParent(self) child_id = child.getId() self[child_id] = child self.setModified(True, False) def delChild(self, child_id): """Remove a child.""" child = self.get(child_id) if child: child.delParent() self.setModified(True, False) def getChildren(self, ids=None, **kwargs): """Return a copy of the list of children.""" if ids: children = [] for (id_, child) in self.iteritems(**kwargs): if id_ in ids: children.append(child) else: children = self.values(**kwargs) return children def _setTable(self, entity): """Recursive helper method for ``entity.setTable``.""" for e in entity.itervalues(): self.table[e.getLevel()].update({e.getFull_id():e}) self._setTable(e) def setTable(self, force=True, unmodify=True): """Populate the children table (``table``) recursively with all children grouped into hierarchy levels. If forced is ``True`` the table will be updated even if the ``Entity`` instance is not modified. If unmodify is ``True`` the ``Entity`` modified flag (``modified``) will be set ``False`` afterwards.""" if self.modified or force: # a table is accurate as long as the contents of a dictionary do not # change. self.delTable() self._setTable(self) if unmodify: self.setUnmodified() def delTable(self): """Delete all children from the children-table (``table``). This does not modify the hierarchy.""" self.table = dict([(level, {}) for level in HIERARCHY[self.index + 1:]]) self.modified = True def getTable(self, level): """Return children of given level from the children-table (``table``).""" return self.table[level] def updateIds(self): """Update self with children ids.""" ids = [] for (id_, child) in self.iteritems(): new_id = child.getId() if id_ != new_id: ids.append((id_, new_id)) for (old_id, new_id) in ids: child = self.pop(old_id) self.update(((new_id, child),)) def getData(self, attr, xtra=False, method=False, forgiving=True, sorted=False): """Get data from children attributes, methods and xtra dicts as a list. If is ``True`` forgiving remove ``None`` values from the output. ``Nones`` are place-holders if a child does not have the requested data. If xtra is True the xtra dictionary (``xtra``) will be searched, if method is ``True`` the child attribute will be called.""" values = self.sortedvalues() if sorted else self.values() if xtra: # looking inside the xtra of children data = [child.xtra.get(attr) for child in values] # could get None else: # looking at attributes data = [] for child in values: try: if not method: data.append(getattr(child, attr)) else: data.append(getattr(child, attr)()) except AttributeError: # data.append(None) if forgiving: # remove Nones data = [point for point in data if point is not None] return data def propagateData(self, function, level, attr, **kwargs): """Propagate data from child level to this ``Entity`` instance. The function defines how children data should be transformed to become the parents data e.g. summed.""" if self.index <= HIERARCHY.index(level) - 2: for child in self.itervalues(): child.propagateData(function, level, attr, **kwargs) datas = self.getData(attr, **kwargs) transformed_datas = function(datas) if kwargs.get('xtra'): self.xtra[attr] = transformed_datas else: setattr(self, attr, transformed_datas) return transformed_datas def countChildren(self, *args, **kwargs): """Count children based on ``getData``. Additional arguments and keyworded arguments are passed to the ``getData`` method.""" data = self.getData(*args, **kwargs) children = defaultdict(int) # by default returns 0 for d in data: children[d] += 1 return children def freqChildren(self, *args, **kwargs): """Frequency of children based on ``countChildren``. Additional arguments and keyworded arguments are passed to the ``countChildren`` method.""" children_count = self.countChildren(*args, **kwargs) lenght = float(len(self)) # it could be len(children_count)? for (key_, value_) in children_count.iteritems(): children_count[key_] = value_ / lenght return children_count def splitChildren(self, *args, **kwargs): """Splits children into groups children based on ``getData``. Additional arguments and keyworded arguments are passed to the ``getData`` method.""" kwargs['forgiving'] = False data = self.getData(*args, **kwargs) clusters = defaultdict(dict) # by default returns {} for (key, (id_, child)) in izip(data, self.iteritems()): clusters[key].update({id_:child}) return clusters def selectChildren(self, value, operator, *args, **kwargs): """Generic method to select children, based on ``getData``. Returns a dictionary of children indexed by ids. Compares the data item for each child using the operator name e.g. "eq" and a value e.g. "H_HOH". Additional arguments and keyworded arguments are passed to the ``getData`` method.""" kwargs['forgiving'] = False data = self.getData(*args, **kwargs) children = {} for (got, (id_, child)) in izip(data, self.iteritems()): if isinstance(operator, basestring): operator = eval(operator) if operator(value, got): children.update({id_:child}) return children def ornamentChildren(self, *args, **kwargs): """Return a list of (ornament, (id, child)) tuples, based on ``getData``. Useful for sorting see: Schwartzian transform. Forgiving is set False. Additional arguments and keyworded arguments are passed to the ``getData`` method.""" kwargs['forgiving'] = False data = self.getData(*args, **kwargs) children = [] for (got, (id_, child)) in izip(data, self.iteritems()): children.append((got, (id_, child))) return children def ornamentdictChildren(self, *args, **kwargs): """Return a dictionary of ornaments indexed by child ids, based on ``getData``. Forgiving is set False. Additional arguments and keyworded arguments are passed to the ``getData`` method.""" kwargs['forgiving'] = False data = self.getData(*args, **kwargs) propertydict = {} for (got, id_) in izip(data, self.iterkeys()): propertydict.update(((id_, got),)) return propertydict def stripChildren(self, *args, **kwargs): """Strips children based on selection criteria. See: ``selectChildren``. Additional arguments and keyworded arguments are passed to the ``selectChildren`` method.""" children_ids = self.selectChildren(*args, **kwargs).keys() for id_ in children_ids: self.delChild(id_) def maskChildren(self, *args, **kwargs): """Mask children based on selection criteria. See: ``selectChildren``. Additional arguments and keyworded arguments are passed to the ``selectChildren`` method.""" children = self.selectChildren(*args, **kwargs).itervalues() for child in children: child.setMasked() # child.setModified child.parent.setModified def unmaskChildren(self, *args, **kwargs): """Unmask children based on selection criteria. See: ``selectChildren``. Additional arguments and keyworded arguments are passed to the ``selectChildren`` method.""" children = self.selectChildren(*args, **kwargs).itervalues() for child in children: child.setUnmasked() # child.setModified child.parent.setModified def moveRecursively(self, origin): """Move ``Entity`` instance recursively to the origin.""" for child in self.itervalues(): try: child.moveRecursively(origin) except: # Atoms do not have this child.move(origin) pass self.setCoords() def setCoordsRecursively(self): """Set coordinates (``coords``) recursively. Useful if any child had its coordinates changed.""" for child in self.itervalues(): try: child.setCoordsRecursively() except: #Atoms do not have this pass self.setCoords() def setCoords(self, *args, **kwargs): """Set coordinates (``coords``) as a centroid of children coordinates. A subset of children can be selected for the calculation. See: ``Entity.selectChildren``. Additional arguments and keyworded arguments are passed to the ``getData`` method.""" # select only some children if args or kwargs: children = self.selectChildren(*args, **kwargs).values() else: children = self coords = [] for child in children: coords.append(child.getCoords()) self.coords = mean(coords, axis=0) def getCoords(self): """Returns the current coordinates (``coords``). Raises an ``AttributeError`` if not set.""" try: return self.coords except AttributeError: raise AttributeError, "Entity has coordinates not set." def dispatch(self, method, *args, **kwargs): """Calls a method of all children with given arguments and keyworded arguments.""" for child in self.itervalues(): getattr(child, method)(*args, **kwargs) class Structure(MultiEntity): """The ``Structure`` instance contains ``Model`` instances.""" def __init__(self, id, *args, **kwargs): self.level = 'S' MultiEntity.__init__(self, id, *args, **kwargs) def __repr__(self): return '' % self.getId() def removeAltmodels(self): """Remove all models with an id != 0""" self.stripChildren((0,), 'ne', 'id', forgiving=False) def getDict(self): """See: ``Entity.getDict``.""" return {'structure':self.getId()[0]} class Model(MultiEntity): """The ``Model`` instance contains ``Chain`` instances.""" def __init__(self, id, *args, **kwargs): self.level = 'M' MultiEntity.__init__(self, id, *args, **kwargs) def __repr__(self): return "" % self.getId() def getDict(self): """See: ``Entity.getDict``.""" try: from_parent = self.parent.getDict() except AttributeError: # we are allowed to silence this becaus a structure id is not # required to write a proper pdb line. from_parent = {} from_parent.update({'model':self.getId()[0]}) return from_parent class Chain(MultiEntity): """The ``Chain`` instance contains ``Residue`` instances.""" def __init__(self, id, *args, **kwargs): self.level = 'C' MultiEntity.__init__(self, id, *args, **kwargs) def __repr__(self): return "" % self.getId() def removeHetero(self): """Remove residues with the hetero flag.""" self.stripChildren('H', 'eq', 'h_flag', forgiving=False) def removeWater(self): """Remove water residues.""" self.stripChildren('H_HOH', 'eq', 'name', forgiving=False) def residueCount(self): """Count residues based on ``name``.""" return self.countChildren('name') def residueFreq(self): """Calculate residue frequency (based on ``name``).""" return self.freqChildren('name') def getDict(self): """See: ``Entity.getDict``.""" from_parent = self.parent.getDict() from_parent.update({'chain_id':self.getId()[0]}) return from_parent class Residue(MultiEntity): """The ``Residue`` instance contains ``Atom`` instances.""" def __init__(self, res_long_id, h_flag, seg_id, *args, **kwargs): self.level = 'R' self.seg_id = seg_id self.h_flag = h_flag self.res_id = res_long_id[1] #ID number self.res_ic = res_long_id[2] #ID long NAME MultiEntity.__init__(self, res_long_id, res_long_id[0], *args, **kwargs) def __repr__(self): res_name, res_id, res_ic = self.getId()[0] full_name = (res_name, res_id, res_ic) return "" % full_name def _getId(self): """Return the residue full id. ``(name, res_id, res_ic)``.""" return ((self.name, self.res_id, self.res_ic),) def _setId(self, id): """Set the residue id ``res_id``, name ``name`` and insertion code ``res_ic`` from a full id.""" (self.name, self.res_id, self.res_ic) = id[0] def removeHydrogens(self): """Remove hydrogen atoms.""" self.stripChildren(' H', 'eq', 'element', forgiving=False) def getSeg_id(self): """Return the segment id.""" return self.seg_id def setSeg_id(self, seg_id): """Set the segment id. This does not change the id.""" self.seg_id = seg_id def getIc(self): """Return the insertion code.""" return self.res_ic def setIc(self, res_ic): """Set the insertion code.""" self.res_ic = res_ic self.setId() def getRes_id(self): """Get the id.""" return self.res_id def setRes_id(self, res_id): """Set the id.""" self.res_id = res_id self.setId() def getH_flag(self): """Return the hetero flag.""" return self.h_flag def setH_flag(self, h_flag): """Sets the hetero flag. A valid flag is ' ' or 'H'. If 'H' the flag becomes part of the residue name i.e. H_XXX.""" if not h_flag in (' ', 'H'): raise AttributeError, "Only ' ' and 'H' hetero flags allowed." if len(self.name) == 3: self.name = "%s_%s" % (h_flag, self.name) elif len(self.name) == 5: self.name = "%s_%s" % (h_flag, self.name[2:]) else: raise ValueError, 'Non-standard residue name' self.h_flag = h_flag self.setId() def getDict(self): """See: ``Entity.getDict``.""" from_parent = self.parent.getDict() if self.h_flag != ' ': at_type = 'HETATM' else: at_type = 'ATOM ' from_parent.update({'at_type': at_type, 'h_flag': self.h_flag, 'res_name': self.name, 'res_long_id': self.getId()[0], 'res_id': self.res_id, 'res_ic': self.res_ic, 'seg_id': self.seg_id, }) return from_parent class Atom(Entity): """The ``Atom`` class contains no children.""" def __init__(self, at_long_id, at_name, ser_num, coords, occupancy, bfactor, element): self.level = 'A' self.index = HIERARCHY.index(self.level) self.coords = coords self.bfactor = bfactor self.occupancy = occupancy self.ser_num = ser_num self.at_id = at_long_id[0] self.alt_loc = at_long_id[1] self.table = dict([(level, {}) for level in HIERARCHY[self.index + 1:]]) self.element = element Entity.__init__(self, at_long_id, at_name) def __nonzero__(self): return bool(self.id) def __repr__(self): return "" % self.getId() def _getId(self): """Return the full id. The id of an atom is not its ' XX ' name but this string after left/right spaces striping. The full id is ``(at_id, alt_loc)``.""" return ((self.at_id, self.alt_loc),) def _setId(self, id): """Set the atom id ``at_id`` and alternate location ``alt_loc`` from a full id. See: ``_getId``.""" (self.at_id, self.alt_loc) = id[0] def setElement(self, element): """Set the atom element ``element``.""" self.element = element def setName(self, name): """Set name and update the id.""" self.name = name self.setAt_id(name.strip()) def setAt_id(self, at_id): """Set id. An atom id should be derived from the atom name. See: ``_getId``.""" self.at_id = at_id self.setId() def setAlt_loc(self, alt_loc): """Set alternate location identifier.""" self.alt_loc = alt_loc self.setId() def setSer_num(self, n): """Set serial number.""" self.ser_num = n def setBfactor(self, bfactor): """Set B-factor.""" self.bfactor = bfactor def setOccupancy(self, occupancy): """Set occupancy.""" self.occupancy = occupancy def setRadius(self, radius=None, radius_type=AREAIMOL_VDW_RADII, \ default_radius=DEFAULT_AREAIMOL_VDW_RADIUS): """Set radius, defaults to the AreaIMol VdW radius.""" if radius: self.radius = radius else: try: self.radius = radius_type[(self.parent.name, self.name)] except KeyError: self.radius = default_radius def getSer_num(self): """Return the serial number.""" return self.ser_num def getBfactor(self): """Return the B-factor.""" return self.bfactor def getOccupancy(self): """Return the occupancy.""" return self.occupancy def getRadius(self): """Return the radius.""" return self.radius def getDict(self): """See: ``Entity.getDict``.""" from_parent = self.parent.getDict() from_parent.update({'at_name': self.name, 'ser_num': self.ser_num, 'coords': self.coords, 'occupancy': self.occupancy, 'bfactor': self.bfactor, 'alt_loc': self.alt_loc, 'at_long_id': self.getId()[0], 'at_id': self.at_id, 'element': self.element}) return from_parent class Holder(MultiEntity): """The ``Holder`` instance exists outside the SMCRA hierarchy. Elements in a ``Holder`` instance are indexed by the full id.""" def __init__(self, name, *args): if not hasattr(self, 'level'): self.level = name MultiEntity.__init__(self, name, name, *args) def __repr__(self): return '' % (self.level, self.getName()) def addChild(self, child): """Add a child.""" child_id = child.getFull_id() self[child_id] = child def delChild(self, child_id): """Remove a child.""" self.pop(child_id) def updateIds(self): """Update self with children long ids.""" ids = [] for (id_, child) in self.iteritems(): new_id = child.getFull_id() if id_ != new_id: ids.append((id_, new_id)) for (old_id, new_id) in ids: child = self.pop(old_id) self.update(((new_id, child),)) class StructureHolder(Holder): """The ``StructureHolder`` contains ``Structure`` instances. See: ``Holder``.""" def __init__(self, *args): self.level = 'H' Holder.__init__(self, *args) def __repr__(self): return "" % self.getName() class ModelHolder(Holder): """The ``ModelHolder`` contains ``Model`` instances. See: ``Holder``.""" def __init__(self, *args): self.level = 'S' Holder.__init__(self, *args) def __repr__(self): return "" % self.getName() class ChainHolder(Holder): """The ``ChainHolder`` contains ``Chain`` instances. See: ``Holder``.""" def __init__(self, *args): self.level = 'M' Holder.__init__(self, *args) def __repr__(self): return "" % self.getName() class ResidueHolder(Holder): """The ``ResidueHolder`` contains ``Residue`` instances. See: ``Holder``.""" def __init__(self, *args): self.level = 'C' Holder.__init__(self, *args) def __repr__(self): return "" % self.getName() class AtomHolder(Holder): """The ``AtomHolder`` contains ``Atom`` instances. See: ``Holder``.""" def __init__(self, *args): self.level = 'R' Holder.__init__(self, *args) def __repr__(self): return "" % self.getName() class StructureBuilder(object): """Constructs a ``Structure`` object. The ``StructureBuilder`` class is used by a parser class to parse a file into a ``Structure`` object. An instance of a ``StructureBuilder`` has methods to create ``Entity`` instances and add them into the SMCRA hierarchy``.""" def __init__(self): self.structure = None def initStructure(self, structure_id): """Initialize a ``Structure`` instance.""" self.structure = Structure(structure_id) def initModel(self, model_id): """Initialize a ``Model`` instance and add it as a child to the ``Structure`` instance. If a model is defined twice a ``ConstructionError`` is raised.""" if not (model_id,) in self.structure: self.model = Model(model_id) self.model.junk = AtomHolder('junk') self.structure._initChild(self.model) else: raise ConstructionError def initChain(self, chain_id): """Initialize a ``Chain`` instance and add it as a child to the ``Model`` instance. If a chain is defined twice a ``ConstructionWarning`` is raised. This means that the model is not continuous.""" if not (chain_id,) in self.model: self.chain = Chain(chain_id) self.model._initChild(self.chain) else: self.chain = self.model[(chain_id,)] raise ConstructionWarning, "Chain %s is not continous" % chain_id def initSeg(self, seg_id): """Does not create an ``Entity`` instance, but updates the segment id, ``seg_id`` which is used to initialize ``Residue`` instances.""" self.seg_id = seg_id def initResidue(self, res_long_id, res_name): """Initialize a ``Residue`` instance and add it as a child to the ``Chain`` instance. If a residue is defined twice a ``ConstructionWarning`` is raised. This means that the chain is not continuous.""" if not (res_long_id,) in self.chain: self.residue = Residue(res_long_id, res_name, self.seg_id) self.chain._initChild(self.residue) else: self.residue = self.chain[(res_long_id,)] raise ConstructionWarning, "Residue %s%s%s is not continuous" % \ res_long_id def initAtom(self, at_long_id, at_name, ser_num, coord, occupancy, \ bfactor, element): """Initialize an ``Atom`` instance and add is as child to the ``Residue`` instance. If an atom is defined twice a ``ConstructionError`` is raised and the ``Atom`` instance is added to the ``structure.model.junk`` ``Holder`` instance.""" if not (at_long_id,) in self.residue: self.atom = Atom(at_long_id, at_name, ser_num, coord, occupancy, \ bfactor, element) self.residue._initChild(self.atom) else: full_id = (tuple(self.residue[(at_long_id,)].getFull_id()), \ ser_num) self.model.junk._initChild(Atom(full_id, at_name, ser_num, coord, \ occupancy, bfactor, element)) raise ConstructionError, 'Atom %s%s is defined twice.' % at_long_id def getStructure(self): """Update coordinates (``coords``), set the children-table (``table``) and return the ``Structure`` instance.""" self.structure.setTable() self.structure.setCoordsRecursively() return self.structure