Add FlashMD energy rescaling to fix metatomic

examples/PACKAGES/metatomic/in.nve_norescale

@@ -0,0 +1,24 @@
+units metal
+atom_style atomic
+boundary p p p
+
+lattice fcc 4.05
+region box block 0 3 0 3 0 3
+create_box 1 box
+create_atoms 1 box
+
+mass 1 26.9815386
+
+velocity all create 800.0 12345 mom yes rot yes dist gaussian
+
+pair_style metatomic flashmd-energy-model.pt device cuda uncertainty_threshold off
+pair_coeff * * 13
+
+timestep 0.016
+
+fix 0 all metatomic flashmd-dynamics-model.pt types 13 device cuda
+
+thermo 5
+thermo_style custom step temp pe ke etotal
+
+run 50

examples/PACKAGES/metatomic/in.nve_rescale

@@ -0,0 +1,24 @@
+units metal
+atom_style atomic
+boundary p p p
+
+lattice fcc 4.05
+region box block 0 3 0 3 0 3
+create_box 1 box
+create_atoms 1 box
+
+mass 1 26.9815386
+
+velocity all create 800.0 12345 mom yes rot yes dist gaussian
+
+pair_style metatomic flashmd-energy-model.pt device cuda uncertainty_threshold off
+pair_coeff * * 13
+
+timestep 0.016
+
+fix 0 all metatomic flashmd-dynamics-model.pt types 13 device cuda rescale_energy on
+
+thermo 5
+thermo_style custom step temp pe ke etotal
+
+run 50

src/KOKKOS/fix_metatomic_kokkos.cpp

@@ -69,6 +69,10 @@ FixMetatomicKokkos<DeviceType>::~FixMetatomicKokkos() {}
 
 template<class DeviceType>
 void FixMetatomicKokkos<DeviceType>::init() {
+    if (this->rescale_energy) {
+        error->all(FLERR, "fix metatomic/kk does not support 'rescale_energy' yet "
+                          "(the Kokkos initial_integrate/post_force overrides bypass the rescale)");
+    }
     FixMetatomic::init();
 
     auto request = neighbor->find_request(this);

src/ML-METATOMIC/fix_metatomic.cpp

@@ -31,6 +31,7 @@
 #include "atom.h"
 #include "memory.h"
 #include "modify.h"
+#include "compute.h"
 #include "error.h"
 #include "group.h"
 #include "force.h"
@@ -42,6 +43,7 @@
 
 #include <vector>
 #include <algorithm>
+#include <cmath>
 
 #include <metatomic/torch.hpp>
 #include <metatensor/torch.hpp>
@@ -93,6 +95,10 @@ FixMetatomic::FixMetatomic(LAMMPS *lmp, int narg, char **arg): Fix(lmp, narg, ar
     this->model_path = arg[3];
     this->requested_device = std::nullopt;
     this->extensions_directory = std::nullopt;
+    this->rescale_energy = false;
+    this->pe_compute = nullptr;
+    this->rescale_U_old = 0.0;
+    this->rescale_K_before = 0.0;
     std::vector<int> parsed_types;
 
     this->mta_data = new FixMetatomicData(std::move(length_unit));
@@ -162,10 +168,23 @@ FixMetatomic::FixMetatomic(LAMMPS *lmp, int narg, char **arg): Fix(lmp, narg, ar
             }
             this->extensions_directory = std::string(arg[iarg + 1]);
             iarg += 2;
+        } else if (strcmp(arg[iarg], "rescale_energy") == 0) {
+            iarg += 1;
+            if (iarg == narg) {
+                error->one(FLERR, "expected <on/off> after 'rescale_energy' in fix metatomic, got nothing");
+            } else if (strcmp(arg[iarg], "on") == 0) {
+                this->rescale_energy = true;
+                iarg += 1;
+            } else if (strcmp(arg[iarg], "off") == 0) {
+                this->rescale_energy = false;
+                iarg += 1;
+            } else {
+                error->one(FLERR, "expected <on/off> after 'rescale_energy' in fix metatomic, got '{}'", arg[iarg]);
+            }
         } else {
             error->all(FLERR,
                 "Illegal fix metatomic command: unrecognized option '{}' (expected "
-                "'types', 'device', `extensions_directory`, or `check_consistency`)",
+                "'types', 'device', `extensions_directory`, `check_consistency`, or `rescale_energy`)",
                 arg[iarg]
             );
         }
@@ -209,10 +228,21 @@ FixMetatomic::FixMetatomic(LAMMPS *lmp, int narg, char **arg): Fix(lmp, narg, ar
     const char* v = std::getenv("LAMMPS_METATOMIC_DISABLE_TORCH_JIT_OPTIMIZATION");
     const bool disable = (v != nullptr) && (std::strcmp(v, "1") == 0);
     if (disable) torch::jit::setGraphExecutorOptimize(false);
+
+    // For energy rescaling we need the potential energy every step; let an
+    // internal "compute pe" tally it (the same mechanism fix npt uses for the
+    // virial via an internal compute pressure).
+    if (this->rescale_energy) {
+        this->pe_compute_id = std::string(this->id) + "_pe";
+        this->pe_compute = modify->add_compute(this->pe_compute_id + " all pe");
+    }
 }
 
 FixMetatomic::~FixMetatomic() {
     memory->destroy(type_mapping);
+    if (this->rescale_energy && modify) {
+        modify->delete_compute(this->pe_compute_id);
+    }
 }
 
 /* ---------------------------------------------------------------------- */
@@ -347,6 +377,27 @@ void FixMetatomic::init() {
         neighbor->binsizeflag = 1;
     }
     // END HACK
+
+    if (this->rescale_energy) {
+        this->pe_compute = modify->get_compute_by_id(this->pe_compute_id);
+        if (!this->pe_compute || this->pe_compute->peflag == 0) {
+            error->all(FLERR, "fix metatomic internal error: potential energy compute unavailable");
+        }
+        if (strcmp(force->pair_style, "none") == 0) {
+            error->all(FLERR,
+                "fix metatomic rescale_energy requires a pair_style providing the potential energy"
+            );
+        }
+    }
+}
+
+void FixMetatomic::setup(int /*vflag*/) {
+    // Seed U(q) for the first step and request the energy tally one step ahead.
+    if (this->rescale_energy) {
+        this->pe_compute->compute_scalar();
+        this->rescale_U_old = this->pe_compute->scalar;
+        this->pe_compute->addstep(update->ntimestep + 1);
+    }
 }
 
 void FixMetatomic::pick_device(c10::Device& device, const char* requested) {
@@ -516,6 +567,7 @@ void FixMetatomic::initial_integrate(int /*vflag*/) {
     std::array<double, 3> com_old = {0.0, 0.0, 0.0};
     std::array<double, 3> com_velocity_old = {0.0, 0.0, 0.0};
     double total_mass = 0.0;
+    double k_before = 0.0;  // kinetic energy of the momenta fed to the model
     for (int i = 0; i < nlocal; i++) {
         if (mask[i] & groupbit) {
             double m_i = rmass ? rmass[i] : mass[type[i]];
@@ -526,8 +578,10 @@ void FixMetatomic::initial_integrate(int /*vflag*/) {
             com_velocity_old[1] += v[i][1] * m_i;
             com_velocity_old[2] += v[i][2] * m_i;
             total_mass += m_i;
+            k_before += m_i * (v[i][0]*v[i][0] + v[i][1]*v[i][1] + v[i][2]*v[i][2]);
         }
     }
+    this->rescale_K_before = 0.5 * force->mvv2e * k_before;
     if (total_mass > 0.0) {
         com_old[0] /= total_mass;
         com_old[1] /= total_mass;
@@ -623,6 +677,45 @@ void FixMetatomic::post_force(int /*vflag*/) {
             f[i][2] = 0.0;
         }
     }
+
+    if (!this->rescale_energy) return;
+
+    // FlashMD energy rescaling (App. C): the pair_style energy model has just
+    // evaluated U(q') for the freshly predicted positions; rescale the predicted
+    // momenta so the total energy matches its value before this step. This runs
+    // after the FlashMD map and before final_integrate's thermostat half-step.
+    double **v = atom->v;
+    double *rmass = atom->rmass;
+    double *mass = atom->mass;
+    int *type = atom->type;
+
+    this->pe_compute->compute_scalar();
+    double U_new = this->pe_compute->scalar;
+
+    double k_new = 0.0;
+    for (int i = 0; i < nlocal; i++) {
+        if (mask[i] & groupbit) {
+            double m_i = rmass ? rmass[i] : mass[type[i]];
+            k_new += m_i * (v[i][0]*v[i][0] + v[i][1]*v[i][1] + v[i][2]*v[i][2]);
+        }
+    }
+    k_new *= 0.5 * force->mvv2e;
+
+    double dE = (U_new + k_new) - (this->rescale_U_old + this->rescale_K_before);
+    double radicand = 1.0 - dE / k_new;
+    if (k_new > 0.0 && radicand > 0.0) {
+        double alpha = std::sqrt(radicand);
+        for (int i = 0; i < nlocal; i++) {
+            if (mask[i] & groupbit) {
+                v[i][0] *= alpha;
+                v[i][1] *= alpha;
+                v[i][2] *= alpha;
+            }
+        }
+    }
+
+    this->rescale_U_old = U_new;  // U(q') becomes U(q) for the next step
+    this->pe_compute->addstep(update->ntimestep + 1);
 }
 
 void FixMetatomic::final_integrate() {

src/ML-METATOMIC/fix_metatomic.h

@@ -38,6 +38,7 @@ namespace at {
 namespace LAMMPS_NS {
 class MetatomicSystemAdaptor;
 class FixMetatomicData;
+class Compute;
 
 class FixMetatomic : public Fix {
 public:
@@ -46,6 +47,7 @@ class FixMetatomic : public Fix {
 
     int setmask() override;
     void init() override;
+    void setup(int) override;
 
     // Integration methods for ML-driven dynamics
     void initial_integrate(int) override;  // ML prediction of positions/momenta
@@ -70,6 +72,15 @@ class FixMetatomic : public Fix {
     // Mapping from LAMMPS atom types to metatomic model types
     int32_t *type_mapping;
 
+    // FlashMD energy rescaling (paper App. C): p' <- alpha p', with
+    // alpha = sqrt(1 - (E'-E)/K'). The potential energy comes from a
+    // pair_style on top of the fix, read via an internal "compute pe".
+    bool rescale_energy;
+    std::string pe_compute_id;
+    Compute *pe_compute;
+    double rescale_U_old;       // U(q) before the current step
+    double rescale_K_before;    // K(p) fed to the model this step
+
     // Helper class to convert between LAMMPS and metatomic representations
     std::unique_ptr<MetatomicSystemAdaptor> system_adaptor;
 };