def random_ellipsoid(box, size, min_axis):
# return a random ellipsoid record of the form:
# [ ELLIPSOID, x_pos, y_pos, z_pos, size, x0, y0, z0, x1, y1, z2, x2, y2, z2 ]
# where the xyz vectors are orthogonal and of length 1.0 or less.
box = box - size
tmp0 = cpv.random_vector()
tmp1 = cpv.random_vector()
tmp2 = cpv.cross_product(tmp1, tmp0)
tmp3 = cpv.cross_product(tmp1, tmp2)
tmp4 = cpv.cross_product(tmp2, tmp3)
tmp2 = cpv.normalize(tmp2)
tmp3 = cpv.normalize(tmp3)
tmp4 = cpv.normalize(tmp4)
primary = cpv.scale(tmp2, random())
secondary = cpv.scale(tmp3,random())
tertiary = cpv.scale(tmp4,random())
factor = 1.0 / max( cpv.length(primary), cpv.length(secondary), cpv.length(tertiary))
primary = cpv.scale(primary, factor)
secodary = cpv.scale(secondary, factor)
tertiary = cpv.scale(tertiary, factor)
return [ ELLIPSOID,
size + random() * box, size + random() * box, size + random() * box,
max(random() * size, min_axis),
] + primary + secondary + tertiary
def plane_orientation(selection, state=-1, visualize=1, quiet=1):
'''
DESCRIPTION
Fit plane (for example beta-sheet). Can also be used with
angle_between_helices (even though this does not fit helices).
Returns center and normal vector of plane.
'''
try:
import numpy
except ImportError:
print ' Error: numpy not available'
raise CmdException
state, visualize, quiet = int(state), int(visualize), int(quiet)
coords = list()
cmd.iterate_state(state, '(%s) and guide' % (selection),
'coords.append([x,y,z])', space=locals())
if len(coords) < 3:
print 'not enough guide atoms in selection'
raise CmdException
x = numpy.array(coords)
U,s,Vh = numpy.linalg.svd(x - x.mean(0))
# normal vector of plane is 3rd principle component
vec = cpv.normalize(Vh[2])
if cpv.dot_product(vec, x[-1] - x[0]) < 0:
vec = cpv.negate(vec)
center = x.mean(0).tolist()
_common_orientation(selection, center, vec, visualize, 4.0, quiet)
# plane visualize
if visualize:
from pymol import cgo
dir1 = cpv.normalize(Vh[0])
dir2 = cpv.normalize(Vh[1])
sx = [max(i/4.0, 2.0) for i in s]
obj = [ cgo.BEGIN, cgo.TRIANGLES, cgo.COLOR, 0.5, 0.5, 0.5 ]
for vertex in [
cpv.scale(dir1, sx[0]),
cpv.scale(dir2, sx[1]),
cpv.scale(dir2, -sx[1]),
cpv.scale(dir1, -sx[0]),
cpv.scale(dir2, -sx[1]),
cpv.scale(dir2, sx[1]),
]:
obj.append(cgo.VERTEX)
obj.extend(cpv.add(center, vertex))
obj.append(cgo.END)
cmd.load_cgo(obj, cmd.get_unused_name('planeFit'))
return center, vec
def planeFromPoints(point1, point2, point3, facetSize):
v1 = cpv.normalize(cpv.sub(point2, point1))
v2 = cpv.normalize(cpv.sub(point3, point1))
normal = cpv.cross_product(v1, v2)
v2 = cpv.cross_product(normal, v1)
x = cpv.scale(v1, facetSize)
y = cpv.scale(v2, facetSize)
center = point2
corner1 = cpv.add(cpv.add(center, x), y)
corner2 = cpv.sub(cpv.add(center, x), y)
corner3 = cpv.sub(cpv.sub(center, x), y)
corner4 = cpv.add(cpv.sub(center, x), y)
return plane(corner1, corner2, corner3, corner4, normal)
def loop_orientation(selection, visualize=1, quiet=0):
'''
DESCRIPTION
Get the center and approximate direction of a peptide. Works for any
secondary structure.
Averages direction of N(i)->C(i) pseudo bonds.
USAGE
loop_orientation selection [, visualize]
SEE ALSO
helix_orientation
'''
visualize, quiet = int(visualize), int(quiet)
stored.x = dict()
cmd.iterate_state(STATE, '(%s) and name N+C' % (selection),
'stored.x.setdefault(chain + resi, dict())[name] = x,y,z')
vec = cpv.get_null()
count = 0
for x in stored.x.itervalues():
if 'C' in x and 'N' in x:
vec = cpv.add(vec, cpv.sub(x['C'], x['N']))
count += 1
if count == 0:
print 'warning: count == 0'
raise CmdException
vec = cpv.normalize(vec)
return _common_orientation(selection, vec, visualize, quiet)
def cafit_orientation(selection, visualize=1, quiet=0):
'''
DESCRIPTION
Get the center and direction of a peptide by least squares
linear fit on CA atoms.
USAGE
cafit_orientation selection [, visualize]
NOTES
Requires python module "numpy".
SEE ALSO
helix_orientation
'''
visualize, quiet = int(visualize), int(quiet)
import numpy
stored.x = list()
cmd.iterate_state(STATE, '(%s) and name CA' % (selection),
'stored.x.append([x,y,z])')
x = numpy.array(stored.x)
U, s, Vh = numpy.linalg.svd(x - x.mean(0))
vec = cpv.normalize(Vh[0])
if cpv.dot_product(vec, x[-1] - x[0]) < 0:
vec = cpv.negate(vec)
return _common_orientation(selection, vec, visualize, quiet)
def helix_orientation(selection, visualize=1, sigma_cutoff=1.5, quiet=0):
"""
DESCRIPTION
Get the center and direction of a helix as vectors. Will only work
for helices and gives slightly different results than loop_orientation.
Averages direction of C(i)->O(i) bonds.
USAGE
helix_orientation selection [, visualize [, sigma_cutoff]]
ARGUMENTS
selection = string: atom selection of helix
visualize = 0 or 1: show fitted vector as arrow {default: 1}
sigma_cutoff = float: drop outliers outside
(standard_deviation * sigma_cutoff) {default: 1.5}
SEE ALSO
angle_between_helices, helix_orientation_hbond, loop_orientation, cafit_orientation
"""
visualize, quiet, sigma_cutoff = int(visualize), int(quiet), float(sigma_cutoff)
stored.x = dict()
cmd.iterate_state(
STATE, "(%s) and name C+O" % (selection), "stored.x.setdefault(chain + resi, dict())[name] = x,y,z"
)
vec_list = []
count = 0
for x in stored.x.values():
if "C" in x and "O" in x:
vec_list.append(cpv.sub(x["O"], x["C"]))
count += 1
if count == 0:
print("warning: count == 0")
raise CmdException
vec = _vec_sum(vec_list)
if count > 2 and sigma_cutoff > 0:
angle_list = [cpv.get_angle(vec, x) for x in vec_list]
angle_mu, angle_sigma = _mean_and_std(angle_list)
vec_list = [
vec_list[i] for i in range(len(vec_list)) if abs(angle_list[i] - angle_mu) < angle_sigma * sigma_cutoff
]
if not quiet:
print("Dropping %d outlier(s)" % (len(angle_list) - len(vec_list)))
vec = _vec_sum(vec_list)
vec = cpv.normalize(vec)
return _common_orientation(selection, vec, visualize, quiet)
def __init__(self, p1, p2, radius, color1,
color2=None, color3=None,
hlength=None, hradius=None,
hlength_scale=3.0, hradius_scale=0.6):
if hlength is None:
hlength = radius * hlength_scale
if hradius is None:
hradius = hlength * hradius_scale
normal = cpv.normalize(cpv.sub(p1, p2))
pM = cpv.add(cpv.scale(normal, hlength), p2)
line = Cylinder(p1, pM, radius, color1, color2)
cone = Cone(
pM, p2, hradius, color2 or color1, radius2=0, color2=color3
)
self._primitive = (line + cone).primitive
def cgo_arrow(atom1='pk1',
atom2='pk2',
radius=0.07,
gap=0.0,
hlength=-1,
hradius=-1,
color='blue red',
name=''):
from chempy import cpv
radius, gap = float(radius), float(gap)
hlength, hradius = float(hlength), float(hradius)
try:
color1, color2 = color.split()
except:
color1 = color2 = color
color1 = list(cmd.get_color_tuple(color1))
color2 = list(cmd.get_color_tuple(color2))
def get_coord(v):
if not isinstance(v, str):
return v
if v.startswith('['):
return cmd.safe_list_eval(v)
return cmd.get_atom_coords(v)
xyz1 = get_coord(atom1)
xyz2 = get_coord(atom2)
normal = cpv.normalize(cpv.sub(xyz1, xyz2))
if hlength < 0:
hlength = radius * 3.0
if hradius < 0:
hradius = hlength * 0.6
if gap:
diff = cpv.scale(normal, gap)
xyz1 = cpv.sub(xyz1, diff)
xyz2 = cpv.add(xyz2, diff)
xyz3 = cpv.add(cpv.scale(normal, hlength), xyz2)
obj = [cgo.CYLINDER] + xyz1 + xyz3 + [radius] + color1 + color2 + [
cgo.CONE
] + xyz3 + xyz2 + [hradius, 0.0] + color2 + color2 + [1.0, 0.0]
return obj
def loop_orientation(selection, state=STATE, visualize=1, quiet=1):
'''
DESCRIPTION
Get the center and approximate direction of a peptide. Works for any
secondary structure.
Averages direction of N(i)->C(i) pseudo bonds.
USAGE
loop_orientation selection [, visualize ]
SEE ALSO
helix_orientation
'''
state, visualize, quiet = int(state), int(visualize), int(quiet)
coords = dict()
cmd.iterate_state(state,
'(%s) and name N+C' % (selection),
'coords.setdefault(chain + resi, {})[name] = x,y,z',
space=locals())
vec = cpv.get_null()
center = cpv.get_null()
count = 0
for x in coords.values():
if 'C' in x and 'N' in x:
vec = cpv.add(vec, cpv.sub(x['C'], x['N']))
for coord in x.values():
center = cpv.add(center, coord)
count += 1
if count == 0:
print('warning: count == 0')
raise CmdException
vec = cpv.normalize(vec)
center = cpv.scale(center, 1. / count)
_common_orientation(selection, center, vec, visualize, 2.0 * len(coords),
quiet)
return center, vec
def _cgo_quad(pos, normal, radius):
'''Return a CGO list specifying a quad.'''
v1 = cpv.normalize(_perp_vec(normal))
v2 = cpv.cross_product(normal, v1)
v1 = cpv.scale(v1, radius)
v2 = cpv.scale(v2, radius)
obj = [cgo.BEGIN, cgo.TRIANGLE_STRIP, cgo.NORMAL]
obj.extend(normal)
obj.append(cgo.VERTEX)
obj.extend(cpv.add(pos, v1))
obj.append(cgo.VERTEX)
obj.extend(cpv.add(pos, v2))
obj.append(cgo.VERTEX)
obj.extend(cpv.sub(pos, v2))
obj.append(cgo.VERTEX)
obj.extend(cpv.sub(pos, v1))
obj.append(cgo.END)
return obj
def draw_arrow(xyz1,
xyz2,
radius=0.5,
gap=0.0,
hlength=-1,
hradius=-1,
color='blue red',
name=''):
'''
Draw an arrow; borrows heavily from cgi arrows.
'''
radius, gap = float(radius), float(gap)
hlength, hradius = float(hlength), float(hradius)
xyz1 = list(xyz1)
xyz2 = list(xyz2)
try:
color1, color2 = color.split()
except:
color1 = color2 = color
color1 = list(cmd.get_color_tuple(color1))
color2 = list(cmd.get_color_tuple(color2))
normal = cpv.normalize(cpv.sub(xyz1, xyz2))
if hlength < 0:
hlength = radius * 3.0
if hradius < 0:
hradius = hlength * 0.6
if gap:
diff = cpv.scale(normal, gap)
xyz1 = cpv.sub(xyz1, diff)
xyz2 = cpv.add(xyz2, diff)
xyz3 = cpv.add(cpv.scale(normal, hlength), xyz2)
obj = [cgo.CYLINDER] + xyz1 + xyz3 + [radius] + color1 + color2 + \
[cgo.CONE] + xyz3 + xyz2 + [hradius, 0.0] + color2 + color2 + \
[1.0, 0.0]
if not name:
name = cmd.get_unused_name('arrow')
cmd.load_cgo(obj, name)
def cafit_orientation(selection, state=STATE, visualize=1, guide=1, quiet=1):
'''
DESCRIPTION
Get the center and direction of a peptide by least squares
linear fit on CA atoms.
USAGE
cafit_orientation selection [, visualize ]
NOTES
Requires python module "numpy".
SEE ALSO
helix_orientation
'''
try:
import numpy
except ImportError:
print ' Error: numpy not available'
raise CmdException
state, visualize, quiet = int(state), int(visualize), int(quiet)
if int(guide):
selection = '(%s) and guide' % (selection)
coords = []
cmd.iterate_state(state, selection,
'coords.append([x,y,z])', space=locals())
x = numpy.array(coords)
center = x.mean(0).tolist()
U,s,Vh = numpy.linalg.svd(x - center)
vec = cpv.normalize(Vh[0])
if cpv.dot_product(vec, x[-1] - x[0]) < 0:
vec = cpv.negate(vec)
_common_orientation(selection, center, vec, visualize, s[0], quiet)
return center, vec
def get_cgo(self, dot_mode=0, dot_radius=0.03):
"""Generate a CGO list for a dot."""
cgolist = []
# COLOR
cgolist.extend(_cgo_color(self.color))
if dot_mode == 0: # spheres
logger.debug("Adding dot to cgolist...")
cgolist.extend(_cgo_sphere(self.coords, dot_radius))
logger.debug("Finished adding dot to cgolist.")
if dot_mode == 1: # quads
logger.debug("Adding quad to cgolist...")
normal = cpv.normalize(cpv.sub(self.coords, self.atom['coords']))
cgolist.extend(_cgo_quad(self.coords, normal, dot_radius * 1.5))
logger.debug("Finished adding quad to cgolist.")
return cgolist
def get_cgo(self, dot_mode=0, dot_radius=0.03):
"""Generate a CGO list for a dot."""
cgolist = []
# COLOR
cgolist.extend(_cgo_color(self.color))
if dot_mode == 0: # spheres
logger.debug("Adding dot to cgolist...")
cgolist.extend(_cgo_sphere(self.coords, dot_radius))
logger.debug("Finished adding dot to cgolist.")
if dot_mode == 1: # quads
logger.debug("Adding quad to cgolist...")
normal = cpv.normalize(cpv.sub(self.coords, self.atom['coords']))
cgolist.extend(_cgo_quad(self.coords, normal, dot_radius * 1.5))
logger.debug("Finished adding quad to cgolist.")
return cgolist
def loop_orientation(selection, state=STATE, visualize=1, quiet=1):
'''
DESCRIPTION
Get the center and approximate direction of a peptide. Works for any
secondary structure.
Averages direction of N(i)->C(i) pseudo bonds.
USAGE
loop_orientation selection [, visualize ]
SEE ALSO
helix_orientation
'''
state, visualize, quiet = int(state), int(visualize), int(quiet)
coords = dict()
cmd.iterate_state(state, '(%s) and name N+C' % (selection),
'coords.setdefault(chain + resi, {})[name] = x,y,z', space=locals())
vec = cpv.get_null()
center = cpv.get_null()
count = 0
for x in coords.itervalues():
if 'C' in x and 'N' in x:
vec = cpv.add(vec, cpv.sub(x['C'], x['N']))
for coord in x.itervalues():
center = cpv.add(center, coord)
count += 1
if count == 0:
print 'warning: count == 0'
raise CmdException
vec = cpv.normalize(vec)
center = cpv.scale(center, 1./count)
_common_orientation(selection, center, vec, visualize, 2.0*len(coords), quiet)
return center, vec
def _cgo_quad(pos, normal, radius):
'''Return a CGO list specifying a quad.'''
v1 = cpv.normalize(_perp_vec(normal))
v2 = cpv.cross_product(normal, v1)
v1 = cpv.scale(v1, radius)
v2 = cpv.scale(v2, radius)
obj = [ cgo.BEGIN,
cgo.TRIANGLE_STRIP,
cgo.NORMAL]
obj.extend(normal)
obj.append(cgo.VERTEX)
obj.extend(cpv.add(pos, v1))
obj.append(cgo.VERTEX)
obj.extend(cpv.add(pos, v2))
obj.append(cgo.VERTEX)
obj.extend(cpv.sub(pos, v2))
obj.append(cgo.VERTEX)
obj.extend(cpv.sub(pos, v1))
obj.append(cgo.END)
return obj
def helix_orientation_hbond(selection, visualize=1, cutoff=3.5, quiet=0):
'''
DESCRIPTION
Get the center and direction of a helix as vectors. Will only work
for alpha helices and gives slightly different results than
helix_orientation. Averages direction of O(i)->N(i+4) hydrogen bonds.
USAGE
helix_orientation selection [, visualize [, cutoff]]
ARGUMENTS
cutoff = float: maximal hydrogen bond distance {default: 3.5}
SEE ALSO
helix_orientation
'''
visualize, quiet, cutoff = int(visualize), int(quiet), float(cutoff)
stored.x = dict()
cmd.iterate_state(STATE, '(%s) and name N+O' % (selection),
'stored.x.setdefault(resv, dict())[name] = x,y,z')
vec_list = []
for resi in stored.x:
resi_other = resi + 4
if 'O' in stored.x[resi] and resi_other in stored.x:
if 'N' in stored.x[resi_other]:
vec = cpv.sub(stored.x[resi_other]['N'], stored.x[resi]['O'])
if cpv.length(vec) < cutoff:
vec_list.append(vec)
if len(vec_list) == 0:
print 'warning: count == 0'
raise CmdException
vec = _vec_sum(vec_list)
vec = cpv.normalize(vec)
return _common_orientation(selection, vec, visualize, quiet)
def helix_orientation_hbond(selection, visualize=1, cutoff=3.5, quiet=0):
'''
DESCRIPTION
Get the center and direction of a helix as vectors. Will only work
for alpha helices and gives slightly different results than
helix_orientation. Averages direction of O(i)->N(i+4) hydrogen bonds.
USAGE
helix_orientation selection [, visualize [, cutoff]]
ARGUMENTS
cutoff = float: maximal hydrogen bond distance {default: 3.5}
SEE ALSO
helix_orientation
'''
visualize, quiet, cutoff = int(visualize), int(quiet), float(cutoff)
stored.x = dict()
cmd.iterate_state(STATE, '(%s) and name N+O' % (selection),
'stored.x.setdefault(resv, dict())[name] = x,y,z')
vec_list = []
for resi in stored.x:
resi_other = resi + 4
if 'O' in stored.x[resi] and resi_other in stored.x:
if 'N' in stored.x[resi_other]:
vec = cpv.sub(stored.x[resi_other]['N'], stored.x[resi]['O'])
if cpv.length(vec) < cutoff:
vec_list.append(vec)
if len(vec_list) == 0:
print 'warning: count == 0'
raise CmdException
vec = _vec_sum(vec_list)
vec = cpv.normalize(vec)
return _common_orientation(selection, vec, visualize, quiet)
def planeFromPoints(p1, p2, p3, vm1=None, vm2=None, center=True, settings={}):
v1 = cpv.sub(p1, p2)
v2 = cpv.sub(p3, p2)
normal = cpv.cross_product(v1, v2)
if 'translate' in settings:
vtran = cpv.scale(cpv.normalize(normal), settings['translate'])
p1_t = cpv.sub(p1, vtran)
p2_t = cpv.sub(p2, vtran)
p3_t = cpv.sub(p3, vtran)
print("New coordinates are:")
print_info("New", p1_t, p2_t, p3_t)
print("New coordinates are for normalized plane:")
v1_t = cpv.normalize(cpv.sub(p1_t, p2_t))
v2_t = cpv.normalize(cpv.sub(p3_t, p2_t))
normal_t = cpv.normalize(cpv.cross_product(v1_t, v2_t))
v2_t = cpv.normalize(cpv.cross_product(normal_t, v1_t))
p1_t2 = cpv.add(v1_t, p2_t)
p3_t2 = cpv.add(v2_t, p2_t)
print_info("Newnormal", p1_t2, p2_t, p3_t2)
if vm1 != None:
v1 = cpv.scale(cpv.normalize(v1), vm1)
if vm2 != None:
v2 = cpv.scale(cpv.normalize(v2), vm2)
centrum = p2
if center:
corner1 = cpv.add(cpv.add(centrum, v1), v2)
corner2 = cpv.sub(cpv.add(centrum, v1), v2)
corner3 = cpv.sub(cpv.sub(centrum, v1), v2)
corner4 = cpv.add(cpv.sub(centrum, v1), v2)
else:
corner1 = cpv.add(cpv.add(centrum, v1), v2)
corner2 = cpv.add(centrum, v1)
corner3 = centrum
corner4 = cpv.add(centrum, v2)
return plane(corner1, corner2, corner3, corner4, normal, settings)
def _cgo_arrow(position=(0., 0., 0),
direction=(1., 0., 0.),
length=1.,
radius=.2,
offset=0.,
color='red'):
'''
Generate and return an arrow CGO from (position + offset * direction) to
(position + ((length + offset) * direction)).
'''
if isinstance(color, str):
color = list(pymol.cmd.get_color_tuple(color))
hlength = radius * 2.5
hradius = radius * 2.0
normal = cpv.normalize(direction)
xyz1 = cpv.add(position, cpv.scale(normal, offset))
xyz2 = cpv.add(position, cpv.scale(normal, length + offset))
xyz3 = cpv.add(xyz2, cpv.scale(normal, hlength))
return [cgo.CYLINDER] + xyz1 + xyz2 + [radius] + color + color + \
[cgo.CONE] + xyz2 + xyz3 + [hradius, 0.0] + color + color + \
[1.0, 0.0]
def helix_orientation(selection, state=STATE, visualize=1, cutoff=3.5, quiet=1):
'''
DESCRIPTION
Get the center and direction of a helix as vectors. Will only work
for alpha helices and gives slightly different results than
cafit_orientation. Averages direction of C(i)->O(i)->N(i+4).
USAGE
helix_orientation selection [, visualize [, cutoff ]]
ARGUMENTS
selection = string: atom selection of helix
visualize = 0 or 1: show fitted vector as arrow {default: 1}
cutoff = float: maximal hydrogen bond distance {default: 3.5}
SEE ALSO
angle_between_helices, loop_orientation, cafit_orientation
'''
state, visualize, quiet = int(state), int(visualize), int(quiet)
cutoff = float(cutoff)
atoms = {'C': dict(), 'O': dict(), 'N': dict()}
cmd.iterate_state(state, '(%s) and name N+O+C' % (selection),
'atoms[name][resv] = x,y,z', space={'atoms': atoms})
vec_list = []
for resi in atoms['C']:
resi_other = resi + 4
try:
aC = atoms['C'][resi]
aO = atoms['O'][resi]
aN = atoms['N'][resi_other]
except KeyError:
continue
dist = cpv.distance(aN, aO)
dist_weight = 1. - (2.8 - dist)
angle = cpv.get_angle_formed_by(aC, aO, aN)
angle_weight = 1. - (3.1 - angle)
if dist_weight > 0.0 and angle_weight > 0.0:
if not quiet:
print ' weight:', angle_weight * dist_weight
vec = cpv.scale(cpv.sub(aN, aC), angle_weight * dist_weight)
vec_list.append(vec)
if len(vec_list) == 0:
print 'warning: count == 0'
raise CmdException
center = cpv.scale(_vec_sum(atoms['O'].itervalues()), 1./len(atoms['O']))
vec = _vec_sum(vec_list)
vec = cpv.normalize(vec)
_common_orientation(selection, center, vec, visualize, 1.5*len(vec_list), quiet)
return center, vec
def torus(center=(0., 0., 0.), normal=(0., 0., 1.), radius=1., color='',
cradius=.25, samples=20, csamples=20):
'''
Generate and return a torus CGO with given center, normal
and ring radius.
'''
from math import cos, sin, pi
if color and isinstance(color, str):
color = list(cmd.get_color_tuple(color))
obj = []
axis = cpv.cross_product(normal, (0., 0., 1.))
angle = -cpv.get_angle(normal, (0., 0., 1.))
matrix = cpv.rotation_matrix(angle, cpv.normalize(axis))
obj_vertex = lambda x, y, z: obj.extend([VERTEX] + cpv.add(center,
cpv.transform(matrix, [x, y, z])))
obj_normal = lambda x, y, z: obj.extend([NORMAL] +
cpv.transform(matrix, [x, y, z]))
r = radius
cr = cradius
rr = 1.5 * cr
dv = 2 * pi / csamples
dw = 2 * pi / samples
v = 0.0
w = 0.0
while w < 2 * pi:
v = 0.0
c_w = cos(w)
s_w = sin(w)
c_wdw = cos(w + dw)
s_wdw = sin(w + dw)
obj.append(BEGIN)
obj.append(TRIANGLE_STRIP)
if color:
obj.append(COLOR)
obj.extend(color)
while v < 2 * pi + dv:
c_v = cos(v)
s_v = sin(v)
c_vdv = cos(v + dv)
s_vdv = sin(v + dv)
obj_normal(
(r + rr * c_v) * c_w - (r + cr * c_v) * c_w,
(r + rr * c_v) * s_w - (r + cr * c_v) * s_w,
(rr * s_v - cr * s_v))
obj_vertex(
(r + cr * c_v) * c_w,
(r + cr * c_v) * s_w,
cr * s_v)
obj_normal(
(r + rr * c_vdv) * c_wdw - (r + cr * c_vdv) * c_wdw,
(r + rr * c_vdv) * s_wdw - (r + cr * c_vdv) * s_wdw,
rr * s_vdv - cr * s_vdv)
obj_vertex(
(r + cr * c_vdv) * c_wdw,
(r + cr * c_vdv) * s_wdw,
cr * s_vdv)
v += dv
obj.append(END)
w += dw
return obj
def rebuild(self) -> None:
"""
Rebuilds torus
"""
obj = []
axis = cpv.cross_product(self.normal.array(), (0., 0., 1.))
angle = -cpv.get_angle(self.normal.array(), (0., 0., 1.))
matrix = cpv.rotation_matrix(angle, cpv.normalize(axis))
def obj_vertex(x, y, z):
return [cgo.VERTEX] + cpv.add(self.center.array(),
cpv.transform(matrix, [x, y, z]))
def obj_normal(x, y, z):
return [cgo.NORMAL] + cpv.transform(matrix, [x, y, z])
r = self.radius
cr = self.cradius
rr = 1.5 * cr
dv = 2 * math.pi / self.csamples
dw = 2 * math.pi / self.samples
v = 0.0
w = 0.0
while w < 2 * math.pi:
v = 0.0
c_w = math.cos(w)
s_w = math.sin(w)
c_wdw = math.cos(w + dw)
s_wdw = math.sin(w + dw)
obj.append(cgo.BEGIN)
obj.append(cgo.TRIANGLE_STRIP)
obj.append(cgo.COLOR)
obj.extend(self.color.array())
while v < 2 * math.pi + dv:
c_v = math.cos(v)
s_v = math.sin(v)
c_vdv = math.cos(v + dv)
s_vdv = math.sin(v + dv)
obj.extend(
obj_normal((r + rr * c_v) * c_w - (r + cr * c_v) * c_w,
(r + rr * c_v) * s_w - (r + cr * c_v) * s_w,
(rr * s_v - cr * s_v)))
obj.extend(
obj_vertex((r + cr * c_v) * c_w, (r + cr * c_v) * s_w,
cr * s_v))
obj.extend(
obj_normal(
(r + rr * c_vdv) * c_wdw - (r + cr * c_vdv) * c_wdw,
(r + rr * c_vdv) * s_wdw - (r + cr * c_vdv) * s_wdw,
rr * s_vdv - cr * s_vdv))
obj.extend(
obj_vertex((r + cr * c_vdv) * c_wdw,
(r + cr * c_vdv) * s_wdw, cr * s_vdv))
v += dv
obj.append(cgo.END)
w += dw
self._data = obj
def helix_orientation(selection,
state=STATE,
visualize=1,
cutoff=3.5,
quiet=1):
'''
DESCRIPTION
Get the center and direction of a helix as vectors. Will only work
for alpha helices and gives slightly different results than
cafit_orientation. Averages direction of C(i)->O(i)->N(i+4).
USAGE
helix_orientation selection [, visualize [, cutoff ]]
ARGUMENTS
selection = string: atom selection of helix
visualize = 0 or 1: show fitted vector as arrow {default: 1}
cutoff = float: maximal hydrogen bond distance {default: 3.5}
SEE ALSO
angle_between_helices, loop_orientation, cafit_orientation
'''
state, visualize, quiet = int(state), int(visualize), int(quiet)
cutoff = float(cutoff)
atoms = {'C': dict(), 'O': dict(), 'N': dict()}
cmd.iterate_state(state,
'(%s) and name N+O+C' % (selection),
'atoms[name][resv] = x,y,z',
space={'atoms': atoms})
vec_list = []
for resi in atoms['C']:
resi_other = resi + 4
try:
aC = atoms['C'][resi]
aO = atoms['O'][resi]
aN = atoms['N'][resi_other]
except KeyError:
continue
dist = cpv.distance(aN, aO)
dist_weight = 1. - (2.8 - dist)
angle = cpv.get_angle_formed_by(aC, aO, aN)
angle_weight = 1. - (3.1 - angle)
if dist_weight > 0.0 and angle_weight > 0.0:
if not quiet:
print(' weight:', angle_weight * dist_weight)
vec = cpv.scale(cpv.sub(aN, aC), angle_weight * dist_weight)
vec_list.append(vec)
if len(vec_list) == 0:
print('warning: count == 0')
raise CmdException
center = cpv.scale(_vec_sum(atoms['O'].values()), 1. / len(atoms['O']))
vec = _vec_sum(vec_list)
vec = cpv.normalize(vec)
_common_orientation(selection, center, vec, visualize, 1.5 * len(vec_list),
quiet)
return center, vec
def torus(center=(0., 0., 0.), normal=(0., 0., 1.), radius=1., color='',
cradius=.25, samples=20, csamples=20, _self=cmd):
'''
Generate and return a torus CGO with given center, normal
and ring radius.
'''
from math import cos, sin, pi
if color and isinstance(color, str):
color = list(_self.get_color_tuple(color))
obj = []
axis = cpv.cross_product(normal, (0., 0., 1.))
angle = -cpv.get_angle(normal, (0., 0., 1.))
matrix = cpv.rotation_matrix(angle, cpv.normalize(axis))
obj_vertex = lambda x, y, z: obj.extend([VERTEX] + cpv.add(center,
cpv.transform(matrix, [x, y, z])))
obj_normal = lambda x, y, z: obj.extend([NORMAL] +
cpv.transform(matrix, [x, y, z]))
r = radius
cr = cradius
rr = 1.5 * cr
dv = 2 * pi / csamples
dw = 2 * pi / samples
v = 0.0
w = 0.0
while w < 2 * pi:
v = 0.0
c_w = cos(w)
s_w = sin(w)
c_wdw = cos(w + dw)
s_wdw = sin(w + dw)
obj.append(BEGIN)
obj.append(TRIANGLE_STRIP)
if color:
obj.append(COLOR)
obj.extend(color)
while v < 2 * pi + dv:
c_v = cos(v)
s_v = sin(v)
c_vdv = cos(v + dv)
s_vdv = sin(v + dv)
obj_normal(
(r + rr * c_v) * c_w - (r + cr * c_v) * c_w,
(r + rr * c_v) * s_w - (r + cr * c_v) * s_w,
(rr * s_v - cr * s_v))
obj_vertex(
(r + cr * c_v) * c_w,
(r + cr * c_v) * s_w,
cr * s_v)
obj_normal(
(r + rr * c_vdv) * c_wdw - (r + cr * c_vdv) * c_wdw,
(r + rr * c_vdv) * s_wdw - (r + cr * c_vdv) * s_wdw,
rr * s_vdv - cr * s_vdv)
obj_vertex(
(r + cr * c_vdv) * c_wdw,
(r + cr * c_vdv) * s_wdw,
cr * s_vdv)
v += dv
obj.append(END)
w += dw
return obj
def plane_orientation(selection, state=STATE, visualize=1, guide=0, quiet=1):
'''
DESCRIPTION
Fit plane (for example beta-sheet). Can also be used with
angle_between_helices (even though this does not fit helices).
Returns center and normal vector of plane.
'''
try:
import numpy
except ImportError:
print(' Error: numpy not available')
raise CmdException
state, visualize, quiet = int(state), int(visualize), int(quiet)
if int(guide):
selection = '(%s) and guide' % (selection)
coords = list()
cmd.iterate_state(state,
selection,
'coords.append([x,y,z])',
space=locals())
if len(coords) < 3:
print('not enough guide atoms in selection')
raise CmdException
x = numpy.array(coords)
U, s, Vh = numpy.linalg.svd(x - x.mean(0))
# normal vector of plane is 3rd principle component
vec = cpv.normalize(Vh[2])
if cpv.dot_product(vec, x[-1] - x[0]) < 0:
vec = cpv.negate(vec)
center = x.mean(0).tolist()
_common_orientation(selection, center, vec, visualize, 4.0, quiet)
# plane visualize
if visualize:
from pymol import cgo
dir1 = cpv.normalize(Vh[0])
dir2 = cpv.normalize(Vh[1])
sx = [max(i / 4.0, 2.0) for i in s]
obj = [cgo.BEGIN, cgo.TRIANGLES, cgo.COLOR, 0.5, 0.5, 0.5]
for vertex in [
cpv.scale(dir1, sx[0]),
cpv.scale(dir2, sx[1]),
cpv.scale(dir2, -sx[1]),
cpv.scale(dir1, -sx[0]),
cpv.scale(dir2, -sx[1]),
cpv.scale(dir2, sx[1]),
]:
obj.append(cgo.VERTEX)
obj.extend(cpv.add(center, vertex))
obj.append(cgo.END)
cmd.load_cgo(obj, get_unused_name('planeFit'))
return center, vec
def read_moestr(contents,
object,
state=0,
finish=1,
discrete=1,
quiet=1,
zoom=-1,
_self=cmd):
moestr = contents
name = object
import sys
if sys.version_info[0] > 2 and isinstance(moestr, bytes):
moestr = moestr.decode()
cmd = _self
mr = MOEReader()
mr.appendFromStr(moestr)
split_chains = cmd.get_setting_int("moe_separate_chains")
cmd.group(name)
if hasattr(mr, 'system'):
have_valences = 0
chain_count = 0
cmd.set_color("_aseg0", [1.0, 1.0, 1.0])
aseg_color = cmd.get_color_index("_aseg0")
aseg_flag = 0
aseg_rep = {}
model = Indexed()
molecule = mr.system['molecule']
if 'atoms' in molecule:
n_atom = molecule['atoms']
model.atom = [Atom() for x in range(n_atom)]
residues = {}
chains = {}
for columns, data in molecule['attr']:
for row in data:
cur_atom = None
for key, value in zip(columns, row):
key = key[0]
if key == 'ID':
ID = value
else:
aProp = _atom_prop_map.get(key, None)
if aProp != None:
setattr(model.atom[ID - 1], aProp, value)
else:
xyz = _atom_coord_map.get(key, None)
if xyz != None:
coord = list(model.atom[ID - 1].coord)
coord[xyz] = value
model.atom[ID - 1].coord = coord
elif key in _atom_vis_map:
atom = model.atom[ID - 1]
if hasattr(atom, 'visible'):
visible = atom.visible
else:
visible = _default_visible
if key == 'aBondLook':
if value == 'cylinder':
atom.visible = 0x00000001 | visible
elif value == 'line':
atom.visible = 0x00000080 | visible
elif value == 'none':
atom.visible = -129 & Visible # 0xFFFFFF7F
elif key == 'aNucleusLook':
if value == 'sphere':
atom.visible = 0x00000002 | visible
elif value == 'small-sphere': # need to set sphere_scale=0.2 for these atoms
atom.visible = 0x00000002 | visible
atom.sphere_scale = 0.2
elif value == 'point': # nonbonded
atom.visible = 0x00000800 | visible
elif value == 'none':
atom.visible = -2067 & visible # 0xFFFFF7ED
elif key == 'aHidden':
atom.visible = 0
atom.hidden = 1
if hasattr(
atom, 'hidden'
): # be sure that hidden atoms aren't shown
atom.visible = 0
elif key in _atom_color_map:
if key == 'aRGB':
model.atom[ID - 1].trgb = value
elif key == 'aColorBy':
model.atom[ID - 1].aColorBy = value
elif key in _atom_label_map:
atom = model.atom[ID - 1]
if hasattr(atom, 'label_dict'):
atom.label_dict[key] = None
else:
atom.label_dict = {key: None}
elif key in _residue_prop_map:
resi_dict = residues.get(ID, {})
resi_dict[key] = value
residues[ID] = resi_dict
elif key in _chain_prop_map:
chain_dict = chains.get(ID, {})
if ID not in chains:
chain_count = chain_count + 1
chain_dict['count'] = chain_count
chain_dict[key] = value
chains[ID] = chain_dict
chn_keys = list(chains.keys())
chn_keys.sort()
res_keys = list(residues.keys())
res_keys.sort()
# map chain properties onto residues
chn_resi = 0
ch_colors = copy.deepcopy(_ch_colors)
unique_chain_names = {}
for chn_idx in chn_keys:
chain_dict = chains[chn_idx]
cName = make_valid_name(chain_dict.get('cName', ''))
segi = cName[0:4]
chain = cName[-1:]
if not len(cName):
if 'count' in chain_dict:
cName = "chain_" + str(chain_dict['count'])
else:
cName = str(chn_idx)
if cName not in unique_chain_names:
unique_chain_names[cName] = cName
else:
cnt = 2
while (cName + "_" + str(cnt)) in unique_chain_names:
cnt = cnt + 1
newCName = cName + "_" + str(cnt)
unique_chain_names[newCName] = cName
cName = newCName
chain_dict['chain_color'] = ch_colors[0]
ch_colors = ch_colors[1:] + [ch_colors[0]]
cResCount = chain_dict.get('cResidueCount', 0)
for res_idx in range(chn_resi, chn_resi + cResCount):
resi_dict = residues[res_keys[res_idx]]
resi_dict['chain'] = chain
resi_dict['segi'] = segi
resi_dict['cName'] = cName
resi_dict['chain_dict'] = chain_dict
chn_resi = chn_resi + cResCount
# map residue properties onto atoms
res_atom = 0
for res_idx in res_keys:
resi_dict = residues[res_idx]
rRibbonMode = resi_dict.get('rRibbonMode', 'none')
rAtomCount = resi_dict['rAtomCount']
rType = resi_dict.get('rType', '')
if rAtomCount > 0:
for at_idx in range(res_atom, res_atom + rAtomCount):
atom = model.atom[at_idx]
setattr(
atom, 'resi',
string.strip(
str(resi_dict.get('rUID', '')) +
resi_dict.get('rINS', '')))
setattr(atom, 'resn', resi_dict.get('rName', ''))
setattr(atom, 'chain', resi_dict.get('chain', ''))
setattr(atom, 'segi', resi_dict.get('segi', ''))
setattr(atom, 'custom', resi_dict.get('cName', ''))
# add labels
if hasattr(atom, 'label_dict'):
label = ''
label_dict = atom.label_dict
if 'aLabelElement' in label_dict:
label = atom.symbol
if 'aLabelRes' in label_dict:
if len(label):
label = label + ","
label = label + atom.resn + "_" + atom.resi
if 'aLabelName' in label_dict:
if len(label):
label = label + "."
label = label + atom.name
atom.label = label
if rType not in ['none', 'heme']:
atom.hetatm = 0 # all normal atoms
else:
atom.flags = 0x02000000 # hetatom or ligand -- ignore when surfacing
if rRibbonMode != 'none':
if hasattr(atom, 'visible'):
visible = atom.visible
else:
visible = _default_visible
rRibbonColorBy = resi_dict['rRibbonColorBy']
repeat = 1
while repeat:
repeat = 0
if rRibbonColorBy in ['rgb', 'r:rgb'
]: # handled automatically
pass
elif rRibbonColorBy == 'chain':
chain_dict = resi_dict['chain_dict']
cColorBy = chain_dict['cColorBy']
if cColorBy == 'rgb':
cColorBy = 'c:rgb'
rRibbonColorBy = cColorBy
repeat = 1
rRibbon_color = 0
rRibbon_trgb = 0xFFFFFF # default -- white
# now color ribbon
if rRibbonColorBy == 'r:rgb':
rRibbon_trgb = resi_dict.get('rRGB', 0xFFFFFF)
elif rRibbonColorBy == 'rgb':
rRibbon_trgb = resi_dict.get(
'rRibbonRGB', 0xFFFFFF)
elif rRibbonColorBy == 'c:rgb':
chain_dict = resi_dict['chain_dict']
rRibbon_trgb = chain_dict.get('cRGB', 0xFFFFFF)
elif rRibbonColorBy == 'r:aseg':
rRibbon_trgb = None
rRibbon_color = aseg_color
aseg_flag = 1
elif rRibbonColorBy == 'tempfactor':
pass # TO DO
elif rRibbonColorBy == 'c:number': # per chain colors
rRibbon_trgb = chain_dict['chain_color']
if rRibbonMode in ['line', 'trace']:
atom.visible = 0x00000040 | visible # PyMOL ribbon
if rRibbon_trgb != None:
atom.ribbon_trgb = rRibbon_trgb
else:
atom.ribbon_color = rRibbon_color
aseg_rep['ribbon'] = 1
else:
atom.visible = 0x00000020 | visible # PyMOL cartoon
if rRibbon_trgb != None:
atom.cartoon_trgb = rRibbon_trgb
else:
atom.cartoon_color = rRibbon_color
aseg_rep['cartoon'] = 1
if hasattr(atom, 'label'):
if hasattr(atom, 'visible'):
visible = atom.visible
else:
visible = _default_visible
atom.visible = 0x00000028 | visible # labels
if not hasattr(atom, 'aColorBy'):
atom.aColorBy = 'element'
if hasattr(atom, 'aColorBy'):
aColorBy = atom.aColorBy
repeat = 1
while repeat:
repeat = 0
if aColorBy == 'ribbon':
aColorBy = resi_dict.get('rRibbonColorBy')
if aColorBy == 'rgb':
aColorBy = 'rib:rgb'
else:
repeat = 1
# TO DO still need to handle "cartoon_color", "ribbon_color"
elif aColorBy == 'element':
if hasattr(atom, 'trgb'):
del atom.trgb
elif aColorBy in ['rgb', 'a:rgb'
]: # handled automatically
pass
elif aColorBy == 'residue':
rColorBy = resi_dict.get('rColorBy')
if rColorBy == 'rgb':
rColorBy = 'r:rgb'
aColorBy = rColorBy
repeat = 1
elif aColorBy == 'chain':
chain_dict = resi_dict['chain_dict']
cColorBy = chain_dict['cColorBy']
if cColorBy == 'rgb':
cColorBy = 'c:rgb'
aColorBy = cColorBy
repeat = 1
# now color atom...
if aColorBy == 'r:rgb':
atom.trgb = resi_dict.get('rRGB', 0xFFFFFF)
elif aColorBy == 'rib:rgb':
atom.trgb = resi_dict.get('rRibbonRGB', 0xFFFFFF)
elif aColorBy == 'c:rgb':
chain_dict = resi_dict['chain_dict']
atom.trgb = chain_dict.get('cRGB', 0xFFFFFF)
elif aColorBy == 'r:aseg':
pass # TO DO
elif aColorBy == 'tempfactor':
pass # TO DO
elif aColorBy == 'c:number': # per chain colors
atom.trgb = chain_dict['chain_color']
res_atom = res_atom + rAtomCount
bond_list = molecule.get('bond', [])
for bond in bond_list:
new_bond = Bond()
new_bond.index = [bond[0] - 1, bond[1] - 1]
if len(bond) > 2:
new_bond.order = bond[2]
if bond[2] == 2: # work around .MOE bug with triple bonds
if model.atom[new_bond.index[0]].hybridization == 'sp':
if model.atom[new_bond.index[1]].hybridization == 'sp':
new_bond.order = 3
have_valences = 1
model.bond.append(new_bond)
if 'ViewOrientationY' in mr.system:
vy = mr.system['ViewOrientationY']
vz = mr.system['ViewOrientationZ']
pos = mr.system['ViewLookAt']
scale = mr.system['ViewScale']
vx = cpv.cross_product(vy, vz)
m = [cpv.normalize(vx), cpv.normalize(vy), cpv.normalize(vz)]
m = cpv.transpose(m)
asp_rat = 0.8 # MOE default (versus 0.75 for PyMOL)
cmd.set("field_of_view", 25.0)
fov = float(cmd.get("field_of_view"))
window_height = scale * asp_rat
dist = (0.5 * window_height) / math.atan(3.1415 *
(0.5 * fov) / 180.0)
new_view = tuple(m[0] + m[1] + m[2] + [0.0, 0.0, -dist] + pos +
[dist * 0.5, dist * 1.5, 0.0])
cmd.set_view(new_view)
zoom = 0
cmd.set("auto_color", 0)
cmd.set_color("carbon", [0.5, 0.5, 0.5]) # default MOE grey
obj_name = name + ".system"
if split_chains < 0: # by default, don't split chains if over 50 objects would be created
if len(unique_chain_names) > 50:
split_chains = 0
if not split_chains:
cmd.load_model(model,
obj_name,
state=state,
finish=finish,
discrete=discrete,
quiet=quiet,
zoom=zoom)
obj_name_list = [obj_name]
else:
cmd.load_model(model,
obj_name,
state=state,
finish=finish,
discrete=discrete,
quiet=1,
zoom=zoom)
obj_name_list = []
system_name = obj_name
for chain in unique_chain_names.keys():
obj_name = name + "." + chain
obj_name_list.append(obj_name)
cmd.select("_moe_ext_tmp",
"custom %s" % chain,
domain=system_name)
cmd.extract(obj_name, "_moe_ext_tmp", quiet=quiet, zoom=0)
# cmd.extract(obj_name,system_name+" and text_type %s"%chain,quiet=quiet)
cmd.delete("_moe_ext_tmp")
if not cmd.count_atoms(system_name):
cmd.delete(system_name)
else:
obj_name_list.append(system_name)
cmd.order(name + ".*", sort=1)
for obj_name in obj_name_list:
cmd.set("stick_radius", 0.1, obj_name)
cmd.set("line_width", 2.0, obj_name)
cmd.set("label_color", "white", obj_name)
cmd.set("nonbonded_size", 0.05,
obj_name) # temporary workaround...
if have_valences: # if this MOE file has valences, then show em!
cmd.set("valence", 1, obj_name)
cmd.set("stick_valence_scale", 1.25, obj_name)
if aseg_flag:
cmd.dss(obj_name)
if 'cartoon' in aseg_rep:
cmd.set("cartoon_color", "red",
obj_name + " and cartoon_color _aseg0 and ss h")
cmd.set("cartoon_color", "yellow",
obj_name + " and cartoon_color _aseg0 and ss s")
cmd.set(
"cartoon_color", "cyan",
obj_name + " and cartoon_color _aseg0 and not ss h+s")
if 'ribbon' in aseg_rep:
cmd.set("ribbon_color", "red",
obj_name + " and ribbon_color _aseg0 and ss h"
) # need selection op ribbon_color
cmd.set("ribbon_color", "yellow",
obj_name + " and ribbon_color _aseg0 and ss s")
cmd.set(
"ribbon_color", "cyan",
obj_name + " and ribbon_color _aseg0 and not ss h+s")
if 'ViewZFront' in mr.system:
moe_front = mr.system['ViewZFront']
moe_width = mr.system['ViewZWidth']
extent = cmd.get_extent(
name) # will this work with groups? TO DO: check!
dx = (extent[0][0] - extent[1][0])
dy = (extent[0][1] - extent[1][1])
dz = (extent[0][2] - extent[1][2])
half_width = 0.5 * math.sqrt(dx * dx + dy * dy + dz * dz)
cmd.clip("atoms", 0.0, name)
cur_view = cmd.get_view()
center = (cur_view[-3] +
cur_view[-2]) * 0.5 # center of clipping slab
front = center - half_width
back = center + half_width
width = half_width * 2
new_view = tuple(
list(cur_view[0:15]) + [
front + width * moe_front, front + width *
(moe_front + moe_width), 0.0
])
cmd.set_view(new_view)
if 'graphics' in mr.system:
cgo_cnt = 1
lab_cnt = 1
unique_cgo_names = {}
for graphics in mr.system['graphics']:
cgo = []
for gvertex in graphics.get('gvertex', []):
vrt = gvertex[0]
seg_list = gvertex[1]['seg']
idx = gvertex[1]['idx']
len_idx = len(idx)
if not cmd.is_list(seg_list):
seg_list = [seg_list] * (len_idx / seg_list)
last_seg = None
ix_start = 0
for seg in seg_list:
if seg != last_seg:
if last_seg != None:
cgo.append(END)
if seg == 3:
cgo.extend([BEGIN, TRIANGLES])
elif seg == 2:
cgo.extend([BEGIN, LINES])
elif seg == 1:
cgo.extend([BEGIN, POINTS])
ix_stop = seg + ix_start
if seg == 3:
for s in idx[ix_start:ix_stop]:
v = vrt[s - 1]
cgo.extend([
COLOR, (0xFF & (v[0] >> 16)) / 255.0,
(0xFF & (v[0] >> 8)) / 255.0,
(0xFF & (v[0])) / 255.0
])
if len(v) > 4:
cgo.extend([NORMAL, v[4], v[5], v[6]])
cgo.extend([VERTEX, v[1], v[2], v[3]])
elif seg == 2:
for s in idx[ix_start:ix_stop]:
v = vrt[s - 1]
cgo.extend([
COLOR, (0xFF & (v[0] >> 16)) / 255.0,
(0xFF & (v[0] >> 8)) / 255.0,
(0xFF & (v[0])) / 255.0
])
if len(v) > 4:
cgo.extend([NORMAL, v[4], v[5], v[6]])
cgo.extend([VERTEX, v[1], v[2], v[3]])
elif seg == 1:
for s in idx[ix_start:ix_stop]:
v = vrt[s - 1]
cgo.extend([
COLOR, (0xFF & (v[0] >> 16)) / 255.0,
(0xFF & (v[0] >> 8)) / 255.0,
(0xFF & (v[0])) / 255.0
])
if len(v) > 4:
cgo.extend([NORMAL, v[4], v[5], v[6]])
cgo.extend([VERTEX, v[1], v[2], v[3]])
ix_start = ix_stop
last_seg = seg
if last_seg != None:
cgo.append(END)
for gtext in graphics.get('gtext', []):
lab_name = name + ".label_%02d" % lab_cnt
exists = 0
for entry in gtext:
exists = 1
cmd.pseudoatom(lab_name,
pos=[
float(entry[1]),
float(entry[2]),
float(entry[3])
],
color="0x%06x" % entry[0],
label=entry[4])
if exists:
cmd.set('label_color', -1, lab_name)
lab_cnt = lab_cnt + 1
# TO DO -- via CGO's?
if len(cgo):
cgo_name = name + "." + make_valid_name(
graphics.get('GTitle', 'graphics'))
if cgo_name not in unique_cgo_names:
unique_cgo_names[cgo_name] = cgo_name
else:
cnt = 2
while cgo_name + "_" + str(cnt) in unique_cgo_names:
cnt = cnt + 1
new_name = cgo_name + "_" + str(cnt)
unique_cgo_names[new_name] = new_name
cgo_name = new_name
cmd.load_cgo(cgo,
cgo_name,
state=state,
quiet=quiet,
zoom=0)
cgo_cnt = cgo_cnt + 1
cmd.set("two_sided_lighting", 1) # global setting...
cmd.set("cgo_line_width", 2, cgo_name)
if 'GTransparency' in graphics:
g_trans = graphics['GTransparency']
if len(g_trans) >= 2:
if g_trans[0] != 0:
cmd.set('cgo_transparency',
'%1.6f' % (g_trans[0] / 255.0),
cgo_name)
cmd.set('transparency_global_sort')
if 'meter' in mr.system:
meter_name = name + ".meter"
exists = 0
for meter_block in mr.system['meter']:
if meter_block[0][0:2] == ['type', 'atoms']:
for meter in meter_block[1]:
(type, atoms) = meter[0:2]
arg = tuple([meter_name] + list(
map(lambda x, o=name: o + " and id " + str(x - 1),
atoms)))
getattr(cmd, type)(*arg)
exists = 1
if exists:
cmd.color("green", meter_name)
# print mr.system['meter']
elif hasattr(mr, 'feature'):
model = Indexed()
cols = mr.feature[0]
rows = mr.feature[1]
col = {}
cnt = 0
for a in cols:
col[a] = cnt
cnt = cnt + 1
for row in rows:
atom = Atom()
atom.coord = [row[col['x']], row[col['y']], row[col['z']]]
atom.vdw = row[col['r']]
atom.custom = row[col['expr']]
model.atom.append(atom)
obj_name = name + ".feature"
cmd.load_model(model,
obj_name,
state=state,
finish=finish,
discrete=discrete,
quiet=quiet,
zoom=zoom)
rank = 1
for row in rows:
cmd.color("0x%06x" % row[col['color']],
obj_name + " & id %d" % (rank - 1))
rank = rank + 1
cmd.show("mesh", obj_name)
cmd.set("min_mesh_spacing", 0.55, obj_name)
cmd.label(obj_name, "' '+custom")
cmd.set("label_color", "yellow", obj_name)
else:
print(dir(mr))
def cgo_arrow(atom1='pk1',
atom2='pk2',
radius=0.5,
gap=0.0,
hlength=-1,
hradius=-1,
color='blue red',
name=''):
'''
DESCRIPTION
Create a CGO arrow between two picked atoms.
ARGUMENTS
atom1 = string: single atom selection or list of 3 floats {default: pk1}
atom2 = string: single atom selection or list of 3 floats {default: pk2}
radius = float: arrow radius {default: 0.5}
gap = float: gap between arrow tips and the two atoms {default: 0.0}
hlength = float: length of head
hradius = float: radius of head
color = string: one or two color names {default: blue red}
name = string: name of CGO object
'''
from chempy import cpv
radius, gap = float(radius), float(gap)
hlength, hradius = float(hlength), float(hradius)
try:
color1, color2 = color.split()
except:
color1 = color2 = color
color1 = list(cmd.get_color_tuple(color1))
color2 = list(cmd.get_color_tuple(color2))
def get_coord(v):
if not isinstance(v, str):
return v
if v.startswith('['):
return cmd.safe_list_eval(v)
return cmd.get_atom_coords(v)
xyz1 = get_coord(atom1)
xyz2 = get_coord(atom2)
normal = cpv.normalize(cpv.sub(xyz1, xyz2))
if hlength < 0:
hlength = radius * 3.0
if hradius < 0:
hradius = hlength * 0.6
if gap:
diff = cpv.scale(normal, gap)
xyz1 = cpv.sub(xyz1, diff)
xyz2 = cpv.add(xyz2, diff)
xyz3 = cpv.add(cpv.scale(normal, hlength), xyz2)
obj = [cgo.CYLINDER] + xyz1 + xyz3 + [radius] + color1 + color2 + \
[cgo.CONE] + xyz3 + xyz2 + [hradius, 0.0] + color2 + color2 + \
[1.0, 0.0]
#obj = [cgo.CYLINDER] + xyz1 + xyz3 + [radius] + color1 + color2 + \
#[cgo.CONE] + xyz3 + xyz2 + [hradius, 0.0] + color2 + color2 + \
#[1.0, 0.0]
if not name:
name = cmd.get_unused_name('arrow')
cmd.load_cgo(obj, name)
def cgo_arrow(atom1='pk1', atom2='pk2', radius=0.5, gap=0.0, hlength=-1, hradius=-1,
color='black black', name=''):
'''
#I modify the color the line just before
DESCRIPTION
Create a CGO arrow between two picked atoms.
ARGUMENTS
atom1 = string: single atom selection or list of 3 floats {default: pk1}
atom2 = string: single atom selection or list of 3 floats {default: pk2}
radius = float: arrow radius {default: 0.5}
gap = float: gap between arrow tips and the two atoms {default: 0.0}
hlength = float: length of head
hradius = float: radius of head
color = string: one or two color names {default: blue red}
name = string: name of CGO object
'''
from chempy import cpv
radius, gap = float(radius), float(gap)
hlength, hradius = float(hlength), float(hradius)
try:
color1, color2 = color.split()
except:
color1 = color2 = color
color1 = list(cmd.get_color_tuple(color1))
color2 = list(cmd.get_color_tuple(color2))
def get_coord(v):
if not isinstance(v, str):
return v
if v.startswith('['):
return cmd.safe_list_eval(v)
return cmd.get_atom_coords(v)
xyz1 = get_coord(atom1)
xyz2 = get_coord(atom2)
normal = cpv.normalize(cpv.sub(xyz1, xyz2))
if hlength < 0:
hlength = radius * 3.0
if hradius < 0:
hradius = hlength * 0.6
if gap:
diff = cpv.scale(normal, gap)
xyz1 = cpv.sub(xyz1, diff)
xyz2 = cpv.add(xyz2, diff)
xyz3 = cpv.add(cpv.scale(normal, hlength), xyz2)
obj = [cgo.CYLINDER] + xyz1 + xyz3 + [radius] + color1 + color2 + \
[cgo.CONE] + xyz3 + xyz2 + [hradius, 0.0] + color2 + color2 + \
[1.0, 0.0]
if not name:
name = cmd.get_unused_name('arrow')
cmd.load_cgo(obj, name)
def calculateNewPoint(p1, p2, distance):
v1 = cpv.normalize(cpv.sub(p1, p2))
return cpv.add(p1, cpv.scale(v1, distance))
[[0.0, 1.0, 0.0], [-1.0, 0.0, 0.0], 0],
[[-1.0, 0.0, 0.0], [0.0, -1.0, 0.0], 0],
[[0.0, -1.0, 0.0], [1.0, 0.0, 0.0], 0],
[[1.0, 0.0, 0.0], [0.0, 1.0, 0.0], 1],
[[0.0, 1.0, 0.0], [-1.0, 0.0, 0.0], 1],
[[-1.0, 0.0, 0.0], [0.0, -1.0, 0.0], 1],
[[0.0, -1.0, 0.0], [1.0, 0.0, 0.0], 1],
]:
hand = basis[2]
if hand:
if z1 == 0.0:
(x, y) = basis[0:2]
else:
(y, x) = basis[0:2]
normal = normalize(cross_product(x, y))
else:
if z1 == 0.0:
(y, x) = basis[0:2]
else:
(x, y) = basis[0:2]
normal = normalize(cross_product(y, x))
obj.extend([BEGIN, TRIANGLE_STRIP] + [COLOR, 1.0, 1.0, 1.0] +
[NORMAL] + normal)
for i in range(sampling + 1):
x1 = edge
y1 = edge * i / sampling
vlen = sqrt(x1 * x1 + y1 * y1)
x0 = radius * x1 / vlen
obj = []
for basis in [ [ [ 1.0, 0.0, 0.0], [ 0.0, 1.0, 0.0], 0],
[ [ 0.0, 1.0, 0.0], [ -1.0, 0.0, 0.0], 0],
[ [-1.0, 0.0, 0.0], [ 0.0, -1.0, 0.0], 0],
[ [ 0.0, -1.0, 0.0], [ 1.0, 0.0, 0.0], 0],
[ [ 1.0, 0.0, 0.0], [ 0.0, 1.0, 0.0], 1],
[ [ 0.0, 1.0, 0.0], [ -1.0, 0.0, 0.0], 1],
[ [-1.0, 0.0, 0.0], [ 0.0, -1.0, 0.0], 1],
[ [ 0.0, -1.0, 0.0], [ 1.0, 0.0, 0.0], 1] ]:
hand = basis[2]
if hand:
(x,y) = basis[0:2]
normal = normalize(cross_product(x,y))
else:
(y,x) = basis[0:2]
normal = normalize(cross_product(y,x))
obj.extend( [BEGIN, TRIANGLE_STRIP] +
[COLOR, 1.0, 1.0, 1.0] +
[NORMAL] + normal )
for i in range(sampling+1):
x1 = edge
y1 = edge*i/sampling
vlen = sqrt(x1*x1+y1*y1)
x0 = radius*x1/vlen
y0 = radius*y1/vlen
v0 = add( scale(x, x0 ), scale(y,y0) )
def bbPlane(selection='(all)', color='gray', transp=0.3, state=-1, name=None, quiet=1):
"""
DESCRIPTION
Draws a plane across the backbone for a selection
ARGUMENTS
selection = string: protein object or selection {default: (all)}
color = string: color name or number {default: white}
transp = float: transparency component (0.0--1.0) {default: 0.0}
state = integer: object state, 0 for all states {default: 1}
NOTES
You need to pass in an object or selection with at least two
amino acids. The plane spans CA_i, O_i, N-H_(i+1), and CA_(i+1)
"""
from pymol.cgo import BEGIN, TRIANGLES, COLOR, VERTEX, END
from pymol import cgo
from chempy import cpv
# format input
transp = float(transp)
state, quiet = int(state), int(quiet)
if name is None:
name = cmd.get_unused_name("backbonePlane")
if state < 0:
state = cmd.get_state()
elif state == 0:
for state in range(1, cmd.count_states(selection) + 1):
bbPlane(selection, color, transp, state, name, quiet)
return
AAs = []
coords = dict()
# need hydrogens on peptide nitrogen
cmd.h_add('(%s) and n. N' % selection)
# get the list of residue ids
for obj in cmd.get_object_list(selection):
sel = obj + " and (" + selection + ")"
for a in cmd.get_model(sel + " and n. CA", state).atom:
key = '/%s/%s/%s/%s' % (obj, a.segi, a.chain, a.resi)
AAs.append(key)
coords[key] = [a.coord, None, None]
for a in cmd.get_model(sel + " and n. O", state).atom:
key = '/%s/%s/%s/%s' % (obj, a.segi, a.chain, a.resi)
if key in coords:
coords[key][1] = a.coord
for a in cmd.get_model(sel + " and ((n. N extend 1 and e. H) or (r. PRO and n. CD))", state).atom:
key = '/%s/%s/%s/%s' % (obj, a.segi, a.chain, a.resi)
if key in coords:
coords[key][2] = a.coord
# need at least two amino acids
if len(AAs) <= 1:
print("ERROR: Please provide at least two amino acids, the alpha-carbon on the 2nd is needed.")
return
# prepare the cgo
obj = [
BEGIN, TRIANGLES,
COLOR,
]
obj.extend(cmd.get_color_tuple(color))
for res in range(0, len(AAs) - 1):
curIdx, nextIdx = str(AAs[res]), str(AAs[res + 1])
# populate the position array
pos = [coords[curIdx][0], coords[curIdx][1], coords[nextIdx][2], coords[nextIdx][0]]
# if the data are incomplete for any residues, ignore
if None in pos:
if not quiet:
print(' bbPlane: peptide bond %s -> %s incomplete' % (curIdx, nextIdx))
continue
if cpv.distance(pos[0], pos[3]) > 4.0:
if not quiet:
print(' bbPlane: %s and %s not adjacent' % (curIdx, nextIdx))
continue
normal = cpv.normalize(cpv.cross_product(
cpv.sub(pos[1], pos[0]),
cpv.sub(pos[2], pos[0])))
obj.append(cgo.NORMAL)
obj.extend(normal)
# need to order vertices to generate correct triangles for plane
if cpv.dot_product(cpv.sub(pos[0], pos[1]), cpv.sub(pos[2], pos[3])) < 0:
vorder = [0, 1, 2, 2, 3, 0]
else:
vorder = [0, 1, 2, 3, 2, 1]
# fill in the vertex data for the triangles;
for i in vorder:
obj.append(VERTEX)
obj.extend(pos[i])
# finish the CGO
obj.append(END)
# update the UI
cmd.load_cgo(obj, name, state, zoom=0)
cmd.set("cgo_transparency", transp, name)
def cgo_modevec(atom1='pk1', atom2='pk2', radius=0.05, gap=0.0, hlength=-1, hradius=-1,
color='green', name='',scalefactor=10.0, cutoff=0.6, transparency=1.0): #was 0.6cut 12 scale
## scalefactor was 10.0
'''
DESCRIPTION
Create a CGO mode vector starting at atom1 and pointint in atom2 displacement
ARGUMENTS
atom1 = string: single atom selection or list of 3 floats {default: pk1}
atom2 = string: displacement to atom1 for modevec
radius = float: arrow radius {default: 0.5}
gap = float: gap between arrow tips and the two atoms {default: 0.0}
hlength = float: length of head
hradius = float: radius of head
color = string: one or two color names {default: blue red}
name = string: name of CGO object
scalefactor = scale how big of an arrow to make. Default 5
transparency = 0.0 ~ 1.0, default=1.0 means being totally opaque
'''
from chempy import cpv
radius, gap = float(radius), float(gap)
hlength, hradius = float(hlength), float(hradius)
scalefactor, cutoff = float(scalefactor), float(cutoff)
transparency = float(transparency)
try:
color1, color2 = color.split()
except:
color1 = color2 = color
color1 = list(cmd.get_color_tuple(color1))
color2 = list(cmd.get_color_tuple(color2))
def get_coord(v):
if not isinstance(v, str):
return v
if v.startswith('['):
return cmd.safe_list_eval(v)
return cmd.get_atom_coords(v)
xyz1 = get_coord(atom1)
xyz2 = get_coord(atom2)
newxyz2 = cpv.scale(xyz2, scalefactor)
newxyz2 = cpv.add(newxyz2, xyz1)
xyz2 = newxyz2
# xyz2 = xyz2[0]*scalefactor, xyz2[1]*scalefactor, xyz2[2]*scalefactor
normal = cpv.normalize(cpv.sub(xyz1, xyz2))
if hlength < 0:
hlength = radius * 3.0
if hradius < 0:
hradius = hlength * 0.6
if gap:
diff = cpv.scale(normal, gap)
xyz1 = cpv.sub(xyz1, diff)
xyz2 = cpv.add(xyz2, diff)
xyz3 = cpv.add(cpv.scale(normal, hlength), xyz2)
# dont draw arrow if distance is too small
distance = cpv.distance(xyz1, xyz2)
if distance <= cutoff:
return
#### generate transparent arrows;
#### The original codes are the next block
#### --Ran
obj = [25.0, transparency, 9.0] + xyz1 + xyz3 + [radius] + color1 + color2 + \
[25.0, transparency, 27.0] + xyz3 + xyz2 + [hradius, 0.0] + color2 + color2 + \
[1.0, 0.0]
# obj = [cgo.CYLINDER] + xyz1 + xyz3 + [radius] + color1 + color2 + \
# [cgo.CONE] + xyz3 + xyz2 + [hradius, 0.0] + color2 + color2 + \
# [1.0, 0.0]
if not name:
name = cmd.get_unused_name('arrow')
cmd.load_cgo(obj, name)
def update_box(self):
if self.points_name in self.cmd.get_names():
model = self.cmd.get_model(self.points_name)
self.coord = (
model.atom[0].coord,
model.atom[1].coord,
model.atom[2].coord,
model.atom[3].coord,
)
p = self.coord[0]
d10 = sub(self.coord[1], p)
d20 = sub(self.coord[2], p)
d30 = sub(self.coord[3], p)
x10_20 = cross_product(d10,d20)
if self.mode != 'quad':
if dot_product(d30,x10_20)<0.0:
p = model.atom[1].coord
d10 = sub(self.coord[0], p)
d20 = sub(self.coord[2], p)
d30 = sub(self.coord[3], p)
n10_20 = normalize(x10_20)
n10 = normalize(d10)
d100 = d10
d010 = remove_component(d20, n10)
if self.mode != 'quad':
d001 = project(d30, n10_20)
else:
d001 = n10_20
n100 = normalize(d100)
n010 = normalize(d010)
n001 = normalize(d001)
f100 = reverse(n100)
f010 = reverse(n010)
f001 = reverse(n001)
if self.mode == 'quad':
p000 = p
p100 = add(p, remove_component(d10,n001))
p010 = add(p, remove_component(d20,n001))
p001 = add(p, remove_component(d30,n001))
else:
p000 = p
p100 = add(p,d100)
p010 = add(p,d010)
p001 = add(p,d001)
p110 = add(p100, d010)
p011 = add(p010, d001)
p101 = add(p100, d001)
p111 = add(p110, d001)
obj = []
if self.mode == 'box': # standard box
obj.extend([ BEGIN, TRIANGLE_STRIP ])
obj.append(NORMAL); obj.extend(f001)
obj.append(VERTEX); obj.extend(p000)
obj.append(VERTEX); obj.extend(p010)
obj.append(VERTEX); obj.extend(p100)
obj.append(VERTEX); obj.extend(p110)
obj.append(END)
obj.extend([ BEGIN, TRIANGLE_STRIP ])
obj.append(NORMAL); obj.extend(n001)
obj.append(VERTEX); obj.extend(p001)
obj.append(VERTEX); obj.extend(p101)
obj.append(VERTEX); obj.extend(p011)
obj.append(VERTEX); obj.extend(p111)
obj.append(END)
obj.extend([ BEGIN, TRIANGLE_STRIP ])
obj.append(NORMAL); obj.extend(f010)
obj.append(VERTEX); obj.extend(p000)
obj.append(VERTEX); obj.extend(p100)
obj.append(VERTEX); obj.extend(p001)
obj.append(VERTEX); obj.extend(p101)
obj.append(END)
obj.extend([ BEGIN, TRIANGLE_STRIP ])
obj.append(NORMAL); obj.extend(n010)
obj.append(VERTEX); obj.extend(p010)
obj.append(VERTEX); obj.extend(p011)
obj.append(VERTEX); obj.extend(p110)
obj.append(VERTEX); obj.extend(p111)
obj.append(END)
obj.extend([ BEGIN, TRIANGLE_STRIP ])
obj.append(NORMAL); obj.extend(f100)
obj.append(VERTEX); obj.extend(p000)
obj.append(VERTEX); obj.extend(p001)
obj.append(VERTEX); obj.extend(p010)
obj.append(VERTEX); obj.extend(p011)
obj.append(END)
obj.extend([ BEGIN, TRIANGLE_STRIP ])
obj.append(NORMAL); obj.extend(n100)
obj.append(VERTEX); obj.extend(p100)
obj.append(VERTEX); obj.extend(p110)
obj.append(VERTEX); obj.extend(p101)
obj.append(VERTEX); obj.extend(p111)
obj.append(END)
model.atom[0].coord = p000
model.atom[1].coord = p100
model.atom[2].coord = add(p010, scale(d100,0.5))
model.atom[3].coord = add(add(p001, scale(d010,0.5)),d100)
elif self.mode=='walls':
obj.extend([ BEGIN, TRIANGLE_STRIP ])
obj.append(NORMAL); obj.extend(n001)
obj.append(VERTEX); obj.extend(p000)
obj.append(VERTEX); obj.extend(p100)
obj.append(VERTEX); obj.extend(p010)
obj.append(VERTEX); obj.extend(p110)
obj.append(END)
obj.extend([ BEGIN, TRIANGLE_STRIP ])
obj.append(NORMAL); obj.extend(n010)
obj.append(VERTEX); obj.extend(p000)
obj.append(VERTEX); obj.extend(p001)
obj.append(VERTEX); obj.extend(p100)
obj.append(VERTEX); obj.extend(p101)
obj.append(END)
obj.extend([ BEGIN, TRIANGLE_STRIP ])
obj.append(NORMAL); obj.extend(n100)
obj.append(VERTEX); obj.extend(p000)
obj.append(VERTEX); obj.extend(p010)
obj.append(VERTEX); obj.extend(p001)
obj.append(VERTEX); obj.extend(p011)
obj.append(END)
model.atom[0].coord = p000
model.atom[1].coord = p100
model.atom[2].coord = p010
model.atom[3].coord = p001
elif self.mode=='plane':
obj.extend([ BEGIN, TRIANGLE_STRIP ])
obj.append(NORMAL); obj.extend(n001)
obj.append(VERTEX); obj.extend(p000)
obj.append(VERTEX); obj.extend(p100)
obj.append(VERTEX); obj.extend(p010)
obj.append(VERTEX); obj.extend(p110)
obj.append(END)
model.atom[0].coord = p000
model.atom[1].coord = p100
model.atom[2].coord = p010
model.atom[3].coord = add(add(p001, scale(d010,0.5)),scale(d100,0.5))
elif self.mode=='quad':
obj.extend([ BEGIN, TRIANGLE_STRIP ])
obj.append(NORMAL); obj.extend(n001)
obj.append(VERTEX); obj.extend(p000)
obj.append(VERTEX); obj.extend(p100)
obj.append(VERTEX); obj.extend(p010)
obj.append(VERTEX); obj.extend(p001)
obj.append(END)
model.atom[0].coord = p000
model.atom[1].coord = p100
model.atom[2].coord = p010
model.atom[3].coord = p001
self.cmd.load_model(model, '_tmp', zoom=0)
self.cmd.update(self.points_name,"_tmp")
self.cmd.delete("_tmp")
# then we load it into PyMOL
self.cmd.delete(self.cgo_name)
self.cmd.load_cgo(obj,self.cgo_name,zoom=0)
self.cmd.order(self.cgo_name+" "+self.points_name,sort=1,location='bottom')
self.cmd.set("nonbonded_size",math.sqrt(dot_product(d10,d10))/10,self.points_name)
def cgo_arrow(origin, endpoint, color='blue', radius=0.10, gap=0.0, hlength=-1, hradius=-1,
type='electric', name='', scaling = 7):
'''
:param origin: List representing origin point of vector to be drawn
:type origin: List of floats
:param endpoint: List representing endpoint of vector to be drawn
:type endpoint: List of floats
:param color: Color of arrow
:type color: String, optional - default blue
:param radius: Radius of cylinder portion of arrow
:type radius: Float, optional - default .1
:param gap: Specifies a gap between the head and body of the arrow, if desired
:type gap: Float, optional - default 0
:param hlength: Length of the head of the arrow
:type hlength: Float, optional - default -1
:param hradius: Radius of the head of the arrow
:type hradius: Float, optional - default -1
:param type: Type of vector being drawn, electric or magnetic
:type type: String, optional - default electric
:param name: Name to be shown in PyMol for the cgo object
:type name: String, optional - default blank
:param scaling: Scaling factor that is passed to scale_endpoint function
:type scaling: Int, optional - default 7
:return: None
'''
from chempy import cpv
#converting parameters to floats
radius, gap = float(radius), float(gap)
hlength, hradius = float(hlength), float(hradius)
if type == 'electric':
color = 'red'
name = 'electric'+name
if type == 'magnetic':
color = 'blue'
name = 'magnetic'+name
try:
color1, color2 = color.split()
except:
color1 = color2 = color
color1 = list(cmd.get_color_tuple(color1))
color2 = list(cmd.get_color_tuple(color2))
if origin == 'sele':
xyz1 = cmd.get_coords('sele', 1)
xyz1 = xyz1.flatten()
xyz1 = xyz1.tolist()
length=np.linalg.norm(np.array(endpoint))
xyz2 = scale_endpoint(endpoint,scaling)
xyz2 = shift_vectors(xyz2, xyz1)
else:
xyz1 = origin
length=np.linalg.norm(np.array(endpoint)-np.array(xyz1))
xyz2 = scale_endpoint(endpoint,scaling)
normal = cpv.normalize(cpv.sub(xyz1, xyz2))
if hlength < 0:
hlength = radius * 3.0
if hradius < 0:
hradius = hlength * 0.6
if gap:
diff = cpv.scale(normal, gap)
xyz1 = cpv.sub(xyz1, diff)
xyz2 = cpv.add(xyz2, diff)
##Location where cylinder switches to cone
xyz3 = cpv.add(cpv.scale(normal, hlength), xyz2)
obj = [cgo.CYLINDER] + xyz1 + xyz3 + [radius] + color1 + color1 + \
[cgo.CONE] + xyz3 + xyz2 + [hradius, 0.0] + color1 + color2 + \
[1.0, 0.0]
print(obj)
##Place pseudoatom with label at midpoint of vector
v0=np.array(xyz1)
v1=np.array(xyz2)
loc=(v0+v1)/2
if not name:
name = cmd.get_unused_name('arrow')
cmd.load_cgo(obj, f"vec_{name}")
##For some reason pos fails, but it just happens to put it where I want
cmd.pseudoatom(f"lab_{name}",name="lab_"+name,label=f"{length:.2f}")#,pos=loc
cmd.group(name,members=f"lab_{name} vec_{name}")