gmx::TorchModel Class Reference

#include <gromacs/applied_forces/nnpot/torchmodel.h>

Inheritance diagram for gmx::TorchModel:

Collaboration diagram for gmx::TorchModel:

Description

Class responsible for loading and evaluating a TorchScript-compiled neural network model. Inherits from NNPotModel.

Public Member Functions
	TorchModel (const std::string &filename, NNPotEmbedding embedding, const MDLogger &logger, const MpiComm &mpiComm)
	Constructor for TorchModel. More...
void	evaluateModel (gmx_enerdata_t *enerd, ArrayRef< RVec > forces, ArrayRef< const int > indexLookup, ArrayRef< const int > mmIndices, ArrayRef< const std::string > inputs, ArrayRef< RVec > positions, ArrayRef< int > atomNumbers, ArrayRef< int > atomPairs, ArrayRef< RVec > pairShifts, ArrayRef< RVec > positionsMM, ArrayRef< real > chargesMM, real nnpCharge, ArrayRef< const LinkFrontierAtom > linkFrontier, matrix *box=nullptr, PbcType *pbcType=nullptr) override
bool	outputsForces () const override
	helper function to check if model outputs forces

Constructor & Destructor Documentation

gmx::TorchModel::TorchModel	(	const std::string &	filename,
		NNPotEmbedding	embedding,
		const MDLogger &	logger,
		const MpiComm &	mpiComm
	)

Constructor for TorchModel.

Parameters

[in]	filename	path to the TorchScript model file
[in]	embedding	embedding scheme used for NNP/MM interaction
[in]	logger	handle to MDLogger
[in]	mpiComm	handle to MpiComm

Member Function Documentation

void gmx::TorchModel::evaluateModel ( gmx_enerdata_t *  enerd, ArrayRef< RVec >  forces, ArrayRef< const int >  indexLookup, ArrayRef< const int >  mmIndices, ArrayRef< const std::string >  inputs, ArrayRef< RVec >  positions, ArrayRef< int >  atomNumbers, ArrayRef< int >  atomPairs, ArrayRef< RVec >  pairShifts, ArrayRef< RVec >  positionsMM, ArrayRef< real >  chargesMM, real  nnpCharge, ArrayRef< const LinkFrontierAtom >  linkFrontier, matrix *  box = nullptr, PbcType *  pbcType = nullptr  )

overridevirtual

Call inference on NN model and retrieve outputs

Parameters

[out]	enerd	energy data struct
[out]	forces	forces on atoms
[in]	indexLookup	lookup table for atom indices
[in]	mmIndices	indices of MM atoms
[in]	inputs	list of strings specifying input data
[in]	positions	atom positions
[in]	atomNumbers	atom numbers
[in]	atomPairs	list of all input atom pairs within cutoff
[in]	pairShifts	list of periodic shift vectors corresponding to atom pairs
[in]	positionsMM	MM atom positions
[in]	chargesMM	MM atom charges
[in]	nnpCharge	total charge of NNP region
[in]	linkFrontier	link frontier atoms
[in]	box	simulation box
[in]	pbcType	periodic boundary conditions

Implements gmx::INNPotModel.

---

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

• src/gromacs/applied_forces/nnpot/torchmodel.h
• src/gromacs/applied_forces/nnpot/torchmodel.cpp