Committee error

mace/calculators/mace.py

@@ -11,6 +11,9 @@
 
 from mace import data
 from mace.tools import torch_geometric, torch_tools, utils
+from typing import Union
+from glob import glob
+import numpy as np
 
 
 class MACECalculator(Calculator):
@@ -25,7 +28,7 @@ def __init__(
         energy_units_to_eV: float = 1.0,
         length_units_to_A: float = 1.0,
         default_dtype="float64",
-        **kwargs
+        **kwargs,
     ):
         Calculator.__init__(self, **kwargs)
         self.results = {}
@@ -106,7 +109,7 @@ def __init__(
         length_units_to_A: float = 1.0,
         default_dtype="float64",
         charges_key="Qs",
-        **kwargs
+        **kwargs,
     ):
         """
         :param charges_key: str, Array field of atoms object where atomic charges are stored
@@ -180,7 +183,7 @@ def __init__(
         length_units_to_A: float = 1.0,
         default_dtype="float64",
         charges_key="Qs",
-        **kwargs
+        **kwargs,
     ):
         """
         :param charges_key: str, Array field of atoms object where atomic charges are stored
@@ -250,3 +253,112 @@ def calculate(self, atoms=None, properties=None, system_changes=all_changes):
             self.results["stress"] = (
                 stress * (self.energy_units_to_eV / self.length_units_to_A**3)
             )[0]
+
+
+class MACECommitteeCalculator(Calculator):
+    """MACE ASE Committee Calculator
+
+    This calculator can be used to obtain energy and force errors from a committee of
+    MACE models. To load multiple committees, either pass a list of file names or a
+    single string with a wildcard to the model_paths argument. This calculator can also
+    be used to load single models.
+    """
+
+    implemented_properties = ["energy", "free_energy", "forces", "stress"]
+
+    def __init__(
+        self,
+        model_paths: Union[list, str],
+        device: str,
+        energy_units_to_eV: float = 1.0,
+        length_units_to_A: float = 1.0,
+        default_dtype="float64",
+        **kwargs,
+    ):
+        Calculator.__init__(self, **kwargs)
+        self.results = {}
+
+        if type(model_paths) == str:
+            # Find all models that staisfy the wildcard (e.g. mace_model_*.pt)
+            model_paths_glob = glob(model_paths)
+
+            if len(model_paths_glob) == 0:
+                raise ValueError(f"Couldn't find MACE model files: {model_paths}")
+            else:
+                model_paths = model_paths_glob
+        if len(model_paths) == 0:
+            raise ValueError(f"No mace file neames supplied")
+        elif len(model_paths) > 1:
+            print(f"Running committee mace with {len(model_paths)} models")
+
+        # Load models
+        self.models = [
+            torch.load(f=model_path, map_location=device) for model_path in model_paths
+        ]
+
+        r_maxs = [model.r_max.cpu() for model in self.models]
+        r_maxs = np.array(r_maxs)
+        assert np.all(
+            r_maxs == r_maxs[0]
+        ), f"committee r_max are not all the same {' '.join(r_maxs)}"
+        self.r_max = r_maxs[0]
+
+        self.device = torch_tools.init_device(device)
+        self.energy_units_to_eV = energy_units_to_eV
+        self.length_units_to_A = length_units_to_A
+        self.z_table = utils.AtomicNumberTable(
+            [int(z) for z in self.models[0].atomic_numbers]
+        )
+
+        torch_tools.set_default_dtype(default_dtype)
+
+    # pylint: disable=dangerous-default-value
+    def calculate(self, atoms=None, properties=None, system_changes=all_changes):
+        """
+        Calculate properties.
+        :param atoms: ase.Atoms object
+        :param properties: [str], properties to be computed, used by ASE internally
+        :param system_changes: [str], system changes since last calculation, used by
+        ASE internally
+        :return:
+        """
+        # call to base-class to set atoms attribute
+        Calculator.calculate(self, atoms)
+
+        # prepare data
+        config = data.config_from_atoms(atoms)
+        data_loader = torch_geometric.dataloader.DataLoader(
+            dataset=[
+                data.AtomicData.from_config(
+                    config, z_table=self.z_table, cutoff=self.r_max
+                )
+            ],
+            batch_size=1,
+            shuffle=False,
+            drop_last=False,
+        )
+        batch = next(iter(data_loader)).to(self.device)
+
+        # predict + extract data
+        energies, forces = [], []
+        for i, model in enumerate(self.models):
+            # Otherwise: RuntimeError: you can only change requires_grad flags of leaf variables.
+            batch = next(iter(data_loader)).to(self.device)
+            out = model(batch, compute_stress=True)
+            energies.append(out["energy"].detach().cpu().item())
+            forces.append(out["forces"].detach().cpu().numpy())
+
+        # convert_units
+        energies = np.array(energies) * self.energy_units_to_eV
+        # force has units eng / len:
+        forces = np.array(forces) * self.energy_units_to_eV / self.length_units_to_A
+        # store results
+        self.results = {
+            "energies": energies,
+            "forcess": forces,
+            "energy": np.mean(energies),
+            "free_energy": np.mean(energies),
+            "energy_var": np.var(energies),
+            "forces": np.mean(forces, axis=0),
+            "forces_var": np.var(forces, axis=0),
+        }

scripts/active_learning_md.py

@@ -0,0 +1,181 @@
+# /home/lls34/rds/hpc-work/GitHub/CarbonCaptureNMR/systems/02_KHCO3/01-First-Develop/01-Bulk/MACE/hkco3-it0-1e-5-3.model
+"""Demonstrates molecular dynamics with constant temperature."""
+import argparse
+from ase.md.langevin import Langevin
+from ase.md.velocitydistribution import MaxwellBoltzmannDistribution
+from ase import units
+
+from mace.calculators.mace import MACECommitteeCalculator
+import ase.io
+import os
+import numpy as np
+import time
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--config", help="path to XYZ configurations", required=True)
+    parser.add_argument(
+        "--config_index", help="index of configuration", type=int, default=-1
+    )
+    parser.add_argument("--error_threshold", help="error threshold", default=0.1)
+    parser.add_argument("--temperature_K", help="temperature", type=float, default=300)
+    parser.add_argument("--friction", help="friction", type=float, default=0.01)
+    parser.add_argument("--timestep", help="timestep", type=float, default=1)
+    parser.add_argument("--nsteps", help="number of steps", type=int, default=1000)
+    parser.add_argument(
+        "--nprint", help="number of steps between prints", type=int, default=10
+    )
+    parser.add_argument(
+        "--nsave", help="number of steps between saves", type=int, default=10
+    )
+    parser.add_argument(
+        "--ncheckerror", help="number of steps between saves", type=int, default=10
+    )
+
+    parser.add_argument(
+        "--model",
+        help="path to model. Use wildcards to add multiple models as committe eg "
+        "(`mace_*.model` to load mace_1.model, mace_2.model) ",
+        required=True,
+    )
+    parser.add_argument("--output", help="output path", required=True)
+    parser.add_argument(
+        "--device",
+        help="select device",
+        type=str,
+        choices=["cpu", "cuda"],
+        default="cuda",
+    )
+    parser.add_argument(
+        "--default_dtype",
+        help="set default dtype",
+        type=str,
+        choices=["float32", "float64"],
+        default="float64",
+    )
+    parser.add_argument(
+        "--compute_stress",
+        help="compute stress",
+        action="store_true",
+        default=False,
+    )
+    parser.add_argument(
+        "--info_prefix",
+        help="prefix for energy, forces and stress keys",
+        type=str,
+        default="MACE_",
+    )
+    return parser.parse_args()
+
+
+def printenergy(dyn, start_time=None):  # store a reference to atoms in the definition.
+    """Function to print the potential, kinetic and total energy."""
+    a = dyn.atoms
+    epot = a.get_potential_energy() / len(a)
+    ekin = a.get_kinetic_energy() / len(a)
+    if start_time is None:
+        elapsed_time = 0
+    else:
+        elapsed_time = time.time() - start_time
+    print(
+        "%.1fs: Energy per atom: Epot = %.3feV  Ekin = %.3feV (T=%3.0fK)  "
+        "Etot = %.3feV t=%.1ffs Eerr = %.3feV Ferr = %.3feV/A"
+        % (
+            elapsed_time,
+            epot,
+            ekin,
+            ekin / (1.5 * units.kB),
+            epot + ekin,
+            dyn.get_time() / units.fs,
+            a.calc.results["energy_var"],
+            np.max(np.linalg.norm(a.calc.results["forces_var"], axis=1)),
+        ),
+        flush=True,
+    )
+
+
+def save_config(dyn, fname):
+    atomsi = dyn.atoms
+    ens = atomsi.get_potential_energy()
+    frcs = atomsi.get_forces()
+
+    atomsi.info.update(
+        {
+            "mlff_energy": ens,
+            "time": np.round(dyn.get_time() / units.fs, 5),
+            "mlff_energy_var": atomsi.calc.results["energy_var"],
+        }
+    )
+    atomsi.arrays.update(
+        {"mlff_forces": frcs, "mlff_forces_var": atomsi.calc.results["forces_var"]}
+    )
+
+    ase.io.write(fname, atomsi, append=True)
+
+
+def stop_error(dyn, threshold, reg=0.2):
+    atomsi = dyn.atoms
+    force_var = atomsi.calc.results["forces_var"]
+    force = atomsi.get_forces()
+    ferr = np.sqrt(np.sum(force_var, axis=1))
+    ferr_rel = ferr / (np.linalg.norm(force, axis=1) + reg)
+
+    if np.max(ferr_rel) > threshold:
+        print(
+            "Error too large {:.3}. Stopping t={:.2} fs.".format(
+                np.max(ferr_rel), dyn.get_time() / units.fs
+            ),
+            flush=True,
+        )
+        dyn.max_steps = 0
+
+
+def main():
+    args = parse_args()
+
+    mace_fname = args.model
+    atoms_fname = args.config
+    atoms_index = args.config_index
+
+    mace_calc = MACECommitteeCalculator(
+        mace_fname,
+        args.device,
+        default_dtype=args.default_dtype,
+    )
+
+    NSTEPS = args.nsteps
+
+    if os.path.exists(args.output):
+        print("Trajectory exists. Continuing from last step.")
+        atoms = ase.io.read(args.output, index=-1)
+        len_save = len(ase.io.read(args.output, ":"))
+        print("Last step: ", atoms.info["time"], "Number of configs: ", len_save)
+        NSTEPS -= len_save * args.nsave
+    else:
+        atoms = ase.io.read(atoms_fname, index=atoms_index)
+        MaxwellBoltzmannDistribution(atoms, temperature_K=args.temperature_K)
+
+    atoms.calc = mace_calc
+
+    # We want to run MD with constant energy using the Langevin algorithm
+    # with a time step of 5 fs, the temperature T and the friction
+    # coefficient to 0.02 atomic units.
+    dyn = Langevin(
+        atoms=atoms,
+        timestep=args.timestep * units.fs,
+        temperature_K=args.temperature_K,
+        friction=args.friction,
+    )
+
+    dyn.attach(printenergy, interval=args.nsave, dyn=dyn, start_time=time.time())
+    dyn.attach(save_config, interval=args.nsave, dyn=dyn, fname=args.output)
+    dyn.attach(
+        stop_error, interval=args.ncheckerror, dyn=dyn, threshold=args.error_threshold
+    )
+    # Now run the dynamics
+    dyn.run(NSTEPS)
+
+
+if __name__ == "__main__":
+    main()