merge changes from main

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
 

mace/__version__.py

@@ -1 +1 @@
-__version__ = "0.3.5"
+__version__ = "0.3.6"

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,

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,15 +620,15 @@ 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(
                 energy_weight=args.swa_energy_weight,
                 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))

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",

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(
                 [

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: