Description

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

Public Member Functions
	TorchModel (const std::string &filename, const MDLogger &logger)
	Constructor for TorchModel. More...

void	evaluateModel (gmx_enerdata_t enerd, const ArrayRef< RVec > &forces, ArrayRef< const int > &indexLookup, ArrayRef< const std::string > &inputs, ArrayRef< RVec > &positions, ArrayRef< int > &atomNumbers, matrix box=nullptr, PbcType *pbcType=nullptr) override

void	setCommRec (const t_commrec *cr) override
	Set communication record for possible communication of input/output data between ranks.

bool	outputsForces () const override
	helper function to check if model outputs forces

Constructor & Destructor Documentation

Constructor for TorchModel.

Parameters

[in]	filename	path to the TorchScript model file
[in]	logger	handle to MDLogger

void gmx::TorchModel::evaluateModel	(	gmx_enerdata_t *	enerd,
		const ArrayRef< RVec > &	forces,
		ArrayRef< const int > &	indexLookup,
		ArrayRef< const std::string > &	inputs,
		ArrayRef< RVec > &	positions,
		ArrayRef< int > &	atomNumbers,
		matrix *	box = `nullptr`,
		PbcType *	pbcType = `nullptr`
	)

overridevirtual

Call inference on NN model and retrieve outputs

Parameters

[in]	indexLookup	lookup table for local atom indices
[out]	enerd	energy data struct
[out]	forces	forces on atoms
[in]	indexLookup	lookup table for atom indices
[in]	inputs	list of strings specifying input data
[in]	positions	atom positions
[in]	atomNumbers	atom numbers
[in]	box	simulation box
[in]	pbcType	periodic boundary conditions

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