LabelLib_pymol.py

from pymol import cmd
import chempy
import numpy as np
import LabelLib as ll

# Usage example:
# fetch 1BNA, async=0
# remove solvent
# genAV('1BNA', '/1BNA/B/B/19/C5', allowed_sphere_radius=1.5)
## genAV('1BNA', '/1BNA/B/B/19/C5', linker_length=22.0, linker_diameter=3.0, dye_radius=4.0, disc_step=0.7, allowed_sphere_radius=2.0)
def genAV(obstacles, attachment, linker_length=20.0, linker_diameter=2.0, dye_radius=3.5, disc_step=0.9, name=None, state=1, stripsc=True, allowed_sphere_radius=0.0, smoothSurf=True):

  source = np.array(cmd.get_model(attachment, state).get_coord_list())
  if (source.shape[0]!=1):
    print('attachment selection must contain exactly one atom, selected: {}'.format(source.shape[0]))
    return
  source=source.reshape(3)

  srcAt = cmd.get_model(attachment, state).atom[0]
  srcModelName = cmd.get_names('objects',0,attachment)[0]

  obstacles = '(' + obstacles + ') and not (' + attachment + ')'
  if stripsc and isAA(srcAt.resn):
    obstacles += ' and not (' + srcModelName
    if len(srcAt.chain)>0:
      obstacles += ' and chain ' + srcAt.chain
    obstacles+=' and resi '+srcAt.resi+' and sidechain'+')'
  if allowed_sphere_radius > 0.0:
    obstacles+=' and not (({}) around {})'.format(attachment, allowed_sphere_radius)

  xyzRT=np.zeros((1,4))
  nAtoms=cmd.count_atoms(obstacles)
  if nAtoms>0:
    atoms=cmd.get_model(obstacles, state).atom
    nAtoms=len(atoms)
    xyzRT=np.zeros((nAtoms,4))
    for i,at in enumerate(atoms):
      xyzRT[i]=[at.coord[0],at.coord[1],at.coord[2],at.vdw]

  av1=ll.dyeDensityAV1(xyzRT.T,source,linker_length, linker_diameter, dye_radius, disc_step)
  m=avToModel(av1)
  if len(m.atom)==0:
    print('Failed: Empty AV. Is attachment position buried?')
    return
  if name is None:
    name = srcModelName + '_'
    if len(srcAt.chain)>0:
      name += srcAt.chain + '-'
    name +=  srcAt.resi + '-' + srcAt.name
  cmd.load_model(m, name)

  if smoothSurf:
    surfName=name+'_surf'
    mapName=name+'_map'
    gRes=cmd.get('gaussian_resolution')
    cmd.set('gaussian_resolution',3.0)
    cmd.map_new(mapName,'gaussian', 1.0, name, 6)
    cmd.isosurface(surfName,mapName,0.9)
    cmd.set('gaussian_resolution',gRes)
    cmd.disable(name)

def isAA(resn):
  names=['ALA', 'ARG', 'ASN', 'ASP', 'ASX',
	 'CYS', 'GLU', 'GLN', 'GLX', 'GLY',
	 'HIS', 'ILE', 'LEU', 'LYS', 'MET',
	 'PHE', 'PRO', 'SER', 'THR', 'TRP', 'TYR', 'VAL']
  if resn in names:
    return True
  return False

def makeAtom(index, xyz, vdw, name='AV', q=1.0):
    atom = chempy.Atom()
    atom.index = index
    atom.name = name
    atom.symbol = 'He'
    atom.resn = 'AV'
    atom.chain = 'A'
    atom.resi = 1
    atom.resi_number = 1
    atom.coord = xyz
    atom.vdw=vdw
    atom.hetatm = False
    atom.b = 100
    atom.q = q
    return atom

def avToModel(av):
  m = chempy.models.Indexed()
  points = av.points()
  r = av.discStep * 0.5

  for i, p in enumerate(points.T):
    x, y, z, w = p
    m.add_atom(makeAtom(i+1, [x, y, z], r, q=w))

  MP = np.average(points[:3,:],axis=1,weights=points[3])
  if points.shape[1]>0:
    m.add_atom(makeAtom(points.shape[1]+1,list(MP),2.0,'AVmp'))

  m.update_index()
  return m


cmd.extend("genAV", genAV)