class DataTransformation(object): """ Args: source ([dict]): list of records config (Configuration): configuration to apply to each record """ def __init__(self, source, config, single_process=False, processes=cpu_count()): if single_process: self.results = imap(transform, izip(source, repeat(config))) else: self.pool = Pool(processes) self.results = self.pool.imap(transform, izip(source, repeat(config))) def __iter__(self): return self def next(self): return self.results.next() def __del__(self): try: self.pool.close() except: pass
def main(args): dens = EMData(args.map) star = parse_star(args.input, keep_index=False) star[["ImageNumber", "ImageName"]] = star['rlnImageName'].str.split("@", expand=True) grouped = star.groupby("ImageName") pool = None if args.nproc > 1: pool = Pool(processes=args.nproc) results = pool.imap(lambda x: project_stack(x, dens, args.dest), (group for name, group in grouped)) else: results = (project_stack(group, dens, args.dest) for name, group in grouped) i = 0 t = 0 for r in results: i += 1 t += r sys.stdout.write("\rProjected %d particles in %d stacks" % (t, i)) sys.stdout.flush() if pool is not None: pool.close() pool.join() sys.stdout.write('\n') sys.stdout.flush() return 0
class DataFusion(object): """ Args: source ([tuple]): list of tuples, each containing the records to merge config (Configuration): configuration to apply to each tuple """ def __init__(self, source, config, processes=cpu_count()): self._pool = Pool(processes) self._results = self._pool.imap(_merge_records, izip(source, repeat(config))) def __iter__(self): return self def next(self): return self._results.next() def __del__(self): self._pool.close()
def main(options): """ Projection subtraction program entry point. :param options: Command-line arguments parsed by ArgumentParser.parse_args() :return: Exit status """ rchop = lambda x, y: x if not x.endswith(y) or len(y) == 0 else x[:-len(y)] options.output = rchop(options.output, ".star") options.suffix = rchop(options.suffix, ".mrc") options.suffix = rchop(options.suffix, ".mrcs") star = StarFile(options.input) npart = len(star['rlnImageName']) sub_dens = EMData(options.submap) if options.wholemap is not None: dens = EMData(options.wholemap) else: print "Reference map is required." return 1 # Write star header for output.star. top_header = "\ndata_\n\nloop_\n" headings = star.keys() output_star = open("{0}.star".format(options.output), 'w') output_star.write(top_header) for i, heading in enumerate(headings): output_star.write("_{0} #{1}\n".format(heading, i + 1)) if options.recenter: # Compute difference vector between new and old mass centers. if options.wholemap is None: print "Reference map required for recentering." return 1 new_dens = dens - sub_dens # Note the sign of the shift in coordinate frame is opposite the shift in the CoM. recenter = Vec3f(*dens.phase_cog()[:3]) - Vec3f( *new_dens.phase_cog()[:3]) else: recenter = None pool = None if options.nproc > 1: # Compute subtraction in parallel. pool = Pool(processes=options.nproc) results = pool.imap( lambda x: subtract(x, dens, sub_dens, recenter=recenter, no_frc=options.no_frc, low_cutoff=options.low_cutoff, high_cutoff=options.high_cutoff), particles(star), chunksize=min(npart / options.nproc, options.maxchunk)) else: # Use serial generator. results = (subtract(x, dens, sub_dens, recenter=recenter, no_frc=options.no_frc, low_cutoff=options.low_cutoff, high_cutoff=options.high_cutoff) for x in particles(star)) # Write subtraction results to .mrcs and .star files. i = 0 nfile = 1 starpath = None mrcs = None mrcs_orig = None for r in results: if i % options.maxpart == 0: mrcsuffix = options.suffix + "_%d" % nfile nfile += 1 starpath = "{0}.mrcs".format( os.path.sep.join( os.path.relpath(mrcsuffix, options.output).split(os.path.sep)[1:])) mrcs = "{0}.mrcs".format(mrcsuffix) mrcs_orig = "{0}_original.mrcs".format(mrcsuffix) if os.path.exists(mrcs): os.remove(mrcs) if os.path.exists(mrcs_orig): os.remove(mrcs_orig) r.ptcl_norm_sub.append_image(mrcs) if options.original: r.ptcl.append_image(mrcs_orig) if logger.getEffectiveLevel( ) == logging.DEBUG: # Write additional debug output. ptcl_sub_img = r.ptcl.process("math.sub.optimal", { "ref": r.ctfproj, "actual": r.ctfproj_sub, "return_subim": True }) ptcl_lowpass = r.ptcl.process("filter.lowpass.gauss", { "apix": 1.22, "cutoff_freq": 0.05 }) ptcl_sub_lowpass = r.ptcl_norm_sub.process("filter.lowpass.gauss", { "apix": 1.22, "cutoff_freq": 0.05 }) ptcl_sub_img.write_image("poreclass_subimg.mrcs", -1) ptcl_lowpass.write_image("poreclass_lowpass.mrcs", -1) ptcl_sub_lowpass.write_image("poreclass_sublowpass.mrcs", -1) r.ctfproj.write_image("poreclass_ctfproj.mrcs", -1) r.ctfproj_sub.write_image("poreclass_ctfprojsub.mrcs", -1) assert r.meta.i == i # Assert particle order is preserved. star['rlnImageName'][i] = "{0:06d}@{1}".format( i % options.maxpart + 1, starpath) # Set new image name. r.meta.update(star) # Update StarFile with altered fields. line = ' '.join(str(star[key][i]) for key in headings) output_star.write("{0}\n".format(line)) i += 1 output_star.close() if pool is not None: pool.close() pool.join() return 0