From 574fcb5ac8b8985f108654a7d55da7f5887c6c86 Mon Sep 17 00:00:00 2001 From: Yusuke KONISHI Date: Wed, 20 Mar 2024 14:31:19 +0900 Subject: [PATCH 1/2] update for Nequip, Allegro, MLIP-3 --- .../latgas_abinitio_interface/__init__.py | 4 + abics/scripts/train.py | 122 ++++++++++++------ .../en/source/inputfiles/parameter_solver.rst | 2 +- .../en/source/inputfiles/parameter_train.rst | 2 +- docs/sphinx/en/source/tutorial/index.rst | 1 + .../ja/source/inputfiles/parameter_solver.rst | 12 ++ .../ja/source/inputfiles/parameter_train.rst | 3 +- pyproject.toml | 2 + 8 files changed, 104 insertions(+), 44 deletions(-) diff --git a/abics/applications/latgas_abinitio_interface/__init__.py b/abics/applications/latgas_abinitio_interface/__init__.py index 5e9ff8cb..9af0e399 100644 --- a/abics/applications/latgas_abinitio_interface/__init__.py +++ b/abics/applications/latgas_abinitio_interface/__init__.py @@ -17,10 +17,14 @@ # from .default_observer import * from .map2perflat import * from .aenet_trainer import * +from .nequip_trainer import * +from .mlip_3_trainer import * from .vasp import VASPSolver from .qe import QESolver from .aenet import AenetSolver from .aenet_pylammps import AenetPyLammpsSolver +from .nequip import NequipSolver +from .mlip_3 import MLIP_3_Solver from .openmx import OpenMXSolver from .user_function_solver import UserFunctionSolver diff --git a/abics/scripts/train.py b/abics/scripts/train.py index a028da65..d6e803b7 100644 --- a/abics/scripts/train.py +++ b/abics/scripts/train.py @@ -15,28 +15,28 @@ # along with this program. If not, see http://www.gnu.org/licenses/. from __future__ import annotations -from typing import MutableMapping, Any -import sys, datetime - -import os, sys +import datetime import itertools +import logging +import os +import sys +from typing import Any, MutableMapping -import numpy as np import networkx as nx +import numpy as np from pymatgen.core import Structure -from abics import __version__ -from abics.applications.latgas_abinitio_interface.params import DFTParams, TrainerParams -from abics.applications.latgas_abinitio_interface import aenet_trainer -from abics.applications.latgas_abinitio_interface import map2perflat -from abics.applications.latgas_abinitio_interface.defect import ( - defect_config, - DFTConfigParams, +from .. import __version__, loggers +from ..applications.latgas_abinitio_interface import ( + aenet_trainer, map2perflat, nequip_trainer, mlip_3_trainer +) +from ..applications.latgas_abinitio_interface.defect import ( + DFTConfigParams, defect_config +) +from ..applications.latgas_abinitio_interface.params import ( + DFTParams, TrainerParams ) - -import logging -import abics.loggers as loggers logger = logging.getLogger("main") @@ -74,7 +74,7 @@ def main_impl(params_root: MutableMapping): species = config.structure.symbol_set dummy_sts = {sp: config.dummy_structure_sp(sp) for sp in species} - if trainer_type != "aenet": + if trainer_type not in ["aenet", "allegro", "nequip", "mlip_3"]: logger.error("Unknown trainer: ", trainer_type) sys.exit(1) @@ -100,7 +100,6 @@ def main_impl(params_root: MutableMapping): logger.info("--Done") - logger.info("-Mapping relaxed structures in AL* to on-lattice model...") # val_map is a list of list [[sp0, vac0], [sp1, vac1], ...] @@ -120,7 +119,7 @@ def main_impl(params_root: MutableMapping): for pair in itertools.combinations(sp_list, 2): G.add_edge(*pair) sp_groups = nx.connected_components(G) - dummy_sts_share : list[tuple[Structure, list]] = [] + dummy_sts_share: list[tuple[Structure, list]] = [] for c in nx.connected_components(G): # merge dummy structures for species that share sublattices sps = list(c) @@ -148,7 +147,9 @@ def main_impl(params_root: MutableMapping): step_ids.append(int(words[2])) for step_id, energy in zip(step_ids, energies_ref): if os.path.exists(f"structure.{step_id}_mapped.vasp"): - structures.append(Structure.from_file(f"structure.{step_id}_mapped.vasp")) + structures.append( + Structure.from_file(f"structure.{step_id}_mapped.vasp") + ) energies.append(energy) rpl += 1 os.chdir(rootdir) @@ -182,7 +183,9 @@ def main_impl(params_root: MutableMapping): st_tmp.remove_species(["X"]) mapped_sts.append(st_tmp) if num_sp != len(st_tmp): - logger.info(f"--mapping failed for structure {step_id} in replica {rpl}") + logger.info( + f"--mapping failed for structure {step_id} in replica {rpl}" + ) mapping_success = False for sts in mapped_sts[1:]: @@ -192,19 +195,23 @@ def main_impl(params_root: MutableMapping): mapped_sts[0].remove_species(ignore_species) if mapping_success: structures.append(mapped_sts[0]) - mapped_sts[0].to(filename=f"structure.{step_id}_mapped.vasp", fmt="POSCAR") + mapped_sts[0].to( + filename=f"structure.{step_id}_mapped.vasp", fmt="POSCAR" + ) energies.append(energy) rpl += 1 os.chdir(rootdir) logger.info("--Finished mapping") - + generate_input_dirs = [] train_input_dirs = [] predict_input_dirs = [] if dftparams.ensemble: if len(trainer_input_dirs) != len(base_input_dir): - logger.error("You must set the number of trainer input dirs equal to baseinput dirs for ensemble NNP") + logger.error( + "You must set the number of trainer input dirs equal to baseinput dirs for ensemble NNP" + ) sys.exit(1) for d in trainer_input_dirs: generate_input_dirs.append(os.path.join(d, "generate")) @@ -215,18 +222,46 @@ def main_impl(params_root: MutableMapping): train_exe = trainer_commands[1] trainers = [] - for i in range(len(trainer_input_dirs)): - trainers.append( - aenet_trainer( - structures, - energies, - generate_input_dirs[i], - train_input_dirs[i], - predict_input_dirs[i], - generate_exe, - train_exe, + if trainer_type == "aenet": + for i in range(len(trainer_input_dirs)): + trainers.append( + aenet_trainer( + structures, + energies, + generate_input_dirs[i], + train_input_dirs[i], + predict_input_dirs[i], + generate_exe, + train_exe, + ) + ) + elif trainer_type == "allegro" or trainer_type == "nequip": + for i in range(len(trainer_input_dirs)): + trainers.append( + nequip_trainer( + structures, + energies, + generate_input_dirs[i], + train_input_dirs[i], + predict_input_dirs[i], + generate_exe, + train_exe, + trainer_type, + ) + ) + elif trainer_type == "mlip_3": + for i in range(len(trainer_input_dirs)): + trainers.append( + mlip_3_trainer( + structures, + energies, + generate_input_dirs[i], + train_input_dirs[i], + predict_input_dirs[i], + generate_exe, + train_exe, + ) ) - ) trainers[0].prepare() # We use threads to parallelize generate.x over ensemble members @@ -248,15 +283,18 @@ def main_impl(params_root: MutableMapping): for trainer in trainers: trainer.generate_wait() logger.info(f"--Finished generate run(s)") - + # We use MPI version of train.x so no need to write parallel code here for i, trainer in enumerate(trainers): logger.info(f"-Training run in train{i}") trainer.train(train_dir="train{}".format(i)) logger.info(f"--Training run finished in train{i}") logger.info(f"-Preparing NN model for abics_sampling in {base_input_dir[i]}") - trainer.new_baseinput(base_input_dir[i]) - logger.info(f"--Success.") + if trainer_type == "aenet": + trainer.new_baseinput(base_input_dir[i]) + if trainer_type in ["allegro","nequip","mlip_3"]: + trainer.new_baseinput(base_input_dir[i], train_dir=f"train{i}") + logger.info("--Success.") with open("ALloop.progress", "a") as fi: logger.info("-Writing ALloop.progress") @@ -272,14 +310,16 @@ def main(): now = datetime.datetime.now() import toml + tomlfile = sys.argv[1] if len(sys.argv) > 1 else "input.toml" params_root = toml.load(tomlfile) loggers.set_log_handles( - app_name = "train", - level = logging.INFO, - console = "serial", - params=params_root.get("log", {})) + app_name="train", + level=logging.INFO, + console="serial", + params=params_root.get("log", {}), + ) logger.info(f"Running abics_train (abICS v{__version__}) on {now}") logger.info(f"-Reading input from: {tomlfile}") diff --git a/docs/sphinx/en/source/inputfiles/parameter_solver.rst b/docs/sphinx/en/source/inputfiles/parameter_solver.rst index 36bb88fb..29932819 100644 --- a/docs/sphinx/en/source/inputfiles/parameter_solver.rst +++ b/docs/sphinx/en/source/inputfiles/parameter_solver.rst @@ -53,7 +53,7 @@ Keywords **Format :** str **Description :** - The solver type (``OpenMX, QE, VASP, aenet, aenetPyLammps, potts``). + The solver type (``OpenMX, QE, VASP, aenet, aenetPyLammps, nequip, allegro, mlip_3, potts``). When ``potts``, the following parameters are not used. - ``path`` diff --git a/docs/sphinx/en/source/inputfiles/parameter_train.rst b/docs/sphinx/en/source/inputfiles/parameter_train.rst index c9372636..8013eae6 100644 --- a/docs/sphinx/en/source/inputfiles/parameter_train.rst +++ b/docs/sphinx/en/source/inputfiles/parameter_train.rst @@ -33,7 +33,7 @@ Key words **Format :** str - **Description :** The trainer to generate the neural network potential (currently only 'aenet'). + **Description :** The trainer to generate the neural network potential (currently 'aenet', 'nequip', 'allegro', and 'mlip_3'). - ``base_input_dir`` diff --git a/docs/sphinx/en/source/tutorial/index.rst b/docs/sphinx/en/source/tutorial/index.rst index 2cede1fb..a4f559fa 100644 --- a/docs/sphinx/en/source/tutorial/index.rst +++ b/docs/sphinx/en/source/tutorial/index.rst @@ -10,3 +10,4 @@ Input files are provided in ``examples/active_learning_qe/`` . :maxdepth: 2 aenet + other_models diff --git a/docs/sphinx/ja/source/inputfiles/parameter_solver.rst b/docs/sphinx/ja/source/inputfiles/parameter_solver.rst index df483652..6b8dcce0 100644 --- a/docs/sphinx/ja/source/inputfiles/parameter_solver.rst +++ b/docs/sphinx/ja/source/inputfiles/parameter_solver.rst @@ -77,6 +77,18 @@ - LAMMPS を経由してaenetを利用します. + - ``nequip`` + + - NequIPを利用します. + + - ``allegro`` + + - Allegroを利用します. + + - ``mlip_3`` + + - MLIP-3を利用します. + - ``user`` - ユーザー定義のソルバーを利用します. diff --git a/docs/sphinx/ja/source/inputfiles/parameter_train.rst b/docs/sphinx/ja/source/inputfiles/parameter_train.rst index 3ad2c4a6..81079cf3 100644 --- a/docs/sphinx/ja/source/inputfiles/parameter_train.rst +++ b/docs/sphinx/ja/source/inputfiles/parameter_train.rst @@ -34,7 +34,8 @@ **形式 :** str型 **説明 :** - 訓練データから配置エネルギー予測モデルを学習する学習器の設定を行います.現在のところ、abICSではaenetのみに対応しています. + 訓練データから配置エネルギー予測モデルを学習する学習器の設定を行います. + 現在のところ、abICSではaenet, nequip, allegro, mlip_3に対応しています. - ``base_input_dir`` diff --git a/pyproject.toml b/pyproject.toml index 62ca8462..bce4e7b6 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -20,8 +20,10 @@ scipy = "^1" mpi4py = "^3" pymatgen = ">=2019.12.3 <2023.5.8" qe_tools = "^1.1" +nequip = {version=">=0.5.6", optional=true} [tool.poetry.extras] +nequip = ["nequip"] [tool.poetry.dev-dependencies] Sphinx = "^4.5.0" From ac8522cbd466db05a6b0d22ade7b64277d968143 Mon Sep 17 00:00:00 2001 From: Yusuke KONISHI Date: Wed, 20 Mar 2024 14:36:36 +0900 Subject: [PATCH 2/2] add files for Nequip, Allegro, MLIP-3 --- .../latgas_abinitio_interface/mlip_3.py | 384 ++++++++++++++++++ .../mlip_3_trainer.py | 187 +++++++++ .../latgas_abinitio_interface/nequip.py | 257 ++++++++++++ .../nequip_trainer.py | 205 ++++++++++ .../en/source/tutorial/other_models.rst | 297 ++++++++++++++ docs/sphinx/ja/source/tutorial/index.rst | 1 + .../ja/source/tutorial/other_models.rst | 295 ++++++++++++++ 7 files changed, 1626 insertions(+) create mode 100644 abics/applications/latgas_abinitio_interface/mlip_3.py create mode 100644 abics/applications/latgas_abinitio_interface/mlip_3_trainer.py create mode 100644 abics/applications/latgas_abinitio_interface/nequip.py create mode 100644 abics/applications/latgas_abinitio_interface/nequip_trainer.py create mode 100644 docs/sphinx/en/source/tutorial/other_models.rst create mode 100644 docs/sphinx/ja/source/tutorial/other_models.rst diff --git a/abics/applications/latgas_abinitio_interface/mlip_3.py b/abics/applications/latgas_abinitio_interface/mlip_3.py new file mode 100644 index 00000000..5ce01d31 --- /dev/null +++ b/abics/applications/latgas_abinitio_interface/mlip_3.py @@ -0,0 +1,384 @@ +# ab-Initio Configuration Sampling tool kit (abICS) +# Copyright (C) 2019- The University of Tokyo +# +# MLIP-3 solver +# Masashi Noda, Yusuke Konishi (Academeia Co., Ltd.) 2024 +# +# This program is free software: you can redistribute it and/or modify +# it under the terms of the GNU General Public License as published by +# the Free Software Foundation, either version 3 of the License, or +# (at your option) any later version. +# +# This program is distributed in the hope that it will be useful, +# but WITHOUT ANY WARRANTY; without even the implied warranty of +# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +# GNU General Public License for more details. +# +# You should have received a copy of the GNU General Public License +# along with this program. If not, see http://www.gnu.org/licenses/. + +""" +Adapted from pymatgen.io.xcrysden distributed under the MIT License +# Copyright (c) Pymatgen Development Team. +# Distributed under the terms of the MIT License. +""" + +from __future__ import annotations + +import io +import os +import shutil +import sys +from collections import namedtuple + +import numpy as np +from pymatgen.core import Structure + +from .base_solver import SolverBase, register_solver +from .params import ALParams, DFTParams + +def map_species_to_sequential_numbers(original_list): + """ + Maps a list of species to sequential numbers, starting from 1. + + Parameters + ---------- + original_list : list + List of species. + + Returns + ------- + list + List of sequential numbers. + """ + # Use a dictionary to map each unique element to a new number + mapping = {} + current_number = 1 + + for item in original_list: + if item not in mapping: + mapping[item] = current_number + current_number += 1 + + # Map each element of the original list to the new number + return [mapping[item] for item in original_list] + +def to_CFG(structure: Structure, energy, write_force_zero=False): + """ + Returns a string with the structure in CFG format + CFG format is a format used in input of MLIP-3 + + Parameters + ---------- + structure : pymatgen.Structure + Atomic structure + energy : float + Total energy + write_force_zero : bool + If True, the forces are written as zeros. + If False, the forces are written as the forces in the structure object. + + Returns + ------- + str + String with the structure in CFG format + """ + + lines = [] + app = lines.append + + app("BEGIN_CFG") + app(" Size") + + cart_coords = structure.cart_coords + app("%6d" % len(cart_coords)) + + cell = structure.lattice.matrix + app(" Supercell") + for i in range(3): + app("%16.6f%16.6f%16.6f" % tuple(cell[i])) + + species = structure.species + mapped_species = map_species_to_sequential_numbers(species) + + site_properties = structure.site_properties + if "forces" not in site_properties.keys(): + write_force_zero = True + else: + forces = site_properties["forces"] + + app(" AtomData: id type cartes_x cartes_y cartes_z fx fy fz") + if write_force_zero: + for a in range(len(cart_coords)): + app("%14d" % int(str(a+1)) + + "%5d" % mapped_species[a] + + "%15.6f%14.6f%14.6f" % tuple(cart_coords[a]) + + "%13.6f%12.6f%12.6f" % tuple([0.0, 0.0, 0.0]) + ) + else: + for a in range(len(cart_coords)): + app("%14d" % int(str(a+1)) + + "%5d" % mapped_species[a] + + "%15.6f%14.6f%14.6f" % tuple(cart_coords[a]) + + "%13.6f%12.6f%12.6f" % tuple(forces[a]) + ) + + app(" Energy") + app("%26.12f" % energy) + app(" PlusStress: xx yy zz yz xz xy") + app("%16.5f%12.5f%12.5f%12.5f%12.5f%12.5f" % tuple([0.0, 0.0, 0.0, 0.0, 0.0, 0.0])) + app(" Feature EFS_by VASP") + app("END_CFG") + app("") + app("") + + return "\n".join(lines) + +def read_CFG(input_string: str): + """ + Reads a string with the structure in CFG format and returns a dictionary + + Parameters + ---------- + input_string : str + String with the structure in CFG format + + Returns + ------- + dict + Dictionary with the structure in CFG format + """ + cfg_dic = {} + size = 0 + lines = input_string.split('\n') + for i, line in enumerate(lines): + if 'Size' in line: + size = int(lines[i+1]) + if 'Supercell' in line: + supercell = [] + for j in range(3): + supercell.append([float(x) for x in lines[i+j+1].split()]) + cfg_dic['supercell'] = supercell + if 'AtomData' in line: + atom_data = [] + atom_type = [] + for j in range(size): + atom_data.append([float(x) for x in lines[i+j+1].split()[2:5]]) + atom_type.append(int(lines[i+j+1].split()[1])) + cfg_dic['atom_data'] = atom_data + cfg_dic['atom_type'] = atom_type + + return cfg_dic + +def from_CFG(input_string: str, species): + """ + Returns a Structure object from a string with the structure in CFG format + + Parameters + ---------- + input_string : str + String with the structure in CFG format + species : list + List of species + + Returns + ------- + pymatgen.Structure + Atomic structure + """ + cfg_dic = read_CFG(input_string) + list_species = [species[i-1] for i in cfg_dic['atom_type']] + s = Structure(cfg_dic['supercell'], list_species, cfg_dic['atom_data'], coords_are_cartesian=True) + return s + +# Need to mod for MLIP-3 +class MLIP_3_Solver(SolverBase): + """ + This class defines the MLIP-3 solver. + """ + + def __init__( + self, path_to_solver: os.PathLike, ignore_species=None, + run_scheme="subprocess" + ): + """ + Initialize the solver. + + Parameters + ---------- + path_to_solver : str + Path to the solver. + """ + super().__init__(path_to_solver) + self.path_to_solver = path_to_solver + self.species = None + self.input = MLIP_3_Solver.Input(self, ignore_species, run_scheme) + self.output = MLIP_3_Solver.Output(self) + + def name(self): + return "mlip_3" + + class Input(object): + def __init__( + self, mlip3_solver, ignore_species: str | None, run_scheme="subprocess" + ): + self.mlip3_solver = mlip3_solver + self.base_info = None + self.pos_info = None + self.pot_info = None + self.ignore_species = ignore_species + self.species = None + self.run_scheme = run_scheme + + def from_directory(self, base_input_dir: os.PathLike): + """ + Initialize information from files in base_input_dir. + + Parameters + ---------- + base_input_dir : str + Path to the directory including base input files. + """ + + # set information of base_input and + # pos_info from files in base_input_dir + self.base_info = os.path.abspath(base_input_dir) + # self.pos_info = open( + # '{}/structure.xsf'.format(base_input_dir), 'r' + # ).read() + + def update_info_by_structure(self, structure: Structure): + """ + Update information by atomic structure. + + Parameters + ---------- + structure : pymatgen.Structure + Atomic structure + """ + if self.ignore_species is not None: + structure = structure.copy() + structure.remove_species(self.ignore_species) + self.mlip3_solver.species = [] + seen = set() + for specie in structure.species: + if specie not in seen: + self.mlip3_solver.species.append(str(specie)) + seen.add(specie) + self.pos_info = to_CFG(structure, 0.0) + + def update_info_from_files(self, output_dir, rerun): + """ + Do nothing. + """ + print("rerun not implemented. Something has gone wrong") + sys.exit(1) + + def write_input(self, output_dir: os.PathLike): + """ + Generate input files of the solver program. + + Parameters + ---------- + output_dir : os.PathLike + Path to working directory. + """ + # Write input files + if self.base_info is None: + raise AttributeError("Fail to set base_info.") + os.makedirs(output_dir, exist_ok=True) + for fname in os.listdir(self.base_info): + shutil.copy(os.path.join(self.base_info, fname), output_dir) + with open(os.path.join(output_dir, "structure.cfg"), "w") as f: + f.write(self.pos_info) + + def cl_args(self, nprocs, nthreads, output_dir): + """ + Generate command line arguments of the solver program. + + Parameters + ---------- + nprocs : int + The number of processes. + nthreads : int + The number of threads. + output_dir : str + Path to the working directory. + + Returns + ------- + args : list[str] + Arguments of command + """ + # Specify command line arguments + if self.run_scheme == "mpi_spawn_ready": + return [ + "calculate_efs", + os.path.join(output_dir, "pot.almtp"), + os.path.join(output_dir, "structure.cfg"), + output_dir, + ] + elif self.run_scheme == "subprocess": + return [ + "calculate_efs", + os.path.join(output_dir, "pot.almtp"), + os.path.join(output_dir, "structure.cfg"), + ] + + class Output(object): + def __init__(self, mlip3_solver): + self.mlip3_solver = mlip3_solver + + def get_results(self, output_dir): + """ + Get energy and structure obtained by the solver program. + + Parameters + ---------- + output_dir: str + Path to the working directory. + + Returns + ------- + phys : named_tuple("energy", "structure") + Total energy and atomic structure. + The energy is measured in the units of eV + and coordinates is measured in the units of Angstrom. + + """ + # Read results from files in output_dir and calculate values + Phys = namedtuple("PhysValues", ("energy", "structure")) + with open(os.path.join(output_dir, "structure.cfg")) as f: + lines = f.read() + structure = from_CFG(lines, self.mlip3_solver.species) + fi_io = io.StringIO(lines) + line = fi_io.readline() + #if "optimized" in lines: + # while "optimized" not in line: + # line = fi_io.readline() + # for i in range(4): + # fi_io.readline() + # for i in range(len(structure)): + # xyz = [float(x) for x in fi_io.readline().split()[1:4]] + # structure.replace( + # i, structure[i].species, coords=xyz, + # coords_are_cartesian=True + # ) + while "Energy" not in line: + line = fi_io.readline() + line = fi_io.readline() + energy = line + return Phys(np.float64(energy), structure) + + def solver_run_schemes(self): + return ("subprocess", "mpi_spawn_ready") + + @classmethod + def create(cls, params: ALParams | DFTParams): + path = params.path + ignore_species = params.ignore_species + run_scheme = params.solver_run_scheme + return cls(path, ignore_species, run_scheme) + + +register_solver("mlip_3", MLIP_3_Solver) diff --git a/abics/applications/latgas_abinitio_interface/mlip_3_trainer.py b/abics/applications/latgas_abinitio_interface/mlip_3_trainer.py new file mode 100644 index 00000000..71beed4e --- /dev/null +++ b/abics/applications/latgas_abinitio_interface/mlip_3_trainer.py @@ -0,0 +1,187 @@ +from __future__ import annotations + +import os +import pathlib +import shlex +import shutil +import subprocess +import time +from typing import Sequence + +import ase +from ase.calculators.singlepoint import SinglePointCalculator +from nequip.utils import Config +from pymatgen.core import Structure + +from ...util import expand_cmd_path +from . import mlip_3 + +def to_XSF(structure: Structure, write_force_zero=False): + """ + Returns a string with the structure in XSF format + See http://www.xcrysden.org/doc/XSF.html + """ + lines = [] + app = lines.append + + app("CRYSTAL") + app("# Primitive lattice vectors in Angstrom") + app("PRIMVEC") + cell = structure.lattice.matrix + for i in range(3): + app(" %.14f %.14f %.14f" % tuple(cell[i])) + + cart_coords = structure.cart_coords + app("# Cartesian coordinates in Angstrom.") + app("PRIMCOORD") + app(" %d 1" % len(cart_coords)) + species = structure.species + site_properties = structure.site_properties + if "forces" not in site_properties.keys(): + write_force_zero = True + else: + forces = site_properties["forces"] + + if write_force_zero: + for a in range(len(cart_coords)): + app( + str(species[a]) + + " %20.14f %20.14f %20.14f" % tuple(cart_coords[a]) + + " 0.0 0.0 0.0" + ) + else: + for a in range(len(cart_coords)): + app( + str(species[a]) + + " %20.14f %20.14f %20.14f" % tuple(cart_coords[a]) + + " %20.14f %20.14f %20.14f" % tuple(forces[a]) + ) + + return "\n".join(lines) + +class mlip_3_trainer: + def __init__( + self, + structures: Sequence[Structure], + energies: Sequence[float], + generate_inputdir: os.PathLike, + train_inputdir: os.PathLike, + predict_inputdir: os.PathLike, + generate_exe: str, + train_exe: str, + ): + self.structures = structures + self.energies = energies + self.generate_inputdir = generate_inputdir + self.train_inputdir = train_inputdir + self.predict_inputdir = predict_inputdir + self.train_exe = [ + expand_cmd_path(e) for e in shlex.split(train_exe) + ] + self.train_exe += ["train", "input.almtp", "input.cfg", "--save_to=./pot.almtp", "--iteration_limit=100", "--al_mode=nbh"] + assert len(self.structures) == len(self.energies) + self.numdata = len(self.structures) + self.is_prepared = False + self.is_trained = False + self.generate_outputdir = None + self.latgas_mode = True + + def prepare(self, latgas_mode=True, st_dir="mlip_3_XSF"): + rootdir = os.getcwd() + xsfdir = os.path.join(rootdir, st_dir) + + # prepare XSF files + os.makedirs(xsfdir, exist_ok=True) + os.chdir(xsfdir) + xsfdir = os.getcwd() + if latgas_mode: + for i, st in enumerate(self.structures): + xsf_string = to_XSF(st, write_force_zero=False) + xsf_string =\ + f"# total energy = {self.energies[i]} eV\n\n{xsf_string}" + with open(f"structure.{i}.xsf", "w") as fi: + fi.write(xsf_string) + else: + for i, st in enumerate(self.structures): + xsf_string = to_XSF(st, write_force_zero=False) + xsf_string =\ + f"# total energy = {self.energies[i]} eV\n\n{xsf_string}" + with open(f"structure.{i}.xsf", "w") as fi: + fi.write(xsf_string) + + self.latgas_mode = latgas_mode + os.chdir(rootdir) + + def generate_run(self, xsfdir="mlip_3_XSF", generate_dir="generate"): + # prepare generate + cfgdir = str(pathlib.Path(xsfdir).resolve()) + if os.path.exists(generate_dir): + shutil.rmtree(generate_dir) + shutil.copytree(xsfdir, generate_dir) + os.chdir(generate_dir) + + # prepare CFG file for MLIP-3 + if self.latgas_mode: + cfg_string = "" + for i, st in enumerate(self.structures): + lines = mlip_3.to_CFG(st, self.energies[i], write_force_zero=False) + cfg_string = cfg_string + lines + with open(f"input.cfg", "w") as fi: + fi.write(cfg_string) + else: + cfg_string = "" + for i, st in enumerate(self.structures): + lines = mlip_3.to_CFG(st, self.energies[i], write_force_zero=False) + cfg_string = cfg_string + lines + with open(f"input.cfg", "w") as fi: + fi.write(cfg_string) + + self.generate_outputdir = os.getcwd() + os.chdir(pathlib.Path(os.getcwd()).parent) + + def generate_wait(self): + interval = 0.1 # sec + #self.is_prepared = False + #if os.path.exists( + # os.path.join(self.generate_outputdir, "input.cfg") + #): + # self.is_prepared = True + self.is_prepared = True + time.sleep(interval) + if not self.is_prepared: + raise RuntimeError(f"{self.generate_outputdir}") + + def train(self, train_dir="train"): + if not self.is_prepared: + raise RuntimeError( + "you have to prepare the trainer before training!" + ) + if os.path.exists(train_dir): + shutil.rmtree(train_dir) + shutil.copytree(self.generate_outputdir, train_dir) + shutil.copy(os.path.join(self.train_inputdir, "input.almtp"), train_dir) + os.chdir(train_dir) + #command = self.train_exe + " train input.almtp input.cfg --save_to=./out/pot.mtp --interaction_limit=100 --al_mode=nbh" + command = self.train_exe + #print(os.getcwd()) + #print(command) + #print(os.path.exists("input.cfg")) + + with open(os.path.join(os.getcwd(), "stdout"), "w") as fi: + subprocess.run( + self.train_exe, stdout=fi, stderr=subprocess.STDOUT, check=True + ) + os.chdir(pathlib.Path(os.getcwd()).parent) + self.is_trained = True + + def new_baseinput(self, baseinput_dir, train_dir="train"): + try: + assert self.is_trained + except AssertionError as e: + e.args += "you have to train before getting results!" + + baseinput = str(pathlib.Path(baseinput_dir).resolve()) + os.makedirs(baseinput, exist_ok=True) + shutil.copy(os.path.join(train_dir, "input.cfg"), baseinput) + shutil.copy(os.path.join(train_dir, "pot.almtp"), baseinput) + diff --git a/abics/applications/latgas_abinitio_interface/nequip.py b/abics/applications/latgas_abinitio_interface/nequip.py new file mode 100644 index 00000000..eaffcc71 --- /dev/null +++ b/abics/applications/latgas_abinitio_interface/nequip.py @@ -0,0 +1,257 @@ +# ab-Initio Configuration Sampling tool kit (abICS) +# Copyright (C) 2019- The University of Tokyo +# +# NequIP solver +# Munehiro Kobayashi, Yusuke Konishi (Academeia Co., Ltd.) 2024 +# +# This program is free software: you can redistribute it and/or modify +# it under the terms of the GNU General Public License as published by +# the Free Software Foundation, either version 3 of the License, or +# (at your option) any later version. +# +# This program is distributed in the hope that it will be useful, +# but WITHOUT ANY WARRANTY; without even the implied warranty of +# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +# GNU General Public License for more details. +# +# You should have received a copy of the GNU General Public License +# along with this program. If not, see http://www.gnu.org/licenses/. + +""" +energy calculator using nequip python interface +""" + +from __future__ import annotations + +import os.path +from collections import namedtuple +import numpy as np +from pymatgen.core import Structure +import torch +from ase import Atoms +from nequip.data import AtomicDataDict, AtomicData +from nequip.utils import Config + +from .base_solver import SolverBase, register_solver +from .params import ALParams, DFTParams + +def to_XSF(structure: Structure, write_force_zero=False): + """ + Returns a string with the structure in XSF format + See http://www.xcrysden.org/doc/XSF.html + """ + lines = [] + app = lines.append + + app("CRYSTAL") + app("# Primitive lattice vectors in Angstrom") + app("PRIMVEC") + cell = structure.lattice.matrix + for i in range(3): + app(" %.14f %.14f %.14f" % tuple(cell[i])) + + cart_coords = structure.cart_coords + app("# Cartesian coordinates in Angstrom.") + app("PRIMCOORD") + app(" %d 1" % len(cart_coords)) + species = structure.species + site_properties = structure.site_properties + if "forces" not in site_properties.keys(): + write_force_zero = True + else: + forces = site_properties["forces"] + + if write_force_zero: + for a in range(len(cart_coords)): + app( + str(species[a]) + + " %20.14f %20.14f %20.14f" % tuple(cart_coords[a]) + + " 0.0 0.0 0.0" + ) + else: + for a in range(len(cart_coords)): + app( + str(species[a]) + + " %20.14f %20.14f %20.14f" % tuple(cart_coords[a]) + + " %20.14f %20.14f %20.14f" % tuple(forces[a]) + ) + + return "\n".join(lines) + +class NequipSolver(SolverBase): + """ + Nequip solver + + Attributes + ---------- + path_to_solver : str + Path to the solver + input : NequipSolver.Input + Input manager + output : NequipSolver.Output + Output manager + """ + + def __init__(self, ignore_species): + """ + Initialize the solver. + + """ + + super(NequipSolver, self).__init__("") + self.path_to_solver = self.calculate_energy + self.input = NequipSolver.Input(ignore_species) + self.output = NequipSolver.Output() + + def name(self): + return "nequip" + + def calculate_energy(self, fi, output_dir): + st = self.input.st + symbols = [site.specie.symbol for site in st] + positions = [site.coords for site in st] + pbc = (True, True, True) + cell = st.lattice.matrix + atoms = Atoms(symbols=symbols, positions=positions, pbc=pbc, cell=cell) + + atom_types = torch.tensor([self.input.element_list.index(atom) for atom in symbols], dtype=torch.long) + + data = AtomicData.from_ase(atoms, r_max=self.input.r_max) + data[AtomicDataDict.ATOM_TYPE_KEY] = atom_types + + self.input.model.eval() + with torch.no_grad(): + # Convert AtomicData to dictionary + data_dict = data.to_dict() + predicted = self.input.model(data_dict) + + # Get predicted energy + ene = predicted['total_energy'].item() + + self.output.st = st + self.output.ene = ene + + class Input(object): + """ + Input manager for Mock + + Attributes + ---------- + st : pymatgen.Structure + structure + """ + + st: Structure + + def __init__(self, ignore_species=None): + self.ignore_species = ignore_species + # self.st = Structure() + + def from_directory(self, base_input_dir): + """ + + Parameters + ---------- + base_input_dir : str + Path to the directory including base input files. + """ + self.base_input_dir = base_input_dir + self.model = torch.jit.load(os.path.join(base_input_dir, "deployed.pth")) + yaml_file = os.path.join(base_input_dir, "input.yaml") + yaml_dic = Config.from_file(yaml_file) + self.element_list = yaml_dic["chemical_symbols"] + self.r_max = yaml_dic["r_max"] + + def update_info_by_structure(self, structure): + """ + Update information by structure file + + Parameters + ---------- + structure : pymatgen.Structure + Atomic structure + """ + self.st = structure.copy() + if self.ignore_species is not None: + self.st.remove_species(self.ignore_species) + + def update_info_from_files(self, workdir, rerun): + """ + Do nothing + """ + pass + + def write_input(self, output_dir): + """ + Generate input files of the solver program. + + Parameters + ---------- + workdir : str + Path to working directory. + """ + if not os.path.exists(output_dir): + import shutil + + shutil.copytree(self.base_input_dir, output_dir) + + # self.st.to("POSCAR", os.path.join(output_dir, "structure.vasp")) + + def cl_args(self, nprocs, nthreads, workdir): + """ + Generate command line argument of the solver program. + + Parameters + ---------- + nprocs : int + The number of processes. + nthreads : int + The number of threads. + workdir : str + Path to the working directory. + + Returns + ------- + args : list[str] + Arguments of command + """ + return [workdir] + + class Output(object): + """ + Output manager. + """ + + def __init__(self): + pass + + def get_results(self, workdir): + """ + Get energy and structure obtained by the solver program. + + Parameters + ---------- + workdir : str + Path to the working directory. + + Returns + ------- + phys : named_tuple("energy", "structure") + Total energy and atomic structure. + The energy is measured in the units of eV + and coodinates is measured in the units of Angstrom. + """ + Phys = namedtuple("PhysVaules", ("energy", "structure")) + return Phys(self.ene, self.st) + + def solver_run_schemes(self): + return ("function",) + + @classmethod + def create(cls, params: ALParams | DFTParams): + ignore_species = params.ignore_species + return cls(ignore_species) + + +register_solver("nequip", NequipSolver) +register_solver("allegro", NequipSolver) diff --git a/abics/applications/latgas_abinitio_interface/nequip_trainer.py b/abics/applications/latgas_abinitio_interface/nequip_trainer.py new file mode 100644 index 00000000..6f6f7c0b --- /dev/null +++ b/abics/applications/latgas_abinitio_interface/nequip_trainer.py @@ -0,0 +1,205 @@ +from __future__ import annotations +from typing import Sequence + +import numpy as np +import os, pathlib, shutil, subprocess, shlex +import time + +from pymatgen.core import Structure + +from abics.util import expand_cmd_path +from abics.applications.latgas_abinitio_interface import nequip + +import ase +from ase import io +from ase.calculators.singlepoint import SinglePointCalculator +from nequip.utils import Config + +def to_XSF(structure: Structure, write_force_zero=False): + """ + Returns a string with the structure in XSF format + See http://www.xcrysden.org/doc/XSF.html + """ + lines = [] + app = lines.append + + app("CRYSTAL") + app("# Primitive lattice vectors in Angstrom") + app("PRIMVEC") + cell = structure.lattice.matrix + for i in range(3): + app(" %.14f %.14f %.14f" % tuple(cell[i])) + + cart_coords = structure.cart_coords + app("# Cartesian coordinates in Angstrom.") + app("PRIMCOORD") + app(" %d 1" % len(cart_coords)) + species = structure.species + site_properties = structure.site_properties + if "forces" not in site_properties.keys(): + write_force_zero = True + else: + forces = site_properties["forces"] + + if write_force_zero: + for a in range(len(cart_coords)): + app( + str(species[a]) + + " %20.14f %20.14f %20.14f" % tuple(cart_coords[a]) + + " 0.0 0.0 0.0" + ) + else: + for a in range(len(cart_coords)): + app( + str(species[a]) + + " %20.14f %20.14f %20.14f" % tuple(cart_coords[a]) + + " %20.14f %20.14f %20.14f" % tuple(forces[a]) + ) + + return "\n".join(lines) + +def xsf_to_ase(xsf): + ase_xsf = ase.io.read(xsf) + with open(xsf) as f: + lines = f.readlines() + + tot_energy = float(lines[0].split()[4]) + ase_xsf.calc = SinglePointCalculator(energy=tot_energy, atoms=ase_xsf) + return ase_xsf + +class nequip_trainer: + def __init__( + self, + structures: Sequence[Structure], + energies: Sequence[float], + generate_inputdir: os.PathLike, + train_inputdir: os.PathLike, + predict_inputdir: os.PathLike, + generate_exe: str, + train_exe: str, + trainer_type: str, + ): + self.structures = structures + self.energies = energies + self.generate_inputdir = generate_inputdir + self.train_inputdir = train_inputdir + self.predict_inputdir = predict_inputdir + self.generate_exe = [expand_cmd_path(e) for e in shlex.split(generate_exe)] + self.generate_exe.append("generate.in") + self.train_exe = [expand_cmd_path(e) for e in shlex.split(train_exe)] + self.train_exe.append("input.yaml") + # self.generate_exe = generate_exe + # self.train_exe = train_exe + assert len(self.structures) == len(self.energies) + self.numdata = len(self.structures) + self.is_prepared = False + self.is_trained = False + self.generate_outputdir = None + self.trainer_type = trainer_type + + def prepare(self, latgas_mode = True, st_dir = "nequipXSF"): + rootdir = os.getcwd() + xsfdir = os.path.join(rootdir, st_dir) + + # prepare XSF files for nequip + os.makedirs(xsfdir, exist_ok=True) + os.chdir(xsfdir) + xsfdir = os.getcwd() + if latgas_mode: + for i, st in enumerate(self.structures): + xsf_string = nequip.to_XSF(st, write_force_zero=False) + xsf_string = ( + "# total energy = {} eV\n\n".format(self.energies[i]) + xsf_string + ) + with open("structure.{}.xsf".format(i), "w") as fi: + fi.write(xsf_string) + else: + for i, st in enumerate(self.structures): + xsf_string = nequip.to_XSF(st, write_force_zero=False) + xsf_string = ( + "# total energy = {} eV\n\n".format(self.energies[i]) + xsf_string + ) + with open("structure.{}.xsf".format(i), "w") as fi: + fi.write(xsf_string) + + os.chdir(rootdir) + + def generate_run(self, xsfdir="nequipXSF", generate_dir="generate"): + # prepare generate + xsfdir = str(pathlib.Path(xsfdir).resolve()) + if os.path.exists(generate_dir): + shutil.rmtree(generate_dir) + shutil.copytree(self.generate_inputdir, generate_dir) + self.generate_dir = generate_dir + os.chdir(generate_dir) + xsf_paths = [ + os.path.join(xsfdir, "structure.{}.xsf".format(i)) + for i in range(self.numdata) + ] + ases = [xsf_to_ase(xsf) for xsf in xsf_paths] + #generate structure.xyz + ase.io.write("structure.xyz", ases) + self.generate_outputdir = os.getcwd() + os.chdir(pathlib.Path(os.getcwd()).parent) + + def generate_wait(self): + interval = 0.1 # sec + self.is_prepared = False + if os.path.exists(os.path.join(self.generate_outputdir, "structure.xyz")): + self.is_prepared = True + time.sleep(interval) + if not self.is_prepared: + raise RuntimeError(f"{self.generate_outputdir}") + + def train(self, train_dir = "train"): + if not self.is_prepared: + raise RuntimeError("you have to prepare the trainer before training!") + if os.path.exists(train_dir): + shutil.rmtree(train_dir) + shutil.copytree(self.train_inputdir, train_dir) + os.chdir(train_dir) + + yaml_dic = Config.from_file("input.yaml") + is_allegro = "allegro.model.Allegro" in yaml_dic["model_builders"] + if self.trainer_type == "nequip": + if is_allegro: + print("Warning: trainer_type=='nequip', but Allegro model is in input.yaml.") + else: + if not is_allegro: + print("Warning: trainer_type=='allegro', but Allegro model is not in input.yaml.") + + os.rename( + os.path.join(self.generate_outputdir, "structure.xyz"), + os.path.join(os.getcwd(), "structure.xyz"), + ) + # command = self.train_exe + " train.in" + # print(os.getcwd()) + # print(command) + # print(os.path.exists("train.in")) + + with open(os.path.join(os.getcwd(), "stdout"), "w") as fi: + subprocess.run( + self.train_exe, stdout=fi, stderr=subprocess.STDOUT, check=True + ) + os.chdir(pathlib.Path(os.getcwd()).parent) + self.is_trained = True + + def new_baseinput(self, baseinput_dir, train_dir = "train"): + try: + assert self.is_trained + except AssertionError as e: + e.args += "you have to train before getting results!" + + baseinput = str(pathlib.Path(baseinput_dir).resolve()) + os.makedirs(baseinput, exist_ok=True) + shutil.copy(os.path.join(train_dir,"input.yaml"),baseinput) + os.chdir(train_dir) + yaml_dic = Config.from_file("input.yaml") + root = yaml_dic["root"] + runname = yaml_dic["run_name"] + nequip_deploy = ["nequip-deploy","build","--train-dir",os.path.join(root,runname),os.path.join(baseinput,"deployed.pth")] + with open("nequip-deploy.out", "w") as fi: + subprocess.run( + nequip_deploy, stdout=fi, stderr=subprocess.STDOUT, check=True + ) + os.chdir(pathlib.Path(os.getcwd()).parent) diff --git a/docs/sphinx/en/source/tutorial/other_models.rst b/docs/sphinx/en/source/tutorial/other_models.rst new file mode 100644 index 00000000..960506c9 --- /dev/null +++ b/docs/sphinx/en/source/tutorial/other_models.rst @@ -0,0 +1,297 @@ +.. _tutorial_nequip: + +*********************************************** +Sampling using other machine learning models +*********************************************** + +In abICS, in addition to the aenet, it is possible to perform sampling using +other machine learning models such as NequIP, Allegro, and MLIP-3. +This section explains how to train and sample using each model. + +Sampling with NequIP +---------------------------------------------- + +Installation of NequIP +~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +To use ``nequip``, you need to install NequIP. + +Install it with the following command. + +.. code-block:: bash + + $ pip3 install wandb + $ pip3 install nequip + +Also, when installing abICS, you can install NequIP by specifying the [nequip] option. + +.. code-block:: bash + + $ cd /path/to/abics + $ pip3 install .[nequip] + +Preparation of input files +~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +First, prepare input_nequip.toml and set the parameters required to run NequIP. +Below, we extract [sampling.solver] and [train] with changes from the aenet input. + +.. code-block:: toml + + [sampling.solver] + type = 'nequip' + base_input_dir = './baseinput_nequip' + perturb = 0.0 + #run_scheme = 'subprocess' #'mpi_spawn_ready' + ignore_species = ["O"] + + [train] + type = 'nequip' + base_input_dir = './nequip_train_input' + exe_command = ['', 'nequip-train'] + ignore_species = ["O"] + vac_map = [] + restart = false + +Also, create the NequIP input file input.yaml in the nequip_train_input/train directory. + +.. code-block:: yaml + + root: results/spinel + run_name: run + seed: 123 + dataset_seed: 456 + + # network + num_basis: 8 + BesselBasis_trainable: true + PolynomialCutoff_p: 6 + l_max: 1 + r_max: 8.0 + parity: true + num_layers: 3 + num_features: 16 + + nonlinearity_type: gate + + nonlinearity_scalars: + e: silu + o: tanh + + nonlinearity_gates: + e: silu + o: tanh + + model_builders: + - SimpleIrrepsConfig + - EnergyModel + - PerSpeciesRescale + - RescaleEnergyEtc + + + dataset: ase + dataset_file_name: structure.xyz + chemical_symbols: + - Mg + - Al + + # logging + wandb: false + # verbose: debug + + # training + n_train: 70 + n_val: 10 + batch_size: 5 + train_val_split: random + #shuffle: true + metrics_key: validation_loss + use_ema: true + ema_decay: 0.99 + ema_use_num_updates: true + max_epochs: 100 + learning_rate: 0.01 + # loss function + loss_coeffs: total_energy + +The procedure of model learning and sampling is the same as aenet. + + +Sampling with Allegro +---------------------------------------------- + +Installation of Allegro +~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +To use ``allegro``, you need to install Allegro. + +Install it with the following command. + +.. code-block:: bash + + $ git clone --depth 1 https://github.com/mir-group/allegro.git + $ cd allegro + $ pip3 install . + + +Preparation of input files +~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +First, prepare input_allegro.toml and set the parameters required to run Allegro. +Below, we extract [sampling.solver] and [train] with changes from the aenet input. + +.. code-block:: toml + + [sampling.solver] + type = 'allegro' + base_input_dir = './baseinput_allegro' + perturb = 0.0 + #run_scheme = 'subprocess' #'mpi_spawn_ready' + ignore_species = ["O"] + + [train] + type = 'allegro' + base_input_dir = './allegro_train_input' + exe_command = ['', 'nequip-train'] + ignore_species = ["O"] + vac_map = [] + restart = false + +Also, create the Allegro input file input.yaml in the allegro_train_input/train directory. + +.. code-block:: yaml + + root: results/spinel + run_name: run + seed: 123 + dataset_seed: 456 + + # network + num_basis: 8 + BesselBasis_trainable: true + PolynomialCutoff_p: 6 + l_max: 1 + r_max: 8.0 + parity: o3_full + num_layers: 2 + num_features: 16 + + env_embed_multiplicity: 16 + embed_initial_edge: true + two_body_latent_mlp_latent_dimensions: [32, 64] + two_body_latent_mlp_nonlinearity: silu + latent_mlp_latent_dimensions: [64, 64] + latent_mlp_nonlinearity: silu + latent_mlp_initialization: uniform + latent_resnet: true + env_embed_mlp_latent_dimensions: [] + env_embed_mlp_nonlinearity: null + env_embed_mlp_initialization: uniform + edge_eng_mlp_latent_dimensions: [16] + edge_eng_mlp_nonlinearity: null + edge_eng_mlp_initialization: uniform + + model_builders: + - allegro.model.Allegro + - PerSpeciesRescale + - RescaleEnergyEtc + + + dataset: ase + dataset_file_name: structure.xyz + chemical_symbols: + - Mg + - Al + + # logging + wandb: false + # verbose: debug + + # training + n_train: 70 + n_val: 10 + batch_size: 5 + train_val_split: random + #shuffle: true + metrics_key: validation_loss + use_ema: true + ema_decay: 0.99 + ema_use_num_updates: true + max_epochs: 100 + learning_rate: 0.01 + # loss function + loss_coeffs: total_energy + +The procedure of model learning and sampling is the same as aenet. + + +Sampling with MLIP-3 +---------------------------------------------- + +Installation of MLIP-3 +~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +To use ``mlip-3``, you need to install MLIP-3. + +Install it with the following command. + +.. code-block:: bash + + $ git clone https://gitlab.com/ashapeev/mlip-3.git + $ cd mlip-3 + $ ./configure --no-mpi + $ make mlp + + +Preparation of input files +~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +First, prepare input_mlip3.toml and set the parameters required to run MLIP-3. +Below, we extract [sampling.solver] and [train] with changes from the aenet input. + +.. code-block:: toml + + [sampling.solver] + type = 'mlip_3' + path= '~/github/mlip-3/bin/mlp' + base_input_dir = './baseinput' + perturb = 0.0 + run_scheme = 'subprocess' #'mpi_spawn_ready' + ignore_species = ["O"] + + [train] + type = 'mlip_3' + base_input_dir = './mlip_3_train_input' + exe_command = ['~/github/mlip-3/bin/mlp','~/github/mlip-3/bin/mlp'] + ignore_species = ["O"] + vac_map = [] + restart = false + +In the above, the path in [sampling.solver] and the exe_command list in [train] +specify the path to the MLIP-3 executable file mlp. +Please change them according to your environment. + +Also, create the MLIP-3 input file input.almtp in the mlip_3_train_input/train directory. + +.. code-block:: none + + MTP + version = 1.1.0 + potential_name = MTP1m + species_count = 3 + potential_tag = + radial_basis_type = RBChebyshev + min_dist = 2.3 + max_dist = 5 + radial_basis_size = 8 + radial_funcs_count = 2 + alpha_moments_count = 8 + alpha_index_basic_count = 5 + alpha_index_basic = {{0, 0, 0, 0}, {0, 1, 0, 0}, {0, 0, 1, 0}, {0, 0, 0, 1}, {1, 0, 0, 0}} + alpha_index_times_count = 5 + alpha_index_times = {{0, 0, 1, 5}, {1, 1, 1, 6}, {2, 2, 1, 6}, {3, 3, 1, 6}, {0, 5, 1, 7}} + alpha_scalar_moments = 5 + alpha_moment_mapping = {0, 4, 5, 6, 7} + + +The procedure of model learning and sampling is the same as aenet. \ No newline at end of file diff --git a/docs/sphinx/ja/source/tutorial/index.rst b/docs/sphinx/ja/source/tutorial/index.rst index 7edce6dd..525d099f 100644 --- a/docs/sphinx/ja/source/tutorial/index.rst +++ b/docs/sphinx/ja/source/tutorial/index.rst @@ -9,3 +9,4 @@ :maxdepth: 2 aenet + other_models diff --git a/docs/sphinx/ja/source/tutorial/other_models.rst b/docs/sphinx/ja/source/tutorial/other_models.rst new file mode 100644 index 00000000..77dd2954 --- /dev/null +++ b/docs/sphinx/ja/source/tutorial/other_models.rst @@ -0,0 +1,295 @@ +.. _tutorial_nequip: + +************************************* +他のモデルを利用したサンプリング +************************************* + +abICSでは、機械学習モデルとして、aenet以外にも、 +NequIP, Allegro, MLIP-3を利用したサンプリングが可能となっています。 +本項では、それぞれのモデルの学習およびサンプリングの方法について説明します。 + +NequIPを利用したサンプリング +---------------------------------------------- + +NequIP のインストール +~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +``nequip`` の利用には、 NequIPのインストールが必要です。 + +下記コマンドにてインストールします。 + +.. code-block:: bash + + $ pip3 install wandb + $ pip3 install nequip + +また、abICSインストール時に[nequip]オプションを指定すれば、NequIPもインストールされます。 + +.. code-block:: bash + + $ cd /path/to/abics + $ pip3 install .abics[nequip] + +インプットファイルの準備 +~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +まず、input_nequip.tomlを準備し、NequIPの実行に必要なパラメータを設定します。 +下では、aenetのインプットから変更のある[sampling.solver]と[train]を抜粋しています。 + +.. code-block:: toml + + [sampling.solver] + type = 'nequip' + base_input_dir = './baseinput_nequip' + perturb = 0.0 + #run_scheme = 'subprocess' #'mpi_spawn_ready' + ignore_species = ["O"] + + [train] + type = 'nequip' + base_input_dir = './nequip_train_input' + exe_command = ['', 'nequip-train'] + ignore_species = ["O"] + vac_map = [] + restart = false + +また、NequIPのインプットファイルinput.yamlをnequip_train_input/trainディレクトリに作成します。 + +.. code-block:: yaml + + root: results/spinel + run_name: run + seed: 123 + dataset_seed: 456 + + # network + num_basis: 8 + BesselBasis_trainable: true + PolynomialCutoff_p: 6 + l_max: 1 + r_max: 8.0 + parity: true + num_layers: 3 + num_features: 16 + + nonlinearity_type: gate + + nonlinearity_scalars: + e: silu + o: tanh + + nonlinearity_gates: + e: silu + o: tanh + + model_builders: + - SimpleIrrepsConfig + - EnergyModel + - PerSpeciesRescale + - RescaleEnergyEtc + + + dataset: ase + dataset_file_name: structure.xyz + chemical_symbols: + - Mg + - Al + + # logging + wandb: false + # verbose: debug + + # training + n_train: 70 + n_val: 10 + batch_size: 5 + train_val_split: random + #shuffle: true + metrics_key: validation_loss + use_ema: true + ema_decay: 0.99 + ema_use_num_updates: true + max_epochs: 100 + learning_rate: 0.01 + # loss function + loss_coeffs: total_energy + +モデル学習、サンプリングの方法に関してはaenetと同様です。 + + +Allegroを利用したサンプリング +---------------------------------------------- + +Allegro のインストール +~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +``allegro`` の利用には、Allegroのインストールが必要です。 + +下記コマンドにてインストールします。 + +.. code-block:: bash + + $ git clone --depth 1 https://github.com/mir-group/allegro.git + $ cd allegro + $ pip3 install . + + +インプットファイルの準備 +~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +まず、input_allegro.tomlを準備し、Allegroの実行に必要なパラメータを設定します。 +下では、aenetのインプットから変更のある[sampling.solver]と[train]を抜粋しています。 + +.. code-block:: toml + + [sampling.solver] + type = 'allegro' + base_input_dir = './baseinput_allegro' + perturb = 0.0 + #run_scheme = 'subprocess' #'mpi_spawn_ready' + ignore_species = ["O"] + + [train] + type = 'allegro' + base_input_dir = './allegro_train_input' + exe_command = ['', 'nequip-train'] + ignore_species = ["O"] + vac_map = [] + restart = false + +また、Allegroのインプットファイルinput.yamlをallegro_train_input/trainディレクトリに作成します。 + +.. code-block:: yaml + + root: results/spinel + run_name: run + seed: 123 + dataset_seed: 456 + + # network + num_basis: 8 + BesselBasis_trainable: true + PolynomialCutoff_p: 6 + l_max: 1 + r_max: 8.0 + parity: o3_full + num_layers: 2 + num_features: 16 + + env_embed_multiplicity: 16 + embed_initial_edge: true + two_body_latent_mlp_latent_dimensions: [32, 64] + two_body_latent_mlp_nonlinearity: silu + latent_mlp_latent_dimensions: [64, 64] + latent_mlp_nonlinearity: silu + latent_mlp_initialization: uniform + latent_resnet: true + env_embed_mlp_latent_dimensions: [] + env_embed_mlp_nonlinearity: null + env_embed_mlp_initialization: uniform + edge_eng_mlp_latent_dimensions: [16] + edge_eng_mlp_nonlinearity: null + edge_eng_mlp_initialization: uniform + + model_builders: + - allegro.model.Allegro + - PerSpeciesRescale + - RescaleEnergyEtc + + + dataset: ase + dataset_file_name: structure.xyz + chemical_symbols: + - Mg + - Al + + # logging + wandb: false + # verbose: debug + + # training + n_train: 70 + n_val: 10 + batch_size: 5 + train_val_split: random + #shuffle: true + metrics_key: validation_loss + use_ema: true + ema_decay: 0.99 + ema_use_num_updates: true + max_epochs: 100 + learning_rate: 0.01 + # loss function + loss_coeffs: total_energy + +モデル学習、サンプリングの方法に関してはaenetと同様です。 + + +MLIP-3を利用したサンプリング +---------------------------------------------- + +MLIP-3 のインストール +~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +``mlip-3`` の利用には、 MLIP-3のインストールが必要です。 + +下記コマンドにてインストールします。 + +.. code-block:: bash + + $ git clone https://gitlab.com/ashapeev/mlip-3.git + $ cd mlip-3 + $ ./configure --no-mpi + $ make mlp + + +インプットファイルの準備 +~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +まず、input_mlip3.tomlを準備し、mlip-3の実行に必要なパラメータを設定します。 +下では、aenetのインプットから変更のある[sampling.solver]と[train]を抜粋しています。 + +.. code-block:: toml + + [sampling.solver] + type = 'mlip_3' + path= '~/github/mlip-3/bin/mlp' + base_input_dir = './baseinput' + perturb = 0.0 + run_scheme = 'subprocess' #'mpi_spawn_ready' + ignore_species = ["O"] + + [train] + type = 'mlip_3' + base_input_dir = './mlip_3_train_input' + exe_command = ['~/github/mlip-3/bin/mlp','~/github/mlip-3/bin/mlp'] + ignore_species = ["O"] + vac_map = [] + restart = false + +上記の内、[sampling.solver]のpathと[train]のexe_commandのリストでは +MLIP-3の実行ファイルmlpのパスを指定します。お使いの環境に合わせて変更してください。 + +また、MLIP-3のインプットファイルinput.almtpをmlip_3_train_input/trainディレクトリに作成します。 + +.. code-block:: none + + MTP + version = 1.1.0 + potential_name = MTP1m + species_count = 3 + potential_tag = + radial_basis_type = RBChebyshev + min_dist = 2.3 + max_dist = 5 + radial_basis_size = 8 + radial_funcs_count = 2 + alpha_moments_count = 8 + alpha_index_basic_count = 5 + alpha_index_basic = {{0, 0, 0, 0}, {0, 1, 0, 0}, {0, 0, 1, 0}, {0, 0, 0, 1}, {1, 0, 0, 0}} + alpha_index_times_count = 5 + alpha_index_times = {{0, 0, 1, 5}, {1, 1, 1, 6}, {2, 2, 1, 6}, {3, 3, 1, 6}, {0, 5, 1, 7}} + alpha_scalar_moments = 5 + alpha_moment_mapping = {0, 4, 5, 6, 7} + +モデル学習、サンプリングの方法に関してはaenetと同様です。 \ No newline at end of file