Implement umbrella integration (SSAGESLabs#68)

Co-authored-by: Pablo Zubieta <[email protected]>

.github/workflows/docker-ci.yml

@@ -55,7 +55,14 @@ jobs:
  name: Load and run test
  run: |
  docker load --input /tmp/pysages.tar
- docker run -t pysages bash -c "cd PySAGES/examples/harmonic_bias/ && ./run.sh"
+ docker run -v/tmp:/tmp -t pysages bash -c "cd PySAGES/examples/harmonic_bias/ && ./run.sh && mv hist.pdf /tmp/"
+ -
+ name: Upload artifact
+ uses: actions/upload-artifact@v2
+ with:
+ name: harmonic-hist.pdf
+ path: /tmp/hist.pdf
+
 
  abf-alanine-dipeptide-openmm:
  runs-on: ubuntu-latest
@@ -75,3 +82,28 @@ jobs:
  run: |
  docker load --input /tmp/pysages.tar
  docker run -t pysages bash -c "cd PySAGES/examples/abf/ && python3 ./alanine-dipeptide_openmm.py"
+
+ umbrella-integration-hoomd:
+ runs-on: ubuntu-latest
+ needs: build
+ steps:
+ -
+ name: Set up Docker Buildx
+ uses: docker/setup-buildx-action@v1
+ -
+ name: Download artifact
+ uses: actions/download-artifact@v2
+ with:
+ name: pysages
+ path: /tmp
+ -
+ name: Load and run test
+ run: |
+ docker load --input /tmp/pysages.tar
+ docker run -v /tmp:/tmp -t pysages bash -c "cd PySAGES/examples/umbrella_integration/ && python3 ./gen_gsd.py && python3 integration.py --N-replica=5 --time-steps=1000 --discard-equi=0 && mkdir /tmp/plots && mv *.pdf /tmp/plots/"
+ -
+ name: Upload artifacts
+ uses: actions/upload-artifact@v2
+ with:
+ name: umbrella-integration-plots
+ path: /tmp/plots

examples/harmonic_bias/harmonic_bias.py

@@ -12,25 +12,7 @@
 
 from pysages.collective_variables import Component
 
-from pysages.methods import HarmonicBias
-
-class HistogramLogger:
- def __init__(self, period):
- self.period = period
- self.data = []
-
- def __call__(self, snapshot, state, timestep):
- if timestep % self.period == 0:
- self.data.append(state.xi)
-
- def get_histograms(self, bins, lim):
- data = np.asarray(self.data)
- data = data.reshape(data.shape[0], data.shape[2])
- histograms = []
- for i in range(data.shape[1]):
- histograms.append(np.histogram(data[:,i], bins=bins, range=lim, density=True)[0])
- return histograms
-
+from pysages.methods import HarmonicBias, HistogramLogger
 
 def plot(xi_hist, target_hist, lim):
  fig, ax = plt.subplots()
@@ -98,9 +80,14 @@ def main():
  target_hist = []
  for i in range(len(center_cv)):
  target_hist.append(get_target_dist(center_cv[i], k, (-Lmax/2, Lmax/2), bins))
- hist = callback.get_histograms(bins, (-Lmax/2, Lmax/2))
- plot(hist, target_hist, (-Lmax/2, Lmax/2))
- validate_hist(hist, target_hist)
+ lims = [(-Lmax/2, Lmax/2) for i in range(3)]
+ hist, edges = callback.get_histograms(bins=bins, range=lims)
+ hist_list = [ np.sum(hist, axis=(1,2))/(Lmax**2),
+ np.sum(hist, axis=(0,2))/(Lmax**2),
+ np.sum(hist, axis=(0,1))/(Lmax**2),
+ ]
+ plot(hist_list, target_hist, (-Lmax/2, Lmax/2))
+ validate_hist(hist_list, target_hist)
 
 
 if __name__ == "__main__":

examples/umbrella_integration/gen_gsd.py

@@ -0,0 +1,58 @@
+#!/usr/bin/env python3
+import sys
+import numpy as np
+import gsd
+import gsd.hoomd
+
+class System:
+ def __init__(self):
+ self.L = 5
+ self.N = 200
+
+
+def post_process_pos(snapshot):
+ snapshot.particles.image = np.rint(snapshot.particles.position/snapshot.configuration.box[:3])
+ snapshot.particles.position -= snapshot.particles.image * snapshot.configuration.box[:3]
+ return snapshot
+
+def get_snap(system):
+ snapshot = gsd.hoomd.Snapshot()
+ snapshot.configuration.box = [system.L, system.L, system.L, 0, 0, 0]
+
+ snapshot.particles.N = system.N
+
+ snapshot.particles.types = ["A", "B"]
+ snapshot.particles.position = np.zeros((snapshot.particles.N, 3))
+ snapshot.particles.velocity = np.random.standard_normal((snapshot.particles.N, 3))
+ snapshot.particles.image = np.zeros((snapshot.particles.N, 3), dtype=np.int)
+ snapshot.particles.typeid = np.zeros(snapshot.particles.N, dtype=np.int)
+ snapshot.particles.typeid += 1
+ snapshot.particles.typeid[0] = 0
+
+ rng = np.random.default_rng()
+ for particle in range(system.N):
+ snapshot.particles.position[particle, 0] = rng.random() * system.L - system.L/2
+ snapshot.particles.position[particle, 1] = rng.random() * system.L - system.L/2
+ snapshot.particles.position[particle, 2] = rng.random() * system.L - system.L/2
+
+ snapshot.particles.position[0 , 0] = -np.pi
+ snapshot.particles.position[0 , 1] = -np.pi
+ snapshot.particles.position[0 , 1] = -np.pi
+
+ return snapshot
+
+
+def main(argv):
+ if len(argv) != 1:
+ print("Usage: ./gen_gsd.py")
+
+ system = System()
+ snap = get_snap(system)
+ snap = post_process_pos(snap)
+ snap.particles.validate()
+ with gsd.hoomd.open("start.gsd", "wb") as f:
+ f.append(snap)
+
+
+if __name__ == "__main__":
+ main(sys.argv)

examples/umbrella_integration/integration.py

@@ -0,0 +1,118 @@
+#!/usr/bin/env python3
+import sys
+import argparse
+import numpy as np
+import matplotlib.pyplot as plt
+
+import hoomd
+import hoomd.md as md
+import hoomd.dlext
+
+import pysages
+from pysages.collective_variables import Component
+from pysages.methods import UmbrellaIntegration
+
+
+param1 = {"A": 0.5, "w": 0.2, "p": 2}
+
+def generate_context(**kwargs):
+ hoomd.context.initialize("")
+ context = hoomd.context.SimulationContext()
+ with context:
+ print("Operating replica {0}".format(kwargs.get("replica_num")))
+ system = hoomd.init.read_gsd("start.gsd")
+
+ hoomd.md.integrate.nve(group=hoomd.group.all())
+ hoomd.md.integrate.mode_standard(dt=0.01)
+
+ nl = hoomd.md.nlist.cell()
+ dpd = hoomd.md.pair.dpd(r_cut=1, nlist=nl, seed=42, kT=1.)
+ dpd.pair_coeff.set("A", "A", A=5., gamma=1.0)
+ dpd.pair_coeff.set("A", "B", A=5., gamma=1.0)
+ dpd.pair_coeff.set("B", "B", A=5., gamma=1.0)
+
+ periodic = hoomd.md.external.periodic()
+ periodic.force_coeff.set('A', A=param1["A"], i=0, w=param1["w"], p=param1["p"])
+ periodic.force_coeff.set('B', A=0.0, i=0, w=0.02, p=1)
+ return context
+
+
+def plot_hist(result, bins=50):
+ fig, ax = plt.subplots(2, 2)
+
+ # ax.set_xlabel("CV")
+ # ax.set_ylabel("p(CV)")
+
+ counter = 0
+ hist_per = len(result["center"])//4+1
+ for x in range(2):
+ for y in range(2):
+ for i in range(hist_per):
+ if counter+i < len(result["center"]):
+ center = np.asarray(result["center"][counter+i])
+ histo, edges = result["histogram"][counter+i].get_histograms(bins=bins)
+ edges = np.asarray(edges)[0]
+ edges = (edges[1:] + edges[:-1]) / 2
+ ax[x,y].plot(edges, histo, label="center {0}".format(center))
+ ax[x,y].legend(loc="best", fontsize="xx-small")
+ ax[x,y].set_yscale("log")
+ counter += hist_per
+ while counter < len(result["center"]):
+ center = np.asarray(result["center"][counter])
+ histo, edges = result["histogram"][counter].get_histograms(bins=bins)
+ edges = np.asarray(edges)[0]
+ edges = (edges[1:] + edges[:-1]) / 2
+ ax[1,1].plot(edges, histo, label="center {0}".format(center))
+ counter += 1
+
+ fig.savefig("hist.pdf")
+
+def external_field(r, A, p, w):
+ return A*np.tanh(1/(2*np.pi*p*w)*np.cos(p*r))
+
+def plot_energy(result):
+ fig, ax = plt.subplots()
+
+ ax.set_xlabel("CV")
+ ax.set_ylabel("Free energy $[\epsilon]$")
+ center = np.asarray(result["center"])
+ A = np.asarray(result["A"])
+ offset = np.min(A)
+ ax.plot(center, A-offset, color="teal")
+
+ x = np.linspace(-3, 3, 50)
+ data = external_field(x, **param1)
+ offset = np.min(data)
+ ax.plot(x, data-offset, label="test")
+
+ fig.savefig("energy.pdf")
+
+
+def get_args(argv):
+ parser = argparse.ArgumentParser(description="Example script to run umbrella integration")
+ parser.add_argument("--k-spring", "-k", type=float, default=50., help="spring constant for each replica")
+ parser.add_argument("--N-replica", "-N", type=int, default=25, help="Number of replica along the path")
+ parser.add_argument("--start-path", "-s", type=float, default=-1.5, help="Start point of the path")
+ parser.add_argument("--end-path", "-e", type=float, default=1.5, help="Start point of the path")
+ parser.add_argument("--time-steps", "-t", type=int, default=int(1e5), help="Number of simulation steps for each replica")
+ parser.add_argument("--log-period", "-l", type=int, default=int(50), help="Frequency of logging the CVS for histogram")
+ parser.add_argument("--discard-equi", "-d", type=int, default=int(1e4), help="Discard timesteps before logging for equilibration")
+ args = parser.parse_args(argv)
+ return args
+
+def main(argv):
+
+ args = get_args(argv)
+
+ cvs = [Component([0], 0),]
+ method = UmbrellaIntegration(cvs)
+
+ centers = list(np.linspace(args.start_path, args.end_path, args.N_replica))
+ result = method.run(generate_context, args.time_steps, centers, args.k_spring, args.log_period, args.discard_equi)
+
+ plot_energy(result)
+ plot_hist(result)
+
+
+if __name__ == "__main__":
+ main(sys.argv[1:])

pysages/methods/__init__.py

@@ -9,3 +9,5 @@
 from .abf import ABF
 from .funn import FUNN
 from .harmonic_bias import HarmonicBias
+from .umbrella_integration import UmbrellaIntegration
+from .utils import HistogramLogger

pysages/methods/harmonic_bias.py

@@ -40,10 +40,15 @@ def __init__(self, cvs, kspring, center, *args, **kwargs):
  """
  super().__init__(cvs, args, kwargs)
  self._N = len(cvs)
- self.set_kspring(kspring)
- self.set_center(center)
+ self.kspring = kspring
+ self.center = center
 
- def set_kspring(self, kspring):
+ @property
+ def kspring(self):
+ return self._kspring
+
+ @kspring.setter
+ def kspring(self, kspring):
  """
  Set new spring constant.
 
@@ -72,35 +77,36 @@ def set_kspring(self, kspring):
  raise RuntimeError(f"Wrong kspring size, expected 1 or {N}, got {n}.")
 
  self._kspring = np.identity(N) * kspring
-
  return self._kspring
 
- def get_kspring(self):
- return self._kspring
+ @property
+ def center(self):
+ return self._center
 
- def set_center(self, center):
+ @center.setter
+ def center(self, center):
  center = np.asarray(center)
+ if center.shape == ():
+ center = center.reshape(1)
  if len(center.shape) !=1 or center.shape[0] != self._N:
  raise RuntimeError(f"Invalid center shape expected {self._N} got {center.shape}.")
  self._center = center
 
- def get_center(self):
- return self._center
-
  def build(self, snapshot, helpers):
  return _harmonic_bias(self, snapshot, helpers)
 
 
 def _harmonic_bias(method, snapshot, helpers):
  cv = method.cv
- center = method.get_center()
- kspring = method.get_kspring()
+ center = method.center
+ kspring = method.kspring
  natoms = np.size(snapshot.positions, 0)
 
  def initialize():
  bias = np.zeros((natoms, 3))
  return HarmonicBiasState(bias, None)
 
+
  def update(state, data):
  xi, Jxi = cv(data)
  D = kspring @ (xi - center).flatten()