def update(self, **kwargs): """Rebuild hierarchical view of atoms. This method is called at instantiation, but can be used to rebuild the hierarchical view when attributes of atoms change.""" array = np.array atoms = self._atoms if isinstance(atoms, AtomGroup): ag = atoms _indices = np.arange(ag._n_atoms) selstr = False else: ag = atoms.getAtomGroup() _indices = atoms._indices selstr = atoms.getSelstr() acsi = self._atoms.getACSIndex() n_atoms = len(ag) self._dict = _dict = dict() self._chains = _chains = list() self._residues = _residues = list() self._segments = _segments = list() segindex = -1 segindices = np.zeros(n_atoms, int) chindex = -1 chindices = np.zeros(n_atoms, int) resindex = -1 resindices = np.zeros(n_atoms, int) sgnms = ag._getSegnames() if sgnms is None: _segments = None else: if selstr: sgnms = sgnms[_indices] unique = np.unique(sgnms) s = sgnms[0] if len(unique) == 1: if s != '': segment = Segment(ag, _indices, acsi=acsi, unique=True, selstr=selstr) _dict[s] = segment _segments.append(segment) LOGGER.info('Hierarchical view contains segments.') else: _segments = None else: ps = None for i, s in enumerate(sgnms): if s == ps or s in _dict: continue ps = s segindex += 1 idx = _indices[i:][sgnms[i:] == s] segment = Segment(ag, idx, acsi=acsi, unique=True, selstr=selstr) segindices[idx] = segindex _dict[s] = segment _segments.append(segment) LOGGER.info('Hierarchical view contains segments.') chids = ag._getChids() if chids is None: _chains = None else: if selstr: chids = chids[_indices] if _segments is None: if len(np.unique(chids)) == 1: chain = Chain(ag, _indices, acsi=acsi, unique=True) _dict[(None, chids[0] or None)] = chain _chains.append(chain) else: pc = None for i, c in enumerate(chids): if c == pc or (None, c) in _dict: continue pc = c chindex += 1 idx = _indices[i:][chids[i:] == c] chain = Chain(ag, idx, acsi=acsi, unique=True) chindices[idx] = chindex _dict[(None, c)] = chain _chains.append(chain) else: pc = chids[0] ps = sgnms[0] _i = 0 for i, c in enumerate(chids): s = sgnms[i] if c == pc and s == ps: continue s_c = (ps, pc or None) chain = _dict.get(s_c) if chain is None: segment = _dict[ps] chindex += 1 idx = _indices[_i:i] chain = Chain(ag, idx, acsi=acsi, segment=segment, unique=True) chindices[idx] = chindex _dict[s_c] = chain segment._dict[pc] = len(segment._list) segment._list.append(chain) _chains.append(chain) else: idx = _indices[_i:i] chindices[idx] = chain._indices[0] chain._indices = np.concatenate((chain._indices, idx)) pc = c ps = s _i = i s_c = (ps, pc or None) chain = _dict.get(s_c) idx = _indices[_i:] if chain is None: segment = _dict[ps] chindex += 1 chindices[idx] = chindex chain = Chain(ag, idx, acsi=acsi, segment=segment, unique=True) _dict[s_c] = chain segment._dict[pc] = len(segment._list) segment._list.append(chain) _chains.append(chain) else: chindices[idx] = chain._indices[0] chain._indices = np.concatenate((chain._indices, idx)) if kwargs.get('chain') == True: return rnums = ag._getResnums() if rnums is None: raise ValueError('resnums are not set') if selstr: rnums = rnums[_indices] nones = None if _segments is None: if nones is None: nones = [None] * len(rnums) sgnms = nones if _chains is None: if nones is None: nones = [None] * len(rnums) chids = nones icods = ag._getIcodes() if icods is None: if nones is None: nones = [None] * len(rnums) icods = nones elif selstr: icods = icods[_indices] pr = rnums[0] pi = icods[0] or None pc = chids[0] ps = sgnms[0] or None _j = 0 for j, r in enumerate(rnums): i = icods[j] or None c = chids[j] or None s = sgnms[j] if r != pr or i != pi or c != pc or s != ps: s_c_r_i = (ps, pc, pr, pi) res = _dict.get(s_c_r_i) if res is None: chain = _dict.get((ps, pc)) resindex += 1 idx = _indices[_j:j] res = Residue(ag, idx, acsi=acsi, chain=chain, unique=True, selstr=selstr) resindices[idx] = resindex if chain is not None: chain._dict[(pr, pi)] = len(chain._list) chain._list.append(res) _residues.append(res) _dict[s_c_r_i] = res else: res._indices = np.concatenate((res._indices, _indices[_j:j])) ps = s pc = c pr = r pi = i _j = j s_c_r_i = (ps, pc, pr, pi) res = _dict.get(s_c_r_i) idx = _indices[_j:] if res is None: chain = _dict.get((ps, pc)) resindex += 1 res = Residue(ag, idx, acsi=acsi, chain=chain, unique=True, selstr=selstr) resindices[idx] = resindex if chain is not None: chain._dict[(pr, pi)] = len(chain._list) chain._list.append(res) _residues.append(res) _dict[s_c_r_i] = res else: resindices[idx] = res._indices[0] res._indices = np.concatenate((res._indices, idx)) ag._data['segindices'] = segindices ag._data['chindices'] = chindices ag._data['resindices'] = resindices
def update(self, **kwargs): """Update (or build) hierarchical view of atoms. This method is called at instantiation, but can be used to rebuild the hierarchical view when attributes of atoms change.""" atoms = self._atoms if isinstance(atoms, AtomGroup): ag = atoms selstr = False _indices = arange(atoms._n_atoms) else: ag = atoms.getAtomGroup() _indices = atoms._getIndices() selstr = atoms.getSelstr() set_indices = False if atoms == ag: set_indices = True _dict = dict() _residues = list() _segments = list() _chains = list() # also set the attributes, so that when no segments or chains # are available, num/iter methods won't raise exceptions self._dict = _dict self._residues = _residues self._segments = _segments self._chains = _chains acsi = atoms.getACSIndex() n_atoms = len(ag) # identify segments if set_indices: segindex = -1 segindices = zeros(n_atoms, int) sgnms = ag._getSegnames() if sgnms is None: _segments = None else: if selstr: sgnms = sgnms[_indices] s = sgnms[0] if len(unique(sgnms)) == 1: if s != '': segment = Segment(ag, _indices, acsi=acsi, unique=True, selstr=selstr) _dict[s] = segment _segments.append(segment) LOGGER.info('Hierarchical view contains segments.') else: _segments = None else: ps = None # previous segment name for i, s in enumerate(sgnms): if s == ps or s in _dict: continue ps = s idx = _indices[i:][sgnms[i:] == s] segment = Segment(ag, idx, acsi=acsi, unique=True, selstr=selstr) if set_indices: segindex += 1 segindices[idx] = segindex _dict[s] = segment _segments.append(segment) LOGGER.info('Hierarchical view contains segments.') if set_indices: ag._data['segindices'] = segindices chindex = -1 chindices = zeros(n_atoms, int) # identify chains chids = ag._getChids() if chids is None: _chains = None else: if selstr: chids = chids[_indices] if _segments is None: if len(unique(chids)) == 1: chain = Chain(ag, _indices, acsi=acsi, unique=True, selstr=selstr) _dict[(None, chids[0] or None)] = chain _chains.append(chain) else: pc = None for i, c in enumerate(chids): if c == pc or (None, c) in _dict: continue pc = c idx = _indices[i:][chids[i:] == c] chain = Chain(ag, idx, acsi=acsi, unique=True, selstr=selstr) if set_indices: chindex += 1 chindices[idx] = chindex _dict[(None, c)] = chain _chains.append(chain) else: pc = chids[0] ps = sgnms[0] _i = 0 for i, c in enumerate(chids): s = sgnms[i] if c == pc and s == ps: continue s_c = (ps, pc or None) chain = _dict.get(s_c) if chain is None: segment = _dict[ps] idx = _indices[_i:i] chain = Chain(ag, idx, acsi=acsi, segment=segment, unique=True, selstr=selstr) if set_indices: chindex += 1 chindices[idx] = chindex _dict[s_c] = chain _chains.append(chain) else: idx = _indices[_i:i] chindices[idx] = chain._indices[0] chain._indices = concatenate((chain._indices, idx)) pc = c ps = s _i = i s_c = (ps, pc or None) chain = _dict.get(s_c) idx = _indices[_i:] if chain is None: segment = _dict[ps] if set_indices: chindex += 1 chindices[idx] = chindex chain = Chain(ag, idx, acsi=acsi, segment=segment, unique=True, selstr=selstr) _dict[s_c] = chain _chains.append(chain) else: if set_indices: chindices[idx] = chain._indices[0] chain._indices = concatenate((chain._indices, idx)) if set_indices: ag._data['chindices'] = chindices if kwargs.get('chain') == True: return if set_indices: resindex = -1 resindices = zeros(n_atoms, int) rnums = ag._getResnums() if rnums is None: raise ValueError('resnums are not set') if selstr: rnums = rnums[_indices] nones = None if _segments is None: if nones is None: nones = [None] * len(rnums) sgnms = nones if _chains is None: if nones is None: nones = [None] * len(rnums) chids = nones icods = ag._getIcodes() if icods is None: if nones is None: nones = [None] * len(rnums) icods = nones elif selstr: icods = icods[_indices] pr = rnums[0] pi = icods[0] or None pc = chids[0] ps = sgnms[0] _j = 0 _append = _residues.append _get = _dict.get _set = _dict.__setitem__ for j, r in enumerate(rnums): i = icods[j] or None c = chids[j] s = sgnms[j] if r != pr or i != pi or c != pc or s != ps: s_c_r_i = (ps, pc, pr, pi) res = _get(s_c_r_i) if res is None: chain = _get((ps, pc)) idx = _indices[_j:j] res = Residue(ag, idx, acsi=acsi, chain=chain, unique=True, selstr=selstr) if set_indices: resindex += 1 resindices[idx] = resindex _append(res) _set(s_c_r_i, res) else: res._indices = concatenate((res._indices, _indices[_j:j])) ps = s pc = c pr = r pi = i _j = j s_c_r_i = (ps, pc, pr, pi) res = _get(s_c_r_i) idx = _indices[_j:] if res is None: chain = _get((ps, pc)) res = Residue(ag, idx, acsi=acsi, chain=chain, unique=True, selstr=selstr) if set_indices: resindex += 1 resindices[idx] = resindex _append(res) _set(s_c_r_i, res) else: if set_indices: resindices[idx] = res._indices[0] res._indices = concatenate((res._indices, idx)) if set_indices: ag._data['resindices'] = resindices