diff --git a/CMakeLists.txt b/CMakeLists.txt index 05a4af8533..a75380191a 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -579,6 +579,11 @@ include(gmxManageCP2K) ######################################################################## include(gmxManageNNPot) +######################################################################## +# Process metatomic settings +######################################################################## +include(gmxManageMetatomic) + ######################################################################## #Process shared/static library settings ######################################################################## @@ -630,7 +635,7 @@ if (GMX_HWLOC) if (HWLOC_FIND_QUIETLY_AFTER_FIRST_RUN) set(HWLOC_FIND_QUIETLY TRUE) endif() - find_package(HWLOC 1.5) + find_package(HWLOC 1.5) if (HWLOC_FOUND) if (HWLOC_LIBRARIES MATCHES ".a$") @@ -716,7 +721,7 @@ if(GMX_GPU) elseif(${_gmx_gpu_uppercase} STREQUAL "HIP") include(gmxManageHip) elseif(${_gmx_gpu_uppercase} STREQUAL "OPENCL") - message(STATUS "GPU support with OpenCL is deprecated. It is still fully supported (and " + message(STATUS "GPU support with OpenCL is deprecated. It is still fully supported (and " "recommended for Apple GPUs). Please use CUDA for running on NVIDIA GPUs and " "SYCL for running on Intel and AMD GPUs.") include(gmxManageOpenCL) diff --git a/api/legacy/include/gromacs/topology/ifunc.h b/api/legacy/include/gromacs/topology/ifunc.h index 21316a401d..a43bdd9e57 100644 --- a/api/legacy/include/gromacs/topology/ifunc.h +++ b/api/legacy/include/gromacs/topology/ifunc.h @@ -198,6 +198,7 @@ enum F_DENSITYFITTING, F_EQM, F_ENNPOT, + F_EMETATOMICPOT, F_EPOT, F_EKIN, F_ETOT, diff --git a/api/nblib/listed_forces/conversionscommon.h b/api/nblib/listed_forces/conversionscommon.h index 50e5fe5c0b..2df7f3425b 100644 --- a/api/nblib/listed_forces/conversionscommon.h +++ b/api/nblib/listed_forces/conversionscommon.h @@ -156,6 +156,7 @@ using GmxToNblibMapping = Unimplemented, // F_DENSITYFITTING, Unimplemented, // F_EQM, Unimplemented, // F_ENNPOT, + Unimplemented, // F_EMETATOMICPOT, Unimplemented, // F_EPOT, Unimplemented, // F_EKIN, Unimplemented, // F_ETOT, diff --git a/cmake/gmxManageMetatomic.cmake b/cmake/gmxManageMetatomic.cmake new file mode 100644 index 0000000000..4136d759b4 --- /dev/null +++ b/cmake/gmxManageMetatomic.cmake @@ -0,0 +1,144 @@ +# +# This file is part of the GROMACS molecular simulation package. +# +# Copyright 2024- The GROMACS Authors +# and the project initiators Erik Lindahl, Berk Hess and David van der Spoel. +# Consult the AUTHORS/COPYING files and https://www.gromacs.org for details. +# +# GROMACS is free software; you can redistribute it and/or +# modify it under the terms of the GNU Lesser General Public License +# as published by the Free Software Foundation; either version 2.1 +# of the License, or (at your option) any later version. +# +# GROMACS is distributed in the hope that it will be useful, +# but WITHOUT ANY WARRANTY; without even the implied warranty of +# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU +# Lesser General Public License for more details. +# +# You should have received a copy of the GNU Lesser General Public +# License along with GROMACS; if not, see +# https://www.gnu.org/licenses, or write to the Free Software Foundation, +# Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA. +# +# If you want to redistribute modifications to GROMACS, please +# consider that scientific software is very special. Version +# control is crucial - bugs must be traceable. We will be happy to +# consider code for inclusion in the official distribution, but +# derived work must not be called official GROMACS. Details are found +# in the README & COPYING files - if they are missing, get the +# official version at https://www.gromacs.org. +# +# To help us fund GROMACS development, we humbly ask that you cite +# the research papers on the package. Check out https://www.gromacs.org. + + +option(GMX_METATOMIC "Enable interface to metatomic atomistic models" OFF) + +# if(TORCH_ALREADY_SEARCHED) +# set(FIND_TORCH_QUIETLY ON) +# endif() + +if(GMX_METATOMIC) + # Bring the `torch` target in scope to allow evaluation + # of cmake generator expression from `metatensor_torch` + find_package(Torch REQUIRED) + + ################ definition of metatensor and metatomic targets ################ + + set(METATENSOR_CORE_VERSION "0.1.17") + set(METATENSOR_CORE_SHA256 "42119e11908239915ccc187d7ca65449b461f1d4b5af4d6df1fb613d687da76a") + + set(METATENSOR_TORCH_VERSION "0.8.0") + set(METATENSOR_TORCH_SHA256 "61d383ce958deafe0e3916088185527680c9118588722b17ec5c39cfbaa6da55") + + set(METATOMIC_TORCH_VERSION "0.1.4") + set(METATOMIC_TORCH_SHA256 "385ec8b8515d674b6a9f093f724792b2469e7ea2365ca596f574b64e38494f94") + + set(VESIN_VERSION "0.4.1") + set(VESIN_GIT_TAG "87dcad999fec47b29ab21be9662ef283edc7530b") + + set(DOWNLOAD_VESIN_DEFAULT ON) + find_package(vesin ${VESIN_VERSION} QUIET) + if (vesin_FOUND) + set(DOWNLOAD_VESIN_DEFAULT OFF) + endif() + + set(DOWNLOAD_METATENSOR_DEFAULT ON) + find_package(metatensor_torch ${METATENSOR_TORCH_VERSION} QUIET) + if (metatensor_torch_FOUND) + set(DOWNLOAD_METATENSOR_DEFAULT OFF) + endif() + + set(DOWNLOAD_METATOMIC_DEFAULT ON) + find_package(metatomic_torch ${METATOMIC_TORCH_VERSION} QUIET) + if (metatomic_torch_FOUND) + set(DOWNLOAD_METATOMIC_DEFAULT OFF) + endif() + + + option(DOWNLOAD_METATENSOR "Download metatensor package instead of using an already installed one" ${DOWNLOAD_METATENSOR_DEFAULT}) + option(DOWNLOAD_METATOMIC "Download metatomic package instead of using an already installed one" ${DOWNLOAD_METATOMIC_DEFAULT}) + + if (DOWNLOAD_METATENSOR) + include(FetchContent) + + set(URL_BASE "https://github.com/metatensor/metatensor/releases/download") + FetchContent_Declare(metatensor + URL ${URL_BASE}/metatensor-core-v${METATENSOR_CORE_VERSION}/metatensor-core-cxx-${METATENSOR_CORE_VERSION}.tar.gz + URL_HASH SHA256=${METATENSOR_CORE_SHA256} + ) + + message(STATUS "Fetching metatensor v${METATENSOR_CORE_VERSION} from github") + FetchContent_MakeAvailable(metatensor) + + FetchContent_Declare(metatensor-torch + URL ${URL_BASE}/metatensor-torch-v${METATENSOR_TORCH_VERSION}/metatensor-torch-cxx-${METATENSOR_TORCH_VERSION}.tar.gz + URL_HASH SHA256=${METATENSOR_TORCH_SHA256} + ) + + message(STATUS "Fetching metatensor-torch v${METATENSOR_TORCH_VERSION} from github") + FetchContent_MakeAvailable(metatensor-torch) + else() + # make sure to fail the configuration if cmake can not find metatensor-torch + find_package(metatensor_torch REQUIRED ${METATENSOR_TORCH_VERSION}) + endif() + + if (DOWNLOAD_METATOMIC) + include(FetchContent) + + set(URL_BASE "https://github.com/metatensor/metatomic/releases/download") + FetchContent_Declare(metatomic-torch + URL ${URL_BASE}/metatomic-torch-v${METATOMIC_TORCH_VERSION}/metatomic-torch-cxx-${METATOMIC_TORCH_VERSION}.tar.gz + URL_HASH SHA256=${METATOMIC_TORCH_SHA256} + ) + + message(STATUS "Fetching metatomic-torch v${METATOMIC_TORCH_VERSION} from github") + FetchContent_MakeAvailable(metatomic-torch) + else() + # make sure to fail the configuration if cmake can not find metatomic-torch + find_package(metatomic_torch REQUIRED ${METATOMIC_TORCH_VERSION}) + endif() + + if (DOWNLOAD_VESIN) + include(FetchContent) + + FetchContent_Declare( + vesin + GIT_REPOSITORY https://github.com/Luthaf/vesin.git + GIT_TAG ${VESIN_GIT_TAG} + ) + + FetchContent_MakeAvailable(vesin) + else() + # make sure to fail the configuration if cmake can not find vesin + find_package(vesin REQUIRED ${VESIN_VERSION}) + endif() + + list(APPEND GMX_COMMON_LIBRARIES + vesin + metatensor + metatomic_torch + metatensor_torch + ) + +endif() diff --git a/cmake/gmxManageOpenMP.cmake b/cmake/gmxManageOpenMP.cmake index 7962bd3abe..8855b35088 100644 --- a/cmake/gmxManageOpenMP.cmake +++ b/cmake/gmxManageOpenMP.cmake @@ -38,7 +38,7 @@ if(GMX_OPENMP) # We should do OpenMP detection if we get here # OpenMP check must come before other CFLAGS! - find_package(OpenMP) + find_package(OpenMP REQUIRED C CXX) if(NOT OPENMP_FOUND) if(CMAKE_CXX_COMPILER_ID MATCHES "AppleClang") message(FATAL_ERROR "The compiler you are using does not support OpenMP parallelism, " diff --git a/src/config.h.cmakein b/src/config.h.cmakein index dbae940a8d..09dc4fb4e4 100644 --- a/src/config.h.cmakein +++ b/src/config.h.cmakein @@ -482,6 +482,9 @@ /* Define if we are building with an external FMM implementation */ #cmakedefine01 GMX_USE_EXT_FMM +/* Define if we are building with Metatomic */ +#cmakedefine01 GMX_METATOMIC + /* Build using clang analyzer */ #cmakedefine01 GMX_CLANG_ANALYZER diff --git a/src/external/muparser/CMakeLists.txt b/src/external/muparser/CMakeLists.txt index 3132ad81f3..9e67abf5bd 100644 --- a/src/external/muparser/CMakeLists.txt +++ b/src/external/muparser/CMakeLists.txt @@ -25,7 +25,7 @@ option(ENABLE_WIDE_CHAR "Enable wide character support" OFF) option(BUILD_SHARED_LIBS "Build shared/static libs" ON) if(ENABLE_OPENMP) - find_package(OpenMP REQUIRED) + find_package(OpenMP REQUIRED C CXX) endif() # Credit: https://stackoverflow.com/questions/2368811/how-to-set-warning-level-in-cmake/3818084 diff --git a/src/external/tng_io/src/tests/CMakeLists.txt b/src/external/tng_io/src/tests/CMakeLists.txt index 1de38e033d..4abd8e4768 100644 --- a/src/external/tng_io/src/tests/CMakeLists.txt +++ b/src/external/tng_io/src/tests/CMakeLists.txt @@ -50,7 +50,7 @@ if(TNG_BUILD_TEST) endif() if(TNG_BUILD_EXAMPLES) - find_package(OpenMP) + find_package(OpenMP REQUIRED C CXX) if(OPENMP_FOUND) set(CMAKE_C_FLAGS "${CMAKE_C_FLAGS} ${OpenMP_C_FLAGS}") set(CMAKE_EXE_LINKER_FLAGS "${CMAKE_EXE_LINKER_FLAGS} ${OpenMP_EXE_LINKER_FLAGS}") diff --git a/src/gromacs/applied_forces/CMakeLists.txt b/src/gromacs/applied_forces/CMakeLists.txt index be4d0eb71f..a8fc268a4f 100644 --- a/src/gromacs/applied_forces/CMakeLists.txt +++ b/src/gromacs/applied_forces/CMakeLists.txt @@ -69,6 +69,7 @@ add_subdirectory(qmmm) add_subdirectory(colvars) add_subdirectory(plumed) add_subdirectory(nnpot) +add_subdirectory(metatomic) if (BUILD_TESTING) add_subdirectory(tests) diff --git a/src/gromacs/applied_forces/metatomic/CMakeLists.txt b/src/gromacs/applied_forces/metatomic/CMakeLists.txt new file mode 100644 index 0000000000..735f457640 --- /dev/null +++ b/src/gromacs/applied_forces/metatomic/CMakeLists.txt @@ -0,0 +1,52 @@ +# +# This file is part of the GROMACS molecular simulation package. +# +# Copyright 2024- The GROMACS Authors +# and the project initiators Erik Lindahl, Berk Hess and David van der Spoel. +# Consult the AUTHORS/COPYING files and https://www.gromacs.org for details. +# +# GROMACS is free software; you can redistribute it and/or +# modify it under the terms of the GNU Lesser General Public License +# as published by the Free Software Foundation; either version 2.1 +# of the License, or (at your option) any later version. +# +# GROMACS is distributed in the hope that it will be useful, +# but WITHOUT ANY WARRANTY; without even the implied warranty of +# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU +# Lesser General Public License for more details. +# +# You should have received a copy of the GNU Lesser General Public +# License along with GROMACS; if not, see +# https://www.gnu.org/licenses, or write to the Free Software Foundation, +# Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA. +# +# If you want to redistribute modifications to GROMACS, please +# consider that scientific software is very special. Version +# control is crucial - bugs must be traceable. We will be happy to +# consider code for inclusion in the official distribution, but +# derived work must not be called official GROMACS. Details are found +# in the README & COPYING files - if they are missing, get the +# official version at https://www.gromacs.org. +# +# To help us fund GROMACS development, we humbly ask that you cite +# the research papers on the package. Check out https://www.gromacs.org. + +gmx_add_libgromacs_sources( + metatomic_options.cpp + metatomic_mdmodule.cpp + metatomic_topologypreprocessor.cpp +) + +if (GMX_METATOMIC) + gmx_add_libgromacs_sources( + metatomic_forceprovider.cpp + ) +else() + gmx_add_libgromacs_sources( + metatomic_forceprovider_stub.cpp + ) +endif() + +if (BUILD_TESTING) + add_subdirectory(tests) +endif() diff --git a/src/gromacs/applied_forces/metatomic/metatomic_forceprovider.cpp b/src/gromacs/applied_forces/metatomic/metatomic_forceprovider.cpp new file mode 100644 index 0000000000..e8da7a8af6 --- /dev/null +++ b/src/gromacs/applied_forces/metatomic/metatomic_forceprovider.cpp @@ -0,0 +1,564 @@ +/* + * This file is part of the GROMACS molecular simulation package. + * + * Copyright 2024- The GROMACS Authors + * and the project initiators Erik Lindahl, Berk Hess and David van der Spoel. + * Consult the AUTHORS/COPYING files and https://www.gromacs.org for details. + * + * GROMACS is free software; you can redistribute it and/or + * modify it under the terms of the GNU Lesser General Public License + * as published by the Free Software Foundation; either version 2.1 + * of the License, or (at your option) any later version. + * + * GROMACS is distributed in the hope that it will be useful, + * but WITHOUT ANY WARRANTY; without even the implied warranty of + * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU + * Lesser General Public License for more details. + * + * You should have received a copy of the GNU Lesser General Public + * License along with GROMACS; if not, see + * https://www.gnu.org/licenses, or write to the Free Software Foundation, + * Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA. + * + * If you want to redistribute modifications to GROMACS, please + * consider that scientific software is very special. Version + * control is crucial - bugs must be traceable. We will be happy to + * consider code for inclusion in the official distribution, but + * derived work must not be called official GROMACS. Details are found + * in the README & COPYING files - if they are missing, get the + * official version at https://www.gromacs.org. + * + * To help us fund GROMACS development, we humbly ask that you cite + * the research papers on the package. Check out https://www.gromacs.org. + */ +/*! \internal \file + * \brief + * Implements the Metatomic Force Provider class + * + * \author Metatensor developers + * \ingroup module_applied_forces + */ +#include "gmxpre.h" + +#include "metatomic_forceprovider.h" + +#include + +#include + +#include + +#ifndef DIM +# define DIM 3 +#endif +#include "gromacs/domdec/localatomset.h" +#include "gromacs/mdlib/broadcaststructs.h" +#include "gromacs/mdtypes/enerdata.h" +#include "gromacs/mdtypes/forceoutput.h" +#include "gromacs/utility/arrayref.h" +#include "gromacs/utility/exceptions.h" +#include "gromacs/utility/gmxassert.h" +#include "gromacs/utility/logger.h" +#include "gromacs/utility/mpicomm.h" + +using namespace std::string_literals; // For ""s +namespace gmx +{ + +MetatomicForceProvider::MetatomicForceProvider(const MetatomicOptions& options, + const MDLogger& logger, + const MpiComm& mpiComm) : + options_(options), + logger_(logger), + mpiComm_(mpiComm), + device_(torch::Device(torch::kCPU)), + box_{ { 0.0, 0.0, 0.0 }, { 0.0, 0.0, 0.0 }, { 0.0, 0.0, 0.0 } } +{ + // NOTE: do NOT return early on non-main ranks. Only perform file I/O / GPU init + // on the main rank, but still run the remainder of the constructor on all ranks + // so that all ranks call the same MPI collectives. + if (mpiComm_.isMainRank()) + { + GMX_LOG(logger_.info).asParagraph().appendText("Initializing MetatomicForceProvider..."); + + if (!std::filesystem::exists(options_.params_.modelPath_)) + { + GMX_THROW(FileIOError("Metatomic model file does not exist: " + options_.params_.modelPath_)); + } + + // Load the model + try + { + torch::optional extensions_directory = torch::nullopt; + if (!options_.params_.extensionsDirectory.empty()) + { + extensions_directory = options_.params_.extensionsDirectory; + } + + model_ = metatomic_torch::load_atomistic_model(options_.params_.modelPath_, extensions_directory); + } + catch (const std::exception& e) + { + GMX_THROW(APIError("Failed to load metatomic model: " + std::string(e.what()))); + } + + // Query model capabilities + capabilities_ = + model_.run_method("capabilities").toCustomClass(); + auto requests_ivalue = model_.run_method("requested_neighbor_lists"); + for (const auto& request_ivalue : requests_ivalue.toList()) + { + nl_requests_.push_back( + request_ivalue.get().toCustomClass()); + } + + // TODO(rg): determine device + model_.to(device_); + + // Set data type locally on main rank + if (capabilities_->dtype() == "float64") + { + dtype_ = torch::kFloat64; + } + else if (capabilities_->dtype() == "float32") + { + dtype_ = torch::kFloat32; + } + else + { + GMX_THROW(APIError("Unsupported dtype from model: " + capabilities_->dtype())); + } + + // 5. Set up evaluation options (only on main rank) + evaluations_options_ = torch::make_intrusive(); + evaluations_options_->set_length_unit("nm"); + + auto outputs = capabilities_->outputs(); + if (!outputs.contains("energy")) + { + GMX_THROW(APIError("Metatomic model must provide an 'energy' output.")); + } + + auto requested_output = torch::make_intrusive(); + requested_output->per_atom = false; + requested_output->explicit_gradients = {}; // Use autograd for forces + + evaluations_options_->outputs.insert("energy", requested_output); + } + else + { + // non-main ranks do not load model or query capabilities, but they must still participate + // in collectives and must know dtype_ afterwards; dtype_ will be broadcast below. + } + + // Initialize ML atom group data structures (must run on all ranks) + const auto& mtaIndices = options_.params_.mtaIndices_; + const int n_atoms = static_cast(mtaIndices.size()); + + // Resize vectors that will be used on all ranks. + positions_.resize(n_atoms); + atomNumbers_.resize(n_atoms); + idxLookup_.resize(n_atoms); // Filled by gatherAtomNumbersIndices + + // Atom numbers are static. We will populate them on the main rank + // and "broadcast" using sumReduce. + + // Ensure non-main ranks have zeroed atomNumbers_ so a sumReduce acts like a broadcast. + std::fill(atomNumbers_.begin(), atomNumbers_.end(), 0); + + // Now, only the main rank populates its version of the data. + if (mpiComm_.isMainRank()) + { + for (int i = 0; i < n_atoms; ++i) + { + const int gIdx = mtaIndices[i]; + + if (gIdx >= options_.params_.numAtoms_) + { + GMX_THROW(APIError("Metatomic atom index " + std::to_string(gIdx) + " is out of bounds for topology numAtoms=" + + std::to_string(options_.params_.numAtoms_))); + } + atomNumbers_[i] = options_.params_.atoms_.atom[gIdx].atomnumber; + } + } + + // "Broadcast" the static atom numbers to all ranks: all ranks must call this. + if (mpiComm_.isParallel()) + { + // resize/allocate on non-main ranks and broadcast raw bytes of vector + nblock_abc(mpiComm_.isMainRank(), mpiComm_.comm(), static_cast(n_atoms), &atomNumbers_); + } + + // Broadcast dtype info (so all ranks know whether the model uses float32 or float64). + // 0 = float64, 1 = float32 + int dtype_code_local = 0; + if (mpiComm_.isMainRank()) + { + dtype_code_local = (dtype_ == torch::kFloat32) ? 1 : 0; + } + // comm() may be MPI_COMM_NULL for SingleRank so.. + if (mpiComm_.isParallel()) + { + block_bc(mpiComm_.comm(), dtype_code_local); + } + // Set dtype_ on all ranks + dtype_ = (dtype_code_local == 1) ? torch::kFloat32 : torch::kFloat64; + + // Initialize the lookup table. It will be populated correctly + // by the first call from the SimulationRunNotifier. + std::fill(idxLookup_.begin(), idxLookup_.end(), -1); + + if (mpiComm_.isMainRank()) + { + GMX_LOG(logger_.info) + .asParagraph() + .appendText("MetatomicForceProvider initialization complete."); + } +} + +void MetatomicForceProvider::gatherAtomNumbersIndices() +{ + // This function is called on domain decomposition. + // Its ONLY job is to update the `idxLookup_` map, which maps + // the contiguous ML atom index (0..n_atoms-1) to the + // sparse GROMACS local index. + // + // `atomNumbers_` is static and was set/broadcast in the constructor. + + const auto& mtaIndices = options_.params_.mtaIndices_; + const int n_atoms = static_cast(mtaIndices.size()); + + // Reset the lookup table to -1 (sentinel for "not local") + std::fill(idxLookup_.begin(), idxLookup_.end(), -1); + + // Build a reverse lookup map for efficiently finding ML atoms. + // Map: (Global GROMACS Index) -> (ML Group Index, 0..n_atoms-1) + std::unordered_map globalToMlIndex; + globalToMlIndex.reserve(n_atoms); + for (int i = 0; i < n_atoms; ++i) + { + globalToMlIndex[mtaIndices[i]] = i; + } + + // Iterate over this rank's LOCAL atoms and populate the lookup table if a + // local atom is part of our ML group. + const auto* mtaAtoms = options_.params_.mtaAtoms_.get(); + for (size_t i = 0; i < mtaAtoms->numAtomsLocal(); ++i) + { + const int lIdx = mtaAtoms->localIndex()[i]; + const int gIdx = mtaAtoms->globalIndex()[mtaAtoms->collectiveIndex()[i]]; + + // Check if this local atom is one of our ML atoms + if (auto it = globalToMlIndex.find(gIdx); it != globalToMlIndex.end()) + { + const int mlIdx = it->second; // This is the index (0..n_atoms-1) + idxLookup_[mlIdx] = lIdx; + } + } + + // Done. No MPI is needed. `gatherAtomPositions` will now correctly + // use this rank-local `idxLookup_` to contribute its positions. +} + +MetatomicForceProvider::~MetatomicForceProvider() = default; + +void MetatomicForceProvider::gatherAtomPositions(ArrayRef globalPositions) +{ + const int n_atoms = static_cast(options_.params_.mtaIndices_.size()); + // Start with a zeroed vector on all ranks + positions_.assign(n_atoms, RVec{ 0.0, 0.0, 0.0 }); + + // Each rank fills in the positions for its local atoms + for (int i = 0; i < n_atoms; ++i) + { + if (idxLookup_[i] != -1) // local + { + positions_[i] = globalPositions[idxLookup_[i]]; + } + } + + // All ranks need the complete, contiguous list of positions. + if (mpiComm_.isParallel()) + { + real* data_ptr = reinterpret_cast(positions_.data()); + const size_t n_reals = static_cast(n_atoms) * 3; + mpiComm_.sumReduce(gmx::ArrayRef(data_ptr, data_ptr + n_reals)); + } +} + + +void MetatomicForceProvider::calculateForces(const ForceProviderInput& inputs, ForceProviderOutput* outputs) +{ + const int n_atoms = static_cast(options_.params_.mtaIndices_.size()); + + // Gather positions + this->gatherAtomPositions(inputs.x_); + copy_mat(inputs.box_, box_); + torch::Tensor forceTensor; + + auto gromacs_scalar_type = torch::kFloat32; + if (std::is_same_v) + { + gromacs_scalar_type = torch::kFloat64; + } + // This blob_options is for reading GROMACS memory + auto blob_options = torch::TensorOptions().dtype(gromacs_scalar_type).device(torch::kCPU); + + if (mpiComm_.isMainRank()) + { + auto coerced_positions = makeArrayRef(positions_); + + auto torch_positions = + torch::from_blob(coerced_positions.data()->as_vec(), { n_atoms, 3 }, blob_options) + .to(this->dtype_) + .to(this->device_) + .set_requires_grad(true); + + auto torch_cell = + torch::from_blob(&box_, { 3, 3 }, blob_options).to(this->dtype_).to(this->device_); + + auto torch_pbc = torch::tensor(std::vector{ 1, 1, 1 }, + torch::TensorOptions().dtype(torch::kBool).device(this->device_)); + auto torch_types = + torch::tensor(atomNumbers_, torch::TensorOptions().dtype(torch::kInt32)).to(this->device_); + + auto system = torch::make_intrusive( + torch_types, torch_positions, torch_cell, torch_pbc); + + bool periodic = torch::all(torch_pbc).item(); + // Compute and add neighbor lists + for (const auto& request : nl_requests_) + { + auto neighbors = computeNeighbors( + request, n_atoms, coerced_positions.data()->as_vec(), box_, periodic); + metatomic_torch::register_autograd_neighbors(system, neighbors, false); + system->add_neighbor_list(request, neighbors); + } + + // Run the model + metatensor_torch::TensorMap output_map; + try + { + auto ivalue_output = this->model_.forward({ + std::vector{ system }, + evaluations_options_, + this->check_consistency_, + }); + auto dict_output = ivalue_output.toGenericDict(); + output_map = dict_output.at("energy").toCustomClass(); + } + catch (const std::exception& e) + { + GMX_THROW(APIError("[MetatomicPotential] Model evaluation failed: "s + e.what())); + } + + // Extract energy + auto energy_block = metatensor_torch::TensorMapHolder::block_by_id(output_map, 0); + auto energy_tensor = energy_block->values(); + + if (energy_tensor.sizes().vec() != std::vector{ 1, 1 }) + { + GMX_THROW(APIError("Model did not return a single scalar energy value."s)); + } + + // Set energy output (GROMACS sums this over ranks) + outputs->enerd_.term[F_EMETATOMICPOT] = energy_tensor.item(); + + // Compute gradients + energy_tensor.backward(); + auto grad = system->positions().grad(); + + // Populate the outer-scoped tensor + forceTensor = -grad.to(torch::kCPU).to(dtype_); + } // --- End of main rank block --- + + + // Broadcast forces and scatter + const size_t total = static_cast(n_atoms) * 3; + MPI_Comm comm = mpiComm_.comm(); + + if (dtype_ == torch::kFloat64) + { + std::vector global_force(total, 0.0); + if (mpiComm_.isMainRank()) + { + // Use the correct accessor for the model's dtype + auto accessor = forceTensor.accessor(); + for (int i = 0; i < n_atoms; ++i) + { + global_force[3 * i + 0] = accessor[i][0]; + global_force[3 * i + 1] = accessor[i][1]; + global_force[3 * i + 2] = accessor[i][2]; + } + } + + if (mpiComm_.isParallel()) + { + nblock_abc(mpiComm_.isMainRank(), comm, total, &global_force); + } + + // Scatter (double -> double) + for (int i = 0; i < n_atoms; ++i) + { + const int localIndex = idxLookup_[i]; + if (localIndex != -1) + { + outputs->forceWithVirial_.force_[localIndex][0] += global_force[3 * i + 0]; + outputs->forceWithVirial_.force_[localIndex][1] += global_force[3 * i + 1]; + outputs->forceWithVirial_.force_[localIndex][2] += global_force[3 * i + 2]; + } + } + } + else if (dtype_ == torch::kFloat32) + { + std::vector global_force(total, 0.0f); + if (mpiComm_.isMainRank()) + { + // Use the correct accessor for the model's dtype + auto accessor = forceTensor.accessor(); + for (int i = 0; i < n_atoms; ++i) + { + global_force[3 * i + 0] = accessor[i][0]; + global_force[3 * i + 1] = accessor[i][1]; + global_force[3 * i + 2] = accessor[i][2]; + } + } + + if (mpiComm_.isParallel()) + { + nblock_abc(mpiComm_.isMainRank(), comm, total, &global_force); + } + + // Scatter (float -> double) + for (int i = 0; i < n_atoms; ++i) + { + const int localIndex = idxLookup_[i]; + if (localIndex != -1) + { + outputs->forceWithVirial_.force_[localIndex][0] += + static_cast(global_force[3 * i + 0]); + outputs->forceWithVirial_.force_[localIndex][1] += + static_cast(global_force[3 * i + 1]); + outputs->forceWithVirial_.force_[localIndex][2] += + static_cast(global_force[3 * i + 2]); + } + } + } + else + { + GMX_THROW(APIError("Unsupported dtype, only float32 and float64 are supported for now")); + } + // TODO(rg): Virial calculation. For now, GROMACS will (incorrectly) calculate it from forces if + // needed. This is the same behavior as nnpot +} + + +metatensor_torch::TensorBlock MetatomicForceProvider::computeNeighbors(metatomic_torch::NeighborListOptions request, + long n_atoms, + const float* positions, + const matrix box, + bool periodic) +{ + auto cutoff = request->engine_cutoff("nm"); + + 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(); + // .............................. gromacs likes floats, vesin does not + double double_box[3][3]; + + for (int i = 0; i < 3; i++) + { + for (int j = 0; j < 3; j++) + { + double_box[i][j] = static_cast(box[i][j]); + } + } + + const size_t total_elements = static_cast(n_atoms) * 3; + std::vector double_positions(total_elements); + + for (size_t i = 0; i < total_elements; i++) + { + double_positions[i] = static_cast(positions[i]); + } + const double* positions_ptr = double_positions.data(); + + VesinDevice cpu{ VesinCPU, 0 }; + const char* error_message = nullptr; + int status = vesin_neighbors(reinterpret_cast(positions_ptr), + static_cast(n_atoms), + double_box, + &periodic, + cpu, + options, + vesin_neighbor_list, + &error_message); + + + if (status != EXIT_SUCCESS) + { + std::string err_str = "vesin_neighbors failed: "; + if (error_message) + { + err_str += error_message; + } + delete vesin_neighbor_list; + GMX_THROW(APIError(err_str)); + } + + auto n_pairs = static_cast(vesin_neighbor_list->length); + + // Build samples tensor (first_atom, second_atom, cell_shifts) + auto pair_samples_values = torch::empty({ n_pairs, 5 }, torch::TensorOptions().dtype(torch::kInt32)); + auto pair_samples_ptr = pair_samples_values.accessor(); + for (int64_t i = 0; i < n_pairs; i++) + { + pair_samples_ptr[i][0] = static_cast(vesin_neighbor_list->pairs[i][0]); + pair_samples_ptr[i][1] = static_cast(vesin_neighbor_list->pairs[i][1]); + pair_samples_ptr[i][2] = vesin_neighbor_list->shifts[i][0]; + pair_samples_ptr[i][3] = vesin_neighbor_list->shifts[i][1]; + pair_samples_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)); + pair_vectors.to(this->dtype_); + + auto neighbor_samples = torch::make_intrusive( + std::vector{ + "first_atom", "second_atom", "cell_shift_a", "cell_shift_b", "cell_shift_c" }, + pair_samples_values.to(device_)); + + auto neighbor_component = torch::make_intrusive( + std::vector{ "xyz" }, + torch::tensor({ 0, 1, 2 }, torch::TensorOptions().dtype(torch::kInt32).device(device_)) + .reshape({ 3, 1 })); + + auto neighbor_properties = torch::make_intrusive( + std::vector{ "distance" }, + torch::zeros({ 1, 1 }, torch::TensorOptions().dtype(torch::kInt32).device(device_))); + + return torch::make_intrusive( + pair_vectors.to(this->dtype_).to(this->device_), + neighbor_samples, + std::vector{ neighbor_component }, + neighbor_properties); +} + + +} // namespace gmx diff --git a/src/gromacs/applied_forces/metatomic/metatomic_forceprovider.h b/src/gromacs/applied_forces/metatomic/metatomic_forceprovider.h new file mode 100644 index 0000000000..30460cdfdd --- /dev/null +++ b/src/gromacs/applied_forces/metatomic/metatomic_forceprovider.h @@ -0,0 +1,117 @@ +/* + * This file is part of the GROMACS molecular simulation package. + * + * Copyright 2024- The GROMACS Authors + * and the project initiators Erik Lindahl, Berk Hess and David van der Spoel. + * Consult the AUTHORS/COPYING files and https://www.gromacs.org for details. + * + * GROMACS is free software; you can redistribute it and/or + * modify it under the terms of the GNU Lesser General Public License + * as published by the Free Software Foundation; either version 2.1 + * of the License, or (at your option) any later version. + * + * GROMACS is distributed in the hope that it will be useful, + * but WITHOUT ANY WARRANTY; without even the implied warranty of + * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU + * Lesser General Public License for more details. + * + * You should have received a copy of the GNU Lesser General Public + * License along with GROMACS; if not, see + * https://www.gnu.org/licenses, or write to the Free Software Foundation, + * Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA. + * + * If you want to redistribute modifications to GROMACS, please + * consider that scientific software is very special. Version + * control is crucial - bugs must be traceable. We will be happy to + * consider code for inclusion in the official distribution, but + * derived work must not be called official GROMACS. Details are found + * in the README & COPYING files - if they are missing, get the + * official version at https://www.gromacs.org. + * + * To help us fund GROMACS development, we humbly ask that you cite + * the research papers on the package. Check out https://www.gromacs.org. + */ +/*! \internal \file + * \brief + * Declares the Metatomic Force Provider class + * + * \author Metatensor developers + * \ingroup module_applied_forces + */ + +#pragma once + +#include "gromacs/mdtypes/iforceprovider.h" + +#include "metatensor.hpp" +// both gromacs and torch define `DIM`, which result in a conflict. We don't need either, +// so we undef before including headers. +#ifdef DIM +# undef DIM +#endif +#include "metatensor/torch.hpp" +#include "metatomic/torch.hpp" +#ifdef DIM +// XXX(rg): ask gromacs folks to do GMX_DIM +# undef DIM +#endif + +#include "metatomic_options.h" + +namespace gmx +{ + +struct MetatomicParameters; + +class MDLogger; +class MpiComm; + +/*! \brief \internal + * TODO + */ +class MetatomicForceProvider final : public IForceProvider +{ +public: + MetatomicForceProvider(const MetatomicOptions&, const MDLogger&, const MpiComm&); + ~MetatomicForceProvider(); + + /*! TODO + */ + void calculateForces(const ForceProviderInput& inputs, ForceProviderOutput* outputs) override; + void updateLocalAtoms(); + void gatherAtomPositions(ArrayRef globalPositions); + void gatherAtomNumbersIndices(); + + metatensor_torch::TensorBlock computeNeighbors(metatomic_torch::NeighborListOptions request, + long n_atoms, + const float* positions, + const matrix box, + bool periodic); + +private: + /// From NNPot + const MetatomicOptions& options_; + const MDLogger& logger_; + const MpiComm& mpiComm_; + torch::Device device_; + //! vector storing all atom positions + std::vector positions_; + + //! vector storing all atomic numbers + std::vector atomNumbers_; + + //! global index lookup table to map indices from model input to global atom indices + std::vector idxLookup_; + + //! local copy of simulation box + matrix box_; + /// From EON + torch::jit::Module model_; + metatomic_torch::ModelCapabilities capabilities_; + std::vector nl_requests_; + metatomic_torch::ModelEvaluationOptions evaluations_options_; + torch::ScalarType dtype_; + bool check_consistency_; +}; + +} // namespace gmx diff --git a/src/gromacs/applied_forces/metatomic/metatomic_forceprovider_stub.cpp b/src/gromacs/applied_forces/metatomic/metatomic_forceprovider_stub.cpp new file mode 100644 index 0000000000..669720f51e --- /dev/null +++ b/src/gromacs/applied_forces/metatomic/metatomic_forceprovider_stub.cpp @@ -0,0 +1,74 @@ +/* + * This file is part of the GROMACS molecular simulation package. + * + * Copyright 2024- The GROMACS Authors + * and the project initiators Erik Lindahl, Berk Hess and David van der Spoel. + * Consult the AUTHORS/COPYING files and https://www.gromacs.org for details. + * + * GROMACS is free software; you can redistribute it and/or + * modify it under the terms of the GNU Lesser General Public License + * as published by the Free Software Foundation; either version 2.1 + * of the License, or (at your option) any later version. + * + * GROMACS is distributed in the hope that it will be useful, + * but WITHOUT ANY WARRANTY; without even the implied warranty of + * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU + * Lesser General Public License for more details. + * + * You should have received a copy of the GNU Lesser General Public + * License along with GROMACS; if not, see + * https://www.gnu.org/licenses, or write to the Free Software Foundation, + * Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA. + * + * If you want to redistribute modifications to GROMACS, please + * consider that scientific software is very special. Version + * control is crucial - bugs must be traceable. We will be happy to + * consider code for inclusion in the official distribution, but + * derived work must not be called official GROMACS. Details are found + * in the README & COPYING files - if they are missing, get the + * official version at https://www.gromacs.org. + * + * To help us fund GROMACS development, we humbly ask that you cite + * the research papers on the package. Check out https://www.gromacs.org. + */ +/*! \internal \file + * \brief + * Stub implementation of the Metatomic Force Provider class. + * Compiled in case Libtorch/Metatomic backend is not linked. + * + * \author Metatensor developers + * \ingroup module_applied_forces + */ + +#include "gmxpre.h" + +#include "gromacs/utility/basedefinitions.h" +#include "gromacs/utility/exceptions.h" + +#include "metatomic_forceprovider.h" + +namespace gmx +{ + +CLANG_DIAGNOSTIC_IGNORE("-Wmissing-noreturn") + +MetatomicForceProvider::MetatomicForceProvider(const MetatomicParameters& params, + const MDLogger& logger, + const MpiComm& mpiComm) : + params_(params), logger_(logger), mpiComm_(mpiComm) +{ + GMX_THROW( + InternalError("This version of GROMACS can not use metatomic atomistic models. " + "Please reconfigure with `-DGMX_METATOMIC=ON` to enable this interface")); +} + +MetatomicForceProvider::~MetatomicForceProvider() {} + +void MetatomicForceProvider::calculateForces(const ForceProviderInput& /*inputs*/, + ForceProviderOutput* /*outputs*/) +{ +} + +CLANG_DIAGNOSTIC_RESET + +} // namespace gmx diff --git a/src/gromacs/applied_forces/metatomic/metatomic_mdmodule.cpp b/src/gromacs/applied_forces/metatomic/metatomic_mdmodule.cpp new file mode 100644 index 0000000000..0de93fb886 --- /dev/null +++ b/src/gromacs/applied_forces/metatomic/metatomic_mdmodule.cpp @@ -0,0 +1,221 @@ +/* + * This file is part of the GROMACS molecular simulation package. + * + * Copyright 2024- The GROMACS Authors + * and the project initiators Erik Lindahl, Berk Hess and David van der Spoel. + * Consult the AUTHORS/COPYING files and https://www.gromacs.org for details. + * + * GROMACS is free software; you can redistribute it and/or + * modify it under the terms of the GNU Lesser General Public License + * as published by the Free Software Foundation; either version 2.1 + * of the License, or (at your option) any later version. + * + * GROMACS is distributed in the hope that it will be useful, + * but WITHOUT ANY WARRANTY; without even the implied warranty of + * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU + * Lesser General Public License for more details. + * + * You should have received a copy of the GNU Lesser General Public + * License along with GROMACS; if not, see + * https://www.gnu.org/licenses, or write to the Free Software Foundation, + * Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA. + * + * If you want to redistribute modifications to GROMACS, please + * consider that scientific software is very special. Version + * control is crucial - bugs must be traceable. We will be happy to + * consider code for inclusion in the official distribution, but + * derived work must not be called official GROMACS. Details are found + * in the README & COPYING files - if they are missing, get the + * official version at https://www.gromacs.org. + * + * To help us fund GROMACS development, we humbly ask that you cite + * the research papers on the package. Check out https://www.gromacs.org. + */ +/*! \internal \file + * \brief + * Implements Metatomic MDModule class + * + * \author Metatensor developers + * \ingroup module_applied_forces + */ + + +#include "gmxpre.h" + +#include "metatomic_mdmodule.h" + +#include "gromacs/domdec/localatomset.h" +#include "gromacs/domdec/localatomsetmanager.h" +#include "gromacs/mdrunutility/mdmodulesnotifiers.h" +#include "gromacs/mdtypes/imdmodule.h" +#include "gromacs/utility/keyvaluetreebuilder.h" + +#include "metatomic_forceprovider.h" +#include "metatomic_options.h" + +namespace gmx +{ + +namespace +{ + +/*! \internal + * \brief Metatomic Module + * + * Class that implements the metatomic MDModule. + */ +class MetatomicMDModule final : public IMDModule +{ +public: + explicit MetatomicMDModule() = default; + + /*! \brief Requests to be notified during preprocessing. + * + * \param[in] notifiers allows the module to subscribe to notifications from MdModules. + * + * The Metatomic module subscribes to the following notifications: + * - The atom groups and their names from the index file (to specify the ML atoms) + * by taking a const IndexGroupsAndNames& as a parameter. + * - The system topology, which might be modified (e.g., to remove classical interactions). + * by taking a gmx_mtop_t* as a parameter. + * - Writing the module parameters to the KVT for storage in the .tpr file + * by taking a KeyValueTreeObjectBuilder as a parameter. + * - Accessing the MDLogger to log messages + * by taking a const MDLogger& as a parameter. + * - Accessing the WarningHandler to output warnings + * by taking a WarningHandler* as a parameter. + */ + void subscribeToPreProcessingNotifications(MDModulesNotifiers* notifiers) override + { + if (!options_.isActive()) + { + return; + } + + const auto setInputGroupIndicesFunction = [this](const IndexGroupsAndNames& indexGroupsAndNames) + { options_.setInputGroupIndices(indexGroupsAndNames); }; + notifiers->preProcessingNotifier_.subscribe(setInputGroupIndicesFunction); + + const auto modifyTopologyFunction = [this](gmx_mtop_t* top) { options_.modifyTopology(top); }; + notifiers->preProcessingNotifier_.subscribe(modifyTopologyFunction); + + const auto writeParamsToKvtFunction = [this](KeyValueTreeObjectBuilder kvt) + { options_.writeParamsToKvt(kvt); }; + notifiers->preProcessingNotifier_.subscribe(writeParamsToKvtFunction); + + const auto setLoggerFunction = [this](const MDLogger& logger) { options_.setLogger(logger); }; + notifiers->preProcessingNotifier_.subscribe(setLoggerFunction); + + const auto setWarningFunction = [this](WarningHandler* wi) + { options_.setWarningHandler(wi); }; + notifiers->preProcessingNotifier_.subscribe(setWarningFunction); + } + + /*! \brief Requests to be notified during simulation setup. + * + * \param[in] notifiers allows the module to subscribe to notifications from MdModules. + * + * The Metatomic module subscribes to the following notifications: + * - Reading the module parameters from the KVT + * by taking a const KeyValueTreeObject& as a parameter. + * - The local atom set manager to construct a local atom set for the ML atoms. + * by taking a LocalAtomSetManager* as a parameter. + * - The system topology + * by taking a const gmx_mtop_t& as a parameter. + * - The PBC type + * by taking a PbcType as a parameter. + * - The MPI communicator + * by taking a const MpiComm& as a parameter. + * - The MDLogger to log messages + * by taking a const MDLogger& as a parameter. + */ + void subscribeToSimulationSetupNotifications(MDModulesNotifiers* notifiers) override + { + if (!options_.isActive()) + { + return; + } + + const auto readParamsFromKvtFunction = [this](const KeyValueTreeObject& kvt) + { options_.readParamsFromKvt(kvt); }; + notifiers->simulationSetupNotifier_.subscribe(readParamsFromKvtFunction); + + const auto setLocalAtomSetFunction = [this](LocalAtomSetManager* localAtomSetManager) + { + // TODO(rg): why are these separate functions, they just create unique pointers.. + LocalAtomSet atomSet1 = localAtomSetManager->add(options_.parameters().mtaIndices_); + options_.setLocalInputAtomSet(atomSet1); + LocalAtomSet atomSet2 = localAtomSetManager->add(options_.parameters().mmIndices_); + options_.setLocalgmxMMAtomSet(atomSet2); + }; + notifiers->simulationSetupNotifier_.subscribe(setLocalAtomSetFunction); + + const auto setTopologyFunction = [this](const gmx_mtop_t& top) { options_.setTopology(top); }; + notifiers->simulationSetupNotifier_.subscribe(setTopologyFunction); + + const auto setPbcTypeFunction = [this](const PbcType& pbc) { options_.setPbcType(pbc); }; + notifiers->simulationSetupNotifier_.subscribe(setPbcTypeFunction); + + const auto setCommFunction = [this](const MpiComm& mpiComm) { options_.setComm(mpiComm); }; + notifiers->simulationSetupNotifier_.subscribe(setCommFunction); + + const auto setLoggerFunction = [this](const MDLogger& logger) { options_.setLogger(logger); }; + notifiers->simulationSetupNotifier_.subscribe(setLoggerFunction); + + // Request that GROMACS adds an energy term for our potential to the .edr file + const auto requestEnergyOutput = + [](MDModulesEnergyOutputToMetatomicPotRequestChecker* energyOutputRequest) + { energyOutputRequest->energyOutputToMetatomicPot_ = true; }; + notifiers->simulationSetupNotifier_.subscribe(requestEnergyOutput); + } + + /*! \brief Requests to be notified during the simulation. + * + * \param[in] notifiers allows the module to subscribe to notifications from MdModules. + * + * The Metatomic module subscribes to the following notifications: + * - Atom redistribution due to domain decomposition + * by taking a const MDModulesAtomsRedistributedSignal as a parameter. + */ + void subscribeToSimulationRunNotifications(MDModulesNotifiers* notifiers) override + { + if (!options_.isActive()) + { + return; + } + + // After domain decomposition, the force provider needs to know which atoms are local. + const auto notifyDDFunction = [this](const MDModulesAtomsRedistributedSignal& /*signal*/) + { force_provider_->gatherAtomNumbersIndices(); }; + notifiers->simulationRunNotifier_.subscribe(notifyDDFunction); + } + + void initForceProviders(ForceProviders* forceProviders) override + { + if (!options_.isActive()) + { + return; + } + + force_provider_ = std::make_unique( + options_, options_.logger(), options_.mpiComm()); + forceProviders->addForceProvider(force_provider_.get(), "Metatomic"); + } + + IMdpOptionProvider* mdpOptionProvider() override { return &options_; } + IMDOutputProvider* outputProvider() override { return nullptr; } + +private: + MetatomicOptions options_; + + std::unique_ptr force_provider_; +}; + +} // end anonymous namespace + +std::unique_ptr MetatomicModuleInfo::create() +{ + return std::make_unique(); +} + +} // end namespace gmx diff --git a/src/gromacs/applied_forces/metatomic/metatomic_mdmodule.h b/src/gromacs/applied_forces/metatomic/metatomic_mdmodule.h new file mode 100644 index 0000000000..dd0e6bfc72 --- /dev/null +++ b/src/gromacs/applied_forces/metatomic/metatomic_mdmodule.h @@ -0,0 +1,67 @@ +/* + * This file is part of the GROMACS molecular simulation package. + * + * Copyright 2024- The GROMACS Authors + * and the project initiators Erik Lindahl, Berk Hess and David van der Spoel. + * Consult the AUTHORS/COPYING files and https://www.gromacs.org for details. + * + * GROMACS is free software; you can redistribute it and/or + * modify it under the terms of the GNU Lesser General Public License + * as published by the Free Software Foundation; either version 2.1 + * of the License, or (at your option) any later version. + * + * GROMACS is distributed in the hope that it will be useful, + * but WITHOUT ANY WARRANTY; without even the implied warranty of + * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU + * Lesser General Public License for more details. + * + * You should have received a copy of the GNU Lesser General Public + * License along with GROMACS; if not, see + * https://www.gnu.org/licenses, or write to the Free Software Foundation, + * Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA. + * + * If you want to redistribute modifications to GROMACS, please + * consider that scientific software is very special. Version + * control is crucial - bugs must be traceable. We will be happy to + * consider code for inclusion in the official distribution, but + * derived work must not be called official GROMACS. Details are found + * in the README & COPYING files - if they are missing, get the + * official version at https://www.gromacs.org. + * + * To help us fund GROMACS development, we humbly ask that you cite + * the research papers on the package. Check out https://www.gromacs.org. + */ +/*! \internal \file + * \brief + * Declares factory structure for Metatomic MDModule class + * + * \author Metatensor developers + * \ingroup module_applied_forces + */ +#pragma once + +#include +#include + +namespace gmx +{ + +class IMDModule; + +/*! \internal + \brief Information about the Metatomic module. + * + * Provides name and method to create a metatomic potential module. + */ +struct MetatomicModuleInfo +{ + /*! \brief + * Creates a module for applying forces derived from a metatomic potential. + */ + static std::unique_ptr create(); + //! The name of the module + static constexpr std::string_view sc_name = "metatomic"; +}; + + +} // namespace gmx diff --git a/src/gromacs/applied_forces/metatomic/metatomic_options.cpp b/src/gromacs/applied_forces/metatomic/metatomic_options.cpp new file mode 100644 index 0000000000..733c2c8f42 --- /dev/null +++ b/src/gromacs/applied_forces/metatomic/metatomic_options.cpp @@ -0,0 +1,261 @@ +/* + * This file is part of the GROMACS molecular simulation package. + * + * Copyright 2024- The GROMACS Authors + * and the project initiators Erik Lindahl, Berk Hess and David van der Spoel. + * Consult the AUTHORS/COPYING files and https://www.gromacs.org for details. + * + * GROMACS is free software; you can redistribute it and/or + * modify it under the terms of the GNU Lesser General Public License + * as published by the Free Software Foundation; either version 2.1 + * of the License, or (at your option) any later version. + * + * GROMACS is distributed in the hope that it will be useful, + * but WITHOUT ANY WARRANTY; without even the implied warranty of + * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU + * Lesser General Public License for more details. + * + * You should have received a copy of the GNU Lesser General Public + * License along with GROMACS; if not, see + * https://www.gnu.org/licenses, or write to the Free Software Foundation, + * Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA. + * + * If you want to redistribute modifications to GROMACS, please + * consider that scientific software is very special. Version + * control is crucial - bugs must be traceable. We will be happy to + * consider code for inclusion in the official distribution, but + * derived work must not be called official GROMACS. Details are found + * in the README & COPYING files - if they are missing, get the + * official version at https://www.gromacs.org. + * + * To help us fund GROMACS development, we humbly ask that you cite + * the research papers on the package. Check out https://www.gromacs.org. + */ +/*! \internal \file + * \brief + * Implements the options for NNPot MDModule class. + * + * \author Metatensor developers + * \ingroup module_applied_forces + */ +// TODO(rg): Figure out how to insert the model into the .tpr file +#include "gmxpre.h" + +#include "metatomic_options.h" + +#include "gromacs/domdec/localatomset.h" +#include "gromacs/fileio/warninp.h" +#include "gromacs/options/basicoptions.h" +#include "gromacs/mdtypes/imdpoptionprovider_helpers.h" +#include "gromacs/options/optionsection.h" +#include "gromacs/selection/indexutil.h" +#include "gromacs/topology/mtop_util.h" +#include "gromacs/utility/keyvaluetreebuilder.h" +#include "gromacs/utility/keyvaluetreetransform.h" +#include "gromacs/utility/logger.h" +#include "gromacs/utility/mpicomm.h" +#include "gromacs/utility/strconvert.h" +#include "gromacs/utility/stringutil.h" + +#include "metatomic_topologypreprocessor.h" + +namespace gmx +{ + +static const std::string METATOMIC_MODULE_NAME = "metatomic"; + +/*! \brief Following Tags denotes names of parameters from .mdp file + * \note Changing this strings will break .tpr backwards compatibility + */ + +static const std::string ACTIVE_TAG = "active"; +static const std::string INPUT_GROUP_TAG = "input_group"; + +static const std::string MODEL_PATH_TAG = "model"; +static const std::string EXTENSIONS_DIRECTORY_TAG = "extensions"; +static const std::string CHECK_CONSISTENCY_TAG = "check_consistency"; +static const std::string DEVICE_TAG = "device"; + + +void MetatomicOptions::initMdpTransform(IKeyValueTreeTransformRules* rules) +{ + const auto& stringIdentityTransform = [](std::string s) { return s; }; + addMdpTransformFromString(rules, &fromStdString, METATOMIC_MODULE_NAME, ACTIVE_TAG); + addMdpTransformFromString( + rules, stringIdentityTransform, METATOMIC_MODULE_NAME, INPUT_GROUP_TAG); + + addMdpTransformFromString( + rules, stringIdentityTransform, METATOMIC_MODULE_NAME, MODEL_PATH_TAG); + addMdpTransformFromString( + rules, stringIdentityTransform, METATOMIC_MODULE_NAME, EXTENSIONS_DIRECTORY_TAG); + addMdpTransformFromString( + rules, &fromStdString, METATOMIC_MODULE_NAME, CHECK_CONSISTENCY_TAG); + addMdpTransformFromString(rules, stringIdentityTransform, METATOMIC_MODULE_NAME, DEVICE_TAG); +} + +void MetatomicOptions::initMdpOptions(IOptionsContainerWithSections* options) +{ + auto section = options->addSection(OptionSection(METATOMIC_MODULE_NAME.c_str())); + section.addOption(BooleanOption(ACTIVE_TAG.c_str()).store(¶ms_.active)); + section.addOption(StringOption(INPUT_GROUP_TAG.c_str()).store(¶ms_.inputGroup)); + + section.addOption(StringOption(MODEL_PATH_TAG.c_str()).store(¶ms_.modelPath_)); + section.addOption(StringOption(EXTENSIONS_DIRECTORY_TAG.c_str()).store(¶ms_.extensionsDirectory)); + section.addOption(StringOption(DEVICE_TAG.c_str()).store(¶ms_.device)); + section.addOption(BooleanOption(CHECK_CONSISTENCY_TAG.c_str()).store(¶ms_.checkConsistency)); +} + +void MetatomicOptions::buildMdpOutput(KeyValueTreeObjectBuilder* builder) const +{ + // new empty line before writing mdp values + // Use helper functions for MDP output + addMdpOutputComment(builder, METATOMIC_MODULE_NAME, "empty-line", ""); + + addMdpOutputComment(builder, + METATOMIC_MODULE_NAME, + "module", + "; Machine learning potential using metatomic"); + addMdpOutputValue(builder, METATOMIC_MODULE_NAME, ACTIVE_TAG, params_.active); + + if (params_.active) + { + addMdpOutputValue(builder, METATOMIC_MODULE_NAME, INPUT_GROUP_TAG, params_.inputGroup); + + addMdpOutputValue(builder, METATOMIC_MODULE_NAME, MODEL_PATH_TAG, params_.modelPath_); + addMdpOutputValue( + builder, METATOMIC_MODULE_NAME, EXTENSIONS_DIRECTORY_TAG, params_.extensionsDirectory); + addMdpOutputValue(builder, METATOMIC_MODULE_NAME, DEVICE_TAG, params_.device); + addMdpOutputValue( + builder, METATOMIC_MODULE_NAME, CHECK_CONSISTENCY_TAG, params_.checkConsistency); + } +} + +const MetatomicParameters& MetatomicOptions::parameters() +{ + return params_; +} + +bool MetatomicOptions::isActive() const +{ + return params_.active; +} + + +void MetatomicOptions::setInputGroupIndices(const IndexGroupsAndNames& indexGroupsAndNames) +{ + if (!params_.active) + { + return; + } + params_.mtaIndices_ = indexGroupsAndNames.indices(params_.inputGroup); + + if (params_.mtaIndices_.empty()) + { + GMX_THROW(InconsistentInputError(formatString( + "Group %s defining metatomic potential input atoms should not be empty.", + params_.inputGroup.c_str()))); + } +} + +void MetatomicOptions::modifyTopology(gmx_mtop_t* top) +{ + if (!params_.active) + { + return; + } + + MetatomicTopologyPreprocessor topPrep(params_.mtaIndices_); + topPrep.preprocess(top, logger(), wi_); +} + +void MetatomicOptions::writeParamsToKvt(KeyValueTreeObjectBuilder treeBuilder) +{ + if (!params_.active) + { + return; + } + + auto GroupIndexAdder = + treeBuilder.addUniformArray(METATOMIC_MODULE_NAME + "-" + INPUT_GROUP_TAG); + for (const auto& indexValue : params_.mtaIndices_) + { + GroupIndexAdder.addValue(indexValue); + } +} + +void MetatomicOptions::readParamsFromKvt(const KeyValueTreeObject& tree) +{ + if (!params_.active) + { + return; + } + + std::string key = METATOMIC_MODULE_NAME + "-" + INPUT_GROUP_TAG; + if (!tree.keyExists(key)) + { + GMX_THROW(InconsistentInputError( + "Cannot find input atoms index vector required for metatomic potential.\n" + "This could be caused by incompatible or corrupted tpr input file.")); + } + + auto kvtIndexArray = tree[key].asArray().values(); + params_.mtaIndices_.resize(kvtIndexArray.size()); + std::transform(std::begin(kvtIndexArray), + std::end(kvtIndexArray), + std::begin(params_.mtaIndices_), + [](const KeyValueTreeValue& val) { return val.cast(); }); +} + + +void MetatomicOptions::setLogger(const MDLogger& logger) +{ + logger_ = &logger; +} + +void MetatomicOptions::setWarningHandler(WarningHandler* wi) +{ + wi_ = wi; +} + +void MetatomicOptions::setTopology(const gmx_mtop_t& top) +{ + params_.atoms_ = gmx_mtop_global_atoms(top); + params_.numAtoms_ = params_.atoms_.nr; +} + +void MetatomicOptions::setPbcType(const PbcType& pbcType) +{ + params_.pbcType_ = std::make_unique(pbcType); +} + +void MetatomicOptions::setComm(const MpiComm& mpiComm) +{ + mpiComm_ = &mpiComm; +} + + +const MDLogger& MetatomicOptions::logger() const +{ + GMX_RELEASE_ASSERT(logger_, "Logger not set for MetatomicOptions."); + return *logger_; +} + +const MpiComm& MetatomicOptions::mpiComm() const +{ + GMX_RELEASE_ASSERT(mpiComm_, "MPI communicator not set for MetatomicOptions."); + return *mpiComm_; +} + + +void MetatomicOptions::setLocalInputAtomSet(const LocalAtomSet& localInputAtomSet) +{ + params_.mtaAtoms_ = std::make_unique(localInputAtomSet); +} + +void MetatomicOptions::setLocalgmxMMAtomSet(const LocalAtomSet& localMMAtomSet) +{ + params_.gmxMMAtoms_ = std::make_unique(localMMAtomSet); +} + + +} // namespace gmx diff --git a/src/gromacs/applied_forces/metatomic/metatomic_options.h b/src/gromacs/applied_forces/metatomic/metatomic_options.h new file mode 100644 index 0000000000..1b747ffe8c --- /dev/null +++ b/src/gromacs/applied_forces/metatomic/metatomic_options.h @@ -0,0 +1,137 @@ +/* + * This file is part of the GROMACS molecular simulation package. + * + * Copyright 2024- The GROMACS Authors + * and the project initiators Erik Lindahl, Berk Hess and David van der Spoel. + * Consult the AUTHORS/COPYING files and https://www.gromacs.org for details. + * + * GROMACS is free software; you can redistribute it and/or + * modify it under the terms of the GNU Lesser General Public License + * as published by the Free Software Foundation; either version 2.1 + * of the License, or (at your option) any later version. + * + * GROMACS is distributed in the hope that it will be useful, + * but WITHOUT ANY WARRANTY; without even the implied warranty of + * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU + * Lesser General Public License for more details. + * + * You should have received a copy of the GNU Lesser General Public + * License along with GROMACS; if not, see + * https://www.gnu.org/licenses, or write to the Free Software Foundation, + * Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA. + * + * If you want to redistribute modifications to GROMACS, please + * consider that scientific software is very special. Version + * control is crucial - bugs must be traceable. We will be happy to + * consider code for inclusion in the official distribution, but + * derived work must not be called official GROMACS. Details are found + * in the README & COPYING files - if they are missing, get the + * official version at https://www.gromacs.org. + * + * To help us fund GROMACS development, we humbly ask that you cite + * the research papers on the package. Check out https://www.gromacs.org. + */ +/*! \internal \file + * \brief + * Declares the options for Metatomic MDModule class, + * set during pre-processing in the .mdp-file. + * + * \author Metatensor developers + * \ingroup module_applied_forces + */ +#pragma once + +#include +#include +#include + +#include "gromacs/mdtypes/imdpoptionprovider.h" +#include "gromacs/topology/atoms.h" +#include "gromacs/utility/vectypes.h" + +// some forward declarations +struct gmx_mtop_t; +class WarningHandler; +enum class PbcType; + + +namespace gmx +{ +class MDLogger; +class IOptionsContainerWithSections; +class IKeyValueTreeTransformRules; +class KeyValueTreeObjectBuilder; +class KeyValueTreeObject; +class IndexGroupsAndNames; +class LocalAtomSet; +class MpiComm; + + +//! TODO +struct MetatomicParameters +{ + //! Is the metatomic force provider enabled? + bool active = false; + + //! path to the exported metatomic model file + std::string modelPath_; + //! path to a directory where extensions will be present at MD time + std::string extensionsDirectory; + //! should metatomic run additional checks on the models inputs & outputs? + bool checkConsistency = true; + //! Torch device to use to run the model. If left empty, this is defined + //! based on the model declared preferences + std::string device; + + // TODO: how should we translate atomic types? + + //! stores atom group name for which metatomic should compute the energy + //! (default whole System) + std::string inputGroup = "System"; + + std::vector mtaIndices_; + std::vector mmIndices_; + std::unique_ptr mtaAtoms_; + std::unique_ptr gmxMMAtoms_; + t_atoms atoms_; + int numAtoms_; + std::unique_ptr pbcType_; +}; + +class MetatomicOptions final : public IMdpOptionProvider +{ +public: + MetatomicParameters params_; + // ^---------- lol + void initMdpTransform(IKeyValueTreeTransformRules* rules) override; + void initMdpOptions(IOptionsContainerWithSections* options) override; + void buildMdpOutput(KeyValueTreeObjectBuilder* builder) const override; + + bool isActive() const; + void setInputGroupIndices(const IndexGroupsAndNames&); + void modifyTopology(gmx_mtop_t*); + void writeParamsToKvt(KeyValueTreeObjectBuilder); + void readParamsFromKvt(const KeyValueTreeObject&); + void setLogger(const MDLogger&); + void setWarningHandler(WarningHandler*); + void setTopology(const gmx_mtop_t&); + void setPbcType(const PbcType&); + void setComm(const MpiComm&); + //! set local atom set for Metatomic input during simulation setup + void setLocalInputAtomSet(const LocalAtomSet&); + + //! set local MM atom set during simulation setup + void setLocalgmxMMAtomSet(const LocalAtomSet&); + + const MetatomicParameters& parameters(); + const MDLogger& logger() const; + const MpiComm& mpiComm() const; + + +private: + const MDLogger* logger_ = nullptr; + const MpiComm* mpiComm_ = nullptr; + WarningHandler* wi_ = nullptr; +}; + +} // namespace gmx diff --git a/src/gromacs/applied_forces/metatomic/metatomic_topologypreprocessor.cpp b/src/gromacs/applied_forces/metatomic/metatomic_topologypreprocessor.cpp new file mode 100644 index 0000000000..4564088724 --- /dev/null +++ b/src/gromacs/applied_forces/metatomic/metatomic_topologypreprocessor.cpp @@ -0,0 +1,96 @@ +/* + * This file is part of the GROMACS molecular simulation package. + * + * Copyright 2024- The GROMACS Authors + * and the project initiators Erik Lindahl, Berk Hess and David van der Spoel. + * Consult the AUTHORS/COPYING files and https://www.gromacs.org for details. + * + * GROMACS is free software; you can redistribute it and/or + * modify it under the terms of the GNU Lesser General Public License + * as published by the Free Software Foundation; either version 2.1 + * of the License, or (at your option) any later version. + * + * GROMACS is distributed in the hope that it will be useful, + * but WITHOUT ANY WARRANTY; without even the implied warranty of + * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU + * Lesser General Public License for more details. + * + * You should have received a copy of the GNU Lesser General Public + * License along with GROMACS; if not, see + * https://www.gnu.org/licenses, or write to the Free Software Foundation, + * Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA. + * + * If you want to redistribute modifications to GROMACS, please + * consider that scientific software is very special. Version + * control is crucial - bugs must be traceable. We will be happy to + * consider code for inclusion in the official distribution, but + * derived work must not be called official GROMACS. Details are found + * in the README & COPYING files - if they are missing, get the + * official version at https://www.gromacs.org. + * + * To help us fund GROMACS development, we humbly ask that you cite + * the research papers on the package. Check out https://www.gromacs.org. + */ +/*! \internal \file + * \brief + * Implements MetatomicTopologyPrepocessor. + * + * \author Metatensor developers + * \ingroup module_applied_forces + */ + +#include "gromacs/fileio/warninp.h" +#include "gromacs/topology/topology.h" +#include "gromacs/utility/logger.h" + +#include "metatomic_topologypreprocessor.h" + +namespace gmx +{ + +MetatomicTopologyPreprocessor::MetatomicTopologyPreprocessor(ArrayRef inputIndices) : + QMMMTopologyPreprocessor(inputIndices), mtaIndices_(inputIndices.begin(), inputIndices.end()) +{ +} + +void MetatomicTopologyPreprocessor::preprocess(gmx_mtop_t* mtop, const MDLogger& logger, WarningHandler* wi) +{ + // We're re-using the topology-modifying functions from QMMM module for now, + // since they contain the same modifications as needed for ML/MM with metatomic. This should be + // refactored in the future. + GMX_LOG(logger.info) + .appendText("Metatomic potential Interface is active, topology was modified!"); + + GMX_LOG(logger.info) + .appendTextFormatted( + "Number of embedded metatomic atoms: %td\nNumber of regular atoms: %td\n", + gmx::ssize(mtaIndices_), + mtop->natoms - gmx::ssize(mtaIndices_)); + + // 1) Split molecules containing Metatomic input atoms from other molecules in blocks + std::vector isMTABlock = splitQMBlocks(mtop, mtaIndices_); + + // 2) Exclude LJ interactions between MTA atoms + // this also excludes coulomb interactions + addQMLJExclusions(mtop, mtaIndices_, logger); + + // 3) Build atomNumbers vector with atomic numbers of all atoms + buildQMMMAtomNumbers(*mtop); + + // 4) Make F_CONNBOND between atoms within MTA region + modifyQMMMTwoCenterInteractions(mtop, mtaIndices_, isMTABlock, logger); + + // 5) Remove angles and settles containing 2 or more MTA atoms + modifyQMMMThreeCenterInteractions(mtop, mtaIndices_, isMTABlock, logger); + + // 6) Remove dihedrals containing 3 or more MTA atoms + modifyQMMMFourCenterInteractions(mtop, mtaIndices_, isMTABlock, logger); + + // 7) Check for constrained bonds in QM subsystem + checkConstrainedBonds(mtop, mtaIndices_, isMTABlock, wi); + + // finalize topology + mtop->finalize(); +} + +} // namespace gmx diff --git a/src/gromacs/applied_forces/metatomic/metatomic_topologypreprocessor.h b/src/gromacs/applied_forces/metatomic/metatomic_topologypreprocessor.h new file mode 100644 index 0000000000..e8e97aa825 --- /dev/null +++ b/src/gromacs/applied_forces/metatomic/metatomic_topologypreprocessor.h @@ -0,0 +1,95 @@ +/* + * This file is part of the GROMACS molecular simulation package. + * + * Copyright 2024- The GROMACS Authors + * and the project initiators Erik Lindahl, Berk Hess and David van der Spoel. + * Consult the AUTHORS/COPYING files and https://www.gromacs.org for details. + * + * GROMACS is free software; you can redistribute it and/or + * modify it under the terms of the GNU Lesser General Public License + * as published by the Free Software Foundation; either version 2.1 + * of the License, or (at your option) any later version. + * + * GROMACS is distributed in the hope that it will be useful, + * but WITHOUT ANY WARRANTY; without even the implied warranty of + * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU + * Lesser General Public License for more details. + * + * You should have received a copy of the GNU Lesser General Public + * License along with GROMACS; if not, see + * https://www.gnu.org/licenses, or write to the Free Software Foundation, + * Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA. + * + * If you want to redistribute modifications to GROMACS, please + * consider that scientific software is very special. Version + * control is crucial - bugs must be traceable. We will be happy to + * consider code for inclusion in the official distribution, but + * derived work must not be called official GROMACS. Details are found + * in the README & COPYING files - if they are missing, get the + * official version at https://www.gromacs.org. + * + * To help us fund GROMACS development, we humbly ask that you cite + * the research papers on the package. Check out https://www.gromacs.org. + */ +/*! \internal \file + * \brief + * MetatomicTopologyPrepocessor class responsible for + * all modificatios of the topology during input pre-processing. + * Inherits from QMMMTopologyPreprocessor. + * + * \author Metatensor developers + * \ingroup module_applied_forces + */ +#pragma once + +#include "gromacs/applied_forces/qmmm/qmmmtopologypreprocessor.h" +#include "gromacs/utility/arrayref.h" + +struct gmx_mtop_t; + +namespace gmx +{ + +/*! \brief + * Class responsible for all modifications of the topology during input pre-processing. + * + * Inherits from QMMMTopologyPreprocessor and performs the following modifications + * for a scheme analoguous to mechanical embedding in QMMM: + * 1) Split molecules containing Metatomic input atoms from other molecules in blocks + * 2) Exclude non-bonded interactions between Metatomic atoms. + * 3) Build vector with atomic numbers of all atoms. + * 4) Make F_CONNBOND between atoms within Metatomic region. + * 5) Remove angles and settles containing 2 or more Metatomic atoms. + * 6) Remove dihedrals containing 3 or more Metatomic atoms. + * + * Is a clone of the NNPot preprocessor + */ +class MetatomicTopologyPreprocessor : public QMMMTopologyPreprocessor +{ +public: + /*! \brief Constructor for MetatomicTopologyPreprocessor from its parameters + * + * \param[in] inputIndices Array with global indices of NN input atoms + */ + MetatomicTopologyPreprocessor(ArrayRef inputIndices); + + /*! \brief Process mtop topology + * containing information about topology modifications. + * + * \param[in,out] mtop Topology that needs to be modified + * \param[in] logger MDLogger for logging info about modifications + * \param[in] wi WarningHandler for handling warnings + */ + void preprocess(gmx_mtop_t* mtop, const MDLogger& logger, WarningHandler* wi); + +private: + /*! \brief Global indices of metatomic atoms; + * The dominant operation is search and we also expect the set of metatomic atoms to be very small + * relative to the rest, so set should outperform unordered set. + */ + std::set mtaIndices_; + //! Vector with atom numbers for the whole system + std::vector atomNumbers_; +}; + +} // namespace gmx diff --git a/src/gromacs/applied_forces/metatomic/tests/CMakeLists.txt b/src/gromacs/applied_forces/metatomic/tests/CMakeLists.txt new file mode 100644 index 0000000000..5d306d64ed --- /dev/null +++ b/src/gromacs/applied_forces/metatomic/tests/CMakeLists.txt @@ -0,0 +1,40 @@ +# +# This file is part of the GROMACS molecular simulation package. +# +# Copyright 2024- The GROMACS Authors +# and the project initiators Erik Lindahl, Berk Hess and David van der Spoel. +# Consult the AUTHORS/COPYING files and https://www.gromacs.org for details. +# +# GROMACS is free software; you can redistribute it and/or +# modify it under the terms of the GNU Lesser General Public License +# as published by the Free Software Foundation; either version 2.1 +# of the License, or (at your option) any later version. +# +# GROMACS is distributed in the hope that it will be useful, +# but WITHOUT ANY WARRANTY; without even the implied warranty of +# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU +# Lesser General Public License for more details. +# +# You should have received a copy of the GNU Lesser General Public +# License along with GROMACS; if not, see +# https://www.gnu.org/licenses, or write to the Free Software Foundation, +# Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA. +# +# If you want to redistribute modifications to GROMACS, please +# consider that scientific software is very special. Version +# control is crucial - bugs must be traceable. We will be happy to +# consider code for inclusion in the official distribution, but +# derived work must not be called official GROMACS. Details are found +# in the README & COPYING files - if they are missing, get the +# official version at https://www.gromacs.org. +# +# To help us fund GROMACS development, we humbly ask that you cite +# the research papers on the package. Check out https://www.gromacs.org. + +gmx_add_unit_test(MTAAppliedForcesUnitTest mta_applied_forces-test + CPP_SOURCE_FILES + # metatomic_forceprovider.cpp + metatomic_options.cpp +) + +target_link_libraries(mta_applied_forces-test PRIVATE applied_forces math topology) diff --git a/src/gromacs/applied_forces/metatomic/tests/metatomic_options.cpp b/src/gromacs/applied_forces/metatomic/tests/metatomic_options.cpp new file mode 100644 index 0000000000..5c816fe120 --- /dev/null +++ b/src/gromacs/applied_forces/metatomic/tests/metatomic_options.cpp @@ -0,0 +1,221 @@ + +/* + * This file is part of the GROMACS molecular simulation package. + * + * Copyright 2024- The GROMACS Authors + * and the project initiators Erik Lindahl, Berk Hess and David van der Spoel. + * Consult the AUTHORS/COPYING files and https://www.gromacs.org for details. + * + * GROMACS is free software; you can redistribute it and/or + * modify it under the terms of the GNU Lesser General Public License + * as published by the Free Software Foundation; either version 2.1 + * of the License, or (at your option) any later version. + * + * GROMACS is distributed in the hope that it will be useful, + * but WITHOUT ANY WARRANTY; without even the implied warranty of + * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU + * Lesser General Public License for more details. + * + * You should have received a copy of the GNU Lesser General Public + * License along with GROMACS; if not, see + * https://www.gnu.org/licenses, or write to the Free Software Foundation, + * Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA. + * + * If you want to redistribute modifications to GROMACS, please + * consider that scientific software is very special. Version + * control is crucial - bugs must be traceable. We will be happy to + * consider code for inclusion in the official distribution, but + * derived work must not be called official GROMACS. Details are found + * in the README & COPYING files - if they are missing, get the + * official version at https://www.gromacs.org. + * + * To help us fund GROMACS development, we humbly ask that you cite + * the research papers on the package. Check out https://www.gromacs.org. + */ +/*! \internal \file + * \brief + * Tests for functionality of the NNPotOptions + * + * \author Lukas Müllender + * \ingroup module_applied_forces + */ + +#include "gmxpre.h" + +#include "gromacs/applied_forces/metatomic/metatomic_options.h" + +#include + +#include "gromacs/applied_forces/nnpot/nnpot.h" +#include "gromacs/domdec/localatomset.h" +#include "gromacs/fileio/warninp.h" +#include "gromacs/mdrunutility/mdmodulesnotifiers.h" +#include "gromacs/mdtypes/imdpoptionprovider_test_helper.h" +#include "gromacs/selection/indexutil.h" +#include "gromacs/topology/index.h" +#include "gromacs/utility/arrayref.h" +#include "gromacs/utility/keyvaluetree.h" +#include "gromacs/utility/keyvaluetreebuilder.h" +#include "gromacs/utility/keyvaluetreemdpwriter.h" +#include "gromacs/utility/logger.h" +#include "gromacs/utility/stringstream.h" +#include "gromacs/utility/textwriter.h" + +#include "testutils/refdata.h" +#include "testutils/testasserts.h" +#include "testutils/testfilemanager.h" +#include "testutils/testmatchers.h" + +namespace gmx +{ + +namespace test +{ + +// Module name from metatomic_options.cpp +static const std::string METATOMIC_MODULE_NAME = "metatomic"; + +class MetatomicOptionsTest : public ::testing::Test +{ +public: + static KeyValueTreeObject metatomicBuildDefaultMdpValues() + { + // Prepare MDP inputs + KeyValueTreeBuilder mdpValueBuilder; + mdpValueBuilder.rootObject().addValue(METATOMIC_MODULE_NAME + "-active", std::string("true")); + return mdpValueBuilder.build(); + } + + static KeyValueTreeObject metatomicBuildInputMdpValues() + { + // Prepare MDP inputs + KeyValueTreeBuilder mdpValueBuilder; + mdpValueBuilder.rootObject().addValue(METATOMIC_MODULE_NAME + "-active", std::string("true")); + mdpValueBuilder.rootObject().addValue( + METATOMIC_MODULE_NAME + "-model", + gmx::test::TestFileManager::getInputFilePath("model.pt").string()); + mdpValueBuilder.rootObject().addValue(METATOMIC_MODULE_NAME + "-input_group", + std::string("System")); + mdpValueBuilder.rootObject().addValue(METATOMIC_MODULE_NAME + "-extensions", std::string("./ext")); + mdpValueBuilder.rootObject().addValue(METATOMIC_MODULE_NAME + "-device", std::string("cpu")); + mdpValueBuilder.rootObject().addValue(METATOMIC_MODULE_NAME + "-check_consistency", + std::string("true")); + return mdpValueBuilder.build(); + } + + static IndexGroupsAndNames indexGroupsAndNamesGeneric() + { + // System group is default + std::vector indexGroups; + indexGroups.push_back({ "A", { 1 } }); + indexGroups.push_back({ "System", { 1, 2, 3 } }); + indexGroups.push_back({ "C", { 2, 3 } }); + + return IndexGroupsAndNames(indexGroups); + } +}; + +TEST_F(MetatomicOptionsTest, DefaultParameters) +{ + MetatomicOptions metatomicOptions; + const MetatomicParameters& defaultParams = metatomicOptions.parameters(); + gmx::test::TestReferenceData data; + gmx::test::TestReferenceChecker checker(data.rootChecker()); + + checker.checkBoolean(defaultParams.active, "active"); + checker.checkString(defaultParams.inputGroup, "inputGroup"); + checker.checkString(defaultParams.modelPath_, "modelPath"); + checker.checkString(defaultParams.extensionsDirectory, "extensionsDirectory"); + checker.checkString(defaultParams.device, "device"); + checker.checkBoolean(defaultParams.checkConsistency, "checkConsistency"); +} + +TEST_F(MetatomicOptionsTest, OptionSetsActive) +{ + MetatomicOptions metatomicOptions; + EXPECT_FALSE(metatomicOptions.parameters().active); + test::fillOptionsFromMdpValues(metatomicBuildDefaultMdpValues(), &metatomicOptions); + EXPECT_TRUE(metatomicOptions.parameters().active); +} + +TEST_F(MetatomicOptionsTest, OutputNoDefaultValuesWhenInactive) +{ + // Transform module data into a flat key-value tree for output. + StringOutputStream stream; + KeyValueTreeBuilder builder; + KeyValueTreeObjectBuilder builderObject = builder.rootObject(); + + MetatomicOptions metatomicOptions; + metatomicOptions.buildMdpOutput(&builderObject); + { + TextWriter writer(&stream); + writeKeyValueTreeAsMdp(&writer, builder.build()); + } + stream.close(); + + gmx::test::TestReferenceData data; + gmx::test::TestReferenceChecker checker(data.rootChecker()); + + checker.checkString(stream.toString(), "Mdp output"); +} + +TEST_F(MetatomicOptionsTest, OutputDefaultValuesWhenActive) +{ + // Set metatomic-active = true + MetatomicOptions metatomicOptions; + test::fillOptionsFromMdpValues(metatomicBuildDefaultMdpValues(), &metatomicOptions); + + // Transform module data into a flat key-value tree for output. + StringOutputStream stream; + KeyValueTreeBuilder builder; + KeyValueTreeObjectBuilder builderObject = builder.rootObject(); + + metatomicOptions.buildMdpOutput(&builderObject); + { + TextWriter writer(&stream); + writeKeyValueTreeAsMdp(&writer, builder.build()); + } + stream.close(); + + gmx::test::TestReferenceData data; + gmx::test::TestReferenceChecker checker(data.rootChecker()); + + checker.checkString(stream.toString(), "Mdp output"); +} + +TEST_F(MetatomicOptionsTest, InternalsToKvtAndBack) +{ + // Set metatomic-active = true and other params + MetatomicOptions metatomicOptions; + fillOptionsFromMdpValues(metatomicBuildInputMdpValues(), &metatomicOptions); + + // Set indices + const IndexGroupsAndNames indexGroupAndNames = indexGroupsAndNamesGeneric(); + metatomicOptions.setInputGroupIndices(indexGroupAndNames); + + // Set dummy logger and warning handler + MDLogger logger; + metatomicOptions.setLogger(logger); + WarningHandler warninp(true, 0); + metatomicOptions.setWarningHandler(&warninp); + + // Copy internal parameters + const MetatomicParameters& params = metatomicOptions.parameters(); + auto mtaIndicesBefore = params.mtaIndices_; + + KeyValueTreeBuilder builder; + // MetatomicOptions::writeParamsToKvt doesn't have external dependencies like + // GMX_TORCH, so we can call it directly. + EXPECT_NO_THROW(metatomicOptions.writeParamsToKvt(builder.rootObject())); + const auto inputTree = builder.build(); + + EXPECT_NO_THROW(metatomicOptions.readParamsFromKvt(inputTree)); + + // Check Internal parameters taken back from KVT + const MetatomicParameters& params2 = metatomicOptions.parameters(); + EXPECT_EQ(mtaIndicesBefore, params2.mtaIndices_); +} + +} // namespace test + +} // namespace gmx diff --git a/src/gromacs/applied_forces/metatomic/tests/refdata/MetatomicOptionsTest_DefaultParameters.xml b/src/gromacs/applied_forces/metatomic/tests/refdata/MetatomicOptionsTest_DefaultParameters.xml new file mode 100644 index 0000000000..37e4ccb557 --- /dev/null +++ b/src/gromacs/applied_forces/metatomic/tests/refdata/MetatomicOptionsTest_DefaultParameters.xml @@ -0,0 +1,11 @@ + + + + false + model.pt + System + + + + + diff --git a/src/gromacs/applied_forces/metatomic/tests/refdata/MetatomicOptionsTest_OutputDefaultValuesWhenActive.xml b/src/gromacs/applied_forces/metatomic/tests/refdata/MetatomicOptionsTest_OutputDefaultValuesWhenActive.xml new file mode 100644 index 0000000000..1635c668d5 --- /dev/null +++ b/src/gromacs/applied_forces/metatomic/tests/refdata/MetatomicOptionsTest_OutputDefaultValuesWhenActive.xml @@ -0,0 +1,13 @@ + + + + +; Machine learning potential using metatomic +metatomic-active = true +metatomic-input_group = System +metatomic-model = +metatomic-extensions = +metatomic-device = +metatomic-check_consistency = true + + diff --git a/src/gromacs/applied_forces/metatomic/tests/refdata/MetatomicOptionsTest_OutputNoDefaultValuesWhenInactive.xml b/src/gromacs/applied_forces/metatomic/tests/refdata/MetatomicOptionsTest_OutputNoDefaultValuesWhenInactive.xml new file mode 100644 index 0000000000..c5c4bbdd90 --- /dev/null +++ b/src/gromacs/applied_forces/metatomic/tests/refdata/MetatomicOptionsTest_OutputNoDefaultValuesWhenInactive.xml @@ -0,0 +1,8 @@ + + + + +; Machine learning potential using metatomic +metatomic-active = false + + diff --git a/src/gromacs/fileio/tpxio.cpp b/src/gromacs/fileio/tpxio.cpp index 5b75ff774c..be98628614 100644 --- a/src/gromacs/fileio/tpxio.cpp +++ b/src/gromacs/fileio/tpxio.cpp @@ -199,6 +199,7 @@ enum tpxv tpxv_RefScaleMultipleCOMs, /**< Add multiple COM groups for refcoord-scale */ tpxv_InputHistogramCounts, /**< Provide input histogram counts for current expanded ensemble state */ tpxv_NNPotIFuncType, /**< Add interaction function type for neural network potential */ + tpxv_MetatomicPotIFuncType, /**< Add interaction function type metatomic potentials */ tpxv_Count /**< the total number of tpxv versions */ }; @@ -294,6 +295,7 @@ static const t_ftupd ftupd[] = { { tpxv_VSite2FD, F_VSITE2FD }, { tpxv_GenericInternalParameters, F_DENSITYFITTING }, { tpxv_NNPotIFuncType, F_ENNPOT }, + { tpxv_MetatomicPotIFuncType, F_EMETATOMICPOT }, { tpxv_Pre96Version69, F_VTEMP_NOLONGERUSED }, { tpxv_Pre96Version66, F_PDISPCORR }, { tpxv_Pre96Version79, F_DVDL_COUL }, diff --git a/src/gromacs/listed_forces/bonded.cpp b/src/gromacs/listed_forces/bonded.cpp index 75991c04fd..27a8ce5fbe 100644 --- a/src/gromacs/listed_forces/bonded.cpp +++ b/src/gromacs/listed_forces/bonded.cpp @@ -4066,6 +4066,7 @@ constexpr std::array c_bondedInteractionFunctions = { BondedInteractions{ unimplemented, -1 }, // F_DENSITYFITTING BondedInteractions{ unimplemented, -1 }, // F_EQM BondedInteractions{ unimplemented, -1 }, // F_ENNPOT + BondedInteractions{ unimplemented, -1 }, // F_EMETATOMICPOT BondedInteractions{ unimplemented, -1 }, // F_EPOT BondedInteractions{ unimplemented, -1 }, // F_EKIN BondedInteractions{ unimplemented, -1 }, // F_ETOT diff --git a/src/gromacs/mdlib/energyoutput.cpp b/src/gromacs/mdlib/energyoutput.cpp index 22e775fa32..f4444678ef 100644 --- a/src/gromacs/mdlib/energyoutput.cpp +++ b/src/gromacs/mdlib/energyoutput.cpp @@ -282,6 +282,11 @@ EnergyOutput::EnergyOutput(ener_file* fp_ene, bEner_[F_ENNPOT] = mdModulesAddOutputToNNPotFieldRequest.energyOutputToNNPot_; + MDModulesEnergyOutputToMetatomicPotRequestChecker mdModulesAddOutputToMetatomicPotFieldRequest; + mdModulesNotifiers.simulationSetupNotifier_.notify(&mdModulesAddOutputToMetatomicPotFieldRequest); + + bEner_[F_EMETATOMICPOT] = mdModulesAddOutputToMetatomicPotFieldRequest.energyOutputToMetatomicPot_; + // Counting the energy terms that will be printed and saving their names f_nre_ = 0; for (i = 0; i < F_NRE; i++) diff --git a/src/gromacs/mdrun/mdmodules.cpp b/src/gromacs/mdrun/mdmodules.cpp index 911c0d4d8b..edac2e39bf 100644 --- a/src/gromacs/mdrun/mdmodules.cpp +++ b/src/gromacs/mdrun/mdmodules.cpp @@ -44,6 +44,7 @@ #include "gromacs/applied_forces/colvars/colvarsMDModule.h" #include "gromacs/applied_forces/densityfitting/densityfitting.h" #include "gromacs/applied_forces/electricfield.h" +#include "gromacs/applied_forces/metatomic/metatomic_mdmodule.h" #include "gromacs/applied_forces/nnpot/nnpot.h" #include "gromacs/applied_forces/plumed/plumedMDModule.h" #include "gromacs/applied_forces/qmmm/qmmm.h" @@ -82,10 +83,11 @@ class MDModules::Impl : public IMDOutputProvider InteractiveMolecularDynamicsModuleInfo::create(); modules_[std::string(QMMMModuleInfo::sc_name)] = QMMMModuleInfo::create(); modules_[std::string(SwapCoordinatesModuleInfo::sc_name)] = SwapCoordinatesModuleInfo::create(); - modules_[std::string(ColvarsModuleInfo::sc_name)] = ColvarsModuleInfo::create(); - modules_[std::string(PlumedModuleInfo::sc_name)] = PlumedModuleInfo::create(); - modules_[std::string(NNPotModuleInfo::sc_name)] = NNPotModuleInfo::create(); - modules_[std::string(FmmModuleInfo::sc_name)] = FmmModuleInfo::create(); + modules_[std::string(ColvarsModuleInfo::sc_name)] = ColvarsModuleInfo::create(); + modules_[std::string(PlumedModuleInfo::sc_name)] = PlumedModuleInfo::create(); + modules_[std::string(NNPotModuleInfo::sc_name)] = NNPotModuleInfo::create(); + modules_[std::string(FmmModuleInfo::sc_name)] = FmmModuleInfo::create(); + modules_[std::string(MetatomicModuleInfo::sc_name)] = MetatomicModuleInfo::create(); } void makeModuleOptions(Options* options) const @@ -102,7 +104,8 @@ class MDModules::Impl : public IMDOutputProvider DensityFittingModuleInfo::sc_name, QMMMModuleInfo::sc_name, ColvarsModuleInfo::sc_name, - NNPotModuleInfo::sc_name }) + NNPotModuleInfo::sc_name, + MetatomicModuleInfo::sc_name }) { IMDModule* module = modules_.at(std::string(moduleName)).get(); IMdpOptionProvider* mdpOptionProvider = module->mdpOptionProvider(); @@ -167,7 +170,8 @@ void MDModules::initMdpTransform(IKeyValueTreeTransformRules* rules) DensityFittingModuleInfo::sc_name, QMMMModuleInfo::sc_name, ColvarsModuleInfo::sc_name, - NNPotModuleInfo::sc_name }) + NNPotModuleInfo::sc_name, + MetatomicModuleInfo::sc_name }) { IMDModule* module = impl_->modules_.at(std::string(moduleName)).get(); IMdpOptionProvider* mdpOptionProvider = module->mdpOptionProvider(); @@ -191,7 +195,8 @@ void MDModules::buildMdpOutput(KeyValueTreeObjectBuilder* builder) DensityFittingModuleInfo::sc_name, QMMMModuleInfo::sc_name, ColvarsModuleInfo::sc_name, - NNPotModuleInfo::sc_name }) + NNPotModuleInfo::sc_name, + MetatomicModuleInfo::sc_name }) { IMDModule* module = impl_->modules_.at(std::string(moduleName)).get(); const IMdpOptionProvider* mdpOptionProvider = module->mdpOptionProvider(); diff --git a/src/gromacs/mdrunutility/mdmodulesnotifiers.h b/src/gromacs/mdrunutility/mdmodulesnotifiers.h index 7cbfb780cd..0a7a253ece 100644 --- a/src/gromacs/mdrunutility/mdmodulesnotifiers.h +++ b/src/gromacs/mdrunutility/mdmodulesnotifiers.h @@ -141,6 +141,16 @@ struct MDModulesEnergyOutputToNNPotRequestChecker bool energyOutputToNNPot_ = false; }; +/*! \libinternal \brief Check if the Metatmoic module outputs energy to a specific field. + * + * Ensures that energy is output for NNPot module. + */ +struct MDModulesEnergyOutputToMetatomicPotRequestChecker +{ + //! Trigger output to Metatomic energy field + bool energyOutputToMetatomicPot_ = false; +}; + /*! \libinternal * \brief Collect errors for the energy calculation frequency. * @@ -404,6 +414,9 @@ struct MDModulesNotifiers * \tparam MDModulesEnergyOutputToNNPotRequestChecker* * Enables NNPot module to report if it wants to write its energy * output to the "NN Potential" field in the energy files + * \tparam MDModulesEnergyOutputToMetatomicPotRequestChecker* + * Enables the Metatomic module to report if it wants to write its energy + * output to the "Metatomic Potential" field in the energy files * \tparam SeparatePmeRanksPermitted* * Enables modules to report if they want to disable dedicated * PME ranks @@ -434,6 +447,7 @@ struct MDModulesNotifiers MDModulesEnergyOutputToDensityFittingRequestChecker*, MDModulesEnergyOutputToQMMMRequestChecker*, MDModulesEnergyOutputToNNPotRequestChecker*, + MDModulesEnergyOutputToMetatomicPotRequestChecker*, SeparatePmeRanksPermitted*, const PbcType&, const SimulationTimeStep&, diff --git a/src/gromacs/topology/ifunc.cpp b/src/gromacs/topology/ifunc.cpp index 3f568f2dbf..6f1cc7c3d2 100644 --- a/src/gromacs/topology/ifunc.cpp +++ b/src/gromacs/topology/ifunc.cpp @@ -209,6 +209,7 @@ const t_interaction_function interaction_function[F_NRE] = { def_nofc("DENSITYFIT", "Dens. fitting"), def_nofc("EQM", "Quantum En."), def_nofc("ENNPOT", "NN Potential"), + def_nofc("EMETATOMICPOT", "Metatomic Potential"), def_nofc("EPOT", "Potential"), def_nofc("EKIN", "Kinetic En."), def_nofc("ETOT", "Total Energy"),