lammps restart, image_array args

Author: mkphuthi

asimtools/asimmodules/lammps/lammps.py

@@ -4,7 +4,7 @@
 
 Author: mkphuthi@github.com
 '''
-from typing import Dict, Optional
+from typing import Dict, Optional, Sequence
 import sys
 from pathlib import Path
 from numpy.random import randint
@@ -24,6 +24,8 @@ def lammps(
     masses: bool = True,
     velocities: bool = False,
     seed: Optional[int] = None,
+    restart_template: Optional[str] = None,
+    specorder: Sequence[str] = None,
 ) -> Dict:
     """Runs a lammps script based on a specified template, variables can be 
     specified as arguments to be defined in the final LAMMPS input file if 
@@ -55,6 +57,12 @@ def lammps(
         seed to be placed, if seed=None, a random one is generated, 
         defaults to None
     :type seed: int, optional
+    :param restart_template: Optional lammps input template to be used to
+        generate a restart.lammps file, defaults to None
+    :type restart_template: str, optional
+    :param specorder: Optional list of atomic species in the order they
+        should appear in the LAMMPS data input file, defaults to None
+    :type specorder: Sequence[str], optional
     :return: LAMMPS out file names
     :rtype: Dict
     """
@@ -73,6 +81,7 @@ def lammps(
                     atom_style=atom_style,
                     masses=masses,
                     velocities=velocities,
+                    specorder=specorder,
                 )
             except ValueError as te:
                 err_txt = 'Need ASE version >=3.23 to support writing '
@@ -92,6 +101,7 @@ def lammps(
         variables['IMAGE_FILE'] = 'image_input.lmpdat'
 
     lmp_txt = ''
+    
     for variable, value in variables.items():
         lmp_txt += f'variable {variable} equal {value}\n'
 
@@ -103,6 +113,18 @@ def lammps(
     if placeholders is None:
         placeholders = {}
 
+    if restart_template is not None:
+        restart_txt = lmp_txt
+        with open(restart_template, 'r', encoding='utf-8') as f:
+            restart_lines = f.readlines()
+        for rline in restart_lines:
+            if placeholders is not None:
+                for placeholder in placeholders:
+                    rline = rline.replace(placeholder, str(placeholders[placeholder]))
+            restart_txt += rline
+        with open('restart.lammps', 'w', encoding='utf-8') as f:
+            f.write(restart_txt)
+
     for line in lines:
         if 'SEED' in line and 'SEED' not in placeholders:
             if seed is None:

asimtools/asimmodules/workflows/image_array.py

@@ -90,6 +90,8 @@ def image_array(
         secondary_array_values=secondary_array_values,
     )
 
+    if template_sim_input is not None:
+        subsim_input = template_sim_input
     if key_sequence is None:
         key_sequence = ['args', 'image']
         # For backwards compatibility where we don't have to specify image

asimtools/utils.py

@@ -245,6 +245,7 @@ def get_atoms(
     builder: Optional[str] = 'bulk',
     atoms: Optional[Atoms] = None,
     repeat: Optional[Tuple[int, int, int]] = None,
+    repeat_to_N_args: Optional[Dict] = None,
     rattle_stdev: Optional[float] = None,
     mp_id: Optional[str] = None,
     user_api_key: Optional[str] = None,
@@ -329,14 +330,16 @@ def get_atoms(
     >>> get_atoms(image_file='molecules.xyz', index=0) # Pick out one structure using indexing
     Atoms(symbols='OH2', pbc=False)
 
-    You can also make supercells and rattle the atoms
+    You can also make supercells and rattle the atoms or repeat to a target
 
     >>> li_bulk = get_atoms(name='Li')
     >>> li_bulk.write('POSCAR', format='vasp')
     >>> get_atoms(image_file='POSCAR', repeat=[3,3,3])
     Atoms(symbols='Li27', pbc=True, cell=[[-5.235, 5.235, 5.235], [5.235, -5.235, 5.235], [5.235, 5.235, -5.235]])
     >>> get_atoms(builder='bulk', name='Li', repeat=[2,2,2], rattle_stdev=0.01)
     Atoms(symbols='Li8', pbc=True, cell=[[-3.49, 3.49, 3.49], [3.49, -3.49, 3.49], [3.49, 3.49, -3.49]])
+    >>> get_atoms(builder='bulk', name='Li', repeat_to_N_args={'N': 16, 'max_dim': 10.0})
+    Atoms(symbols='Li16', pbc=True, cell=[[-6.98, 6.98, 6.98], [6.98, -6.98, 6.98], [6.98, 6.98, -6.98]])
 
     Mostly for internal use and use in asimmodules, one can specify atoms
     directly
@@ -439,6 +442,13 @@ def get_atoms(
     elif rattle_stdev is not None and interface == 'pymatgen':
         struct.perturb(distance=rattle_stdev, min_distance=0)
 
+    if repeat_to_N_args is not None and interface == 'ase':
+        atoms = repeat_to_N(atoms, **repeat_to_N_args)
+    elif repeat_to_N_args is not None and interface == 'pymatgen':
+        raise NotImplementedError(
+            'repeat_to_N_args is only implemented for ASE interface'
+        )
+
     if constraints is not None and interface == 'ase':
         consts = []
         for constraint_args in constraints:
@@ -491,6 +501,47 @@ def parse_slice(value: str, bash: bool = False) -> slice:
             parts.append('1')
         return f'$(seq {parts[0]} {parts[2]} {parts[1]})'
 
+def repeat_to_N(
+    atoms: Atoms,
+    N: int,
+    max_dim: float = 50.0,
+) -> Atoms:
+    """Scale a structure to have approximately N atoms in the unit cell.
+    The function repeats the shortest axis of the unit cell until the
+    number of atoms >= N or the longest axis of the unit cell is > max_dim.
+
+    :param atoms: Input atoms object
+    :type atoms: Atoms
+    :param N: Target number of atoms
+    :type N: int
+    :param max_dim: Maximum length of the longest axis of the unit cell,
+        defaults to 50.0
+    :type max_dim: float, optional
+    :raises ValueError: If it fails to scale the structure to N atoms
+    :return: Scaled atoms object
+    :rtype: Atoms
+    """
+    cell = atoms.get_cell()
+    lengths = [np.linalg.norm(vec) for vec in cell]
+    num_atoms = len(atoms)
+    new_atoms = atoms.copy()
+
+    while num_atoms < N and np.max(lengths) < max_dim:
+        shortest_axis = np.argmin(lengths)
+        repeat_vec = [1, 1, 1]
+        repeat_vec[shortest_axis] += 1
+        new_atoms = new_atoms.repeat(repeat_vec)
+        cell = new_atoms.get_cell()
+        lengths = [np.linalg.norm(vec) for vec in cell]
+        num_atoms = len(new_atoms)
+
+    if num_atoms < N:
+        raise ValueError(
+            f'Failed to scale structure to {N} atoms without exceeding \
+                max_dim of {max_dim} Angstroms'
+        )
+    return new_atoms
+
 def get_images(
     image_file: str = None,
     pattern: str = None,

tests/unit/test_utils.py

@@ -24,6 +24,7 @@
     get_str_btn,
     expand_wildcards,
     write_atoms,
+    repeat_to_N,
 )
 import ase.build
 
@@ -132,6 +133,7 @@ def test_join_names(test_input, expected):
 ])
 def test_get_atoms(test_input, expected):
     ''' Test getting atoms from different inputs '''
+    print('e', expected)
     assert get_atoms(**test_input) == expected
 
 def test_get_atoms_constraints(tmp_path):
@@ -373,4 +375,16 @@ def test_expand_wildcards(test_input, expected, tmp_path):
             f.write('')
     print(f'Found paths in {os.getcwd()}: {[f for f in Path(tmp_path).glob("*")]}')
 
-    assert expand_wildcards(test_input, root_path=tmp_path) == expected
+    assert expand_wildcards(test_input, root_path=tmp_path) == expected
+
+def test_repeat_to_N():
+    ''' Test repeating unit cell to at least N atoms '''
+    atoms = ase.build.bulk('Cu', crystalstructure='fcc', cubic=True, a=2.0)
+    repeated_atoms = repeat_to_N(atoms, 16)
+    assert len(repeated_atoms) == 16
+    assert np.abs(repeated_atoms.get_cell()[0][0] - 2*2.0) < 1e-6
+    assert np.abs(repeated_atoms.get_cell()[1][1] - 2*2.0) < 1e-6
+    assert np.abs(repeated_atoms.get_cell()[2][2] - 1*2.0) < 1e-6
+    assert len(repeat_to_N(atoms, 15)) == 16
+    with pytest.raises(ValueError):
+        repeat_to_N(atoms, 16, max_dim=4)