From 07510b5b5ef2ccd0ab6c1b961d02c970cc033dfc Mon Sep 17 00:00:00 2001 From: Siva Dasetty Date: Thu, 12 May 2022 08:58:35 -0500 Subject: [PATCH 01/29] created template for pbmetad --- pysages/methods/pbmetad.py | 289 +++++++++++++++++++++++++++++++++++++ 1 file changed, 289 insertions(+) create mode 100644 pysages/methods/pbmetad.py diff --git a/pysages/methods/pbmetad.py b/pysages/methods/pbmetad.py new file mode 100644 index 00000000..29c7739b --- /dev/null +++ b/pysages/methods/pbmetad.py @@ -0,0 +1,289 @@ +# SPDX-License-Identifier: MIT +# Copyright (c) 2020-2021: PySAGES contributors +# See LICENSE.md and CONTRIBUTORS.md at https://github.com/SSAGESLabs/PySAGES + +""" +Implementation of Standard and Well-tempered Metadynamics both with optional support for grids. +""" + +from typing import NamedTuple, Optional + +from jax import numpy as np, grad, jit, value_and_grad, 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 SamplingMethod, generalize +from pysages.utils import JaxArray, gaussian, identity +from pysages.grids import build_indexer + + +class MetadynamicsState(NamedTuple): + """ + Attributes + ---------- + + bias: JaxArray + Array of metadynamics bias forces for each particle in the simulation. + + xi: JaxArray + Collective variable value in the last simulation step. + + heights: JaxArray + Height values for all accumulated gaussians (zeros for not yet added gaussians). + + centers: JaxArray + Centers of the accumulated gaussians. + + sigmas: JaxArray + Widths of the accumulated gaussians. + + grid_potential: Optional[JaxArray] + Array of metadynamics bias potentials stored on a grid. + + grid_gradient: Optional[JaxArray] + Array of metadynamics bias gradients for each particle in the simulation stored on a grid. + + idx: int + Index of the next gaussian to be deposited. + + nstep: int + Counts the number of times `method.update` has been called. + """ + + bias: JaxArray + xi: JaxArray + heights: JaxArray + centers: JaxArray + sigmas: JaxArray + grid_potential: Optional[JaxArray] + grid_gradient: Optional[JaxArray] + idx: int + nstep: int + + def __repr__(self): + return repr("PySAGES" + type(self).__name__) + + +class PartialMetadynamicsState(NamedTuple): + """ + Helper intermediate Metadynamics state + """ + + xi: JaxArray + heights: JaxArray + centers: JaxArray + sigmas: JaxArray + grid_potential: Optional[JaxArray] + grid_gradient: Optional[JaxArray] + idx: int + grid_idx: Optional[JaxArray] + + +class Metadynamics(SamplingMethod): + """ + Implementation of Standard and Well-tempered Metadynamics as described in + [PNAS 99.20, 12562-6 (2002)](https://doi.org/10.1073/pnas.202427399) and + [Phys. Rev. Lett. 100, 020603 (2008)](https://doi.org/10.1103/PhysRevLett.100.020603) + """ + + snapshot_flags = {"positions", "indices"} + + def __init__(self, cvs, height, sigma, stride, ngaussians, *args, deltaT=None, **kwargs): + """ + Arguments + --------- + + cvs: + Set of user selected collective variable. + + height: + Initial height of the deposited Gaussians. + + sigma: + Initial standard deviation of the to-be-deposit Gaussians. + + stride: int + Bias potential deposition frequency. + + ngaussians: int + Total number of expected gaussians (timesteps // stride + 1). + + Keyword arguments + ----------------- + + deltaT: Optional[float] = None + Well-tempered metadynamics $\\Delta T$ parameter + (if `None` standard metadynamics is used). + + grid: Optional[Grid] = None + If provided, it will be used to accelerate the computation by + approximating the bias potential and its gradient over its centers. + + kB: Optional[float] + Boltzmann constant. Must be provided for well-tempered metadynamics + simulations and should match the internal units of the backend. + """ + + if deltaT is not None and "kB" not in kwargs: + raise KeyError( + "For well-tempered metadynamics a keyword argument `kB` for " + "the value of the Boltzmann constant (that matches the " + "internal units of the backend) must be provided." + ) + + super().__init__(cvs, args, kwargs) + + self.height = height + self.sigma = sigma + self.stride = stride + self.ngaussians = ngaussians # NOTE: infer from timesteps and stride + self.deltaT = deltaT + + self.kB = kwargs.get("kB", None) + self.grid = kwargs.get("grid", None) + + def build(self, snapshot, helpers, *args, **kwargs): + return _metadynamics(self, snapshot, helpers) + + +def _metadynamics(method, snapshot, helpers): + # Initialization and update of biasing forces. Interface 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, 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 + grid_potential = np.zeros((*shape,), dtype=np.float64) if method.deltaT else None + grid_gradient = np.zeros((*shape, 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=True) + + +def build_gaussian_accumulator(method: Metadynamics): + """ + 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 = method.height + deltaT = method.deltaT + grid = method.grid + kB = method.kB + + if deltaT is None: + next_height = jit(lambda *args: height_0) + else: # if well-tempered + if grid is None: + evaluate_potential = jit(lambda pstate: sum_of_gaussians(*pstate[:4], periods)) + else: + evaluate_potential = jit(lambda pstate: pstate.grid_potential[pstate.grid_idx]) + + def next_height(pstate): + V = evaluate_potential(pstate) + return height_0 * np.exp(-V / (deltaT * kB)) + + if grid is None: + get_grid_index = jit(lambda arg: None) + update_grids = jit(lambda *args: (None, None)) + 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)) + + if deltaT is None: + transform = grad + pack = jit(lambda x: (x,)) + # No need to accumulate values for the potential (V is None) + update = jit(lambda V, dV, vals: (V, accum(dV, vals))) + else: + transform = value_and_grad + pack = identity + update = jit(lambda V, dV, vals, grads: (accum(V, vals), accum(dV, grads))) + + def update_grids(pstate, height, xi, sigma): + # We use sum_of_gaussians since it already takes care of the wrapping + current_gaussian = jit(lambda x: sum_of_gaussians(x, height, xi, sigma, periods)) + # Evaluate gradient of bias (and bias potential for WT version) + grid_values = pack(vmap(transform(current_gaussian))(grid_mesh)) + return update(pstate.grid_potential, pstate.grid_gradient, *grid_values) + + def deposit_gaussian(pstate): + xi, idx = pstate.xi, pstate.idx + current_height = next_height(pstate) + heights = pstate.heights.at[idx].set(current_height) + 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): + pstate = PartialMetadynamicsState(xi, *state[2:-1], get_grid_index(xi)) + return cond(in_deposition_step, deposit_gaussian, identity, pstate) + + return _deposit_gaussian + + +def build_bias_grad_evaluator(method: Metadynamics): + """ + Returns a function that given the deposited gaussians parameters, computes the + gradient of the biasing potential with respect to the CVs. + """ + if method.grid is None: + periods = get_periods(method.cvs) + evaluate_bias_grad = jit(lambda pstate: grad(sum_of_gaussians)(*pstate[:4], periods)) + else: + evaluate_bias_grad = jit(lambda pstate: pstate.grid_gradient[pstate.grid_idx]) + + return evaluate_bias_grad + + +# Helper function to evaluate bias potential -- may be moved to analysis part +def sum_of_gaussians(xi, heights, centers, sigmas, periods): + """ + Sum of n-dimensional gaussians potential. + """ + delta_x = wrap(xi - centers, periods) + return gaussian(heights, sigmas, delta_x).sum() From 2d3f7a7f9ed3fe5d24b8e397ad91a5ca8823f644 Mon Sep 17 00:00:00 2001 From: Siva Dasetty Date: Tue, 19 Jul 2022 12:00:47 -0500 Subject: [PATCH 02/29] init modif of metad to pbmetad --- pysages/methods/pbmetad.py | 76 ++++---------------------------------- 1 file changed, 7 insertions(+), 69 deletions(-) diff --git a/pysages/methods/pbmetad.py b/pysages/methods/pbmetad.py index 29c7739b..7c3539e9 100644 --- a/pysages/methods/pbmetad.py +++ b/pysages/methods/pbmetad.py @@ -3,7 +3,7 @@ # See LICENSE.md and CONTRIBUTORS.md at https://github.com/SSAGESLabs/PySAGES """ -Implementation of Standard and Well-tempered Metadynamics both with optional support for grids. +Implementation of Parallel Bias Well-tempered Metadynamics with optional support for grids. """ from typing import NamedTuple, Optional @@ -16,75 +16,13 @@ from pysages.methods.core import SamplingMethod, generalize from pysages.utils import JaxArray, gaussian, identity from pysages.grids import build_indexer +from pysages.methods.metad import MetadynamicsState, PartialMetadynamicsState -class MetadynamicsState(NamedTuple): +class ParallelBiasMetadynamics(SamplingMethod): """ - Attributes - ---------- - - bias: JaxArray - Array of metadynamics bias forces for each particle in the simulation. - - xi: JaxArray - Collective variable value in the last simulation step. - - heights: JaxArray - Height values for all accumulated gaussians (zeros for not yet added gaussians). - - centers: JaxArray - Centers of the accumulated gaussians. - - sigmas: JaxArray - Widths of the accumulated gaussians. - - grid_potential: Optional[JaxArray] - Array of metadynamics bias potentials stored on a grid. - - grid_gradient: Optional[JaxArray] - Array of metadynamics bias gradients for each particle in the simulation stored on a grid. - - idx: int - Index of the next gaussian to be deposited. - - nstep: int - Counts the number of times `method.update` has been called. - """ - - bias: JaxArray - xi: JaxArray - heights: JaxArray - centers: JaxArray - sigmas: JaxArray - grid_potential: Optional[JaxArray] - grid_gradient: Optional[JaxArray] - idx: int - nstep: int - - def __repr__(self): - return repr("PySAGES" + type(self).__name__) - - -class PartialMetadynamicsState(NamedTuple): - """ - Helper intermediate Metadynamics state - """ - - xi: JaxArray - heights: JaxArray - centers: JaxArray - sigmas: JaxArray - grid_potential: Optional[JaxArray] - grid_gradient: Optional[JaxArray] - idx: int - grid_idx: Optional[JaxArray] - - -class Metadynamics(SamplingMethod): - """ - Implementation of Standard and Well-tempered Metadynamics as described in - [PNAS 99.20, 12562-6 (2002)](https://doi.org/10.1073/pnas.202427399) and - [Phys. Rev. Lett. 100, 020603 (2008)](https://doi.org/10.1103/PhysRevLett.100.020603) + Implementation of Parallel Bias Metadynamics as described in + [J. Chem. Theory Comput. 11, 5062–5067 (2015)](https://doi.org/10.1021/acs.jctc.5b00846) """ snapshot_flags = {"positions", "indices"} @@ -144,10 +82,10 @@ def __init__(self, cvs, height, sigma, stride, ngaussians, *args, deltaT=None, * self.grid = kwargs.get("grid", None) def build(self, snapshot, helpers, *args, **kwargs): - return _metadynamics(self, snapshot, helpers) + return _parallelbiasmetadynamics(self, snapshot, helpers) -def _metadynamics(method, snapshot, helpers): +def _parallelbiasmetadynamics(method, snapshot, helpers): # Initialization and update of biasing forces. Interface expected for methods. cv = method.cv stride = method.stride From 29f5e424df8498f8ed9291e98b5f310d63787b3f Mon Sep 17 00:00:00 2001 From: Siva Dasetty Date: Wed, 20 Jul 2022 14:35:50 -0500 Subject: [PATCH 03/29] update - separate heights for each CV --- pysages/methods/pbmetad.py | 115 ++++++++++++++++++++----------------- 1 file changed, 62 insertions(+), 53 deletions(-) diff --git a/pysages/methods/pbmetad.py b/pysages/methods/pbmetad.py index 7c3539e9..3864f21f 100644 --- a/pysages/methods/pbmetad.py +++ b/pysages/methods/pbmetad.py @@ -14,7 +14,7 @@ from pysages.approxfun import compute_mesh from pysages.collective_variables import get_periods, wrap from pysages.methods.core import SamplingMethod, generalize -from pysages.utils import JaxArray, gaussian, identity +from pysages.utils import JaxArray, identity from pysages.grids import build_indexer from pysages.methods.metad import MetadynamicsState, PartialMetadynamicsState @@ -23,6 +23,9 @@ class ParallelBiasMetadynamics(SamplingMethod): """ Implementation of Parallel Bias Metadynamics 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 total bias potential expression and height at which + bias is deposited for each CV is different in parallel bias metadynamics. """ snapshot_flags = {"positions", "indices"} @@ -47,29 +50,22 @@ def __init__(self, cvs, height, sigma, stride, ngaussians, *args, deltaT=None, * ngaussians: int Total number of expected gaussians (timesteps // stride + 1). + deltaT: float = None + 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 ----------------- - - deltaT: Optional[float] = None - Well-tempered metadynamics $\\Delta T$ parameter - (if `None` standard metadynamics is used). - + grid: Optional[Grid] = None If provided, it will be used to accelerate the computation by approximating the bias potential and its gradient over its centers. - kB: Optional[float] - Boltzmann constant. Must be provided for well-tempered metadynamics - simulations and should match the internal units of the backend. """ - if deltaT is not None and "kB" not in kwargs: - raise KeyError( - "For well-tempered metadynamics a keyword argument `kB` for " - "the value of the Boltzmann constant (that matches the " - "internal units of the backend) must be provided." - ) - super().__init__(cvs, args, kwargs) self.height = height @@ -77,8 +73,8 @@ def __init__(self, cvs, height, sigma, stride, ngaussians, *args, deltaT=None, * self.stride = stride self.ngaussians = ngaussians # NOTE: infer from timesteps and stride self.deltaT = deltaT - self.kB = kwargs.get("kB", None) + self.grid = kwargs.get("grid", None) def build(self, snapshot, helpers, *args, **kwargs): @@ -100,7 +96,7 @@ def initialize(): xi, _ = cv(helpers.query(snapshot)) # NOTE: for restart; use hills file to initialize corresponding arrays. - heights = np.zeros(ngaussians, dtype=np.float64) + 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) @@ -109,7 +105,7 @@ def initialize(): grid_potential = grid_gradient = None else: shape = method.grid.shape - grid_potential = np.zeros((*shape,), dtype=np.float64) if method.deltaT else None + grid_potential = np.zeros((*shape,), dtype=np.float64) grid_gradient = np.zeros((*shape, shape.size), dtype=np.float64) return MetadynamicsState( @@ -148,43 +144,35 @@ def build_gaussian_accumulator(method: Metadynamics): grid = method.grid kB = method.kB - if deltaT is None: - next_height = jit(lambda *args: height_0) - else: # if well-tempered - if grid is None: - evaluate_potential = jit(lambda pstate: sum_of_gaussians(*pstate[:4], periods)) - else: - evaluate_potential = jit(lambda pstate: pstate.grid_potential[pstate.grid_idx]) + if grid is None: + evaluate_potential = jit(lambda pstate: sum_of_gaussians(*pstate[:4], periods)) + #else: + # evaluate_potential = jit(lambda pstate: pstate.grid_potential[pstate.grid_idx]) - def next_height(pstate): - V = evaluate_potential(pstate) - return height_0 * np.exp(-V / (deltaT * kB)) + def next_height(pstate): + V = evaluate_potential(pstate) + return height_0 * np.exp(-V / (deltaT * kB)) if grid is None: get_grid_index = jit(lambda arg: None) update_grids = jit(lambda *args: (None, None)) - 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)) - - if deltaT is None: - transform = grad - pack = jit(lambda x: (x,)) - # No need to accumulate values for the potential (V is None) - update = jit(lambda V, dV, vals: (V, accum(dV, vals))) - else: - transform = value_and_grad - pack = identity - update = jit(lambda V, dV, vals, grads: (accum(V, vals), accum(dV, grads))) - - def update_grids(pstate, height, xi, sigma): - # We use sum_of_gaussians since it already takes care of the wrapping - current_gaussian = jit(lambda x: sum_of_gaussians(x, height, xi, sigma, periods)) - # Evaluate gradient of bias (and bias potential for WT version) - grid_values = pack(vmap(transform(current_gaussian))(grid_mesh)) - return update(pstate.grid_potential, pstate.grid_gradient, *grid_values) + #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)) +# +# + # transform = value_and_grad + # pack = identity + # update = jit(lambda V, dV, vals, grads: (accum(V, vals), accum(dV, grads))) +# + # def update_grids(pstate, height, xi, sigma): + # # We use sum_of_gaussians since it already takes care of the wrapping + # current_gaussian = jit(lambda x: sum_of_gaussians(x, height, xi, sigma, periods)) + # # Evaluate gradient of bias (and bias potential for WT version) + # grid_values = pack(vmap(transform(current_gaussian))(grid_mesh)) + # return update(pstate.grid_potential, pstate.grid_gradient, *grid_values) def deposit_gaussian(pstate): xi, idx = pstate.xi, pstate.idx @@ -212,8 +200,8 @@ def build_bias_grad_evaluator(method: Metadynamics): if method.grid is None: periods = get_periods(method.cvs) evaluate_bias_grad = jit(lambda pstate: grad(sum_of_gaussians)(*pstate[:4], periods)) - else: - evaluate_bias_grad = jit(lambda pstate: pstate.grid_gradient[pstate.grid_idx]) + #else: + # evaluate_bias_grad = jit(lambda pstate: pstate.grid_gradient[pstate.grid_idx]) return evaluate_bias_grad @@ -225,3 +213,24 @@ def sum_of_gaussians(xi, heights, centers, sigmas, periods): """ delta_x = wrap(xi - centers, periods) return gaussian(heights, sigmas, delta_x).sum() + + +def parallel_bias(x): + """ + Sum array `x` along each of its row (`axis = 1`), + Parallel bias potential + V_PB = + """ + # calculate epx( (-1/kT)*Gaussian ) for each CV using VMAP + vmap(gaussian_singleCV)(x) + + # sum the result and return it + + return v_pb + + +def gaussian_singleCV(a, sigma, cv_diff): + """ + N-dimensional origin-centered gaussian with height `a` and standard deviation `sigma`. + """ + return a * np.exp(-(cv_diff / sigma) ** 2 / 2) \ No newline at end of file From c4fac86a20f651354bfef24e2f5a33381bb67f68 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Wed, 20 Jul 2022 19:36:23 +0000 Subject: [PATCH 04/29] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- pysages/methods/pbmetad.py | 28 ++++++++++++++-------------- 1 file changed, 14 insertions(+), 14 deletions(-) diff --git a/pysages/methods/pbmetad.py b/pysages/methods/pbmetad.py index 3864f21f..e89179c6 100644 --- a/pysages/methods/pbmetad.py +++ b/pysages/methods/pbmetad.py @@ -23,7 +23,7 @@ class ParallelBiasMetadynamics(SamplingMethod): """ Implementation of Parallel Bias Metadynamics 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 total bias potential expression and height at which bias is deposited for each CV is different in parallel bias metadynamics. """ @@ -53,13 +53,13 @@ def __init__(self, cvs, height, sigma, stride, ngaussians, *args, deltaT=None, * deltaT: float = None 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 and its gradient over its centers. @@ -74,7 +74,7 @@ def __init__(self, cvs, height, sigma, stride, ngaussians, *args, deltaT=None, * self.ngaussians = ngaussians # NOTE: infer from timesteps and stride self.deltaT = deltaT self.kB = kwargs.get("kB", None) - + self.grid = kwargs.get("grid", None) def build(self, snapshot, helpers, *args, **kwargs): @@ -146,7 +146,7 @@ def build_gaussian_accumulator(method: Metadynamics): if grid is None: evaluate_potential = jit(lambda pstate: sum_of_gaussians(*pstate[:4], periods)) - #else: + # else: # evaluate_potential = jit(lambda pstate: pstate.grid_potential[pstate.grid_idx]) def next_height(pstate): @@ -156,17 +156,17 @@ def next_height(pstate): if grid is None: get_grid_index = jit(lambda arg: None) update_grids = jit(lambda *args: (None, None)) - #else: + # 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)) -# -# + # + # # transform = value_and_grad # pack = identity # update = jit(lambda V, dV, vals, grads: (accum(V, vals), accum(dV, grads))) -# + # # def update_grids(pstate, height, xi, sigma): # # We use sum_of_gaussians since it already takes care of the wrapping # current_gaussian = jit(lambda x: sum_of_gaussians(x, height, xi, sigma, periods)) @@ -200,7 +200,7 @@ def build_bias_grad_evaluator(method: Metadynamics): if method.grid is None: periods = get_periods(method.cvs) evaluate_bias_grad = jit(lambda pstate: grad(sum_of_gaussians)(*pstate[:4], periods)) - #else: + # else: # evaluate_bias_grad = jit(lambda pstate: pstate.grid_gradient[pstate.grid_idx]) return evaluate_bias_grad @@ -219,13 +219,13 @@ def parallel_bias(x): """ Sum array `x` along each of its row (`axis = 1`), Parallel bias potential - V_PB = + V_PB = """ # calculate epx( (-1/kT)*Gaussian ) for each CV using VMAP vmap(gaussian_singleCV)(x) - + # sum the result and return it - + return v_pb @@ -233,4 +233,4 @@ def gaussian_singleCV(a, sigma, cv_diff): """ N-dimensional origin-centered gaussian with height `a` and standard deviation `sigma`. """ - return a * np.exp(-(cv_diff / sigma) ** 2 / 2) \ No newline at end of file + return a * np.exp(-((cv_diff / sigma) ** 2) / 2) From 277381969ebfedb566ab64f318fd2cd12e7efe7e Mon Sep 17 00:00:00 2001 From: Siva Dasetty Date: Thu, 21 Jul 2022 00:02:04 -0500 Subject: [PATCH 05/29] draft of pbmetad without grids support --- pysages/methods/pbmetad.py | 74 ++++++++++++++++++++++---------------- 1 file changed, 43 insertions(+), 31 deletions(-) diff --git a/pysages/methods/pbmetad.py b/pysages/methods/pbmetad.py index 3864f21f..5cc84ed9 100644 --- a/pysages/methods/pbmetad.py +++ b/pysages/methods/pbmetad.py @@ -30,7 +30,7 @@ class ParallelBiasMetadynamics(SamplingMethod): snapshot_flags = {"positions", "indices"} - def __init__(self, cvs, height, sigma, stride, ngaussians, *args, deltaT=None, **kwargs): + def __init__(self, cvs, height, sigma, stride, ngaussians, *args, T, deltaT, kB, **kwargs): """ Arguments --------- @@ -72,6 +72,7 @@ def __init__(self, cvs, height, sigma, stride, ngaussians, *args, deltaT=None, * self.sigma = sigma self.stride = stride self.ngaussians = ngaussians # NOTE: infer from timesteps and stride + self.T = T self.deltaT = deltaT self.kB = kwargs.get("kB", None) @@ -87,6 +88,11 @@ def _parallelbiasmetadynamics(method, snapshot, helpers): stride = method.stride ngaussians = method.ngaussians natoms = np.size(snapshot.positions, 0) + T = method.T + deltaT = method.deltaT + kB = method.kB + beta = 1/(kB * T) + kB_deltaT = kB * deltaT deposit_gaussian = build_gaussian_accumulator(method) evaluate_bias_grad = build_bias_grad_evaluator(method) @@ -119,7 +125,7 @@ def update(state, 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) + partial_state = deposit_gaussian(xi, state, in_deposition_step, beta, kB_deltaT) # Evaluate gradient of biasing potential (or generalized force) generalized_force = evaluate_bias_grad(partial_state) @@ -140,18 +146,24 @@ def build_gaussian_accumulator(method: Metadynamics): """ periods = get_periods(method.cvs) height_0 = method.height + T = method.T deltaT = method.deltaT grid = method.grid kB = method.kB + beta = 1/(kB * T) if grid is None: - evaluate_potential = jit(lambda pstate: sum_of_gaussians(*pstate[:4], periods)) + evaluate_potential_cv = jit(lambda pstate: get_bias_each_cv(*pstate[:4], periods)) + evaluate_potential = jit(lambda pstate: parallelbias_potential(*pstate[:4], beta, periods)) #else: # evaluate_potential = jit(lambda pstate: pstate.grid_potential[pstate.grid_idx]) - def next_height(pstate): - V = evaluate_potential(pstate) - return height_0 * np.exp(-V / (deltaT * kB)) + def next_height(pstate, beta, kB_deltaT): + V = evaluate_potential_cv(pstate) + w = height_0 * np.exp(-V / (kB_deltaT)) + switching_probability_sum = np.sum(np.exp(-beta*V)) + switching_probability = np.exp(-beta*V)/switching_probability_sum + return w * switching_probability if grid is None: get_grid_index = jit(lambda arg: None) @@ -169,14 +181,15 @@ def next_height(pstate): # # def update_grids(pstate, height, xi, sigma): # # We use sum_of_gaussians since it already takes care of the wrapping + # may be replace sum_of_gaussians with parallelbias_potential # current_gaussian = jit(lambda x: sum_of_gaussians(x, height, xi, sigma, periods)) # # Evaluate gradient of bias (and bias potential for WT version) # grid_values = pack(vmap(transform(current_gaussian))(grid_mesh)) # return update(pstate.grid_potential, pstate.grid_gradient, *grid_values) - def deposit_gaussian(pstate): + def deposit_gaussian(pstate, beta, kB_deltaT): xi, idx = pstate.xi, pstate.idx - current_height = next_height(pstate) + current_height = next_height(pstate, beta, kB_deltaT) heights = pstate.heights.at[idx].set(current_height) centers = pstate.centers.at[idx].set(xi.flatten()) sigmas = pstate.sigmas @@ -199,38 +212,37 @@ def build_bias_grad_evaluator(method: Metadynamics): """ if method.grid is None: periods = get_periods(method.cvs) - evaluate_bias_grad = jit(lambda pstate: grad(sum_of_gaussians)(*pstate[:4], periods)) + T = method.T + kB = method.kB + beta = 1/(kB * T) + evaluate_bias_grad = jit(lambda pstate: grad(parallelbias_potential)(*pstate[:4], beta, periods)) #else: # evaluate_bias_grad = jit(lambda pstate: pstate.grid_gradient[pstate.grid_idx]) return evaluate_bias_grad -# Helper function to evaluate bias potential -- may be moved to analysis part -def sum_of_gaussians(xi, heights, centers, sigmas, periods): - """ - Sum of n-dimensional gaussians potential. - """ - delta_x = wrap(xi - centers, periods) - return gaussian(heights, sigmas, delta_x).sum() - - -def parallel_bias(x): +# Helper function to evaluate parallel bias potential +def parallelbias_potential(xi, heights, centers, sigmas, beta, periods): """ - Sum array `x` along each of its row (`axis = 1`), - Parallel bias potential - V_PB = + Evaluate parallel bias potential Eq. 7 in + [J. Chem. Theory Comput. 11, 5062–5067 (2015)](https://doi.org/10.1021/acs.jctc.5b00846) """ - # calculate epx( (-1/kT)*Gaussian ) for each CV using VMAP - vmap(gaussian_singleCV)(x) - - # sum the result and return it + bias_each_cv = get_bias_each_cv((xi, heights, centers, sigmas, periods) + exp_sum_gaussian = np.exp(-beta*sum_gaussian) - return v_pb + return -(1/beta) * np.log(np.sum(exp_sum_gaussian)) - -def gaussian_singleCV(a, sigma, cv_diff): +# Helper function to evaluate bias potential along each CV +def get_bias_each_cv(xi, heights, centers, sigmas, periods): """ - N-dimensional origin-centered gaussian with height `a` and standard deviation `sigma`. + Evaluate parallel bias potential along each CV. """ - return a * np.exp(-(cv_diff / sigma) ** 2 / 2) \ No newline at end of file + delta_x_each_cv = wrap(xi - centers, periods) + gaussian_each_cv = a * np.exp(-((delta_x_each_cv / sigma) ** 2) / 2) + bias_each_cv = np.sum(gaussian_each_cv, axis=0) + + return bias_each_cv + + + \ No newline at end of file From 63da8b4b979ae7c12defad4b86d62d0754778037 Mon Sep 17 00:00:00 2001 From: Siva Dasetty Date: Wed, 27 Jul 2022 12:28:31 -0500 Subject: [PATCH 06/29] Crude state pbmetad without grids support --- examples/openmm/pbmetad/alanine-dipeptide.py | 169 ++++++++++++++++ pysages/methods/pbmetad.py | 199 +++++++++++++------ pysages/methods/utils.py | 2 +- 3 files changed, 310 insertions(+), 60 deletions(-) create mode 100644 examples/openmm/pbmetad/alanine-dipeptide.py diff --git a/examples/openmm/pbmetad/alanine-dipeptide.py b/examples/openmm/pbmetad/alanine-dipeptide.py new file mode 100644 index 00000000..74d41190 --- /dev/null +++ b/examples/openmm/pbmetad/alanine-dipeptide.py @@ -0,0 +1,169 @@ +#!/usr/bin/env python3 + +""" +Parallel bias metadynamics simulation of Alanine Dipeptide 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 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_val = 298.15 +T = T_val * 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 = [ + ("well-tempered", "w", bool, 0, "Whether to use well-tempered metadynamics"), + ("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) + + 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) + grid = grid if args.use_grids else None + + # Method + method = ParallelBiasMetadynamics(cvs, height, sigma, stride, ngauss, T_val, deltaT, kB, grid=None) + + # 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_grid = pysages.Grid(lower=(-pi, -pi), upper=(pi, pi), shape=(64, 64), periodic=True) + #### xi = (compute_mesh(plot_grid) + 1) / 2 * plot_grid.size + plot_grid.lower +#### + #### # determine bias factor depending on method (for standard = 1 and for well-tempered = (T+deltaT)/deltaT) + #### alpha = ( + #### 1 + #### if method.deltaT is None + #### else (T.value_in_unit(unit.kelvin) + method.deltaT) / method.deltaT + #### ) + #### kT = kB * T.value_in_unit(unit.kelvin) +#### + #### # extract metapotential function from result + #### result = pysages.analyze(run_result) + #### metapotential = result["metapotential"] +#### + #### # report in kT and set min free energy to zero + #### A = metapotential(xi) * -alpha / kT + #### A = A - A.min() + #### A = A.reshape(plot_grid.shape) +#### + #### # plot and save free energy to a PNG file + #### fig, ax = plt.subplots(dpi=120) +#### + #### im = ax.imshow(A, interpolation="bicubic", origin="lower", extent=[-pi, pi, -pi, pi]) + #### ax.contour(A, levels=12, linewidths=0.75, colors="k", extent=[-pi, pi, -pi, pi]) + #### ax.set_xlabel(r"$\phi$") + #### ax.set_ylabel(r"$\psi$") +#### + #### cbar = plt.colorbar(im) + #### cbar.ax.set_ylabel(r"$A~[k_{B}T]$", rotation=270, labelpad=20) +#### + #### fig.savefig("adp-fe.png", dpi=fig.dpi) + + return run_result + + +# %% +if __name__ == "__main__": + main(sys.argv[1:]) diff --git a/pysages/methods/pbmetad.py b/pysages/methods/pbmetad.py index 5cc84ed9..63620b84 100644 --- a/pysages/methods/pbmetad.py +++ b/pysages/methods/pbmetad.py @@ -13,36 +13,39 @@ from pysages.approxfun import compute_mesh from pysages.collective_variables import get_periods, wrap -from pysages.methods.core import SamplingMethod, generalize +from pysages.methods.core import Result, GriddedSamplingMethod, generalize from pysages.utils import JaxArray, identity from pysages.grids import build_indexer +from pysages.utils import dispatch from pysages.methods.metad import MetadynamicsState, PartialMetadynamicsState -class ParallelBiasMetadynamics(SamplingMethod): +class ParallelBiasMetadynamics(GriddedSamplingMethod): """ Implementation of Parallel Bias Metadynamics 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 total bias potential expression and height at which - bias is deposited for each CV is different in parallel bias metadynamics. + Compared to well-tempered metadynamics, the height of Gaussian bias deposited along + each CV is different in parallel bias metadynamics. In addition, the total bias potential + has a slightly different expression involving the sum of exponential of Gaussians (see Eq. 8) + in the paper. """ snapshot_flags = {"positions", "indices"} - def __init__(self, cvs, height, sigma, stride, ngaussians, *args, T, deltaT, kB, **kwargs): + def __init__(self, cvs, height, sigma, stride, ngaussians, T, deltaT, kB, *args, **kwargs): """ Arguments --------- cvs: - Set of user selected collective variable. + Set of user selected collective variables. height: - Initial height of the deposited Gaussians. + Initial height of the deposited Gaussians along each CV. sigma: - Initial standard deviation of the to-be-deposit Gaussians. + Initial standard deviation of the to-be-deposit Gaussians along each CV. stride: int Bias potential deposition frequency. @@ -66,7 +69,8 @@ def __init__(self, cvs, height, sigma, stride, ngaussians, *args, T, deltaT, kB, """ - super().__init__(cvs, args, kwargs) + kwargs["grid"] = kwargs.get("grid", None) + super().__init__(cvs, **kwargs) self.height = height self.sigma = sigma @@ -74,9 +78,7 @@ def __init__(self, cvs, height, sigma, stride, ngaussians, *args, T, deltaT, kB, self.ngaussians = ngaussians # NOTE: infer from timesteps and stride self.T = T self.deltaT = deltaT - self.kB = kwargs.get("kB", None) - - self.grid = kwargs.get("grid", None) + self.kB = kB def build(self, snapshot, helpers, *args, **kwargs): return _parallelbiasmetadynamics(self, snapshot, helpers) @@ -88,11 +90,6 @@ def _parallelbiasmetadynamics(method, snapshot, helpers): stride = method.stride ngaussians = method.ngaussians natoms = np.size(snapshot.positions, 0) - T = method.T - deltaT = method.deltaT - kB = method.kB - beta = 1/(kB * T) - kB_deltaT = kB * deltaT deposit_gaussian = build_gaussian_accumulator(method) evaluate_bias_grad = build_bias_grad_evaluator(method) @@ -122,10 +119,10 @@ def update(state, data): # Compute the collective variable and its jacobian xi, Jxi = cv(data) - # Deposit gaussian depending on the stride + # 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, beta, kB_deltaT) + partial_state = deposit_gaussian(xi, state, in_deposition_step) # Evaluate gradient of biasing potential (or generalized force) generalized_force = evaluate_bias_grad(partial_state) @@ -136,61 +133,74 @@ def update(state, data): return MetadynamicsState(bias, *partial_state[:-1], nstep + 1) - return snapshot, initialize, generalize(update, helpers, jit_compile=True) + return snapshot, initialize, generalize(update, helpers, jit_compile=False) -def build_gaussian_accumulator(method: Metadynamics): +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. + stores the next Gaussian that is added to the biasing potential. """ periods = get_periods(method.cvs) - height_0 = method.height + 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: + if grid is None: evaluate_potential_cv = jit(lambda pstate: get_bias_each_cv(*pstate[:4], periods)) - evaluate_potential = jit(lambda pstate: parallelbias_potential(*pstate[:4], beta, periods)) + evaluate_potential = jit(lambda pstate: parallelbias(*pstate[:4], beta, periods)) + #else: # evaluate_potential = jit(lambda pstate: pstate.grid_potential[pstate.grid_idx]) - def next_height(pstate, beta, kB_deltaT): + def next_height(pstate): V = evaluate_potential_cv(pstate) - w = height_0 * np.exp(-V / (kB_deltaT)) + w = height_0 * np.exp(-V / kB_deltaT) switching_probability_sum = np.sum(np.exp(-beta*V)) switching_probability = np.exp(-beta*V)/switching_probability_sum + return w * 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)) -# -# - # transform = value_and_grad - # pack = identity - # update = jit(lambda V, dV, vals, grads: (accum(V, vals), accum(dV, grads))) -# + + # if deltaT is None: + # transform = grad + # pack = jit(lambda x: (x,)) + # # No need to accumulate values for the potential (V is None) + # update = jit(lambda V, dV, vals: (V, accum(dV, vals))) + # else: + # transform = value_and_grad + # pack = identity + # update = jit(lambda V, dV, vals, grads: (accum(V, vals), accum(dV, grads))) + # def update_grids(pstate, height, xi, sigma): - # # We use sum_of_gaussians since it already takes care of the wrapping - # may be replace sum_of_gaussians with parallelbias_potential + # # We use `sum_of_gaussians` since it already takes care of the wrapping # current_gaussian = jit(lambda x: sum_of_gaussians(x, height, xi, sigma, periods)) # # Evaluate gradient of bias (and bias potential for WT version) # grid_values = pack(vmap(transform(current_gaussian))(grid_mesh)) # return update(pstate.grid_potential, pstate.grid_gradient, *grid_values) - def deposit_gaussian(pstate, beta, kB_deltaT): + # 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, beta, kB_deltaT) - heights = pstate.heights.at[idx].set(current_height) + 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) @@ -199,50 +209,121 @@ def deposit_gaussian(pstate, beta, kB_deltaT): ) def _deposit_gaussian(xi, state, in_deposition_step): - pstate = PartialMetadynamicsState(xi, *state[2:-1], get_grid_index(xi)) - return cond(in_deposition_step, deposit_gaussian, identity, pstate) + 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: Metadynamics): +def build_bias_grad_evaluator(method: ParallelBiasMetadynamics): """ - Returns a function that given the deposited gaussians parameters, computes the + Returns a function that given the deposited Gaussians parameters, computes the gradient of the biasing potential with respect to the CVs. """ - if method.grid is None: + grid = method.grid + T = method.T + kB = method.kB + beta = 1/(kB * T) + + if grid is None: periods = get_periods(method.cvs) - T = method.T - kB = method.kB - beta = 1/(kB * T) - evaluate_bias_grad = jit(lambda pstate: grad(parallelbias_potential)(*pstate[:4], beta, periods)) - #else: - # evaluate_bias_grad = jit(lambda pstate: pstate.grid_gradient[pstate.grid_idx]) + evaluate_bias_grad = jit(lambda pstate: grad(parallelbias)(*pstate[:4], beta, periods)) + else: + + def zero_force(_): + return np.zeros(grid.shape.size) + + def get_force(pstate): + return pstate.grid_gradient[pstate.grid_idx] + + 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 -def parallelbias_potential(xi, heights, centers, sigmas, beta, periods): +def parallelbias(xi, heights, centers, sigmas, beta, periods): """ - Evaluate parallel bias potential Eq. 7 in + 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 = get_bias_each_cv((xi, heights, centers, sigmas, periods) - exp_sum_gaussian = np.exp(-beta*sum_gaussian) + bias_each_cv = get_bias_each_cv(xi, heights, centers, sigmas, periods) + exp_sum_gaussian = np.exp(-beta*bias_each_cv) return -(1/beta) * np.log(np.sum(exp_sum_gaussian)) -# Helper function to evaluate bias potential along each CV +# Helper function to evaluate parallel bias potential along each CV def get_bias_each_cv(xi, heights, centers, sigmas, periods): """ Evaluate parallel bias potential along each CV. """ - delta_x_each_cv = wrap(xi - centers, periods) - gaussian_each_cv = a * np.exp(-((delta_x_each_cv / sigma) ** 2) / 2) + delta_xi_each_cv = wrap(xi - centers, periods) + gaussian_each_cv = heights * np.exp(-((delta_xi_each_cv / sigmas) ** 2) / 2) bias_each_cv = np.sum(gaussian_each_cv, axis=0) return bias_each_cv - - - \ No newline at end of file + + +@dispatch +def analyze(result: Result[ParallelBiasMetadynamics]): + """ + Helper for calculating the free energy from the final state of a `Metadynamics` run. + + Parameters + ---------- + + result: Result[Metadynamics]: + Result bundle containing method, final metadynamics state, and callback. + + Returns + ------- + + dict: + A dictionary with the following keys: + + heights: JaxArray + Height of the Gaussian bias potential during the simulation. + + metapotential: Callable + Maps a user-provided array of CV values to the corresponding deposited bias + potential. For standard metadynamics, the free energy along user-provided CV + range is the same as `metapotential(cv)`. + In the case of well-tempered metadynamics, the free energy is equal to + `(T + deltaT) / deltaT * metapotential(cv)`, where `T` is the simulation + temperature and `deltaT` is the user-defined parameter in + well-tempered metadynamics. + """ + method = result.method + states = result.states + + P = get_periods(method.cvs) + + if len(states) == 1: + heights = states[0].heights + centers = states[0].centers + sigmas = states[0].sigmas + + metapotential = jit(vmap(lambda x: sum_of_gaussians(x, heights, centers, sigmas, P))) + + return dict(heights=heights, metapotential=metapotential) + + # For multiple-replicas runs we return a list heights and functions + # (one for each replica) + + def build_metapotential(heights, centers, sigmas): + return jit(vmap(lambda x: sum_of_gaussians(x, heights, centers, sigmas, P))) + + heights = [] + metapotentials = [] + + for s in states: + heights.append(s.heights) + metapotentials.append(build_metapotential(s.heights, s.centers, s.sigmas)) + + return dict(heights=heights, metapotential=metapotentials) + + diff --git a/pysages/methods/utils.py b/pysages/methods/utils.py index 382c72c3..b1ae71e0 100644 --- a/pysages/methods/utils.py +++ b/pysages/methods/utils.py @@ -147,7 +147,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): """ From ab9b6cc4083a5757878198c0908ba95c4d35cb43 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Wed, 27 Jul 2022 17:29:44 +0000 Subject: [PATCH 07/29] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- examples/openmm/pbmetad/alanine-dipeptide.py | 21 +++++----- pysages/methods/pbmetad.py | 41 ++++++++++---------- 2 files changed, 31 insertions(+), 31 deletions(-) diff --git a/examples/openmm/pbmetad/alanine-dipeptide.py b/examples/openmm/pbmetad/alanine-dipeptide.py index 74d41190..4b8a86df 100644 --- a/examples/openmm/pbmetad/alanine-dipeptide.py +++ b/examples/openmm/pbmetad/alanine-dipeptide.py @@ -107,14 +107,15 @@ def main(argv=[]): 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) grid = grid if args.use_grids else None # Method - method = ParallelBiasMetadynamics(cvs, height, sigma, stride, ngauss, T_val, deltaT, kB, grid=None) + method = ParallelBiasMetadynamics( + cvs, height, sigma, stride, ngauss, T_val, deltaT, kB, grid=None + ) # Logging hills_file = "hills.dat" @@ -126,11 +127,11 @@ def main(argv=[]): 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_grid = pysages.Grid(lower=(-pi, -pi), upper=(pi, pi), shape=(64, 64), periodic=True) #### xi = (compute_mesh(plot_grid) + 1) / 2 * plot_grid.size + plot_grid.lower -#### + #### #### # determine bias factor depending on method (for standard = 1 and for well-tempered = (T+deltaT)/deltaT) #### alpha = ( #### 1 @@ -138,27 +139,27 @@ def main(argv=[]): #### else (T.value_in_unit(unit.kelvin) + method.deltaT) / method.deltaT #### ) #### kT = kB * T.value_in_unit(unit.kelvin) -#### + #### #### # extract metapotential function from result #### result = pysages.analyze(run_result) #### metapotential = result["metapotential"] -#### + #### #### # report in kT and set min free energy to zero #### A = metapotential(xi) * -alpha / kT #### A = A - A.min() #### A = A.reshape(plot_grid.shape) -#### + #### #### # plot and save free energy to a PNG file #### fig, ax = plt.subplots(dpi=120) -#### + #### #### im = ax.imshow(A, interpolation="bicubic", origin="lower", extent=[-pi, pi, -pi, pi]) #### ax.contour(A, levels=12, linewidths=0.75, colors="k", extent=[-pi, pi, -pi, pi]) #### ax.set_xlabel(r"$\phi$") #### ax.set_ylabel(r"$\psi$") -#### + #### #### cbar = plt.colorbar(im) #### cbar.ax.set_ylabel(r"$A~[k_{B}T]$", rotation=270, labelpad=20) -#### + #### #### fig.savefig("adp-fe.png", dpi=fig.dpi) return run_result diff --git a/pysages/methods/pbmetad.py b/pysages/methods/pbmetad.py index 63620b84..684f96a2 100644 --- a/pysages/methods/pbmetad.py +++ b/pysages/methods/pbmetad.py @@ -24,7 +24,7 @@ class ParallelBiasMetadynamics(GriddedSamplingMethod): """ Implementation of Parallel Bias Metadynamics 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 height of Gaussian bias deposited along each CV is different in parallel bias metadynamics. In addition, the total bias potential has a slightly different expression involving the sum of exponential of Gaussians (see Eq. 8) @@ -56,13 +56,13 @@ def __init__(self, cvs, height, sigma, stride, ngaussians, T, deltaT, kB, *args, deltaT: float = None 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 and its gradient over its centers. @@ -147,30 +147,30 @@ def build_gaussian_accumulator(method: ParallelBiasMetadynamics): deltaT = method.deltaT grid = method.grid kB = method.kB - beta = 1/(kB * T) + beta = 1 / (kB * T) kB_deltaT = kB * deltaT - if grid is None: + if grid is None: evaluate_potential_cv = jit(lambda pstate: get_bias_each_cv(*pstate[:4], periods)) evaluate_potential = jit(lambda pstate: parallelbias(*pstate[:4], beta, periods)) - - #else: + + # else: # evaluate_potential = jit(lambda pstate: pstate.grid_potential[pstate.grid_idx]) def next_height(pstate): V = evaluate_potential_cv(pstate) - w = height_0 * np.exp(-V / kB_deltaT) - switching_probability_sum = np.sum(np.exp(-beta*V)) - switching_probability = np.exp(-beta*V)/switching_probability_sum - + w = height_0 * np.exp(-V / kB_deltaT) + switching_probability_sum = np.sum(np.exp(-beta * V)) + switching_probability = np.exp(-beta * V) / switching_probability_sum + return w * 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: + + # else: # grid_mesh = compute_mesh(grid) * (grid.size / 2) # get_grid_index = build_indexer(grid) # # Reshape so the dimensions are compatible @@ -225,7 +225,7 @@ def build_bias_grad_evaluator(method: ParallelBiasMetadynamics): grid = method.grid T = method.T kB = method.kB - beta = 1/(kB * T) + beta = 1 / (kB * T) if grid is None: periods = get_periods(method.cvs) @@ -248,13 +248,14 @@ def evaluate_bias_grad(pstate): # Helper function to evaluate parallel bias potential def parallelbias(xi, heights, centers, sigmas, beta, periods): """ - Evaluate parallel bias potential according to Eq. 8 in + 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 = get_bias_each_cv(xi, heights, centers, sigmas, periods) - exp_sum_gaussian = np.exp(-beta*bias_each_cv) - - return -(1/beta) * np.log(np.sum(exp_sum_gaussian)) + exp_sum_gaussian = np.exp(-beta * bias_each_cv) + + return -(1 / beta) * np.log(np.sum(exp_sum_gaussian)) + # Helper function to evaluate parallel bias potential along each CV def get_bias_each_cv(xi, heights, centers, sigmas, periods): @@ -264,7 +265,7 @@ def get_bias_each_cv(xi, heights, centers, sigmas, periods): delta_xi_each_cv = wrap(xi - centers, periods) gaussian_each_cv = heights * np.exp(-((delta_xi_each_cv / sigmas) ** 2) / 2) bias_each_cv = np.sum(gaussian_each_cv, axis=0) - + return bias_each_cv @@ -325,5 +326,3 @@ def build_metapotential(heights, centers, sigmas): metapotentials.append(build_metapotential(s.heights, s.centers, s.sigmas)) return dict(heights=heights, metapotential=metapotentials) - - From 107fd49dc54ce3d545503ebea775a05a23c979ef Mon Sep 17 00:00:00 2001 From: Siva Dasetty Date: Wed, 27 Jul 2022 12:29:50 -0500 Subject: [PATCH 08/29] Crude state added pbmetad to init --- pysages/methods/__init__.py | 1 + 1 file changed, 1 insertion(+) diff --git a/pysages/methods/__init__.py b/pysages/methods/__init__.py index 5ae6b6a6..83b19994 100644 --- a/pysages/methods/__init__.py +++ b/pysages/methods/__init__.py @@ -70,6 +70,7 @@ from .harmonic_bias import HarmonicBias from .umbrella_integration import UmbrellaIntegration from .metad import Metadynamics +from .pbmetad import ParallelBiasMetadynamics from .spectral_abf import SpectralABF from .utils import ( HistogramLogger, From 5d41d8d8c1128d695cbf142c89a7f9f9ab46955c Mon Sep 17 00:00:00 2001 From: Siva Dasetty Date: Wed, 27 Jul 2022 16:20:05 -0500 Subject: [PATCH 09/29] Crude nogrids - updated ADP example --- examples/openmm/pbmetad/alanine-dipeptide.py | 135 +++++++++++++------ pysages/methods/pbmetad.py | 62 ++++----- 2 files changed, 126 insertions(+), 71 deletions(-) diff --git a/examples/openmm/pbmetad/alanine-dipeptide.py b/examples/openmm/pbmetad/alanine-dipeptide.py index 4b8a86df..fcd6d0ef 100644 --- a/examples/openmm/pbmetad/alanine-dipeptide.py +++ b/examples/openmm/pbmetad/alanine-dipeptide.py @@ -16,6 +16,7 @@ import numpy import pysages +from jax import numpy as np from pysages.colvars import DihedralAngle from pysages.methods import ParallelBiasMetadynamics, MetaDLogger @@ -34,8 +35,7 @@ kB = unit.BOLTZMANN_CONSTANT_kB * unit.AVOGADRO_CONSTANT_NA kB = kB.value_in_unit(unit.kilojoules_per_mole / unit.kelvin) -T_val = 298.15 -T = T_val * 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") @@ -114,7 +114,7 @@ def main(argv=[]): # Method method = ParallelBiasMetadynamics( - cvs, height, sigma, stride, ngauss, T_val, deltaT, kB, grid=None + cvs, height, sigma, stride, ngauss, T.value_in_unit(unit.kelvin), deltaT, kB, grid=None ) # Logging @@ -126,43 +126,98 @@ def main(argv=[]): 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_grid = pysages.Grid(lower=(-pi, -pi), upper=(pi, pi), shape=(64, 64), periodic=True) - #### xi = (compute_mesh(plot_grid) + 1) / 2 * plot_grid.size + plot_grid.lower - #### - #### # determine bias factor depending on method (for standard = 1 and for well-tempered = (T+deltaT)/deltaT) - #### alpha = ( - #### 1 - #### if method.deltaT is None - #### else (T.value_in_unit(unit.kelvin) + method.deltaT) / method.deltaT - #### ) - #### kT = kB * T.value_in_unit(unit.kelvin) - #### - #### # extract metapotential function from result - #### result = pysages.analyze(run_result) - #### metapotential = result["metapotential"] - #### - #### # report in kT and set min free energy to zero - #### A = metapotential(xi) * -alpha / kT - #### A = A - A.min() - #### A = A.reshape(plot_grid.shape) - #### - #### # plot and save free energy to a PNG file - #### fig, ax = plt.subplots(dpi=120) - #### - #### im = ax.imshow(A, interpolation="bicubic", origin="lower", extent=[-pi, pi, -pi, pi]) - #### ax.contour(A, levels=12, linewidths=0.75, colors="k", extent=[-pi, pi, -pi, pi]) - #### ax.set_xlabel(r"$\phi$") - #### ax.set_ylabel(r"$\psi$") - #### - #### cbar = plt.colorbar(im) - #### cbar.ax.set_ylabel(r"$A~[k_{B}T]$", rotation=270, labelpad=20) - #### - #### fig.savefig("adp-fe.png", dpi=fig.dpi) - - return run_result + # 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) + plot_2Dgrid = pysages.Grid(lower=(-pi, -pi), upper=(pi, pi), shape=(64, 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)) + xi_2D = (compute_mesh(plot_2Dgrid) + 1) / 2 * plot_2Dgrid.size + plot_2Dgrid.lower + + # 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_potential = result["pbmetad_potential"] + + # report in kT and set min free energy to zero + A_cv = pbmetad_potential_cv(xi_1D_2CVs) * -alpha / kT + 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) + + #A = pbmetad_potential(xi_2D, beta) * -alpha / kT + #A = A - A.min() + #A = A.reshape(plot_2Dgrid.shape) + + # plot and save free energy along each CV to a PNG file + fig = plt.figure(figsize=(16,12),dpi=120) + fig.subplots_adjust(hspace=0.5,wspace=0.4) + + # plot centers along phi + ax = fig.add_subplot(3, 3, 1) + ax.scatter(np.arange(np.shape(centers)[0]), centers[:,0], s=20, color='blue') + ax.set_xlabel(r"Step") + ax.set_ylabel(r"$\phi$") + ax.set_yticks(np.arange(-pi, pi+pi/2, step=(pi/2)), ['-π','-π/2','0','π/2','π']) + + # plot centers along psi + ax = fig.add_subplot(3, 3, 2) + ax.scatter(np.arange(np.shape(centers)[0]), centers[:,1], s=20, color='red') + ax.set_xlabel(r"Step") + ax.set_ylabel(r"$\psi$") + ax.set_yticks(np.arange(-pi, pi+pi/2, step=(pi/2)), ['-π','-π/2','0','π/2','π']) + + # plot height along phi + ax = fig.add_subplot(3, 3, 4) + ax.scatter(np.arange(np.shape(heights)[0]), heights[:,0], s=20, color='blue') + ax.set_xlabel(r"Step") + ax.set_ylabel(r"$W(\phi) ~[k_{B}T]$") + + # plot height along psi + ax = fig.add_subplot(3, 3, 5) + ax.scatter(np.arange(np.shape(heights)[0]), heights[:,1], s=20, color='red') + ax.set_xlabel(r"Step") + ax.set_ylabel(r"$W(\psi) ~[k_{B}T]$") + + # plot free energy along phi + ax = fig.add_subplot(3, 3, 7) + ax.plot(xi_1D, A_cv1, lw=3, color='blue') + ax.set_xticks(np.arange(-pi, pi+pi/2, step=(pi/2)), ['-π','-π/2','0','π/2','π']) + ax.set_xlabel(r"$\phi$") + ax.set_ylabel(r"$A~[k_{B}T]$") + + # plot free energy along psi + ax = fig.add_subplot(3, 3, 8) + ax.plot(xi_1D, A_cv2, lw=3, color='red') + ax.set_xticks(np.arange(-pi, pi+pi/2, step=(pi/2)), ['-π','-π/2','0','π/2','π']) + ax.set_xlabel(r"$\psi$") + ax.set_ylabel(r"$A~[k_{B}T]$") + + ##### plot 2D free energy along phi-psi + ####ax = fig.add_subplot(3, 3, 9) + ####im = ax.imshow(A, interpolation="bicubic", origin="lower", extent=[-pi, pi, -pi, pi]) + ####ax.contour(A, levels=12, linewidths=0.75, colors="k", extent=[-pi, pi, -pi, pi]) + ####ax.set_xticks(np.arange(-pi, pi+pi/2, step=(pi/2)), ['-π','-π/2','0','π/2','π']) + ####ax.set_yticks(np.arange(-pi, pi+pi/2, step=(pi/2)), ['-π','-π/2','0','π/2','π']) + ####ax.set_xlabel(r"$\phi$") + ####ax.set_ylabel(r"$\psi$") + + ####cbar = plt.colorbar(im) + ####cbar.ax.set_ylabel(r"$A~[k_{B}T]$", rotation=270, labelpad=20) + + fig.savefig("adp-fe.png", dpi=fig.dpi) + + return result # %% diff --git a/pysages/methods/pbmetad.py b/pysages/methods/pbmetad.py index 684f96a2..18bc891d 100644 --- a/pysages/methods/pbmetad.py +++ b/pysages/methods/pbmetad.py @@ -152,9 +152,8 @@ def build_gaussian_accumulator(method: ParallelBiasMetadynamics): if grid is None: evaluate_potential_cv = jit(lambda pstate: get_bias_each_cv(*pstate[:4], periods)) - evaluate_potential = jit(lambda pstate: parallelbias(*pstate[:4], beta, periods)) - - # else: + #evaluate_potential = jit(lambda pstate: parallelbias(*pstate[:4], beta, periods)) + #else: # evaluate_potential = jit(lambda pstate: pstate.grid_potential[pstate.grid_idx]) def next_height(pstate): @@ -170,21 +169,15 @@ def next_height(pstate): update_grids = jit(lambda *args: (None, None)) should_deposit = jit(lambda pred, _: pred) - # else: + #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)) - # if deltaT is None: - # transform = grad - # pack = jit(lambda x: (x,)) - # # No need to accumulate values for the potential (V is None) - # update = jit(lambda V, dV, vals: (V, accum(dV, vals))) - # else: - # transform = value_and_grad - # pack = identity - # update = jit(lambda V, dV, vals, grads: (accum(V, vals), accum(dV, grads))) + # transform = value_and_grad + # pack = identity + # update = jit(lambda V, dV, vals, grads: (accum(V, vals), accum(dV, grads))) # def update_grids(pstate, height, xi, sigma): # # We use `sum_of_gaussians` since it already takes care of the wrapping @@ -272,13 +265,13 @@ def get_bias_each_cv(xi, heights, centers, sigmas, periods): @dispatch def analyze(result: Result[ParallelBiasMetadynamics]): """ - Helper for calculating the free energy from the final state of a `Metadynamics` run. + Helper for calculating the free energy from the final state of a `Parallel Bias Metadynamics` run. Parameters ---------- - result: Result[Metadynamics]: - Result bundle containing method, final metadynamics state, and callback. + result: Result[ParallelBiasMetadynamics]: + Result bundle containing method, final parallel bias metadynamics state, and callback. Returns ------- @@ -287,16 +280,15 @@ def analyze(result: Result[ParallelBiasMetadynamics]): A dictionary with the following keys: heights: JaxArray - Height of the Gaussian bias potential during the simulation. + Height of the Gaussian bias potential along each CV during the simulation. - metapotential: Callable + parallelbias_metapotential: Callable Maps a user-provided array of CV values to the corresponding deposited bias - potential. For standard metadynamics, the free energy along user-provided CV - range is the same as `metapotential(cv)`. - In the case of well-tempered metadynamics, the free energy is equal to - `(T + deltaT) / deltaT * metapotential(cv)`, where `T` is the simulation + 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 - well-tempered metadynamics. + parallel bias metadynamics. """ method = result.method states = result.states @@ -308,21 +300,29 @@ def analyze(result: Result[ParallelBiasMetadynamics]): centers = states[0].centers sigmas = states[0].sigmas - metapotential = jit(vmap(lambda x: sum_of_gaussians(x, heights, centers, sigmas, P))) - - return dict(heights=heights, metapotential=metapotential) + pbmetad_potential_cv = jit(vmap(lambda x: get_bias_each_cv(x, heights, centers, sigmas, P))) + pbmetad_potential = jit(vmap(lambda x, beta: parallelbias(x, heights, centers, sigmas, beta, P), in_axes=(0, None))) + + return dict(centers=centers, heights=heights, pbmetad_potential_cv=pbmetad_potential_cv, pbmetad_potential=pbmetad_potential) # For multiple-replicas runs we return a list heights and functions # (one for each replica) - def build_metapotential(heights, centers, sigmas): - return jit(vmap(lambda x: sum_of_gaussians(x, heights, centers, sigmas, P))) + def build_pbmetapotential_cv(heights, centers, sigmas): + return jit(vmap(lambda x: get_bias_each_cv(x, heights, centers, sigmas, P))) + + def build_pbmetapotential(heights, centers, sigmas): + return jit(vmap(lambda x, beta: parallelbias(x, heights, centers, sigmas, beta, P), in_axes=(0, None))) heights = [] - metapotentials = [] + centers = [] + pbmetad_potentials_cv = [] + pbmetad_potentials = [] for s in states: + centers.append(s.centers) heights.append(s.heights) - metapotentials.append(build_metapotential(s.heights, s.centers, s.sigmas)) + pbmetad_potentials_cv.append(build_pbmetapotential_cv(s.heights, s.centers, s.sigmas)) + pbmetad_potentials.append(build_pbmetapotential(s.heights, s.centers, s.sigmas)) - return dict(heights=heights, metapotential=metapotentials) + return dict(centers=centers, heights=heights, pbmetad_potential_cv=pbmetad_potentials_cv, pbmetad_potential=pbmetad_potentials) From a862abc960be3cbf3ba4cfa31f68eea88f196920 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Wed, 27 Jul 2022 21:20:32 +0000 Subject: [PATCH 10/29] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- examples/openmm/pbmetad/alanine-dipeptide.py | 70 ++++++++++---------- pysages/methods/pbmetad.py | 38 ++++++++--- 2 files changed, 64 insertions(+), 44 deletions(-) diff --git a/examples/openmm/pbmetad/alanine-dipeptide.py b/examples/openmm/pbmetad/alanine-dipeptide.py index fcd6d0ef..791f681b 100644 --- a/examples/openmm/pbmetad/alanine-dipeptide.py +++ b/examples/openmm/pbmetad/alanine-dipeptide.py @@ -127,79 +127,79 @@ def main(argv=[]): 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) plot_2Dgrid = pysages.Grid(lower=(-pi, -pi), upper=(pi, pi), shape=(64, 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)) xi_2D = (compute_mesh(plot_2Dgrid) + 1) / 2 * plot_2Dgrid.size + plot_2Dgrid.lower - + # 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 + 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_potential = result["pbmetad_potential"] # report in kT and set min free energy to zero A_cv = pbmetad_potential_cv(xi_1D_2CVs) * -alpha / kT 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) - - #A = pbmetad_potential(xi_2D, beta) * -alpha / kT - #A = A - A.min() - #A = A.reshape(plot_2Dgrid.shape) - + A_cv1 = A_cv[:, 0].reshape(plot_1Dgrid.shape) + A_cv2 = A_cv[:, 1].reshape(plot_1Dgrid.shape) + + # A = pbmetad_potential(xi_2D, beta) * -alpha / kT + # A = A - A.min() + # A = A.reshape(plot_2Dgrid.shape) + # plot and save free energy along each CV to a PNG file - fig = plt.figure(figsize=(16,12),dpi=120) - fig.subplots_adjust(hspace=0.5,wspace=0.4) + fig = plt.figure(figsize=(16, 12), dpi=120) + fig.subplots_adjust(hspace=0.5, wspace=0.4) # plot centers along phi - ax = fig.add_subplot(3, 3, 1) - ax.scatter(np.arange(np.shape(centers)[0]), centers[:,0], s=20, color='blue') + ax = fig.add_subplot(3, 3, 1) + ax.scatter(np.arange(np.shape(centers)[0]), centers[:, 0], s=20, color="blue") ax.set_xlabel(r"Step") ax.set_ylabel(r"$\phi$") - ax.set_yticks(np.arange(-pi, pi+pi/2, step=(pi/2)), ['-π','-π/2','0','π/2','π']) - + ax.set_yticks(np.arange(-pi, pi + pi / 2, step=(pi / 2)), ["-π", "-π/2", "0", "π/2", "π"]) + # plot centers along psi - ax = fig.add_subplot(3, 3, 2) - ax.scatter(np.arange(np.shape(centers)[0]), centers[:,1], s=20, color='red') + ax = fig.add_subplot(3, 3, 2) + ax.scatter(np.arange(np.shape(centers)[0]), centers[:, 1], s=20, color="red") ax.set_xlabel(r"Step") ax.set_ylabel(r"$\psi$") - ax.set_yticks(np.arange(-pi, pi+pi/2, step=(pi/2)), ['-π','-π/2','0','π/2','π']) - + ax.set_yticks(np.arange(-pi, pi + pi / 2, step=(pi / 2)), ["-π", "-π/2", "0", "π/2", "π"]) + # plot height along phi - ax = fig.add_subplot(3, 3, 4) - ax.scatter(np.arange(np.shape(heights)[0]), heights[:,0], s=20, color='blue') + ax = fig.add_subplot(3, 3, 4) + ax.scatter(np.arange(np.shape(heights)[0]), heights[:, 0], s=20, color="blue") ax.set_xlabel(r"Step") ax.set_ylabel(r"$W(\phi) ~[k_{B}T]$") - + # plot height along psi - ax = fig.add_subplot(3, 3, 5) - ax.scatter(np.arange(np.shape(heights)[0]), heights[:,1], s=20, color='red') + ax = fig.add_subplot(3, 3, 5) + ax.scatter(np.arange(np.shape(heights)[0]), heights[:, 1], s=20, color="red") ax.set_xlabel(r"Step") ax.set_ylabel(r"$W(\psi) ~[k_{B}T]$") - + # plot free energy along phi - ax = fig.add_subplot(3, 3, 7) - ax.plot(xi_1D, A_cv1, lw=3, color='blue') - ax.set_xticks(np.arange(-pi, pi+pi/2, step=(pi/2)), ['-π','-π/2','0','π/2','π']) + ax = fig.add_subplot(3, 3, 7) + ax.plot(xi_1D, A_cv1, lw=3, color="blue") + ax.set_xticks(np.arange(-pi, pi + pi / 2, step=(pi / 2)), ["-π", "-π/2", "0", "π/2", "π"]) ax.set_xlabel(r"$\phi$") ax.set_ylabel(r"$A~[k_{B}T]$") - + # plot free energy along psi - ax = fig.add_subplot(3, 3, 8) - ax.plot(xi_1D, A_cv2, lw=3, color='red') - ax.set_xticks(np.arange(-pi, pi+pi/2, step=(pi/2)), ['-π','-π/2','0','π/2','π']) + ax = fig.add_subplot(3, 3, 8) + ax.plot(xi_1D, A_cv2, lw=3, color="red") + ax.set_xticks(np.arange(-pi, pi + pi / 2, step=(pi / 2)), ["-π", "-π/2", "0", "π/2", "π"]) ax.set_xlabel(r"$\psi$") ax.set_ylabel(r"$A~[k_{B}T]$") diff --git a/pysages/methods/pbmetad.py b/pysages/methods/pbmetad.py index 18bc891d..32faf647 100644 --- a/pysages/methods/pbmetad.py +++ b/pysages/methods/pbmetad.py @@ -152,8 +152,8 @@ def build_gaussian_accumulator(method: ParallelBiasMetadynamics): if grid is None: evaluate_potential_cv = jit(lambda pstate: get_bias_each_cv(*pstate[:4], periods)) - #evaluate_potential = jit(lambda pstate: parallelbias(*pstate[:4], beta, periods)) - #else: + # evaluate_potential = jit(lambda pstate: parallelbias(*pstate[:4], beta, periods)) + # else: # evaluate_potential = jit(lambda pstate: pstate.grid_potential[pstate.grid_idx]) def next_height(pstate): @@ -169,7 +169,7 @@ def next_height(pstate): update_grids = jit(lambda *args: (None, None)) should_deposit = jit(lambda pred, _: pred) - #else: + # else: # grid_mesh = compute_mesh(grid) * (grid.size / 2) # get_grid_index = build_indexer(grid) # # Reshape so the dimensions are compatible @@ -301,18 +301,33 @@ def analyze(result: Result[ParallelBiasMetadynamics]): sigmas = states[0].sigmas pbmetad_potential_cv = jit(vmap(lambda x: get_bias_each_cv(x, heights, centers, sigmas, P))) - pbmetad_potential = jit(vmap(lambda x, beta: parallelbias(x, heights, centers, sigmas, beta, P), in_axes=(0, None))) - - return dict(centers=centers, heights=heights, pbmetad_potential_cv=pbmetad_potential_cv, pbmetad_potential=pbmetad_potential) + pbmetad_potential = jit( + vmap( + lambda x, beta: parallelbias(x, heights, centers, sigmas, beta, P), + in_axes=(0, None), + ) + ) + + return dict( + centers=centers, + heights=heights, + pbmetad_potential_cv=pbmetad_potential_cv, + pbmetad_potential=pbmetad_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: get_bias_each_cv(x, heights, centers, sigmas, P))) - + def build_pbmetapotential(heights, centers, sigmas): - return jit(vmap(lambda x, beta: parallelbias(x, heights, centers, sigmas, beta, P), in_axes=(0, None))) + return jit( + vmap( + lambda x, beta: parallelbias(x, heights, centers, sigmas, beta, P), + in_axes=(0, None), + ) + ) heights = [] centers = [] @@ -325,4 +340,9 @@ def build_pbmetapotential(heights, centers, sigmas): pbmetad_potentials_cv.append(build_pbmetapotential_cv(s.heights, s.centers, s.sigmas)) pbmetad_potentials.append(build_pbmetapotential(s.heights, s.centers, s.sigmas)) - return dict(centers=centers, heights=heights, pbmetad_potential_cv=pbmetad_potentials_cv, pbmetad_potential=pbmetad_potentials) + return dict( + centers=centers, + heights=heights, + pbmetad_potential_cv=pbmetad_potentials_cv, + pbmetad_potential=pbmetad_potentials, + ) From 9210fea8403c9d85e30bd4b6d1542ae0a99226df Mon Sep 17 00:00:00 2001 From: Siva Dasetty Date: Sat, 30 Jul 2022 10:19:47 -0500 Subject: [PATCH 11/29] added support for grids --- examples/openmm/pbmetad/alanine-dipeptide.py | 131 +++++------- pysages/approxfun/core.py | 8 +- pysages/grids.py | 12 +- pysages/methods/pbmetad.py | 199 +++++++++++-------- 4 files changed, 183 insertions(+), 167 deletions(-) diff --git a/examples/openmm/pbmetad/alanine-dipeptide.py b/examples/openmm/pbmetad/alanine-dipeptide.py index 791f681b..e512b86a 100644 --- a/examples/openmm/pbmetad/alanine-dipeptide.py +++ b/examples/openmm/pbmetad/alanine-dipeptide.py @@ -1,9 +1,10 @@ #!/usr/bin/env python3 """ -Parallel bias metadynamics simulation of Alanine Dipeptide in vacuum with OpenMM and PySAGES. +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` +Example command to run the simulation `python3 alanine-dipeptide.py --time_steps 1000` For other supported commandline parameters, check `python3 alanine-dipeptide.py --help` """ @@ -84,10 +85,9 @@ def generate_simulation(pdb_filename=adp_pdb, T=T, dt=dt): # %% def get_args(argv): available_args = [ - ("well-tempered", "w", bool, 0, "Whether to use well-tempered metadynamics"), - ("use-grids", "g", bool, 0, "Whether to use grid acceleration"), + ("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"), + ("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: @@ -98,6 +98,7 @@ def get_args(argv): # %% def main(argv=[]): args = get_args(argv) + print(args) cvs = [DihedralAngle([4, 6, 8, 14]), DihedralAngle([6, 8, 14, 16])] @@ -109,12 +110,13 @@ def main(argv=[]): 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) + 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=None + cvs, height, sigma, stride, ngauss, T.value_in_unit(unit.kelvin), deltaT, kB, grid=grid ) # Logging @@ -127,93 +129,68 @@ def main(argv=[]): 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) - plot_2Dgrid = pysages.Grid(lower=(-pi, -pi), upper=(pi, pi), shape=(64, 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)) - xi_2D = (compute_mesh(plot_2Dgrid) + 1) / 2 * plot_2Dgrid.size + plot_2Dgrid.lower - + # 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 + 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_potential = result["pbmetad_potential"] + pbmetad_net_potential = result["pbmetad_net_potential"] - # report in kT and set min free energy to zero + # 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) - - # A = pbmetad_potential(xi_2D, beta) * -alpha / kT - # A = A - A.min() - # A = A.reshape(plot_2Dgrid.shape) - + 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=(16, 12), dpi=120) - fig.subplots_adjust(hspace=0.5, wspace=0.4) - - # plot centers along phi - ax = fig.add_subplot(3, 3, 1) - ax.scatter(np.arange(np.shape(centers)[0]), centers[:, 0], s=20, color="blue") - ax.set_xlabel(r"Step") - ax.set_ylabel(r"$\phi$") - ax.set_yticks(np.arange(-pi, pi + pi / 2, step=(pi / 2)), ["-π", "-π/2", "0", "π/2", "π"]) - - # plot centers along psi - ax = fig.add_subplot(3, 3, 2) - ax.scatter(np.arange(np.shape(centers)[0]), centers[:, 1], s=20, color="red") - ax.set_xlabel(r"Step") - ax.set_ylabel(r"$\psi$") - ax.set_yticks(np.arange(-pi, pi + pi / 2, step=(pi / 2)), ["-π", "-π/2", "0", "π/2", "π"]) - - # plot height along phi - ax = fig.add_subplot(3, 3, 4) - ax.scatter(np.arange(np.shape(heights)[0]), heights[:, 0], s=20, color="blue") - ax.set_xlabel(r"Step") - ax.set_ylabel(r"$W(\phi) ~[k_{B}T]$") - - # plot height along psi - ax = fig.add_subplot(3, 3, 5) - ax.scatter(np.arange(np.shape(heights)[0]), heights[:, 1], s=20, color="red") - ax.set_xlabel(r"Step") - ax.set_ylabel(r"$W(\psi) ~[k_{B}T]$") - - # plot free energy along phi - ax = fig.add_subplot(3, 3, 7) - ax.plot(xi_1D, A_cv1, lw=3, color="blue") - ax.set_xticks(np.arange(-pi, pi + pi / 2, step=(pi / 2)), ["-π", "-π/2", "0", "π/2", "π"]) - ax.set_xlabel(r"$\phi$") - ax.set_ylabel(r"$A~[k_{B}T]$") - - # plot free energy along psi - ax = fig.add_subplot(3, 3, 8) - ax.plot(xi_1D, A_cv2, lw=3, color="red") - ax.set_xticks(np.arange(-pi, pi + pi / 2, step=(pi / 2)), ["-π", "-π/2", "0", "π/2", "π"]) - ax.set_xlabel(r"$\psi$") - ax.set_ylabel(r"$A~[k_{B}T]$") - - ##### plot 2D free energy along phi-psi - ####ax = fig.add_subplot(3, 3, 9) - ####im = ax.imshow(A, interpolation="bicubic", origin="lower", extent=[-pi, pi, -pi, pi]) - ####ax.contour(A, levels=12, linewidths=0.75, colors="k", extent=[-pi, pi, -pi, pi]) - ####ax.set_xticks(np.arange(-pi, pi+pi/2, step=(pi/2)), ['-π','-π/2','0','π/2','π']) - ####ax.set_yticks(np.arange(-pi, pi+pi/2, step=(pi/2)), ['-π','-π/2','0','π/2','π']) - ####ax.set_xlabel(r"$\phi$") - ####ax.set_ylabel(r"$\psi$") - - ####cbar = plt.colorbar(im) - ####cbar.ax.set_ylabel(r"$A~[k_{B}T]$", rotation=270, labelpad=20) + 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) diff --git a/pysages/approxfun/core.py b/pysages/approxfun/core.py index c4630a47..1565d272 100644 --- a/pysages/approxfun/core.py +++ b/pysages/approxfun/core.py @@ -108,13 +108,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 @@ -122,6 +123,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): @@ -129,7 +131,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 410e3981..55c76893 100644 --- a/pysages/grids.py +++ b/pysages/grids.py @@ -46,13 +46,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_parameter(self) 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 or "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.") @@ -117,7 +119,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) @@ -134,7 +136,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) @@ -152,6 +154,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/pbmetad.py b/pysages/methods/pbmetad.py index 32faf647..6dd4b520 100644 --- a/pysages/methods/pbmetad.py +++ b/pysages/methods/pbmetad.py @@ -22,13 +22,19 @@ class ParallelBiasMetadynamics(GriddedSamplingMethod): """ - Implementation of Parallel Bias Metadynamics as described in + 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 height of Gaussian bias deposited along - each CV is different in parallel bias metadynamics. In addition, the total bias potential - has a slightly different expression involving the sum of exponential of Gaussians (see Eq. 8) - in the paper. + Compared to well-tempered metadynamics, the Gaussian bias deposited along + each CV have different heights in PBMetaD. In addition, the total bias potential + involves the log of sum of exponential of bias potential (see Eq. 8 in the paper) + compared to just sum of Gaussians in well-tempered metadynamics. + + Because the method requires sampling along each CV separately, only the diagonal points + of the grids are required for storing potential along each CV and the net gradient of bias in + PBMetaD. To activate this, the keyword ``parallelbias`` can be used when defining + grids to create the grid centers for each CV separately. Currently, only same number of bins + for each CV is supported. """ snapshot_flags = {"positions", "indices"} @@ -41,10 +47,10 @@ def __init__(self, cvs, height, sigma, stride, ngaussians, T, deltaT, kB, *args, cvs: Set of user selected collective variables. - height: + height: JaxArray Initial height of the deposited Gaussians along each CV. - sigma: + sigma: JaxArray Initial standard deviation of the to-be-deposit Gaussians along each CV. stride: int @@ -53,7 +59,7 @@ def __init__(self, cvs, height, sigma, stride, ngaussians, T, deltaT, kB, *args, ngaussians: int Total number of expected gaussians (timesteps // stride + 1). - deltaT: float = None + deltaT: float Well-tempered metadynamics $\\Delta T$ parameter to set the energy scale for sampling. @@ -65,8 +71,8 @@ def __init__(self, cvs, height, sigma, stride, ngaussians, T, deltaT, kB, *args, grid: Optional[Grid] = None If provided, it will be used to accelerate the computation by - approximating the bias potential and its gradient over its centers. - + 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) @@ -85,7 +91,7 @@ def build(self, snapshot, helpers, *args, **kwargs): def _parallelbiasmetadynamics(method, snapshot, helpers): - # Initialization and update of biasing forces. Interface expected for methods. + # Initialization and update of biasing forces. Interface as expected for methods. cv = method.cv stride = method.stride ngaussians = method.ngaussians @@ -108,8 +114,12 @@ def initialize(): grid_potential = grid_gradient = None else: shape = method.grid.shape - grid_potential = np.zeros((*shape,), dtype=np.float64) - grid_gradient = np.zeros((*shape, shape.size), dtype=np.float64) + # NOTE: for now, we assume, number of bins defined by shape along each CV are same. + # This need not be the case for PBMetaD as it generate free energy along each CV separately. + # PySAGES will throw an concatenation error if bins or shape of each CV is different. + # So, we use shape[0] to define the size of grids as all bins are expected to be same. + 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 @@ -151,44 +161,50 @@ def build_gaussian_accumulator(method: ParallelBiasMetadynamics): kB_deltaT = kB * deltaT if grid is None: - evaluate_potential_cv = jit(lambda pstate: get_bias_each_cv(*pstate[:4], periods)) - # evaluate_potential = jit(lambda pstate: parallelbias(*pstate[:4], beta, periods)) - # else: - # evaluate_potential = jit(lambda pstate: pstate.grid_potential[pstate.grid_idx]) + evaluate_potential_each_cv = jit(lambda pstate: parallelbias_each_cv(*pstate[:4], periods)) + else: + # each index in pstate.grid_idx correpsonds to different CV. + # so, we extract it using np.choose + 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_cv(pstate) + V = evaluate_potential_each_cv(pstate) w = height_0 * np.exp(-V / kB_deltaT) - switching_probability_sum = np.sum(np.exp(-beta * V)) - switching_probability = np.exp(-beta * V) / switching_probability_sum + cv_switching_probability_sum = np.sum(np.exp(-beta * V)) + cv_switching_probability = np.exp(-beta * V) / cv_switching_probability_sum - return w * switching_probability + 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)) - - # transform = value_and_grad - # pack = identity - # update = jit(lambda V, dV, vals, grads: (accum(V, vals), accum(dV, grads))) - - # def update_grids(pstate, height, xi, sigma): - # # We use `sum_of_gaussians` since it already takes care of the wrapping - # current_gaussian = jit(lambda x: sum_of_gaussians(x, height, xi, sigma, periods)) - # # Evaluate gradient of bias (and bias potential for WT version) - # grid_values = pack(vmap(transform(current_gaussian))(grid_mesh)) - # return update(pstate.grid_potential, pstate.grid_gradient, *grid_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 + 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)) + transform = grad + 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. Total bias potential is not required. + current_parallelbias_each_cv = jit(lambda x: parallelbias_each_cv_grids(x, height, xi, sigma, periods)) + + # We need bias gradient obtained by grad of total bias potential along each CV. + current_parallelbias = jit(lambda x: parallelbias_grids(x, height, xi, sigma, beta, periods)) + + grid_potential_values = vmap(current_parallelbias_each_cv)(grid_mesh) + grid_grad_values= vmap(transform(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) == np.array([int(grid.shape[0]), grid.shape.size]))) + return in_deposition_step & in_bounds + def deposit_gaussian(pstate): xi, idx = pstate.xi, pstate.idx @@ -229,10 +245,12 @@ def zero_force(_): return np.zeros(grid.shape.size) def get_force(pstate): - return pstate.grid_gradient[pstate.grid_idx] + # 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 + ob = np.any(np.array(pstate.grid_idx) == np.array([int(grid.shape[0]), grid.shape.size])) # out of bounds return cond(ob, zero_force, get_force, pstate) return evaluate_bias_grad @@ -244,14 +262,14 @@ def parallelbias(xi, heights, centers, sigmas, beta, periods): 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 = get_bias_each_cv(xi, heights, centers, sigmas, periods) + bias_each_cv = parallelbias_each_cv(xi, heights, centers, sigmas, periods) exp_sum_gaussian = np.exp(-beta * bias_each_cv) return -(1 / beta) * np.log(np.sum(exp_sum_gaussian)) # Helper function to evaluate parallel bias potential along each CV -def get_bias_each_cv(xi, heights, centers, sigmas, periods): +def parallelbias_each_cv(xi, heights, centers, sigmas, periods): """ Evaluate parallel bias potential along each CV. """ @@ -261,6 +279,29 @@ def get_bias_each_cv(xi, heights, centers, sigmas, periods): return bias_each_cv +# Helper function to evaluate parallel bias potential +def parallelbias_grids(xi, heights, centers, sigmas, beta, periods): + """ + 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_grids(xi, heights, centers, sigmas, periods) + exp_sum_gaussian = np.exp(-beta * bias_each_cv) + + return -(1 / beta) * np.log(np.sum(exp_sum_gaussian)) + + +# Helper function to evaluate parallel bias potential along each CV +def parallelbias_each_cv_grids(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 + @dispatch def analyze(result: Result[ParallelBiasMetadynamics]): @@ -278,17 +319,26 @@ def analyze(result: Result[ParallelBiasMetadynamics]): 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. - parallelbias_metapotential: 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 + pbmetad_potential_cv: Callable + Maps a user-provided array of CV values and step 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 @@ -300,49 +350,34 @@ def analyze(result: Result[ParallelBiasMetadynamics]): centers = states[0].centers sigmas = states[0].sigmas - pbmetad_potential_cv = jit(vmap(lambda x: get_bias_each_cv(x, heights, centers, sigmas, P))) - pbmetad_potential = jit( - vmap( - lambda x, beta: parallelbias(x, heights, centers, sigmas, beta, P), - in_axes=(0, None), - ) - ) - - return dict( - centers=centers, - heights=heights, - pbmetad_potential_cv=pbmetad_potential_cv, - pbmetad_potential=pbmetad_potential, - ) + pbmetad_potential_cv = jit(vmap(lambda x: parallelbias_each_cv(x, heights, centers, sigmas, P))) + pbmetad_net_potential = jit(vmap(lambda x, beta: parallelbias(x, heights, centers, sigmas, beta, P), + 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: get_bias_each_cv(x, heights, centers, sigmas, P))) - + return jit(vmap(lambda x: parallelbias_each_cv(x, heights, centers, sigmas, P))) + def build_pbmetapotential(heights, centers, sigmas): - return jit( - vmap( - lambda x, beta: parallelbias(x, heights, centers, sigmas, beta, P), - in_axes=(0, None), - ) - ) + return jit(vmap(lambda x, beta: parallelbias(x, heights, centers, sigmas, beta, P), in_axes=(0, None))) heights = [] centers = [] pbmetad_potentials_cv = [] - pbmetad_potentials = [] + 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_potentials.append(build_pbmetapotential(s.heights, s.centers, s.sigmas)) - - return dict( - centers=centers, - heights=heights, - pbmetad_potential_cv=pbmetad_potentials_cv, - pbmetad_potential=pbmetad_potentials, - ) + 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) From 51b221f0cffe2e9aa3f1110661e50af91ef61eed Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Sat, 30 Jul 2022 15:20:43 +0000 Subject: [PATCH 12/29] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- examples/openmm/pbmetad/alanine-dipeptide.py | 58 +++++++------ pysages/approxfun/core.py | 2 +- pysages/methods/pbmetad.py | 88 +++++++++++++------- 3 files changed, 90 insertions(+), 58 deletions(-) diff --git a/examples/openmm/pbmetad/alanine-dipeptide.py b/examples/openmm/pbmetad/alanine-dipeptide.py index e512b86a..e54fb999 100644 --- a/examples/openmm/pbmetad/alanine-dipeptide.py +++ b/examples/openmm/pbmetad/alanine-dipeptide.py @@ -110,8 +110,10 @@ def main(argv=[]): 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 = pysages.Grid( + lower=(-pi, -pi), upper=(pi, pi), shape=(50, 50), periodic=True, parallelbias=True + ) + grid = grid if args.use_grids else None # Method @@ -129,23 +131,23 @@ def main(argv=[]): 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 + 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"] @@ -153,42 +155,46 @@ def main(argv=[]): 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) - + 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 + 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' + 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.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','π']) - + 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' + 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.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' + 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_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]$") diff --git a/pysages/approxfun/core.py b/pysages/approxfun/core.py index 1565d272..6bb13933 100644 --- a/pysages/approxfun/core.py +++ b/pysages/approxfun/core.py @@ -114,7 +114,7 @@ def compute_mesh(grid: Grid): o = -1 + h / 2 nodes = o + h * np.hstack([np.arange(i).reshape(-1, 1) for i in grid.shape]) - + return nodes if grid.parallelbias else _compute_mesh(nodes) diff --git a/pysages/methods/pbmetad.py b/pysages/methods/pbmetad.py index 6dd4b520..767cef0f 100644 --- a/pysages/methods/pbmetad.py +++ b/pysages/methods/pbmetad.py @@ -27,9 +27,9 @@ class ParallelBiasMetadynamics(GriddedSamplingMethod): Compared to well-tempered metadynamics, the Gaussian bias deposited along each CV have different heights in PBMetaD. In addition, the total bias potential - involves the log of sum of exponential of bias potential (see Eq. 8 in the paper) + involves the log of sum of exponential of bias potential (see Eq. 8 in the paper) compared to just sum of Gaussians in well-tempered metadynamics. - + Because the method requires sampling along each CV separately, only the diagonal points of the grids are required for storing potential along each CV and the net gradient of bias in PBMetaD. To activate this, the keyword ``parallelbias`` can be used when defining @@ -165,8 +165,9 @@ def build_gaussian_accumulator(method: ParallelBiasMetadynamics): else: # each index in pstate.grid_idx correpsonds to different CV. # so, we extract it using np.choose - evaluate_potential_each_cv = jit(lambda pstate: np.choose(np.array(pstate.grid_idx), - pstate.grid_potential, mode='clip')) + 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) @@ -191,20 +192,25 @@ def next_height(pstate): def update_grids(pstate, height, xi, sigma): # We need bias potential along each CV to update the heights. Total bias potential is not required. - current_parallelbias_each_cv = jit(lambda x: parallelbias_each_cv_grids(x, height, xi, sigma, periods)) + current_parallelbias_each_cv = jit( + lambda x: parallelbias_each_cv_grids(x, height, xi, sigma, periods) + ) # We need bias gradient obtained by grad of total bias potential along each CV. - current_parallelbias = jit(lambda x: parallelbias_grids(x, height, xi, sigma, beta, periods)) - + current_parallelbias = jit( + lambda x: parallelbias_grids(x, height, xi, sigma, beta, periods) + ) + grid_potential_values = vmap(current_parallelbias_each_cv)(grid_mesh) - grid_grad_values= vmap(transform(current_parallelbias))(grid_mesh) - - return update(pstate.grid_potential, pstate.grid_gradient, grid_potential_values, grid_grad_values) - + grid_grad_values = vmap(transform(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) == np.array([int(grid.shape[0]), grid.shape.size]))) return in_deposition_step & in_bounds - def deposit_gaussian(pstate): xi, idx = pstate.xi, pstate.idx @@ -247,10 +253,12 @@ def zero_force(_): 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') + 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) == np.array([int(grid.shape[0]), grid.shape.size])) # out of bounds + ob = np.any( + np.array(pstate.grid_idx) == np.array([int(grid.shape[0]), grid.shape.size]) + ) # out of bounds return cond(ob, zero_force, get_force, pstate) return evaluate_bias_grad @@ -279,6 +287,7 @@ def parallelbias_each_cv(xi, heights, centers, sigmas, periods): return bias_each_cv + # Helper function to evaluate parallel bias potential def parallelbias_grids(xi, heights, centers, sigmas, beta, periods): """ @@ -319,7 +328,7 @@ def analyze(result: Result[ParallelBiasMetadynamics]): dict: A dictionary with the following keys: - + centers: JaxArray Centers of the CVs used for depositing Gaussian bias potential during the simulation. @@ -329,12 +338,12 @@ def analyze(result: Result[ParallelBiasMetadynamics]): pbmetad_potential_cv: Callable Maps a user-provided array of CV values and step 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)`, + + 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 @@ -350,22 +359,36 @@ def analyze(result: Result[ParallelBiasMetadynamics]): centers = states[0].centers sigmas = states[0].sigmas - pbmetad_potential_cv = jit(vmap(lambda x: parallelbias_each_cv(x, heights, centers, sigmas, P))) - pbmetad_net_potential = jit(vmap(lambda x, beta: parallelbias(x, heights, centers, sigmas, beta, P), - in_axes=(0, None))) - - return dict(centers=centers, heights=heights, - pbmetad_potential_cv=pbmetad_potential_cv, - pbmetad_net_potential=pbmetad_net_potential) + pbmetad_potential_cv = jit( + vmap(lambda x: parallelbias_each_cv(x, heights, centers, sigmas, P)) + ) + pbmetad_net_potential = jit( + vmap( + lambda x, beta: parallelbias(x, heights, centers, sigmas, beta, P), + 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: parallelbias_each_cv(x, heights, centers, sigmas, P))) - + def build_pbmetapotential(heights, centers, sigmas): - return jit(vmap(lambda x, beta: parallelbias(x, heights, centers, sigmas, beta, P), in_axes=(0, None))) + return jit( + vmap( + lambda x, beta: parallelbias(x, heights, centers, sigmas, beta, P), + in_axes=(0, None), + ) + ) heights = [] centers = [] @@ -378,6 +401,9 @@ def build_pbmetapotential(heights, centers, sigmas): 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) + return dict( + centers=centers, + heights=heights, + pbmetad_potential_cv=pbmetad_potentials_cv, + pbmetad_net_potential=pbmetad_net_potentials, + ) From f05a812db5a36625052fd9c9faa2f88213d85d21 Mon Sep 17 00:00:00 2001 From: Siva Dasetty Date: Sat, 30 Jul 2022 10:36:40 -0500 Subject: [PATCH 13/29] updated test_pickle to include pbmetad --- tests/test_pickle.py | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/tests/test_pickle.py b/tests/test_pickle.py index abff2fb8..65e4c097 100644 --- a/tests/test_pickle.py +++ b/tests/test_pickle.py @@ -55,6 +55,17 @@ "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), + }, "SpectralABF": { "cvs": [pysages.colvars.Component([0], 0), pysages.colvars.Component([0], 1)], "grid": pysages.Grid(lower=(1, 1), upper=(5, 5), shape=(32, 32)), From 4f3b44fd129e2d4232caf2d3064e7e3926335af9 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Sat, 30 Jul 2022 15:37:14 +0000 Subject: [PATCH 14/29] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- tests/test_pickle.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_pickle.py b/tests/test_pickle.py index 65e4c097..abc2514b 100644 --- a/tests/test_pickle.py +++ b/tests/test_pickle.py @@ -61,7 +61,7 @@ "sigma": 5.0, "stride": 7.0, "ngaussians": 5, - "T":300, + "T": 300, "deltaT": 0.1, "kB": 614.0, "grid": pysages.Grid(lower=(-pi), upper=(pi), shape=(32), periodic=True), From 51b3c4745b5cd9fa8204ede8918e9ac147302909 Mon Sep 17 00:00:00 2001 From: Siva Dasetty Date: Sat, 30 Jul 2022 10:47:49 -0500 Subject: [PATCH 15/29] updated grids for pbmetad --- pysages/grids.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pysages/grids.py b/pysages/grids.py index 55c76893..1b0ac26d 100644 --- a/pysages/grids.py +++ b/pysages/grids.py @@ -52,7 +52,7 @@ def __check_init_invariants__(self, **kwargs): T = type_parameter(self) if not (issubclass(type(T), type) and issubclass(T, GridType)): raise TypeError("Type parameter must be a subclass of GridType.") - if len(kwargs) > 1: + if len(kwargs) >= 1: if "periodic" not in kwargs or "parallelbias" not in kwargs: raise ValueError("Invalid keyword arguments") periodic = kwargs.get("periodic", T is Periodic) From a1c433313fbeab780f74edb28a28017b63a5ba1a Mon Sep 17 00:00:00 2001 From: Siva Dasetty Date: Sat, 30 Jul 2022 10:52:41 -0500 Subject: [PATCH 16/29] updated for pbmetad --- tests/test_pickle.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_pickle.py b/tests/test_pickle.py index abc2514b..59231597 100644 --- a/tests/test_pickle.py +++ b/tests/test_pickle.py @@ -64,7 +64,7 @@ "T": 300, "deltaT": 0.1, "kB": 614.0, - "grid": pysages.Grid(lower=(-pi), upper=(pi), shape=(32), periodic=True), + "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)], From f497b8122e58779c7e4a25cc38a3adf7e5a09a37 Mon Sep 17 00:00:00 2001 From: Siva Dasetty Date: Sat, 30 Jul 2022 11:43:53 -0500 Subject: [PATCH 17/29] fixed grids kwargs --- pysages/grids.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pysages/grids.py b/pysages/grids.py index 1b0ac26d..41fedf35 100644 --- a/pysages/grids.py +++ b/pysages/grids.py @@ -53,7 +53,7 @@ def __check_init_invariants__(self, **kwargs): if not (issubclass(type(T), type) and issubclass(T, GridType)): raise TypeError("Type parameter must be a subclass of GridType.") if len(kwargs) >= 1: - if "periodic" not in kwargs or "parallelbias" not in kwargs: + 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: From b26d81cb4b9904df96bb0659734441d828a9809c Mon Sep 17 00:00:00 2001 From: Siva Dasetty Date: Sat, 30 Jul 2022 12:13:48 -0500 Subject: [PATCH 18/29] added pbmetad to ci --- .github/workflows/ci.yml | 2 ++ 1 file changed, 2 insertions(+) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index a7f4dfb9..3e8294ca 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -66,6 +66,7 @@ jobs: pylint pysages/methods/__init__.py pylint pysages/methods/abf.py pylint pysages/methods/metad.py + pylint pysages/methods/pbmetad.py pylint pysages/methods/utils.py pylint pysages/methods/harmonic_bias.py pylint pysages/methods/umbrella_integration.py @@ -73,6 +74,7 @@ jobs: flake8 pysages/methods/__init__.py flake8 pysages/methods/abf.py flake8 pysages/methods/metad.py + flake8 pysages/methods/pbmetad.py flake8 pysages/methods/utils.py flake8 pysages/methods/harmonic_bias.py flake8 pysages/methods/umbrella_integration.py From 77af9a4ea82ab31b25e1a9e43664352b3e28ffee Mon Sep 17 00:00:00 2001 From: Siva Dasetty Date: Sat, 30 Jul 2022 12:17:57 -0500 Subject: [PATCH 19/29] added pbmetad to docker-ci --- .github/workflows/docker-ci.yml | 16 ++++++++++++++++ 1 file changed, 16 insertions(+) diff --git a/.github/workflows/docker-ci.yml b/.github/workflows/docker-ci.yml index 45ab97ff..441c6e73 100644 --- a/.github/workflows/docker-ci.yml +++ b/.github/workflows/docker-ci.yml @@ -122,6 +122,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 From 86b326ed970be610b8d3d3ab9caca59505bc80f4 Mon Sep 17 00:00:00 2001 From: Siva Dasetty Date: Sat, 30 Jul 2022 12:27:02 -0500 Subject: [PATCH 20/29] removed unused imports and vars --- pysages/methods/pbmetad.py | 8 +++----- 1 file changed, 3 insertions(+), 5 deletions(-) diff --git a/pysages/methods/pbmetad.py b/pysages/methods/pbmetad.py index 767cef0f..2ea47534 100644 --- a/pysages/methods/pbmetad.py +++ b/pysages/methods/pbmetad.py @@ -6,15 +6,13 @@ Implementation of Parallel Bias Well-tempered Metadynamics with optional support for grids. """ -from typing import NamedTuple, Optional - -from jax import numpy as np, grad, jit, value_and_grad, vmap +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 JaxArray, identity +from pysages.utils import identity from pysages.grids import build_indexer from pysages.utils import dispatch from pysages.methods.metad import MetadynamicsState, PartialMetadynamicsState @@ -39,7 +37,7 @@ class ParallelBiasMetadynamics(GriddedSamplingMethod): snapshot_flags = {"positions", "indices"} - def __init__(self, cvs, height, sigma, stride, ngaussians, T, deltaT, kB, *args, **kwargs): + def __init__(self, cvs, height, sigma, stride, ngaussians, T, deltaT, kB, **kwargs): """ Arguments --------- From 4106b1e5573ff5e9125a6606c83e0f73e54706d1 Mon Sep 17 00:00:00 2001 From: Siva Dasetty Date: Sat, 30 Jul 2022 12:37:15 -0500 Subject: [PATCH 21/29] fixed len of lines --- pysages/methods/pbmetad.py | 33 +++++++++++++++++++-------------- 1 file changed, 19 insertions(+), 14 deletions(-) diff --git a/pysages/methods/pbmetad.py b/pysages/methods/pbmetad.py index 2ea47534..9ab72b62 100644 --- a/pysages/methods/pbmetad.py +++ b/pysages/methods/pbmetad.py @@ -28,11 +28,11 @@ class ParallelBiasMetadynamics(GriddedSamplingMethod): involves the log of sum of exponential of bias potential (see Eq. 8 in the paper) compared to just sum of Gaussians in well-tempered metadynamics. - Because the method requires sampling along each CV separately, only the diagonal points - of the grids are required for storing potential along each CV and the net gradient of bias in - PBMetaD. To activate this, the keyword ``parallelbias`` can be used when defining - grids to create the grid centers for each CV separately. Currently, only same number of bins - for each CV is supported. + 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 implementing 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"} @@ -113,7 +113,8 @@ def initialize(): else: shape = method.grid.shape # NOTE: for now, we assume, number of bins defined by shape along each CV are same. - # This need not be the case for PBMetaD as it generate free energy along each CV separately. + # This need not be the case for PBMetaD as it generates free energy along each + # CV separately. # PySAGES will throw an concatenation error if bins or shape of each CV is different. # So, we use shape[0] to define the size of grids as all bins are expected to be same. grid_potential = np.zeros((shape[0], shape.size), dtype=np.float64) @@ -189,7 +190,8 @@ def next_height(pstate): 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. Total bias potential is not required. + # We need bias potential along each CV to update the heights. + # Total bias potential is required only for storing gradient of bias. current_parallelbias_each_cv = jit( lambda x: parallelbias_each_cv_grids(x, height, xi, sigma, periods) ) @@ -313,7 +315,8 @@ def parallelbias_each_cv_grids(xi, heights, centers, sigmas, periods): @dispatch def analyze(result: Result[ParallelBiasMetadynamics]): """ - Helper for calculating the free energy from the final state of a `Parallel Bias Metadynamics` run. + Helper for calculating the free energy from the final state of a + `Parallel Bias Metadynamics` run. Parameters ---------- @@ -337,15 +340,17 @@ def analyze(result: Result[ParallelBiasMetadynamics]): Maps a user-provided array of CV values and step 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. + 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. + 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 From c67ce3e5e57884983f6af2b1ace40d2c42a5095d Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Sat, 30 Jul 2022 17:37:38 +0000 Subject: [PATCH 22/29] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- pysages/methods/pbmetad.py | 24 ++++++++++++------------ 1 file changed, 12 insertions(+), 12 deletions(-) diff --git a/pysages/methods/pbmetad.py b/pysages/methods/pbmetad.py index 9ab72b62..b9369113 100644 --- a/pysages/methods/pbmetad.py +++ b/pysages/methods/pbmetad.py @@ -28,10 +28,10 @@ class ParallelBiasMetadynamics(GriddedSamplingMethod): involves the log of sum of exponential of bias potential (see Eq. 8 in the 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 implementing this, the keyword - ``parallelbias`` is added to define grids for each CV separately. Currently, only + 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 implementing 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. """ @@ -113,7 +113,7 @@ def initialize(): else: shape = method.grid.shape # NOTE: for now, we assume, number of bins defined by shape along each CV are same. - # This need not be the case for PBMetaD as it generates free energy along each + # This need not be the case for PBMetaD as it generates free energy along each # CV separately. # PySAGES will throw an concatenation error if bins or shape of each CV is different. # So, we use shape[0] to define the size of grids as all bins are expected to be same. @@ -190,7 +190,7 @@ def next_height(pstate): 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. + # We need bias potential along each CV to update the heights. # Total bias potential is required only for storing gradient of bias. current_parallelbias_each_cv = jit( lambda x: parallelbias_each_cv_grids(x, height, xi, sigma, periods) @@ -315,7 +315,7 @@ def parallelbias_each_cv_grids(xi, heights, centers, sigmas, periods): @dispatch def analyze(result: Result[ParallelBiasMetadynamics]): """ - Helper for calculating the free energy from the final state of a + Helper for calculating the free energy from the final state of a `Parallel Bias Metadynamics` run. Parameters @@ -340,16 +340,16 @@ def analyze(result: Result[ParallelBiasMetadynamics]): Maps a user-provided array of CV values and step 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 + 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 + 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, + 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 From 09699b32271668cd871dcc55fbf4041a84249e81 Mon Sep 17 00:00:00 2001 From: Siva Dasetty Date: Sat, 30 Jul 2022 13:41:52 -0500 Subject: [PATCH 23/29] final fixes --- pysages/methods/pbmetad.py | 129 ++++++++++++++++--------------------- 1 file changed, 56 insertions(+), 73 deletions(-) diff --git a/pysages/methods/pbmetad.py b/pysages/methods/pbmetad.py index b9369113..91089e1c 100644 --- a/pysages/methods/pbmetad.py +++ b/pysages/methods/pbmetad.py @@ -25,13 +25,13 @@ class ParallelBiasMetadynamics(GriddedSamplingMethod): Compared to well-tempered metadynamics, the Gaussian bias deposited along each CV have different heights in PBMetaD. In addition, the total bias potential - involves the log of sum of exponential of bias potential (see Eq. 8 in the paper) + involves the 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 implementing this, the keyword - ``parallelbias`` is added to define grids for each CV separately. Currently, only + 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 implementing 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. """ @@ -112,9 +112,11 @@ def initialize(): grid_potential = grid_gradient = None else: shape = method.grid.shape - # NOTE: for now, we assume, number of 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. + # 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 efficient. + # PySAGES will throw an concatenation error if bins or shape of each CV is different. # So, we use shape[0] to define the size of grids as all bins are expected to be same. grid_potential = np.zeros((shape[0], shape.size), dtype=np.float64) @@ -160,13 +162,20 @@ def build_gaussian_accumulator(method: ParallelBiasMetadynamics): kB_deltaT = kB * deltaT if grid is None: - evaluate_potential_each_cv = jit(lambda pstate: parallelbias_each_cv(*pstate[:4], periods)) + 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 - evaluate_potential_each_cv = jit( - lambda pstate: np.choose(np.array(pstate.grid_idx), pstate.grid_potential, mode="clip") - ) + # 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) @@ -190,15 +199,14 @@ def next_height(pstate): 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. - # Total bias potential is required only for storing gradient of bias. + # We need bias potential along each CV to update the heights. current_parallelbias_each_cv = jit( - lambda x: parallelbias_each_cv_grids(x, height, xi, sigma, periods) + lambda x: parallelbias_each_cv(x, height, xi, sigma, periods) ) - # We need bias gradient obtained by grad of total bias potential along each CV. + # Total bias potential is required only for calculating and storing gradient of bias. current_parallelbias = jit( - lambda x: parallelbias_grids(x, height, xi, sigma, beta, periods) + lambda x: parallelbias(x, height, xi, sigma, beta, periods, grid) ) grid_potential_values = vmap(current_parallelbias_each_cv)(grid_mesh) @@ -209,7 +217,7 @@ def update_grids(pstate, height, xi, sigma): ) def should_deposit(in_deposition_step, I_xi): - in_bounds = ~(np.any(np.array(I_xi) == np.array([int(grid.shape[0]), grid.shape.size]))) + in_bounds = ~(np.any(np.array(I_xi) == grid.shape)) return in_deposition_step & in_bounds def deposit_gaussian(pstate): @@ -244,7 +252,8 @@ def build_bias_grad_evaluator(method: ParallelBiasMetadynamics): if grid is None: periods = get_periods(method.cvs) - evaluate_bias_grad = jit(lambda pstate: grad(parallelbias)(*pstate[:4], beta, periods)) + evaluate_bias_grad = jit(lambda pstate: grad(parallelbias)(*pstate[:4], beta, + periods, grid)) else: def zero_force(_): @@ -256,26 +265,12 @@ def get_force(pstate): 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) == np.array([int(grid.shape[0]), grid.shape.size]) - ) # out of bounds + 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 -def parallelbias(xi, heights, centers, sigmas, beta, periods): - """ - 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) - exp_sum_gaussian = np.exp(-beta * bias_each_cv) - - return -(1 / beta) * np.log(np.sum(exp_sum_gaussian)) - - # Helper function to evaluate parallel bias potential along each CV def parallelbias_each_cv(xi, heights, centers, sigmas, periods): """ @@ -283,39 +278,27 @@ def parallelbias_each_cv(xi, heights, centers, sigmas, periods): """ delta_xi_each_cv = wrap(xi - centers, periods) gaussian_each_cv = heights * np.exp(-((delta_xi_each_cv / sigmas) ** 2) / 2) - bias_each_cv = np.sum(gaussian_each_cv, axis=0) + bias_each_cv = gaussian_each_cv return bias_each_cv - # Helper function to evaluate parallel bias potential -def parallelbias_grids(xi, heights, centers, sigmas, beta, periods): +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_grids(xi, heights, centers, sigmas, periods) + 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)) -# Helper function to evaluate parallel bias potential along each CV -def parallelbias_each_cv_grids(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 - - @dispatch def analyze(result: Result[ParallelBiasMetadynamics]): """ - Helper for calculating the free energy from the final state of a + Helper for calculating the free energy from the final state of a `Parallel Bias Metadynamics` run. Parameters @@ -340,37 +323,38 @@ def analyze(result: Result[ParallelBiasMetadynamics]): Maps a user-provided array of CV values and step 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 + 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 + 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, + 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: parallelbias_each_cv(x, heights, centers, sigmas, P)) - ) - pbmetad_net_potential = jit( - vmap( - lambda x, beta: parallelbias(x, heights, centers, sigmas, beta, P), - in_axes=(0, None), - ) - ) + 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, @@ -383,15 +367,14 @@ def analyze(result: Result[ParallelBiasMetadynamics]): # (one for each replica) def build_pbmetapotential_cv(heights, centers, sigmas): - return jit(vmap(lambda x: parallelbias_each_cv(x, heights, centers, sigmas, P))) + 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), - in_axes=(0, None), - ) - ) + return jit(vmap(lambda x, beta: parallelbias(x, heights, centers, sigmas, + beta, P, grid), in_axes=(0, None)) + ) heights = [] centers = [] From 389c9006a9997551b92360d5a810755d9fccf959 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Sat, 30 Jul 2022 18:42:21 +0000 Subject: [PATCH 24/29] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- pysages/methods/pbmetad.py | 84 +++++++++++++++++++------------------- 1 file changed, 42 insertions(+), 42 deletions(-) diff --git a/pysages/methods/pbmetad.py b/pysages/methods/pbmetad.py index 91089e1c..a6356c9c 100644 --- a/pysages/methods/pbmetad.py +++ b/pysages/methods/pbmetad.py @@ -28,10 +28,10 @@ class ParallelBiasMetadynamics(GriddedSamplingMethod): involves the 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 implementing this, the keyword - ``parallelbias`` is added to define grids for each CV separately. Currently, only + 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 implementing 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. """ @@ -112,11 +112,11 @@ def initialize(): 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 + # 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 efficient. - + # PySAGES will throw an concatenation error if bins or shape of each CV is different. # So, we use shape[0] to define the size of grids as all bins are expected to be same. grid_potential = np.zeros((shape[0], shape.size), dtype=np.float64) @@ -162,20 +162,15 @@ def build_gaussian_accumulator(method: ParallelBiasMetadynamics): 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) - ) + 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") - ) + 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) @@ -199,7 +194,7 @@ def next_height(pstate): 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. + # 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) ) @@ -252,8 +247,9 @@ def build_bias_grad_evaluator(method: ParallelBiasMetadynamics): if grid is None: periods = get_periods(method.cvs) - evaluate_bias_grad = jit(lambda pstate: grad(parallelbias)(*pstate[:4], beta, - periods, grid)) + evaluate_bias_grad = jit( + lambda pstate: grad(parallelbias)(*pstate[:4], beta, periods, grid) + ) else: def zero_force(_): @@ -282,6 +278,7 @@ def parallelbias_each_cv(xi, heights, centers, sigmas, periods): return bias_each_cv + # Helper function to evaluate parallel bias potential def parallelbias(xi, heights, centers, sigmas, beta, periods, grid): """ @@ -298,7 +295,7 @@ def parallelbias(xi, heights, centers, sigmas, beta, periods, grid): @dispatch def analyze(result: Result[ParallelBiasMetadynamics]): """ - Helper for calculating the free energy from the final state of a + Helper for calculating the free energy from the final state of a `Parallel Bias Metadynamics` run. Parameters @@ -323,16 +320,16 @@ def analyze(result: Result[ParallelBiasMetadynamics]): Maps a user-provided array of CV values and step 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 + 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 + 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, + 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 @@ -346,15 +343,15 @@ def analyze(result: Result[ParallelBiasMetadynamics]): 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)) - ) + 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, @@ -367,14 +364,17 @@ def analyze(result: Result[ParallelBiasMetadynamics]): # (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) )) + 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)) - ) + return jit( + vmap( + lambda x, beta: parallelbias(x, heights, centers, sigmas, beta, P, grid), + in_axes=(0, None), + ) + ) heights = [] centers = [] From e15f77de82da08fbb7bbc53813f90ea930931b70 Mon Sep 17 00:00:00 2001 From: Siva Dasetty Date: Sat, 30 Jul 2022 13:58:32 -0500 Subject: [PATCH 25/29] minor fmt edits --- examples/openmm/pbmetad/alanine-dipeptide.py | 4 +- pysages/methods/pbmetad.py | 85 ++++++++++---------- 2 files changed, 45 insertions(+), 44 deletions(-) diff --git a/examples/openmm/pbmetad/alanine-dipeptide.py b/examples/openmm/pbmetad/alanine-dipeptide.py index e54fb999..b2e881eb 100644 --- a/examples/openmm/pbmetad/alanine-dipeptide.py +++ b/examples/openmm/pbmetad/alanine-dipeptide.py @@ -85,9 +85,9 @@ def generate_simulation(pdb_filename=adp_pdb, T=T, dt=dt): # %% def get_args(argv): available_args = [ - ("use_grids", "g", bool, 0, "Whether to use grid acceleration"), + ("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"), + ("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: diff --git a/pysages/methods/pbmetad.py b/pysages/methods/pbmetad.py index a6356c9c..aae7f118 100644 --- a/pysages/methods/pbmetad.py +++ b/pysages/methods/pbmetad.py @@ -28,10 +28,10 @@ class ParallelBiasMetadynamics(GriddedSamplingMethod): involves the 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 implementing this, the keyword - ``parallelbias`` is added to define grids for each CV separately. Currently, only + 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 implementing 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. """ @@ -112,11 +112,11 @@ def initialize(): 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 + # 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 efficient. - + # PySAGES will throw an concatenation error if bins or shape of each CV is different. # So, we use shape[0] to define the size of grids as all bins are expected to be same. grid_potential = np.zeros((shape[0], shape.size), dtype=np.float64) @@ -162,15 +162,20 @@ def build_gaussian_accumulator(method: ParallelBiasMetadynamics): 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) - ) + 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") - ) + 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) @@ -194,7 +199,7 @@ def next_height(pstate): 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. + # 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) ) @@ -247,9 +252,8 @@ def build_bias_grad_evaluator(method: ParallelBiasMetadynamics): if grid is None: periods = get_periods(method.cvs) - evaluate_bias_grad = jit( - lambda pstate: grad(parallelbias)(*pstate[:4], beta, periods, grid) - ) + evaluate_bias_grad = jit(lambda pstate: grad(parallelbias)(*pstate[:4], beta, + periods, grid)) else: def zero_force(_): @@ -278,7 +282,6 @@ def parallelbias_each_cv(xi, heights, centers, sigmas, periods): return bias_each_cv - # Helper function to evaluate parallel bias potential def parallelbias(xi, heights, centers, sigmas, beta, periods, grid): """ @@ -295,7 +298,7 @@ def parallelbias(xi, heights, centers, sigmas, beta, periods, grid): @dispatch def analyze(result: Result[ParallelBiasMetadynamics]): """ - Helper for calculating the free energy from the final state of a + Helper for calculating the free energy from the final state of a `Parallel Bias Metadynamics` run. Parameters @@ -320,16 +323,16 @@ def analyze(result: Result[ParallelBiasMetadynamics]): Maps a user-provided array of CV values and step 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 + 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 + 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, + 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 @@ -343,15 +346,15 @@ def analyze(result: Result[ParallelBiasMetadynamics]): 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), - ) - ) + 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, @@ -364,17 +367,14 @@ def analyze(result: Result[ParallelBiasMetadynamics]): # (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)) - ) + 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), - ) - ) + return jit(vmap(lambda x, beta: parallelbias(x, heights, centers, sigmas, + beta, P, grid), in_axes=(0, None)) + ) heights = [] centers = [] @@ -393,3 +393,4 @@ def build_pbmetapotential(heights, centers, sigmas): pbmetad_potential_cv=pbmetad_potentials_cv, pbmetad_net_potential=pbmetad_net_potentials, ) + From 4d5b522ae6535bfa91a4a20e19d336860a23961e Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Sat, 30 Jul 2022 18:58:57 +0000 Subject: [PATCH 26/29] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- pysages/methods/pbmetad.py | 85 +++++++++++++++++++------------------- 1 file changed, 42 insertions(+), 43 deletions(-) diff --git a/pysages/methods/pbmetad.py b/pysages/methods/pbmetad.py index aae7f118..a6356c9c 100644 --- a/pysages/methods/pbmetad.py +++ b/pysages/methods/pbmetad.py @@ -28,10 +28,10 @@ class ParallelBiasMetadynamics(GriddedSamplingMethod): involves the 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 implementing this, the keyword - ``parallelbias`` is added to define grids for each CV separately. Currently, only + 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 implementing 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. """ @@ -112,11 +112,11 @@ def initialize(): 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 + # 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 efficient. - + # PySAGES will throw an concatenation error if bins or shape of each CV is different. # So, we use shape[0] to define the size of grids as all bins are expected to be same. grid_potential = np.zeros((shape[0], shape.size), dtype=np.float64) @@ -162,20 +162,15 @@ def build_gaussian_accumulator(method: ParallelBiasMetadynamics): 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) - ) + 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") - ) + 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) @@ -199,7 +194,7 @@ def next_height(pstate): 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. + # 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) ) @@ -252,8 +247,9 @@ def build_bias_grad_evaluator(method: ParallelBiasMetadynamics): if grid is None: periods = get_periods(method.cvs) - evaluate_bias_grad = jit(lambda pstate: grad(parallelbias)(*pstate[:4], beta, - periods, grid)) + evaluate_bias_grad = jit( + lambda pstate: grad(parallelbias)(*pstate[:4], beta, periods, grid) + ) else: def zero_force(_): @@ -282,6 +278,7 @@ def parallelbias_each_cv(xi, heights, centers, sigmas, periods): return bias_each_cv + # Helper function to evaluate parallel bias potential def parallelbias(xi, heights, centers, sigmas, beta, periods, grid): """ @@ -298,7 +295,7 @@ def parallelbias(xi, heights, centers, sigmas, beta, periods, grid): @dispatch def analyze(result: Result[ParallelBiasMetadynamics]): """ - Helper for calculating the free energy from the final state of a + Helper for calculating the free energy from the final state of a `Parallel Bias Metadynamics` run. Parameters @@ -323,16 +320,16 @@ def analyze(result: Result[ParallelBiasMetadynamics]): Maps a user-provided array of CV values and step 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 + 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 + 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, + 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 @@ -346,15 +343,15 @@ def analyze(result: Result[ParallelBiasMetadynamics]): 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)) - ) + 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, @@ -367,14 +364,17 @@ def analyze(result: Result[ParallelBiasMetadynamics]): # (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) )) + 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)) - ) + return jit( + vmap( + lambda x, beta: parallelbias(x, heights, centers, sigmas, beta, P, grid), + in_axes=(0, None), + ) + ) heights = [] centers = [] @@ -393,4 +393,3 @@ def build_pbmetapotential(heights, centers, sigmas): pbmetad_potential_cv=pbmetad_potentials_cv, pbmetad_net_potential=pbmetad_net_potentials, ) - From de156e9876046cb1a75c2ab4a125ae47b3a59501 Mon Sep 17 00:00:00 2001 From: Siva Dasetty Date: Sat, 30 Jul 2022 22:01:02 -0500 Subject: [PATCH 27/29] minor fmt fixes --- pysages/methods/pbmetad.py | 17 ++++++++--------- 1 file changed, 8 insertions(+), 9 deletions(-) diff --git a/pysages/methods/pbmetad.py b/pysages/methods/pbmetad.py index a6356c9c..88b15842 100644 --- a/pysages/methods/pbmetad.py +++ b/pysages/methods/pbmetad.py @@ -25,12 +25,12 @@ class ParallelBiasMetadynamics(GriddedSamplingMethod): Compared to well-tempered metadynamics, the Gaussian bias deposited along each CV have different heights in PBMetaD. In addition, the total bias potential - involves the log of sum of exponential of bias potential (see Eq. 8 in the PBMetaD paper) + 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 implementing this, the keyword + 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. """ @@ -46,7 +46,7 @@ def __init__(self, cvs, height, sigma, stride, ngaussians, T, deltaT, kB, **kwar Set of user selected collective variables. height: JaxArray - Initial height of the deposited Gaussians along each CV. + Initial heights of the deposited Gaussians along each CV. sigma: JaxArray Initial standard deviation of the to-be-deposit Gaussians along each CV. @@ -115,10 +115,10 @@ def initialize(): # 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 efficient. + # but that is not supported yet. - # PySAGES will throw an concatenation error if bins or shape of each CV is different. - # So, we use shape[0] to define the size of grids as all bins are expected to be same. + # 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) @@ -190,7 +190,6 @@ def next_height(pstate): get_grid_index = build_indexer(grid) # Reshape so the dimensions are compatible accum = jit(lambda total, val: total + val.reshape(total.shape)) - transform = grad update = jit(lambda V_each_cv, dV, vals, grads: (accum(V_each_cv, vals), accum(dV, grads))) def update_grids(pstate, height, xi, sigma): @@ -205,7 +204,7 @@ def update_grids(pstate, height, xi, sigma): ) grid_potential_values = vmap(current_parallelbias_each_cv)(grid_mesh) - grid_grad_values = vmap(transform(current_parallelbias))(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 @@ -317,7 +316,7 @@ def analyze(result: Result[ParallelBiasMetadynamics]): Height of the Gaussian bias potential along each CV during the simulation. pbmetad_potential_cv: Callable - Maps a user-provided array of CV values and step to the corresponding deposited bias + 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 From 0aa475d10076b426c621ab4ea936b3b130a55656 Mon Sep 17 00:00:00 2001 From: Siva Dasetty Date: Sun, 31 Jul 2022 00:23:18 -0500 Subject: [PATCH 28/29] updated docs to include pbmetad --- docs/source/module-pysages-methods-pbmetad.rst | 13 +++++++++++++ docs/source/module-pysages-methods.rst | 2 ++ docs/source/pysages_wordlist.txt | 1 + 3 files changed, 16 insertions(+) create mode 100644 docs/source/module-pysages-methods-pbmetad.rst 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 1ab96678..7aaf6493 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 @@ -40,6 +41,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-utils module-pysages-methods-core diff --git a/docs/source/pysages_wordlist.txt b/docs/source/pysages_wordlist.txt index 4bebbd17..c08ba359 100644 --- a/docs/source/pysages_wordlist.txt +++ b/docs/source/pysages_wordlist.txt @@ -28,6 +28,7 @@ Rev Lett equilibration Metadynamics +ParallelBiasMetadynamics backend backends colab From 321de8667e3e14869ea86cdebca2c68af2873db8 Mon Sep 17 00:00:00 2001 From: Siva Dasetty Date: Sun, 31 Jul 2022 00:28:55 -0500 Subject: [PATCH 29/29] added missing spell --- docs/source/pysages_wordlist.txt | 3 +++ 1 file changed, 3 insertions(+) diff --git a/docs/source/pysages_wordlist.txt b/docs/source/pysages_wordlist.txt index c08ba359..847ce5ab 100644 --- a/docs/source/pysages_wordlist.txt +++ b/docs/source/pysages_wordlist.txt @@ -24,10 +24,13 @@ anisotropy convolving Jacobian Phys +Comput Rev Lett equilibration Metadynamics +metadynamics +Eq ParallelBiasMetadynamics backend backends