Add custom force module and otf helper functions; remove xyz logs

flare/md.py

@@ -1,4 +1,5 @@
 import numpy as np
+from typing import List
 
 
 def update_positions(dt, noa, structure):
@@ -34,3 +35,64 @@ def calculate_temperature(new_pos, structure, dt, noa):
     temperature = 2 * KE / ((3 * noa - 3) * kb)
 
     return KE, temperature, velocities
+
+
+def get_random_velocities(noa: int, temperature: float, mass: float):
+    """Draw velocities from the Maxwell-Boltzmann distribution, assuming a
+    fixed mass for all particles in amu.
+
+    Args:
+        noa (int): Number of atoms in the system.
+        temperature (float): Temperature of the system.
+        mass (float): Mass of each particle in amu.
+
+    Returns:
+        np.ndarray: Particle velocities, corrected to give zero center of mass motion.
+    """
+
+    # Use FLARE mass units (time = ps, length = A, energy = eV)
+    mass_md = mass * 0.000103642695727
+    kb = 0.0000861733034
+    std = np.sqrt(kb * temperature / mass_md)
+    velocities = np.random.normal(scale=std, size=(noa, 3))
+
+    # Remove center-of-mass motion
+    vel_sum = np.sum(velocities, axis=0)
+    corrected_velocities = velocities - vel_sum / noa
+
+    return corrected_velocities
+
+
+def multicomponent_velocities(temperature: float, masses: List[float]):
+    """Draw velocities from the Maxwell-Boltzmann distribution for particles of
+    varying mass.
+
+    Args:
+        temperature (float): Temperature of the system.
+        masses (List[float]): Particle masses in amu.
+
+    Returns:
+        np.ndarray: Particle velocities, corrected to give zero center of mass motion.
+    """
+
+    noa = len(masses)
+    velocities = np.zeros((noa, 3))
+    kb = 0.0000861733034
+    mom_tot = np.array([0., 0., 0.])
+
+    for n, mass in enumerate(masses):
+        # Convert to FLARE mass units (time = ps, length = A, energy = eV)
+        mass_md = mass * 0.000103642695727
+        std = np.sqrt(kb * temperature / mass_md)
+        rand_vel = np.random.normal(scale=std, size=(3))
+        velocities[n] = rand_vel
+        mom_curr = rand_vel * mass_md
+        mom_tot += mom_curr
+
+    # Correct momentum, remove center of mass motion
+    mom_corr = mom_tot / noa
+    for n, mass in enumerate(masses):
+        mass_md = mass * 0.000103642695727
+        velocities[n] -= mom_corr / mass_md
+
+    return velocities

flare/otf.py

@@ -74,6 +74,12 @@ class OTF:
             single file name, or a list of several. Copied files will be
             prepended with the date and time with the format
             'Year.Month.Day:Hour:Minute:Second:'.
+        custom_module: A custom module that can be used in place of the DFT
+            modules available in the FLARE package. The module must contain
+            two functions: parse_dft_input, which takes a file name (in string
+            format) as input and returns the positions, species, cell, and
+            masses of a structure of atoms; and run_dft_par, which takes a
+            number of DFT related inputs and returns the forces on all atoms.
     """
     def __init__(self, dft_input: str, dt: float, number_of_steps: int,
                  gp: gp.GaussianProcess, dft_loc: str,
@@ -83,10 +89,10 @@ def __init__(self, dft_input: str, dt: float, number_of_steps: int,
                  calculate_energy: bool = False, output_name: str = 'otf_run',
                  max_atoms_added: int = 1, freeze_hyps: int = 10,
                  rescale_steps: List[int] = [], rescale_temps: List[int] = [],
-                 dft_softwarename: str = "qe",
-                 no_cpus: int = 1, npool: int = None, mpi: str = "srun",
-                 dft_kwargs=None,
-                 store_dft_output: Tuple[Union[str,List[str]],str] = None):
+                 dft_softwarename: str = "qe", no_cpus: int = 1,
+                 npool: int = None, mpi: str = "srun", dft_kwargs=None,
+                 store_dft_output: Tuple[Union[str, List[str]], str] = None,
+                 custom_module=None):
 
         self.dft_input = dft_input
         self.dt = dt
@@ -97,13 +103,16 @@ def __init__(self, dft_input: str, dt: float, number_of_steps: int,
         self.skip = skip
         self.dft_step = True
         self.freeze_hyps = freeze_hyps
-        self.dft_module = dft_software[dft_softwarename]
+
+        if custom_module is not None:
+            self.dft_module = custom_module
+        else:
+            self.dft_module = dft_software[dft_softwarename]
 
         # parse input file
         positions, species, cell, masses = \
             self.dft_module.parse_dft_input(self.dft_input)
 
-
         self.structure = struc.Structure(cell=cell, species=species,
                                          positions=positions,
                                          mass_dict=masses,
@@ -169,7 +178,6 @@ def run(self):
         self.start_time = time.time()
 
         while self.curr_step < self.number_of_steps:
-            print('curr_step:', self.curr_step)
             # run DFT and train initial model if first step and DFT is on
             if self.curr_step == 0 and self.std_tolerance != 0:
                 # call dft and update positions
@@ -228,7 +236,6 @@ def run(self):
                     if (self.dft_count-1) < self.freeze_hyps:
                         self.train_gp()
 
-
             # write gp forces
             if counter >= self.skip and not self.dft_step:
                 self.update_temperature(new_pos)
@@ -243,7 +250,7 @@ def run(self):
 
     def run_dft(self):
         """Calculates DFT forces on atoms in the current structure.
-        
+
         If OTF has store_dft_output set, then the specified DFT files will
         be copied with the current date and time prepended in the format
         'Year.Month.Day:Hour:Minute:Second:'.
@@ -266,7 +273,7 @@ def run_dft(self):
         time_curr = time.time() - self.start_time
         self.output.write_to_log('number of DFT calls: %i \n' % self.dft_count)
         self.output.write_to_log('wall time from start: %.2f s \n' % time_curr)
-        
+
         # Store DFT outputs in another folder if desired
         # specified in self.store_dft_output
         if self.store_dft_output is not None:
@@ -304,9 +311,6 @@ def train_gp(self):
         """Optimizes the hyperparameters of the current GP model."""
 
         self.gp.train(self.output)
-        self.output.write_hyps(self.gp.hyp_labels, self.gp.hyps,
-                               self.start_time,
-                               self.gp.likelihood, self.gp.likelihood_gradient)
 
     def update_positions(self, new_pos: 'ndarray'):
         """Performs a Verlet update of the atomic positions.
@@ -332,8 +336,8 @@ def update_temperature(self, new_pos: 'ndarray'):
             new_pos (np.ndarray): Positions of atoms in the next MD frame.
         """
         KE, temperature, velocities = \
-                md.calculate_temperature(new_pos, self.structure, self.dt,
-                                         self.noa)
+            md.calculate_temperature(new_pos, self.structure, self.dt,
+                                     self.noa)
         self.KE = KE
         self.temperature = temperature
         self.velocities = velocities
@@ -344,5 +348,3 @@ def record_state(self):
                                     self.local_energies, self.start_time,
                                     self.dft_step,
                                     self.velocities)
-        self.output.write_xyz_config(self.curr_step, self.structure,
-                                     self.dft_step)

flare/output.py

@@ -36,9 +36,7 @@ def __init__(self, basename: str = 'otf_run',
         """
         self.basename = f"{basename}"
         self.outfiles = {}
-        filesuffix = {'log': '.out', 'xyz': '.xyz',
-                      'fxyz': '-f.xyz', 'hyps': "-hyps.dat",
-                      'std': '-std.xyz', 'stat': '-stat.dat'}
+        filesuffix = {'log': '.out', 'hyps': '-hyps.dat'}
 
         for filetype in filesuffix.keys():
             self.open_new_log(filetype, filesuffix[filetype])
@@ -151,9 +149,9 @@ def write_header(self, cutoffs, kernel_name: str,
         # report previous positions
         headerstring += '\nprevious positions (A):\n'
         for i in range(len(structure.positions)):
-            headerstring += f'{structure.species_labels[i]} '
+            headerstring += f'{structure.species_labels[i]:5}'
             for j in range(3):
-                headerstring += f'{structure.prev_positions[i][j]:10.4} '
+                headerstring += f'{structure.prev_positions[i][j]:8.4f}'
             headerstring += '\n'
         headerstring += '-' * 80 + '\n'
 
@@ -162,7 +160,6 @@ def write_header(self, cutoffs, kernel_name: str,
         if self.always_flush:
             f.flush()
 
-    # TO DO: this module should be removed in the future
     def write_md_config(self, dt, curr_step, structure,
                         temperature, KE, local_energies,
                         start_time, dft_step, velocities):
@@ -195,36 +192,36 @@ def write_md_config(self, dt, curr_step, structure,
         string += f'\nSimulation Time: {(dt * curr_step):.3} ps \n'
 
         # Construct Header line
-        string += 'El  Position (A) \t\t\t\t '
+        n_space = 24
+        string += str.ljust('El', 5)
+        string += str.center('Position (A)', n_space)
+        string += ' ' * 4
         if not dft_step:
-            string += 'GP Force (ev/A) '
+            string += str.center('GP Force (ev/A)', n_space)
+            string += ' ' * 4
         else:
-            string += 'DFT Force (ev/A) '
-        string += '\t\t\t\t Std. Dev (ev/A) \t'
-        string += '\t\t\t\t Velocities (A/ps) \n'
+            string += str.center('DFT Force (ev/A)', n_space)
+            string += ' ' * 4
+        string += str.center('Std. Dev (ev/A)', n_space) + ' ' * 4
+        string += str.center('Velocities (A/ps)', n_space) + '\n'
 
         # Construct atom-by-atom description
         for i in range(len(structure.positions)):
-            string += f'{structure.species_labels[i]} '
+            string += f'{structure.species_labels[i]:5}'
+            # string += '\t'
             for j in range(3):
-                string += f"{structure.positions[i][j]:8.4} "
-            string += '\t'
+                string += f'{structure.positions[i][j]:8.4f}'
+            string += ' ' * 4
             for j in range(3):
-                string += f"{structure.forces[i][j]:8.4} "
-            string += '\t'
+                string += f'{structure.forces[i][j]:8.4f}'
+            string += ' ' * 4
             for j in range(3):
-                string += f'{structure.stds[i][j]:8.4} '
-            string += '\t'
+                string += f'{structure.stds[i][j]:8.4f}'
+            string += ' ' * 4
             for j in range(3):
-                string += f'{velocities[i][j]:8.4} '
+                string += f'{velocities[i][j]:8.4f}'
             string += '\n'
 
-        print(curr_step)
-        print(structure.species_labels)
-        self.write_xyz_config(curr_step, structure, dft_step)
-        self.write_xyz(curr_step, structure.stds, structure.species_labels,
-                       "std", header)
-
         string += '\n'
         string += f'temperature: {temperature:.2f} K \n'
         string += f'kinetic energy: {KE:.6f} eV \n'
@@ -334,16 +331,16 @@ def write_hyps(self, hyp_labels, hyps, start_time, like, like_grad,
 
         if hyp_labels is not None:
             for i, label in enumerate(hyp_labels):
-                f.write(f'Hyp{i} : {label} = {hyps[i]}\n')
+                f.write(f'Hyp{i} : {label} = {hyps[i]:.4f}\n')
         else:
             for i, hyp in enumerate(hyps):
-                f.write(f'Hyp{i} : {hyp}\n')
+                f.write(f'Hyp{i} : {hyp:.4f}\n')
 
-        f.write(f'likelihood: {like}\n')
+        f.write(f'likelihood: {like:.4f}\n')
         f.write(f'likelihood gradient: {like_grad}\n')
         if start_time:
             time_curr = time.time() - start_time
-            f.write(f'wall time from start: {time_curr:.2} s \n')
+            f.write(f'wall time from start: {time_curr:.2f} s \n')
 
         if self.always_flush:
             f.flush()
@@ -393,9 +390,9 @@ def write_gp_dft_comparison(self, curr_step, frame,
 
         string += '\n'
 
-        self.write_xyz_config(curr_step, frame, forces=frame.forces,
-                              stds=frame.stds, forces_2=dft_forces,
-                              dft_step=True)
+        # self.write_xyz_config(curr_step, frame, forces=frame.forces,
+        #                       stds=frame.stds, forces_2=dft_forces,
+        #                       dft_step=True)
 
         mae = np.mean(error) * 1000
         mac = np.mean(np.abs(dft_forces)) * 1000
@@ -436,7 +433,7 @@ def write_gp_dft_comparison(self, curr_step, frame,
         stat += f' {dt}\n'
 
         self.outfiles['log'].write(string)
-        self.outfiles['stat'].write(stat)
+        # self.outfiles['stat'].write(stat)
 
         if self.always_flush:
             self.outfiles['log'].flush()

flare/util.py

@@ -8,6 +8,67 @@
 
 import numpy as np
 
+
+def get_supercell_positions(sc_size: int, cell: np.ndarray,
+                            positions: np.ndarray):
+    """Returns the positions of a supercell of atoms, with the number of cells
+    in each direction fixed.
+
+    Args:
+        sc_size (int): Size of the supercell.
+        cell (np.ndarray): 3x3 array of cell vectors.
+        positions (np.ndarray): Positions of atoms in the unit cell.
+
+    Returns:
+        np.ndarray: Positions of atoms in the supercell.
+    """
+
+    sc_positions = []
+    for m in range(sc_size):
+        vec1 = m * cell[0]
+        for n in range(sc_size):
+            vec2 = n * cell[1]
+            for p in range(sc_size):
+                vec3 = p * cell[2]
+
+                # append translated positions
+                for pos in positions:
+                    sc_positions.append(pos+vec1+vec2+vec3)
+
+    return np.array(sc_positions)
+
+
+def supercell_custom(cell: np.ndarray, positions: np.ndarray,
+                     size1: int, size2: int, size3: int):
+    """Returns the positions of a supercell of atoms with a chosen number of
+    cells in each direction.
+
+    Args:
+        cell (np.ndarray): 3x3 array of cell vectors.
+        positions (np.ndarray): Positions of atoms in the unit cell.
+        size1 (int): Number of cells along the first cell vector.
+        size2 (int): Number of cells along the second cell vector.
+        size3 (int): Number of cells along the third cell vector.
+
+    Returns:
+        np.ndarray: Positions of atoms in the supercell.
+    """
+
+    sc_positions = []
+    for m in range(size1):
+        vec1 = m * cell[0]
+        for n in range(size2):
+            vec2 = n * cell[1]
+            for p in range(size3):
+                vec3 = p * cell[2]
+
+                # append translated positions
+                for pos in positions:
+                    sc_positions.append(pos+vec1+vec2+vec3)
+
+    return np.array(sc_positions)
+
+
 # Dictionary mapping elements to their atomic number (Z)
 _element_to_Z = {'H': 1,
                  'He': 2,