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