diff --git a/README.md b/README.md index 83a02a01e..31d466af4 100644 --- a/README.md +++ b/README.md @@ -50,64 +50,33 @@ A partial documentation is available at: https://mace-docs.readthedocs.io ## Installation -Requirements: +### 1. Requirements: -- Python >= 3.7 -- [PyTorch](https://pytorch.org/) >= 1.12 **(training with float64 is not supported with PyTorch 2.1 but is supported with 2.2 and later.)**. +- 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. -### 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. - -### 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 +### 2b. Installation from source -# (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) +```sh 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 -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 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, diff --git a/mace/cli/run_train.py b/mace/cli/run_train.py index aecc3f713..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, @@ -585,19 +588,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 +620,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 +628,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 +810,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 + "_stagetwo.model") else: model_path = Path(args.checkpoints_dir) / (tag + ".model") logging.info(f"Saving model to {model_path}") @@ -821,10 +824,12 @@ 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 + "_stagetwo.model") + ) try: path_complied = Path(args.model_dir) / ( - args.name + "_swa_compiled.model" + args.name + "_stagetwo_compiled.model" ) logging.info(f"Compiling model, saving metadata {path_complied}") model_compiled = jit.compile(deepcopy(model)) diff --git a/mace/tools/arg_parser.py b/mace/tools/arg_parser.py index 893203aa5..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", @@ -389,45 +390,55 @@ def build_default_arg_parser() -> argparse.ArgumentParser: ) parser.add_argument( "--swa_forces_weight", - help="weight of forces loss after starting swa", + "--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", + "--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", + "--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", + "--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", + "--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 +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", 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 @@ -490,15 +506,19 @@ def build_default_arg_parser() -> argparse.ArgumentParser: ) 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", + "--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_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/scripts_utils.py b/mace/tools/scripts_utils.py index cc7b39291..106cb9b03 100644 --- a/mace/tools/scripts_utils.py +++ b/mace/tools/scripts_utils.py @@ -304,11 +304,18 @@ 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}") + with open(E0s, "r", encoding="utf-8") as f: + atomic_energies_dict = json.load(f) + 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" @@ -454,6 +461,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", @@ -558,6 +573,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 7ebf3ce15..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 @@ -163,7 +183,7 @@ 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 @@ -233,7 +253,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: 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