Class DeepPotJAX#
Defined in File DeepPotJAX.h
Inheritance Relationships#
Base Type#
public deepmd::DeepPotBackend
(Class DeepPotBackend)
Class Documentation#
- class DeepPotJAX : public deepmd::DeepPotBackend#
TensorFlow implementation for Deep Potential.
Public Functions
- DeepPotJAX()#
DP constructor without initialization.
- virtual ~DeepPotJAX()#
- DeepPotJAX(const std::string &model, const int &gpu_rank = 0, const std::string &file_content = "")#
DP constructor with initialization.
- Parameters:
model – [in] The name of the frozen model file.
gpu_rank – [in] The GPU rank. Default is 0. If < 0, use CPU.
file_content – [in] The content of the model file. If it is not empty, DP will read from the string instead of the file.
- virtual void init(const std::string &model, const int &gpu_rank = 0, const std::string &file_content = "")#
Initialize the DP.
- Parameters:
model – [in] The name of the frozen model file.
gpu_rank – [in] The GPU rank. Default is 0. If < 0, use CPU.
file_content – [in] The content of the model file. If it is not empty, DP will read from the string instead of the file.
- inline virtual double cutoff() const#
Get the cutoff radius.
- Returns:
The cutoff radius.
- inline virtual int numb_types() const#
Get the number of types.
- Returns:
The number of types.
- inline virtual int numb_types_spin() const#
Get the number of types with spin.
- Returns:
The number of types with spin.
- inline virtual int dim_fparam() const#
Get the dimension of the frame parameter.
- Returns:
The dimension of the frame parameter.
- inline virtual int dim_aparam() const#
Get the dimension of the atomic parameter.
- Returns:
The dimension of the atomic parameter.
- virtual void get_type_map(std::string &type_map)#
Get the type map (element name of the atom types) of this model.
- Parameters:
type_map – [out] The type map of this model.
- inline virtual bool is_aparam_nall() const#
Get whether the atom dimension of aparam is nall instead of fparam.
- Parameters:
aparam_nall – [out] whether the atom dimension of aparam is nall instead of fparam.
- virtual void computew(std::vector<double> &ener, std::vector<double> &force, std::vector<double> &virial, std::vector<double> &atom_energy, std::vector<double> &atom_virial, const std::vector<double> &coord, const std::vector<int> &atype, const std::vector<double> &box, const std::vector<double> &fparam, const std::vector<double> &aparam, const bool atomic)#
Evaluate the energy, force, virial, atomic energy, and atomic virial by using this DP.
Note
The double precision interface is used by i-PI, GROMACS, ABACUS, and CP2k.
- Parameters:
ener – [out] The system energy.
force – [out] The force on each atom.
virial – [out] The virial.
atom_energy – [out] The atomic energy.
atom_virial – [out] The atomic virial.
coord – [in] The coordinates of atoms. The array should be of size nframes x natoms x 3.
atype – [in] The atom types. The list should contain natoms ints.
box – [in] The cell of the region. The array should be of size nframes x 9.
fparam – [in] The frame parameter. The array can be of size : nframes x dim_fparam. dim_fparam. Then all frames are assumed to be provided with the same fparam.
aparam – [in] The atomic parameter The array can be of size : nframes x natoms x dim_aparam. natoms x dim_aparam. Then all frames are assumed to be provided with the same aparam.
atomic – [in] Request atomic energy and virial if atomic is true.
- virtual void computew(std::vector<double> &ener, std::vector<float> &force, std::vector<float> &virial, std::vector<float> &atom_energy, std::vector<float> &atom_virial, const std::vector<float> &coord, const std::vector<int> &atype, const std::vector<float> &box, const std::vector<float> &fparam, const std::vector<float> &aparam, const bool atomic)#
- virtual void computew(std::vector<double> &ener, std::vector<double> &force, std::vector<double> &virial, std::vector<double> &atom_energy, std::vector<double> &atom_virial, const std::vector<double> &coord, const std::vector<int> &atype, const std::vector<double> &box, const int nghost, const InputNlist &inlist, const int &ago, const std::vector<double> &fparam, const std::vector<double> &aparam, const bool atomic)#
Evaluate the energy, force, virial, atomic energy, and atomic virial by using this DP.
Note
The double precision interface is used by LAMMPS and AMBER.
- Parameters:
ener – [out] The system energy.
force – [out] The force on each atom.
virial – [out] The virial.
atom_energy – [out] The atomic energy.
atom_virial – [out] The atomic virial.
coord – [in] The coordinates of atoms. The array should be of size nframes x natoms x 3.
atype – [in] The atom types. The list should contain natoms ints.
box – [in] The cell of the region. The array should be of size nframes x 9.
nghost – [in] The number of ghost atoms.
lmp_list – [in] The input neighbour list.
ago – [in] Update the internal neighbour list if ago is 0.
fparam – [in] The frame parameter. The array can be of size : nframes x dim_fparam. dim_fparam. Then all frames are assumed to be provided with the same fparam.
aparam – [in] The atomic parameter The array can be of size : nframes x natoms x dim_aparam. natoms x dim_aparam. Then all frames are assumed to be provided with the same aparam.
atomic – [in] Request atomic energy and virial if atomic is true.
- virtual void computew(std::vector<double> &ener, std::vector<float> &force, std::vector<float> &virial, std::vector<float> &atom_energy, std::vector<float> &atom_virial, const std::vector<float> &coord, const std::vector<int> &atype, const std::vector<float> &box, const int nghost, const InputNlist &inlist, const int &ago, const std::vector<float> &fparam, const std::vector<float> &aparam, const bool atomic)#
- virtual void computew_mixed_type(std::vector<double> &ener, std::vector<double> &force, std::vector<double> &virial, std::vector<double> &atom_energy, std::vector<double> &atom_virial, const int &nframes, const std::vector<double> &coord, const std::vector<int> &atype, const std::vector<double> &box, const std::vector<double> &fparam, const std::vector<double> &aparam, const bool atomic)#
Evaluate the energy, force, and virial with the mixed type by using this DP.
Note
At this time, no external program uses this interface.
- Parameters:
ener – [out] The system energy.
force – [out] The force on each atom.
virial – [out] The virial.
atom_energy – [out] The atomic energy.
atom_virial – [out] The atomic virial.
nframes – [in] The number of frames.
coord – [in] The coordinates of atoms. The array should be of size nframes x natoms x 3.
atype – [in] The atom types. The array should be of size nframes x natoms.
box – [in] The cell of the region. The array should be of size nframes x 9.
fparam – [in] The frame parameter. The array can be of size : nframes x dim_fparam. dim_fparam. Then all frames are assumed to be provided with the same fparam.
aparam – [in] The atomic parameter The array can be of size : nframes x natoms x dim_aparam. natoms x dim_aparam. Then all frames are assumed to be provided with the same aparam.
atomic – [in] Request atomic energy and virial if atomic is true.
- virtual void computew_mixed_type(std::vector<double> &ener, std::vector<float> &force, std::vector<float> &virial, std::vector<float> &atom_energy, std::vector<float> &atom_virial, const int &nframes, const std::vector<float> &coord, const std::vector<int> &atype, const std::vector<float> &box, const std::vector<float> &fparam, const std::vector<float> &aparam, const bool atomic)#
- DeepPotJAX()#