Better PLUMED interface (i-pi#330)

Lots of small & medium fixes to the PLUMED interface. Most notably, it is now possible to fetch CV values from PLUMED and bring them into i-PI so they can be printed using the extras infrastructure. Also took this opportunity to clean up a bit extras, components access and miscellaneous advanced output features. 

* Added an example of PLUMED bias. Did is so that it does not run in the CI as we don't have PLUMED there
* Added regtest for metadynamics. Again, excluded from CI and full tests, must run manually
* Implemented direct output of CVs and other internal PLUMED variables. Use plumed cmd, and extras to accumulate them
* Added an example of propagating CV to i-PI outside of the regtests
* Renamed force_extras, and added outputs for force components
* Raw outputs for potential and force components
* Moved some relevant examples from temp to features 

---------

Co-authored-by: Venkat Kapil <[email protected]>

examples/features/committee_models/weighted_baseline/input.xml

@@ -43,8 +43,8 @@
     </initialize>
     <forces>
         <force forcefield='committee' nbeads='1'> 
-            <!-- adding these to force_extras will make the committee extras be treated properly with ring-polymer contraction, ect. -->
-            <force_extras> [committee_pot, committee_force, committee_virial, baseline_pot, baseline_force, baseline_virial, wb_mixing] </force_extras>
+            <!-- adding these to interpolate_extras will make the committee extras be treated properly with ring-polymer contraction, ect. -->
+            <interpolate_extras> [committee_pot, committee_force, committee_virial, baseline_pot, baseline_force, baseline_virial, wb_mixing] </interpolate_extras>
         </force>
     </forces>
     <motion mode='dynamics'>

examples/features/committee_models/weighted_baseline_extra/input.xml

@@ -40,8 +40,8 @@
     </initialize>
     <forces>
         <force forcefield='committee' nbeads='1'> 
-            <!-- adding these to force_extras will make the committee extras be treated properly with ring-polymer contraction, ect. -->
-            <force_extras> [committee_pot, committee_force, committee_virial, baseline_pot, baseline_force, baseline_virial, wb_mixing] </force_extras>
+            <!-- adding these to interpolate_extras will make the committee extras be treated properly with ring-polymer contraction, ect. -->
+            <interpolate_extras> [committee_pot, committee_force, committee_virial, baseline_pot, baseline_force, baseline_virial, wb_mixing] </interpolate_extras>
         </force>
     </forces>
     <motion mode='dynamics'>

examples/features/metadynamics/README.md

@@ -15,3 +15,7 @@ i-pi--driver -u -h zundel -m zundel
 ```
 
 However, they require having PLUMED libraries and Python bindings in the appropriate system paths.
+
+Note also that the example in `pimd_metadynamics` provides a demonstration of how to retrieve quantities
+computed PLUMED-side (such as CV values) back into i-PI as `extras` that can be associated with 
+individual structures and printed out alongside the other i-PI outputs.

examples/features/metadynamics/pimd_metadynamics_zundel/input.xml

@@ -9,11 +9,13 @@
    <ffplumed name="plumed">
          <file mode="xyz">./h5o2+.xyz</file>
          <plumeddat> plumed/plumed.dat </plumeddat>
+         <plumedextras> [doo, co1.lessthan, co2.lessthan, mtd.bias ] </plumedextras>
    </ffplumed>
    <total_steps>400</total_steps>
    <output prefix="data">
       <trajectory stride="40" filename="pos" cell_units="angstrom">positions{angstrom}</trajectory>
       <trajectory stride="20" filename="xc" format="xyz">x_centroid{angstrom}</trajectory>
+      <trajectory stride="20" filename="doo" bead="0" extra_type="doo"> extras_bias </trajectory>
       <properties stride="2"> [ step, time, conserved, temperature{kelvin}, kinetic_cv, 
             potential, kinetic_cv(H), kinetic_cv(O), ensemble_bias ] </properties>
    </output>
@@ -33,7 +35,9 @@
       <ensemble>
          <temperature units="kelvin"> 300.0 </temperature>
          <bias>
-            <force forcefield="plumed" nbeads="1"></force>
+           <force forcefield="plumed" nbeads="1">
+             <interpolate_extras> [ doo, mtd.bias ] </interpolate_extras>
+           </force>
          </bias>
       </ensemble>
       <motion mode="dynamics">

examples/temp/pes/pswater/combine-extras/input.xml → examples/features/outputs/combine-extras/input.xml

@@ -7,8 +7,8 @@
     <trajectory filename='kin' format='xyz' stride='2'> kinetic_cv </trajectory>
     <trajectory filename='dip' stride='2' extra_type='dipole'> extras </trajectory>
     <trajectory filename='raw' stride='2' extra_type='raw'> extras </trajectory>    
-    <trajectory filename='dip0' stride='2' extra_type='dipole'> extras_component(0) </trajectory>
-    <trajectory filename='dip1' stride='2' extra_type='dipole'> extras_component(1) </trajectory>
+    <trajectory filename='dip0' stride='2' extra_type='dipole'> extras_component_raw(0) </trajectory>
+    <trajectory filename='dip1' stride='2' extra_type='dipole'> extras_component_raw(1) </trajectory>
     <checkpoint stride='100' filename='chk' overwrite='true'/>
     <checkpoint stride='2000' filename='restart' overwrite='false'/>
   </output>
@@ -39,10 +39,10 @@
     </initialize>
     <forces>
       <force forcefield='driver' weight="1.0">
-          <force_extras> ['dipole'] </force_extras>
+          <interpolate_extras> ['dipole'] </interpolate_extras>
       </force>
       <force forcefield='driver2' weight="2.0" >
-           <force_extras> ['dipole'] </force_extras>  
+           <interpolate_extras> ['dipole'] </interpolate_extras>  
       </force>
     </forces>
     <ensemble>

examples/temp/pes/pswater/contract-extras/input.xml → examples/features/outputs/contract-extras/input.xml

@@ -6,6 +6,8 @@
     <trajectory filename='for' format='xyz' stride='2'> forces </trajectory>
     <trajectory filename='kin' format='xyz' stride='2'> kinetic_cv </trajectory>
     <trajectory filename='dip' stride='2' extra_type='dipole'> extras </trajectory>
+    <!-- these two output the extras for the contracted beads -->
+    <trajectory filename='dip_raw' stride='2' extra_type='dipole'> extras_component_raw(0) </trajectory>
     <checkpoint stride='100' filename='chk' overwrite='true'/>
     <checkpoint stride='2000' filename='restart' overwrite='false'/>
   </output>
@@ -31,7 +33,7 @@
         <!-- The line below is the one that tells i-PI that the 'dipole' extras quantity will be 
         treated as a physical quantity like forces. It can then be interpolated, and if there are
         several force components, it will be automatically summed with the respective weights -->
-        <force_extras> ['dipole'] </force_extras>
+        <interpolate_extras> [dipole] </interpolate_extras>
       </force>
     </forces>
     <ensemble>

examples/features/path_integral_molecular_dynamics/README.md

@@ -2,9 +2,9 @@ Path integral molecular dynamics.
 =================================================
 `standard_pimd`: path integral molecular dynamics using the PILE thermostat and MTS for propagating the ring polymer modes.
 `cayley_pimd`: path integral molecular dynamics using the PILE thermostat and Cayley integrator for propagating the ring polymer modes.
-`constrained_centroid`: constant-temperature PIMD, with the centroid held fixed in the initial position. useful to compute the centroid potential of mean force.
+`constrained_centroid`: constant-temperature PIMD, with the centroid held fixed in the initial position. Useful to compute the centroid potential of mean force.
 `piglet`: path integral molecular dynamics using the PIGLET thermostat
-`pimd+mts`:  path integral molecular dynamics with MTS algorithm to integrate short and long ranged forces with different timesteps.
+`pimd+mts`:  path integral molecular dynamics with MTS algorithm to integrate short and long ranged forces with different timesteps. Includes option to print out the slow and/or fast force components.
 `rpc`:  path integral molecular dynamics with RPC algorithm to use different number of replicas for short and long ranged forces.
 scpimd:  path integral molecular dynamics with the Suzuki-Chin splitting.
 `water_remd`: path integral molecular dynamics with replica-exchange spanning different temperatures and pressures.

examples/features/path_integral_molecular_dynamics/pimd+mts_liquid_water/input.xml

@@ -1,7 +1,11 @@
 <simulation mode='md' verbosity='high'>
   <output prefix='simulation'>
-    <properties stride='1' filename='out'>  [ step, time{picosecond}, conserved, temperature{kelvin}, kinetic_cv, potential, pressure_cv{megapascal} ] </properties>
-    <trajectory filename='pos' stride='1'> positions </trajectory>
+    <properties stride='1' filename='out'>  [ step, time{picosecond}, conserved, temperature{kelvin}, kinetic_cv, potential, pressure_cv{megapascal}, pot_component_raw(0), pot_component_raw(1) ] </properties>
+    <properties stride='10' filename='pot_lr'>  [ pot_component(1;0), pot_component(1;1), pot_component(1;2), pot_component(1;3) ] </properties> 
+    <properties stride='10' filename='pot_lr-raw'>  [ pot_component_raw(1;0), pot_component_raw(1;1) ] </properties>
+    <trajectory filename='pos' stride='10'> positions </trajectory>
+    <trajectory filename='f_sr_c' stride='10'> forces_component_raw(0) </trajectory>
+    <trajectory filename='f_lr-raw' stride='10'> forces_component_raw(0;0) </trajectory>
     <checkpoint stride='200'/>
   </output>
   <total_steps>100</total_steps>
@@ -21,14 +25,16 @@
     </initialize>
     <forces>
        <!--
-         MTS setup - apply the fast (short-range) force in the inner loop and the slow (long-range) force in the outer loop. 
-         Note that if the outer loop contains a *correction* to the inner loop the weights should be 
-         [-1,1] (fast force) and [1,0] (slow force)
+           MTS setup - apply the fast (short-range) force in the inner loop and the slow (long-range) force in the 
+           outer loop. Also does ring-polymer contraction, by computing the long-range force on just two beads. 
+
+           Note that if the outer loop contains a *correction* to the inner loop the weights should be 
+           [-1,1] (fast force) and [1,0] (slow force)         
        -->
       <force forcefield='short_range'>
         <mts_weights>[0,1]</mts_weights>
       </force>
-      <force forcefield='long_range'>
+      <force forcefield='long_range' nbeads='2'>
         <mts_weights>[1,0]</mts_weights>
       </force>
     </forces>

examples/temp/pes/pswater/nvt/input.xml

@@ -7,7 +7,7 @@
     <trajectory filename='kin' format='xyz' stride='2'> kinetic_cv </trajectory>
     <trajectory filename='dip' stride='2' extra_type='dipole'> extras </trajectory>
     <trajectory filename='raw' stride='2' extra_type='raw'> extras </trajectory>    
-    <trajectory filename='dip0' stride='2' extra_type='dipole'> extras_component(0) </trajectory>
+    <trajectory filename='dip0' stride='2' extra_type='dipole'> extras_component_raw(0) </trajectory>
     <checkpoint stride='100' filename='chk' overwrite='true'/>
     <checkpoint stride='2000' filename='restart' overwrite='false'/>
   </output>

ipi/engine/forcefields.py

@@ -671,6 +671,7 @@ def __init__(
         init_file="",
         plumeddat="",
         plumedstep=0,
+        plumedextras=[],
     ):
         """Initialises FFPlumed.
 
@@ -689,6 +690,7 @@ def __init__(
         self.plumed = plumed.Plumed()
         self.plumeddat = plumeddat
         self.plumedstep = plumedstep
+        self.plumedextras = plumedextras
         self.init_file = init_file
 
         if self.init_file.mode == "xyz":
@@ -718,6 +720,20 @@ def __init__(
             self.plumedrestart = True
             self.plumed.cmd("setRestart", 1)
         self.plumed.cmd("init")
+
+        self.plumed_data = {}
+        for x in plumedextras:
+            rank = np.zeros(1, dtype=np.int_)
+            self.plumed.cmd(f"getDataRank {x}", rank)
+            if rank[0] > 1:
+                raise ValueError("Cannot retrieve varibles with rank > 1")
+            shape = np.zeros(rank[0], dtype=np.int_)
+            if shape[0] > 1:
+                raise ValueError("Cannot retrieve varibles with size > 1")
+            self.plumed.cmd(f"getDataShape {x}", shape)
+            self.plumed_data[x] = np.zeros(shape, dtype=np.double)
+            self.plumed.cmd(f"setMemoryForData {x}", self.plumed_data[x])
+
         self.charges = dstrip(myatoms.q) * 0.0
         self.masses = dstrip(myatoms.m)
         self.lastq = np.zeros(3 * self.natoms)
@@ -771,8 +787,12 @@ def evaluate(self, r):
             f[:] -= dstrip(self.system_force.f).flatten()
             vir[:] -= -dstrip(self.system_force.vir)
 
+        extras = {"raw": ""}
+        for x in self.plumed_data:
+            extras[str(x)] = self.plumed_data[x].copy()
+
         # nb: the virial is a symmetric tensor, so we don't need to transpose
-        r["result"] = [v, f, vir, {"raw": ""}]
+        r["result"] = [v, f, vir, extras]
         r["status"] = "Done"
 
     def mtd_update(self, pos, cell):

ipi/engine/forces.py

@@ -301,7 +301,7 @@ def __init__(
         weight=1.0,
         name="",
         mts_weights=None,
-        force_extras=None,
+        interpolate_extras=None,
         epsilon=-0.001,
     ):
         """Initializes ForceComponent
@@ -316,7 +316,7 @@ def __init__(
               will be weighted by this factor. The combination is a weighted sum.
            name: The name of the forcefield.
            mts_weights: Weight of forcefield at each mts level.
-           force_extras: A list of properties that should be treated as physical quantities,
+           interpolate_extras: A list of properties that should be treated as physical quantities,
               converted to numpy arrays and treated with ring polymer contraction. If
               different force components have this field, they will also be summed with
               the respective weight like a forces object.
@@ -330,10 +330,10 @@ def __init__(
             self.mts_weights = np.asarray([])
         else:
             self.mts_weights = np.asarray(mts_weights)
-        if force_extras is None:
-            self.force_extras = []
+        if interpolate_extras is None:
+            self.interpolate_extras = []
         else:
-            self.force_extras = force_extras
+            self.interpolate_extras = interpolate_extras
         self.epsilon = epsilon
 
     def bind(self, beads, cell, fflist, output_maker):
@@ -400,7 +400,7 @@ def bind(self, beads, cell, fflist, output_maker):
         )
         self._extras = depend_value(
             name="extras",
-            value=np.zeros(self.nbeads, float),
+            value={},
             func=self.extra_gather,
             dependencies=[self._forces[b]._extra for b in range(self.nbeads)],
         )
@@ -464,21 +464,21 @@ def extra_gather(self):
                     )
                 fc_extra[e].append(b.extra[e])
 
-        # force_extras should be numerical, thus can be converted to numpy arrays.
+        # interpolate_extras should be numerical, thus can be converted to numpy arrays.
         # we enforce the type and numpy will raise an error if not.
-        for e in self.force_extras:
+        for e in self.interpolate_extras:
             try:
                 fc_extra[e] = np.asarray(fc_extra[e], dtype=float)
             except KeyError:
                 raise KeyError(
-                    "force_extras required "
+                    "interpolate_extras required "
                     + e
                     + " to promote, but was not found among extras "
                     + str(list(fc_extra.keys()))
                 )
             except:
                 raise Exception(
-                    "force_extras has to be numerical to be treated as a physical quantity. It is not -- check the quantity that is being passed."
+                    "interpolate_extras has to be numerical to be treated as a physical quantity. It is not -- check the quantity that is being passed."
                 )
         return fc_extra
 
@@ -530,6 +530,7 @@ class ScaledForceComponent:
     def __init__(self, baseforce, scaling=1):
         self.bf = baseforce
         self.name = baseforce.name
+        self.nbeads = baseforce.nbeads
         self.ffield = baseforce.ffield
         self._scaling = depend_value(name="scaling", value=scaling)
         self._f = depend_array(
@@ -554,10 +555,10 @@ def __init__(self, baseforce, scaling=1):
             value=np.zeros((self.bf.nbeads, 3, 3)),
             dependencies=[self.bf._virs, self._scaling],
         )
-        self._extras = depend_array(
+        self._extras = depend_value(
             name="extras",
             func=lambda: self.bf.extras,
-            value=np.zeros(self.bf.nbeads),
+            value={},
             dependencies=[self.bf._extras],
         )
 
@@ -576,7 +577,7 @@ def __init__(self, baseforce, scaling=1):
         self._weight = depend_value(name="weight", value=0)
         dpipe(self.bf._weight, self._weight)
         self.mts_weights = self.bf.mts_weights
-        self.force_extras = self.bf.force_extras
+        self.interpolate_extras = self.bf.interpolate_extras
 
     def get_pots(self):
         return (
@@ -812,7 +813,7 @@ def make_rpc(rpc, beads):
                 nbeads=newb,
                 weight=fc.weight,
                 mts_weights=fc.mts_weights,
-                force_extras=fc.force_extras,
+                interpolate_extras=fc.interpolate_extras,
                 epsilon=fc.epsilon,
             )
             newbeads = Beads(beads.natoms, newb)
@@ -1160,41 +1161,44 @@ def get_vir(self):
             vir += v
         return vir
 
-    def pots_component(self, index, weighted=True):
+    def pots_component(self, index, weighted=True, interpolate=True):
         """Fetches the index^th component of the total potential."""
-        if weighted:
-            if self.mforces[index].weight != 0:
-                return self.mforces[index].weight * self.mrpc[index].b2tob1(
-                    self.mforces[index].pots
-                )
+
+        if weighted and self.mforces[index].weight == 0:
+            if interpolate:
+                return np.zeros(self.mrpc[index].nbeads1)
             else:
-                return 0
+                return np.zeros(self.mrpc[index].nbeads2)
         else:
-            return self.mrpc[index].b2tob1(self.mforces[index].pots)
-
-    def forces_component(self, index, weighted=True):
+            pots = dstrip(self.mforces[index].pots).copy()
+            if weighted:
+                pots *= self.mforces[index].weight
+            if interpolate:
+                pots = self.mrpc[index].b2tob1(pots)
+            return pots
+
+    def forces_component(self, index, weighted=True, interpolate=True):
         """Fetches the index^th component of the total force."""
-        if weighted:
-            if self.mforces[index].weight != 0:
-                return self.mforces[index].weight * self.mrpc[index].b2tob1(
-                    dstrip(self.mforces[index].f)
-                )
+
+        if weighted and self.mforces[index].weight == 0:
+            if interpolate:
+                return np.zeros((self.mrpc[index].nbeads1, 3 * self.natoms))
             else:
-                return np.zeros((self.nbeads, self.natoms * 3), float)
+                return np.zeros((self.mrpc[index].nbeads2, 3 * self.natoms))
         else:
-            return self.mrpc[index].b2tob1(dstrip(self.mforces[index].f))
+            forces = dstrip(self.mforces[index].f).copy()
+            if weighted:
+                forces *= self.mforces[index].weight
+            if interpolate:
+                forces = self.mrpc[index].b2tob1(forces)
+            return forces
 
     def extras_component(self, index):
-        """Fetches extras that are computed for one specific force component."""
+        """Fetches extras that are computed for one specific force component.
+
+        Does not attempt to apply weights or interpolate, always returns raw stuff.
+        """
 
-        if self.nbeads != self.mforces[index].nbeads:
-            raise ValueError(
-                "Cannot fetch extras for a component when using ring polymer contraction"
-            )
-        if self.mforces[index].weight == 0:
-            raise ValueError(
-                "Cannot fetch extras for a component that has not been computed because of zero weight"
-            )
         return self.mforces[index].extras
 
     def forcesvirs_4th_order(self, index):
@@ -1449,8 +1453,8 @@ def extra_combine(self):
         for k in range(self.nforces):
             # combines the extras from the different force components
             for e, v in self.mforces[k].extras.items():
-                if e in self.mforces[k].force_extras:
-                    # extras that are tagged as force_extras are treated exactly as if they were an energy/force/stress
+                if e in self.mforces[k].interpolate_extras:
+                    # extras that are tagged as interpolate_extras are treated exactly as if they were an energy/force/stress
                     v = (
                         self.mforces[k].weight
                         * self.mforces[k].mts_weights.sum()