Line data Source code
1 : /* +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
2 : Copyright (c) 2023-2026 of Amr Alhossary and the KEnRef Authors.
3 :
4 : The kenref module is free software: you can redistribute it and/or modify
5 : it under the terms of the GNU Lesser General Public License as published by
6 : the Free Software Foundation, either version 3 of the License, or
7 : (at your option) any later version.
8 :
9 : The kenref module is distributed in the hope that it will be useful,
10 : but WITHOUT ANY WARRANTY; without even the implied warranty of
11 : MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
12 : GNU Lesser General Public License for more details.
13 :
14 : You should have received a copy of the GNU Lesser General Public License
15 : along with plumed. If not, see <http://www.gnu.org/licenses/>.
16 :
17 : --------------------------------------------------------------------------
18 :
19 : This module is only compiled when PLUMED is configured with --enable-kenref,
20 : in which case it links against the external KEnRef core library
21 : (https://github.com/Smith-Group/KEnRef), which is distributed separately under
22 : the BSD 3-Clause License. That licence is permissive and imposes no additional
23 : restriction on this module or on PLUMED; see LICENSE_CORE.txt in the KEnRef
24 : distribution for its text.
25 : +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++ */
26 : /*
27 : * KEnRefBias.cpp — PLUMED kenref module (frozen frame)
28 : *
29 : * KEnRefBias is a PLUMED Bias/ActionAtomistic that restrains kinetic-ensemble observables (interproton
30 : * relaxation / order parameters) computed by the KEnRef library. It implements kenref::EngineAdapter
31 : * and delegates the whole per-step refinement pipeline (no-jump -> Kabsch fit -> replica gather ->
32 : * energy/derivative -> scatter -> inverse-fit -> force) to kenref::KEnRefDriver, implementing here only
33 : * the handful of PLUMED-specific callbacks (positions, box, force application, inter-replica MPI).
34 : *
35 : * This file holds the STABLE parts of the action (declaration is in KEnRefBias.h via the KEnRef repo):
36 : * registerKeywords, the PLUMED<->Eigen glue, calculate(), the lock/unlock boilerplate, and the action
37 : * registration. The one-time constructor (model + sub-indexing + driver setup) evolves with the KEnRef
38 : * model abstraction, so it is hosted in the KEnRef repository and compiled in via KEnRefBias_setup.cpp.
39 : */
40 :
41 : // The whole module is optional: PLUMED can be configured with --enable-modules=all on a machine that
42 : // has no KEnRef installed, in which case configure leaves __PLUMED_HAS_KENREF undefined and this
43 : // translation unit compiles to nothing. Same pattern as src/pytorch and the ttsketch/ITensor module.
44 : #ifdef __PLUMED_HAS_KENREF
45 :
46 : #include "core/ActionRegister.h"
47 : #include "core/PlumedMain.h"
48 : #include "tools/Vector.h"
49 : #include "tools/Tensor.h"
50 : #include "tools/Communicator.h"
51 : #include "tools/OpenMP.h"
52 :
53 : #include "core/ModelRegistry.h" // registerKeywords iterates the model registry's schemas
54 :
55 : // Unprefixed: kenref_plumed.pc puts the directory holding this header directly on the include path.
56 : // A "plumedinterface/..." form would look to plumedcheck like a cross-module include, and
57 : // plumedinterface is a KEnRef directory, not a PLUMED module.
58 : #include "KEnRefBias.h"
59 :
60 : #include <chrono>
61 : #include <set>
62 :
63 : // Written as two nested namespaces rather than `namespace PLMD::kenref` because plumedcheck matches
64 : // each of them line-by-line ("^ *namespace *PLMD" and "^ *namespace *kenref"); the C++17 nested form
65 : // satisfies only the first. This is also the form every other module uses.
66 : namespace PLMD {
67 : namespace kenref {
68 :
69 : //+PLUMEDOC BIAS KENREF
70 : /*
71 : Restrains an ensemble of replicas against NMR observables (kinetic ensemble refinement).
72 :
73 : KENREF computes an NMR observable as an average over the whole ensemble of replicas, compares it with
74 : experimental data, and applies the resulting restraint force to the atoms. Unlike most biases it acts on
75 : coordinates rather than on collective variables, so it takes **no** `ARG`.
76 :
77 : The energy model is selected by name with `MODEL`:
78 :
79 : | `MODEL` | restrains |
80 : |:-----------|:---------------------------------------------------|
81 : | `SIGMA` | cross-relaxation rates (interproton NOE build-up) |
82 : | `PLATEAUS` | NOE plateau values |
83 : | `RELAX` | longitudinal / transverse relaxation |
84 :
85 : Because the observable is an ensemble average, `RELAX` needs at least two replicas (run with `--multi`);
86 : `SIGMA` and `PLATEAUS` also work with a single replica.
87 :
88 : `EXP_DATA_FOLDER` holds the experimental data and the atom-pair definitions. `ATOMNAME_MAPPING` is a PDB
89 : that maps atom names to indices, and is also used as the reference structure when `REF` is omitted. When
90 : `FIT_TO_REFERENCE` is set the coordinates are Kabsch-fitted onto the reference, using `GUIDE_ATOMS`,
91 : before the observable is computed, and the derivatives are rotated back afterwards. `SATURATE_FORCES`
92 : clamps the resulting force to `MAX_FORCE`.
93 :
94 : The restraint needs whole molecules: the fit uses all `GUIDE_ATOMS` at once and the atom-pair distances
95 : have no cutoff. At every step KENREF therefore moves each atom it reads to the periodic image nearest its
96 : neighbour in a spanning tree built from the reference structure. The reference must be whole itself; if
97 : one edge of that tree is not shorter than half the smallest box width, KENREF stops with an error.
98 : Coordinates that are already whole are left unchanged.
99 :
100 : The action outputs its restraint energy as `.energy` (equal to `.bias`) and, as a diagnostic, the RMSD
101 : from the reference as `.rmsd`.
102 :
103 : This module needs the external KEnRef library and is not compiled by default; see the module
104 : documentation for how to build it.
105 :
106 : ## Examples
107 :
108 : Restraining a fragment of GB3 against cross-relaxation rates:
109 :
110 : ```plumed
111 : #SETTINGS MOLFILE=regtest/kenref/rt-kenref-sigma/gb3_frag.pdb
112 : kenref: KENREF ...
113 : MODEL=SIGMA
114 : K=1.0
115 : N=0.25
116 : PROTON_MHZ=700.0
117 : EXP_DATA_FOLDER=regtest/kenref/rt-kenref-sigma/
118 : REF=regtest/kenref/rt-kenref-sigma/gb3_frag.pdb
119 : ATOMNAME_MAPPING=regtest/kenref/rt-kenref-sigma/gb3_frag.pdb
120 : GUIDE_ATOMS=1,5,18,20,22,35,37,39,56,58,60,78,80,82,97
121 : MAX_FORCE=999
122 : FIT_TO_REFERENCE
123 : SATURATE_FORCES
124 : ...
125 :
126 : PRINT ARG=kenref.bias,kenref.energy,kenref.rmsd FILE=colvar
127 : ```
128 :
129 : */
130 : //+ENDPLUMEDOC
131 :
132 : PLUMED_REGISTER_ACTION(KEnRefBias, "KENREF")
133 :
134 : // ============================================================
135 : // registerKeywords
136 : // ============================================================
137 19 : void KEnRefBias::registerKeywords(Keywords &keys) {
138 19 : Bias::registerKeywords(keys);
139 19 : ActionAtomistic::registerKeywords(keys);
140 38 : keys.setDisplayName("KENREF");
141 :
142 : // KENREF is atom-based: it biases on coordinates, never on a CV, so it takes no ARG. PLUMED
143 : // 2.11-dev made ARG compulsory for every Bias, which forced input files to carry an empty
144 : // `ARG=` on master while 2.10 required it to be ABSENT -- the same input could not serve both.
145 : // Dropping the keyword makes one input file work everywhere. Guarded by exists() so this stays
146 : // correct on branches whose Bias never added ARG in the first place.
147 19 : if (keys.exists("ARG")) {
148 38 : keys.remove("ARG");
149 : }
150 :
151 : // ---- General (framework-owned) keywords ----
152 19 : keys.add("compulsory", "MODEL", "SIGMA", "The energy model to use (e.g. SIGMA or PLATEAUS)");
153 19 : keys.add("compulsory", "K", "1.0", "Force constant");
154 19 : keys.add("compulsory", "N", "0.25", "Power scaling factor");
155 19 : keys.add("optional", "MAX_FORCE", "Maximum force magnitude (default 9999)");
156 19 : keys.add("atoms", "GUIDE_ATOMS", "Atoms used for alignment to reference");
157 19 : keys.add("optional", "REF", "Reference structure PDB for alignment");
158 19 : keys.add("compulsory", "ATOMNAME_MAPPING", "undefined",
159 : "PDB file with atom-name -> atom-index mapping (also used as reference if REF is omitted)");
160 19 : keys.addFlag("FIT_TO_REFERENCE", false, "Fit coordinates to reference before calculating energy");
161 19 : keys.addFlag("SATURATE_FORCES", false, "Clamp forces to MAX_FORCE");
162 :
163 : // ---- Model-specific keywords (the union across every registered KEnRef model, optional; the
164 : // selected model validates required-ness in buildCache). Deduplicated by key. ----
165 19 : ::kenref::bootstrapModels();
166 : std::set<std::string> seen;
167 76 : for (const auto &[name, schema]: ::kenref::ModelRegistry<KEnRef_Real_t>::allSchemas()) {
168 171 : for (const auto &spec: schema.specs()) {
169 114 : if (seen.insert(spec.key).second) {
170 76 : keys.add("optional", spec.key, spec.doc);
171 : }
172 : }
173 19 : }
174 :
175 : // ---- Output components ----
176 38 : keys.addOutputComponent("energy", "default", "scalar", "Total KEnRef restraint energy");
177 38 : keys.addOutputComponent("rmsd", "default", "scalar", "RMSD from reference (after fitting)");
178 :
179 19 : keys.use("RESTART");
180 19 : keys.use("UPDATE_FROM");
181 19 : keys.use("UPDATE_UNTIL");
182 19 : }
183 :
184 : // ============================================================
185 : // Helpers: current PLUMED positions (nm) -> Angstrom Eigen
186 : // ============================================================
187 30 : CoordsMatrixType<KEnRef_Real_t> KEnRefBias::getGuideAtomsX() const {
188 : // Periodically repaired positions, not getPositions(): see makeRequestedAtomsWhole().
189 : const std::vector<Vector> &pos = wholePositions_;
190 30 : const int nGuide = static_cast<int>(guideAtoms_.size());
191 30 : CoordsMatrixType<KEnRef_Real_t> gx(nGuide, 3);
192 480 : for (int i = 0; i < nGuide; ++i) {
193 450 : const int li = serial_to_localIdx_.at(guideAtoms_[i].serial());
194 450 : gx(i, 0) = to<KEnRef_Real_t>(pos[li][0]) * 10;
195 450 : gx(i, 1) = to<KEnRef_Real_t>(pos[li][1]) * 10;
196 450 : gx(i, 2) = to<KEnRef_Real_t>(pos[li][2]) * 10;
197 : }
198 30 : return gx;
199 : }
200 :
201 15 : void KEnRefBias::fillSubAtomsX(CoordsMatrixType<KEnRef_Real_t> &out) const {
202 : // Periodically repaired positions, not getPositions(): see makeRequestedAtomsWhole().
203 : const std::vector<Vector> &pos = wholePositions_;
204 495 : for (int i = 0; i < out.rows(); ++i) {
205 480 : const int serial = sub0Id_to_global1Serial_[i];
206 480 : const int li = serial_to_localIdx_.at(serial);
207 480 : out(i, 0) = to<KEnRef_Real_t>(pos[li][0]) * 10; // nm -> Angstrom
208 480 : out(i, 1) = to<KEnRef_Real_t>(pos[li][1]) * 10;
209 480 : out(i, 2) = to<KEnRef_Real_t>(pos[li][2]) * 10;
210 : }
211 15 : }
212 :
213 : // ============================================================
214 : // kenref::EngineAdapter implementation (PLUMED-specific I/O + replica MPI)
215 : // ============================================================
216 39 : std::optional<std::string> KEnRefBias::getRawParam(const std::string &key) const {
217 : const auto it = modelParams_.find(key);
218 39 : if (it == modelParams_.end() || it->second.empty()) {
219 15 : return std::nullopt;
220 : }
221 24 : return it->second;
222 : }
223 :
224 15 : int KEnRefBias::numOmpThreads() const {
225 15 : return static_cast<int>(OpenMP::getNumThreads());
226 : }
227 :
228 15 : void KEnRefBias::getLocalModelX(int /*localModel*/, CoordsMatrixType<KEnRef_Real_t> &guideX,
229 : CoordsMatrixType<KEnRef_Real_t> &subX,
230 : Eigen::Matrix<KEnRef_Real_t, 3, 3> &box) const {
231 : /* Repair first, for the WHOLE requested set, before a single coordinate is read. KEnRef restrains a
232 : * whole molecule -- global Kabsch fit, pair distances with no cutoff -- so a molecule split across a
233 : * periodic boundary would give a different energy and different forces. */
234 15 : makeRequestedAtomsWhole();
235 15 : guideX = getGuideAtomsX();
236 15 : subX.resize(numSubAtoms_, 3);
237 15 : fillSubAtomsX(subX);
238 : // PLUMED box Tensor (nm) -> Eigen 3x3 (raw; the driver scales it to Angstrom via toAngstrom).
239 15 : const Tensor &b = getPbc().getBox();
240 60 : for (int i = 0; i < 3; ++i)
241 180 : for (int j = 0; j < 3; ++j) {
242 135 : box(i, j) = to<KEnRef_Real_t>(b[i][j]);
243 : }
244 15 : }
245 :
246 15 : void KEnRefBias::addLocalModelDerivatives(int /*localModel*/, const CoordsMatrixType<KEnRef_Real_t> &derivs) {
247 : // derivs are already inverse-fitted, unit-scaled and saturated by the driver. KEnRef returns
248 : // dE/dx; the applied force is -dE/dx, so we negate.
249 495 : for (int i = 0; i < derivs.rows(); ++i) {
250 480 : const int serial = sub0Id_to_global1Serial_[i];
251 480 : const int li = serial_to_localIdx_.at(serial);
252 : const Vector d(static_cast<double>(derivs(i, 0)),
253 : static_cast<double>(derivs(i, 1)),
254 480 : static_cast<double>(derivs(i, 2)));
255 480 : addForce(getValueIndices(getAbsoluteIndex(li)), -d);
256 : }
257 15 : }
258 :
259 15 : void KEnRefBias::gatherFittedSubAtomsX(const std::vector<CoordsMatrixType<KEnRef_Real_t>> &localFitted,
260 : std::vector<CoordsMatrixType<KEnRef_Real_t>> &all) const {
261 15 : if (!isMultiSim_) {
262 3 : all = localFitted; // single replica owns the whole ensemble
263 3 : return;
264 : }
265 : /* ONLY a replica's main rank may touch multi_sim_comm: PLUMED populates it under
266 : * `if (MAIN(cr))`, and elsewhere it is MPI_COMM_SELF, whose Get_size() silently returns 1. A
267 : * collective over it from a non-main rank is meaningless rather than diagnosable (PLUMED #1443;
268 : * #1484 makes debug builds say so). Every rank of a replica fits the same atoms -- the request
269 : * list, and therefore getPositions(), is complete on every rank -- so localFitted is identical
270 : * across the replica and gathering it from the main rank alone loses nothing.
271 : *
272 : * Returning early is safe because `all` is read only where simulationIndex() == 0, which is a
273 : * main rank by construction. */
274 12 : if (!isReplicaMainRank_) {
275 : return;
276 : }
277 : const auto &fitted = localFitted[0];
278 8 : const int subSize = static_cast<int>(fitted.size()); // rows*3
279 8 : multi_sim_comm.Gather(const_cast<KEnRef_Real_t *>(fitted.data()), subSize,
280 : allSimulationsSubAtomsX_buffer_.data(), subSize, 0);
281 8 : if (simulationIndex() == 0) {
282 : const Eigen::Index nSub = fitted.rows();
283 4 : all.resize(numSimulations_);
284 12 : for (int m = 0; m < numSimulations_; ++m) { // zero-copy view materialised by the assignment
285 8 : all[m] = CoordsMapTypeConst<KEnRef_Real_t>(&allSimulationsSubAtomsX_buffer_[m * nSub * 3], nSub, 3);
286 : }
287 : }
288 : }
289 :
290 15 : void KEnRefBias::scatterModelDerivatives(const std::vector<CoordsMatrixType<KEnRef_Real_t>> &allPerModel,
291 : std::vector<CoordsMatrixType<KEnRef_Real_t>> &localPerModel) const {
292 : /* EVERY RANK NEEDS THE FINISHED DERIVATIVES, even though only one rank computed them.
293 : * addLocalModelDerivatives() walks every restrained atom and calls addForce for it; PLUMED then
294 : * writes those forces through g2l, so a rank's contribution lands only on the atoms it homes
295 : * (DomainDecomposition keeps `unique` local-only, and apply() indexes through g2l). That makes
296 : * the redundant call on every rank harmless -- but it also means a rank with no derivatives
297 : * silently contributes nothing for its own atoms. Hence the broadcast down `comm` in both
298 : * branches below: it is what makes the forces complete, not an optimisation. */
299 15 : const int subSize = numSubAtoms_ * 3;
300 :
301 15 : if (!isMultiSim_) {
302 : /* One replica, but possibly several ranks over it. Only simulationIndex() == 0 -- this
303 : * replica's main rank -- ran the model, so allPerModel is populated there and EMPTY elsewhere;
304 : * copying it would leave the other ranks indexing an empty vector. */
305 3 : localPerModel.resize(1);
306 3 : if (simulationIndex() == 0) {
307 3 : std::copy_n(allPerModel[0].data(), subSize, derivatives_buffer_.data());
308 : }
309 3 : if (comm.Get_size() > 1) {
310 0 : comm.Bcast(derivatives_buffer_.data(), subSize, 0);
311 : }
312 3 : localPerModel[0] = CoordsMapType<KEnRef_Real_t>(derivatives_buffer_.data(), numSubAtoms_, 3);
313 3 : return;
314 : }
315 :
316 : /* Several replicas. Only the replica main ranks are in multi_sim_comm, so only they may pack and
317 : * scatter; the ensemble master is the one that also has replicaIndex_ == 0. */
318 12 : if (isReplicaMainRank_) {
319 8 : if (simulationIndex() == 0)
320 12 : for (int m = 0; m < numSimulations_; ++m) {
321 8 : std::copy_n(allPerModel[m].data(), subSize, &allDerivatives_buffer_[m * subSize]);
322 : }
323 8 : multi_sim_comm.Scatter(allDerivatives_buffer_.data(), subSize, derivatives_buffer_.data(), subSize, 0);
324 : }
325 : // Push this replica's block from its main rank down to the rest of the replica's ranks.
326 12 : if (comm.Get_size() > 1) {
327 8 : comm.Bcast(derivatives_buffer_.data(), subSize, 0);
328 : }
329 12 : localPerModel.resize(1);
330 12 : localPerModel[0] = CoordsMapType<KEnRef_Real_t>(derivatives_buffer_.data(), numSubAtoms_, 3);
331 : }
332 :
333 : // ============================================================
334 : // calculate()
335 : // ============================================================
336 15 : void KEnRefBias::calculate() {
337 15 : const auto begin = std::chrono::high_resolution_clock::now();
338 15 : const long step = getStep();
339 :
340 : /* The replica state is NOT refreshed here any more. It used to be read from multi_sim_comm every
341 : * step, which is wrong on every rank that was never given that communicator: unset means
342 : * MPI_COMM_SELF, so Get_size() returns 1 and Get_rank() 0 in silence, and each non-main rank
343 : * restrained its own domain against a one-replica ensemble (PLUMED #1443). It now comes from
344 : * getNumberOfReplicas()/getReplicaIndex() at setup -- correct on every rank, communication-free,
345 : * and fixed for the whole run, so there is nothing to refresh. */
346 :
347 : // Run the whole engine-agnostic pipeline (no-jump -> fit -> gather -> compute -> scatter ->
348 : // inverse-fit -> unit-scale -> saturate -> addForce) through the adapter callbacks above.
349 15 : const KEnRef_Real_t energy = driver_->step(*this, step % 10 == 0);
350 :
351 : // ---- RMSD (crude diagnostic; current guide atoms vs reference, Angstrom -> nm) ----
352 : double rmsd = 0.0;
353 15 : if (guideAtomsReferenceCoords_.rows() > 0) {
354 30 : const CoordsMatrixType<KEnRef_Real_t> diff = getGuideAtomsX() - guideAtomsReferenceCoords_;
355 15 : rmsd = std::sqrt(diff.rowwise().squaredNorm().mean()) * 0.1;
356 : }
357 :
358 : // ---- output components ----
359 : const double bias_value = energy;
360 15 : getPntrToComponent("bias")->set(bias_value);
361 15 : getPntrToComponent("energy")->set(bias_value);
362 15 : getPntrToComponent("rmsd")->set(rmsd);
363 15 : setBias(bias_value);
364 :
365 15 : const auto end = std::chrono::high_resolution_clock::now();
366 15 : calculate_time_ += std::chrono::duration_cast<std::chrono::nanoseconds>(end - begin).count();
367 : // One line per job, not one per domain: simulationIndex() is 0 on exactly one rank in the run.
368 11 : if (simulationIndex() == 0 && step % 10 == 0) {
369 14 : std::cout << " Step " << step << " KEnRef energy = " << energy << std::endl;
370 : }
371 15 : }
372 :
373 : // ============================================================
374 : // Ambiguity resolution (lock both Atoms and Arguments/CVs)
375 : // ============================================================
376 15 : void KEnRefBias::lockRequests() {
377 : ActionAtomistic::lockRequests();
378 : Bias::lockRequests();
379 15 : }
380 :
381 15 : void KEnRefBias::unlockRequests() {
382 : ActionAtomistic::unlockRequests();
383 : Bias::unlockRequests();
384 15 : }
385 :
386 0 : void KEnRefBias::calculateNumericalDerivatives(ActionWithValue *a) {
387 0 : ActionAtomistic::calculateNumericalDerivatives(a);
388 0 : }
389 :
390 : } // namespace kenref
391 : } // namespace PLMD
392 :
393 : #endif // __PLUMED_HAS_KENREF
|
