def writePQRStream(stream, atoms, **kwargs):
if isinstance(atoms, Atom):
atoms = Selection(atoms.getAtomGroup(), [atoms.getIndex()],
atoms.getACSIndex(),
'index ' + str(atoms.getIndex()))
n_atoms = atoms.numAtoms()
atomnames = atoms.getNames()
if atomnames is None:
raise RuntimeError('atom names are not set')
for i, an in enumerate(atomnames):
lenan = len(an)
if lenan < 4:
atomnames[i] = ' ' + an
elif lenan > 4:
atomnames[i] = an[:4]
s_or_u = np.array(['a']).dtype.char
resnames = atoms._getResnames()
if resnames is None:
resnames = ['UNK'] * n_atoms
resnums = atoms._getResnums()
if resnums is None:
resnums = np.ones(n_atoms, int)
chainids = atoms._getChids()
if chainids is None:
chainids = np.zeros(n_atoms, s_or_u + '1')
charges = atoms._getCharges()
if charges is None:
charges = np.zeros(n_atoms, float)
radii = atoms._getRadii()
if radii is None:
radii = np.zeros(n_atoms, float)
icodes = atoms._getIcodes()
if icodes is None:
icodes = np.zeros(n_atoms, s_or_u + '1')
hetero = ['ATOM'] * n_atoms
heteroflags = atoms._getFlags('hetatm')
if heteroflags is None:
heteroflags = atoms._getFlags('hetero')
if heteroflags is not None:
hetero = np.array(hetero, s_or_u + '6')
hetero[heteroflags] = 'HETATM'
altlocs = atoms._getAltlocs()
if altlocs is None:
altlocs = np.zeros(n_atoms, s_or_u + '1')
format = ('{0:6s} {1:5d} {2:4s} {3:1s}' +
'{4:4s} {5:1s} {6:4d} {7:1s} ' +
'{8:8.3f} {9:8.3f} {10:8.3f}' + '{11:8.4f} {12:7.4f}\n').format
coords = atoms._getCoords()
write = stream.write
for i, xyz in enumerate(coords):
write(
format(hetero[i], i + 1,
atomnames[i], altlocs[i], resnames[i], chainids[i],
int(resnums[i]), icodes[i], xyz[0], xyz[1], xyz[2],
charges[i], radii[i]))
def sliceVector(vector, atoms, select):
"""Return part of the *vector* for *atoms* matching *select*. Note that
returned :class:`.Vector` instance is not normalized.
:arg vector: vector instance to be sliced
:type vector: :class:`.VectorBase`
:arg atoms: atoms for which *vector* describes a deformation, motion, etc.
:type atoms: :class:`.Atomic`
:arg select: an atom selection or a selection string
:type select: :class:`.Selection`, str
:returns: (:class:`.Vector`, :class:`.Selection`)"""
if not isinstance(vector, VectorBase):
raise TypeError('vector must be a VectorBase instance, not {0}'
.format(type(vector)))
if not isinstance(atoms, Atomic):
raise TypeError('atoms must be an Atomic instance, not {0}'
.format(type(atoms)))
if atoms.numAtoms() != vector.numAtoms():
raise ValueError('number of atoms in model and atoms must be equal')
if isinstance(select, str):
selstr = select
if isinstance(atoms, AtomGroup):
sel = atoms.select(selstr)
which = sel._getIndices()
else:
which = SELECT.getIndices(atoms, selstr)
sel = Selection(atoms.getAtomGroup(), atoms.getIndices()[which],
selstr, atoms.getACSIndex())
elif isinstance(select, AtomSubset):
sel = select
if isinstance(atoms, AtomGroup):
if sel.getAtomGroup() != atoms:
raise ValueError('select and atoms do not match')
which = sel._getIndices()
else:
if atoms.getAtomGroup() != sel.getAtomGroup():
raise ValueError('select and atoms do not match')
elif not sel in atoms:
raise ValueError('select is not a subset of atoms')
idxset = set(atoms._getIndices())
which = np.array([idx in idxset for idx in sel._getIndices()])
which = which.nonzero()[0]
selstr = sel.getSelstr()
else:
raise TypeError('select must be a string or a Selection instance')
vec = Vector(vector.getArrayNx3()[
which, :].flatten(),
'{0} slice {1}'.format(str(vector), repr(selstr)),
vector.is3d())
return (vec, sel)
def __call__(self, radius, center):
"""Select atoms radius *radius* (Å) of *center*, which can be point(s)
in 3-d space (:class:`numpy.ndarray` with shape ``(n_atoms, 3)``) or a
set of atoms, e.g. :class:`.Selection`."""
try:
center = center._getCoords()
except AttributeError:
try:
ndim, shape = center.ndim, center.shape
except AttributeError:
raise TypeError('center must be an Atomic instance or a'
'coordinate array')
else:
if shape == (3, ):
center = [center]
elif not ndim == 2 and shape[1] == 3:
raise ValueError('center.shape must be (n_atoms, 3) or'
'(3,)')
else:
if center is None:
raise ValueError('center does not have coordinate data')
search = self._kdtree.search
get_indices = self._kdtree.getIndices
get_count = self._kdtree.getCount
indices = set()
update = indices.update
radius = float(radius)
for xyz in center:
search(radius, xyz)
if get_count():
update(get_indices())
indices = list(indices)
if indices:
indices.sort()
if self._ag is None:
return array(indices)
else:
if self._indices is not None:
indices = self._indices[indices]
return Selection(self._ag,
array(indices),
'index ' + rangeString(indices),
acsi=self._acsi,
unique=True)
def writePQR(filename, atoms):
"""Write *atoms* in PQR format to a file with name *filename*. Only
current coordinate set is written. Returns *filename* upon success. If
*filename* ends with :file:`.gz`, a compressed file will be written."""
if not isinstance(atoms, Atomic):
raise TypeError('atoms does not have a valid type')
if isinstance(atoms, Atom):
atoms = Selection(atoms.getAtomGroup(), [atoms.getIndex()],
atoms.getACSIndex(),
'index ' + str(atoms.getIndex()))
stream = openFile(filename, 'w')
n_atoms = atoms.numAtoms()
atomnames = atoms.getNames()
if atomnames is None:
raise RuntimeError('atom names are not set')
for i, an in enumerate(atomnames):
lenan = len(an)
if lenan < 4:
atomnames[i] = ' ' + an
elif lenan > 4:
atomnames[i] = an[:4]
s_or_u = np.array(['a']).dtype.char
resnames = atoms._getResnames()
if resnames is None:
resnames = ['UNK'] * n_atoms
resnums = atoms._getResnums()
if resnums is None:
resnums = np.ones(n_atoms, int)
chainids = atoms._getChids()
if chainids is None:
chainids = np.zeros(n_atoms, s_or_u + '1')
charges = atoms._getCharges()
if charges is None:
charges = np.zeros(n_atoms, float)
radii = atoms._getRadii()
if radii is None:
radii = np.zeros(n_atoms, float)
icodes = atoms._getIcodes()
if icodes is None:
icodes = np.zeros(n_atoms, s_or_u + '1')
hetero = ['ATOM'] * n_atoms
heteroflags = atoms._getFlags('hetatm')
if heteroflags is None:
heteroflags = atoms._getFlags('hetero')
if heteroflags is not None:
hetero = np.array(hetero, s_or_u + '6')
hetero[heteroflags] = 'HETATM'
altlocs = atoms._getAltlocs()
if altlocs is None:
altlocs = np.zeros(n_atoms, s_or_u + '1')
format = ('{0:6s}{1:5d} {2:4s}{3:1s}' + '{4:4s}{5:1s}{6:4d}{7:1s} ' +
'{8:8.3f}{9:8.3f}{10:8.3f}' + '{11:8.4f}{12:7.4f}\n').format
coords = atoms._getCoords()
write = stream.write
for i, xyz in enumerate(coords):
write(
format(hetero[i], i + 1,
atomnames[i], altlocs[i], resnames[i], chainids[i],
int(resnums[i]), icodes[i], xyz[0], xyz[1], xyz[2],
charges[i], radii[i]))
write('TER\nEND')
stream.close()
return filename
def reduceModel(model, atoms, select):
"""Return reduced NMA model. Reduces a :class:`.NMA` model to a subset of
*atoms* matching *select*. This function behaves differently depending on
the type of the *model* argument. For :class:`.ANM` and :class:`.GNM` or
other :class:`.NMA` models, force constant matrix for system of interest
(specified by the *select*) is derived from the force constant matrix for
the *model* by assuming that for any given displacement of the system of
interest, other atoms move along in such a way as to minimize the potential
energy. This is based on the formulation in [KH00]_. For :class:`.PCA`
models, this function simply takes the sub-covariance matrix for selection.
.. [KH00] Konrad H, Andrei-Jose P, Serge D, Marie-Claire BF, Gerald RK.
Harmonicity in slow protein dynamics. *Chem Phys* **2000** 261:25-37.
:arg model: dynamics model
:type model: :class:`.ANM`, :class:`.GNM`, or :class:`.PCA`
:arg atoms: atoms that were used to build the model
:type atoms: :class:`.Atomic`
:arg select: an atom selection or a selection string
:type select: :class:`.Selection`, str
:returns: (:class:`.NMA`, :class:`.Selection`)"""
linalg = importLA()
if not isinstance(model, NMA):
raise TypeError('model must be an NMA instance, not {0}'
.format(type(model)))
if not isinstance(atoms, (AtomGroup, AtomSubset, AtomMap)):
raise TypeError('atoms type is not valid')
if len(atoms) <= 1:
raise TypeError('atoms must contain more than 1 atoms')
if isinstance(model, GNM):
matrix = model._kirchhoff
elif isinstance(model, ANM):
matrix = model._hessian
elif isinstance(model, PCA):
matrix = model._cov
else:
raise TypeError('model does not have a valid type derived from NMA')
if matrix is None:
raise ValueError('model matrix (Hessian/Kirchhoff/Covariance) is not '
'built')
if isinstance(select, str):
system = SELECT.getBoolArray(atoms, select)
n_sel = sum(system)
if n_sel == 0:
raise ValueError('select matches 0 atoms')
if len(atoms) == n_sel:
raise ValueError('select matches all atoms')
if isinstance(atoms, AtomGroup):
ag = atoms
which = np.arange(len(atoms))[system]
else:
ag = atoms.getAtomGroup()
which = atoms._getIndices()[system]
sel = Selection(ag, which, select, atoms.getACSIndex())
elif isinstance(select, AtomSubset):
sel = select
if isinstance(atoms, AtomGroup):
if sel.getAtomGroup() != atoms:
raise ValueError('select and atoms do not match')
system = np.zeros(len(atoms), bool)
system[sel._getIndices()] = True
else:
if atoms.getAtomGroup() != sel.getAtomGroup():
raise ValueError('select and atoms do not match')
elif not sel in atoms:
raise ValueError('select is not a subset of atoms')
idxset = set(atoms._getIndices())
system = np.array([idx in idxset for idx in sel._getIndices()])
else:
raise TypeError('select must be a string or a Selection instance')
other = np.invert(system)
if model.is3d():
system = np.tile(system, (3, 1)).transpose().flatten()
other = np.tile(other, (3, 1)).transpose().flatten()
ss = matrix[system, :][:, system]
if isinstance(model, PCA):
eda = PCA(model.getTitle() + ' reduced')
eda.setCovariance(ss)
return eda, system
so = matrix[system, :][:, other]
os = matrix[other, :][:, system]
oo = matrix[other, :][:, other]
matrix = ss - np.dot(so, np.dot(linalg.inv(oo), os))
if isinstance(model, GNM):
gnm = GNM(model.getTitle() + ' reduced')
gnm.setKirchhoff(matrix)
return gnm, sel
elif isinstance(model, ANM):
anm = ANM(model.getTitle() + ' reduced')
anm.setHessian(matrix)
return anm, sel
elif isinstance(model, PCA):
eda = PCA(model.getTitle() + ' reduced')
eda.setCovariance(matrix)
return eda, sel
def sliceModel(model, atoms, select):
"""Return a part of the *model* for *atoms* matching *select*. Note that
normal modes (eigenvectors) are not normalized.
:arg mode: NMA model instance to be sliced
:type mode: :class:`.NMA`
:arg atoms: atoms for which the *model* was built
:type atoms: :class:`.Atomic`
:arg select: an atom selection or a selection string
:type select: :class:`.Selection`, str
:returns: (:class:`.NMA`, :class:`.Selection`)"""
if not isinstance(model, NMA):
raise TypeError('mode must be a NMA instance, not {0}'
.format(type(model)))
if not isinstance(atoms, Atomic):
raise TypeError('atoms must be an Atomic instance, not {0}'
.format(type(atoms)))
if atoms.numAtoms() != model.numAtoms():
raise ValueError('number of atoms in model and atoms must be equal')
array = model._getArray()
if isinstance(select, str):
selstr = select
if isinstance(atoms, AtomGroup):
sel = atoms.select(selstr)
which = sel.getIndices()
else:
which = SELECT.getIndices(atoms, selstr)
sel = Selection(atoms.getAtomGroup(), atoms.getIndices()[which],
selstr, atoms.getACSIndex())
elif isinstance(select, AtomSubset):
sel = select
if isinstance(atoms, AtomGroup):
if sel.getAtomGroup() != atoms:
raise ValueError('select and atoms do not match')
which = sel._getIndices()
else:
if atoms.getAtomGroup() != sel.getAtomGroup():
raise ValueError('select and atoms do not match')
elif not sel in atoms:
raise ValueError('select is not a subset of atoms')
idxset = set(atoms._getIndices())
which = np.array([idx in idxset for idx in sel._getIndices()])
which = which.nonzero()[0]
selstr = sel.getSelstr()
else:
raise TypeError('select must be a string or a Selection instance')
nma = type(model)('{0} slice {1}'
.format(model.getTitle(), repr(selstr)))
if model.is3d():
which = [which.reshape((len(which), 1))*3]
which.append(which[0]+1)
which.append(which[0]+2)
which = np.concatenate(which, 1).flatten()
nma.setEigens(array[which, :], model.getEigvals())
return (nma, sel)
def sliceMode(mode, atoms, select):
"""Return part of the *mode* for *atoms* matching *select*. This works
slightly different from :func:`.sliceVector`. Mode array (eigenvector) is
multiplied by square-root of the variance along the mode. If mode is from
an elastic network model, variance is defined as the inverse of the
eigenvalue. Note that returned :class:`.Vector` instance is not
normalized.
:arg mode: mode instance to be sliced
:type mode: :class:`.Mode`
:arg atoms: atoms for which *mode* describes a deformation, motion, etc.
:type atoms: :class:`.Atomic`
:arg select: an atom selection or a selection string
:type select: :class:`.Selection`, str
:returns: (:class:`.Vector`, :class:`.Selection`)"""
if not isinstance(mode, Mode):
raise TypeError('mode must be a Mode instance, not {0}'
.format(type(mode)))
if not isinstance(atoms, Atomic):
raise TypeError('atoms must be an Atomic instance, not {0}'
.format(type(atoms)))
if atoms.numAtoms() != mode.numAtoms():
raise ValueError('number of atoms in model and atoms must be equal')
if isinstance(select, str):
selstr = select
if isinstance(atoms, AtomGroup):
sel = atoms.select(selstr)
which = sel._getIndices()
else:
which = SELECT.getIndices(atoms, selstr)
sel = Selection(atoms.getAtomGroup(), atoms.getIndices()[which],
selstr, atoms.getACSIndex())
elif isinstance(select, AtomSubset):
sel = select
if isinstance(atoms, AtomGroup):
if sel.getAtomGroup() != atoms:
raise ValueError('select and atoms do not match')
which = sel._getIndices()
else:
if atoms.getAtomGroup() != sel.getAtomGroup():
raise ValueError('select and atoms do not match')
elif not sel in atoms:
raise ValueError('select is not a subset of atoms')
idxset = set(atoms._getIndices())
which = np.array([idx in idxset for idx in sel._getIndices()])
which = which.nonzero()[0]
selstr = sel.getSelstr()
else:
raise TypeError('select must be a string or a Selection instance')
vec = Vector(mode.getArrayNx3()[which, :].flatten() *
mode.getVariance()**0.5,
'{0} slice {1}'.format(str(mode), repr(selstr)), mode.is3d())
return (vec, sel)
def trimPDBEnsemble(pdb_ensemble, **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.
**Occupancy**
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``
:type occupancy: float
:arg selstr: the function will trim residues that are NOT specified by
the selection string.
:type selstr: str
:arg hard: hard trimming or soft trimming. If set to `False`, *pdb_ensemble*
will be trimmed by selection. This is useful for example when one uses
:func:`calcEnsembleENMs` and :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
"""
atoms = pdb_ensemble.getAtoms()
selstr = kwargs.pop('selstr', None)
occupancy = kwargs.pop('occupancy', None)
hard = kwargs.pop('hard', False) or atoms is None
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
elif selstr is not None:
atoms = pdb_ensemble.getAtoms()
assert atoms is not None, 'atoms are empty'
selector = Select()
torf = selector.getBoolArray(atoms, selstr)
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]
if type(atoms) is Chain:
trim_atoms = Chain(atoms.copy(), trim_atoms_idx, atoms._hv)
elif type(atoms) is AtomGroup:
trim_atoms = AtomMap(atoms, trim_atoms_idx)
else:
trim_atoms = AtomMap(atoms.copy(), 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()
trimmed.addCoordset(confs[:, torf], weights[:, torf], labels)
else:
indices = np.where(torf)
selids = pdb_ensemble._indices
if selids is not None:
ag = atoms.getAtomGroup()
indices = selids[indices]
else:
ag = atoms.copy()
selstr = '' if selstr is None else selstr
select = Selection(ag, indices, selstr, ag._acsi)
trimmed.setAtoms(ag)
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()
trimmed.addCoordset(confs, weights, labels)
trimmed.setAtoms(select)
return trimmed