MetatomicPotential Class Reference

A potential class that uses a metatomic model for energy and force calculations. More...

#include <MetatomicPotential.h>

Inheritance diagram for MetatomicPotential:

Classes
struct	CloneTag

Public Member Functions
	MetatomicPotential (const eonc::Parameters &params)
	Constructor for the MetatomicPotential.
	~MetatomicPotential () override=default
	Destructor.
std::shared_ptr< eonc::Potential >	clonePotential () const override
	Independent instance that does not reload from disk.
void	force (long nAtoms, const double *positions, const int *atomicNrs, double *forces, double *energy, double *variance, const double *box) override
	Calculates the energy and forces for a given atomic configuration.
bool	isThreadSafe () const noexcept override
	Single shared instance, serialized via mutex.
bool	needsPerImageInstance () const noexcept override
	Whether NEB should create separate Potential instances per image for true parallel force evaluation.
bool	supportsBatchEvaluation () const noexcept override
	Batched evaluation via single shared instance.
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
Public Member Functions inherited from eonc::Potential
	Potential (PotType a_ptype)
	Production default: construction-scope registry, else PotRegistry::get().
	Potential (PotType a_ptype, IPotRegistry &registry)
	Test seam: injected registry, no process-default get() counters.
	Potential (PotType a_ptype, const Parameters &p)
	Potential (const Parameters &a_params)
virtual	~Potential ()
void	force (std::span< const double > positions, std::span< const int > atomicNrs, std::span< double > forces, double *energy, double *variance, std::span< const double > box)
	C++ call site: size-checked view over the raw FFI force().
virtual void	setFixedMask (long nAtoms, const double *isFixed)
	Optional frozen-atom mask (nAtoms*3, 1.0 = fixed).
std::tuple< double, AtomMatrix >	get_ef (const AtomMatrix &pos, const VectorXi &atmnrs, const Matrix3d &box)
PotType	getType () const
virtual double	finiteCutoff () const noexcept
	Finite interaction range in position length units.
virtual bool	isSurrogate () const noexcept
	Whether this is a surrogate (GP) potential.
virtual bool	requiresIsolatedMoleculeLayout () const noexcept
	True for molecular QM / non-PBC backends (NWChem socket, ASE ORCA/NWChem, …).
virtual unsigned	layoutFlags () const noexcept
virtual bool	isSharedInstanceThreadSafe () const noexcept
	Conservative gate for sharing one Potential instance across threads.
virtual bool	computesStress () const noexcept
	True when force() leaves a Cauchy stress that cauchyStress() can read until the next force() on this instance.
virtual Matrix3d	cauchyStress () const
	Cauchy stress in eV/Angstrom^3.
virtual void	forceBatchOwned (long nSystems, long nAtoms, const double *const *positions, const int *const *atomicNrs, double *const *forces, double *energies, double *variances, const double *const *boxes, const long *owners)
	Evaluate forces for N systems in a single call.

Private Member Functions
metatensor_torch::TensorBlock	computeNeighbors (metatomic_torch::NeighborListOptions request, long nAtoms, const double *positions, const double *box, const bool periodic[3])
	Computes neighbor list using the vesin library.
	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)

Private Attributes
eonc::log::Scoped	m_log
eonc::Parameters::metatomic_options_t	m_metatomic_opts
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_
std::string	energy_key_
std::string	energy_uncertainty_key_
std::string	nc_forces_key_
bool	non_conservative_ {false}
bool	random_rotation_ {false}
long	n_symmetry_rotations_ {0}
double	uncertainty_threshold_ {-1.0}
std::mutex	inference_mutex_

Additional Inherited Members
Public Types inherited from eonc::Potential
enum class	PotLayout : unsigned { InProcess = 1u << 0 , NeedsWorkingDirectory = 1u << 1 , Subprocess = 1u << 2 }
	How the pot is executed. Combine with bitwise or. More...
Public Attributes inherited from eonc::Potential
std::atomic< size_t >	forceCallCounter
Protected Attributes inherited from eonc::Potential
PotType	ptype

Detailed Description

A potential class that uses a metatomic model for energy and force calculations.

This class loads a pre-trained atomistic model (via metatomic/PyTorch) and uses it to compute the potential energy and corresponding forces on atoms. It uses the vesin library to compute neighbor lists required by the model.

Definition at line 54 of file MetatomicPotential.h.

Constructor & Destructor Documentation

MetatomicPotential() [1/2]

MetatomicPotential::MetatomicPotential

(

const eonc::Parameters &

params

)

Constructor for the MetatomicPotential.

Parameters

params

A shared pointer to the simulation parameters object. This object should contain the settings needed for the metatomic model.

Definition at line 56 of file MetatomicPotential.cpp.

    : eonc::Potential(eonc::PotType::METATOMIC),
      m_metatomic_opts{params.metatomic_options()},
      model_(torch::jit::Module()),
      device_type_(c10::DeviceType::CPU),
      device_(torch::Device(device_type_)) {
 
  // Determinism knobs (see
  // https://rgoswami.me/snippets/pytorch-deterministic-regression/): JIT
  // profiling specializes graphs after the first few forwards, so a fresh model
  // instance can differ at ULP level from a warm one — bad for parallel NEB.
  // cuBLAS / index_add_ on CUDA are likewise nondeterministic unless forced.
  // [Metatomic] deterministic=true (default) applies the safe defaults;
  // deterministic_strict=true requires CUBLAS_WORKSPACE_CONFIG (e.g. :4096:8)
  // and fails on nondeterministic ops instead of warning.
  if (m_metatomic_opts.deterministic) {
    torch::jit::getProfilingMode() = false;
    const bool strict = m_metatomic_opts.deterministic_strict ||
                        (std::getenv("CUBLAS_WORKSPACE_CONFIG") != nullptr);
    at::globalContext().setDeterministicAlgorithms(true,
                                                   /*warn_only=*/!strict);
    at::globalContext().setBenchmarkCuDNN(false);
    QUILL_LOG_INFO(m_log,
                   "[MetatomicPotential] Deterministic algorithms enabled "
                   "(strict={})",
                   strict);
  } else {
    QUILL_LOG_INFO(m_log,
                   "[MetatomicPotential] Deterministic algorithms disabled "
                   "(faster, may diverge across runs / NEB images)");
  }
 
  eonc::FPEHandler fpeh;
  fpeh.eat_fpe();
 
  QUILL_LOG_INFO(m_log, "[MetatomicPotential] Initializing...");
 
  // 1. Load the model from the path specified in parameters
  torch::optional<std::string> extensions_directory = torch::nullopt;
  if (!m_metatomic_opts.extensions_directory.empty()) {
    extensions_directory = m_metatomic_opts.extensions_directory;
  }
 
  try {
    QUILL_LOG_INFO(m_log, "[MetatomicPotential] Loading model from '{}'",
                   m_metatomic_opts.model_path);
    this->model_ = metatomic_torch::load_atomistic_model(
        m_metatomic_opts.model_path, extensions_directory);
  } catch (const std::exception &e) {
    QUILL_LOG_ERROR(m_log, "[MetatomicPotential] Failed to load model: {}",
                    e.what());
    throw;
  }
 
  // 2. Extract capabilities and neighbor list requests from the model
  this->capabilities_ =
      this->model_.run_method("capabilities")
          .toCustomClass<metatomic_torch::ModelCapabilitiesHolder>();
  auto requests_ivalue = this->model_.run_method("requested_neighbor_lists");
  for (const auto &request_ivalue : requests_ivalue.toList()) {
    auto request =
        request_ivalue.get()
            .toCustomClass<metatomic_torch::NeighborListOptionsHolder>();
    this->nl_requests_.push_back(request);
  }
 
  // 3. Determine and set up the device (CPU/CUDA/MPS)
  torch::optional<std::string> desired = torch::nullopt;
  if (!m_metatomic_opts.device.empty()) {
    desired = m_metatomic_opts.device;
  }
  device_type_ = metatomic_torch::pick_device(
      this->capabilities_->supported_devices, desired);
 
  device_ = torch::Device(device_type_);
  QUILL_LOG_INFO(m_log, "[MetatomicPotential] Using device: {}", device_.str());
 
  move_atomistic_model(this->model_, this->device_);
 
  // 4. Set data type (float32/float64) based on model capabilities
  if (this->capabilities_->dtype() == "float64") {
    this->dtype_ = torch::kFloat64;
  } else if (this->capabilities_->dtype() == "float32") {
    this->dtype_ = torch::kFloat32;
  } else {
    throw std::runtime_error("Unsupported dtype: " +
                             this->capabilities_->dtype());
  }
  QUILL_LOG_INFO(m_log, "[MetatomicPotential] Using dtype: {}",
                 this->capabilities_->dtype().c_str());
 
  // 5. Resolve energy / force output keys: explicit keys (#215) or variants
  // (#296 for non_conservative_force)
  auto outputs = this->capabilities_->outputs();
 
  auto v_base = normalize_variant(m_metatomic_opts.variant.base);
  auto v_energy = m_metatomic_opts.variant.energy.empty()
                      ? v_base
                      : normalize_variant(m_metatomic_opts.variant.energy);
  auto v_energy_uq =
      m_metatomic_opts.variant.energy_uncertainty.empty()
          ? v_energy
          : normalize_variant(m_metatomic_opts.variant.energy_uncertainty);
  auto v_force = m_metatomic_opts.variant.force.empty()
                     ? v_energy
                     : normalize_variant(m_metatomic_opts.variant.force);
 
  if (!m_metatomic_opts.energy_output.empty()) {
    this->energy_key_ = m_metatomic_opts.energy_output;
  } else {
    this->energy_key_ =
        metatomic_torch::pick_output("energy", outputs, v_energy);
  }
 
  this->non_conservative_ = m_metatomic_opts.non_conservative;
  this->random_rotation_ = m_metatomic_opts.random_rotation;
  this->n_symmetry_rotations_ = m_metatomic_opts.n_symmetry_rotations;
  if (this->non_conservative_) {
    if (!m_metatomic_opts.force_output.empty()) {
      this->nc_forces_key_ = m_metatomic_opts.force_output;
      if (!outputs.contains(this->nc_forces_key_)) {
        throw std::runtime_error(
            "Missing explicit force_output in metatomic model: " +
            this->nc_forces_key_);
      }
    } else {
      this->nc_forces_key_ = metatomic_torch::pick_output(
          "non_conservative_force", outputs, v_force);
    }
    QUILL_LOG_INFO(m_log,
                   "[MetatomicPotential] Non-conservative forces from '{}'",
                   this->nc_forces_key_);
  }
  if (this->n_symmetry_rotations_ > 0) {
    QUILL_LOG_INFO(m_log,
                   "[MetatomicPotential] Symmetry averaging over {} rotations",
                   this->n_symmetry_rotations_);
  } else if (this->random_rotation_) {
    QUILL_LOG_INFO(
        m_log, "[MetatomicPotential] Per-call random SO(3) rotation enabled");
  }
  if ((this->random_rotation_ || this->n_symmetry_rotations_ > 0) &&
      params.main_options().randomSeed > 0) {
    torch::manual_seed(static_cast<uint64_t>(params.main_options().randomSeed));
    QUILL_LOG_INFO(m_log,
                   "[MetatomicPotential] torch RNG seeded from "
                   "main.randomSeed={}",
                   params.main_options().randomSeed);
  }
 
  if (!outputs.contains(this->energy_key_)) {
    QUILL_LOG_ERROR(
        m_log,
        "[MetatomicPotential] The model does not provide an '{}' output.",
        this->energy_key_);
    throw std::runtime_error("Missing energy output in metatomic model");
  }
 
  // 6. Set up evaluation options to request total energy
  this->evaluations_options_ =
      torch::make_intrusive<metatomic_torch::ModelEvaluationOptionsHolder>();
  evaluations_options_->set_length_unit(m_metatomic_opts.length_unit);
 
  auto model_output = outputs.at(this->energy_key_);
  auto requested_output =
      torch::make_intrusive<metatomic_torch::ModelOutputHolder>();
 
  // Per-atom granularity is sample_kind == "atom" (get/set_per_atom removed).
  requested_output->set_sample_kind(model_output->sample_kind());
  requested_output->explicit_gradients = {};
  requested_output->set_unit("eV");
  evaluations_options_->outputs.insert(this->energy_key_, requested_output);
 
  // Request non-conservative forces when enabled (#296)
  if (this->non_conservative_ && !this->nc_forces_key_.empty()) {
    auto nc_info = outputs.at(this->nc_forces_key_);
    auto requested_nc =
        torch::make_intrusive<metatomic_torch::ModelOutputHolder>();
    requested_nc->set_sample_kind(nc_info->sample_kind());
    requested_nc->explicit_gradients = {};
    requested_nc->set_unit("eV/Angstrom");
    evaluations_options_->outputs.insert(this->nc_forces_key_, requested_nc);
  }
 
  // 7. Optionally request energy uncertainty if threshold is positive
  if (m_metatomic_opts.uncertainty_threshold > 0) {
    this->uncertainty_threshold_ = m_metatomic_opts.uncertainty_threshold;
    const bool explicit_uq_key =
        !m_metatomic_opts.energy_uncertainty_output.empty();
    if (explicit_uq_key) {
      // User-specified key: hard fail if missing (not soft-disabled).
      this->energy_uncertainty_key_ =
          m_metatomic_opts.energy_uncertainty_output;
      if (!outputs.contains(this->energy_uncertainty_key_)) {
        QUILL_LOG_ERROR(m_log,
                        "[MetatomicPotential] energy_uncertainty_output '{}' "
                        "is not provided by the model.",
                        this->energy_uncertainty_key_);
        throw std::runtime_error(
            "Missing explicit energy_uncertainty_output in metatomic model: " +
            this->energy_uncertainty_key_);
      }
    } else {
      try {
        this->energy_uncertainty_key_ = metatomic_torch::pick_output(
            "energy_uncertainty", outputs, v_energy_uq);
      } catch (const std::exception &e) {
        QUILL_LOG_DEBUG(
            m_log, "[MetatomicPotential] No uncertainty output available: {}",
            e.what());
        this->uncertainty_threshold_ = -1.0;
      }
    }
 
    if (this->uncertainty_threshold_ > 0) {
      auto uncertainty_info = outputs.at(this->energy_uncertainty_key_);
      if (uncertainty_info->sample_kind() == "atom") {
        auto requested_uncertainty =
            torch::make_intrusive<metatomic_torch::ModelOutputHolder>();
        requested_uncertainty->set_sample_kind("atom");
        requested_uncertainty->explicit_gradients = {};
        requested_uncertainty->set_unit("eV");
        evaluations_options_->outputs.insert(this->energy_uncertainty_key_,
                                             requested_uncertainty);
        QUILL_LOG_INFO(m_log,
                       "[MetatomicPotential] Requested per-atom "
                       "'{}' from model (threshold = {})",
                       this->energy_uncertainty_key_,
                       this->uncertainty_threshold_);
      } else {
        QUILL_LOG_DEBUG(m_log,
                        "[MetatomicPotential] Model provides '{}' "
                        "but sample_kind is not \"atom\"; skipping uncertainty "
                        "checks.",
                        this->energy_uncertainty_key_);
        this->uncertainty_threshold_ = -1.0;
      }
    }
  }
 
  this->check_consistency_ = m_metatomic_opts.check_consistency;
  QUILL_LOG_INFO(m_log, "[MetatomicPotential] Initialization complete.");
 
  fpeh.restore_fpe();
}

~MetatomicPotential()

MetatomicPotential::~MetatomicPotential ( )

overridedefault

Destructor.

MetatomicPotential() [2/2]

MetatomicPotential::MetatomicPotential ( const MetatomicPotential & src, CloneTag  )

private

Definition at line 302 of file MetatomicPotential.cpp.

    : eonc::Potential(eonc::PotType::METATOMIC),
      m_metatomic_opts{src.m_metatomic_opts},
      model_(src.model_.clone(/*inplace=*/false)),
      capabilities_{src.capabilities_},
      nl_requests_{src.nl_requests_},
      evaluations_options_{src.evaluations_options_},
      dtype_{src.dtype_},
      device_type_{src.device_type_},
      device_{src.device_},
      check_consistency_{src.check_consistency_},
      energy_key_{src.energy_key_},
      energy_uncertainty_key_{src.energy_uncertainty_key_},
      nc_forces_key_{src.nc_forces_key_},
      non_conservative_{src.non_conservative_},
      random_rotation_{src.random_rotation_},
      n_symmetry_rotations_{src.n_symmetry_rotations_},
      uncertainty_threshold_{src.uncertainty_threshold_} {
  QUILL_LOG_INFO(m_log, "[MetatomicPotential] Cloned loaded model (no disk)");
}

Member Function Documentation

clonePotential()

std::shared_ptr< eonc::Potential > MetatomicPotential::clonePotential ( ) const

nodiscardoverridevirtual

Independent instance that does not reload from disk.

nullptr means the caller should use makePotential().

Reimplemented from eonc::Potential.

Definition at line 323 of file MetatomicPotential.cpp.

                                                                      {
  return std::shared_ptr<eonc::Potential>(
      new MetatomicPotential(*this, CloneTag{}));
}

computeNeighbors()

metatensor_torch::TensorBlock MetatomicPotential::computeNeighbors ( metatomic_torch::NeighborListOptions request, long nAtoms, const double * positions, const double * box, const bool periodic[3] )

private

Computes neighbor list using the vesin library.

This function calls vesin_neighbors to build a neighbor list based on the model's requirements (cutoff, full/half list) and converts it into a metatensor::TensorBlock suitable for metatomic.

Parameters

request	The neighbor list options requested by the model.
nAtoms	The number of atoms in the system.
positions	Pointer to the atomic positions array.
box	Pointer to the simulation box matrix.

Returns

A metatensor_torch::TensorBlock containing the neighbor list.

Definition at line 571 of file MetatomicPotential.cpp.

                                                                        {
 
  auto cutoff = request->engine_cutoff(m_metatomic_opts.length_unit);
 
  // Zero-init so vesin 0.6 skin/n_threads stay 0 when compiling against 0.6
  // headers (must match linked libvesin — pin vesin>=0.6 for metatomic builds).
  VesinOptions options{};
  options.cutoff = cutoff;
  options.full = request->full_list();
  options.sorted = false;
  options.return_shifts = true;
  options.return_distances = false; // we don't need distances
  options.return_vectors = true;    // metatomic uses vectors for autograd
 
  VesinNeighborList *vesin_neighbor_list = new VesinNeighborList();
 
  VesinDevice cpu{VesinCPU, 0};
  const char *error_message = nullptr;
  int status = vesin_neighbors(reinterpret_cast<const double (*)[3]>(positions),
                               static_cast<size_t>(nAtoms),
                               reinterpret_cast<const double (*)[3]>(box),
                               const_cast<bool *>(periodic), cpu, options,
                               vesin_neighbor_list, &error_message);
 
  if (status != EXIT_SUCCESS) {
    std::string err_str = "vesin_neighbors failed";
    if (error_message != nullptr) {
      err_str += ": " + std::string(error_message);
    } else {
      err_str += " (no message; vesin header/lib ABI mismatch? need vesin>=0.6 "
                 "with matching engine)";
    }
    delete vesin_neighbor_list;
    throw std::runtime_error(err_str);
  }
 
  // Convert from vesin to metatomic format
  auto n_pairs = static_cast<int64_t>(vesin_neighbor_list->length);
  auto labels_options_cpu =
      torch::TensorOptions().dtype(torch::kInt32).device(torch::kCPU);
 
  auto pair_samples_values = torch::empty({n_pairs, 5}, labels_options_cpu);
  auto pair_samples_values_ptr = pair_samples_values.accessor<int32_t, 2>();
  for (int64_t i = 0; i < n_pairs; i++) {
    pair_samples_values_ptr[i][0] =
        static_cast<int32_t>(vesin_neighbor_list->pairs[i][0]);
    pair_samples_values_ptr[i][1] =
        static_cast<int32_t>(vesin_neighbor_list->pairs[i][1]);
    pair_samples_values_ptr[i][2] = vesin_neighbor_list->shifts[i][0];
    pair_samples_values_ptr[i][3] = vesin_neighbor_list->shifts[i][1];
    pair_samples_values_ptr[i][4] = vesin_neighbor_list->shifts[i][2];
  }
 
  // Custom deleter to free vesin's memory when the torch tensor is destroyed
  auto deleter = [=](void *) {
    vesin_free(vesin_neighbor_list);
    delete vesin_neighbor_list;
  };
 
  auto pair_vectors = torch::from_blob(
      vesin_neighbor_list->vectors, {n_pairs, 3, 1}, deleter,
      torch::TensorOptions().dtype(torch::kFloat64).device(torch::kCPU));
 
  auto neighbor_samples = torch::make_intrusive<metatensor_torch::LabelsHolder>(
      std::vector<std::string>{"first_atom", "second_atom", "cell_shift_a",
                               "cell_shift_b", "cell_shift_c"},
      pair_samples_values.to(this->device_));
 
  auto labels_options_dev =
      torch::TensorOptions().dtype(torch::kInt32).device(this->device_);
  auto neighbor_component =
      torch::make_intrusive<metatensor_torch::LabelsHolder>(
          "xyz", torch::tensor({0, 1, 2}, labels_options_dev).reshape({3, 1}));
  auto neighbor_properties =
      torch::make_intrusive<metatensor_torch::LabelsHolder>(
          "distance", torch::zeros({1, 1}, labels_options_dev));
 
  return torch::make_intrusive<metatensor_torch::TensorBlockHolder>(
      pair_vectors.to(this->dtype_).to(this->device_), neighbor_samples,
      std::vector<metatensor_torch::Labels>{neighbor_component},
      neighbor_properties);
}

force()

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

overridevirtual

Calculates the energy and forces for a given atomic configuration.

This is the core method of the potential. It takes the current atomic positions, builds the necessary data structures for metatomic, executes the model to get the potential energy, and uses PyTorch's autograd to compute forces.

Parameters

nAtoms	Number of atoms.
positions	Flat array of atomic positions (size nAtoms * 3).
atomicNrs	Flat array of atomic numbers (size nAtoms), used for consistency checks.
forces	Flat array to store the calculated forces (size nAtoms * 3).
energy	Pointer to a double to store the calculated potential energy.
variance	Pointer to a double to store the variance of the energy (currently unused, set to NULL).
box	The simulation box vectors (3x3 matrix).

Implements eonc::Potential.

Definition at line 352 of file MetatomicPotential.cpp.

                                                  {
  // Serialize concurrent calls -- PyTorch model inference on the same
  // Module instance is not thread-safe
  std::lock_guard<std::mutex> lock(inference_mutex_);
 
  eonc::FPEHandler fpeh;
  fpeh.eat_fpe();
 
  if (!atomicNrs) {
    throw std::runtime_error(
        "[MetatomicPotential] `atomicNrs` must be provided.");
  }
 
  const bool use_rotation =
      this->random_rotation_ || this->n_symmetry_rotations_ > 0;
  // n_symmetry_rotations averages; random_rotation alone is one rotated eval
  const long n_passes =
      this->n_symmetry_rotations_ > 0 ? this->n_symmetry_rotations_ : 1;
 
  auto f64_options =
      torch::TensorOptions().dtype(torch::kFloat64).device(torch::kCPU);
  std::vector<int32_t> types_vec(atomicNrs, atomicNrs + nAtoms);
  auto atomic_types_cpu =
      torch::tensor(types_vec, torch::TensorOptions().dtype(torch::kInt32));
 
  double energy_acc = 0.0;
  auto forces_acc = torch::zeros({nAtoms, 3}, f64_options);
  bool variance_set = false;
 
  for (long i_pass = 0; i_pass < n_passes; ++i_pass) {
    torch::Tensor R = torch::eye(
        3, torch::TensorOptions().dtype(this->dtype_).device(this->device_));
    if (use_rotation) {
      R = random_so3(this->device_, this->dtype_);
    }
    // R is applied to row vectors: pos' = pos @ R^T  (equiv. R @ pos for cols)
    auto R_cpu = R.to(torch::kCPU).to(torch::kFloat64);
    auto R_T = R.transpose(0, 1);
 
    auto pos_cpu = torch::from_blob(const_cast<double *>(positions),
                                    {nAtoms, 3}, f64_options)
                       .clone();
    auto cell_cpu =
        torch::from_blob(const_cast<double *>(box), {3, 3}, f64_options)
            .clone();
    if (use_rotation) {
      pos_cpu = pos_cpu.matmul(R_cpu.transpose(0, 1));
      // Rotate cell vectors (rows) the same way
      cell_cpu = cell_cpu.matmul(R_cpu.transpose(0, 1));
    }
 
    std::vector<double> pos_buf(static_cast<size_t>(nAtoms) * 3);
    std::vector<double> cell_buf(9);
    std::memcpy(pos_buf.data(), pos_cpu.contiguous().data_ptr<double>(),
                pos_buf.size() * sizeof(double));
    std::memcpy(cell_buf.data(), cell_cpu.contiguous().data_ptr<double>(),
                9 * sizeof(double));
 
    auto torch_positions =
        torch::from_blob(pos_buf.data(), {nAtoms, 3}, f64_options)
            .to(this->dtype_)
            .to(this->device_)
            .set_requires_grad(!this->non_conservative_);
 
    auto torch_cell = torch::from_blob(cell_buf.data(), {3, 3}, f64_options)
                          .to(this->dtype_)
                          .to(this->device_);
 
    auto cell_norms = torch::norm(torch_cell, 2, /*dim=*/1);
    auto torch_pbc = cell_norms.abs() > 1e-9;
    bool periodic[3] = {torch_pbc[0].item<bool>(), torch_pbc[1].item<bool>(),
                        torch_pbc[2].item<bool>()};
 
    auto atomic_types = atomic_types_cpu.to(this->device_);
 
    auto system = torch::make_intrusive<metatomic_torch::SystemHolder>(
        atomic_types, torch_positions, torch_cell, torch_pbc);
 
    for (const auto &request : this->nl_requests_) {
      auto neighbors = this->computeNeighbors(request, nAtoms, pos_buf.data(),
                                              cell_buf.data(), periodic);
      metatomic_torch::register_autograd_neighbors(system, neighbors,
                                                   this->check_consistency_);
      system->add_neighbor_list(request, neighbors);
    }
 
    torch::Tensor forces_tensor;
    try {
      auto ivalue_output = this->model_.forward({
          std::vector<metatomic_torch::System>{system},
          evaluations_options_,
          this->check_consistency_,
      });
      auto dict_output = ivalue_output.toGenericDict();
      auto output_map = dict_output.at(this->energy_key_)
                            .toCustomClass<metatensor_torch::TensorMapHolder>();
 
      if (this->uncertainty_threshold_ > 0 && i_pass == 0) {
        try {
          if (dict_output.contains(this->energy_uncertainty_key_)) {
            auto uncertainty_map =
                dict_output.at(this->energy_uncertainty_key_)
                    .toCustomClass<metatensor_torch::TensorMapHolder>();
            auto uncertainty_block =
                metatensor_torch::TensorMapHolder::block_by_id(uncertainty_map,
                                                               0);
            auto flat_uncertainty =
                uncertainty_block->values().reshape({-1}).to(torch::kCPU);
            if (variance != nullptr && flat_uncertainty.numel() > 0) {
              try {
                *variance =
                    flat_uncertainty.to(torch::kFloat64).mean().item<double>();
                variance_set = true;
              } catch (...) {
                QUILL_LOG_DEBUG(m_log,
                                "[MetatomicPotential] Failed to compute mean "
                                "uncertainty for variance.");
              }
            }
            auto atoms_above_threshold =
                flat_uncertainty > this->uncertainty_threshold_;
            if (torch::any(atoms_above_threshold).item<bool>()) {
              auto samples = uncertainty_block->samples();
              auto atom_indices_all = samples->column("atom").to(torch::kCPU);
              auto atom_indices_above =
                  atom_indices_all.index({atoms_above_threshold});
              std::ostringstream ss;
              ss << "atoms at index [";
              auto n_report = std::min<int64_t>(10, atom_indices_above.size(0));
              for (int64_t i = 0; i < n_report; ++i) {
                if (i > 0)
                  ss << ", ";
                ss << atom_indices_above[i].item<int32_t>();
              }
              ss << "]";
              if (atom_indices_above.size(0) > n_report) {
                ss << " and " << (atom_indices_above.size(0) - n_report)
                   << " more";
              }
              QUILL_LOG_WARNING(
                  m_log,
                  "[MetatomicPotential] The uncertainty on atomic energies for "
                  "{} are larger than the threshold of {}. (Key: {}) Be "
                  "careful "
                  "when analyzing the results, and consider retraining the "
                  "model to better describe these configurations.",
                  ss.str(), this->uncertainty_threshold_,
                  this->energy_uncertainty_key_);
            }
          }
        } catch (const std::exception &e) {
          QUILL_LOG_WARNING(m_log,
                            "[MetatomicPotential] Failed to check {}: {}",
                            this->energy_uncertainty_key_, e.what());
        }
      }
 
      auto energy_block =
          metatensor_torch::TensorMapHolder::block_by_id(output_map, 0);
      auto energy_tensor = energy_block->values();
      energy_acc += energy_tensor.sum().item<double>();
 
      if (this->non_conservative_ && !this->nc_forces_key_.empty()) {
        auto nc_map = dict_output.at(this->nc_forces_key_)
                          .toCustomClass<metatensor_torch::TensorMapHolder>();
        auto nc_block =
            metatensor_torch::TensorMapHolder::block_by_id(nc_map, 0);
        auto nc_vals = nc_block->values();
        if (nc_vals.numel() != nAtoms * 3) {
          throw std::runtime_error("[MetatomicPotential] NC force block has " +
                                   std::to_string(nc_vals.numel()) +
                                   " values, expected " +
                                   std::to_string(nAtoms * 3));
        }
        auto nc_samples = nc_block->samples();
        if (nc_samples->size() > 0 && nc_samples->names().size() > 1) {
          auto atom_col = nc_samples->column("atom").to(torch::kCPU);
          if (atom_col.size(0) != nAtoms) {
            throw std::runtime_error(
                "[MetatomicPotential] NC force samples atom count mismatch");
          }
        }
        forces_tensor =
            nc_vals.reshape({nAtoms, 3}).to(torch::kCPU).to(torch::kFloat64);
      } else {
        energy_tensor.backward(torch::ones_like(energy_tensor));
        auto positions_grad = system->positions().grad();
        forces_tensor = (-positions_grad).to(torch::kCPU).to(torch::kFloat64);
      }
    } catch (const std::exception &e) {
      QUILL_LOG_ERROR(m_log, "[MetatomicPotential] Model evaluation failed: {}",
                      e.what());
      throw;
    }
 
    // Rotate forces back to original frame: F = F' @ R  (since pos' = pos @
    // R^T)
    if (use_rotation) {
      forces_tensor = forces_tensor.matmul(R_cpu);
    }
    forces_acc += forces_tensor;
  }
 
  const double inv_n = 1.0 / static_cast<double>(n_passes);
  *energy = energy_acc * inv_n;
  forces_acc = forces_acc * inv_n;
  (void)variance_set;
 
  std::memcpy(forces, forces_acc.contiguous().data_ptr<double>(),
              nAtoms * 3 * sizeof(double));
 
  fpeh.restore_fpe();
}

forceBatch()

void MetatomicPotential::forceBatch ( long nSystems, long nAtoms, const double *const * positions, const int *const * atomicNrs, double *const * forces, double * energies, double * variances, const double *const * boxes )

overridevirtual

Reimplemented from eonc::Potential.

Definition at line 662 of file MetatomicPotential.cpp.

                                                                {
  if (nSystems > 1) {
    try {
      forceBatchNative(nSystems, nAtoms, positions, atomicNrs, forces, energies,
                       variances, boxes);
      forceCallCounter += nSystems;
      eonc::PotRegistry::get().on_force_call(ptype);
      return;
    } catch (const std::exception &e) {
      QUILL_LOG_WARNING(m_log,
                        "[MetatomicPotential] batched forward failed ({}); "
                        "falling back to sequential force()",
                        e.what());
    }
  }
  for (long s = 0; s < nSystems; s++) {
    double var = 0;
    force(nAtoms, positions[s], atomicNrs[s], forces[s], &energies[s], &var,
          boxes[s]);
    if (variances)
      variances[s] = var;
    forceCallCounter++;
    eonc::PotRegistry::get().on_force_call(ptype);
  }
}

forceBatchNative()

void MetatomicPotential::forceBatchNative ( long nSystems, long nAtoms, const double *const * positions, const int *const * atomicNrs, double *const * forces, double * energies, double * variances, const double *const * boxes )

private

Definition at line 693 of file MetatomicPotential.cpp.

                                                                      {
  std::lock_guard<std::mutex> lock(inference_mutex_);
 
  eonc::FPEHandler fpeh;
  fpeh.eat_fpe();
 
  auto f64_options =
      torch::TensorOptions().dtype(torch::kFloat64).device(torch::kCPU);
 
  std::vector<metatomic_torch::System> systems;
  std::vector<torch::Tensor> pos_tensors;
  systems.reserve(static_cast<size_t>(nSystems));
  pos_tensors.reserve(static_cast<size_t>(nSystems));
 
  for (long s = 0; s < nSystems; s++) {
    auto torch_positions = torch::from_blob(const_cast<double *>(positions[s]),
                                            {nAtoms, 3}, f64_options)
                               .to(this->dtype_)
                               .to(this->device_)
                               .set_requires_grad(true);
    pos_tensors.push_back(torch_positions);
 
    auto torch_cell =
        torch::from_blob(const_cast<double *>(boxes[s]), {3, 3}, f64_options)
            .to(this->dtype_)
            .to(this->device_);
 
    auto cell_norms = torch::norm(torch_cell, 2, /*dim=*/1);
    auto torch_pbc = cell_norms.abs() > 1e-9;
    bool periodic[3] = {torch_pbc[0].item<bool>(), torch_pbc[1].item<bool>(),
                        torch_pbc[2].item<bool>()};
 
    if (!atomicNrs[s]) {
      throw std::runtime_error(
          "[MetatomicPotential] `atomicNrs` must be provided.");
    }
    std::vector<int32_t> types_vec(atomicNrs[s], atomicNrs[s] + nAtoms);
    auto atomic_types =
        torch::tensor(types_vec, torch::TensorOptions().dtype(torch::kInt32))
            .to(this->device_);
 
    auto system = torch::make_intrusive<metatomic_torch::SystemHolder>(
        atomic_types, torch_positions, torch_cell, torch_pbc);
 
    // Compute and register neighbor lists for this system
    for (const auto &request : this->nl_requests_) {
      auto neighbors = this->computeNeighbors(request, nAtoms, positions[s],
                                              boxes[s], periodic);
      metatomic_torch::register_autograd_neighbors(system, neighbors,
                                                   this->check_consistency_);
      system->add_neighbor_list(request, neighbors);
    }
 
    systems.push_back(system);
  }
 
  // Single batched forward pass
  metatensor_torch::TensorMap output_map;
  try {
    auto ivalue_output = this->model_.forward({
        systems,
        evaluations_options_,
        this->check_consistency_,
    });
    auto dict_output = ivalue_output.toGenericDict();
    output_map = dict_output.at(this->energy_key_)
                     .toCustomClass<metatensor_torch::TensorMapHolder>();
  } catch (const std::exception &e) {
    QUILL_LOG_ERROR(m_log,
                    "[MetatomicPotential] Batched model evaluation failed: {}",
                    e.what());
    throw;
  }
 
  // Extract per-system energies from the output TensorMap.
  // For per-atom output: samples have ["system", "atom"] dimensions.
  // For system-level output: samples have ["system"] dimension.
  // In both cases, sum over all non-system dimensions to get per-system energy.
  auto energy_block =
      metatensor_torch::TensorMapHolder::block_by_id(output_map, 0);
  auto energy_values = energy_block->values();
  auto samples = energy_block->samples();
 
  // Check if this is per-atom or per-system output
  bool per_atom = samples->size() > 0 && samples->names().size() > 1;
 
  if (per_atom) {
    // Per-atom output: sum energies by system index
    auto system_col = samples->column("system").to(torch::kCPU);
    auto flat_energies = energy_values.reshape({-1}).to(torch::kCPU);
 
    // Sum per-system energies for output
    for (long s = 0; s < nSystems; s++) {
      auto mask = (system_col == s);
      energies[s] =
          flat_energies.index({mask}).sum().to(torch::kFloat64).item<double>();
    }
 
    // Backward: sum ALL energies, single backward call.
    // Each system's positions.grad() gets only its own contribution
    // because energy_i depends only on positions_i.
    energy_values.sum().backward();
  } else {
    // System-level output: values shape is (nSystems, 1) or similar
    auto cpu_energies = energy_values.to(torch::kCPU).to(torch::kFloat64);
    for (long s = 0; s < nSystems; s++) {
      energies[s] = cpu_energies[s].sum().item<double>();
    }
    energy_values.backward(torch::ones_like(energy_values));
  }
 
  // Extract per-system forces from position gradients
  for (long s = 0; s < nSystems; s++) {
    auto positions_grad = pos_tensors[s].grad();
    auto forces_tensor = -positions_grad.to(torch::kCPU).to(torch::kFloat64);
    std::memcpy(forces[s], forces_tensor.contiguous().data_ptr<double>(),
                nAtoms * 3 * sizeof(double));
  }
 
  // Variances: not yet supported in batched path
  if (variances) {
    for (long s = 0; s < nSystems; s++) {
      variances[s] = 0.0;
    }
  }
 
  fpeh.restore_fpe();
}

isThreadSafe()

bool MetatomicPotential::isThreadSafe ( ) const

inlinenodiscardoverridevirtualnoexcept

Single shared instance, serialized via mutex.

Sequential evaluation through computePotential() ensures correct force counting, removeNetForce, and PotRegistry bookkeeping. JIT profiling is disabled so all calls on the same instance produce deterministic results.

Reimplemented from eonc::Potential.

Definition at line 140 of file MetatomicPotential.h.

{ return false; }

needsPerImageInstance()

bool MetatomicPotential::needsPerImageInstance ( ) const

inlinenodiscardoverridevirtualnoexcept

Whether NEB should create separate Potential instances per image for true parallel force evaluation.

When true, NEB calls makePotential() once per image instead of sharing one instance. Override in potentials that use internal mutexes (e.g. MetatomicPotential).

Reimplemented from eonc::Potential.

Definition at line 141 of file MetatomicPotential.h.

                                                                     {
    return false;
  }

supportsBatchEvaluation()

bool MetatomicPotential::supportsBatchEvaluation ( ) const

inlinenodiscardoverridevirtualnoexcept

Batched evaluation via single shared instance.

Processes each system sequentially through the same model (identical to N force() calls but bypasses Matter::computePotential overhead). Forces, energies, and forceCallCounter are handled correctly.

Reimplemented from eonc::Potential.

Definition at line 149 of file MetatomicPotential.h.

                                                                       {
    return true;
  }

Member Data Documentation

capabilities_

metatomic_torch::ModelCapabilities MetatomicPotential::capabilities_

private

Definition at line 60 of file MetatomicPotential.h.

check_consistency_

bool MetatomicPotential::check_consistency_

private

Definition at line 67 of file MetatomicPotential.h.

device_

torch::Device MetatomicPotential::device_

private

Definition at line 66 of file MetatomicPotential.h.

device_type_

c10::DeviceType MetatomicPotential::device_type_

private

Definition at line 65 of file MetatomicPotential.h.

dtype_

torch::ScalarType MetatomicPotential::dtype_

private

Definition at line 64 of file MetatomicPotential.h.

energy_key_

std::string MetatomicPotential::energy_key_

private

Definition at line 69 of file MetatomicPotential.h.

energy_uncertainty_key_

std::string MetatomicPotential::energy_uncertainty_key_

private

Definition at line 70 of file MetatomicPotential.h.

evaluations_options_

metatomic_torch::ModelEvaluationOptions MetatomicPotential::evaluations_options_

private

Definition at line 62 of file MetatomicPotential.h.

inference_mutex_

std::mutex MetatomicPotential::inference_mutex_

mutableprivate

Definition at line 165 of file MetatomicPotential.h.

m_log

eonc::log::Scoped MetatomicPotential::m_log

private

Definition at line 56 of file MetatomicPotential.h.

m_metatomic_opts

eonc::Parameters::metatomic_options_t MetatomicPotential::m_metatomic_opts

private

Definition at line 57 of file MetatomicPotential.h.

model_

metatensor_torch::Module MetatomicPotential::model_

private

Definition at line 59 of file MetatomicPotential.h.

n_symmetry_rotations_

long MetatomicPotential::n_symmetry_rotations_ {0}

private

Definition at line 74 of file MetatomicPotential.h.

{0};

nc_forces_key_

std::string MetatomicPotential::nc_forces_key_

private

Definition at line 71 of file MetatomicPotential.h.

nl_requests_

std::vector<metatomic_torch::NeighborListOptions> MetatomicPotential::nl_requests_

private

Definition at line 61 of file MetatomicPotential.h.

non_conservative_

bool MetatomicPotential::non_conservative_ {false}

private

Definition at line 72 of file MetatomicPotential.h.

{false};

random_rotation_

bool MetatomicPotential::random_rotation_ {false}

private

Definition at line 73 of file MetatomicPotential.h.

{false};

uncertainty_threshold_

double MetatomicPotential::uncertainty_threshold_ {-1.0}

private

Definition at line 77 of file MetatomicPotential.h.

{-1.0};

---

The documentation for this class was generated from the following files:

• /home/runner/work/eOn/eOn/include/eon/potentials/Metatomic/MetatomicPotential.h
• /home/runner/work/eOn/eOn/client/potentials/Metatomic/MetatomicPotential.cpp