enable use of CVs defined by PyTorch neural network models

[zwpku]

This branch implements a class called torchANN, which allows to define cv components by loading pretrained PyTorch neural network models.

Installation Steps

1. Download LibTorch. This package is required in order to enable the torchann class. First, download the code and unzip it.

        wget https://download.pytorch.org/libtorch/nightly/cpu/libtorch-cxx11-abi-shared-with-deps-latest.zip
        unzip libtorch-cxx11-abi-shared-with-deps-latest.zip
   

   In this way, the library is uncompressed under the current directory. Let's say it is located at /path/to/libtorch.

2. Patch MD engine. This step is done as usual using the script update-colvars-code.sh. Enter the source code of Colvars package, and run:

        ./update-colvars-code.sh /path/to/md-engine        
   

3. Compilation. This step depends on the engine to be compiled.

   • NAMD: add "--with-colvars-torch --torch-prefix path/to/libtorch" to the argument of ./config

     Assume packages that are required to build NAMD, e.g. charm, tcl/tcl-threaded, are already prepared.
     Then, one can compile the NAMD package with the following commands:

       ./config Linux-x86_64-g++ --charm-arch multicore-linux-x86_64 --with-colvars-torch    \
             --torch-prefix /path/to/libtorch  --with-fftw3 --fftw-prefix /path/to/fftw
       cd Linux-x86_64-g++
       make 
   

   • GROMACS: add "-DTorch_DIR=/path/to/libtorch/share/cmake/Torch" when running cmake

   An example of the command is:

       cmake .. -DCMAKE_INSTALL_PREFIX=/home/username/local/gromacs  \
                       -DFFTWF_LIBRARY=/home/username/mambaforge/lib/libfftw3f.so  \
                       -DFFTWF_INCLUDE_DIR=/home/username/mambaforge/include \
                       -DTorch_DIR=/path/to/libtorch/share/cmake/Torch/  \
                       -DCMAKE_CXX_COMPILER=/usr/bin/mpicxx \
                       -DOpenMP_gomp_LIBRARY=/home/username/mambaforge/lib/libgomp.so
   

   • LAMMPS: only installation by cmake is supported. In the directory of LAMMPS source code, run

        mkdir build && cd build
        cmake ../cmake -D PKG_COLVARS=yes -D COLVARS_TORCH=yes 
   

   and set the variable Torch_DIR in the file CMakeCache.txt. When a cpu version of libtorch library is used, it may
   also be necessary to set MKL path to empty:

        MKL_INCLUDE_DIR:PATH=
   

   Alternatively, one could combine these steps in one command:

        cmake ../cmake -D PKG_COLVARS=yes -D COLVARS_TORCH=yes      \ 
            -D  Torch_DIR=/path/to/libtorch/share/cmake/Torch -D MKL_INCLUDE_DIR=
   

   After that, run make and make install to compile and install the package.

The class has only been tested using simple neural network models (i.e. an autoencoder on alanine dipeptide), under NAMD and GROMACS engines. Feedbacks are welcome!

A (trivial) example

1. Create a PyTorch model

import torch

class MyModel(torch.nn.Module):
    def __init__(self):
        super().__init__()
    def forward(self, x):
        return x

model = MyModel()
scripted_cv_filename = f'./identity.pt'
torch.jit.script(model).save(scripted_cv_filename)

This Python script simply creates a model which is an identity map and save it to a file named identity.pt.

2. Define the COLVARS config file

This file defines two CVs using torchann class taking other cv components (here dihedral angles) as inputs.

colvarsTrajFrequency    10000
colvarsRestartFrequency 10000

colvar {
  name nn_0
  lowerBoundary -180.0
  upperBoundary 180
  width 5.0
  extendedLagrangian on
  extendedFluctuation 5.0
  extendedTimeConstant 200

  torchann {
    modelFile identity.pt
    m_output_index 0
    period 360

    dihedral {
      group1 { 
	atomnumbers 5
      }
      group2 { 
	atomnumbers 7
      }
      group3 { 
	atomnumbers 9
      }
      group4 { 
	atomnumbers 15
      }
    }

    dihedral {
      group1 { 
	atomnumbers 7
      }
      group2 { 
	atomnumbers 9
      }
      group3 { 
	atomnumbers 15
      }
      group4 { 
	atomnumbers 17
      }
    }

  }
}

colvar {
  name nn_1
  lowerBoundary -180.0
  upperBoundary 180
  width 5.0
  extendedLagrangian on
  extendedFluctuation 5.0
  extendedTimeConstant 200

  torchann {
    modelFile identity.pt
    m_output_index 1
    period 360

    dihedral {
      group1 { 
	atomnumbers 5
      }
      group2 { 
	atomnumbers 7
      }
      group3 { 
	atomnumbers 9
      }
      group4 { 
	atomnumbers 15
      }
    }

    dihedral {
      group1 { 
	atomnumbers 7
      }
      group2 { 
	atomnumbers 9
      }
      group3 { 
	atomnumbers 15
      }
      group4 { 
	atomnumbers 17
      }
    }
  }
}

abf {
  colvars nn_0 nn_1
  fullSamples	200
}

[giacomofiorin]

Hi, since we have one CI test working for the library, my only remaining concern would be avoiding routine conflicts with the build systems of the MD engines (see #570 (comment)).

[giacomofiorin]

The latest GROMACS test error is unrelated to Colvars: https://gitlab.com/gromacs/gromacs/-/issues/5204

[giacomofiorin]

Hi there! GROMACS 2025 runs without errors the torchann input from this branch, but there are differences from the reference files generated by @zwpku and updated by @jhenin.

See the outputs here:
https://github.com/Colvars/colvars/actions/runs/11747796165/artifacts/2164700677

[zwpku]

Hi there! GROMACS 2025 runs without errors the torchann input from this branch, but there are differences from the reference files generated by @zwpku and updated by @jhenin.

See the outputs here: https://github.com/Colvars/colvars/actions/runs/11747796165/artifacts/2164700677

@giacomofiorin thanks for the work!
Any hint on what the reason might be? I don't have a good understanding on the regression test... Could it be that the files under gromacs/tests/library/000_torchann/AutoDiff are outdated?

[giacomofiorin]

Hi there! GROMACS 2025 runs without errors the torchann input from this branch, but there are differences from the reference files generated by @zwpku and updated by @jhenin.
See the outputs here: https://github.com/Colvars/colvars/actions/runs/11747796165/artifacts/2164700677

@giacomofiorin thanks for the work! Any hint on what the reason might be? I don't have a good understanding on the regression test... Could it be that the files under gromacs/tests/library/000_torchann/AutoDiff are outdated?

@zwpku Yes, the reference files currently in that folder were produced came from another build, with a different version of libTorch. Would you expect this kind of difference? It is small, but it did exceed our threshold (1.0e-6 relative error).

[zwpku]

Hi there! GROMACS 2025 runs without errors the torchann input from this branch, but there are differences from the reference files generated by @zwpku and updated by @jhenin.
See the outputs here: https://github.com/Colvars/colvars/actions/runs/11747796165/artifacts/2164700677

@giacomofiorin thanks for the work! Any hint on what the reason might be? I don't have a good understanding on the regression test... Could it be that the files under gromacs/tests/library/000_torchann/AutoDiff are outdated?

@zwpku Yes, the reference files currently in that folder were produced came from another build, with a different version of libTorch. Would you expect this kind of difference? It is small, but it did exceed our threshold (1.0e-6 relative error).

@giacomofiorin If I see correctly, the torch model in that test is simply the identity map and the CV is a dihedral angle. So I expect there should be little difference due to different versions of libTorch. Could it also be caused by some changes in the source code or in the config files of that test? I can try to build and examine the test on my local machine.