diff --git a/.github/workflows/docker-ci.yml b/.github/workflows/docker-ci.yml index 00b1b0ad..255baf49 100644 --- a/.github/workflows/docker-ci.yml +++ b/.github/workflows/docker-ci.yml @@ -123,6 +123,22 @@ jobs: docker load --input /tmp/pysages.tar docker run -t pysages bash -c "cd PySAGES/examples/openmm/metad/ && python3 ./alanine-dipeptide.py --time-steps=25" + pbmetad-alanine-dipeptide-openmm: + 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 -t pysages bash -c "cd PySAGES/examples/openmm/pbmetad/ && python3 ./alanine-dipeptide.py --time-steps=25" + alanine-dipeptide-openmm-mpi: runs-on: ubuntu-latest needs: build diff --git a/docs/source/module-pysages-methods-pbmetad.rst b/docs/source/module-pysages-methods-pbmetad.rst new file mode 100644 index 00000000..b0f400f0 --- /dev/null +++ b/docs/source/module-pysages-methods-pbmetad.rst @@ -0,0 +1,13 @@ +pysages.methods.pbmetad +----------------------- + +.. rubric:: Overview + +.. autosummary:: + pysages.methods.pbmetad.ParallelBiasMetadynamics + +.. rubric:: Details + +.. automodule:: pysages.methods.pbmetad + :synopsis: Python classes for Parallel Bias Well-tempered Metadynamics. + :members: diff --git a/docs/source/module-pysages-methods.rst b/docs/source/module-pysages-methods.rst index 0182c4b5..97234eb6 100644 --- a/docs/source/module-pysages-methods.rst +++ b/docs/source/module-pysages-methods.rst @@ -12,6 +12,7 @@ Methods available in PySAGES pysages.methods.ffs.FFS pysages.methods.funn.FUNN pysages.methods.metad.Metadynamics + pysages.methods.pbmetad.ParallelBiasMetadynamics pysages.methods.umbrella_integration.UmbrellaIntegration pysages.methods.harmonic_bias.HarmonicBias pysages.methods.spline_string.SplineString @@ -41,6 +42,7 @@ Abstract base classes module-pysages-methods-funn module-pysages-methods-harmonic_bias module-pysages-methods-metad + module-pysages-methods-pbmetad module-pysages-methods-umbrella module-pysages-methods-string module-pysages-methods-utils diff --git a/docs/source/pysages_wordlist.txt b/docs/source/pysages_wordlist.txt index 2a946024..feedbdaa 100644 --- a/docs/source/pysages_wordlist.txt +++ b/docs/source/pysages_wordlist.txt @@ -52,6 +52,7 @@ Metaclass metad metadynamics Metadynamics +ParallelBiasMetadynamics metapotential monocyclic multi diff --git a/examples/openmm/pbmetad/alanine-dipeptide.py b/examples/openmm/pbmetad/alanine-dipeptide.py new file mode 100644 index 00000000..b2e881eb --- /dev/null +++ b/examples/openmm/pbmetad/alanine-dipeptide.py @@ -0,0 +1,208 @@ +#!/usr/bin/env python3 + +""" +Parallel Bias Well-tempered Metadynamics (PBMetaD) simulation of +Alanine Dipeptide along backbone dihedrals in vacuum with OpenMM and PySAGES. + +Example command to run the simulation `python3 alanine-dipeptide.py --time_steps 1000` +For other supported commandline parameters, check `python3 alanine-dipeptide.py --help` +""" + + +# %% +import argparse +import os +import sys +import time + +import numpy +import pysages +from jax import numpy as np + +from pysages.colvars import DihedralAngle +from pysages.methods import ParallelBiasMetadynamics, MetaDLogger +from pysages.utils import try_import +from pysages.approxfun import compute_mesh + +import matplotlib.pyplot as plt + +openmm = try_import("openmm", "simtk.openmm") +unit = try_import("openmm.unit", "simtk.unit") +app = try_import("openmm.app", "simtk.openmm.app") + + +# %% +pi = numpy.pi +kB = unit.BOLTZMANN_CONSTANT_kB * unit.AVOGADRO_CONSTANT_NA +kB = kB.value_in_unit(unit.kilojoules_per_mole / unit.kelvin) + +T = 298.15 * unit.kelvin +dt = 2.0 * unit.femtoseconds +adp_pdb = os.path.join(os.pardir, os.pardir, "inputs", "alanine-dipeptide", "adp-vacuum.pdb") + + +# %% +def generate_simulation(pdb_filename=adp_pdb, T=T, dt=dt): + pdb = app.PDBFile(pdb_filename) + + ff = app.ForceField("amber99sb.xml") + cutoff_distance = 1.0 * unit.nanometer + topology = pdb.topology + + system = ff.createSystem( + topology, constraints=app.HBonds, nonbondedMethod=app.PME, nonbondedCutoff=cutoff_distance + ) + + # Set dispersion correction use. + forces = {} + for i in range(system.getNumForces()): + force = system.getForce(i) + forces[force.__class__.__name__] = force + + forces["NonbondedForce"].setUseDispersionCorrection(True) + forces["NonbondedForce"].setEwaldErrorTolerance(1.0e-5) + + positions = pdb.getPositions(asNumpy=True) + + integrator = openmm.LangevinIntegrator(T, 1 / unit.picosecond, dt) + + integrator.setRandomNumberSeed(42) + + # platform = openmm.Platform.getPlatformByName(platform) + # simulation = app.Simulation(topology, system, integrator, platform) + simulation = app.Simulation(topology, system, integrator) + simulation.context.setPositions(positions) + simulation.minimizeEnergy() + + simulation.reporters.append(app.PDBReporter("output.pdb", 1000)) + simulation.reporters.append( + app.StateDataReporter("log.dat", 1000, step=True, potentialEnergy=True, temperature=True) + ) + + return simulation + + +# %% +def get_args(argv): + available_args = [ + ("use-grids", "g", bool, 0, "Whether to use grid acceleration"), + ("log", "l", bool, 0, "Whether to use a callback to log data into a file"), + ("time-steps", "t", int, 5e5, "Number of simulation steps"), + ] + parser = argparse.ArgumentParser(description="Example script to run metadynamics") + for (name, short, T, val, doc) in available_args: + parser.add_argument("--" + name, "-" + short, type=T, default=T(val), help=doc) + return parser.parse_args(argv) + + +# %% +def main(argv=[]): + args = get_args(argv) + print(args) + + cvs = [DihedralAngle([4, 6, 8, 14]), DihedralAngle([6, 8, 14, 16])] + + height = [1.2, 1.2] # kJ/mol + sigma = [0.35, 0.35] # radians + deltaT = 5000 + stride = 500 # frequency for depositing gaussians + timesteps = args.time_steps + ngauss = timesteps // stride + 1 # total number of gaussians + + # Grid for storing bias potential and its gradient + grid = pysages.Grid( + lower=(-pi, -pi), upper=(pi, pi), shape=(50, 50), periodic=True, parallelbias=True + ) + + grid = grid if args.use_grids else None + + # Method + method = ParallelBiasMetadynamics( + cvs, height, sigma, stride, ngauss, T.value_in_unit(unit.kelvin), deltaT, kB, grid=grid + ) + + # Logging + hills_file = "hills.dat" + callback = MetaDLogger(hills_file, stride) if args.log else None + + tic = time.perf_counter() + run_result = pysages.run(method, generate_simulation, timesteps, callback) + toc = time.perf_counter() + print(f"Completed the simulation in {toc - tic:0.4f} seconds.") + + # Analysis: Calculate free energy using the deposited bias potential + + # generate CV values on a grid to evaluate bias potential + plot_1Dgrid = pysages.Grid(lower=(-pi), upper=(pi), shape=(64), periodic=True) + + xi_1D = (compute_mesh(plot_1Dgrid) + 1) / 2 * plot_1Dgrid.size + plot_1Dgrid.lower + xi_1D_2CVs = np.hstack((xi_1D, xi_1D)) + + # determine bias factor = (T+deltaT)/deltaT) + alpha = (T.value_in_unit(unit.kelvin) + method.deltaT) / method.deltaT + kT = kB * T.value_in_unit(unit.kelvin) + beta = 1 / kT + + # extract metapotential function from result + result = pysages.analyze(run_result) + centers = result["centers"] + heights = result["heights"] + + pbmetad_potential_cv = result["pbmetad_potential_cv"] + pbmetad_net_potential = result["pbmetad_net_potential"] + + # calculate free energy and report in kT at the end of simulation. + A_cv = pbmetad_potential_cv(xi_1D_2CVs) * -alpha / kT + # set min free energy to zero + A_cv = A_cv - A_cv.min(axis=0) + A_cv1 = A_cv[:, 0].reshape(plot_1Dgrid.shape) + A_cv2 = A_cv[:, 1].reshape(plot_1Dgrid.shape) + + # plot and save free energy along each CV to a PNG file + fig = plt.figure(figsize=(8, 8), dpi=120) + fig.subplots_adjust(hspace=0.25, wspace=0.25) + + range_to_time = stride * dt.value_in_unit(unit.femtoseconds) * 1e-6 + + # plot centers along phi and psi to monitor sampling + for i in range(2): + ax = fig.add_subplot(3, 2, i + 1) + color = "blue" if i == 0 else "red" + ylabel = r"$\phi$" if i == 0 else r"$\psi$" + ax.scatter( + np.arange(np.shape(centers)[0]) * range_to_time, centers[:, i], s=20, color=color + ) + ax.set_xlabel(r"Time [ns]") + ax.set_ylabel(ylabel) + ax.set_yticks(np.arange(-pi, pi + pi / 2, step=(pi / 2)), ["-π", "-π/2", "0", "π/2", "π"]) + + # plot height along phi and psi to monitor sampling + for i in range(2): + ax = fig.add_subplot(3, 2, i + 3) + color = "blue" if i == 0 else "red" + ylabel = r"$W(\phi) ~[k_{B}T]$" if i == 0 else r"$W(\psi) ~[k_{B}T]$" + ax.scatter( + np.arange(np.shape(heights)[0]) * range_to_time, heights[:, i], s=20, color=color + ) + ax.set_xlabel(r"Time [ns]") + ax.set_ylabel(ylabel) + + # plot free energy along phi and psi at the end of simulation + for i in range(2): + ax = fig.add_subplot(3, 2, i + 5) + color = "blue" if i == 0 else "red" + xlabel = r"$\phi$" if i == 0 else r"$\phi$" + y = A_cv1 if i == 0 else A_cv2 + ax.plot(xi_1D, y, lw=3, color=color) + ax.set_xticks(np.arange(-pi, pi + pi / 2, step=(pi / 2)), ["-π", "-π/2", "0", "π/2", "π"]) + ax.set_xlabel(xlabel) + ax.set_ylabel(r"$A~[k_{B}T]$") + + fig.savefig("adp-fe.png", dpi=fig.dpi) + + return result + + +# %% +if __name__ == "__main__": + main(sys.argv[1:]) diff --git a/pysages/approxfun/core.py b/pysages/approxfun/core.py index fd76ba79..7075cfb0 100644 --- a/pysages/approxfun/core.py +++ b/pysages/approxfun/core.py @@ -109,13 +109,14 @@ def compute_mesh(grid: Grid): """ Returns a dense mesh with the same shape as `grid`, but on the hypercube [-1, 1]ⁿ, where `n` is the dimensionality of `grid`. + For parallelbias, only the diagonal elements are returned. """ h = 2 / grid.shape o = -1 + h / 2 nodes = o + h * np.hstack([np.arange(i).reshape(-1, 1) for i in grid.shape]) - return _compute_mesh(nodes) + return nodes if grid.parallelbias else _compute_mesh(nodes) @dispatch @@ -123,6 +124,7 @@ def compute_mesh(grid: Grid[Chebyshev]): # noqa: F811 # pylint: disable=C0116,E """ Returns a Chebyshev-distributed dense mesh with the same shape as `grid`, but on the hypercube [-1, 1]ⁿ, where n is the dimensionality of `grid`. + For parallelbias, only the diagonal elements are returned. """ def transform(n): @@ -130,7 +132,7 @@ def transform(n): nodes = np.hstack([transform(i)(np.arange(i).reshape(-1, 1)) for i in grid.shape]) - return _compute_mesh(nodes) + return nodes if grid.parallelbias else _compute_mesh(nodes) def _compute_mesh(nodes): diff --git a/pysages/grids.py b/pysages/grids.py index d8fdb63e..27d0f876 100644 --- a/pysages/grids.py +++ b/pysages/grids.py @@ -47,13 +47,15 @@ def __init__(self, lower, upper, shape, **kwargs): self.upper = np.asarray(upper).reshape(n) self.shape = shape.reshape(n) self.size = self.upper - self.lower + self.parallelbias = kwargs.get("parallelbias", False) def __check_init_invariants__(self, **kwargs): T = type(self).type_parameter if not (issubclass(type(T), type) and issubclass(T, GridType)): raise TypeError("Type parameter must be a subclass of GridType.") - if len(kwargs) > 1 or (len(kwargs) == 1 and "periodic" not in kwargs): - raise ValueError("Invalid keyword argument") + if len(kwargs) >= 1: + if ("periodic" not in kwargs) and ("parallelbias" not in kwargs): + raise ValueError("Invalid keyword arguments") periodic = kwargs.get("periodic", T is Periodic) if type(periodic) is not bool: raise TypeError("`periodic` must be a bool.") @@ -118,7 +120,7 @@ def get_index(x): h = grid.size / grid.shape idx = (x.flatten() - grid.lower) // h idx = np.where((idx < 0) | (idx > grid.shape), grid.shape, idx) - return (*np.flip(np.uint32(idx)),) + return np.uint32(idx) if grid.parallelbias else (*np.flip(np.uint32(idx)),) return jit(get_index) @@ -135,7 +137,7 @@ def get_index(x): h = grid.size / grid.shape idx = (x.flatten() - grid.lower) // h idx = idx % grid.shape - return (*np.flip(np.uint32(idx)),) + return np.uint32(idx) if grid.parallelbias else (*np.flip(np.uint32(idx)),) return jit(get_index) @@ -153,6 +155,6 @@ def get_index(x): x = 2 * (grid.lower - x.flatten()) / grid.size + 1 idx = (grid.shape * np.arccos(x)) // np.pi idx = np.nan_to_num(idx, nan=grid.shape) - return (*np.flip(np.uint32(idx)),) + return np.uint32(idx) if grid.parallelbias else (*np.flip(np.uint32(idx)),) return jit(get_index) diff --git a/pysages/methods/__init__.py b/pysages/methods/__init__.py index dbeeab15..25ab8097 100644 --- a/pysages/methods/__init__.py +++ b/pysages/methods/__init__.py @@ -68,6 +68,7 @@ from .funn import FUNN from .harmonic_bias import HarmonicBias from .metad import Metadynamics +from .pbmetad import ParallelBiasMetadynamics from .restraints import CVRestraints from .spectral_abf import SpectralABF from .spline_string import SplineString diff --git a/pysages/methods/pbmetad.py b/pysages/methods/pbmetad.py new file mode 100644 index 00000000..88b15842 --- /dev/null +++ b/pysages/methods/pbmetad.py @@ -0,0 +1,394 @@ +# SPDX-License-Identifier: MIT +# Copyright (c) 2020-2021: PySAGES contributors +# See LICENSE.md and CONTRIBUTORS.md at https://github.com/SSAGESLabs/PySAGES + +""" +Implementation of Parallel Bias Well-tempered Metadynamics with optional support for grids. +""" + +from jax import numpy as np, grad, jit, vmap +from jax.lax import cond + +from pysages.approxfun import compute_mesh +from pysages.collective_variables import get_periods, wrap +from pysages.methods.core import Result, GriddedSamplingMethod, generalize +from pysages.utils import identity +from pysages.grids import build_indexer +from pysages.utils import dispatch +from pysages.methods.metad import MetadynamicsState, PartialMetadynamicsState + + +class ParallelBiasMetadynamics(GriddedSamplingMethod): + """ + Implementation of Parallel Bias Well-tempered Metadynamics (PBMetaD) as described in + [J. Chem. Theory Comput. 11, 5062–5067 (2015)](https://doi.org/10.1021/acs.jctc.5b00846) + + Compared to well-tempered metadynamics, the Gaussian bias deposited along + each CV have different heights in PBMetaD. In addition, the total bias potential + involves a log of sum of exponential of bias potential (see Eq. 8 in the PBMetaD paper) + compared to just sum of Gaussians in well-tempered metadynamics. + + Because the method requires sampling along each CV separately, only the diagonal center + points of the grids are required for storing potential along each CV and to store the + net gradient of bias in PBMetaD. For this, the keyword + ``parallelbias`` is added to define grids for each CV separately. Currently, only + same number of bins for each CV is supported, which is the default. + """ + + snapshot_flags = {"positions", "indices"} + + def __init__(self, cvs, height, sigma, stride, ngaussians, T, deltaT, kB, **kwargs): + """ + Arguments + --------- + + cvs: + Set of user selected collective variables. + + height: JaxArray + Initial heights of the deposited Gaussians along each CV. + + sigma: JaxArray + Initial standard deviation of the to-be-deposit Gaussians along each CV. + + stride: int + Bias potential deposition frequency. + + ngaussians: int + Total number of expected gaussians (timesteps // stride + 1). + + deltaT: float + Well-tempered metadynamics $\\Delta T$ parameter to set the energy + scale for sampling. + + kB: float + Boltzmann constant. It should match the internal units of the backend. + + Keyword arguments + ----------------- + + grid: Optional[Grid] = None + If provided, it will be used to accelerate the computation by + approximating the bias potential along each CV and the gradient + of the total paralle bias potential over the grid centers. + """ + + kwargs["grid"] = kwargs.get("grid", None) + super().__init__(cvs, **kwargs) + + self.height = height + self.sigma = sigma + self.stride = stride + self.ngaussians = ngaussians # NOTE: infer from timesteps and stride + self.T = T + self.deltaT = deltaT + self.kB = kB + + def build(self, snapshot, helpers, *args, **kwargs): + return _parallelbiasmetadynamics(self, snapshot, helpers) + + +def _parallelbiasmetadynamics(method, snapshot, helpers): + # Initialization and update of biasing forces. Interface as expected for methods. + cv = method.cv + stride = method.stride + ngaussians = method.ngaussians + natoms = np.size(snapshot.positions, 0) + + deposit_gaussian = build_gaussian_accumulator(method) + evaluate_bias_grad = build_bias_grad_evaluator(method) + + def initialize(): + bias = np.zeros((natoms, 3), dtype=np.float64) + xi, _ = cv(helpers.query(snapshot)) + + # NOTE: for restart; use hills file to initialize corresponding arrays. + heights = np.zeros((ngaussians, xi.size), dtype=np.float64) + centers = np.zeros((ngaussians, xi.size), dtype=np.float64) + sigmas = np.array(method.sigma, dtype=np.float64, ndmin=2) + + # Arrays to store forces and bias potential on a grid. + if method.grid is None: + grid_potential = grid_gradient = None + else: + shape = method.grid.shape + # NOTE: for now, we assume, number of grid bins defined by shape along each + # CV are same. This need not be the case for PBMetaD as it generates + # free energy along each CV separately. We can define separate grids for each CV + # but that is not supported yet. + + # Given bins or shape of each CV is same, + # we use shape[0] to define the size of grids. + grid_potential = np.zeros((shape[0], shape.size), dtype=np.float64) + grid_gradient = np.zeros((shape[0], shape.size), dtype=np.float64) + + return MetadynamicsState( + bias, xi, heights, centers, sigmas, grid_potential, grid_gradient, 0, 0 + ) + + def update(state, data): + # Compute the collective variable and its jacobian + xi, Jxi = cv(data) + + # Deposit Gaussian depending on the stride + nstep = state.nstep + in_deposition_step = (nstep > 0) & (nstep % stride == 0) + partial_state = deposit_gaussian(xi, state, in_deposition_step) + + # Evaluate gradient of biasing potential (or generalized force) + generalized_force = evaluate_bias_grad(partial_state) + + # Calculate biasing forces + bias = -Jxi.T @ generalized_force.flatten() + bias = bias.reshape(state.bias.shape) + + return MetadynamicsState(bias, *partial_state[:-1], nstep + 1) + + return snapshot, initialize, generalize(update, helpers, jit_compile=False) + + +def build_gaussian_accumulator(method: ParallelBiasMetadynamics): + """ + Returns a function that given a `MetadynamicsState`, and the value of the CV, + stores the next Gaussian that is added to the biasing potential. + """ + periods = get_periods(method.cvs) + height_0 = np.array(method.height, dtype=np.float64) + T = method.T + deltaT = method.deltaT + grid = method.grid + kB = method.kB + beta = 1 / (kB * T) + kB_deltaT = kB * deltaT + + if grid is None: + evaluate_potential_each_cv = jit( + lambda pstate: np.sum(parallelbias_each_cv(*pstate[:4], periods), axis=0) + ) + else: + # each index in pstate.grid_idx correpsonds to different CV. + # so, we extract it using np.choose. mode='clip' is added for jit compilation. + evaluate_potential_each_cv = jit( + lambda pstate: np.choose(np.array(pstate.grid_idx), pstate.grid_potential, mode="clip") + ) + + def next_height(pstate): + V = evaluate_potential_each_cv(pstate) + w = height_0 * np.exp(-V / kB_deltaT) + cv_switching_probability_sum = np.sum(np.exp(-beta * V)) + cv_switching_probability = np.exp(-beta * V) / cv_switching_probability_sum + + return w * cv_switching_probability + + if grid is None: + get_grid_index = jit(lambda arg: None) + update_grids = jit(lambda *args: (None, None)) + should_deposit = jit(lambda pred, _: pred) + + else: + grid_mesh = compute_mesh(grid) * (grid.size / 2) + get_grid_index = build_indexer(grid) + # Reshape so the dimensions are compatible + accum = jit(lambda total, val: total + val.reshape(total.shape)) + update = jit(lambda V_each_cv, dV, vals, grads: (accum(V_each_cv, vals), accum(dV, grads))) + + def update_grids(pstate, height, xi, sigma): + # We need bias potential along each CV to update the heights. + current_parallelbias_each_cv = jit( + lambda x: parallelbias_each_cv(x, height, xi, sigma, periods) + ) + + # Total bias potential is required only for calculating and storing gradient of bias. + current_parallelbias = jit( + lambda x: parallelbias(x, height, xi, sigma, beta, periods, grid) + ) + + grid_potential_values = vmap(current_parallelbias_each_cv)(grid_mesh) + grid_grad_values = vmap(grad(current_parallelbias))(grid_mesh) + + return update( + pstate.grid_potential, pstate.grid_gradient, grid_potential_values, grid_grad_values + ) + + def should_deposit(in_deposition_step, I_xi): + in_bounds = ~(np.any(np.array(I_xi) == grid.shape)) + return in_deposition_step & in_bounds + + def deposit_gaussian(pstate): + xi, idx = pstate.xi, pstate.idx + current_height = next_height(pstate) + heights = pstate.heights.at[idx].set(current_height.flatten()) + centers = pstate.centers.at[idx].set(xi.flatten()) + sigmas = pstate.sigmas + grid_potential, grid_gradient = update_grids(pstate, current_height, xi, sigmas) + return PartialMetadynamicsState( + xi, heights, centers, sigmas, grid_potential, grid_gradient, idx + 1, pstate.grid_idx + ) + + def _deposit_gaussian(xi, state, in_deposition_step): + I_xi = get_grid_index(xi) + pstate = PartialMetadynamicsState(xi, *state[2:-1], I_xi) + predicate = should_deposit(in_deposition_step, I_xi) + return cond(predicate, deposit_gaussian, identity, pstate) + + return _deposit_gaussian + + +def build_bias_grad_evaluator(method: ParallelBiasMetadynamics): + """ + Returns a function that given the deposited Gaussians parameters, computes the + gradient of the biasing potential with respect to the CVs. + """ + grid = method.grid + T = method.T + kB = method.kB + beta = 1 / (kB * T) + + if grid is None: + periods = get_periods(method.cvs) + evaluate_bias_grad = jit( + lambda pstate: grad(parallelbias)(*pstate[:4], beta, periods, grid) + ) + else: + + def zero_force(_): + return np.zeros(grid.shape.size) + + def get_force(pstate): + # each index in pstate.grid_idx correpsonds to different CV. + # so, we extract it using np.choose + return np.choose(np.array(pstate.grid_idx), pstate.grid_gradient, mode="clip") + + def evaluate_bias_grad(pstate): + ob = np.any(np.array(pstate.grid_idx) == grid.shape) # out of bounds + return cond(ob, zero_force, get_force, pstate) + + return evaluate_bias_grad + + +# Helper function to evaluate parallel bias potential along each CV +def parallelbias_each_cv(xi, heights, centers, sigmas, periods): + """ + Evaluate parallel bias potential along each CV. + """ + delta_xi_each_cv = wrap(xi - centers, periods) + gaussian_each_cv = heights * np.exp(-((delta_xi_each_cv / sigmas) ** 2) / 2) + bias_each_cv = gaussian_each_cv + + return bias_each_cv + + +# Helper function to evaluate parallel bias potential +def parallelbias(xi, heights, centers, sigmas, beta, periods, grid): + """ + Evaluate parallel bias potential according to Eq. 8 in + [J. Chem. Theory Comput. 11, 5062–5067 (2015)](https://doi.org/10.1021/acs.jctc.5b00846) + """ + bias_each_cv = parallelbias_each_cv(xi, heights, centers, sigmas, periods) + bias_each_cv = bias_each_cv if grid else np.sum(bias_each_cv, axis=0) + exp_sum_gaussian = np.exp(-beta * bias_each_cv) + + return -(1 / beta) * np.log(np.sum(exp_sum_gaussian)) + + +@dispatch +def analyze(result: Result[ParallelBiasMetadynamics]): + """ + Helper for calculating the free energy from the final state of a + `Parallel Bias Metadynamics` run. + + Parameters + ---------- + + result: Result[ParallelBiasMetadynamics]: + Result bundle containing method, final parallel bias metadynamics state, and callback. + + Returns + ------- + + dict: + A dictionary with the following keys: + + centers: JaxArray + Centers of the CVs used for depositing Gaussian bias potential during the simulation. + + heights: JaxArray + Height of the Gaussian bias potential along each CV during the simulation. + + pbmetad_potential_cv: Callable + Maps a user-provided array of CV values to the corresponding deposited bias + potential. + + The free energy along each user-provided CV range is similar to well-tempered + metadynamics i.e., the free energy is equal to + `(T + deltaT) / deltaT * parallelbias_metapotential(cv)`, + where `T` is the simulation temperature and `deltaT` is the user-defined parameter + in parallel bias metadynamics. + + pbmetad_net_potential: Callable + Maps a user-provided array of CV values to the total parallel bias well-tempered + potential. Ideally, this can be used for obtaining multi-dimensional free energy + landscape using umbrella sampling like reweighting technique can be applied, + which is not yet supported. + """ + method = result.method + states = result.states + + P = get_periods(method.cvs) + grid = method.grid + + if len(states) == 1: + heights = states[0].heights + centers = states[0].centers + sigmas = states[0].sigmas + + pbmetad_potential_cv = jit( + vmap(lambda x: np.sum(parallelbias_each_cv(x, heights, centers, sigmas, P), axis=0)) + ) + pbmetad_net_potential = jit( + vmap( + lambda x, beta: parallelbias(x, heights, centers, sigmas, beta, P, grid), + in_axes=(0, None), + ) + ) + + return dict( + centers=centers, + heights=heights, + pbmetad_potential_cv=pbmetad_potential_cv, + pbmetad_net_potential=pbmetad_net_potential, + ) + + # For multiple-replicas runs we return a list heights and functions + # (one for each replica) + + def build_pbmetapotential_cv(heights, centers, sigmas): + return jit( + vmap(lambda x: np.sum(parallelbias_each_cv(x, heights, centers, sigmas, P), axis=0)) + ) + + def build_pbmetapotential(heights, centers, sigmas): + return jit( + vmap( + lambda x, beta: parallelbias(x, heights, centers, sigmas, beta, P, grid), + in_axes=(0, None), + ) + ) + + heights = [] + centers = [] + pbmetad_potentials_cv = [] + pbmetad_net_potentials = [] + + for s in states: + centers.append(s.centers) + heights.append(s.heights) + pbmetad_potentials_cv.append(build_pbmetapotential_cv(s.heights, s.centers, s.sigmas)) + pbmetad_net_potentials.append(build_pbmetapotential(s.heights, s.centers, s.sigmas)) + + return dict( + centers=centers, + heights=heights, + pbmetad_potential_cv=pbmetad_potentials_cv, + pbmetad_net_potential=pbmetad_net_potentials, + ) diff --git a/pysages/methods/utils.py b/pysages/methods/utils.py index a52894c1..b8912628 100644 --- a/pysages/methods/utils.py +++ b/pysages/methods/utils.py @@ -156,7 +156,7 @@ def save_hills(self, xi, sigma, height): f.write(str(self.counter) + "\t") f.write("\t".join(map(str, xi.flatten())) + "\t") f.write("\t".join(map(str, sigma.flatten())) + "\t") - f.write(str(height) + "\n") + f.write("\t".join(map(str, height.flatten())) + "\n") def __call__(self, snapshot, state, timestep): """ diff --git a/tests/test_pickle.py b/tests/test_pickle.py index efd797b6..38cc8a20 100644 --- a/tests/test_pickle.py +++ b/tests/test_pickle.py @@ -77,6 +77,17 @@ def build_neighbor_list(box_size, positions, r_cutoff, capacity_multiplier): "grid": pysages.Grid(lower=(-pi), upper=(pi), shape=(32), periodic=True), "kB": 614.0, }, + "ParallelBiasMetadynamics": { + "cvs": [pysages.colvars.Component([0], 0)], + "height": 1.0, + "sigma": 5.0, + "stride": 7.0, + "ngaussians": 5, + "T": 300, + "deltaT": 0.1, + "kB": 614.0, + "grid": pysages.Grid(lower=(-pi), upper=(pi), shape=(32), periodic=True, parallelbias=True), + }, "SpectralABF": { "cvs": [pysages.colvars.Component([0], 0), pysages.colvars.Component([0], 1)], "grid": pysages.Grid(lower=(1, 1), upper=(5, 5), shape=(32, 32)),