colvar_neuralnetworkcompute.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 NEURALNETWORKCOMPUTE_H
#define NEURALNETWORKCOMPUTE_H
 
#include <vector>
#include <functional>
#include <string>
#include <cmath>
#include <memory>
#include <map>
 
#ifdef LEPTON
#include "Lepton.h"
#endif
 
class colvarmodule; // Forward declaration
 
namespace neuralnetworkCV {
extern std::map<std::string, std::pair<std::function<double(double)>, std::function<double(double)>>> activation_function_map;
 
#ifdef LEPTON
// allow to define a custom activation function
class customActivationFunction {
public:
    customActivationFunction();
    customActivationFunction(const std::string& expression_string, colvarmodule *cvmodule);
    customActivationFunction(const customActivationFunction& source);
    customActivationFunction& operator=(const customActivationFunction& source);
    void setExpression(const std::string& expression_string);
    std::string getExpression() const;
    double evaluate(double x) const;
    double derivative(double x) const;
private:
    std::string expression;
    std::unique_ptr<Lepton::CompiledExpression> value_evaluator;
    std::unique_ptr<Lepton::CompiledExpression> gradient_evaluator;
    double* input_reference;
    double* derivative_reference;
    colvarmodule *cvmodule;
};
#endif
 
class denseLayer {
private:
    size_t m_input_size;
    size_t m_output_size;
    std::function<double(double)> m_activation_function;
    std::function<double(double)> m_activation_function_derivative;
    colvarmodule *cvmodule;
#ifdef LEPTON
    bool m_use_custom_activation;
    customActivationFunction m_custom_activation_function;
#else
    static const bool m_use_custom_activation = false;
#endif
    std::vector<std::vector<double>> m_weights;
    std::vector<double> m_biases;
public:
    denseLayer() {}
    denseLayer(const std::string& weights_file,
               const std::string& biases_file,
               const std::function<double(double)>& f,
               const std::function<double(double)>& df,
               colvarmodule *cvmodule);
#ifdef LEPTON
    denseLayer(const std::string& weights_file,
               const std::string& biases_file,
               const std::string& custom_activation_expression,
               colvarmodule *cvmodule);
#endif
    void readFromFile(const std::string& weights_file, const std::string& biases_file);
    void setActivationFunction(const std::function<double(double)>& f, const std::function<double(double)>& df);
    void compute(const std::vector<double>& input, std::vector<double>& output) const;
    double computeGradientElement(const std::vector<double>& input, const size_t i, const size_t j) const;
    void computeGradient(const std::vector<double>& input, std::vector<std::vector<double>>& output_grad) const;
    size_t getInputSize() const {
        return m_input_size;
    }
    size_t getOutputSize() const {
        return m_output_size;
    }
    double getWeight(size_t i, size_t j) const {
        return m_weights[i][j];
    }
    double getBias(size_t i) const {
        return m_biases[i];
    }
    ~denseLayer() {}
};
 
class neuralNetworkCompute {
private:
    std::vector<denseLayer> m_dense_layers;
    std::vector<double> m_input;
    std::vector<std::vector<double>> m_layers_output;
    std::vector<std::vector<std::vector<double>>> m_grads_tmp;
    std::vector<std::vector<double>> m_chained_grad;
private:
    static std::vector<std::vector<double>> multiply_matrix(const std::vector<std::vector<double>>& A, const std::vector<std::vector<double>>& B);
public:
    neuralNetworkCompute(): m_dense_layers(0), m_layers_output(0) {}
    neuralNetworkCompute(const std::vector<denseLayer>& dense_layers);
    bool addDenseLayer(const denseLayer& layer);
    // for faster computation
    const std::vector<double>& input() const {return m_input;}
    std::vector<double>& input() {return m_input;}
    void compute();
    double getOutput(const size_t i) const {return m_layers_output.back()[i];}
    double getGradient(const size_t i, const size_t j) const {return m_chained_grad[i][j];}
    const denseLayer& getLayer(const size_t i) const {return m_dense_layers[i];}
    size_t getNumberOfLayers() const {return m_dense_layers.size();}
};
 
}
#endif