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