MetatomicPotential.h

Go to the documentation of this file.

/*
** This file is part of eOn.
**
** SPDX-License-Identifier: BSD-3-Clause
**
** Copyright (c) 2010--present, eOn Development Team
** All rights reserved.
**
** Repo:
** https://github.com/TheochemUI/eOn
*/
#pragma once
 
#include "eon/Parameters.h"
#include "eon/Potential.h"
#include <mutex>
 
// Metatomic and torch headers
// These pragmas are included to suppress warnings from third-party libraries
#pragma GCC diagnostic push
#pragma GCC diagnostic ignored "-Wpedantic"
#pragma GCC diagnostic ignored "-Wunused-parameter"
#pragma GCC diagnostic ignored "-Wfloat-equal"
#pragma GCC diagnostic ignored "-Wfloat-conversion"
#pragma GCC diagnostic ignored "-Wsign-conversion"
#pragma GCC diagnostic ignored "-Wold-style-cast"
 
#include <torch/script.h>
#include <torch/version.h>
#if __has_include(<torch/cuda.h>)
#include <torch/cuda.h>
#endif
#if __has_include(<torch/mps.h>)
#include <torch/mps.h>
#endif
 
#include "metatensor/torch.hpp"
#include "metatensor/torch/module.hpp"
#include "metatomic/torch.hpp"
 
#pragma GCC diagnostic pop

class MetatomicPotential : public eonc::Potential {
private:
  eonc::log::Scoped m_log;
  eonc::Parameters::metatomic_options_t m_metatomic_opts;
  // --- Metatomic and Torch members ---
  metatensor_torch::Module model_;
  metatomic_torch::ModelCapabilities capabilities_;
  std::vector<metatomic_torch::NeighborListOptions> nl_requests_;
  metatomic_torch::ModelEvaluationOptions evaluations_options_;
 
  torch::ScalarType dtype_;
  c10::DeviceType device_type_;
  torch::Device device_;
  bool check_consistency_;
  // -- Variants / output keys
  std::string energy_key_;
  std::string energy_uncertainty_key_;
  std::string nc_forces_key_; // empty => conservative (autograd) forces
  bool non_conservative_{false};
  bool random_rotation_{false};
  long n_symmetry_rotations_{0};
  // --- Uncertainty handling ---
  // If non-positive, uncertainty checks are effectively disabled.
  double uncertainty_threshold_{-1.0};

  metatensor_torch::TensorBlock
  computeNeighbors(metatomic_torch::NeighborListOptions request, long nAtoms,
                   const double *positions, const double *box,
                   const bool periodic[3]);
 
public:
  MetatomicPotential(const eonc::Parameters &params);

  ~MetatomicPotential() override = default;
 
  [[nodiscard]] std::shared_ptr<eonc::Potential>
  clonePotential() const override;

  void force(long nAtoms, const double *positions, const int *atomicNrs,
             double *forces, double *energy, double *variance,
             const double *box) override;

  [[nodiscard]] bool isThreadSafe() const noexcept override { return false; }
  [[nodiscard]] bool needsPerImageInstance() const noexcept override {
    return false;
  }

  [[nodiscard]] bool supportsBatchEvaluation() const noexcept override {
    return true;
  }
  void forceBatch(long nSystems, long nAtoms, const double *const *positions,
                  const int *const *atomicNrs, double *const *forces,
                  double *energies, double *variances,
                  const double *const *boxes) override;
 
private:
  struct CloneTag {};
  MetatomicPotential(const MetatomicPotential &src, CloneTag);
  void forceBatchNative(long nSystems, long nAtoms,
                        const double *const *positions,
                        const int *const *atomicNrs, double *const *forces,
                        double *energies, double *variances,
                        const double *const *boxes);
  mutable std::mutex inference_mutex_;
};