colvarcomp_torchann.h

// -*- c++ -*-
 
// This file is part of the Collective Variables module (Colvars).
// The original version of Colvars and its updates are located at:
// https://github.com/Colvars/colvars
// Please update all Colvars source files before making any changes.
// If you wish to distribute your changes, please submit them to the
// Colvars repository at GitHub.
//
#ifndef COLVARCOMP_TORCH_H
#define COLVARCOMP_TORCH_H
 
// Declaration of torchann 
 
#include <memory>
 
#include "colvar.h"
#include "colvarcomp.h"
#include "colvarmodule.h"
 
#ifdef COLVARS_TORCH 
 
#include <torch/torch.h>
#include <torch/script.h>
 
class colvar::torchANN
  : public colvar::linearCombination
{
protected:
    torch::jit::script::Module nn;
    size_t m_output_index;
    bool use_double_input;
    //bool use_gpu;
    // 1d tensor, concatenation of values of sub-cvcs
    torch::Tensor input_tensor;
    torch::Tensor nn_outputs;
    torch::Tensor input_grad;
    // record the initial index of of sub-cvcs in input_tensor
    std::vector<int> cvc_indices;
public:
    torchANN();
    virtual ~torchANN();
    virtual int init(std::string const &conf);
    virtual void calc_value();
    virtual void calc_gradients();
    virtual void apply_force(colvarvalue const &force);
};
 
#else
 
class colvar::torchANN
  : public colvar::cvc
{
public:
    torchANN();
    virtual ~torchANN();
    virtual int init(std::string const &conf);
    virtual void calc_value();
};
#endif // COLVARS_TORCH checking
       
#endif