Major update to how the plumed module collects hills work

This avoids a lot of unnecessary computations, and one can switch off
the calculation of total work to save even more.
Note that this change makes calculations somewhat less robust, as the
assumption mtd_update is called right after having updated the bias
i-pi-side is even more important.
However this was an assumption all along, so overall this is just being
more consequential about it.

demos/ensemble-deltamu/README.md

@@ -130,7 +130,7 @@ The relevant section in the `input.xml` file is
 ```
   <ffplumed name="plumed">
       <file mode="xyz">init.xyz</file>
-      <plumeddat>plumed.dat </plumeddat>
+      <plumed_dat>plumed.dat </plumed_dat>
   </ffplumed>
 ```
 

demos/ensemble-deltamu/input_n2p2.xml

@@ -26,7 +26,7 @@
   </ffcommittee>
   <ffplumed name="plumed">                                                    
       <file mode="xyz">init_ase.xyz</file>
-      <plumeddat>plumed.dat </plumeddat>                              
+      <plumed_dat>plumed.dat </plumed_dat>                              
   </ffplumed> 
    <system>
       <initialize nbeads="1">

demos/ensemble-deltamu/input_rascaline.xml

@@ -17,7 +17,7 @@
   </ffcommittee>
   <ffplumed name="plumed">                                                    
       <file mode="xyz">init_ase.xyz</file>
-      <plumeddat>plumed.dat </plumeddat>
+      <plumed_dat>plumed.dat </plumed_dat>
   </ffplumed> 
    <system>
       <initialize nbeads="1">

examples/features/metadynamics/pimd_metadynamics_zundel/input.xml

@@ -8,7 +8,7 @@
    </ffsocket>
    <ffplumed name="plumed">
          <file mode="xyz">./h5o2+.xyz</file>
-         <plumeddat> plumed/plumed.dat </plumeddat>
+         <plumed_dat> plumed/plumed.dat </plumed_dat>
          <plumed_extras> [doo, co1.lessthan, co2.lessthan, mtd.bias ] </plumed_extras>
    </ffplumed>
    <total_steps>400</total_steps>

examples/features/metadynamics/remd_metadynamics_zundel/input.xml

@@ -8,15 +8,15 @@
    </ffsocket>
    <ffplumed name="plumed-200K">
          <file mode="xyz">./h5o2+.xyz</file>
-         <plumeddat> plumed_200K/plumed.dat </plumeddat>
+         <plumed_dat> plumed_200K/plumed.dat </plumed_dat>
    </ffplumed>
    <ffplumed name="plumed-250K">
          <file mode="xyz">./h5o2+.xyz</file>
-         <plumeddat> plumed_250K/plumed.dat </plumeddat>
+         <plumed_dat> plumed_250K/plumed.dat </plumed_dat>
    </ffplumed>
    <ffplumed name="plumed-320K">
          <file mode="xyz">./h5o2+.xyz</file>
-         <plumeddat> plumed_320K/plumed.dat </plumeddat>
+         <plumed_dat> plumed_320K/plumed.dat </plumed_dat>
    </ffplumed>
    <total_steps>20000</total_steps>
    <output prefix="data">

examples/features/metadynamics/wte_metadynamics_zundel/input.xml

@@ -8,7 +8,7 @@
    </ffsocket>
    <ffplumed name="plumed">
          <file mode="xyz">./h5o2+.xyz</file>
-         <plumeddat> plumed/plumed.dat </plumeddat>
+         <plumed_dat> plumed/plumed.dat </plumed_dat>
    </ffplumed>
    <total_steps>400</total_steps>
    <output prefix="data">

ipi/engine/forcefields.py

@@ -720,8 +720,9 @@ def __init__(
         dopbc=False,
         threaded=False,
         init_file="",
-        plumeddat="",
-        plumedstep=0,
+        compute_work=True,
+        plumed_dat="",
+        plumed_step=0,
         plumed_extras=[],
     ):
         """Initialises FFPlumed.
@@ -742,9 +743,10 @@ def __init__(
             latency, offset, name, pars, dopbc=False, threaded=threaded
         )
         self.plumed = plumed.Plumed()
-        self.plumeddat = plumeddat
-        self.plumedstep = plumedstep
+        self.plumed_dat = plumed_dat
+        self.plumed_step = plumed_step
         self.plumed_extras = plumed_extras
+        self.compute_work = compute_work
         self.init_file = init_file
 
         if self.init_file.mode == "xyz":
@@ -758,7 +760,7 @@ def __init__(
         self.natoms = myatoms.natoms
         self.plumed.cmd("setRealPrecision", 8)  # i-PI uses double precision
         self.plumed.cmd("setMDEngine", "i-pi")
-        self.plumed.cmd("setPlumedDat", self.plumeddat)
+        self.plumed.cmd("setPlumedDat", self.plumed_dat)
         self.plumed.cmd("setNatoms", self.natoms)
         timeunit = 2.4188843e-05  # atomic time to ps
         self.plumed.cmd("setMDTimeUnits", timeunit)
@@ -773,7 +775,7 @@ def __init__(
             "setMDLengthUnits", 0.052917721
         )  # Pass a pointer to the conversion factor between the length unit used in your code and nm
         self.plumedrestart = False
-        if self.plumedstep > 0:
+        if self.plumed_step > 0:
             # we are restarting, signal that PLUMED should continue
             self.plumedrestart = True
             self.plumed.cmd("setRestart", 1)
@@ -821,7 +823,7 @@ def evaluate(self, r):
         # linking with the current value in simulations is non-trivial, as masses
         # are not expected to be the force evaluator's business, and charges are not
         # i-PI's business.
-        self.plumed.cmd("setStep", self.plumedstep)
+        self.plumed.cmd("setStep", self.plumed_step)
         self.plumed.cmd("setCharges", self.charges)
         self.plumed.cmd("setMasses", self.masses)
 
@@ -863,27 +865,31 @@ def mtd_update(self, pos, cell):
         """Makes updates to the potential that only need to be triggered
         upon completion of a time step."""
 
-        self.plumedstep += 1
-        f = np.zeros((self.natoms, 3))
-        vir = np.zeros((3, 3))
+        # NB - this assumes this is called at the end of a step,
+        # when the bias has already been computed to integrate MD
+        # unexpected behavior will happen if it's called when the
+        # bias force is not "freshly computed"
 
-        self.plumed.cmd("setStep", self.plumedstep)
-        self.plumed.cmd("setCharges", self.charges)
-        self.plumed.cmd("setMasses", self.masses)
-        rpos = pos.reshape((-1, 3))
-        self.plumed.cmd("setPositions", rpos)
-        self.plumed.cmd("setBox", cell.T.copy())
-        if self.system_force is not None:
-            f[:] = dstrip(self.system_force.f).reshape((-1, 3))
-            vir[:] = -dstrip(self.system_force.vir)
-            self.plumed.cmd("setEnergy", dstrip(self.system_force.pot))
-        self.plumed.cmd("setForces", f)
-        self.plumed.cmd("setVirial", vir)
-        self.plumed.cmd("prepareCalc")
-        self.plumed.cmd("performCalcNoUpdate")
+        self.plumed_step += 1
+
+        bias_before = np.zeros(1, float)
+        bias_after = np.zeros(1, float)
+
+        if self.compute_work:
+            self.plumed.cmd("getBias", bias_before)
+
+        # sets the step and does the actual update
+        self.plumed.cmd("setStep", self.plumed_step)
         self.plumed.cmd("update")
 
-        return True
+        # recompute the bias so we can compute the work
+        if self.compute_work:
+            self.plumed.cmd("performCalcNoForces")
+            self.plumed.cmd("getBias", bias_after)
+
+        work = (bias_before - bias_after).item()
+
+        return work
 
 
 class FFYaff(FFEval):

ipi/engine/smotion/metad.py

@@ -71,6 +71,7 @@ def step(self, step=None):
                 k = bc.ffield
                 if k not in self.metaff:
                     continue  # only does metad for the indicated forcefield
+
                 f = s.ensemble.bias.ff[k]
                 if not hasattr(
                     f, "mtd_update"
@@ -83,16 +84,19 @@ def step(self, step=None):
                 if s.ensemble.bweights[ik] == 0:
                     continue  # do not put metad bias on biases with zero weights (useful to do remd+metad!)
 
-                meta_pot_before = s.ensemble.bias.pot
+                mtd_work = f.mtd_update(pos=s.beads.qc, cell=s.cell.h)
+                # updates the conserved quantity with the change in bias so that
+                # we remove the shift due to added hills
+                s.ensemble.eens += mtd_work
 
-                fmtd = f.mtd_update(pos=s.beads.qc, cell=s.cell.h)
-                if fmtd:  # if metadyn has updated, then we must recompute forces.
-                    # hacky but cannot think of a better way: we must manually taint *just* that component
+                if mtd_work != 0:
+                    # hacky but cannot think of a better way: we must manually taint *just* that component.
+                    # we also use the fact that the bias force from a hill is zero when it's added so we
+                    # don't need changes to the forces, only to the bias
                     for fc in s.ensemble.bias.mforces:
                         if fc.ffield == k:
                             for fb in fc._forces:
-                                fb._ufvx.taint()
-                    meta_pot_after = s.ensemble.bias.pot
-                    # updates the conserved quantity with the change in bias so that
-                    # we remove the shift due to added hills
-                    s.ensemble.eens += meta_pot_before - meta_pot_after
+                                # this open-heart surgery on a depend object is fugly
+                                # but can't se a better way
+                                fb._ufvx._value[0] -= mtd_work
+                                fb._ufvx.taint(taintme=False)

ipi/inputs/forcefields.py

@@ -535,18 +535,29 @@ class InputFFPlumed(InputForceField):
                 "help": "This describes the location to read the reference structure file from.",
             },
         ),
-        "plumeddat": (
+        "plumed_dat": (
             InputValue,
             {"dtype": str, "default": "plumed.dat", "help": "The PLUMED input file"},
         ),
-        "plumedstep": (
+        "plumed_step": (
             InputValue,
             {
                 "dtype": int,
                 "default": 0,
                 "help": "The current step counter for PLUMED calls",
             },
         ),
+        "compute_work": (
+            InputValue,
+            {
+                "dtype": bool,
+                "default": True,
+                "help": """Compute the work done by the metadynamics bias 
+                (to correct the conserved quantity). Note that this might 
+                require multiple evaluations of the conserved quantities,
+                and can add some overhead.""",
+            },
+        ),
         "plumed_extras": (
             InputArray,
             {
@@ -572,11 +583,9 @@ class InputFFPlumed(InputForceField):
 
     def store(self, ff):
         super(InputFFPlumed, self).store(ff)
-        self.plumeddat.store(ff.plumeddat)
-        # pstep = ff.plumedstep
-        # if pstep > 0: pstep -= 1 # roll back plumed step before writing a restart
-        # self.plumedstep.store(pstep)
-        self.plumedstep.store(ff.plumedstep)
+        self.plumed_dat.store(ff.plumed_dat)
+        self.plumed_step.store(ff.plumed_step)
+        self.compute_work.store(ff.compute_work)
         self.file.store(ff.init_file)
         self.plumed_extras.store(np.array(list(ff.plumed_data.keys())))
 
@@ -589,8 +598,9 @@ def fetch(self):
             offset=self.offset.fetch(),
             dopbc=self.pbc.fetch(),
             threaded=self.threaded.fetch(),
-            plumeddat=self.plumeddat.fetch(),
-            plumedstep=self.plumedstep.fetch(),
+            compute_work=self.compute_work.fetch(),
+            plumed_dat=self.plumed_dat.fetch(),
+            plumed_step=self.plumed_step.fetch(),
             plumed_extras=self.plumed_extras.fetch(),
             init_file=self.file.fetch(),
         )

ipi_tests/regression_tests/tests/PLUMED/BIAS.NOAUTO/input.xml

@@ -15,7 +15,7 @@
    </ffsocket>
    <ffplumed name="plumed">
        <file mode="xyz">./init.xyz</file>
-       <plumeddat> plumed.dat </plumeddat>
+       <plumed_dat> plumed.dat </plumed_dat>
        <plumedextras> [dhh, rest.bias] </plumedextras>
    </ffplumed>
    <system >

ipi_tests/regression_tests/tests/PLUMED/METAD.NOAUTO/input.xml

@@ -14,7 +14,7 @@
    </ffsocket>
    <ffplumed name="plumed">
        <file mode="xyz">./init.xyz</file>
-       <plumeddat> plumed.dat </plumeddat>
+       <plumed_dat> plumed.dat </plumed_dat>
        <plumedextras> [dhh, metad.bias] </plumedextras>
    </ffplumed>
    <system >