diff --git a/asimtools/asimmodules/matgl/train_matgl.py b/asimtools/asimmodules/matgl/train_matgl.py new file mode 100644 index 0000000..3ec8fc3 --- /dev/null +++ b/asimtools/asimmodules/matgl/train_matgl.py @@ -0,0 +1,50 @@ +#!/usr/bin/env python +''' +Asimmodule for training MACE models + +Author: mkphuthi@github.com +''' +from typing import Dict, Optional, Union +import sys +from pathlib import Path +import logging +import warnings +warnings.filterwarnings("ignore") +import json +from numpy.random import randint +from mace.cli.run_train import main as mace_run_train_main +from mace.cli.create_lammps_model import main as create_lammps_model + +def train_mace( + config: Union[Dict,str], + randomize_seed: bool = False, + compile_lammps: bool = False, +) -> Dict: + """Runs MACE training + + :param config: MACE config dictionary or path to config file + :type config: Union[Dict,str] + :param randomize_seed: Whether to randomize the seed, defaults to False + :type randomize_seed: bool + :return: Dictionary of results + :rtype: Dict + """ + + if isinstance(config, str): + with open(config, 'r') as fp: + config = json.load(fp) + + if randomize_seed: + config['seed'] = randint(0, 1000000) + + config_file_path = str(Path("mace_config.yaml").resolve()) + with open(config_file_path, "w") as f: + json.dump(config, f, indent=2) + + logging.getLogger().handlers.clear() + sys.argv = ["program", "--config", config_file_path] + mace_run_train_main() + + if compile_lammps: + create_lammps_model('mace_test_compiled.model') + return {} diff --git a/asimtools/asimmodules/transformations/delete_atoms.py b/asimtools/asimmodules/transformations/delete_atoms.py new file mode 100644 index 0000000..9c49976 --- /dev/null +++ b/asimtools/asimmodules/transformations/delete_atoms.py @@ -0,0 +1,79 @@ +''' +Produce a set of images with unit cells scaled compared to the input + +author: mkphuthi@github.com +''' + +from typing import Dict, Optional, Sequence +import numpy as np +from ase.io import write +from asimtools.utils import ( + get_atoms, +) + +def apply_scale(old_atoms, scale): + ''' Applies a scaling factor to a unit cell ''' + atoms = old_atoms.copy() + new_cell = atoms.get_cell() * scale + atoms.set_cell(new_cell, scale_atoms=True) + atoms.info['scale'] = f'{scale:.3f}' + return atoms + +def scale_unit_cells( + image: Dict, + scales: Optional[Sequence] = None, + logspace: Optional[Sequence] = None, + linspace: Optional[Sequence] = None, + scale_by: str = 'a', +) -> Dict: + """Produce a set of images with unit cells scaled compared to the input + + :param image: Image specification, see :func:`asimtools.utils.get_atoms` + :type image: Dict + :param scales: Scaling values by which to scale cell, defaults to None + :type scales: Optional[Sequence], optional + :param logspace: Parameters to pass to np.logspace for scaling values, + defaults to None + :type logspace: Optional[Sequence], optional + :param linspace: Parameters to pass to np.linspace for scaling values, + defaults to None + :type linspace: Optional[Sequence], optional + :param scale_by: Scale either "volume" or "a" which is lattice parameter, + defaults to 'a' + :type scale_by: str, optional + :raises ValueError: If more than one of scales, linspace, logspace are + provided + :return: Path to xyz file + :rtype: Dict + """ + + assert scale_by in ['volume', 'a'], \ + 'Only scaling by "a" and "volume" allowed' + + if (scales is None and linspace is None and logspace is not None): + scales = np.logspace(*logspace) + elif (scales is None and linspace is not None and logspace is None): + scales = np.linspace(*linspace) + elif (scales is not None and linspace is None and logspace is None): + pass + else: + raise ValueError( + 'Provide only one of factors, factor_logspacem factor_linspace' + ) + + atoms = get_atoms(**image) + + scales = np.array(scales) + if scale_by == 'volume': + scales = scales**(1/3) + + # Make a database of structures with the volumes scaled appropriately + scaled_images = [] + for scale in scales: + new_atoms = apply_scale(atoms, scale) + scaled_images.append(new_atoms) + + scaled_images_file = 'scaled_unitcells_output.xyz' + write(scaled_images_file, scaled_images, format='extxyz') + + return {'files': {'images': scaled_images_file}} diff --git a/asimtools/job.py b/asimtools/job.py index dddc52a..46f4bbc 100644 --- a/asimtools/job.py +++ b/asimtools/job.py @@ -44,6 +44,7 @@ def __init__( sim_input: Dict, env_input: Union[Dict,None] = None, calc_input: Union[Dict,None] = None, + asimrun_mode: bool = False, ) -> None: if env_input is None: env_input = get_env_input() @@ -63,7 +64,7 @@ def __init__( self.sim_input['src_dir'] = self.launchdir self.env_id = self.sim_input.get('env_id', None) - if self.env_id is not None and self.env_input is not None: + if self.env_id is not None and not asimrun_mode: self.env = self.env_input[self.env_id] else: self.env = { @@ -1039,7 +1040,7 @@ def submit(self, dependency: Union[List,None] = None, debug: bool = False) -> Li return job_ids -def load_job_from_directory(workdir: os.PathLike) -> Job: +def load_job_from_directory(workdir: os.PathLike, asimrun_mode=False) -> Job: ''' Loads a job from a given directory ''' workdir = Path(workdir) assert workdir.exists(), f'Work directory "{workdir}" does not exist' @@ -1066,6 +1067,7 @@ def load_job_from_directory(workdir: os.PathLike) -> Job: sim_input=sim_input, env_input=env_input, calc_input=calc_input, + asimrun_mode=asimrun_mode, ) # This makes sure that wherever we may be loading the job from, we refer diff --git a/asimtools/scripts/asim_run.py b/asimtools/scripts/asim_run.py index 4fa90a4..8a8432a 100755 --- a/asimtools/scripts/asim_run.py +++ b/asimtools/scripts/asim_run.py @@ -132,7 +132,7 @@ def main(args=None) -> None: sim_func = getattr(sim_module, func_name) cwd = Path('.').resolve() - job = load_job_from_directory(cwd) + job = load_job_from_directory(cwd, asimrun_mode=True) job.start() try: