diff --git a/examples/PACKAGES/metatomic/in.nve_norescale b/examples/PACKAGES/metatomic/in.nve_norescale new file mode 100644 index 00000000000..d3b839652e0 --- /dev/null +++ b/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 diff --git a/examples/PACKAGES/metatomic/in.nve_rescale b/examples/PACKAGES/metatomic/in.nve_rescale new file mode 100644 index 00000000000..a23c18d9f56 --- /dev/null +++ b/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 diff --git a/src/KOKKOS/fix_metatomic_kokkos.cpp b/src/KOKKOS/fix_metatomic_kokkos.cpp index 7307bafe84c..5729bc9725c 100644 --- a/src/KOKKOS/fix_metatomic_kokkos.cpp +++ b/src/KOKKOS/fix_metatomic_kokkos.cpp @@ -69,6 +69,10 @@ FixMetatomicKokkos::~FixMetatomicKokkos() {} template void FixMetatomicKokkos::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); diff --git a/src/ML-METATOMIC/fix_metatomic.cpp b/src/ML-METATOMIC/fix_metatomic.cpp index 1bb01b7f344..3d4ea0233cf 100644 --- a/src/ML-METATOMIC/fix_metatomic.cpp +++ b/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 #include +#include #include #include @@ -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 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 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 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 com_old = {0.0, 0.0, 0.0}; std::array 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() { diff --git a/src/ML-METATOMIC/fix_metatomic.h b/src/ML-METATOMIC/fix_metatomic.h index 5c7e87048e4..6ff88594f28 100644 --- a/src/ML-METATOMIC/fix_metatomic.h +++ b/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 system_adaptor; };