def main():
arglist = []
for arg in sys.argv:
arglist.append(arg)
progname = os.path.basename(arglist[0])
usage = progname + """ firstvolume secondvolume maskfile outputfile --wn --step --cutoff --radius --fsc --res_overall --out_ang_res --apix --MPI
Compute local resolution in real space within area outlined by the maskfile and within regions wn x wn x wn
"""
parser = optparse.OptionParser(usage, version=global_def.SPARXVERSION)
parser.add_option(
"--wn",
type="int",
default=7,
help=
"Size of window within which local real-space FSC is computed. (default 7)"
)
parser.add_option(
"--step",
type="float",
default=1.0,
help="Shell step in Fourier size in pixels. (default 1.0)")
parser.add_option("--cutoff",
type="float",
default=0.5,
help="Resolution cut-off for FSC. (default 0.5)")
parser.add_option(
"--radius",
type="int",
default=-1,
help=
"If there is no maskfile, sphere with r=radius will be used. By default, the radius is nx/2-wn (default -1)"
)
parser.add_option(
"--fsc",
type="string",
default=None,
help=
"Save overall FSC curve (might be truncated). By default, the program does not save the FSC curve. (default none)"
)
parser.add_option(
"--res_overall",
type="float",
default=-1.0,
help=
"Overall resolution at the cutoff level estimated by the user [abs units]. (default None)"
)
parser.add_option(
"--out_ang_res",
action="store_true",
default=False,
help=
"Additionally creates a local resolution file in Angstroms. (default False)"
)
parser.add_option(
"--apix",
type="float",
default=1.0,
help=
"Pixel size in Angstrom. Effective only with --out_ang_res options. (default 1.0)"
)
parser.add_option("--MPI",
action="store_true",
default=False,
help="Use MPI version.")
(options, args) = parser.parse_args(arglist[1:])
if len(args) < 3 or len(args) > 4:
print("See usage " + usage)
sys.exit()
if global_def.CACHE_DISABLE:
utilities.disable_bdb_cache()
res_overall = options.res_overall
if options.MPI:
sys.argv = mpi.mpi_init(len(sys.argv), sys.argv)
number_of_proc = mpi.mpi_comm_size(mpi.MPI_COMM_WORLD)
myid = mpi.mpi_comm_rank(mpi.MPI_COMM_WORLD)
main_node = 0
global_def.MPI = True
cutoff = options.cutoff
nk = int(options.wn)
if (myid == main_node):
#print sys.argv
vi = utilities.get_im(sys.argv[1])
ui = utilities.get_im(sys.argv[2])
nx = vi.get_xsize()
ny = vi.get_ysize()
nz = vi.get_zsize()
dis = [nx, ny, nz]
else:
dis = [0, 0, 0, 0]
global_def.BATCH = True
dis = 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 = utilities.model_blank(nx, ny, nz)
ui = utilities.model_blank(nx, ny, nz)
if len(args) == 3:
m = utilities.model_circle((min(nx, ny, nz) - nk) // 2, nx, ny, nz)
outvol = args[2]
elif len(args) == 4:
if (myid == main_node):
m = morphology.binarize(utilities.get_im(args[2]), 0.5)
else:
m = utilities.model_blank(nx, ny, nz)
outvol = args[3]
utilities.bcast_EMData_to_all(m, myid, main_node)
"""Multiline Comment0"""
freqvol, resolut = statistics.locres(vi, ui, m, nk, cutoff,
options.step, myid, main_node,
number_of_proc)
if (myid == 0):
# Remove outliers based on the Interquartile range
output_volume(freqvol, resolut, options.apix, outvol, options.fsc,
options.out_ang_res, nx, ny, nz, res_overall)
mpi.mpi_finalize()
else:
cutoff = options.cutoff
vi = utilities.get_im(args[0])
ui = utilities.get_im(args[1])
nn = vi.get_xsize()
nk = int(options.wn)
if len(args) == 3:
m = utilities.model_circle((nn - nk) // 2, nn, nn, nn)
outvol = args[2]
elif len(args) == 4:
m = morphology.binarize(utilities.get_im(args[2]), 0.5)
outvol = args[3]
mc = utilities.model_blank(nn, nn, nn, 1.0) - m
vf = fundamentals.fft(vi)
uf = fundamentals.fft(ui)
"""Multiline Comment1"""
lp = int(nn / 2 / options.step + 0.5)
step = 0.5 / lp
freqvol = utilities.model_blank(nn, nn, nn)
resolut = []
for i in range(1, lp):
fl = step * i
fh = fl + step
#print(lp,i,step,fl,fh)
v = fundamentals.fft(filter.filt_tophatb(vf, fl, fh))
u = fundamentals.fft(filter.filt_tophatb(uf, fl, fh))
tmp1 = EMAN2_cppwrap.Util.muln_img(v, v)
tmp2 = EMAN2_cppwrap.Util.muln_img(u, u)
do = EMAN2_cppwrap.Util.infomask(
morphology.square_root(
morphology.threshold(
EMAN2_cppwrap.Util.muln_img(tmp1, tmp2))), m, True)[0]
tmp3 = EMAN2_cppwrap.Util.muln_img(u, v)
dp = EMAN2_cppwrap.Util.infomask(tmp3, m, True)[0]
resolut.append([i, (fl + fh) / 2.0, dp / do])
tmp1 = EMAN2_cppwrap.Util.box_convolution(tmp1, nk)
tmp2 = EMAN2_cppwrap.Util.box_convolution(tmp2, nk)
tmp3 = EMAN2_cppwrap.Util.box_convolution(tmp3, nk)
EMAN2_cppwrap.Util.mul_img(tmp1, tmp2)
tmp1 = morphology.square_root(morphology.threshold(tmp1))
EMAN2_cppwrap.Util.mul_img(tmp1, m)
EMAN2_cppwrap.Util.add_img(tmp1, mc)
EMAN2_cppwrap.Util.mul_img(tmp3, m)
EMAN2_cppwrap.Util.add_img(tmp3, mc)
EMAN2_cppwrap.Util.div_img(tmp3, tmp1)
EMAN2_cppwrap.Util.mul_img(tmp3, m)
freq = (fl + fh) / 2.0
bailout = True
for x in range(nn):
for y in range(nn):
for z in range(nn):
if (m.get_value_at(x, y, z) > 0.5):
if (freqvol.get_value_at(x, y, z) == 0.0):
if (tmp3.get_value_at(x, y, z) < cutoff):
freqvol.set_value_at(x, y, z, freq)
bailout = False
else:
bailout = False
if (bailout): break
#print(len(resolut))
# remove outliers
output_volume(freqvol, resolut, options.apix, outvol, options.fsc,
options.out_ang_res, nx, ny, nz, res_overall)
def main():
def params_3D_2D_NEW(phi, theta, psi, s2x, s2y, mirror):
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=15. --aa=0.01 --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=10,
help="number of neighbouring projections")
parser.add_option("--no_norm",
action="store_true",
default=False,
help="do not use 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")
parser.add_option("--sym", type="string", default="c1", help="symmetry")
parser.add_option(
"--fl",
type="float",
default=0.0,
help=
"cutoff freqency in absolute frequency (0.0-0.5). (Default - no filtration)"
)
parser.add_option(
"--aa",
type="float",
default=0.0,
help=
"fall off of the filter. Put 0.01 if user has no clue about falloff (Default - no filtration)"
)
parser.add_option("--CTF",
action="store_true",
default=False,
help="use CFT correction")
parser.add_option("--VERBOSE",
action="store_true",
default=False,
help="Long output for debugging")
#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 on input consists of 2D variances (Default False)")
parser.add_option(
"--decimate",
type="float",
default=1.0,
help=
"image decimate rate, a number larger (expand image) or less (shrink image) than 1. default is 1"
)
parser.add_option(
"--window",
type="int",
default=0,
help=
"reduce images to a small image size without changing pixel_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)")
(options, args) = parser.parse_args()
#####
from mpi import mpi_init, 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 applications import MPI_start_end
from reconstruction import recons3d_em, recons3d_em_MPI
from reconstruction import recons3d_4nn_MPI, recons3d_4nn_ctf_MPI
from utilities import print_begin_msg, print_end_msg, print_msg
from utilities import read_text_row, get_image, get_im
from utilities import bcast_EMData_to_all, bcast_number_to_all
from 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 global_def.CACHE_DISABLE:
from utilities import disable_bdb_cache
disable_bdb_cache()
# Set up global variables related to ERROR function
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:
global_def.MPI = True
if options.symmetrize:
if RUNNING_UNDER_MPI:
try:
sys.argv = mpi_init(len(sys.argv), sys.argv)
try:
number_of_proc = mpi_comm_size(MPI_COMM_WORLD)
if (number_of_proc > 1):
ERROR(
"Cannot use more than one CPU for symmetry prepration",
"sx3dvariability", 1)
except:
pass
except:
pass
if options.output_dir != "./" and not os.path.exists(
options.output_dir):
os.mkdir(options.output_dir)
# Input
#instack = "Clean_NORM_CTF_start_wparams.hdf"
#instack = "bdb:data"
from 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(options.output_dir, "./")
instack = args[0]
sym = options.sym.lower()
if (sym == "c1"):
ERROR("There is no need to symmetrize stack for C1 symmetry",
"sx3dvariability", 1)
line = ""
for a in sys.argv:
line += " " + a
log_main.add(line)
if (instack[:4] != "bdb:"):
if output_dir == "./": stack = "bdb:data"
else: stack = "bdb:" + options.output_dir + "/data"
delete_bdb(stack)
junk = cmdexecute("sxcpy.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 xrange(ks):
#Qfile = "Q%1d"%k
if options.output_dir != "./":
Qfile = os.path.join(options.output_dir, "Q%1d" % k)
else:
Qfile = os.path.join(options.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 xrange(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")
else: delete_bdb("bdb:" + options.output_dir + "/" + "sdata")
#junk = cmdexecute("e2bdb.py . --makevstack=bdb:sdata --filt=Q")
sdata = "bdb:" + options.output_dir + "/" + "sdata"
print(sdata)
junk = cmdexecute("e2bdb.py " + options.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:
sys.argv = mpi_init(len(sys.argv), sys.argv)
myid = mpi_comm_rank(MPI_COMM_WORLD)
number_of_proc = mpi_comm_size(MPI_COMM_WORLD)
main_node = 0
if len(args) == 1:
stack = args[0]
else:
print(("usage: " + usage))
print(("Please run '" + progname + " -h' for detailed options"))
return 1
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",
"sx3dvariability", 1, myid)
if options.VAR and options.SND:
ERROR("Only one of var and SND can be set!", "sx3dvariability",
myid)
exit()
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",
"sx3dvariability", 1, myid)
exit()
#if options.SND and (options.ave2D or options.ave3D):
# ERROR("When SND is set, the program cannot output ave2D or ave3D", "sx3dvariability", 1, myid)
# exit()
if options.nvec > 0:
ERROR("PCA option not implemented", "sx3dvariability", 1, myid)
exit()
if options.nvec > 0 and options.ave3D == None:
ERROR("When doing PCA analysis, one must set ave3D",
"sx3dvariability",
myid=myid)
exit()
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 options.output_dir != "./" and not os.path.exists(
options.output_dir):
os.mkdir(options.output_dir)
img_per_grp = options.img_per_grp
nvec = options.nvec
radiuspca = options.radiuspca
from 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(options.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("instack :" + stack)
log_main.add("output_dir :" + options.output_dir)
log_main.add("var3d :" + options.var3D)
if myid == main_node:
line = strftime("%Y-%m-%d_%H:%M:%S", localtime()) + " =>"
#print_begin_msg("sx3dvariability")
msg = "sx3dvariability"
log_main.add(msg)
print(line, msg)
msg = ("%-70s: %s\n" % ("Input stack", stack))
log_main.add(msg)
print(line, msg)
symbaselen = 0
if myid == main_node:
nima = EMUtil.get_image_count(stack)
img = get_image(stack)
nx = img.get_xsize()
ny = img.get_ysize()
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",
"sx3dvariability",
myid=myid)
except:
ERROR(
"Input stack is not prepared for symmetry, please follow instructions",
"sx3dvariability",
myid=myid)
from 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",
"sx3dvariability",
myid=myid)
symbaselen = nima / i
else:
symbaselen = nima
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)
Tracker = {}
Tracker["total_stack"] = nima
if options.decimate == 1.:
if options.window != 0:
nx = options.window
ny = options.window
else:
if options.window == 0:
nx = int(nx * options.decimate)
ny = int(ny * options.decimate)
else:
nx = int(options.window * options.decimate)
ny = nx
Tracker["nx"] = nx
Tracker["ny"] = ny
Tracker["nz"] = nx
symbaselen = bcast_number_to_all(symbaselen)
if radiuspca == -1: radiuspca = nx / 2 - 2
if myid == main_node:
line = strftime("%Y-%m-%d_%H:%M:%S", localtime()) + " =>"
msg = "%-70s: %d\n" % ("Number of projection", nima)
log_main.add(msg)
print(line, msg)
img_begin, img_end = MPI_start_end(nima, number_of_proc, myid)
"""
if options.SND:
from projection import prep_vol, prgs
from statistics import im_diff
from utilities import get_im, model_circle, get_params_proj, set_params_proj
from utilities import get_ctf, generate_ctf
from 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:
#varList = EMData.read_images(stack, range(img_begin, img_end))
varList = []
this_image = EMData()
for index_of_particle in xrange(img_begin, img_end):
this_image.read_image(stack, index_of_particle)
varList.append(
image_decimate_window_xform_ctf(this_image,
options.decimate,
options.window,
options.CTF))
else:
from utilities import bcast_number_to_all, bcast_list_to_all, send_EMData, recv_EMData
from utilities import set_params_proj, get_params_proj, params_3D_2D, get_params2D, set_params2D, compose_transform2
from utilities import model_blank, nearest_proj, model_circle
from applications import pca
from statistics import avgvar, avgvar_ctf, ccc
from filter import filt_tanl
from morphology import threshold, square_root
from projection import project, prep_vol, prgs
from sets import Set
if myid == main_node:
t1 = time()
proj_angles = []
aveList = []
tab = EMUtil.get_all_attributes(stack, 'xform.projection')
for i in xrange(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()
line = strftime("%Y-%m-%d_%H:%M:%S", localtime()) + " =>"
msg = "%-70s: %d\n" % ("Number of neighboring projections",
img_per_grp)
log_main.add(msg)
print(line, msg)
msg = "...... Finding neighboring projections\n"
log_main.add(msg)
print(line, msg)
if options.VERBOSE:
msg = "Number of images per group: %d" % img_per_grp
log_main.add(msg)
print(line, msg)
msg = "Now grouping projections"
log_main.add(msg)
print(line, msg)
proj_angles.sort()
proj_angles_list = [0.0] * (nima * 4)
if myid == main_node:
for i in xrange(nima):
proj_angles_list[i * 4] = proj_angles[i][1]
proj_angles_list[i * 4 + 1] = proj_angles[i][2]
proj_angles_list[i * 4 + 2] = proj_angles[i][3]
proj_angles_list[i * 4 + 3] = proj_angles[i][4]
proj_angles_list = bcast_list_to_all(proj_angles_list, myid,
main_node)
proj_angles = []
for i in xrange(nima):
proj_angles.append([
proj_angles_list[i * 4], proj_angles_list[i * 4 + 1],
proj_angles_list[i * 4 + 2],
int(proj_angles_list[i * 4 + 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)
if options.VERBOSE:
print("On node %2d, number of images needed to be read = %5d" %
(myid, len(all_proj)))
index = {}
for i in xrange(len(all_proj)):
index[all_proj[i]] = i
mpi_barrier(MPI_COMM_WORLD)
if myid == main_node:
line = strftime("%Y-%m-%d_%H:%M:%S", localtime()) + " =>"
msg = ("%-70s: %.2f\n" %
("Finding neighboring projections lasted [s]",
time() - t2))
log_main.add(msg)
print(msg)
msg = ("%-70s: %d\n" %
("Number of groups processed on the main node",
len(proj_list)))
log_main.add(msg)
print(line, msg)
if options.VERBOSE:
print("Grouping projections took: ", (time() - t2) / 60,
"[min]")
print("Number of groups on main node: ", len(proj_list))
mpi_barrier(MPI_COMM_WORLD)
if myid == main_node:
line = strftime("%Y-%m-%d_%H:%M:%S", localtime()) + " =>"
msg = ("...... calculating the stack of 2D variances \n")
log_main.add(msg)
print(line, msg)
if options.VERBOSE:
print("Now calculating the stack of 2D variances")
proj_params = [0.0] * (nima * 5)
aveList = []
varList = []
if nvec > 0:
eigList = [[] for i in xrange(nvec)]
if options.VERBOSE:
print("Begin to read images on processor %d" % (myid))
ttt = time()
#imgdata = EMData.read_images(stack, all_proj)
imgdata = []
for index_of_proj in xrange(len(all_proj)):
#img = EMData()
#img.read_image(stack, all_proj[index_of_proj])
dmg = image_decimate_window_xform_ctf(
get_im(stack, all_proj[index_of_proj]), options.decimate,
options.window, options.CTF)
#print dmg.get_xsize(), "init"
imgdata.append(dmg)
if options.VERBOSE:
print("Reading images on processor %d done, time = %.2f" %
(myid, time() - ttt))
print("On processor %d, we got %d images" %
(myid, len(imgdata)))
mpi_barrier(MPI_COMM_WORLD)
'''
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 applications import prepare_2d_forPCA
from utilities import model_blank
for i in xrange(len(proj_list)):
ki = proj_angles[proj_list[i][0]][3]
if ki >= symbaselen: continue
mi = index[ki]
phiM, thetaM, psiM, s2xM, s2yM = get_params_proj(imgdata[mi])
grp_imgdata = []
for j in xrange(img_per_grp):
mj = index[proj_angles[proj_list[i][j]][3]]
phi, theta, psi, s2x, s2y = get_params_proj(imgdata[mj])
alpha, sx, sy, mirror = params_3D_2D_NEW(
phi, theta, psi, s2x, s2y, mirror_list[i][j])
if thetaM <= 90:
if mirror == 0:
alpha, sx, sy, scale = compose_transform2(
alpha, sx, sy, 1.0, phiM - phi, 0.0, 0.0, 1.0)
else:
alpha, sx, sy, scale = compose_transform2(
alpha, sx, sy, 1.0, 180 - (phiM - phi), 0.0,
0.0, 1.0)
else:
if mirror == 0:
alpha, sx, sy, scale = compose_transform2(
alpha, sx, sy, 1.0, -(phiM - phi), 0.0, 0.0,
1.0)
else:
alpha, sx, sy, scale = compose_transform2(
alpha, sx, sy, 1.0, -(180 - (phiM - phi)), 0.0,
0.0, 1.0)
set_params2D(imgdata[mj], [alpha, sx, sy, mirror, 1.0])
grp_imgdata.append(imgdata[mj])
#print grp_imgdata[j].get_xsize(), imgdata[mj].get_xsize()
if not options.no_norm:
#print grp_imgdata[j].get_xsize()
mask = model_circle(nx / 2 - 2, nx, nx)
for k in xrange(img_per_grp):
ave, std, minn, maxx = Util.infomask(
grp_imgdata[k], mask, False)
grp_imgdata[k] -= ave
grp_imgdata[k] /= std
del mask
if options.fl > 0.0:
from filter import filt_ctf, filt_table
from fundamentals import fft, window2d
nx2 = 2 * nx
ny2 = 2 * ny
if options.CTF:
from utilities import pad
for k in xrange(img_per_grp):
grp_imgdata[k] = window2d(
fft(
filt_tanl(
filt_ctf(
fft(
pad(grp_imgdata[k], nx2, ny2,
1, 0.0)),
grp_imgdata[k].get_attr("ctf"),
binary=1), options.fl,
options.aa)), nx, ny)
#grp_imgdata[k] = window2d(fft( filt_table( filt_tanl( filt_ctf(fft(pad(grp_imgdata[k], nx2, ny2, 1,0.0)), grp_imgdata[k].get_attr("ctf"), binary=1), options.fl, options.aa), fifi) ),nx,ny)
#grp_imgdata[k] = filt_tanl(grp_imgdata[k], options.fl, options.aa)
else:
for k in xrange(img_per_grp):
grp_imgdata[k] = filt_tanl(grp_imgdata[k],
options.fl, options.aa)
#grp_imgdata[k] = window2d(fft( filt_table( filt_tanl( filt_ctf(fft(pad(grp_imgdata[k], nx2, ny2, 1,0.0)), grp_imgdata[k].get_attr("ctf"), binary=1), options.fl, options.aa), fifi) ),nx,ny)
#grp_imgdata[k] = filt_tanl(grp_imgdata[k], options.fl, options.aa)
else:
from utilities import pad, read_text_file
from filter import filt_ctf, filt_table
from fundamentals import fft, window2d
nx2 = 2 * nx
ny2 = 2 * ny
if options.CTF:
from utilities import pad
for k in xrange(img_per_grp):
grp_imgdata[k] = window2d(
fft(
filt_ctf(fft(
pad(grp_imgdata[k], nx2, ny2, 1, 0.0)),
grp_imgdata[k].get_attr("ctf"),
binary=1)), nx, ny)
#grp_imgdata[k] = window2d(fft( filt_table( filt_tanl( filt_ctf(fft(pad(grp_imgdata[k], nx2, ny2, 1,0.0)), grp_imgdata[k].get_attr("ctf"), binary=1), options.fl, options.aa), fifi) ),nx,ny)
#grp_imgdata[k] = filt_tanl(grp_imgdata[k], options.fl, options.aa)
'''
if i < 10 and myid == main_node:
for k in xrange(10):
grp_imgdata[k].write_image("grp%03d.hdf"%i, k)
'''
"""
if myid == main_node and i==0:
for pp in xrange(len(grp_imgdata)):
grp_imgdata[pp].write_image("pp.hdf", pp)
"""
ave, grp_imgdata = prepare_2d_forPCA(grp_imgdata)
"""
if myid == main_node and i==0:
for pp in xrange(len(grp_imgdata)):
grp_imgdata[pp].write_image("qq.hdf", pp)
"""
var = model_blank(nx, ny)
for q in grp_imgdata:
Util.add_img2(var, q)
Util.mul_scalar(var, 1.0 / (len(grp_imgdata) - 1))
# Switch to std dev
var = square_root(threshold(var))
#if options.CTF: ave, var = avgvar_ctf(grp_imgdata, mode="a")
#else: ave, var = avgvar(grp_imgdata, mode="a")
"""
if myid == main_node:
ave.write_image("avgv.hdf",i)
var.write_image("varv.hdf",i)
"""
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)
if options.VERBOSE:
print("%5.2f%% done on processor %d" %
(i * 100.0 / len(proj_list), myid))
if nvec > 0:
eig = pca(input_stacks=grp_imgdata,
subavg="",
mask_radius=radiuspca,
nvec=nvec,
incore=True,
shuffle=False,
genbuf=True)
for k in xrange(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)
"""
del imgdata
# To this point, all averages, variances, and eigenvectors are computed
if options.ave2D:
from fundamentals import fpol
if myid == main_node:
km = 0
for i in xrange(number_of_proc):
if i == main_node:
for im in xrange(len(aveList)):
aveList[im].write_image(
os.path.join(options.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 xrange(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, Tracker["nx"],
Tracker["nx"], 1)
tmpvol.write_image(
os.path.join(options.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 xrange(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 options.ave3D:
from fundamentals import fpol
if options.VERBOSE:
print("Reconstructing 3D average volume")
ave3D = recons3d_4nn_MPI(myid,
aveList,
symmetry=options.sym,
npad=options.npad)
bcast_EMData_to_all(ave3D, myid)
if myid == main_node:
line = strftime("%Y-%m-%d_%H:%M:%S", localtime()) + " =>"
ave3D = fpol(ave3D, Tracker["nx"], Tracker["nx"],
Tracker["nx"])
ave3D.write_image(
os.path.join(options.output_dir, options.ave3D))
msg = ("%-70s: %s\n" % (
"Writing to the disk volume reconstructed from averages as",
options.ave3D))
log_main.add(msg)
print(line, msg)
del ave, var, proj_list, stack, phi, theta, psi, s2x, s2y, alpha, sx, sy, mirror, aveList
if nvec > 0:
for k in xrange(nvec):
if options.VERBOSE:
print("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 xrange(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:
line = strftime("%Y-%m-%d_%H:%M:%S",
localtime()) + " =>"
u = int(u[0])
msg = (" Eigenvector: ", k, " number changed ",
int(icont[0]))
log_main.add(msg)
print(line, msg)
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 fundamentals import fpol
if myid == main_node:
km = 0
for i in xrange(number_of_proc):
if i == main_node:
for im in xrange(len(varList)):
tmpvol = fpol(varList[im], Tracker["nx"],
Tracker["nx"], 1)
tmpvol.write_image(
os.path.join(options.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 xrange(nl):
ave = recv_EMData(i, im + i + 70000)
tmpvol = fpol(ave, Tracker["nx"],
Tracker["nx"], 1)
tmpvol.write_image(
os.path.join(options.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 xrange(len(varList)):
send_EMData(varList[im], main_node, im + myid +
70000) # What with the attributes??
mpi_barrier(MPI_COMM_WORLD)
if options.var3D:
if myid == main_node and options.VERBOSE:
line = strftime("%Y-%m-%d_%H:%M:%S", localtime()) + " =>"
msg = ("Reconstructing 3D variability volume")
log_main.add(msg)
print(line, msg)
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 fundamentals import fpol
res = fpol(res, Tracker["nx"], Tracker["nx"], Tracker["nx"])
res.write_image(os.path.join(options.output_dir,
options.var3D))
if myid == main_node:
line = strftime("%Y-%m-%d_%H:%M:%S", localtime()) + " =>"
msg = ("%-70s: %.2f\n" %
("Reconstructing 3D variability took [s]", time() - t6))
log_main.add(msg)
print(line, msg)
if options.VERBOSE:
print("Reconstruction took: %.2f [min]" %
((time() - t6) / 60))
if myid == main_node:
line = strftime("%Y-%m-%d_%H:%M:%S", localtime()) + " =>"
msg = ("%-70s: %.2f\n" %
("Total time for these computations [s]", time() - t0))
print(line, msg)
log_main.add(msg)
if options.VERBOSE:
print("Total time for these computations: %.2f [min]" %
((time() - t0) / 60))
line = strftime("%Y-%m-%d_%H:%M:%S", localtime()) + " =>"
msg = ("sx3dvariability")
print(line, msg)
log_main.add(msg)
from mpi import mpi_finalize
mpi_finalize()
if RUNNING_UNDER_MPI:
global_def.MPI = False
global_def.BATCH = False
def helicalshiftali_MPI(stack,
maskfile=None,
maxit=100,
CTF=False,
snr=1.0,
Fourvar=False,
search_rng=-1):
from applications import MPI_start_end
from utilities import model_circle, model_blank, get_image, peak_search, get_im, pad
from utilities import reduce_EMData_to_root, bcast_EMData_to_all, send_attr_dict, file_type, bcast_number_to_all, bcast_list_to_all
from pap_statistics import varf2d_MPI
from fundamentals import fft, ccf, rot_shift3D, rot_shift2D, fshift
from utilities import get_params2D, set_params2D, chunks_distribution
from utilities import print_msg, print_begin_msg, print_end_msg
import os
import sys
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
from mpi import MPI_SUM, MPI_FLOAT, MPI_INT
from time import time
from pixel_error import ordersegments
from math import sqrt, atan2, tan, pi
nproc = mpi_comm_size(MPI_COMM_WORLD)
myid = mpi_comm_rank(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), "ehelix_MPI", 1, 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)
print("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 filter import filt_ctf
from 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 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', "helicalshiftali_MPI", 1,
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_reduce(sx_sum, 1, MPI_FLOAT, MPI_SUM, main_node,
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_barrier(MPI_COMM_WORLD)
par_str = ["xform.align2d", "ID"]
if myid == main_node:
from utilities import file_type
if (file_type(stack) == "bdb"):
from utilities import recv_attr_dict_bdb
recv_attr_dict_bdb(main_node, stack, data, par_str, 0, ldata,
nproc)
else:
from 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 ali3d_MPI(stack, ref_vol, outdir, maskfile = None, ir = 1, ou = -1, rs = 1,
xr = "4 2 2 1", yr = "-1", ts = "1 1 0.5 0.25", delta = "10 6 4 4", an = "-1",
center = 0, maxit = 5, term = 95, CTF = False, fourvar = False, snr = 1.0, ref_a = "S", sym = "c1",
sort=True, cutoff=999.99, pix_cutoff="0", two_tail=False, model_jump="1 1 1 1 1", restart=False, save_half=False,
protos=None, oplane=None, lmask=-1, ilmask=-1, findseam=False, vertstep=None, hpars="-1", hsearch="73.0 170.0",
full_output = False, compare_repro = False, compare_ref_free = "-1", ref_free_cutoff= "-1 -1 -1 -1",
wcmask = None, debug = False, recon_pad = 4):
from alignment import Numrinit, prepare_refrings
from utilities import model_circle, get_image, drop_image, get_input_from_string
from utilities import bcast_list_to_all, bcast_number_to_all, reduce_EMData_to_root, bcast_EMData_to_all
from utilities import send_attr_dict
from utilities import get_params_proj, file_type
from fundamentals import rot_avg_image
import os
import types
from utilities import print_begin_msg, print_end_msg, print_msg
from mpi import mpi_bcast, mpi_comm_size, mpi_comm_rank, MPI_FLOAT, MPI_COMM_WORLD, mpi_barrier, mpi_reduce
from mpi import mpi_reduce, MPI_INT, MPI_SUM, mpi_finalize
from filter import filt_ctf
from projection import prep_vol, prgs
from statistics import hist_list, varf3d_MPI, fsc_mask
from numpy import array, bincount, array2string, ones
number_of_proc = mpi_comm_size(MPI_COMM_WORLD)
myid = mpi_comm_rank(MPI_COMM_WORLD)
main_node = 0
if myid == main_node:
if os.path.exists(outdir): ERROR('Output directory exists, please change the name and restart the program', "ali3d_MPI", 1)
os.mkdir(outdir)
mpi_barrier(MPI_COMM_WORLD)
if debug:
from time import sleep
while not os.path.exists(outdir):
print "Node ",myid," waiting..."
sleep(5)
info_file = os.path.join(outdir, "progress%04d"%myid)
finfo = open(info_file, 'w')
else:
finfo = None
mjump = get_input_from_string(model_jump)
xrng = get_input_from_string(xr)
if yr == "-1": yrng = xrng
else : yrng = get_input_from_string(yr)
step = get_input_from_string(ts)
delta = get_input_from_string(delta)
ref_free_cutoff = get_input_from_string(ref_free_cutoff)
pix_cutoff = get_input_from_string(pix_cutoff)
lstp = min(len(xrng), len(yrng), len(step), len(delta))
if an == "-1":
an = [-1] * lstp
else:
an = get_input_from_string(an)
# make sure pix_cutoff is set for all iterations
if len(pix_cutoff)<lstp:
for i in xrange(len(pix_cutoff),lstp):
pix_cutoff.append(pix_cutoff[-1])
# don't waste time on sub-pixel alignment for low-resolution ang incr
for i in range(len(step)):
if (delta[i] > 4 or delta[i] == -1) and step[i] < 1:
step[i] = 1
first_ring = int(ir)
rstep = int(rs)
last_ring = int(ou)
max_iter = int(maxit)
center = int(center)
nrefs = EMUtil.get_image_count( ref_vol )
nmasks = 0
if maskfile:
# read number of masks within each maskfile (mc)
nmasks = EMUtil.get_image_count( maskfile )
# open masks within maskfile (mc)
maskF = EMData.read_images(maskfile, xrange(nmasks))
vol = EMData.read_images(ref_vol, xrange(nrefs))
nx = vol[0].get_xsize()
## make sure box sizes are the same
if myid == main_node:
im=EMData.read_images(stack,[0])
bx = im[0].get_xsize()
if bx!=nx:
print_msg("Error: Stack box size (%i) differs from initial model (%i)\n"%(bx,nx))
sys.exit()
del im,bx
# for helical processing:
helicalrecon = False
if protos is not None or hpars != "-1" or findseam is True:
helicalrecon = True
# if no out-of-plane param set, use 5 degrees
if oplane is None:
oplane=5.0
if protos is not None:
proto = get_input_from_string(protos)
if len(proto) != nrefs:
print_msg("Error: insufficient protofilament numbers supplied")
sys.exit()
if hpars != "-1":
hpars = get_input_from_string(hpars)
if len(hpars) != 2*nrefs:
print_msg("Error: insufficient helical parameters supplied")
sys.exit()
## create helical parameter file for helical reconstruction
if helicalrecon is True and myid == main_node:
from hfunctions import createHpar
# create initial helical parameter files
dp=[0]*nrefs
dphi=[0]*nrefs
vdp=[0]*nrefs
vdphi=[0]*nrefs
for iref in xrange(nrefs):
hpar = os.path.join(outdir,"hpar%02d.spi"%(iref))
params = False
if hpars != "-1":
# if helical parameters explicitly given, set twist & rise
params = [float(hpars[iref*2]),float(hpars[(iref*2)+1])]
dp[iref],dphi[iref],vdp[iref],vdphi[iref] = createHpar(hpar,proto[iref],params,vertstep)
# get values for helical search parameters
hsearch = get_input_from_string(hsearch)
if len(hsearch) != 2:
print_msg("Error: specify outer and inner radii for helical search")
sys.exit()
if last_ring < 0 or last_ring > int(nx/2)-2 : last_ring = int(nx/2) - 2
if myid == main_node:
# import user_functions
# user_func = user_functions.factory[user_func_name]
print_begin_msg("ali3d_MPI")
print_msg("Input stack : %s\n"%(stack))
print_msg("Reference volume : %s\n"%(ref_vol))
print_msg("Output directory : %s\n"%(outdir))
if nmasks > 0:
print_msg("Maskfile (number of masks) : %s (%i)\n"%(maskfile,nmasks))
print_msg("Inner radius : %i\n"%(first_ring))
print_msg("Outer radius : %i\n"%(last_ring))
print_msg("Ring step : %i\n"%(rstep))
print_msg("X search range : %s\n"%(xrng))
print_msg("Y search range : %s\n"%(yrng))
print_msg("Translational step : %s\n"%(step))
print_msg("Angular step : %s\n"%(delta))
print_msg("Angular search range : %s\n"%(an))
print_msg("Maximum iteration : %i\n"%(max_iter))
print_msg("Center type : %i\n"%(center))
print_msg("CTF correction : %s\n"%(CTF))
print_msg("Signal-to-Noise Ratio : %f\n"%(snr))
print_msg("Reference projection method : %s\n"%(ref_a))
print_msg("Symmetry group : %s\n"%(sym))
print_msg("Fourier padding for 3D : %i\n"%(recon_pad))
print_msg("Number of reference models : %i\n"%(nrefs))
print_msg("Sort images between models : %s\n"%(sort))
print_msg("Allow images to jump : %s\n"%(mjump))
print_msg("CC cutoff standard dev : %f\n"%(cutoff))
print_msg("Two tail cutoff : %s\n"%(two_tail))
print_msg("Termination pix error : %f\n"%(term))
print_msg("Pixel error cutoff : %s\n"%(pix_cutoff))
print_msg("Restart : %s\n"%(restart))
print_msg("Full output : %s\n"%(full_output))
print_msg("Compare reprojections : %s\n"%(compare_repro))
print_msg("Compare ref free class avgs : %s\n"%(compare_ref_free))
print_msg("Use cutoff from ref free : %s\n"%(ref_free_cutoff))
if protos:
print_msg("Protofilament numbers : %s\n"%(proto))
print_msg("Using helical search range : %s\n"%hsearch)
if findseam is True:
print_msg("Using seam-based reconstruction\n")
if hpars != "-1":
print_msg("Using hpars : %s\n"%hpars)
if vertstep != None:
print_msg("Using vertical step : %.2f\n"%vertstep)
if save_half is True:
print_msg("Saving even/odd halves\n")
for i in xrange(100) : print_msg("*")
print_msg("\n\n")
if maskfile:
if type(maskfile) is types.StringType: mask3D = get_image(maskfile)
else: mask3D = maskfile
else: mask3D = model_circle(last_ring, nx, nx, nx)
numr = Numrinit(first_ring, last_ring, rstep, "F")
mask2D = model_circle(last_ring,nx,nx) - model_circle(first_ring,nx,nx)
fscmask = model_circle(last_ring,nx,nx,nx)
if CTF:
from filter import filt_ctf
from reconstruction_rjh import rec3D_MPI_noCTF
if myid == main_node:
active = EMUtil.get_all_attributes(stack, 'active')
list_of_particles = []
for im in xrange(len(active)):
if active[im]: list_of_particles.append(im)
del active
nima = len(list_of_particles)
else:
nima = 0
total_nima = bcast_number_to_all(nima, source_node = main_node)
if myid != main_node:
list_of_particles = [-1]*total_nima
list_of_particles = bcast_list_to_all(list_of_particles, source_node = main_node)
image_start, image_end = MPI_start_end(total_nima, number_of_proc, myid)
# create a list of images for each node
list_of_particles = list_of_particles[image_start: image_end]
nima = len(list_of_particles)
if debug:
finfo.write("image_start, image_end: %d %d\n" %(image_start, image_end))
finfo.flush()
data = EMData.read_images(stack, list_of_particles)
t_zero = Transform({"type":"spider","phi":0,"theta":0,"psi":0,"tx":0,"ty":0})
transmulti = [[t_zero for i in xrange(nrefs)] for j in xrange(nima)]
for iref,im in ((iref,im) for iref in xrange(nrefs) for im in xrange(nima)):
if nrefs == 1:
transmulti[im][iref] = data[im].get_attr("xform.projection")
else:
# if multi models, keep track of eulers for all models
try:
transmulti[im][iref] = data[im].get_attr("eulers_txty.%i"%iref)
except:
data[im].set_attr("eulers_txty.%i"%iref,t_zero)
scoremulti = [[0.0 for i in xrange(nrefs)] for j in xrange(nima)]
pixelmulti = [[0.0 for i in xrange(nrefs)] for j in xrange(nima)]
ref_res = [0.0 for x in xrange(nrefs)]
apix = data[0].get_attr('apix_x')
# for oplane parameter, create cylindrical mask
if oplane is not None and myid == main_node:
from hfunctions import createCylMask
cmaskf=os.path.join(outdir, "mask3D_cyl.mrc")
mask3D = createCylMask(data,ou,lmask,ilmask,cmaskf)
# if finding seam of helix, create wedge masks
if findseam is True:
wedgemask=[]
for pf in xrange(nrefs):
wedgemask.append(EMData())
# wedgemask option
if wcmask is not None:
wcmask = get_input_from_string(wcmask)
if len(wcmask) != 3:
print_msg("Error: wcmask option requires 3 values: x y radius")
sys.exit()
# determine if particles have helix info:
try:
data[0].get_attr('h_angle')
original_data = []
boxmask = True
from hfunctions import createBoxMask
except:
boxmask = False
# prepare particles
for im in xrange(nima):
data[im].set_attr('ID', list_of_particles[im])
data[im].set_attr('pix_score', int(0))
if CTF:
# only phaseflip particles, not full CTF correction
ctf_params = data[im].get_attr("ctf")
st = Util.infomask(data[im], mask2D, False)
data[im] -= st[0]
data[im] = filt_ctf(data[im], ctf_params, sign = -1, binary=1)
data[im].set_attr('ctf_applied', 1)
# for window mask:
if boxmask is True:
h_angle = data[im].get_attr("h_angle")
original_data.append(data[im].copy())
bmask = createBoxMask(nx,apix,ou,lmask,h_angle)
data[im]*=bmask
del bmask
if debug:
finfo.write( '%d loaded \n' % nima )
finfo.flush()
if myid == main_node:
# initialize data for the reference preparation function
ref_data = [ mask3D, max(center,0), None, None, None, None ]
# for method -1, switch off centering in user function
from time import time
# this is needed for gathering of pixel errors
disps = []
recvcount = []
disps_score = []
recvcount_score = []
for im in xrange(number_of_proc):
if( im == main_node ):
disps.append(0)
disps_score.append(0)
else:
disps.append(disps[im-1] + recvcount[im-1])
disps_score.append(disps_score[im-1] + recvcount_score[im-1])
ib, ie = MPI_start_end(total_nima, number_of_proc, im)
recvcount.append( ie - ib )
recvcount_score.append((ie-ib)*nrefs)
pixer = [0.0]*nima
cs = [0.0]*3
total_iter = 0
volodd = EMData.read_images(ref_vol, xrange(nrefs))
voleve = EMData.read_images(ref_vol, xrange(nrefs))
if restart:
# recreate initial volumes from alignments stored in header
itout = "000_00"
for iref in xrange(nrefs):
if(nrefs == 1):
modout = ""
else:
modout = "_model_%02d"%(iref)
if(sort):
group = iref
for im in xrange(nima):
imgroup = data[im].get_attr('group')
if imgroup == iref:
data[im].set_attr('xform.projection',transmulti[im][iref])
else:
group = int(999)
for im in xrange(nima):
data[im].set_attr('xform.projection',transmulti[im][iref])
fscfile = os.path.join(outdir, "fsc_%s%s"%(itout,modout))
vol[iref], fscc, volodd[iref], voleve[iref] = rec3D_MPI_noCTF(data, sym, fscmask, fscfile, myid, main_node, index = group, npad = recon_pad)
if myid == main_node:
if helicalrecon:
from hfunctions import processHelicalVol
vstep=None
if vertstep is not None:
vstep=(vdp[iref],vdphi[iref])
print_msg("Old rise and twist for model %i : %8.3f, %8.3f\n"%(iref,dp[iref],dphi[iref]))
hvals=processHelicalVol(vol[iref],voleve[iref],volodd[iref],iref,outdir,itout,
dp[iref],dphi[iref],apix,hsearch,findseam,vstep,wcmask)
(vol[iref],voleve[iref],volodd[iref],dp[iref],dphi[iref],vdp[iref],vdphi[iref])=hvals
print_msg("New rise and twist for model %i : %8.3f, %8.3f\n"%(iref,dp[iref],dphi[iref]))
# get new FSC from symmetrized half volumes
fscc = fsc_mask( volodd[iref], voleve[iref], mask3D, rstep, fscfile)
else:
vol[iref].write_image(os.path.join(outdir, "vol_%s.hdf"%itout),-1)
if save_half is True:
volodd[iref].write_image(os.path.join(outdir, "volodd_%s.hdf"%itout),-1)
voleve[iref].write_image(os.path.join(outdir, "voleve_%s.hdf"%itout),-1)
if nmasks > 1:
# Read mask for multiplying
ref_data[0] = maskF[iref]
ref_data[2] = vol[iref]
ref_data[3] = fscc
# call user-supplied function to prepare reference image, i.e., center and filter it
vol[iref], cs,fl = ref_ali3d(ref_data)
vol[iref].write_image(os.path.join(outdir, "volf_%s.hdf"%(itout)),-1)
if (apix == 1):
res_msg = "Models filtered at spatial frequency of:\t"
res = fl
else:
res_msg = "Models filtered at resolution of: \t"
res = apix / fl
ares = array2string(array(res), precision = 2)
print_msg("%s%s\n\n"%(res_msg,ares))
bcast_EMData_to_all(vol[iref], myid, main_node)
# write out headers, under MPI writing has to be done sequentially
mpi_barrier(MPI_COMM_WORLD)
# projection matching
for N_step in xrange(lstp):
terminate = 0
Iter = -1
while(Iter < max_iter-1 and terminate == 0):
Iter += 1
total_iter += 1
itout = "%03g_%02d" %(delta[N_step], Iter)
if myid == main_node:
print_msg("ITERATION #%3d, inner iteration #%3d\nDelta = %4.1f, an = %5.2f, xrange = %5.2f, yrange = %5.2f, step = %5.2f\n\n"%(N_step, Iter, delta[N_step], an[N_step], xrng[N_step],yrng[N_step],step[N_step]))
for iref in xrange(nrefs):
if myid == main_node: start_time = time()
volft,kb = prep_vol( vol[iref] )
## constrain projections to out of plane parameter
theta1 = None
theta2 = None
if oplane is not None:
theta1 = 90-oplane
theta2 = 90+oplane
refrings = prepare_refrings( volft, kb, nx, delta[N_step], ref_a, sym, numr, MPI=True, phiEqpsi = "Minus", initial_theta=theta1, delta_theta=theta2)
del volft,kb
if myid== main_node:
print_msg( "Time to prepare projections for model %i: %s\n" % (iref, legibleTime(time()-start_time)) )
start_time = time()
for im in xrange( nima ):
data[im].set_attr("xform.projection", transmulti[im][iref])
if an[N_step] == -1:
t1, peak, pixer[im] = proj_ali_incore(data[im],refrings,numr,xrng[N_step],yrng[N_step],step[N_step],finfo)
else:
t1, peak, pixer[im] = proj_ali_incore_local(data[im],refrings,numr,xrng[N_step],yrng[N_step],step[N_step],an[N_step],finfo)
#data[im].set_attr("xform.projection"%iref, t1)
if nrefs > 1: data[im].set_attr("eulers_txty.%i"%iref,t1)
scoremulti[im][iref] = peak
from pixel_error import max_3D_pixel_error
# t1 is the current param, t2 is old
t2 = transmulti[im][iref]
pixelmulti[im][iref] = max_3D_pixel_error(t1,t2,numr[-3])
transmulti[im][iref] = t1
if myid == main_node:
print_msg("Time of alignment for model %i: %s\n"%(iref, legibleTime(time()-start_time)))
start_time = time()
# gather scoring data from all processors
from mpi import mpi_gatherv
scoremultisend = sum(scoremulti,[])
pixelmultisend = sum(pixelmulti,[])
tmp = mpi_gatherv(scoremultisend,len(scoremultisend),MPI_FLOAT, recvcount_score, disps_score, MPI_FLOAT, main_node,MPI_COMM_WORLD)
tmp1 = mpi_gatherv(pixelmultisend,len(pixelmultisend),MPI_FLOAT, recvcount_score, disps_score, MPI_FLOAT, main_node,MPI_COMM_WORLD)
tmp = mpi_bcast(tmp,(total_nima * nrefs), MPI_FLOAT,0, MPI_COMM_WORLD)
tmp1 = mpi_bcast(tmp1,(total_nima * nrefs), MPI_FLOAT,0, MPI_COMM_WORLD)
tmp = map(float,tmp)
tmp1 = map(float,tmp1)
score = array(tmp).reshape(-1,nrefs)
pixelerror = array(tmp1).reshape(-1,nrefs)
score_local = array(scoremulti)
mean_score = score.mean(axis=0)
std_score = score.std(axis=0)
cut = mean_score - (cutoff * std_score)
cut2 = mean_score + (cutoff * std_score)
res_max = score_local.argmax(axis=1)
minus_cc = [0.0 for x in xrange(nrefs)]
minus_pix = [0.0 for x in xrange(nrefs)]
minus_ref = [0.0 for x in xrange(nrefs)]
#output pixel errors
if(myid == main_node):
from statistics import hist_list
lhist = 20
pixmin = pixelerror.min(axis=1)
region, histo = hist_list(pixmin, lhist)
if(region[0] < 0.0): region[0] = 0.0
print_msg("Histogram of pixel errors\n ERROR number of particles\n")
for lhx in xrange(lhist):
print_msg(" %10.3f %7d\n"%(region[lhx], histo[lhx]))
# Terminate if 95% within 1 pixel error
im = 0
for lhx in xrange(lhist):
if(region[lhx] > 1.0): break
im += histo[lhx]
print_msg( "Percent of particles with pixel error < 1: %f\n\n"% (im/float(total_nima)*100))
term_cond = float(term)/100
if(im/float(total_nima) > term_cond):
terminate = 1
print_msg("Terminating internal loop\n")
del region, histo
terminate = mpi_bcast(terminate, 1, MPI_INT, 0, MPI_COMM_WORLD)
terminate = int(terminate[0])
for im in xrange(nima):
if(sort==False):
data[im].set_attr('group',999)
elif (mjump[N_step]==1):
data[im].set_attr('group',int(res_max[im]))
pix_run = data[im].get_attr('pix_score')
if (pix_cutoff[N_step]==1 and (terminate==1 or Iter == max_iter-1)):
if (pixelmulti[im][int(res_max[im])] > 1):
data[im].set_attr('pix_score',int(777))
if (score_local[im][int(res_max[im])]<cut[int(res_max[im])]) or (two_tail and score_local[im][int(res_max[im])]>cut2[int(res_max[im])]):
data[im].set_attr('group',int(888))
minus_cc[int(res_max[im])] = minus_cc[int(res_max[im])] + 1
if(pix_run == 777):
data[im].set_attr('group',int(777))
minus_pix[int(res_max[im])] = minus_pix[int(res_max[im])] + 1
if (compare_ref_free != "-1") and (ref_free_cutoff[N_step] != -1) and (total_iter > 1):
id = data[im].get_attr('ID')
if id in rejects:
data[im].set_attr('group',int(666))
minus_ref[int(res_max[im])] = minus_ref[int(res_max[im])] + 1
minus_cc_tot = mpi_reduce(minus_cc,nrefs,MPI_FLOAT,MPI_SUM,0,MPI_COMM_WORLD)
minus_pix_tot = mpi_reduce(minus_pix,nrefs,MPI_FLOAT,MPI_SUM,0,MPI_COMM_WORLD)
minus_ref_tot = mpi_reduce(minus_ref,nrefs,MPI_FLOAT,MPI_SUM,0,MPI_COMM_WORLD)
if (myid == main_node):
if(sort):
tot_max = score.argmax(axis=1)
res = bincount(tot_max)
else:
res = ones(nrefs) * total_nima
print_msg("Particle distribution: \t\t%s\n"%(res*1.0))
afcut1 = res - minus_cc_tot
afcut2 = afcut1 - minus_pix_tot
afcut3 = afcut2 - minus_ref_tot
print_msg("Particle distribution after cc cutoff:\t\t%s\n"%(afcut1))
print_msg("Particle distribution after pix cutoff:\t\t%s\n"%(afcut2))
print_msg("Particle distribution after ref cutoff:\t\t%s\n\n"%(afcut3))
res = [0.0 for i in xrange(nrefs)]
for iref in xrange(nrefs):
if(center == -1):
from utilities import estimate_3D_center_MPI, rotate_3D_shift
dummy=EMData()
cs[0], cs[1], cs[2], dummy, dummy = estimate_3D_center_MPI(data, total_nima, myid, number_of_proc, main_node)
cs = mpi_bcast(cs, 3, MPI_FLOAT, main_node, MPI_COMM_WORLD)
cs = [-float(cs[0]), -float(cs[1]), -float(cs[2])]
rotate_3D_shift(data, cs)
if(sort):
group = iref
for im in xrange(nima):
imgroup = data[im].get_attr('group')
if imgroup == iref:
data[im].set_attr('xform.projection',transmulti[im][iref])
else:
group = int(999)
for im in xrange(nima):
data[im].set_attr('xform.projection',transmulti[im][iref])
if(nrefs == 1):
modout = ""
else:
modout = "_model_%02d"%(iref)
fscfile = os.path.join(outdir, "fsc_%s%s"%(itout,modout))
vol[iref], fscc, volodd[iref], voleve[iref] = rec3D_MPI_noCTF(data, sym, fscmask, fscfile, myid, main_node, index=group, npad=recon_pad)
if myid == main_node:
print_msg("3D reconstruction time for model %i: %s\n"%(iref, legibleTime(time()-start_time)))
start_time = time()
# Compute Fourier variance
if fourvar:
outvar = os.path.join(outdir, "volVar_%s.hdf"%(itout))
ssnr_file = os.path.join(outdir, "ssnr_%s"%(itout))
varf = varf3d_MPI(data, ssnr_text_file=ssnr_file, mask2D=None, reference_structure=vol[iref], ou=last_ring, rw=1.0, npad=1, CTF=None, sign=1, sym=sym, myid=myid)
if myid == main_node:
print_msg("Time to calculate 3D Fourier variance for model %i: %s\n"%(iref, legibleTime(time()-start_time)))
start_time = time()
varf = 1.0/varf
varf.write_image(outvar,-1)
else: varf = None
if myid == main_node:
if helicalrecon:
from hfunctions import processHelicalVol
vstep=None
if vertstep is not None:
vstep=(vdp[iref],vdphi[iref])
print_msg("Old rise and twist for model %i : %8.3f, %8.3f\n"%(iref,dp[iref],dphi[iref]))
hvals=processHelicalVol(vol[iref],voleve[iref],volodd[iref],iref,outdir,itout,
dp[iref],dphi[iref],apix,hsearch,findseam,vstep,wcmask)
(vol[iref],voleve[iref],volodd[iref],dp[iref],dphi[iref],vdp[iref],vdphi[iref])=hvals
print_msg("New rise and twist for model %i : %8.3f, %8.3f\n"%(iref,dp[iref],dphi[iref]))
# get new FSC from symmetrized half volumes
fscc = fsc_mask( volodd[iref], voleve[iref], mask3D, rstep, fscfile)
print_msg("Time to search and apply helical symmetry for model %i: %s\n\n"%(iref, legibleTime(time()-start_time)))
start_time = time()
else:
vol[iref].write_image(os.path.join(outdir, "vol_%s.hdf"%(itout)),-1)
if save_half is True:
volodd[iref].write_image(os.path.join(outdir, "volodd_%s.hdf"%(itout)),-1)
voleve[iref].write_image(os.path.join(outdir, "voleve_%s.hdf"%(itout)),-1)
if nmasks > 1:
# Read mask for multiplying
ref_data[0] = maskF[iref]
ref_data[2] = vol[iref]
ref_data[3] = fscc
ref_data[4] = varf
# call user-supplied function to prepare reference image, i.e., center and filter it
vol[iref], cs,fl = ref_ali3d(ref_data)
vol[iref].write_image(os.path.join(outdir, "volf_%s.hdf"%(itout)),-1)
if (apix == 1):
res_msg = "Models filtered at spatial frequency of:\t"
res[iref] = fl
else:
res_msg = "Models filtered at resolution of: \t"
res[iref] = apix / fl
del varf
bcast_EMData_to_all(vol[iref], myid, main_node)
if compare_ref_free != "-1": compare_repro = True
if compare_repro:
outfile_repro = comp_rep(refrings, data, itout, modout, vol[iref], group, nima, nx, myid, main_node, outdir)
mpi_barrier(MPI_COMM_WORLD)
if compare_ref_free != "-1":
ref_free_output = os.path.join(outdir,"ref_free_%s%s"%(itout,modout))
rejects = compare(compare_ref_free, outfile_repro,ref_free_output,yrng[N_step], xrng[N_step], rstep,nx,apix,ref_free_cutoff[N_step], number_of_proc, myid, main_node)
# retrieve alignment params from all processors
par_str = ['xform.projection','ID','group']
if nrefs > 1:
for iref in xrange(nrefs):
par_str.append('eulers_txty.%i'%iref)
if myid == main_node:
from 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:
ares = array2string(array(res), precision = 2)
print_msg("%s%s\n\n"%(res_msg,ares))
dummy = EMData()
if full_output:
nimat = EMUtil.get_image_count(stack)
output_file = os.path.join(outdir, "paramout_%s"%itout)
foutput = open(output_file, 'w')
for im in xrange(nimat):
# save the parameters for each of the models
outstring = ""
dummy.read_image(stack,im,True)
param3d = dummy.get_attr('xform.projection')
g = dummy.get_attr("group")
# retrieve alignments in EMAN-format
pE = param3d.get_params('eman')
outstring += "%f\t%f\t%f\t%f\t%f\t%i\n" %(pE["az"], pE["alt"], pE["phi"], pE["tx"], pE["ty"],g)
foutput.write(outstring)
foutput.close()
del dummy
mpi_barrier(MPI_COMM_WORLD)
# mpi_finalize()
if myid == main_node: print_end_msg("ali3d_MPI")
def resample( prjfile, outdir, bufprefix, nbufvol, nvol, seedbase,\
delta, d, snr, CTF, npad,\
MPI, myid, ncpu, verbose = 0 ):
from utilities import even_angles
from random import seed, jumpahead, shuffle
import os
from sys import exit
nprj = EMUtil.get_image_count(prjfile)
if MPI:
from mpi import mpi_barrier, MPI_COMM_WORLD
if myid == 0:
if os.path.exists(outdir): nx = 1
else: nx = 0
else: nx = 0
ny = bcast_number_to_all(nx, source_node=0)
if ny == 1:
ERROR(
'Output directory exists, please change the name and restart the program',
"resample", 1, myid)
mpi_barrier(MPI_COMM_WORLD)
if myid == 0:
os.mkdir(outdir)
mpi_barrier(MPI_COMM_WORLD)
else:
if os.path.exists(outdir):
ERROR(
'Output directory exists, please change the name and restart the program',
"resample", 1, 0)
os.mkdir(outdir)
if (verbose == 1):
finfo = open(os.path.join(outdir, "progress%04d.txt" % myid), "w")
else:
finfo = None
#print " before evenangles",myid
from utilities import getvec
from numpy import array, reshape
refa = even_angles(delta)
nrefa = len(refa)
refnormal = zeros((nrefa, 3), 'float32')
tetref = [0.0] * nrefa
for i in range(nrefa):
tr = getvec(refa[i][0], refa[i][1])
for j in range(3):
refnormal[i][j] = tr[j]
tetref[i] = refa[i][1]
del refa
vct = array([0.0] * (3 * nprj), 'float32')
if myid == 0:
print(" will read ", myid)
tr = EMUtil.get_all_attributes(prjfile, 'xform.projection')
tetprj = [0.0] * nprj
for i in range(nprj):
temp = tr[i].get_params("spider")
tetprj[i] = temp["theta"]
if (tetprj[i] > 90.0): tetprj[i] = 180.0 - tetprj[i]
vct[3 * i + 0] = tr[i].at(2, 0)
vct[3 * i + 1] = tr[i].at(2, 1)
vct[3 * i + 2] = tr[i].at(2, 2)
del tr
else:
tetprj = [0.0] * nprj
#print " READ ",myid
if MPI:
#print " will bcast",myid
from mpi import mpi_bcast, MPI_FLOAT, MPI_COMM_WORLD
vct = mpi_bcast(vct, len(vct), MPI_FLOAT, 0, MPI_COMM_WORLD)
from utilities import bcast_list_to_all
tetprj = bcast_list_to_all(tetprj, myid, 0)
#print " reshape ",myid
vct = reshape(vct, (nprj, 3))
assignments = [[] for i in range(nrefa)]
dspn = 1.25 * delta
for k in range(nprj):
best_s = -1.0
best_i = -1
for i in range(nrefa):
if (abs(tetprj[k] - tetref[i]) <= dspn):
s = abs(refnormal[i][0] * vct[k][0] +
refnormal[i][1] * vct[k][1] +
refnormal[i][2] * vct[k][2])
if s > best_s:
best_s = s
best_i = i
assignments[best_i].append(k)
am = len(assignments[0])
mufur = 1.0 / am
for i in range(1, len(assignments)):
ti = len(assignments[i])
am = min(am, ti)
if (ti > 0): mufur += 1.0 / ti
del tetprj, tetref
dp = 1.0 - d # keep that many in each direction
keep = int(am * dp + 0.5)
mufur = keep * nrefa / (1.0 - mufur * keep / float(nrefa))
if myid == 0:
print(" Number of projections ", nprj,
". Number of reference directions ", nrefa,
", multiplicative factor for the variance ", mufur)
print(" Minimum number of assignments ", am,
" Number of projections used per stratum ", keep,
" Number of projections in resampled structure ",
int(am * dp + 0.5) * nrefa)
if am < 2 or am == keep:
print("incorrect settings")
exit() # FIX
if (seedbase < 1):
seed()
jumpahead(17 * myid + 123)
else:
seed(seedbase)
jumpahead(17 * myid + 123)
volfile = os.path.join(outdir, "bsvol%04d.hdf" % myid)
from random import randint
niter = nvol / ncpu / nbufvol
for kiter in range(niter):
if (verbose == 1):
finfo.write("Iteration %d: \n" % kiter)
finfo.flush()
iter_start = time()
# the following has to be converted to resample mults=1 means take given projection., mults=0 means omit
mults = [[0] * nprj for i in range(nbufvol)]
for i in range(nbufvol):
for l in range(nrefa):
mass = assignments[l][:]
shuffle(mass)
mass = mass[:keep]
mass.sort()
#print l, " * ",mass
for k in range(keep):
mults[i][mass[k]] = 1
'''
lout = []
for l in xrange(len(mults[i])):
if mults[i][l] == 1: lout.append(l)
write_text_file(lout, os.path.join(outdir, "list%04d_%03d.txt" %(i, myid)))
del lout
'''
del mass
rectors, fftvols, wgtvols = resample_prepare(prjfile, nbufvol, snr,
CTF, npad)
resample_insert(bufprefix, fftvols, wgtvols, mults, CTF, npad, finfo)
del mults
resample_finish(rectors, fftvols, wgtvols, volfile, kiter, nprj, finfo)
rectors = None
fftvols = None
wgtvols = None
if (verbose == 1):
finfo.write("time for iteration: %10.3f\n" % (time() - iter_start))
finfo.flush()
def ali3d_MPI(stack,
ref_vol,
outdir,
maskfile=None,
ir=1,
ou=-1,
rs=1,
xr="4 2 2 1",
yr="-1",
ts="1 1 0.5 0.25",
delta="10 6 4 4",
an="-1",
center=0,
maxit=5,
term=95,
CTF=False,
fourvar=False,
snr=1.0,
ref_a="S",
sym="c1",
sort=True,
cutoff=999.99,
pix_cutoff="0",
two_tail=False,
model_jump="1 1 1 1 1",
restart=False,
save_half=False,
protos=None,
oplane=None,
lmask=-1,
ilmask=-1,
findseam=False,
vertstep=None,
hpars="-1",
hsearch="0.0 50.0",
full_output=False,
compare_repro=False,
compare_ref_free="-1",
ref_free_cutoff="-1 -1 -1 -1",
wcmask=None,
debug=False,
recon_pad=4,
olmask=75):
from alignment import Numrinit, prepare_refrings
from utilities import model_circle, get_image, drop_image, get_input_from_string
from utilities import bcast_list_to_all, bcast_number_to_all, reduce_EMData_to_root, bcast_EMData_to_all
from utilities import send_attr_dict
from utilities import get_params_proj, file_type
from fundamentals import rot_avg_image
import os
import types
from utilities import print_begin_msg, print_end_msg, print_msg
from mpi import mpi_bcast, mpi_comm_size, mpi_comm_rank, MPI_FLOAT, MPI_COMM_WORLD, mpi_barrier, mpi_reduce
from mpi import mpi_reduce, MPI_INT, MPI_SUM, mpi_finalize
from filter import filt_ctf
from projection import prep_vol, prgs
from statistics import hist_list, varf3d_MPI, fsc_mask
from numpy import array, bincount, array2string, ones
number_of_proc = mpi_comm_size(MPI_COMM_WORLD)
myid = mpi_comm_rank(MPI_COMM_WORLD)
main_node = 0
if myid == main_node:
if os.path.exists(outdir):
ERROR(
'Output directory exists, please change the name and restart the program',
"ali3d_MPI", 1)
os.mkdir(outdir)
mpi_barrier(MPI_COMM_WORLD)
if debug:
from time import sleep
while not os.path.exists(outdir):
print "Node ", myid, " waiting..."
sleep(5)
info_file = os.path.join(outdir, "progress%04d" % myid)
finfo = open(info_file, 'w')
else:
finfo = None
mjump = get_input_from_string(model_jump)
xrng = get_input_from_string(xr)
if yr == "-1": yrng = xrng
else: yrng = get_input_from_string(yr)
step = get_input_from_string(ts)
delta = get_input_from_string(delta)
ref_free_cutoff = get_input_from_string(ref_free_cutoff)
pix_cutoff = get_input_from_string(pix_cutoff)
lstp = min(len(xrng), len(yrng), len(step), len(delta))
if an == "-1":
an = [-1] * lstp
else:
an = get_input_from_string(an)
# make sure pix_cutoff is set for all iterations
if len(pix_cutoff) < lstp:
for i in xrange(len(pix_cutoff), lstp):
pix_cutoff.append(pix_cutoff[-1])
# don't waste time on sub-pixel alignment for low-resolution ang incr
for i in range(len(step)):
if (delta[i] > 4 or delta[i] == -1) and step[i] < 1:
step[i] = 1
first_ring = int(ir)
rstep = int(rs)
last_ring = int(ou)
max_iter = int(maxit)
center = int(center)
nrefs = EMUtil.get_image_count(ref_vol)
nmasks = 0
if maskfile:
# read number of masks within each maskfile (mc)
nmasks = EMUtil.get_image_count(maskfile)
# open masks within maskfile (mc)
maskF = EMData.read_images(maskfile, xrange(nmasks))
vol = EMData.read_images(ref_vol, xrange(nrefs))
nx = vol[0].get_xsize()
## make sure box sizes are the same
if myid == main_node:
im = EMData.read_images(stack, [0])
bx = im[0].get_xsize()
if bx != nx:
print_msg(
"Error: Stack box size (%i) differs from initial model (%i)\n"
% (bx, nx))
sys.exit()
del im, bx
# for helical processing:
helicalrecon = False
if protos is not None or hpars != "-1" or findseam is True:
helicalrecon = True
# if no out-of-plane param set, use 5 degrees
if oplane is None:
oplane = 5.0
if protos is not None:
proto = get_input_from_string(protos)
if len(proto) != nrefs:
print_msg("Error: insufficient protofilament numbers supplied")
sys.exit()
if hpars != "-1":
hpars = get_input_from_string(hpars)
if len(hpars) != 2 * nrefs:
print_msg("Error: insufficient helical parameters supplied")
sys.exit()
## create helical parameter file for helical reconstruction
if helicalrecon is True and myid == main_node:
from hfunctions import createHpar
# create initial helical parameter files
dp = [0] * nrefs
dphi = [0] * nrefs
vdp = [0] * nrefs
vdphi = [0] * nrefs
for iref in xrange(nrefs):
hpar = os.path.join(outdir, "hpar%02d.spi" % (iref))
params = False
if hpars != "-1":
# if helical parameters explicitly given, set twist & rise
params = [float(hpars[iref * 2]), float(hpars[(iref * 2) + 1])]
dp[iref], dphi[iref], vdp[iref], vdphi[iref] = createHpar(
hpar, proto[iref], params, vertstep)
# get values for helical search parameters
hsearch = get_input_from_string(hsearch)
if len(hsearch) != 2:
print_msg("Error: specify outer and inner radii for helical search")
sys.exit()
if last_ring < 0 or last_ring > int(nx / 2) - 2:
last_ring = int(nx / 2) - 2
if myid == main_node:
# import user_functions
# user_func = user_functions.factory[user_func_name]
print_begin_msg("ali3d_MPI")
print_msg("Input stack : %s\n" % (stack))
print_msg("Reference volume : %s\n" % (ref_vol))
print_msg("Output directory : %s\n" % (outdir))
if nmasks > 0:
print_msg("Maskfile (number of masks) : %s (%i)\n" %
(maskfile, nmasks))
print_msg("Inner radius : %i\n" % (first_ring))
print_msg("Outer radius : %i\n" % (last_ring))
print_msg("Ring step : %i\n" % (rstep))
print_msg("X search range : %s\n" % (xrng))
print_msg("Y search range : %s\n" % (yrng))
print_msg("Translational step : %s\n" % (step))
print_msg("Angular step : %s\n" % (delta))
print_msg("Angular search range : %s\n" % (an))
print_msg("Maximum iteration : %i\n" % (max_iter))
print_msg("Center type : %i\n" % (center))
print_msg("CTF correction : %s\n" % (CTF))
print_msg("Signal-to-Noise Ratio : %f\n" % (snr))
print_msg("Reference projection method : %s\n" % (ref_a))
print_msg("Symmetry group : %s\n" % (sym))
print_msg("Fourier padding for 3D : %i\n" % (recon_pad))
print_msg("Number of reference models : %i\n" % (nrefs))
print_msg("Sort images between models : %s\n" % (sort))
print_msg("Allow images to jump : %s\n" % (mjump))
print_msg("CC cutoff standard dev : %f\n" % (cutoff))
print_msg("Two tail cutoff : %s\n" % (two_tail))
print_msg("Termination pix error : %f\n" % (term))
print_msg("Pixel error cutoff : %s\n" % (pix_cutoff))
print_msg("Restart : %s\n" % (restart))
print_msg("Full output : %s\n" % (full_output))
print_msg("Compare reprojections : %s\n" % (compare_repro))
print_msg("Compare ref free class avgs : %s\n" % (compare_ref_free))
print_msg("Use cutoff from ref free : %s\n" % (ref_free_cutoff))
if protos:
print_msg("Protofilament numbers : %s\n" % (proto))
print_msg("Using helical search range : %s\n" % hsearch)
if findseam is True:
print_msg("Using seam-based reconstruction\n")
if hpars != "-1":
print_msg("Using hpars : %s\n" % hpars)
if vertstep != None:
print_msg("Using vertical step : %.2f\n" % vertstep)
if save_half is True:
print_msg("Saving even/odd halves\n")
for i in xrange(100):
print_msg("*")
print_msg("\n\n")
if maskfile:
if type(maskfile) is types.StringType: mask3D = get_image(maskfile)
else: mask3D = maskfile
else: mask3D = model_circle(last_ring, nx, nx, nx)
numr = Numrinit(first_ring, last_ring, rstep, "F")
mask2D = model_circle(last_ring, nx, nx) - model_circle(first_ring, nx, nx)
fscmask = model_circle(last_ring, nx, nx, nx)
if CTF:
from filter import filt_ctf
from reconstruction_rjh import rec3D_MPI_noCTF
if myid == main_node:
active = EMUtil.get_all_attributes(stack, 'active')
list_of_particles = []
for im in xrange(len(active)):
if active[im]: list_of_particles.append(im)
del active
nima = len(list_of_particles)
else:
nima = 0
total_nima = bcast_number_to_all(nima, source_node=main_node)
if myid != main_node:
list_of_particles = [-1] * total_nima
list_of_particles = bcast_list_to_all(list_of_particles,
source_node=main_node)
image_start, image_end = MPI_start_end(total_nima, number_of_proc, myid)
# create a list of images for each node
list_of_particles = list_of_particles[image_start:image_end]
nima = len(list_of_particles)
if debug:
finfo.write("image_start, image_end: %d %d\n" %
(image_start, image_end))
finfo.flush()
data = EMData.read_images(stack, list_of_particles)
t_zero = Transform({
"type": "spider",
"phi": 0,
"theta": 0,
"psi": 0,
"tx": 0,
"ty": 0
})
transmulti = [[t_zero for i in xrange(nrefs)] for j in xrange(nima)]
for iref, im in ((iref, im) for iref in xrange(nrefs)
for im in xrange(nima)):
if nrefs == 1:
transmulti[im][iref] = data[im].get_attr("xform.projection")
else:
# if multi models, keep track of eulers for all models
try:
transmulti[im][iref] = data[im].get_attr("eulers_txty.%i" %
iref)
except:
data[im].set_attr("eulers_txty.%i" % iref, t_zero)
scoremulti = [[0.0 for i in xrange(nrefs)] for j in xrange(nima)]
pixelmulti = [[0.0 for i in xrange(nrefs)] for j in xrange(nima)]
ref_res = [0.0 for x in xrange(nrefs)]
apix = data[0].get_attr('apix_x')
# for oplane parameter, create cylindrical mask
if oplane is not None and myid == main_node:
from hfunctions import createCylMask
cmaskf = os.path.join(outdir, "mask3D_cyl.mrc")
mask3D = createCylMask(data, olmask, lmask, ilmask, cmaskf)
# if finding seam of helix, create wedge masks
if findseam is True:
wedgemask = []
for pf in xrange(nrefs):
wedgemask.append(EMData())
# wedgemask option
if wcmask is not None:
wcmask = get_input_from_string(wcmask)
if len(wcmask) != 3:
print_msg(
"Error: wcmask option requires 3 values: x y radius")
sys.exit()
# determine if particles have helix info:
try:
data[0].get_attr('h_angle')
original_data = []
boxmask = True
from hfunctions import createBoxMask
except:
boxmask = False
# prepare particles
for im in xrange(nima):
data[im].set_attr('ID', list_of_particles[im])
data[im].set_attr('pix_score', int(0))
if CTF:
# only phaseflip particles, not full CTF correction
ctf_params = data[im].get_attr("ctf")
st = Util.infomask(data[im], mask2D, False)
data[im] -= st[0]
data[im] = filt_ctf(data[im], ctf_params, sign=-1, binary=1)
data[im].set_attr('ctf_applied', 1)
# for window mask:
if boxmask is True:
h_angle = data[im].get_attr("h_angle")
original_data.append(data[im].copy())
bmask = createBoxMask(nx, apix, ou, lmask, h_angle)
data[im] *= bmask
del bmask
if debug:
finfo.write('%d loaded \n' % nima)
finfo.flush()
if myid == main_node:
# initialize data for the reference preparation function
ref_data = [mask3D, max(center, 0), None, None, None, None]
# for method -1, switch off centering in user function
from time import time
# this is needed for gathering of pixel errors
disps = []
recvcount = []
disps_score = []
recvcount_score = []
for im in xrange(number_of_proc):
if (im == main_node):
disps.append(0)
disps_score.append(0)
else:
disps.append(disps[im - 1] + recvcount[im - 1])
disps_score.append(disps_score[im - 1] + recvcount_score[im - 1])
ib, ie = MPI_start_end(total_nima, number_of_proc, im)
recvcount.append(ie - ib)
recvcount_score.append((ie - ib) * nrefs)
pixer = [0.0] * nima
cs = [0.0] * 3
total_iter = 0
volodd = EMData.read_images(ref_vol, xrange(nrefs))
voleve = EMData.read_images(ref_vol, xrange(nrefs))
if restart:
# recreate initial volumes from alignments stored in header
itout = "000_00"
for iref in xrange(nrefs):
if (nrefs == 1):
modout = ""
else:
modout = "_model_%02d" % (iref)
if (sort):
group = iref
for im in xrange(nima):
imgroup = data[im].get_attr('group')
if imgroup == iref:
data[im].set_attr('xform.projection',
transmulti[im][iref])
else:
group = int(999)
for im in xrange(nima):
data[im].set_attr('xform.projection', transmulti[im][iref])
fscfile = os.path.join(outdir, "fsc_%s%s" % (itout, modout))
vol[iref], fscc, volodd[iref], voleve[iref] = rec3D_MPI_noCTF(
data,
sym,
fscmask,
fscfile,
myid,
main_node,
index=group,
npad=recon_pad)
if myid == main_node:
if helicalrecon:
from hfunctions import processHelicalVol
vstep = None
if vertstep is not None:
vstep = (vdp[iref], vdphi[iref])
print_msg(
"Old rise and twist for model %i : %8.3f, %8.3f\n"
% (iref, dp[iref], dphi[iref]))
hvals = processHelicalVol(vol[iref], voleve[iref],
volodd[iref], iref, outdir,
itout, dp[iref], dphi[iref],
apix, hsearch, findseam, vstep,
wcmask)
(vol[iref], voleve[iref], volodd[iref], dp[iref],
dphi[iref], vdp[iref], vdphi[iref]) = hvals
print_msg(
"New rise and twist for model %i : %8.3f, %8.3f\n"
% (iref, dp[iref], dphi[iref]))
# get new FSC from symmetrized half volumes
fscc = fsc_mask(volodd[iref], voleve[iref], mask3D, rstep,
fscfile)
else:
vol[iref].write_image(
os.path.join(outdir, "vol_%s.hdf" % itout), -1)
if save_half is True:
volodd[iref].write_image(
os.path.join(outdir, "volodd_%s.hdf" % itout), -1)
voleve[iref].write_image(
os.path.join(outdir, "voleve_%s.hdf" % itout), -1)
if nmasks > 1:
# Read mask for multiplying
ref_data[0] = maskF[iref]
ref_data[2] = vol[iref]
ref_data[3] = fscc
# call user-supplied function to prepare reference image, i.e., center and filter it
vol[iref], cs, fl = ref_ali3d(ref_data)
vol[iref].write_image(
os.path.join(outdir, "volf_%s.hdf" % (itout)), -1)
if (apix == 1):
res_msg = "Models filtered at spatial frequency of:\t"
res = fl
else:
res_msg = "Models filtered at resolution of: \t"
res = apix / fl
ares = array2string(array(res), precision=2)
print_msg("%s%s\n\n" % (res_msg, ares))
bcast_EMData_to_all(vol[iref], myid, main_node)
# write out headers, under MPI writing has to be done sequentially
mpi_barrier(MPI_COMM_WORLD)
# projection matching
for N_step in xrange(lstp):
terminate = 0
Iter = -1
while (Iter < max_iter - 1 and terminate == 0):
Iter += 1
total_iter += 1
itout = "%03g_%02d" % (delta[N_step], Iter)
if myid == main_node:
print_msg(
"ITERATION #%3d, inner iteration #%3d\nDelta = %4.1f, an = %5.2f, xrange = %5.2f, yrange = %5.2f, step = %5.2f\n\n"
% (N_step, Iter, delta[N_step], an[N_step], xrng[N_step],
yrng[N_step], step[N_step]))
for iref in xrange(nrefs):
if myid == main_node: start_time = time()
volft, kb = prep_vol(vol[iref])
## constrain projections to out of plane parameter
theta1 = None
theta2 = None
if oplane is not None:
theta1 = 90 - oplane
theta2 = 90 + oplane
refrings = prepare_refrings(volft,
kb,
nx,
delta[N_step],
ref_a,
sym,
numr,
MPI=True,
phiEqpsi="Minus",
initial_theta=theta1,
delta_theta=theta2)
del volft, kb
if myid == main_node:
print_msg(
"Time to prepare projections for model %i: %s\n" %
(iref, legibleTime(time() - start_time)))
start_time = time()
for im in xrange(nima):
data[im].set_attr("xform.projection", transmulti[im][iref])
if an[N_step] == -1:
t1, peak, pixer[im] = proj_ali_incore(
data[im], refrings, numr, xrng[N_step],
yrng[N_step], step[N_step], finfo)
else:
t1, peak, pixer[im] = proj_ali_incore_local(
data[im], refrings, numr, xrng[N_step],
yrng[N_step], step[N_step], an[N_step], finfo)
#data[im].set_attr("xform.projection"%iref, t1)
if nrefs > 1:
data[im].set_attr("eulers_txty.%i" % iref, t1)
scoremulti[im][iref] = peak
from pixel_error import max_3D_pixel_error
# t1 is the current param, t2 is old
t2 = transmulti[im][iref]
pixelmulti[im][iref] = max_3D_pixel_error(t1, t2, numr[-3])
transmulti[im][iref] = t1
if myid == main_node:
print_msg("Time of alignment for model %i: %s\n" %
(iref, legibleTime(time() - start_time)))
start_time = time()
# gather scoring data from all processors
from mpi import mpi_gatherv
scoremultisend = sum(scoremulti, [])
pixelmultisend = sum(pixelmulti, [])
tmp = mpi_gatherv(scoremultisend, len(scoremultisend), MPI_FLOAT,
recvcount_score, disps_score, MPI_FLOAT,
main_node, MPI_COMM_WORLD)
tmp1 = mpi_gatherv(pixelmultisend, len(pixelmultisend), MPI_FLOAT,
recvcount_score, disps_score, MPI_FLOAT,
main_node, MPI_COMM_WORLD)
tmp = mpi_bcast(tmp, (total_nima * nrefs), MPI_FLOAT, 0,
MPI_COMM_WORLD)
tmp1 = mpi_bcast(tmp1, (total_nima * nrefs), MPI_FLOAT, 0,
MPI_COMM_WORLD)
tmp = map(float, tmp)
tmp1 = map(float, tmp1)
score = array(tmp).reshape(-1, nrefs)
pixelerror = array(tmp1).reshape(-1, nrefs)
score_local = array(scoremulti)
mean_score = score.mean(axis=0)
std_score = score.std(axis=0)
cut = mean_score - (cutoff * std_score)
cut2 = mean_score + (cutoff * std_score)
res_max = score_local.argmax(axis=1)
minus_cc = [0.0 for x in xrange(nrefs)]
minus_pix = [0.0 for x in xrange(nrefs)]
minus_ref = [0.0 for x in xrange(nrefs)]
#output pixel errors
if (myid == main_node):
from statistics import hist_list
lhist = 20
pixmin = pixelerror.min(axis=1)
region, histo = hist_list(pixmin, lhist)
if (region[0] < 0.0): region[0] = 0.0
print_msg(
"Histogram of pixel errors\n ERROR number of particles\n"
)
for lhx in xrange(lhist):
print_msg(" %10.3f %7d\n" % (region[lhx], histo[lhx]))
# Terminate if 95% within 1 pixel error
im = 0
for lhx in xrange(lhist):
if (region[lhx] > 1.0): break
im += histo[lhx]
print_msg("Percent of particles with pixel error < 1: %f\n\n" %
(im / float(total_nima) * 100))
term_cond = float(term) / 100
if (im / float(total_nima) > term_cond):
terminate = 1
print_msg("Terminating internal loop\n")
del region, histo
terminate = mpi_bcast(terminate, 1, MPI_INT, 0, MPI_COMM_WORLD)
terminate = int(terminate[0])
for im in xrange(nima):
if (sort == False):
data[im].set_attr('group', 999)
elif (mjump[N_step] == 1):
data[im].set_attr('group', int(res_max[im]))
pix_run = data[im].get_attr('pix_score')
if (pix_cutoff[N_step] == 1
and (terminate == 1 or Iter == max_iter - 1)):
if (pixelmulti[im][int(res_max[im])] > 1):
data[im].set_attr('pix_score', int(777))
if (score_local[im][int(res_max[im])] < cut[int(
res_max[im])]) or (two_tail and score_local[im][int(
res_max[im])] > cut2[int(res_max[im])]):
data[im].set_attr('group', int(888))
minus_cc[int(res_max[im])] = minus_cc[int(res_max[im])] + 1
if (pix_run == 777):
data[im].set_attr('group', int(777))
minus_pix[int(
res_max[im])] = minus_pix[int(res_max[im])] + 1
if (compare_ref_free != "-1") and (ref_free_cutoff[N_step] !=
-1) and (total_iter > 1):
id = data[im].get_attr('ID')
if id in rejects:
data[im].set_attr('group', int(666))
minus_ref[int(
res_max[im])] = minus_ref[int(res_max[im])] + 1
minus_cc_tot = mpi_reduce(minus_cc, nrefs, MPI_FLOAT, MPI_SUM, 0,
MPI_COMM_WORLD)
minus_pix_tot = mpi_reduce(minus_pix, nrefs, MPI_FLOAT, MPI_SUM, 0,
MPI_COMM_WORLD)
minus_ref_tot = mpi_reduce(minus_ref, nrefs, MPI_FLOAT, MPI_SUM, 0,
MPI_COMM_WORLD)
if (myid == main_node):
if (sort):
tot_max = score.argmax(axis=1)
res = bincount(tot_max)
else:
res = ones(nrefs) * total_nima
print_msg("Particle distribution: \t\t%s\n" % (res * 1.0))
afcut1 = res - minus_cc_tot
afcut2 = afcut1 - minus_pix_tot
afcut3 = afcut2 - minus_ref_tot
print_msg("Particle distribution after cc cutoff:\t\t%s\n" %
(afcut1))
print_msg("Particle distribution after pix cutoff:\t\t%s\n" %
(afcut2))
print_msg("Particle distribution after ref cutoff:\t\t%s\n\n" %
(afcut3))
res = [0.0 for i in xrange(nrefs)]
for iref in xrange(nrefs):
if (center == -1):
from utilities import estimate_3D_center_MPI, rotate_3D_shift
dummy = EMData()
cs[0], cs[1], cs[2], dummy, dummy = estimate_3D_center_MPI(
data, total_nima, myid, number_of_proc, main_node)
cs = mpi_bcast(cs, 3, MPI_FLOAT, main_node, MPI_COMM_WORLD)
cs = [-float(cs[0]), -float(cs[1]), -float(cs[2])]
rotate_3D_shift(data, cs)
if (sort):
group = iref
for im in xrange(nima):
imgroup = data[im].get_attr('group')
if imgroup == iref:
data[im].set_attr('xform.projection',
transmulti[im][iref])
else:
group = int(999)
for im in xrange(nima):
data[im].set_attr('xform.projection',
transmulti[im][iref])
if (nrefs == 1):
modout = ""
else:
modout = "_model_%02d" % (iref)
fscfile = os.path.join(outdir, "fsc_%s%s" % (itout, modout))
vol[iref], fscc, volodd[iref], voleve[iref] = rec3D_MPI_noCTF(
data,
sym,
fscmask,
fscfile,
myid,
main_node,
index=group,
npad=recon_pad)
if myid == main_node:
print_msg("3D reconstruction time for model %i: %s\n" %
(iref, legibleTime(time() - start_time)))
start_time = time()
# Compute Fourier variance
if fourvar:
outvar = os.path.join(outdir, "volVar_%s.hdf" % (itout))
ssnr_file = os.path.join(outdir, "ssnr_%s" % (itout))
varf = varf3d_MPI(data,
ssnr_text_file=ssnr_file,
mask2D=None,
reference_structure=vol[iref],
ou=last_ring,
rw=1.0,
npad=1,
CTF=None,
sign=1,
sym=sym,
myid=myid)
if myid == main_node:
print_msg(
"Time to calculate 3D Fourier variance for model %i: %s\n"
% (iref, legibleTime(time() - start_time)))
start_time = time()
varf = 1.0 / varf
varf.write_image(outvar, -1)
else:
varf = None
if myid == main_node:
if helicalrecon:
from hfunctions import processHelicalVol
vstep = None
if vertstep is not None:
vstep = (vdp[iref], vdphi[iref])
print_msg(
"Old rise and twist for model %i : %8.3f, %8.3f\n"
% (iref, dp[iref], dphi[iref]))
hvals = processHelicalVol(vol[iref], voleve[iref],
volodd[iref], iref, outdir,
itout, dp[iref], dphi[iref],
apix, hsearch, findseam,
vstep, wcmask)
(vol[iref], voleve[iref], volodd[iref], dp[iref],
dphi[iref], vdp[iref], vdphi[iref]) = hvals
print_msg(
"New rise and twist for model %i : %8.3f, %8.3f\n"
% (iref, dp[iref], dphi[iref]))
# get new FSC from symmetrized half volumes
fscc = fsc_mask(volodd[iref], voleve[iref], mask3D,
rstep, fscfile)
print_msg(
"Time to search and apply helical symmetry for model %i: %s\n\n"
% (iref, legibleTime(time() - start_time)))
start_time = time()
else:
vol[iref].write_image(
os.path.join(outdir, "vol_%s.hdf" % (itout)), -1)
if save_half is True:
volodd[iref].write_image(
os.path.join(outdir, "volodd_%s.hdf" % (itout)),
-1)
voleve[iref].write_image(
os.path.join(outdir, "voleve_%s.hdf" % (itout)),
-1)
if nmasks > 1:
# Read mask for multiplying
ref_data[0] = maskF[iref]
ref_data[2] = vol[iref]
ref_data[3] = fscc
ref_data[4] = varf
# call user-supplied function to prepare reference image, i.e., center and filter it
vol[iref], cs, fl = ref_ali3d(ref_data)
vol[iref].write_image(
os.path.join(outdir, "volf_%s.hdf" % (itout)), -1)
if (apix == 1):
res_msg = "Models filtered at spatial frequency of:\t"
res[iref] = fl
else:
res_msg = "Models filtered at resolution of: \t"
res[iref] = apix / fl
del varf
bcast_EMData_to_all(vol[iref], myid, main_node)
if compare_ref_free != "-1": compare_repro = True
if compare_repro:
outfile_repro = comp_rep(refrings, data, itout, modout,
vol[iref], group, nima, nx, myid,
main_node, outdir)
mpi_barrier(MPI_COMM_WORLD)
if compare_ref_free != "-1":
ref_free_output = os.path.join(
outdir, "ref_free_%s%s" % (itout, modout))
rejects = compare(compare_ref_free, outfile_repro,
ref_free_output, yrng[N_step],
xrng[N_step], rstep, nx, apix,
ref_free_cutoff[N_step],
number_of_proc, myid, main_node)
# retrieve alignment params from all processors
par_str = ['xform.projection', 'ID', 'group']
if nrefs > 1:
for iref in xrange(nrefs):
par_str.append('eulers_txty.%i' % iref)
if myid == main_node:
from 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:
ares = array2string(array(res), precision=2)
print_msg("%s%s\n\n" % (res_msg, ares))
dummy = EMData()
if full_output:
nimat = EMUtil.get_image_count(stack)
output_file = os.path.join(outdir, "paramout_%s" % itout)
foutput = open(output_file, 'w')
for im in xrange(nimat):
# save the parameters for each of the models
outstring = ""
dummy.read_image(stack, im, True)
param3d = dummy.get_attr('xform.projection')
g = dummy.get_attr("group")
# retrieve alignments in EMAN-format
pE = param3d.get_params('eman')
outstring += "%f\t%f\t%f\t%f\t%f\t%i\n" % (
pE["az"], pE["alt"], pE["phi"], pE["tx"], pE["ty"],
g)
foutput.write(outstring)
foutput.close()
del dummy
mpi_barrier(MPI_COMM_WORLD)
# mpi_finalize()
if myid == main_node: print_end_msg("ali3d_MPI")
def helicalshiftali_MPI(stack, maskfile=None, maxit=100, CTF=False, snr=1.0, Fourvar=False, search_rng=-1):
from applications import MPI_start_end
from utilities import model_circle, model_blank, get_image, peak_search, get_im, pad
from utilities import reduce_EMData_to_root, bcast_EMData_to_all, send_attr_dict, file_type, bcast_number_to_all, bcast_list_to_all
from statistics import varf2d_MPI
from fundamentals import fft, ccf, rot_shift3D, rot_shift2D, fshift
from utilities import get_params2D, set_params2D, chunks_distribution
from utilities import print_msg, print_begin_msg, print_end_msg
import os
import sys
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
from mpi import MPI_SUM, MPI_FLOAT, MPI_INT
from time import time
from pixel_error import ordersegments
from math import sqrt, atan2, tan, pi
nproc = mpi_comm_size(MPI_COMM_WORLD)
myid = mpi_comm_rank(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 xrange(total_nfils): inidl[i] = len(filaments[i])
linidl = sum(inidl)
nima = linidl
tfilaments = []
for i in xrange(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 xrange(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), "ehelix_MPI", 1,myid)
# balanced load
temp = chunks_distribution([[len(filaments[i]), i] for i in xrange(len(filaments))], nproc)[myid:myid+1][0]
filaments = [filaments[temp[i][1]] for i in xrange(len(temp))]
nfils = len(filaments)
#filaments = [[0,1]]
#print "filaments",filaments
list_of_particles = []
indcs = []
k = 0
for i in xrange(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)
print "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 xrange(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 filter import filt_ctf
from 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 utilities import info
for im in xrange(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', "helicalshiftali_MPI", 1,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 xrange(0,nx+2,2):
for j in xrange(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 xrange(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 xrange(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 xrange(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 xrange(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 xrange(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_reduce(sx_sum, 1, MPI_FLOAT, MPI_SUM, main_node, 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 xrange(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 xrange(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_barrier(MPI_COMM_WORLD)
par_str = ["xform.align2d", "ID"]
if myid == main_node:
from utilities import file_type
if(file_type(stack) == "bdb"):
from utilities import recv_attr_dict_bdb
recv_attr_dict_bdb(main_node, stack, data, par_str, 0, ldata, nproc)
else:
from 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 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 utilities import bcast_number_to_all, bcast_list_to_all, model_blank, bcast_EMData_to_all, reduce_EMData_to_root from morphology import threshold_outside from filter import filt_tanl from 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 xrange(nx): for y in xrange(ny): for z in xrange(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 xrange(myid, nz, number_of_proc): for x in xrange(nx): for y in xrange(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 main():
def params_3D_2D_NEW(phi, theta, psi, s2x, s2y, mirror):
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=0.2 --aa=0.1 --sym=symmetry --CTF"
parser = OptionParser(usage, version=SPARXVERSION)
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=10 , help="number of neighbouring projections")
parser.add_option("--no_norm", action="store_true", default=False, help="do not use normalization")
parser.add_option("--radiusvar", type="int" , default=-1 , help="radius for 3D var" )
parser.add_option("--npad", type="int" , default=2 , help="number of time to pad the original images")
parser.add_option("--sym" , type="string" , default="c1" , help="symmetry")
parser.add_option("--fl", type="float" , default=0.0 , help="stop-band frequency (Default - no filtration)")
parser.add_option("--aa", type="float" , default=0.0 , help="fall off of the filter (Default - no filtration)")
parser.add_option("--CTF", action="store_true", default=False, help="use CFT correction")
parser.add_option("--VERBOSE", action="store_true", default=False, help="Long output for debugging")
#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 on input consists of 2D variances (Default False)")
parser.add_option("--decimate", type="float", default=1.0, help="image decimate rate, a number large than 1. default is 1")
parser.add_option("--window", type="int", default=0, help="reduce images to a small image size without changing pixel_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)")
(options,args) = parser.parse_args()
#####
from mpi import mpi_init, mpi_comm_rank, mpi_comm_size, mpi_recv, MPI_COMM_WORLD, MPI_TAG_UB
from mpi import mpi_barrier, mpi_reduce, mpi_bcast, mpi_send, MPI_FLOAT, MPI_SUM, MPI_INT, MPI_MAX
from applications import MPI_start_end
from reconstruction import recons3d_em, recons3d_em_MPI
from reconstruction import recons3d_4nn_MPI, recons3d_4nn_ctf_MPI
from utilities import print_begin_msg, print_end_msg, print_msg
from utilities import read_text_row, get_image, get_im
from utilities import bcast_EMData_to_all, bcast_number_to_all
from 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
if options.symmetrize :
try:
sys.argv = mpi_init(len(sys.argv), sys.argv)
try:
number_of_proc = mpi_comm_size(MPI_COMM_WORLD)
if( number_of_proc > 1 ):
ERROR("Cannot use more than one CPU for symmetry prepration","sx3dvariability",1)
except:
pass
except:
pass
# Input
#instack = "Clean_NORM_CTF_start_wparams.hdf"
#instack = "bdb:data"
instack = args[0]
sym = options.sym
if( sym == "c1" ):
ERROR("Thre is no need to symmetrize stack for C1 symmetry","sx3dvariability",1)
if(instack[:4] !="bdb:"):
stack = "bdb:data"
delete_bdb(stack)
cmdexecute("sxcpy.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 xrange(ks):
delete_bdb("bdb:Q%1d"%k)
cmdexecute("e2bdb.py "+stack+" --makevstack=bdb:Q%1d"%k)
DB = db_open_dict("bdb:Q%1d"%k)
for i in xrange(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)
#cmdexecute("e2bdb.py "+stack+" --makevstack=bdb:Q%1d"%k)
#cmdexecute("sxheader.py bdb:Q%1d --params=xform.projection --import=ptsma%1d.txt"%(k,k))
DB.close()
delete_bdb("bdb:sdata")
cmdexecute("e2bdb.py . --makevstack=bdb:sdata --filt=Q")
#cmdexecute("ls EMAN2DB/sdata*")
a = get_im("bdb:sdata")
a.set_attr("variabilitysymmetry",sym)
a.write_image("bdb:sdata")
else:
sys.argv = mpi_init(len(sys.argv), sys.argv)
myid = mpi_comm_rank(MPI_COMM_WORLD)
number_of_proc = mpi_comm_size(MPI_COMM_WORLD)
main_node = 0
if len(args) == 1:
stack = args[0]
else:
print( "usage: " + usage)
print( "Please run '" + progname + " -h' for detailed options")
return 1
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", "sx3dvariability", 1, myid)
if options.VAR and options.SND:
ERROR("Only one of var and SND can be set!", "sx3dvariability", myid)
exit()
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", "sx3dvariability", 1, myid)
exit()
#if options.SND and (options.ave2D or options.ave3D):
# ERROR("When SND is set, the program cannot output ave2D or ave3D", "sx3dvariability", 1, myid)
# exit()
if options.nvec > 0 :
ERROR("PCA option not implemented", "sx3dvariability", 1, myid)
exit()
if options.nvec > 0 and options.ave3D == None:
ERROR("When doing PCA analysis, one must set ave3D", "sx3dvariability", myid=myid)
exit()
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:
print_begin_msg("sx3dvariability")
print_msg("%-70s: %s\n"%("Input stack", stack))
img_per_grp = options.img_per_grp
nvec = options.nvec
radiuspca = options.radiuspca
symbaselen = 0
if myid == main_node:
nima = EMUtil.get_image_count(stack)
img = get_image(stack)
nx = img.get_xsize()
ny = img.get_ysize()
if options.sym != "c1" :
imgdata = get_im(stack)
try:
i = imgdata.get_attr("variabilitysymmetry")
if(i != options.sym):
ERROR("The symmetry provided does not agree with the symmetry of the input stack", "sx3dvariability", myid=myid)
except:
ERROR("Input stack is not prepared for symmetry, please follow instructions", "sx3dvariability", myid=myid)
from 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", "sx3dvariability", myid=myid)
symbaselen = nima/i
else: symbaselen = nima
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)
Tracker ={}
Tracker["nx"] =nx
Tracker["ny"] =ny
Tracker["total_stack"]=nima
if options.decimate==1.:
if options.window !=0:
nx = options.window
ny = options.window
else:
if options.window ==0:
nx = int(nx/options.decimate)
ny = int(ny/options.decimate)
else:
nx = int(options.window/options.decimate)
ny = nx
symbaselen = bcast_number_to_all(symbaselen)
if radiuspca == -1: radiuspca = nx/2-2
if myid == main_node:
print_msg("%-70s: %d\n"%("Number of projection", nima))
img_begin, img_end = MPI_start_end(nima, number_of_proc, myid)
"""
if options.SND:
from projection import prep_vol, prgs
from statistics import im_diff
from utilities import get_im, model_circle, get_params_proj, set_params_proj
from utilities import get_ctf, generate_ctf
from 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:
#varList = EMData.read_images(stack, range(img_begin, img_end))
varList = []
this_image = EMData()
for index_of_particle in xrange(img_begin,img_end):
this_image.read_image(stack,index_of_particle)
varList.append(image_decimate_window_xform_ctf(img,options.decimate,options.window,options.CTF))
else:
from utilities import bcast_number_to_all, bcast_list_to_all, send_EMData, recv_EMData
from utilities import set_params_proj, get_params_proj, params_3D_2D, get_params2D, set_params2D, compose_transform2
from utilities import model_blank, nearest_proj, model_circle
from applications import pca
from statistics import avgvar, avgvar_ctf, ccc
from filter import filt_tanl
from morphology import threshold, square_root
from projection import project, prep_vol, prgs
from sets import Set
if myid == main_node:
t1 = time()
proj_angles = []
aveList = []
tab = EMUtil.get_all_attributes(stack, 'xform.projection')
for i in xrange(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()
print_msg("%-70s: %d\n"%("Number of neighboring projections", img_per_grp))
print_msg("...... Finding neighboring projections\n")
if options.VERBOSE:
print "Number of images per group: ", img_per_grp
print "Now grouping projections"
proj_angles.sort()
proj_angles_list = [0.0]*(nima*4)
if myid == main_node:
for i in xrange(nima):
proj_angles_list[i*4] = proj_angles[i][1]
proj_angles_list[i*4+1] = proj_angles[i][2]
proj_angles_list[i*4+2] = proj_angles[i][3]
proj_angles_list[i*4+3] = proj_angles[i][4]
proj_angles_list = bcast_list_to_all(proj_angles_list, myid, main_node)
proj_angles = []
for i in xrange(nima):
proj_angles.append([proj_angles_list[i*4], proj_angles_list[i*4+1], proj_angles_list[i*4+2], int(proj_angles_list[i*4+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)
if options.VERBOSE:
print "On node %2d, number of images needed to be read = %5d"%(myid, len(all_proj))
index = {}
for i in xrange(len(all_proj)): index[all_proj[i]] = i
mpi_barrier(MPI_COMM_WORLD)
if myid == main_node:
print_msg("%-70s: %.2f\n"%("Finding neighboring projections lasted [s]", time()-t2))
print_msg("%-70s: %d\n"%("Number of groups processed on the main node", len(proj_list)))
if options.VERBOSE:
print "Grouping projections took: ", (time()-t2)/60 , "[min]"
print "Number of groups on main node: ", len(proj_list)
mpi_barrier(MPI_COMM_WORLD)
if myid == main_node:
print_msg("...... calculating the stack of 2D variances \n")
if options.VERBOSE:
print "Now calculating the stack of 2D variances"
proj_params = [0.0]*(nima*5)
aveList = []
varList = []
if nvec > 0:
eigList = [[] for i in xrange(nvec)]
if options.VERBOSE: print "Begin to read images on processor %d"%(myid)
ttt = time()
#imgdata = EMData.read_images(stack, all_proj)
img = EMData()
imgdata = []
for index_of_proj in xrange(len(all_proj)):
img.read_image(stack, all_proj[index_of_proj])
dmg = image_decimate_window_xform_ctf(img,options.decimate,options.window,options.CTF)
#print dmg.get_xsize(), "init"
imgdata.append(dmg)
if options.VERBOSE:
print "Reading images on processor %d done, time = %.2f"%(myid, time()-ttt)
print "On processor %d, we got %d images"%(myid, len(imgdata))
mpi_barrier(MPI_COMM_WORLD)
'''
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 applications import prepare_2d_forPCA
from utilities import model_blank
for i in xrange(len(proj_list)):
ki = proj_angles[proj_list[i][0]][3]
if ki >= symbaselen: continue
mi = index[ki]
phiM, thetaM, psiM, s2xM, s2yM = get_params_proj(imgdata[mi])
grp_imgdata = []
for j in xrange(img_per_grp):
mj = index[proj_angles[proj_list[i][j]][3]]
phi, theta, psi, s2x, s2y = get_params_proj(imgdata[mj])
alpha, sx, sy, mirror = params_3D_2D_NEW(phi, theta, psi, s2x, s2y, mirror_list[i][j])
if thetaM <= 90:
if mirror == 0: alpha, sx, sy, scale = compose_transform2(alpha, sx, sy, 1.0, phiM-phi, 0.0, 0.0, 1.0)
else: alpha, sx, sy, scale = compose_transform2(alpha, sx, sy, 1.0, 180-(phiM-phi), 0.0, 0.0, 1.0)
else:
if mirror == 0: alpha, sx, sy, scale = compose_transform2(alpha, sx, sy, 1.0, -(phiM-phi), 0.0, 0.0, 1.0)
else: alpha, sx, sy, scale = compose_transform2(alpha, sx, sy, 1.0, -(180-(phiM-phi)), 0.0, 0.0, 1.0)
set_params2D(imgdata[mj], [alpha, sx, sy, mirror, 1.0])
grp_imgdata.append(imgdata[mj])
#print grp_imgdata[j].get_xsize(), imgdata[mj].get_xsize()
if not options.no_norm:
#print grp_imgdata[j].get_xsize()
mask = model_circle(nx/2-2, nx, nx)
for k in xrange(img_per_grp):
ave, std, minn, maxx = Util.infomask(grp_imgdata[k], mask, False)
grp_imgdata[k] -= ave
grp_imgdata[k] /= std
del mask
if options.fl > 0.0:
from filter import filt_ctf, filt_table
from fundamentals import fft, window2d
nx2 = 2*nx
ny2 = 2*ny
if options.CTF:
from utilities import pad
for k in xrange(img_per_grp):
grp_imgdata[k] = window2d(fft( filt_tanl( filt_ctf(fft(pad(grp_imgdata[k], nx2, ny2, 1,0.0)), grp_imgdata[k].get_attr("ctf"), binary=1), options.fl, options.aa) ),nx,ny)
#grp_imgdata[k] = window2d(fft( filt_table( filt_tanl( filt_ctf(fft(pad(grp_imgdata[k], nx2, ny2, 1,0.0)), grp_imgdata[k].get_attr("ctf"), binary=1), options.fl, options.aa), fifi) ),nx,ny)
#grp_imgdata[k] = filt_tanl(grp_imgdata[k], options.fl, options.aa)
else:
for k in xrange(img_per_grp):
grp_imgdata[k] = filt_tanl( grp_imgdata[k], options.fl, options.aa)
#grp_imgdata[k] = window2d(fft( filt_table( filt_tanl( filt_ctf(fft(pad(grp_imgdata[k], nx2, ny2, 1,0.0)), grp_imgdata[k].get_attr("ctf"), binary=1), options.fl, options.aa), fifi) ),nx,ny)
#grp_imgdata[k] = filt_tanl(grp_imgdata[k], options.fl, options.aa)
else:
from utilities import pad, read_text_file
from filter import filt_ctf, filt_table
from fundamentals import fft, window2d
nx2 = 2*nx
ny2 = 2*ny
if options.CTF:
from utilities import pad
for k in xrange(img_per_grp):
grp_imgdata[k] = window2d( fft( filt_ctf(fft(pad(grp_imgdata[k], nx2, ny2, 1,0.0)), grp_imgdata[k].get_attr("ctf"), binary=1) ) , nx,ny)
#grp_imgdata[k] = window2d(fft( filt_table( filt_tanl( filt_ctf(fft(pad(grp_imgdata[k], nx2, ny2, 1,0.0)), grp_imgdata[k].get_attr("ctf"), binary=1), options.fl, options.aa), fifi) ),nx,ny)
#grp_imgdata[k] = filt_tanl(grp_imgdata[k], options.fl, options.aa)
'''
if i < 10 and myid == main_node:
for k in xrange(10):
grp_imgdata[k].write_image("grp%03d.hdf"%i, k)
'''
"""
if myid == main_node and i==0:
for pp in xrange(len(grp_imgdata)):
grp_imgdata[pp].write_image("pp.hdf", pp)
"""
ave, grp_imgdata = prepare_2d_forPCA(grp_imgdata)
"""
if myid == main_node and i==0:
for pp in xrange(len(grp_imgdata)):
grp_imgdata[pp].write_image("qq.hdf", pp)
"""
var = model_blank(nx,ny)
for q in grp_imgdata: Util.add_img2( var, q )
Util.mul_scalar( var, 1.0/(len(grp_imgdata)-1))
# Switch to std dev
var = square_root(threshold(var))
#if options.CTF: ave, var = avgvar_ctf(grp_imgdata, mode="a")
#else: ave, var = avgvar(grp_imgdata, mode="a")
"""
if myid == main_node:
ave.write_image("avgv.hdf",i)
var.write_image("varv.hdf",i)
"""
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)
if options.VERBOSE:
print "%5.2f%% done on processor %d"%(i*100.0/len(proj_list), myid)
if nvec > 0:
eig = pca(input_stacks=grp_imgdata, subavg="", mask_radius=radiuspca, nvec=nvec, incore=True, shuffle=False, genbuf=True)
for k in xrange(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)
"""
del imgdata
# To this point, all averages, variances, and eigenvectors are computed
if options.ave2D:
from fundamentals import fpol
if myid == main_node:
km = 0
for i in xrange(number_of_proc):
if i == main_node :
for im in xrange(len(aveList)):
aveList[im].write_image(options.ave2D, km)
km += 1
else:
nl = mpi_recv(1, MPI_INT, i, MPI_TAG_UB, MPI_COMM_WORLD)
nl = int(nl[0])
for im in xrange(nl):
ave = recv_EMData(i, im+i+70000)
"""
nm = mpi_recv(1, MPI_INT, i, MPI_TAG_UB, MPI_COMM_WORLD)
nm = int(nm[0])
members = mpi_recv(nm, MPI_INT, i, MPI_TAG_UB, MPI_COMM_WORLD)
ave.set_attr('members', map(int, members))
members = mpi_recv(nm, MPI_FLOAT, i, MPI_TAG_UB, MPI_COMM_WORLD)
ave.set_attr('pix_err', map(float, members))
members = mpi_recv(3, MPI_FLOAT, i, MPI_TAG_UB, MPI_COMM_WORLD)
ave.set_attr('refprojdir', map(float, members))
"""
tmpvol=fpol(ave, Tracker["nx"],Tracker["nx"],Tracker["nx"])
tmpvol.write_image(options.ave2D, km)
km += 1
else:
mpi_send(len(aveList), 1, MPI_INT, main_node, MPI_TAG_UB, MPI_COMM_WORLD)
for im in xrange(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, MPI_TAG_UB, MPI_COMM_WORLD)
mpi_send(members, len(members), MPI_INT, main_node, MPI_TAG_UB, MPI_COMM_WORLD)
members = aveList[im].get_attr('pix_err')
mpi_send(members, len(members), MPI_FLOAT, main_node, MPI_TAG_UB, MPI_COMM_WORLD)
try:
members = aveList[im].get_attr('refprojdir')
mpi_send(members, 3, MPI_FLOAT, main_node, MPI_TAG_UB, MPI_COMM_WORLD)
except:
mpi_send([-999.0,-999.0,-999.0], 3, MPI_FLOAT, main_node, MPI_TAG_UB, MPI_COMM_WORLD)
"""
if options.ave3D:
from fundamentals import fpol
if options.VERBOSE:
print "Reconstructing 3D average volume"
ave3D = recons3d_4nn_MPI(myid, aveList, symmetry=options.sym, npad=options.npad)
bcast_EMData_to_all(ave3D, myid)
if myid == main_node:
ave3D=fpol(ave3D,Tracker["nx"],Tracker["nx"],Tracker["nx"])
ave3D.write_image(options.ave3D)
print_msg("%-70s: %s\n"%("Writing to the disk volume reconstructed from averages as", options.ave3D))
del ave, var, proj_list, stack, phi, theta, psi, s2x, s2y, alpha, sx, sy, mirror, aveList
if nvec > 0:
for k in xrange(nvec):
if options.VERBOSE:
print "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("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 xrange(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])
print " 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 fundamentals import fpol
if myid == main_node:
km = 0
for i in xrange(number_of_proc):
if i == main_node :
for im in xrange(len(varList)):
tmpvol=fpol(varList[im], Tracker["nx"], Tracker["nx"],1)
tmpvol.write_image(options.var2D, km)
km += 1
else:
nl = mpi_recv(1, MPI_INT, i, MPI_TAG_UB, MPI_COMM_WORLD)
nl = int(nl[0])
for im in xrange(nl):
ave = recv_EMData(i, im+i+70000)
tmpvol=fpol(ave, Tracker["nx"], Tracker["nx"],1)
tmpvol.write_image(options.var2D, km)
km += 1
else:
mpi_send(len(varList), 1, MPI_INT, main_node, MPI_TAG_UB, MPI_COMM_WORLD)
for im in xrange(len(varList)):
send_EMData(varList[im], main_node, im+myid+70000)# What with the attributes??
mpi_barrier(MPI_COMM_WORLD)
if options.var3D:
if myid == main_node and options.VERBOSE:
print "Reconstructing 3D variability volume"
t6 = time()
radiusvar = options.radiusvar
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 fundamentals import fpol
res =fpol(res, Tracker["nx"], Tracker["nx"], Tracker["nx"])
res.write_image(options.var3D)
if myid == main_node:
print_msg("%-70s: %.2f\n"%("Reconstructing 3D variability took [s]", time()-t6))
if options.VERBOSE:
print "Reconstruction took: %.2f [min]"%((time()-t6)/60)
if myid == main_node:
print_msg("%-70s: %.2f\n"%("Total time for these computations [s]", time()-t0))
if options.VERBOSE:
print "Total time for these computations: %.2f [min]"%((time()-t0)/60)
print_end_msg("sx3dvariability")
global_def.BATCH = False
from mpi import mpi_finalize
mpi_finalize()
def resample( prjfile, outdir, bufprefix, nbufvol, nvol, seedbase,\
delta, d, snr, CTF, npad,\
MPI, myid, ncpu, verbose = 0 ):
from utilities import even_angles
from random import seed, jumpahead, shuffle
import os
from sys import exit
nprj = EMUtil.get_image_count( prjfile )
if MPI:
from mpi import mpi_barrier, MPI_COMM_WORLD
if myid == 0:
if os.path.exists(outdir): nx = 1
else: nx = 0
else: nx = 0
ny = bcast_number_to_all(nx, source_node = 0)
if ny == 1: ERROR('Output directory exists, please change the name and restart the program', "resample", 1,myid)
mpi_barrier(MPI_COMM_WORLD)
if myid == 0:
os.mkdir(outdir)
mpi_barrier(MPI_COMM_WORLD)
else:
if os.path.exists(outdir):
ERROR('Output directory exists, please change the name and restart the program', "resample", 1,0)
os.mkdir(outdir)
if(verbose == 1): finfo=open( os.path.join(outdir, "progress%04d.txt" % myid), "w" )
else: finfo = None
#print " before evenangles",myid
from utilities import getvec
from numpy import array, reshape
refa = even_angles(delta)
nrefa = len(refa)
refnormal = zeros((nrefa,3),'float32')
tetref = [0.0]*nrefa
for i in xrange(nrefa):
tr = getvec( refa[i][0], refa[i][1] )
for j in xrange(3): refnormal[i][j] = tr[j]
tetref[i] = refa[i][1]
del refa
vct = array([0.0]*(3*nprj),'float32')
if myid == 0:
print " will read ",myid
tr = EMUtil.get_all_attributes(prjfile,'xform.projection')
tetprj = [0.0]*nprj
for i in xrange(nprj):
temp = tr[i].get_params("spider")
tetprj[i] = temp["theta"]
if(tetprj[i] > 90.0): tetprj[i] = 180.0 - tetprj[i]
vct[3*i+0] = tr[i].at(2,0)
vct[3*i+1] = tr[i].at(2,1)
vct[3*i+2] = tr[i].at(2,2)
del tr
else:
tetprj = [0.0]*nprj
#print " READ ",myid
if MPI:
#print " will bcast",myid
from mpi import mpi_bcast, MPI_FLOAT, MPI_COMM_WORLD
vct = mpi_bcast(vct,len(vct),MPI_FLOAT,0,MPI_COMM_WORLD)
from utilities import bcast_list_to_all
tetprj = bcast_list_to_all(tetprj, myid, 0)
#print " reshape ",myid
vct = reshape(vct,(nprj,3))
assignments = [[] for i in xrange(nrefa)]
dspn = 1.25*delta
for k in xrange(nprj):
best_s = -1.0
best_i = -1
for i in xrange( nrefa ):
if(abs(tetprj[k] - tetref[i]) <= dspn):
s = abs(refnormal[i][0]*vct[k][0] + refnormal[i][1]*vct[k][1] + refnormal[i][2]*vct[k][2])
if s > best_s:
best_s = s
best_i = i
assignments[best_i].append(k)
am = len(assignments[0])
mufur = 1.0/am
for i in xrange(1,len(assignments)):
ti = len(assignments[i])
am = min(am, ti)
if(ti>0): mufur += 1.0/ti
del tetprj,tetref
dp = 1.0 - d # keep that many in each direction
keep = int(am*dp +0.5)
mufur = keep*nrefa/(1.0 - mufur*keep/float(nrefa))
if myid == 0:
print " Number of projections ",nprj,". Number of reference directions ",nrefa,", multiplicative factor for the variance ",mufur
print " Minimum number of assignments ",am," Number of projections used per stratum ", keep," Number of projections in resampled structure ",int(am*dp +0.5)*nrefa
if am <2 or am == keep:
print "incorrect settings"
exit() # FIX
if(seedbase < 1):
seed()
jumpahead(17*myid+123)
else:
seed(seedbase)
jumpahead(17*myid+123)
volfile = os.path.join(outdir, "bsvol%04d.hdf" % myid)
from random import randint
niter = nvol/ncpu/nbufvol
for kiter in xrange(niter):
if(verbose == 1):
finfo.write( "Iteration %d: \n" % kiter )
finfo.flush()
iter_start = time()
# the following has to be converted to resample mults=1 means take given projection., mults=0 means omit
mults = [ [0]*nprj for i in xrange(nbufvol) ]
for i in xrange(nbufvol):
for l in xrange(nrefa):
mass = assignments[l][:]
shuffle(mass)
mass = mass[:keep]
mass.sort()
#print l, " * ",mass
for k in xrange(keep):
mults[i][mass[k]] = 1
'''
lout = []
for l in xrange(len(mults[i])):
if mults[i][l] == 1: lout.append(l)
write_text_file(lout, os.path.join(outdir, "list%04d_%03d.txt" %(i, myid)))
del lout
'''
del mass
rectors, fftvols, wgtvols = resample_prepare( prjfile, nbufvol, snr, CTF, npad )
resample_insert( bufprefix, fftvols, wgtvols, mults, CTF, npad, finfo )
del mults
resample_finish( rectors, fftvols, wgtvols, volfile, kiter, nprj, finfo )
rectors = None
fftvols = None
wgtvols = None
if(verbose == 1):
finfo.write( "time for iteration: %10.3f\n" % (time() - iter_start) )
finfo.flush()
