예제 #1
0
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
예제 #2
0
파일: atomgroup.py 프로젝트: crosvera/ProDy
 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
예제 #3
0
 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