MolecularMachineLearningModel.h

Go to the documentation of this file.

#ifndef SWOOSE_MOLECULARMACHINELEARNINGMODEL_H
#define SWOOSE_MOLECULARMACHINELEARNINGMODEL_H

#include <Utils/Math/MachineLearning/Regression/KernelRidgeRegression.h>
#include <Eigen/Dense>
#include <vector>

namespace Scine {

// Forward declarations
namespace Utils {
class MolecularTrajectory;
class AtomCollection;
} // namespace Utils

namespace Swoose {
using ForcesCollection = Eigen::Matrix<double, Eigen::Dynamic, 3, Eigen::RowMajor>;

namespace MachineLearning {

class MolecularMachineLearningModel {
 public:
  MolecularMachineLearningModel() = default;
  void setReferenceData(const Utils::MolecularTrajectory& structures, const std::vector<ForcesCollection>& refForces);
  void trainEnergyModel();
  void trainForcesModel();
  double predictEnergy(const Utils::AtomCollection& structure);
  ForcesCollection predictForces(const Utils::AtomCollection& structure);
  std::pair<double, double> evaluateEnergyModel(int k);
  std::pair<double, double> evaluateForcesModel(int k, bool pooledVariance = true);
  Utils::MachineLearning::KernelRidgeRegression& energyPredictor();
  Utils::MachineLearning::KernelRidgeRegression& forcePredictor(int atomIndex);

 private:
  // Sets the default values of the model hyperparameters for all of the predictors.
  void setDefaultHyperparameters();
  // Trains a single forces model for a given atom index
  void trainSingleForceModel(int atomIndex);
  // Evaluates a single force model for a given atom index with k-fold cross-validation.
  std::pair<double, double> evaluateSingleForceModel(int atomIndex, int k);
  // Returns the energy feature matrix
  Eigen::MatrixXd getEnergyFeatures();
  // Returns the energy target vector
  Eigen::VectorXd getEnergyTargets();
  // Returns the feature matrix for the atomic force of the atom with index 'atomIndex'
  Eigen::MatrixXd getSingleForceFeatures(int atomIndex);
  // Returns the target matrix for the atomic force of the atom with index 'atomIndex'
  Eigen::MatrixXd getSingleForceTargets(int atomIndex);
  // The KRR model that learns and predicts the energy of the molecular system
  Utils::MachineLearning::KernelRidgeRegression energyPredictor_;
  // The KRR models that learn and predict the atomic forces of the molecular system.
  // It exists one model for each atom.
  std::vector<Utils::MachineLearning::KernelRidgeRegression> forcePredictors_;
  // The Coulomb matrix features for all structures (reference data)
  std::vector<Eigen::VectorXd> coulombMatrixFeatures_;
  // The features for learning atomic forces for all structures
  std::vector<std::vector<Eigen::VectorXd>> forcesFeatures_;
  // Reference energy targets
  std::vector<double> refEnergies_;
  // Reference atomic forces targets
  std::vector<ForcesCollection> refForces_;
};

} // namespace MachineLearning
} // namespace Swoose
} // namespace Scine

#endif // SWOOSE_MOLECULARMACHINELEARNINGMODEL_H