def saveAtoms(atoms, filename=None, **kwargs):
"""Save *atoms* in ProDy internal format. All atomic classes are accepted
as *atoms* argument. This function saves user set atomic data as well.
Note that title of the AtomGroup instance is used as the filename when
*atoms* is not an AtomGroup. To avoid overwriting an existing file with
the same name, specify a *filename*."""
if not isinstance(atoms, Atomic):
raise TypeError('atoms must be Atomic instance, not {0:s}'
.format(type(atoms)))
if isinstance(atoms, AtomGroup):
ag = atoms
title = ag.getTitle()
SKIP = SAVE_SKIP_ATOMGROUP
else:
ag = atoms.getAtomGroup()
title = str(atoms)
SKIP = SAVE_SKIP_POINTER
if filename is None:
filename = ag.getTitle().replace(' ', '_')
filename += '.ag.npz'
attr_dict = {'title': title}
attr_dict['n_atoms'] = atoms.numAtoms()
attr_dict['n_csets'] = atoms.numCoordsets()
attr_dict['cslabels'] = atoms.getCSLabels()
coords = atoms._getCoordsets()
if coords is not None:
attr_dict['coordinates'] = coords
bonds = ag._bonds
bmap = ag._bmap
if bonds is not None and bmap is not None:
if isinstance(atoms, AtomGroup):
attr_dict['bonds'] = bonds
attr_dict['bmap'] = bmap
attr_dict['numbonds'] = ag._data['numbonds']
frags = ag._data['fragindices']
if frags is not None:
attr_dict['fragindices'] = frags
else:
bonds = trimBonds(bonds, atoms._getIndices())
attr_dict['bonds'] = bonds
attr_dict['bmap'], attr_dict['numbonds'] = \
evalBonds(bonds, len(atoms))
for key, data in ag._data.iteritems():
if key in SKIP:
continue
if data is not None:
attr_dict[key] = data
ostream = openFile(filename, 'wb', **kwargs)
savez(ostream, **attr_dict)
ostream.close()
return filename
def copy(self, which=None):
"""Return a copy of atoms indicated *which* as a new AtomGroup
instance.
*which* may be:
* ``None``, make a copy of the AtomGroup
* a Selection, Residue, Chain, or Atom instance
* a list or an array of indices
* a selection string"""
title = self._title
if which is None:
indices = None
newmol = AtomGroup('{0:s}'.format(title))
newmol.setCoords(self._coords.copy())
elif isinstance(which, int):
indices = [which]
newmol = AtomGroup('{0:s} index {1:d}'.format(title, which))
elif isinstance(which, str):
indices = SELECT.getIndices(self, which)
if len(indices) == 0:
return None
newmol = AtomGroup('{0:s} selection "{1:s}"'.format(title, which))
elif isinstance(which, (list, np.ndarray)):
if isinstance(which, list):
indices = np.array(which)
elif which.ndim != 1:
raise ValueError('which must be a 1d array')
else:
indices = which
newmol = AtomGroup('{0:s} subset'.format(title))
else:
if isinstance(which, Atom):
indices = [which.getIndex()]
elif isinstance(which, (AtomSubset, AtomMap)):
indices = which.getIndices()
else:
raise TypeError('{0:s} is not a valid type'.format(
type(which)))
newmol = AtomGroup('{0:s} selection "{1:s}"'.format(title,
str(which)))
if indices is not None:
newmol.setCoords(self._coords[:, indices])
for key, array in self._data.iteritems():
if key == 'numbonds':
continue
if array is not None:
if indices is None:
newmol._data[key] = array.copy()
else:
newmol._data[key] = array[indices]
newmol._cslabels = list(self._cslabels)
bonds = self._bonds
bmap = self._bmap
if bonds is not None and bmap is not None:
if indices is None:
newmol._bonds = bonds.copy()
newmol._bmap = bmap.copy()
newmol._data['numbonds'] = self._data['numbonds'].copy()
else:
bonds = trimBonds(bonds, indices)
if bonds is not None:
newmol.setBonds(bonds)
return newmol
def copy(self, which=None):
"""Return a copy of atomic data indicated by *which* in a new
:class:`AtomGroup` instance. *which* may be one of:
* ``None``, make a copy of the :class:`AtomGroup` instance
* a :class:`.Selection`, :class:`.Residue`, :class:`.Chain`,
:class:`.Atom`, :class:`.Segment`, :class:`.AtomMap` instance
* a list or an array of indices
* a selection string"""
title = self._title
atommap = False
copy_coords = self._coords is not None
if which is None:
indices = None
newmol = AtomGroup('{0:s}'.format(title))
if copy_coords:
newmol.setCoords(self._coords.copy())
elif isinstance(which, int):
indices = [which]
newmol = AtomGroup('{0:s} index {1:d}'.format(title, which))
elif isinstance(which, str):
indices = SELECT.getIndices(self, which)
if len(indices) == 0:
return None
newmol = AtomGroup('{0:s} selection {1:s}'
.format(title, repr(which)))
elif isinstance(which, (list, np.ndarray)):
if isinstance(which, list):
indices = np.array(which)
elif which.ndim != 1:
raise ValueError('which must be a 1d array')
else:
indices = which
newmol = AtomGroup('{0:s} subset'.format(title))
else:
if isinstance(which, Atom):
idx = which.getIndex()
indices = [idx]
newmol = AtomGroup('{0:s} index {1:d}'.format(title, idx))
elif isinstance(which, AtomSubset):
indices = which._getIndices()
newmol = AtomGroup('{0:s} selection {1:s}'
.format(title, repr(which.getSelstr())))
elif isinstance(which, AtomMap):
indices = np.unique(which._getIndices())
newmol = AtomGroup('{0:s} mapping {1:s}'
.format(title, repr(which.getTitle())))
atommap = True
else:
raise TypeError('{0:s} is not a valid type'.format(
type(which)))
if indices is not None:
if atommap:
if copy_coords:
newmol.setCoords(which.getCoords())
newmol.setData('dummy', which.getDummyFlags())
newmol.setData('mapped', which.getMappedFlags())
elif copy_coords:
newmol.setCoords(self._coords[:, indices])
if not copy_coords:
newmol._n_atoms = len(indices)
for key, array in self._data.iteritems():
if key in READONLY:
continue
if array is not None:
if atommap:
if key in ATOMIC_ATTRIBUTES:
newmol._data[key] = getattr(which, 'get' +
ATOMIC_ATTRIBUTES[key].meth_pl)()
else:
newmol._data[key] = which.getData(key)
else:
if indices is None:
newmol._data[key] = array.copy()
else:
newmol._data[key] = array[indices]
newmol._cslabels = list(self._cslabels)
bonds = self._bonds
bmap = self._bmap
if bonds is not None and bmap is not None:
if indices is None:
newmol._bonds = bonds.copy()
newmol._bmap = bmap.copy()
newmol._data['numbonds'] = self._data['numbonds'].copy()
if self._data['fragindices'] is not None:
newmol._data['fragindices'
] = self._data['fragindices'].copy()
else:
bonds = trimBonds(bonds, indices)
if bonds is not None:
newmol.setBonds(bonds)
newmol._data['fragindices'] = None
return newmol
