Merge pull request #76 from mir-group/feature/jon/npool

Feature/jon/npool

flare/dft_interface/cp2k_util.py

@@ -6,29 +6,34 @@
 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)
 
     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)

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):
     """

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,

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}

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}

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()

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()

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()