Skip to content

Commit

Permalink
Move code for vesin in a separate folder
Browse files Browse the repository at this point in the history 
This allows building it in the usual way inside PLUMED
while not having to modify the code from upstream
  • Loading branch information
@Luthaf
Luthaf committed May 27, 2024
1 parent 6ef31c6 commit b6bf849
Show file tree
Hide file tree
Showing 8 changed files with 242 additions and 115 deletions.
7 changes: 4 additions & 3 deletions src/maketools/codecheck
View file
Original file line number Diff line number Diff line change
Expand Up @@ -65,7 +65,7 @@ CLEAN_FILE="$TMPDIR/codecheck.clean"

echo "++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++"
echo " In the following analysis we exclude:"
echo " src/molfile src/lapack src/blas src/lepton src/asmjit src/xdrfile src/small_vector"
echo " src/molfile src/lapack src/blas src/lepton src/asmjit src/xdrfile src/small_vector src/metatensor/external"
echo " since they are full of false positives"
cat codecheck.log |
gawk '{
Expand All @@ -77,6 +77,7 @@ cat codecheck.log |
if(match($0,"[[]lepton/")) next;
if(match($0,"[[]asmjit/")) next;
if(match($0,"[[]xdrfile/")) next;
if(match($0,"[[]metatensor/external/")) next;
} else {
# with cppcheck we exclude molfile lapack and blas
if(match($0,"[[]molfile/")) next;
Expand All @@ -86,6 +87,7 @@ cat codecheck.log |
if(match($0,"[[]asmjit/")) next;
if(match($0,"[[]xdrfile/")) next;
if(match($0,"[[]small_vector/")) next;
if(match($0,"[[]metatensor/external/")) next;
}
print
}' > "${CLEAN_FILE}"
Expand All @@ -102,7 +104,7 @@ echo "++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
## (the "true" command is necessary so that the script does not fail if some string is not found)
grep -v "(style)" "${CLEAN_FILE}" | grep -v "(information)" | grep -v :duplInheritedMember: >> "${FATAL_FILE}" || true
count=$(cat "${FATAL_FILE}" | wc -l)
echo
echo
echo "++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++"
if ((count==0))
then
Expand All @@ -121,4 +123,3 @@ echo "++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
echo

exit 1

2 changes: 1 addition & 1 deletion src/metatensor/Makefile
View file
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
USE=core
USE=core tools

#generic makefile
include ../maketools/make.module
7 changes: 7 additions & 0 deletions src/metatensor/external/README.md
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1,7 @@
This directory contains the code for the library used to compute neighbor lists
in the metatensor interface to PLUMED. This library lives in
https://github.com/Luthaf/vesin.

If you want to update the library after a bug fix or improvement upstream, you
should run the `create-single-cpp.py` script in this repo, and copy over the
`include/vesin.h` file and `vesin-single-build.cpp` file.
1 change: 0 additions & 1 deletion src/metatensor/vesin_single_build.cpp → ...etatensor/external/vesin-single-build.cpp
View file
Original file line number Diff line number Diff line change
@@ -1,4 +1,3 @@
/*INDENT-OFF*/
#include <cassert>
#include <cstdlib>
#include <cstring>
Expand Down
29 changes: 25 additions & 4 deletions src/metatensor/vesin.h → src/metatensor/external/vesin.h
View file
Original file line number Diff line number Diff line change
@@ -1,10 +1,31 @@
/*INDENT-OFF*/
#ifndef VESIN_H
#define VESIN_H

#include <stddef.h>
#include <stdint.h>

#if defined(VESIN_SHARED)
#if defined(VESIN_EXPORTS)
#if defined(__clang__) || defined(__GNUC__) || defined(__GNUG__)
#define VESIN_API __attribute__((visibility("default")))
#elif defined(_MSC_VER)
#define VESIN_API __declspec(dllexport)
#else
#define VESIN_API
#endif
#else
#if defined(__clang__) || defined(__GNUC__) || defined(__GNUG__)
#define VESIN_API __attribute__((visibility("default")))
#elif defined(_MSC_VER)
#define VESIN_API __declspec(dllimport)
#else
#define VESIN_API
#endif
#endif
#else
#define VESIN_API
#endif

#ifdef __cplusplus
extern "C" {
#endif
Expand Down Expand Up @@ -55,7 +76,7 @@ enum VesinDevice {
///
/// Under periodic boundary conditions, two atoms can be part of multiple pairs,
/// each pair having a different periodic shift.
typedef struct VesinNeighborList {
typedef struct VESIN_API VesinNeighborList {
#ifdef __cplusplus
VesinNeighborList():
length(0),
Expand Down Expand Up @@ -91,7 +112,7 @@ typedef struct VesinNeighborList {

/// Free all allocated memory inside a `VesinNeighborList`, according the it's
/// `device`.
void vesin_free(VesinNeighborList* neighbors);
void VESIN_API vesin_free(VesinNeighborList* neighbors);

/// Compute a neighbor list.
///
Expand All @@ -114,7 +135,7 @@ void vesin_free(VesinNeighborList* neighbors);
/// @param error_message Pointer to a `char*` that wil be set to the error
/// message if this function fails. This does not need to be freed when no
/// longer needed.
int vesin_neighbors(
int VESIN_API vesin_neighbors(
const double (*points)[3],
size_t n_points,
const double box[3][3],
Expand Down
123 changes: 17 additions & 106 deletions src/metatensor/metatensor.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -159,8 +159,6 @@ class MetatensorPlumedAction: public ActionAtomistic, public ActionWithValue {

#else

#include <type_traits>

#pragma GCC diagnostic push
#pragma GCC diagnostic ignored "-Wpedantic"
#pragma GCC diagnostic ignored "-Wunused-parameter"
Expand All @@ -184,7 +182,7 @@ class MetatensorPlumedAction: public ActionAtomistic, public ActionWithValue {
#include <metatensor/torch.hpp>
#include <metatensor/torch/atomistic.hpp>

#include "vesin.h"
#include "neighbors.h"


namespace PLMD {
Expand Down Expand Up @@ -212,12 +210,6 @@ class MetatensorPlumedAction: public ActionAtomistic, public ActionWithValue {
private:
// fill this->system_ according to the current PLUMED data
void createSystem();
// compute a neighbor list following metatensor format, using data from PLUMED
metatensor_torch::TorchTensorBlock computeNeighbors(
metatensor_torch::NeighborListOptions request,
const std::vector<PLMD::Vector>& positions,
const PLMD::Tensor& cell
);

// execute the model for the given system
metatensor_torch::TorchTensorBlock executeModel(metatensor_torch::System system);
Expand Down Expand Up @@ -498,7 +490,14 @@ MetatensorPlumedAction::MetatensorPlumedAction(const ActionOptions& options):
model_length_unit.c_str()
);

auto neighbors = this->computeNeighbors(request, {PLMD::Vector(0, 0, 0)}, PLMD::Tensor(0, 0, 0, 0, 0, 0, 0, 0, 0));
auto neighbors = PLMD::metatensor::compute_neighbors(
request,
{PLMD::Vector(0, 0, 0)},
PLMD::Tensor(0, 0, 0, 0, 0, 0, 0, 0, 0),
this->dtype_,
this->device_,
this->getUnits().getLengthString()
);
metatensor_torch::register_autograd_neighbors(dummy_system, neighbors, this->check_consistency_);
dummy_system->add_neighbor_list(request, neighbors);
}
Expand Down Expand Up @@ -589,15 +588,11 @@ void MetatensorPlumedAction::createSystem() {
plumed_merror(oss.str());
}

// this->getTotAtoms()

const auto& cell = this->getPbc().getBox();

auto cpu_f64_tensor = torch::TensorOptions().dtype(torch::kFloat64).device(torch::kCPU);
auto torch_cell = torch::zeros({3, 3}, cpu_f64_tensor);

// TODO: check if cell is stored in row or column major order
// TODO: check if cell is zero for non-periodic systems
torch_cell[0][0] = cell(0, 0);
torch_cell[0][1] = cell(0, 1);
torch_cell[0][2] = cell(0, 2);
Expand Down Expand Up @@ -642,103 +637,20 @@ void MetatensorPlumedAction::createSystem() {
// compute the neighbors list requested by the model, and register them with
// the system
for (auto request: this->nl_requests_) {
auto neighbors = this->computeNeighbors(request, positions, cell);
auto neighbors = PLMD::metatensor::compute_neighbors(
request,
positions,
cell,
this->dtype_,
this->device_,
this->getUnits().getLengthString()
);
metatensor_torch::register_autograd_neighbors(this->system_, neighbors, this->check_consistency_);
this->system_->add_neighbor_list(request, neighbors);
}
}


metatensor_torch::TorchTensorBlock MetatensorPlumedAction::computeNeighbors(
metatensor_torch::NeighborListOptions request,
const std::vector<PLMD::Vector>& positions,
const PLMD::Tensor& cell
) {
auto labels_options = 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).reshape({3, 1})
);
auto neighbor_properties = torch::make_intrusive<metatensor_torch::LabelsHolder>(
"distance", torch::zeros({1, 1}, labels_options)
);

auto cutoff = request->engine_cutoff(this->getUnits().getLengthString());

auto non_periodic = (
cell(0, 0) == 0.0 && cell(0, 1) == 0.0 && cell(0, 2) == 0.0 &&
cell(1, 0) == 0.0 && cell(1, 1) == 0.0 && cell(1, 2) == 0.0 &&
cell(2, 0) == 0.0 && cell(2, 2) == 0.0 && cell(2, 2) == 0.0
);

// use https://github.com/Luthaf/vesin to compute the requested neighbor
// lists since we can not get these from PLUMED
VesinOptions options;
options.cutoff = cutoff;
options.full = request->full_list();
options.return_shifts = true;
options.return_distances = false;
options.return_vectors = true;

VesinNeighborList* vesin_neighbor_list = new VesinNeighborList();
memset(vesin_neighbor_list, 0, sizeof(VesinNeighborList));

const char* error_message = NULL;
int status = vesin_neighbors(
reinterpret_cast<const double (*)[3]>(positions.data()),
positions.size(),
reinterpret_cast<const double (*)[3]>(&cell(0, 0)),
!non_periodic,
VesinCPU,
options,
vesin_neighbor_list,
&error_message
);

if (status != EXIT_SUCCESS) {
plumed_merror(
"failed to compute neighbor list (cutoff=" + std::to_string(cutoff) +
", full=" + (request->full_list() ? "true" : "false") + "): " + error_message
);
}

// transform from vesin to metatensor format
auto n_pairs = static_cast<int64_t>(vesin_neighbor_list->length);

auto pair_vectors = torch::from_blob(
vesin_neighbor_list->vectors,
{n_pairs, 3, 1},
/*deleter*/ [=](void*) {
vesin_free(vesin_neighbor_list);
delete vesin_neighbor_list;
},
torch::TensorOptions().dtype(torch::kFloat64).device(torch::kCPU)
);

auto pair_samples_values = torch::zeros({n_pairs, 5}, labels_options.device(torch::kCPU));
for (unsigned i=0; i<n_pairs; i++) {
pair_samples_values[i][0] = static_cast<int32_t>(vesin_neighbor_list->pairs[i][0]);
pair_samples_values[i][1] = static_cast<int32_t>(vesin_neighbor_list->pairs[i][1]);
pair_samples_values[i][2] = vesin_neighbor_list->shifts[i][0];
pair_samples_values[i][3] = vesin_neighbor_list->shifts[i][1];
pair_samples_values[i][4] = vesin_neighbor_list->shifts[i][2];
}

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 neighbors = torch::make_intrusive<metatensor_torch::TensorBlockHolder>(
pair_vectors.to(this->dtype_).to(this->device_),
neighbor_samples,
std::vector<metatensor_torch::TorchLabels>{neighbor_component},
neighbor_properties
);

return neighbors;
}

metatensor_torch::TorchTensorBlock MetatensorPlumedAction::executeModel(metatensor_torch::System system) {
try {
auto ivalue_output = this->model_.forward({
Expand Down Expand Up @@ -833,7 +745,6 @@ void MetatensorPlumedAction::calculate() {
}
}


void MetatensorPlumedAction::apply() {
auto* value = this->getPntrToComponent(0);
if (!value->forcesWereAdded()) {
Expand Down
Loading

0 comments on commit b6bf849

Please sign in to comment.