def test_random_projection_generation(self):
""" Generation of random projection from a PDB file with em2d images"""
testfile = self.get_input_file_name("opencv_test.spi")
if os.path.isfile(testfile):
# delete the file to check
os.remove(testfile)
smodel = IMP.Model()
ssel = IMP.atom.ATOMPDBSelector()
prot = IMP.atom.read_pdb(self.get_input_file_name("1z5s.pdb"), smodel, ssel)
IMP.atom.add_radii(prot)
particles = IMP.core.get_leaves(prot)
rows = 80
cols = 80
resolution = 1
apix = 1.5
img = em2d.Image()
img.set_size(rows, cols)
srw = em2d.SpiderImageReaderWriter()
rr = em2d.RegistrationResult()
rr.set_random_registration(0, 5)
options = em2d.ProjectingOptions(apix, resolution)
options.srw = srw
em2d.get_projection(img, particles, rr, options)
img.write(testfile, srw)
self.assertTrue(os.path.isfile(testfile), "Projection image not generated")
os.remove(testfile)
def test_rigid_body_image_fit_restraint(self):
"""Test scoring with RigidBodiesImageFitRestraint"""
m = IMP.Model()
# read full complex
fn = self.get_input_file_name("1z5s.pdb")
prot = atom.read_pdb(fn, m, IMP.atom.ATOMPDBSelector())
# read components
names = ["1z5sA", "1z5sB", "1z5sC", "1z5sD"]
fn_pdbs = [self.get_input_file_name(name + ".pdb") for name in names]
components = [atom.read_pdb(fn, m, IMP.atom.ATOMPDBSelector())
for fn in fn_pdbs]
components_rbs = [atom.create_rigid_body(c) for c in components]
# img
R = alg.get_identity_rotation_3d()
reg = em2d.RegistrationResult(R)
img = em2d.Image()
img.set_size(80, 80)
srw = em2d.SpiderImageReaderWriter()
resolution = 5
pixel_size = 1.5
options = em2d.ProjectingOptions(pixel_size, resolution)
ls = core.get_leaves(prot)
em2d.get_projection(img, ls, reg, options)
# img.write("rbfit_test_image.spi",srw)
# set restraint
score_function = em2d.EM2DScore()
rb_fit = em2d.RigidBodiesImageFitRestraint(score_function,
components_rbs, img)
pp = em2d.ProjectingParameters(pixel_size, resolution)
rb_fit.set_projecting_parameters(pp)
# set the trivial case:
n_masks = 1
for rb in components_rbs:
# set as the only possible orientation the one that the rigid
# body already has
rb_fit.set_orientations(rb,
[rb.get_reference_frame().get_transformation_to().get_rotation()])
self.assertEqual(rb_fit.get_number_of_masks(rb), n_masks,
"Incorrect number rigid body masks")
# Calculate the positions of the rigid bodies respect to the centroid
# of the entire molecule
ls = core.get_leaves(prot)
xyzs = core.XYZs(ls)
centroid = core.get_centroid(xyzs)
coords = [rb.get_coordinates() - centroid for rb in components_rbs]
for rb, coord in zip(components_rbs, coords):
rb.set_coordinates(coord)
# Check that the value is a perfect registration
score = rb_fit.evaluate(False)
# print "score ...", score
# It seems that projecting with the masks is slightly less accurate
# I have to establish a tolerance of 0.03
self.assertAlmostEqual(score, 0, delta=0.03,
msg="Wrong value for the score %f " % (score))
def test_random_projection_generation(self):
"""Generation of random projection from a PDB file with em2d images"""
testfile = "opencv_test.spi"
if os.path.isfile(testfile):
# delete the file to check
os.remove(testfile)
smodel = IMP.Model()
ssel = IMP.atom.ATOMPDBSelector()
prot = IMP.atom.read_pdb(
self.get_input_file_name("1z5s.pdb"),
smodel,
ssel)
IMP.atom.add_radii(prot)
particles = IMP.core.get_leaves(prot)
rows = 80
cols = 80
resolution = 1
apix = 1.5
img = em2d.Image()
img.set_size(rows, cols)
srw = em2d.SpiderImageReaderWriter()
rr = em2d.RegistrationResult()
rr.set_random_registration(0, 5)
options = em2d.ProjectingOptions(apix, resolution)
options.srw = srw
em2d.get_projection(img, particles, rr, options)
img.write(testfile, srw)
self.assertTrue(os.path.isfile(testfile),
"Projection image not generated")
os.remove(testfile)
def test_registration(self):
"""Test the registration of 3 subjects from 1gyt.pdb at 0.5 SNR"""
# Get model from PDB file
smodel = IMP.Model()
ssel = IMP.atom.ATOMPDBSelector()
fn_model = self.get_input_file_name("1gyt.pdb")
prot = IMP.atom.read_pdb(fn_model, smodel, ssel)
particles = IMP.core.get_leaves(prot)
# Read subject images
srw = em2d.SpiderImageReaderWriter()
selection_file = self.get_input_file_name("1gyt-subjects-0.5-SNR.sel")
images_to_read_names = em2d.read_selection_file(selection_file)
for i in range(0, len(images_to_read_names)):
images_to_read_names[i] = self.get_input_file_name(
images_to_read_names[i])
subjects = em2d.read_images(images_to_read_names, srw)
self.assertEqual(len(subjects), 3, "Problem reading subject images")
# Generate 20 evenly distributed projections from the PDB file
n_projections = 20
proj_params = em2d.get_evenly_distributed_registration_results(
n_projections)
rows = 128
cols = 128
pixel_size = 1.5
# for generating projections, use a very high resolution
resolution = 8.5
options = em2d.ProjectingOptions(pixel_size, resolution)
projections = em2d.get_projections(particles, proj_params, rows, cols,
options)
self.assertEqual(len(projections), n_projections,
"Problem generating projections")
# Prepare registration
# IMP.set_log_level(IMP.VERBOSE)
finder = em2d.ProjectionFinder()
score_function = em2d.EM2DScore()
params = em2d.Em2DRestraintParameters(pixel_size, resolution,
n_projections)
params.save_match_images = False
params.coarse_registration_method = em2d.ALIGN2D_PREPROCESSING
params.optimization_steps = 30
params.simplex_initial_length = 0.1
params.simplex_minimum_size = 0.01
finder.setup(score_function, params)
finder.set_model_particles(particles)
finder.set_subjects(subjects)
finder.set_projections(projections)
finder.set_fast_mode(2)
finder.get_complete_registration()
# Recover the registration results:
registration_parameters = finder.get_registration_results()
fn_registration_results = "my_1gyt_registration.params"
em2d.write_registration_results(fn_registration_results,
registration_parameters)
# Read the correct registration results:
correct_parameters = em2d.read_registration_results(
self.get_input_file_name("1gyt-subjects-0.5-SNR.params"))
print("determined: ")
for r in registration_parameters:
print(r.get_rotation(), r.get_shift())
print("correct: ")
for r in correct_parameters:
print(r.get_rotation(), r.get_shift())
for i in range(0, len(registration_parameters)):
# Generate the registered projection
imgx = em2d.Image()
imgx.set_size(rows, cols)
em2d.get_projection(imgx, particles, registration_parameters[i],
options)
ccc = em2d.get_cross_correlation_coefficient(
subjects[i].get_data(), imgx.get_data())
print(i, "ccc", ccc)
snr = 0.5
theoretical_ccc = (snr / (1. + snr))**.5
self.assertAlmostEqual(
ccc,
theoretical_ccc,
delta=0.02,
msg="Error in registration of subject %d: ccc %8.3f "
"theoretical_ccc %8.3f " % (i, ccc, theoretical_ccc))
os.remove(fn_registration_results)
def test_registration(self):
"""Test the registration of 3 subjects from 1gyt.pdb at 0.5 SNR"""
# Get model from PDB file
smodel = IMP.Model()
ssel = IMP.atom.ATOMPDBSelector()
fn_model = self.get_input_file_name("1gyt.pdb")
prot = IMP.atom.read_pdb(fn_model, smodel, ssel)
particles = IMP.core.get_leaves(prot)
# Read subject images
srw = em2d.SpiderImageReaderWriter()
selection_file = self.get_input_file_name("1gyt-subjects-0.5-SNR.sel")
images_to_read_names = em2d.read_selection_file(selection_file)
for i in range(0, len(images_to_read_names)):
images_to_read_names[i] = self.get_input_file_name(
images_to_read_names[i])
subjects = em2d.read_images(images_to_read_names, srw)
self.assertEqual(len(subjects), 3, "Problem reading subject images")
# Generate 20 evenly distributed projections from the PDB file
n_projections = 20
proj_params = em2d.get_evenly_distributed_registration_results(
n_projections)
rows = 128
cols = 128
pixel_size = 1.5
# for generating projections, use a very high resolution
resolution = 1
options = em2d.ProjectingOptions(pixel_size, resolution)
projections = em2d.get_projections(particles, proj_params,
rows, cols, options)
self.assertEqual(len(projections), n_projections,
"Problem generating projections")
# Prepare registration
# IMP.set_log_level(IMP.VERBOSE)
finder = em2d.ProjectionFinder()
score_function = em2d.EM2DScore()
params = em2d.Em2DRestraintParameters(pixel_size,
resolution, n_projections)
params.save_match_images = False
params.coarse_registration_method = em2d.ALIGN2D_PREPROCESSING
params.optimization_steps = 30
params.simplex_initial_length = 0.1
params.simplex_minimum_size = 0.01
finder.setup(score_function, params)
finder.set_model_particles(particles)
finder.set_subjects(subjects)
finder.set_projections(projections)
finder.set_fast_mode(2)
finder.get_complete_registration()
# Recover the registration results:
registration_parameters = finder.get_registration_results()
fn_registration_results = "my_1gyt_registration.params"
em2d.write_registration_results(fn_registration_results,
registration_parameters)
# Read the correct registration results:
correct_parameters = em2d.read_registration_results(
self.get_input_file_name("1gyt-subjects-0.5-SNR.params"))
print("determined: ")
for r in registration_parameters:
print(r.get_rotation(), r.get_shift())
print("correct: ")
for r in correct_parameters:
print(r.get_rotation(), r.get_shift())
for i in range(0, len(registration_parameters)):
# Generate the registered projection
imgx = em2d.Image()
imgx.set_size(rows, cols)
em2d.get_projection(imgx, particles,
registration_parameters[i], options)
ccc = em2d.get_cross_correlation_coefficient(
subjects[i].get_data(),
imgx.get_data())
print(i, "ccc", ccc)
snr = 0.5
theoretical_ccc = (snr / (1. + snr)) ** .5
self.assertAlmostEqual(ccc, theoretical_ccc, delta=0.02,
msg="Error in registration of subject %d: ccc %8.3f "
"theoretical_ccc %8.3f " % (i, ccc, theoretical_ccc))
os.remove(fn_registration_results)
def test_rigid_body_image_fit_restraint(self):
"""Test scoring with RigidBodiesImageFitRestraint"""
m = IMP.kernel.Model()
# read full complex
fn = self.get_input_file_name("1z5s.pdb")
prot = atom.read_pdb(fn, m, IMP.atom.ATOMPDBSelector())
# read components
names = ["1z5sA", "1z5sB", "1z5sC", "1z5sD"]
fn_pdbs = [self.get_input_file_name(name + ".pdb") for name in names]
components = [
atom.read_pdb(fn, m, IMP.atom.ATOMPDBSelector()) for fn in fn_pdbs
]
components_rbs = [atom.create_rigid_body(c) for c in components]
# img
R = alg.get_identity_rotation_3d()
reg = em2d.RegistrationResult(R)
img = em2d.Image()
img.set_size(80, 80)
srw = em2d.SpiderImageReaderWriter()
resolution = 5
pixel_size = 1.5
options = em2d.ProjectingOptions(pixel_size, resolution)
ls = core.get_leaves(prot)
em2d.get_projection(img, ls, reg, options)
# img.write("rbfit_test_image.spi",srw)
# set restraint
score_function = em2d.EM2DScore()
rb_fit = em2d.RigidBodiesImageFitRestraint(score_function,
components_rbs, img)
pp = em2d.ProjectingParameters(pixel_size, resolution)
rb_fit.set_projecting_parameters(pp)
# set the trivial case:
n_masks = 1
for rb in components_rbs:
# set as the only possible orientation the one that the rigid
# body already has
rb_fit.set_orientations(rb, [
rb.get_reference_frame().get_transformation_to().get_rotation(
)
])
self.assertEqual(rb_fit.get_number_of_masks(rb), n_masks,
"Incorrect number rigid body masks")
# Calculate the positions of the rigid bodies respect to the centroid
# of the entire molecule
ls = core.get_leaves(prot)
xyzs = core.XYZs(ls)
centroid = core.get_centroid(xyzs)
coords = [rb.get_coordinates() - centroid for rb in components_rbs]
for rb, coord in zip(components_rbs, coords):
rb.set_coordinates(coord)
# Check that the value is a perfect registration
m.add_restraint(rb_fit)
score = rb_fit.evaluate(False)
# print "score ...", score
# It seems that projecting with the masks is slightly less accurate
# I have to establish a tolerance of 0.03
self.assertAlmostEqual(score,
0,
delta=0.03,
msg="Wrong value for the score %f " % (score))
