The PYCV module for PLUMED 2

The PYCV module enables PLUMED2 Collective Variables (CVs) and arbitrary functions to be defined and auto-differentiated in the Python language.

Advantages of using PYCV over standard development of CVs in C++ are:

functions may be prototyped in high-level code, using extensive mathematical libraries, and no boilerplate;
just-in-time compilation occurs transparently: there are no compilation and link delays for code changes;
CVs may be automatically differentiated in common cases.

[!WARNING]
This page and the associated paper describe the first implementation of PYCV. Since PLUMED 2.9, a newer variant of PYCV is distributed with PLUMED, albeit with a different syntax and installation procedure. Please see tutorial 24.015. I recommend adapting your code to use the “official” PYCV, if possible.

The code is organized as a standard PLUMED module: development occurs in a fork of the original repository. All code is in the src/pycv and regtest/pycv directories.

Documentation

The PYCV module defines the following actions:

PYTHONCV, to implement single- and multi-component CVs in Python;
PYTHONFUNCTION, to implement arbitrary functions.

In both cases, a Python interpreter is first started; the Python file indicated in the IMPORT= keyword is then imported; from it, an user-chosen function (FUNCTION=) is called to perform the computations at each timestep. Modules can be shared for multiple functions, and contain one-time initialization.

Transparent auto-differentiation, JIT compilation, and vectorization are available through Google’s JAX library (recommended).

Please see the generated documentation and regression tests.

Prerequisites

Python 3 is required. Under OSX, Homebrew’s Python 3 package works well. Conda also seems to work. Make sure you have a Python version recent enough that this command works:

python3-config --embed --ldflags

You will also need to install numpy, either via pip3 install numpy or your distribution’s packages.

Warning If numpy is not available, PYCV will be disabled in the configure stage.

Automatic differentiation examples require the JAX library: install it with pip3 install jax jaxlib. That can be installed later, after building.

Installation

Follow the usual PLUMED 2 configuration procedure, i.e. check-out the PYCV repository and then:

git clone https://github.com/giorginolab/plumed2-pycv.git
cd plumed2-pycv
./configure --enable-modules=+pycv 
make -j4
make install   # optional

Please inspect the configure messages to see if any missing dependency prevents PYCV from actually being enabled.

In case your system has multiple python versions, you may select the one to be embedded adding the argument pycv_python3_config=... to configure. For example: ./configure --enable-modules=+pycv pycv_python3_config=/bin/python3.10-config.

Conda

Installation under Conda seems unexplicably complicated. For one, Conda lacks a symbolic link that you can recreate with

ln -s $CONDA_PREFIX/lib/libpython* $(python3-config --embed --configdir)

But other problems remain, and they are being investigated. On a clean miniconda environment with the conda-forge channel, the following combination worked at least partially:

conda create -n pycv -c conda-forge numpy python=3.10
conda activate pycv
conda install -c conda-forge compilers make
ln -s $CONDA_PREFIX/lib/libpython* $(python3-config --embed --configdir)
./configure --enable-modules=+pycv
make -j4
make -k

Which generates plumed-runtime and plumed-static, but not plumed. In practice:

pip install jax jaxlib
PLUMED_PROGRAM_NAME=$PWD/src/lib/plumed-static make -C regtest/pycv/
PLUMED_KERNEL=$PWD/src/lib/libplumedKernel.so PLUMED_PROGRAM_NAME=$PWD/src/lib/plumed-runtime make -C regtest/pycv/

Test

Verify the successful installation with e.g.

./src/lib/plumed-static manual --action PYTHONCV

which should return the manual for PYTHONCV. (The plumed-static executable should work even without installation.)

Next, try the regression tests:

PLUMED_PROGRAM_NAME=$PWD/src/lib/plumed-static make -C regtest/pycv/

The following Colab notebook should demonstrate all of the above (but may become outdated): PyCV build and test

Note (It is also possible to compile the module as a LOAD-able dynamic object. Once in the src/pycv directory, issue make PYCV.so (PYCV.dylib under OSX). The compilation step should pick Python-specific flags as long as the correct python3-config executable is in your path.)

Quickstart

Here’s a quick example to whet your appetite, following the regression test rt-jax2.

export PLUMED_PROGRAM_NAME=$PWD/src/lib/plumed-static 
make -C regtest/pycv/rt-jax2

The files in regtest/pycv/rt-jax2 are reproduced here:

File plumed.dat

cv1:  PYTHONCV ATOMS=1,4,3 IMPORT=jaxcv FUNCTION=cv1

File jaxcv.py

# Import the JAX library
import jax.numpy as np
from jax import grad, jit, vmap

# Implementation of the angle function. @jit really improves speed
@jit
def angle(x):
    r1 = x[0,:]-x[1,:]
    r2 = x[2,:]-x[1,:]

    costheta = np.dot(r1,r2) / np.sqrt(np.dot(r1,r1) * np.dot(r2,r2))
    # OR: 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
grad_angle = grad(angle)

# The CV function actually called
def cv1(x):
    return angle(x), grad_angle(x)

Demos and regression tests

A number of PLUMED-style regression tests are provided to test and demonstrate the package features.

Name	Feature
`rt-1`	Basic test (distance of two atoms), pure Python, no gradient
`rt-2`	CV and gradient explicitly coded, pure Python
`rt-3`	Multiple PYTHON actions in the same file
`rt-f1`	Function, pure Python
`rt-f2`	Function, with JIT and auto-gradient, with MATHEVAL equivalent
`rt-jax1`	Distance CV with JAX reverse-mode differentiation
`rt-jax2`	Angle CV with JAX reverse-mode differentiation and JIT decorator
`rt-jax3`	Radius of curvature, as in doi:10.1016/j.cpc.2018.02.017
`rt-multi1`	Multi-component CV, pure Python
`rt-multi2`	Multi-component CV, with JAX reverse-mode differentiation

To run:

cd [path_to_repository]
export PLUMED_PROGRAM_NAME=$PWD/src/lib/plumed
cd regtest/pycv
make

Common errors

Missing modules at run time: you may need to set the PYTHONHOME environment variable.
uncaught exception [...] The interpreter is already running This occurs if the module is loaded twice. The most common cause is when it is both built-in in the current kernel, and loaded as a module.

Limitations

Behavior with MPI, and JAX’s GPU/TPU offloading are untested.

Author

Toni Giorgino toni.giorgino@gmail.com

Contributing

Please report bugs and ideas via this repository’s Issues.

Citation

Giorgino T. PYCV: a PLUMED 2 Module Enabling the Rapid Prototyping of Collective Variables in Python. The Journal of Open Source Software 4(42):1773

Copyright

PYCV is distributed under the LGPL terms: see COPYRIGHT.

PYCV includes the pybind11 library, distributed under its own license terms (BSD-3).