def testGetSymmetry(self):
cmd.load(self.datafile('1rx1.pdb'))
sym = cmd.get_symmetry('1rx1')
self.assertArrayEqual(sym[:6],
[34.455, 45.370, 98.701, 90.0, 90.0, 90.0],
delta=1e-3)
self.assertEqual(sym[6], 'P 21 21 21')
cmd.fragment('gly', 'm1')
self.assertTrue(cmd.get_symmetry('m1') is None)
def test_set_symmetry(self):
sym = [68.7, 126.8, 184.0, 90.0, 90.0, 90.0, 'P 21 21 21']
cmd.pseudoatom('m1')
cmd.set_symmetry('m1', *sym)
v = cmd.get_symmetry('m1')
self.assertEqual(v[-1], sym[-1])
self.assertArrayEqual(v[:-1], sym[:-1], 1e-4)
cmd.pseudoatom('m2')
cmd.symmetry_copy('m1', 'm2')
v = cmd.get_symmetry('m2')
self.assertEqual(v[-1], sym[-1])
self.assertArrayEqual(v[:-1], sym[:-1], 1e-4)
def test_set_symmetry(self):
sym = [68.7, 126.8, 184.0, 90.0, 90.0, 90.0, 'P 21 21 21']
cmd.pseudoatom('m1')
cmd.set_symmetry('m1', *sym)
v = cmd.get_symmetry('m1')
self.assertEqual(v[-1], sym[-1])
self.assertArrayEqual(v[:-1], sym[:-1], 1e-4)
cmd.pseudoatom('m2')
cmd.symmetry_copy('m1', 'm2')
v = cmd.get_symmetry('m2')
self.assertEqual(v[-1], sym[-1])
self.assertArrayEqual(v[:-1], sym[:-1], 1e-4)
def draw_symops(obj, radius=0.2, extension=0):
"""
From pymol issue the "run draw_symops_cctbx.py" command to load the script,
then issue the "draw_symops(object,<optional radius>,<optional extension>)" command
to actually run it and create the cgo object.
e.g. load 1avv.pdb
run draw_symops_cctbx.py
draw_symops 1avv, 0.5, .2
or draw_symops('1avv',.5,.2)
or draw_symops 1avv, radius=.5, extension=.2
The different axis types appear as different objects on the PyMOL menu so they can be turned
on and off individually.
See also help(draw_symops_param) to draw operators by specifying the space group
and cell dimensions directly (i.e. not loaded from a pdb file)
The 'extension' parameter is a fractional increase in the length of each symmetry
operator axis drawn. i.e. a value of 0 is the default and a value of .2 increases
the length by 20% at each end
"""
radius = float(radius)
extension = float(extension)
cell_info = cmd.get_symmetry(obj)
draw_symops_param(cell_info[0:6], cell_info[6], radius, extension)
def test_CCDC(self):
cmd.load(self.datafile('CAFINE.cif'), 'm1')
self.assertEqual(cmd.count_atoms(), 25)
sym = cmd.get_symmetry('m1')
self.assertEqual(sym[6], 'P 21/a')
self.assertArrayEqual(sym[:6], [14.8, 16.7, 3.97, 90., 97., 90.],
delta=1e-4)
def cell_shift(object, dX, dY, dZ, rename=1):
if rename:
oldName = object.split("_")
oldPre = oldName[0]
oldX = int(oldName[1])
oldY = int(oldName[2])
oldZ = int(oldName[3])
newX = "_" + str(int(dX) + oldX)
newY = "_" + str(int(dY) + oldY)
newZ = "_" + str(int(dZ) + oldZ)
newName = oldPre + newX + newY + newZ
#if cmd.get_names().find(newName) != -1:
# print "Symmetry partner already exists in destination position!"
# quit()
cmd.set_name(object, newName)
object = newName
stored.shift = [float(dX), float(dY), float(dZ)]
stored.sgInfo = cmd.get_symmetry(object)
a, b, c, alpha, beta, gamma = stored.sgInfo[0:6]
ca = math.cos(math.radians(alpha))
cb = math.cos(math.radians(beta))
cg = math.cos(math.radians(gamma))
sb = math.sin(math.radians(beta))
sg = math.sin(math.radians(gamma))
stored.fracToOrt = N.array(
[[a, b * cg, c * cb], [0.0, b * sg, c * (ca - cb * cg) / sg],
[0.0, 0.0,
c * sb * math.sqrt(1.0 - ((cb * cg - ca) / (sb * sg))**2)]])
stored.fracToOrt = stored.fracToOrt.transpose()
stored.ortToFrac = inv(stored.fracToOrt)
cmd.cell_shift_helper = cell_shift_helper
cmd.alter_state(1, object,
"x,y,z = cmd.cell_shift_helper([x,y,z],stored.shift)")
def get_grostr(selection="all", state=-1, *, _self=cmd):
"""
DESCRIPTION
Save a Gromacs .gro file.
"""
from chempy import cpv
buffer = ["Created with PSICO", None]
total = [0]
_self.iterate_state(
state,
selection, "total[0] += 1;"
"buffer.append(f'{resv:5}{resn:5}{name:>5}{total[0]:5}{x/10:8.3f}{y/10:8.3f}{z/10:8.3f}')",
space={
"buffer": buffer,
"total": total
})
buffer[1] = f"{total[0]:5}"
sym = cmd.get_symmetry(f"first ({selection})")
a, b, c = sym[:3] if sym else cpv.sub(
*reversed(_self.get_extent(selection)))
buffer.append(f"{(a)/10:10.5f} {(b)/10:9.5f} {(c)/10:9.5f}")
buffer.append("")
return "\n".join(buffer)
def draw_symops(obj,radius=0.2,extension=0):
"""
From pymol issue the "run draw_symops_cctbx.py" command to load the script,
then issue the "draw_symops(object,<optional radius>,<optional extension>)" command
to actually run it and create the cgo object.
e.g. load 1avv.pdb
run draw_symops_cctbx.py
draw_symops 1avv, 0.5, .2
or draw_symops('1avv',.5,.2)
or draw_symops 1avv, radius=.5, extension=.2
The different axis types appear as different objects on the PyMOL menu so they can be turned
on and off individually.
See also help(draw_symops_param) to draw operators by specifying the space group
and cell dimensions directly (i.e. not loaded from a pdb file)
The 'extension' parameter is a fractional increase in the length of each symmetry
operator axis drawn. i.e. a value of 0 is the default and a value of .2 increases
the length by 20% at each end
"""
radius=float(radius)
extension=float(extension)
cell_info=cmd.get_symmetry(obj)
draw_symops_param(cell_info[0:6],cell_info[6],radius,extension)
def cell_shift(object, dX, dY, dZ, rename = 1):
if rename:
oldName = object.split("_")
oldPre = oldName[0]
oldX = int(oldName[1])
oldY = int(oldName[2])
oldZ = int(oldName[3])
newX = "_" + str(int(dX) + oldX)
newY = "_" + str(int(dY) + oldY)
newZ = "_" + str(int(dZ) + oldZ)
newName = oldPre + newX + newY + newZ
#if cmd.get_names().find(newName) != -1:
# print "Symmetry partner already exists in destination position!"
# quit()
cmd.set_name(object, newName)
object = newName
stored.shift = [float(dX),float(dY),float(dZ)]
stored.sgInfo = cmd.get_symmetry(object)
a,b,c,alpha,beta,gamma = stored.sgInfo[0:6]
ca = math.cos(math.radians(alpha))
cb = math.cos(math.radians(beta))
cg = math.cos(math.radians(gamma))
sb = math.sin(math.radians(beta))
sg = math.sin(math.radians(gamma))
stored.fracToOrt = N.array([[a, b * cg, c * cb],
[0.0, b * sg, c * (ca - cb * cg) / sg],
[0.0, 0.0, c * sb * math.sqrt(1.0 - ((cb * cg - ca) / (sb * sg))**2)]])
stored.fracToOrt = stored.fracToOrt.transpose()
stored.ortToFrac = inv(stored.fracToOrt)
cmd.cell_shift_helper = cell_shift_helper
cmd.alter_state(1, object, "x,y,z = cmd.cell_shift_helper([x,y,z],stored.shift)")
def testLoadPDBML(self):
cmd.load(self.datafile("1ubq.xml.gz"))
self.assertEqual(660, cmd.count_atoms())
symmetry = cmd.get_symmetry()
self.assertArrayEqual(symmetry[:6],
[50.84, 42.77, 28.95, 90.0, 90.0, 90.0],
delta=1e-4)
self.assertEqual(symmetry[6], 'P 21 21 21')
def testLoadMMTF(self):
cmd.load(self.datafile("3njw.mmtf.gz"))
self.assertEqual(169, cmd.count_atoms())
self.assertEqual(36, cmd.count_atoms('ss S'))
self.assertEqual(25, cmd.count_atoms('solvent'))
symmetry = cmd.get_symmetry()
self.assertArrayEqual(symmetry[:6], [19.465, 21.432, 29.523, 90.0, 90.0, 90.0], delta=1e-4)
self.assertEqual(symmetry[6], 'P 21 21 21')
def symexpcell(prefix='mate', object=None, a=0, b=0, c=0):
'''
DESCRIPTION
Creates all symmetry-related objects for the specified object that
occur with their bounding box center within the unit cell.
USAGE
symexpcell prefix, object, [a, b, c]
ARGUMENTS
prefix = string: prefix of new objects
object = string: object for which to create symmetry mates
a, b, c = integer: create neighboring cell {default: 0,0,0}
SEE ALSO
symexp, http://www.pymolwiki.org/index.php/SuperSym
'''
if object is None:
object = cmd.get_object_list()[0]
sym = cmd.get_symmetry(object)
cell_edges = sym[0:3]
cell_angles = sym[3:6]
spacegroup = sym[6]
basis = cellbasis(cell_angles, cell_edges)
basis = numpy.matrix(basis)
extent = cmd.get_extent(object)
center = sum(numpy.array(extent)) * 0.5
center = numpy.matrix(center.tolist() + [1.0]).T
center_cell = basis.I * center
extra_shift = [[float(i)] for i in (a,b,c)]
i = 0
matrices = xray.sg_sym_to_mat_list(spacegroup)
for mat in matrices:
i += 1
mat = numpy.matrix(mat)
shift = numpy.floor(mat * center_cell)
mat[0:3,3] -= shift[0:3,0]
mat[0:3,3] += extra_shift
mat = basis * mat * basis.I
mat_list = list(mat.flat)
name = '%s%d' % (prefix, i)
cmd.create(name, object)
cmd.transform_object(name, mat_list)
cmd.color(i+1, name)
def symexpcell(prefix='mate', object=None, a=0, b=0, c=0):
"""
DESCRIPTION
Creates all symmetry-related objects for the specified object that
occur with their bounding box center within the unit cell.
USAGE
symexpcell prefix, object, [a, b, c]
ARGUMENTS
prefix = string: prefix of new objects
object = string: object for which to create symmetry mates
a, b, c = integer: create neighboring cell {default: 0,0,0}
SEE ALSO
symexp, http://www.pymolwiki.org/index.php/SuperSym
"""
if object is None:
object = cmd.get_object_list()[0]
sym = cmd.get_symmetry(object)
cell_edges = sym[0:3]
cell_angles = sym[3:6]
spacegroup = sym[6]
basis = cellbasis(cell_angles, cell_edges)
basis = numpy.matrix(basis)
extent = cmd.get_extent(object)
center = sum(numpy.array(extent)) * 0.5
center = numpy.matrix(center.tolist() + [1.0]).T
center_cell = basis.I * center
extra_shift = [[float(i)] for i in (a,b,c)]
i = 0
matrices = xray.sg_sym_to_mat_list(spacegroup)
for mat in matrices:
i += 1
mat = numpy.matrix(mat)
shift = numpy.floor(mat * center_cell)
mat[0:3,3] -= shift[0:3,0]
mat[0:3,3] += extra_shift
mat = basis * mat * basis.I
mat_list = list(mat.flat)
name = '%s%d' % (prefix, i)
cmd.create(name, object)
cmd.transform_object(name, mat_list, 0)
cmd.color(i+1, name)
def find_cell(self):
if self._obj != None:
cell_param = cmd.get_symmetry(self._obj)
self._a = cell_param[0]
self._b = cell_param[1]
self._c = cell_param[2]
self._alpha = cell_param[3]
self._beta = cell_param[4]
self._gamma = cell_param[5]
def testLoadMMTF(self):
cmd.load(self.datafile("3njw.mmtf.gz"))
self.assertEqual(169, cmd.count_atoms())
self.assertEqual(36, cmd.count_atoms('ss S'))
self.assertEqual(25, cmd.count_atoms('solvent'))
symmetry = cmd.get_symmetry()
self.assertArrayEqual(symmetry[:6],
[19.465, 21.432, 29.523, 90.0, 90.0, 90.0],
delta=1e-4)
self.assertEqual(symmetry[6], 'P 21 21 21')
def get_orthogonalization_matrix(object, quiet = 0):
a,b,c,alpha,beta,gamma = cmd.get_symmetry(object)[0:6]
ca = math.cos(math.radians(alpha))
cb = math.cos(math.radians(beta))
cg = math.cos(math.radians(gamma))
sb = math.sin(math.radians(beta))
sg = math.sin(math.radians(gamma))
fracToOrt = N.array([[a, b * cg, c * cb],
[0.0, b * sg, c * (ca - cb * cg) / sg],
[0.0, 0.0, c * sb * math.sqrt(1.0 - ((cb * cg - ca) / (sb * sg))**2)]])
if not quiet:
print fracToOrt
print inv(fracToOrt)
return fracToOrt
def get_orthogonalization_matrix(object, quiet=0):
a, b, c, alpha, beta, gamma = cmd.get_symmetry(object)[0:6]
ca = math.cos(math.radians(alpha))
cb = math.cos(math.radians(beta))
cg = math.cos(math.radians(gamma))
sb = math.sin(math.radians(beta))
sg = math.sin(math.radians(gamma))
fracToOrt = N.array(
[[a, b * cg, c * cb], [0.0, b * sg, c * (ca - cb * cg) / sg],
[0.0, 0.0,
c * sb * math.sqrt(1.0 - ((cb * cg - ca) / (sb * sg))**2)]])
if not quiet:
print(fracToOrt)
print(inv(fracToOrt))
return fracToOrt
def get_box_params(selection):
try:
a, b, c, alpha, beta, gamma, spacegroup = cmd.get_symmetry(selection)
except TypeError:
# this happens if get_symmetry() is None
return 0, 0, 0, 0, 0, 0, 0, 0, 0
alpha *= math.pi / 180.
beta *= math.pi / 180.
gamma *= math.pi / 180.
a /= 10
b /= 10
c /= 10
v1x = a
v2x = b * math.cos(gamma)
v2y = b * math.sin(gamma)
v3x = c * math.cos(beta)
v3y = c * (math.cos(alpha) -
math.cos(gamma) * math.cos(beta)) / (math.sin(gamma))
v3z = (c**2 - v3x**2 - v3y**2)**0.5
return (v1x, v2y, v3z, 0, 0, v2x, 0, v3x, v3y)
def __init__(self, map_name, level, radius, name, sym_source):
self.level = level
self.radius = radius
self.map_name = map_name
self.name = name
self.center_name = cmd.get_unused_name('_center')
self.callback_name = cmd.get_unused_name('_cb')
cmd.set("auto_zoom", 0)
cmd.pseudoatom(self.center_name)
cmd.hide("everything", self.center_name)
symmetry = cmd.get_symmetry(sym_source or map_name)
if symmetry:
cmd.set("map_auto_expand_sym", 1)
cmd.set_symmetry(self.center_name, *symmetry)
cmd.set_key("pgup", self.contour_plus)
cmd.set_key("pgdn", self.contour_minus)
self.update()
def draw_cell(obj, radius=1.0, mode=0):
"""
From pymol issue the "run draw_cell.py" command to load the script,
then issue the "draw_cell(object,<optional radius>)" command
to actually run it and create the cgo object showing the unit cell
border for the space group specified by molecular object 'object'.
e.g. load 1avv.pdb
run draw_cell.py
draw_cell 1avv 0.5 (or draw_cell('1avv',.5))
see also help(draw_cell_param) to draw the cell border for
user-defined cell dimensions (i.e. not loaded from a pdb file)
See also "help(draw_cell_param) to draw the cell border by
specifying the unit cell parameters directly (i.e. not loaded from
a pdb file).
"""
radius = float(radius)
cell_info = cmd.get_symmetry(obj)
draw_cell_param(cell_info[0:6], radius, mode)
def draw_cell(obj,radius=0.2):
"""
From pymol issue the "run draw_cell.py" command to load the script,
then issue the "draw_cell(object,<optional radius>)" command
to actually run it and create the cgo object showing the unit cell
border for the space group specified by molecular object 'object'.
e.g. load 1avv.pdb
run draw_cell.py
draw_cell 1avv 0.5 (or draw_cell('1avv',.5))
see also help(draw_cell_param) to draw the cell border for
user-defined cell dimensions (i.e. not loaded from a pdb file)
See also "help(draw_cell_param) to draw the cell border by
specifying the unit cell parameters directly (i.e. not loaded from
a pdb file).
"""
radius=float(radius)
cell_info=cmd.get_symmetry(obj)
draw_cell_param(cell_info[0:6],radius)
def __init__(self, map_name, level, radius, name, sym_source):
self.level = level
self.radius = radius
self.map_name = map_name
self.name = name
self.center_name = cmd.get_unused_name('_center')
self.callback_name = cmd.get_unused_name('_cb')
cmd.set("auto_zoom", 0)
cmd.pseudoatom(self.center_name)
cmd.hide("everything", self.center_name)
symmetry = cmd.get_symmetry(sym_source or map_name)
if symmetry:
cmd.set("map_auto_expand_sym", 1)
cmd.set_symmetry(self.center_name, *symmetry)
cmd.set_key("pgup", self.contour_plus)
cmd.set_key("pgdn", self.contour_minus)
self.update()
def get_box(obj):
dimensions = cmd.get_symmetry(obj)[:6]
B = np.zeros((3, 3), dtype=np.float32)
x, y, z, a, b, c = dimensions
if np.all(dimensions[:3] == 0):
return B
B[0][0] = x
if a == 90.0 and b == 90.0 and c == 90.0:
B[1][1] = y
B[2][2] = z
else:
a = np.deg2rad(a)
b = np.deg2rad(b)
c = np.deg2rad(c)
B[1][0] = y * np.cos(c)
B[1][1] = y * np.sin(c)
B[2][0] = z * np.cos(b)
B[2][1] = z * (np.cos(a) - np.cos(b) * np.cos(c)) / np.sin(c)
B[2][2] = np.sqrt(z * z - B[2][0] ** 2 - B[2][1] ** 2)
return B
def get_box(obj):
dimensions = cmd.get_symmetry(obj)[:6]
B = np.zeros((3, 3), dtype=np.float32)
x, y, z, a, b, c = dimensions
if np.all(dimensions[:3] == 0):
return B
B[0][0] = x
if a == 90. and b == 90. and c == 90.:
B[1][1] = y
B[2][2] = z
else:
a = np.deg2rad(a)
b = np.deg2rad(b)
c = np.deg2rad(c)
B[1][0] = y * np.cos(c)
B[1][1] = y * np.sin(c)
B[2][0] = z * np.cos(b)
B[2][1] = z * (np.cos(a) - np.cos(b) * np.cos(c)) / np.sin(c)
B[2][2] = np.sqrt(z * z - B[2][0] ** 2 - B[2][1] ** 2)
return B
def supercell(a=1,
b=1,
c=1,
object=None,
color='green',
name='supercell',
withmates=1):
'''
DESCRIPTION
Draw a supercell, as requested by Nicolas Bock on the pymol-users
mailing list (Subject: [PyMOL] feature request: supercell construction
Date: 04/12/2010 10:12:17 PM (Mon, 12 Apr 2010 14:12:17 -0600))
USAGE
supercell a, b, c [, object [, color [, name [, withmates]]]]
ARGUMENTS
a, b, c = integer: repeat cell in x,y,z direction a,b,c times
{default: 1,1,1}
object = string: name of object to take cell definition from
color = string: color of cell {default: blue}
name = string: name of the cgo object to create {default: supercell}
withmates = bool: also create symmetry mates in displayed cells
{default: 1}
SEE ALSO
show cell
'''
import numpy
from pymol import cgo
if object is None:
object = cmd.get_object_list()[0]
withmates = int(withmates)
sym = cmd.get_symmetry(object)
cell_edges = sym[0:3]
cell_angles = sym[3:6]
basis = cellbasis(cell_angles, cell_edges)
assert isinstance(basis, numpy.ndarray)
ts = list()
for i in range(int(a)):
for j in range(int(b)):
for k in range(int(c)):
ts.append([i, j, k])
obj = [
cgo.BEGIN,
cgo.LINES,
]
for t in ts:
shift = basis[0:3, 0:3] * t
shift = shift[:, 0] + shift[:, 1] + shift[:, 2]
for i in range(3):
vi = basis[0:3, i]
vj = [
numpy.array([0., 0., 0.]), basis[0:3, (i + 1) % 3],
basis[0:3, (i + 2) % 3],
basis[0:3, (i + 1) % 3] + basis[0:3, (i + 2) % 3]
]
for j in range(4):
obj.append(cgo.VERTEX)
obj.extend((shift + vj[j]).tolist())
obj.append(cgo.VERTEX)
obj.extend((shift + vj[j] + vi).tolist())
if withmates:
groupname = 'm%d%d%d' % tuple(t)
symexpcell(groupname + '_', object, *t)
cmd.group(groupname, groupname + '_*')
obj.append(cgo.END)
cmd.delete(name)
cmd.load_cgo(obj, name)
cmd.color(color, name)
def supercell(a=1, b=1, c=1, object=None, color='blue', name='supercell', withmates=1,prefix="m",center=0,transformation=None,cutoff=None):
'''
DESCRIPTION
Draw a supercell, as requested by Nicolas Bock on the pymol-users
mailing list (Subject: [PyMOL] feature request: supercell construction
Date: 04/12/2010 10:12:17 PM (Mon, 12 Apr 2010 14:12:17 -0600))
USAGE
supercell a, b, c [, object [, color [, name [, withmates]]]]
ARGUMENTS
a, b, c = integer: repeat cell in x,y,z direction a,b,c times
{default: 1,1,1}
object = string: name of object to take cell definition from
color = string: color of cell {default: blue}
name = string: name of the cgo object to create {default: supercell}
withmates = bool: also create symmetry mates in displayed cells
{default: 1}
prefix = string: prefix for the symmetry mates {default: m}
center = boolean: If 1, indicates that the lattice should be centered on the
origin, as opposed to having the corner at the origin cell. {default: 0}
transformation = list: a 16-element list giving the 4x4 transformation
matrix, as described in get_object_matrix() {default: identity matrix}
cutoff = int: restrict symmetry mates to within cutoff angstroms of the origin.
Use 0 to generate all symmetry mates. {default: 0}
SEE ALSO
show cell
cmd
'''
if object is None:
object = cmd.get_object_list()[0]
withmates = int(withmates)
sym = cmd.get_symmetry(object)
cell_edges = sym[0:3]
cell_angles = sym[3:6]
basis = cellbasis(cell_angles, cell_edges)
if transformation is not None:
transmat = transformation_to_numpy(transformation)
assert isinstance(basis, numpy.ndarray)
ts = list()
a = int(a)
b = int(b)
c = int(c)
if int(center) == 0:
astart = 0
bstart = 0
cstart = 0
else:
#TODO Maybe would be more useful to center at the asymmetric unit?
# For now, center on the origin cell.
astart = (1-a)/2
bstart = (1-b)/2
cstart = (1-c)/2
for i in range( astart,astart+a ):
for j in range( bstart,bstart+b ):
for k in range( cstart,cstart+c ):
ts.append([i,j,k])
obj = [
cgo.BEGIN,
cgo.LINES,
cgo.COLOR,
]
obj.extend(cmd.get_color_tuple(color))
for t in ts:
# draw bounding box around cell t
shift = basis[0:3,0:3] * t
shift = shift[:,0] + shift[:,1] + shift[:,2]
for i in range(3):
# vi is direction of the edges to draw
vi = basis[0:3,i]
# vj are starting points for the four edges in that direction
vj = [
numpy.array([0.,0.,0.]),
basis[0:3,(i+1)%3],
basis[0:3,(i+2)%3],
basis[0:3,(i+1)%3] + basis[0:3,(i+2)%3]
]
for j in range(4):
start = shift + vj[j]
end = start + vi
if transformation is not None:
start = numpy.dot(transmat, numpy.append(start,1))[:3]
end = numpy.dot(transmat, numpy.append(end ,1))[:3]
obj.append(cgo.VERTEX)
obj.extend(start.tolist())
obj.append(cgo.VERTEX)
obj.extend(end.tolist())
if withmates:
symexpcell('%s%d%d%d_' % (prefix,t[0]-astart,t[1]-bstart,t[2]-cstart), object, *t,transformation=transformation,cutoff=cutoff)
obj.append(cgo.END)
cmd.delete(name)
cmd.load_cgo(obj, name)
def save_pdb(filename, selection='(all)', state=-1, symm=1, ss=1, aniso=0, seqres=0, quiet=1):
'''
DESCRIPTION
Save the coordinates of a selection as pdb including the
secondary structure information and, if possible, the unit
cell. The latter requires the selction of a single object
USAGE
save_pdb filename, selection [, state [, symm [, ss [, aniso ]]]]
ARGUMENTS
filename = string: file path to be written
selection = string: atoms to save {default: (all)}
Note: to include the unit cell information you
need to select a single object
state = integer: state to save {default: -1 (current state)}
symm = 0 or 1: save symmetry info if possible {default: 1}
ss = 0 or 1: save secondary structure info {default: 1}
aniso = 0 or 1: save ANISO records {default: 0}
SEE ALSO
save
'''
_assert_package_import()
from . import pymol_version
selection = selector.process(selection)
state, quiet = int(state), int(quiet)
symm, ss = int(symm), int(ss)
filename = cmd.exp_path(filename)
f = open(filename, 'w')
print('REMARK 200 Generated with PyMOL and psico'.ljust(80), file=f)
# Write sequence
if int(seqres):
f.write(get_pdb_seqres(selection))
# Write the CRYST1 line if possible
if symm:
try:
obj1 = cmd.get_object_list(selection)[0]
sym = cmd.get_symmetry(obj1)
if len(sym) != 7:
raise
f.write("CRYST1%9.3f%9.3f%9.3f%7.2f%7.2f%7.2f %-10s%4d\n" % tuple(sym + [1]))
if not quiet:
print(' Info: Wrote unit cell and space group info')
except:
if not quiet:
print(' Info: No crystal information')
# Write secondary structure
if ss:
try:
sss = get_pdb_sss(selection, state, quiet)
if not sss:
raise
f.write(sss)
if not quiet:
print(' Info: Wrote secondary structure info')
except:
if not quiet:
print(' Info: No secondary structure information')
# Write coordinates of selection
pdbstr = cmd.get_pdbstr(selection, state)
# fix END records
if state == 0 and pymol_version < 1.6:
pdbstr = '\n'.join(line for line in pdbstr.splitlines() if line != 'END') + '\nEND\n'
# anisotropic b-factors
if int(aniso) and pymol_version < 1.6 and \
cmd.get_model('first (%s)' % selection).atom[0].u_aniso[0] != 0.0:
def mergeaniso():
atom_it = iter(cmd.get_model(selection, state).atom)
for line in pdbstr.splitlines(True):
yield line
if line[:6] in ['ATOM ', 'HETATM']:
yield 'ANISOU' + line[6:28] + \
''.join('%7.0f' % (u*1e4) for u in next(atom_it).u_aniso) + \
line[70:]
pdbstr = ''.join(mergeaniso())
f.write(pdbstr)
f.close()
if not quiet:
print('Wrote PDB to \''+filename+'\'')
def test_CCDC(self):
cmd.load(self.datafile('CAFINE.cif'), 'm1')
self.assertEqual(cmd.count_atoms(), 25)
sym = cmd.get_symmetry('m1')
self.assertEqual(sym[6], 'P 21/a')
self.assertArrayEqual(sym[:6], [14.8, 16.7, 3.97, 90., 97., 90.], delta=1e-4)
def save_gro(filename, selection='(all)', bx=None, by=None, bz=None):
"""
DESCRIPTION
"save_gro" saves content to a .gro file in a similar fashion as "save"
USAGE
save_gro filename [, selection [, bx [, by [, bz ]]]]
ARGUMENTS
filename = string: file path to be written
selection = string: atoms to save {default: (all)}
bx = float: rectangular box size X dimension in Angströms
by = float: rectangular box size Y dimension in Angströms
bz = float: rectangular box size Z dimension in Angströms
NOTES
In contrast to "save", only the current state can be saved
as the GRO format is a single-structure format.
If any of bx, by and bz is not supplied, the box will be
determined from the get_symmetry() command.
Atoms are written with increasing ID.
SEE ALSO
load, save
"""
with open(filename, 'wt') as f:
f.write('{}\n'.format(selection))
f.write('{:>5d}\n'.format(cmd.count_atoms(selection)))
space = {'l': []}
cmd.iterate_state(-1, selection, 'l.append((ID, resi,resn,name,index,x,y,z))', space=space)
for ID, resi, resn, name, index, x, y, z in sorted(space['l'], key=lambda a: (a[1], a[4])):
f.write('{:>5d}{:<5s}{:>5s}{:>5d}{:>8.3f}{:>8.3f}{:>8.3f}\n'.format(int(resi), resn, name, int(ID) + 1,
float(x) * 0.1, float(y) * 0.1,
float(z) * 0.1))
if bx is None and by is None and bz is None:
try:
a, b, c, alpha, beta, gamma, spacegroup = cmd.get_symmetry(selection)
alpha *= math.pi / 180.
beta *= math.pi / 180.
gamma *= math.pi / 180.
a /= 10
b /= 10
c /= 10
v1x = a
v2x = b * math.cos(gamma)
v2y = b * math.sin(gamma)
v3x = c * math.cos(beta)
v3y = c * (math.cos(alpha) - math.cos(gamma) * math.cos(beta)) / (math.sin(gamma))
v3z = (c ** 2 - v3x ** 2 - v3y ** 2) ** 0.5
f.write(
' {:>9.5f} {:>9.5f} {:>9.5f} {:>9.5f} {:>9.5f} {:>9.5f} {:>9.5f} {:>9.5f} {:>9.5f}\n'.format(v1x,
v2y,
v3z, 0,
0, v2x,
0, v3x,
v3y))
except TypeError:
f.write(' 0 0 0\n')
else:
f.write(' {} {} {}\n'.format(bx, by, bz))
print('Saved {} atoms.'.format(len(space['l'])))
def save_pdb(filename, selection="(all)", state=-1, symm=1, ss=1, aniso=0, quiet=1):
"""
DESCRIPTION
Save the coordinates of a selection as pdb including the
secondary structure information and, if possible, the unit
cell. The latter requires the selction of a single object
USAGE
save_pdb filename, selection [, state [, symm [, ss [, aniso ]]]]
ARGUMENTS
filename = string: file path to be written
selection = string: atoms to save {default: (all)}
Note: to include the unit cell information you
need to select a single object
state = integer: state to save {default: -1 (current state)}
symm = 0 or 1: save symmetry info if possible {default: 1}
ss = 0 or 1: save secondary structure info {default: 1}
aniso = 0 or 1: save ANISO records {default: 0}
SEE ALSO
save
"""
selection = selector.process(selection)
state, quiet = int(state), int(quiet)
symm, ss = int(symm), int(ss)
filename = cmd.exp_path(filename)
f = open(filename, "w")
print >> f, "REMARK 200 Generated with PyMOL and psico".ljust(80)
# Write the CRYST1 line if possible
if symm:
try:
obj1 = cmd.get_object_list(selection)[0]
sym = cmd.get_symmetry(obj1)
if len(sym) != 7:
raise
f.write("CRYST1%9.3f%9.3f%9.3f%7.2f%7.2f%7.2f %-10s%4d\n" % tuple(sym + [1]))
if not quiet:
print " Info: Wrote unit cell and space group info"
except:
if not quiet:
print " Info: No crystal information"
# Write secondary structure
if ss:
try:
sss = get_pdb_sss(selection, state, quiet)
if not sss:
raise
f.write(sss)
if not quiet:
print " Info: Wrote secondary structure info"
except:
if not quiet:
print " Info: No secondary structure information"
# Write coordinates of selection
pdbstr = cmd.get_pdbstr(selection, state)
# fix END records
if state == 0 and cmd.get_version()[1] < 1.6:
pdbstr = "\n".join(line for line in pdbstr.splitlines() if line != "END") + "\nEND\n"
# anisotropic b-factors
if int(aniso) and cmd.get_model("first (%s)" % selection).atom[0].u_aniso[0] != 0.0:
def mergeaniso():
atom_it = iter(cmd.get_model(selection, state).atom)
for line in pdbstr.splitlines(True):
yield line
if line[:6] in ["ATOM ", "HETATM"]:
yield "ANISOU" + line[6:28] + "".join("%7.0f" % (u * 1e4) for u in atom_it.next().u_aniso) + line[
70:
]
pdbstr = "".join(mergeaniso())
f.write(pdbstr)
f.close()
if not quiet:
print "Wrote PDB to '" + filename + "'"
def symexpcell(prefix='mate', object=None, a=0, b=0, c=0,transformation=None,cutoff=None):
'''
DESCRIPTION
Creates all symmetry-related objects for the specified object that
occur with their bounding box center within the unit cell.
USAGE
symexpcell prefix, object, [a, b, c]
ARGUMENTS
prefix = string: prefix of new objects
object = string: object for which to create symmetry mates
a, b, c = integer: create neighboring cell {default: 0,0,0}
transformation = list: list of 16 floats giving the transformation matrix
to apply to the generated symmetry mates {default: identity matrix}
cutoff = int: restrict symmetry mates to within cutoff angstroms of the origin.
Use 0 to generate all symmetry mates. {default: 0}
SEE ALSO
symexp, http://www.pymolwiki.org/index.php/SuperSym
'''
#print "symexpcell %s,%s,%d,%d,%d,%s"%(prefix,object,int(a),int(b),int(c),transformation)
if object is None:
object = cmd.get_object_list()[0]
if cutoff is not None:
cutoff = int(cutoff)
if cutoff <= 0: cutoff = None
sym = cmd.get_symmetry(object)
cell_edges = sym[0:3]
cell_angles = sym[3:6]
spacegroup = sym[6]
basis = cellbasis(cell_angles, cell_edges)
extent = cmd.get_extent(object)
center = sum(numpy.array(extent)) * 0.5
center = numpy.append(center,1.0).reshape(4,1)
center_cell = numpy.linalg.inv(basis) * center
extra_shift = numpy.array([[float(i)] for i in (a,b,c)])
origin = numpy.array([[0,0,0,1]]).T
if transformation is not None:
transmat = transformation_to_numpy(transformation)
#print "%s\n*\n%s\n=\n%s\n" % (origin,transmat,
# numpy.dot(numpy.linalg.inv(transmat),origin) )
origin = numpy.dot(numpy.linalg.inv(transmat),origin)
i = 0
matrices = xray.sg_sym_to_mat_list(spacegroup)
for mat in matrices:
i += 1
mat = numpy.array(mat)
shift = numpy.floor(numpy.dot(mat, center_cell))
mat[0:3,3] -= shift[0:3,0]
mat[0:3,3] += extra_shift[0:3,0]
mat = numpy.dot(numpy.dot(basis, mat), numpy.linalg.inv(basis) )
mat_list = list(mat.flat)
new_center = numpy.dot(mat,center)
#print "%s\n* (%d)\n%s\n=\n%s\n" % (center,i,mat, new_center)
if cutoff is not None:
dist = new_center - origin
dist = numpy.dot(dist.T,dist)
if dist > cutoff**2:
#print "Skipping %d%d%d_%d at distance %f" % (a,b,c,i,sqrt(dist))
continue
name = '%s%d' % (prefix, i)
cmd.create(name, object)
cmd.transform_object(name, mat_list)
# Apply extra transformation afterwards
if transformation is not None:
cmd.transform_object(name, transformation)
cmd.color(i+1, name)
def supercell(a=1, b=1, c=1, object=None, color='green', name='supercell', withmates=1):
'''
DESCRIPTION
Draw a supercell, as requested by Nicolas Bock on the pymol-users
mailing list (Subject: [PyMOL] feature request: supercell construction
Date: 04/12/2010 10:12:17 PM (Mon, 12 Apr 2010 14:12:17 -0600))
USAGE
supercell a, b, c [, object [, color [, name [, withmates]]]]
ARGUMENTS
a, b, c = integer: repeat cell in x,y,z direction a,b,c times
{default: 1,1,1}
object = string: name of object to take cell definition from
color = string: color of cell {default: blue}
name = string: name of the cgo object to create {default: supercell}
withmates = bool: also create symmetry mates in displayed cells
{default: 1}
SEE ALSO
show cell
'''
import numpy
from pymol import cgo
if object is None:
object = cmd.get_object_list()[0]
withmates = int(withmates)
sym = cmd.get_symmetry(object)
cell_edges = sym[0:3]
cell_angles = sym[3:6]
basis = cellbasis(cell_angles, cell_edges)
assert isinstance(basis, numpy.ndarray)
ts = list()
for i in range(int(a)):
for j in range(int(b)):
for k in range(int(c)):
ts.append([i,j,k])
obj = [
cgo.BEGIN,
cgo.LINES,
cgo.COLOR,
]
obj.extend(cmd.get_color_tuple(color))
for t in ts:
shift = basis[0:3,0:3] * t
shift = shift[:,0] + shift[:,1] + shift[:,2]
for i in range(3):
vi = basis[0:3,i]
vj = [
numpy.array([0.,0.,0.]),
basis[0:3,(i+1)%3],
basis[0:3,(i+2)%3],
basis[0:3,(i+1)%3] + basis[0:3,(i+2)%3]
]
for j in range(4):
obj.append(cgo.VERTEX)
obj.extend((shift + vj[j]).tolist())
obj.append(cgo.VERTEX)
obj.extend((shift + vj[j] + vi).tolist())
if withmates:
groupname = 'm%d%d%d' % tuple(t)
symexpcell(groupname + '_', object, *t)
cmd.group(groupname, groupname + '_*')
obj.append(cgo.END)
cmd.delete(name)
cmd.load_cgo(obj, name)
def testCopyMap(self):
cmd.load(self.datafile('emd_1155.ccp4'), 'map1')
cmd.copy('map2', 'map1')
self.assertEqual(cmd.get_symmetry('map1'), cmd.get_symmetry('map2'))
self.assertArrayEqual(cmd.get_volume_field('map1'), cmd.get_volume_field('map2'))
def testLoadPDBML(self):
cmd.load(self.datafile("1ubq.xml.gz"))
self.assertEqual(660, cmd.count_atoms())
symmetry = cmd.get_symmetry()
self.assertArrayEqual(symmetry[:6], [50.84, 42.77, 28.95, 90.0, 90.0, 90.0], delta=1e-4)
self.assertEqual(symmetry[6], 'P 21 21 21')
def cell_axes(object=None, a_color='0xd95f02', b_color='0x1b9e77', c_color='0x7570b3', name='cell_axes'):
'''
DESCRIPTION
Draw arrows corresponding to the crystallographic axes.
The default color palette is colorblind friendly but close to the familiar
red, green, and blue for the a, b, and, c axes respectively.
(See https://colorbrewer2.org/#type=qualitative&scheme=Dark2&n=3 for details)
USAGE
cell_axes [ object, [, a_color [, b_color [, c_color, [ name]]]]
ARGUMENTS
object = string: name of object to take cell definition from
a_color = string: color of a-axis {default: '0xd95f02'}
b_color = string: color of b-axis {default: '0x1b9e77'}
c_color = string: color of c-axis {default: '0x7570b3'}
name = string: name of the cgo object to create {default: cell_axes}
'''
from pymol import cgo
if object is None:
object = cmd.get_object_list()[0]
sym = cmd.get_symmetry(object)
cell_edges = sym[0:3]
cell_angles = sym[3:6]
basis = cellbasis(cell_angles, cell_edges)
cmd.set('auto_zoom', 0)
w = 0.06 # cylinder width
l = 0.75 # cylinder length
h = 0.25 # cone height
d = w * 1.618 # cone base diameter
obj = []
import numpy as np
A,B,C = basis[:3,:3].T
r = A
r = np.where(np.isclose(r, 0., atol=1e-5), 0., r)
r = r / np.linalg.norm(r)
rgb = cmd.get_color_tuple(a_color)
obj.extend([
cgo.CYLINDER, 0.0, 0.0, 0.0, l*r[0], l*r[1], l*r[2], w, *rgb, *rgb,
cgo.CONE, l*r[0], l*r[1], l*r[2], (h+l)*r[0], (h+l)*r[1], (h+l)*r[2], d, 0.0, *rgb, *rgb, 1.0, 1.0
])
def get_cgo_vector_list(r, color_name, eps=1e-5):
rgb = cmd.get_color_tuple(color_name)
r = r / np.linalg.norm(r)
r = np.where(np.isclose(r, 0.0, atol=eps), 0., r) #see https://github.com/schrodinger/pymol-open-source/issues/220
cgo_list = [
cgo.CYLINDER, 0.0, 0.0, 0.0, l*r[0], l*r[1], l*r[2], w, *rgb, *rgb,
cgo.CONE, l*r[0], l*r[1], l*r[2], (h+l)*r[0], (h+l)*r[1], (h+l)*r[2], d, 0.0, *rgb, *rgb, 1.0, 1.0
]
return cgo_list
obj = get_cgo_vector_list(A, a_color) + \
get_cgo_vector_list(B, b_color) + \
get_cgo_vector_list(C, c_color)
PutCenterCallback(name, 1).load()
cmd.load_cgo(obj, name)
def procD2dat(lfile=None,biod="/data/biounit",outd=None):
N = 4
if lfile is None: lfile='/Users/sheffler/project/sym_comp/meta/D2.list'
if outd is None: outd="/Users/sheffler/project/sym_comp/D2"
print outd
Nnobio=0; Nok=0; Ncontact=0; Nbig=0; Nnsym=0; Nnomxatm=0; Nhomogen=0
for fn in open(lfile).readlines():
fn = fn.strip()
pdb = fn[3:7]
bnum = int(fn[-1:])
if os.path.exists(outd+"/"+pdb+"_"+str(bnum)+"_sub1.pdb") or os.path.exists(outd+"/"+pdb+"_"+str(bnum)+"_sub1.pdb.gz"):
Nok += 1
continue
fname = biod+"/"+fn
if not os.path.exists(fname):
fname += ".gz"
if not os.path.exists(fname):
Nnobio += 1
continue
#print pdb,bnum,fname
cmd.delete("all")
cmd.load(fname,'m')
cmd.remove("resn HOH")
cmd.remove('not alt a+""')
#hf = cmd.select("(HET and not resn MSE+CSW)",state=1) / cmd.select("ALL",state=1)
#if hf > 0.1: continue
#cmd.remove("(HET and not resn MSE+CSW)")
if cmd.select('all',state=4) != 0:
for i in range(1,N+1):
cmd.create("sub%i"%i,"m and not (HET and not resn MSE+CSW)",i,1)
elif cmd.select('all',state=2) != 0:
cc = chaincount('m')
if len(cc) < 2:
Nnsym += 1
continue
cmd.create("sub1","m and chain %s and not (HET and not resn MSE+CSW)"%(cc[0][1]),1,1)
cmd.create("sub2","m and chain %s and not (HET and not resn MSE+CSW)"%(cc[1][1]),1,1)
cmd.create("sub3","m and chain %s and not (HET and not resn MSE+CSW)"%(cc[0][1]),2,1)
cmd.create("sub4","m and chain %s and not (HET and not resn MSE+CSW)"%(cc[1][1]),2,1)
else:
cc = chaincount("m")
if len(cc) < N:
sym = cmd.get_symmetry("m")
if sym[6] == "I 2 2 2":
trans('m',Vec(0,-sym[1],0))
print pid
#elif sym[6] == "P 21 21 2" and len(cc)==2:
# trans('m',Vec(-sym[0]/2.0,-sym[1]/2.0,0))
# cmd.create("sub1","m and chain %s and not (HET and not resn MSE+CSW)"%(cc[0][1]),1,1)
# cmd.create("sub2","m and chain %s and not (HET and not resn MSE+CSW)"%(cc[1][1]),1,1)
# cmd.create("sub3","m and chain %s and not (HET and not resn MSE+CSW)"%(cc[0][1]),1,1)
# cmd.create("sub4","m and chain %s and not (HET and not resn MSE+CSW)"%(cc[1][1]),1,1)
# rot("sub3",Vec(0,0,1),180,Vec(0,0,0))
# rot("sub4",Vec(0,0,1),180,Vec(0,0,0))
elif sym[6] in ('C 1 2 1','P 21 21 21','P 62 2 2','P 64 2 2','P 65 2 2','P 63 2 2','P 61 2 2','C 2 2 21'):
Nnsym += 1
#if pid != "1y2k_2": return
continue
else:
Nnsym += 1
continue#return
cmd.save(outd+"/"+pdb+"_"+str(bnum)+"_sub1.pdb","m")
Nok += 1
continue
else:
for i in range(1,N+1):
cmd.create("sub%i"%i,"m and chain %s and not (HET and not resn MSE+CSW)"%(cc[-i][1]),1,1)
for i in range(1,N+1):
if iscontig("sub%i"%i):
cmd.create("mxatm","sub%i"%i)
break
if cmd.select("name CA and mxatm") > 250:
Nbig += 1
continue
if cmd.select("mxatm") < 50:
Nnomxatm += 1
continue
chains = ["sub%i"%i for i in range(1,N+1)]
done = False
count = 0
while not done and count < 50:
done = True
random.shuffle(chains)
for i in range(len(chains)):
for j in range(i+1,len(chains)):
done = done and homogenizechains(chains[i],chains[j])
count += 1
if count >= 50:
Nhomogen += 1
continue
if cmd.select("sub1") < 50:
Nnomxatm += 1
continue
cm = com("sub*")
for i in range(1,N+1):
trans("sub%i"%i,-cm)
a = [cmd.get_model("sub%i and name CA"%i).atom for i in range(1,N+1)]
a1 = Vec(0,0,0)
for i in range(len(a[0])):
axis1 = Vec(a[0][i].coord) + Vec(a[1][i].coord)
if axis1.length() > 0.0001 and a1.dot(axis1) < 0: axis1 *= -1
a1 += axis1
a1.normalize()
for i in range(1,N+1):
alignaxis("sub%i"%i,Vec(1,0,0),a1,Vec(0,0,0))
a = [cmd.get_model("sub%i and name CA"%i).atom for i in range(1,N+1)]
a1 = Vec(0,0,0)
for i in range(len(a[0])):
axis1 = Vec(a[0][i].coord) + Vec(a[2][i].coord)
if axis1.length() > 0.0001 and a1.dot(axis1) < 0: axis1 *= -1
a1 += axis1
a1.normalize()
for i in range(1,N+1):
alignaxis("sub%i"%i,Vec(0,1,0),a1,Vec(0,0,0))
cmd.align("mxatm","sub1")
cmd.create("final2","mxatm")
cmd.create("final3","mxatm")
cmd.create("final4","mxatm")
rot('final2',Vec(1,0,0),180,Vec(0,0,0))
rot('final3',Vec(0,1,0),180,Vec(0,0,0))
rot('final4',Vec(0,0,1),180,Vec(0,0,0))
n1 = cmd.select('mxatm within 4 of final2')
n2 = cmd.select('mxatm within 4 of final3')
n3 = cmd.select('mxatm within 4 of final4')
if n1 < 10 and n2 < 10 and n3 < 10:
Ncontact += 1
continue
if not os.path.exists(outd): os.mkdir(outd)
cmd.save(outd+"/"+pdb+"_"+str(bnum)+"_sub1.pdb","mxatm")
Nok += 1
#return
print Nok, Nbig, Nnsym, Ncontact, Nnobio, Nnomxatm, Nhomogen
def save_traj(filename, selection="(all)", format="", box=0, quiet=1):
"""
DESCRIPTION
Save coordinates of a multi-state object to a trajectory file (DCD OR CRD).
Based on http://pymolwiki.org/index.php/Save2traj by Sean Law
USAGE
save_traj filename [, selection [, format ]]
ARGUMENTS
filename = string: file path to be written
selection = string: atoms to save
format = string: 'dcd' or 'crd' (alias 'charmm' or 'amber') {default:
determined from filename extension)
"""
box = int(box)
# Get NATOMS, NSTATES
NATOMS = cmd.count_atoms(selection)
NSTATES = cmd.count_states(selection)
# Determine Trajectory Format
if format == "" and "." in filename:
format = filename.rsplit(".", 1)[1]
format = format.lower()
if format in ["charmm", "dcd"]:
format = "charmm"
elif format in ["amber", "trj", "crd"]:
format = "amber"
else:
print "Unknown format:", format
raise CmdException
f = open(filename, "wb")
if format == "charmm":
outfile = DCDOutfile(filename, NSTATES, NATOMS)
elif format == "amber":
# size of periodic box
if box:
try:
boxdim = cmd.get_symmetry(selection)[0:3]
except:
boxdim = [0, 0, 0]
else:
boxdim = None
outfile = CRDOutfile(filename, NSTATES, NATOMS, box=boxdim)
else:
raise Exception, "This should not happen"
# Write Trajectory Coordinates
for state in range(1, NSTATES + 1):
xyz = cmd.get_model(selection, state).get_coord_list()
outfile.writeCoordSet(xyz)
outfile.close()
if not int(quiet):
fmt = "Wrote trajectory in %s format with %d atoms and %d frames to file %s"
print fmt % (format, NATOMS, NSTATES, filename)
def save_pdb(filename,
selection='(all)',
state=-1,
symm=1,
ss=1,
aniso=0,
seqres=0,
quiet=1):
'''
DESCRIPTION
Save the coordinates of a selection as pdb including the
secondary structure information and, if possible, the unit
cell. The latter requires the selction of a single object
USAGE
save_pdb filename, selection [, state [, symm [, ss [, aniso ]]]]
ARGUMENTS
filename = string: file path to be written
selection = string: atoms to save {default: (all)}
Note: to include the unit cell information you
need to select a single object
state = integer: state to save {default: -1 (current state)}
symm = 0 or 1: save symmetry info if possible {default: 1}
ss = 0 or 1: save secondary structure info {default: 1}
aniso = 0 or 1: save ANISO records {default: 0}
SEE ALSO
save
'''
selection = selector.process(selection)
state, quiet = int(state), int(quiet)
symm, ss = int(symm), int(ss)
filename = cmd.exp_path(filename)
f = open(filename, 'w')
print('REMARK 200 Generated with PyMOL and psico'.ljust(80), file=f)
# Write sequence
if int(seqres):
f.write(get_pdb_seqres(selection))
# Write the CRYST1 line if possible
if symm:
try:
obj1 = cmd.get_object_list(selection)[0]
sym = cmd.get_symmetry(obj1)
if len(sym) != 7:
raise
f.write("CRYST1%9.3f%9.3f%9.3f%7.2f%7.2f%7.2f %-10s%4d\n" %
tuple(sym + [1]))
if not quiet:
print(' Info: Wrote unit cell and space group info')
except:
if not quiet:
print(' Info: No crystal information')
# Write secondary structure
if ss:
try:
sss = get_pdb_sss(selection, state, quiet)
if not sss:
raise
f.write(sss)
if not quiet:
print(' Info: Wrote secondary structure info')
except:
if not quiet:
print(' Info: No secondary structure information')
# Write coordinates of selection
pdbstr = cmd.get_pdbstr(selection, state)
# fix END records
if state == 0 and cmd.get_version()[1] < 1.6:
pdbstr = '\n'.join(
line for line in pdbstr.splitlines() if line != 'END') + '\nEND\n'
# anisotropic b-factors
from . import pymol_version
if int(aniso) and pymol_version < 1.6 and \
cmd.get_model('first (%s)' % selection).atom[0].u_aniso[0] != 0.0:
def mergeaniso():
atom_it = iter(cmd.get_model(selection, state).atom)
for line in pdbstr.splitlines(True):
yield line
if line[:6] in ['ATOM ', 'HETATM']:
yield 'ANISOU' + line[6:28] + \
''.join('%7.0f' % (u*1e4) for u in atom_it.next().u_aniso) + \
line[70:]
pdbstr = ''.join(mergeaniso())
f.write(pdbstr)
f.close()
if not quiet:
print('Wrote PDB to \'' + filename + '\'')
def symset(prefix="sym", object=-1, x=0, y=0, z=0, opList=[]):
if object == -1:
object = cmd.get_names()[0]
cell = [float(x), float(y), float(z)]
view = cmd.get_view()
cmd.show("lines", object)
sgInfo = cmd.get_symmetry(object)
raw_ops = []
for s in sgtbx.space_group_info(sgInfo[6]).group():
raw_ops.append(str(s))
if len(opList) == 0:
for i in range(len(raw_ops)):
opList.append(i)
opMatrices = []
vars = ["x", "y", "z"]
#CREATE 4X4 MATRICES FOR SYMMETRY OPERATORS
for i, raw_op in enumerate(raw_ops):
ops = raw_op.split(",")
matrix = [[0, 0, 0, 0], [0, 0, 0, 0], [0, 0, 0, 0], [0, 0, 0, 1]]
for j in range(len(ops)):
for k in range(len(vars)):
index = ops[j].find(vars[k])
if index != -1:
if index == 0:
matrix[k][j] = 1
elif ops[j][index - 1] == "-":
matrix[k][j] = -1
else:
matrix[k][j] = 1
index = ops[j].find("/")
if index != -1:
matrix[3][j] = float(ops[j][index - 1]) / float(
ops[j][index + 1])
opMatrices.append(matrix)
a, b, c, alpha, beta, gamma = sgInfo[0:6]
ca = math.cos(math.radians(alpha))
cb = math.cos(math.radians(beta))
cg = math.cos(math.radians(gamma))
sb = math.sin(math.radians(beta))
sg = math.sin(math.radians(gamma))
fracToOrt = N.array([[a, b * cg, c * cb, 0.0],
[0.0, b * sg, c * (ca - cb * cg) / sg, 0.0],
[
0.0, 0.0,
c * sb * math.sqrt(1.0 - ((cb * cg - ca) /
(sb * sg))**2), 0.0
], [0.0, 0.0, 0.0, 1.0]])
fracToOrt = fracToOrt.transpose()
ortToFrac = inv(fracToOrt)
stored.atoms = []
cmd.iterate_state(1, object, "stored.atoms.append([x,y,z,1])")
stored.atoms = N.array(stored.atoms)
fracCoords = N.dot(stored.atoms, ortToFrac)
for i in opList:
try:
op = opMatrices[i]
except:
print("Bad symmetry partner numbers. Try again.")
quit()
copy = "%s%02d_%d_%d_%d" % (prefix, i, x, y, z)
cmd.copy(copy, object)
newCoordsFrac = N.dot(fracCoords, op)
stored.newCoords = N.dot(newCoordsFrac, fracToOrt)
stored.j = 0
cmd.alter_state(
1, copy,
"x,y,z = stored.newCoords[stored.j][0], stored.newCoords[stored.j][1], stored.newCoords[stored.j][2]; stored.j = stored.j + 1"
)
xSum = ySum = zSum = 0.0
for a, b, c in newCoordsFrac:
xSum += a
ySum += b
zSum += c
center = N.array([xSum, ySum, zSum])
center = center / len(stored.newCoords)
shift = [
cell[0] - math.floor(center[0]), cell[1] - math.floor(center[1]),
cell[2] - math.floor(center[2])
]
cell_shift(copy, shift[0], shift[1], shift[2], 0)
'''
#COPIES COORDINATES OF EACH ATOM TO CORRESPONDING ONE IN GIVEN SYMMETRY PARTNER
#cmd.alter_state(1, copy, "x,y,z = cmd.sym_partner([x,y,z], stored.tmpOp)")
#MOVES SYMMETRY PARTNER TO PROPER LATTICE COORDINATES AND CORRECTS FOR NATIVE LATTICE POSITION ERROR
#stored.xSum,stored.ySum,stored.zSum = 0.0,0.0,0.0
#atoms = cmd.count_atoms(copy)
#cmd.iterate_state(1, copy, "stored.xSum = stored.xSum + x; stored.ySum = stored.ySum + y; stored.zSum = stored.zSum + z")
#xMean = stored.xSum / atoms
#yMean = stored.ySum / atoms
#zMean = stored.zSum / atoms
#xError, yError, zError = N.dot(N.array([xMean,yMean,zMean]), stored.ortToFrac)
#dX,dY,dZ = cell[0]-math.floor(xError), cell[1]-math.floor(yError), cell[2]-math.floor(zError)
#cell_shift(copy,dX,dY,dZ, 0)
'''
cmd.hide("everything", object)
cmd.set_view(view)
def symexpcell(prefix='mate', object=None, a=0, b=0, c=0):
'''
DESCRIPTION
Creates all symmetry-related objects for the specified object that
occur with their bounding box center within the unit cell.
USAGE
symexpcell prefix, object, [a, b, c]
ARGUMENTS
prefix = string: prefix of new objects
object = string: object for which to create symmetry mates
a, b, c = integer: create neighboring cell {default: 0,0,0}
SEE ALSO
symexp
'''
import numpy
from pymol import xray
if object is None:
object = cmd.get_object_list()[0]
sym = cmd.get_symmetry(object)
cell_edges = sym[0:3]
cell_angles = sym[3:6]
spacegroup = sym[6]
basis = cellbasis(cell_angles, cell_edges)
basis = numpy.matrix(basis)
extent = cmd.get_extent(object)
center = sum(numpy.array(extent)) * 0.5
center = numpy.matrix(center.tolist() + [1.0]).T
center_cell = basis.I * center
spacegroup = xray.space_group_map.get(spacegroup, spacegroup)
i = 0
matrices = xray.sg_sym_to_mat_list(spacegroup)
for mat in matrices:
i += 1
mat = numpy.matrix(mat)
shift = -numpy.floor(numpy.array(mat * center_cell)[0:3, 0])
shift = shift.flatten().astype(int)
shift += [a, b, c]
mat[0:3, 3] += shift.reshape((3, 1))
mat = basis * mat * basis.I
mat_list = list(mat.flat)
name = '%s%d' % (prefix, i)
cmd.create(name, object)
cmd.transform_object(name, mat_list, 0)
cmd.set_title(name, 1, "{}_{}{}{}".format(i, *(shift + 5).tolist()))
if len(matrices) > 1:
cmd.color(i + 1, name)
def get_operations(object):
raw_ops = []
sgInfo = cmd.get_symmetry(object)
for s in sgtbx.space_group_info(sgInfo[6]).group():
raw_ops.append(str(s))
return raw_ops
def testLoadCryst1(self):
cmd.load(self.datafile('cryst1.mae'), 'm1')
s = cmd.get_symmetry('m1')
self.assertArrayEqual(s[:6], [50.84, 42.77, 28.95, 90., 90., 90.], delta=0.01)
self.assertEqual(s[6], 'P 21 21 21')
def save_traj(filename, selection='(all)', format='', box=0, quiet=1):
'''
DESCRIPTION
Save coordinates of a multi-state object to a trajectory file (DCD OR CRD).
Based on http://pymolwiki.org/index.php/Save2traj by Sean Law
USAGE
save_traj filename [, selection [, format ]]
ARGUMENTS
filename = string: file path to be written
selection = string: atoms to save
format = string: 'dcd' or 'crd' (alias 'charmm' or 'amber') {default:
determined from filename extension)
'''
box = int(box)
# Get NATOMS, NSTATES
NATOMS = cmd.count_atoms(selection)
NSTATES = cmd.count_states(selection)
# Determine Trajectory Format
if format == '' and '.' in filename:
format = filename.rsplit('.', 1)[1]
format = format.lower()
if format in ['charmm', 'dcd']:
Outfile = DCDOutfile
elif format in ['amber', 'trj', 'crd']:
Outfile = CRDOutfile
elif format in ['rst', 'rst7']:
Outfile = RSTOutfile
else:
print('Unknown format:', format)
raise CmdException
f = open(filename, 'wb')
# size of periodic box
if box:
try:
boxdim = cmd.get_symmetry(selection)[0:3]
except:
boxdim = [0, 0, 0]
else:
boxdim = None
outfile = Outfile(filename, NSTATES, NATOMS, box=boxdim)
# Write Trajectory Coordinates
for state in range(1, NSTATES + 1):
xyz = cmd.get_model(selection, state).get_coord_list()
outfile.writeCoordSet(xyz)
outfile.close()
if not int(quiet):
fmt = 'Wrote trajectory in %s format with %d atoms and %d frames to file %s'
print(fmt % (format, NATOMS, NSTATES, filename))
def save_traj(filename, selection='(all)', format='', box=0, quiet=1):
'''
DESCRIPTION
Save coordinates of a multi-state object to a trajectory file (DCD OR CRD).
Based on http://pymolwiki.org/index.php/Save2traj by Sean Law
USAGE
save_traj filename [, selection [, format ]]
ARGUMENTS
filename = string: file path to be written
selection = string: atoms to save
format = string: 'dcd' or 'crd' (alias 'charmm' or 'amber') {default:
determined from filename extension)
'''
box = int(box)
# Get NATOMS, NSTATES
NATOMS = cmd.count_atoms(selection)
NSTATES = cmd.count_states(selection)
# Determine Trajectory Format
if format == '' and '.' in filename:
format = filename.rsplit('.', 1)[1]
format = format.lower()
if format in ['charmm', 'dcd']:
Outfile = DCDOutfile
elif format in ['amber', 'trj', 'crd']:
Outfile = CRDOutfile
elif format in ['rst', 'rst7']:
Outfile = RSTOutfile
else:
print 'Unknown format:', format
raise CmdException
f = open(filename, 'wb')
# size of periodic box
if box:
try:
boxdim = cmd.get_symmetry(selection)[0:3]
except:
boxdim = [0,0,0]
else:
boxdim = None
outfile = Outfile(filename, NSTATES, NATOMS, box=boxdim)
# Write Trajectory Coordinates
for state in range(1, NSTATES+1):
xyz = cmd.get_model(selection, state).get_coord_list()
outfile.writeCoordSet(xyz)
outfile.close()
if not int(quiet):
fmt = 'Wrote trajectory in %s format with %d atoms and %d frames to file %s'
print fmt % (format, NATOMS, NSTATES, filename)
def symset(prefix = "sym", object = -1, x=0,y=0,z=0, opList = []):
if object == -1:
object = cmd.get_names()[0]
cell = [float(x),float(y),float(z)]
view = cmd.get_view()
cmd.show("lines", object)
sgInfo = cmd.get_symmetry(object)
raw_ops = []
for s in sgtbx.space_group_info(sgInfo[6]).group():
raw_ops.append(str(s))
if (len(opList) == 0):
for i in range(len(raw_ops)):
opList.append(i)
opMatrices = []
vars = ["x","y","z"]
i = 0
j = 0
k = 0
#CREATE 4X4 MATRICES FOR SYMMETRY OPERATORS
for raw_op in raw_ops:
ops = raw_op.split(",")
matrix = [[0,0,0,0],[0,0,0,0],[0,0,0,0],[0,0,0,1]]
for j in range(len(ops)):
for k in range(len(vars)):
index = ops[j].find(vars[k])
if index != -1:
if index == 0:
matrix[k][j] = 1
elif ops[j][index - 1] == "-":
matrix[k][j] = -1
else:
matrix[k][j] = 1
index = ops[j].find("/")
if index != -1:
matrix[3][j] = float(ops[j][index - 1]) / float(ops[j][index + 1])
opMatrices.append(matrix)
i=i+1
a,b,c,alpha,beta,gamma = sgInfo[0:6]
ca = math.cos(math.radians(alpha))
cb = math.cos(math.radians(beta))
cg = math.cos(math.radians(gamma))
sb = math.sin(math.radians(beta))
sg = math.sin(math.radians(gamma))
fracToOrt = N.array([[a, b * cg, c * cb, 0.0],
[0.0, b * sg, c * (ca - cb * cg) / sg, 0.0],
[0.0, 0.0, c * sb * math.sqrt(1.0 - ((cb * cg - ca) / (sb * sg))**2), 0.0],
[0.0,0.0,0.0,1.0]])
fracToOrt = fracToOrt.transpose()
ortToFrac = inv(fracToOrt)
stored.atoms = []
cmd.iterate_state(1,object,"stored.atoms.append([x,y,z,1])")
stored.atoms = N.array(stored.atoms)
fracCoords = N.dot(stored.atoms,ortToFrac)
for i in opList:
try:
op = opMatrices[i]
except:
print "Bad symmetry partner numbers. Try again."
quit()
if i > 9:
copy = prefix + str(i) + "_" + str(x) + "_" + str(y) + "_" + str(z)
else:
copy = prefix + "0" + str(i) + "_" + str(x) + "_" + str(y) + "_" + str(z)
cmd.copy(copy, object)
newCoordsFrac = N.dot(fracCoords, op)
stored.newCoords = N.dot(newCoordsFrac, fracToOrt)
stored.j = 0
cmd.alter_state(1,copy,"x,y,z = stored.newCoords[stored.j][0], stored.newCoords[stored.j][1], stored.newCoords[stored.j][2]; stored.j = stored.j + 1")
xSum=ySum=zSum=0.0
for k in range(len(newCoordsFrac)):
xSum = newCoordsFrac[k][0] + xSum
ySum = newCoordsFrac[k][1] + ySum
zSum = newCoordsFrac[k][2] + zSum
center = N.array([xSum,ySum,zSum])
center = center/len(stored.newCoords)
shift = [-math.floor(center[0]) + cell[0], -math.floor(center[1]) + cell[1], -math.floor(center[2]) + cell[2]]
cell_shift(copy,shift[0],shift[1],shift[2],0)
'''
#COPIES COORDINATES OF EACH ATOM TO CORRESPONDING ONE IN GIVEN SYMMETRY PARTNER
#cmd.alter_state(1, copy, "x,y,z = cmd.sym_partner([x,y,z], stored.tmpOp)")
#MOVES SYMMETRY PARTNER TO PROPER LATTICE COORDINATES AND CORRECTS FOR NATIVE LATTICE POSITION ERROR
#stored.xSum,stored.ySum,stored.zSum = 0.0,0.0,0.0
#atoms = cmd.count_atoms(copy)
#cmd.iterate_state(1, copy, "stored.xSum = stored.xSum + x; stored.ySum = stored.ySum + y; stored.zSum = stored.zSum + z")
#xMean = (stored.xSum / atoms)
#yMean = (stored.ySum / atoms)
#zMean = (stored.zSum / atoms)
#xError, yError, zError = N.dot(N.array([xMean,yMean,zMean]), stored.ortToFrac)
#dX,dY,dZ = -math.floor(xError) + cell[0], -math.floor(yError) + cell[1], -math.floor(zError) + cell[2]
#cell_shift(copy,dX,dY,dZ, 0)
'''
cmd.hide("everything", object)
cmd.set_view(view)
def get_operations(object):
raw_ops = []
sgInfo = cmd.get_symmetry(object)
for s in sgtbx.space_group_info(sgInfo[6]).group():
raw_ops.append(str(s))
return raw_ops
