Python Molecule.coords方法代码示例

# 需要导入模块: from htmd.molecule.molecule import Molecule [as 别名]
# 或者: from htmd.molecule.molecule.Molecule import coords [as 别名]
    def _viewStatesNGL(self, states, statetype, protein, ligand, mols, numsamples):
        if states is None:
            states = range(self.macronum)
        if isinstance(states, int):
            states = [states]
        if mols is None:
            mols = self.getStates(states, statetype, numsamples=min(numsamples, 15))
        colors = [0, 1, 3, 4, 5, 6, 7, 9]
        if protein is None and ligand is None:
            raise NameError('Please provide either the "protein" or "ligand" parameter for viewStates.')
        if protein:
            mol = Molecule()
        if ligand:
            mol = mols[0].copy()
            mol.remove(ligand, _logger=False)
            mol.coords = np.atleast_3d(mol.coords[:, :, 0])
            mol.reps.add(sel='protein', style='NewCartoon', color='Secondary Structure')
        for i, s in enumerate(states):
            if protein:
                mol.reps.add(sel='segid ST{}'.format(s), style='NewCartoon', color='Index')
            if ligand:
                mol.reps.add(sel='segid ST{}'.format(s), style='Licorice', color=colors[np.mod(i, len(colors))])
                mols[i].filter(ligand, _logger=False)

            mols[i].set('segid', 'ST{}'.format(s))
            tmpcoo = mols[i].coords
            for j in range(mols[i].numFrames):
                mols[i].coords = np.atleast_3d(tmpcoo[:, :, j])
                mol.append(mols[i])

        w = mol.view(viewer='ngl')
        self._nglButtons(w, statetype, states)
        return w

# 需要导入模块: from htmd.molecule.molecule import Molecule [as 别名]
# 或者: from htmd.molecule.molecule.Molecule import coords [as 别名]
def _filterPDBPSF(sim, outfolder, filtsel):
    try:
        mol = Molecule(sim.molfile)
    except IOError as e:
        raise NameError('simFilter: {}. Cannot create filtered.pdb due to problematic pdb: {}'.format(e, sim.molfile))

    if not path.isfile(path.join(outfolder, 'filtered.pdb')):
        if mol.coords.size == 0:  # If we read for example psf or prmtop which have no coords, just add 0s everywhere
            mol.coords = np.zeros((mol.numAtoms, 3, 1), dtype=np.float32)
        mol.write(path.join(outfolder, 'filtered.pdb'), filtsel)

# 需要导入模块: from htmd.molecule.molecule import Molecule [as 别名]
# 或者: from htmd.molecule.molecule.Molecule import coords [as 别名]
def _fillMolecule(name, resname, chain, resid, insertion, coords, segid, elements):
    numAtoms = len(name)
    mol = Molecule()
    mol.empty(numAtoms)

    mol.name = np.array(name, dtype=mol._pdb_fields['name'])
    mol.resname = np.array(resname, dtype=mol._pdb_fields['resname'])
    mol.chain = np.array(chain, dtype=mol._pdb_fields['chain'])
    mol.resid = np.array(resid, dtype=mol._pdb_fields['resid'])
    mol.insertion = np.array(insertion, dtype=mol._pdb_fields['insertion'])
    mol.coords = np.array(np.atleast_3d(np.vstack(coords)), dtype=mol._pdb_fields['coords'])
    mol.segid = np.array(segid, dtype=mol._pdb_fields['segid'])
    mol.element = np.array(elements, dtype=mol._pdb_fields['element'])
    return mol

# 需要导入模块: from htmd.molecule.molecule import Molecule [as 别名]
# 或者: from htmd.molecule.molecule.Molecule import coords [as 别名]
def _fillMolecule(name, resname, chain, resid, insertion, coords, segid, element,
                  occupancy, beta, charge, record):
    numAtoms = len(name)
    mol = Molecule()
    mol.empty(numAtoms)

    mol.name = np.array(name, dtype=mol._dtypes['name'])
    mol.resname = np.array(resname, dtype=mol._dtypes['resname'])
    mol.chain = np.array(chain, dtype=mol._dtypes['chain'])
    mol.resid = np.array(resid, dtype=mol._dtypes['resid'])
    mol.insertion = np.array(insertion, dtype=mol._dtypes['insertion'])
    mol.coords = np.array(np.atleast_3d(np.vstack(coords)), dtype=mol._dtypes['coords'])
    mol.segid = np.array(segid, dtype=mol._dtypes['segid'])
    mol.element = np.array(element, dtype=mol._dtypes['element'])
    mol.occupancy = np.array(occupancy, dtype=mol._dtypes['occupancy'])
    mol.beta = np.array(beta, dtype=mol._dtypes['beta'])
    # mol.charge = np.array(charge, dtype=mol._dtypes['charge'])
    # mol.record = np.array(record, dtype=mol._dtypes['record'])
    return mol

# 需要导入模块: from htmd.molecule.molecule import Molecule [as 别名]
# 或者: from htmd.molecule.molecule.Molecule import coords [as 别名]
    def _viewStatesNGL(self, states, statetype, protein, ligand, mols, numsamples, gui=False):
        from htmd.builder.builder import sequenceID
        if states is None:
            states = range(self.macronum)
        if isinstance(states, int):
            states = [states]
        if mols is None:
            mols = self.getStates(states, statetype, numsamples=min(numsamples, 15))
        colors = [0, 1, 3, 4, 5, 6, 7, 9]
        hexcolors = {0: '#0000ff', 1: '#ff0000', 2: '#333333', 3: '#ff6600', 4: '#ffff00', 5: '#4c4d00', 6: '#b2b2cc',
                     7: '#33cc33', 8: '#ffffff', 9: '#ff3399', 10: '#33ccff'}
        if protein is None and ligand is None:
            raise NameError('Please provide either the "protein" or "ligand" parameter for viewStates.')
        k = 0
        from nglview import NGLWidget, HTMDTrajectory
        view = NGLWidget(gui=gui)
        ref = mols[0].copy()
        for i, s in enumerate(states):
            if protein:
                mol = Molecule()
            if ligand:
                mol = ref.copy()
                mol.remove(ligand, _logger=False)
                mol.coords = np.atleast_3d(mol.coords[:, :, 0])
                mols[i].filter(ligand, _logger=False)
            mols[i].set('chain', '{}'.format(s))
            tmpcoo = mols[i].coords
            for j in range(mols[i].numFrames):
                mols[i].coords = np.atleast_3d(tmpcoo[:, :, j])
                if ligand:
                    mols[i].set('segid', sequenceID(mols[i].resid)+k)
                    k = int(mols[i].segid[-1])
                mol.append(mols[i])
            view.add_trajectory(HTMDTrajectory(mol))
            # Setting up representations
            if ligand:
                view[i].add_cartoon('protein', color='sstruc')
                view[i].add_hyperball(':{}'.format(s), color=hexcolors[np.mod(i, len(hexcolors))])
            if protein:
                view[i].add_cartoon('protein', color='residueindex')

        self._nglButtons(view, statetype, states)
        return view

# 需要导入模块: from htmd.molecule.molecule import Molecule [as 别名]
# 或者: from htmd.molecule.molecule.Molecule import coords [as 别名]
def _filterTopology(sim, outfolder, filtsel):
    from htmd.util import ensurelist
    try:
        from htmd.molecule.molecule import Molecule
        mol = Molecule(sim.molfile)
    except IOError as e:
        raise RuntimeError('simFilter: {}. Cannot read topology file {}'.format(e, sim.molfile))

    if mol.coords.size == 0:  # If we read for example psf or prmtop which have no coords, just add 0s everywhere
        mol.coords = np.zeros((mol.numAtoms, 3, 1), dtype=np.float32)

    extensions = ['pdb',]  # Adding pdb to make sure it's always written
    for m in ensurelist(sim.molfile):
        extensions.append(os.path.splitext(m)[1][1:])

    for ext in list(set(extensions)):
        filttopo = path.join(outfolder, 'filtered.{}'.format(ext))
        if not path.isfile(filttopo):
            try:
                mol.write(filttopo, filtsel)
            except Exception as e:
                logger.warning('Filtering was not able to write {} due to error: {}'.format(filttopo, e))

# 需要导入模块: from htmd.molecule.molecule import Molecule [as 别名]
# 或者: from htmd.molecule.molecule.Molecule import coords [as 别名]
    def viewStates(self, protein=None, ligand=None, nsamples=20):
        from htmd.projections.metric import _singleMolfile
        from htmd.molecule.molecule import Molecule
        from htmd.vmdviewer import getCurrentViewer
        (single, molfile) = _singleMolfile(self.data.simlist)
        if not single:
            raise RuntimeError('Can''t visualize states without unique molfile')

        viewer = getCurrentViewer()
        colors = [0, 1, 3, 4, 5, 6, 7, 9]

        print('Active set includes macrostates: {}'.format(self.hmm.active_set))

        # dtraj = np.vstack(self.hmm.discrete_trajectories_full)
        res = self.hmm.sample_by_observation_probabilities(nsamples)
        refmol = Molecule(molfile)

        for i, s in enumerate(self.hmm.active_set):
            mol = Molecule(molfile)
            mol.coords = []
            mol.box = []
            # idx = np.where(dtraj == i)[0]
            # samples = np.random.choice(idx, 20)
            # frames = self.data.abs2sim(samples)

            frames = self.data.rel2sim(res[i])
            for f in frames:
                mol._readTraj(f.sim.trajectory[f.piece], frames=[f.frame], append=True)
            mol.wrap('protein')
            mol.align('protein', refmol=refmol)
            viewer.loadMol(mol, name='hmm macro ' + str(s))
            if ligand is not None:
                viewer.rep('ligand', sel=ligand, color=colors[np.mod(i, len(colors))])
            if protein is not None:
                viewer.rep('protein')
            viewer.send('start_sscache')

# 需要导入模块: from htmd.molecule.molecule import Molecule [as 别名]
# 或者: from htmd.molecule.molecule.Molecule import coords [as 别名]
def _filterTopology(sim, outfolder, filtsel):
    try:
        from htmd.molecule.molecule import Molecule
        mol = Molecule(sim.molfile)
    except IOError as e:
        raise RuntimeError('simFilter: {}. Cannot read topology file {}'.format(e, sim.molfile))

    ext = os.path.splitext(sim.molfile)[1][1:]
    filttopo = path.join(outfolder, 'filtered.{}'.format(ext))
    filtpdb = path.join(outfolder, 'filtered.pdb')

    if not path.isfile(filttopo):
        try:
            mol.write(filttopo, filtsel)
        except Exception as e:
            logger.warning('Was not able to write {} due to error: {}'.format(filttopo, e))

        if mol.coords.size == 0:  # If we read for example psf or prmtop which have no coords, just add 0s everywhere
            mol.coords = np.zeros((mol.numAtoms, 3, 1), dtype=np.float32)

        try:
            mol.write(filtpdb, filtsel)
        except Exception as e:
            logger.warning('Was not able to write {} due to error: {}'.format(filtpdb, e))

# 需要导入模块: from htmd.molecule.molecule import Molecule [as 别名]
# 或者: from htmd.molecule.molecule.Molecule import coords [as 别名]
def _viewSphere(spherecoor):
    spheremol = Molecule()
    spheremol.empty(spherecoor.shape[0])
    spheremol.coords = np.atleast_3d(spherecoor)
    spheremol.view(guessbonds=False)

# 需要导入模块: from htmd.molecule.molecule import Molecule [as 别名]
# 或者: from htmd.molecule.molecule.Molecule import coords [as 别名]
def reconstructAdaptiveTraj(simlist, trajID):
    """ Reconstructs a long trajectory out of short adaptive runs.

    Parameters
    ----------
    simlist : numpy.ndarray of :class:`Sim <htmd.simlist.Sim>` objects
        A simulation list generated by the :func:`simlist <htmd.simlist.simlist>` function
    trajID : int
        The id of the trajectory from which to start going back.

    Returns
    -------
    mol : :class:`Molecule <htmd.molecule.molecule.Molecule>` object
        A Molecule object containing the reconstructed trajectory
    chain : np.ndarray
        The simulation IDs of all simulations involved
    pathlist : np.ndarray of str
        The names of all simulations involved.

    Examples
    --------
    >>> mol, chain, pathlist = reconstructAdaptiveTraj(data.simlist, 52)
    """

    sim = None
    for s in simlist:
        if s.simid == trajID:
            sim = s
            break
    if sim is None:
        raise NameError('Could not find sim with ID {} in the simlist.'.format(trajID))

    pathlist = []
    pathlist.append(sim.trajectory[0])
    chain = []
    chain.append((sim, -1, -1))

    epo = None
    while epo != 1:
        [sim, piece, frame, epo] = _findprevioustraj(simlist, _simName(sim.trajectory[0]))
        pathlist.append(sim.trajectory[piece])
        chain.append((sim, piece, frame))
    pathlist = pathlist[::-1]
    chain = chain[::-1]

    mol = Molecule(sim.molfile)
    mol.coords = np.zeros((mol.numAtoms, 3, 0), dtype=np.float32)
    mol.fileloc = []
    mol.box = np.zeros((3, 0))
    for i, c in enumerate(chain):
        tmpmol = Molecule(sim.molfile)
        tmpmol.read(c[0].trajectory)
        endpiece = c[1]
        fileloc = np.vstack(tmpmol.fileloc)
        filenames = fileloc[:, 0]
        pieces = np.unique(filenames)
        firstpieceframe = np.where(filenames == pieces[endpiece])[0][0]
        endFrame = firstpieceframe + c[2]
        if endFrame != -1:
            tmpmol.coords = tmpmol.coords[:, :, 0:endFrame + 1]  # Adding the actual respawned frame (+1) since the respawned sim doesn't include it in the xtc
            tmpmol.fileloc = tmpmol.fileloc[0:endFrame + 1]
            tmpmol.box = tmpmol.box[:, 0:endFrame + 1]
        mol.coords = np.concatenate((mol.coords, tmpmol.coords), axis=2)
        mol.box = np.concatenate((mol.box, tmpmol.box), axis=1)
        mol.fileloc += tmpmol.fileloc
    #mol.fileloc[:, 1] = range(np.size(mol.fileloc, 0))

    return mol, chain, pathlist

# 需要导入模块: from htmd.molecule.molecule import Molecule [as 别名]
# 或者: from htmd.molecule.molecule.Molecule import coords [as 别名]
def ionizePlace(mol, anion_resname, cation_resname, anion_name, cation_name, nanion, ncation, dfrom=5, dbetween=5, segname=None):
    """Place a given number of negative and positive ions in the solvent.

    Replaces water molecules al long as they respect the given distance criteria.

    Parameters
    ----------
    mol : :class:`Molecule <htmd.molecule.molecule.Molecule>` object
        The Molecule object
    anion_resname : str
        Resname of the added anions
    cation_resname : str
        Resname of the added cations
    anion_name : str
        Name of the added anions
    cation_name : str
        Name of the added cations
    nanion : int
        Number of anions to add
    ncation : int
        Number of cations to add
    dfrom : float
        Min distance of ions from molecule
    dbetween : float
        Min distance between ions
    segname : str
        Segment name to add
        
    Returns
    -------
    mol : :class:`Molecule <htmd.molecule.molecule.Molecule>` object
        The molecule with the ions added
    """

    newmol = mol.copy()

    logger.debug('Min distance of ions from molecule: ' + str(dfrom) + 'A')
    logger.debug('Min distance between ions: ' + str(dbetween) + 'A')
    logger.debug('Placing {:d} anions and {:d} cations.'.format(nanion,ncation))

    if (nanion + ncation) == 0:
        return newmol

    nions = nanion + ncation

    betabackup = newmol.beta.copy()
    newmol.set('beta', sequenceID((newmol.resid, newmol.insertion, newmol.segid)))

    # Find water oxygens to replace with ions
    ntries = 0
    maxtries = 10
    while True:
        ionlist = []
        watindex = newmol.atomselect('noh and water and not (within ' + str(dfrom) + ' of not water)', indexes=True)
        watsize = len(watindex)

        if watsize == 0:
            raise NameError('No waters could be found further than ' + str(dfrom) + ' from other molecules to be replaced by ions. You might need to solvate with a bigger box or disable the ionize property when building.')

        while len(ionlist) < nions:
            if len(watindex) == 0:
                break
            randwat = np.random.randint(len(watindex))
            thision = watindex[randwat]
            addit = True
            if len(ionlist) != 0:  # Check for distance from precious ions
                ionspos = newmol.get('coords', sel=ionlist)
                thispos = newmol.get('coords', sel=thision)
                dists = distance.cdist(np.atleast_2d(ionspos), np.atleast_2d(thispos), metric='euclidean')

                if np.any(dists < dbetween):
                    addit = False
            if addit:
                ionlist.append(thision)
                watindex = np.delete(watindex, randwat)
        if len(ionlist) == nions:
            break

        ntries += 1
        if ntries == maxtries:
            raise NameError('Failed to add ions after ' + str(maxtries) + ' attempts. Try decreasing the ''from'' and ''between'' parameters, decreasing ion concentration or making a larger water box.')

    # Delete waters but keep their coordinates
    waterpos = np.atleast_2d(newmol.get('coords', ionlist))
    betasel = np.zeros(newmol.numAtoms, dtype=bool)
    for b in newmol.beta[ionlist]:
        betasel |= newmol.beta == b
    atmrem = np.sum(betasel)
    atmput = 3 * len(ionlist)
    # assert atmrem == atmput, 'Removing {} atoms instead of {}. Report this bug.'.format(atmrem, atmput)
    sel = np.where(betasel)[0]
    newmol.remove(sel, _logger=False)
    # assert np.size(sel) == atmput, 'Removed {} atoms instead of {}. Report this bug.'.format(np.size(sel), atmput)
    betabackup = np.delete(betabackup, sel)

    # Add the ions
    randidx = np.random.permutation(np.size(waterpos, 0))
    atom = Molecule()
    atom.empty(1)
    atom.set('chain', 'I')
#.........这里部分代码省略.........

# 需要导入模块: from htmd.molecule.molecule import Molecule [as 别名]
# 或者: from htmd.molecule.molecule.Molecule import coords [as 别名]
def _applyProteinCaps(mol, caps):

    # AMBER capping
    # =============
    # This is the (horrible) way of adding caps in tleap:
    # For now, this is hardwired for ACE and NME
    # 1. Change one of the hydrogens of the N terminal (H[T]?[123]) to the ACE C atom, giving it a new resid
    # 1a. If no hydrogen present, create the ACE C atom.
    # 2. Change one of the oxygens of the C terminal ({O,OT1,OXT}) to the NME N atom, giving it a new resid
    # 2a. If no oxygen present, create the NME N atom.
    # 3. Reorder to put the new atoms first and last
    # 4. Remove the lingering hydrogens of the N terminal and oxygens of the C terminal.

    # Define the atoms to be replaced (0 and 1 corresponds to N- and C-terminal caps)
    terminalatoms = {'ACE': 'H1 H2 H3 HT1 HT2 HT3', 'NME': 'OXT OT1 O'}  # XPLOR names for H[123] and OXT are HT[123]
                                                                         # and OT1, respectively.
    capresname = ['ACE', 'NME']
    capatomtype = ['C', 'N']

    # For each caps definition
    for seg in caps:
        # Get the segment
        segment = mol.atomselect('segid {}'.format(seg), indexes=True)
        # Test segment
        if len(segment) == 0:
            raise RuntimeError('There is no segment {} in the molecule.'.format(seg))
        if len(mol.atomselect('protein and segid {}'.format(seg), indexes=True)) == 0:
            raise RuntimeError(
                'Segment {} is not protein. Capping for non-protein segments is not supported.'.format(seg))
        # For each cap
        for i, cap in enumerate(caps[seg]):
            if cap is None or (isinstance(cap, str) and cap == 'none'):
                continue
            # Get info on segment and its terminals
            segment = mol.atomselect('segid {}'.format(seg), indexes=True)
            resids = np.unique(mol.get('resid', sel=segment))
            terminalids = [segment[0], segment[-1]]
            terminalresids = [np.min(resids), np.max(resids)]
            if i == 0:
                orig_terminalresids = [np.min(resids), np.max(resids)]
            # In case there is no cap defined
            if cap is None or cap == '':
                logger.warning(
                    'No cap provided for resid {} on segment {}. Did not apply it.'.format(terminalresids[i], seg))
                continue
            # If it is defined, test if supported
            elif cap not in capresname:
                raise RuntimeError(
                    'In segment {}, the {} cap is not supported. Try using {} instead.'.format(seg, cap, capresname))
            # Test if cap is already applied
            testcap = mol.atomselect('segid {} and resid "{}" and resname {}'.format(seg, terminalresids[i], cap),
                                     indexes=True)
            if len(testcap) != 0:
                logger.warning('Cap {} already exists on segment {}. Did not re-apply it.'.format(cap, seg))
                continue
            # Test if the atom to change exists
            termatomsids = mol.atomselect('segid {} and resid "{}" and name {}'.format(seg,
                                                                                       terminalresids[i],
                                                                                       terminalatoms[cap]),
                                          indexes=True)
            if len(termatomsids) == 0:
                # Create new atom
                termcaid = mol.atomselect('segid {} and resid "{}" and name CA'.format(seg, terminalresids[i]),
                                        indexes=True)
                termcenterid = mol.atomselect('segid {} and resid "{}" and name {}'.format(seg, terminalresids[i],
                                                                                           capatomtype[-i+1]),
                                        indexes=True)  # if i=0 => capatomtype[1]; i=1 => capatomtype[0]
                atom = Molecule()
                atom.empty(1)
                atom.record = 'ATOM'
                atom.name = capatomtype[i]
                atom.resid = terminalresids[i]-1+2*i
                atom.resname = cap
                atom.segid = seg
                atom.element = capatomtype[i]
                atom.chain = np.unique(mol.get('chain', sel='segid {}'.format(seg)))
                atom.coords = mol.coords[termcenterid] + 0.33 * np.subtract(mol.coords[termcenterid],
                                                                               mol.coords[termcaid])
                mol.insert(atom, terminalids[i])
                # newatom = mol.numAtoms - 1
                logger.info('In segment {}, resid {} had none of these atoms: {}. Capping was performed by creating '
                            'a new atom for cap construction by tleap.'.format(seg, terminalresids[i],
                                                                              terminalatoms[cap]))
            else:
                # Select atom to change, do changes to cap, and change resid
                newatom = np.max(termatomsids)
                mol.set('resname', cap, sel=newatom)
                mol.set('name', capatomtype[i], sel=newatom)
                mol.set('element', capatomtype[i], sel=newatom)
                mol.set('resid', terminalresids[i]-1+2*i, sel=newatom)  # if i=0 => resid-1; i=1 => resid+1

                # Reorder
                neworder = np.arange(mol.numAtoms)
                neworder[newatom] = terminalids[i]
                neworder[terminalids[i]] = newatom
                _reorderMol(mol, neworder)

        # For each cap
        for i, cap in enumerate(caps[seg]):
            if cap is None or (isinstance(cap, str) and cap == 'none'):
#.........这里部分代码省略.........

# 需要导入模块: from htmd.molecule.molecule import Molecule [as 别名]
# 或者: from htmd.molecule.molecule.Molecule import coords [as 别名]

for ext in _MDTRAJ_SAVERS:
    if ext not in _WRITERS:
        _WRITERS[ext] = MDTRAJwrite


if __name__ == '__main__':
    from htmd.home import home
    from htmd.molecule.molecule import Molecule, mol_equal
    from htmd.util import tempname
    import numpy as np
    import os
    testfolder = home(dataDir='metricdistance')
    mol = Molecule(os.path.join(testfolder, 'filtered.pdb'))
    mol.coords = np.tile(mol.coords, (1, 1, 2))
    mol.filter('protein and resid 1 to 20')
    mol.boxangles = np.ones((3, 2), dtype=np.float32) * 90
    mol.box = np.ones((3, 2), dtype=np.float32) * 15
    mol.step = np.arange(2)
    mol.time = np.arange(2) * 1E5

    for ext in _WRITERS:
        tmp = tempname(suffix='.'+ext)
        mol.write(tmp)
        print('Can write {} files'.format(ext))


    # from difflib import Differ
    # d = Differ()
    #

# 需要导入模块: from htmd.molecule.molecule import Molecule [as 别名]
# 或者: from htmd.molecule.molecule.Molecule import coords [as 别名]
def ionizePlace(mol, anion, cation, anionatom, cationatom, nanion, ncation, dfrom=5, dbetween=5, segname=None):
    newmol = mol.copy()

    logger.info('Min distance of ions from molecule: ' + str(dfrom) + 'A')
    logger.info('Min distance between ions: ' + str(dbetween) + 'A')
    logger.info('Placing ' + str(nanion+ncation) + ' ions.')

    if (nanion + ncation) == 0:
        return newmol

    segname = _getSegname(newmol, segname)
    nions = nanion + ncation

    betabackup = newmol.beta
    newmol.set('beta', sequenceID((newmol.resid, newmol.segid)))

    # Find water oxygens to replace with ions
    ntries = 0
    maxtries = 10
    while True:
        ionlist = np.empty(0, dtype=int)
        watindex = newmol.atomselect('noh and water and not (within ' + str(dfrom) + ' of not water)', indexes=True)
        watsize = len(watindex)

        if watsize == 0:
            raise NameError('No waters could be found further than ' + str(dfrom) + ' from other molecules to be replaced by ions. You might need to solvate with a bigger box or disable the ionize property when building.')

        while len(ionlist) < nions:
            if len(watindex) == 0:
                break
            randwat = np.random.randint(len(watindex))
            thision = watindex[randwat]
            addit = True
            if len(ionlist) != 0:  # Check for distance from precious ions
                ionspos = newmol.get('coords', sel=ionlist)
                thispos = newmol.get('coords', sel=thision)
                dists = distance.cdist(np.atleast_2d(ionspos), np.atleast_2d(thispos), metric='euclidean')

                if np.any(dists < dbetween):
                    addit = False
            if addit:
                ionlist = np.append(ionlist, thision)
                watindex = np.delete(watindex, randwat)
        if len(ionlist) == nions:
            break

        ntries += 1
        if ntries == maxtries:
            raise NameError('Failed to add ions after ' + str(maxtries) + ' attempts. Try decreasing the ''from'' and ''between'' parameters, decreasing ion concentration or making a larger water box.')

    # Delete waters but keep their coordinates
    waterpos = np.atleast_2d(newmol.get('coords', ionlist))
    stringsel = 'beta'
    for x in newmol.beta[ionlist]:
        stringsel += ' ' + str(int(x))
    atmrem = np.sum(newmol.atomselect(stringsel))
    atmput = 3 * len(ionlist)
    # assert atmrem == atmput, 'Removing {} atoms instead of {}. Report this bug.'.format(atmrem, atmput)
    sel = newmol.atomselect(stringsel, indexes=True)
    newmol.remove(sel, _logger=False)
    # assert np.size(sel) == atmput, 'Removed {} atoms instead of {}. Report this bug.'.format(np.size(sel), atmput)
    betabackup = np.delete(betabackup, sel)

    # Add the ions
    randidx = np.random.permutation(np.size(waterpos, 0))
    atom = Molecule()
    atom.empty(1)
    atom.set('chain', 'I')
    atom.set('segid', 'I')

    for i in range(nanion):
        atom.set('name', anionatom)
        atom.set('resname', anion)
        atom.set('resid', newmol.resid[-1] + 1)
        atom.coords = waterpos[randidx[i], :]
        newmol.insert(atom, len(newmol.name))
    for i in range(ncation):
        atom.set('name', cationatom)
        atom.set('resname', cation)
        atom.set('resid', newmol.resid[-1] + 1)
        atom.coords = waterpos[randidx[i+nanion], :]
        newmol.insert(atom, len(newmol.name))

    # Restoring the original betas
    sel = np.ones(len(betabackup) + nions, dtype=bool)
    sel[len(betabackup)::] = False
    newmol.set('beta', betabackup, sel=sel)
    return newmol

# 需要导入模块: from htmd.molecule.molecule import Molecule [as 别名]
# 或者: from htmd.molecule.molecule.Molecule import coords [as 别名]
def _applyProteinCaps(mol, caps):

    # AMBER capping
    # =============
    # This is the (horrible) way of adding caps in tleap:
    # For now, this is hardwired for ACE and NME
    # 1. Change one of the hydrogens of the N terminal (H[T]?[123]) to the ACE C atom, giving it a new resid
    # 1a. If no hydrogen present, create the ACE C atom.
    # 2. Change one of the oxygens of the C terminal ({O,OT1,OXT}) to the NME N atom, giving it a new resid
    # 2a. If no oxygen present, create the NME N atom.
    # 3. Reorder to put the new atoms first and last
    # 4. Remove the lingering hydrogens of the N terminal and oxygens of the C terminal.

    # Define the atoms to be replaced (0 and 1 corresponds to N- and C-terminal caps)
    terminalatoms = {'ACE': ['H1', 'H2', 'H3', 'HT1', 'HT2', 'HT3'], 'NME': ['OXT', 'OT1', 'O']}  # XPLOR names for H[123] and OXT are HT[123]
                                                                         # and OT1, respectively.
    capresname = ['ACE', 'NME']
    capatomtype = ['C', 'N']

    # For each caps definition
    for seg in caps:
        prot = mol.atomselect('protein')  # Can't move this out since we remove atoms in this loop
        # Get the segment
        segment = np.where(mol.segid == seg)[0]
        # Test segment
        if len(segment) == 0:
            raise RuntimeError('There is no segment {} in the molecule.'.format(seg))
        if not np.any(prot & (mol.segid == seg)):
            raise RuntimeError('Segment {} is not protein. Capping for non-protein segments is not supported.'.format(seg))
        # For each cap
        passed = False
        for i, cap in enumerate(caps[seg]):
            if cap is None or (isinstance(cap, str) and cap == 'none'):
                continue
            # Get info on segment and its terminals
            segidm = mol.segid == seg  # Mask for segid
            segididx = np.where(segidm)[0]
            resids = mol.resid[segididx]
            terminalids = [segididx[0], segididx[-1]]
            terminalresids = [resids[0], resids[-1]]
            residm = mol.resid == terminalresids[i]  # Mask for resid

            if not passed:
                orig_terminalresids = terminalresids
                passed = True

            if cap is None or cap == '':  # In case there is no cap defined
                logger.warning('No cap provided for resid {} on segment {}. Did not apply it.'.format(terminalresids[i], seg))
                continue
            elif cap not in capresname:  # If it is defined, test if supported
                raise RuntimeError('In segment {}, the {} cap is not supported. Try using {} instead.'.format(seg, cap, capresname))

            # Test if cap is already applied
            testcap = np.where(segidm & residm & (mol.resname == cap))[0]
            if len(testcap) != 0:
                logger.warning('Cap {} already exists on segment {}. Did not re-apply it.'.format(cap, seg))
                continue

            # Test if the atom to change exists
            termatomsids = np.zeros(residm.shape, dtype=bool)
            for atm in terminalatoms[cap]:
                termatomsids |= mol.name == atm
            termatomsids = np.where(termatomsids & segidm & residm)[0]

            if len(termatomsids) == 0:
                # Create new atom
                termcaid = np.where(segidm & residm & (mol.name == 'CA'))[0]
                termcenterid = np.where(segidm & residm & (mol.name == capatomtype[1-i]))[0]
                atom = Molecule()
                atom.empty(1)
                atom.record = np.array(['ATOM'], dtype=Molecule._dtypes['record'])
                atom.name = np.array([capatomtype[i]], dtype=Molecule._dtypes['name'])
                atom.resid = np.array([terminalresids[i]-1+2*i], dtype=Molecule._dtypes['resid'])
                atom.resname = np.array([cap], dtype=Molecule._dtypes['resname'])
                atom.segid = np.array([seg], dtype=Molecule._dtypes['segid'])
                atom.element = np.array([capatomtype[i]], dtype=Molecule._dtypes['element'])
                atom.chain = np.array([np.unique(mol.chain[segidm])], dtype=Molecule._dtypes['chain'])  # TODO: Assumption of single chain in a segment might be wrong
                atom.coords = mol.coords[termcenterid] + 0.33 * np.subtract(mol.coords[termcenterid],
                                                                            mol.coords[termcaid])
                mol.insert(atom, terminalids[i])
                # logger.info('In segment {}, resid {} had none of these atoms: {}. Capping was performed by creating '
                #             'a new atom for cap construction by tleap.'.format(seg, terminalresids[i],
                #                                                                ' '.join(terminalatoms[cap])))

            else:
                # Select atom to change, do changes to cap, and change resid
                newatom = np.max(termatomsids)
                mol.set('resname', cap, sel=newatom)
                mol.set('name', capatomtype[i], sel=newatom)
                mol.set('element', capatomtype[i], sel=newatom)
                mol.set('resid', terminalresids[i]-1+2*i, sel=newatom)  # if i=0 => resid-1; i=1 => resid+1

                # Reorder
                neworder = np.arange(mol.numAtoms)
                neworder[newatom] = terminalids[i]
                neworder[terminalids[i]] = newatom
                _reorderMol(mol, neworder)

        # For each cap
        for i, cap in enumerate(caps[seg]):
#.........这里部分代码省略.........