def addCoordset(self, coords):
"""Add coordinate set(s) to the ensemble. *coords* must be a Numpy
array with suitable data type, shape and dimensionality, or an object
with :meth:`getCoordsets` method."""
n_atoms = self._n_atoms
n_select = self.numSelected()
try:
if self._coords is not None:
if isinstance(coords, Ensemble):
coords = coords._getCoordsets(selected=False)
elif hasattr(coords, '_getCoordsets'):
coords = coords._getCoordsets()
elif hasattr(coords, '_getCoords'):
coords = coords._getCoords()
coords = array(coords)
else:
if isinstance(coords, Ensemble):
coords = coords.getCoordsets(selected=False)
elif hasattr(coords, 'getCoordsets'):
coords = coords.getCoordsets()
elif hasattr(coords, 'getCoords'):
coords = coords.getCoords()
coords = array(coords)
except AttributeError:
pass
else:
if coords is None:
raise ValueError('coordinates are not set')
try:
checkCoords(coords, csets=True, natoms=n_atoms)
except:
try:
checkCoords(coords, csets=True, natoms=n_select)
except TypeError:
raise TypeError('coords must be a numpy array or an object '
'with `getCoords` method')
if coords.ndim == 2:
n_nodes, _ = coords.shape
coords = coords.reshape((1, n_nodes, 3))
n_confs = 1
else:
n_confs, n_nodes, _ = coords.shape
if not n_atoms:
self._n_atoms = n_atoms = n_nodes
if n_nodes == n_select and self.isSelected():
full_coords = repeat(self._coords[newaxis, :, :], n_confs, axis=0)
full_coords[:, self._indices, :] = coords
coords = full_coords
if self._confs is None:
self._confs = coords
else:
self._confs = concatenate((self._confs, coords), axis=0)
self._n_csets += n_confs
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 calcCenter(atoms, weights=None):
"""Returns geometric center of *atoms*. If *weights* is given it must
be a flat array with length equal to number of atoms. Mass center of
atoms can be calculated by setting weights equal to atom masses, i.e.
``weights=atoms.getMasses()``."""
try:
coords = atoms._getCoords()
except AttributeError:
try:
coords = atoms.getCoords()
except AttributeError:
coords = atoms
try:
checkCoords(coords, csets=True, dtype=None, name='atoms')
except TypeError:
raise TypeError('atoms must be an Atomic instance')
if weights is not None:
try:
ndim, shape = weights.ndim, weights.shape
except AttributeError:
raise TypeError('weights must be a numpy array')
else:
if shape[0] != coords.shape[-2]:
raise ValueError('weights.shape[0] must be equal to number of '
'atoms')
if ndim != 2:
try:
weights = weights.reshape((shape[0], 1))
except ValueError:
raise ValueError('weights.shape must be a (n_atoms, 1)')
return getCenter(coords, weights)
def setCoords(self, coords):
"""Set *coords* as the trajectory reference coordinate set. *coords*
must be an object with :meth:`getCoords` method, or a Numpy array with
suitable data type, shape, and dimensionality."""
atoms = coords
try:
coords = atoms.getCoords()
except AttributeError:
pass
else:
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
def addCoordset(self, coords, weights=None, label=None, **kwargs):
"""Add coordinate set(s) to the ensemble. *coords* must be a Numpy
array with suitable shape and dimensionality, or an object with
:meth:`getCoordsets`. *weights* is an optional argument.
If provided, its length must match number of atoms. Weights of
missing (not resolved) atoms must be ``0`` and weights of those
that are resolved can be anything greater than ``0``. If not
provided, weights of all atoms for this coordinate set will be
set equal to ``1``. *label*, which may be a PDB identifier or a
list of identifiers, is used to label conformations."""
degeneracy = kwargs.pop('degeneracy', False)
atoms = coords
n_atoms = self._n_atoms
n_select = self.numSelected()
n_confs = self.numCoordsets()
try:
if degeneracy:
if self._coords is not None:
if isinstance(coords, Ensemble):
coords = coords._getCoords(selected=False)
elif hasattr(coords, '_getCoords'):
coords = coords._getCoords()
else:
if isinstance(coords, Ensemble):
coords = coords.getCoords(selected=False)
elif hasattr(coords, 'getCoords'):
coords = coords.getCoords()
else:
if self._coords is not None:
if isinstance(coords, Ensemble):
coords = coords._getCoordsets(selected=False)
elif hasattr(coords, '_getCoordsets'):
coords = coords._getCoordsets()
else:
if isinstance(coords, Ensemble):
coords = coords.getCoordsets(selected=False)
elif hasattr(coords, 'getCoordsets'):
coords = coords.getCoordsets()
except AttributeError:
label = label or 'Unknown'
else:
if coords is None:
raise ValueError('coordinates are not set')
elif label is None and isinstance(atoms, Atomic):
if not isinstance(atoms, AtomGroup):
ag = atoms.getAtomGroup()
else:
ag = atoms
label = ag.getTitle()
if coords.shape[0] < ag.numCoordsets():
label += '_m' + str(atoms.getACSIndex())
else:
label = label or 'Unknown'
# check coordinates
try:
checkCoords(coords, csets=True, natoms=n_atoms)
except:
try:
checkCoords(coords, csets=True, natoms=n_select)
except TypeError:
raise TypeError('coords must be a numpy array or an object '
'with `getCoords` method')
if coords.ndim == 2:
n_nodes, _ = coords.shape
coords = coords.reshape((1, n_nodes, 3))
n_csets = 1
else:
n_csets, n_nodes, _ = coords.shape
if degeneracy:
coords = coords[:1]
n_repeats = 1 if degeneracy else n_csets
if not n_atoms:
self._n_atoms = n_nodes
if n_nodes == n_select and self.isSelected():
full_coords = np.repeat(self._coords[np.newaxis, :, :],
n_csets,
axis=0)
full_coords[:, self._indices, :] = coords
coords = full_coords
# check weights
if weights is None:
weights = np.ones((n_csets, n_atoms, 1), dtype=float)
else:
weights = checkWeights(weights, n_atoms, n_csets)
if degeneracy:
weights = weights[:1]
# assign new values
# update labels
if n_csets > 1 and not degeneracy:
if isinstance(label, str):
labels = [
'{0}_m{1}'.format(label, i + 1) for i in range(n_csets)
]
else:
if len(label) != n_csets:
raise ValueError('length of label and number of '
'coordinate sets must be the same')
labels = label
else:
labels = [label] if np.isscalar(label) else label
self._labels.extend(labels)
# update coordinates
if self._confs is None and self._weights is None:
self._confs = coords
self._weights = weights
self._n_csets = n_repeats
elif self._confs is not None and self._weights is not None:
self._confs = np.concatenate((self._confs, coords), axis=0)
self._weights = np.concatenate((self._weights, weights), axis=0)
self._n_csets += n_repeats
else:
raise RuntimeError('_confs and _weights must be set or None at '
'the same time')
def write(self, coords, unitcell=None, **kwargs):
"""Write *coords* to a file open in 'a' or 'w' mode. *coords* may be
a NUmpy array or a ProDy object that stores or points to coordinate
data. Note that all coordinate sets of ProDy object will be written.
Number of atoms will be determined from the file or based on the size
of the first coordinate set written. If *unitcell* is provided for
the first coordinate set, it will be expected for the following
coordinate sets as well. If *coords* is an :class:`~.Atomic` or
:class:`~.Ensemble` all coordinate sets will be written.
Following keywords are used when writing the first coordinate set:
:arg timestep: timestep used for integration, default is 1
:arg firsttimestep: number of the first timestep, default is 0
:arg framefreq: number of timesteps between frames, default is 1"""
if self._closed:
raise ValueError('I/O operation on closed file')
if self._mode == 'r':
raise IOError('File not open for writing')
try:
coords = coords._getCoordsets()
except AttributeError:
try:
coords = coords._getCoords()
except AttributeError:
checkCoords(coords, csets=True, dtype=None)
else:
if unitcell is None:
try:
coords = coords.getUnitcell()
except AttributeError:
pass
if coords.dtype != float32:
coords = coords.astype(float32)
n_atoms = coords.shape[-2]
if self._n_atoms == 0:
self._n_atoms = n_atoms
elif self._n_atoms != n_atoms:
raise ValueError('coords does not have correct number of atoms')
if coords.ndim == 2:
coords = [coords]
dcd = self._file
pack_i_4N = pack('i', self._n_atoms * 4)
if self._n_csets == 0:
if unitcell is None:
self._unitcell = False
else:
self._unitcell = True
timestep = float(kwargs.get('timestep', 1.0))
first_ts = int(kwargs.get('firsttimestep', 0))
framefreq = int(kwargs.get('framefreq', 1))
n_fixed = 0
pack_i_0 = pack(b'i', 0)
pack_ix4_0x4 = pack(b'i' * 4, 0, 0, 0, 0)
pack_i_1 = pack(b'i', 1)
pack_i_2 = pack(b'i', 2)
pack_i_4 = pack(b'i', 4)
pack_i_84 = pack(b'i', 84)
pack_i_164 = pack(b'i', 164)
dcd.write(pack_i_84)
dcd.write(b'CORD')
dcd.write(pack_i_0) # 0 Number of frames in file, none written yet
dcd.write(pack(b'i', first_ts)) # 1 Starting timestep
dcd.write(pack(b'i', framefreq)) # 2 Timesteps between frames
dcd.write(pack_i_0) # 3 Number of timesteps in simulation
dcd.write(pack_i_0) # 4 NAMD writes NSTEP or ISTART - NSAVC here?
dcd.write(pack_ix4_0x4) # 5, 6, 7, 8
dcd.write(pack('f', timestep)) # 9 timestep
dcd.write(pack('i', int(self._unitcell))) # 10 with unitcell
dcd.write(pack_ix4_0x4) # 11, 12, 13, 14
dcd.write(pack_ix4_0x4) # 15, 16, 17, 18
dcd.write(pack('i', 24)) # 19 Pretend to be CHARMM version 24
dcd.write(pack_i_84)
dcd.write(pack_i_164)
dcd.write(pack_i_2)
dcd.write(b'Created by ProDy'.ljust(80))
temp = now().strftime('%d %B, %Y at %H:%M')
try:
temp = bytes(temp, encoding='utf-8')
except TypeError:
pass
dcd.write((b'REMARKS Created ' + temp).ljust(80))
dcd.write(pack_i_164)
dcd.write(pack_i_4)
dcd.write(pack(b'i', n_atoms))
dcd.write(pack_i_4)
self._first_byte = dcd.tell()
if self._unitcell:
if unitcell is None:
raise TypeError('unitcell data is expected')
else:
uc = unitcell
uc[3:] = np.sin((PISQUARE / 90) * (90 - uc[3:]))
uc = uc[[0, 3, 1, 4, 5, 2]]
pack_i_48 = pack('i', 48)
dcd.seek(0, 2)
for xyz in coords:
if self._unitcell:
dcd.write(pack_i_48)
uc.tofile(dcd)
dcd.write(pack_i_48)
xyz = xyz.T
dcd.write(pack_i_4N)
xyz[0].tofile(dcd)
dcd.write(pack_i_4N)
dcd.write(pack_i_4N)
xyz[1].tofile(dcd)
dcd.write(pack_i_4N)
dcd.write(pack_i_4N)
xyz[2].tofile(dcd)
dcd.write(pack_i_4N)
self._n_csets += 1
dcd.seek(8, 0)
dcd.write(pack('i', self._n_csets))
dcd.seek(0, 2)
self._nfi = self._n_csets