def __getitem__(self, index): """Returns a conformation at given index.""" labels = self._labels if isinstance(index, Integral): return self.getConformation(index) elif isinstance(index, slice): ens = PDBEnsemble('{0} ({1[0]}:{1[1]}:{1[2]})'.format( self._title, index.indices(len(self)))) ens.setCoords(copy(self._coords)) ens.addCoordset(self._confs[index].copy(), self._weights[index].copy(), label=self._labels[index]) if self._trans is not None: ens._trans = self._trans[index] ens.setAtoms(self._atoms) ens._indices = self._indices for key in self._data.keys(): ens._data[key] = self._data[key][index].copy() return ens elif isinstance(index, (list, np.ndarray)): index2 = list(index) for i in range(len(index)): if isinstance(index[i], str): try: index2[i] = labels.index(index[i]) except ValueError: raise IndexError('invalid label: %s' % index[i]) ens = PDBEnsemble('{0}'.format(self._title)) ens.setCoords(copy(self._coords)) labels = list(np.array(self._labels)[index2]) ens.addCoordset(self._confs[index2].copy(), self._weights[index2].copy(), label=labels) if self._trans is not None: ens._trans = self._trans[index2] ens.setAtoms(self._atoms) ens._indices = self._indices for key in self._data.keys(): ens._data[key] = self._data[key][index].copy() return ens elif isinstance(index, str): try: i = labels.index(index) return self.getConformation(i) except ValueError: raise IndexError('invalid label: %s' % index) else: raise IndexError('invalid index')
def setCoords(self, coords): """Set *coords* as the ensemble reference coordinate set. *coords* may be an array with suitable data type, shape, and dimensionality, or an object with :meth:`getCoords` method.""" atoms = coords try: if isinstance(coords, Ensemble): coords = copy(coords._coords) else: coords = coords.getCoords() except AttributeError: pass finally: if coords is None: raise ValueError('coordinates of {0} are not set'.format( str(atoms))) try: checkCoords(coords, natoms=self._n_atoms) except TypeError: raise TypeError('coords must be a numpy array or an object ' 'with `getCoords` method') self._coords = coords self._n_atoms = coords.shape[0] if isinstance(atoms, Ensemble): self._indices = atoms._indices self._atoms = atoms._atoms
def __getitem__(self, index): """Returns a conformation at given index.""" if self._confs is None: return None # use sth like follows # which = arange(self._n_csets)[index].nonzero()[0] # if len(which) == 1: # return getConf... # else: # return SubEnsemble if isinstance(index, Integral): return self.getConformation(index) elif isinstance(index, slice): ens = Ensemble('{0} ({1[0]}:{1[1]}:{1[2]})'.format( self._title, index.indices(len(self)))) ens.setCoords(copy(self._coords)) for key in self._data.keys(): ens._data[key] = self._data[key][index].copy() ens.addCoordset(self._confs[index].copy()) if self._weights is not None: ens.setWeights(self._weights.copy()) ens.setAtoms(self._atoms) ens._indices = self._indices return ens elif isinstance(index, (list, ndarray)): ens = Ensemble('{0}'.format(self._title)) ens.setCoords(copy(self._coords)) for key in self._data.keys(): ens._data[key] = self._data[key][index].copy() ens.addCoordset(self._confs[index].copy()) if self._weights is not None: ens.setWeights(self._weights.copy()) ens.setAtoms(self._atoms) ens._indices = self._indices return ens else: raise IndexError('invalid index')
def __add__(self, other): """Concatenate two ensembles. The reference coordinates of *self* is used in the result.""" if not isinstance(other, Ensemble): raise TypeError('an Ensemble instance cannot be added to an {0} ' 'instance'.format(type(other))) elif self.numAtoms() != other.numAtoms(): raise ValueError('Ensembles must have same number of atoms.') ensemble = PDBEnsemble('{0} + {1}'.format(self.getTitle(), other.getTitle())) ensemble.setCoords(copy(self._coords)) weights = copy(self._weights) if self._confs is not None: ensemble.addCoordset(copy(self._confs), weights=weights, label=self.getLabels()) other_weights = copy(other._weights) ensemble.addCoordset(copy(other._confs), weights=other_weights, label=other.getLabels()) if self._atoms is not None: ensemble.setAtoms(self._atoms) ensemble._indices = self._indices else: ensemble.setAtoms(other._atoms) ensemble._indices = other._indices all_keys = list(self._data.keys()) + list(other._data.keys()) for key in all_keys: if key in self._data and key in other._data: self_data = self._data[key] other_data = other._data[key] elif key in self._data: self_data = self._data[key] other_data = np.zeros(other.numConfs(), dtype=self_data.dtype) elif key in other._data: other_data = other._data[key] self_data = np.zeros(other.numConfs(), dtype=other_data.dtype) ensemble._data[key] = np.concatenate((self_data, other_data), axis=0) return ensemble
def getIndices(self): """Returns a copy of indices of selected columns""" return copy(self._indices)
def iterpose(self, rmsd=0.0001): confs = copy(self._confs) Ensemble.iterpose(self, rmsd) self._confs = confs LOGGER.info('Final superposition to calculate transformations.') self.superpose()
def trimPDBEnsemble(pdb_ensemble, occupancy=None, **kwargs): """Returns a new PDB ensemble obtained by trimming given *pdb_ensemble*. This function helps selecting atoms in a pdb ensemble based on one of the following criteria, and returns them in a new :class:`.PDBEnsemble` instance. Resulting PDB ensemble will contain atoms whose occupancies are greater or equal to *occupancy* keyword argument. Occupancies for atoms will be calculated using ``calcOccupancies(pdb_ensemble, normed=True)``. :arg occupancy: occupancy for selecting atoms, must satisfy ``0 < occupancy <= 1``. If set to *None* then *hard* trimming will be performed. :type occupancy: float :arg hard: Whether to perform hard trimming. Default is **False** If set to **True**, atoms will be completely removed from *pdb_ensemble*. If set to **False**, a soft trimming of *pdb_ensemble* will be done where atoms will be removed from the active selection. This is useful, for example, when one uses :func:`calcEnsembleENMs` together with :func:`sliceModel` or :func:`reduceModel` to calculate the modes from the remaining part while still taking the removed part into consideration (e.g. as the environment). :type hard: bool """ hard = kwargs.pop('hard', False) or pdb_ensemble._atoms is None \ or occupancy is None atoms = pdb_ensemble.getAtoms(selected=hard) if not isinstance(pdb_ensemble, PDBEnsemble): raise TypeError('pdb_ensemble argument must be a PDBEnsemble') if pdb_ensemble.numConfs() == 0 or pdb_ensemble.numAtoms() == 0: raise ValueError('pdb_ensemble must have conformations') if occupancy is not None: occupancy = float(occupancy) assert 0 < occupancy <= 1, ('occupancy is not > 0 and <= 1: ' '{0}'.format(repr(occupancy))) n_confs = pdb_ensemble.numConfs() assert n_confs > 0, 'pdb_ensemble does not contain any conformations' occupancies = calcOccupancies(pdb_ensemble, normed=True) #assert weights is not None, 'weights must be set for pdb_ensemble' #weights = weights.flatten() #mean_weights = weights / n_confs torf = occupancies >= occupancy else: n_atoms = pdb_ensemble.getCoords().shape[0] torf = np.ones(n_atoms, dtype=bool) trimmed = PDBEnsemble(pdb_ensemble.getTitle()) if hard: if atoms is not None: trim_atoms_idx = [n for n, t in enumerate(torf) if t] trim_atoms = atoms[trim_atoms_idx] trimmed.setAtoms(trim_atoms) coords = pdb_ensemble.getCoords() if coords is not None: trimmed.setCoords(coords[torf]) confs = pdb_ensemble.getCoordsets() if confs is not None: weights = pdb_ensemble.getWeights() labels = pdb_ensemble.getLabels() msa = pdb_ensemble.getMSA() if msa: msa = msa[:, torf] trimmed.addCoordset(confs[:, torf], weights[:, torf], labels, sequence=msa) else: indices = np.where(torf)[0] selids = pdb_ensemble._indices if selids is not None: indices = selids[indices] select = atoms[indices] trimmed.setAtoms(atoms) trimmed.setAtoms(select) coords = copy(pdb_ensemble._coords) if coords is not None: trimmed.setCoords(coords) confs = copy(pdb_ensemble._confs) if confs is not None: weights = copy(pdb_ensemble._weights) labels = pdb_ensemble.getLabels() msa = pdb_ensemble._msa trimmed.addCoordset(confs, weights, labels, sequence=msa) trimmed.setAtoms(select) trimmed._data = pdb_ensemble._data return trimmed