def prgq(volft, kb, nx, delta, ref_a, sym, MPI=False):
"""
Generate set of projections based on even angles
The command returns list of ffts of projections
"""
from sp_projection import prep_vol, prgs
from sp_applications import MPI_start_end
from sp_utilities import even_angles, model_blank
from sp_fundamentals import fft
# generate list of Eulerian angles for reference projections
# phi, theta, psi
mode = "F"
ref_angles = even_angles(delta,
symmetry=sym,
method=ref_a,
phiEqpsi="Minus")
cnx = nx // 2 + 1
cny = nx // 2 + 1
num_ref = len(ref_angles)
if MPI:
from mpi import mpi_comm_rank, mpi_comm_size, MPI_COMM_WORLD
myid = mpi_comm_rank(MPI_COMM_WORLD)
ncpu = mpi_comm_size(MPI_COMM_WORLD)
else:
ncpu = 1
myid = 0
from sp_applications import MPI_start_end
ref_start, ref_end = MPI_start_end(num_ref, ncpu, myid)
prjref = [
] # list of (image objects) reference projections in Fourier representation
for i in range(num_ref):
prjref.append(model_blank(
nx,
nx)) # I am not sure why is that necessary, why not put None's??
for i in range(ref_start, ref_end):
prjref[i] = prgs(
volft, kb,
[ref_angles[i][0], ref_angles[i][1], ref_angles[i][2], 0.0, 0.0])
if MPI:
from sp_utilities import bcast_EMData_to_all
for i in range(num_ref):
for j in range(ncpu):
ref_start, ref_end = MPI_start_end(num_ref, ncpu, j)
if i >= ref_start and i < ref_end: rootid = j
bcast_EMData_to_all(prjref[i], myid, rootid)
for i in range(len(ref_angles)):
prjref[i].set_attr_dict({
"phi": ref_angles[i][0],
"theta": ref_angles[i][1],
"psi": ref_angles[i][2]
})
return prjref
def getalldata(stack, myid, nproc): if(myid == 0): ndata = EMUtil.get_image_count(stack) else: ndata = 0 ndata = bcast_number_to_all(ndata) if( ndata < nproc): if(myid<ndata): image_start = myid image_end = myid+1 else: image_start = 0 image_end = 1 else: image_start, image_end = MPI_start_end(ndata, nproc, myid) data = EMData.read_images(stack, list(range(image_start, image_end))) return data
def getindexdata(stack, partids, partstack, myid, nproc): # The function will read from stack a subset of images specified in partids # and assign to them parameters from partstack # So, the lengths of partids and partstack are the same. # The read data is properly distributed among MPI threads. lpartids = list(map(int, read_text_file(partids) )) ndata = len(lpartids) partstack = read_text_row(partstack) if( ndata < nproc): if(myid<ndata): image_start = myid image_end = myid+1 else: image_start = 0 image_end = 1 else: image_start, image_end = MPI_start_end(ndata, nproc, myid) lpartids = lpartids[image_start:image_end] partstack = partstack[image_start:image_end] data = EMData.read_images(stack, lpartids) for i in range(len(partstack)): set_params_proj(data[i], partstack[i]) return data
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 main():
progname = os.path.basename(sys.argv[0])
usage = progname + " proj_stack output_averages --MPI"
parser = OptionParser(usage, version=SPARXVERSION)
parser.add_option("--img_per_group",
type="int",
default=100,
help="number of images per group")
parser.add_option("--radius",
type="int",
default=-1,
help="radius for alignment")
parser.add_option(
"--xr",
type="string",
default="2 1",
help="range for translation search in x direction, search is +/xr")
parser.add_option(
"--yr",
type="string",
default="-1",
help=
"range for translation search in y direction, search is +/yr (default = same as xr)"
)
parser.add_option(
"--ts",
type="string",
default="1 0.5",
help=
"step size of the translation search in both directions, search is -xr, -xr+ts, 0, xr-ts, xr, can be fractional"
)
parser.add_option(
"--iter",
type="int",
default=30,
help="number of iterations within alignment (default = 30)")
parser.add_option(
"--num_ali",
type="int",
default=5,
help="number of alignments performed for stability (default = 5)")
parser.add_option("--thld_err",
type="float",
default=1.0,
help="threshold of pixel error (default = 1.732)")
parser.add_option(
"--grouping",
type="string",
default="GRP",
help=
"do grouping of projections: PPR - per projection, GRP - different size groups, exclusive (default), GEV - grouping equal size"
)
parser.add_option(
"--delta",
type="float",
default=-1.0,
help="angular step for reference projections (required for GEV method)"
)
parser.add_option(
"--fl",
type="float",
default=0.3,
help="cut-off frequency of hyperbolic tangent low-pass Fourier filter")
parser.add_option(
"--aa",
type="float",
default=0.2,
help="fall-off of hyperbolic tangent low-pass Fourier filter")
parser.add_option("--CTF",
action="store_true",
default=False,
help="Consider CTF correction during the alignment ")
parser.add_option("--MPI",
action="store_true",
default=False,
help="use MPI version")
(options, args) = parser.parse_args()
myid = mpi.mpi_comm_rank(MPI_COMM_WORLD)
number_of_proc = mpi.mpi_comm_size(MPI_COMM_WORLD)
main_node = 0
if len(args) == 2:
stack = args[0]
outdir = args[1]
else:
sp_global_def.ERROR("Incomplete list of arguments",
"sxproj_stability.main",
1,
myid=myid)
return
if not options.MPI:
sp_global_def.ERROR("Non-MPI not supported!",
"sxproj_stability.main",
1,
myid=myid)
return
if sp_global_def.CACHE_DISABLE:
from sp_utilities import disable_bdb_cache
disable_bdb_cache()
sp_global_def.BATCH = True
img_per_grp = options.img_per_group
radius = options.radius
ite = options.iter
num_ali = options.num_ali
thld_err = options.thld_err
xrng = get_input_from_string(options.xr)
if options.yr == "-1":
yrng = xrng
else:
yrng = get_input_from_string(options.yr)
step = get_input_from_string(options.ts)
if myid == main_node:
nima = EMUtil.get_image_count(stack)
img = get_image(stack)
nx = img.get_xsize()
ny = img.get_ysize()
else:
nima = 0
nx = 0
ny = 0
nima = bcast_number_to_all(nima)
nx = bcast_number_to_all(nx)
ny = bcast_number_to_all(ny)
if radius == -1: radius = nx / 2 - 2
mask = model_circle(radius, nx, nx)
st = time()
if options.grouping == "GRP":
if myid == main_node:
sxprint(" A ", myid, " ", time() - st)
proj_attr = EMUtil.get_all_attributes(stack, "xform.projection")
proj_params = []
for i in range(nima):
dp = proj_attr[i].get_params("spider")
phi, theta, psi, s2x, s2y = dp["phi"], dp["theta"], dp[
"psi"], -dp["tx"], -dp["ty"]
proj_params.append([phi, theta, psi, s2x, s2y])
# Here is where the grouping is done, I didn't put enough annotation in the group_proj_by_phitheta,
# So I will briefly explain it here
# proj_list : Returns a list of list of particle numbers, each list contains img_per_grp particle numbers
# except for the last one. Depending on the number of particles left, they will either form a
# group or append themselves to the last group
# angle_list : Also returns a list of list, each list contains three numbers (phi, theta, delta), (phi,
# theta) is the projection angle of the center of the group, delta is the range of this group
# mirror_list: Also returns a list of list, each list contains img_per_grp True or False, which indicates
# whether it should take mirror position.
# In this program angle_list and mirror list are not of interest.
proj_list_all, angle_list, mirror_list = group_proj_by_phitheta(
proj_params, img_per_grp=img_per_grp)
del proj_params
sxprint(" B number of groups ", myid, " ", len(proj_list_all),
time() - st)
mpi_barrier(MPI_COMM_WORLD)
# Number of groups, actually there could be one or two more groups, since the size of the remaining group varies
# we will simply assign them to main node.
n_grp = nima / img_per_grp - 1
# Divide proj_list_all equally to all nodes, and becomes proj_list
proj_list = []
for i in range(n_grp):
proc_to_stay = i % number_of_proc
if proc_to_stay == main_node:
if myid == main_node: proj_list.append(proj_list_all[i])
elif myid == main_node:
mpi_send(len(proj_list_all[i]), 1, MPI_INT, proc_to_stay,
SPARX_MPI_TAG_UNIVERSAL, MPI_COMM_WORLD)
mpi_send(proj_list_all[i], len(proj_list_all[i]), MPI_INT,
proc_to_stay, SPARX_MPI_TAG_UNIVERSAL, MPI_COMM_WORLD)
elif myid == proc_to_stay:
img_per_grp = mpi_recv(1, MPI_INT, main_node,
SPARX_MPI_TAG_UNIVERSAL, MPI_COMM_WORLD)
img_per_grp = int(img_per_grp[0])
temp = mpi_recv(img_per_grp, MPI_INT, main_node,
SPARX_MPI_TAG_UNIVERSAL, MPI_COMM_WORLD)
proj_list.append(list(map(int, temp)))
del temp
mpi_barrier(MPI_COMM_WORLD)
sxprint(" C ", myid, " ", time() - st)
if myid == main_node:
# Assign the remaining groups to main_node
for i in range(n_grp, len(proj_list_all)):
proj_list.append(proj_list_all[i])
del proj_list_all, angle_list, mirror_list
# Compute stability per projection projection direction, equal number assigned, thus overlaps
elif options.grouping == "GEV":
if options.delta == -1.0:
ERROR(
"Angular step for reference projections is required for GEV method"
)
return
from sp_utilities import even_angles, nearestk_to_refdir, getvec
refproj = even_angles(options.delta)
img_begin, img_end = MPI_start_end(len(refproj), number_of_proc, myid)
# Now each processor keeps its own share of reference projections
refprojdir = refproj[img_begin:img_end]
del refproj
ref_ang = [0.0] * (len(refprojdir) * 2)
for i in range(len(refprojdir)):
ref_ang[i * 2] = refprojdir[0][0]
ref_ang[i * 2 + 1] = refprojdir[0][1] + i * 0.1
sxprint(" A ", myid, " ", time() - st)
proj_attr = EMUtil.get_all_attributes(stack, "xform.projection")
# the solution below is very slow, do not use it unless there is a problem with the i/O
"""
for i in xrange(number_of_proc):
if myid == i:
proj_attr = EMUtil.get_all_attributes(stack, "xform.projection")
mpi_barrier(MPI_COMM_WORLD)
"""
sxprint(" B ", myid, " ", time() - st)
proj_ang = [0.0] * (nima * 2)
for i in range(nima):
dp = proj_attr[i].get_params("spider")
proj_ang[i * 2] = dp["phi"]
proj_ang[i * 2 + 1] = dp["theta"]
sxprint(" C ", myid, " ", time() - st)
asi = Util.nearestk_to_refdir(proj_ang, ref_ang, img_per_grp)
del proj_ang, ref_ang
proj_list = []
for i in range(len(refprojdir)):
proj_list.append(asi[i * img_per_grp:(i + 1) * img_per_grp])
del asi
sxprint(" D ", myid, " ", time() - st)
#from sys import exit
#exit()
# Compute stability per projection
elif options.grouping == "PPR":
sxprint(" A ", myid, " ", time() - st)
proj_attr = EMUtil.get_all_attributes(stack, "xform.projection")
sxprint(" B ", myid, " ", time() - st)
proj_params = []
for i in range(nima):
dp = proj_attr[i].get_params("spider")
phi, theta, psi, s2x, s2y = dp["phi"], dp["theta"], dp[
"psi"], -dp["tx"], -dp["ty"]
proj_params.append([phi, theta, psi, s2x, s2y])
img_begin, img_end = MPI_start_end(nima, number_of_proc, myid)
sxprint(" C ", myid, " ", time() - st)
from sp_utilities import nearest_proj
proj_list, mirror_list = nearest_proj(
proj_params, img_per_grp,
list(range(img_begin, img_begin + 1))) #range(img_begin, img_end))
refprojdir = proj_params[img_begin:img_end]
del proj_params, mirror_list
sxprint(" D ", myid, " ", time() - st)
else:
ERROR("Incorrect projection grouping option")
return
###########################################################################################################
# Begin stability test
from sp_utilities import get_params_proj, read_text_file
#if myid == 0:
# from utilities import read_text_file
# proj_list[0] = map(int, read_text_file("lggrpp0.txt"))
from sp_utilities import model_blank
aveList = [model_blank(nx, ny)] * len(proj_list)
if options.grouping == "GRP":
refprojdir = [[0.0, 0.0, -1.0]] * len(proj_list)
for i in range(len(proj_list)):
sxprint(" E ", myid, " ", time() - st)
class_data = EMData.read_images(stack, proj_list[i])
#print " R ",myid," ",time()-st
if options.CTF:
from sp_filter import filt_ctf
for im in range(len(class_data)): # MEM LEAK!!
atemp = class_data[im].copy()
btemp = filt_ctf(atemp, atemp.get_attr("ctf"), binary=1)
class_data[im] = btemp
#class_data[im] = filt_ctf(class_data[im], class_data[im].get_attr("ctf"), binary=1)
for im in class_data:
try:
t = im.get_attr(
"xform.align2d") # if they are there, no need to set them!
except:
try:
t = im.get_attr("xform.projection")
d = t.get_params("spider")
set_params2D(im, [0.0, -d["tx"], -d["ty"], 0, 1.0])
except:
set_params2D(im, [0.0, 0.0, 0.0, 0, 1.0])
#print " F ",myid," ",time()-st
# Here, we perform realignment num_ali times
all_ali_params = []
for j in range(num_ali):
if (xrng[0] == 0.0 and yrng[0] == 0.0):
avet = ali2d_ras(class_data,
randomize=True,
ir=1,
ou=radius,
rs=1,
step=1.0,
dst=90.0,
maxit=ite,
check_mirror=True,
FH=options.fl,
FF=options.aa)
else:
avet = within_group_refinement(class_data, mask, True, 1,
radius, 1, xrng, yrng, step,
90.0, ite, options.fl,
options.aa)
ali_params = []
for im in range(len(class_data)):
alpha, sx, sy, mirror, scale = get_params2D(class_data[im])
ali_params.extend([alpha, sx, sy, mirror])
all_ali_params.append(ali_params)
#aveList[i] = avet
#print " G ",myid," ",time()-st
del ali_params
# We determine the stability of this group here.
# stable_set contains all particles deemed stable, it is a list of list
# each list has two elements, the first is the pixel error, the second is the image number
# stable_set is sorted based on pixel error
#from utilities import write_text_file
#write_text_file(all_ali_params, "all_ali_params%03d.txt"%myid)
stable_set, mir_stab_rate, average_pix_err = multi_align_stability(
all_ali_params, 0.0, 10000.0, thld_err, False, 2 * radius + 1)
#print " H ",myid," ",time()-st
if (len(stable_set) > 5):
stable_set_id = []
members = []
pix_err = []
# First put the stable members into attr 'members' and 'pix_err'
for s in stable_set:
# s[1] - number in this subset
stable_set_id.append(s[1])
# the original image number
members.append(proj_list[i][s[1]])
pix_err.append(s[0])
# Then put the unstable members into attr 'members' and 'pix_err'
from sp_fundamentals import rot_shift2D
avet.to_zero()
if options.grouping == "GRP":
aphi = 0.0
atht = 0.0
vphi = 0.0
vtht = 0.0
l = -1
for j in range(len(proj_list[i])):
# Here it will only work if stable_set_id is sorted in the increasing number, see how l progresses
if j in stable_set_id:
l += 1
avet += rot_shift2D(class_data[j], stable_set[l][2][0],
stable_set[l][2][1],
stable_set[l][2][2],
stable_set[l][2][3])
if options.grouping == "GRP":
phi, theta, psi, sxs, sy_s = get_params_proj(
class_data[j])
if (theta > 90.0):
phi = (phi + 540.0) % 360.0
theta = 180.0 - theta
aphi += phi
atht += theta
vphi += phi * phi
vtht += theta * theta
else:
members.append(proj_list[i][j])
pix_err.append(99999.99)
aveList[i] = avet.copy()
if l > 1:
l += 1
aveList[i] /= l
if options.grouping == "GRP":
aphi /= l
atht /= l
vphi = (vphi - l * aphi * aphi) / l
vtht = (vtht - l * atht * atht) / l
from math import sqrt
refprojdir[i] = [
aphi, atht,
(sqrt(max(vphi, 0.0)) + sqrt(max(vtht, 0.0))) / 2.0
]
# Here more information has to be stored, PARTICULARLY WHAT IS THE REFERENCE DIRECTION
aveList[i].set_attr('members', members)
aveList[i].set_attr('refprojdir', refprojdir[i])
aveList[i].set_attr('pixerr', pix_err)
else:
sxprint(" empty group ", i, refprojdir[i])
aveList[i].set_attr('members', [-1])
aveList[i].set_attr('refprojdir', refprojdir[i])
aveList[i].set_attr('pixerr', [99999.])
del class_data
if myid == main_node:
km = 0
for i in range(number_of_proc):
if i == main_node:
for im in range(len(aveList)):
aveList[im].write_image(args[1], km)
km += 1
else:
nl = mpi_recv(1, MPI_INT, i, SPARX_MPI_TAG_UNIVERSAL,
MPI_COMM_WORLD)
nl = int(nl[0])
for im in range(nl):
ave = recv_EMData(i, im + i + 70000)
nm = mpi_recv(1, MPI_INT, i, SPARX_MPI_TAG_UNIVERSAL,
MPI_COMM_WORLD)
nm = int(nm[0])
members = mpi_recv(nm, MPI_INT, i, SPARX_MPI_TAG_UNIVERSAL,
MPI_COMM_WORLD)
ave.set_attr('members', list(map(int, members)))
members = mpi_recv(nm, MPI_FLOAT, i,
SPARX_MPI_TAG_UNIVERSAL, MPI_COMM_WORLD)
ave.set_attr('pixerr', list(map(float, members)))
members = mpi_recv(3, MPI_FLOAT, i,
SPARX_MPI_TAG_UNIVERSAL, MPI_COMM_WORLD)
ave.set_attr('refprojdir', list(map(float, members)))
ave.write_image(args[1], km)
km += 1
else:
mpi_send(len(aveList), 1, MPI_INT, main_node, SPARX_MPI_TAG_UNIVERSAL,
MPI_COMM_WORLD)
for im in range(len(aveList)):
send_EMData(aveList[im], main_node, im + myid + 70000)
members = aveList[im].get_attr('members')
mpi_send(len(members), 1, MPI_INT, main_node,
SPARX_MPI_TAG_UNIVERSAL, MPI_COMM_WORLD)
mpi_send(members, len(members), MPI_INT, main_node,
SPARX_MPI_TAG_UNIVERSAL, MPI_COMM_WORLD)
members = aveList[im].get_attr('pixerr')
mpi_send(members, len(members), MPI_FLOAT, main_node,
SPARX_MPI_TAG_UNIVERSAL, MPI_COMM_WORLD)
try:
members = aveList[im].get_attr('refprojdir')
mpi_send(members, 3, MPI_FLOAT, main_node,
SPARX_MPI_TAG_UNIVERSAL, MPI_COMM_WORLD)
except:
mpi_send([-999.0, -999.0, -999.0], 3, MPI_FLOAT, main_node,
SPARX_MPI_TAG_UNIVERSAL, MPI_COMM_WORLD)
sp_global_def.BATCH = False
mpi_barrier(MPI_COMM_WORLD)
def main():
def params_3D_2D_NEW(phi, theta, psi, s2x, s2y, mirror):
# the final ali2d parameters already combine shifts operation first and rotation operation second for parameters converted from 3D
if mirror:
m = 1
alpha, sx, sy, scalen = compose_transform2(0, s2x, s2y, 1.0,
540.0 - psi, 0, 0, 1.0)
else:
m = 0
alpha, sx, sy, scalen = compose_transform2(0, s2x, s2y, 1.0,
360.0 - psi, 0, 0, 1.0)
return alpha, sx, sy, m
progname = os.path.basename(sys.argv[0])
usage = progname + " prj_stack --ave2D= --var2D= --ave3D= --var3D= --img_per_grp= --fl= --aa= --sym=symmetry --CTF"
parser = OptionParser(usage, version=SPARXVERSION)
parser.add_option("--output_dir",
type="string",
default="./",
help="Output directory")
parser.add_option("--ave2D",
type="string",
default=False,
help="Write to the disk a stack of 2D averages")
parser.add_option("--var2D",
type="string",
default=False,
help="Write to the disk a stack of 2D variances")
parser.add_option("--ave3D",
type="string",
default=False,
help="Write to the disk reconstructed 3D average")
parser.add_option("--var3D",
type="string",
default=False,
help="Compute 3D variability (time consuming!)")
parser.add_option(
"--img_per_grp",
type="int",
default=100,
help="Number of neighbouring projections.(Default is 100)")
parser.add_option(
"--no_norm",
action="store_true",
default=False,
help="Do not use normalization.(Default is to apply normalization)")
#parser.add_option("--radius", type="int" , default=-1 , help="radius for 3D variability" )
parser.add_option(
"--npad",
type="int",
default=2,
help=
"Number of time to pad the original images.(Default is 2 times padding)"
)
parser.add_option("--sym",
type="string",
default="c1",
help="Symmetry. (Default is no symmetry)")
parser.add_option(
"--fl",
type="float",
default=0.0,
help=
"Low pass filter cutoff in absolute frequency (0.0 - 0.5) and is applied to decimated images. (Default - no filtration)"
)
parser.add_option(
"--aa",
type="float",
default=0.02,
help=
"Fall off of the filter. Use default value if user has no clue about falloff (Default value is 0.02)"
)
parser.add_option("--CTF",
action="store_true",
default=False,
help="Use CFT correction.(Default is no CTF correction)")
#parser.add_option("--MPI" , action="store_true", default=False, help="use MPI version")
#parser.add_option("--radiuspca", type="int" , default=-1 , help="radius for PCA" )
#parser.add_option("--iter", type="int" , default=40 , help="maximum number of iterations (stop criterion of reconstruction process)" )
#parser.add_option("--abs", type="float" , default=0.0 , help="minimum average absolute change of voxels' values (stop criterion of reconstruction process)" )
#parser.add_option("--squ", type="float" , default=0.0 , help="minimum average squared change of voxels' values (stop criterion of reconstruction process)" )
parser.add_option(
"--VAR",
action="store_true",
default=False,
help="Stack of input consists of 2D variances (Default False)")
parser.add_option(
"--decimate",
type="float",
default=0.25,
help="Image decimate rate, a number less than 1. (Default is 0.25)")
parser.add_option(
"--window",
type="int",
default=0,
help=
"Target image size relative to original image size. (Default value is zero.)"
)
#parser.add_option("--SND", action="store_true", default=False, help="compute squared normalized differences (Default False)")
#parser.add_option("--nvec", type="int" , default=0 , help="Number of eigenvectors, (Default = 0 meaning no PCA calculated)")
parser.add_option(
"--symmetrize",
action="store_true",
default=False,
help="Prepare input stack for handling symmetry (Default False)")
parser.add_option("--overhead",
type="float",
default=0.5,
help="python overhead per CPU.")
(options, args) = parser.parse_args()
#####
from mpi import mpi_comm_rank, mpi_comm_size, mpi_recv, MPI_COMM_WORLD
from mpi import mpi_barrier, mpi_reduce, mpi_bcast, mpi_send, MPI_FLOAT, MPI_SUM, MPI_INT, MPI_MAX
#from mpi import *
from sp_applications import MPI_start_end
from sp_reconstruction import recons3d_em, recons3d_em_MPI
from sp_reconstruction import recons3d_4nn_MPI, recons3d_4nn_ctf_MPI
from sp_utilities import print_begin_msg, print_end_msg, print_msg
from sp_utilities import read_text_row, get_image, get_im, wrap_mpi_send, wrap_mpi_recv
from sp_utilities import bcast_EMData_to_all, bcast_number_to_all
from sp_utilities import get_symt
# This is code for handling symmetries by the above program. To be incorporated. PAP 01/27/2015
from EMAN2db import db_open_dict
# Set up global variables related to bdb cache
if sp_global_def.CACHE_DISABLE:
from sp_utilities import disable_bdb_cache
disable_bdb_cache()
# Set up global variables related to ERROR function
sp_global_def.BATCH = True
# detect if program is running under MPI
RUNNING_UNDER_MPI = "OMPI_COMM_WORLD_SIZE" in os.environ
if RUNNING_UNDER_MPI: sp_global_def.MPI = True
if options.output_dir == "./":
current_output_dir = os.path.abspath(options.output_dir)
else:
current_output_dir = options.output_dir
if options.symmetrize:
if mpi.mpi_comm_size(MPI_COMM_WORLD) > 1:
ERROR("Cannot use more than one CPU for symmetry preparation")
if not os.path.exists(current_output_dir):
os.makedirs(current_output_dir)
sp_global_def.write_command(current_output_dir)
from sp_logger import Logger, BaseLogger_Files
if os.path.exists(os.path.join(current_output_dir, "log.txt")):
os.remove(os.path.join(current_output_dir, "log.txt"))
log_main = Logger(BaseLogger_Files())
log_main.prefix = os.path.join(current_output_dir, "./")
instack = args[0]
sym = options.sym.lower()
if (sym == "c1"):
ERROR("There is no need to symmetrize stack for C1 symmetry")
line = ""
for a in sys.argv:
line += " " + a
log_main.add(line)
if (instack[:4] != "bdb:"):
#if output_dir =="./": stack = "bdb:data"
stack = "bdb:" + current_output_dir + "/data"
delete_bdb(stack)
junk = cmdexecute("sp_cpy.py " + instack + " " + stack)
else:
stack = instack
qt = EMUtil.get_all_attributes(stack, 'xform.projection')
na = len(qt)
ts = get_symt(sym)
ks = len(ts)
angsa = [None] * na
for k in range(ks):
#Qfile = "Q%1d"%k
#if options.output_dir!="./": Qfile = os.path.join(options.output_dir,"Q%1d"%k)
Qfile = os.path.join(current_output_dir, "Q%1d" % k)
#delete_bdb("bdb:Q%1d"%k)
delete_bdb("bdb:" + Qfile)
#junk = cmdexecute("e2bdb.py "+stack+" --makevstack=bdb:Q%1d"%k)
junk = cmdexecute("e2bdb.py " + stack + " --makevstack=bdb:" +
Qfile)
#DB = db_open_dict("bdb:Q%1d"%k)
DB = db_open_dict("bdb:" + Qfile)
for i in range(na):
ut = qt[i] * ts[k]
DB.set_attr(i, "xform.projection", ut)
#bt = ut.get_params("spider")
#angsa[i] = [round(bt["phi"],3)%360.0, round(bt["theta"],3)%360.0, bt["psi"], -bt["tx"], -bt["ty"]]
#write_text_row(angsa, 'ptsma%1d.txt'%k)
#junk = cmdexecute("e2bdb.py "+stack+" --makevstack=bdb:Q%1d"%k)
#junk = cmdexecute("sxheader.py bdb:Q%1d --params=xform.projection --import=ptsma%1d.txt"%(k,k))
DB.close()
#if options.output_dir =="./": delete_bdb("bdb:sdata")
delete_bdb("bdb:" + current_output_dir + "/" + "sdata")
#junk = cmdexecute("e2bdb.py . --makevstack=bdb:sdata --filt=Q")
sdata = "bdb:" + current_output_dir + "/" + "sdata"
sxprint(sdata)
junk = cmdexecute("e2bdb.py " + current_output_dir +
" --makevstack=" + sdata + " --filt=Q")
#junk = cmdexecute("ls EMAN2DB/sdata*")
#a = get_im("bdb:sdata")
a = get_im(sdata)
a.set_attr("variabilitysymmetry", sym)
#a.write_image("bdb:sdata")
a.write_image(sdata)
else:
from sp_fundamentals import window2d
myid = mpi_comm_rank(MPI_COMM_WORLD)
number_of_proc = mpi_comm_size(MPI_COMM_WORLD)
main_node = 0
shared_comm = mpi_comm_split_type(MPI_COMM_WORLD, MPI_COMM_TYPE_SHARED,
0, MPI_INFO_NULL)
myid_on_node = mpi_comm_rank(shared_comm)
no_of_processes_per_group = mpi_comm_size(shared_comm)
masters_from_groups_vs_everything_else_comm = mpi_comm_split(
MPI_COMM_WORLD, main_node == myid_on_node, myid_on_node)
color, no_of_groups, balanced_processor_load_on_nodes = get_colors_and_subsets(main_node, MPI_COMM_WORLD, myid, \
shared_comm, myid_on_node, masters_from_groups_vs_everything_else_comm)
overhead_loading = options.overhead * number_of_proc
#memory_per_node = options.memory_per_node
#if memory_per_node == -1.: memory_per_node = 2.*no_of_processes_per_group
keepgoing = 1
current_window = options.window
current_decimate = options.decimate
if len(args) == 1: stack = args[0]
else:
sxprint("Usage: " + usage)
sxprint("Please run \'" + progname + " -h\' for detailed options")
ERROR(
"Invalid number of parameters used. Please see usage information above."
)
return
t0 = time()
# obsolete flags
options.MPI = True
#options.nvec = 0
options.radiuspca = -1
options.iter = 40
options.abs = 0.0
options.squ = 0.0
if options.fl > 0.0 and options.aa == 0.0:
ERROR("Fall off has to be given for the low-pass filter",
myid=myid)
#if options.VAR and options.SND:
# ERROR( "Only one of var and SND can be set!",myid=myid )
if options.VAR and (options.ave2D or options.ave3D or options.var2D):
ERROR(
"When VAR is set, the program cannot output ave2D, ave3D or var2D",
myid=myid)
#if options.SND and (options.ave2D or options.ave3D):
# ERROR( "When SND is set, the program cannot output ave2D or ave3D", myid=myid )
#if options.nvec > 0 :
# ERROR( "PCA option not implemented", myid=myid )
#if options.nvec > 0 and options.ave3D == None:
# ERROR( "When doing PCA analysis, one must set ave3D", myid=myid )
if current_decimate > 1.0 or current_decimate < 0.0:
ERROR("Decimate rate should be a value between 0.0 and 1.0",
myid=myid)
if current_window < 0.0:
ERROR("Target window size should be always larger than zero",
myid=myid)
if myid == main_node:
img = get_image(stack, 0)
nx = img.get_xsize()
ny = img.get_ysize()
if (min(nx, ny) < current_window): keepgoing = 0
keepgoing = bcast_number_to_all(keepgoing, main_node, MPI_COMM_WORLD)
if keepgoing == 0:
ERROR(
"The target window size cannot be larger than the size of decimated image",
myid=myid)
import string
options.sym = options.sym.lower()
# if global_def.CACHE_DISABLE:
# from utilities import disable_bdb_cache
# disable_bdb_cache()
# global_def.BATCH = True
if myid == main_node:
if not os.path.exists(current_output_dir):
os.makedirs(current_output_dir
) # Never delete output_dir in the program!
img_per_grp = options.img_per_grp
#nvec = options.nvec
radiuspca = options.radiuspca
from sp_logger import Logger, BaseLogger_Files
#if os.path.exists(os.path.join(options.output_dir, "log.txt")): os.remove(os.path.join(options.output_dir, "log.txt"))
log_main = Logger(BaseLogger_Files())
log_main.prefix = os.path.join(current_output_dir, "./")
if myid == main_node:
line = ""
for a in sys.argv:
line += " " + a
log_main.add(line)
log_main.add("-------->>>Settings given by all options<<<-------")
log_main.add("Symmetry : %s" % options.sym)
log_main.add("Input stack : %s" % stack)
log_main.add("Output_dir : %s" % current_output_dir)
if options.ave3D:
log_main.add("Ave3d : %s" % options.ave3D)
if options.var3D:
log_main.add("Var3d : %s" % options.var3D)
if options.ave2D:
log_main.add("Ave2D : %s" % options.ave2D)
if options.var2D:
log_main.add("Var2D : %s" % options.var2D)
if options.VAR: log_main.add("VAR : True")
else: log_main.add("VAR : False")
if options.CTF: log_main.add("CTF correction : True ")
else: log_main.add("CTF correction : False ")
log_main.add("Image per group : %5d" % options.img_per_grp)
log_main.add("Image decimate rate : %4.3f" % current_decimate)
log_main.add("Low pass filter : %4.3f" % options.fl)
current_fl = options.fl
if current_fl == 0.0: current_fl = 0.5
log_main.add(
"Current low pass filter is equivalent to cutoff frequency %4.3f for original image size"
% round((current_fl * current_decimate), 3))
log_main.add("Window size : %5d " % current_window)
log_main.add("sx3dvariability begins")
symbaselen = 0
if myid == main_node:
nima = EMUtil.get_image_count(stack)
img = get_image(stack)
nx = img.get_xsize()
ny = img.get_ysize()
nnxo = nx
nnyo = ny
if options.sym != "c1":
imgdata = get_im(stack)
try:
i = imgdata.get_attr("variabilitysymmetry").lower()
if (i != options.sym):
ERROR(
"The symmetry provided does not agree with the symmetry of the input stack",
myid=myid)
except:
ERROR(
"Input stack is not prepared for symmetry, please follow instructions",
myid=myid)
from sp_utilities import get_symt
i = len(get_symt(options.sym))
if ((nima / i) * i != nima):
ERROR(
"The length of the input stack is incorrect for symmetry processing",
myid=myid)
symbaselen = nima / i
else:
symbaselen = nima
else:
nima = 0
nx = 0
ny = 0
nnxo = 0
nnyo = 0
nima = bcast_number_to_all(nima)
nx = bcast_number_to_all(nx)
ny = bcast_number_to_all(ny)
nnxo = bcast_number_to_all(nnxo)
nnyo = bcast_number_to_all(nnyo)
if current_window > max(nx, ny):
ERROR("Window size is larger than the original image size")
if current_decimate == 1.:
if current_window != 0:
nx = current_window
ny = current_window
else:
if current_window == 0:
nx = int(nx * current_decimate + 0.5)
ny = int(ny * current_decimate + 0.5)
else:
nx = int(current_window * current_decimate + 0.5)
ny = nx
symbaselen = bcast_number_to_all(symbaselen)
# check FFT prime number
from sp_fundamentals import smallprime
is_fft_friendly = (nx == smallprime(nx))
if not is_fft_friendly:
if myid == main_node:
log_main.add(
"The target image size is not a product of small prime numbers"
)
log_main.add("Program adjusts the input settings!")
### two cases
if current_decimate == 1.:
nx = smallprime(nx)
ny = nx
current_window = nx # update
if myid == main_node:
log_main.add("The window size is updated to %d." %
current_window)
else:
if current_window == 0:
nx = smallprime(int(nx * current_decimate + 0.5))
current_decimate = float(nx) / nnxo
ny = nx
if (myid == main_node):
log_main.add("The decimate rate is updated to %f." %
current_decimate)
else:
nx = smallprime(
int(current_window * current_decimate + 0.5))
ny = nx
current_window = int(nx / current_decimate + 0.5)
if (myid == main_node):
log_main.add("The window size is updated to %d." %
current_window)
if myid == main_node:
log_main.add("The target image size is %d" % nx)
if radiuspca == -1: radiuspca = nx / 2 - 2
if myid == main_node:
log_main.add("%-70s: %d\n" % ("Number of projection", nima))
img_begin, img_end = MPI_start_end(nima, number_of_proc, myid)
"""
if options.SND:
from sp_projection import prep_vol, prgs
from sp_statistics import im_diff
from sp_utilities import get_im, model_circle, get_params_proj, set_params_proj
from sp_utilities import get_ctf, generate_ctf
from sp_filter import filt_ctf
imgdata = EMData.read_images(stack, range(img_begin, img_end))
if options.CTF:
vol = recons3d_4nn_ctf_MPI(myid, imgdata, 1.0, symmetry=options.sym, npad=options.npad, xysize=-1, zsize=-1)
else:
vol = recons3d_4nn_MPI(myid, imgdata, symmetry=options.sym, npad=options.npad, xysize=-1, zsize=-1)
bcast_EMData_to_all(vol, myid)
volft, kb = prep_vol(vol)
mask = model_circle(nx/2-2, nx, ny)
varList = []
for i in xrange(img_begin, img_end):
phi, theta, psi, s2x, s2y = get_params_proj(imgdata[i-img_begin])
ref_prj = prgs(volft, kb, [phi, theta, psi, -s2x, -s2y])
if options.CTF:
ctf_params = get_ctf(imgdata[i-img_begin])
ref_prj = filt_ctf(ref_prj, generate_ctf(ctf_params))
diff, A, B = im_diff(ref_prj, imgdata[i-img_begin], mask)
diff2 = diff*diff
set_params_proj(diff2, [phi, theta, psi, s2x, s2y])
varList.append(diff2)
mpi_barrier(MPI_COMM_WORLD)
"""
if options.VAR: # 2D variance images have no shifts
#varList = EMData.read_images(stack, range(img_begin, img_end))
from EMAN2 import Region
for index_of_particle in range(img_begin, img_end):
image = get_im(stack, index_of_proj)
if current_window > 0:
varList.append(
fdecimate(
window2d(image, current_window, current_window),
nx, ny))
else:
varList.append(fdecimate(image, nx, ny))
else:
from sp_utilities import bcast_number_to_all, bcast_list_to_all, send_EMData, recv_EMData
from sp_utilities import set_params_proj, get_params_proj, params_3D_2D, get_params2D, set_params2D, compose_transform2
from sp_utilities import model_blank, nearest_proj, model_circle, write_text_row, wrap_mpi_gatherv
from sp_applications import pca
from sp_statistics import avgvar, avgvar_ctf, ccc
from sp_filter import filt_tanl
from sp_morphology import threshold, square_root
from sp_projection import project, prep_vol, prgs
from sets import Set
from sp_utilities import wrap_mpi_recv, wrap_mpi_bcast, wrap_mpi_send
import numpy as np
if myid == main_node:
t1 = time()
proj_angles = []
aveList = []
tab = EMUtil.get_all_attributes(stack, 'xform.projection')
for i in range(nima):
t = tab[i].get_params('spider')
phi = t['phi']
theta = t['theta']
psi = t['psi']
x = theta
if x > 90.0: x = 180.0 - x
x = x * 10000 + psi
proj_angles.append([x, t['phi'], t['theta'], t['psi'], i])
t2 = time()
log_main.add(
"%-70s: %d\n" %
("Number of neighboring projections", img_per_grp))
log_main.add("...... Finding neighboring projections\n")
log_main.add("Number of images per group: %d" % img_per_grp)
log_main.add("Now grouping projections")
proj_angles.sort()
proj_angles_list = np.full((nima, 4), 0.0, dtype=np.float32)
for i in range(nima):
proj_angles_list[i][0] = proj_angles[i][1]
proj_angles_list[i][1] = proj_angles[i][2]
proj_angles_list[i][2] = proj_angles[i][3]
proj_angles_list[i][3] = proj_angles[i][4]
else:
proj_angles_list = 0
proj_angles_list = wrap_mpi_bcast(proj_angles_list, main_node,
MPI_COMM_WORLD)
proj_angles = []
for i in range(nima):
proj_angles.append([
proj_angles_list[i][0], proj_angles_list[i][1],
proj_angles_list[i][2],
int(proj_angles_list[i][3])
])
del proj_angles_list
proj_list, mirror_list = nearest_proj(proj_angles, img_per_grp,
range(img_begin, img_end))
all_proj = Set()
for im in proj_list:
for jm in im:
all_proj.add(proj_angles[jm][3])
all_proj = list(all_proj)
index = {}
for i in range(len(all_proj)):
index[all_proj[i]] = i
mpi_barrier(MPI_COMM_WORLD)
if myid == main_node:
log_main.add("%-70s: %.2f\n" %
("Finding neighboring projections lasted [s]",
time() - t2))
log_main.add("%-70s: %d\n" %
("Number of groups processed on the main node",
len(proj_list)))
log_main.add("Grouping projections took: %12.1f [m]" %
((time() - t2) / 60.))
log_main.add("Number of groups on main node: ", len(proj_list))
mpi_barrier(MPI_COMM_WORLD)
if myid == main_node:
log_main.add("...... Calculating the stack of 2D variances \n")
# Memory estimation. There are two memory consumption peaks
# peak 1. Compute ave, var;
# peak 2. Var volume reconstruction;
# proj_params = [0.0]*(nima*5)
aveList = []
varList = []
#if nvec > 0: eigList = [[] for i in range(nvec)]
dnumber = len(
all_proj) # all neighborhood set for assigned to myid
pnumber = len(proj_list) * 2. + img_per_grp # aveList and varList
tnumber = dnumber + pnumber
vol_size2 = nx**3 * 4. * 8 / 1.e9
vol_size1 = 2. * nnxo**3 * 4. * 8 / 1.e9
proj_size = nnxo * nnyo * len(
proj_list) * 4. * 2. / 1.e9 # both aveList and varList
orig_data_size = nnxo * nnyo * 4. * tnumber / 1.e9
reduced_data_size = nx * nx * 4. * tnumber / 1.e9
full_data = np.full((number_of_proc, 2), -1., dtype=np.float16)
full_data[myid] = orig_data_size, reduced_data_size
if myid != main_node:
wrap_mpi_send(full_data, main_node, MPI_COMM_WORLD)
if myid == main_node:
for iproc in range(number_of_proc):
if iproc != main_node:
dummy = wrap_mpi_recv(iproc, MPI_COMM_WORLD)
full_data[np.where(dummy > -1)] = dummy[np.where(
dummy > -1)]
del dummy
mpi_barrier(MPI_COMM_WORLD)
full_data = wrap_mpi_bcast(full_data, main_node, MPI_COMM_WORLD)
# find the CPU with heaviest load
minindx = np.argsort(full_data, 0)
heavy_load_myid = minindx[-1][1]
total_mem = sum(full_data)
if myid == main_node:
if current_window == 0:
log_main.add(
"Nx: current image size = %d. Decimated by %f from %d"
% (nx, current_decimate, nnxo))
else:
log_main.add(
"Nx: current image size = %d. Windowed to %d, and decimated by %f from %d"
% (nx, current_window, current_decimate, nnxo))
log_main.add("Nproj: number of particle images.")
log_main.add("Navg: number of 2D average images.")
log_main.add("Nvar: number of 2D variance images.")
log_main.add(
"Img_per_grp: user defined image per group for averaging = %d"
% img_per_grp)
log_main.add(
"Overhead: total python overhead memory consumption = %f"
% overhead_loading)
log_main.add("Total memory) = 4.0*nx^2*(nproj + navg +nvar+ img_per_grp)/1.0e9 + overhead: %12.3f [GB]"%\
(total_mem[1] + overhead_loading))
del full_data
mpi_barrier(MPI_COMM_WORLD)
if myid == heavy_load_myid:
log_main.add(
"Begin reading and preprocessing images on processor. Wait... "
)
ttt = time()
#imgdata = EMData.read_images(stack, all_proj)
imgdata = [None for im in range(len(all_proj))]
for index_of_proj in range(len(all_proj)):
#image = get_im(stack, all_proj[index_of_proj])
if (current_window > 0):
imgdata[index_of_proj] = fdecimate(
window2d(get_im(stack, all_proj[index_of_proj]),
current_window, current_window), nx, ny)
else:
imgdata[index_of_proj] = fdecimate(
get_im(stack, all_proj[index_of_proj]), nx, ny)
if (current_decimate > 0.0 and options.CTF):
ctf = imgdata[index_of_proj].get_attr("ctf")
ctf.apix = ctf.apix / current_decimate
imgdata[index_of_proj].set_attr("ctf", ctf)
if myid == heavy_load_myid and index_of_proj % 100 == 0:
log_main.add(" ...... %6.2f%% " %
(index_of_proj / float(len(all_proj)) * 100.))
mpi_barrier(MPI_COMM_WORLD)
if myid == heavy_load_myid:
log_main.add("All_proj preprocessing cost %7.2f m" %
((time() - ttt) / 60.))
log_main.add("Wait untill reading on all CPUs done...")
'''
imgdata2 = EMData.read_images(stack, range(img_begin, img_end))
if options.fl > 0.0:
for k in xrange(len(imgdata2)):
imgdata2[k] = filt_tanl(imgdata2[k], options.fl, options.aa)
if options.CTF:
vol = recons3d_4nn_ctf_MPI(myid, imgdata2, 1.0, symmetry=options.sym, npad=options.npad, xysize=-1, zsize=-1)
else:
vol = recons3d_4nn_MPI(myid, imgdata2, symmetry=options.sym, npad=options.npad, xysize=-1, zsize=-1)
if myid == main_node:
vol.write_image("vol_ctf.hdf")
print_msg("Writing to the disk volume reconstructed from averages as : %s\n"%("vol_ctf.hdf"))
del vol, imgdata2
mpi_barrier(MPI_COMM_WORLD)
'''
from sp_applications import prepare_2d_forPCA
from sp_utilities import model_blank
from EMAN2 import Transform
if not options.no_norm:
mask = model_circle(nx / 2 - 2, nx, nx)
if options.CTF:
from sp_utilities import pad
from sp_filter import filt_ctf
from sp_filter import filt_tanl
if myid == heavy_load_myid:
log_main.add("Start computing 2D aveList and varList. Wait...")
ttt = time()
inner = nx // 2 - 4
outer = inner + 2
xform_proj_for_2D = [None for i in range(len(proj_list))]
for i in range(len(proj_list)):
ki = proj_angles[proj_list[i][0]][3]
if ki >= symbaselen: continue
mi = index[ki]
dpar = Util.get_transform_params(imgdata[mi],
"xform.projection", "spider")
phiM, thetaM, psiM, s2xM, s2yM = dpar["phi"], dpar[
"theta"], dpar[
"psi"], -dpar["tx"] * current_decimate, -dpar[
"ty"] * current_decimate
grp_imgdata = []
for j in range(img_per_grp):
mj = index[proj_angles[proj_list[i][j]][3]]
cpar = Util.get_transform_params(imgdata[mj],
"xform.projection",
"spider")
alpha, sx, sy, mirror = params_3D_2D_NEW(
cpar["phi"], cpar["theta"], cpar["psi"],
-cpar["tx"] * current_decimate,
-cpar["ty"] * current_decimate, mirror_list[i][j])
if thetaM <= 90:
if mirror == 0:
alpha, sx, sy, scale = compose_transform2(
alpha, sx, sy, 1.0, phiM - cpar["phi"], 0.0,
0.0, 1.0)
else:
alpha, sx, sy, scale = compose_transform2(
alpha, sx, sy, 1.0, 180 - (phiM - cpar["phi"]),
0.0, 0.0, 1.0)
else:
if mirror == 0:
alpha, sx, sy, scale = compose_transform2(
alpha, sx, sy, 1.0, -(phiM - cpar["phi"]), 0.0,
0.0, 1.0)
else:
alpha, sx, sy, scale = compose_transform2(
alpha, sx, sy, 1.0,
-(180 - (phiM - cpar["phi"])), 0.0, 0.0, 1.0)
imgdata[mj].set_attr(
"xform.align2d",
Transform({
"type": "2D",
"alpha": alpha,
"tx": sx,
"ty": sy,
"mirror": mirror,
"scale": 1.0
}))
grp_imgdata.append(imgdata[mj])
if not options.no_norm:
for k in range(img_per_grp):
ave, std, minn, maxx = Util.infomask(
grp_imgdata[k], mask, False)
grp_imgdata[k] -= ave
grp_imgdata[k] /= std
if options.fl > 0.0:
for k in range(img_per_grp):
grp_imgdata[k] = filt_tanl(grp_imgdata[k], options.fl,
options.aa)
# Because of background issues, only linear option works.
if options.CTF:
ave, var = aves_wiener(grp_imgdata,
SNR=1.0e5,
interpolation_method="linear")
else:
ave, var = ave_var(grp_imgdata)
# Switch to std dev
# threshold is not really needed,it is just in case due to numerical accuracy something turns out negative.
var = square_root(threshold(var))
set_params_proj(ave, [phiM, thetaM, 0.0, 0.0, 0.0])
set_params_proj(var, [phiM, thetaM, 0.0, 0.0, 0.0])
aveList.append(ave)
varList.append(var)
xform_proj_for_2D[i] = [phiM, thetaM, 0.0, 0.0, 0.0]
'''
if nvec > 0:
eig = pca(input_stacks=grp_imgdata, subavg="", mask_radius=radiuspca, nvec=nvec, incore=True, shuffle=False, genbuf=True)
for k in range(nvec):
set_params_proj(eig[k], [phiM, thetaM, 0.0, 0.0, 0.0])
eigList[k].append(eig[k])
"""
if myid == 0 and i == 0:
for k in xrange(nvec):
eig[k].write_image("eig.hdf", k)
"""
'''
if (myid == heavy_load_myid) and (i % 100 == 0):
log_main.add(" ......%6.2f%% " %
(i / float(len(proj_list)) * 100.))
del imgdata, grp_imgdata, cpar, dpar, all_proj, proj_angles, index
if not options.no_norm: del mask
if myid == main_node: del tab
# At this point, all averages and variances are computed
mpi_barrier(MPI_COMM_WORLD)
if (myid == heavy_load_myid):
log_main.add("Computing aveList and varList took %12.1f [m]" %
((time() - ttt) / 60.))
xform_proj_for_2D = wrap_mpi_gatherv(xform_proj_for_2D, main_node,
MPI_COMM_WORLD)
if (myid == main_node):
write_text_row([str(entry) for entry in xform_proj_for_2D],
os.path.join(current_output_dir, "params.txt"))
del xform_proj_for_2D
mpi_barrier(MPI_COMM_WORLD)
if options.ave2D:
from sp_fundamentals import fpol
from sp_applications import header
if myid == main_node:
log_main.add("Compute ave2D ... ")
km = 0
for i in range(number_of_proc):
if i == main_node:
for im in range(len(aveList)):
aveList[im].write_image(
os.path.join(current_output_dir,
options.ave2D), km)
km += 1
else:
nl = mpi_recv(1, MPI_INT, i,
SPARX_MPI_TAG_UNIVERSAL,
MPI_COMM_WORLD)
nl = int(nl[0])
for im in range(nl):
ave = recv_EMData(i, im + i + 70000)
"""
nm = mpi_recv(1, MPI_INT, i, SPARX_MPI_TAG_UNIVERSAL, MPI_COMM_WORLD)
nm = int(nm[0])
members = mpi_recv(nm, MPI_INT, i, SPARX_MPI_TAG_UNIVERSAL, MPI_COMM_WORLD)
ave.set_attr('members', map(int, members))
members = mpi_recv(nm, MPI_FLOAT, i, SPARX_MPI_TAG_UNIVERSAL, MPI_COMM_WORLD)
ave.set_attr('pix_err', map(float, members))
members = mpi_recv(3, MPI_FLOAT, i, SPARX_MPI_TAG_UNIVERSAL, MPI_COMM_WORLD)
ave.set_attr('refprojdir', map(float, members))
"""
tmpvol = fpol(ave, nx, nx, 1)
tmpvol.write_image(
os.path.join(current_output_dir,
options.ave2D), km)
km += 1
else:
mpi_send(len(aveList), 1, MPI_INT, main_node,
SPARX_MPI_TAG_UNIVERSAL, MPI_COMM_WORLD)
for im in range(len(aveList)):
send_EMData(aveList[im], main_node, im + myid + 70000)
"""
members = aveList[im].get_attr('members')
mpi_send(len(members), 1, MPI_INT, main_node, SPARX_MPI_TAG_UNIVERSAL, MPI_COMM_WORLD)
mpi_send(members, len(members), MPI_INT, main_node, SPARX_MPI_TAG_UNIVERSAL, MPI_COMM_WORLD)
members = aveList[im].get_attr('pix_err')
mpi_send(members, len(members), MPI_FLOAT, main_node, SPARX_MPI_TAG_UNIVERSAL, MPI_COMM_WORLD)
try:
members = aveList[im].get_attr('refprojdir')
mpi_send(members, 3, MPI_FLOAT, main_node, SPARX_MPI_TAG_UNIVERSAL, MPI_COMM_WORLD)
except:
mpi_send([-999.0,-999.0,-999.0], 3, MPI_FLOAT, main_node, SPARX_MPI_TAG_UNIVERSAL, MPI_COMM_WORLD)
"""
if myid == main_node:
header(os.path.join(current_output_dir, options.ave2D),
params='xform.projection',
fimport=os.path.join(current_output_dir,
"params.txt"))
mpi_barrier(MPI_COMM_WORLD)
if options.ave3D:
from sp_fundamentals import fpol
t5 = time()
if myid == main_node: log_main.add("Reconstruct ave3D ... ")
ave3D = recons3d_4nn_MPI(myid,
aveList,
symmetry=options.sym,
npad=options.npad)
bcast_EMData_to_all(ave3D, myid)
if myid == main_node:
if current_decimate != 1.0:
ave3D = resample(ave3D, 1. / current_decimate)
ave3D = fpol(ave3D, nnxo, nnxo,
nnxo) # always to the orignal image size
set_pixel_size(ave3D, 1.0)
ave3D.write_image(
os.path.join(current_output_dir, options.ave3D))
log_main.add("Ave3D reconstruction took %12.1f [m]" %
((time() - t5) / 60.0))
log_main.add("%-70s: %s\n" %
("The reconstructed ave3D is saved as ",
options.ave3D))
mpi_barrier(MPI_COMM_WORLD)
del ave, var, proj_list, stack, alpha, sx, sy, mirror, aveList
'''
if nvec > 0:
for k in range(nvec):
if myid == main_node:log_main.add("Reconstruction eigenvolumes", k)
cont = True
ITER = 0
mask2d = model_circle(radiuspca, nx, nx)
while cont:
#print "On node %d, iteration %d"%(myid, ITER)
eig3D = recons3d_4nn_MPI(myid, eigList[k], symmetry=options.sym, npad=options.npad)
bcast_EMData_to_all(eig3D, myid, main_node)
if options.fl > 0.0:
eig3D = filt_tanl(eig3D, options.fl, options.aa)
if myid == main_node:
eig3D.write_image(os.path.join(options.outpout_dir, "eig3d_%03d.hdf"%(k, ITER)))
Util.mul_img( eig3D, model_circle(radiuspca, nx, nx, nx) )
eig3Df, kb = prep_vol(eig3D)
del eig3D
cont = False
icont = 0
for l in range(len(eigList[k])):
phi, theta, psi, s2x, s2y = get_params_proj(eigList[k][l])
proj = prgs(eig3Df, kb, [phi, theta, psi, s2x, s2y])
cl = ccc(proj, eigList[k][l], mask2d)
if cl < 0.0:
icont += 1
cont = True
eigList[k][l] *= -1.0
u = int(cont)
u = mpi_reduce([u], 1, MPI_INT, MPI_MAX, main_node, MPI_COMM_WORLD)
icont = mpi_reduce([icont], 1, MPI_INT, MPI_SUM, main_node, MPI_COMM_WORLD)
if myid == main_node:
u = int(u[0])
log_main.add(" Eigenvector: ",k," number changed ",int(icont[0]))
else: u = 0
u = bcast_number_to_all(u, main_node)
cont = bool(u)
ITER += 1
del eig3Df, kb
mpi_barrier(MPI_COMM_WORLD)
del eigList, mask2d
'''
if options.ave3D: del ave3D
if options.var2D:
from sp_fundamentals import fpol
from sp_applications import header
if myid == main_node:
log_main.add("Compute var2D...")
km = 0
for i in range(number_of_proc):
if i == main_node:
for im in range(len(varList)):
tmpvol = fpol(varList[im], nx, nx, 1)
tmpvol.write_image(
os.path.join(current_output_dir,
options.var2D), km)
km += 1
else:
nl = mpi_recv(1, MPI_INT, i,
SPARX_MPI_TAG_UNIVERSAL,
MPI_COMM_WORLD)
nl = int(nl[0])
for im in range(nl):
ave = recv_EMData(i, im + i + 70000)
tmpvol = fpol(ave, nx, nx, 1)
tmpvol.write_image(
os.path.join(current_output_dir,
options.var2D), km)
km += 1
else:
mpi_send(len(varList), 1, MPI_INT, main_node,
SPARX_MPI_TAG_UNIVERSAL, MPI_COMM_WORLD)
for im in range(len(varList)):
send_EMData(varList[im], main_node, im + myid +
70000) # What with the attributes??
mpi_barrier(MPI_COMM_WORLD)
if myid == main_node:
from sp_applications import header
header(os.path.join(current_output_dir, options.var2D),
params='xform.projection',
fimport=os.path.join(current_output_dir,
"params.txt"))
mpi_barrier(MPI_COMM_WORLD)
if options.var3D:
if myid == main_node: log_main.add("Reconstruct var3D ...")
t6 = time()
# radiusvar = options.radius
# if( radiusvar < 0 ): radiusvar = nx//2 -3
res = recons3d_4nn_MPI(myid,
varList,
symmetry=options.sym,
npad=options.npad)
#res = recons3d_em_MPI(varList, vol_stack, options.iter, radiusvar, options.abs, True, options.sym, options.squ)
if myid == main_node:
from sp_fundamentals import fpol
if current_decimate != 1.0:
res = resample(res, 1. / current_decimate)
res = fpol(res, nnxo, nnxo, nnxo)
set_pixel_size(res, 1.0)
res.write_image(os.path.join(current_output_dir,
options.var3D))
log_main.add(
"%-70s: %s\n" %
("The reconstructed var3D is saved as ", options.var3D))
log_main.add("Var3D reconstruction took %f12.1 [m]" %
((time() - t6) / 60.0))
log_main.add("Total computation time %f12.1 [m]" %
((time() - t0) / 60.0))
log_main.add("sx3dvariability finishes")
if RUNNING_UNDER_MPI:
sp_global_def.MPI = False
sp_global_def.BATCH = False
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")