From 7e1dedc8b10773cc882ae1043cffc464422b544c Mon Sep 17 00:00:00 2001 From: Eszter Varga-Umbrich Date: Mon, 15 Jul 2024 00:05:47 +0100 Subject: [PATCH 01/11] remove conda install steps --- README.md | 21 +-------------------- 1 file changed, 1 insertion(+), 20 deletions(-) diff --git a/README.md b/README.md index 83a02a01e..6313454b8 100644 --- a/README.md +++ b/README.md @@ -71,26 +71,7 @@ pip install mace-torch For CPU or MPS (Apple Silicon) installation, use `pip install torch torchvision torchaudio` instead. -### conda installation from source - -To install from source using `conda`, follow the steps below: -```sh -# Create a virtual environment and activate it -conda create --name mace_env -conda activate mace_env - -# Install PyTorch -conda install pytorch torchvision torchaudio pytorch-cuda=11.6 -c pytorch -c nvidia - -# (optional) Install MACE's dependencies from Conda as well -conda install numpy scipy matplotlib ase opt_einsum prettytable pandas e3nn - -# Clone and install MACE (and all required packages) -git clone https://github.com/ACEsuit/mace.git -pip install ./mace -``` -For the Pytorch version, use the appropriate version for your CUDA version. -### pip installation from source +### Installation from source To install via `pip`, follow the steps below: From 9bab980b0d117a7f409ab7580ab440b195ec95eb Mon Sep 17 00:00:00 2001 From: Eszter Varga-Umbrich Date: Mon, 15 Jul 2024 16:38:09 +0100 Subject: [PATCH 02/11] reorder and rephrase installation guide --- README.md | 27 ++++++++------------------- 1 file changed, 8 insertions(+), 19 deletions(-) diff --git a/README.md b/README.md index 6313454b8..9a71c5ecb 100644 --- a/README.md +++ b/README.md @@ -50,45 +50,34 @@ A partial documentation is available at: https://mace-docs.readthedocs.io ## Installation -Requirements: +### 1. Requirements: -- Python >= 3.7 +- Python >= 3.7 (for openMM, use Python = 3.9) - [PyTorch](https://pytorch.org/) >= 1.12 **(training with float64 is not supported with PyTorch 2.1 but is supported with 2.2 and later.)**. -(for openMM, use Python = 3.9) +**Make sure to install PyTorch.** Please refer to the [official PyTorch installation](https://pytorch.org/get-started/locally/) for the installation instructions. Select the appropriate options for your system. For GPU installation, make sure to select pip + the appropriate CUDA version for your system. For recent GPUs, the latest cuda version is usually the best choice. -### pip installation +### 2a. Installation from PyPI This is the recommended way to install MACE. -**First, make sure to install PyTorch.** Please refer to the [official PyTorch installation](https://pytorch.org/get-started/locally/) for the installation instructions. Select the appropriate options for your system. For GPU installation, make sure to select pip + the appropriate CUDA version for your system. For recent GPUs, the latest cuda version is usually the best choice. - -To install via `pip`, follow the steps below: - ```sh pip install --upgrade pip pip install mace-torch ``` +**Note:** The homonymous package on [PyPI](https://pypi.org/project/MACE/) has nothing to do with this one. -For CPU or MPS (Apple Silicon) installation, use `pip install torch torchvision torchaudio` instead. -### Installation from source +### 2b. Installation from source -To install via `pip`, follow the steps below: ```sh -# Create a virtual environment and activate it -python -m venv mace-venv -source mace-venv/bin/activate - -# Install PyTorch (for example, for CUDA 11.6 [cu116]) -pip3 install torch==2.0.1 torchvision==0.15.2 torchaudio==2.0.2 --index-url https://download.pytorch.org/whl/cu118 - # Clone and install MACE (and all required packages) git clone https://github.com/ACEsuit/mace.git pip install ./mace ``` -**Note:** The homonymous package on [PyPI](https://pypi.org/project/MACE/) has nothing to do with this one. + + ## Usage From 134d4aacdf26535865f908043967c397f7c29146 Mon Sep 17 00:00:00 2001 From: Eszter Varga-Umbrich Date: Tue, 16 Jul 2024 00:07:35 +0100 Subject: [PATCH 03/11] Simplify installation steps in the readme --- README.md | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/README.md b/README.md index 9a71c5ecb..31d466af4 100644 --- a/README.md +++ b/README.md @@ -53,9 +53,9 @@ A partial documentation is available at: https://mace-docs.readthedocs.io ### 1. Requirements: - Python >= 3.7 (for openMM, use Python = 3.9) -- [PyTorch](https://pytorch.org/) >= 1.12 **(training with float64 is not supported with PyTorch 2.1 but is supported with 2.2 and later.)**. +- [PyTorch](https://pytorch.org/) >= 1.12 **(training with float64 is not supported with PyTorch 2.1 but is supported with 2.2 and later)** -**Make sure to install PyTorch.** Please refer to the [official PyTorch installation](https://pytorch.org/get-started/locally/) for the installation instructions. Select the appropriate options for your system. For GPU installation, make sure to select pip + the appropriate CUDA version for your system. For recent GPUs, the latest cuda version is usually the best choice. +**Make sure to install PyTorch.** Please refer to the [official PyTorch installation](https://pytorch.org/get-started/locally/) for the installation instructions. Select the appropriate options for your system. ### 2a. Installation from PyPI This is the recommended way to install MACE. @@ -71,7 +71,6 @@ pip install mace-torch ```sh -# Clone and install MACE (and all required packages) git clone https://github.com/ACEsuit/mace.git pip install ./mace ``` From 1700caef7dcaaaea3f6d6635ffc603e0fd52de9c Mon Sep 17 00:00:00 2001 From: Ilyes Batatia <48651863+ilyes319@users.noreply.github.com> Date: Tue, 16 Jul 2024 11:35:44 +0100 Subject: [PATCH 04/11] remove Literal for python >3.9 --- mace/__version__.py | 2 +- mace/calculators/foundations_models.py | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/mace/__version__.py b/mace/__version__.py index a8d4557d2..d7b30e121 100644 --- a/mace/__version__.py +++ b/mace/__version__.py @@ -1 +1 @@ -__version__ = "0.3.5" +__version__ = "0.3.6" diff --git a/mace/calculators/foundations_models.py b/mace/calculators/foundations_models.py index 1ee7659a6..bc9ed6548 100644 --- a/mace/calculators/foundations_models.py +++ b/mace/calculators/foundations_models.py @@ -1,7 +1,7 @@ import os import urllib.request from pathlib import Path -from typing import Literal, Union +from typing import Union import torch from ase import units @@ -20,7 +20,7 @@ def mace_mp( device: str = "", default_dtype: str = "float32", dispersion: bool = False, - damping: Literal["zero", "bj", "zerom", "bjm"] = "bj", + damping: str = "bj", # choices: ["zero", "bj", "zerom", "bjm"] dispersion_xc: str = "pbe", dispersion_cutoff: float = 40.0 * units.Bohr, **kwargs, From 442928201bb7272329d07ce142d509472d5195ff Mon Sep 17 00:00:00 2001 From: vue1999 Date: Wed, 17 Jul 2024 01:08:29 +0100 Subject: [PATCH 05/11] Add two stage flags --- mace/tools/arg_parser.py | 37 ++++++++++++++++++++++--------------- mace/tools/train.py | 2 ++ 2 files changed, 24 insertions(+), 15 deletions(-) diff --git a/mace/tools/arg_parser.py b/mace/tools/arg_parser.py index 893203aa5..37d0ce8de 100644 --- a/mace/tools/arg_parser.py +++ b/mace/tools/arg_parser.py @@ -388,46 +388,51 @@ def build_default_arg_parser() -> argparse.ArgumentParser: "--forces_weight", help="weight of forces loss", type=float, default=100.0 ) parser.add_argument( - "--swa_forces_weight", - help="weight of forces loss after starting swa", + "--swa_forces_weight","--stage_two_forces_weight", + help="weight of forces loss after starting Stage Two (previously called swa)", type=float, default=100.0, + dest="swa_forces_weight", ) parser.add_argument( "--energy_weight", help="weight of energy loss", type=float, default=1.0 ) parser.add_argument( - "--swa_energy_weight", - help="weight of energy loss after starting swa", + "--swa_energy_weight","--stage_two_energy_weight", + help="weight of energy loss after starting Stage Two (previously called swa)", type=float, default=1000.0, + dest="swa_energy_weight", ) parser.add_argument( "--virials_weight", help="weight of virials loss", type=float, default=1.0 ) parser.add_argument( - "--swa_virials_weight", - help="weight of virials loss after starting swa", + "--swa_virials_weight", "--stage_two_virials_weight", + help="weight of virials loss after starting Stage Two (previously called swa)", type=float, default=10.0, + dest="swa_virials_weight", ) parser.add_argument( "--stress_weight", help="weight of virials loss", type=float, default=1.0 ) parser.add_argument( - "--swa_stress_weight", - help="weight of stress loss after starting swa", + "--swa_stress_weight", "--stage_two_stress_weight", + help="weight of stress loss after starting Stage Two (previously called swa)", type=float, default=10.0, + dest="swa_stress_weight", ) parser.add_argument( "--dipole_weight", help="weight of dipoles loss", type=float, default=1.0 ) parser.add_argument( - "--swa_dipole_weight", - help="weight of dipoles after starting swa", + "--swa_dipole_weight","--stage_two_dipole_weight", + help="weight of dipoles after starting Stage Two (previously called swa)", type=float, default=1.0, + dest="swa_dipole_weight", ) parser.add_argument( "--config_type_weights", @@ -462,7 +467,7 @@ def build_default_arg_parser() -> argparse.ArgumentParser: "--lr", help="Learning rate of optimizer", type=float, default=0.01 ) parser.add_argument( - "--swa_lr", help="Learning rate of optimizer in swa", type=float, default=1e-3 + "--swa_lr", "--stage_two_lr", help="Learning rate of optimizer in Stage Two (previously called swa)", type=float, default=1e-3, dest="swa_lr" ) parser.add_argument( "--weight_decay", help="weight decay (L2 penalty)", type=float, default=5e-7 @@ -489,16 +494,18 @@ def build_default_arg_parser() -> argparse.ArgumentParser: default=0.9993, ) parser.add_argument( - "--swa", - help="use Stochastic Weight Averaging, which decreases the learning rate and increases the energy weight at the end of the training to help converge them", + "--swa", "--stage_two", + help="use Stage Two loss weight, which decreases the learning rate and increases the energy weight at the end of the training to help converge them", action="store_true", default=False, + dest="swa", ) parser.add_argument( - "--start_swa", - help="Number of epochs before switching to swa", + "--start_swa","--start_stage_two", + help="Number of epochs before changing to Stage Two loss weights", type=int, default=None, + dest="start_swa", ) parser.add_argument( "--ema", diff --git a/mace/tools/train.py b/mace/tools/train.py index 7ebf3ce15..22e2c982b 100644 --- a/mace/tools/train.py +++ b/mace/tools/train.py @@ -168,6 +168,8 @@ def train( swa_start = False keep_last = True loss_fn = swa.loss_fn + print("EsztiBUG: ") + print(swa.model) swa.model.update_parameters(model) if epoch > start_epoch: swa.scheduler.step() From af3af99411ef46240b2d1a5f3bad3ccc07113ec1 Mon Sep 17 00:00:00 2001 From: vue1999 Date: Wed, 17 Jul 2024 16:54:20 +0100 Subject: [PATCH 06/11] changed loggings of swa to Stage Two --- mace/cli/run_train.py | 18 +++++++++--------- mace/tools/train.py | 6 ++---- 2 files changed, 11 insertions(+), 13 deletions(-) diff --git a/mace/cli/run_train.py b/mace/cli/run_train.py index aecc3f713..63b6d7089 100644 --- a/mace/cli/run_train.py +++ b/mace/cli/run_train.py @@ -585,19 +585,19 @@ def run(args: argparse.Namespace) -> None: swa: Optional[tools.SWAContainer] = None swas = [False] if args.swa: - assert dipole_only is False, "swa for dipole fitting not implemented" + assert dipole_only is False, "Stage Two for dipole fitting not implemented" swas.append(True) if args.start_swa is None: args.start_swa = max(1, args.max_num_epochs // 4 * 3) else: if args.start_swa > args.max_num_epochs: logging.info( - f"Start swa must be less than max_num_epochs, got {args.start_swa} > {args.max_num_epochs}" + f"Start Stage Two must be less than max_num_epochs, got {args.start_swa} > {args.max_num_epochs}" ) args.start_swa = max(1, args.max_num_epochs // 4 * 3) - logging.info(f"Setting start swa to {args.start_swa}") + logging.info(f"Setting start Stage Two to {args.start_swa}") if args.loss == "forces_only": - raise ValueError("Can not select swa with forces only loss.") + raise ValueError("Can not select Stage Two with forces only loss.") if args.loss == "virials": loss_fn_energy = modules.WeightedEnergyForcesVirialsLoss( energy_weight=args.swa_energy_weight, @@ -617,7 +617,7 @@ def run(args: argparse.Namespace) -> None: dipole_weight=args.swa_dipole_weight, ) logging.info( - f"Using stochastic weight averaging (after {args.start_swa} epochs) with energy weight : {args.swa_energy_weight}, forces weight : {args.swa_forces_weight}, dipole weight : {args.swa_dipole_weight} and learning rate : {args.swa_lr}" + f"Stage Two (after {args.start_swa} epochs) with energy weight : {args.swa_energy_weight}, forces weight : {args.swa_forces_weight}, dipole weight : {args.swa_dipole_weight} and learning rate : {args.swa_lr}" ) else: loss_fn_energy = modules.WeightedEnergyForcesLoss( @@ -625,7 +625,7 @@ def run(args: argparse.Namespace) -> None: forces_weight=args.swa_forces_weight, ) logging.info( - f"Using stochastic weight averaging (after {args.start_swa} epochs) with energy weight : {args.swa_energy_weight}, forces weight : {args.swa_forces_weight} and learning rate : {args.swa_lr}" + f"Stage Two (after {args.start_swa} epochs) with energy weight : {args.swa_energy_weight}, forces weight : {args.swa_forces_weight} and learning rate : {args.swa_lr}" ) swa = tools.SWAContainer( model=AveragedModel(model), @@ -807,7 +807,7 @@ def run(args: argparse.Namespace) -> None: if rank == 0: # Save entire model if swa_eval: - model_path = Path(args.checkpoints_dir) / (tag + "_swa.model") + model_path = Path(args.checkpoints_dir) / (tag + "_s2.model") else: model_path = Path(args.checkpoints_dir) / (tag + ".model") logging.info(f"Saving model to {model_path}") @@ -821,10 +821,10 @@ def run(args: argparse.Namespace) -> None: ), } if swa_eval: - torch.save(model, Path(args.model_dir) / (args.name + "_swa.model")) + torch.save(model, Path(args.model_dir) / (args.name + "_s2.model")) try: path_complied = Path(args.model_dir) / ( - args.name + "_swa_compiled.model" + args.name + "_s2_compiled.model" ) logging.info(f"Compiling model, saving metadata {path_complied}") model_compiled = jit.compile(deepcopy(model)) diff --git a/mace/tools/train.py b/mace/tools/train.py index 22e2c982b..575fb02dc 100644 --- a/mace/tools/train.py +++ b/mace/tools/train.py @@ -163,13 +163,11 @@ def train( ) # Can break if exponential LR, TODO fix that! else: if swa_start: - logging.info("Changing loss based on SWA") + logging.info("Changing loss based on Stage Two Weights") lowest_loss = np.inf swa_start = False keep_last = True loss_fn = swa.loss_fn - print("EsztiBUG: ") - print(swa.model) swa.model.update_parameters(model) if epoch > start_epoch: swa.scheduler.step() @@ -235,7 +233,7 @@ def train( patience_counter += 1 if patience_counter >= patience and epoch < swa.start: logging.info( - f"Stopping optimization after {patience_counter} epochs without improvement and starting swa" + f"Stopping optimization after {patience_counter} epochs without improvement and starting Stage Two" ) epoch = swa.start elif patience_counter >= patience and epoch >= swa.start: From d8f6ae1befddb9affd2d1830504138c4134b5260 Mon Sep 17 00:00:00 2001 From: vue1999 Date: Wed, 17 Jul 2024 23:50:10 +0100 Subject: [PATCH 07/11] changed s2 to stagetwo in saved model names --- mace/cli/run_train.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/mace/cli/run_train.py b/mace/cli/run_train.py index 63b6d7089..725c5d215 100644 --- a/mace/cli/run_train.py +++ b/mace/cli/run_train.py @@ -807,7 +807,7 @@ def run(args: argparse.Namespace) -> None: if rank == 0: # Save entire model if swa_eval: - model_path = Path(args.checkpoints_dir) / (tag + "_s2.model") + model_path = Path(args.checkpoints_dir) / (tag + "_stagetwo.model") else: model_path = Path(args.checkpoints_dir) / (tag + ".model") logging.info(f"Saving model to {model_path}") @@ -821,10 +821,10 @@ def run(args: argparse.Namespace) -> None: ), } if swa_eval: - torch.save(model, Path(args.model_dir) / (args.name + "_s2.model")) + torch.save(model, Path(args.model_dir) / (args.name + "_stagetwo.model")) try: path_complied = Path(args.model_dir) / ( - args.name + "_s2_compiled.model" + args.name + "_stagetwo_compiled.model" ) logging.info(f"Compiling model, saving metadata {path_complied}") model_compiled = jit.compile(deepcopy(model)) From 0d2c32d5d751b8fda67bd0f3e4f527ebde7d33ee Mon Sep 17 00:00:00 2001 From: Ilyes Batatia <48651863+ilyes319@users.noreply.github.com> Date: Tue, 23 Jul 2024 16:01:20 +0100 Subject: [PATCH 08/11] remove mypy from checks --- scripts/run_checks.sh | 1 - 1 file changed, 1 deletion(-) diff --git a/scripts/run_checks.sh b/scripts/run_checks.sh index 4a01da4b0..bd1214a40 100755 --- a/scripts/run_checks.sh +++ b/scripts/run_checks.sh @@ -4,7 +4,6 @@ python -m isort . # Check python -m pylint --rcfile=pyproject.toml mace tests scripts -python -m mypy --config-file=.mypy.ini mace tests scripts # Tests python -m pytest tests From b76a2a947576d8b10fd5f72ccf2053aebaed96f7 Mon Sep 17 00:00:00 2001 From: Ilyes Batatia <48651863+ilyes319@users.noreply.github.com> Date: Thu, 8 Aug 2024 13:06:22 +0100 Subject: [PATCH 09/11] add option to load E0s from json --- mace/tools/scripts_utils.py | 16 +++++++++++----- 1 file changed, 11 insertions(+), 5 deletions(-) diff --git a/mace/tools/scripts_utils.py b/mace/tools/scripts_utils.py index cc7b39291..5261c124e 100644 --- a/mace/tools/scripts_utils.py +++ b/mace/tools/scripts_utils.py @@ -304,11 +304,17 @@ def get_atomic_energies(E0s, train_collection, z_table) -> dict: f"Could not compute average E0s if no training xyz given, error {e} occured" ) from e else: - try: - atomic_energies_dict = ast.literal_eval(E0s) - assert isinstance(atomic_energies_dict, dict) - except Exception as e: - raise RuntimeError(f"E0s specified invalidly, error {e} occured") from e + if E0s.endswith(".json"): + logging.info(f"Loading atomic energies from {E0s}") + atomic_energies_dict = json.load(open(E0s, "r")) + else: + try: + atomic_energies_dict = ast.literal_eval(E0s) + assert isinstance(atomic_energies_dict, dict) + except Exception as e: + raise RuntimeError( + f"E0s specified invalidly, error {e} occured" + ) from e else: raise RuntimeError( "E0s not found in training file and not specified in command line" From 562f76096508e5e5233cc094cf349caa3569195e Mon Sep 17 00:00:00 2001 From: Ilyes Batatia <48651863+ilyes319@users.noreply.github.com> Date: Mon, 12 Aug 2024 10:42:26 +0100 Subject: [PATCH 10/11] fix MAE table with universal loss --- mace/cli/run_train.py | 9 +++++++-- mace/tools/arg_parser.py | 29 +++++++++++++++++++++-------- mace/tools/scripts_utils.py | 34 ++++++++++++++++++++++++++++++++++ mace/tools/train.py | 20 ++++++++++++++++++++ 4 files changed, 82 insertions(+), 10 deletions(-) diff --git a/mace/cli/run_train.py b/mace/cli/run_train.py index 725c5d215..a4b2aebe6 100644 --- a/mace/cli/run_train.py +++ b/mace/cli/run_train.py @@ -357,7 +357,10 @@ def run(args: argparse.Namespace) -> None: if args.loss in ("stress", "virials", "huber", "universal"): compute_virials = True args.compute_stress = True - args.error_table = "PerAtomRMSEstressvirials" + if "MAE" in args.error_table: + args.error_table = "PerAtomMAEstressvirials" + else: + args.error_table = "PerAtomRMSEstressvirials" output_args = { "energy": compute_energy, @@ -821,7 +824,9 @@ def run(args: argparse.Namespace) -> None: ), } if swa_eval: - torch.save(model, Path(args.model_dir) / (args.name + "_stagetwo.model")) + torch.save( + model, Path(args.model_dir) / (args.name + "_stagetwo.model") + ) try: path_complied = Path(args.model_dir) / ( args.name + "_stagetwo_compiled.model" diff --git a/mace/tools/arg_parser.py b/mace/tools/arg_parser.py index 37d0ce8de..38034335d 100644 --- a/mace/tools/arg_parser.py +++ b/mace/tools/arg_parser.py @@ -80,6 +80,7 @@ def build_default_arg_parser() -> argparse.ArgumentParser: "PerAtomRMSE", "TotalRMSE", "PerAtomRMSEstressvirials", + "PerAtomMAEstressvirials", "PerAtomMAE", "TotalMAE", "DipoleRMSE", @@ -388,7 +389,8 @@ def build_default_arg_parser() -> argparse.ArgumentParser: "--forces_weight", help="weight of forces loss", type=float, default=100.0 ) parser.add_argument( - "--swa_forces_weight","--stage_two_forces_weight", + "--swa_forces_weight", + "--stage_two_forces_weight", help="weight of forces loss after starting Stage Two (previously called swa)", type=float, default=100.0, @@ -398,7 +400,8 @@ def build_default_arg_parser() -> argparse.ArgumentParser: "--energy_weight", help="weight of energy loss", type=float, default=1.0 ) parser.add_argument( - "--swa_energy_weight","--stage_two_energy_weight", + "--swa_energy_weight", + "--stage_two_energy_weight", help="weight of energy loss after starting Stage Two (previously called swa)", type=float, default=1000.0, @@ -408,7 +411,8 @@ def build_default_arg_parser() -> argparse.ArgumentParser: "--virials_weight", help="weight of virials loss", type=float, default=1.0 ) parser.add_argument( - "--swa_virials_weight", "--stage_two_virials_weight", + "--swa_virials_weight", + "--stage_two_virials_weight", help="weight of virials loss after starting Stage Two (previously called swa)", type=float, default=10.0, @@ -418,7 +422,8 @@ def build_default_arg_parser() -> argparse.ArgumentParser: "--stress_weight", help="weight of virials loss", type=float, default=1.0 ) parser.add_argument( - "--swa_stress_weight", "--stage_two_stress_weight", + "--swa_stress_weight", + "--stage_two_stress_weight", help="weight of stress loss after starting Stage Two (previously called swa)", type=float, default=10.0, @@ -428,7 +433,8 @@ def build_default_arg_parser() -> argparse.ArgumentParser: "--dipole_weight", help="weight of dipoles loss", type=float, default=1.0 ) parser.add_argument( - "--swa_dipole_weight","--stage_two_dipole_weight", + "--swa_dipole_weight", + "--stage_two_dipole_weight", help="weight of dipoles after starting Stage Two (previously called swa)", type=float, default=1.0, @@ -467,7 +473,12 @@ def build_default_arg_parser() -> argparse.ArgumentParser: "--lr", help="Learning rate of optimizer", type=float, default=0.01 ) parser.add_argument( - "--swa_lr", "--stage_two_lr", help="Learning rate of optimizer in Stage Two (previously called swa)", type=float, default=1e-3, dest="swa_lr" + "--swa_lr", + "--stage_two_lr", + help="Learning rate of optimizer in Stage Two (previously called swa)", + type=float, + default=1e-3, + dest="swa_lr", ) parser.add_argument( "--weight_decay", help="weight decay (L2 penalty)", type=float, default=5e-7 @@ -494,14 +505,16 @@ def build_default_arg_parser() -> argparse.ArgumentParser: default=0.9993, ) parser.add_argument( - "--swa", "--stage_two", + "--swa", + "--stage_two", help="use Stage Two loss weight, which decreases the learning rate and increases the energy weight at the end of the training to help converge them", action="store_true", default=False, dest="swa", ) parser.add_argument( - "--start_swa","--start_stage_two", + "--start_swa", + "--start_stage_two", help="Number of epochs before changing to Stage Two loss weights", type=int, default=None, diff --git a/mace/tools/scripts_utils.py b/mace/tools/scripts_utils.py index 5261c124e..98bbc08ae 100644 --- a/mace/tools/scripts_utils.py +++ b/mace/tools/scripts_utils.py @@ -460,6 +460,14 @@ def create_error_table( "relative F RMSE %", "RMSE Stress (Virials) / meV / A (A^3)", ] + elif table_type == "PerAtomMAEstressvirials": + table.field_names = [ + "config_type", + "MAE E / meV / atom", + "MAE F / meV / A", + "relative F MAE %", + "MAE Stress (Virials) / meV / A (A^3)", + ] elif table_type == "TotalMAE": table.field_names = [ "config_type", @@ -564,6 +572,32 @@ def create_error_table( f"{metrics['rmse_virials'] * 1000:.1f}", ] ) + elif ( + table_type == "PerAtomMAEstressvirials" + and metrics["mae_stress"] is not None + ): + table.add_row( + [ + name, + f"{metrics['mae_e_per_atom'] * 1000:.1f}", + f"{metrics['mae_f'] * 1000:.1f}", + f"{metrics['rel_mae_f']:.2f}", + f"{metrics['mae_stress'] * 1000:.1f}", + ] + ) + elif ( + table_type == "PerAtomMAEstressvirials" + and metrics["mae_virials"] is not None + ): + table.add_row( + [ + name, + f"{metrics['mae_e_per_atom'] * 1000:.1f}", + f"{metrics['mae_f'] * 1000:.1f}", + f"{metrics['rel_mae_f']:.2f}", + f"{metrics['mae_virials'] * 1000:.1f}", + ] + ) elif table_type == "TotalMAE": table.add_row( [ diff --git a/mace/tools/train.py b/mace/tools/train.py index 575fb02dc..20256cec9 100644 --- a/mace/tools/train.py +++ b/mace/tools/train.py @@ -71,6 +71,26 @@ def valid_err_log(valid_loss, eval_metrics, logger, log_errors, epoch=None): logging.info( f"Epoch {epoch}: loss={valid_loss:.4f}, RMSE_E_per_atom={error_e:.1f} meV, RMSE_F={error_f:.1f} meV / A, RMSE_virials_per_atom={error_virials:.1f} meV" ) + elif ( + log_errors == "PerAtomMAEstressvirials" + and eval_metrics["mae_stress_per_atom"] is not None + ): + error_e = eval_metrics["mae_e_per_atom"] * 1e3 + error_f = eval_metrics["mae_f"] * 1e3 + error_stress = eval_metrics["mae_stress"] * 1e3 + logging.info( + f"Epoch {epoch}: loss={valid_loss:.4f}, MAE_E_per_atom={error_e:.1f} meV, MAE_F={error_f:.1f} meV / A, MAE_stress={error_stress:.1f} meV / A^3" + ) + elif ( + log_errors == "PerAtomMAEstressvirials" + and eval_metrics["mae_virials_per_atom"] is not None + ): + error_e = eval_metrics["mae_e_per_atom"] * 1e3 + error_f = eval_metrics["mae_f"] * 1e3 + error_virials = eval_metrics["mae_virials"] * 1e3 + logging.info( + f"Epoch {epoch}: loss={valid_loss:.4f}, MAE_E_per_atom={error_e:.1f} meV, MAE_F={error_f:.1f} meV / A, MAE_virials={error_virials:.1f} meV" + ) elif log_errors == "TotalRMSE": error_e = eval_metrics["rmse_e"] * 1e3 error_f = eval_metrics["rmse_f"] * 1e3 From 81c02b46412524a65b0d39ee054832bb43820b49 Mon Sep 17 00:00:00 2001 From: Ilyes Batatia <48651863+ilyes319@users.noreply.github.com> Date: Mon, 12 Aug 2024 10:50:57 +0100 Subject: [PATCH 11/11] fix linter json E0s --- mace/tools/scripts_utils.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/mace/tools/scripts_utils.py b/mace/tools/scripts_utils.py index 98bbc08ae..106cb9b03 100644 --- a/mace/tools/scripts_utils.py +++ b/mace/tools/scripts_utils.py @@ -306,7 +306,8 @@ def get_atomic_energies(E0s, train_collection, z_table) -> dict: else: if E0s.endswith(".json"): logging.info(f"Loading atomic energies from {E0s}") - atomic_energies_dict = json.load(open(E0s, "r")) + with open(E0s, "r", encoding="utf-8") as f: + atomic_energies_dict = json.load(f) else: try: atomic_energies_dict = ast.literal_eval(E0s)