PYCVINTERFACE

Define collective variables in the Python language.

Plumed will import the module chosen wiht the IMPORT keyword. The module can be a simple py file or a directory with the standard module configuration (a directory that contains at least an init.py, see the rt-persistentData test for an example).

In the module there must be a function that will be used in caclulation and the definition of the component(s) (meaning period and if derivatives will be returned) of your cv.

PYCVINTERFACE will call two or more elements from the module. Each element or function can be selected with its dedicated keyword:

• CALCULATE will select the function to be called at each step (defaults to "plumedCalculate")
• INIT will select the function (or the dict, see down) to be called during the initilization (defaults to "plumedInit")
• UPDATE will select the function to be called at each step, during the update() invocation
• PREPARE will select the function to be called at each step, during the prepare() invocation All the function called will need a single argument of type plumedCommunications.PythonCVInterface

Getting started

The minimal plumed input is:

Click on the labels of the actions for more information on what each action computes

cv1: PYTHONCV 
IMPORT
 could not find this keyword 
=hello 
ATOMS
 could not find this keyword 
=1  You cannot view the components that are calculated by each action for this input file. Sorry 
PRINT 
FILE
the name of the file on which to output these quantities 
=colvar.out 
ARG
the input for this action is the scalar output from one or more other actions. 
=*  You cannot view the components that are calculated by each action for this input file. Sorry 

this should be paired with the hello.py file:

import plumedCommunications as PLMD
 
plumedInit={"Value":PLMD.defaults.COMPONENT_NODEV}
 
def plumedCompute(action: PLMD.PythonCVInterface):
    action.log("Hello, world, from calculate!")
    return 0.0

If INIT is not specified, plumed will search for an object "plumedInit", that can be either a function that returns a dict or a dict This dict MUST contain at least the informations about the presence of the derivatives and on the periodicity of the variable. We will refer to this dict as "the init dict" from now.

If CALCULATE is not specified, plumed will search for a function named "plumedCalculate" plumed will read the variable returned accordingly to what it was specified in the initialization dict.

The init dict will tell plumed how many components the calculate function will return and how they shall behave. Along this the dict can contain all the keyword that are compatible with PYCVINTERFACE. Mind that if the same keyword is specified both in the init dict and in the plumed file the calculation will be aborted to avoid unwated settings confict. In case of flags the dict entry must be a bool, differently from the standard plumed input.

The only keyword that can only be specified in python is COMPONENTS. The COMPONENTS key must point to a dict that has as keys the names of the componensts. Each component dictionary must have two keys:

• "period": None of a list of two values, min and max (like [0,1] or also strings like ["0.5*pi","2*pi"])
• "derivative": True or False If you want to use a single component you can create the "COMPONENTS" dict with as single key, the name will be ignored. In the previous example the key "Value" is used instead of "COMPONENTS": it is a shorter form for "COMPONENTS":{"any":{...}}. To avoid confusion you cannot specify both "COMPONENTS" and "Value" in the same dict.

To speed up the declarations of the components the plumedCommunications module contains a submodule defaults with the default dictionaries already set up:

• plumedCommunications.defaults.COMPONENT={"period":None, "derivative":True}
• plumedCommunications.defaults.COMPONENT_NODEV={"period":None, "derivative":False}

The calculate function

The calculate funtion must, as all the other functions accept a PLMD.PythonCVInterface object as only input.

The calculate function must either return a float or a tuple or, in the case of multiple components, a dict whose keys are the name of the components, whose elements are either float or tuple.

Plumed will assign automatically assign the result to the CV (to the key named element), if the name of the component is missing the calculation will be interrupted with an error message. If derivatives are disabled it will expect a float(or a double). In case of activated derivatives it will interrupt the calculation if the return value would not be a tuple. The tuple should be (float, ndArray(nat,3),ndArray(3,3)) with the first elements the value, the second the atomic derivatives and the third the box derivative (that can also have shape(9), with format (x_x, x_y, x_z, y_x, y_y, y_z, z_x, z_y, z_z)), if the box derivative are not present a WARNING will be raised, but the calculation won't be interrupted.

The prepare and update functions and the "data" attribute

If the PREPARE keyword is used, the defined function will be called at prepare time, before calculate. The prepare dictionary can contain a "setAtomRequest" key with a parseable ATOM string, like in the input (or a list of indexes, 0 based).

#this , with "PREPARE=changeAtom" in the plumed file will select a new atom at each new step
def changeAtom(plmdAction: plumedCommunications.PythonCVInterface):
    toret = {"setAtomRequest": f"1, {int(plmdAction.getStep()) + 2}"}
    if plmdAction.getStep() == 3:
        toret["setAtomRequest"] = "1,2"
    return toret

If the UPDATE keyword is used, the defined function will be called at update time, after calculate. As now plumed will ignore the return of this function (but it stills need to return a dict) and it is intended to accumulate things or post process data afer calculate

In the example plmdAction.data["pycv"]=0 is intialized in pyinit and its value is updated in calculate.

Click on the labels of the actions for more information on what each action computes

cv1: PYCVINTERFACE ...
   
ATOMS
 could not find this keyword 
=@mdatoms 
   
IMPORT
 could not find this keyword 
=pycvPersistentData 
   
CALCULATE
 could not find this keyword 
=pydist 
   
INIT
 could not find this keyword 
=pyinit 
... You cannot view the components that are calculated by each action for this input file. Sorry 
PRINT 
FILE
the name of the file on which to output these quantities 
=colvar.out 
ARG
the input for this action is the scalar output from one or more other actions. 
=*  You cannot view the components that are calculated by each action for this input file. Sorry 

import plumedCommunications as PLMD
from plumedCommunications.defaults import COMPONENT_NODEV
 
def pyinit(plmdAction: PLMD.PythonCVInterface):
    plmdAction.data["pycv"]=0
    print(f"{plmdAction.data=}", file=log)
    return {"Value":COMPONENT_NODEV}
 
def pydist(plmdAction: PLMD.PythonCVInterface):
    plmdAction.data["pycv"]+=plmdAction.getStep()
    d=plmdAction.data["pycv"]
    return d

The plumedCommunications.PythonCVInterface has a data attribute that is a dictionary and can be used to store data during the calculations

Getting the manual

You can obtain the manual of the module (and of its components) by running plumed driver with this plumed file

Click on the labels of the actions for more information on what each action computes

LOAD 
GLOBAL
 could not find this keyword 
 
FILE
compulsory keyword 
file to be loaded 
=PythonCVInterface.so  You cannot view the components that are calculated by each action for this input file. Sorry 
cvdist: PYCVINTERFACE 
IMPORT
 could not find this keyword 
=pyhelp  You cannot view the components that are calculated by each action for this input file. Sorry 
PRINT 
FILE
the name of the file on which to output these quantities 
=colvar.out 
ARG
the input for this action is the scalar output from one or more other actions. 
=*  You cannot view the components that are calculated by each action for this input file. Sorry 

and this py module

import plumedCommunications
import pydoc
 
def plumedInit(_):
    with open('PythonCVInterface.help.txt', 'w') as f:
        h = pydoc.Helper(output=f)
        h(plumedCommunications.PythonCVInterface)
    with open('plumedCommunications.help.txt', 'w') as f:
        h = pydoc.Helper(output=f)
        h(plumedCommunications)
    with open('plumedCommunications.defaults.help.txt', 'w') as f:
        h = pydoc.Helper(output=f)
        h(plumedCommunications.defaults)
    return {"Value":plumedCommunications.defaults.COMPONENT_NODEV, "ATOMS":"1"}
 
def plumedCalculate(_):
    return 0

Tips and tricks and examples

Automatic differentiation and transparent compilation (including to GPU) can be performed via Google's JAX library: see the example below.

The following input tells PLUMED to print the distance between atoms 1 and 4.

Click on the labels of the actions for more information on what each action computes

cv1: PYTHONCV 
ATOMS
 could not find this keyword 
=1,4 
IMPORT
 could not find this keyword 
=distcv 
CALCULATE
 could not find this keyword 
=cv  You cannot view the components that are calculated by each action for this input file. Sorry 
PRINT 
FILE
the name of the file on which to output these quantities 
=colvar.out 
ARG
the input for this action is the scalar output from one or more other actions. 
=*  You cannot view the components that are calculated by each action for this input file. Sorry 

The file distcv.py should contain something as follows.

import numpy as np
import plumedCommunications as PLMD
 
plumedInit={"Value":PLMD.defaults.COMPONENT}
# Define the distance function
def dist_f(x):
    r = x[0,:]-x[1,:]
    d2 = np.dot(r,r)
    return np.sqrt(d2)
 
def grad_dist(x):
    d = dist_f(x)
    r = x[0,:]-x[1,:]
    g = r/d
    return np.array([g,-g])
 
# The CV function actually called
def cv(action:PLMD.PythonCVInterface):
    return dist_f(action.getPositions()), grad_dist(action.getPositions())

JAX for automatic differentiation and compilation

Automatic differentiation and transparent compilation (including to GPU) can be performed via Google's JAX library. In a nutshell, it's sufficient to replace numpy with jax.numpy. See the following example.

Click on the labels of the actions for more information on what each action computes

cv1: PYTHONCV 
ATOMS
 could not find this keyword 
=1,2,3 
IMPORT
 could not find this keyword 
=jaxcv 
CALCULATE
 could not find this keyword 
=angle  You cannot view the components that are calculated by each action for this input file. Sorry 

PRINT FILE=colvar.out ARG=*

And, in jaxcv.py...

# Import the JAX library
import jax.numpy as np
from jax import grad, jit, vmap
import plumedCommunications as PLMD
 
plumedInit={"Value":PLMD.defaults.COMPONENT}
 
# Implementation of the angle function
def angle_f(x):
    r1 = x[0,:]-x[1,:]
    r2 = x[2,:]-x[1,:]
 
    costheta = np.dot(r1,r2) / np.linalg.norm(r1) / np.linalg.norm(r2)
    theta = np.arccos(costheta)
    return theta
 
# Use JAX to auto-gradient it
angle_grad = grad(angle_f)
 
def cv(action:PLMD.PythonCVInterface):
    return angle_f(action.getPositions()), angle_grad(action.getPositions())

There are however limitations, the most notable of which is that indexed assignments such as x[i]=y are not allowed, and should instead be replaced by functional equivalents such as x=jax.ops.index_update(x, jax.ops.index[i], y).

Multiple components

It is possible to return multiple components at a time. This may be useful e.g. if they reuse part of the computation. To do so, pass the declare the "COMPONENTS" key in the init dict, and assign to each key a dict with the same rules of the key "Value". In this case, the function must return a dict with the component names as keys, see the "calculate" paragraph. Inside PLUMED, component names will be prefixed by py-.

Note that you can use JAX's Jacobian function jax.jacrev() to conveniently compute the gradients all at once (see regtests). For example:

Click on the labels of the actions for more information on what each action computes

cv1: PYTHONCV 
ATOMS
 could not find this keyword 
=1,3,4 
IMPORT
 could not find this keyword 
=distcv 
FUNCTION
 could not find this keyword 
=cv 
COMPONENTS
 could not find this keyword 
=d12,d13  You cannot view the components that are calculated by each action for this input file. Sorry 

import jax.numpy as np
from jax import jacrev, jit
import plumedCommunications as PLMD
 
plumedInit = dict(
    COMPONENTS=dict(d12=PLMD.defaults.COMPONENT, d13=PLMD.defaults.COMPONENT)
)
 
# Define the distance function
@jit
def dist_f(X):
    return {
     'd12': np.linalg.norm( X[0,:]-X[1,:] ),
     'd13': np.linalg.norm( X[0,:]-X[2,:] )
     }
 
dist_grad=jacrev(dist_f)
 
def cv(action: PLMD.PythonCVInterface):
    toret = {}
    d=dist_f(action.getPositions())
    g=dist_grad(action.getPositions())
    for key in ["d12","d13"]:
        toret[key] = (d[key],g[key])
    return  toret

Installation

A use of an virtual environment or equivalent is recomended. To compile pycv you just need numpy and pybind11, jax is not necessary for compilation and installation.

To compile the shared object library you need also plumed in your path. You need to export the following environmental variables:

export PLUMED_MKLIB_CFLAGS="$(python3-config --cflags --embed) $(python -m pybind11 --includes)"
export PLUMED_MKLIB_LDFLAGS="$(python3-config --ldflags --embed)"

and then compile the shared object:

plumed mklib PythonCVInterface.cpp ActionWithPython.cpp PlumedPythonEmbeddedModule.cpp

If you are on linux you can use pycv only with a plumed version that is compatible with the GLOBAL keyword for the action LOAD

Glossary of keywords and components