def main(args):
log = logging.getLogger('root')
hdlr = logging.StreamHandler(sys.stdout)
log.addHandler(hdlr)
log.setLevel(logging.getLevelName(args.loglevel.upper()))
# apix = args.apix = hdr["xlen"] / hdr["nx"]
for fn in args.input:
if not (fn.endswith(".star") or fn.endswith(".mrcs")
or fn.endswith(".mrc")):
log.error("Only .star, .mrc, and .mrcs files supported")
return 1
first_ptcl = 0
dfs = []
with mrc.ZSliceWriter(args.output) as writer:
for fn in args.input:
if fn.endswith(".star"):
df = star.parse_star(fn, keep_index=False)
star.augment_star_ucsf(df)
df = df.sort_values([
star.UCSF.IMAGE_ORIGINAL_PATH,
star.UCSF.IMAGE_ORIGINAL_INDEX
])
gb = df.groupby(star.UCSF.IMAGE_ORIGINAL_PATH)
for name, g in gb:
with mrc.ZSliceReader(name) as reader:
for i in g[star.UCSF.IMAGE_ORIGINAL_INDEX].values:
writer.write(reader.read(i))
else:
with mrc.ZSliceReader(fn) as reader:
for img in reader:
writer.write(img)
df = pd.DataFrame(
{star.UCSF.IMAGE_ORIGINAL_INDEX: np.arange(reader.nz)})
df[star.UCSF.IMAGE_ORIGINAL_PATH] = fn
if args.star is not None:
df[star.UCSF.IMAGE_INDEX] = np.arange(first_ptcl,
first_ptcl + df.shape[0])
df[star.UCSF.IMAGE_PATH] = writer.path
df["index"] = df[star.UCSF.IMAGE_INDEX]
star.simplify_star_ucsf(df)
dfs.append(df)
first_ptcl += df.shape[0]
if args.star is not None:
df = pd.concat(dfs, join="inner")
# df = pd.concat(dfs)
# df = df.dropna(df, axis=1, how="any")
star.write_star(args.star, df, reindex=True)
return 0
def read_first_mrc(particles):
'''
Read first mrc in a star file
'''
zreader = mrc.ZSliceReader(particles[star.UCSF.IMAGE_ORIGINAL_PATH].iloc[0])
first_ptcl = particles.iloc[0]
p1r = zreader.read(first_ptcl[star.UCSF.IMAGE_ORIGINAL_INDEX])
return p1r
def producer(pool,
queue,
submap_ft,
refmap_ft,
fname,
particles,
idx,
stack,
sx,
sy,
s,
a,
apix,
def1,
def2,
angast,
phase,
kv,
ac,
cs,
az,
el,
sk,
xshift,
yshift,
new_idx,
new_stack,
coefs_method,
r,
nr,
fftthreads=1):
log = logging.getLogger('root')
log.debug("Producing %s" % fname)
zreader = mrc.ZSliceReader(stack[particles[0]])
for i in particles:
log.debug("Produce %d@%s" % (idx[i], stack[i]))
# p1r = mrc.read_imgs(stack[i], idx[i] - 1, compat="relion")
p1r = zreader.read(idx[i] - 1)
log.debug("Apply")
ri = pool.apply_async(
subtract_outer,
(p1r, submap_ft, refmap_ft, sx, sy, s, a, apix, def1[i], def2[i],
angast[i], phase[i], kv[i], ac[i], cs[i], az[i], el[i], sk[i],
xshift[i], yshift[i], coefs_method, r, nr),
{"fftthreads": fftthreads})
log.debug("Put")
queue.put((new_idx[i], ri), block=True)
log.debug("Queue for %s is size %d" % (stack[i], queue.qsize()))
zreader.close()
# Either the poison-pill-put blocks, we have multiple queues and
# consumers, or the consumer knows maps results to multiple files.
log.debug("Put poison pill")
queue.put((-1, None), block=False)
def producer(pool, queue, submap_ft, refmap_ft, fname, particles,
sx, sy, s, a, apix, coefs_method, r, nr, fftthreads=1, crop=None, pfac=2):
log = logging.getLogger('root')
log.debug("Producing %s" % fname)
zreader = mrc.ZSliceReader(particles[star.UCSF.IMAGE_ORIGINAL_PATH].iloc[0])
for i, ptcl in particles.iterrows():
log.debug("Produce %d@%s" % (ptcl[star.UCSF.IMAGE_ORIGINAL_INDEX], ptcl[star.UCSF.IMAGE_ORIGINAL_PATH]))
# p1r = mrc.read_imgs(stack[i], idx[i] - 1, compat="relion")
p1r = zreader.read(ptcl[star.UCSF.IMAGE_ORIGINAL_INDEX])
log.debug("Apply")
ri = pool.apply_async(
subtract_outer,
(p1r, ptcl, submap_ft, refmap_ft, sx, sy, s, a, apix, coefs_method, r, nr),
{"fftthreads": fftthreads, "crop": crop, "pfac": pfac})
log.debug("Put")
queue.put((ptcl[star.UCSF.IMAGE_INDEX], ri), block=True)
log.debug("Queue for %s is size %d" % (ptcl[star.UCSF.IMAGE_ORIGINAL_PATH], queue.qsize()))
zreader.close()
log.debug("Put poison pill")
queue.put((-1, None), block=True)
def producer(pool, queue, submap_ft, refmap_ft, fname, particles,
sx, sy, s, a, apix, coefs_method, r, nr, fftthreads=1):
log = logging.getLogger('root')
log.debug("Producing %s" % fname)
zreader = mrc.ZSliceReader(particles[star.UCSF.IMAGE_ORIGINAL_PATH].iloc[0])
for i, ptcl in particles.iterrows():
log.debug("Produce %d@%s" % (ptcl[star.UCSF.IMAGE_ORIGINAL_INDEX], ptcl[star.UCSF.IMAGE_ORIGINAL_PATH]))
# p1r = mrc.read_imgs(stack[i], idx[i] - 1, compat="relion")
p1r = zreader.read(ptcl[star.UCSF.IMAGE_ORIGINAL_INDEX])
log.debug("Apply")
ri = pool.apply_async(
subtract_outer,
(p1r, ptcl, submap_ft, refmap_ft, sx, sy, s, a, apix, coefs_method, r, nr),
{"fftthreads": fftthreads})
log.debug("Put")
queue.put((ptcl[star.UCSF.IMAGE_INDEX], ri), block=True)
log.debug("Queue for %s is size %d" % (ptcl[star.UCSF.IMAGE_ORIGINAL_PATH], queue.qsize()))
zreader.close()
# Either the poison-pill-put blocks, we have multiple queues and
# consumers, or the consumer knows maps results to multiple files.
log.debug("Put poison pill")
queue.put((-1, None), block=False)
def particle_xcorr(ptcl, refmap_ft):
r = util.euler2rot(*np.deg2rad(ptcl[star.Relion.ANGLES]))
proj = vop.interpolate_slice_numba(refmap_ft, r)
c = ctf.eval_ctf(s / apix,
a,
def1[i],
def2[i],
angast[i],
phase[i],
kv[i],
ac[i],
cs[i],
bf=0,
lp=2 * apix)
pshift = np.exp(-2 * np.pi * 1j * (-xshift[i] * sx + -yshift * sy))
proj_ctf = proj * pshift * c
with mrc.ZSliceReader(ptcl[star.Relion.IMAGE_NAME]) as f:
exp_image_fft = rfft2(fftshift(f.read(i)))
xcor_fft = exp_image_fft * proj_ctf
xcor = fftshift(irfft2(xcor_fft))
return xcor
def main(args):
log = logging.getLogger('root')
hdlr = logging.StreamHandler(sys.stdout)
log.addHandler(hdlr)
log.setLevel(logging.getLevelName(args.loglevel.upper()))
df = star.parse_star(args.input, keep_index=False)
star.augment_star_ucsf(df)
maxshift = np.round(np.max(np.abs(df[star.Relion.ORIGINS].values)))
if args.map is not None:
if args.map.endswith(".npy"):
log.info("Reading precomputed 3D FFT of volume")
f3d = np.load(args.map)
log.info("Finished reading 3D FFT of volume")
if args.size is None:
args.size = (f3d.shape[0] - 3) // args.pfac
else:
vol = mrc.read(args.map, inc_header=False, compat="relion")
if args.mask is not None:
mask = mrc.read(args.mask, inc_header=False, compat="relion")
vol *= mask
if args.size is None:
args.size = vol.shape[0]
if args.crop is not None and args.size // 2 < maxshift + args.crop // 2:
log.error(
"Some shifts are too large to crop (maximum crop is %d)" %
(args.size - 2 * maxshift))
return 1
log.info("Preparing 3D FFT of volume")
f3d = vop.vol_ft(vol, pfac=args.pfac, threads=args.threads)
log.info("Finished 3D FFT of volume")
else:
log.error("Please supply a map")
return 1
sz = (f3d.shape[0] - 3) // args.pfac
apix = star.calculate_apix(df) * np.double(args.size) / sz
sx, sy = np.meshgrid(np.fft.rfftfreq(sz), np.fft.fftfreq(sz))
s = np.sqrt(sx**2 + sy**2)
a = np.arctan2(sy, sx)
log.info("Projection size is %d, unpadded volume size is %d" %
(args.size, sz))
log.info("Effective pixel size is %f A/px" % apix)
if args.subtract and args.size != sz:
log.error("Volume and projections must be same size when subtracting")
return 1
if args.crop is not None and args.size // 2 < maxshift + args.crop // 2:
log.error("Some shifts are too large to crop (maximum crop is %d)" %
(args.size - 2 * maxshift))
return 1
ift = None
with mrc.ZSliceWriter(args.output, psz=apix) as zsw:
for i, p in df.iterrows():
f2d = project(f3d,
p,
s,
sx,
sy,
a,
pfac=args.pfac,
apply_ctf=args.ctf,
size=args.size,
flip_phase=args.flip)
if ift is None:
ift = irfft2(f2d.copy(),
threads=args.threads,
planner_effort="FFTW_ESTIMATE",
auto_align_input=True,
auto_contiguous=True)
proj = fftshift(
ift(f2d.copy(),
np.zeros(ift.output_shape, dtype=ift.output_dtype)))
log.debug("%f +/- %f" % (np.mean(proj), np.std(proj)))
if args.subtract:
with mrc.ZSliceReader(p["ucsfImagePath"]) as zsr:
img = zsr.read(p["ucsfImageIndex"])
log.debug("%f +/- %f" % (np.mean(img), np.std(img)))
proj = img - proj
if args.crop is not None:
orihalf = args.size // 2
newhalf = args.crop // 2
x = orihalf - np.int(np.round(p[star.Relion.ORIGINX]))
y = orihalf - np.int(np.round(p[star.Relion.ORIGINY]))
proj = proj[y - newhalf:y + newhalf, x - newhalf:x + newhalf]
zsw.write(proj)
log.debug(
"%d@%s: %d/%d" %
(p["ucsfImageIndex"], p["ucsfImagePath"], i + 1, df.shape[0]))
if args.star is not None:
log.info("Writing output .star file")
if args.crop is not None:
df = star.recenter(df, inplace=True)
if args.subtract:
df[star.UCSF.IMAGE_ORIGINAL_PATH] = df[star.UCSF.IMAGE_PATH]
df[star.UCSF.IMAGE_ORIGINAL_INDEX] = df[star.UCSF.IMAGE_INDEX]
df[star.UCSF.IMAGE_PATH] = args.output
df[star.UCSF.IMAGE_INDEX] = np.arange(df.shape[0])
star.simplify_star_ucsf(df)
star.write_star(args.star, df)
return 0
def main(args):
log = logging.getLogger('root')
hdlr = logging.StreamHandler(sys.stdout)
log.addHandler(hdlr)
log.setLevel(logging.getLevelName(args.loglevel.upper()))
# apix = args.apix = hdr["xlen"] / hdr["nx"]
for fn in args.input:
if not (fn.endswith(".star") or fn.endswith(".mrcs") or
fn.endswith(".mrc") or fn.endswith(".par")):
log.error("Only .star, .mrc, .mrcs, and .par files supported")
return 1
first_ptcl = 0
dfs = []
with mrc.ZSliceWriter(args.output) as writer:
for fn in args.input:
if fn.endswith(".star"):
df = star.parse_star(fn, augment=True)
if args.cls is not None:
df = star.select_classes(df, args.cls)
star.set_original_fields(df, inplace=True)
if args.resort:
df = df.sort_values([star.UCSF.IMAGE_ORIGINAL_PATH,
star.UCSF.IMAGE_ORIGINAL_INDEX])
for idx, row in df.iterrows():
if args.stack_path is not None:
input_stack_path = os.path.join(args.stack_path, row[star.UCSF.IMAGE_ORIGINAL_PATH])
else:
input_stack_path = row[star.UCSF.IMAGE_ORIGINAL_PATH]
with mrc.ZSliceReader(input_stack_path) as reader:
i = row[star.UCSF.IMAGE_ORIGINAL_INDEX]
writer.write(reader.read(i))
elif fn.endswith(".par"):
if args.stack_path is None:
log.error(".par file input requires --stack-path")
return 1
df = metadata.par2star(metadata.parse_fx_par(fn), data_path=args.stack_path)
# star.set_original_fields(df, inplace=True) # Redundant.
star.augment_star_ucsf(df)
elif fn.endswith(".csv"):
return 1
elif fn.endswith(".cs"):
return 1
else:
if fn.endswith(".mrcs"):
with mrc.ZSliceReader(fn) as reader:
for img in reader:
writer.write(img)
df = pd.DataFrame(
{star.UCSF.IMAGE_ORIGINAL_INDEX: np.arange(reader.nz)})
df[star.UCSF.IMAGE_ORIGINAL_PATH] = fn
else:
print("Unrecognized input file type")
return 1
if args.star is not None:
df[star.UCSF.IMAGE_INDEX] = np.arange(first_ptcl,
first_ptcl + df.shape[0])
if args.abs_path:
df[star.UCSF.IMAGE_PATH] = writer.path
else:
df[star.UCSF.IMAGE_PATH] = os.path.relpath(writer.path, os.path.dirname(args.star))
df["index"] = df[star.UCSF.IMAGE_INDEX]
star.simplify_star_ucsf(df)
dfs.append(df)
first_ptcl += df.shape[0]
if args.star is not None:
df = pd.concat(dfs, join="inner")
# df = pd.concat(dfs)
# df = df.dropna(df, axis=1, how="any")
if not args.relion2: # Relion 3.1 style output.
df = star.remove_deprecated_relion2(df, inplace=True)
star.write_star(args.star, df, resort_records=False, optics=True)
else:
df = star.remove_new_relion31(df, inplace=True)
star.write_star(args.star, df, resort_records=False, optics=False)
return 0
def main(args):
log = logging.getLogger('root')
hdlr = logging.StreamHandler(sys.stdout)
log.addHandler(hdlr)
log.setLevel(logging.getLevelName(args.loglevel.upper()))
df = star.parse_star(args.input, keep_index=False)
star.augment_star_ucsf(df)
if args.map is not None:
vol = mrc.read(args.map, inc_header=False, compat="relion")
if args.mask is not None:
mask = mrc.read(args.mask, inc_header=False, compat="relion")
vol *= mask
else:
print("Please supply a map")
return 1
f3d = vop.vol_ft(vol, pfac=args.pfac, threads=args.threads)
sz = f3d.shape[0] // 2 - 1
sx, sy = np.meshgrid(np.fft.rfftfreq(sz), np.fft.fftfreq(sz))
s = np.sqrt(sx**2 + sy**2)
a = np.arctan2(sy, sx)
ift = None
with mrc.ZSliceWriter(args.output) as zsw:
for i, p in df.iterrows():
f2d = project(f3d,
p,
s,
sx,
sy,
a,
apply_ctf=args.ctf,
size=args.size)
if ift is None:
ift = irfft2(f2d.copy(),
threads=cpu_count(),
planner_effort="FFTW_ESTIMATE",
auto_align_input=True,
auto_contiguous=True)
proj = fftshift(
ift(f2d.copy(), np.zeros(vol.shape[:-1], dtype=vol.dtype)))
log.debug("%f +/- %f" % (np.mean(proj), np.std(proj)))
if args.subtract:
with mrc.ZSliceReader(p["ucsfImagePath"]) as zsr:
img = zsr.read(p["ucsfImageIndex"])
log.debug("%f +/- %f" % (np.mean(img), np.std(img)))
proj = img - proj
zsw.write(proj)
log.info(
"%d@%s: %d/%d" %
(p["ucsfImageIndex"], p["ucsfImagePath"], i + 1, df.shape[0]))
if args.star is not None:
if args.subtract:
df[star.UCSF.IMAGE_ORIGINAL_PATH] = df[star.UCSF.IMAGE_PATH]
df[star.UCSF.IMAGE_ORIGINAL_INDEX] = df[star.UCSF.IMAGE_INDEX]
df[star.UCSF.IMAGE_PATH] = args.output
df[star.UCSF.IMAGE_INDEX] = np.arange(df.shape[0])
star.simplify_star_ucsf(df)
star.write_star(args.star, df)
return 0