mirror of
https://github.com/schrodinger/pymol-open-source.git
synced 2026-06-04 20:04:21 +08:00
466 lines
17 KiB
Python
466 lines
17 KiB
Python
'''
|
|
unit tests for pymol.exporting
|
|
'''
|
|
|
|
import os
|
|
import sys
|
|
import tempfile
|
|
import Image
|
|
import unittest
|
|
|
|
import pymol.exporting
|
|
from pymol import cmd, testing, stored
|
|
|
|
class TestExporting(testing.PyMOLTestCase):
|
|
|
|
def testCache(self):
|
|
for action in pymol.exporting.cache_action_dict:
|
|
cmd.cache(action)
|
|
|
|
def testCopyImage(self):
|
|
cmd.copy_image
|
|
self.skipTest("TODO")
|
|
|
|
def testGetFastastr(self):
|
|
seq, name = 'ACD', 'm1'
|
|
cmd.fab(seq, name)
|
|
s = cmd.get_fastastr()
|
|
lines = s.split()
|
|
self.assertTrue(lines[0] in (
|
|
'>' + name,
|
|
'>' + name + '_',
|
|
))
|
|
self.assertEqual(lines[1:], [seq])
|
|
|
|
def testGetPdbstr(self):
|
|
cmd.pseudoatom()
|
|
lines = cmd.get_pdbstr().splitlines()
|
|
self.assertTrue(lines[0].startswith('HETATM'))
|
|
self.assertTrue(lines[1].startswith('END'))
|
|
|
|
def testGetSession(self):
|
|
cmd.fragment('ala')
|
|
x = cmd.count_atoms()
|
|
s = cmd.get_session()
|
|
cmd.reinitialize()
|
|
cmd.set_session(s)
|
|
self.assertEqual(['ala'], cmd.get_names())
|
|
self.assertEqual(x, cmd.count_atoms())
|
|
|
|
@testing.requires_version('2.5')
|
|
def testGetSession25(self):
|
|
cmd.set('pse_binary_dump', 0)
|
|
cmd.set('pse_export_version', 0)
|
|
s = cmd.get_session(binary=-1, version=-1)
|
|
self.assertEqual(cmd.get_setting_int('pse_binary_dump'), 0)
|
|
self.assertEqual(cmd.get_setting_float('pse_export_version'), 0.0)
|
|
cmd.set_session(s)
|
|
self.assertEqual(cmd.get_setting_int('pse_binary_dump'), 0)
|
|
self.assertEqual(cmd.get_setting_float('pse_export_version'), 0.0)
|
|
s = cmd.get_session(binary=True, version=1.2)
|
|
self.assertEqual(cmd.get_setting_int('pse_binary_dump'), 0)
|
|
self.assertEqual(cmd.get_setting_float('pse_export_version'), 0.0)
|
|
cmd.set_session(s)
|
|
self.assertEqual(cmd.get_setting_int('pse_binary_dump'), 1)
|
|
self.assertEqual(cmd.get_setting_float('pse_export_version'), 1.2)
|
|
|
|
@testing.requires_version('1.8.4')
|
|
def testMultisave(self):
|
|
names = ['ala', 'gly', 'his', 'arg']
|
|
|
|
for name in names:
|
|
cmd.fragment(name)
|
|
|
|
sym1 = (1.2, 3.4, 5.6, 60.0, 70.0, 80.0)
|
|
sym2 = (2.3, 4.5, 6.7, 90.0, 90.0, 90.0)
|
|
sym3 = (3.4, 5.6, 7.8, 70.0, 70.0, 70.0)
|
|
cmd.set_symmetry("gly", *sym1)
|
|
cmd.set_symmetry("his", *sym2)
|
|
cmd.set_symmetry("arg", *sym3)
|
|
|
|
with testing.mktemp('.pdb') as filename:
|
|
cmd.multisave(filename, names[0]) # new
|
|
cmd.multisave(filename, names[1]) # new (overwrite)
|
|
cmd.multisave(filename, ' '.join(names[2:]), append=1) # append
|
|
cmd.delete('*')
|
|
cmd.load(filename)
|
|
|
|
self.assertEqual(cmd.get_object_list(), names[1:])
|
|
|
|
self.assertArrayEqual(cmd.get_symmetry("gly")[:6], sym1, delta=1e-4)
|
|
self.assertArrayEqual(cmd.get_symmetry("his")[:6], sym2, delta=1e-4)
|
|
self.assertArrayEqual(cmd.get_symmetry("arg")[:6], sym3, delta=1e-4)
|
|
|
|
@testing.requires_version('2.1')
|
|
def testMultifilesave(self):
|
|
import glob
|
|
|
|
for name in ['ala', 'gly', 'his', 'arg']:
|
|
cmd.fragment(name)
|
|
|
|
# multistate
|
|
for i in range(2, 11):
|
|
cmd.create('ala', 'ala', 1, i)
|
|
|
|
for fmt in ['{}-{:02}.cif', '{}-{state}.cif']:
|
|
with testing.mkdtemp() as dirname:
|
|
cmd.multifilesave(os.path.join(dirname, fmt), 'ala', 0)
|
|
filenames = [os.path.basename(p) for p in glob.glob(os.path.join(dirname, '*.cif'))]
|
|
self.assertEqual(len(filenames), 10)
|
|
self.assertTrue('ala-03.cif' in filenames)
|
|
|
|
with testing.mkdtemp() as dirname:
|
|
cmd.multifilesave(os.path.join(dirname, '{}.pdb'), 'a* g*')
|
|
filenames_full = sorted(glob.glob(os.path.join(dirname, '*.pdb')))
|
|
filenames = [os.path.basename(p) for p in filenames_full]
|
|
self.assertEqual(filenames, ['ala.pdb', 'arg.pdb', 'gly.pdb'])
|
|
|
|
cmd.delete('*')
|
|
cmd.load(filenames_full[0])
|
|
self.assertEqual(cmd.count_atoms(), 10)
|
|
|
|
cmd.delete('*')
|
|
cmd.load(filenames_full[1])
|
|
self.assertEqual(cmd.count_atoms(), 24)
|
|
|
|
@testing.foreach(
|
|
('0.5in', '0.25in', 200, (100, 50)), # inch
|
|
('2.54cm', '1.27cm', 100, (100, 50)), # centimeter
|
|
(100, 0, -1, (100, 75)), # px, only width
|
|
(0, 75, -1, (100, 75)), # px, only height
|
|
)
|
|
@testing.requires('no_edu') # ray
|
|
@testing.requires_version('1.6.0')
|
|
def testPng(self, w, h, dpi, size):
|
|
with testing.mktemp('.png') as filename:
|
|
cmd.png(filename, w, h, dpi, ray=1)
|
|
self.assertEqual(size, Image.open(filename).size)
|
|
|
|
@testing.requires_version('2.5')
|
|
def testPngNoFile(self):
|
|
self.ambientOnly()
|
|
cmd.fragment('gly')
|
|
cmd.show_as('spheres')
|
|
cmd.color('yellow')
|
|
cmd.zoom(complete=1)
|
|
|
|
cmd.refresh() # TODO should not be necessary, see also PYMOL-3328
|
|
|
|
ncol, nrow = 120, 80
|
|
buf = cmd.png(None, ncol, nrow)
|
|
self.assertTrue(buf.startswith(b'\x89PNG'))
|
|
import io
|
|
img = self.get_imagearray(Image.open(io.BytesIO(buf)))
|
|
self.assertEqual(img.shape[:2], (nrow, ncol))
|
|
self.assertImageHasColor('yellow', img)
|
|
|
|
# not supported in older versions: xyz (no ref)
|
|
@testing.foreach('pdb', 'sdf', 'mol', 'mol2')
|
|
def testSaveRef(self, format):
|
|
# for rms_cur (not all formats save all identifiers)
|
|
m = -1
|
|
cmd.set('retain_order')
|
|
|
|
cmd.fragment('ala', 'm1')
|
|
cmd.copy('m2', 'm1')
|
|
cmd.copy('m3', 'm1')
|
|
|
|
cmd.rotate('y', 90, 'm2')
|
|
cmd.align('m3', 'm2')
|
|
|
|
# with ref=m3
|
|
with testing.mktemp('.' + format) as filename:
|
|
cmd.save(filename, 'm2', ref='m3')
|
|
cmd.load(filename, 'm4')
|
|
self.assertAlmostEqual(cmd.rms_cur('m4', 'm1', matchmaker=m), 0.00, delta=1e-2)
|
|
self.assertAlmostEqual(cmd.rms_cur('m4', 'm2', matchmaker=m), 1.87, delta=1e-2)
|
|
|
|
# without ref
|
|
with testing.mktemp('.' + format) as filename:
|
|
cmd.save(filename, 'm2')
|
|
cmd.load(filename, 'm5')
|
|
self.assertAlmostEqual(cmd.rms_cur('m5', 'm2', matchmaker=m), 0.00, delta=1e-2)
|
|
self.assertAlmostEqual(cmd.rms_cur('m5', 'm1', matchmaker=m), 1.87, delta=1e-2)
|
|
|
|
@testing.foreach(
|
|
('pdb', 1.2),
|
|
('sdf', 1.7),
|
|
('mol2', 1.7),
|
|
('xyz', 1.8),
|
|
('mol', 1.831),
|
|
('mae', 1.831),
|
|
('cif', 1.8),
|
|
)
|
|
def testSaveState(self, format, pymol_version):
|
|
if pymol_version > testing.PYMOL_VERSION[1]:
|
|
self.skipTest("version %f" % (pymol_version))
|
|
|
|
# for rms_cur (not all formats save all identifiers)
|
|
m = -1
|
|
cmd.set('retain_order')
|
|
|
|
# create a multistate object
|
|
cmd.fragment('ala', 'm1')
|
|
cmd.create('m1', 'm1', 1, 2)
|
|
cmd.create('m1', 'm1', 1, 3)
|
|
cmd.translate([5, 0, 0], 'm1', state=2)
|
|
cmd.translate([0, 5, 0], 'm1', state=3)
|
|
n_states = cmd.count_states('m1')
|
|
|
|
with testing.mktemp('.' + format) as filename:
|
|
if format == 'mae' and not pymol.invocation.options.incentive_product:
|
|
format = 'cms'
|
|
|
|
# explicit
|
|
for state in range(1, n_states + 1):
|
|
cmd.delete('m2')
|
|
cmd.save(filename, 'm1', state=state)
|
|
cmd.load(filename, 'm2', format=format)
|
|
rms = cmd.rms_cur('m1', 'm2', state, 1, matchmaker=m)
|
|
self.assertAlmostEqual(rms, 0.00, delta=1e-2)
|
|
|
|
# current state
|
|
for state in range(1, n_states + 1):
|
|
cmd.frame(state)
|
|
cmd.delete('m2')
|
|
cmd.save(filename, 'm1')
|
|
cmd.load(filename, 'm2', 1, format=format)
|
|
rms = cmd.rms_cur('m1', 'm2', state, 1, matchmaker=m)
|
|
self.assertAlmostEqual(rms, 0.00, delta=1e-2)
|
|
|
|
if format in ('mol', 'cms'):
|
|
# no multi support
|
|
return
|
|
|
|
# all states
|
|
cmd.delete('m2')
|
|
cmd.save(filename, 'm1', state=0)
|
|
cmd.load(filename, 'm2', 1, discrete=1, multiplex=0)
|
|
self.assertEqual(cmd.count_states('m2'), n_states)
|
|
for state in range(1, n_states + 1):
|
|
rms = cmd.rms_cur('m1', 'm2 and state %d' % state,
|
|
state, state, matchmaker=m)
|
|
self.assertAlmostEqual(rms, 0.00, delta=1e-2)
|
|
|
|
@testing.foreach(
|
|
('pdb', 1.2),
|
|
('sdf', 1.7),
|
|
('mol2', 1.7),
|
|
('xyz', 1.7),
|
|
('mol', 1.7),
|
|
('mae', 1.831),
|
|
('cif', 1.8),
|
|
)
|
|
def testSaveSelection(self, format, pymol_version):
|
|
if pymol_version > testing.PYMOL_VERSION[1]:
|
|
self.skipTest("version %f" % (pymol_version))
|
|
|
|
cmd.fragment('trp', 'm1')
|
|
cmd.fragment('glu', 'm2')
|
|
|
|
n_O = cmd.count_atoms('elem O')
|
|
n_N = cmd.count_atoms('elem N')
|
|
|
|
with testing.mktemp('.' + format) as filename:
|
|
if testing.PYMOL_VERSION_TUPLE < (1, 8):
|
|
cmd.set('raise_exceptions', 0) # 1.7.6 save xyz doesn't set r=DEFAULT_SUCCESS
|
|
cmd.save(filename, 'elem O+N')
|
|
if testing.PYMOL_VERSION_TUPLE < (1, 8):
|
|
cmd.set('raise_exceptions', 1)
|
|
|
|
if format == 'mae' and not pymol.invocation.options.incentive_product:
|
|
format = 'cms'
|
|
|
|
cmd.delete('*')
|
|
cmd.load(filename, 'm2', discrete=1, format=format) # avoid merging of atoms
|
|
|
|
self.assertEqual(n_O, cmd.count_atoms('elem O'))
|
|
self.assertEqual(n_N, cmd.count_atoms('elem N'))
|
|
self.assertEqual(n_O + n_N, cmd.count_atoms())
|
|
|
|
@testing.foreach('pdb', 'cif', 'mmtf')
|
|
@testing.requires_version('2.5')
|
|
def testSave_symmetry(self, format):
|
|
sym1 = (1.2, 3.4, 5.6, 60.0, 70.0, 80.0, "P 3")
|
|
cmd.fragment("gly", "m1")
|
|
cmd.set_symmetry("m1", *sym1)
|
|
|
|
with testing.mktemp('.' + format) as filename:
|
|
cmd.save(filename)
|
|
cmd.delete('*')
|
|
cmd.load(filename, "m2")
|
|
|
|
sym = cmd.get_symmetry("m2")
|
|
self.assertArrayEqual(sym[:6], sym1[:6], delta=1e-4)
|
|
self.assertEqual(sym[6], sym1[6])
|
|
|
|
# cmp_atom : compares all fields in Atom (see chempy/__init__.py)
|
|
# except the id (which is unique to the instance)
|
|
def cmp_atom(self, selfobj,other):
|
|
cmp = lambda a, b: a != b # py3k
|
|
return \
|
|
cmp(type(selfobj), type(other)) or \
|
|
cmp(selfobj.segi, other.segi) or \
|
|
cmp(selfobj.chain, other.chain) or \
|
|
cmp(selfobj.resi_number, other.resi_number) or \
|
|
cmp(selfobj.resi, other.resi) or \
|
|
cmp(selfobj.resn, other.resn) or \
|
|
cmp(selfobj.symbol, other.symbol) or \
|
|
cmp(selfobj.name, other.name)
|
|
|
|
# cmp_bond : compares all fields in Bond (see chempy/__init__.py)
|
|
def cmp_bond(self, selfobj,other):
|
|
cmp = lambda a, b: a != b # py3k
|
|
return \
|
|
cmp(selfobj.order, other.order) or \
|
|
cmp(selfobj.index, other.index)
|
|
|
|
def assertModelsAreSame(self, m1, m2):
|
|
self.assertTrue(len(m1.atom) == len(m2.atom))
|
|
idx = 0
|
|
for m1atomidx in m1.atom:
|
|
self.assertTrue(self.cmp_atom(m1atomidx, m2.atom[idx]) == 0)
|
|
idx = idx + 1
|
|
idx = 0
|
|
for m1bondidx in m1.bond:
|
|
self.assertTrue(self.cmp_bond(m1bondidx, m2.bond[idx]) == 0)
|
|
idx = idx + 1
|
|
|
|
@testing.requires_version('1.7.6')
|
|
def testPSEBulkImport(self):
|
|
cmd.load(self.datafile('1rx1_1766_bulk.pse.gz'))
|
|
m1 = cmd.get_model()
|
|
cmd.load(self.datafile('1rx1_176.pse.gz'))
|
|
m2 = cmd.get_model()
|
|
self.assertModelsAreSame(m1, m2)
|
|
|
|
@testing.foreach.product((0, 1.7, 1.76, 1.8, 1.82, 1.9), (0, 1))
|
|
@testing.requires_version('1.7.6.5')
|
|
def testPSEBulkExportImport(self, pse_export_version, pse_binary_dump):
|
|
with testing.mktemp('.pse') as filename:
|
|
cmd.load(self.datafile("1oky-frag.pdb"))
|
|
m1 = cmd.get_model()
|
|
cmd.set("pse_export_version", pse_export_version)
|
|
cmd.set("pse_binary_dump", pse_binary_dump)
|
|
cmd.save(filename)
|
|
cmd.reinitialize()
|
|
cmd.load(filename)
|
|
m2 = cmd.get_model()
|
|
self.assertModelsAreSame(m1, m2)
|
|
|
|
def testGetModelObjectName(self):
|
|
cmd.load(self.datafile("1oky-frag.pdb"))
|
|
cmd.load(self.datafile('1rna.cif'))
|
|
|
|
m1 = cmd.get_model()
|
|
cnt = cmd.count_atoms('%1oky-frag')
|
|
|
|
self.assertEqual(m1.atom[0].model, '1oky-frag')
|
|
self.assertEqual(m1.atom[-1].model, '1rna')
|
|
self.assertEqual(m1.atom[cnt].model, '1rna')
|
|
|
|
|
|
@testing.requires_version('2.1')
|
|
def testMMTF(self):
|
|
'''Styled MMTF export/import'''
|
|
S = 0b10 # 1 << 1 spheres
|
|
D = 0b1000000000 # 1 << 9 dots
|
|
B = 2 # blue
|
|
R = 4 # red
|
|
|
|
cmd.fragment('gly')
|
|
cmd.color(B)
|
|
cmd.color(R, 'elem C')
|
|
cmd.show_as('spheres')
|
|
cmd.show_as('dots', 'elem C')
|
|
|
|
with testing.mktemp('.mmtf') as filename:
|
|
cmd.save(filename)
|
|
cmd.delete('*')
|
|
cmd.load(filename)
|
|
|
|
color_list = []
|
|
reps_list = []
|
|
|
|
cmd.iterate('*', 'color_list.append(color)', space=locals())
|
|
cmd.iterate('*', 'reps_list.append(reps)', space=locals())
|
|
|
|
self.assertEqual(color_list, [B, R, R, B, B, B, B])
|
|
self.assertEqual(reps_list, [S, D, D, S, S, S, S])
|
|
|
|
@testing.requires_version('2.1')
|
|
def testMMTFExportSele(self):
|
|
cmd.fab('ACDE')
|
|
|
|
with testing.mktemp('.mmtf') as filename:
|
|
cmd.save(filename, 'resn CYS+ASP')
|
|
cmd.delete('*')
|
|
cmd.load(filename)
|
|
|
|
self.assertEqual(cmd.count_atoms(), 23)
|
|
self.assertEqual(cmd.count_atoms('bound_to ASP/CG'), 3)
|
|
self.assertEqual(cmd.get_model('ASP/CG ASP/OD1').bond[0].order, 2)
|
|
self.assertEqual(cmd.get_model('ASP/CG ASP/OD2').bond[0].order, 1)
|
|
self.assertEqual(cmd.get_model('CYS/C ASP/N').bond[0].order, 1)
|
|
|
|
@testing.requires_version('2.1')
|
|
def testMMTFExportEmpty(self):
|
|
with testing.mktemp('.mmtf') as filename:
|
|
cmd.save(filename)
|
|
cmd.load(filename, 'm1')
|
|
|
|
self.assertEqual(cmd.count_atoms(), 0)
|
|
self.assertEqual(cmd.get_names(), ['m1'])
|
|
|
|
@testing.requires_version('2.4')
|
|
def testglTF(self):
|
|
'''glTF export'''
|
|
cmd.fragment('gly')
|
|
|
|
with testing.mktemp('.gltf') as filename:
|
|
self.assertEqual(cmd.save(filename), 0)
|
|
cmd.delete('*')
|
|
|
|
def testSaveAln(self):
|
|
cmd.fab('ACDEFGH', 'm1')
|
|
cmd.fab('ACDFGH', 'm2')
|
|
cmd.align('m1', 'm2', cycles=0, object='aln')
|
|
with testing.mktemp('.aln') as filename:
|
|
cmd.save(filename)
|
|
with open(filename) as handle:
|
|
lines = list(handle)
|
|
self.assertEqual(lines[0].split(), ['CLUSTAL'])
|
|
self.assertEqual(lines[1].split(), [])
|
|
self.assertEqual(lines[2].split(), ['m1', 'ACDEFGH'])
|
|
self.assertEqual(lines[3].split(), ['m2', 'ACD-FGH'])
|
|
|
|
@testing.requires_version('2.3')
|
|
def testSaveAlnNucleic(self):
|
|
cmd.load(self.datafile('1rna.cif'))
|
|
cmd.create('m1', 'chain A and not resi 6-7')
|
|
cmd.create('m2', 'chain B')
|
|
cmd.alter('m1 and resi 1', 'resn = "DT"') # mimic DNA
|
|
cmd.alter('m2 and resi 20', 'resn = "UNK"') # mimic nonstd residue
|
|
cmd.align('m1', 'm2', cycles=0, object='aln')
|
|
with testing.mktemp('.aln') as filename:
|
|
cmd.save(filename)
|
|
with open(filename) as handle:
|
|
lines = list(handle)
|
|
self.assertEqual(lines[0].split(), ['CLUSTAL'])
|
|
self.assertEqual(lines[1].split(), [])
|
|
self.assertEqual(lines[2].split(), ['m1', 'TUAUA--UAUAUAA'])
|
|
self.assertEqual(lines[3].split(), ['m2', 'UUAUA?AUAUAUAA'])
|
|
|
|
@testing.requires_version('2.0')
|
|
def testAssingAtomTypes(self):
|
|
cmd.fragment('his')
|
|
cmd.alter('all', 'text_type = "none"')
|
|
mytypes = {}
|
|
cmd.iterate('all', 'mytypes[name] = text_type', space=locals())
|
|
self.assertEqual(mytypes['NE2'], 'none')
|
|
pymol.exporting.assign_atom_types('all', 'mol2')
|
|
cmd.iterate('all', 'mytypes[name] = text_type', space=locals())
|
|
self.assertEqual(mytypes['NE2'], 'N.pl3')
|
|
self.assertEqual(mytypes['N'], 'N.3')
|
|
self.assertEqual(mytypes['CG'], 'C.2')
|