def filterlocal(ui, vi, m, falloff, myid, main_node, number_of_proc):
if myid == main_node:
nx = vi.get_xsize()
ny = vi.get_ysize()
nz = vi.get_zsize()
# Round all resolution numbers to two digits
for x in range(nx):
for y in range(ny):
for z in range(nz):
ui.set_value_at_fast(x, y, z, round(ui.get_value_at(x, y, z), 2))
dis = [nx, ny, nz]
else:
falloff = 0.0
radius = 0
dis = [0, 0, 0]
falloff = sp_utilities.bcast_number_to_all(falloff, main_node)
dis = sp_utilities.bcast_list_to_all(dis, myid, source_node=main_node)
if myid != main_node:
nx = int(dis[0])
ny = int(dis[1])
nz = int(dis[2])
vi = sp_utilities.model_blank(nx, ny, nz)
ui = sp_utilities.model_blank(nx, ny, nz)
sp_utilities.bcast_EMData_to_all(vi, myid, main_node)
sp_utilities.bcast_EMData_to_all(ui, myid, main_node)
sp_fundamentals.fftip(vi) # volume to be filtered
st = EMAN2_cppwrap.Util.infomask(ui, m, True)
filteredvol = sp_utilities.model_blank(nx, ny, nz)
cutoff = max(st[2] - 0.01, 0.0)
while cutoff < st[3]:
cutoff = round(cutoff + 0.01, 2)
# if(myid == main_node): print cutoff,st
pt = EMAN2_cppwrap.Util.infomask(
sp_morphology.threshold_outside(ui, cutoff - 0.00501, cutoff + 0.005),
m,
True,
) # Ideally, one would want to check only slices in question...
if pt[0] != 0.0:
# print cutoff,pt[0]
vovo = sp_fundamentals.fft(filt_tanl(vi, cutoff, falloff))
for z in range(myid, nz, number_of_proc):
for x in range(nx):
for y in range(ny):
if m.get_value_at(x, y, z) > 0.5:
if round(ui.get_value_at(x, y, z), 2) == cutoff:
filteredvol.set_value_at_fast(
x, y, z, vovo.get_value_at(x, y, z)
)
mpi.mpi_barrier(mpi.MPI_COMM_WORLD)
sp_utilities.reduce_EMData_to_root(filteredvol, myid, main_node, mpi.MPI_COMM_WORLD)
return filteredvol
def helicalshiftali_MPI(stack,
maskfile=None,
maxit=100,
CTF=False,
snr=1.0,
Fourvar=False,
search_rng=-1):
nproc = mpi.mpi_comm_size(mpi.MPI_COMM_WORLD)
myid = mpi.mpi_comm_rank(mpi.MPI_COMM_WORLD)
main_node = 0
ftp = file_type(stack)
if myid == main_node:
print_begin_msg("helical-shiftali_MPI")
max_iter = int(maxit)
if (myid == main_node):
infils = EMUtil.get_all_attributes(stack, "filament")
ptlcoords = EMUtil.get_all_attributes(stack, 'ptcl_source_coord')
filaments = ordersegments(infils, ptlcoords)
total_nfils = len(filaments)
inidl = [0] * total_nfils
for i in range(total_nfils):
inidl[i] = len(filaments[i])
linidl = sum(inidl)
nima = linidl
tfilaments = []
for i in range(total_nfils):
tfilaments += filaments[i]
del filaments
else:
total_nfils = 0
linidl = 0
total_nfils = bcast_number_to_all(total_nfils, source_node=main_node)
if myid != main_node:
inidl = [-1] * total_nfils
inidl = bcast_list_to_all(inidl, myid, source_node=main_node)
linidl = bcast_number_to_all(linidl, source_node=main_node)
if myid != main_node:
tfilaments = [-1] * linidl
tfilaments = bcast_list_to_all(tfilaments, myid, source_node=main_node)
filaments = []
iendi = 0
for i in range(total_nfils):
isti = iendi
iendi = isti + inidl[i]
filaments.append(tfilaments[isti:iendi])
del tfilaments, inidl
if myid == main_node:
print_msg("total number of filaments: %d" % total_nfils)
if total_nfils < nproc:
ERROR(
'number of CPUs (%i) is larger than the number of filaments (%i), please reduce the number of CPUs used'
% (nproc, total_nfils),
myid=myid)
# balanced load
temp = chunks_distribution([[len(filaments[i]), i]
for i in range(len(filaments))],
nproc)[myid:myid + 1][0]
filaments = [filaments[temp[i][1]] for i in range(len(temp))]
nfils = len(filaments)
#filaments = [[0,1]]
#print "filaments",filaments
list_of_particles = []
indcs = []
k = 0
for i in range(nfils):
list_of_particles += filaments[i]
k1 = k + len(filaments[i])
indcs.append([k, k1])
k = k1
data = EMData.read_images(stack, list_of_particles)
ldata = len(data)
sxprint("ldata=", ldata)
nx = data[0].get_xsize()
ny = data[0].get_ysize()
if maskfile == None:
mrad = min(nx, ny) // 2 - 2
mask = pad(model_blank(2 * mrad + 1, ny, 1, 1.0), nx, ny, 1, 0.0)
else:
mask = get_im(maskfile)
# apply initial xform.align2d parameters stored in header
init_params = []
for im in range(ldata):
t = data[im].get_attr('xform.align2d')
init_params.append(t)
p = t.get_params("2d")
data[im] = rot_shift2D(data[im], p['alpha'], p['tx'], p['ty'],
p['mirror'], p['scale'])
if CTF:
from sp_filter import filt_ctf
from sp_morphology import ctf_img
ctf_abs_sum = EMData(nx, ny, 1, False)
ctf_2_sum = EMData(nx, ny, 1, False)
else:
ctf_2_sum = None
ctf_abs_sum = None
from sp_utilities import info
for im in range(ldata):
data[im].set_attr('ID', list_of_particles[im])
st = Util.infomask(data[im], mask, False)
data[im] -= st[0]
if CTF:
ctf_params = data[im].get_attr("ctf")
qctf = data[im].get_attr("ctf_applied")
if qctf == 0:
data[im] = filt_ctf(fft(data[im]), ctf_params)
data[im].set_attr('ctf_applied', 1)
elif qctf != 1:
ERROR('Incorrectly set qctf flag', myid=myid)
ctfimg = ctf_img(nx, ctf_params, ny=ny)
Util.add_img2(ctf_2_sum, ctfimg)
Util.add_img_abs(ctf_abs_sum, ctfimg)
else:
data[im] = fft(data[im])
del list_of_particles
if CTF:
reduce_EMData_to_root(ctf_2_sum, myid, main_node)
reduce_EMData_to_root(ctf_abs_sum, myid, main_node)
if CTF:
if myid != main_node:
del ctf_2_sum
del ctf_abs_sum
else:
temp = EMData(nx, ny, 1, False)
tsnr = 1. / snr
for i in range(0, nx + 2, 2):
for j in range(ny):
temp.set_value_at(i, j, tsnr)
temp.set_value_at(i + 1, j, 0.0)
#info(ctf_2_sum)
Util.add_img(ctf_2_sum, temp)
#info(ctf_2_sum)
del temp
total_iter = 0
shift_x = [0.0] * ldata
for Iter in range(max_iter):
if myid == main_node:
start_time = time()
print_msg("Iteration #%4d\n" % (total_iter))
total_iter += 1
avg = EMData(nx, ny, 1, False)
for im in range(ldata):
Util.add_img(avg, fshift(data[im], shift_x[im]))
reduce_EMData_to_root(avg, myid, main_node)
if myid == main_node:
if CTF: tavg = Util.divn_filter(avg, ctf_2_sum)
else: tavg = Util.mult_scalar(avg, 1.0 / float(nima))
else:
tavg = model_blank(nx, ny)
if Fourvar:
bcast_EMData_to_all(tavg, myid, main_node)
vav, rvar = varf2d_MPI(myid, data, tavg, mask, "a", CTF)
if myid == main_node:
if Fourvar:
tavg = fft(Util.divn_img(fft(tavg), vav))
vav_r = Util.pack_complex_to_real(vav)
# normalize and mask tavg in real space
tavg = fft(tavg)
stat = Util.infomask(tavg, mask, False)
tavg -= stat[0]
Util.mul_img(tavg, mask)
tavg.write_image("tavg.hdf", Iter)
# For testing purposes: shift tavg to some random place and see if the centering is still correct
#tavg = rot_shift3D(tavg,sx=3,sy=-4)
if Fourvar: del vav
bcast_EMData_to_all(tavg, myid, main_node)
tavg = fft(tavg)
sx_sum = 0.0
nxc = nx // 2
for ifil in range(nfils):
"""
# Calculate filament average
avg = EMData(nx, ny, 1, False)
filnima = 0
for im in xrange(indcs[ifil][0], indcs[ifil][1]):
Util.add_img(avg, data[im])
filnima += 1
tavg = Util.mult_scalar(avg, 1.0/float(filnima))
"""
# Calculate 1D ccf between each segment and filament average
nsegms = indcs[ifil][1] - indcs[ifil][0]
ctx = [None] * nsegms
pcoords = [None] * nsegms
for im in range(indcs[ifil][0], indcs[ifil][1]):
ctx[im - indcs[ifil][0]] = Util.window(ccf(tavg, data[im]), nx,
1)
pcoords[im - indcs[ifil][0]] = data[im].get_attr(
'ptcl_source_coord')
#ctx[im-indcs[ifil][0]].write_image("ctx.hdf",im-indcs[ifil][0])
#print " CTX ",myid,im,Util.infomask(ctx[im-indcs[ifil][0]], None, True)
# search for best x-shift
cents = nsegms // 2
dst = sqrt(
max((pcoords[cents][0] - pcoords[0][0])**2 +
(pcoords[cents][1] - pcoords[0][1])**2,
(pcoords[cents][0] - pcoords[-1][0])**2 +
(pcoords[cents][1] - pcoords[-1][1])**2))
maxincline = atan2(ny // 2 - 2 - float(search_rng), dst)
kang = int(dst * tan(maxincline) + 0.5)
#print " settings ",nsegms,cents,dst,search_rng,maxincline,kang
# ## C code for alignment. @ming
results = [0.0] * 3
results = Util.helixshiftali(ctx, pcoords, nsegms, maxincline,
kang, search_rng, nxc)
sib = int(results[0])
bang = results[1]
qm = results[2]
#print qm, sib, bang
# qm = -1.e23
#
# for six in xrange(-search_rng, search_rng+1,1):
# q0 = ctx[cents].get_value_at(six+nxc)
# for incline in xrange(kang+1):
# qt = q0
# qu = q0
# if(kang>0): tang = tan(maxincline/kang*incline)
# else: tang = 0.0
# for kim in xrange(cents+1,nsegms):
# dst = sqrt((pcoords[cents][0] - pcoords[kim][0])**2 + (pcoords[cents][1] - pcoords[kim][1])**2)
# xl = dst*tang+six+nxc
# ixl = int(xl)
# dxl = xl - ixl
# #print " A ", ifil,six,incline,kim,xl,ixl,dxl
# qt += (1.0-dxl)*ctx[kim].get_value_at(ixl) + dxl*ctx[kim].get_value_at(ixl+1)
# xl = -dst*tang+six+nxc
# ixl = int(xl)
# dxl = xl - ixl
# qu += (1.0-dxl)*ctx[kim].get_value_at(ixl) + dxl*ctx[kim].get_value_at(ixl+1)
# for kim in xrange(cents):
# dst = sqrt((pcoords[cents][0] - pcoords[kim][0])**2 + (pcoords[cents][1] - pcoords[kim][1])**2)
# xl = -dst*tang+six+nxc
# ixl = int(xl)
# dxl = xl - ixl
# qt += (1.0-dxl)*ctx[kim].get_value_at(ixl) + dxl*ctx[kim].get_value_at(ixl+1)
# xl = dst*tang+six+nxc
# ixl = int(xl)
# dxl = xl - ixl
# qu += (1.0-dxl)*ctx[kim].get_value_at(ixl) + dxl*ctx[kim].get_value_at(ixl+1)
# if( qt > qm ):
# qm = qt
# sib = six
# bang = tang
# if( qu > qm ):
# qm = qu
# sib = six
# bang = -tang
#if incline == 0: print "incline = 0 ",six,tang,qt,qu
#print qm,six,sib,bang
#print " got results ",indcs[ifil][0], indcs[ifil][1], ifil,myid,qm,sib,tang,bang,len(ctx),Util.infomask(ctx[0], None, True)
for im in range(indcs[ifil][0], indcs[ifil][1]):
kim = im - indcs[ifil][0]
dst = sqrt((pcoords[cents][0] - pcoords[kim][0])**2 +
(pcoords[cents][1] - pcoords[kim][1])**2)
if (kim < cents): xl = -dst * bang + sib
else: xl = dst * bang + sib
shift_x[im] = xl
# Average shift
sx_sum += shift_x[indcs[ifil][0] + cents]
# #print myid,sx_sum,total_nfils
sx_sum = mpi.mpi_reduce(sx_sum, 1, mpi.MPI_FLOAT, mpi.MPI_SUM,
main_node, mpi.MPI_COMM_WORLD)
if myid == main_node:
sx_sum = float(sx_sum[0]) / total_nfils
print_msg("Average shift %6.2f\n" % (sx_sum))
else:
sx_sum = 0.0
sx_sum = 0.0
sx_sum = bcast_number_to_all(sx_sum, source_node=main_node)
for im in range(ldata):
shift_x[im] -= sx_sum
#print " %3d %6.3f"%(im,shift_x[im])
#exit()
# combine shifts found with the original parameters
for im in range(ldata):
t1 = Transform()
##import random
##shix=random.randint(-10, 10)
##t1.set_params({"type":"2D","tx":shix})
t1.set_params({"type": "2D", "tx": shift_x[im]})
# combine t0 and t1
tt = t1 * init_params[im]
data[im].set_attr("xform.align2d", tt)
# write out headers and STOP, under MPI writing has to be done sequentially
mpi.mpi_barrier(mpi.MPI_COMM_WORLD)
par_str = ["xform.align2d", "ID"]
if myid == main_node:
from sp_utilities import file_type
if (file_type(stack) == "bdb"):
from sp_utilities import recv_attr_dict_bdb
recv_attr_dict_bdb(main_node, stack, data, par_str, 0, ldata,
nproc)
else:
from sp_utilities import recv_attr_dict
recv_attr_dict(main_node, stack, data, par_str, 0, ldata, nproc)
else:
send_attr_dict(main_node, data, par_str, 0, ldata)
if myid == main_node: print_end_msg("helical-shiftali_MPI")
def shiftali_MPI(stack,
maskfile=None,
maxit=100,
CTF=False,
snr=1.0,
Fourvar=False,
search_rng=-1,
oneDx=False,
search_rng_y=-1):
number_of_proc = mpi.mpi_comm_size(mpi.MPI_COMM_WORLD)
myid = mpi.mpi_comm_rank(mpi.MPI_COMM_WORLD)
main_node = 0
ftp = file_type(stack)
if myid == main_node:
print_begin_msg("shiftali_MPI")
max_iter = int(maxit)
if myid == main_node:
if ftp == "bdb":
from EMAN2db import db_open_dict
dummy = db_open_dict(stack, True)
nima = EMUtil.get_image_count(stack)
else:
nima = 0
nima = bcast_number_to_all(nima, source_node=main_node)
list_of_particles = list(range(nima))
image_start, image_end = MPI_start_end(nima, number_of_proc, myid)
list_of_particles = list_of_particles[image_start:image_end]
# read nx and ctf_app (if CTF) and broadcast to all nodes
if myid == main_node:
ima = EMData()
ima.read_image(stack, list_of_particles[0], True)
nx = ima.get_xsize()
ny = ima.get_ysize()
if CTF: ctf_app = ima.get_attr_default('ctf_applied', 2)
del ima
else:
nx = 0
ny = 0
if CTF: ctf_app = 0
nx = bcast_number_to_all(nx, source_node=main_node)
ny = bcast_number_to_all(ny, source_node=main_node)
if CTF:
ctf_app = bcast_number_to_all(ctf_app, source_node=main_node)
if ctf_app > 0:
ERROR("data cannot be ctf-applied", myid=myid)
if maskfile == None:
mrad = min(nx, ny)
mask = model_circle(mrad // 2 - 2, nx, ny)
else:
mask = get_im(maskfile)
if CTF:
from sp_filter import filt_ctf
from sp_morphology import ctf_img
ctf_abs_sum = EMData(nx, ny, 1, False)
ctf_2_sum = EMData(nx, ny, 1, False)
else:
ctf_2_sum = None
from sp_global_def import CACHE_DISABLE
if CACHE_DISABLE:
data = EMData.read_images(stack, list_of_particles)
else:
for i in range(number_of_proc):
if myid == i:
data = EMData.read_images(stack, list_of_particles)
if ftp == "bdb": mpi.mpi_barrier(mpi.MPI_COMM_WORLD)
for im in range(len(data)):
data[im].set_attr('ID', list_of_particles[im])
st = Util.infomask(data[im], mask, False)
data[im] -= st[0]
if CTF:
ctf_params = data[im].get_attr("ctf")
ctfimg = ctf_img(nx, ctf_params, ny=ny)
Util.add_img2(ctf_2_sum, ctfimg)
Util.add_img_abs(ctf_abs_sum, ctfimg)
if CTF:
reduce_EMData_to_root(ctf_2_sum, myid, main_node)
reduce_EMData_to_root(ctf_abs_sum, myid, main_node)
else:
ctf_2_sum = None
if CTF:
if myid != main_node:
del ctf_2_sum
del ctf_abs_sum
else:
temp = EMData(nx, ny, 1, False)
for i in range(0, nx, 2):
for j in range(ny):
temp.set_value_at(i, j, snr)
Util.add_img(ctf_2_sum, temp)
del temp
total_iter = 0
# apply initial xform.align2d parameters stored in header
init_params = []
for im in range(len(data)):
t = data[im].get_attr('xform.align2d')
init_params.append(t)
p = t.get_params("2d")
data[im] = rot_shift2D(data[im],
p['alpha'],
sx=p['tx'],
sy=p['ty'],
mirror=p['mirror'],
scale=p['scale'])
# fourier transform all images, and apply ctf if CTF
for im in range(len(data)):
if CTF:
ctf_params = data[im].get_attr("ctf")
data[im] = filt_ctf(fft(data[im]), ctf_params)
else:
data[im] = fft(data[im])
sx_sum = 0
sy_sum = 0
sx_sum_total = 0
sy_sum_total = 0
shift_x = [0.0] * len(data)
shift_y = [0.0] * len(data)
ishift_x = [0.0] * len(data)
ishift_y = [0.0] * len(data)
for Iter in range(max_iter):
if myid == main_node:
start_time = time()
print_msg("Iteration #%4d\n" % (total_iter))
total_iter += 1
avg = EMData(nx, ny, 1, False)
for im in data:
Util.add_img(avg, im)
reduce_EMData_to_root(avg, myid, main_node)
if myid == main_node:
if CTF:
tavg = Util.divn_filter(avg, ctf_2_sum)
else:
tavg = Util.mult_scalar(avg, 1.0 / float(nima))
else:
tavg = EMData(nx, ny, 1, False)
if Fourvar:
bcast_EMData_to_all(tavg, myid, main_node)
vav, rvar = varf2d_MPI(myid, data, tavg, mask, "a", CTF)
if myid == main_node:
if Fourvar:
tavg = fft(Util.divn_img(fft(tavg), vav))
vav_r = Util.pack_complex_to_real(vav)
# normalize and mask tavg in real space
tavg = fft(tavg)
stat = Util.infomask(tavg, mask, False)
tavg -= stat[0]
Util.mul_img(tavg, mask)
# For testing purposes: shift tavg to some random place and see if the centering is still correct
#tavg = rot_shift3D(tavg,sx=3,sy=-4)
tavg = fft(tavg)
if Fourvar: del vav
bcast_EMData_to_all(tavg, myid, main_node)
sx_sum = 0
sy_sum = 0
if search_rng > 0: nwx = 2 * search_rng + 1
else: nwx = nx
if search_rng_y > 0: nwy = 2 * search_rng_y + 1
else: nwy = ny
not_zero = 0
for im in range(len(data)):
if oneDx:
ctx = Util.window(ccf(data[im], tavg), nwx, 1)
p1 = peak_search(ctx)
p1_x = -int(p1[0][3])
ishift_x[im] = p1_x
sx_sum += p1_x
else:
p1 = peak_search(Util.window(ccf(data[im], tavg), nwx, nwy))
p1_x = -int(p1[0][4])
p1_y = -int(p1[0][5])
ishift_x[im] = p1_x
ishift_y[im] = p1_y
sx_sum += p1_x
sy_sum += p1_y
if not_zero == 0:
if (not (ishift_x[im] == 0.0)) or (not (ishift_y[im] == 0.0)):
not_zero = 1
sx_sum = mpi.mpi_reduce(sx_sum, 1, mpi.MPI_INT, mpi.MPI_SUM, main_node,
mpi.MPI_COMM_WORLD)
if not oneDx:
sy_sum = mpi.mpi_reduce(sy_sum, 1, mpi.MPI_INT, mpi.MPI_SUM,
main_node, mpi.MPI_COMM_WORLD)
if myid == main_node:
sx_sum_total = int(sx_sum[0])
if not oneDx:
sy_sum_total = int(sy_sum[0])
else:
sx_sum_total = 0
sy_sum_total = 0
sx_sum_total = bcast_number_to_all(sx_sum_total, source_node=main_node)
if not oneDx:
sy_sum_total = bcast_number_to_all(sy_sum_total,
source_node=main_node)
sx_ave = round(float(sx_sum_total) / nima)
sy_ave = round(float(sy_sum_total) / nima)
for im in range(len(data)):
p1_x = ishift_x[im] - sx_ave
p1_y = ishift_y[im] - sy_ave
params2 = {
"filter_type": Processor.fourier_filter_types.SHIFT,
"x_shift": p1_x,
"y_shift": p1_y,
"z_shift": 0.0
}
data[im] = Processor.EMFourierFilter(data[im], params2)
shift_x[im] += p1_x
shift_y[im] += p1_y
# stop if all shifts are zero
not_zero = mpi.mpi_reduce(not_zero, 1, mpi.MPI_INT, mpi.MPI_SUM,
main_node, mpi.MPI_COMM_WORLD)
if myid == main_node:
not_zero_all = int(not_zero[0])
else:
not_zero_all = 0
not_zero_all = bcast_number_to_all(not_zero_all, source_node=main_node)
if myid == main_node:
print_msg("Time of iteration = %12.2f\n" % (time() - start_time))
start_time = time()
if not_zero_all == 0: break
#for im in xrange(len(data)): data[im] = fft(data[im]) This should not be required as only header information is used
# combine shifts found with the original parameters
for im in range(len(data)):
t0 = init_params[im]
t1 = Transform()
t1.set_params({
"type": "2D",
"alpha": 0,
"scale": t0.get_scale(),
"mirror": 0,
"tx": shift_x[im],
"ty": shift_y[im]
})
# combine t0 and t1
tt = t1 * t0
data[im].set_attr("xform.align2d", tt)
# write out headers and STOP, under MPI writing has to be done sequentially
mpi.mpi_barrier(mpi.MPI_COMM_WORLD)
par_str = ["xform.align2d", "ID"]
if myid == main_node:
from sp_utilities import file_type
if (file_type(stack) == "bdb"):
from sp_utilities import recv_attr_dict_bdb
recv_attr_dict_bdb(main_node, stack, data, par_str, image_start,
image_end, number_of_proc)
else:
from sp_utilities import recv_attr_dict
recv_attr_dict(main_node, stack, data, par_str, image_start,
image_end, number_of_proc)
else:
send_attr_dict(main_node, data, par_str, image_start, image_end)
if myid == main_node: print_end_msg("shiftali_MPI")
def filterlocal(ui, vi, m, falloff, myid, main_node, number_of_proc): from mpi import mpi_init, mpi_comm_size, mpi_comm_rank, MPI_COMM_WORLD from mpi import mpi_reduce, mpi_bcast, mpi_barrier, mpi_gatherv, mpi_send, mpi_recv from mpi import MPI_SUM, MPI_FLOAT, MPI_INT from sp_utilities import bcast_number_to_all, bcast_list_to_all, model_blank, bcast_EMData_to_all, reduce_EMData_to_root from sp_morphology import threshold_outside from sp_filter import filt_tanl from sp_fundamentals import fft, fftip if(myid == main_node): nx = vi.get_xsize() ny = vi.get_ysize() nz = vi.get_zsize() # Round all resolution numbers to two digits for x in range(nx): for y in range(ny): for z in range(nz): ui.set_value_at_fast( x,y,z, round(ui.get_value_at(x,y,z), 2) ) dis = [nx,ny,nz] else: falloff = 0.0 radius = 0 dis = [0,0,0] falloff = bcast_number_to_all(falloff, main_node) dis = bcast_list_to_all(dis, myid, source_node = main_node) if(myid != main_node): nx = int(dis[0]) ny = int(dis[1]) nz = int(dis[2]) vi = model_blank(nx,ny,nz) ui = model_blank(nx,ny,nz) bcast_EMData_to_all(vi, myid, main_node) bcast_EMData_to_all(ui, myid, main_node) fftip(vi) # volume to be filtered st = Util.infomask(ui, m, True) filteredvol = model_blank(nx,ny,nz) cutoff = max(st[2] - 0.01,0.0) while(cutoff < st[3] ): cutoff = round(cutoff + 0.01, 2) #if(myid == main_node): print cutoff,st pt = Util.infomask( threshold_outside(ui, cutoff - 0.00501, cutoff + 0.005), m, True) # Ideally, one would want to check only slices in question... if(pt[0] != 0.0): #print cutoff,pt[0] vovo = fft( filt_tanl(vi, cutoff, falloff) ) for z in range(myid, nz, number_of_proc): for x in range(nx): for y in range(ny): if(m.get_value_at(x,y,z) > 0.5): if(round(ui.get_value_at(x,y,z),2) == cutoff): filteredvol.set_value_at_fast(x,y,z,vovo.get_value_at(x,y,z)) mpi_barrier(MPI_COMM_WORLD) reduce_EMData_to_root(filteredvol, myid, main_node, MPI_COMM_WORLD) return filteredvol
def recons3d_4nn_ctf_MPI(
myid,
prjlist,
snr=1.0,
sign=1,
symmetry="c1",
finfo=None,
npad=2,
xysize=-1,
zsize=-1,
mpi_comm=None,
smearstep=0.0,
):
"""
recons3d_4nn_ctf - calculate CTF-corrected 3-D reconstruction from a set of projections using three Eulerian angles, two shifts, and CTF settings for each projeciton image
Input
stack: name of the stack file containing projection data, projections have to be squares
list_proj: list of projections to be included in the reconstruction or image iterator
snr: Signal-to-Noise Ratio of the data
sign: sign of the CTF
symmetry: point-group symmetry to be enforced, each projection will enter the reconstruction in all symmetry-related directions.
"""
if mpi_comm == None:
mpi_comm = mpi.MPI_COMM_WORLD
if type(prjlist) == list:
prjlist = sp_utilities.iterImagesList(prjlist)
if not prjlist.goToNext():
sp_global_def.ERROR("empty input list", "recons3d_4nn_ctf_MPI", 1)
imgsize = prjlist.image().get_xsize()
if prjlist.image().get_ysize() != imgsize:
imgsize = max(imgsize, prjlist.image().get_ysize())
dopad = True
else:
dopad = False
prjlist.goToPrev()
fftvol = EMAN2_cppwrap.EMData()
if smearstep > 0.0:
# if myid == 0: print " Setting smear in prepare_recons_ctf"
ns = 1
smear = []
for j in range(-ns, ns + 1):
if j != 0:
for i in range(-ns, ns + 1):
for k in range(-ns, ns + 1):
smear += [
i * smearstep, j * smearstep, k * smearstep, 1.0
]
# Deal with theta = 0.0 cases
prj = []
for i in range(-ns, ns + 1):
for k in range(-ns, ns + 1):
prj.append(i + k)
for i in range(-2 * ns, 2 * ns + 1, 1):
smear += [i * smearstep, 0.0, 0.0, float(prj.count(i))]
# if myid == 0: print " Smear ",smear
fftvol.set_attr("smear", smear)
weight = EMAN2_cppwrap.EMData()
if xysize == -1 and zsize == -1:
params = {
"size": imgsize,
"npad": npad,
"snr": snr,
"sign": sign,
"symmetry": symmetry,
"fftvol": fftvol,
"weight": weight,
}
r = EMAN2_cppwrap.Reconstructors.get("nn4_ctf", params)
else:
if xysize != -1 and zsize != -1:
rx = old_div(float(xysize), imgsize)
ry = old_div(float(xysize), imgsize)
rz = old_div(float(zsize), imgsize)
elif xysize != -1:
rx = old_div(float(xysize), imgsize)
ry = old_div(float(xysize), imgsize)
rz = 1.0
else:
rx = 1.0
ry = 1.0
rz = old_div(float(zsize), imgsize)
# There is an error here with sizeprojection PAP 10/22/2014
params = {
"size": sizeprojection,
"npad": npad,
"snr": snr,
"sign": sign,
"symmetry": symmetry,
"fftvol": fftvol,
"weight": weight,
"xratio": rx,
"yratio": ry,
"zratio": rz,
}
r = EMAN2_cppwrap.Reconstructors.get("nn4_ctf_rect", params)
r.setup()
# if not (finfo is None):
nimg = 0
while prjlist.goToNext():
prj = prjlist.image()
if dopad:
prj = sp_utilities.pad(prj, imgsize, imgsize, 1, "circumference")
# if params:
insert_slices(r, prj)
if not (finfo is None):
nimg += 1
finfo.write(" %4d inserted\n" % (nimg))
finfo.flush()
del sp_utilities.pad
if not (finfo is None):
finfo.write("begin reduce\n")
finfo.flush()
sp_utilities.reduce_EMData_to_root(fftvol, myid, comm=mpi_comm)
sp_utilities.reduce_EMData_to_root(weight, myid, comm=mpi_comm)
if not (finfo is None):
finfo.write("after reduce\n")
finfo.flush()
if myid == 0:
dummy = r.finish(True)
else:
if xysize == -1 and zsize == -1:
fftvol = sp_utilities.model_blank(imgsize, imgsize, imgsize)
else:
if zsize == -1:
fftvol = sp_utilities.model_blank(xysize, xysize, imgsize)
elif xysize == -1:
fftvol = sp_utilities.model_blank(imgsize, imgsize, zsize)
else:
fftvol = sp_utilities.model_blank(xysize, xysize, zsize)
return fftvol
def recons3d_trl_struct_MPI(
myid,
main_node,
prjlist,
paramstructure,
refang,
rshifts_shrank,
delta,
upweighted=True,
mpi_comm=None,
CTF=True,
target_size=-1,
avgnorm=1.0,
norm_per_particle=None,
):
"""
recons3d_4nn_ctf - calculate CTF-corrected 3-D reconstruction from a set of projections using three Eulerian angles, two shifts, and CTF settings for each projeciton image
Input
list_of_prjlist: list of lists of projections to be included in the reconstruction
"""
if mpi_comm == None:
mpi_comm = mpi.MPI_COMM_WORLD
refvol = sp_utilities.model_blank(target_size)
refvol.set_attr("fudge", 1.0)
if CTF:
do_ctf = 1
else:
do_ctf = 0
fftvol = EMAN2_cppwrap.EMData()
weight = EMAN2_cppwrap.EMData()
params = {
"size": target_size,
"npad": 2,
"snr": 1.0,
"sign": 1,
"symmetry": "c1",
"refvol": refvol,
"fftvol": fftvol,
"weight": weight,
"do_ctf": do_ctf,
}
r = EMAN2_cppwrap.Reconstructors.get("nn4_ctfw", params)
r.setup()
if prjlist:
if norm_per_particle == None:
norm_per_particle = len(prjlist) * [1.0]
nnx = prjlist[0].get_xsize()
nny = prjlist[0].get_ysize()
nshifts = len(rshifts_shrank)
for im in range(len(prjlist)):
# parse projection structure, generate three lists:
# [ipsi+iang], [ishift], [probability]
# Number of orientations for a given image
numbor = len(paramstructure[im][2])
ipsiandiang = [
old_div(paramstructure[im][2][i][0], 1000)
for i in range(numbor)
]
allshifts = [
paramstructure[im][2][i][0] % 1000 for i in range(numbor)
]
probs = [paramstructure[im][2][i][1] for i in range(numbor)]
# Find unique projection directions
tdir = list(set(ipsiandiang))
bckgn = prjlist[im].get_attr("bckgnoise")
ct = prjlist[im].get_attr("ctf")
# For each unique projection direction:
data = [None] * nshifts
for ii in range(len(tdir)):
# Find the number of times given projection direction appears on the list, it is the number of different shifts associated with it.
lshifts = sp_utilities.findall(tdir[ii], ipsiandiang)
toprab = 0.0
for ki in range(len(lshifts)):
toprab += probs[lshifts[ki]]
recdata = EMAN2_cppwrap.EMData(nny, nny, 1, False)
recdata.set_attr("is_complex", 0)
for ki in range(len(lshifts)):
lpt = allshifts[lshifts[ki]]
if data[lpt] == None:
data[lpt] = sp_fundamentals.fshift(
prjlist[im], rshifts_shrank[lpt][0],
rshifts_shrank[lpt][1])
data[lpt].set_attr("is_complex", 0)
EMAN2_cppwrap.Util.add_img(
recdata,
EMAN2_cppwrap.Util.mult_scalar(
data[lpt], old_div(probs[lshifts[ki]], toprab)),
)
recdata.set_attr_dict({
"padffted": 1,
"is_fftpad": 1,
"is_fftodd": 0,
"is_complex_ri": 1,
"is_complex": 1,
})
if not upweighted:
recdata = sp_filter.filt_table(recdata, bckgn)
recdata.set_attr_dict({"bckgnoise": bckgn, "ctf": ct})
ipsi = tdir[ii] % 100000
iang = old_div(tdir[ii], 100000)
r.insert_slice(
recdata,
EMAN2_cppwrap.Transform({
"type":
"spider",
"phi":
refang[iang][0],
"theta":
refang[iang][1],
"psi":
refang[iang][2] + ipsi * delta,
}),
old_div(toprab * avgnorm, norm_per_particle[im]),
)
# clean stuff
del bckgn, recdata, tdir, ipsiandiang, allshifts, probs
sp_utilities.reduce_EMData_to_root(fftvol, myid, main_node, comm=mpi_comm)
sp_utilities.reduce_EMData_to_root(weight, myid, main_node, comm=mpi_comm)
if myid == main_node:
dummy = r.finish(True)
mpi.mpi_barrier(mpi_comm)
if myid == main_node:
return fftvol, weight, refvol
else:
return None, None, None
def recons3d_4nnw_MPI(
myid,
prjlist,
bckgdata,
snr=1.0,
sign=1,
symmetry="c1",
finfo=None,
npad=2,
xysize=-1,
zsize=-1,
mpi_comm=None,
smearstep=0.0,
fsc=None,
):
"""
recons3d_4nn_ctf - calculate CTF-corrected 3-D reconstruction from a set of projections using three Eulerian angles, two shifts, and CTF settings for each projeciton image
Input
stack: name of the stack file containing projection data, projections have to be squares
prjlist: list of projections to be included in the reconstruction or image iterator
bckgdata = [get_im("tsd.hdf"),read_text_file("data_stamp.txt")]
snr: Signal-to-Noise Ratio of the data
sign: sign of the CTF
symmetry: point-group symmetry to be enforced, each projection will enter the reconstruction in all symmetry-related directions.
"""
pass # IMPORTIMPORTIMPORT from sp_utilities import reduce_EMData_to_root, pad
pass # IMPORTIMPORTIMPORT from EMAN2 import Reconstructors
pass # IMPORTIMPORTIMPORT from sp_utilities import iterImagesList, set_params_proj, model_blank
pass # IMPORTIMPORTIMPORT from mpi import MPI_COMM_WORLD
pass # IMPORTIMPORTIMPORT import types
if mpi_comm == None:
mpi_comm = mpi.MPI_COMM_WORLD
if type(prjlist) == list:
prjlist = sp_utilities.iterImagesList(prjlist)
if not prjlist.goToNext():
sp_global_def.ERROR("empty input list", "recons3d_4nnw_MPI", 1)
imgsize = prjlist.image().get_xsize()
if prjlist.image().get_ysize() != imgsize:
imgsize = max(imgsize, prjlist.image().get_ysize())
dopad = True
else:
dopad = False
prjlist.goToPrev()
# Do the FSC shtick.
bnx = old_div(imgsize * npad, 2) + 1
if fsc:
pass # IMPORTIMPORTIMPORT from math import sqrt
pass # IMPORTIMPORTIMPORT from sp_utilities import reshape_1d
t = [0.0] * len(fsc)
for i in range(len(fsc)):
t[i] = min(max(fsc[i], 0.0), 0.999)
t = sp_utilities.reshape_1d(t, len(t), npad * len(t))
refvol = sp_utilities.model_blank(bnx, 1, 1, 0.0)
for i in range(len(fsc)):
refvol.set_value_at(i, t[i])
else:
refvol = sp_utilities.model_blank(bnx, 1, 1, 1.0)
refvol.set_attr("fudge", 1.0)
fftvol = EMAN2_cppwrap.EMData()
weight = EMAN2_cppwrap.EMData()
if smearstep > 0.0:
# if myid == 0: print " Setting smear in prepare_recons_ctf"
ns = 1
smear = []
for j in range(-ns, ns + 1):
if j != 0:
for i in range(-ns, ns + 1):
for k in range(-ns, ns + 1):
smear += [
i * smearstep, j * smearstep, k * smearstep, 1.0
]
# Deal with theta = 0.0 cases
prj = []
for i in range(-ns, ns + 1):
for k in range(-ns, ns + 1):
prj.append(i + k)
for i in range(-2 * ns, 2 * ns + 1, 1):
smear += [i * smearstep, 0.0, 0.0, float(prj.count(i))]
# if myid == 0: print " Smear ",smear
fftvol.set_attr("smear", smear)
if xysize == -1 and zsize == -1:
params = {
"size": imgsize,
"npad": npad,
"snr": snr,
"sign": sign,
"symmetry": symmetry,
"refvol": refvol,
"fftvol": fftvol,
"weight": weight,
}
r = EMAN2_cppwrap.Reconstructors.get("nn4_ctfw", params)
else:
if xysize != -1 and zsize != -1:
rx = old_div(float(xysize), imgsize)
ry = old_div(float(xysize), imgsize)
rz = old_div(float(zsize), imgsize)
elif xysize != -1:
rx = old_div(float(xysize), imgsize)
ry = old_div(float(xysize), imgsize)
rz = 1.0
else:
rx = 1.0
ry = 1.0
rz = old_div(float(zsize), imgsize)
# There is an error here with sizeprojection PAP 10/22/2014
params = {
"size": sizeprojection,
"npad": npad,
"snr": snr,
"sign": sign,
"symmetry": symmetry,
"fftvol": fftvol,
"weight": weight,
"xratio": rx,
"yratio": ry,
"zratio": rz,
}
r = EMAN2_cppwrap.Reconstructors.get("nn4_ctf_rect", params)
r.setup()
# from utilities import model_blank, get_im, read_text_file
# bckgdata = [get_im("tsd.hdf"),read_text_file("data_stamp.txt")]
nnx = bckgdata[0].get_xsize()
nny = bckgdata[0].get_ysize()
bckgnoise = []
for i in range(nny):
prj = sp_utilities.model_blank(nnx)
for k in range(nnx):
prj[k] = bckgdata[0].get_value_at(k, i)
bckgnoise.append(prj)
datastamp = bckgdata[1]
if not (finfo is None):
nimg = 0
while prjlist.goToNext():
prj = prjlist.image()
try:
stmp = old_div(nnx, 0)
stmp = prj.get_attr("ptcl_source_image")
except:
try:
stmp = prj.get_attr("ctf")
stmp = round(stmp.defocus, 4)
except:
sp_global_def.ERROR(
"Either ptcl_source_image or ctf has to be present in the header.",
"recons3d_4nnw_MPI",
1,
myid,
)
try:
indx = datastamp.index(stmp)
except:
sp_global_def.ERROR("Problem with indexing ptcl_source_image.",
"recons3d_4nnw_MPI", 1, myid)
if dopad:
prj = sp_utilities.pad(prj, imgsize, imgsize, 1, "circumference")
prj.set_attr("bckgnoise", bckgnoise[indx])
insert_slices(r, prj)
if not (finfo is None):
nimg += 1
finfo.write(" %4d inserted\n" % (nimg))
finfo.flush()
del sp_utilities.pad
if not (finfo is None):
finfo.write("begin reduce\n")
finfo.flush()
sp_utilities.reduce_EMData_to_root(fftvol, myid, comm=mpi_comm)
sp_utilities.reduce_EMData_to_root(weight, myid, comm=mpi_comm)
if not (finfo is None):
finfo.write("after reduce\n")
finfo.flush()
if myid == 0:
dummy = r.finish(True)
else:
pass # IMPORTIMPORTIMPORT from sp_utilities import model_blank
if xysize == -1 and zsize == -1:
fftvol = sp_utilities.model_blank(imgsize, imgsize, imgsize)
else:
if zsize == -1:
fftvol = sp_utilities.model_blank(xysize, xysize, imgsize)
elif xysize == -1:
fftvol = sp_utilities.model_blank(imgsize, imgsize, zsize)
else:
fftvol = sp_utilities.model_blank(xysize, xysize, zsize)
return fftvol
def recons3d_4nn_MPI(
myid,
prjlist,
symmetry="c1",
finfo=None,
snr=1.0,
npad=2,
xysize=-1,
zsize=-1,
mpi_comm=None,
):
if mpi_comm == None:
mpi_comm = mpi.MPI_COMM_WORLD
if type(prjlist) == list:
prjlist = sp_utilities.iterImagesList(prjlist)
if not prjlist.goToNext():
sp_global_def.ERROR("empty input list", "recons3d_4nn_MPI", 1)
imgsize = prjlist.image().get_xsize()
if prjlist.image().get_ysize() != imgsize:
imgsize = max(imgsize, prjlist.image().get_ysize())
dopad = True
else:
dopad = False
prjlist.goToPrev()
fftvol = EMAN2_cppwrap.EMData()
weight = EMAN2_cppwrap.EMData()
if xysize == -1 and zsize == -1:
params = {
"size": imgsize,
"npad": npad,
"symmetry": symmetry,
"fftvol": fftvol,
"weight": weight,
"snr": snr,
}
r = EMAN2_cppwrap.Reconstructors.get("nn4", params)
else:
if xysize != -1 and zsize != -1:
rx = old_div(float(xysize), imgsize)
ry = old_div(float(xysize), imgsize)
rz = old_div(float(zsize), imgsize)
elif xysize != -1:
rx = old_div(float(xysize), imgsize)
ry = old_div(float(xysize), imgsize)
rz = 1.0
else:
rx = 1.0
ry = 1.0
rz = old_div(float(zsize), imgsize)
params = {
"sizeprojection": imgsize,
"npad": npad,
"symmetry": symmetry,
"fftvol": fftvol,
"weight": weight,
"xratio": rx,
"yratio": ry,
"zratio": rz,
}
r = EMAN2_cppwrap.Reconstructors.get("nn4_rect", params)
r.setup()
if not (finfo is None):
nimg = 0
while prjlist.goToNext():
prj = prjlist.image()
if dopad:
prj = sp_utilities.pad(prj, imgsize, imgsize, 1, "circumference")
insert_slices(r, prj)
if not (finfo is None):
nimg += 1
finfo.write("Image %4d inserted.\n" % (nimg))
finfo.flush()
if not (finfo is None):
finfo.write("Begin reducing ...\n")
finfo.flush()
sp_utilities.reduce_EMData_to_root(fftvol, myid, comm=mpi_comm)
sp_utilities.reduce_EMData_to_root(weight, myid, comm=mpi_comm)
if myid == 0:
dummy = r.finish(True)
else:
if xysize == -1 and zsize == -1:
fftvol = sp_utilities.model_blank(imgsize, imgsize, imgsize)
else:
if zsize == -1:
fftvol = sp_utilities.model_blank(xysize, xysize, imgsize)
elif xysize == -1:
fftvol = sp_utilities.model_blank(imgsize, imgsize, zsize)
else:
fftvol = sp_utilities.model_blank(xysize, xysize, zsize)
return fftvol
def mref_ali2d_MPI(stack,
refim,
outdir,
maskfile=None,
ir=1,
ou=-1,
rs=1,
xrng=0,
yrng=0,
step=1,
center=1,
maxit=10,
CTF=False,
snr=1.0,
user_func_name="ref_ali2d",
rand_seed=1000):
# 2D multi-reference alignment using rotational ccf in polar coordinates and quadratic interpolation
from sp_utilities import model_circle, combine_params2, inverse_transform2, drop_image, get_image, get_im
from sp_utilities import reduce_EMData_to_root, bcast_EMData_to_all, bcast_number_to_all
from sp_utilities import send_attr_dict
from sp_utilities import center_2D
from sp_statistics import fsc_mask
from sp_alignment import Numrinit, ringwe, search_range
from sp_fundamentals import rot_shift2D, fshift
from sp_utilities import get_params2D, set_params2D
from random import seed, randint
from sp_morphology import ctf_2
from sp_filter import filt_btwl, filt_params
from numpy import reshape, shape
from sp_utilities import print_msg, print_begin_msg, print_end_msg
import os
import sys
import shutil
from sp_applications import MPI_start_end
from mpi import mpi_comm_size, mpi_comm_rank, MPI_COMM_WORLD
from mpi import mpi_reduce, mpi_bcast, mpi_barrier, mpi_recv, mpi_send
from mpi import MPI_SUM, MPI_FLOAT, MPI_INT
number_of_proc = mpi_comm_size(MPI_COMM_WORLD)
myid = mpi_comm_rank(MPI_COMM_WORLD)
main_node = 0
# create the output directory, if it does not exist
if os.path.exists(outdir):
ERROR(
'Output directory exists, please change the name and restart the program',
"mref_ali2d_MPI ", 1, myid)
mpi_barrier(MPI_COMM_WORLD)
import sp_global_def
if myid == main_node:
os.mkdir(outdir)
sp_global_def.LOGFILE = os.path.join(outdir, sp_global_def.LOGFILE)
print_begin_msg("mref_ali2d_MPI")
nima = EMUtil.get_image_count(stack)
image_start, image_end = MPI_start_end(nima, number_of_proc, myid)
nima = EMUtil.get_image_count(stack)
ima = EMData()
ima.read_image(stack, image_start)
first_ring = int(ir)
last_ring = int(ou)
rstep = int(rs)
max_iter = int(maxit)
if max_iter == 0:
max_iter = 10
auto_stop = True
else:
auto_stop = False
if myid == main_node:
print_msg("Input stack : %s\n" % (stack))
print_msg("Reference stack : %s\n" % (refim))
print_msg("Output directory : %s\n" % (outdir))
print_msg("Maskfile : %s\n" % (maskfile))
print_msg("Inner radius : %i\n" % (first_ring))
nx = ima.get_xsize()
# default value for the last ring
if last_ring == -1: last_ring = nx / 2 - 2
if myid == main_node:
print_msg("Outer radius : %i\n" % (last_ring))
print_msg("Ring step : %i\n" % (rstep))
print_msg("X search range : %f\n" % (xrng))
print_msg("Y search range : %f\n" % (yrng))
print_msg("Translational step : %f\n" % (step))
print_msg("Center type : %i\n" % (center))
print_msg("Maximum iteration : %i\n" % (max_iter))
print_msg("CTF correction : %s\n" % (CTF))
print_msg("Signal-to-Noise Ratio : %f\n" % (snr))
print_msg("Random seed : %i\n\n" % (rand_seed))
print_msg("User function : %s\n" % (user_func_name))
import sp_user_functions
user_func = sp_user_functions.factory[user_func_name]
if maskfile:
import types
if type(maskfile) is bytes: mask = get_image(maskfile)
else: mask = maskfile
else: mask = model_circle(last_ring, nx, nx)
# references, do them on all processors...
refi = []
numref = EMUtil.get_image_count(refim)
# IMAGES ARE SQUARES! center is in SPIDER convention
cnx = nx / 2 + 1
cny = cnx
mode = "F"
#precalculate rings
numr = Numrinit(first_ring, last_ring, rstep, mode)
wr = ringwe(numr, mode)
# prepare reference images on all nodes
ima.to_zero()
for j in range(numref):
# even, odd, numer of even, number of images. After frc, totav
refi.append([get_im(refim, j), ima.copy(), 0])
# for each node read its share of data
data = EMData.read_images(stack, list(range(image_start, image_end)))
for im in range(image_start, image_end):
data[im - image_start].set_attr('ID', im)
if myid == main_node: seed(rand_seed)
a0 = -1.0
again = True
Iter = 0
ref_data = [mask, center, None, None]
while Iter < max_iter and again:
ringref = []
mashi = cnx - last_ring - 2
for j in range(numref):
refi[j][0].process_inplace("normalize.mask", {
"mask": mask,
"no_sigma": 1
}) # normalize reference images to N(0,1)
cimage = Util.Polar2Dm(refi[j][0], cnx, cny, numr, mode)
Util.Frngs(cimage, numr)
Util.Applyws(cimage, numr, wr)
ringref.append(cimage)
# zero refi
refi[j][0].to_zero()
refi[j][1].to_zero()
refi[j][2] = 0
assign = [[] for i in range(numref)]
# begin MPI section
for im in range(image_start, image_end):
alpha, sx, sy, mirror, scale = get_params2D(data[im - image_start])
# Why inverse? 07/11/2015 PAP
alphai, sxi, syi, scalei = inverse_transform2(alpha, sx, sy)
# normalize
data[im - image_start].process_inplace("normalize.mask", {
"mask": mask,
"no_sigma": 0
}) # subtract average under the mask
# If shifts are outside of the permissible range, reset them
if (abs(sxi) > mashi or abs(syi) > mashi):
sxi = 0.0
syi = 0.0
set_params2D(data[im - image_start], [0.0, 0.0, 0.0, 0, 1.0])
ny = nx
txrng = search_range(nx, last_ring, sxi, xrng, "mref_ali2d_MPI")
txrng = [txrng[1], txrng[0]]
tyrng = search_range(ny, last_ring, syi, yrng, "mref_ali2d_MPI")
tyrng = [tyrng[1], tyrng[0]]
# align current image to the reference
[angt, sxst, syst, mirrort, xiref,
peakt] = Util.multiref_polar_ali_2d(data[im - image_start],
ringref, txrng, tyrng, step,
mode, numr, cnx + sxi,
cny + syi)
iref = int(xiref)
# combine parameters and set them to the header, ignore previous angle and mirror
[alphan, sxn, syn,
mn] = combine_params2(0.0, -sxi, -syi, 0, angt, sxst, syst,
(int)(mirrort))
set_params2D(data[im - image_start],
[alphan, sxn, syn, int(mn), scale])
data[im - image_start].set_attr('assign', iref)
# apply current parameters and add to the average
temp = rot_shift2D(data[im - image_start], alphan, sxn, syn, mn)
it = im % 2
Util.add_img(refi[iref][it], temp)
assign[iref].append(im)
#assign[im] = iref
refi[iref][2] += 1.0
del ringref
# end MPI section, bring partial things together, calculate new reference images, broadcast them back
for j in range(numref):
reduce_EMData_to_root(refi[j][0], myid, main_node)
reduce_EMData_to_root(refi[j][1], myid, main_node)
refi[j][2] = mpi_reduce(refi[j][2], 1, MPI_FLOAT, MPI_SUM,
main_node, MPI_COMM_WORLD)
if (myid == main_node): refi[j][2] = int(refi[j][2][0])
# gather assignements
for j in range(numref):
if myid == main_node:
for n in range(number_of_proc):
if n != main_node:
import sp_global_def
ln = mpi_recv(1, MPI_INT, n,
sp_global_def.SPARX_MPI_TAG_UNIVERSAL,
MPI_COMM_WORLD)
lis = mpi_recv(ln[0], MPI_INT, n,
sp_global_def.SPARX_MPI_TAG_UNIVERSAL,
MPI_COMM_WORLD)
for l in range(ln[0]):
assign[j].append(int(lis[l]))
else:
import sp_global_def
mpi_send(len(assign[j]), 1, MPI_INT, main_node,
sp_global_def.SPARX_MPI_TAG_UNIVERSAL, MPI_COMM_WORLD)
mpi_send(assign[j], len(assign[j]), MPI_INT, main_node,
sp_global_def.SPARX_MPI_TAG_UNIVERSAL, MPI_COMM_WORLD)
if myid == main_node:
# replace the name of the stack with reference with the current one
refim = os.path.join(outdir, "aqm%03d.hdf" % Iter)
a1 = 0.0
ave_fsc = []
for j in range(numref):
if refi[j][2] < 4:
#ERROR("One of the references vanished","mref_ali2d_MPI",1)
# if vanished, put a random image (only from main node!) there
assign[j] = []
assign[j].append(
randint(image_start, image_end - 1) - image_start)
refi[j][0] = data[assign[j][0]].copy()
#print 'ERROR', j
else:
#frsc = fsc_mask(refi[j][0], refi[j][1], mask, 1.0, os.path.join(outdir,"drm%03d%04d"%(Iter, j)))
from sp_statistics import fsc
frsc = fsc(
refi[j][0], refi[j][1], 1.0,
os.path.join(outdir, "drm%03d%04d.txt" % (Iter, j)))
Util.add_img(refi[j][0], refi[j][1])
Util.mul_scalar(refi[j][0], 1.0 / float(refi[j][2]))
if ave_fsc == []:
for i in range(len(frsc[1])):
ave_fsc.append(frsc[1][i])
c_fsc = 1
else:
for i in range(len(frsc[1])):
ave_fsc[i] += frsc[1][i]
c_fsc += 1
#print 'OK', j, len(frsc[1]), frsc[1][0:5], ave_fsc[0:5]
#print 'sum', sum(ave_fsc)
if sum(ave_fsc) != 0:
for i in range(len(ave_fsc)):
ave_fsc[i] /= float(c_fsc)
frsc[1][i] = ave_fsc[i]
for j in range(numref):
ref_data[2] = refi[j][0]
ref_data[3] = frsc
refi[j][0], cs = user_func(ref_data)
# write the current average
TMP = []
for i_tmp in range(len(assign[j])):
TMP.append(float(assign[j][i_tmp]))
TMP.sort()
refi[j][0].set_attr_dict({'ave_n': refi[j][2], 'members': TMP})
del TMP
refi[j][0].process_inplace("normalize.mask", {
"mask": mask,
"no_sigma": 1
})
refi[j][0].write_image(refim, j)
Iter += 1
msg = "ITERATION #%3d %d\n\n" % (Iter, again)
print_msg(msg)
for j in range(numref):
msg = " group #%3d number of particles = %7d\n" % (
j, refi[j][2])
print_msg(msg)
Iter = bcast_number_to_all(Iter, main_node) # need to tell all
if again:
for j in range(numref):
bcast_EMData_to_all(refi[j][0], myid, main_node)
# clean up
del assign
# write out headers and STOP, under MPI writing has to be done sequentially (time-consumming)
mpi_barrier(MPI_COMM_WORLD)
if CTF and data_had_ctf == 0:
for im in range(len(data)):
data[im].set_attr('ctf_applied', 0)
par_str = ['xform.align2d', 'assign', 'ID']
if myid == main_node:
from sp_utilities import file_type
if (file_type(stack) == "bdb"):
from sp_utilities import recv_attr_dict_bdb
recv_attr_dict_bdb(main_node, stack, data, par_str, image_start,
image_end, number_of_proc)
else:
from sp_utilities import recv_attr_dict
recv_attr_dict(main_node, stack, data, par_str, image_start,
image_end, number_of_proc)
else:
send_attr_dict(main_node, data, par_str, image_start, image_end)
if myid == main_node:
print_end_msg("mref_ali2d_MPI")