Select Atoms within a Region

GetResiduesWithinNang.py

import numpy as np
import mdtraj as md
import matplotlib
import pandas as pd

def GetResiduesWithinNang(traj,radius=20):
    top = traj.topology
    ligand = top.select('resname JLJ')
    prot = top.select("protein and mass>2")
    #def SelectWithinNangs(prot,ligand):
    ap = []
    for i in prot:
        for j in ligand: ap.append([i,j])
    ap = np.array(ap)
    print(ap.shape)
    distance = md.compute_distances(traj,ap)
    df = pd.DataFrame(ap)
    df.columns = ['PROT','LIGA']
    df['DIST']=distance[0]
    reg = df[df.DIST<=radius*0.1]
    reg['RES_CODE'] = [top.atom(i).residue.index for i in reg.PROT]
    ## Expand atomselection to complete residues
    all_atoms = []
    for res in reg.RES_CODE.unique():
        for at in top.residue(res).atoms: all_atoms.append(at.index)
    return all_atoms
traj = md.load_pdb('first.pdb')
print(GetResiduesWithinNang(traj))