Merge branch 'mkphuthi/develop' of https://github.com/BattModels/asimtools into mkphuthi/develop

Author: mkphuthi

asimtools/asimmodules/ase_md/ase_md.py

@@ -290,6 +290,8 @@ def ase_md(
         )
 
 
+    if plot_args is None:
+        plot_args = {}
     if plot:
         plot_thermo(
             images={'image_file': 'output.traj'},

asimtools/asimmodules/data/collect_images.py

@@ -61,6 +61,7 @@ def collect_images(
     :rtype: Dict
     
     """
+
     write_kwargs = {}
     if fnames == (1):
         fnames = [f'{fnames}-{i:03d}' for i in range(len(splits))]
@@ -125,6 +126,7 @@ def collect_images(
         selected_atoms = [selected_atoms[i] for i in selected_inds]
 
     if shuffle:
+        assert not sort_by_energy_per_atom, 'Either sort or shuffle, not both'
         np.random.shuffle(selected_atoms)
     elif sort_by_energy_per_atom:
         assert not shuffle, 'Either sort or shuffle, not both'

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/vasp/vasp.py

@@ -73,9 +73,24 @@ def vasp(
     :param image: Initial image for VASP calculation. Image specification,
         see :func:`asimtools.utils.get_atoms`
     :type image: Dict
-    :param vaspinput_args: Dictionary of pymatgen's VaspInput arguments. 
+    :param user_incar_settings: Dictionary of INCAR settings to override
+        defaults or MP settings, defaults to None
+    :type user_incar_settings: Dict, optional
+    :param user_kpoints_settings: Dictionary of KPOINTS settings to override
+        defaults or MP settings, defaults to None
+    :type user_kpoints_settings: Dict, optional
+    :param user_potcar_functional: Potcar functional to use in case of MP
+        settings, defaults to 'PBE_64'
+    :type user_potcar_functional: str, optional
+    :param potcar: Dictionary specifying Potcar settings, see
+        :class:`pymatgen.io.vasp.inputs.Potcar`, defaults to None
+    :type potcar: Dict, optional
+    :param prev_calc: Path to previous VASP calculation to use as starting
+        point for MP settings, defaults to None
+    :type prev_calc: os.PathLike, optional
+    :param vaspinput_kwargs: Dictionary of pymatgen's VaspInput arguments. 
         See :class:`pymatgen.io.vasp.inputs.VaspInput`
-    :type vaspinput_args: Dict
+    :type vaspinput_kwargs: Dict
     :param command: Command with which to run VASP, defaults to 'vasp_std'
     :type command: str, optional
     :param mpset: Materials Project VASP set to use see 
@@ -84,6 +99,9 @@ def vasp(
     :param write_image_output: Whether to write output image in standard 
         asimtools format to file, defaults to False
     :type write_image_output: bool, optional
+    :param run_vasp: Whether to run VASP after writing input files,
+        defaults to True
+    :type run_vasp: bool, optional
     """
 
     if vaspinput_kwargs is None:

asimtools/asimmodules/workflows/image_array.py

@@ -14,7 +14,8 @@
 
 def image_array(
     images: Dict,
-    subsim_input: Dict,
+    subsim_input: Optional[Dict] = None,
+    template_sim_input: Optional[Dict] = None, 
     calc_input: Optional[Dict] = None,
     env_input: Optional[Dict] = None,
     array_max: Optional[int] = None,
@@ -33,8 +34,14 @@ def image_array(
 
     :param images: Images specification, see :func:`asimtools.utils.get_images`
     :type images: Dict
-    :param subsim_input: sim_input of asimmodule to be run
-    :type subsim_input: Dict
+    :param subsim_input: sim_input of asimmodule to be run, included for backward
+        compatibility, please use template_sim_input instead,
+        defaults to None
+    :type subsim_input: Optional[Dict], optional
+    :param template_sim_input: sim_input of asimmodule to be run, defaults to None
+    :type template_sim_input: Optional[Dict], optional
+    :param str_btn_args: args to pass to :func:`asimtools.utils.get_str_btn`
+    :type str_btn_args: Optional[Sequence], optional
     :param calc_input: calc_input to override global file, defaults to None
     :type calc_input: Optional[Dict], optional
     :param env_input: env_input to override global file, defaults to None
@@ -83,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/calculators.py

@@ -38,7 +38,7 @@ def load_calc(
             calc_params = calc_input[calc_id]
         except KeyError as exc:
             msg = f'Calculator with calc_id: {calc_id} not found in'
-            msg += f'calc_input {calc_input}'
+            msg += f'calc_input {[c for c in calc_input]}'
             raise KeyError(msg) from exc
         except AttributeError as exc:
             raise AttributeError('No calc_input found') from exc
@@ -400,6 +400,28 @@ def load_ase_dftd3(calc_params):
 
     return calc
 
+def load_aqcat(calc_params):
+    """Load AQCat Calculator
+    :param calc_params: args to pass to loader, including checkpoint_path
+    :type calc_params: Dict
+    :return: AQCat calculator
+    :rtype: :class:`fairchem.core.common.relaxation.ase_utils.patched_calc`
+    """
+    from fairchem.core.common.relaxation.ase_utils import patched_calc
+
+    CHECKPOINT_PATH = "aqcat25/checkpoints_aqcat_ev2/ev2-in+midFiLM-AQCat25+OC20-20M_20251008_223220.pt"
+
+    try:
+        calc = patched_calc(**calc_params)
+    except Exception:
+        logging.error(
+            "Failed to load AQCat FAIRChemCalculator with parameters:\n %s", \
+                calc_params
+        )
+        raise
+
+    return calc
+
 external_calcs = {
     'NequIP': load_nequip,
     'Allegro': load_nequip,
@@ -415,4 +437,5 @@ def load_ase_dftd3(calc_params):
     'fairchemV2': load_fairchemV2,
     'fairchem': load_fairchemV2,
     'ASEDFTD3': load_ase_dftd3,
+    'AQCat': load_aqcat
 }

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)