Python pool.Pool类代码示例

def main():

	options = options_desc.parse_args(sys.argv)[0]
	pool = Pool()

	needy_nodes = pool.where("state == 'merge-able'").multilock()

	if(len(needy_nodes) == 0):
		return
	
	# find out about trr time step
	dt = 0	
	nodeDir = needy_nodes[0].dir.split('/')[-1]
	for fn in os.listdir(needy_nodes[0].dir):
		if re.match("^"+nodeDir+".+run\d+\.trr", fn):
			trr = TrrFile(needy_nodes[0].dir+"/"+fn)			
			dt = trr.first_frame.next().t - trr.first_frame.t
			trr.close()
			break

	# dt is sometimes noisy in the final digits (three digits is femtosecond step = enough)
	dt = np.around(dt, decimals=3)
	for n in needy_nodes:

		if(options.trr):
			# merge sampling trajectories
			trr_fns = sorted([ fn for fn in os.listdir(n.dir) if re.match("[^#].+run\d+.trr", fn) ])
			cmd = ["trjcat", "-f"]
			cmd += trr_fns
			cmd += ["-o", "../../"+n.trr_fn, "-cat"]
			print("Calling: %s"%" ".join(cmd))
			check_call(cmd, cwd=n.dir)

		if(options.edr):
			# merge edr files
			# get list of edr-files
			edr_fnames = sorted([n.dir+"/"+fn for fn in os.listdir(n.dir) if re.match("[^#].+run\d+.edr", fn)])
			assert( len(edr_fnames) ==  n.extensions_counter+1 )
			assert( len(edr_fnames) ==  n.extensions_max+1 )

			time_offset = n.sampling_length+dt

			for edr_fn in edr_fnames[1:]:	
				# adapt edr starting times
				cmd = ["eneconv", "-f", edr_fn, "-o", edr_fn, "-settime"]
				print("Calling: "+(" ".join(cmd)))
				p = Popen(cmd, stdin=PIPE)
				p.communicate(input=(str(time_offset)+"\n"))
				assert(p.wait() == 0)

				time_offset += n.extensions_length+dt

			# concatenate edr files with adapted starting times
			cmd = ["eneconv", "-f"] + edr_fnames + ["-o", n.dir+"/ener.edr"]
			print("Calling: "+(" ".join(cmd)))
			p = Popen(cmd)
			retcode = p.wait()
			assert(retcode == 0)

	needy_nodes.unlock()

def main(argv=None):
    if argv == None:
        argv = sys.argv
        options = options_desc.parse_args(argv)[0]

    pool = Pool()

    found_parents = [n for n in pool if n.name == options.parent_node]
    assert len(found_parents) == 1
    parent = found_parents[0]

    chosen_idx = np.linspace(start=0, stop=parent.trajectory.n_frames - 1, num=options.numnodes).astype(int)

    print "choosen_idx: ", chosen_idx

    for i in chosen_idx:
        n = Node()
        n.parent_frame_num = i
        n.parent = parent
        n.state = "created"
        n.extensions_counter = 0
        n.extensions_max = 0
        n.extensions_length = 0
        n.sampling_length = parent.sampling_length * 3
        n.internals = parent.trajectory.getframe(i)
        pool.append(n)
        n.save()

def main():
	options = options_desc.parse_args(sys.argv)[0]
	
	pool = Pool()
	
	choice = "state in ('converged', 'refined')"
	if(options.ignore_convergence):
		choice = "state in ('converged', 'not-converged', 'refined')"	

	needy_nodes = NodeList([n for n in pool.where(choice) if not n == pool.root]) # we won't touch the root

	if not(len(needy_nodes)):
		sys.exit("Nothing to do.")

	if not(userinput("Once the solvent has been removed, further refinement of the pool is not possible. This includes the generation of unrestrained transition nodes! Continue?", "bool")):
		sys.exit("Quit by user.")
		
	assert(len(needy_nodes) == len(needy_nodes.multilock())) # make sure we lock ALL nodes

	try:
		for n in needy_nodes:	
			discard_solvent(n, "pdb")
			discard_solvent(n, "trr")

		for n in needy_nodes:
			n.unlock()
	except:
		traceback.print_exc()

def main(argv=None):
	if(argv==None):
		argv = sys.argv
	options = options_desc.parse_args(argv)[0]

	assert(not(options.refine_all and options.extend_all)) 
	
	pool = Pool()
	needy_nodes = pool.where("isa_partition and is_sampled").multilock()
	
	# 1. Trying to detect fake convergence
	for n in pool.where("state == 'converged'"):
		means = kmeans(n.trajectory, k=2)
		d = (means[0] - means[1]).norm2()
		if(d > 2.0 and (options.refine_all or userinput("%s has converged but appears to have a bimodal distribution.\nDo you want to refine?"%n.name, "bool"))): #TODO decide upon threshold (per coordinate?)
			refine(n, options)
	
	# 2. Dealing with not-converged nodes
	for n in pool.where("state == 'not-converged'"):
		if(not(options.refine_all or options.extend_all)):
			choice = userchoice("%s has not converged. What do you want to do?"%n.name, ['_refine', '_extend', '_ignore'])
		if(options.refine_all or choice=="r"):
			refine(n, options)
		elif(options.extend_all or choice=="e"):
			extend(n)
		elif(choice=="i"):
			continue
	
	for n in needy_nodes:
		n.save()
		n.unlock()
			
	zgf_setup_nodes.main()
	zgf_grompp.main()
	zgf_cleanup.main()	

def main():
	pool = Pool()

	if(len(pool.where("state == 'grompp-able'")) > 0):
		call_grompp(pool)
		
	if(len(pool.where("state == 'em-grompp-able'")) > 0):
		assert(path.exists("em.mdp")) #TODO that's not super nice yet
		call_grompp(pool, mode='em')

def main():
	options = options_desc.parse_args(sys.argv)[0]
	
	zgf_cleanup.main()
	
	pool = Pool()

	#not_reweightable = "state not in ('refined','converged')"
	not_reweightable = "isa_partition and state!='converged'"
	if options.ignore_convergence:
		not_reweightable = "isa_partition and state not in ('converged','not-converged')"

	if pool.where(not_reweightable):
		print "Pool can not be reweighted due to the following nodes:"		
		for bad_guy in pool.where(not_reweightable):
			print "Node %s with state %s."%(bad_guy.name, bad_guy.state)
		sys.exit("Aborting.")
		
	active_nodes = pool.where("isa_partition")
	assert(len(active_nodes) == len(active_nodes.multilock())) # make sure we lock ALL nodes

	for n in active_nodes:
		check_restraint_energy(n)

	# find out about number of energygrps
	mdp_file = gromacs.read_mdp_file(pool.mdp_fn)
	energygrps = [str(egrp) for egrp in re.findall('[\S]+', mdp_file["energygrps"])]
	moi_energies = True	
	if len(energygrps) < 2:
		moi_energies = False # Gromacs energies are named differently when there are less than two energygrps :(

	if(options.method == "direct"): 
		reweight_direct(active_nodes, moi_energies, options.sol_energy, options.save_refpoints)
	elif(options.method == "entropy"):
		reweight_entropy(active_nodes, moi_energies, options.sol_energy, options.save_refpoints)
	elif(options.method == "presampling"):
		reweight_presampling(active_nodes, options.presamp_temp, moi_energies, options.sol_energy)
	else:
		raise(Exception("Method unkown: "+options.method))
	
	weight_sum = np.sum([n.tmp['weight'] for n in active_nodes])
	
	print "Thermodynamic weights calculated by method '%s' (sol-energy=%s):"%(options.method, options.sol_energy)
	for n in active_nodes:
		n.obs.weight_direct = n.tmp['weight'] / weight_sum
		if(options.method == "direct"):
			print("  %s with mean_V: %f [kJ/mol], %d refpoints and weight: %f" % (n.name, n.obs.mean_V, n.tmp['n_refpoints'], n.obs.weight_direct))
		else:
			print("  %s with A: %f [kJ/mol] and weight: %f" % (n.name, n.obs.A, n.obs.weight_direct))

	for n in active_nodes:
		n.save()

	active_nodes.unlock()

def main():
	pool = Pool()
	needy_nodes = pool.where("state == 'grompp-able'").multilock()
		
	try:
		for n in needy_nodes:
			call_grompp(n)			
	except:
		traceback.print_exc()

	for n in needy_nodes:
		n.unlock()

def main():
	options = options_desc.parse_args(sys.argv)[0]
	
	zgf_cleanup.main()
	
	pool = Pool()

	not_reweightable = "isa_partition and state not in ('converged'"
	if(options.ignore_convergence):
		not_reweightable += ",'not-converged'"
	if(options.ignore_failed):
		not_reweightable += ",'mdrun-failed'"
	not_reweightable += ")"

	if pool.where(not_reweightable):
		print "Pool can not be reweighted due to the following nodes:"		
		for bad_guy in pool.where(not_reweightable):
			print "Node %s with state %s."%(bad_guy.name, bad_guy.state)
		sys.exit("Aborting.")

	active_nodes = pool.where("isa_partition and state != 'mdrun-failed'")
	assert(len(active_nodes) == len(active_nodes.multilock())) # make sure we lock ALL nodes

	if(options.check_restraint):
		for n in active_nodes:
			check_restraint_energy(n)

	if(options.method == "direct"):
		reweight_direct(active_nodes, options)
	elif(options.method == "entropy"):
		reweight_entropy(active_nodes, options)
	elif(options.method == "presampling"):
		reweight_presampling(active_nodes, options)
	else:
		raise(Exception("Method unkown: "+options.method))
	
	weight_sum = np.sum([n.tmp['weight'] for n in active_nodes])
	
	print "Thermodynamic weights calculated by method '%s':"%options.method
	for n in active_nodes:
		n.obs.weight_direct = n.tmp['weight'] / weight_sum
		if(options.method == "direct"):
			print("  %s with mean_V: %f [kJ/mol], %d refpoints and weight: %f" % (n.name, n.obs.mean_V, n.tmp['n_refpoints'], n.obs.weight_direct))
		else:
			print("  %s with A: %f [kJ/mol] and weight: %f" % (n.name, n.obs.A, n.obs.weight_direct))
	print "The above weighting uses bonded energies='%s' and nonbonded energies='%s'."%(options.e_bonded, options.e_nonbonded)

	for n in active_nodes:
		n.save()

	active_nodes.unlock()

def main():
	
	pool = Pool()
	needy_nodes = pool.where("state == 'created'")
	assert(len(needy_nodes) == len(needy_nodes.multilock())) # make sure we lock ALL nodes 
	
	extract_frames(pool)
	generate_topology(pool)
	generate_mdp(pool)
	
	for n in needy_nodes:
		n.state = "grompp-able"
		n.save()
		n.unlock()

def main():
	options = options_desc.parse_args(sys.argv)[0]
	
	zgf_cleanup.main()

	pool = Pool()

	needy_nodes = pool.where("state == '%s'"%options.current_state).multilock()

	for n in needy_nodes:
		print "Recovering node %s with state %s to state %s ..."%(n.name, n.state, options.recover_state)
		n.state = options.recover_state
		n.save()		
		n.unlock()

def main():
	options = options_desc.parse_args(sys.argv)[0]
	
	pool = Pool()
	needy_nodes = pool.where("state == 'em-mdrun-able'")
	assert(len(needy_nodes) == len(needy_nodes.multilock())) # make sure we lock ALL nodes


	# add ions to simulation boxes
	call_genion(pool, options.np, options.pname, options.nn, options.nname, options.random_seed)

	
	for n in needy_nodes:
		n.state = "em-grompp-able"
		zgf_grompp.call_grompp(n, mdp_file=options.grompp, final_state="em-mdrun-able") # re-grompp to get a tpr for energy minimization
		n.unlock()

def main():
	options = options_desc.parse_args(sys.argv)[0]

	outfile = open(options.outfile,"w")
	
	pool = Pool()
	needy_nodes = pool.where("isa_partition and state not in ('refined','mdrun-failed')")
	
	for n in needy_nodes:
		outfile.write("%s, state: '%s':\n"%(n.name,n.state))
		outfile.write(str(n.internals.array)+"\n")
		outfile.write("mean pot.: %f, std pot.: %f, free energy estimate: %f\n"%(n.obs.mean_V,n.obs.std_V,n.obs.A))
		outfile.write("#========================================================================#\n")

	outfile.close()
	print "Pool info was written to %s."%options.outfile

def main():
	options = options_desc.parse_args(sys.argv)[0]
	
	pool = Pool()
	needy_nodes = pool.where("state == 'grompp-able'")
	assert(len(needy_nodes) == len(needy_nodes.multilock())) # make sure we lock ALL nodes

	if(options.solv_model == "tip3p"):
		solv_box = "spc216.gro"
		solv_fn = "tip3p.itp"
	elif(options.solv_model == "tip4p"):
		solv_box = "tip4p.gro"
		solv_fn = "tip4p.itp"
	elif(options.solv_model == "tip4pew"):
		solv_box = "tip4p.gro"
		solv_fn = "tip4pew.itp"
	elif(options.solv_model == "tip5"):
		solv_box = "tip5p.gro"
		solv_fn = "tip5p.itp"
	elif(options.solv_model == "spc"):
		solv_box = "spc216.gro"
		solv_fn = "spc.itp"
	elif(options.solv_model == "spce"):
		solv_box = "spc216.gro"
		solv_fn = "spce.itp"
	elif(options.solv_model == "acetonitrile"): # TODO one might change this one to "custom" and let user enter name of template box
		solv_box = "acetonitrile.pdb"
		msg = "Topology update for acetonitrile is not supported. Proceed?"
		if not(userinput(msg, "bool")):
			for n in needy_nodes:
				n.unlock()
			return("Quit by user.")
	
	# determine maximum length of linears, if any
	max_linear = query_linear_length(pool)

	# make box and fill with solvent
	genbox(pool, max_linear, options.bt, (options.box_x, options.box_y, options.box_z), solv_box)

	# update topology files (add solvent model and ions includes)
	if not(options.solv_model == "acetonitrile"):
		update_tops(pool, solv_fn)

	for n in needy_nodes:
		n.state = "em-grompp-able"
		zgf_grompp.call_grompp(n, mdp_file=options.grompp, final_state="em-mdrun-able") # re-grompp to get a tpr for energy minimization
		n.unlock()

def main():
	options = options_desc.parse_args(sys.argv)[0]
	
	#TODO put somehow into Options, e.g. min_value=1 or required=True
	if(not options.doomed_nodes):
		sys.exit("Option --doomed_nodes is required.")
		
	pool = Pool()
	old_pool_size = len(pool)
	old_alpha = pool.alpha

	doomed_nodes = NodeList()
	
	#TODO: maybe this code should go into ZIBMolPy.ui 
	for name in options.doomed_nodes.split(","):
		found = [n for n in pool if n.name == name]
		if(len(found) != 1):
			sys.exit("Coult not find node '%s'"%(name))
		doomed_nodes.append(found[0])
	
	for n in doomed_nodes:
		if(n == pool.root):
			sys.exit("Node %s is the root. Removal not allowed."%(n.name))		
		#if(len(n.children) > 0):
		#	sys.exit("Node %s has children. Removal not allowed."%(n.name)) #TODO why should we forbid this?

	if not(userinput("The selected node(s) will be removed permanently. Continue?", "bool")):
		sys.exit("Quit by user.")

	assert(len(doomed_nodes) == len(doomed_nodes.multilock()))
	for n in doomed_nodes:
		print("Removing directory: "+n.dir)
		shutil.rmtree(n.dir)

	pool.reload_nodes()
	
	#TODO: this code-block also exists in zgf_create_node
	if(len(pool.where("isa_partition")) < 2):
		pool.alpha = None
	elif(options.methodalphas == "theta"):
		pool.alpha = zgf_create_nodes.calc_alpha_theta(pool)
	elif(options.methodalphas == "user"):
		pool.alpha = userinput("Please enter a value for alpha", "float")
	else:
		raise(Exception("Method unkown: "+options.methodalphas))

	pool.history.append({'removed_nodes': [(n.name, n.state) for n in doomed_nodes], 'size':old_pool_size, 'alpha':old_alpha, 'timestamp':datetime.now()})
	pool.save()

	#TODO: deal with analysis dir and dependencies
	zgf_cleanup.main()	

def main():
	options = options_desc.parse_args(sys.argv)[0]
	
	pool = Pool()
	needy_nodes = pool.where("state == 'em-mdrun-able'")
	assert(len(needy_nodes) == len(needy_nodes.multilock())) # make sure we lock ALL nodes


	# add ions to simulation boxes
	call_genion(pool, options.np, options.pname, options.nn, options.nname, options.random_seed)

	
	for n in needy_nodes:
		n.state = "em-grompp-able"
		n.save()
		n.unlock()

	zgf_grompp.main()