Line data Source code
1 : /* +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
2 : Copyright (c) 2024 The METATENSOR code team
3 : (see the PEOPLE-METATENSOR file at the root of this folder for a list of names)
4 :
5 : See https://docs.metatensor.org/latest/ for more information about the
6 : metatensor package that this module allows you to call from PLUMED.
7 :
8 : This file is part of METATENSOR-PLUMED module.
9 :
10 : The METATENSOR-PLUMED module is free software: you can redistribute it and/or modify
11 : it under the terms of the GNU Lesser General Public License as published by
12 : the Free Software Foundation, either version 3 of the License, or
13 : (at your option) any later version.
14 :
15 : The METATENSOR-PLUMED module is distributed in the hope that it will be useful,
16 : but WITHOUT ANY WARRANTY; without even the implied warranty of
17 : MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
18 : GNU Lesser General Public License for more details.
19 :
20 : You should have received a copy of the GNU Lesser General Public License
21 : along with the METATENSOR-PLUMED module. If not, see <http://www.gnu.org/licenses/>.
22 : +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++ */
23 :
24 : #include "core/ActionAtomistic.h"
25 : #include "core/ActionWithValue.h"
26 : #include "core/ActionRegister.h"
27 : #include "core/PlumedMain.h"
28 :
29 : //+PLUMEDOC METATENSORMOD_COLVAR METATENSOR
30 : /*
31 : Use arbitrary machine learning models as collective variables.
32 :
33 : Note that this action requires the metatensor-torch library. Check the
34 : instructions in the \ref METATENSORMOD page to enable this module.
35 :
36 : This action enables the use of fully custom machine learning models — based on
37 : the [metatensor atomistic models][mts_models] interface — as collective
38 : variables in PLUMED. Such machine learning model are typically written and
39 : customized using Python code, and then exported to run within PLUMED as
40 : [TorchScript], which is a subset of Python that can be executed by the C++ torch
41 : library.
42 :
43 : Metatensor offers a way to define such models and pass data from PLUMED (or any
44 : other simulation engine) to the model and back. For more information on how to
45 : define such model, have a look at the [corresponding tutorials][mts_tutorials],
46 : or at the code in `regtest/metatensor/`. Each of the Python scripts in this
47 : directory defines a custom machine learning CV that can be used with PLUMED.
48 :
49 : \par Examples
50 :
51 : The following input shows how you can call metatensor and evaluate the model
52 : that is described in the file `custom_cv.pt` from PLUMED.
53 :
54 : \plumedfile metatensor_cv: METATENSOR ... MODEL=custom_cv.pt
55 :
56 : SPECIES1=1-26
57 : SPECIES2=27-62
58 : SPECIES3=63-76
59 : SPECIES_TO_TYPES=6,1,8
60 : ...
61 : \endplumedfile
62 :
63 : The numbered `SPECIES` labels are used to indicate the list of atoms that belong
64 : to each atomic species in the system. The `SPECIES_TO_TYPE` keyword then
65 : provides information on the atom type for each species. The first number here is
66 : the atomic type of the atoms that have been specified using the `SPECIES1` flag,
67 : the second number is the atomic number of the atoms that have been specified
68 : using the `SPECIES2` flag and so on.
69 :
70 : `METATENSOR` action also accepts the following options:
71 :
72 : - `EXTENSIONS_DIRECTORY` should be the path to a directory containing
73 : TorchScript extensions (as shared libraries) that are required to load and
74 : execute the model. This matches the `collect_extensions` argument to
75 : `MetatensorAtomisticModel.export` in Python.
76 : - `CHECK_CONSISTENCY` can be used to enable internal consistency checks;
77 : - `SELECTED_ATOMS` can be used to signal the metatensor models that it should
78 : only run its calculation for the selected subset of atoms. The model still
79 : need to know about all the atoms in the system (through the `SPECIES`
80 : keyword); but this can be used to reduce the calculation cost. Note that the
81 : indices of the selected atoms should start at 1 in the PLUMED input file, but
82 : they will be translated to start at 0 when given to the model (i.e. in
83 : Python/TorchScript, the `forward` method will receive a `selected_atoms` which
84 : starts at 0)
85 :
86 : Here is another example with all the possible keywords:
87 :
88 : \plumedfile soap: METATENSOR ... MODEL=soap.pt EXTENSION_DIRECTORY=extensions
89 : CHECK_CONSISTENCY
90 :
91 : SPECIES1=1-10
92 : SPECIES2=11-20
93 : SPECIES_TO_TYPES=8,13
94 :
95 : # only run the calculation for the Aluminium (type 13) atoms, but
96 : # include the Oxygen (type 8) as potential neighbors.
97 : SELECTED_ATOMS=11-20
98 : ...
99 : \endplumedfile
100 :
101 : \par Collective variables and metatensor models
102 :
103 : PLUMED can use the [`"features"` output][features_output] of metatensor
104 : atomistic models as a collective variables. Alternatively, the code also accepts
105 : an output named `"plumed::cv"`, with the same metadata structure as the
106 : `"features"` output.
107 :
108 : */ /*
109 :
110 : [TorchScript]: https://pytorch.org/docs/stable/jit.html
111 : [mts_models]: https://docs.metatensor.org/latest/atomistic/index.html
112 : [mts_tutorials]: https://docs.metatensor.org/latest/examples/atomistic/index.html
113 : [mts_block]: https://docs.metatensor.org/latest/torch/reference/block.html
114 : [features_output]: https://docs.metatensor.org/latest/examples/atomistic/outputs/features.html
115 : */
116 : //+ENDPLUMEDOC
117 :
118 : /*INDENT-OFF*/
119 : #if !defined(__PLUMED_HAS_METATENSOR) || !defined(__PLUMED_HAS_LIBTORCH)
120 :
121 : namespace PLMD { namespace metatensor {
122 : class MetatensorPlumedAction: public ActionAtomistic, public ActionWithValue {
123 : public:
124 : static void registerKeywords(Keywords& keys);
125 0 : explicit MetatensorPlumedAction(const ActionOptions& options):
126 : Action(options),
127 : ActionAtomistic(options),
128 0 : ActionWithValue(options)
129 : {
130 0 : throw std::runtime_error(
131 : "Can not use metatensor action without the corresponding libraries. \n"
132 : "Make sure to configure with `--enable-metatensor --enable-libtorch` "
133 : "and that the corresponding libraries are found"
134 0 : );
135 0 : }
136 :
137 0 : void calculate() override {}
138 0 : void apply() override {}
139 0 : unsigned getNumberOfDerivatives() override {return 0;}
140 : };
141 :
142 : }} // namespace PLMD::metatensor
143 :
144 : #else
145 :
146 : #include <type_traits>
147 :
148 : #pragma GCC diagnostic push
149 : #pragma GCC diagnostic ignored "-Wpedantic"
150 : #pragma GCC diagnostic ignored "-Wunused-parameter"
151 : #pragma GCC diagnostic ignored "-Wfloat-equal"
152 : #pragma GCC diagnostic ignored "-Wfloat-conversion"
153 : #pragma GCC diagnostic ignored "-Wimplicit-float-conversion"
154 : #pragma GCC diagnostic ignored "-Wimplicit-int-conversion"
155 : #pragma GCC diagnostic ignored "-Wshorten-64-to-32"
156 : #pragma GCC diagnostic ignored "-Wsign-conversion"
157 : #pragma GCC diagnostic ignored "-Wold-style-cast"
158 :
159 : #include <torch/script.h>
160 : #include <torch/version.h>
161 : #include <torch/cuda.h>
162 : #if TORCH_VERSION_MAJOR >= 2
163 : #include <torch/mps.h>
164 : #endif
165 :
166 : #pragma GCC diagnostic pop
167 :
168 : #include <metatensor/torch.hpp>
169 : #include <metatensor/torch/atomistic.hpp>
170 :
171 : #include "vesin.h"
172 :
173 :
174 : namespace PLMD {
175 : namespace metatensor {
176 :
177 : // We will cast Vector/Tensor to pointers to arrays and doubles, so let's make
178 : // sure this is legal to do
179 : static_assert(std::is_standard_layout<PLMD::Vector>::value);
180 : static_assert(sizeof(PLMD::Vector) == sizeof(std::array<double, 3>));
181 : static_assert(alignof(PLMD::Vector) == alignof(std::array<double, 3>));
182 :
183 : static_assert(std::is_standard_layout<PLMD::Tensor>::value);
184 : static_assert(sizeof(PLMD::Tensor) == sizeof(std::array<std::array<double, 3>, 3>));
185 : static_assert(alignof(PLMD::Tensor) == alignof(std::array<std::array<double, 3>, 3>));
186 :
187 : class MetatensorPlumedAction: public ActionAtomistic, public ActionWithValue {
188 : public:
189 : static void registerKeywords(Keywords& keys);
190 : explicit MetatensorPlumedAction(const ActionOptions&);
191 :
192 : void calculate() override;
193 : void apply() override;
194 : unsigned getNumberOfDerivatives() override;
195 :
196 : private:
197 : // fill this->system_ according to the current PLUMED data
198 : void createSystem();
199 : // compute a neighbor list following metatensor format, using data from PLUMED
200 : metatensor_torch::TorchTensorBlock computeNeighbors(
201 : metatensor_torch::NeighborListOptions request,
202 : const std::vector<PLMD::Vector>& positions,
203 : const PLMD::Tensor& cell
204 : );
205 :
206 : // execute the model for the given system
207 : metatensor_torch::TorchTensorBlock executeModel(metatensor_torch::System system);
208 :
209 : torch::jit::Module model_;
210 :
211 : metatensor_torch::ModelCapabilities capabilities_;
212 : std::string model_output_;
213 :
214 : // neighbor lists requests made by the model
215 : std::vector<metatensor_torch::NeighborListOptions> nl_requests_;
216 :
217 : // dtype/device to use to execute the model
218 : torch::ScalarType dtype_;
219 : torch::Device device_;
220 :
221 : torch::Tensor atomic_types_;
222 : // store the strain to be able to compute the virial with autograd
223 : torch::Tensor strain_;
224 :
225 : metatensor_torch::System system_;
226 : metatensor_torch::ModelEvaluationOptions evaluations_options_;
227 : bool check_consistency_;
228 :
229 : metatensor_torch::TorchTensorMap output_;
230 : // shape of the output of this model
231 : unsigned n_samples_;
232 : unsigned n_properties_;
233 : };
234 :
235 :
236 : MetatensorPlumedAction::MetatensorPlumedAction(const ActionOptions& options):
237 : Action(options),
238 : ActionAtomistic(options),
239 : ActionWithValue(options),
240 : device_(torch::kCPU)
241 : {
242 : if (metatensor_torch::version().find("0.5.") != 0) {
243 : this->error(
244 : "this code requires version 0.5.x of metatensor-torch, got version " +
245 : metatensor_torch::version()
246 : );
247 : }
248 :
249 : // first, load the model
250 : std::string extensions_directory_str;
251 : this->parse("EXTENSIONS_DIRECTORY", extensions_directory_str);
252 :
253 : torch::optional<std::string> extensions_directory = torch::nullopt;
254 : if (!extensions_directory_str.empty()) {
255 : extensions_directory = std::move(extensions_directory_str);
256 : }
257 :
258 : std::string model_path;
259 : this->parse("MODEL", model_path);
260 :
261 : try {
262 : this->model_ = metatensor_torch::load_atomistic_model(model_path, extensions_directory);
263 : } catch (const std::exception& e) {
264 : this->error("failed to load model at '" + model_path + "': " + e.what());
265 : }
266 :
267 : // extract information from the model
268 : auto metadata = this->model_.run_method("metadata").toCustomClass<metatensor_torch::ModelMetadataHolder>();
269 : this->capabilities_ = this->model_.run_method("capabilities").toCustomClass<metatensor_torch::ModelCapabilitiesHolder>();
270 : auto requests_ivalue = this->model_.run_method("requested_neighbor_lists");
271 : for (auto request_ivalue: requests_ivalue.toList()) {
272 : auto request = request_ivalue.get().toCustomClass<metatensor_torch::NeighborListOptionsHolder>();
273 : this->nl_requests_.push_back(request);
274 : }
275 :
276 : log.printf("\n%s\n", metadata->print().c_str());
277 : // add the model references to PLUMED citation handling mechanism
278 : for (const auto& it: metadata->references) {
279 : for (const auto& ref: it.value()) {
280 : this->cite(ref);
281 : }
282 : }
283 :
284 : // parse the atomic types from the input file
285 : std::vector<int32_t> atomic_types;
286 : std::vector<int32_t> species_to_types;
287 : this->parseVector("SPECIES_TO_TYPES", species_to_types);
288 : bool has_custom_types = !species_to_types.empty();
289 :
290 : std::vector<AtomNumber> all_atoms;
291 : this->parseAtomList("SPECIES", all_atoms);
292 :
293 : size_t n_species = 0;
294 : if (all_atoms.empty()) {
295 : std::vector<AtomNumber> t;
296 : int i = 0;
297 : while (true) {
298 : i += 1;
299 : this->parseAtomList("SPECIES", i, t);
300 : if (t.empty()) {
301 : break;
302 : }
303 :
304 : int32_t type = i;
305 : if (has_custom_types) {
306 : if (species_to_types.size() < static_cast<size_t>(i)) {
307 : this->error(
308 : "SPECIES_TO_TYPES is too small, it should have one entry "
309 : "for each species (we have at least " + std::to_string(i) +
310 : " species and " + std::to_string(species_to_types.size()) +
311 : "entries in SPECIES_TO_TYPES)"
312 : );
313 : }
314 :
315 : type = species_to_types[static_cast<size_t>(i - 1)];
316 : }
317 :
318 : log.printf(" atoms with type %d are: ", type);
319 : for(unsigned j=0; j<t.size(); j++) {
320 : log.printf("%d ", t[j]);
321 : all_atoms.push_back(t[j]);
322 : atomic_types.push_back(type);
323 : }
324 : log.printf("\n"); t.resize(0);
325 :
326 : n_species += 1;
327 : }
328 : } else {
329 : n_species = 1;
330 :
331 : int32_t type = 1;
332 : if (has_custom_types) {
333 : type = species_to_types[0];
334 : }
335 : atomic_types.resize(all_atoms.size(), type);
336 : }
337 :
338 : if (has_custom_types && species_to_types.size() != n_species) {
339 : this->warning(
340 : "SPECIES_TO_TYPES contains more entries (" +
341 : std::to_string(species_to_types.size()) +
342 : ") than there where species (" + std::to_string(n_species) + ")"
343 : );
344 : }
345 :
346 : this->atomic_types_ = torch::tensor(std::move(atomic_types));
347 : this->requestAtoms(all_atoms);
348 :
349 : this->check_consistency_ = false;
350 : this->parseFlag("CHECK_CONSISTENCY", this->check_consistency_);
351 : if (this->check_consistency_) {
352 : log.printf(" checking for internal consistency of the model\n");
353 : }
354 :
355 : // create evaluation options for the model. These won't change during the
356 : // simulation, so we initialize them once here.
357 : evaluations_options_ = torch::make_intrusive<metatensor_torch::ModelEvaluationOptionsHolder>();
358 : evaluations_options_->set_length_unit(getUnits().getLengthString());
359 :
360 : auto outputs = this->capabilities_->outputs();
361 : if (outputs.contains("features")) {
362 : this->model_output_ = "features";
363 : }
364 :
365 : if (outputs.contains("plumed::cv")) {
366 : if (outputs.contains("features")) {
367 : this->warning(
368 : "this model exposes both 'features' and 'plumed::cv' outputs, "
369 : "we will use 'features'. 'plumed::cv' is deprecated, please "
370 : "remove it from your models"
371 : );
372 : } else {
373 : this->warning(
374 : "this model is using 'plumed::cv' output, which is deprecated. "
375 : "Please replace it with a 'features' output"
376 : );
377 : this->model_output_ = "plumed::cv";
378 : }
379 : }
380 :
381 :
382 : if (this->model_output_.empty()) {
383 : auto existing_outputs = std::vector<std::string>();
384 : for (const auto& it: this->capabilities_->outputs()) {
385 : existing_outputs.push_back(it.key());
386 : }
387 :
388 : this->error(
389 : "expected 'features' or 'plumed::cv' in the capabilities of the model, "
390 : "could not find it. the following outputs exist: " + torch::str(existing_outputs)
391 : );
392 : }
393 :
394 : auto output = torch::make_intrusive<metatensor_torch::ModelOutputHolder>();
395 : // this output has no quantity or unit to set
396 :
397 : output->per_atom = this->capabilities_->outputs().at(this->model_output_)->per_atom;
398 : // we are using torch autograd system to compute gradients,
399 : // so we don't need any explicit gradients.
400 : output->explicit_gradients = {};
401 : evaluations_options_->outputs.insert(this->model_output_, output);
402 :
403 : // Determine which device we should use based on user input, what the model
404 : // supports and what's available
405 : auto available_devices = std::vector<torch::Device>();
406 : for (const auto& device: this->capabilities_->supported_devices) {
407 : if (device == "cpu") {
408 : available_devices.push_back(torch::kCPU);
409 : } else if (device == "cuda") {
410 : if (torch::cuda::is_available()) {
411 : available_devices.push_back(torch::Device("cuda"));
412 : }
413 : } else if (device == "mps") {
414 : #if TORCH_VERSION_MAJOR >= 2
415 : if (torch::mps::is_available()) {
416 : available_devices.push_back(torch::Device("mps"));
417 : }
418 : #endif
419 : } else {
420 : this->warning(
421 : "the model declared support for unknown device '" + device +
422 : "', it will be ignored"
423 : );
424 : }
425 : }
426 :
427 : if (available_devices.empty()) {
428 : this->error(
429 : "failed to find a valid device for the model at '" + model_path + "': "
430 : "the model supports " + torch::str(this->capabilities_->supported_devices) +
431 : ", none of these where available"
432 : );
433 : }
434 :
435 : std::string requested_device;
436 : this->parse("DEVICE", requested_device);
437 : if (requested_device.empty()) {
438 : // no user request, pick the device the model prefers
439 : this->device_ = available_devices[0];
440 : } else {
441 : bool found_requested_device = false;
442 : for (const auto& device: available_devices) {
443 : if (device.is_cpu() && requested_device == "cpu") {
444 : this->device_ = device;
445 : found_requested_device = true;
446 : break;
447 : } else if (device.is_cuda() && requested_device == "cuda") {
448 : this->device_ = device;
449 : found_requested_device = true;
450 : break;
451 : } else if (device.is_mps() && requested_device == "mps") {
452 : this->device_ = device;
453 : found_requested_device = true;
454 : break;
455 : }
456 : }
457 :
458 : if (!found_requested_device) {
459 : this->error(
460 : "failed to find requested device (" + requested_device + "): it is either "
461 : "not supported by this model or not available on this machine"
462 : );
463 : }
464 : }
465 :
466 : this->model_.to(this->device_);
467 : this->atomic_types_ = this->atomic_types_.to(this->device_);
468 :
469 : log.printf(
470 : " running model on %s device with %s data\n",
471 : this->device_.str().c_str(),
472 : this->capabilities_->dtype().c_str()
473 : );
474 :
475 : if (this->capabilities_->dtype() == "float64") {
476 : this->dtype_ = torch::kFloat64;
477 : } else if (this->capabilities_->dtype() == "float32") {
478 : this->dtype_ = torch::kFloat32;
479 : } else {
480 : this->error(
481 : "the model requested an unsupported dtype '" + this->capabilities_->dtype() + "'"
482 : );
483 : }
484 :
485 : auto tensor_options = torch::TensorOptions().dtype(this->dtype_).device(this->device_);
486 : this->strain_ = torch::eye(3, tensor_options.requires_grad(true));
487 :
488 : // determine how many properties there will be in the output by running the
489 : // model once on a dummy system
490 : auto dummy_system = torch::make_intrusive<metatensor_torch::SystemHolder>(
491 : /*types = */ torch::zeros({0}, tensor_options.dtype(torch::kInt32)),
492 : /*positions = */ torch::zeros({0, 3}, tensor_options),
493 : /*cell = */ torch::zeros({3, 3}, tensor_options)
494 : );
495 :
496 : log.printf(" the following neighbor lists have been requested:\n");
497 : auto length_unit = this->getUnits().getLengthString();
498 : auto model_length_unit = this->capabilities_->length_unit();
499 : for (auto request: this->nl_requests_) {
500 : log.printf(" - %s list, %g %s cutoff (requested %g %s)\n",
501 : request->full_list() ? "full" : "half",
502 : request->engine_cutoff(length_unit),
503 : length_unit.c_str(),
504 : request->cutoff(),
505 : model_length_unit.c_str()
506 : );
507 :
508 : auto neighbors = this->computeNeighbors(request, {PLMD::Vector(0, 0, 0)}, PLMD::Tensor(0, 0, 0, 0, 0, 0, 0, 0, 0));
509 : metatensor_torch::register_autograd_neighbors(dummy_system, neighbors, this->check_consistency_);
510 : dummy_system->add_neighbor_list(request, neighbors);
511 : }
512 :
513 : this->n_properties_ = static_cast<unsigned>(
514 : this->executeModel(dummy_system)->properties()->count()
515 : );
516 :
517 : // parse and handle atom sub-selection. This is done AFTER determining the
518 : // output size, since the selection might not be valid for the dummy system
519 : std::vector<int32_t> selected_atoms;
520 : this->parseVector("SELECTED_ATOMS", selected_atoms);
521 : if (!selected_atoms.empty()) {
522 : auto selection_value = torch::zeros(
523 : {static_cast<int64_t>(selected_atoms.size()), 2},
524 : torch::TensorOptions().dtype(torch::kInt32).device(this->device_)
525 : );
526 :
527 : for (unsigned i=0; i<selected_atoms.size(); i++) {
528 : auto n_atoms = static_cast<int32_t>(this->atomic_types_.size(0));
529 : if (selected_atoms[i] <= 0 || selected_atoms[i] > n_atoms) {
530 : this->error(
531 : "Values in metatensor's SELECTED_ATOMS should be between 1 "
532 : "and the number of atoms (" + std::to_string(n_atoms) + "), "
533 : "got " + std::to_string(selected_atoms[i]));
534 : }
535 : // PLUMED input uses 1-based indexes, but metatensor wants 0-based
536 : selection_value[i][1] = selected_atoms[i] - 1;
537 : }
538 :
539 : evaluations_options_->set_selected_atoms(
540 : torch::make_intrusive<metatensor_torch::LabelsHolder>(
541 : std::vector<std::string>{"system", "atom"}, selection_value
542 : )
543 : );
544 : }
545 :
546 : // Now that we now both n_samples and n_properties, we can setup the
547 : // PLUMED-side storage for the computed CV
548 : if (output->per_atom) {
549 : if (selected_atoms.empty()) {
550 : this->n_samples_ = static_cast<unsigned>(this->atomic_types_.size(0));
551 : } else {
552 : this->n_samples_ = static_cast<unsigned>(selected_atoms.size());
553 : }
554 : } else {
555 : this->n_samples_ = 1;
556 : }
557 :
558 : if (n_samples_ == 1 && n_properties_ == 1) {
559 : log.printf(" the output of this model is a scalar\n");
560 :
561 : this->addValue();
562 : } else if (n_samples_ == 1) {
563 : log.printf(" the output of this model is 1x%d vector\n", n_properties_);
564 :
565 : this->addValue({this->n_properties_});
566 : this->getPntrToComponent(0)->buildDataStore();
567 : } else if (n_properties_ == 1) {
568 : log.printf(" the output of this model is %dx1 vector\n", n_samples_);
569 :
570 : this->addValue({this->n_samples_});
571 : this->getPntrToComponent(0)->buildDataStore();
572 : } else {
573 : log.printf(" the output of this model is a %dx%d matrix\n", n_samples_, n_properties_);
574 :
575 : this->addValue({this->n_samples_, this->n_properties_});
576 : this->getPntrToComponent(0)->buildDataStore();
577 : this->getPntrToComponent(0)->reshapeMatrixStore(n_properties_);
578 : }
579 :
580 : this->setNotPeriodic();
581 : }
582 :
583 : unsigned MetatensorPlumedAction::getNumberOfDerivatives() {
584 : // gradients w.r.t. positions (3 x N values) + gradients w.r.t. strain (9 values)
585 : return 3 * this->getNumberOfAtoms() + 9;
586 : }
587 :
588 :
589 : void MetatensorPlumedAction::createSystem() {
590 : if (this->getTotAtoms() != static_cast<unsigned>(this->atomic_types_.size(0))) {
591 : std::ostringstream oss;
592 : oss << "METATENSOR action needs to know about all atoms in the system. ";
593 : oss << "There are " << this->getTotAtoms() << " atoms overall, ";
594 : oss << "but we only have atomic types for " << this->atomic_types_.size(0) << " of them.";
595 : plumed_merror(oss.str());
596 : }
597 :
598 : // this->getTotAtoms()
599 :
600 : const auto& cell = this->getPbc().getBox();
601 :
602 : auto cpu_f64_tensor = torch::TensorOptions().dtype(torch::kFloat64).device(torch::kCPU);
603 : auto torch_cell = torch::zeros({3, 3}, cpu_f64_tensor);
604 :
605 : torch_cell[0][0] = cell(0, 0);
606 : torch_cell[0][1] = cell(0, 1);
607 : torch_cell[0][2] = cell(0, 2);
608 :
609 : torch_cell[1][0] = cell(1, 0);
610 : torch_cell[1][1] = cell(1, 1);
611 : torch_cell[1][2] = cell(1, 2);
612 :
613 : torch_cell[2][0] = cell(2, 0);
614 : torch_cell[2][1] = cell(2, 1);
615 : torch_cell[2][2] = cell(2, 2);
616 :
617 : const auto& positions = this->getPositions();
618 :
619 : auto torch_positions = torch::from_blob(
620 : const_cast<PLMD::Vector*>(positions.data()),
621 : {static_cast<int64_t>(positions.size()), 3},
622 : cpu_f64_tensor
623 : );
624 :
625 : torch_positions = torch_positions.to(this->dtype_).to(this->device_);
626 : torch_cell = torch_cell.to(this->dtype_).to(this->device_);
627 :
628 : // setup torch's automatic gradient tracking
629 : if (!this->doNotCalculateDerivatives()) {
630 : torch_positions.requires_grad_(true);
631 :
632 : // pretend to scale positions/cell by the strain so that it enters the
633 : // computational graph.
634 : torch_positions = torch_positions.matmul(this->strain_);
635 : torch_positions.retain_grad();
636 :
637 : torch_cell = torch_cell.matmul(this->strain_);
638 : }
639 :
640 : this->system_ = torch::make_intrusive<metatensor_torch::SystemHolder>(
641 : this->atomic_types_,
642 : torch_positions,
643 : torch_cell
644 : );
645 :
646 : // compute the neighbors list requested by the model, and register them with
647 : // the system
648 : for (auto request: this->nl_requests_) {
649 : auto neighbors = this->computeNeighbors(request, positions, cell);
650 : metatensor_torch::register_autograd_neighbors(this->system_, neighbors, this->check_consistency_);
651 : this->system_->add_neighbor_list(request, neighbors);
652 : }
653 : }
654 :
655 :
656 : metatensor_torch::TorchTensorBlock MetatensorPlumedAction::computeNeighbors(
657 : metatensor_torch::NeighborListOptions request,
658 : const std::vector<PLMD::Vector>& positions,
659 : const PLMD::Tensor& cell
660 : ) {
661 : auto labels_options = torch::TensorOptions().dtype(torch::kInt32).device(this->device_);
662 : auto neighbor_component = torch::make_intrusive<metatensor_torch::LabelsHolder>(
663 : "xyz",
664 : torch::tensor({0, 1, 2}, labels_options).reshape({3, 1})
665 : );
666 : auto neighbor_properties = torch::make_intrusive<metatensor_torch::LabelsHolder>(
667 : "distance", torch::zeros({1, 1}, labels_options)
668 : );
669 :
670 : auto cutoff = request->engine_cutoff(this->getUnits().getLengthString());
671 :
672 : auto non_periodic = (
673 : cell(0, 0) == 0.0 && cell(0, 1) == 0.0 && cell(0, 2) == 0.0 &&
674 : cell(1, 0) == 0.0 && cell(1, 1) == 0.0 && cell(1, 2) == 0.0 &&
675 : cell(2, 0) == 0.0 && cell(2, 2) == 0.0 && cell(2, 2) == 0.0
676 : );
677 :
678 : // use https://github.com/Luthaf/vesin to compute the requested neighbor
679 : // lists since we can not get these from PLUMED
680 : vesin::VesinOptions options;
681 : options.cutoff = cutoff;
682 : options.full = request->full_list();
683 : options.return_shifts = true;
684 : options.return_distances = false;
685 : options.return_vectors = true;
686 :
687 : vesin::VesinNeighborList* vesin_neighbor_list = new vesin::VesinNeighborList();
688 : memset(vesin_neighbor_list, 0, sizeof(vesin::VesinNeighborList));
689 :
690 : const char* error_message = NULL;
691 : int status = vesin_neighbors(
692 : reinterpret_cast<const double (*)[3]>(positions.data()),
693 : positions.size(),
694 : reinterpret_cast<const double (*)[3]>(&cell(0, 0)),
695 : !non_periodic,
696 : vesin::VesinCPU,
697 : options,
698 : vesin_neighbor_list,
699 : &error_message
700 : );
701 :
702 : if (status != EXIT_SUCCESS) {
703 : plumed_merror(
704 : "failed to compute neighbor list (cutoff=" + std::to_string(cutoff) +
705 : ", full=" + (request->full_list() ? "true" : "false") + "): " + error_message
706 : );
707 : }
708 :
709 : // transform from vesin to metatensor format
710 : auto n_pairs = static_cast<int64_t>(vesin_neighbor_list->length);
711 :
712 : auto pair_vectors = torch::from_blob(
713 : vesin_neighbor_list->vectors,
714 : {n_pairs, 3, 1},
715 : /*deleter*/ [=](void*) {
716 : vesin_free(vesin_neighbor_list);
717 : delete vesin_neighbor_list;
718 : },
719 : torch::TensorOptions().dtype(torch::kFloat64).device(torch::kCPU)
720 : );
721 :
722 : auto pair_samples_values = torch::zeros({n_pairs, 5}, labels_options.device(torch::kCPU));
723 : for (unsigned i=0; i<n_pairs; i++) {
724 : pair_samples_values[i][0] = static_cast<int32_t>(vesin_neighbor_list->pairs[i][0]);
725 : pair_samples_values[i][1] = static_cast<int32_t>(vesin_neighbor_list->pairs[i][1]);
726 : pair_samples_values[i][2] = vesin_neighbor_list->shifts[i][0];
727 : pair_samples_values[i][3] = vesin_neighbor_list->shifts[i][1];
728 : pair_samples_values[i][4] = vesin_neighbor_list->shifts[i][2];
729 : }
730 :
731 : auto neighbor_samples = torch::make_intrusive<metatensor_torch::LabelsHolder>(
732 : std::vector<std::string>{"first_atom", "second_atom", "cell_shift_a", "cell_shift_b", "cell_shift_c"},
733 : pair_samples_values.to(this->device_)
734 : );
735 :
736 : auto neighbors = torch::make_intrusive<metatensor_torch::TensorBlockHolder>(
737 : pair_vectors.to(this->dtype_).to(this->device_),
738 : neighbor_samples,
739 : std::vector<metatensor_torch::TorchLabels>{neighbor_component},
740 : neighbor_properties
741 : );
742 :
743 : return neighbors;
744 : }
745 :
746 : metatensor_torch::TorchTensorBlock MetatensorPlumedAction::executeModel(metatensor_torch::System system) {
747 : try {
748 : auto ivalue_output = this->model_.forward({
749 : std::vector<metatensor_torch::System>{system},
750 : evaluations_options_,
751 : this->check_consistency_,
752 : });
753 :
754 : auto dict_output = ivalue_output.toGenericDict();
755 : auto cv = dict_output.at(this->model_output_);
756 : this->output_ = cv.toCustomClass<metatensor_torch::TensorMapHolder>();
757 : } catch (const std::exception& e) {
758 : plumed_merror("failed to evaluate the model: " + std::string(e.what()));
759 : }
760 :
761 : plumed_massert(this->output_->keys()->count() == 1, "output should have a single block");
762 : auto block = metatensor_torch::TensorMapHolder::block_by_id(this->output_, 0);
763 : plumed_massert(block->components().empty(), "components are not yet supported in the output");
764 :
765 : return block;
766 : }
767 :
768 :
769 : void MetatensorPlumedAction::calculate() {
770 : this->createSystem();
771 :
772 : auto block = this->executeModel(this->system_);
773 : auto torch_values = block->values().to(torch::kCPU).to(torch::kFloat64);
774 :
775 : if (static_cast<unsigned>(torch_values.size(0)) != this->n_samples_) {
776 : plumed_merror(
777 : "expected the model to return a TensorBlock with " +
778 : std::to_string(this->n_samples_) + " samples, got " +
779 : std::to_string(torch_values.size(0)) + " instead"
780 : );
781 : } else if (static_cast<unsigned>(torch_values.size(1)) != this->n_properties_) {
782 : plumed_merror(
783 : "expected the model to return a TensorBlock with " +
784 : std::to_string(this->n_properties_) + " properties, got " +
785 : std::to_string(torch_values.size(1)) + " instead"
786 : );
787 : }
788 :
789 : Value* value = this->getPntrToComponent(0);
790 : // reshape the plumed `Value` to hold the data returned by the model
791 : if (n_samples_ == 1) {
792 : if (n_properties_ == 1) {
793 : value->set(torch_values.item<double>());
794 : } else {
795 : // we have multiple CV describing a single thing (atom or full system)
796 : for (unsigned i=0; i<n_properties_; i++) {
797 : value->set(i, torch_values[0][i].item<double>());
798 : }
799 : }
800 : } else {
801 : auto samples = block->samples();
802 : plumed_assert((samples->names() == std::vector<std::string>{"system", "atom"}));
803 :
804 : auto samples_values = samples->values().to(torch::kCPU);
805 : auto selected_atoms = this->evaluations_options_->get_selected_atoms();
806 :
807 : // handle the possibility that samples are returned in
808 : // a non-sorted order.
809 : auto get_output_location = [&](unsigned i) {
810 : if (selected_atoms.has_value()) {
811 : // If the users picked some selected atoms, then we store the
812 : // output in the same order as the selection was given
813 : auto sample = samples_values.index({static_cast<int64_t>(i), torch::indexing::Slice()});
814 : auto position = selected_atoms.value()->position(sample);
815 : plumed_assert(position.has_value());
816 : return static_cast<unsigned>(position.value());
817 : } else {
818 : return static_cast<unsigned>(samples_values[i][1].item<int32_t>());
819 : }
820 : };
821 :
822 : if (n_properties_ == 1) {
823 : // we have a single CV describing multiple things (i.e. atoms)
824 : for (unsigned i=0; i<n_samples_; i++) {
825 : auto output_i = get_output_location(i);
826 : value->set(output_i, torch_values[i][0].item<double>());
827 : }
828 : } else {
829 : // the CV is a matrix
830 : for (unsigned i=0; i<n_samples_; i++) {
831 : auto output_i = get_output_location(i);
832 : for (unsigned j=0; j<n_properties_; j++) {
833 : value->set(output_i * n_properties_ + j, torch_values[i][j].item<double>());
834 : }
835 : }
836 : }
837 : }
838 : }
839 :
840 :
841 : void MetatensorPlumedAction::apply() {
842 : const auto* value = this->getPntrToComponent(0);
843 : if (!value->forcesWereAdded()) {
844 : return;
845 : }
846 :
847 : auto block = metatensor_torch::TensorMapHolder::block_by_id(this->output_, 0);
848 : auto torch_values = block->values().to(torch::kCPU).to(torch::kFloat64);
849 :
850 : auto output_grad = torch::zeros_like(torch_values);
851 : if (n_samples_ == 1) {
852 : if (n_properties_ == 1) {
853 : output_grad[0][0] = value->getForce();
854 : } else {
855 : for (unsigned i=0; i<n_properties_; i++) {
856 : output_grad[0][i] = value->getForce(i);
857 : }
858 : }
859 : } else {
860 : auto samples = block->samples();
861 : plumed_assert((samples->names() == std::vector<std::string>{"system", "atom"}));
862 :
863 : auto samples_values = samples->values().to(torch::kCPU);
864 : auto selected_atoms = this->evaluations_options_->get_selected_atoms();
865 :
866 : // see above for an explanation of why we use this function
867 : auto get_output_location = [&](unsigned i) {
868 : if (selected_atoms.has_value()) {
869 : auto sample = samples_values.index({static_cast<int64_t>(i), torch::indexing::Slice()});
870 : auto position = selected_atoms.value()->position(sample);
871 : plumed_assert(position.has_value());
872 : return static_cast<unsigned>(position.value());
873 : } else {
874 : return static_cast<unsigned>(samples_values[i][1].item<int32_t>());
875 : }
876 : };
877 :
878 : if (n_properties_ == 1) {
879 : for (unsigned i=0; i<n_samples_; i++) {
880 : auto output_i = get_output_location(i);
881 : output_grad[i][0] = value->getForce(output_i);
882 : }
883 : } else {
884 : for (unsigned i=0; i<n_samples_; i++) {
885 : auto output_i = get_output_location(i);
886 : for (unsigned j=0; j<n_properties_; j++) {
887 : output_grad[i][j] = value->getForce(output_i * n_properties_ + j);
888 : }
889 : }
890 : }
891 : }
892 :
893 : this->system_->positions().mutable_grad() = torch::Tensor();
894 : this->strain_.mutable_grad() = torch::Tensor();
895 :
896 : torch_values.backward(output_grad);
897 : auto positions_grad = this->system_->positions().grad();
898 : auto strain_grad = this->strain_.grad();
899 :
900 : positions_grad = positions_grad.to(torch::kCPU).to(torch::kFloat64);
901 : strain_grad = strain_grad.to(torch::kCPU).to(torch::kFloat64);
902 :
903 : plumed_assert(positions_grad.sizes().size() == 2);
904 : plumed_assert(positions_grad.is_contiguous());
905 :
906 : plumed_assert(strain_grad.sizes().size() == 2);
907 : plumed_assert(strain_grad.is_contiguous());
908 :
909 : auto derivatives = std::vector<double>(
910 : positions_grad.data_ptr<double>(),
911 : positions_grad.data_ptr<double>() + 3 * this->system_->size()
912 : );
913 :
914 : // add virials to the derivatives
915 : derivatives.push_back(-strain_grad[0][0].item<double>());
916 : derivatives.push_back(-strain_grad[0][1].item<double>());
917 : derivatives.push_back(-strain_grad[0][2].item<double>());
918 :
919 : derivatives.push_back(-strain_grad[1][0].item<double>());
920 : derivatives.push_back(-strain_grad[1][1].item<double>());
921 : derivatives.push_back(-strain_grad[1][2].item<double>());
922 :
923 : derivatives.push_back(-strain_grad[2][0].item<double>());
924 : derivatives.push_back(-strain_grad[2][1].item<double>());
925 : derivatives.push_back(-strain_grad[2][2].item<double>());
926 :
927 : unsigned index = 0;
928 : this->setForcesOnAtoms(derivatives, index);
929 : }
930 :
931 : } // namespace metatensor
932 : } // namespace PLMD
933 :
934 :
935 : #endif
936 :
937 :
938 : namespace PLMD {
939 : namespace metatensor {
940 :
941 : // use the same implementation for both the actual action and the dummy one
942 : // (when libtorch and libmetatensor could not be found).
943 2 : void MetatensorPlumedAction::registerKeywords(Keywords& keys) {
944 2 : Action::registerKeywords(keys);
945 2 : ActionAtomistic::registerKeywords(keys);
946 2 : ActionWithValue::registerKeywords(keys);
947 :
948 4 : keys.add("compulsory", "MODEL", "path to the exported metatensor model");
949 4 : keys.add("optional", "EXTENSIONS_DIRECTORY", "path to the directory containing TorchScript extensions to load");
950 4 : keys.add("optional", "DEVICE", "Torch device to use for the calculation");
951 :
952 4 : keys.addFlag("CHECK_CONSISTENCY", false, "Should we enable internal consistency of the model");
953 :
954 4 : keys.add("numbered", "SPECIES", "the atoms in each PLUMED species");
955 4 : keys.reset_style("SPECIES", "atoms");
956 :
957 4 : keys.add("optional", "SELECTED_ATOMS", "subset of atoms that should be used for the calculation");
958 4 : keys.reset_style("SELECTED_ATOMS", "atoms");
959 :
960 4 : keys.add("optional", "SPECIES_TO_TYPES", "mapping from PLUMED SPECIES to metatensor's atomic types");
961 :
962 4 : keys.addOutputComponent("outputs", "default", "scalar", "collective variable created by the metatensor model");
963 4 : keys.setValueDescription("scalar/vector/matrix","collective variable created by the metatensor model");
964 2 : }
965 :
966 : PLUMED_REGISTER_ACTION(MetatensorPlumedAction, "METATENSOR")
967 :
968 : } // namespace metatensor
969 : } // namespace PLMD
|
