diff --git a/flare/dft_interface/cp2k_util.py b/flare/dft_interface/cp2k_util.py index ea8e48d99..fbc92d436 100644 --- a/flare/dft_interface/cp2k_util.py +++ b/flare/dft_interface/cp2k_util.py @@ -6,15 +6,19 @@ from flare import struc from typing import List -name="CP2K" +name = "CP2K" -def run_dft_par(dft_input, structure, dft_loc, no_cpus=1, dft_out="dft.out"): +def run_dft_par(dft_input, structure, dft_loc, no_cpus=1, dft_out="dft.out", + npool=None, mpi="mpi"): newfilename = edit_dft_input_positions(dft_input, structure) dft_command = \ - '{} -i {} > {}'.format(dft_loc, newfilename, dft_out) + f'{dft_loc} -i {newfilename} > {dft_out}' if (no_cpus > 1): - dft_command = 'mpirun -np {} {}'.format(no_cpus, dft_command) + if (mpi == "mpi"): + dft_command = f'mpirun -np {no_cpus} {dft_command}' + else: + dft_command = f'srun -n {no_cpus} {dft_command}' # output.write_to_output(dft_command+'\n') call(dft_command, shell=True) os.remove(newfilename) @@ -22,13 +26,14 @@ def run_dft_par(dft_input, structure, dft_loc, no_cpus=1, dft_out="dft.out"): return parse_dft_forces(dft_out) -def run_dft_en_par(dft_input, structure, dft_loc, no_cpus, dft_out="dft.out"): +def run_dft_en_par(dft_input, structure, dft_loc, no_cpus, dft_out="dft.out", + npool=None, mpi="mpi"): newfilename = edit_dft_input_positions(dft_input, structure) dft_command = \ - '{} -i {} > {}'.format(dft_loc, newfilename, dft_out) + f'{dft_loc} -i {newfilename} > {dft_out}' if (no_cpus > 1): - dft_command = 'mpirun -np {} {}'.format(no_cpus, dft_command) + dft_command = f'mpirun -np {no_cpus} {dft_command}' # output.write_to_output(dft_command+'\n') call(dft_command, shell=True) os.remove(newfilename) diff --git a/flare/dft_interface/qe_util.py b/flare/dft_interface/qe_util.py index d289d7e2c..3179e1031 100644 --- a/flare/dft_interface/qe_util.py +++ b/flare/dft_interface/qe_util.py @@ -9,31 +9,40 @@ def run_dft(qe_input, structure, dft_loc): run_qe_path = qe_input edit_dft_input_positions(run_qe_path, structure) - qe_command = '{0} < {1} > {2}'.format(dft_loc, run_qe_path, - 'pwscf.out') + qe_command = f'{dft_loc} < {run_qe_path} > pwscf.out' call(qe_command, shell=True) return parse_dft_forces('pwscf.out') -def run_dft_par(qe_input, structure, dft_loc, no_cpus): - run_qe_path = qe_input - edit_dft_input_positions(run_qe_path, structure) - qe_command = \ - 'mpirun -np {0} {1} < {2} > {3}'.format(no_cpus, dft_loc, run_qe_path, - 'pwscf.out') +def run_dft_par(qe_input, structure, dft_loc, no_cpus=1, dft_out='pwscf.out', + npool=None, mpi="mpi"): + newfilename = edit_dft_input_positions(qe_input, structure) - call(qe_command, shell=True) + if npool is None: + dft_command = \ + f'{dft_loc} -i {newfilename} > {dft_out}' + else: + dft_command = \ + f'{dft_loc} -nk {npool} -i {newfilename} > {dft_out}' - return parse_dft_forces('pwscf.out') + if (no_cpus > 1): + if (mpi == "mpi"): + dft_command = f'mpirun -np {no_cpus} {dft_command}' + else: + dft_command = f'srun -n {no_cpus} --mpi=pmi2 {dft_command}' + + call(dft_command, shell=True) + os.remove(newfilename) + + return parse_dft_forces(dft_out) def run_dft_en_par(qe_input, structure, dft_loc, no_cpus): run_qe_path = qe_input edit_dft_input_positions(run_qe_path, structure) qe_command = \ - 'mpirun -np {0} {1} < {2} > {3}'.format(no_cpus, dft_loc, run_qe_path, - 'pwscf.out') + 'mpirun -np {no_cpus} {dft_loc} < {run_qe_path} > pwscf.out' call(qe_command, shell=True) forces, energy = parse_dft_forces_and_energy('pwscf.out') @@ -193,10 +202,14 @@ def edit_dft_input_positions(qe_input: str, structure): lines[cell_index + 2] = ' '.join([str(x) for x in structure.vec3]) \ + '\n' - with open(qe_input, 'w') as f: + newfilename = qe_input + "_run" + + with open(newfilename, 'w') as f: for line in lines: f.write(line) + return newfilename + def parse_dft_forces(outfile: str): """ diff --git a/flare/otf.py b/flare/otf.py index 57b9a8fa2..26c330a78 100644 --- a/flare/otf.py +++ b/flare/otf.py @@ -23,7 +23,7 @@ def __init__(self, dft_input: str, dt: float, number_of_steps: int, max_atoms_added=1, freeze_hyps=10, rescale_steps=[], rescale_temps=[], dft_softwarename="qe", - no_cpus=1): + no_cpus=1, npool=None, mpi="srun"): self.dft_input = dft_input self.dt = dt @@ -85,8 +85,10 @@ def __init__(self, dft_input: str, dt: float, number_of_steps: int, self.output = Output(output_name, always_flush=True) - # set number of cpus for qe runs + # set number of cpus and npool for qe runs self.no_cpus = no_cpus + self.npool = npool + self.mpi = mpi def run(self): self.output.write_header(self.gp.cutoffs, self.gp.kernel_name, @@ -170,7 +172,10 @@ def run_dft(self): # calculate DFT forces forces = self.dft_module.run_dft_par(self.dft_input, self.structure, - self.dft_loc, self.no_cpus) + self.dft_loc, + no_cpus=self.no_cpus, + npool=self.npool, + mpi=self.mpi) self.structure.forces = forces # write wall time of DFT calculation @@ -192,11 +197,6 @@ def update_gp(self, train_atoms, dft_frcs): self.gp.set_L_alpha() - # if self.curr_step == 0: - # self.gp.set_L_alpha() - # else: - # self.gp.update_L_alpha() - def train_gp(self): self.gp.train(self.output) self.output.write_hyps(self.gp.hyp_labels, self.gp.hyps, diff --git a/tests/pwscf.in_run b/tests/pwscf.in_run new file mode 100644 index 000000000..804e9f73d --- /dev/null +++ b/tests/pwscf.in_run @@ -0,0 +1,34 @@ +&CONTROL + pseudo_dir = './test_files/pseudos' + outdir = '.' + calculation = 'scf' + disk_io = 'low' + tprnfor = .true. + wf_collect = .false. +/ +&SYSTEM + ecutwfc = 50 + ecutrho = 100 + ntyp = 1 +nat = 2 + ibrav = 0 +/ +&ELECTRONS + diagonalization = 'david' + mixing_beta = 0.5 + conv_thr = 1e-07 +/ +&IONS +/ +&CELL +/ +ATOMIC_SPECIES + H 1.0 H.pbe-kjpaw.UPF +CELL_PARAMETERS {angstrom} +5.0 0.0 0.0 +0.0 5.0 0.0 +0.0 0.0 5.0 +ATOMIC_POSITIONS {angstrom} +H 2.3 2.5 2.5 +H 2.8 2.5 2.5 +K_POINTS {gamma} diff --git a/tests/qe_input_1.in_run b/tests/qe_input_1.in_run new file mode 100644 index 000000000..e2debd89e --- /dev/null +++ b/tests/qe_input_1.in_run @@ -0,0 +1,34 @@ +&CONTROL + pseudo_dir = './test_files/pseudos' + outdir = '.' + calculation = 'scf' + disk_io = 'low' + tprnfor = .true. + wf_collect = .false. +/ +&SYSTEM + ecutwfc = 50 + ecutrho = 100 + ntyp = 1 +nat = 2 + ibrav = 0 +/ +&ELECTRONS + diagonalization = 'david' + mixing_beta = 0.5 + conv_thr = 1e-07 +/ +&IONS +/ +&CELL +/ +ATOMIC_SPECIES + H 1.0 H.pbe-kjpaw.UPF +CELL_PARAMETERS {angstrom} +6.222887684575656 -0.2229615903227085 0.380670175839734 +0.0 5.0 0.0 +0.0 0.0 5.0 +ATOMIC_POSITIONS {angstrom} +H 3.488144951197148 2.3560459114463352 2.1211854930486513 +H 2.8 2.5 2.5 +K_POINTS {gamma} diff --git a/tests/test_OTF_cp2k_par.py b/tests/test_OTF_cp2k_par.py index 3a102bf4a..18ec3c455 100644 --- a/tests/test_OTF_cp2k_par.py +++ b/tests/test_OTF_cp2k_par.py @@ -58,7 +58,7 @@ def test_otf_h2_par(): calculate_energy=True, max_atoms_added=1, dft_softwarename="cp2k", no_cpus=2, - par=True, + par=True, mpi="mpi", output_name='h2_otf_cp2k_par') otf.run() diff --git a/tests/test_OTF_qe_par.py b/tests/test_OTF_qe_par.py index 15b0f7894..1046327ee 100644 --- a/tests/test_OTF_qe_par.py +++ b/tests/test_OTF_qe_par.py @@ -36,7 +36,7 @@ def test_otf_h2(): qe_input = './pwscf.in' dt = 0.0001 - number_of_steps = 20 + number_of_steps = 4 cutoffs = np.array([5]) dft_loc = os.environ.get('PWSCF_COMMAND') std_tolerance_factor = -0.1 @@ -62,6 +62,55 @@ def test_otf_h2(): std_tolerance_factor, init_atoms=[0], calculate_energy=True, max_atoms_added=1, no_cpus=2, par=True, + mpi="mpi", + output_name='h2_otf_qe_par') + + otf.run() + os.system('mkdir test_outputs') + os.system('mv h2_otf_qe_par* test_outputs') + cleanup_espresso_run() + +@pytest.mark.skipif(not os.environ.get('PWSCF_COMMAND', + False), reason='PWSCF_COMMAND not found ' + 'in environment: Please install Quantum ' + 'ESPRESSO and set the PWSCF_COMMAND env. ' + 'variable to point to pw.x.') +def test_otf_Al_npool(): + """ + Test that an otf run can survive going for more steps + :return: + """ + os.system('cp ./test_files/qe_input_2.in ./pwscf.in') + + qe_input = './pwscf.in' + dt = 0.0001 + number_of_steps = 4 + cutoffs = np.array([5]) + dft_loc = os.environ.get('PWSCF_COMMAND') + std_tolerance_factor = -0.1 + + # make gp model + kernel = en.two_body + kernel_grad = en.two_body_grad + hyps = np.array([1, 1, 1]) + hyp_labels = ['Signal Std', 'Length Scale', 'Noise Std'] + energy_force_kernel = en.two_body_force_en + + gp = \ + GaussianProcess(kernel=kernel, + kernel_grad=kernel_grad, + hyps=hyps, + cutoffs=cutoffs, + hyp_labels=hyp_labels, + par=True, + energy_force_kernel=energy_force_kernel, + maxiter=50) + + otf = OTF(qe_input, dt, number_of_steps, gp, dft_loc, + std_tolerance_factor, init_atoms=[0], + calculate_energy=True, max_atoms_added=1, + no_cpus=2, par=True, npool=2, + mpi="mpi", output_name='h2_otf_qe_par') otf.run() diff --git a/tests/test_qe_util.py b/tests/test_qe_util.py index 132342899..c554efd8d 100644 --- a/tests/test_qe_util.py +++ b/tests/test_qe_util.py @@ -118,9 +118,9 @@ def test_espresso_input_edit(): structure.vec1 += np.random.randn(3) structure.positions[0] += np.random.randn(3) - edit_dft_input_positions('./qe_input_1.in', structure=structure) + new_file = edit_dft_input_positions('./qe_input_1.in', structure=structure) - positions, species, cell, masses = parse_dft_input('./qe_input_1.in') + positions, species, cell, masses = parse_dft_input(new_file) assert np.equal(positions[0], structure.positions[0]).all() assert np.equal(structure.vec1, cell[0, :]).all()