From 0bb942e78081dedc5fb091faffdc5a4aa88c2fe0 Mon Sep 17 00:00:00 2001 From: YutackPark Date: Thu, 18 Jul 2024 14:50:05 +0900 Subject: [PATCH 01/25] refactor, lint, feature: sevenn_preset, sevenn_patch_lammps, pretrained models as keyword (continue, calculator, sevenn get model) --- .pre-commit-config.yaml | 36 ++-- README.md | 4 +- logo_ascii | 20 --- pyproject.toml | 53 ++++++ setup.py | 7 +- sevenn/_const.py | 35 ++-- sevenn/main/sevenn_get_model.py | 17 +- sevenn/main/sevenn_patch_lammps.py | 34 ++++ sevenn/main/sevenn_preset.py | 30 ++++ sevenn/model_build.py | 45 ++--- sevenn/nn/edge_embedding.py | 31 +--- sevenn/nn/equivariant_gate.py | 18 -- sevenn/nn/force_output.py | 2 +- sevenn/nn/sequential.py | 2 +- .../pair_e3gnn}/comm_brick.cpp | 0 .../pair_e3gnn}/comm_brick.h | 0 .../pair_e3gnn}/pair_e3gnn.cpp | 0 .../pair_e3gnn}/pair_e3gnn.h | 0 .../pair_e3gnn}/pair_e3gnn_parallel.cpp | 0 .../pair_e3gnn}/pair_e3gnn_parallel.h | 0 .../pair_e3gnn/patch_lammps.sh | 17 +- sevenn/parse_input.py | 167 +++++++----------- sevenn/presets/base.yaml | 96 ++++++++++ sevenn/presets/fine_tune.yaml | 80 +++++++++ sevenn/presets/sevennet_0.yaml | 80 +++++++++ .../SevenNet_0__11July2024/README.md | 0 .../checkpoint_sevennet_0.pth | Bin .../SevenNet_0__11July2024/fine_tune.yaml | 0 .../parallel_model/deployed_parallel_0.pt | Bin .../parallel_model/deployed_parallel_1.pt | Bin .../parallel_model/deployed_parallel_2.pt | Bin .../parallel_model/deployed_parallel_3.pt | Bin .../parallel_model/deployed_parallel_4.pt | Bin .../SevenNet_0__11July2024/pre_train.yaml | 0 .../serial_model/deployed_serial.pt | Bin .../SevenNet_0__22May2024/README.md | 0 .../checkpoint_sevennet_0.pth | Bin .../SevenNet_0__22May2024/fine_tune.yaml | 0 .../parallel_model/deployed_parallel_0.pt | Bin .../parallel_model/deployed_parallel_1.pt | Bin .../parallel_model/deployed_parallel_2.pt | Bin .../parallel_model/deployed_parallel_3.pt | Bin .../parallel_model/deployed_parallel_4.pt | Bin .../serial_model/deployed_serial.pt | Bin .../240423_checkpoint_sevennet_0.pth | Bin .../archive/240423_SevenNet_0/README.txt | 0 sevenn/scripts/deploy.py | 3 +- sevenn/scripts/processing_continue.py | 16 +- sevenn/scripts/processing_dataset.py | 28 +-- sevenn/sevenn_logger.py | 11 +- sevenn/sevennet_calculator.py | 59 ++++--- sevenn/train/dataload.py | 19 +- sevenn/train/dataset.py | 17 +- sevenn/train/optim.py | 74 -------- sevenn/util.py | 59 ++++--- 55 files changed, 653 insertions(+), 407 deletions(-) delete mode 100644 logo_ascii create mode 100644 pyproject.toml create mode 100644 sevenn/main/sevenn_patch_lammps.py create mode 100644 sevenn/main/sevenn_preset.py rename {pair_e3gnn => sevenn/pair_e3gnn}/comm_brick.cpp (100%) rename {pair_e3gnn => sevenn/pair_e3gnn}/comm_brick.h (100%) rename {pair_e3gnn => sevenn/pair_e3gnn}/pair_e3gnn.cpp (100%) rename {pair_e3gnn => sevenn/pair_e3gnn}/pair_e3gnn.h (100%) rename {pair_e3gnn => sevenn/pair_e3gnn}/pair_e3gnn_parallel.cpp (100%) rename {pair_e3gnn => sevenn/pair_e3gnn}/pair_e3gnn_parallel.h (100%) rename patch_lammps.sh => sevenn/pair_e3gnn/patch_lammps.sh (82%) mode change 100644 => 100755 create mode 100644 sevenn/presets/base.yaml create mode 100644 sevenn/presets/fine_tune.yaml create mode 100644 sevenn/presets/sevennet_0.yaml rename {pretrained_potentials => sevenn/pretrained_potentials}/SevenNet_0__11July2024/README.md (100%) rename {pretrained_potentials => sevenn/pretrained_potentials}/SevenNet_0__11July2024/checkpoint_sevennet_0.pth (100%) rename {pretrained_potentials => sevenn/pretrained_potentials}/SevenNet_0__11July2024/fine_tune.yaml (100%) rename {pretrained_potentials => sevenn/pretrained_potentials}/SevenNet_0__11July2024/parallel_model/deployed_parallel_0.pt (100%) rename {pretrained_potentials => sevenn/pretrained_potentials}/SevenNet_0__11July2024/parallel_model/deployed_parallel_1.pt (100%) rename {pretrained_potentials => sevenn/pretrained_potentials}/SevenNet_0__11July2024/parallel_model/deployed_parallel_2.pt (100%) rename {pretrained_potentials => sevenn/pretrained_potentials}/SevenNet_0__11July2024/parallel_model/deployed_parallel_3.pt (100%) rename {pretrained_potentials => sevenn/pretrained_potentials}/SevenNet_0__11July2024/parallel_model/deployed_parallel_4.pt (100%) rename {pretrained_potentials => sevenn/pretrained_potentials}/SevenNet_0__11July2024/pre_train.yaml (100%) rename {pretrained_potentials => sevenn/pretrained_potentials}/SevenNet_0__11July2024/serial_model/deployed_serial.pt (100%) rename {pretrained_potentials => sevenn/pretrained_potentials}/SevenNet_0__22May2024/README.md (100%) rename {pretrained_potentials => sevenn/pretrained_potentials}/SevenNet_0__22May2024/checkpoint_sevennet_0.pth (100%) rename {pretrained_potentials => sevenn/pretrained_potentials}/SevenNet_0__22May2024/fine_tune.yaml (100%) rename {pretrained_potentials => sevenn/pretrained_potentials}/SevenNet_0__22May2024/parallel_model/deployed_parallel_0.pt (100%) rename {pretrained_potentials => sevenn/pretrained_potentials}/SevenNet_0__22May2024/parallel_model/deployed_parallel_1.pt (100%) rename {pretrained_potentials => sevenn/pretrained_potentials}/SevenNet_0__22May2024/parallel_model/deployed_parallel_2.pt (100%) rename {pretrained_potentials => sevenn/pretrained_potentials}/SevenNet_0__22May2024/parallel_model/deployed_parallel_3.pt (100%) rename {pretrained_potentials => sevenn/pretrained_potentials}/SevenNet_0__22May2024/parallel_model/deployed_parallel_4.pt (100%) rename {pretrained_potentials => sevenn/pretrained_potentials}/SevenNet_0__22May2024/serial_model/deployed_serial.pt (100%) rename {pretrained_potentials => sevenn/pretrained_potentials}/archive/240423_SevenNet_0/240423_checkpoint_sevennet_0.pth (100%) rename {pretrained_potentials => sevenn/pretrained_potentials}/archive/240423_SevenNet_0/README.txt (100%) diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 6287962..c6034f2 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -41,32 +41,44 @@ repos: rev: 5.13.2 hooks: - id: isort - exclude: 'pair_e3gnn/pair_e3gnn_parallel.h' - exclude: 'pair_e3gnn/comm_brick.h' - exclude: 'pair_e3gnn/comm_brick.cpp' + exclude: 'sevenn/pair_e3gnn/pair_e3gnn.cpp' + exclude: 'sevenn/pair_e3gnn/pair_e3gnn.h' + exclude: 'sevenn/pair_e3gnn/pair_e3gnn_parallel.cpp' + exclude: 'sevenn/pair_e3gnn/pair_e3gnn_parallel.h' + exclude: 'sevenn/pair_e3gnn/comm_brick.h' + exclude: 'sevenn/pair_e3gnn/comm_brick.cpp' - repo: https://github.com/psf/black rev: 23.12.1 hooks: - id: black + exclude: 'sevenn/pair_e3gnn/pair_e3gnn.cpp' + exclude: 'sevenn/pair_e3gnn/pair_e3gnn.h' + exclude: 'sevenn/pair_e3gnn/pair_e3gnn_parallel.cpp' + exclude: 'sevenn/pair_e3gnn/pair_e3gnn_parallel.h' + exclude: 'sevenn/pair_e3gnn/comm_brick.h' + exclude: 'sevenn/pair_e3gnn/comm_brick.cpp' args: ['--skip-string-normalization', '--line-length=79', '--preview'] - exclude: 'pair_e3gnn/pair_e3gnn_parallel.h' - exclude: 'pair_e3gnn/comm_brick.h' - exclude: 'pair_e3gnn/comm_brick.cpp' - repo: https://github.com/pycqa/flake8 rev: 6.1.0 hooks: - id: flake8 + exclude: 'sevenn/pair_e3gnn/pair_e3gnn.cpp' + exclude: 'sevenn/pair_e3gnn/pair_e3gnn.h' + exclude: 'sevenn/pair_e3gnn/pair_e3gnn_parallel.cpp' + exclude: 'sevenn/pair_e3gnn/pair_e3gnn_parallel.h' + exclude: 'sevenn/pair_e3gnn/comm_brick.h' + exclude: 'sevenn/pair_e3gnn/comm_brick.cpp' require_serial: true - exclude: 'pair_e3gnn/pair_e3gnn_parallel.h' - exclude: 'pair_e3gnn/comm_brick.h' - exclude: 'pair_e3gnn/comm_brick.cpp' - repo: https://github.com/pre-commit/mirrors-clang-format rev: 'v17.0.6' hooks: - id: clang-format - exclude: 'pair_e3gnn/pair_e3gnn_parallel.h' - exclude: 'pair_e3gnn/comm_brick.h' - exclude: 'pair_e3gnn/comm_brick.cpp' + exclude: 'sevenn/pair_e3gnn/pair_e3gnn.cpp' + exclude: 'sevenn/pair_e3gnn/pair_e3gnn.h' + exclude: 'sevenn/pair_e3gnn/pair_e3gnn_parallel.cpp' + exclude: 'sevenn/pair_e3gnn/pair_e3gnn_parallel.h' + exclude: 'sevenn/pair_e3gnn/comm_brick.h' + exclude: 'sevenn/pair_e3gnn/comm_brick.cpp' diff --git a/README.md b/README.md index f345232..f68c756 100644 --- a/README.md +++ b/README.md @@ -128,7 +128,7 @@ sevenn_get_model checkpoint_best.pt -p This will create multiple `deployed_parallel_*.pt` files. The number of deployed models equals the number of message-passing layers. These models can be used as lammps potential to run parallel MD simulations with GNN potential using multiple GPU cards. -## Installation for LAMMPS +## Installation for LAMMPS * PyTorch (same version as used for training) * LAMMPS version of 'stable_2Aug2023' [`LAMMPS`](https://github.com/lammps/lammps) @@ -241,7 +241,7 @@ Note that SevenNet-0 parallel model is located in `{PATH TO SEVENNET}/pretrained I recommend using variables to handle file paths for parallel models. ``` -pair_style e3gnn/parallel +pair_style e3gnn/parallel pair_coeff * * 5 ${pre}/deployed_parallel_0.pt ${pre}/deployed_parallel_1.pt ${pre}/deployed_parallel_2.pt ${pre}/deployed_parallel_3.pt ${pre}/deployed_parallel_4.pt {chemical species} ``` diff --git a/logo_ascii b/logo_ascii deleted file mode 100644 index 476d63d..0000000 --- a/logo_ascii +++ /dev/null @@ -1,20 +0,0 @@ - - **** - ******** . - *//////, .. . ,*. - ,,***. .. , ********. ./, - . . .. /////. ., . *///////// /////////. - .&@&/ . .(((((((.. / *//////*. ... *((((((((((. - @@@@@@@@@@* @@@@@@@@@@ @@@@@ *((@@@@@ ( %@@@@@@@@@@ .@@@@@@ ..@@@@. @@@@@@* .(@@@@@(((* - @@@@@. @@@@ @@@@@ . @@@@@ # %@@@@ @@@@@@@@ @@@@(, @@@@@@@@. @@@@@(*. - %@@@@@@@& @@@@@@@@@@ @@@@@ @@@@@ # ., .%@@@@@@@@@ @@@@@@@@@@ @@@@, @@@@@@@@@@ @@@@@ - ,(%@@@@@@@@@ @@@@@@@@@@ @@@@@ @@@@& (//////%@@@@@@@@@ @@@@ @@@@@@ @@@@ . @@@@@ @@@@@.@@@@@ - . @@@@@ @@@@ . . @@@@@@@@% . . ( .////,%@@@@ @@@@ @@@@@@@@@ @@@@@ @@@@@@@@@ - (@@@@@@@@@@@ @@@@@@@@@@**. @@@@@@* *. .%@@@@@@@@@@ @@@@ . @@@@@@@ @@@@@ .@@@@@@@ - @@@@@@@@@. @@@@@@@@@@///, @@@@. . / %@@@@@@@@@@ @@@@***, @@@@@ @@@@@ @@@@@ - . //////////*. / . .*******... . ,. - .&&&&&... ,//////*. ...////. / ,*/. . ,////, .,///// - &@@@@@@ ,(/((, * ,((((((. .***. - ,/@(, .. * ,((((* - , - . diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 0000000..217986b --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,53 @@ +[project] +name = "sevenn" +version = "0.9.2" +authors = [ + { name="Yutack Park", email="parkyutack@snu.ac.kr" }, + { name="Jaesun Kim" }, +] +description = "Scalable EquiVariance Enabled Neural Network" +readme = "README.md" +license = { file = "LICENSE" } +requires-python = ">=3.8" +classifiers = [ + "Programming Language :: Python :: 3", + "License :: OSI Approved :: GNU General Public License v3 (GPLv3)", + "Operating System :: POSIX :: Linux", +] + +[project.scripts] +sevenn = "sevenn.main.sevenn:main" +sevenn_get_model = "sevenn.main.sevenn_get_model:main" +sevenn_graph_build = "sevenn.main.sevenn_graph_build:main" +sevenn_inference = "sevenn.main.sevenn_inference:main" +sevenn_patch_lammps = "sevenn.main.sevenn_patch_lammps:main" +sevenn_preset = "sevenn.main.sevenn_preset:main" + +[project.urls] +Homepage = "https://github.com/MDIL-SNU/SevenNet" +Issues = "https://github.com/MDIL-SNU/SevenNet/issues" + +[build-system] +requires = [ + "setuptools>=61.0", + "ase", + "braceexpand", + "pyyaml", + "e3nn", + "tqdm", +] +build-backend = "setuptools.build_meta" + +[tool.setuptools.package-data] +sevenn = [ + "logo_ascii", + "pair_e3gnn/*.cpp", + "pair_e3gnn/*.h", + "presets/*.yaml", + "pretrained_potentials/SevenNet_0__11July2024/checkpoint_sevennet_0.pth", + "pretrained_potentials/SevenNet_0__22May2024/checkpoint_sevennet_0.pth" +] + +[tool.setuptools.packages.find] +include = ["sevenn*"] +exclude = ["tests*", "example_inputs*", ] diff --git a/setup.py b/setup.py index ce4086b..c87882d 100644 --- a/setup.py +++ b/setup.py @@ -1,5 +1,9 @@ -from setuptools import find_packages, setup +# from setuptools import find_packages, setup +import setuptools +setuptools.setup() + +""" setup( name='sevenn', version='0.9.2', @@ -25,3 +29,4 @@ ] }, ) +""" diff --git a/sevenn/_const.py b/sevenn/_const.py index e84bf95..4bbf3ba 100644 --- a/sevenn/_const.py +++ b/sevenn/_const.py @@ -1,3 +1,4 @@ +import os from enum import Enum from typing import Dict @@ -44,6 +45,14 @@ ACTIVATION_FOR_ODD = {'tanh': torch.tanh, 'abs': torch.abs} ACTIVATION_DICT = {'e': ACTIVATION_FOR_EVEN, 'o': ACTIVATION_FOR_ODD} +_prefix = os.path.abspath(f'{os.path.dirname(__file__)}/pretrained_potentials') +SEVENNET_0_11July2024 = ( + f'{_prefix}/SevenNet_0__11July2024/checkpoint_sevennet_0.pth' +) +SEVENNET_0_22May2024 = ( + f'{_prefix}/SevenNet_0__22May2024/checkpoint_sevennet_0.pth' +) + # to avoid torch script to compile torch_geometry.data AtomGraphDataType = Dict[str, torch.Tensor] @@ -90,12 +99,12 @@ def error_record_condition(x): KEY.NUM_CONVOLUTION: 3, KEY.ACTIVATION_SCARLAR: {'e': 'silu', 'o': 'tanh'}, KEY.ACTIVATION_GATE: {'e': 'silu', 'o': 'tanh'}, - #KEY.AVG_NUM_NEIGH: True, # deprecated - #KEY.TRAIN_AVG_NUM_NEIGH: False, # deprecated + # KEY.AVG_NUM_NEIGH: True, # deprecated + # KEY.TRAIN_AVG_NUM_NEIGH: False, # deprecated KEY.CONV_DENOMINATOR: 'avg_num_neigh', KEY.TRAIN_DENOMINTAOR: False, - KEY.TRAIN_SHIFT_SCALE: False, - #KEY.OPTIMIZE_BY_REDUCE: True, # deprecated, always True + KEY.TRAIN_SHIFT_SCALE: False, + # KEY.OPTIMIZE_BY_REDUCE: True, # deprecated, always True KEY.USE_BIAS_IN_LINEAR: False, KEY.READOUT_AS_FCN: False, # Applied af readout as fcn is True @@ -121,7 +130,10 @@ def error_record_condition(x): }, KEY.CUTOFF: float, KEY.NUM_CONVOLUTION: int, - KEY.CONV_DENOMINATOR: lambda x: isinstance(x, float) or x in ["avg_num_neigh", "sqrt_avg_num_neigh"], + KEY.CONV_DENOMINATOR: lambda x: isinstance(x, float) or x in [ + 'avg_num_neigh', + 'sqrt_avg_num_neigh', + ], KEY.CONVOLUTION_WEIGHT_NN_HIDDEN_NEURONS: list, KEY.TRAIN_SHIFT_SCALE: bool, KEY.TRAIN_DENOMINTAOR: bool, @@ -147,7 +159,6 @@ def model_defaults(config): return defaults - DEFAULT_DATA_CONFIG = { KEY.DTYPE: 'single', KEY.DATA_FORMAT: 'ase', @@ -158,9 +169,9 @@ def model_defaults(config): KEY.RATIO: 0.1, KEY.BATCH_SIZE: 6, KEY.PREPROCESS_NUM_CORES: 1, - #KEY.USE_SPECIES_WISE_SHIFT_SCALE: False, - KEY.SHIFT: "per_atom_energy_mean", - KEY.SCALE: "force_rms", + # KEY.USE_SPECIES_WISE_SHIFT_SCALE: False, + KEY.SHIFT: 'per_atom_energy_mean', + KEY.SCALE: 'force_rms', KEY.DATA_SHUFFLE: True, } @@ -174,7 +185,7 @@ def model_defaults(config): KEY.RATIO: float, KEY.BATCH_SIZE: int, KEY.PREPROCESS_NUM_CORES: int, - #KEY.USE_SPECIES_WISE_SHIFT_SCALE: bool, + # KEY.USE_SPECIES_WISE_SHIFT_SCALE: bool, KEY.SHIFT: lambda x: type(x) in [float, list] or x in IMPLEMENTED_SHIFT, KEY.SCALE: lambda x: type(x) in [float, list] or x in IMPLEMENTED_SCALE, KEY.DATA_SHUFFLE: bool, @@ -194,7 +205,9 @@ def data_defaults(config): KEY.EPOCH: 300, KEY.LOSS: 'mse', KEY.OPTIMIZER: 'adam', + KEY.OPTIM_PARAM: {}, KEY.SCHEDULER: 'exponentiallr', + KEY.SCHEDULER_PARAM: {}, KEY.FORCE_WEIGHT: 0.1, KEY.STRESS_WEIGHT: 1e-6, # SIMPLE-NN default KEY.PER_EPOCH: 5, @@ -228,7 +241,7 @@ def data_defaults(config): KEY.STRESS_WEIGHT: float, KEY.USE_TESTSET: None, # Not used KEY.NUM_WORKERS: None, # Not used - KEY.PER_EPOCH: int, + KEY.PER_EPOCH: int, KEY.CONTINUE: { KEY.CHECKPOINT: str, KEY.RESET_OPTIMIZER: bool, diff --git a/sevenn/main/sevenn_get_model.py b/sevenn/main/sevenn_get_model.py index f3c4947..50d097b 100644 --- a/sevenn/main/sevenn_get_model.py +++ b/sevenn/main/sevenn_get_model.py @@ -1,6 +1,5 @@ import argparse - -import torch +import os import sevenn._const as _const import sevenn.util @@ -10,7 +9,10 @@ f'sevenn version={_const.SEVENN_VERSION}, sevenn_get_model.' + ' Deploy model for LAMMPS from the checkpoint' ) -checkpoint_help = 'checkpoint path' +checkpoint_help = ( + 'path to the checkpoint | SevenNet-0 | 7net-0 |' + ' {SevenNet-0|7net-0}_{11July2024|22May2024}' +) output_name_help = 'filename prefix' get_parallel_help = 'deploy parallel model' @@ -18,14 +20,19 @@ def main(args=None): checkpoint, output_prefix, get_parallel = cmd_parse_get_model(args) get_serial = not get_parallel - cp_file = torch.load(checkpoint, map_location=torch.device('cpu')) if output_prefix is None: output_prefix = ( 'deployed_parallel' if not get_serial else 'deployed_serial' ) - model, config = sevenn.util.model_from_checkpoint(checkpoint) + checkpoint_path = None + if os.path.isfile(checkpoint): + checkpoint_path = checkpoint + else: + checkpoint_path = sevenn.util.pretrained_name_to_path(checkpoint) + + model, config = sevenn.util.model_from_checkpoint(checkpoint_path) stct_dct = model.state_dict() if get_serial: diff --git a/sevenn/main/sevenn_patch_lammps.py b/sevenn/main/sevenn_patch_lammps.py new file mode 100644 index 0000000..8746a0d --- /dev/null +++ b/sevenn/main/sevenn_patch_lammps.py @@ -0,0 +1,34 @@ +import os +import argparse +import subprocess + +# python wrapper of patch_lammps.sh script + +# importlib.resources is correct way to do these things +# but it changes so frequently to use +pair_e3gnn_dir = os.path.abspath(f'{os.path.dirname(__file__)}/../pair_e3gnn') + +from sevenn._const import SEVENN_VERSION + +description = ( + f'sevenn version={SEVENN_VERSION}, patch LAMMPS for pair_e3gnn styles' +) + + +def main(args=None): + lammps_dir = cmd_parse_main(args) + script = f"{pair_e3gnn_dir}/patch_lammps.sh" + cmd = f"{script} {lammps_dir}" + res = subprocess.run(cmd.split()) + return res.returncode # is it meaningless? + + +def cmd_parse_main(args=None): + ag = argparse.ArgumentParser(description=description) + ag.add_argument('lammps_dir', help="Path to LAMMPS source", type=str) + args = ag.parse_args() + return args.lammps_dir + + +if __name__ == '__main__': + main() diff --git a/sevenn/main/sevenn_preset.py b/sevenn/main/sevenn_preset.py new file mode 100644 index 0000000..78a7dd3 --- /dev/null +++ b/sevenn/main/sevenn_preset.py @@ -0,0 +1,30 @@ +import os +import argparse + +import sevenn._const as _const + +description_preset = ( + f'sevenn version={_const.SEVENN_VERSION}, sevenn_preset.' + + ' copy paste preset training yaml file to current directory' + + ' ex) sevennet_preset fine_tune > my_input.yaml' +) + +preset_help = "Name of preset" + + +def main(args=None): + preset = cmd_parse_preset(args) + prefix = os.path.abspath(f'{os.path.dirname(__file__)}/../presets') + + with open(f"{prefix}/{preset}.yaml", "r") as f: + print(f.read()) + + +def cmd_parse_preset(args=None): + ag = argparse.ArgumentParser(description=description_preset) + ag.add_argument( + 'preset', choices=['fine_tune', 'sevennet-0', 'base'], + help = preset_help + ) + args = ag.parse_args() + return args.preset diff --git a/sevenn/model_build.py b/sevenn/model_build.py index 75003eb..badfdf5 100644 --- a/sevenn/model_build.py +++ b/sevenn/model_build.py @@ -5,8 +5,8 @@ import sevenn._const as _const import sevenn._keys as KEY -from sevenn.nn.convolution import IrrepsConvolution import sevenn.util as util +from sevenn.nn.convolution import IrrepsConvolution from sevenn.nn.edge_embedding import ( BesselBasis, EdgeEmbedding, @@ -54,12 +54,12 @@ def init_self_connection(config): def init_radial_basis(config): radial_basis_dct = config[KEY.RADIAL_BASIS] - param = {"cutoff_length": config[KEY.CUTOFF]} + param = {'cutoff_length': config[KEY.CUTOFF]} param.update(radial_basis_dct) del param[KEY.RADIAL_BASIS_NAME] if radial_basis_dct[KEY.RADIAL_BASIS_NAME] == 'bessel': - basis_function = BesselBasis(**param) + basis_function = BesselBasis(**param) return basis_function, basis_function.num_basis raise RuntimeError('something went very wrong...') @@ -67,7 +67,7 @@ def init_radial_basis(config): def init_cutoff_function(config): cutoff_function_dct = config[KEY.CUTOFF_FUNCTION] - param = {"cutoff_length": config[KEY.CUTOFF]} + param = {'cutoff_length': config[KEY.CUTOFF]} param.update(cutoff_function_dct) del param[KEY.CUTOFF_FUNCTION_NAME] @@ -93,16 +93,8 @@ def build_E3_equivariant_model(config: dict, parallel=False): feature_multiplicity = config[KEY.NODE_FEATURE_MULTIPLICITY] lmax = config[KEY.LMAX] - lmax_edge = ( - config[KEY.LMAX_EDGE] - if config[KEY.LMAX_EDGE] >= 0 - else lmax - ) - lmax_node = ( - config[KEY.LMAX_NODE] - if config[KEY.LMAX_NODE] >= 0 - else lmax - ) + lmax_edge = config[KEY.LMAX_EDGE] if config[KEY.LMAX_EDGE] >= 0 else lmax + lmax_node = config[KEY.LMAX_NODE] if config[KEY.LMAX_NODE] >= 0 else lmax num_convolution_layer = config[KEY.NUM_CONVOLUTION] @@ -154,7 +146,9 @@ def build_E3_equivariant_model(config: dict, parallel=False): basis_module=radial_basis_module, cutoff_module=cutoff_function_module, # operate on r/||r|| - spherical_module=SphericalEncoding(lmax_edge, -1 if is_parity else 1, normalize=_normalize_sph), + spherical_module=SphericalEncoding( + lmax_edge, -1 if is_parity else 1, normalize=_normalize_sph + ), ) if not parallel: layers.update({ @@ -214,13 +208,12 @@ def build_E3_equivariant_model(config: dict, parallel=False): weight_nn_layers = [radial_basis_num] + weight_nn_hidden for i in range(num_convolution_layer): - # here, we can infer irreps of x after interaction from lmax and f0_irreps interaction_block = {} - parity_mode = "full" + parity_mode = 'full' if i == num_convolution_layer - 1: lmax_node = 0 - parity_mode = "even" + parity_mode = 'even' # raw irreps out after message(convolution) function tp_irreps_out = util.infer_irreps_out( @@ -243,19 +236,9 @@ def build_E3_equivariant_model(config: dict, parallel=False): else irreps_manual[i + 1] ) - # output irreps of linear 2 & self_connection is determined by Gate - # Gate require extra scalars(or weight) for l>0 features in nequip, - # they make it from linear2. (and self_connection have to fit its dimension) - # Here, initialize gate first and put it later gate_layer = EquivariantGate(true_irreps_out, act_scalar, act_gate) irreps_for_gate_in = gate_layer.get_gate_irreps_in() - # from here, data flow split into self connection part and convolution part - # self connection part is represented as Intro, Outro pair - - # note that this layer does not overwrite x, it calculates tp of in & operand - # and save its results in somewhere to concatenate to new_x at Outro - interaction_block[f'{i}_self_connection_intro'] = sc_intro( irreps_x=irreps_x, irreps_operand=irreps_node_attr, @@ -290,7 +273,11 @@ def build_E3_equivariant_model(config: dict, parallel=False): basis_module=radial_basis_module, cutoff_module=cutoff_function_module, # operate on r/||r|| - spherical_module=SphericalEncoding(lmax_edge, -1 if is_parity else 1, normalize=_normalize_sph), + spherical_module=SphericalEncoding( + lmax_edge, + -1 if is_parity else 1, + normalize=_normalize_sph, + ), ) }) ####################################################### diff --git a/sevenn/nn/edge_embedding.py b/sevenn/nn/edge_embedding.py index 4581558..87f99f4 100644 --- a/sevenn/nn/edge_embedding.py +++ b/sevenn/nn/edge_embedding.py @@ -91,9 +91,9 @@ def __init__( super().__init__() self.num_basis = bessel_basis_num self.prefactor = 2.0 / cutoff_length - self.coeffs = torch.FloatTensor( - [n * math.pi / cutoff_length for n in range(1, bessel_basis_num + 1)] - ) + self.coeffs = torch.FloatTensor([ + n * math.pi / cutoff_length for n in range(1, bessel_basis_num + 1) + ]) if trainable_coeff: self.coeffs = nn.Parameter(self.coeffs) @@ -147,7 +147,6 @@ def __init__( assert self.r_on < self.r_cut def forward(self, r: torch.Tensor) -> torch.Tensor: - # r > r_cut switch is not necessary since edges are already based on cutoff r_sq = r * r r_on_sq = self.r_on * self.r_on r_cut_sq = self.r_cut * self.r_cut @@ -162,26 +161,12 @@ def forward(self, r: torch.Tensor) -> torch.Tensor: @compile_mode('script') class SphericalEncoding(nn.Module): - """ - Calculate spherical harmonics from 0 to lmax - taking displacement vector (EDGE_VEC) as input. - - lmax: maximum angular momentum quantum number used in model - normalization : {'integral', 'component', 'norm'} - normalization of the output tensors - Valid options: - * *component*: :math:`\|Y^l(x)\|^2 = 2l+1, x \in S^2` - * *norm*: :math:`\|Y^l(x)\| = 1, x \in S^2`, ``component / sqrt(2l+1)`` - * *integral*: :math:`\int_{S^2} Y^l_m(x)^2 dx = 1`, ``component / sqrt(4pi)`` - - Returns - ------- - `torch.Tensor` - a tensor of shape ``(..., (lmax+1)^2)`` - """ - def __init__( - self, lmax: int, parity: int = -1, normalization: str = 'component', normalize = True, + self, + lmax: int, + parity: int = -1, + normalization: str = 'component', + normalize=True, ): super().__init__() self.lmax = lmax diff --git a/sevenn/nn/equivariant_gate.py b/sevenn/nn/equivariant_gate.py index d49f5e3..8422405 100644 --- a/sevenn/nn/equivariant_gate.py +++ b/sevenn/nn/equivariant_gate.py @@ -11,21 +11,6 @@ @compile_mode('script') class EquivariantGate(nn.Module): - """ - wrapper of e3nn.nn Gate (equivariant-nonlinear gate for irreps) - required irreps_in for Gate forward is computed after instantiation - of this class - see - https://docs.e3nn.org/en/stable/api/nn/nn_gate.html - in nequip, result of convolution and self-interaction linear2 - is directly used for irreps_gates - - Usage in NequIP - irreps_x: Representation of lmax, fixed multiplicity applied irreps - act_scalar/gate_dict: dictionary of parity and activation function - depends on parity, the activation function is regulated (odd or even function) - """ - def __init__( self, irreps_x: Irreps, @@ -53,9 +38,6 @@ def __init__( irreps_scalars = Irreps(irreps_scalars_elem) irreps_gated = Irreps(irreps_gated_elem) - # determine whether this scalar is odd or even - # in gates, whether scalar is odd or even is not important but it should be - # found in irreps_x to operate irreps_gates_parity = 1 if '0e' in irreps_scalars else -1 irreps_gates = Irreps( [(mul, (0, irreps_gates_parity)) for mul, _ in irreps_gated] diff --git a/sevenn/nn/force_output.py b/sevenn/nn/force_output.py index b2a0094..4f1ff3d 100644 --- a/sevenn/nn/force_output.py +++ b/sevenn/nn/force_output.py @@ -76,7 +76,7 @@ def forward(self, data: AtomGraphDataType) -> AtomGraphDataType: create_graph=self.training, )[0] - # without this 'if', type(grad) is 'Optional[Tensor]' which result in error + # For torchscript if grad is not None: data[self.key_force] = torch.neg(grad) return data diff --git a/sevenn/nn/sequential.py b/sevenn/nn/sequential.py index fc28135..ae8f256 100644 --- a/sevenn/nn/sequential.py +++ b/sevenn/nn/sequential.py @@ -24,7 +24,7 @@ def __init__( cutoff: float = 0.0, type_map: Dict[int, int] = {-1: -1}, ): - if type(modules) != OrderedDict: + if not isinstance(modules, OrderedDict): modules = OrderedDict(modules) self.cutoff = cutoff self.type_map = type_map diff --git a/pair_e3gnn/comm_brick.cpp b/sevenn/pair_e3gnn/comm_brick.cpp similarity index 100% rename from pair_e3gnn/comm_brick.cpp rename to sevenn/pair_e3gnn/comm_brick.cpp diff --git a/pair_e3gnn/comm_brick.h b/sevenn/pair_e3gnn/comm_brick.h similarity index 100% rename from pair_e3gnn/comm_brick.h rename to sevenn/pair_e3gnn/comm_brick.h diff --git a/pair_e3gnn/pair_e3gnn.cpp b/sevenn/pair_e3gnn/pair_e3gnn.cpp similarity index 100% rename from pair_e3gnn/pair_e3gnn.cpp rename to sevenn/pair_e3gnn/pair_e3gnn.cpp diff --git a/pair_e3gnn/pair_e3gnn.h b/sevenn/pair_e3gnn/pair_e3gnn.h similarity index 100% rename from pair_e3gnn/pair_e3gnn.h rename to sevenn/pair_e3gnn/pair_e3gnn.h diff --git a/pair_e3gnn/pair_e3gnn_parallel.cpp b/sevenn/pair_e3gnn/pair_e3gnn_parallel.cpp similarity index 100% rename from pair_e3gnn/pair_e3gnn_parallel.cpp rename to sevenn/pair_e3gnn/pair_e3gnn_parallel.cpp diff --git a/pair_e3gnn/pair_e3gnn_parallel.h b/sevenn/pair_e3gnn/pair_e3gnn_parallel.h similarity index 100% rename from pair_e3gnn/pair_e3gnn_parallel.h rename to sevenn/pair_e3gnn/pair_e3gnn_parallel.h diff --git a/patch_lammps.sh b/sevenn/pair_e3gnn/patch_lammps.sh old mode 100644 new mode 100755 similarity index 82% rename from patch_lammps.sh rename to sevenn/pair_e3gnn/patch_lammps.sh index 582c112..68a060d --- a/patch_lammps.sh +++ b/sevenn/pair_e3gnn/patch_lammps.sh @@ -1,6 +1,7 @@ #!/bin/bash lammps_root=$1 +SCRIPT_DIR=$(dirname "${BASH_SOURCE[0]}") echo "Usage: sh patch_lammps.sh {lammps_root}" # Check if the lammps_root directory exists @@ -10,14 +11,14 @@ if [ ! -d "$lammps_root" ]; then fi # Check if the given directory is the root of LAMMPS source -if [ ! -d "$lammps_root/cmake" ]; then - echo "Given $lammps_root is not root of LAMMPS source" +if [ ! -d "$lammps_root/cmake" ] && [ ! -d "$lammps_root/potentials" ]; then + echo "Given $lammps_root is not a root of LAMMPS source" exit 1 fi # Check if the script is being run from the root of SevenNet -if [ ! -d "./pair_e3gnn" ]; then - echo "Please run this script in the root of SevenNet" +if [ ! -f "${SCRIPT_DIR}/pair_e3gnn.cpp" ]; then + echo "Script executed in a wrong directory" exit 1 fi @@ -47,7 +48,9 @@ cp $lammps_root/src/comm_brick.cpp $backup_dir/ cp $lammps_root/src/comm_brick.h $backup_dir/ # 2. Copy everything inside pair_e3gnn to LAMMPS source -cp ./pair_e3gnn/* $lammps_root/src/ +# script is located in pair_e3gnn folder +cp $SCRIPT_DIR/*.cpp $lammps_root/src/ +cp $SCRIPT_DIR/*.h $lammps_root/src/ # 3. Copy cmake/CMakeLists.txt from original source as backup cp $lammps_root/cmake/CMakeLists.txt $backup_dir/CMakeLists.txt @@ -65,11 +68,11 @@ EOF2 # Check if the command is found and its value is true cuda_support=$(ompi_info --parsable --all | grep mpi_built_with_cuda_support:value) if [[ -z "$cuda_support" ]]; then - echo "OpenMPI not found, parallel performance could be low" + echo "OpenMPI not found, parallel performance is not optimal" elif [[ "$cuda_support" == *"true" ]]; then echo "OpenMPI is CUDA aware" else - echo "This OpenMPI is not CUDA aware, parallel performance could be low" + echo "This system's OpenMPI is not 'CUDA aware', parallel performance is not optimal" fi # ?. Print changes and backup file locations diff --git a/sevenn/parse_input.py b/sevenn/parse_input.py index f762a18..1759b54 100644 --- a/sevenn/parse_input.py +++ b/sevenn/parse_input.py @@ -1,26 +1,21 @@ import glob import os import warnings -from typing import Callable, Any +from typing import Any, Callable import torch import yaml import sevenn._const as _const import sevenn._keys as KEY -from sevenn.train.optim import ( - loss_dict, - loss_param_name_type_dict, - optim_dict, - optim_param_name_type_dict, - scheduler_dict, - scheduler_param_name_type_dict, -) -from sevenn.util import chemical_species_preprocess +from sevenn.util import chemical_species_preprocess, pretrained_name_to_path def config_initialize( - key: str, config: dict, default: Any, conditions, + key: str, + config: dict, + default: Any, + conditions, ): # default value exist & no user input -> return default if key not in config.keys(): @@ -35,41 +30,46 @@ def config_initialize( if type(default) is dict and isinstance(condition, dict): for i_key, val in default.items(): - user_input[i_key] =\ - config_initialize(i_key, user_input, val, condition) + user_input[i_key] = config_initialize( + i_key, user_input, val, condition + ) return user_input - elif isinstance(condition, type): + elif isinstance(condition, type): if isinstance(user_input, condition): return user_input else: try: return condition(user_input) # try type casting except ValueError: - raise ValueError(f"Expect '{user_input}' for '{key}' is {condition}") + raise ValueError( + f"Expect '{user_input}' for '{key}' is {condition}" + ) elif isinstance(condition, Callable) and condition(user_input): return user_input else: - raise ValueError(f"Given input '{user_input}' for '{key}' is not valid") + raise ValueError( + f"Given input '{user_input}' for '{key}' is not valid" + ) def init_model_config(config: dict): - defaults = _const.model_defaults(config) + # defaults = _const.model_defaults(config) model_meta = {} # init complicated ones if KEY.CHEMICAL_SPECIES not in config.keys(): raise ValueError('required key chemical_species not exist') input_chem = config[KEY.CHEMICAL_SPECIES] - if type(input_chem) == str and input_chem.lower() == 'auto': + if isinstance(input_chem, str) and input_chem.lower() == 'auto': model_meta[KEY.CHEMICAL_SPECIES] = 'auto' model_meta[KEY.NUM_SPECIES] = 'auto' model_meta[KEY.TYPE_MAP] = 'auto' else: - if type(input_chem) == list and all( - type(x) == str for x in input_chem + if isinstance(input_chem, list) and all( + isinstance(x, str) for x in input_chem ): pass - elif type(input_chem) == str: + elif isinstance(input_chem, str): input_chem = ( input_chem.replace('-', ',').replace(' ', ',').split(',') ) @@ -78,19 +78,20 @@ def init_model_config(config: dict): raise ValueError(f'given {KEY.CHEMICAL_SPECIES} input is strange') model_meta.update(chemical_species_preprocess(input_chem)) - ######## deprecation warnings ######### + # deprecation warnings if KEY.AVG_NUM_NEIGH in config: warnings.warn( - "key 'avg_num_neigh' is deprecated. Please use 'conv_denominator'. " - "We use the default, the average number of neighbors in the dataset, " - "if not provided.", + "key 'avg_num_neigh' is deprecated. Please use 'conv_denominator'." + " We use the default, the average number of neighbors in the" + " dataset, if not provided.", UserWarning, ) config.pop(KEY.AVG_NUM_NEIGH) if KEY.TRAIN_AVG_NUM_NEIGH in config: warnings.warn( - "key 'train_avg_num_neigh' is deprecated. Please use 'train_denominator'. " - "We overwrite train_denominator as given train_avg_num_neigh", + "key 'train_avg_num_neigh' is deprecated. Please use" + " 'train_denominator'. We overwrite train_denominator as given" + " train_avg_num_neigh", UserWarning, ) config[KEY.TRAIN_DENOMINTAOR] = config[KEY.TRAIN_AVG_NUM_NEIGH] @@ -101,29 +102,31 @@ def init_model_config(config: dict): UserWarning, ) config.pop(KEY.OPTIMIZE_BY_REDUCE) - ######## deprecation warnings ######### # init simpler ones for key, default in _const.DEFAULT_E3_EQUIVARIANT_MODEL_CONFIG.items(): - model_meta[key] =\ - config_initialize(key, config, default, _const.MODEL_CONFIG_CONDITION) + model_meta[key] = config_initialize( + key, config, default, _const.MODEL_CONFIG_CONDITION + ) unknown_keys = [ key for key in config.keys() if key not in model_meta.keys() ] if len(unknown_keys) != 0: - warnings.warn(f"Unexpected model keys: {unknown_keys} will be ignored", UserWarning) + warnings.warn( + f'Unexpected model keys: {unknown_keys} will be ignored', + UserWarning, + ) return model_meta def init_train_config(config: dict): train_meta = {} - defaults = _const.train_defaults(config) + # defaults = _const.train_defaults(config) try: device_input = config[KEY.DEVICE] - # TODO: device input sanity? train_meta[KEY.DEVICE] = torch.device(device_input) except KeyError: train_meta[KEY.DEVICE] = ( @@ -132,89 +135,37 @@ def init_train_config(config: dict): else torch.device('cpu') ) - name_dicts = [optim_dict, scheduler_dict, loss_dict] - name_keys = [KEY.OPTIMIZER, KEY.SCHEDULER, KEY.LOSS] - for idx, type_key in enumerate(name_keys): - if type_key not in config.keys(): - train_meta[type_key] = defaults[type_key] - continue - user_input = config[type_key].lower() - available_keys = name_dicts[idx].keys() - if type(user_input) is not str: - raise ValueError(f'{type_key} should be type: string.') - if user_input not in available_keys: - ava_key_to_str = '' - for i, key in enumerate(available_keys): - if i == 0: - ava_key_to_str += f'{key}' - else: - ava_key_to_str += f', {key}' - raise ValueError(f'{type_key} should be one of {ava_key_to_str}') - train_meta[type_key] = user_input - - param_type_dicts = [ - optim_param_name_type_dict, - scheduler_param_name_type_dict, - loss_param_name_type_dict, - ] - for idx, param_key in enumerate( - [KEY.OPTIM_PARAM, KEY.SCHEDULER_PARAM, KEY.LOSS_PARAM] - ): - if param_key not in config.keys(): - continue - user_input = config[param_key] - type_value = train_meta[name_keys[idx]] - universal_keys = list(param_type_dicts[idx]['universial'].keys()) - available_keys = list(param_type_dicts[idx][type_value].keys()) - available_keys.extend(universal_keys) - for key, value in user_input.items(): - # key = key.lower() # case sensitive detect of param name - if key not in available_keys: - ava_key_to_str = '' - for i, k in enumerate(available_keys): - if i == 0: - ava_key_to_str += f'{k}' - else: - ava_key_to_str += f', {k}' - raise ValueError( - f'{param_key}: {key} should be one of {available_keys}' - ) - if key in universal_keys: - type_ = param_type_dicts[idx]['universial'][key] - else: - type_ = param_type_dicts[idx][type_value][key] - - if type(value) is not type_: - raise ValueError(f'{param_key}: {key} should be type: {type_}') - train_meta[param_key] = user_input + # init simpler ones + for key, default in _const.DEFAULT_TRAINING_CONFIG.items(): + train_meta[key] = config_initialize( + key, config, default, _const.TRAINING_CONFIG_CONDITION + ) if KEY.CONTINUE in config.keys(): cnt_dct = config[KEY.CONTINUE] if KEY.CHECKPOINT not in cnt_dct.keys(): raise ValueError('no checkpoint is given in continue') checkpoint = cnt_dct[KEY.CHECKPOINT] - if type(checkpoint) != str or os.path.isfile(checkpoint) is False: - raise ValueError(f'Checkpoint file:{checkpoint} is not found') - train_meta[KEY.CONTINUE] = {} - train_meta[KEY.CONTINUE][KEY.CHECKPOINT] = checkpoint - - # init simpler ones - for key, default in _const.DEFAULT_TRAINING_CONFIG.items(): - train_meta[key] =\ - config_initialize(key, config, default, _const.TRAINING_CONFIG_CONDITION) + if os.path.isfile(checkpoint): + checkpoint_file = checkpoint + else: + checkpoint_file = pretrained_name_to_path(checkpoint) + train_meta[KEY.CONTINUE].update({KEY.CHECKPOINT: checkpoint_file}) unknown_keys = [ key for key in config.keys() if key not in train_meta.keys() ] if len(unknown_keys) != 0: - warnings.warn(f"Unexpected train keys: {unknown_keys} will be ignored", UserWarning) - + warnings.warn( + f'Unexpected train keys: {unknown_keys} will be ignored', + UserWarning, + ) return train_meta def init_data_config(config: dict): data_meta = {} - defaults = _const.data_defaults(config) + # defaults = _const.data_defaults(config) if KEY.LOAD_DATASET not in config.keys(): raise ValueError('load_dataset_path is not given') @@ -231,21 +182,27 @@ def init_data_config(config: dict): for i in inp: extended.extend(glob.glob(i)) if len(extended) == 0: - raise ValueError(f'Cannot find {inp} for {load_data_key}' - + ' or path is not given') + raise ValueError( + f'Cannot find {inp} for {load_data_key}' + + ' or path is not given' + ) data_meta[load_data_key] = extended else: data_meta[load_data_key] = False for key, default in _const.DEFAULT_DATA_CONFIG.items(): - data_meta[key] =\ - config_initialize(key, config, default, _const.DATA_CONFIG_CONDITION) + data_meta[key] = config_initialize( + key, config, default, _const.DATA_CONFIG_CONDITION + ) unknown_keys = [ key for key in config.keys() if key not in data_meta.keys() ] if len(unknown_keys) != 0: - warnings.warn(f"Unexpected data keys: {unknown_keys} will be ignored", UserWarning) + warnings.warn( + f'Unexpected data keys: {unknown_keys} will be ignored', + UserWarning, + ) return data_meta diff --git a/sevenn/presets/base.yaml b/sevenn/presets/base.yaml new file mode 100644 index 0000000..d85d2f8 --- /dev/null +++ b/sevenn/presets/base.yaml @@ -0,0 +1,96 @@ +# Example input.yaml for training SevenNet. +# '*' signifies default. You can check log.sevenn. + +model: + chemical_species: 'Auto' # Chemical elements present in the dataset, guess them from load_dataset data if 'auto' + cutoff: 5.0 # Cutoff radius in Angstroms. If two atoms are within the cutoff, they are connected. + channel: 32 # The multiplicity(channel) of node features. + lmax: 2 # Maximum order of irreducible representations (rotation order). + num_convolution_layer: 3 # The number of message passing layers. + + #irreps_manual: # Manually set irreps of the model in each layer + #- "128x0e" + #- "128x0e+64x1e+32x2e" + #- "128x0e+64x1e+32x2e" + #- "128x0e+64x1e+32x2e" + #- "128x0e+64x1e+32x2e" + #- "128x0e" + + weight_nn_hidden_neurons: [64, 64] # Hidden neurons in convolution weight neural network + radial_basis: # Function and its parameters to encode radial distance + radial_basis_name: 'bessel' # Only 'bessel' is currently supported + bessel_basis_num: 8 + cutoff_function: # Envelop function, multiplied to radial_basis functions to init edge featrues + cutoff_function_name: 'poly_cut' # {'poly_cut' and 'poly_cut_p_value'} or {'XPLOR' and 'cutoff_on'} + poly_cut_p_value: 6 + + act_gate: {'e': 'silu', 'o': 'tanh'} # Equivalent to 'nonlinearity_gates' in nequip + act_scalar: {'e': 'silu', 'o': 'tanh'} # Equivalent to 'nonlinearity_scalars' in nequip + + is_parity: False # Pairy True (E(3) group) or False (to SE(3) group) + + self_connection_type: 'nequip' # Default is 'nequip'. 'linear' is used for SevenNet-0. + + conv_denominator: "avg_num_neigh" # Valid options are "avg_num_neigh*", "sqrt_avg_num_neigh", or float + train_denominator: False # Enable training for denominator in convolution layer + train_shift_scale: False # Enable training for shift & scale in output layer + +train: + random_seed: 1 + is_train_stress: True # Includes stress in the loss function + epoch: 200 # Ends training after this number of epochs + + #loss: 'Huber' # Default is 'mse' (mean squared error) + #loss_param: + #delta: 0.01 + + # Each optimizer and scheduler have different available parameters. + # You can refer to sevenn/train/optim.py for supporting optimizer & schedulers + optimizer: 'adam' # Options available are 'sgd', 'adagrad', 'adam', 'adamw', 'radam' + optim_param: + lr: 0.005 + scheduler: 'exponentiallr' # 'steplr', 'multisteplr', 'exponentiallr', 'cosineannealinglr', 'reducelronplateau', 'linearlr' + scheduler_param: + gamma: 0.99 + + force_loss_weight: 0.1 # Coefficient for force loss + stress_loss_weight: 1e-06 # Coefficient for stress loss (to kbar unit) + + per_epoch: 10 # Generate checkpoints every this epoch + + # ['target y', 'metric'] + # Target y: TotalEnergy, Energy, Force, Stress, Stress_GPa, TotalLoss + # Metric : RMSE, MAE, or Loss + error_record: + - ['Energy', 'RMSE'] + - ['Force', 'RMSE'] + - ['Stress', 'RMSE'] + - ['TotalLoss', 'None'] + + # Continue training model from given checkpoint, or pre-trained model checkpoint for fine-tuning + #continue: + #checkpoint: 'checkpoint_best.pth' # Checkpoint of pre-trained model or a model want to continue training. + #reset_optimizer: False # Set True for fine-tuning + #reset_scheduler: False # Set True for fine-tuning + #use_statistic_values_of_checkpoint: False # Set True to use shift, scale, and avg_num_neigh from checkpoint or not + +data: + batch_size: 4 # Per GPU batch size. + data_divide_ratio: 0.1 # Split dataset into training and validation sets by this ratio + + shift: 'per_atom_energy_mean' # One of 'per_atom_energy_mean*', 'elemwise_reference_energies', float + scale: 'force_rms' # One of 'force_rms*', 'per_atom_energy_std', 'elemwise_force_rms', float + + # ase.io.read readable data files or structure_list or .sevenn_data files can be used as dataset. + # .sevenn_data is preprocessed data set has edges connected (can be obtained by using sevenn_graph_build or by save_** options below) + data_format: 'ase' # Default is 'ase' + data_format_args: # Paramaters, will be passed to ase.io.read see: https://wiki.fysik.dtu.dk/ase/ase/io/io.html + index: ':' + + # If only load_dataset_path is provided, train/valid set is automatically decided by splitting dataset by divide ratio + # If both load_dataset_path & load_validset_path is provided, use load_dataset_path as training set. + load_dataset_path: ['../data/train.extxyz'] # Example of using ase as data_format, support multiple datasets and expansion(*) + #load_validset_path: ['./valid.sevenn_data'] + + #save_dataset_path: './total' # Save the preprocessed (in load_dataset_path) dataset + #save_by_train_valid: True # Save the preprocessed train.sevenn_data, valid.sevenn_data diff --git a/sevenn/presets/fine_tune.yaml b/sevenn/presets/fine_tune.yaml new file mode 100644 index 0000000..573a50d --- /dev/null +++ b/sevenn/presets/fine_tune.yaml @@ -0,0 +1,80 @@ +# Example input.yaml for fine-tuning sevennet-0 + +model: # model keys should be consistent except for train_* keys + chemical_species: 'Auto' + cutoff: 5.0 + channel: 128 + is_parity: False + lmax: 2 + num_convolution_layer: 5 + irreps_manual: + - "128x0e" + - "128x0e+64x1e+32x2e" + - "128x0e+64x1e+32x2e" + - "128x0e+64x1e+32x2e" + - "128x0e+64x1e+32x2e" + - "128x0e" + + weight_nn_hidden_neurons: [64, 64] + radial_basis: + radial_basis_name: 'bessel' + bessel_basis_num: 8 + cutoff_function: + cutoff_function_name: 'XPLOR' + cutoff_on: 4.5 + self_connection_type: 'linear' + + train_shift_scale: False # customizable (True | False) + train_denominator: False # customizable (True | False) + +train: # Customizable + random_seed: 1 + is_train_stress: True + epoch: 100 + + optimizer: 'adam' + optim_param: + lr: 0.004 + scheduler: 'exponentiallr' + scheduler_param: + gamma: 0.99 + + force_loss_weight: 0.1 + stress_loss_weight: 1e-06 + + per_epoch: 10 # Generate checkpoints every this epoch + + # ['target y', 'metric'] + # Target y: TotalEnergy, Energy, Force, Stress, Stress_GPa, TotalLoss + # Metric : RMSE, MAE, or Loss + error_record: + - ['Energy', 'RMSE'] + - ['Force', 'RMSE'] + - ['Stress', 'RMSE'] + - ['TotalLoss', 'None'] + + continue: + reset_optimizer: True + reset_scheduler: True + reset_epoch: True + checkpoint: 'SevenNet-0_11July2024' + # Set True to use shift, scale, and avg_num_neigh from checkpoint (highly recommended) + use_statistic_values_of_checkpoint: True + +data: # Customizable + batch_size: 4 + data_divide_ratio: 0.1 + + # ase.io.read readable data files or structure_list or .sevenn_data files can be used as dataset. + # .sevenn_data is preprocessed data set has edges connected (can be obtained by using sevenn_graph_build or by save_** options below) + data_format: 'ase' # Default is 'ase' + data_format_args: # Paramaters, will be passed to ase.io.read see: https://wiki.fysik.dtu.dk/ase/ase/io/io.html + index: ':' + + # If only load_dataset_path is provided, train/valid set is automatically decided by splitting dataset by divide ratio + # If both load_dataset_path & load_validset_path is provided, use load_dataset_path as training set. + load_dataset_path: ['fine_tune.extxyz'] # Support multiple files and expansion(*) + #load_validset_path: ['./valid.sevenn_data'] + + #save_dataset_path: './total' # Save the preprocessed (in load_dataset_path) dataset + #save_by_train_valid: True # Save the preprocessed train.sevenn_data, valid.sevenn_data diff --git a/sevenn/presets/sevennet_0.yaml b/sevenn/presets/sevennet_0.yaml new file mode 100644 index 0000000..9e53f7a --- /dev/null +++ b/sevenn/presets/sevennet_0.yaml @@ -0,0 +1,80 @@ +# SevenNet-0 +model: + chemical_species: 'auto' + cutoff: 5.0 + channel: 128 + is_parity: False + lmax: 2 + num_convolution_layer: 5 + irreps_manual: + - "128x0e" + - "128x0e+64x1e+32x2e" + - "128x0e+64x1e+32x2e" + - "128x0e+64x1e+32x2e" + - "128x0e+64x1e+32x2e" + - "128x0e" + + weight_nn_hidden_neurons: [64, 64] + radial_basis: + radial_basis_name: 'bessel' + bessel_basis_num: 8 + cutoff_function: + cutoff_function_name: 'XPLOR' + cutoff_on: 4.5 + + act_gate: {'e': 'silu', 'o': 'tanh'} + act_scalar: {'e': 'silu', 'o': 'tanh'} + + conv_denominator: 'avg_num_neigh' + train_shift_scale: False + train_denominator: False + self_connection_type: 'linear' +train: + train_shuffle: False + random_seed: 1 + is_train_stress : True + epoch: 600 + + loss: 'Huber' + loss_param: + delta: 0.01 + + optimizer: 'adam' + optim_param: + lr: 0.01 + scheduler: 'linearlr' + scheduler_param: + start_factor: 1.0 + total_iters: 600 + end_factor: 0.0001 + + force_loss_weight : 1.00 + stress_loss_weight: 0.01 + + error_record: + - ['Energy', 'RMSE'] + - ['Force', 'RMSE'] + - ['Stress', 'RMSE'] + - ['Energy', 'MAE'] + - ['Force', 'MAE'] + - ['Stress', 'MAE'] + - ['Energy', 'Loss'] + - ['Force', 'Loss'] + - ['Stress', 'Loss'] + - ['TotalLoss', 'None'] + + per_epoch: 10 + # continue: + # checkpoint: './checkpoint_last.pth' + # reset_optimizer: False + # reset_scheduler: False +data: + data_shuffle: False + batch_size: 4 + scale: 'per_atom_energy_std' + shift: 'elemwise_reference_energies' + + data_format: 'ase' + save_by_train_valid: False + load_dataset_path: ["path_to_MPtrj_total.sevenn_data"] + load_validset_path: ["validaset.sevenn_data"] diff --git a/pretrained_potentials/SevenNet_0__11July2024/README.md b/sevenn/pretrained_potentials/SevenNet_0__11July2024/README.md similarity index 100% rename from pretrained_potentials/SevenNet_0__11July2024/README.md rename to sevenn/pretrained_potentials/SevenNet_0__11July2024/README.md diff --git a/pretrained_potentials/SevenNet_0__11July2024/checkpoint_sevennet_0.pth b/sevenn/pretrained_potentials/SevenNet_0__11July2024/checkpoint_sevennet_0.pth similarity index 100% rename from pretrained_potentials/SevenNet_0__11July2024/checkpoint_sevennet_0.pth rename to sevenn/pretrained_potentials/SevenNet_0__11July2024/checkpoint_sevennet_0.pth diff --git a/pretrained_potentials/SevenNet_0__11July2024/fine_tune.yaml b/sevenn/pretrained_potentials/SevenNet_0__11July2024/fine_tune.yaml similarity index 100% rename from pretrained_potentials/SevenNet_0__11July2024/fine_tune.yaml rename to sevenn/pretrained_potentials/SevenNet_0__11July2024/fine_tune.yaml diff --git a/pretrained_potentials/SevenNet_0__11July2024/parallel_model/deployed_parallel_0.pt b/sevenn/pretrained_potentials/SevenNet_0__11July2024/parallel_model/deployed_parallel_0.pt similarity index 100% rename from pretrained_potentials/SevenNet_0__11July2024/parallel_model/deployed_parallel_0.pt rename to sevenn/pretrained_potentials/SevenNet_0__11July2024/parallel_model/deployed_parallel_0.pt diff --git a/pretrained_potentials/SevenNet_0__11July2024/parallel_model/deployed_parallel_1.pt b/sevenn/pretrained_potentials/SevenNet_0__11July2024/parallel_model/deployed_parallel_1.pt similarity index 100% rename from pretrained_potentials/SevenNet_0__11July2024/parallel_model/deployed_parallel_1.pt rename to sevenn/pretrained_potentials/SevenNet_0__11July2024/parallel_model/deployed_parallel_1.pt diff --git a/pretrained_potentials/SevenNet_0__11July2024/parallel_model/deployed_parallel_2.pt b/sevenn/pretrained_potentials/SevenNet_0__11July2024/parallel_model/deployed_parallel_2.pt similarity index 100% rename from pretrained_potentials/SevenNet_0__11July2024/parallel_model/deployed_parallel_2.pt rename to sevenn/pretrained_potentials/SevenNet_0__11July2024/parallel_model/deployed_parallel_2.pt diff --git a/pretrained_potentials/SevenNet_0__11July2024/parallel_model/deployed_parallel_3.pt b/sevenn/pretrained_potentials/SevenNet_0__11July2024/parallel_model/deployed_parallel_3.pt similarity index 100% rename from pretrained_potentials/SevenNet_0__11July2024/parallel_model/deployed_parallel_3.pt rename to sevenn/pretrained_potentials/SevenNet_0__11July2024/parallel_model/deployed_parallel_3.pt diff --git a/pretrained_potentials/SevenNet_0__11July2024/parallel_model/deployed_parallel_4.pt b/sevenn/pretrained_potentials/SevenNet_0__11July2024/parallel_model/deployed_parallel_4.pt similarity index 100% rename from pretrained_potentials/SevenNet_0__11July2024/parallel_model/deployed_parallel_4.pt rename to sevenn/pretrained_potentials/SevenNet_0__11July2024/parallel_model/deployed_parallel_4.pt diff --git a/pretrained_potentials/SevenNet_0__11July2024/pre_train.yaml b/sevenn/pretrained_potentials/SevenNet_0__11July2024/pre_train.yaml similarity index 100% rename from pretrained_potentials/SevenNet_0__11July2024/pre_train.yaml rename to sevenn/pretrained_potentials/SevenNet_0__11July2024/pre_train.yaml diff --git a/pretrained_potentials/SevenNet_0__11July2024/serial_model/deployed_serial.pt b/sevenn/pretrained_potentials/SevenNet_0__11July2024/serial_model/deployed_serial.pt similarity index 100% rename from pretrained_potentials/SevenNet_0__11July2024/serial_model/deployed_serial.pt rename to sevenn/pretrained_potentials/SevenNet_0__11July2024/serial_model/deployed_serial.pt diff --git a/pretrained_potentials/SevenNet_0__22May2024/README.md b/sevenn/pretrained_potentials/SevenNet_0__22May2024/README.md similarity index 100% rename from pretrained_potentials/SevenNet_0__22May2024/README.md rename to sevenn/pretrained_potentials/SevenNet_0__22May2024/README.md diff --git a/pretrained_potentials/SevenNet_0__22May2024/checkpoint_sevennet_0.pth b/sevenn/pretrained_potentials/SevenNet_0__22May2024/checkpoint_sevennet_0.pth similarity index 100% rename from pretrained_potentials/SevenNet_0__22May2024/checkpoint_sevennet_0.pth rename to sevenn/pretrained_potentials/SevenNet_0__22May2024/checkpoint_sevennet_0.pth diff --git a/pretrained_potentials/SevenNet_0__22May2024/fine_tune.yaml b/sevenn/pretrained_potentials/SevenNet_0__22May2024/fine_tune.yaml similarity index 100% rename from pretrained_potentials/SevenNet_0__22May2024/fine_tune.yaml rename to sevenn/pretrained_potentials/SevenNet_0__22May2024/fine_tune.yaml diff --git a/pretrained_potentials/SevenNet_0__22May2024/parallel_model/deployed_parallel_0.pt b/sevenn/pretrained_potentials/SevenNet_0__22May2024/parallel_model/deployed_parallel_0.pt similarity index 100% rename from pretrained_potentials/SevenNet_0__22May2024/parallel_model/deployed_parallel_0.pt rename to sevenn/pretrained_potentials/SevenNet_0__22May2024/parallel_model/deployed_parallel_0.pt diff --git a/pretrained_potentials/SevenNet_0__22May2024/parallel_model/deployed_parallel_1.pt b/sevenn/pretrained_potentials/SevenNet_0__22May2024/parallel_model/deployed_parallel_1.pt similarity index 100% rename from pretrained_potentials/SevenNet_0__22May2024/parallel_model/deployed_parallel_1.pt rename to sevenn/pretrained_potentials/SevenNet_0__22May2024/parallel_model/deployed_parallel_1.pt diff --git a/pretrained_potentials/SevenNet_0__22May2024/parallel_model/deployed_parallel_2.pt b/sevenn/pretrained_potentials/SevenNet_0__22May2024/parallel_model/deployed_parallel_2.pt similarity index 100% rename from pretrained_potentials/SevenNet_0__22May2024/parallel_model/deployed_parallel_2.pt rename to sevenn/pretrained_potentials/SevenNet_0__22May2024/parallel_model/deployed_parallel_2.pt diff --git a/pretrained_potentials/SevenNet_0__22May2024/parallel_model/deployed_parallel_3.pt b/sevenn/pretrained_potentials/SevenNet_0__22May2024/parallel_model/deployed_parallel_3.pt similarity index 100% rename from pretrained_potentials/SevenNet_0__22May2024/parallel_model/deployed_parallel_3.pt rename to sevenn/pretrained_potentials/SevenNet_0__22May2024/parallel_model/deployed_parallel_3.pt diff --git a/pretrained_potentials/SevenNet_0__22May2024/parallel_model/deployed_parallel_4.pt b/sevenn/pretrained_potentials/SevenNet_0__22May2024/parallel_model/deployed_parallel_4.pt similarity index 100% rename from pretrained_potentials/SevenNet_0__22May2024/parallel_model/deployed_parallel_4.pt rename to sevenn/pretrained_potentials/SevenNet_0__22May2024/parallel_model/deployed_parallel_4.pt diff --git a/pretrained_potentials/SevenNet_0__22May2024/serial_model/deployed_serial.pt b/sevenn/pretrained_potentials/SevenNet_0__22May2024/serial_model/deployed_serial.pt similarity index 100% rename from pretrained_potentials/SevenNet_0__22May2024/serial_model/deployed_serial.pt rename to sevenn/pretrained_potentials/SevenNet_0__22May2024/serial_model/deployed_serial.pt diff --git a/pretrained_potentials/archive/240423_SevenNet_0/240423_checkpoint_sevennet_0.pth b/sevenn/pretrained_potentials/archive/240423_SevenNet_0/240423_checkpoint_sevennet_0.pth similarity index 100% rename from pretrained_potentials/archive/240423_SevenNet_0/240423_checkpoint_sevennet_0.pth rename to sevenn/pretrained_potentials/archive/240423_SevenNet_0/240423_checkpoint_sevennet_0.pth diff --git a/pretrained_potentials/archive/240423_SevenNet_0/README.txt b/sevenn/pretrained_potentials/archive/240423_SevenNet_0/README.txt similarity index 100% rename from pretrained_potentials/archive/240423_SevenNet_0/README.txt rename to sevenn/pretrained_potentials/archive/240423_SevenNet_0/README.txt diff --git a/sevenn/scripts/deploy.py b/sevenn/scripts/deploy.py index 111dd91..0ba8519 100644 --- a/sevenn/scripts/deploy.py +++ b/sevenn/scripts/deploy.py @@ -81,8 +81,7 @@ def deploy_parallel(model_state_dct, config, fname): chem_list.strip() # dim of irreps_in of last model convolution is (max)comm_size - # except first one, first of every model is embedding followed by convolution - # TODO: this code is error prone + # TODO: this is error prone comm_size = ( model_list[-1][1].convolution.irreps_in1.dim if len(model_list) > 1 diff --git a/sevenn/scripts/processing_continue.py b/sevenn/scripts/processing_continue.py index 4af6dac..607b3ac 100644 --- a/sevenn/scripts/processing_continue.py +++ b/sevenn/scripts/processing_continue.py @@ -58,15 +58,8 @@ def processing_continue(config): continue_dct = config[KEY.CONTINUE] Logger().write('\nContinue found, loading checkpoint\n') - try: - checkpoint = torch.load( - continue_dct[KEY.CHECKPOINT], map_location='cpu' - ) - config_cp = checkpoint['config'] - except FileNotFoundError: - raise FileNotFoundError( - f'checkpoint file {continue_dct[KEY.CHECKPOINT]} not found' - ) + checkpoint = torch.load(continue_dct[KEY.CHECKPOINT], map_location='cpu') + config_cp = checkpoint['config'] model_cp, config_cp = util.model_from_checkpoint(checkpoint) model_state_dict_cp = model_cp.state_dict() @@ -75,9 +68,8 @@ def processing_continue(config): check_config_compatible(config, config_cp) Logger().write('Checkpoint config is compatible\n') - ################## for backward compat. + # for backward compat. config.update({KEY._NORMALIZE_SPH: config_cp[KEY._NORMALIZE_SPH]}) - ################## for backward compat. from_epoch = checkpoint['epoch'] optimizer_state_dict_cp = ( @@ -104,7 +96,7 @@ def processing_continue(config): ) del model_state_dict_cp[f'{i}_convolution.denominator'] - # these dataset-dependent values should be later handled by processing_dataset.py + # Further handled by processing_dataset.py config.update({ KEY.SHIFT + '_cp': shift_cp, KEY.SCALE + '_cp': scale_cp, diff --git a/sevenn/scripts/processing_dataset.py b/sevenn/scripts/processing_dataset.py index 5f74f16..8a816be 100644 --- a/sevenn/scripts/processing_dataset.py +++ b/sevenn/scripts/processing_dataset.py @@ -71,11 +71,13 @@ def handle_shift_scale(config, train_set, checkpoint_given): use_species_wise_shift_scale = True avg_num_neigh = train_set.get_avg_num_neigh() - Logger().format_k_v('Average # of neighbors', f'{avg_num_neigh:.6f}', write=True) + Logger().format_k_v( + 'Average # of neighbors', f'{avg_num_neigh:.6f}', write=True + ) - if config[KEY.CONV_DENOMINATOR] == "avg_num_neigh": + if config[KEY.CONV_DENOMINATOR] == 'avg_num_neigh': conv_denominator = avg_num_neigh - elif config[KEY.CONV_DENOMINATOR] == "sqrt_avg_num_neigh": + elif config[KEY.CONV_DENOMINATOR] == 'sqrt_avg_num_neigh': conv_denominator = avg_num_neigh ** (0.5) if ( @@ -85,7 +87,7 @@ def handle_shift_scale(config, train_set, checkpoint_given): Logger().writeline( 'Overwrite shift, scale, conv_denominator from model checkpoint' ) - # Values extracted from checkpoint (not config) in processing_continue.py + # Values extracted from checkpoint in processing_continue.py shift = config[KEY.SHIFT + '_cp'] scale = config[KEY.SCALE + '_cp'] # shift & scale would be both array (with same lenght) or scalar @@ -126,9 +128,13 @@ def handle_shift_scale(config, train_set, checkpoint_given): Logger().format_k_v(f'{cstr}', f'{sh:.6f}, {sc:.6f}', write=True) else: Logger().write('Use global shift, scale\n') - Logger().format_k_v("shift, scale", f'{shift:.6f}, {scale:.6f}', write=True) + Logger().format_k_v( + 'shift, scale', f'{shift:.6f}, {scale:.6f}', write=True + ) - Logger().format_k_v("(1st) conv_denominator is", f'{conv_denominator[0]:.6f}', write=True) + Logger().format_k_v( + '(1st) conv_denominator is', f'{conv_denominator[0]:.6f}', write=True + ) config[KEY.USE_SPECIES_WISE_SHIFT_SCALE] = use_species_wise_shift_scale return shift, scale, conv_denominator @@ -268,7 +274,7 @@ def processing_dataset(config, working_dir): Logger().write('Turn off stress training\n') Logger().bar() - # If I did right job, saved data must have atomic numbers as X not one hot idx + # saved data must have atomic numbers as X not one hot idx if config[KEY.SAVE_BY_TRAIN_VALID]: train_set.save(prefix + 'train') valid_set.save(prefix + 'valid') @@ -292,9 +298,11 @@ def processing_dataset(config, working_dir): shift, scale, conv_denominator = handle_shift_scale( config, train_set, checkpoint_given ) - config.update( - {KEY.SHIFT: shift, KEY.SCALE: scale, KEY.CONV_DENOMINATOR: conv_denominator} - ) + config.update({ + KEY.SHIFT: shift, + KEY.SCALE: scale, + KEY.CONV_DENOMINATOR: conv_denominator, + }) data_lists = (train_set.to_list(), valid_set.to_list(), test_set.to_list()) if config[KEY.DATA_SHUFFLE]: diff --git a/sevenn/sevenn_logger.py b/sevenn/sevenn_logger.py index 7e6e4a4..94453ed 100644 --- a/sevenn/sevenn_logger.py +++ b/sevenn/sevenn_logger.py @@ -25,7 +25,9 @@ class Logger(metaclass=Singleton): SCREEN_WIDTH = 120 # half size of my screen / changed due to stress output - def __init__(self, filename: str = "log.sevenn", screen: bool = False, rank: int = 0): + def __init__( + self, filename: str = 'log.sevenn', screen: bool = False, rank: int = 0 + ): self.rank = rank if rank == 0: self.logfile = open(filename, 'w', buffering=1) @@ -93,10 +95,7 @@ def natoms_write(self, natoms): content += self.format_k_v('Total', sum(total_natom.values())) self.write(content) - def statistic_write(self, statistic): - """ - expect statistic is dict(key as label) of dict(key of mean, std, and so on) - """ + def statistic_write(self, statistic: dict): content = '' for label, dct in statistic.items(): dct_new = {} @@ -235,7 +234,7 @@ def format_k_v(key, val, write=False): return content else: Logger().write(content) - return "" + return '' def greeting(self): LOGO_ASCII_FILE = f'{os.path.dirname(__file__)}/logo_ascii' diff --git a/sevenn/sevennet_calculator.py b/sevenn/sevennet_calculator.py index 4c1487b..3c9eb1f 100644 --- a/sevenn/sevennet_calculator.py +++ b/sevenn/sevennet_calculator.py @@ -7,11 +7,13 @@ from ase.calculators.calculator import Calculator, all_changes from ase.data import chemical_symbols +import sevenn._const import sevenn._keys as KEY import sevenn.util torch_script_type = torch.jit._script.RecursiveScriptModule + class SevenNetCalculator(Calculator): """ASE calculator for SevenNet models @@ -32,21 +34,19 @@ def __init__( file_type: str = 'checkpoint', device: Union[torch.device, str] = 'auto', sevennet_config=None, - **kwargs + **kwargs, ): """Initialize the calculator Args: - model (SevenNet): AtomGraphSequential or path to the checkpoint file. + model (SevenNet): path to the checkpoint file, or pretrained device (str, optional): Torch device to use. Defaults to "auto". """ super().__init__(**kwargs) file_type = file_type.lower() if file_type not in ['checkpoint', 'torchscript']: - raise ValueError( - 'file_type should be checkpoint or torchscript' - ) + raise ValueError('file_type should be checkpoint or torchscript') if not isinstance(device, torch.device) and not isinstance( device, str @@ -65,16 +65,13 @@ def __init__( self.device = device if file_type == 'checkpoint': - if model == 'SevenNet-0': # special case loading pre-trained model - checkpoint = os.getenv('SEVENNET_0_CP') - if checkpoint is None: - raise ValueError( - 'Please set env variable SEVENNET_0_CP as checkpoint' - ' path.' - ) - else: + if os.path.isfile(model): checkpoint = model - model_loaded, config = sevenn.util.model_from_checkpoint(checkpoint) + else: + checkpoint = sevenn.util.pretrained_name_to_path(model) + model_loaded, config = sevenn.util.model_from_checkpoint( + checkpoint + ) model_loaded.set_is_batch_data(False) self.type_map = config[KEY.TYPE_MAP] self.cutoff = config[KEY.CUTOFF] @@ -90,24 +87,32 @@ def __init__( 'time': b'', } model_loaded = torch.jit.load( - model, - _extra_files=extra_dict, - map_location=self.device + model, _extra_files=extra_dict, map_location=self.device + ) + chem_symbols = extra_dict['chemical_symbols_to_index'].decode( + 'utf-8' ) - chem_symbols = extra_dict['chemical_symbols_to_index'].decode('utf-8') sym_to_num = {sym: n for n, sym in enumerate(chemical_symbols)} - self.type_map = {sym_to_num[sym]: i for i, sym in enumerate(chem_symbols.split())} + self.type_map = { + sym_to_num[sym]: i + for i, sym in enumerate(chem_symbols.split()) + } self.cutoff = float(extra_dict['cutoff'].decode('utf-8')) else: - raise ValueError(f"Unknown file type") + raise ValueError('Unknown file type') self.model = model_loaded self.model.to(self.device) self.model.eval() - # atomic_energy = energies in ASE, atoms_instance.get_potential_energies() - self.implemented_properties = ['free_energy', 'energy', 'forces', 'stress', 'energies'] + self.implemented_properties = [ + 'free_energy', + 'energy', + 'forces', + 'stress', + 'energies', + ] def calculate( self, atoms=None, properties=None, system_changes=all_changes @@ -123,7 +128,7 @@ def calculate( if isinstance(self.model, torch_script_type): data = data.to_dict() - del data["data_info"] + del data['data_info'] output = self.model(data) energy = output[KEY.PRED_TOTAL_ENERGY].detach().cpu().item() @@ -131,7 +136,13 @@ def calculate( self.results = { 'free_energy': energy, 'energy': energy, - 'energies': output[KEY.ATOMIC_ENERGY].detach().cpu().reshape(len(atoms)).numpy(), + 'energies': ( + output[KEY.ATOMIC_ENERGY] + .detach() + .cpu() + .reshape(len(atoms)) + .numpy() + ), 'forces': output[KEY.PRED_FORCE].detach().cpu().numpy(), 'stress': np.array( (-output[KEY.PRED_STRESS]) diff --git a/sevenn/train/dataload.py b/sevenn/train/dataload.py index 5acf19b..89a5d0d 100644 --- a/sevenn/train/dataload.py +++ b/sevenn/train/dataload.py @@ -36,7 +36,6 @@ def unlabeled_atoms_to_graph(atoms: ase.Atoms, cutoff: float): 'ijDS', atoms.get_pbc(), cell, pos, cutoff, self_interaction=True ) - # remove redundant edges (self interaction) but saves self interaction cross PBC is_zero_idx = np.all(edge_vec == 0, axis=1) is_self_idx = edge_src == edge_dst non_trivials = ~(is_zero_idx & is_self_idx) @@ -83,8 +82,8 @@ def atoms_to_graph( Raises: RuntimeError: if ase atoms are somewhat imperfect - Use free_energy by default (atoms.get_potential_energy(force_consistent=True)) - If it is not available, use energy (atoms.get_potential_energy()) + Use free_energy: atoms.get_potential_energy(force_consistent=True) + If it is not available, use atoms.get_potential_energy() If stress is available, initialize stress tensor Ignore constraints like selective dynamics @@ -114,7 +113,6 @@ def atoms_to_graph( 'ijDS', atoms.get_pbc(), cell, pos, cutoff, self_interaction=True ) - # remove redundant edges (self interaction) but saves self interaction cross PBC is_zero_idx = np.all(edge_vec == 0, axis=1) is_self_idx = edge_src == edge_dst non_trivials = ~(is_zero_idx & is_self_idx) @@ -177,7 +175,6 @@ def graph_build( if not serial: pool = mp.Pool(num_cores) - # this is not strictly correct because it updates for every input not output graph_list = pool.starmap( atoms_to_graph, tqdm.tqdm(inputs, total=len(atoms_list)) ) @@ -204,9 +201,9 @@ def pkl_atoms_reader(fname): """ with open(fname, 'rb') as f: atoms_list = pickle.load(f) - if type(atoms_list) != list: + if not isinstance(atoms_list, list): raise TypeError('The content of the pkl is not list') - if type(atoms_list[0]) != ase.Atoms: + if not isinstance(atoms_list[0], ase.Atoms): raise TypeError('The content of the pkl is not list of ase.Atoms') return atoms_list @@ -313,7 +310,7 @@ def match_reader(reader_name: str, **kwargs): metadata.update({'origin': 'structure_list'}) else: reader = partial(ase_reader, **kwargs) - metadata.update({'origin': f'ase_reader'}) + metadata.update({'origin': 'ase_reader'}) return reader, metadata @@ -332,15 +329,15 @@ def file_to_dataset( # expect label: atoms_list dct or atoms or list of atoms atoms = reader(file) - if type(atoms) == list: + if type(atoms) is list: if label is None: label = KEY.LABEL_NONE atoms_dct = {label: atoms} - elif type(atoms) == ase.Atoms: + elif isinstance(atoms, ase.Atoms): if label is None: label = KEY.LABEL_NONE atoms_dct = {label: [atoms]} - elif type(atoms) == dict: + elif isinstance(atoms, dict): atoms_dct = atoms else: raise TypeError('The return of reader is not list or dict') diff --git a/sevenn/train/dataset.py b/sevenn/train/dataset.py index 1a9cc84..51562d0 100644 --- a/sevenn/train/dataset.py +++ b/sevenn/train/dataset.py @@ -43,7 +43,7 @@ def __init__( dataset: Union[Dict[str, List], List], cutoff: float, metadata: Optional[Dict] = None, - x_is_one_hot_idx: Optional[bool] = False, + x_is_one_hot_idx: bool = False, ): """ Default constructor of AtomGraphDataset @@ -51,14 +51,13 @@ def __init__( dataset (Union[Dict[str, List], List]: dataset as dict or pure list metadata (Dict, Optional): metadata of data cutoff (float): cutoff radius of graphs inside the dataset - x_is_one_hot_idx (bool, Optional): if True, x is one_hot_idx, eles 'Z' + x_is_one_hot_idx (bool): if True, x is one_hot_idx, eles 'Z' 'x' (node feature) of dataset can have 3 states, atomic_numbers, one_hot_idx, or one_hot_vector. - atomic_numbers is the most general one but cannot directly used for input + atomic_numbers is general but cannot directly used for input one_hot_idx is can be input of the model but requires 'type_map' - one_hot_idx to one_hot_vector is done by model, so it should not be used here """ self.cutoff = cutoff self.x_is_one_hot_idx = x_is_one_hot_idx @@ -196,7 +195,6 @@ def divide_dataset( """ def divide(ratio: float, data_list: List, ignore_test=True): - # Get list as input and return list divided by 1-2*ratio : ratio : ratio if ratio > 0.5: raise ValueError('Ratio must not exceed 0.5') data_len = len(data_list) @@ -392,7 +390,6 @@ def augment(self, dataset, validator: Optional[Callable] = None): check consistent data type, float, double, long integer etc """ - assert type(dataset) == AtomGraphDataset def default_validator(db1, db2): cut_consis = db1.cutoff == db2.cutoff @@ -425,19 +422,13 @@ def delete_data_key(self, key): # TODO: this by_label is not straightforward def save(self, path, by_label=False): - """ - with open(path, 'wb') as f: - pickle.dump(self, f) - """ if by_label: for label, data in self.dataset.items(): - to = f'{path}/{label}.sevenn_data' - # torch.save(AtomGraphDataset({label: data}, metadata=self.meta), to) torch.save( AtomGraphDataset( {label: data}, self.cutoff, metadata=self.meta ), - to, + f'{path}/{label}.sevenn_data', ) else: if path.endswith('.sevenn_data') is False: diff --git a/sevenn/train/optim.py b/sevenn/train/optim.py index 05aaefe..68984dc 100644 --- a/sevenn/train/optim.py +++ b/sevenn/train/optim.py @@ -10,48 +10,6 @@ 'radam': optim.RAdam, } -optim_param_name_type_dict = { - 'universial': {'lr': float, 'weight_decay': float}, - 'sgd': { - 'momentum': float, - 'dampening': float, - 'nesterov': bool, - 'maximize': bool, - 'foreach': bool, - }, - 'adagrad': { - 'lr_decay': float, - 'eps': float, - 'foreach': bool, - 'maximize': bool, - }, - 'adam': { - 'betas': tuple, # How to make it tuple[float, float]? - 'eps': float, - 'amsgrad': bool, - 'foreach': bool, - 'maximize': bool, - 'capturable': bool, - 'fused': bool, - }, - 'adamw': { - 'betas': tuple, # How to make it tuple[float, float]? - 'eps': float, - 'amsgrad': bool, - 'maximize': bool, - 'foreach': bool, - 'capturable': bool, - }, - 'radam': { - 'betas': tuple, # How to make it tuple[float, float]? - 'eps': float, - 'foreach': bool, - }, -} - -# Some scheduler use lambda function (e.g. LambdaLR) as input. -# This is not possible using simple yaml configuration. -# TODO: How to implement this? scheduler_dict = { 'steplr': scheduler.StepLR, @@ -62,36 +20,4 @@ 'linearlr': scheduler.LinearLR, } -scheduler_param_name_type_dict = { - 'universial': {'last_epoch': int, 'verbose': bool}, - 'steplr': {'step_size': int, 'gamma': float}, - 'multisteplr': { - 'milestones': list, # How to make it list[int]? - 'gamma': float, - }, - 'exponentiallr': {'gamma': float}, - 'cosineannealinglr': {'T_max': int, 'eta_min': float}, - 'reducelronplateau': { - 'mode': str, - 'factor': float, - 'patience': int, - 'threshold': float, - 'threshold_mode': str, - 'cooldown': int, - 'min_lr': float, - 'eps': float, - }, - 'linearlr': { - 'start_factor': float, - 'end_factor': float, - 'total_iters': int, - 'last_epoch': int, - } -} - loss_dict = {'mse': nn.MSELoss, 'huber': nn.HuberLoss} -loss_param_name_type_dict = { - 'universial': {}, - 'mse': {}, - 'huber': {'delta': float}, # default = 1.0 -} diff --git a/sevenn/util.py b/sevenn/util.py index 310fb50..f6544ab 100644 --- a/sevenn/util.py +++ b/sevenn/util.py @@ -1,9 +1,9 @@ -from typing import Union import warnings +from typing import Union import numpy as np import torch -from e3nn.o3 import Irreps, FullTensorProduct +from e3nn.o3 import FullTensorProduct, Irreps import sevenn._keys as KEY @@ -82,7 +82,7 @@ def error_recorder_from_loss_functions(loss_functions): for loss_function, _ in loss_functions: ref_key = loss_function.ref_key pred_key = loss_function.pred_key - unit = loss_function.unit + # unit = loss_function.unit criterion = loss_function.criterion name = loss_function.name base = None @@ -159,18 +159,19 @@ def onehot_to_chem(one_hot_indicies, type_map): def _patch_old_config(config): # Fixing my old mistakes - if config[KEY.CUTOFF_FUNCTION][KEY.CUTOFF_FUNCTION_NAME] == "XPLOR": - config[KEY.CUTOFF_FUNCTION].pop("poly_cut_p_value", None) - config[KEY.TRAIN_DENOMINTAOR] = config.pop("train_avg_num_neigh", False) - _opt = config.pop("optimize_by_reduce", None) - if _opt == False: - raise ValueError("This checkpoint(optimize_by_reduce: False) is no longer supported") + if config[KEY.CUTOFF_FUNCTION][KEY.CUTOFF_FUNCTION_NAME] == 'XPLOR': + config[KEY.CUTOFF_FUNCTION].pop('poly_cut_p_value', None) + config[KEY.TRAIN_DENOMINTAOR] = config.pop('train_avg_num_neigh', False) + _opt = config.pop('optimize_by_reduce', None) + if _opt is False: + raise ValueError( + 'This checkpoint(optimize_by_reduce: False) is no longer supported' + ) if KEY.CONV_DENOMINATOR not in config: config[KEY.CONV_DENOMINATOR] = 0.0 - #warnings.warn(f'conv denominator will be loaded from state_dict of checkpoint', UserWarning) if KEY._NORMALIZE_SPH not in config: config[KEY._NORMALIZE_SPH] = False - warnings.warn(f'You are loading very old model, normalize_sph is False', UserWarning) + # Warn this in the docs, not here for SevenNet-0 (22May2024) return config @@ -212,11 +213,7 @@ def _map_old_model(old_model_state_dict): def model_from_checkpoint(checkpoint): - from sevenn._const import ( - model_defaults, - data_defaults, - train_defaults, - ) + from sevenn._const import data_defaults, model_defaults, train_defaults from sevenn.model_build import build_E3_equivariant_model if isinstance(checkpoint, str): @@ -237,7 +234,9 @@ def model_from_checkpoint(checkpoint): for k, v in defaults.items(): if k not in config: - warnings.warn(f'{k} not in config, using default value {v}', UserWarning) + warnings.warn( + f'{k} not in config, using default value {v}', UserWarning + ) config[k] = v # expect only non-tensor values in config, if exists, move to cpu @@ -261,8 +260,8 @@ def model_from_checkpoint(checkpoint): def unlabeled_atoms_to_input(atoms, cutoff): - from sevenn.train.dataload import unlabeled_atoms_to_graph from sevenn.atom_graph_data import AtomGraphData + from sevenn.train.dataload import unlabeled_atoms_to_graph atom_graph = AtomGraphData.from_numpy_dict( unlabeled_atoms_to_graph(atoms, cutoff) @@ -319,7 +318,11 @@ def load_model_from_checkpoint(checkpoint): DEFAULT_TRAINING_CONFIG, ) from sevenn.model_build import build_E3_equivariant_model - warnings.warn(f'This method is deprecated, use model_from_checkpoint instead', DeprecationWarning) + + warnings.warn( + 'This method is deprecated, use model_from_checkpoint instead', + DeprecationWarning, + ) if isinstance(checkpoint, str): checkpoint = torch.load(checkpoint, map_location='cpu') @@ -375,9 +378,25 @@ def infer_irreps_out( continue if parity_mode == 'even' and p == -1: continue - elif parity_mode == 'sph' and p != (-1)**l: + elif parity_mode == 'sph' and p != (-1) ** l: continue if fix_multiplicity: elem = (fix_multiplicity, (l, p)) new_irreps_elem.append(elem) return Irreps(new_irreps_elem) + + +def pretrained_name_to_path(name: str) -> str: + import sevenn._const as _const + + name = name.lower() + heads = ['sevennet-0', '7net-0'] + checkpoint_path = None + if name in [f'{n}_11july2024' for n in heads] or name in heads: + checkpoint_path = _const.SEVENNET_0_11July2024 + elif name in [f'{n}_22may2024' for n in heads]: + checkpoint_path = _const.SEVENNET_0_22May2024 + else: + raise ValueError('Not a valid potential') + + return checkpoint_path From fc4b0259c50efe1d298a97160a197c8d03f8b614 Mon Sep 17 00:00:00 2001 From: YutackPark Date: Thu, 18 Jul 2024 15:05:34 +0900 Subject: [PATCH 02/25] bugfix: loss param to const --- sevenn/_const.py | 1 + 1 file changed, 1 insertion(+) diff --git a/sevenn/_const.py b/sevenn/_const.py index 4bbf3ba..7dbe364 100644 --- a/sevenn/_const.py +++ b/sevenn/_const.py @@ -204,6 +204,7 @@ def data_defaults(config): KEY.RANDOM_SEED: 1, KEY.EPOCH: 300, KEY.LOSS: 'mse', + KEY.LOSS_PARAM: {}, KEY.OPTIMIZER: 'adam', KEY.OPTIM_PARAM: {}, KEY.SCHEDULER: 'exponentiallr', From dbb2949e836103cd176d6e4b2388218b710a331f Mon Sep 17 00:00:00 2001 From: YutackPark <111348843+YutackPark@users.noreply.github.com> Date: Thu, 18 Jul 2024 15:47:05 +0900 Subject: [PATCH 03/25] docs: update readme.md with tools --- README.md | 179 +++++++++++++++++------------------------------------- 1 file changed, 56 insertions(+), 123 deletions(-) diff --git a/README.md b/README.md index f68c756..896e60c 100644 --- a/README.md +++ b/README.md @@ -9,8 +9,6 @@ SevenNet (Scalable EquiVariance Enabled Neural Network) is a graph neural networ The project provides parallel molecular dynamics simulations using graph neural network interatomic potentials, which enable large-scale MD simulations or faster MD simulations. -**PLEASE NOTE:** SevenNet is under active development and may not be fully stable. - The installation and usage of SevenNet are split into two parts: training (handled by PyTorch) and molecular dynamics (handled by [`LAMMPS`](https://github.com/lammps/lammps)). The model, once trained with PyTorch, is deployed using TorchScript and is later used to run molecular dynamics simulations via LAMMPS. - [SevenNet](#sevennet) @@ -20,10 +18,10 @@ The installation and usage of SevenNet are split into two parts: training (handl + [SevenNet Calculator for ASE](#sevennet-calculator-for-ase) + [Training sevenn](#training) - [Multi-GPU training](#multi-gpu-training) + + [sevenn_graph build](#sevenn_graph_build) + [sevenn_inference](#sevenn_inference) + [sevenn_get_model](#sevenn_get_model) * [Installation for LAMMPS](#installation-for-lammps) - + [Notes](#notes) * [Usage for LAMMPS](#usage-for-lammps) + [To check installation](#to-check-installation) + [For serial model](#for-serial-model) @@ -49,85 +47,86 @@ pip install . ## Usage ### SevenNet-0 -SevenNet-0 is a general-purpose interatomic potential trained on the [`MPF dataset of M3GNet`](https://figshare.com/articles/dataset/MPF_2021_2_8/19470599) or [`MPtrj dataset of CHGNet`](https://figshare.com/articles/dataset/Materials_Project_Trjectory_MPtrj_Dataset/23713842). You can try SevenNet-0 to your application without any training. If the accuracy is unsatisfactory, SevenNet-0 can be [fine-tuned](#Training). +SevenNet-0 is a general-purpose interatomic potential trained on the [`MPF dataset of M3GNet`](https://figshare.com/articles/dataset/MPF_2021_2_8/19470599) or [`MPtrj dataset of CHGNet`](https://figshare.com/articles/dataset/Materials_Project_Trjectory_MPtrj_Dataset/23713842). You can try SevenNet-0 to your application without any training. If the accuracy is unsatisfactory, SevenNet-0 can be [fine-tuned](#training). #### SevenNet-0 (11July2024) -This model was trained on [`MPtrj`](https://figshare.com/articles/dataset/Materials_Project_Trjectory_MPtrj_Dataset/23713842). We suggest starting with this model as we found that it performs better than the previous SevenNet-0 (22May2024). +This model was trained on [`MPtrj`](https://figshare.com/articles/dataset/Materials_Project_Trjectory_MPtrj_Dataset/23713842). We suggest starting with this model as we found that it performs better than the previous SevenNet-0 (22May2024). Check [`Matbench Discovery leaderborad`](https://matbench-discovery.materialsproject.org/) for this model's performance on materials discovery. + +Whenever the checkpoint path is the input, this model can be loaded via `7net-0 | SevenNet-0 | 7net-0_11July2024 | SevenNet-0_11July2024` keywords. #### SevenNet-0 (22May2024) This model was trained on [`MPF.2021.2.8`](https://figshare.com/articles/dataset/MPF_2021_2_8/19470599). This is the model used in [our paper](https://pubs.acs.org/doi/10.1021/acs.jctc.4c00190). -Checkpoints of SevenNet-0 (for use in ASE or fine-tuning) and deployed potentials (for LAMMPS) are located in `{path_to_SevenNet}/pretrained_potentials/SevenNet_0__{release date}`. - -For its detailed usage, please check [SevenNet Calculator for ASE](#sevennet-calculator-for-ase), [For serial model](#for-serial-model), and [For parallel model](#for-parallel-model) +Whenever the checkpoint path is the input, this model can be loaded via `7net-0_22May2024 | SevenNet-0_22May2024` keywords. ### SevenNet Calculator for ASE +[ASE (Atomic Simulation Environment)](https://wiki.fysik.dtu.dk/ase/) is a set of tools and Python modules for atomistic simulations. SevenNet-0 and SevenNet-trained potentials can be used with ASE for its use in python. + +For pre-trained models, ```python from sevenn.sevennet_calculator import SevenNetCalculator -checkpoint_path = ### PATH TO CHECKPOINT ### -sevenet_cal = SevenNetCalculator(checkpoint_path, device='cpu') +sevenet_0_cal = SevenNetCalculator("7net-0", device='cpu') # 7net-0, SevenNet-0, 7net-0_22May2024, 7net-0_11July2024 ... ``` -If you want to use SevenNet-0, you can do something like below -```bash -echo "export SEVENNET_0_CP={PATH_TO_SEVENNET}/pretrained_potentials/SevenNet_0__11July2024/checkpoint_sevennet_0.pth" >> ~/.bashrc -``` -SevenNetCalculator tries to read the SEVENNET_0_CP environment variable. - +For user trained models, ```python from sevenn.sevennet_calculator import SevenNetCalculator -sevenet_0_cal = SevenNetCalculator(device='cpu') +checkpoint_path = ### PATH TO CHECKPOINT ### +sevenet_cal = SevenNetCalculator(checkpoint_path, device='cpu') ``` ### Training -``` -cd example_inputs/training -sevenn input_full.yaml -s +```bash +sevenn_preset base > input.yaml +sevenn input.yaml -s ``` -Example `input_full.yaml` can be found under `SevenNet/example_inputs`. The `structure_list` file is used to select VASP OUTCARs for training. -To reuse a preprocessed training set, you can specify `${dataset_name}.sevenn_data` to the `load_dataset_path:` in the `input.yaml`. +Other valid preset options are: `base`, `fine_tune`, and `sevennet-0`. +Check comments of `base` yaml for explanations. -Once you initiate training, `log.sevenn` will contain all parsed inputs from `input.yaml`. Any parameters not specified in the input will be automatically assigned as their default values. You can refer to the log to check the default inputs. -Currently, detailed explanations of model hyperparameters can be found at `input_full.yaml`. +To reuse a preprocessed training set, you can specify `${dataset_name}.sevenn_data` to the `load_dataset_path:` in the `input.yaml`. +Once you initiate training, `log.sevenn` will contain all parsed inputs from `input.yaml`. You can refer to the log to check the default inputs. #### Multi-GPU training - We support multi-GPU training features using PyTorch DDP (distributed data parallel). We use one process (CPU core) per GPU. -``` +```bash torchrun --standalone --nnodes={# of nodes} --nproc_per_node {# of GPUs} --no_python sevenn input.yaml -d ``` Please note that `batch_size` in input.yaml indicates `batch_size` per GPU. -### sevenn_inference - -Assuming that you've done temporal training of 10 epochs by above "To start training using 'sevenn'", try below at the same directory -``` -sevenn_inference checkpoint_best.pt ../data/label_1/* +### sevenn_graph_build +```bash +sevenn_graph_build -f ase my_train_data.extxyz 5.0 ``` -This will create dir 'sevenn_infer_result'. It includes .csv files that enumerate prediction/reference results of energy and force on OUTCARs in `data/label_1` directory. -You can try `sevenn_inference --help` for more information on this command. - -### sevenn_get_model +You can preprocess the dataset with `sevenn_graph_build` to obtain `*.sevenn_data` files. The cutoff length should be provided. +See `sevenn_graph_build --help` for more information. -Assuming that you've done temporal training of 10 epochs by above "To start training using 'sevenn'", try below at the same directory -``` -sevenn_get_model checkpoint_best.pt +### sevenn_inference +```bash +sevenn_inference checkpoint_best.pt path_to_my_structures/* ``` +This will create dir `sevenn_infer_result`. It includes .csv files that enumerate prediction/reference results of energy and force. +See `sevenn_inference --help` for more information. -This will create `deployed_serial.pt`, which can be used as lammps potential under `e3gnn` pair_style. Please take a look at the lammps installation process below. +### sevenn_get_model +This command is for deploying lammps potentials from checkpoints. The argument is either the path to checkpoint or the name of pre-trained potential. +```bash +sevenn_get_model 7net-0 +``` +This will create `deployed_serial.pt`, which can be used as lammps potential under `e3gnn` pair_style. The parallel model can be obtained in a similar way +```bash +sevenn_get_model 7net-0 -p ``` -sevenn_get_model checkpoint_best.pt -p -``` - This will create multiple `deployed_parallel_*.pt` files. The number of deployed models equals the number of message-passing layers. These models can be used as lammps potential to run parallel MD simulations with GNN potential using multiple GPU cards. +See `sevenn_inference --help` for more information. + ## Installation for LAMMPS * PyTorch (same version as used for training) @@ -139,20 +138,19 @@ These models can be used as lammps potential to run parallel MD simulations with **PLEASE NOTE:** CUDA-aware OpenMPI does not support NVIDIA Gaming GPUs. Given that the software is closely tied to hardware specifications, please consult with your server administrator if unavailable. Ensure the LAMMPS version (stable_2Aug2023). You can easily switch the version using git. -``` -$ git clone https://github.com/lammps/lammps.git lammps_dir -$ cd lammps_dir -$ git checkout stable_2Aug2023 +```bash +git clone https://github.com/lammps/lammps.git lammps_dir +cd lammps_dir +git checkout stable_2Aug2023 ``` -Run patch_lammps.sh -``` -$ cd {path_to_SevenNet_root} -$ sh patch_lammps.sh {path_to_lammps_dir} +Run sevenn_patch_lammps +```bash +sevenn_patch_lammps {path_to_lammps_dir} ``` +Refer to `sevenn/pair_e3gnn/patch_lammps.sh` for the patch process. Build LAMMPS with cmake (example): - ``` $ cd {path_to_lammps_dir} $ mkdir build @@ -161,97 +159,32 @@ $ cmake ../cmake -DCMAKE_PREFIX_PATH=`python -c 'import torch;print(torch.utils. $ make -j4 ``` -If you prefer a manual patch, see the notes below. - -### Notes - -Note that the following commands will overwrite `comm_brick.cpp` and `comm_brick.h` in the original `LAMMPS`. While it does not affect the original functionality of `LAMMPS`, you may want to back up these files from the source if you're unsure. - -``` -cp {path_to_SevenNet}/pair_e3gnn/* path_to_lammps/src/ -``` - -If you have correctly installed CUDA-aware OpenMPI, the remaining process is identical to [`pair-nequip`](https://github.com/mir-group/pair_nequip). - -Please make the following modifications to lammps/cmake/CMakeLists.txt: -Change `set(CMAKE_CXX_STANDARD 11)` to `set(CMAKE_CXX_STANDARD 14)`. -Then append the following lines in the same file: - -``` -find_package(Torch REQUIRED) -set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} ${TORCH_CXX_FLAGS}") -target_link_libraries(lammps PUBLIC "${TORCH_LIBRARIES}") -``` - -You can check whether your OpenMPI is CUDA-aware by using `ompi_info` command: - -``` -$ ompi_info --parsable --all | grep mpi_built_with_cuda_support:value -mca:mpi:base:param:mpi_built_with_cuda_support:value:true -``` - ## Usage for LAMMPS ### To check installation -For serial MD, -``` -$ cd ${path_to_SevenNetetet}/example_inputs/md_serial_example -$ {lammps_binary} -in in.lmp - -###lammps outputs for 5 MD steps### - -$ grep PairE3GNN log.lammps -PairE3GNN using device : CUDA -``` - -For parallel MD -``` -$ cd ${path_to_SevenNet}/example_inputs/md_serial_example -$ mpirun -np {# of GPUs you want to use} {lammps_binary} -in in.lmp - -###lammps outputs for 5 MD steps### - -$ grep PairE3GNN log.lammps -PairE3GNNParallel using device : CUDA -PairE3GNNParallel cuda-aware mpi : True +```bash +{lmp_binary} -help | grep e3gnn ``` - -Example MD input scripts for `LAMMPS` can be found under `SevenNet/example_inputs`. If you've correctly installed `LAMMPS`, there are two additional pair styles available: `e3gnn` and `e3gnn/parallel`. - -In the `pair_coeff` of the lammps script, you need to provide the path of the trained models (either serial or parallel). For parallel models, you should specify how many segmented models will be used. +You will see `e3gnn` and `e3gnn/parallel` as pair_style. ### For serial model ``` pair_style e3gnn -pair_coeff * * {path to serial model} {chemical species} +pair_coeff * * {path to serial model} {space separated chemical species} ``` -Note that SevenNet-0 serial model is located in `{PATH TO SEVENNET}/pretrained_potentials/SevenNet_0/serial_model/deployed_serial.pt`. - ### For parallel model ``` pair_style e3gnn/parallel -pair_coeff * * {number of segmented parallel models} {space separated paths of segmented parallel models} {chemical species} +pair_coeff * * {number of segmented parallel models} {space separated paths of segmented parallel models} {space separated chemical species} ``` -Note that SevenNet-0 parallel model is located in `{PATH TO SEVENNET}/pretrained_potentials/SevenNet_0/parallel_model/deployed_parallel_*.pt`. -I recommend using variables to handle file paths for parallel models. - -``` -pair_style e3gnn/parallel -pair_coeff * * 5 ${pre}/deployed_parallel_0.pt ${pre}/deployed_parallel_1.pt ${pre}/deployed_parallel_2.pt ${pre}/deployed_parallel_3.pt ${pre}/deployed_parallel_4.pt {chemical species} -``` - -Now, you can execute this LAMMPS script with a prefix for parallel models - -``` -mpirun -np {# of GPUS to use} {LAMMPS_binary} -in {LAMMPS_script} -var pre {PATH TO SEVENNET}/pretrained_potentials/SevenNet_0/parallel_model/ -``` +Check [sevenn_get_model](#sevenn_get_model) for deploying lammps models from checkpoint for both serial and parallel. -Ideally, one GPU per MPI process is expected. If the available GPUs are fewer than the MPI processes, the simulation may run inefficiently. +**PLEASE NOTE:** One GPU per MPI process is expected. If the available GPUs are fewer than the MPI processes, the simulation may run inefficiently. **PLEASE NOTE:** Currently, the parallel version raises an error when there are no atoms in one of the subdomain cells. This issue can be addressed using the `processors` command and, more optimally, the `fix balance` command in LAMMPS. This will be patched in the future. From 5acb2434b4547db511672cacfba4e66e9a71abad Mon Sep 17 00:00:00 2001 From: YutackPark <111348843+YutackPark@users.noreply.github.com> Date: Thu, 18 Jul 2024 16:32:15 +0900 Subject: [PATCH 04/25] docs: prepare for pypi release --- README.md | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/README.md b/README.md index 896e60c..019beae 100644 --- a/README.md +++ b/README.md @@ -38,10 +38,10 @@ The installation and usage of SevenNet are split into two parts: training (handl You can find the installation guides for these packages from the [`PyTorch official`](https://pytorch.org/get-started/locally/), [`TorchGeometric docs`](https://pytorch-geometric.readthedocs.io/en/latest/install/installation.html) and [`pytorch_scatter`](https://github.com/rusty1s/pytorch_scatter). Remember that these packages have dependencies on your CUDA version. -``` -git clone https://github.com/MDIL-SNU/SevenNet.git -cd SevenNet -pip install . +**PLEASE NOTE:** You must install PyTorch, TorchGeometric, and pytorch_scatter before installing SevenNet. They are not marked as dependencies since they are coupled with the CUDA version. + +```bash +pip install sevenn ``` ## Usage From a352dfae21e1d67b750744f48c7227dd215ca681 Mon Sep 17 00:00:00 2001 From: YutackPark <111348843+YutackPark@users.noreply.github.com> Date: Thu, 18 Jul 2024 16:39:48 +0900 Subject: [PATCH 05/25] docs: typo --- README.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/README.md b/README.md index 019beae..c105130 100644 --- a/README.md +++ b/README.md @@ -92,7 +92,7 @@ Once you initiate training, `log.sevenn` will contain all parsed inputs from `in #### Multi-GPU training We support multi-GPU training features using PyTorch DDP (distributed data parallel). We use one process (CPU core) per GPU. ```bash -torchrun --standalone --nnodes={# of nodes} --nproc_per_node {# of GPUs} --no_python sevenn input.yaml -d +torchrun --standalone --nnodes {# of nodes} --nproc_per_node {# of GPUs} --no_python sevenn input.yaml -d ``` Please note that `batch_size` in input.yaml indicates `batch_size` per GPU. From 8c835095da1402f3fc4c7a2b281a3c4b095f697f Mon Sep 17 00:00:00 2001 From: YutackPark Date: Thu, 18 Jul 2024 17:05:22 +0900 Subject: [PATCH 06/25] bugfix: fix issues within precommit configs --- .pre-commit-config.yaml | 13 ------------- 1 file changed, 13 deletions(-) diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index c6034f2..3b372e6 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -22,7 +22,6 @@ repos: - id: check-shebang-scripts-are-executable - id: check-merge-conflict - id: check-vcs-permalinks - - id: check-yaml - id: debug-statements - id: destroyed-symlinks - id: detect-private-key @@ -48,18 +47,6 @@ repos: exclude: 'sevenn/pair_e3gnn/comm_brick.h' exclude: 'sevenn/pair_e3gnn/comm_brick.cpp' - - repo: https://github.com/psf/black - rev: 23.12.1 - hooks: - - id: black - exclude: 'sevenn/pair_e3gnn/pair_e3gnn.cpp' - exclude: 'sevenn/pair_e3gnn/pair_e3gnn.h' - exclude: 'sevenn/pair_e3gnn/pair_e3gnn_parallel.cpp' - exclude: 'sevenn/pair_e3gnn/pair_e3gnn_parallel.h' - exclude: 'sevenn/pair_e3gnn/comm_brick.h' - exclude: 'sevenn/pair_e3gnn/comm_brick.cpp' - args: ['--skip-string-normalization', '--line-length=79', '--preview'] - - repo: https://github.com/pycqa/flake8 rev: 6.1.0 hooks: From 12f4133a0ba5de9ed6416b1fd06f9bf730478769 Mon Sep 17 00:00:00 2001 From: YutackPark <111348843+YutackPark@users.noreply.github.com> Date: Thu, 18 Jul 2024 17:47:22 +0900 Subject: [PATCH 07/25] chore: docs # -> number --- README.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/README.md b/README.md index c105130..cde76bb 100644 --- a/README.md +++ b/README.md @@ -92,7 +92,7 @@ Once you initiate training, `log.sevenn` will contain all parsed inputs from `in #### Multi-GPU training We support multi-GPU training features using PyTorch DDP (distributed data parallel). We use one process (CPU core) per GPU. ```bash -torchrun --standalone --nnodes {# of nodes} --nproc_per_node {# of GPUs} --no_python sevenn input.yaml -d +torchrun --standalone --nnodes {number of nodes} --nproc_per_node {number of GPUs} --no_python sevenn input.yaml -d ``` Please note that `batch_size` in input.yaml indicates `batch_size` per GPU. From ba7ddc370dc4a38277acb8287f1d388f26e7353b Mon Sep 17 00:00:00 2001 From: YutackPark Date: Thu, 18 Jul 2024 21:45:42 +0900 Subject: [PATCH 08/25] bugfix: rename yaml --- sevenn/presets/sevennet-0.yaml | 80 ++++++++++++++++++++++++++++++++++ 1 file changed, 80 insertions(+) create mode 100644 sevenn/presets/sevennet-0.yaml diff --git a/sevenn/presets/sevennet-0.yaml b/sevenn/presets/sevennet-0.yaml new file mode 100644 index 0000000..9e53f7a --- /dev/null +++ b/sevenn/presets/sevennet-0.yaml @@ -0,0 +1,80 @@ +# SevenNet-0 +model: + chemical_species: 'auto' + cutoff: 5.0 + channel: 128 + is_parity: False + lmax: 2 + num_convolution_layer: 5 + irreps_manual: + - "128x0e" + - "128x0e+64x1e+32x2e" + - "128x0e+64x1e+32x2e" + - "128x0e+64x1e+32x2e" + - "128x0e+64x1e+32x2e" + - "128x0e" + + weight_nn_hidden_neurons: [64, 64] + radial_basis: + radial_basis_name: 'bessel' + bessel_basis_num: 8 + cutoff_function: + cutoff_function_name: 'XPLOR' + cutoff_on: 4.5 + + act_gate: {'e': 'silu', 'o': 'tanh'} + act_scalar: {'e': 'silu', 'o': 'tanh'} + + conv_denominator: 'avg_num_neigh' + train_shift_scale: False + train_denominator: False + self_connection_type: 'linear' +train: + train_shuffle: False + random_seed: 1 + is_train_stress : True + epoch: 600 + + loss: 'Huber' + loss_param: + delta: 0.01 + + optimizer: 'adam' + optim_param: + lr: 0.01 + scheduler: 'linearlr' + scheduler_param: + start_factor: 1.0 + total_iters: 600 + end_factor: 0.0001 + + force_loss_weight : 1.00 + stress_loss_weight: 0.01 + + error_record: + - ['Energy', 'RMSE'] + - ['Force', 'RMSE'] + - ['Stress', 'RMSE'] + - ['Energy', 'MAE'] + - ['Force', 'MAE'] + - ['Stress', 'MAE'] + - ['Energy', 'Loss'] + - ['Force', 'Loss'] + - ['Stress', 'Loss'] + - ['TotalLoss', 'None'] + + per_epoch: 10 + # continue: + # checkpoint: './checkpoint_last.pth' + # reset_optimizer: False + # reset_scheduler: False +data: + data_shuffle: False + batch_size: 4 + scale: 'per_atom_energy_std' + shift: 'elemwise_reference_energies' + + data_format: 'ase' + save_by_train_valid: False + load_dataset_path: ["path_to_MPtrj_total.sevenn_data"] + load_validset_path: ["validaset.sevenn_data"] From 2bb02ac26fb38b98834c26970d0ab527387c604a Mon Sep 17 00:00:00 2001 From: YutackPark Date: Thu, 18 Jul 2024 21:46:22 +0900 Subject: [PATCH 09/25] chore --- sevenn/presets/sevennet_0.yaml | 80 ---------------------------------- 1 file changed, 80 deletions(-) delete mode 100644 sevenn/presets/sevennet_0.yaml diff --git a/sevenn/presets/sevennet_0.yaml b/sevenn/presets/sevennet_0.yaml deleted file mode 100644 index 9e53f7a..0000000 --- a/sevenn/presets/sevennet_0.yaml +++ /dev/null @@ -1,80 +0,0 @@ -# SevenNet-0 -model: - chemical_species: 'auto' - cutoff: 5.0 - channel: 128 - is_parity: False - lmax: 2 - num_convolution_layer: 5 - irreps_manual: - - "128x0e" - - "128x0e+64x1e+32x2e" - - "128x0e+64x1e+32x2e" - - "128x0e+64x1e+32x2e" - - "128x0e+64x1e+32x2e" - - "128x0e" - - weight_nn_hidden_neurons: [64, 64] - radial_basis: - radial_basis_name: 'bessel' - bessel_basis_num: 8 - cutoff_function: - cutoff_function_name: 'XPLOR' - cutoff_on: 4.5 - - act_gate: {'e': 'silu', 'o': 'tanh'} - act_scalar: {'e': 'silu', 'o': 'tanh'} - - conv_denominator: 'avg_num_neigh' - train_shift_scale: False - train_denominator: False - self_connection_type: 'linear' -train: - train_shuffle: False - random_seed: 1 - is_train_stress : True - epoch: 600 - - loss: 'Huber' - loss_param: - delta: 0.01 - - optimizer: 'adam' - optim_param: - lr: 0.01 - scheduler: 'linearlr' - scheduler_param: - start_factor: 1.0 - total_iters: 600 - end_factor: 0.0001 - - force_loss_weight : 1.00 - stress_loss_weight: 0.01 - - error_record: - - ['Energy', 'RMSE'] - - ['Force', 'RMSE'] - - ['Stress', 'RMSE'] - - ['Energy', 'MAE'] - - ['Force', 'MAE'] - - ['Stress', 'MAE'] - - ['Energy', 'Loss'] - - ['Force', 'Loss'] - - ['Stress', 'Loss'] - - ['TotalLoss', 'None'] - - per_epoch: 10 - # continue: - # checkpoint: './checkpoint_last.pth' - # reset_optimizer: False - # reset_scheduler: False -data: - data_shuffle: False - batch_size: 4 - scale: 'per_atom_energy_std' - shift: 'elemwise_reference_energies' - - data_format: 'ase' - save_by_train_valid: False - load_dataset_path: ["path_to_MPtrj_total.sevenn_data"] - load_validset_path: ["validaset.sevenn_data"] From 32445958326b0c301a9060d9b2264cbf369c08b4 Mon Sep 17 00:00:00 2001 From: YutackPark Date: Mon, 22 Jul 2024 16:58:51 +0900 Subject: [PATCH 10/25] docs: preset explanation --- sevenn/presets/base.yaml | 12 ++++++------ sevenn/presets/fine_tune.yaml | 9 +++++---- 2 files changed, 11 insertions(+), 10 deletions(-) diff --git a/sevenn/presets/base.yaml b/sevenn/presets/base.yaml index d85d2f8..e758bc5 100644 --- a/sevenn/presets/base.yaml +++ b/sevenn/presets/base.yaml @@ -1,5 +1,5 @@ # Example input.yaml for training SevenNet. -# '*' signifies default. You can check log.sevenn. +# '*' signifies default. You can check log.sevenn for defaults. model: chemical_species: 'Auto' # Chemical elements present in the dataset, guess them from load_dataset data if 'auto' @@ -19,10 +19,10 @@ model: weight_nn_hidden_neurons: [64, 64] # Hidden neurons in convolution weight neural network radial_basis: # Function and its parameters to encode radial distance radial_basis_name: 'bessel' # Only 'bessel' is currently supported - bessel_basis_num: 8 + bessel_basis_num: 8 cutoff_function: # Envelop function, multiplied to radial_basis functions to init edge featrues cutoff_function_name: 'poly_cut' # {'poly_cut' and 'poly_cut_p_value'} or {'XPLOR' and 'cutoff_on'} - poly_cut_p_value: 6 + poly_cut_p_value: 6 act_gate: {'e': 'silu', 'o': 'tanh'} # Equivalent to 'nonlinearity_gates' in nequip act_scalar: {'e': 'silu', 'o': 'tanh'} # Equivalent to 'nonlinearity_scalars' in nequip @@ -83,9 +83,9 @@ data: # ase.io.read readable data files or structure_list or .sevenn_data files can be used as dataset. # .sevenn_data is preprocessed data set has edges connected (can be obtained by using sevenn_graph_build or by save_** options below) - data_format: 'ase' # Default is 'ase' - data_format_args: # Paramaters, will be passed to ase.io.read see: https://wiki.fysik.dtu.dk/ase/ase/io/io.html - index: ':' + data_format: 'ase' # One of 'ase', 'structure_list' (.sevenn_data is always readable) + data_format_args: # if `data_format` is `ase`, args will be passed to `ase.io.read` + index: ':' # see `https://wiki.fysik.dtu.dk/ase/ase/io/io.html` for more valid arguments # If only load_dataset_path is provided, train/valid set is automatically decided by splitting dataset by divide ratio # If both load_dataset_path & load_validset_path is provided, use load_dataset_path as training set. diff --git a/sevenn/presets/fine_tune.yaml b/sevenn/presets/fine_tune.yaml index 573a50d..7fc6faf 100644 --- a/sevenn/presets/fine_tune.yaml +++ b/sevenn/presets/fine_tune.yaml @@ -1,4 +1,5 @@ # Example input.yaml for fine-tuning sevennet-0 +# '*' signifies default. You can check log.sevenn for defaults. model: # model keys should be consistent except for train_* keys chemical_species: 'Auto' @@ -59,7 +60,7 @@ train: # Customizable reset_epoch: True checkpoint: 'SevenNet-0_11July2024' # Set True to use shift, scale, and avg_num_neigh from checkpoint (highly recommended) - use_statistic_values_of_checkpoint: True + use_statistic_values_of_checkpoint: True data: # Customizable batch_size: 4 @@ -67,9 +68,9 @@ data: # Customizable # ase.io.read readable data files or structure_list or .sevenn_data files can be used as dataset. # .sevenn_data is preprocessed data set has edges connected (can be obtained by using sevenn_graph_build or by save_** options below) - data_format: 'ase' # Default is 'ase' - data_format_args: # Paramaters, will be passed to ase.io.read see: https://wiki.fysik.dtu.dk/ase/ase/io/io.html - index: ':' + data_format: 'ase' # One of 'ase', 'structure_list' (.sevenn_data is always readable) + data_format_args: # if `data_format` is `ase`, args will be passed to `ase.io.read` + index: ':' # see `https://wiki.fysik.dtu.dk/ase/ase/io/io.html` for more valid arguments # If only load_dataset_path is provided, train/valid set is automatically decided by splitting dataset by divide ratio # If both load_dataset_path & load_validset_path is provided, use load_dataset_path as training set. From a5d4c6b6be1a0433208048b4178b3bf3d9491d92 Mon Sep 17 00:00:00 2001 From: YutackPark Date: Mon, 22 Jul 2024 17:53:49 +0900 Subject: [PATCH 11/25] refactor: remove torch_scatter dependency --- setup.cfg | 2 +- sevenn/model_build.py | 41 +++++++++------------ sevenn/nn/convolution.py | 35 ++++++++++++++++-- sevenn/nn/force_output.py | 43 ----------------------- sevenn/nn/linear.py | 28 +++++++-------- sevenn/pair_e3gnn/pair_e3gnn_parallel.cpp | 18 ++++------ sevenn/train/dataset.py | 14 +++++--- 7 files changed, 79 insertions(+), 102 deletions(-) diff --git a/setup.cfg b/setup.cfg index 799269e..feeb5aa 100644 --- a/setup.cfg +++ b/setup.cfg @@ -10,5 +10,5 @@ include_trailing_comma=True force_grid_wrap=0 use_parentheses=True line_length=80 -known_third_party=ase,braceexpand,e3nn,numpy,setuptools,sklearn,torch,torch_geometric,torch_scatter,tqdm,yaml +known_third_party=ase,braceexpand,e3nn,numpy,setuptools,sklearn,torch,torch_geometric,tqdm,yaml known_first_party= diff --git a/sevenn/model_build.py b/sevenn/model_build.py index badfdf5..c0f8c7e 100644 --- a/sevenn/model_build.py +++ b/sevenn/model_build.py @@ -16,7 +16,7 @@ XPLORCutoff, ) from sevenn.nn.equivariant_gate import EquivariantGate -from sevenn.nn.force_output import ForceOutputFromEdge, ForceStressOutput +from sevenn.nn.force_output import ForceStressOutput from sevenn.nn.linear import AtomReduce, FCN_e3nn, IrrepsLinear from sevenn.nn.node_embedding import OnehotEmbedding from sevenn.nn.scale import Rescale, SpeciesWiseRescale @@ -32,7 +32,7 @@ 'ignore', message=( "The TorchScript type system doesn't " - "support instance-level annotations" + 'support instance-level annotations' ), ) @@ -61,8 +61,7 @@ def init_radial_basis(config): if radial_basis_dct[KEY.RADIAL_BASIS_NAME] == 'bessel': basis_function = BesselBasis(**param) return basis_function, basis_function.num_basis - - raise RuntimeError('something went very wrong...') + raise RuntimeError('something went wrong...') def init_cutoff_function(config): @@ -80,10 +79,6 @@ def init_cutoff_function(config): # TODO: it gets bigger and bigger. refactor it def build_E3_equivariant_model(config: dict, parallel=False): - """ - IDENTICAL to nequip model - atom embedding is not part of model - """ data_key_weight_input = KEY.EDGE_EMBEDDING # default # parameter initialization @@ -103,6 +98,7 @@ def build_E3_equivariant_model(config: dict, parallel=False): irreps_spherical_harm = Irreps.spherical_harmonics( lmax_edge, -1 if is_parity else 1 ) + layers_list = [] if parallel: layers_list = [OrderedDict() for _ in range(num_convolution_layer)] layers_idx = 0 @@ -150,11 +146,10 @@ def build_E3_equivariant_model(config: dict, parallel=False): lmax_edge, -1 if is_parity else 1, normalize=_normalize_sph ), ) - if not parallel: - layers.update({ - # simple edge preprocessor module with no param - 'edge_preprocess': EdgePreprocess(is_stress=True), - }) + layers.update({ + # simple edge preprocessor module with no param + 'edge_preprocess': EdgePreprocess(is_stress=True), + }) layers.update({ # 'Not' simple edge embedding module @@ -375,22 +370,18 @@ def build_E3_equivariant_model(config: dict, parallel=False): constant=1.0, ), }) - if not parallel: - fso = ForceStressOutput( - data_key_energy=KEY.PRED_TOTAL_ENERGY, - data_key_force=KEY.PRED_FORCE, - data_key_stress=KEY.PRED_STRESS, - ) - fof = ForceOutputFromEdge( - data_key_energy=KEY.PRED_TOTAL_ENERGY, - data_key_force=KEY.PRED_FORCE, - ) - gradient_module = fso if not parallel else fof - layers.update({'force_output': gradient_module}) + gradient_module = ForceStressOutput( + data_key_energy=KEY.PRED_TOTAL_ENERGY, + data_key_force=KEY.PRED_FORCE, + data_key_stress=KEY.PRED_STRESS, + ) + layers.update({'force_output': gradient_module}) # output extraction part type_map = config[KEY.TYPE_MAP] if parallel: + del layers_list[0]['edge_preprocess'] # done in LAMMPS + del layers_list[-1]['force_output'] # done in LAMMPS return [AtomGraphSequential(v, cutoff, type_map) for v in layers_list] else: return AtomGraphSequential(layers, cutoff, type_map) diff --git a/sevenn/nn/convolution.py b/sevenn/nn/convolution.py index 1235605..709eaf2 100644 --- a/sevenn/nn/convolution.py +++ b/sevenn/nn/convolution.py @@ -5,13 +5,44 @@ from e3nn.nn import FullyConnectedNet from e3nn.o3 import Irreps, TensorProduct from e3nn.util.jit import compile_mode -from torch_scatter import scatter import sevenn._keys as KEY from sevenn._const import AtomGraphDataType from sevenn.nn.activation import ShiftedSoftPlus +def _broadcast( + src: torch.Tensor, + other: torch.Tensor, + dim: int +): + if dim < 0: + dim = other.dim() + dim + if src.dim() == 1: + for _ in range(0, dim): + src = src.unsqueeze(0) + for _ in range(src.dim(), other.dim()): + src = src.unsqueeze(-1) + src = src.expand_as(other) + return src + + +def message_gather( + node_features: torch.Tensor, + edge_dst: torch.Tensor, + message: torch.Tensor +): + index = _broadcast(edge_dst, message, 0) + out_shape = [len(node_features)] + list(message.shape[1:]) + out = torch.zeros( + out_shape, + dtype=node_features.dtype, + device=node_features.device + ) + out.scatter_reduce_(0, index, message, reduce='sum') + return out + + @compile_mode('script') class IrrepsConvolution(nn.Module): """ @@ -85,7 +116,7 @@ def forward(self, data: AtomGraphDataType) -> AtomGraphDataType: message = self.convolution(x[edge_src], data[self.key_filter], weight) - x = scatter(message, edge_dst, dim=0, dim_size=len(x)) + x = message_gather(x, edge_dst, message) x = x.div(self.denominator) if self.is_parallel: # NLOCAL is # of atoms in system at 'CPU' diff --git a/sevenn/nn/force_output.py b/sevenn/nn/force_output.py index 4f1ff3d..184a907 100644 --- a/sevenn/nn/force_output.py +++ b/sevenn/nn/force_output.py @@ -1,54 +1,11 @@ import torch import torch.nn as nn from e3nn.util.jit import compile_mode -from torch_scatter import scatter import sevenn._keys as KEY from sevenn._const import AtomGraphDataType -@compile_mode('script') -class ForceOutputFromEdge(nn.Module): - """ - works when edge_vec.requires_grad_ is True - """ - - def __init__( - self, - data_key_edge_vec: str = KEY.EDGE_VEC, - data_key_edge_idx: str = KEY.EDGE_IDX, - data_key_energy: str = KEY.SCALED_ENERGY, - data_key_force: str = KEY.SCALED_FORCE, - data_key_num_atoms: str = KEY.NUM_ATOMS, - ): - super().__init__() - self.key_edge_vec = data_key_edge_vec - self.key_energy = data_key_energy - self.key_force = data_key_force - self.key_edge_idx = data_key_edge_idx - self.key_num_atoms = data_key_num_atoms - - def forward(self, data: AtomGraphDataType) -> AtomGraphDataType: - tot_num = torch.sum(data[self.key_num_atoms]) - edge_idx = data[self.key_edge_idx] - - edge_vec_tensor = [data[self.key_edge_vec]] - energy = [(data[self.key_energy]).sum()] - - dE_dr = torch.autograd.grad( - energy, edge_vec_tensor, create_graph=self.training - )[0] - - if dE_dr is not None: - force = torch.zeros(tot_num, 3) - force = scatter(dE_dr, edge_idx[0], dim=0, dim_size=tot_num) - force -= scatter(dE_dr, edge_idx[1], dim=0, dim_size=tot_num) - - data[self.key_force] = force - - return data - - @compile_mode('script') class ForceOutput(nn.Module): """ diff --git a/sevenn/nn/linear.py b/sevenn/nn/linear.py index db3e7e6..015e147 100644 --- a/sevenn/nn/linear.py +++ b/sevenn/nn/linear.py @@ -1,9 +1,10 @@ +from typing import Callable, List, Optional + import torch import torch.nn as nn from e3nn.nn import FullyConnectedNet from e3nn.o3 import Irreps, Linear from e3nn.util.jit import compile_mode -from torch_scatter import scatter import sevenn._keys as KEY from sevenn._const import AtomGraphDataType @@ -20,7 +21,7 @@ def __init__( irreps_in: Irreps, irreps_out: Irreps, data_key_in: str, - data_key_out: str = None, + data_key_out: Optional[str] = None, **e3nn_linear_params, ): super().__init__() @@ -48,7 +49,7 @@ def __init__( self, data_key_in: str, data_key_out: str, - reduce='sum', + reduce: str = 'sum', constant: float = 1.0, ): super().__init__() @@ -63,16 +64,13 @@ def __init__( def forward(self, data: AtomGraphDataType) -> AtomGraphDataType: if self._is_batch_data: - data[self.key_output] = ( - scatter( - data[self.key_input], - data[KEY.BATCH], - dim=0, - reduce=self.reduce, - ) - * self.constant + src = data[self.key_input].squeeze(1) + size = int(data[KEY.BATCH].max()) + 1 + output = torch.zeros( + (size), dtype=src.dtype, device=src.device, ) - data[self.key_output] = data[self.key_output].squeeze(1) + output.scatter_reduce_(0, data[KEY.BATCH], src, reduce='sum') + data[self.key_output] = output * self.constant else: data[self.key_output] = ( torch.sum(data[self.key_input]) * self.constant @@ -91,10 +89,10 @@ def __init__( self, irreps_in: Irreps, # confirm it is scalar & input size dim_out: int, - hidden_neurons, - activation, + hidden_neurons: List[int], + activation: Callable, data_key_in: str, - data_key_out: str = None, + data_key_out: Optional[str] = None, **e3nn_params, ): super().__init__() diff --git a/sevenn/pair_e3gnn/pair_e3gnn_parallel.cpp b/sevenn/pair_e3gnn/pair_e3gnn_parallel.cpp index c1940e5..5a404a9 100644 --- a/sevenn/pair_e3gnn/pair_e3gnn_parallel.cpp +++ b/sevenn/pair_e3gnn/pair_e3gnn_parallel.cpp @@ -111,8 +111,8 @@ PairE3GNNParallel::PairE3GNNParallel(LAMMPS *lmp) : Pair(lmp) { #else use_cuda_mpi = false; #endif - //use_cuda_mpi = use_gpu && use_cuda_mpi; - //if (use_cuda_mpi) { + // use_cuda_mpi = use_gpu && use_cuda_mpi; + // if (use_cuda_mpi) { if (use_gpu) { device = get_cuda_device(); device_name = "CUDA"; @@ -125,7 +125,8 @@ PairE3GNNParallel::PairE3GNNParallel(LAMMPS *lmp) : Pair(lmp) { if (use_gpu && !use_cuda_mpi) { // GPU device + cuda-'NOT'aware mpi combination. is not supported yet device_comm = torch::kCPU; - fprintf(lmp->screen, "cuda-aware mpi not found, communicate via host device\n"); + fprintf(lmp->screen, + "cuda-aware mpi not found, communicate via host device\n"); } else { device_comm = device; } @@ -138,7 +139,8 @@ PairE3GNNParallel::PairE3GNNParallel(LAMMPS *lmp) : Pair(lmp) { if (use_gpu && !use_cuda_mpi) { // GPU device + cuda-'NOT'aware mpi combination. is not supported yet device_comm = torch::kCPU; - fprintf(lmp->logfile, "cuda-aware mpi not found, communicate via host device\n"); + fprintf(lmp->logfile, + "cuda-aware mpi not found, communicate via host device\n"); } else { device_comm = device; } @@ -488,14 +490,11 @@ void PairE3GNNParallel::compute(int eflag, int vflag) { std::cout << world_rank << " Used/GraphSize: " << Mused / graph_size << "\n" << std::endl; } - // TODO: atomic energy things? eng_vdwl += energy_tensor.item(); // accumulate energy dE_dr = dE_dr.to(torch::kCPU); torch::Tensor force_tensor = torch::zeros({graph_indexer, 3}); - // TODO:where I can find torch_scatter cpp version? I heard this version(torch - // defaults) is slower. force_tensor.scatter_( 0, edge_idx_src_tensor.repeat_interleave(3).view({nedges, 3}), dE_dr, "add"); @@ -522,7 +521,6 @@ void PairE3GNNParallel::compute(int eflag, int vflag) { } } - // clean up comm preprocess variables comm_preprocess_done = false; for (int i = 0; i < 6; i++) { @@ -671,9 +669,7 @@ void PairE3GNNParallel::init_style() { neighbor->add_request(this, NeighConst::REQ_FULL); } -double PairE3GNNParallel::init_one(int i, int j) { - return cutoff; -} +double PairE3GNNParallel::init_one(int i, int j) { return cutoff; } void PairE3GNNParallel::comm_preprocess() { assert(!comm_preprocess_done); diff --git a/sevenn/train/dataset.py b/sevenn/train/dataset.py index 51562d0..c404f7d 100644 --- a/sevenn/train/dataset.py +++ b/sevenn/train/dataset.py @@ -7,7 +7,6 @@ import torch from ase.data import chemical_symbols from sklearn.linear_model import Ridge -from torch_scatter import scatter import sevenn._keys as KEY import sevenn.util as util @@ -329,11 +328,16 @@ def get_species_wise_force_rms(self, num_chem_species): atomx = torch.concat([d[self.DATA_KEY_X] for d in data_list]) force = torch.concat([d[self.DATA_KEY_FORCE] for d in data_list]) - rms = torch.sqrt( - scatter(force.square(), atomx, dim=0, reduce='mean').mean(dim=1) + index = atomx.repeat_interleave(3, 0).reshape(force.shape) + rms = torch.zeros( + (num_chem_species, 3), + dtype=force.dtype, + device=force.device + ) + rms.scatter_reduce_( + 0, index, force.square(), + reduce='mean', include_self=False ) - if len(rms) < num_chem_species: - rms = torch.cat([rms, torch.zeros(num_chem_species - len(rms))]) return rms def get_avg_num_neigh(self): From 49034b77c5c1e2d7646f6becc641d4e4bca2bf11 Mon Sep 17 00:00:00 2001 From: YutackPark Date: Mon, 22 Jul 2024 17:56:12 +0900 Subject: [PATCH 12/25] fix: add torch_geometric dependency --- pyproject.toml | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 217986b..8c5af5a 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -35,16 +35,17 @@ requires = [ "pyyaml", "e3nn", "tqdm", + "torch_geometric" ] build-backend = "setuptools.build_meta" [tool.setuptools.package-data] sevenn = [ - "logo_ascii", - "pair_e3gnn/*.cpp", - "pair_e3gnn/*.h", + "logo_ascii", + "pair_e3gnn/*.cpp", + "pair_e3gnn/*.h", "presets/*.yaml", - "pretrained_potentials/SevenNet_0__11July2024/checkpoint_sevennet_0.pth", + "pretrained_potentials/SevenNet_0__11July2024/checkpoint_sevennet_0.pth", "pretrained_potentials/SevenNet_0__22May2024/checkpoint_sevennet_0.pth" ] From 9559fa12db418a6e1a326ed77db8831ef0b04b4b Mon Sep 17 00:00:00 2001 From: YutackPark Date: Tue, 23 Jul 2024 12:35:54 +0900 Subject: [PATCH 13/25] fix: support torch ver>2.0 with lammps, adopt deprecation warnings from libtorch, correct dependencies --- pyproject.toml | 22 +++++++++++++--------- sevenn/main/sevenn_patch_lammps.py | 22 ++++++++++++++-------- sevenn/pair_e3gnn/pair_e3gnn.cpp | 2 +- sevenn/pair_e3gnn/pair_e3gnn_parallel.cpp | 17 ++++++++--------- sevenn/pair_e3gnn/patch_lammps.sh | 10 +++++++--- 5 files changed, 43 insertions(+), 30 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 8c5af5a..e1ee9ec 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -14,6 +14,17 @@ classifiers = [ "License :: OSI Approved :: GNU General Public License v3 (GPLv3)", "Operating System :: POSIX :: Linux", ] +dependencies = [ + "ase", + "braceexpand", + "pyyaml", + "e3nn", + "tqdm", + "scikit-learn", + "torch_geometric", + "numpy<2.0", +] + [project.scripts] sevenn = "sevenn.main.sevenn:main" @@ -28,22 +39,15 @@ Homepage = "https://github.com/MDIL-SNU/SevenNet" Issues = "https://github.com/MDIL-SNU/SevenNet/issues" [build-system] -requires = [ - "setuptools>=61.0", - "ase", - "braceexpand", - "pyyaml", - "e3nn", - "tqdm", - "torch_geometric" -] build-backend = "setuptools.build_meta" +requires = ["setuptools>=61.0"] [tool.setuptools.package-data] sevenn = [ "logo_ascii", "pair_e3gnn/*.cpp", "pair_e3gnn/*.h", + "pair_e3gnn/patch_lammps.sh", "presets/*.yaml", "pretrained_potentials/SevenNet_0__11July2024/checkpoint_sevennet_0.pth", "pretrained_potentials/SevenNet_0__22May2024/checkpoint_sevennet_0.pth" diff --git a/sevenn/main/sevenn_patch_lammps.py b/sevenn/main/sevenn_patch_lammps.py index 8746a0d..cc1e053 100644 --- a/sevenn/main/sevenn_patch_lammps.py +++ b/sevenn/main/sevenn_patch_lammps.py @@ -1,15 +1,16 @@ -import os import argparse +import os import subprocess -# python wrapper of patch_lammps.sh script +from torch import __version__ + +from sevenn._const import SEVENN_VERSION +# python wrapper of patch_lammps.sh script # importlib.resources is correct way to do these things -# but it changes so frequently to use +# but it changes so frequently to use pair_e3gnn_dir = os.path.abspath(f'{os.path.dirname(__file__)}/../pair_e3gnn') -from sevenn._const import SEVENN_VERSION - description = ( f'sevenn version={SEVENN_VERSION}, patch LAMMPS for pair_e3gnn styles' ) @@ -17,15 +18,20 @@ def main(args=None): lammps_dir = cmd_parse_main(args) - script = f"{pair_e3gnn_dir}/patch_lammps.sh" - cmd = f"{script} {lammps_dir}" + cxx_standard = '17' if __version__.startswith('2') else '14' + if cxx_standard == '17': + print('Torch version >= 2.0 detacted, use CXX STANDARD 17') + else: + print('Torch version < 2.0 detacted, use CXX STANDARD 14') + script = f'{pair_e3gnn_dir}/patch_lammps.sh' + cmd = f'{script} {lammps_dir} {cxx_standard}' res = subprocess.run(cmd.split()) return res.returncode # is it meaningless? def cmd_parse_main(args=None): ag = argparse.ArgumentParser(description=description) - ag.add_argument('lammps_dir', help="Path to LAMMPS source", type=str) + ag.add_argument('lammps_dir', help='Path to LAMMPS source', type=str) args = ag.parse_args() return args.lammps_dir diff --git a/sevenn/pair_e3gnn/pair_e3gnn.cpp b/sevenn/pair_e3gnn/pair_e3gnn.cpp index c1f6038..d1ba87c 100644 --- a/sevenn/pair_e3gnn/pair_e3gnn.cpp +++ b/sevenn/pair_e3gnn/pair_e3gnn.cpp @@ -129,7 +129,7 @@ void PairE3GNN::compute(int eflag, int vflag) { torch::Tensor inp_pos = torch::zeros({nlocal, 3}); torch::Tensor inp_cell_volume = - torch::dot(inp_cell[0], torch::cross(inp_cell[1], inp_cell[2])); + torch::dot(inp_cell[0], torch::cross(inp_cell[1], inp_cell[2], 0)); float pbc_shift_tmp[nedges_upper_bound][3]; diff --git a/sevenn/pair_e3gnn/pair_e3gnn_parallel.cpp b/sevenn/pair_e3gnn/pair_e3gnn_parallel.cpp index 5a404a9..5a3b1b3 100644 --- a/sevenn/pair_e3gnn/pair_e3gnn_parallel.cpp +++ b/sevenn/pair_e3gnn/pair_e3gnn_parallel.cpp @@ -123,7 +123,6 @@ PairE3GNNParallel::PairE3GNNParallel(LAMMPS *lmp) : Pair(lmp) { if (lmp->screen) { if (use_gpu && !use_cuda_mpi) { - // GPU device + cuda-'NOT'aware mpi combination. is not supported yet device_comm = torch::kCPU; fprintf(lmp->screen, "cuda-aware mpi not found, communicate via host device\n"); @@ -137,7 +136,6 @@ PairE3GNNParallel::PairE3GNNParallel(LAMMPS *lmp) : Pair(lmp) { } if (lmp->logfile) { if (use_gpu && !use_cuda_mpi) { - // GPU device + cuda-'NOT'aware mpi combination. is not supported yet device_comm = torch::kCPU; fprintf(lmp->logfile, "cuda-aware mpi not found, communicate via host device\n"); @@ -212,12 +210,12 @@ void PairE3GNNParallel::warning_pressure() { fprintf( lmp->screen, "WARNING: PairE3GNNParallel does not support pressure calculation. " - "Pressure on log is WRONG. Use serial version if you needed\n"); + "Pressure on log is wrong. Use serial version if you needed\n"); if (lmp->logfile) fprintf( lmp->logfile, "WARNING: PairE3GNNParallel does not support pressure calculation. " - "Pressure on log is WRONG. Use serial version if you needed\n"); + "Pressure on log is wrong. Use serial version if you needed\n"); already_did = true; } } @@ -495,12 +493,12 @@ void PairE3GNNParallel::compute(int eflag, int vflag) { dE_dr = dE_dr.to(torch::kCPU); torch::Tensor force_tensor = torch::zeros({graph_indexer, 3}); - force_tensor.scatter_( + force_tensor.scatter_reduce_( 0, edge_idx_src_tensor.repeat_interleave(3).view({nedges, 3}), dE_dr, - "add"); - force_tensor.scatter_( + "sum"); + force_tensor.scatter_reduce_( 0, edge_idx_dst_tensor.repeat_interleave(3).view({nedges, 3}), - torch::neg(dE_dr), "add"); + torch::neg(dE_dr), "sum"); auto forces = force_tensor.accessor(); @@ -639,7 +637,8 @@ void PairE3GNNParallel::coeff(int narg, char **arg) { } } if (!found_flag) { - error->all(FLERR, "Unknown chemical specie is given"); + error->all(FLERR, "Unknown chemical specie is given or the number of " + "potential files is not consistent"); } } diff --git a/sevenn/pair_e3gnn/patch_lammps.sh b/sevenn/pair_e3gnn/patch_lammps.sh index 68a060d..9acd726 100755 --- a/sevenn/pair_e3gnn/patch_lammps.sh +++ b/sevenn/pair_e3gnn/patch_lammps.sh @@ -1,6 +1,7 @@ #!/bin/bash lammps_root=$1 +cxx_standard=$2 SCRIPT_DIR=$(dirname "${BASH_SOURCE[0]}") echo "Usage: sh patch_lammps.sh {lammps_root}" @@ -24,6 +25,10 @@ fi if [ -f "$lammps_root/src/pair_e3gnn.cpp" ]; then echo "Seems like given LAMMPS is already patched." + echo "Example build commends, under LAMMPS root" + echo " mkdir build; cd build" + echo " cmake ../cmake -DCMAKE_PREFIX_PATH=`python -c 'import torch;print(torch.utils.cmake_prefix_path)'`" + echo " make -j 4" exit 0 fi @@ -56,7 +61,7 @@ cp $SCRIPT_DIR/*.h $lammps_root/src/ cp $lammps_root/cmake/CMakeLists.txt $backup_dir/CMakeLists.txt # 4. Patch cmake/CMakeLists.txt -sed -i "s/set(CMAKE_CXX_STANDARD 11)/set(CMAKE_CXX_STANDARD 14)/" $lammps_root/cmake/CMakeLists.txt +sed -i "s/set(CMAKE_CXX_STANDARD 11)/set(CMAKE_CXX_STANDARD $cxx_standard)/" $lammps_root/cmake/CMakeLists.txt cat >> $lammps_root/cmake/CMakeLists.txt << "EOF2" find_package(Torch REQUIRED) @@ -79,7 +84,7 @@ fi echo "Changes made:" echo " - Original LAMMPS files (src/comm_brick.*, cmake/CMakeList.txt) are in {lammps_root}/_backups" echo " - Copied contents of pair_e3gnn to $lammps_root/src/" -echo " - Patched CMakeLists.txt: include LibTorch, CXX_STANDARD 14" +echo " - Patched CMakeLists.txt: include LibTorch, CXX_STANDARD $cxx_standard" # ?. Provide example cmake command to the user echo "Example build commends, under LAMMPS root" @@ -88,4 +93,3 @@ echo " cmake ../cmake -DCMAKE_PREFIX_PATH=`python -c 'import torch;print(torch. echo " make -j 4" exit 0 - From a74e5ac621cd2ecf9e137d20a1d1fa41498e14af Mon Sep 17 00:00:00 2001 From: YutackPark <111348843+YutackPark@users.noreply.github.com> Date: Tue, 23 Jul 2024 17:06:37 +0900 Subject: [PATCH 14/25] docs: readme.md --- README.md | 32 +++++++++++++++++++------------- 1 file changed, 19 insertions(+), 13 deletions(-) diff --git a/README.md b/README.md index cde76bb..5bea03f 100644 --- a/README.md +++ b/README.md @@ -32,14 +32,21 @@ The installation and usage of SevenNet are split into two parts: training (handl ## Installation * Python >= 3.8 -* PyTorch >= 1.11 -* [`TorchGeometric`](https://pytorch-geometric.readthedocs.io/en/latest/install/installation.html) -* [`pytorch_scatter`](https://github.com/rusty1s/pytorch_scatter) +* PyTorch >= 1.12.0 -You can find the installation guides for these packages from the [`PyTorch official`](https://pytorch.org/get-started/locally/), [`TorchGeometric docs`](https://pytorch-geometric.readthedocs.io/en/latest/install/installation.html) and [`pytorch_scatter`](https://github.com/rusty1s/pytorch_scatter). Remember that these packages have dependencies on your CUDA version. +Please install PyTorch from [`PyTorch official`](https://pytorch.org/get-started/locally/) before installing the SevenNet. +Note that for SevenNet, `torchvision` and `torchaudio` are redundant. You can omit them from the command provided in the installation guide. -**PLEASE NOTE:** You must install PyTorch, TorchGeometric, and pytorch_scatter before installing SevenNet. They are not marked as dependencies since they are coupled with the CUDA version. +Matching PyTorch + CUDA versions may cause problems, especially when compiling SevenNet with LAMMPS. Here are the versions we've been using internally, without problems. +- PyTorch/2.2.2 + CUDA/12.1.0 +- PyTorch/1.13.1 + CUDA/12.1.0 +- PyTorch/1.12.0 + CUDA/11.6.2 +Using the newer versions of CUDA with PyTorch is usually not an issue. For example, we were able to compile and use `PyTorch/1.13.1+cu117` with `CUDA/12.1.0`. + +**PLEASE NOTE:** You must install PyTorch before installing SevenNet. They are not marked as dependencies since it is coupled with the CUDA version, and manual installation is safe for this case. + +After the PyTorch installation, simply run ```bash pip install sevenn ``` @@ -54,6 +61,8 @@ This model was trained on [`MPtrj`](https://figshare.com/articles/dataset/Materi Whenever the checkpoint path is the input, this model can be loaded via `7net-0 | SevenNet-0 | 7net-0_11July2024 | SevenNet-0_11July2024` keywords. +Acknowledgments: This potential was developed with the support of the Samsung Advanced Institute of Technology (SAIT) and utilized the resources of the Samsung SSC-21 cluster. + #### SevenNet-0 (22May2024) This model was trained on [`MPF.2021.2.8`](https://figshare.com/articles/dataset/MPF_2021_2_8/19470599). This is the model used in [our paper](https://pubs.acs.org/doi/10.1021/acs.jctc.4c00190). @@ -87,10 +96,9 @@ Other valid preset options are: `base`, `fine_tune`, and `sevennet-0`. Check comments of `base` yaml for explanations. To reuse a preprocessed training set, you can specify `${dataset_name}.sevenn_data` to the `load_dataset_path:` in the `input.yaml`. -Once you initiate training, `log.sevenn` will contain all parsed inputs from `input.yaml`. You can refer to the log to check the default inputs. #### Multi-GPU training -We support multi-GPU training features using PyTorch DDP (distributed data parallel). We use one process (CPU core) per GPU. +We support multi-GPU training features using PyTorch DDP (distributed data parallel). We use one process (or a CPU core) per GPU. ```bash torchrun --standalone --nnodes {number of nodes} --nproc_per_node {number of GPUs} --no_python sevenn input.yaml -d ``` @@ -125,30 +133,28 @@ sevenn_get_model 7net-0 -p This will create multiple `deployed_parallel_*.pt` files. The number of deployed models equals the number of message-passing layers. These models can be used as lammps potential to run parallel MD simulations with GNN potential using multiple GPU cards. -See `sevenn_inference --help` for more information. - ## Installation for LAMMPS * PyTorch (same version as used for training) -* LAMMPS version of 'stable_2Aug2023' [`LAMMPS`](https://github.com/lammps/lammps) +* LAMMPS version of 'stable_2Aug2023_update3' [`LAMMPS`](https://github.com/lammps/lammps) * (Optional) [`CUDA-aware OpenMPI`](https://www.open-mpi.org/faq/?category=buildcuda) for parallel MD **PLEASE NOTE:** CUDA-aware OpenMPI is optional, but recommended for parallel MD. If it is not available, in parallel mode, GPUs will communicate via CPU. It is still faster than using only one GPU, but its efficiency is low. **PLEASE NOTE:** CUDA-aware OpenMPI does not support NVIDIA Gaming GPUs. Given that the software is closely tied to hardware specifications, please consult with your server administrator if unavailable. -Ensure the LAMMPS version (stable_2Aug2023). You can easily switch the version using git. +Ensure the LAMMPS version (stable_2Aug2023_update3). You can easily switch the version using git. ```bash git clone https://github.com/lammps/lammps.git lammps_dir cd lammps_dir -git checkout stable_2Aug2023 +git checkout stable_2Aug2023_update3 ``` Run sevenn_patch_lammps ```bash sevenn_patch_lammps {path_to_lammps_dir} ``` -Refer to `sevenn/pair_e3gnn/patch_lammps.sh` for the patch process. +Refer to `sevenn/pair_e3gnn/patch_lammps.sh` for the detailed patch process. Build LAMMPS with cmake (example): ``` From 2ee524a95e5b22ec11a915dc7e83f8e6f41bc602 Mon Sep 17 00:00:00 2001 From: YutackPark <111348843+YutackPark@users.noreply.github.com> Date: Tue, 23 Jul 2024 17:49:17 +0900 Subject: [PATCH 15/25] docs: update README.md, parallel model example --- README.md | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/README.md b/README.md index 5bea03f..4dbfac7 100644 --- a/README.md +++ b/README.md @@ -188,6 +188,13 @@ pair_style e3gnn/parallel pair_coeff * * {number of segmented parallel models} {space separated paths of segmented parallel models} {space separated chemical species} ``` +For example, +``` +pair_style e3gnn/parallel +pair_coeff * * 4 ./deployed_parallel_0.pt ./deployed_parallel_1.pt ./deployed_parallel_2.pt ./deployed_parallel_3.pt Hf O +``` +The number of segmented `*.pt` files is the same as the number of message-passing layers of the model. + Check [sevenn_get_model](#sevenn_get_model) for deploying lammps models from checkpoint for both serial and parallel. **PLEASE NOTE:** One GPU per MPI process is expected. If the available GPUs are fewer than the MPI processes, the simulation may run inefficiently. From 3af3247f38c64f3b1dbb1fc9fcc7898ed7b4b423 Mon Sep 17 00:00:00 2001 From: YutackPark <111348843+YutackPark@users.noreply.github.com> Date: Tue, 23 Jul 2024 21:36:39 +0900 Subject: [PATCH 16/25] Update README.md --- README.md | 45 ++++++++++++++++++++++++++++++++++----------- 1 file changed, 34 insertions(+), 11 deletions(-) diff --git a/README.md b/README.md index 4dbfac7..1673a50 100644 --- a/README.md +++ b/README.md @@ -37,14 +37,14 @@ The installation and usage of SevenNet are split into two parts: training (handl Please install PyTorch from [`PyTorch official`](https://pytorch.org/get-started/locally/) before installing the SevenNet. Note that for SevenNet, `torchvision` and `torchaudio` are redundant. You can omit them from the command provided in the installation guide. -Matching PyTorch + CUDA versions may cause problems, especially when compiling SevenNet with LAMMPS. Here are the versions we've been using internally, without problems. +Matching PyTorch + CUDA versions may cause problems, especially when compiling SevenNet with LAMMPS. Here are the recommended versions we've been using internally without an issue. - PyTorch/2.2.2 + CUDA/12.1.0 - PyTorch/1.13.1 + CUDA/12.1.0 - PyTorch/1.12.0 + CUDA/11.6.2 -Using the newer versions of CUDA with PyTorch is usually not an issue. For example, we were able to compile and use `PyTorch/1.13.1+cu117` with `CUDA/12.1.0`. +Using the newer versions of CUDA with PyTorch is usually not an issue. For example, you can compile and use `PyTorch/1.13.1+cu117` with `CUDA/12.1.0`. -**PLEASE NOTE:** You must install PyTorch before installing SevenNet. They are not marked as dependencies since it is coupled with the CUDA version, and manual installation is safe for this case. +**PLEASE NOTE:** You must install PyTorch before installing SevenNet. They are not marked as dependencies since it is coupled with the CUDA version, therefore manual installation is safer. After the PyTorch installation, simply run ```bash @@ -139,22 +139,20 @@ These models can be used as lammps potential to run parallel MD simulations with * LAMMPS version of 'stable_2Aug2023_update3' [`LAMMPS`](https://github.com/lammps/lammps) * (Optional) [`CUDA-aware OpenMPI`](https://www.open-mpi.org/faq/?category=buildcuda) for parallel MD +As system-wise requirements, you also need CUDA and Intel MKL. + **PLEASE NOTE:** CUDA-aware OpenMPI is optional, but recommended for parallel MD. If it is not available, in parallel mode, GPUs will communicate via CPU. It is still faster than using only one GPU, but its efficiency is low. **PLEASE NOTE:** CUDA-aware OpenMPI does not support NVIDIA Gaming GPUs. Given that the software is closely tied to hardware specifications, please consult with your server administrator if unavailable. -Ensure the LAMMPS version (stable_2Aug2023_update3). You can easily switch the version using git. +Ensure the LAMMPS version (stable_2Aug2023_update3). You can easily switch the version using git. After switching the version, run `sevenn_patch_lammps` with the lammps directory path as an argument. ```bash git clone https://github.com/lammps/lammps.git lammps_dir cd lammps_dir git checkout stable_2Aug2023_update3 +sevenn_patch_lammps ./ ``` - -Run sevenn_patch_lammps -```bash -sevenn_patch_lammps {path_to_lammps_dir} -``` -Refer to `sevenn/pair_e3gnn/patch_lammps.sh` for the detailed patch process. +You can refer to `sevenn/pair_e3gnn/patch_lammps.sh` for the detailed patch process. Build LAMMPS with cmake (example): ``` @@ -165,6 +163,31 @@ $ cmake ../cmake -DCMAKE_PREFIX_PATH=`python -c 'import torch;print(torch.utils. $ make -j4 ``` + If you're having trouble with the above procedure and it is related to MKL, check the notes below. + +### Note for MKL errors +You may encounter `MKL_INCLUDE_DIR NOT-FOUND` during cmake or see hundreds of `undefined reference to XXX` errors from `libtorch_cpu.so` at the end of the compilation. This typically indicates that either Intel MKL is not installed or its environment variables are not correctly set. + +First, check if there are Intel MKL modules available on your system (these may have names like intel-oneapi, intel-mkl, intel-compiler, etc.). Load the relevant module to automatically configure the necessary settings. + +If the module is not available, I recommend using Conda. Install `mkl` and `mkl-include` with the following command: +```bash +conda install -c intel mkl mkl-include +``` +If you encounter an error, remove `-c intel`. This is a known bug in the recent Conda version. + +Next, configure the `LD_LIBRARY_PATH` for MKL libraries: +```bash +export LD_LIBRARY_PATH=$CONDA_PREFIX/lib:$LD_LIBRARY_PATH +``` + +Finally, append the following to your cmake command: +```bash +-DMKL_INCLUDE_DIR="$CONDA_PREFIX/include" +``` + +Before running the command, I recommend you clean the files generated from the previous `cmake` command. + ## Usage for LAMMPS ### To check installation @@ -197,7 +220,7 @@ The number of segmented `*.pt` files is the same as the number of message-passin Check [sevenn_get_model](#sevenn_get_model) for deploying lammps models from checkpoint for both serial and parallel. -**PLEASE NOTE:** One GPU per MPI process is expected. If the available GPUs are fewer than the MPI processes, the simulation may run inefficiently. +One GPU per MPI process is expected. If the available GPUs are fewer than the MPI processes, the simulation may run inefficiently. **PLEASE NOTE:** Currently, the parallel version raises an error when there are no atoms in one of the subdomain cells. This issue can be addressed using the `processors` command and, more optimally, the `fix balance` command in LAMMPS. This will be patched in the future. From 0d876027ca2f2d89e112f4bc7403e738431fc557 Mon Sep 17 00:00:00 2001 From: YutackPark <111348843+YutackPark@users.noreply.github.com> Date: Tue, 23 Jul 2024 21:51:00 +0900 Subject: [PATCH 17/25] Update README.md --- README.md | 11 +++++++++-- 1 file changed, 9 insertions(+), 2 deletions(-) diff --git a/README.md b/README.md index 1673a50..1e4b96c 100644 --- a/README.md +++ b/README.md @@ -135,11 +135,12 @@ These models can be used as lammps potential to run parallel MD simulations with ## Installation for LAMMPS + * PyTorch (same version as used for training) * LAMMPS version of 'stable_2Aug2023_update3' [`LAMMPS`](https://github.com/lammps/lammps) * (Optional) [`CUDA-aware OpenMPI`](https://www.open-mpi.org/faq/?category=buildcuda) for parallel MD -As system-wise requirements, you also need CUDA and Intel MKL. +As system-wise requirements, you also need CUDA and Intel MKL. Currently, this installation guide is dedicated to Linux. **PLEASE NOTE:** CUDA-aware OpenMPI is optional, but recommended for parallel MD. If it is not available, in parallel mode, GPUs will communicate via CPU. It is still faster than using only one GPU, but its efficiency is low. @@ -163,9 +164,15 @@ $ cmake ../cmake -DCMAKE_PREFIX_PATH=`python -c 'import torch;print(torch.utils. $ make -j4 ``` +If the compilation is successful, you will find the executable at `{path_to_lammps_dir}/build/lmp`. To use this binary easily, for example, create a soft link in your bin directory (which should be included in your $PATH). +```bash +ln -s {absolute_path_to_lammps_dir}/build/lmp $HOME/.local/bin/lmp +``` +This will allow you to run the binary using `lmp -in my_lammps_script.lmp`. + If you're having trouble with the above procedure and it is related to MKL, check the notes below. -### Note for MKL errors +### Note for MKL You may encounter `MKL_INCLUDE_DIR NOT-FOUND` during cmake or see hundreds of `undefined reference to XXX` errors from `libtorch_cpu.so` at the end of the compilation. This typically indicates that either Intel MKL is not installed or its environment variables are not correctly set. First, check if there are Intel MKL modules available on your system (these may have names like intel-oneapi, intel-mkl, intel-compiler, etc.). Load the relevant module to automatically configure the necessary settings. From 0cb7c46945cdfac08e1714ce5d226847f04875b0 Mon Sep 17 00:00:00 2001 From: YutackPark <111348843+YutackPark@users.noreply.github.com> Date: Tue, 23 Jul 2024 21:56:57 +0900 Subject: [PATCH 18/25] Update README.md --- README.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/README.md b/README.md index 1e4b96c..6b9f5b0 100644 --- a/README.md +++ b/README.md @@ -9,7 +9,7 @@ SevenNet (Scalable EquiVariance Enabled Neural Network) is a graph neural networ The project provides parallel molecular dynamics simulations using graph neural network interatomic potentials, which enable large-scale MD simulations or faster MD simulations. -The installation and usage of SevenNet are split into two parts: training (handled by PyTorch) and molecular dynamics (handled by [`LAMMPS`](https://github.com/lammps/lammps)). The model, once trained with PyTorch, is deployed using TorchScript and is later used to run molecular dynamics simulations via LAMMPS. +The installation and usage of SevenNet are split into two parts: training + command-line interface + ASE calculator (handled by Python) and molecular dynamics (handled by [`LAMMPS`](https://github.com/lammps/lammps)). - [SevenNet](#sevennet) * [Installation](#installation) From 3bb11f53570de2a44edb64cd42ebc8b89bbf008e Mon Sep 17 00:00:00 2001 From: YutackPark <111348843+YutackPark@users.noreply.github.com> Date: Wed, 24 Jul 2024 10:32:24 +0900 Subject: [PATCH 19/25] docs: 7net-0 11July2024 acknowledgement --- README.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/README.md b/README.md index 6b9f5b0..00b0717 100644 --- a/README.md +++ b/README.md @@ -61,7 +61,7 @@ This model was trained on [`MPtrj`](https://figshare.com/articles/dataset/Materi Whenever the checkpoint path is the input, this model can be loaded via `7net-0 | SevenNet-0 | 7net-0_11July2024 | SevenNet-0_11July2024` keywords. -Acknowledgments: This potential was developed with the support of the Samsung Advanced Institute of Technology (SAIT) and utilized the resources of the Samsung SSC-21 cluster. +**Acknowledgments**: This work was supported by the Neural Processing Research Center program of Samsung Advanced Institute of Technology, Samsung Electronics Co., Ltd. The computations for training models were carried out using the Samsung SSC-21 cluster. #### SevenNet-0 (22May2024) This model was trained on [`MPF.2021.2.8`](https://figshare.com/articles/dataset/MPF_2021_2_8/19470599). This is the model used in [our paper](https://pubs.acs.org/doi/10.1021/acs.jctc.4c00190). From 4495bba8d96c2e4e4ab0e10e26c98146e8367111 Mon Sep 17 00:00:00 2001 From: YutackPark Date: Wed, 24 Jul 2024 14:29:30 +0900 Subject: [PATCH 20/25] version up to 0.9.3 --- pyproject.toml | 2 +- sevenn/_const.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index e1ee9ec..5912e95 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "sevenn" -version = "0.9.2" +version = "0.9.3" authors = [ { name="Yutack Park", email="parkyutack@snu.ac.kr" }, { name="Jaesun Kim" }, diff --git a/sevenn/_const.py b/sevenn/_const.py index 7dbe364..035061e 100644 --- a/sevenn/_const.py +++ b/sevenn/_const.py @@ -7,7 +7,7 @@ import sevenn._keys as KEY from sevenn.nn.activation import ShiftedSoftPlus -SEVENN_VERSION = '0.9.2' +SEVENN_VERSION = '0.9.3' IMPLEMENTED_RADIAL_BASIS = ['bessel'] IMPLEMENTED_CUTOFF_FUNCTION = ['poly_cut', 'XPLOR'] # TODO: support None. This became difficult because of paralell model From 6cf602d5992aa639cc365a00a4173b23798b1240 Mon Sep 17 00:00:00 2001 From: YutackPark <111348843+YutackPark@users.noreply.github.com> Date: Wed, 24 Jul 2024 16:23:42 +0900 Subject: [PATCH 21/25] docs: update readme.md --- README.md | 47 +++++++++++++++++++---------------------------- 1 file changed, 19 insertions(+), 28 deletions(-) diff --git a/README.md b/README.md index 00b0717..da54c99 100644 --- a/README.md +++ b/README.md @@ -35,18 +35,18 @@ The installation and usage of SevenNet are split into two parts: training + comm * PyTorch >= 1.12.0 Please install PyTorch from [`PyTorch official`](https://pytorch.org/get-started/locally/) before installing the SevenNet. -Note that for SevenNet, `torchvision` and `torchaudio` are redundant. You can omit them from the command provided in the installation guide. +Note that for SevenNet, `torchvision` and `torchaudio` are redundant. You can safely exclude these packages from the installation commands. -Matching PyTorch + CUDA versions may cause problems, especially when compiling SevenNet with LAMMPS. Here are the recommended versions we've been using internally without an issue. +Here are the recommended versions we've been using internally without an issue. - PyTorch/2.2.2 + CUDA/12.1.0 - PyTorch/1.13.1 + CUDA/12.1.0 - PyTorch/1.12.0 + CUDA/11.6.2 -Using the newer versions of CUDA with PyTorch is usually not an issue. For example, you can compile and use `PyTorch/1.13.1+cu117` with `CUDA/12.1.0`. +Using the newer versions of CUDA with PyTorch is usually not a problem. For example, you can compile and use `PyTorch/1.13.1+cu117` with `CUDA/12.1.0`. -**PLEASE NOTE:** You must install PyTorch before installing SevenNet. They are not marked as dependencies since it is coupled with the CUDA version, therefore manual installation is safer. +**PLEASE NOTE:** You must install PyTorch before installing SevenNet. They are not marked as dependencies since it is coupled with the CUDA version. -After the PyTorch installation, simply run +After the PyTorch installation, run ```bash pip install sevenn ``` @@ -135,29 +135,27 @@ These models can be used as lammps potential to run parallel MD simulations with ## Installation for LAMMPS - * PyTorch (same version as used for training) * LAMMPS version of 'stable_2Aug2023_update3' [`LAMMPS`](https://github.com/lammps/lammps) * (Optional) [`CUDA-aware OpenMPI`](https://www.open-mpi.org/faq/?category=buildcuda) for parallel MD +* MKL-include -As system-wise requirements, you also need CUDA and Intel MKL. Currently, this installation guide is dedicated to Linux. +**PLEASE NOTE:** CUDA-aware OpenMPI does not support NVIDIA Gaming GPUs. Given that the software is closely tied to hardware specifications, please consult with your server administrator if unavailable. -**PLEASE NOTE:** CUDA-aware OpenMPI is optional, but recommended for parallel MD. If it is not available, in parallel mode, GPUs will communicate via CPU. It is still faster than using only one GPU, but its efficiency is low. +If your cluster supports the Intel MKL module (often included with Intel OneAPI, Intel Compiler, and other Intel-related modules), load the module. If it is unavailable, read the 'Note for MKL' section before running cmake. -**PLEASE NOTE:** CUDA-aware OpenMPI does not support NVIDIA Gaming GPUs. Given that the software is closely tied to hardware specifications, please consult with your server administrator if unavailable. +CUDA-aware OpenMPI is optional but recommended for parallel MD. If it is not available, in parallel mode, GPUs will communicate via CPU. It is still faster than using only one GPU, but its efficiency is low. Ensure the LAMMPS version (stable_2Aug2023_update3). You can easily switch the version using git. After switching the version, run `sevenn_patch_lammps` with the lammps directory path as an argument. ```bash -git clone https://github.com/lammps/lammps.git lammps_dir -cd lammps_dir -git checkout stable_2Aug2023_update3 -sevenn_patch_lammps ./ +git clone https://github.com/lammps/lammps.git lammps_test --branch stable_2Aug2023_update3 --depth=1 +sevenn_patch_lammps ./lammps_test ``` You can refer to `sevenn/pair_e3gnn/patch_lammps.sh` for the detailed patch process. Build LAMMPS with cmake (example): ``` -$ cd {path_to_lammps_dir} +$ cd ./lammps_test $ mkdir build $ cd build $ cmake ../cmake -DCMAKE_PREFIX_PATH=`python -c 'import torch;print(torch.utils.cmake_prefix_path)'` @@ -170,30 +168,23 @@ ln -s {absolute_path_to_lammps_dir}/build/lmp $HOME/.local/bin/lmp ``` This will allow you to run the binary using `lmp -in my_lammps_script.lmp`. - If you're having trouble with the above procedure and it is related to MKL, check the notes below. - ### Note for MKL -You may encounter `MKL_INCLUDE_DIR NOT-FOUND` during cmake or see hundreds of `undefined reference to XXX` errors from `libtorch_cpu.so` at the end of the compilation. This typically indicates that either Intel MKL is not installed or its environment variables are not correctly set. +You may encounter MKL_INCLUDE_DIR NOT-FOUND during cmake. This usually means the environment variable is not set correctly, or mkl-include is not present on your system. -First, check if there are Intel MKL modules available on your system (these may have names like intel-oneapi, intel-mkl, intel-compiler, etc.). Load the relevant module to automatically configure the necessary settings. - -If the module is not available, I recommend using Conda. Install `mkl` and `mkl-include` with the following command: +Install mkl-include with: ```bash -conda install -c intel mkl mkl-include +conda install -c intel mkl-include ``` If you encounter an error, remove `-c intel`. This is a known bug in the recent Conda version. -Next, configure the `LD_LIBRARY_PATH` for MKL libraries: +Append the following to your cmake command: ```bash -export LD_LIBRARY_PATH=$CONDA_PREFIX/lib:$LD_LIBRARY_PATH +-DMKL_INCLUDE_DIR=$CONDA_PREFIX/include ``` -Finally, append the following to your cmake command: -```bash --DMKL_INCLUDE_DIR="$CONDA_PREFIX/include" -``` +If you see hundreds of `undefined reference to XXX` errors with `libtorch_cpu.so` at the end of compilation, check your `$LD_LIBRARY_PATH`. PyTorch depends on MKL libraries (this is a default backend for torch+CPU), therefore you already have them. For example, if you installed PyTorch using Conda, you may find `libmkl_*.so` files under `$CONDA_PREFIX/lib`. Ensure that `$LD_LIBRARY_PATH` includes `$CONDA_PREFIX/lib`. -Before running the command, I recommend you clean the files generated from the previous `cmake` command. +For other error cases, you might want to check [`pair-nequip`](https://github.com/mir-group/pair_nequip). ## Usage for LAMMPS From c3cb96355274cf2b1be07e79c978ee5edf0014eb Mon Sep 17 00:00:00 2001 From: YutackPark <111348843+YutackPark@users.noreply.github.com> Date: Wed, 24 Jul 2024 16:35:54 +0900 Subject: [PATCH 22/25] docs: linke to SevenNet-0 models --- README.md | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/README.md b/README.md index da54c99..c1c6db7 100644 --- a/README.md +++ b/README.md @@ -57,14 +57,14 @@ pip install sevenn SevenNet-0 is a general-purpose interatomic potential trained on the [`MPF dataset of M3GNet`](https://figshare.com/articles/dataset/MPF_2021_2_8/19470599) or [`MPtrj dataset of CHGNet`](https://figshare.com/articles/dataset/Materials_Project_Trjectory_MPtrj_Dataset/23713842). You can try SevenNet-0 to your application without any training. If the accuracy is unsatisfactory, SevenNet-0 can be [fine-tuned](#training). #### SevenNet-0 (11July2024) -This model was trained on [`MPtrj`](https://figshare.com/articles/dataset/Materials_Project_Trjectory_MPtrj_Dataset/23713842). We suggest starting with this model as we found that it performs better than the previous SevenNet-0 (22May2024). Check [`Matbench Discovery leaderborad`](https://matbench-discovery.materialsproject.org/) for this model's performance on materials discovery. +This model was trained on [`MPtrj`](https://figshare.com/articles/dataset/Materials_Project_Trjectory_MPtrj_Dataset/23713842). We suggest starting with this model as we found that it performs better than the previous SevenNet-0 (22May2024). Check [`Matbench Discovery leaderborad`](https://matbench-discovery.materialsproject.org/) for this model's performance on materials discovery. For more information, click [here](sevenn/pretrained_potentials/SevenNet_0__11July2024). Whenever the checkpoint path is the input, this model can be loaded via `7net-0 | SevenNet-0 | 7net-0_11July2024 | SevenNet-0_11July2024` keywords. **Acknowledgments**: This work was supported by the Neural Processing Research Center program of Samsung Advanced Institute of Technology, Samsung Electronics Co., Ltd. The computations for training models were carried out using the Samsung SSC-21 cluster. #### SevenNet-0 (22May2024) -This model was trained on [`MPF.2021.2.8`](https://figshare.com/articles/dataset/MPF_2021_2_8/19470599). This is the model used in [our paper](https://pubs.acs.org/doi/10.1021/acs.jctc.4c00190). +This model was trained on [`MPF.2021.2.8`](https://figshare.com/articles/dataset/MPF_2021_2_8/19470599). This is the model used in [our paper](https://pubs.acs.org/doi/10.1021/acs.jctc.4c00190). For more information, click [here](sevenn/pretrained_potentials/SevenNet_0__22May2024). Whenever the checkpoint path is the input, this model can be loaded via `7net-0_22May2024 | SevenNet-0_22May2024` keywords. From 708592800ebcb327d11a5f01b54ea526ca1933ad Mon Sep 17 00:00:00 2001 From: YutackPark <111348843+YutackPark@users.noreply.github.com> Date: Wed, 24 Jul 2024 16:41:30 +0900 Subject: [PATCH 23/25] docs: sevennet-0 11july2024 --- .../SevenNet_0__11July2024/README.md | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/sevenn/pretrained_potentials/SevenNet_0__11July2024/README.md b/sevenn/pretrained_potentials/SevenNet_0__11July2024/README.md index 79f5d5f..ea7377f 100644 --- a/sevenn/pretrained_potentials/SevenNet_0__11July2024/README.md +++ b/sevenn/pretrained_potentials/SevenNet_0__11July2024/README.md @@ -15,11 +15,15 @@ It can be directly applied to any system without training and fine-tuned with an - fine_tune.yaml: example input.yaml for fine-tuning. - pre_train.yaml: example input.yaml file used in pre-training. +For SevenNet-0 (11July2024), we did not split the dataset. +| |Energy (eV/atom)|Force (eV/Å)|Stress (GPa)| +|----------------|--------|-------|-------| +|Train|0.011|0.040|0.28| + ### Note You can obtain the same deployed models by running the following command (-p for parallel) ```bash -$ sevenn_get_model {-p} checkpoint_sevennet_0.pth +$ sevenn_get_model {-p} 7net-0_11July2024 ``` Refer to example_inputs/md_{serial/parallel}_example/ for their usage in LAMMPS. Both serial and parallel model gives the same results but the parallel model enables multi-GPU MD simulation. - From 39759ae5e17da81cf0015463961fe75f5b56ed78 Mon Sep 17 00:00:00 2001 From: YutackPark <111348843+YutackPark@users.noreply.github.com> Date: Wed, 24 Jul 2024 19:28:45 +0900 Subject: [PATCH 24/25] docs: final revision --- README.md | 26 +++++++++++++++----------- 1 file changed, 15 insertions(+), 11 deletions(-) diff --git a/README.md b/README.md index c1c6db7..1368e7d 100644 --- a/README.md +++ b/README.md @@ -5,11 +5,11 @@ # SevenNet -SevenNet (Scalable EquiVariance Enabled Neural Network) is a graph neural network interatomic potential package that supports parallel molecular dynamics simulations with [`LAMMPS`](https://github.com/lammps/lammps). Its underlying GNN model is based on [`nequip`](https://github.com/mir-group/nequip). +SevenNet (Scalable EquiVariance Enabled Neural Network) is a graph neural network interatomic potential package that supports parallel molecular dynamics simulations with [`LAMMPS`](https://docs.lammps.org/Manual.html). Its underlying GNN model is based on [`nequip`](https://github.com/mir-group/nequip). The project provides parallel molecular dynamics simulations using graph neural network interatomic potentials, which enable large-scale MD simulations or faster MD simulations. -The installation and usage of SevenNet are split into two parts: training + command-line interface + ASE calculator (handled by Python) and molecular dynamics (handled by [`LAMMPS`](https://github.com/lammps/lammps)). +The installation and usage of SevenNet are split into two parts: training + command-line interface + ASE calculator (handled by Python) and molecular dynamics (handled by [`LAMMPS`](https://docs.lammps.org/Manual.html)). - [SevenNet](#sevennet) * [Installation](#installation) @@ -37,7 +37,7 @@ The installation and usage of SevenNet are split into two parts: training + comm Please install PyTorch from [`PyTorch official`](https://pytorch.org/get-started/locally/) before installing the SevenNet. Note that for SevenNet, `torchvision` and `torchaudio` are redundant. You can safely exclude these packages from the installation commands. -Here are the recommended versions we've been using internally without an issue. +Here are the recommended versions we've been using internally without any issues. - PyTorch/2.2.2 + CUDA/12.1.0 - PyTorch/1.13.1 + CUDA/12.1.0 - PyTorch/1.12.0 + CUDA/11.6.2 @@ -148,14 +148,14 @@ CUDA-aware OpenMPI is optional but recommended for parallel MD. If it is not ava Ensure the LAMMPS version (stable_2Aug2023_update3). You can easily switch the version using git. After switching the version, run `sevenn_patch_lammps` with the lammps directory path as an argument. ```bash -git clone https://github.com/lammps/lammps.git lammps_test --branch stable_2Aug2023_update3 --depth=1 -sevenn_patch_lammps ./lammps_test +git clone https://github.com/lammps/lammps.git lammps_sevenn --branch stable_2Aug2023_update3 --depth=1 +sevenn_patch_lammps ./lammps_sevenn ``` You can refer to `sevenn/pair_e3gnn/patch_lammps.sh` for the detailed patch process. Build LAMMPS with cmake (example): ``` -$ cd ./lammps_test +$ cd ./lammps_sevenn $ mkdir build $ cd build $ cmake ../cmake -DCMAKE_PREFIX_PATH=`python -c 'import torch;print(torch.utils.cmake_prefix_path)'` @@ -169,7 +169,7 @@ ln -s {absolute_path_to_lammps_dir}/build/lmp $HOME/.local/bin/lmp This will allow you to run the binary using `lmp -in my_lammps_script.lmp`. ### Note for MKL -You may encounter MKL_INCLUDE_DIR NOT-FOUND during cmake. This usually means the environment variable is not set correctly, or mkl-include is not present on your system. +You may encounter `MKL_INCLUDE_DIR NOT-FOUND` during cmake. This usually means the environment variable is not set correctly, or mkl-include is not present on your system. Install mkl-include with: ```bash @@ -184,20 +184,22 @@ Append the following to your cmake command: If you see hundreds of `undefined reference to XXX` errors with `libtorch_cpu.so` at the end of compilation, check your `$LD_LIBRARY_PATH`. PyTorch depends on MKL libraries (this is a default backend for torch+CPU), therefore you already have them. For example, if you installed PyTorch using Conda, you may find `libmkl_*.so` files under `$CONDA_PREFIX/lib`. Ensure that `$LD_LIBRARY_PATH` includes `$CONDA_PREFIX/lib`. -For other error cases, you might want to check [`pair-nequip`](https://github.com/mir-group/pair_nequip). +For other error cases, you might want to check [`pair-nequip`](https://github.com/mir-group/pair_nequip), as the `pair-nequip` and SevenNet+LAMMPS shares similar requirements: torch + LAMMPS. ## Usage for LAMMPS ### To check installation ```bash -{lmp_binary} -help | grep e3gnn +{lammps_binary} -help | grep e3gnn ``` You will see `e3gnn` and `e3gnn/parallel` as pair_style. ### For serial model ``` +units metal +atom_style atomic pair_style e3gnn pair_coeff * * {path to serial model} {space separated chemical species} ``` @@ -205,6 +207,8 @@ pair_coeff * * {path to serial model} {space separated chemical species} ### For parallel model ``` +units metal +atom_style atomic pair_style e3gnn/parallel pair_coeff * * {number of segmented parallel models} {space separated paths of segmented parallel models} {space separated chemical species} ``` @@ -214,11 +218,11 @@ For example, pair_style e3gnn/parallel pair_coeff * * 4 ./deployed_parallel_0.pt ./deployed_parallel_1.pt ./deployed_parallel_2.pt ./deployed_parallel_3.pt Hf O ``` -The number of segmented `*.pt` files is the same as the number of message-passing layers of the model. +The number of segmented `*.pt` files is same as the number of message-passing layers of the model. Check [sevenn_get_model](#sevenn_get_model) for deploying lammps models from checkpoint for both serial and parallel. -One GPU per MPI process is expected. If the available GPUs are fewer than the MPI processes, the simulation may run inefficiently. +One GPU per MPI process is expected. The simulation may run inefficiently if the available GPUs are fewer than the MPI processes. **PLEASE NOTE:** Currently, the parallel version raises an error when there are no atoms in one of the subdomain cells. This issue can be addressed using the `processors` command and, more optimally, the `fix balance` command in LAMMPS. This will be patched in the future. From 493c0a677c46c93d6641d4159043731b3bf7490a Mon Sep 17 00:00:00 2001 From: YutackPark Date: Wed, 24 Jul 2024 21:03:07 +0900 Subject: [PATCH 25/25] bugfix: species wise force rms --- sevenn/train/dataset.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/sevenn/train/dataset.py b/sevenn/train/dataset.py index c404f7d..39d2660 100644 --- a/sevenn/train/dataset.py +++ b/sevenn/train/dataset.py @@ -338,7 +338,7 @@ def get_species_wise_force_rms(self, num_chem_species): 0, index, force.square(), reduce='mean', include_self=False ) - return rms + return torch.sqrt(rms.mean(dim=1)) def get_avg_num_neigh(self): n_neigh = []