def test_even_projections(self):
""" Evenly distributed em2d image projections from a PDB file"""
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)
n_projections = 3
rows = 80
cols = 80
resolution = 1
apix = 1.5
srw = em2d.SpiderImageReaderWriter()
registration_values = em2d.get_evenly_distributed_registration_results(
n_projections)
options = em2d.ProjectingOptions(apix, resolution)
projections = em2d.get_projections(particles,
registration_values, rows, cols, options)
# Read the stored projections
stored_projection_names = em2d.create_filenames(
n_projections, "1z5s-fast-projection", "spi")
for n in range(0, n_projections):
stored_projection_names[n] = self.get_input_file_name(
stored_projection_names[n])
stored_projections = em2d.read_images(stored_projection_names, srw)
# check
for n in range(0, n_projections):
for i in range(0, rows):
for j in range(0, cols):
self.assertAlmostEqual(projections[n](i, j),
stored_projections[n](i, j), delta=0.001,
msg="Projections generated and stored are different")
def test_read_and_write_opencv_images(self):
"""Test read/write for the images using OpenCV data storage"""
srw = em2d.SpiderImageReaderWriter()
img = em2d.Image()
img.read(self.get_input_file_name("1gyt-subject-1-0.5-SNR.spi"), srw)
rows = int(img.get_header().get_number_of_rows())
cols = int(img.get_header().get_number_of_columns())
self.assertEqual(rows, 128, "Error reading image")
self.assertEqual(cols, 128, "Error reading image")
temp = "opencv_test_image.spi"
img.write(temp, srw)
img2 = em2d.Image()
img2.read(temp, srw)
for i in range(0, rows):
for j in range(0, cols):
self.assertAlmostEqual(
img(i, j),
img2(i, j),
delta=0.001,
msg="Generated image is different from stored")
os.remove(temp)
def test_do_extend_borders(self):
"""Test that extending the borders of an image is done correctly"""
srw = em2d.SpiderImageReaderWriter()
fn_img1 = self.get_input_file_name("lena-256x256.spi")
img1 = em2d.Image(fn_img1, srw)
img2 = em2d.Image()
border = 10
em2d.do_extend_borders(img1, img2, border)
rows2 = int(img2.get_header().get_number_of_rows())
cols2 = int(img2.get_header().get_number_of_columns())
self.assertEqual(rows2, 256 + 2 * border,
"Border rows are not extended properly")
self.assertEqual(cols2, 256 + 2 * border,
"Border columns are not extended properly")
for i in range(0, rows2):
for j in range(0, cols2):
if (i < border or i > (rows2 - border)):
self.assertAlmostEqual(
img2(i, j),
0,
delta=0.001,
msg="Borders are not zero at row %d col %d" % (i, j))
if (j < border or i > (cols2 - border)):
self.assertAlmostEqual(
img2(i, j),
0,
delta=0.001,
msg="Borders are not zero at row %d col %d" % (i, j))
def test_read_jpg(self):
"""Test of JPGReaderWriter reading"""
srw = em2d.SpiderImageReaderWriter()
jrw = em2d.JPGImageReaderWriter()
fn_jpg_img = self.get_input_file_name("lena-256x256.jpg")
jpg_img = em2d.Image(fn_jpg_img, jrw)
fn_spider_img = self.get_input_file_name("lena-256x256.spi")
spider_img = em2d.Image(fn_spider_img, srw)
rows = int(jpg_img.get_header().get_number_of_rows())
cols = int(jpg_img.get_header().get_number_of_columns())
self.assertEqual(spider_img.get_header().get_number_of_rows(), rows)
self.assertEqual(spider_img.get_header().get_number_of_columns(), cols)
for i in range(0, rows):
for j in range(0, cols):
# due to rounding, integer numbers in the jpg file can vary
# to the next integer. Allow delta 1
self.assertAlmostEqual(
abs(spider_img(i, j) - jpg_img(i, j)),
0,
delta=1,
msg="JPG image is not equal to spider image "
"at pixel (%d,%d)" % (i, j))
def get_em2d_restraint(assembly,
images_selection_file,
restraint_params,
mode="fast",
n_optimized=1):
""" Sets a restraint for comparing the model to a set of EM images
"""
model = assembly.get_model()
# Setup the restraint
sc = em2d.EM2DScore()
r = em2d.Em2DRestraint(model)
r.setup(sc, restraint_params)
names = em2d.read_selection_file(images_selection_file)
names = [base.get_relative_path(images_selection_file, x) for x in names]
log.debug("names of the images %s", names)
srw = em2d.SpiderImageReaderWriter()
imgs = em2d.read_images(names, srw)
r.set_images(imgs)
ps = atom.get_leaves(assembly)
lsc = container.ListSingletonContainer(ps)
r.set_particles(lsc)
if (mode == "coarse"):
r.set_coarse_registration_mode(True)
elif (mode == "fast"):
r.set_fast_mode(n_optimized)
elif (mode == "complete"):
pass
else:
raise ValueError("Em2DRestraint mode not recognized")
return r
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_read_tiff(self):
"""Test of TIFFReaderWriter reading"""
srw = em2d.SpiderImageReaderWriter()
trw = em2d.TIFFImageReaderWriter()
fn_tif_img = self.get_input_file_name("lena-256x256.tif")
tif_img = em2d.Image(fn_tif_img, trw)
fn_spider_img = self.get_input_file_name("lena-256x256.spi")
spider_img = em2d.Image(fn_spider_img, srw)
rows = int(tif_img.get_header().get_number_of_rows())
cols = int(tif_img.get_header().get_number_of_columns())
self.assertEqual(spider_img.get_header().get_number_of_rows(), rows)
self.assertEqual(spider_img.get_header().get_number_of_columns(), cols)
ccc = em2d.get_cross_correlation_coefficient(tif_img.get_data(),
spider_img.get_data())
self.assertAlmostEqual(ccc, 1, delta=0.01, msg="ccc ins not 1")
def test_polar_resampling(self):
"""Test of polar resampling of images"""
srw = em2d.SpiderImageReaderWriter()
fn_input = self.get_input_file_name("1gyt-subject-1-0.5-SNR.spi")
img = em2d.Image(fn_input, srw)
polar_params = em2d.PolarResamplingParameters()
polar = em2d.Image()
em2d.do_resample_polar(img, polar, polar_params)
fn_saved = self.get_input_file_name("1gyt-subject-1-0.5-SNR-polar.spi")
saved = em2d.Image(fn_saved, srw)
rows = int(polar.get_header().get_number_of_rows())
cols = int(polar.get_header().get_number_of_columns())
for i in range(0, rows):
for j in range(0, cols):
self.assertAlmostEqual(saved(i, j), polar(i, j), delta=0.001,
msg="Generated polar image is different from stored"
" row %d col %d" % (i, j))
def test_substract(self):
"""Test subtracting images"""
srw = em2d.SpiderImageReaderWriter()
rows = int(10)
cols = int(5)
img1 = em2d.Image(rows, cols)
img2 = em2d.Image(rows, cols)
result = em2d.Image(rows, cols)
for i in range(0, rows):
for j in range(0, cols):
img1.set_value(i, j, random.uniform(-1, 1))
img2.set_value(i, j, img1(i, j))
em2d.do_subtract_images(img1, img2, result)
for i in range(0, rows):
for j in range(0, cols):
self.assertAlmostEqual(abs(result(i, j)), 0, delta=0.001,
msg="Subtract images error")
def test_write_tiff(self):
"""Test of TIFFReaderWriter writing"""
trw = em2d.TIFFImageReaderWriter()
srw = em2d.SpiderImageReaderWriter()
fn_img1 = self.get_input_file_name("lena-256x256.tif")
img1 = em2d.Image(fn_img1, trw)
fn_img2 = "temp.tif"
img1.write(fn_img2, trw)
img2 = em2d.Image(fn_img2, trw)
# Use the ccc for testing instead of the pixel values. The matrix
# in img2 is transformed from floats to ints son it can be written.
# Values can change, but the ccc has to be very close to 1.
ccc = em2d.get_cross_correlation_coefficient(img1.get_data(),
img2.get_data())
print ccc
self.assertAlmostEqual(ccc, 1, delta=0.01,
msg="Written TIFF image is not equal to read ")
os.remove(fn_img2)
def test_noisy_projections(self):
""" Test the generation of noisy projections"""
smodel = IMP.kernel.Model()
ssel = IMP.atom.ATOMPDBSelector()
fn_model = self.get_input_file_name("1e6v.pdb")
prot = IMP.atom.read_pdb(fn_model, smodel, ssel)
particles = IMP.core.get_leaves(prot)
n_projections = 16
rows = 100
cols = 100
resolution = 1
apix = 1.5
noise_SNR = 0.5
# read the stored noisy images
stored_names = []
srw = em2d.SpiderImageReaderWriter()
for i in range(0, n_projections):
fn_subject = "1e6v-subject-%d-set-%d-%s-apix" \
"-%s-SNR.spi" % (i, n_projections, str(apix), str(noise_SNR))
stored_names.append(self.get_input_file_name(fn_subject))
stored_images = em2d.read_images(stored_names, srw)
# Read registration parameters and generate new images
fn_regs = self.get_input_file_name('1e6v-subjects-0.5.params')
Regs = em2d.read_registration_results(fn_regs)
options = em2d.ProjectingOptions(apix, resolution)
projections = em2d.get_projections(particles, Regs, rows, cols,
options)
# Add noise
for i in range(0, n_projections):
em2d.do_normalize(projections[i], True)
em2d.add_noise(projections[i], 0.0, 1. / (noise_SNR**0.5),
"gaussian", 3)
# theoretical ccc for same images at a level of noise
theoretical_ccc = noise_SNR / (noise_SNR + 1)
for n in range(0, n_projections):
ccc = em2d.get_cross_correlation_coefficient(
projections[n], stored_images[n])
# allow 3% difference in cross-correlation
self.assertAlmostEqual(
theoretical_ccc,
ccc,
delta=0.03,
msg="Noisy projections generated and stored are different")
def test_variance_filter(self):
"""Test that the variance filter is working"""
srw = em2d.SpiderImageReaderWriter()
img = em2d.Image()
img.read(self.get_input_file_name("1z5s-projection-2.spi"), srw)
filtered = em2d.Image()
kernelsize = 7
em2d.apply_variance_filter(img, filtered, kernelsize)
saved = em2d.Image()
saved.read(self.get_input_file_name("filtered_image.spi"), srw)
rows = int(img.get_header().get_number_of_rows())
cols = int(img.get_header().get_number_of_columns())
for i in range(0, rows):
for j in range(0, cols):
self.assertAlmostEqual(
saved(i, j),
filtered(i, j),
delta=0.001,
msg="Generated image is different from stored")
def test_building_an_optimization_problem_with_em2d_restraint(self):
"""Test that an a optimization with em2d restraint is properly built"""
m = IMP.kernel.Model()
prot = IMP.atom.read_pdb(self.get_input_file_name("1z5s.pdb"),
m, IMP.atom.ATOMPDBSelector())
# get the chains
chains = IMP.atom.get_by_type(prot, IMP.atom.CHAIN_TYPE)
# set the chains as rigid bodies
rigid_bodies = []
native_chain_centers = []
for c in chains:
atoms = IMP.core.get_leaves(c)
rbd = IMP.core.RigidBody.setup_particle(c, atoms)
rbd.set_coordinates_are_optimized(True)
rigid_bodies.append(rbd)
native_chain_centers.append(rbd.get_coordinates())
self.assertEqual(
len(rigid_bodies),
4,
"Problem generating rigid bodies")
bb = IMP.algebra.BoundingBox3D(IMP.algebra.Vector3D(-25, -40, -60),
IMP.algebra.Vector3D(25, 40, 60))
# set distance restraints
d01 = IMP.algebra.get_distance(native_chain_centers[0],
native_chain_centers[1])
r01 = IMP.core.DistanceRestraint(IMP.core.Harmonic(d01, 1),
chains[0],
chains[1])
r01.set_name("distance 0-1")
d12 = IMP.algebra.get_distance(native_chain_centers[1],
native_chain_centers[2])
r12 = IMP.core.DistanceRestraint(IMP.core.Harmonic(d12, 1),
chains[1],
chains[2])
r12.set_name("distance 1-2")
d23 = IMP.algebra.get_distance(native_chain_centers[2],
native_chain_centers[3])
r23 = IMP.core.DistanceRestraint(IMP.core.Harmonic(d23, 1),
chains[2],
chains[3])
r23.set_name("distance 2-3")
d30 = IMP.algebra.get_distance(native_chain_centers[3],
native_chain_centers[0])
r30 = IMP.core.DistanceRestraint(IMP.core.Harmonic(d30, 1),
chains[3],
chains[0])
r30.set_name("distance 3-0")
# set distance restraints
for r in [r01, r12, r23, r30]:
m.add_restraint(r)
self.assertEqual(m.get_number_of_restraints(), 4,
"Incorrect number of distance restraints")
# set em2D restraint
srw = em2d.SpiderImageReaderWriter()
selection_file = self.get_input_file_name("all-1z5s-projections.sel")
images_to_read_names = [IMP.base.get_relative_path(selection_file, x)
for x in em2d.read_selection_file(selection_file)]
em_images = em2d.read_images(images_to_read_names, srw)
self.assertEqual(len(em_images), 3, "Incorrect number images read")
apix = 1.5
resolution = 1
n_projections = 20
params = em2d.Em2DRestraintParameters(apix, resolution, n_projections)
params.save_match_images = False
params.coarse_registration_method = em2d.ALIGN2D_PREPROCESSING
score_function = em2d.EM2DScore()
em2d_restraint = em2d.Em2DRestraint(m)
em2d_restraint.setup(score_function, params)
em2d_restraint.set_images(em_images)
em2d_restraint.set_name("em2d restraint")
container = IMP.container.ListSingletonContainer(
IMP.core.get_leaves(prot))
em2d_restraint.set_particles(container)
em2d_restraints_set = IMP.kernel.RestraintSet(m)
em2d_restraints_set.add_restraint(em2d_restraint)
em2d_restraints_set.set_weight(1000) # weight for the em2D restraint
m.add_restraint(em2d_restraints_set)
self.assertEqual(m.get_number_of_restraints(), 5,
"Incorrect number of restraints")
# MONTECARLO OPTIMIZATION
s = IMP.core.MonteCarlo(m)
# Add movers for the rigid bodies
movers = []
for rbd in rigid_bodies:
movers.append(IMP.core.RigidBodyMover(rbd, 5, 2))
s.add_movers(movers)
self.assertEqual(s.get_number_of_movers(), 4,
"Incorrect number of MonteCarlo movers")
# Optimizer state to save intermediate configurations
o_state = IMP.atom.WritePDBOptimizerState(chains,
"intermediate-step-%1%.pdb")
o_state.set_period(11)
s.add_optimizer_state(o_state)
ostate2 = WriteStatisticsOptimizerScore(m)
s.add_optimizer_state(ostate2)
# Perform optimization
temperatures = [200, 100, 60, 40, 20, 5]
optimization_steps = 200
# for temp in temperatures:
# s.optimize(optimization_steps)
# IMP.atom.write_pdb(prot,"solution.pdb")
self.assertTrue(True)
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)
## \example em2d/em_images_conversion.py
# Conversion of Electron Microscopy Images.
#
import IMP
import IMP.em2d as em2d
import os
"""
Conversion of Electron Microscopy Images.
"""
# Read images
fn_selection = em2d.get_example_path("all-1z5s-projections.sel")
srw = em2d.SpiderImageReaderWriter()
trw = em2d.TIFFImageReaderWriter()
fn_images = em2d.read_selection_file(fn_selection)
fn_images = [em2d.get_example_path(x) for x in fn_images]
images = em2d.read_images(fn_images, srw)
# write
fn_saved = em2d.create_filenames(3, "1z5s-projection", "tif")
em2d.save_images(images, fn_saved, trw)
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))
