diff --git a/source/api_cc/include/DeepPot.h b/source/api_cc/include/DeepPot.h index 0e61b03ce9..e8e64e31c0 100644 --- a/source/api_cc/include/DeepPot.h +++ b/source/api_cc/include/DeepPot.h @@ -1,6 +1,8 @@ // SPDX-License-Identifier: LGPL-3.0-or-later #pragma once +#include + #include "common.h" #include "neighbor_list.h" @@ -8,6 +10,499 @@ namespace deepmd { /** * @brief Deep Potential. **/ +class DeepPotBase { + public: + /** + * @brief DP constructor without initialization. + **/ + DeepPotBase(){}; + virtual ~DeepPotBase(){}; + /** + * @brief DP constructor with initialization. + * @param[in] model The name of the frozen model file. + * @param[in] gpu_rank The GPU rank. Default is 0. + * @param[in] file_content The content of the model file. If it is not empty, + *DP will read from the string instead of the file. + **/ + DeepPotBase(const std::string& model, + const int& gpu_rank = 0, + const std::string& file_content = ""); + /** + * @brief Initialize the DP. + * @param[in] model The name of the frozen model file. + * @param[in] gpu_rank The GPU rank. Default is 0. + * @param[in] file_content 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 = "") = 0; + + /** + * @brief Evaluate the energy, force and virial by using this DP. + * @param[out] ener The system energy. + * @param[out] force The force on each atom. + * @param[out] virial The virial. + * @param[in] coord The coordinates of atoms. The array should be of size + *nframes x natoms x 3. + * @param[in] atype The atom types. The list should contain natoms ints. + * @param[in] box The cell of the region. The array should be of size nframes + *x 9. + * @param[in] fparam 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. + * @param[in] aparam 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. + * @{ + **/ + virtual void computew( + double& ener, + std::vector& force, + std::vector& virial, + const std::vector& coord, + const std::vector& atype, + const std::vector& box, + const std::vector& fparam = std::vector(), + const std::vector& aparam = std::vector()) = 0; + + virtual void computew( + double& ener, + std::vector& force, + std::vector& virial, + const std::vector& coord, + const std::vector& atype, + const std::vector& box, + const std::vector& fparam = std::vector(), + const std::vector& aparam = std::vector()) = 0; + virtual void computew( + std::vector& ener, + std::vector& force, + std::vector& virial, + const std::vector& coord, + const std::vector& atype, + const std::vector& box, + const std::vector& fparam = std::vector(), + const std::vector& aparam = std::vector()) = 0; + virtual void computew( + std::vector& ener, + std::vector& force, + std::vector& virial, + const std::vector& coord, + const std::vector& atype, + const std::vector& box, + const std::vector& fparam = std::vector(), + const std::vector& aparam = std::vector()) = 0; + /** @} */ + /** + * @brief Evaluate the energy, force and virial by using this DP. + * @param[out] ener The system energy. + * @param[out] force The force on each atom. + * @param[out] virial The virial. + * @param[in] coord The coordinates of atoms. The array should be of size + *nframes x natoms x 3. + * @param[in] atype The atom types. The list should contain natoms ints. + * @param[in] box The cell of the region. The array should be of size nframes + *x 9. + * @param[in] nghost The number of ghost atoms. + * @param[in] inlist The input neighbour list. + * @param[in] ago Update the internal neighbour list if ago is 0. + * @param[in] fparam 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. + * @param[in] aparam 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. + * @{ + **/ + virtual void computew( + double& ener, + std::vector& force, + std::vector& virial, + const std::vector& coord, + const std::vector& atype, + const std::vector& box, + const int nghost, + const InputNlist& inlist, + const int& ago, + const std::vector& fparam = std::vector(), + const std::vector& aparam = std::vector()) = 0; + + virtual void computew( + double& ener, + std::vector& force, + std::vector& virial, + const std::vector& coord, + const std::vector& atype, + const std::vector& box, + const int nghost, + const InputNlist& inlist, + const int& ago, + const std::vector& fparam = std::vector(), + const std::vector& aparam = std::vector()) = 0; + virtual void computew( + std::vector& ener, + std::vector& force, + std::vector& virial, + const std::vector& coord, + const std::vector& atype, + const std::vector& box, + const int nghost, + const InputNlist& inlist, + const int& ago, + const std::vector& fparam = std::vector(), + const std::vector& aparam = std::vector()) = 0; + virtual void computew( + std::vector& ener, + std::vector& force, + std::vector& virial, + const std::vector& coord, + const std::vector& atype, + const std::vector& box, + const int nghost, + const InputNlist& inlist, + const int& ago, + const std::vector& fparam = std::vector(), + const std::vector& aparam = std::vector()) = 0; + /** @} */ + + /** + * @brief Evaluate the energy, force, virial, atomic energy, and atomic virial + *by using this DP. + * @param[out] ener The system energy. + * @param[out] force The force on each atom. + * @param[out] virial The virial. + * @param[out] atom_energy The atomic energy. + * @param[out] atom_virial The atomic virial. + * @param[in] coord The coordinates of atoms. The array should be of size + *nframes x natoms x 3. + * @param[in] atype The atom types. The list should contain natoms ints. + * @param[in] box The cell of the region. The array should be of size nframes + *x 9. + * @param[in] fparam 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. + * @param[in] aparam 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. + * @{ + **/ + virtual void computew( + double& ener, + std::vector& force, + std::vector& virial, + std::vector& atom_energy, + std::vector& atom_virial, + const std::vector& coord, + const std::vector& atype, + const std::vector& box, + const std::vector& fparam = std::vector(), + const std::vector& aparam = std::vector()) = 0; + + virtual void computew( + double& ener, + std::vector& force, + std::vector& virial, + std::vector& atom_energy, + std::vector& atom_virial, + const std::vector& coord, + const std::vector& atype, + const std::vector& box, + const std::vector& fparam = std::vector(), + const std::vector& aparam = std::vector()) = 0; + virtual void computew( + std::vector& ener, + std::vector& force, + std::vector& virial, + std::vector& atom_energy, + std::vector& atom_virial, + const std::vector& coord, + const std::vector& atype, + const std::vector& box, + const std::vector& fparam = std::vector(), + const std::vector& aparam = std::vector()) = 0; + virtual void computew( + std::vector& ener, + std::vector& force, + std::vector& virial, + std::vector& atom_energy, + std::vector& atom_virial, + const std::vector& coord, + const std::vector& atype, + const std::vector& box, + const std::vector& fparam = std::vector(), + const std::vector& aparam = std::vector()) = 0; + /** @} */ + /** + * @brief Evaluate the energy, force, virial, atomic energy, and atomic virial + *by using this DP. + * @param[out] ener The system energy. + * @param[out] force The force on each atom. + * @param[out] virial The virial. + * @param[out] atom_energy The atomic energy. + * @param[out] atom_virial The atomic virial. + * @param[in] coord The coordinates of atoms. The array should be of size + *nframes x natoms x 3. + * @param[in] atype The atom types. The list should contain natoms ints. + * @param[in] box The cell of the region. The array should be of size nframes + *x 9. + * @param[in] nghost The number of ghost atoms. + * @param[in] lmp_list The input neighbour list. + * @param[in] ago Update the internal neighbour list if ago is 0. + * @param[in] fparam 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. + * @param[in] aparam 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. + * @{ + **/ + virtual void computew( + double& ener, + std::vector& force, + std::vector& virial, + std::vector& atom_energy, + std::vector& atom_virial, + const std::vector& coord, + const std::vector& atype, + const std::vector& box, + const int nghost, + const InputNlist& inlist, + const int& ago, + const std::vector& fparam = std::vector(), + const std::vector& aparam = std::vector()) = 0; + + virtual void computew( + double& ener, + std::vector& force, + std::vector& virial, + std::vector& atom_energy, + std::vector& atom_virial, + const std::vector& coord, + const std::vector& atype, + const std::vector& box, + const int nghost, + const InputNlist& inlist, + const int& ago, + const std::vector& fparam = std::vector(), + const std::vector& aparam = std::vector()) = 0; + virtual void computew( + std::vector& ener, + std::vector& force, + std::vector& virial, + std::vector& atom_energy, + std::vector& atom_virial, + const std::vector& coord, + const std::vector& atype, + const std::vector& box, + const int nghost, + const InputNlist& inlist, + const int& ago, + const std::vector& fparam = std::vector(), + const std::vector& aparam = std::vector()) = 0; + virtual void computew( + std::vector& ener, + std::vector& force, + std::vector& virial, + std::vector& atom_energy, + std::vector& atom_virial, + const std::vector& coord, + const std::vector& atype, + const std::vector& box, + const int nghost, + const InputNlist& inlist, + const int& ago, + const std::vector& fparam = std::vector(), + const std::vector& aparam = std::vector()) = 0; + /** @} */ + /** + * @brief Evaluate the energy, force, and virial with the mixed type + *by using this DP. + * @param[out] ener The system energy. + * @param[out] force The force on each atom. + * @param[out] virial The virial. + * @param[in] nframes The number of frames. + * @param[in] coord The coordinates of atoms. The array should be of size + *nframes x natoms x 3. + * @param[in] atype The atom types. The array should be of size nframes x + *natoms. + * @param[in] box The cell of the region. The array should be of size nframes + *x 9. + * @param[in] fparam 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. + * @param[in] aparam 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. + * @{ + **/ + virtual void computew_mixed_type( + double& ener, + std::vector& force, + std::vector& virial, + const int& nframes, + const std::vector& coord, + const std::vector& atype, + const std::vector& box, + const std::vector& fparam = std::vector(), + const std::vector& aparam = std::vector()) = 0; + virtual void computew_mixed_type( + double& ener, + std::vector& force, + std::vector& virial, + const int& nframes, + const std::vector& coord, + const std::vector& atype, + const std::vector& box, + const std::vector& fparam = std::vector(), + const std::vector& aparam = std::vector()) = 0; + virtual void computew_mixed_type( + std::vector& ener, + std::vector& force, + std::vector& virial, + const int& nframes, + const std::vector& coord, + const std::vector& atype, + const std::vector& box, + const std::vector& fparam = std::vector(), + const std::vector& aparam = std::vector()) = 0; + virtual void computew_mixed_type( + std::vector& ener, + std::vector& force, + std::vector& virial, + const int& nframes, + const std::vector& coord, + const std::vector& atype, + const std::vector& box, + const std::vector& fparam = std::vector(), + const std::vector& aparam = std::vector()) = 0; + /** @} */ + /** + * @brief Evaluate the energy, force, and virial with the mixed type + *by using this DP. + * @param[out] ener The system energy. + * @param[out] force The force on each atom. + * @param[out] virial The virial. + * @param[out] atom_energy The atomic energy. + * @param[out] atom_virial The atomic virial. + * @param[in] nframes The number of frames. + * @param[in] coord The coordinates of atoms. The array should be of size + *nframes x natoms x 3. + * @param[in] atype The atom types. The array should be of size nframes x + *natoms. + * @param[in] box The cell of the region. The array should be of size nframes + *x 9. + * @param[in] fparam 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. + * @param[in] aparam 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. + * @{ + **/ + virtual void computew_mixed_type( + double& ener, + std::vector& force, + std::vector& virial, + std::vector& atom_energy, + std::vector& atom_virial, + const int& nframes, + const std::vector& coord, + const std::vector& atype, + const std::vector& box, + const std::vector& fparam = std::vector(), + const std::vector& aparam = std::vector()) = 0; + virtual void computew_mixed_type( + double& ener, + std::vector& force, + std::vector& virial, + std::vector& atom_energy, + std::vector& atom_virial, + const int& nframes, + const std::vector& coord, + const std::vector& atype, + const std::vector& box, + const std::vector& fparam = std::vector(), + const std::vector& aparam = std::vector()) = 0; + virtual void computew_mixed_type( + std::vector& ener, + std::vector& force, + std::vector& virial, + std::vector& atom_energy, + std::vector& atom_virial, + const int& nframes, + const std::vector& coord, + const std::vector& atype, + const std::vector& box, + const std::vector& fparam = std::vector(), + const std::vector& aparam = std::vector()) = 0; + virtual void computew_mixed_type( + std::vector& ener, + std::vector& force, + std::vector& virial, + std::vector& atom_energy, + std::vector& atom_virial, + const int& nframes, + const std::vector& coord, + const std::vector& atype, + const std::vector& box, + const std::vector& fparam = std::vector(), + const std::vector& aparam = std::vector()) = 0; + /** @} */ + /** + * @brief Get the cutoff radius. + * @return The cutoff radius. + **/ + virtual double cutoff() const = 0; + /** + * @brief Get the number of types. + * @return The number of types. + **/ + virtual int numb_types() const = 0; + /** + * @brief Get the number of types with spin. + * @return The number of types with spin. + **/ + virtual int numb_types_spin() const = 0; + /** + * @brief Get the dimension of the frame parameter. + * @return The dimension of the frame parameter. + **/ + virtual int dim_fparam() const = 0; + /** + * @brief Get the dimension of the atomic parameter. + * @return The dimension of the atomic parameter. + **/ + virtual int dim_aparam() const = 0; + /** + * @brief Get the type map (element name of the atom types) of this model. + * @param[out] type_map The type map of this model. + **/ + virtual void get_type_map(std::string& type_map) = 0; + + /** + * @brief Get whether the atom dimension of aparam is nall instead of fparam. + * @param[out] aparam_nall whether the atom dimension of aparam is nall + *instead of fparam. + **/ + virtual bool is_aparam_nall() const = 0; +}; + +/** + * @brief Deep Potential to automatically switch backends. + **/ class DeepPot { public: /** @@ -35,12 +530,12 @@ class DeepPot { void init(const std::string& model, const int& gpu_rank = 0, const std::string& file_content = ""); + /** * @brief Print the DP summary to the screen. * @param[in] pre The prefix to each line. **/ void print_summary(const std::string& pre) const; - /** * @brief Evaluate the energy, force and virial by using this DP. * @param[out] ener The system energy. @@ -249,42 +744,27 @@ class DeepPot { * @brief Get the cutoff radius. * @return The cutoff radius. **/ - double cutoff() const { - assert(inited); - return rcut; - }; + double cutoff() const; /** * @brief Get the number of types. * @return The number of types. **/ - int numb_types() const { - assert(inited); - return ntypes; - }; + int numb_types() const; /** * @brief Get the number of types with spin. * @return The number of types with spin. **/ - int numb_types_spin() const { - assert(inited); - return ntypes_spin; - }; + int numb_types_spin() const; /** * @brief Get the dimension of the frame parameter. * @return The dimension of the frame parameter. **/ - int dim_fparam() const { - assert(inited); - return dfparam; - }; + int dim_fparam() const; /** * @brief Get the dimension of the atomic parameter. * @return The dimension of the atomic parameter. **/ - int dim_aparam() const { - assert(inited); - return daparam; - }; + int dim_aparam() const; /** * @brief Get the type map (element name of the atom types) of this model. * @param[out] type_map The type map of this model. @@ -296,77 +776,11 @@ class DeepPot { * @param[out] aparam_nall whether the atom dimension of aparam is nall *instead of fparam. **/ - bool is_aparam_nall() const { - assert(inited); - return aparam_nall; - }; + bool is_aparam_nall() const; private: - tensorflow::Session* session; - int num_intra_nthreads, num_inter_nthreads; - tensorflow::GraphDef* graph_def; bool inited; - template - VT get_scalar(const std::string& name) const; - // VALUETYPE get_rcut () const; - // int get_ntypes () const; - double rcut; - int dtype; - double cell_size; - std::string model_type; - std::string model_version; - int ntypes; - int ntypes_spin; - int dfparam; - int daparam; - bool aparam_nall; - /** - * @brief Validate the size of frame and atomic parameters. - * @param[in] nframes The number of frames. - * @param[in] nloc The number of local atoms. - * @param[in] fparam The frame parameter. - * @param[in] aparam The atomic parameter. - * @tparam VALUETYPE The type of the parameters, double or float. - */ - template - void validate_fparam_aparam(const int& nframes, - const int& nloc, - const std::vector& fparam, - const std::vector& aparam) const; - /** - * @brief Tile the frame or atomic parameters if there is only - * a single frame of frame or atomic parameters. - * @param[out] out_param The tiled frame or atomic parameters. - * @param[in] nframes The number of frames. - * @param[in] dparam The dimension of the frame or atomic parameters in a - * frame. - * @param[in] param The frame or atomic parameters. - * @tparam VALUETYPE The type of the parameters, double or float. - */ - template - void tile_fparam_aparam(std::vector& out_param, - const int& nframes, - const int& dparam, - const std::vector& param) const; - template - void compute_inner( - ENERGYVTYPE& ener, - std::vector& force, - std::vector& virial, - const std::vector& coord, - const std::vector& atype, - const std::vector& box, - const int nghost, - const int& ago, - const std::vector& fparam = std::vector(), - const std::vector& aparam = std::vector()); - - // copy neighbor list info from host - bool init_nbor; - std::vector sec_a; - NeighborListData nlist_data; - InputNlist nlist; - AtomMap atommap; + std::shared_ptr dp; }; class DeepPotModelDevi { diff --git a/source/api_cc/include/DeepPotTF.h b/source/api_cc/include/DeepPotTF.h new file mode 100644 index 0000000000..59bd5f476f --- /dev/null +++ b/source/api_cc/include/DeepPotTF.h @@ -0,0 +1,628 @@ +// SPDX-License-Identifier: LGPL-3.0-or-later +#pragma once + +#include "DeepPot.h" +#include "common.h" +#include "neighbor_list.h" + +namespace deepmd { +/** + * @brief TensorFlow implementation for Deep Potential. + **/ +class DeepPotTF : public DeepPotBase { + public: + /** + * @brief DP constructor without initialization. + **/ + DeepPotTF(); + ~DeepPotTF(); + /** + * @brief DP constructor with initialization. + * @param[in] model The name of the frozen model file. + * @param[in] gpu_rank The GPU rank. Default is 0. + * @param[in] file_content The content of the model file. If it is not empty, + *DP will read from the string instead of the file. + **/ + DeepPotTF(const std::string& model, + const int& gpu_rank = 0, + const std::string& file_content = ""); + /** + * @brief Initialize the DP. + * @param[in] model The name of the frozen model file. + * @param[in] gpu_rank The GPU rank. Default is 0. + * @param[in] file_content The content of the model file. If it is not empty, + *DP will read from the string instead of the file. + **/ + void init(const std::string& model, + const int& gpu_rank = 0, + const std::string& file_content = ""); + + /** + * @brief Evaluate the energy, force and virial by using this DP. + * @param[out] ener The system energy. + * @param[out] force The force on each atom. + * @param[out] virial The virial. + * @param[in] coord The coordinates of atoms. The array should be of size + *nframes x natoms x 3. + * @param[in] atype The atom types. The list should contain natoms ints. + * @param[in] box The cell of the region. The array should be of size nframes + *x 9. + * @param[in] fparam 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. + * @param[in] aparam 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. + **/ + template + void compute(ENERGYVTYPE& ener, + std::vector& force, + std::vector& virial, + const std::vector& coord, + const std::vector& atype, + const std::vector& box, + const std::vector& fparam = std::vector(), + const std::vector& aparam = std::vector()); + /** + * @brief Evaluate the energy, force and virial by using this DP. + * @param[out] ener The system energy. + * @param[out] force The force on each atom. + * @param[out] virial The virial. + * @param[in] coord The coordinates of atoms. The array should be of size + *nframes x natoms x 3. + * @param[in] atype The atom types. The list should contain natoms ints. + * @param[in] box The cell of the region. The array should be of size nframes + *x 9. + * @param[in] nghost The number of ghost atoms. + * @param[in] inlist The input neighbour list. + * @param[in] ago Update the internal neighbour list if ago is 0. + * @param[in] fparam 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. + * @param[in] aparam 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. + **/ + template + void compute(ENERGYVTYPE& ener, + std::vector& force, + std::vector& virial, + const std::vector& coord, + const std::vector& atype, + const std::vector& box, + const int nghost, + const InputNlist& inlist, + const int& ago, + const std::vector& fparam = std::vector(), + const std::vector& aparam = std::vector()); + /** + * @brief Evaluate the energy, force, virial, atomic energy, and atomic virial + *by using this DP. + * @param[out] ener The system energy. + * @param[out] force The force on each atom. + * @param[out] virial The virial. + * @param[out] atom_energy The atomic energy. + * @param[out] atom_virial The atomic virial. + * @param[in] coord The coordinates of atoms. The array should be of size + *nframes x natoms x 3. + * @param[in] atype The atom types. The list should contain natoms ints. + * @param[in] box The cell of the region. The array should be of size nframes + *x 9. + * @param[in] fparam 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. + * @param[in] aparam 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. + **/ + template + void compute(ENERGYVTYPE& ener, + std::vector& force, + std::vector& virial, + std::vector& atom_energy, + std::vector& atom_virial, + const std::vector& coord, + const std::vector& atype, + const std::vector& box, + const std::vector& fparam = std::vector(), + const std::vector& aparam = std::vector()); + /** + * @brief Evaluate the energy, force, virial, atomic energy, and atomic virial + *by using this DP. + * @param[out] ener The system energy. + * @param[out] force The force on each atom. + * @param[out] virial The virial. + * @param[out] atom_energy The atomic energy. + * @param[out] atom_virial The atomic virial. + * @param[in] coord The coordinates of atoms. The array should be of size + *nframes x natoms x 3. + * @param[in] atype The atom types. The list should contain natoms ints. + * @param[in] box The cell of the region. The array should be of size nframes + *x 9. + * @param[in] nghost The number of ghost atoms. + * @param[in] lmp_list The input neighbour list. + * @param[in] ago Update the internal neighbour list if ago is 0. + * @param[in] fparam 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. + * @param[in] aparam 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. + **/ + template + void compute(ENERGYVTYPE& ener, + std::vector& force, + std::vector& virial, + std::vector& atom_energy, + std::vector& atom_virial, + const std::vector& coord, + const std::vector& atype, + const std::vector& box, + const int nghost, + const InputNlist& lmp_list, + const int& ago, + const std::vector& fparam = std::vector(), + const std::vector& aparam = std::vector()); + /** + * @brief Evaluate the energy, force, and virial with the mixed type + *by using this DP. + * @param[out] ener The system energy. + * @param[out] force The force on each atom. + * @param[out] virial The virial. + * @param[in] nframes The number of frames. + * @param[in] coord The coordinates of atoms. The array should be of size + *nframes x natoms x 3. + * @param[in] atype The atom types. The array should be of size nframes x + *natoms. + * @param[in] box The cell of the region. The array should be of size nframes + *x 9. + * @param[in] fparam 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. + * @param[in] aparam 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. + **/ + template + void compute_mixed_type( + ENERGYVTYPE& ener, + std::vector& force, + std::vector& virial, + const int& nframes, + const std::vector& coord, + const std::vector& atype, + const std::vector& box, + const std::vector& fparam = std::vector(), + const std::vector& aparam = std::vector()); + /** + * @brief Evaluate the energy, force, and virial with the mixed type + *by using this DP. + * @param[out] ener The system energy. + * @param[out] force The force on each atom. + * @param[out] virial The virial. + * @param[out] atom_energy The atomic energy. + * @param[out] atom_virial The atomic virial. + * @param[in] nframes The number of frames. + * @param[in] coord The coordinates of atoms. The array should be of size + *nframes x natoms x 3. + * @param[in] atype The atom types. The array should be of size nframes x + *natoms. + * @param[in] box The cell of the region. The array should be of size nframes + *x 9. + * @param[in] fparam 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. + * @param[in] aparam 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. + **/ + template + void compute_mixed_type( + ENERGYVTYPE& ener, + std::vector& force, + std::vector& virial, + std::vector& atom_energy, + std::vector& atom_virial, + const int& nframes, + const std::vector& coord, + const std::vector& atype, + const std::vector& box, + const std::vector& fparam = std::vector(), + const std::vector& aparam = std::vector()); + /** + * @brief Get the cutoff radius. + * @return The cutoff radius. + **/ + double cutoff() const { + assert(inited); + return rcut; + }; + /** + * @brief Get the number of types. + * @return The number of types. + **/ + int numb_types() const { + assert(inited); + return ntypes; + }; + /** + * @brief Get the number of types with spin. + * @return The number of types with spin. + **/ + int numb_types_spin() const { + assert(inited); + return ntypes_spin; + }; + /** + * @brief Get the dimension of the frame parameter. + * @return The dimension of the frame parameter. + **/ + int dim_fparam() const { + assert(inited); + return dfparam; + }; + /** + * @brief Get the dimension of the atomic parameter. + * @return The dimension of the atomic parameter. + **/ + int dim_aparam() const { + assert(inited); + return daparam; + }; + /** + * @brief Get the type map (element name of the atom types) of this model. + * @param[out] type_map The type map of this model. + **/ + void get_type_map(std::string& type_map); + + /** + * @brief Get whether the atom dimension of aparam is nall instead of fparam. + * @param[out] aparam_nall whether the atom dimension of aparam is nall + *instead of fparam. + **/ + bool is_aparam_nall() const { + assert(inited); + return aparam_nall; + }; + + // forward to template class + void computew(double& ener, + std::vector& force, + std::vector& virial, + const std::vector& coord, + const std::vector& atype, + const std::vector& box, + const std::vector& fparam = std::vector(), + const std::vector& aparam = std::vector()); + + void computew(double& ener, + std::vector& force, + std::vector& virial, + const std::vector& coord, + const std::vector& atype, + const std::vector& box, + const std::vector& fparam = std::vector(), + const std::vector& aparam = std::vector()); + void computew(std::vector& ener, + std::vector& force, + std::vector& virial, + const std::vector& coord, + const std::vector& atype, + const std::vector& box, + const std::vector& fparam = std::vector(), + const std::vector& aparam = std::vector()); + void computew(std::vector& ener, + std::vector& force, + std::vector& virial, + const std::vector& coord, + const std::vector& atype, + const std::vector& box, + const std::vector& fparam = std::vector(), + const std::vector& aparam = std::vector()); + void computew(double& ener, + std::vector& force, + std::vector& virial, + const std::vector& coord, + const std::vector& atype, + const std::vector& box, + const int nghost, + const InputNlist& inlist, + const int& ago, + const std::vector& fparam = std::vector(), + const std::vector& aparam = std::vector()); + + void computew(double& ener, + std::vector& force, + std::vector& virial, + const std::vector& coord, + const std::vector& atype, + const std::vector& box, + const int nghost, + const InputNlist& inlist, + const int& ago, + const std::vector& fparam = std::vector(), + const std::vector& aparam = std::vector()); + void computew(std::vector& ener, + std::vector& force, + std::vector& virial, + const std::vector& coord, + const std::vector& atype, + const std::vector& box, + const int nghost, + const InputNlist& inlist, + const int& ago, + const std::vector& fparam = std::vector(), + const std::vector& aparam = std::vector()); + void computew(std::vector& ener, + std::vector& force, + std::vector& virial, + const std::vector& coord, + const std::vector& atype, + const std::vector& box, + const int nghost, + const InputNlist& inlist, + const int& ago, + const std::vector& fparam = std::vector(), + const std::vector& aparam = std::vector()); + void computew(double& ener, + std::vector& force, + std::vector& virial, + std::vector& atom_energy, + std::vector& atom_virial, + const std::vector& coord, + const std::vector& atype, + const std::vector& box, + const std::vector& fparam = std::vector(), + const std::vector& aparam = std::vector()); + + void computew(double& ener, + std::vector& force, + std::vector& virial, + std::vector& atom_energy, + std::vector& atom_virial, + const std::vector& coord, + const std::vector& atype, + const std::vector& box, + const std::vector& fparam = std::vector(), + const std::vector& aparam = std::vector()); + void computew(std::vector& ener, + std::vector& force, + std::vector& virial, + std::vector& atom_energy, + std::vector& atom_virial, + const std::vector& coord, + const std::vector& atype, + const std::vector& box, + const std::vector& fparam = std::vector(), + const std::vector& aparam = std::vector()); + void computew(std::vector& ener, + std::vector& force, + std::vector& virial, + std::vector& atom_energy, + std::vector& atom_virial, + const std::vector& coord, + const std::vector& atype, + const std::vector& box, + const std::vector& fparam = std::vector(), + const std::vector& aparam = std::vector()); + void computew(double& ener, + std::vector& force, + std::vector& virial, + std::vector& atom_energy, + std::vector& atom_virial, + const std::vector& coord, + const std::vector& atype, + const std::vector& box, + const int nghost, + const InputNlist& inlist, + const int& ago, + const std::vector& fparam = std::vector(), + const std::vector& aparam = std::vector()); + + void computew(double& ener, + std::vector& force, + std::vector& virial, + std::vector& atom_energy, + std::vector& atom_virial, + const std::vector& coord, + const std::vector& atype, + const std::vector& box, + const int nghost, + const InputNlist& inlist, + const int& ago, + const std::vector& fparam = std::vector(), + const std::vector& aparam = std::vector()); + void computew(std::vector& ener, + std::vector& force, + std::vector& virial, + std::vector& atom_energy, + std::vector& atom_virial, + const std::vector& coord, + const std::vector& atype, + const std::vector& box, + const int nghost, + const InputNlist& inlist, + const int& ago, + const std::vector& fparam = std::vector(), + const std::vector& aparam = std::vector()); + void computew(std::vector& ener, + std::vector& force, + std::vector& virial, + std::vector& atom_energy, + std::vector& atom_virial, + const std::vector& coord, + const std::vector& atype, + const std::vector& box, + const int nghost, + const InputNlist& inlist, + const int& ago, + const std::vector& fparam = std::vector(), + const std::vector& aparam = std::vector()); + void computew_mixed_type( + double& ener, + std::vector& force, + std::vector& virial, + const int& nframes, + const std::vector& coord, + const std::vector& atype, + const std::vector& box, + const std::vector& fparam = std::vector(), + const std::vector& aparam = std::vector()); + void computew_mixed_type( + double& ener, + std::vector& force, + std::vector& virial, + const int& nframes, + const std::vector& coord, + const std::vector& atype, + const std::vector& box, + const std::vector& fparam = std::vector(), + const std::vector& aparam = std::vector()); + void computew_mixed_type( + std::vector& ener, + std::vector& force, + std::vector& virial, + const int& nframes, + const std::vector& coord, + const std::vector& atype, + const std::vector& box, + const std::vector& fparam = std::vector(), + const std::vector& aparam = std::vector()); + void computew_mixed_type( + std::vector& ener, + std::vector& force, + std::vector& virial, + const int& nframes, + const std::vector& coord, + const std::vector& atype, + const std::vector& box, + const std::vector& fparam = std::vector(), + const std::vector& aparam = std::vector()); + void computew_mixed_type( + double& ener, + std::vector& force, + std::vector& virial, + std::vector& atom_energy, + std::vector& atom_virial, + const int& nframes, + const std::vector& coord, + const std::vector& atype, + const std::vector& box, + const std::vector& fparam = std::vector(), + const std::vector& aparam = std::vector()); + void computew_mixed_type( + double& ener, + std::vector& force, + std::vector& virial, + std::vector& atom_energy, + std::vector& atom_virial, + const int& nframes, + const std::vector& coord, + const std::vector& atype, + const std::vector& box, + const std::vector& fparam = std::vector(), + const std::vector& aparam = std::vector()); + void computew_mixed_type( + std::vector& ener, + std::vector& force, + std::vector& virial, + std::vector& atom_energy, + std::vector& atom_virial, + const int& nframes, + const std::vector& coord, + const std::vector& atype, + const std::vector& box, + const std::vector& fparam = std::vector(), + const std::vector& aparam = std::vector()); + void computew_mixed_type( + std::vector& ener, + std::vector& force, + std::vector& virial, + std::vector& atom_energy, + std::vector& atom_virial, + const int& nframes, + const std::vector& coord, + const std::vector& atype, + const std::vector& box, + const std::vector& fparam = std::vector(), + const std::vector& aparam = std::vector()); + + private: + tensorflow::Session* session; + int num_intra_nthreads, num_inter_nthreads; + tensorflow::GraphDef* graph_def; + bool inited; + template + VT get_scalar(const std::string& name) const; + double rcut; + int dtype; + double cell_size; + std::string model_type; + std::string model_version; + int ntypes; + int ntypes_spin; + int dfparam; + int daparam; + bool aparam_nall; + /** + * @brief Validate the size of frame and atomic parameters. + * @param[in] nframes The number of frames. + * @param[in] nloc The number of local atoms. + * @param[in] fparam The frame parameter. + * @param[in] aparam The atomic parameter. + * @tparam VALUETYPE The type of the parameters, double or float. + */ + template + void validate_fparam_aparam(const int& nframes, + const int& nloc, + const std::vector& fparam, + const std::vector& aparam) const; + /** + * @brief Tile the frame or atomic parameters if there is only + * a single frame of frame or atomic parameters. + * @param[out] out_param The tiled frame or atomic parameters. + * @param[in] nframes The number of frames. + * @param[in] dparam The dimension of the frame or atomic parameters in a + * frame. + * @param[in] param The frame or atomic parameters. + * @tparam VALUETYPE The type of the parameters, double or float. + */ + template + void tile_fparam_aparam(std::vector& out_param, + const int& nframes, + const int& dparam, + const std::vector& param) const; + template + void compute_inner( + ENERGYVTYPE& ener, + std::vector& force, + std::vector& virial, + const std::vector& coord, + const std::vector& atype, + const std::vector& box, + const int nghost, + const int& ago, + const std::vector& fparam = std::vector(), + const std::vector& aparam = std::vector()); + + // copy neighbor list info from host + bool init_nbor; + std::vector sec_a; + NeighborListData nlist_data; + InputNlist nlist; + AtomMap atommap; +}; + +} // namespace deepmd diff --git a/source/api_cc/include/common.h b/source/api_cc/include/common.h index 481e09cc89..7982c4f89d 100644 --- a/source/api_cc/include/common.h +++ b/source/api_cc/include/common.h @@ -19,6 +19,8 @@ namespace deepmd { typedef double ENERGYTYPE; +// TODO: currently we only implement TF; reserve for future use +enum DPBackend { TensorFlow, PyTorch, Paddle, Unknown }; struct NeighborListData { /// Array stores the core region atom's index diff --git a/source/api_cc/src/DeepPot.cc b/source/api_cc/src/DeepPot.cc index cd1b571153..d290565c2b 100644 --- a/source/api_cc/src/DeepPot.cc +++ b/source/api_cc/src/DeepPot.cc @@ -1,419 +1,27 @@ // SPDX-License-Identifier: LGPL-3.0-or-later #include "DeepPot.h" +#include "common.h" +// TODO: only include when TF backend is built +#include #include #include "AtomMap.h" +#include "DeepPotTF.h" #include "device.h" -using namespace tensorflow; using namespace deepmd; -// start multiple frames - -template -static void run_model( - std::vector& dener, - std::vector& dforce_, - std::vector& dvirial, - Session* session, - const std::vector>& input_tensors, - const AtomMap& atommap, - const int nframes, - const int nghost = 0) { - unsigned nloc = atommap.get_type().size(); - unsigned nall = nloc + nghost; - dener.resize(nframes); - if (nloc == 0) { - // no backward map needed - // dforce of size nall * 3 - dforce_.resize(nframes * nall * 3); - fill(dforce_.begin(), dforce_.end(), (VALUETYPE)0.0); - // dvirial of size 9 - dvirial.resize(nframes * 9); - fill(dvirial.begin(), dvirial.end(), (VALUETYPE)0.0); - return; - } - - std::vector output_tensors; - check_status(session->Run( - input_tensors, {"o_energy", "o_force", "o_atom_energy", "o_atom_virial"}, - {}, &output_tensors)); - - Tensor output_e = output_tensors[0]; - Tensor output_f = output_tensors[1]; - Tensor output_av = output_tensors[3]; - - auto oe = output_e.flat(); - auto of = output_f.flat(); - auto oav = output_av.flat(); - - std::vector dforce(nframes * 3 * nall); - dvirial.resize(nframes * 9); - for (int ii = 0; ii < nframes; ++ii) { - dener[ii] = oe(ii); - } - for (unsigned ii = 0; ii < nframes * nall * 3; ++ii) { - dforce[ii] = of(ii); - } - // set dvirial to zero, prevent input vector is not zero (#1123) - std::fill(dvirial.begin(), dvirial.end(), (VALUETYPE)0.); - for (int kk = 0; kk < nframes; ++kk) { - for (int ii = 0; ii < nall; ++ii) { - dvirial[kk * 9 + 0] += (VALUETYPE)1.0 * oav(kk * nall * 9 + 9 * ii + 0); - dvirial[kk * 9 + 1] += (VALUETYPE)1.0 * oav(kk * nall * 9 + 9 * ii + 1); - dvirial[kk * 9 + 2] += (VALUETYPE)1.0 * oav(kk * nall * 9 + 9 * ii + 2); - dvirial[kk * 9 + 3] += (VALUETYPE)1.0 * oav(kk * nall * 9 + 9 * ii + 3); - dvirial[kk * 9 + 4] += (VALUETYPE)1.0 * oav(kk * nall * 9 + 9 * ii + 4); - dvirial[kk * 9 + 5] += (VALUETYPE)1.0 * oav(kk * nall * 9 + 9 * ii + 5); - dvirial[kk * 9 + 6] += (VALUETYPE)1.0 * oav(kk * nall * 9 + 9 * ii + 6); - dvirial[kk * 9 + 7] += (VALUETYPE)1.0 * oav(kk * nall * 9 + 9 * ii + 7); - dvirial[kk * 9 + 8] += (VALUETYPE)1.0 * oav(kk * nall * 9 + 9 * ii + 8); - } - } - dforce_ = dforce; - atommap.backward(dforce_.begin(), dforce.begin(), 3, nframes, - nall); -} - -template void run_model( - std::vector& dener, - std::vector& dforce_, - std::vector& dvirial, - Session* session, - const std::vector>& input_tensors, - const AtomMap& atommap, - const int nframes, - const int nghost); - -template void run_model( - std::vector& dener, - std::vector& dforce_, - std::vector& dvirial, - Session* session, - const std::vector>& input_tensors, - const AtomMap& atommap, - const int nframes, - const int nghost); - -template void run_model( - std::vector& dener, - std::vector& dforce_, - std::vector& dvirial, - Session* session, - const std::vector>& input_tensors, - const AtomMap& atommap, - const int nframes, - const int nghost); - -template void run_model( - std::vector& dener, - std::vector& dforce_, - std::vector& dvirial, - Session* session, - const std::vector>& input_tensors, - const AtomMap& atommap, - const int nframes, - const int nghost); - -template -static void run_model( - std::vector& dener, - std::vector& dforce_, - std::vector& dvirial, - std::vector& datom_energy_, - std::vector& datom_virial_, - Session* session, - const std::vector>& input_tensors, - const deepmd::AtomMap& atommap, - const int& nframes, - const int& nghost = 0) { - unsigned nloc = atommap.get_type().size(); - unsigned nall = nloc + nghost; - dener.resize(nframes); - if (nloc == 0) { - // no backward map needed - // dforce of size nall * 3 - dforce_.resize(nframes * nall * 3); - fill(dforce_.begin(), dforce_.end(), (VALUETYPE)0.0); - // dvirial of size 9 - dvirial.resize(nframes * 9); - fill(dvirial.begin(), dvirial.end(), (VALUETYPE)0.0); - // datom_energy_ of size nall - datom_energy_.resize(nframes * nall); - fill(datom_energy_.begin(), datom_energy_.end(), (VALUETYPE)0.0); - // datom_virial_ of size nall * 9 - datom_virial_.resize(nframes * nall * 9); - fill(datom_virial_.begin(), datom_virial_.end(), (VALUETYPE)0.0); - return; - } - std::vector output_tensors; - - check_status(session->Run( - input_tensors, {"o_energy", "o_force", "o_atom_energy", "o_atom_virial"}, - {}, &output_tensors)); - - Tensor output_e = output_tensors[0]; - Tensor output_f = output_tensors[1]; - Tensor output_ae = output_tensors[2]; - Tensor output_av = output_tensors[3]; - - auto oe = output_e.flat(); - auto of = output_f.flat(); - auto oae = output_ae.flat(); - auto oav = output_av.flat(); - - std::vector dforce(nframes * 3 * nall); - std::vector datom_energy(nframes * nall, 0); - std::vector datom_virial(nframes * 9 * nall); - dvirial.resize(nframes * 9); - for (int ii = 0; ii < nframes; ++ii) { - dener[ii] = oe(ii); - } - for (int ii = 0; ii < nframes * nall * 3; ++ii) { - dforce[ii] = of(ii); - } - for (int ii = 0; ii < nframes; ++ii) { - for (int jj = 0; jj < nloc; ++jj) { - datom_energy[ii * nall + jj] = oae(ii * nloc + jj); - } - } - for (int ii = 0; ii < nframes * nall * 9; ++ii) { - datom_virial[ii] = oav(ii); - } - // set dvirial to zero, prevent input vector is not zero (#1123) - std::fill(dvirial.begin(), dvirial.end(), (VALUETYPE)0.); - for (int kk = 0; kk < nframes; ++kk) { - for (int ii = 0; ii < nall; ++ii) { - dvirial[kk * 9 + 0] += - (VALUETYPE)1.0 * datom_virial[kk * nall * 9 + 9 * ii + 0]; - dvirial[kk * 9 + 1] += - (VALUETYPE)1.0 * datom_virial[kk * nall * 9 + 9 * ii + 1]; - dvirial[kk * 9 + 2] += - (VALUETYPE)1.0 * datom_virial[kk * nall * 9 + 9 * ii + 2]; - dvirial[kk * 9 + 3] += - (VALUETYPE)1.0 * datom_virial[kk * nall * 9 + 9 * ii + 3]; - dvirial[kk * 9 + 4] += - (VALUETYPE)1.0 * datom_virial[kk * nall * 9 + 9 * ii + 4]; - dvirial[kk * 9 + 5] += - (VALUETYPE)1.0 * datom_virial[kk * nall * 9 + 9 * ii + 5]; - dvirial[kk * 9 + 6] += - (VALUETYPE)1.0 * datom_virial[kk * nall * 9 + 9 * ii + 6]; - dvirial[kk * 9 + 7] += - (VALUETYPE)1.0 * datom_virial[kk * nall * 9 + 9 * ii + 7]; - dvirial[kk * 9 + 8] += - (VALUETYPE)1.0 * datom_virial[kk * nall * 9 + 9 * ii + 8]; - } - } - dforce_ = dforce; - datom_energy_ = datom_energy; - datom_virial_ = datom_virial; - atommap.backward(dforce_.begin(), dforce.begin(), 3, nframes, - nall); - atommap.backward(datom_energy_.begin(), datom_energy.begin(), 1, - nframes, nall); - atommap.backward(datom_virial_.begin(), datom_virial.begin(), 9, - nframes, nall); -} - -template void run_model( - std::vector& dener, - std::vector& dforce_, - std::vector& dvirial, - std::vector& datom_energy_, - std::vector& datom_virial_, - Session* session, - const std::vector>& input_tensors, - const deepmd::AtomMap& atommap, - const int& nframes, - const int& nghost); - -template void run_model( - std::vector& dener, - std::vector& dforce_, - std::vector& dvirial, - std::vector& datom_energy_, - std::vector& datom_virial_, - Session* session, - const std::vector>& input_tensors, - const deepmd::AtomMap& atommap, - const int& nframes, - const int& nghost); - -template void run_model( - std::vector& dener, - std::vector& dforce_, - std::vector& dvirial, - std::vector& datom_energy_, - std::vector& datom_virial_, - Session* session, - const std::vector>& input_tensors, - const deepmd::AtomMap& atommap, - const int& nframes, - const int& nghost); - -template void run_model( - std::vector& dener, - std::vector& dforce_, - std::vector& dvirial, - std::vector& datom_energy_, - std::vector& datom_virial_, - Session* session, - const std::vector>& input_tensors, - const deepmd::AtomMap& atommap, - const int& nframes, - const int& nghost); - -// end multiple frames - -// start single frame - -template -static void run_model( - ENERGYTYPE& dener, - std::vector& dforce_, - std::vector& dvirial, - Session* session, - const std::vector>& input_tensors, - const AtomMap& atommap, - const int nframes = 1, - const int nghost = 0) { - assert(nframes == 1); - std::vector dener_(1); - // call multi-frame version - run_model(dener_, dforce_, dvirial, session, - input_tensors, atommap, nframes, nghost); - dener = dener_[0]; -} - -template void run_model( - ENERGYTYPE& dener, - std::vector& dforce_, - std::vector& dvirial, - Session* session, - const std::vector>& input_tensors, - const AtomMap& atommap, - const int nframes, - const int nghost); - -template void run_model( - ENERGYTYPE& dener, - std::vector& dforce_, - std::vector& dvirial, - Session* session, - const std::vector>& input_tensors, - const AtomMap& atommap, - const int nframes, - const int nghost); - -template void run_model( - ENERGYTYPE& dener, - std::vector& dforce_, - std::vector& dvirial, - Session* session, - const std::vector>& input_tensors, - const AtomMap& atommap, - const int nframes, - const int nghost); - -template void run_model( - ENERGYTYPE& dener, - std::vector& dforce_, - std::vector& dvirial, - Session* session, - const std::vector>& input_tensors, - const AtomMap& atommap, - const int nframes, - const int nghost); - -template -static void run_model( - ENERGYTYPE& dener, - std::vector& dforce_, - std::vector& dvirial, - std::vector& datom_energy_, - std::vector& datom_virial_, - Session* session, - const std::vector>& input_tensors, - const deepmd::AtomMap& atommap, - const int& nframes = 1, - const int& nghost = 0) { - assert(nframes == 1); - std::vector dener_(1); - // call multi-frame version - run_model(dener_, dforce_, dvirial, datom_energy_, - datom_virial_, session, input_tensors, - atommap, nframes, nghost); - dener = dener_[0]; -} - -template void run_model( - ENERGYTYPE& dener, - std::vector& dforce_, - std::vector& dvirial, - std::vector& datom_energy_, - std::vector& datom_virial_, - Session* session, - const std::vector>& input_tensors, - const deepmd::AtomMap& atommap, - const int& nframes, - const int& nghost); - -template void run_model( - ENERGYTYPE& dener, - std::vector& dforce_, - std::vector& dvirial, - std::vector& datom_energy_, - std::vector& datom_virial_, - Session* session, - const std::vector>& input_tensors, - const deepmd::AtomMap& atommap, - const int& nframes, - const int& nghost); - -template void run_model( - ENERGYTYPE& dener, - std::vector& dforce_, - std::vector& dvirial, - std::vector& datom_energy_, - std::vector& datom_virial_, - Session* session, - const std::vector>& input_tensors, - const deepmd::AtomMap& atommap, - const int& nframes, - const int& nghost); - -template void run_model( - ENERGYTYPE& dener, - std::vector& dforce_, - std::vector& dvirial, - std::vector& datom_energy_, - std::vector& datom_virial_, - Session* session, - const std::vector>& input_tensors, - const deepmd::AtomMap& atommap, - const int& nframes, - const int& nghost); - -// end single frame - -DeepPot::DeepPot() - : inited(false), init_nbor(false), graph_def(new GraphDef()) {} +DeepPot::DeepPot() : inited(false) {} DeepPot::DeepPot(const std::string& model, const int& gpu_rank, const std::string& file_content) - : inited(false), init_nbor(false), graph_def(new GraphDef()) { - try { - init(model, gpu_rank, file_content); - } catch (...) { - // Clean up and rethrow, as the destructor will not be called - delete graph_def; - throw; - } + : inited(false) { + init(model, gpu_rank, file_content); } -DeepPot::~DeepPot() { delete graph_def; } +DeepPot::~DeepPot() {} void DeepPot::init(const std::string& model, const int& gpu_rank, @@ -424,152 +32,25 @@ void DeepPot::init(const std::string& model, << std::endl; return; } - SessionOptions options; - get_env_nthreads(num_intra_nthreads, num_inter_nthreads); - options.config.set_inter_op_parallelism_threads(num_inter_nthreads); - options.config.set_intra_op_parallelism_threads(num_intra_nthreads); - deepmd::load_op_library(); - - if (file_content.size() == 0) { - check_status(ReadBinaryProto(Env::Default(), model, graph_def)); + // TODO: To implement detect_backend + DPBackend backend = deepmd::DPBackend::TensorFlow; + if (deepmd::DPBackend::TensorFlow == backend) { + // TODO: throw errors if TF backend is not built, without mentioning TF + dp = std::make_shared(model, gpu_rank, file_content); + } else if (deepmd::DPBackend::PyTorch == backend) { + throw deepmd::deepmd_exception("PyTorch backend is not supported yet"); + } else if (deepmd::DPBackend::Paddle == backend) { + throw deepmd::deepmd_exception("PaddlePaddle backend is not supported yet"); } else { - (*graph_def).ParseFromString(file_content); - } - int gpu_num = -1; -#if GOOGLE_CUDA || TENSORFLOW_USE_ROCM - DPGetDeviceCount(gpu_num); // check current device environment - if (gpu_num > 0) { - options.config.set_allow_soft_placement(true); - options.config.mutable_gpu_options()->set_per_process_gpu_memory_fraction( - 0.9); - options.config.mutable_gpu_options()->set_allow_growth(true); - DPErrcheck(DPSetDevice(gpu_rank % gpu_num)); - std::string str = "/gpu:"; - str += std::to_string(gpu_rank % gpu_num); - graph::SetDefaultDevice(str, graph_def); - } -#endif // GOOGLE_CUDA || TENSORFLOW_USE_ROCM - check_status(NewSession(options, &session)); - check_status(session->Create(*graph_def)); - try { - model_version = get_scalar("model_attr/model_version"); - } catch (deepmd::tf_exception& e) { - // no model version defined in old models - model_version = "0.0"; - } - if (!model_compatable(model_version)) { - throw deepmd::deepmd_exception( - "incompatable model: version " + model_version + - " in graph, but version " + global_model_version + - " supported " - "See https://deepmd.rtfd.io/compatability/ for details."); - } - dtype = session_get_dtype(session, "descrpt_attr/rcut"); - if (dtype == tensorflow::DT_DOUBLE) { - rcut = get_scalar("descrpt_attr/rcut"); - } else { - rcut = get_scalar("descrpt_attr/rcut"); - } - cell_size = rcut; - ntypes = get_scalar("descrpt_attr/ntypes"); - try { - ntypes_spin = get_scalar("spin_attr/ntypes_spin"); - } catch (const deepmd::deepmd_exception) { - ntypes_spin = 0; - } - dfparam = get_scalar("fitting_attr/dfparam"); - daparam = get_scalar("fitting_attr/daparam"); - if (dfparam < 0) { - dfparam = 0; - } - if (daparam < 0) { - daparam = 0; + throw deepmd::deepmd_exception("Unknown file type"); } - if (daparam > 0) { - try { - aparam_nall = get_scalar("fitting_attr/aparam_nall"); - } catch (const deepmd::deepmd_exception) { - aparam_nall = false; - } - } else { - aparam_nall = false; - } - model_type = get_scalar("model_attr/model_type"); inited = true; - - init_nbor = false; } void DeepPot::print_summary(const std::string& pre) const { deepmd::print_summary(pre); } -template -VT DeepPot::get_scalar(const std::string& name) const { - return session_get_scalar(session, name); -} - -template -void DeepPot::validate_fparam_aparam( - const int& nframes, - const int& nloc, - const std::vector& fparam, - const std::vector& aparam) const { - if (fparam.size() != dfparam && fparam.size() != nframes * dfparam) { - throw deepmd::deepmd_exception( - "the dim of frame parameter provided is not consistent with what the " - "model uses"); - } - - if (aparam.size() != daparam * nloc && - aparam.size() != nframes * daparam * nloc) { - throw deepmd::deepmd_exception( - "the dim of atom parameter provided is not consistent with what the " - "model uses"); - } -} - -template void DeepPot::validate_fparam_aparam( - const int& nframes, - const int& nloc, - const std::vector& fparam, - const std::vector& aparam) const; - -template void DeepPot::validate_fparam_aparam( - const int& nframes, - const int& nloc, - const std::vector& fparam, - const std::vector& aparam) const; - -template -void DeepPot::tile_fparam_aparam(std::vector& out_param, - const int& nframes, - const int& dparam, - const std::vector& param) const { - if (param.size() == dparam) { - out_param.resize(nframes * dparam); - for (int ii = 0; ii < nframes; ++ii) { - std::copy(param.begin(), param.end(), out_param.begin() + ii * dparam); - } - } else if (param.size() == nframes * dparam) { - out_param = param; - } -} - -template void DeepPot::tile_fparam_aparam( - std::vector& out_param, - const int& nframes, - const int& dparam, - const std::vector& param) const; - -template void DeepPot::tile_fparam_aparam( - std::vector& out_param, - const int& nframes, - const int& dparam, - const std::vector& param) const; - -// ENERGYVTYPE: std::vector or ENERGYTYPE - template void DeepPot::compute(ENERGYVTYPE& dener, std::vector& dforce_, @@ -579,37 +60,8 @@ void DeepPot::compute(ENERGYVTYPE& dener, const std::vector& dbox, const std::vector& fparam_, const std::vector& aparam_) { - int nall = datype_.size(); - // if nall==0, unclear nframes, but 1 is ok - int nframes = nall > 0 ? (dcoord_.size() / nall / 3) : 1; - int nloc = nall; - atommap = deepmd::AtomMap(datype_.begin(), datype_.begin() + nloc); - assert(nloc == atommap.get_type().size()); - std::vector fparam; - std::vector aparam; - validate_fparam_aparam(nframes, (aparam_nall ? nall : nloc), fparam_, - aparam_); - tile_fparam_aparam(fparam, nframes, dfparam, fparam_); - tile_fparam_aparam(aparam, nframes, (aparam_nall ? nall : nloc) * daparam, - aparam_); - - std::vector> input_tensors; - - if (dtype == tensorflow::DT_DOUBLE) { - int ret = session_input_tensors(input_tensors, dcoord_, ntypes, - datype_, dbox, cell_size, fparam, - aparam, atommap, "", aparam_nall); - assert(ret == nloc); - run_model(dener, dforce_, dvirial, session, input_tensors, atommap, - nframes); - } else { - int ret = session_input_tensors(input_tensors, dcoord_, ntypes, - datype_, dbox, cell_size, fparam, - aparam, atommap, "", aparam_nall); - assert(ret == nloc); - run_model(dener, dforce_, dvirial, session, input_tensors, atommap, - nframes); - } + dp->computew(dener, dforce_, dvirial, dcoord_, datype_, dbox, fparam_, + aparam_); } template void DeepPot::compute( @@ -664,37 +116,8 @@ void DeepPot::compute(ENERGYVTYPE& dener, const int& ago, const std::vector& fparam_, const std::vector& aparam__) { - int nall = datype_.size(); - // if nall==0, unclear nframes, but 1 is ok - int nframes = nall > 0 ? (dcoord_.size() / nall / 3) : 1; - std::vector fparam; - std::vector aparam_; - validate_fparam_aparam(nframes, (aparam_nall ? nall : (nall - nghost)), - fparam_, aparam__); - tile_fparam_aparam(fparam, nframes, dfparam, fparam_); - tile_fparam_aparam(aparam_, nframes, - (aparam_nall ? nall : (nall - nghost)) * daparam, - aparam__); - - // select real atoms - std::vector dcoord, dforce, aparam; - std::vector datype, fwd_map, bkw_map; - int nghost_real, nall_real, nloc_real; - select_real_atoms_coord(dcoord, datype, aparam, nghost_real, fwd_map, bkw_map, - nall_real, nloc_real, dcoord_, datype_, aparam_, - nghost, ntypes, nframes, daparam, nall, aparam_nall); - - // internal nlist - if (ago == 0) { - nlist_data.copy_from_nlist(lmp_list); - nlist_data.shuffle_exclude_empty(fwd_map); - } - compute_inner(dener, dforce, dvirial, dcoord, datype, dbox, nghost_real, ago, - fparam, aparam); - // bkw map - dforce_.resize(nframes * fwd_map.size() * 3); - select_map(dforce_, dforce, bkw_map, 3, nframes, fwd_map.size(), - bkw_map.size()); + dp->computew(dener, dforce_, dvirial, dcoord_, datype_, dbox, nghost, + lmp_list, ago, fparam_, aparam__); } template void DeepPot::compute( @@ -749,96 +172,6 @@ template void DeepPot::compute>( const std::vector& fparam, const std::vector& aparam_); -template -void DeepPot::compute_inner(ENERGYVTYPE& dener, - std::vector& dforce_, - std::vector& dvirial, - const std::vector& dcoord_, - const std::vector& datype_, - const std::vector& dbox, - const int nghost, - const int& ago, - const std::vector& fparam, - const std::vector& aparam) { - int nall = datype_.size(); - // if nall==0, unclear nframes, but 1 is ok - int nframes = nall > 0 ? (dcoord_.size() / nall / 3) : 1; - int nloc = nall - nghost; - - std::vector> input_tensors; - - // agp == 0 means that the LAMMPS nbor list has been updated - if (ago == 0) { - atommap = deepmd::AtomMap(datype_.begin(), datype_.begin() + nloc); - assert(nloc == atommap.get_type().size()); - nlist_data.shuffle(atommap); - nlist_data.make_inlist(nlist); - } - if (dtype == tensorflow::DT_DOUBLE) { - int ret = session_input_tensors( - input_tensors, dcoord_, ntypes, datype_, dbox, nlist, fparam, aparam, - atommap, nghost, ago, "", aparam_nall); - assert(nloc == ret); - run_model(dener, dforce_, dvirial, session, input_tensors, atommap, - nframes, nghost); - } else { - int ret = session_input_tensors( - input_tensors, dcoord_, ntypes, datype_, dbox, nlist, fparam, aparam, - atommap, nghost, ago, "", aparam_nall); - assert(nloc == ret); - run_model(dener, dforce_, dvirial, session, input_tensors, atommap, - nframes, nghost); - } -} - -template void DeepPot::compute_inner( - ENERGYTYPE& dener, - std::vector& dforce_, - std::vector& dvirial, - const std::vector& dcoord_, - const std::vector& datype_, - const std::vector& dbox, - const int nghost, - const int& ago, - const std::vector& fparam, - const std::vector& aparam); - -template void DeepPot::compute_inner( - ENERGYTYPE& dener, - std::vector& dforce_, - std::vector& dvirial, - const std::vector& dcoord_, - const std::vector& datype_, - const std::vector& dbox, - const int nghost, - const int& ago, - const std::vector& fparam, - const std::vector& aparam); - -template void DeepPot::compute_inner>( - std::vector& dener, - std::vector& dforce_, - std::vector& dvirial, - const std::vector& dcoord_, - const std::vector& datype_, - const std::vector& dbox, - const int nghost, - const int& ago, - const std::vector& fparam, - const std::vector& aparam); - -template void DeepPot::compute_inner>( - std::vector& dener, - std::vector& dforce_, - std::vector& dvirial, - const std::vector& dcoord_, - const std::vector& datype_, - const std::vector& dbox, - const int nghost, - const int& ago, - const std::vector& fparam, - const std::vector& aparam); - template void DeepPot::compute(ENERGYVTYPE& dener, std::vector& dforce_, @@ -850,31 +183,8 @@ void DeepPot::compute(ENERGYVTYPE& dener, const std::vector& dbox, const std::vector& fparam_, const std::vector& aparam_) { - // if datype.size is 0, not clear nframes; but 1 is just ok - int nframes = datype_.size() > 0 ? (dcoord_.size() / 3 / datype_.size()) : 1; - atommap = deepmd::AtomMap(datype_.begin(), datype_.end()); - int nloc = datype_.size(); - std::vector fparam; - std::vector aparam; - validate_fparam_aparam(nframes, nloc, fparam_, aparam_); - tile_fparam_aparam(fparam, nframes, dfparam, fparam_); - tile_fparam_aparam(aparam, nframes, nloc * daparam, aparam_); - - std::vector> input_tensors; - - if (dtype == tensorflow::DT_DOUBLE) { - int nloc = session_input_tensors(input_tensors, dcoord_, ntypes, - datype_, dbox, cell_size, fparam, - aparam, atommap, "", aparam_nall); - run_model(dener, dforce_, dvirial, datom_energy_, datom_virial_, - session, input_tensors, atommap, nframes); - } else { - int nloc = session_input_tensors(input_tensors, dcoord_, ntypes, - datype_, dbox, cell_size, fparam, - aparam, atommap, "", aparam_nall); - run_model(dener, dforce_, dvirial, datom_energy_, datom_virial_, - session, input_tensors, atommap, nframes); - } + dp->computew(dener, dforce_, dvirial, datom_energy_, datom_virial_, dcoord_, + datype_, dbox, fparam_, aparam_); } template void DeepPot::compute( @@ -939,62 +249,8 @@ void DeepPot::compute(ENERGYVTYPE& dener, const int& ago, const std::vector& fparam_, const std::vector& aparam__) { - int nall = datype_.size(); - // if nall==0, unclear nframes, but 1 is ok - int nframes = nall > 0 ? (dcoord_.size() / nall / 3) : 1; - int nloc = nall - nghost; - std::vector fparam; - std::vector aparam_; - validate_fparam_aparam(nframes, (aparam_nall ? nall : nloc), fparam_, - aparam__); - tile_fparam_aparam(fparam, nframes, dfparam, fparam_); - tile_fparam_aparam(aparam_, nframes, (aparam_nall ? nall : nloc) * daparam, - aparam__); - std::vector> input_tensors; - // select real atoms - std::vector dcoord, dforce, aparam, datom_energy, datom_virial; - std::vector datype, fwd_map, bkw_map; - int nghost_real, nall_real, nloc_real; - select_real_atoms_coord(dcoord, datype, aparam, nghost_real, fwd_map, bkw_map, - nall_real, nloc_real, dcoord_, datype_, aparam_, - nghost, ntypes, nframes, daparam, nall, aparam_nall); - - if (ago == 0) { - atommap = deepmd::AtomMap(datype.begin(), datype.begin() + nloc_real); - assert(nloc_real == atommap.get_type().size()); - - nlist_data.copy_from_nlist(lmp_list); - nlist_data.shuffle_exclude_empty(fwd_map); - nlist_data.shuffle(atommap); - nlist_data.make_inlist(nlist); - } - - if (dtype == tensorflow::DT_DOUBLE) { - int ret = session_input_tensors( - input_tensors, dcoord, ntypes, datype, dbox, nlist, fparam, aparam, - atommap, nghost_real, ago, "", aparam_nall); - assert(nloc_real == ret); - run_model(dener, dforce, dvirial, datom_energy, datom_virial, - session, input_tensors, atommap, nframes, nghost_real); - } else { - int ret = session_input_tensors( - input_tensors, dcoord, ntypes, datype, dbox, nlist, fparam, aparam, - atommap, nghost_real, ago, "", aparam_nall); - assert(nloc_real == ret); - run_model(dener, dforce, dvirial, datom_energy, datom_virial, - session, input_tensors, atommap, nframes, nghost_real); - } - - // bkw map - dforce_.resize(nframes * fwd_map.size() * 3); - datom_energy_.resize(nframes * fwd_map.size()); - datom_virial_.resize(nframes * fwd_map.size() * 9); - select_map(dforce_, dforce, bkw_map, 3, nframes, fwd_map.size(), - nall_real); - select_map(datom_energy_, datom_energy, bkw_map, 1, nframes, - fwd_map.size(), nall_real); - select_map(datom_virial_, datom_virial, bkw_map, 9, nframes, - fwd_map.size(), nall_real); + dp->computew(dener, dforce_, dvirial, datom_energy_, datom_virial_, dcoord_, + datype_, dbox, nghost, lmp_list, ago, fparam_, aparam__); } template void DeepPot::compute( @@ -1068,32 +324,8 @@ void DeepPot::compute_mixed_type(ENERGYVTYPE& dener, const std::vector& dbox, const std::vector& fparam_, const std::vector& aparam_) { - int nloc = datype_.size() / nframes; - // here atommap only used to get nloc - atommap = deepmd::AtomMap(datype_.begin(), datype_.begin() + nloc); - std::vector fparam; - std::vector aparam; - validate_fparam_aparam(nframes, nloc, fparam_, aparam_); - tile_fparam_aparam(fparam, nframes, dfparam, fparam_); - tile_fparam_aparam(aparam, nframes, nloc * daparam, aparam_); - - std::vector> input_tensors; - - if (dtype == tensorflow::DT_DOUBLE) { - int ret = session_input_tensors_mixed_type( - input_tensors, nframes, dcoord_, ntypes, datype_, dbox, cell_size, - fparam, aparam, atommap, "", aparam_nall); - assert(ret == nloc); - run_model(dener, dforce_, dvirial, session, input_tensors, atommap, - nframes); - } else { - int ret = session_input_tensors_mixed_type( - input_tensors, nframes, dcoord_, ntypes, datype_, dbox, cell_size, - fparam, aparam, atommap, "", aparam_nall); - assert(ret == nloc); - run_model(dener, dforce_, dvirial, session, input_tensors, atommap, - nframes); - } + dp->computew_mixed_type(dener, dforce_, dvirial, nframes, dcoord_, datype_, + dbox, fparam_, aparam_); } template void DeepPot::compute_mixed_type( @@ -1152,30 +384,8 @@ void DeepPot::compute_mixed_type(ENERGYVTYPE& dener, const std::vector& dbox, const std::vector& fparam_, const std::vector& aparam_) { - int nloc = datype_.size() / nframes; - // here atommap only used to get nloc - atommap = deepmd::AtomMap(datype_.begin(), datype_.begin() + nloc); - std::vector fparam; - std::vector aparam; - validate_fparam_aparam(nframes, nloc, fparam_, aparam_); - tile_fparam_aparam(fparam, nframes, dfparam, fparam_); - tile_fparam_aparam(aparam, nframes, nloc * daparam, aparam_); - - std::vector> input_tensors; - - if (dtype == tensorflow::DT_DOUBLE) { - int nloc = session_input_tensors_mixed_type( - input_tensors, nframes, dcoord_, ntypes, datype_, dbox, cell_size, - fparam, aparam, atommap, "", aparam_nall); - run_model(dener, dforce_, dvirial, datom_energy_, datom_virial_, - session, input_tensors, atommap, nframes); - } else { - int nloc = session_input_tensors_mixed_type( - input_tensors, nframes, dcoord_, ntypes, datype_, dbox, cell_size, - fparam, aparam, atommap, "", aparam_nall); - run_model(dener, dforce_, dvirial, datom_energy_, datom_virial_, - session, input_tensors, atommap, nframes); - } + dp->computew_mixed_type(dener, dforce_, dvirial, datom_energy_, datom_virial_, + nframes, dcoord_, datype_, dbox, fparam_, aparam_); } template void DeepPot::compute_mixed_type( @@ -1230,10 +440,22 @@ template void DeepPot::compute_mixed_type>( const std::vector& fparam, const std::vector& aparam); +double DeepPot::cutoff() const { return dp->cutoff(); } + +int DeepPot::numb_types() const { return dp->numb_types(); } + +int DeepPot::numb_types_spin() const { return dp->numb_types_spin(); } + +int DeepPot::dim_fparam() const { return dp->dim_fparam(); } + +int DeepPot::dim_aparam() const { return dp->dim_aparam(); } + void DeepPot::get_type_map(std::string& type_map) { - type_map = get_scalar("model_attr/tmap"); + dp->get_type_map(type_map); } +bool DeepPot::is_aparam_nall() const { return dp->is_aparam_nall(); } + DeepPotModelDevi::DeepPotModelDevi() : inited(false), numb_models(0) {} DeepPotModelDevi::DeepPotModelDevi( diff --git a/source/api_cc/src/DeepPotTF.cc b/source/api_cc/src/DeepPotTF.cc new file mode 100644 index 0000000000..ed3bede4df --- /dev/null +++ b/source/api_cc/src/DeepPotTF.cc @@ -0,0 +1,1556 @@ +// SPDX-License-Identifier: LGPL-3.0-or-later +#include "DeepPotTF.h" + +#include + +#include "AtomMap.h" +#include "common.h" +#include "device.h" + +using namespace tensorflow; +using namespace deepmd; + +// start multiple frames + +template +static void run_model( + std::vector& dener, + std::vector& dforce_, + std::vector& dvirial, + Session* session, + const std::vector>& input_tensors, + const AtomMap& atommap, + const int nframes, + const int nghost = 0) { + unsigned nloc = atommap.get_type().size(); + unsigned nall = nloc + nghost; + dener.resize(nframes); + if (nloc == 0) { + // no backward map needed + // dforce of size nall * 3 + dforce_.resize(static_cast(nframes) * nall * 3); + fill(dforce_.begin(), dforce_.end(), (VALUETYPE)0.0); + // dvirial of size 9 + dvirial.resize(static_cast(nframes) * 9); + fill(dvirial.begin(), dvirial.end(), (VALUETYPE)0.0); + return; + } + + std::vector output_tensors; + check_status(session->Run( + input_tensors, {"o_energy", "o_force", "o_atom_energy", "o_atom_virial"}, + {}, &output_tensors)); + + Tensor output_e = output_tensors[0]; + Tensor output_f = output_tensors[1]; + Tensor output_av = output_tensors[3]; + + auto oe = output_e.flat(); + auto of = output_f.flat(); + auto oav = output_av.flat(); + + std::vector dforce(static_cast(nframes) * 3 * nall); + dvirial.resize(static_cast(nframes) * 9); + for (int ii = 0; ii < nframes; ++ii) { + dener[ii] = oe(ii); + } + for (size_t ii = 0; ii < static_cast(nframes) * nall * 3; ++ii) { + dforce[ii] = of(ii); + } + // set dvirial to zero, prevent input vector is not zero (#1123) + std::fill(dvirial.begin(), dvirial.end(), (VALUETYPE)0.); + for (int kk = 0; kk < nframes; ++kk) { + for (int ii = 0; ii < nall; ++ii) { + dvirial[kk * 9 + 0] += (VALUETYPE)1.0 * oav(kk * nall * 9 + 9 * ii + 0); + dvirial[kk * 9 + 1] += (VALUETYPE)1.0 * oav(kk * nall * 9 + 9 * ii + 1); + dvirial[kk * 9 + 2] += (VALUETYPE)1.0 * oav(kk * nall * 9 + 9 * ii + 2); + dvirial[kk * 9 + 3] += (VALUETYPE)1.0 * oav(kk * nall * 9 + 9 * ii + 3); + dvirial[kk * 9 + 4] += (VALUETYPE)1.0 * oav(kk * nall * 9 + 9 * ii + 4); + dvirial[kk * 9 + 5] += (VALUETYPE)1.0 * oav(kk * nall * 9 + 9 * ii + 5); + dvirial[kk * 9 + 6] += (VALUETYPE)1.0 * oav(kk * nall * 9 + 9 * ii + 6); + dvirial[kk * 9 + 7] += (VALUETYPE)1.0 * oav(kk * nall * 9 + 9 * ii + 7); + dvirial[kk * 9 + 8] += (VALUETYPE)1.0 * oav(kk * nall * 9 + 9 * ii + 8); + } + } + dforce_ = dforce; + atommap.backward(dforce_.begin(), dforce.begin(), 3, nframes, + nall); +} + +template void run_model( + std::vector& dener, + std::vector& dforce_, + std::vector& dvirial, + Session* session, + const std::vector>& input_tensors, + const AtomMap& atommap, + const int nframes, + const int nghost); + +template void run_model( + std::vector& dener, + std::vector& dforce_, + std::vector& dvirial, + Session* session, + const std::vector>& input_tensors, + const AtomMap& atommap, + const int nframes, + const int nghost); + +template void run_model( + std::vector& dener, + std::vector& dforce_, + std::vector& dvirial, + Session* session, + const std::vector>& input_tensors, + const AtomMap& atommap, + const int nframes, + const int nghost); + +template void run_model( + std::vector& dener, + std::vector& dforce_, + std::vector& dvirial, + Session* session, + const std::vector>& input_tensors, + const AtomMap& atommap, + const int nframes, + const int nghost); + +template +static void run_model( + std::vector& dener, + std::vector& dforce_, + std::vector& dvirial, + std::vector& datom_energy_, + std::vector& datom_virial_, + Session* session, + const std::vector>& input_tensors, + const deepmd::AtomMap& atommap, + const int& nframes, + const int& nghost = 0) { + unsigned nloc = atommap.get_type().size(); + unsigned nall = nloc + nghost; + dener.resize(nframes); + if (nloc == 0) { + // no backward map needed + // dforce of size nall * 3 + dforce_.resize(static_cast(nframes) * nall * 3); + fill(dforce_.begin(), dforce_.end(), (VALUETYPE)0.0); + // dvirial of size 9 + dvirial.resize(static_cast(nframes) * 9); + fill(dvirial.begin(), dvirial.end(), (VALUETYPE)0.0); + // datom_energy_ of size nall + datom_energy_.resize(static_cast(nframes) * nall); + fill(datom_energy_.begin(), datom_energy_.end(), (VALUETYPE)0.0); + // datom_virial_ of size nall * 9 + datom_virial_.resize(static_cast(nframes) * nall * 9); + fill(datom_virial_.begin(), datom_virial_.end(), (VALUETYPE)0.0); + return; + } + std::vector output_tensors; + + check_status(session->Run( + input_tensors, {"o_energy", "o_force", "o_atom_energy", "o_atom_virial"}, + {}, &output_tensors)); + + Tensor output_e = output_tensors[0]; + Tensor output_f = output_tensors[1]; + Tensor output_ae = output_tensors[2]; + Tensor output_av = output_tensors[3]; + + auto oe = output_e.flat(); + auto of = output_f.flat(); + auto oae = output_ae.flat(); + auto oav = output_av.flat(); + + std::vector dforce(static_cast(nframes) * 3 * nall); + std::vector datom_energy(static_cast(nframes) * nall, 0); + std::vector datom_virial(static_cast(nframes) * 9 * nall); + dvirial.resize(static_cast(nframes) * 9); + for (int ii = 0; ii < nframes; ++ii) { + dener[ii] = oe(ii); + } + for (size_t ii = 0; ii < static_cast(nframes) * nall * 3; ++ii) { + dforce[ii] = of(ii); + } + for (int ii = 0; ii < nframes; ++ii) { + for (int jj = 0; jj < nloc; ++jj) { + datom_energy[ii * nall + jj] = oae(ii * nloc + jj); + } + } + for (size_t ii = 0; ii < static_cast(nframes) * nall * 9; ++ii) { + datom_virial[ii] = oav(ii); + } + // set dvirial to zero, prevent input vector is not zero (#1123) + std::fill(dvirial.begin(), dvirial.end(), (VALUETYPE)0.); + for (int kk = 0; kk < nframes; ++kk) { + for (int ii = 0; ii < nall; ++ii) { + dvirial[kk * 9 + 0] += + (VALUETYPE)1.0 * datom_virial[kk * nall * 9 + 9 * ii + 0]; + dvirial[kk * 9 + 1] += + (VALUETYPE)1.0 * datom_virial[kk * nall * 9 + 9 * ii + 1]; + dvirial[kk * 9 + 2] += + (VALUETYPE)1.0 * datom_virial[kk * nall * 9 + 9 * ii + 2]; + dvirial[kk * 9 + 3] += + (VALUETYPE)1.0 * datom_virial[kk * nall * 9 + 9 * ii + 3]; + dvirial[kk * 9 + 4] += + (VALUETYPE)1.0 * datom_virial[kk * nall * 9 + 9 * ii + 4]; + dvirial[kk * 9 + 5] += + (VALUETYPE)1.0 * datom_virial[kk * nall * 9 + 9 * ii + 5]; + dvirial[kk * 9 + 6] += + (VALUETYPE)1.0 * datom_virial[kk * nall * 9 + 9 * ii + 6]; + dvirial[kk * 9 + 7] += + (VALUETYPE)1.0 * datom_virial[kk * nall * 9 + 9 * ii + 7]; + dvirial[kk * 9 + 8] += + (VALUETYPE)1.0 * datom_virial[kk * nall * 9 + 9 * ii + 8]; + } + } + dforce_ = dforce; + datom_energy_ = datom_energy; + datom_virial_ = datom_virial; + atommap.backward(dforce_.begin(), dforce.begin(), 3, nframes, + nall); + atommap.backward(datom_energy_.begin(), datom_energy.begin(), 1, + nframes, nall); + atommap.backward(datom_virial_.begin(), datom_virial.begin(), 9, + nframes, nall); +} + +template void run_model( + std::vector& dener, + std::vector& dforce_, + std::vector& dvirial, + std::vector& datom_energy_, + std::vector& datom_virial_, + Session* session, + const std::vector>& input_tensors, + const deepmd::AtomMap& atommap, + const int& nframes, + const int& nghost); + +template void run_model( + std::vector& dener, + std::vector& dforce_, + std::vector& dvirial, + std::vector& datom_energy_, + std::vector& datom_virial_, + Session* session, + const std::vector>& input_tensors, + const deepmd::AtomMap& atommap, + const int& nframes, + const int& nghost); + +template void run_model( + std::vector& dener, + std::vector& dforce_, + std::vector& dvirial, + std::vector& datom_energy_, + std::vector& datom_virial_, + Session* session, + const std::vector>& input_tensors, + const deepmd::AtomMap& atommap, + const int& nframes, + const int& nghost); + +template void run_model( + std::vector& dener, + std::vector& dforce_, + std::vector& dvirial, + std::vector& datom_energy_, + std::vector& datom_virial_, + Session* session, + const std::vector>& input_tensors, + const deepmd::AtomMap& atommap, + const int& nframes, + const int& nghost); + +// end multiple frames + +// start single frame + +template +static void run_model( + ENERGYTYPE& dener, + std::vector& dforce_, + std::vector& dvirial, + Session* session, + const std::vector>& input_tensors, + const AtomMap& atommap, + const int nframes = 1, + const int nghost = 0) { + assert(nframes == 1); + std::vector dener_(1); + // call multi-frame version + run_model(dener_, dforce_, dvirial, session, + input_tensors, atommap, nframes, nghost); + dener = dener_[0]; +} + +template void run_model( + ENERGYTYPE& dener, + std::vector& dforce_, + std::vector& dvirial, + Session* session, + const std::vector>& input_tensors, + const AtomMap& atommap, + const int nframes, + const int nghost); + +template void run_model( + ENERGYTYPE& dener, + std::vector& dforce_, + std::vector& dvirial, + Session* session, + const std::vector>& input_tensors, + const AtomMap& atommap, + const int nframes, + const int nghost); + +template void run_model( + ENERGYTYPE& dener, + std::vector& dforce_, + std::vector& dvirial, + Session* session, + const std::vector>& input_tensors, + const AtomMap& atommap, + const int nframes, + const int nghost); + +template void run_model( + ENERGYTYPE& dener, + std::vector& dforce_, + std::vector& dvirial, + Session* session, + const std::vector>& input_tensors, + const AtomMap& atommap, + const int nframes, + const int nghost); + +template +static void run_model( + ENERGYTYPE& dener, + std::vector& dforce_, + std::vector& dvirial, + std::vector& datom_energy_, + std::vector& datom_virial_, + Session* session, + const std::vector>& input_tensors, + const deepmd::AtomMap& atommap, + const int& nframes = 1, + const int& nghost = 0) { + assert(nframes == 1); + std::vector dener_(1); + // call multi-frame version + run_model(dener_, dforce_, dvirial, datom_energy_, + datom_virial_, session, input_tensors, + atommap, nframes, nghost); + dener = dener_[0]; +} + +template void run_model( + ENERGYTYPE& dener, + std::vector& dforce_, + std::vector& dvirial, + std::vector& datom_energy_, + std::vector& datom_virial_, + Session* session, + const std::vector>& input_tensors, + const deepmd::AtomMap& atommap, + const int& nframes, + const int& nghost); + +template void run_model( + ENERGYTYPE& dener, + std::vector& dforce_, + std::vector& dvirial, + std::vector& datom_energy_, + std::vector& datom_virial_, + Session* session, + const std::vector>& input_tensors, + const deepmd::AtomMap& atommap, + const int& nframes, + const int& nghost); + +template void run_model( + ENERGYTYPE& dener, + std::vector& dforce_, + std::vector& dvirial, + std::vector& datom_energy_, + std::vector& datom_virial_, + Session* session, + const std::vector>& input_tensors, + const deepmd::AtomMap& atommap, + const int& nframes, + const int& nghost); + +template void run_model( + ENERGYTYPE& dener, + std::vector& dforce_, + std::vector& dvirial, + std::vector& datom_energy_, + std::vector& datom_virial_, + Session* session, + const std::vector>& input_tensors, + const deepmd::AtomMap& atommap, + const int& nframes, + const int& nghost); + +// end single frame + +DeepPotTF::DeepPotTF() + : inited(false), init_nbor(false), graph_def(new GraphDef()) {} + +DeepPotTF::DeepPotTF(const std::string& model, + const int& gpu_rank, + const std::string& file_content) + : inited(false), init_nbor(false), graph_def(new GraphDef()) { + try { + init(model, gpu_rank, file_content); + } catch (...) { + // Clean up and rethrow, as the destructor will not be called + delete graph_def; + throw; + } +} + +DeepPotTF::~DeepPotTF() { delete graph_def; } + +void DeepPotTF::init(const std::string& model, + const int& gpu_rank, + const std::string& file_content) { + if (inited) { + std::cerr << "WARNING: deepmd-kit should not be initialized twice, do " + "nothing at the second call of initializer" + << std::endl; + return; + } + SessionOptions options; + get_env_nthreads(num_intra_nthreads, num_inter_nthreads); + options.config.set_inter_op_parallelism_threads(num_inter_nthreads); + options.config.set_intra_op_parallelism_threads(num_intra_nthreads); + deepmd::load_op_library(); + + if (file_content.size() == 0) { + check_status(ReadBinaryProto(Env::Default(), model, graph_def)); + } else { + (*graph_def).ParseFromString(file_content); + } + int gpu_num = -1; +#if GOOGLE_CUDA || TENSORFLOW_USE_ROCM + DPGetDeviceCount(gpu_num); // check current device environment + if (gpu_num > 0) { + options.config.set_allow_soft_placement(true); + options.config.mutable_gpu_options()->set_per_process_gpu_memory_fraction( + 0.9); + options.config.mutable_gpu_options()->set_allow_growth(true); + DPErrcheck(DPSetDevice(gpu_rank % gpu_num)); + std::string str = "/gpu:"; + str += std::to_string(gpu_rank % gpu_num); + graph::SetDefaultDevice(str, graph_def); + } +#endif // GOOGLE_CUDA || TENSORFLOW_USE_ROCM + check_status(NewSession(options, &session)); + check_status(session->Create(*graph_def)); + try { + model_version = get_scalar("model_attr/model_version"); + } catch (deepmd::tf_exception& e) { + // no model version defined in old models + model_version = "0.0"; + } + if (!model_compatable(model_version)) { + throw deepmd::deepmd_exception( + "incompatable model: version " + model_version + + " in graph, but version " + global_model_version + + " supported " + "See https://deepmd.rtfd.io/compatability/ for details."); + } + dtype = session_get_dtype(session, "descrpt_attr/rcut"); + if (dtype == tensorflow::DT_DOUBLE) { + rcut = get_scalar("descrpt_attr/rcut"); + } else { + rcut = get_scalar("descrpt_attr/rcut"); + } + cell_size = rcut; + ntypes = get_scalar("descrpt_attr/ntypes"); + try { + ntypes_spin = get_scalar("spin_attr/ntypes_spin"); + } catch (const deepmd::deepmd_exception&) { + ntypes_spin = 0; + } + dfparam = get_scalar("fitting_attr/dfparam"); + daparam = get_scalar("fitting_attr/daparam"); + if (dfparam < 0) { + dfparam = 0; + } + if (daparam < 0) { + daparam = 0; + } + if (daparam > 0) { + try { + aparam_nall = get_scalar("fitting_attr/aparam_nall"); + } catch (const deepmd::deepmd_exception&) { + aparam_nall = false; + } + } else { + aparam_nall = false; + } + model_type = get_scalar("model_attr/model_type"); + inited = true; + + init_nbor = false; +} + +template +VT DeepPotTF::get_scalar(const std::string& name) const { + return session_get_scalar(session, name); +} + +template +void DeepPotTF::validate_fparam_aparam( + const int& nframes, + const int& nloc, + const std::vector& fparam, + const std::vector& aparam) const { + if (fparam.size() != dfparam && + fparam.size() != static_cast(nframes) * dfparam) { + throw deepmd::deepmd_exception( + "the dim of frame parameter provided is not consistent with what the " + "model uses"); + } + + if (aparam.size() != static_cast(daparam) * nloc && + aparam.size() != static_cast(nframes) * daparam * nloc) { + throw deepmd::deepmd_exception( + "the dim of atom parameter provided is not consistent with what the " + "model uses"); + } +} + +template void DeepPotTF::validate_fparam_aparam( + const int& nframes, + const int& nloc, + const std::vector& fparam, + const std::vector& aparam) const; + +template void DeepPotTF::validate_fparam_aparam( + const int& nframes, + const int& nloc, + const std::vector& fparam, + const std::vector& aparam) const; + +template +void DeepPotTF::tile_fparam_aparam(std::vector& out_param, + const int& nframes, + const int& dparam, + const std::vector& param) const { + if (param.size() == dparam) { + out_param.resize(static_cast(nframes) * dparam); + for (int ii = 0; ii < nframes; ++ii) { + std::copy(param.begin(), param.end(), + out_param.begin() + static_cast(ii) * dparam); + } + } else if (param.size() == static_cast(nframes) * dparam) { + out_param = param; + } +} + +template void DeepPotTF::tile_fparam_aparam( + std::vector& out_param, + const int& nframes, + const int& dparam, + const std::vector& param) const; + +template void DeepPotTF::tile_fparam_aparam( + std::vector& out_param, + const int& nframes, + const int& dparam, + const std::vector& param) const; + +// ENERGYVTYPE: std::vector or ENERGYTYPE + +template +void DeepPotTF::compute(ENERGYVTYPE& dener, + std::vector& dforce_, + std::vector& dvirial, + const std::vector& dcoord_, + const std::vector& datype_, + const std::vector& dbox, + const std::vector& fparam_, + const std::vector& aparam_) { + int nall = datype_.size(); + // if nall==0, unclear nframes, but 1 is ok + int nframes = nall > 0 ? (dcoord_.size() / nall / 3) : 1; + int nloc = nall; + atommap = deepmd::AtomMap(datype_.begin(), datype_.begin() + nloc); + assert(nloc == atommap.get_type().size()); + std::vector fparam; + std::vector aparam; + validate_fparam_aparam(nframes, (aparam_nall ? nall : nloc), fparam_, + aparam_); + tile_fparam_aparam(fparam, nframes, dfparam, fparam_); + tile_fparam_aparam(aparam, nframes, (aparam_nall ? nall : nloc) * daparam, + aparam_); + + std::vector> input_tensors; + + if (dtype == tensorflow::DT_DOUBLE) { + int ret = session_input_tensors(input_tensors, dcoord_, ntypes, + datype_, dbox, cell_size, fparam, + aparam, atommap, "", aparam_nall); + assert(ret == nloc); + run_model(dener, dforce_, dvirial, session, input_tensors, atommap, + nframes); + } else { + int ret = session_input_tensors(input_tensors, dcoord_, ntypes, + datype_, dbox, cell_size, fparam, + aparam, atommap, "", aparam_nall); + assert(ret == nloc); + run_model(dener, dforce_, dvirial, session, input_tensors, atommap, + nframes); + } +} + +template void DeepPotTF::compute( + ENERGYTYPE& dener, + std::vector& dforce_, + std::vector& dvirial, + const std::vector& dcoord_, + const std::vector& datype_, + const std::vector& dbox, + const std::vector& fparam, + const std::vector& aparam); + +template void DeepPotTF::compute( + ENERGYTYPE& dener, + std::vector& dforce_, + std::vector& dvirial, + const std::vector& dcoord_, + const std::vector& datype_, + const std::vector& dbox, + const std::vector& fparam, + const std::vector& aparam); + +template void DeepPotTF::compute>( + std::vector& dener, + std::vector& dforce_, + std::vector& dvirial, + const std::vector& dcoord_, + const std::vector& datype_, + const std::vector& dbox, + const std::vector& fparam, + const std::vector& aparam); + +template void DeepPotTF::compute>( + std::vector& dener, + std::vector& dforce_, + std::vector& dvirial, + const std::vector& dcoord_, + const std::vector& datype_, + const std::vector& dbox, + const std::vector& fparam, + const std::vector& aparam); + +template +void DeepPotTF::compute(ENERGYVTYPE& dener, + std::vector& dforce_, + std::vector& dvirial, + const std::vector& dcoord_, + const std::vector& datype_, + const std::vector& dbox, + const int nghost, + const InputNlist& lmp_list, + const int& ago, + const std::vector& fparam_, + const std::vector& aparam__) { + int nall = datype_.size(); + // if nall==0, unclear nframes, but 1 is ok + int nframes = nall > 0 ? (dcoord_.size() / nall / 3) : 1; + std::vector fparam; + std::vector aparam_; + validate_fparam_aparam(nframes, (aparam_nall ? nall : (nall - nghost)), + fparam_, aparam__); + tile_fparam_aparam(fparam, nframes, dfparam, fparam_); + tile_fparam_aparam(aparam_, nframes, + (aparam_nall ? nall : (nall - nghost)) * daparam, + aparam__); + + // select real atoms + std::vector dcoord, dforce, aparam; + std::vector datype, fwd_map, bkw_map; + int nghost_real, nall_real, nloc_real; + select_real_atoms_coord(dcoord, datype, aparam, nghost_real, fwd_map, bkw_map, + nall_real, nloc_real, dcoord_, datype_, aparam_, + nghost, ntypes, nframes, daparam, nall, aparam_nall); + + // internal nlist + if (ago == 0) { + nlist_data.copy_from_nlist(lmp_list); + nlist_data.shuffle_exclude_empty(fwd_map); + } + compute_inner(dener, dforce, dvirial, dcoord, datype, dbox, nghost_real, ago, + fparam, aparam); + // bkw map + dforce_.resize(static_cast(nframes) * fwd_map.size() * 3); + select_map(dforce_, dforce, bkw_map, 3, nframes, fwd_map.size(), + bkw_map.size()); +} + +template void DeepPotTF::compute( + ENERGYTYPE& dener, + std::vector& dforce_, + std::vector& dvirial, + const std::vector& dcoord_, + const std::vector& datype_, + const std::vector& dbox, + const int nghost, + const InputNlist& lmp_list, + const int& ago, + const std::vector& fparam, + const std::vector& aparam_); + +template void DeepPotTF::compute( + ENERGYTYPE& dener, + std::vector& dforce_, + std::vector& dvirial, + const std::vector& dcoord_, + const std::vector& datype_, + const std::vector& dbox, + const int nghost, + const InputNlist& lmp_list, + const int& ago, + const std::vector& fparam, + const std::vector& aparam_); + +template void DeepPotTF::compute>( + std::vector& dener, + std::vector& dforce_, + std::vector& dvirial, + const std::vector& dcoord_, + const std::vector& datype_, + const std::vector& dbox, + const int nghost, + const InputNlist& lmp_list, + const int& ago, + const std::vector& fparam, + const std::vector& aparam_); + +template void DeepPotTF::compute>( + std::vector& dener, + std::vector& dforce_, + std::vector& dvirial, + const std::vector& dcoord_, + const std::vector& datype_, + const std::vector& dbox, + const int nghost, + const InputNlist& lmp_list, + const int& ago, + const std::vector& fparam, + const std::vector& aparam_); + +template +void DeepPotTF::compute_inner(ENERGYVTYPE& dener, + std::vector& dforce_, + std::vector& dvirial, + const std::vector& dcoord_, + const std::vector& datype_, + const std::vector& dbox, + const int nghost, + const int& ago, + const std::vector& fparam, + const std::vector& aparam) { + int nall = datype_.size(); + // if nall==0, unclear nframes, but 1 is ok + int nframes = nall > 0 ? (dcoord_.size() / nall / 3) : 1; + int nloc = nall - nghost; + + std::vector> input_tensors; + + // agp == 0 means that the LAMMPS nbor list has been updated + if (ago == 0) { + atommap = deepmd::AtomMap(datype_.begin(), datype_.begin() + nloc); + assert(nloc == atommap.get_type().size()); + nlist_data.shuffle(atommap); + nlist_data.make_inlist(nlist); + } + if (dtype == tensorflow::DT_DOUBLE) { + int ret = session_input_tensors( + input_tensors, dcoord_, ntypes, datype_, dbox, nlist, fparam, aparam, + atommap, nghost, ago, "", aparam_nall); + assert(nloc == ret); + run_model(dener, dforce_, dvirial, session, input_tensors, atommap, + nframes, nghost); + } else { + int ret = session_input_tensors( + input_tensors, dcoord_, ntypes, datype_, dbox, nlist, fparam, aparam, + atommap, nghost, ago, "", aparam_nall); + assert(nloc == ret); + run_model(dener, dforce_, dvirial, session, input_tensors, atommap, + nframes, nghost); + } +} + +template void DeepPotTF::compute_inner( + ENERGYTYPE& dener, + std::vector& dforce_, + std::vector& dvirial, + const std::vector& dcoord_, + const std::vector& datype_, + const std::vector& dbox, + const int nghost, + const int& ago, + const std::vector& fparam, + const std::vector& aparam); + +template void DeepPotTF::compute_inner( + ENERGYTYPE& dener, + std::vector& dforce_, + std::vector& dvirial, + const std::vector& dcoord_, + const std::vector& datype_, + const std::vector& dbox, + const int nghost, + const int& ago, + const std::vector& fparam, + const std::vector& aparam); + +template void DeepPotTF::compute_inner>( + std::vector& dener, + std::vector& dforce_, + std::vector& dvirial, + const std::vector& dcoord_, + const std::vector& datype_, + const std::vector& dbox, + const int nghost, + const int& ago, + const std::vector& fparam, + const std::vector& aparam); + +template void DeepPotTF::compute_inner>( + std::vector& dener, + std::vector& dforce_, + std::vector& dvirial, + const std::vector& dcoord_, + const std::vector& datype_, + const std::vector& dbox, + const int nghost, + const int& ago, + const std::vector& fparam, + const std::vector& aparam); + +template +void DeepPotTF::compute(ENERGYVTYPE& dener, + std::vector& dforce_, + std::vector& dvirial, + std::vector& datom_energy_, + std::vector& datom_virial_, + const std::vector& dcoord_, + const std::vector& datype_, + const std::vector& dbox, + const std::vector& fparam_, + const std::vector& aparam_) { + // if datype.size is 0, not clear nframes; but 1 is just ok + int nframes = datype_.size() > 0 ? (dcoord_.size() / 3 / datype_.size()) : 1; + atommap = deepmd::AtomMap(datype_.begin(), datype_.end()); + int nloc = datype_.size(); + std::vector fparam; + std::vector aparam; + validate_fparam_aparam(nframes, nloc, fparam_, aparam_); + tile_fparam_aparam(fparam, nframes, dfparam, fparam_); + tile_fparam_aparam(aparam, nframes, nloc * daparam, aparam_); + + std::vector> input_tensors; + + if (dtype == tensorflow::DT_DOUBLE) { + int ret = session_input_tensors(input_tensors, dcoord_, ntypes, + datype_, dbox, cell_size, fparam, + aparam, atommap, "", aparam_nall); + run_model(dener, dforce_, dvirial, datom_energy_, datom_virial_, + session, input_tensors, atommap, nframes); + } else { + int ret = session_input_tensors(input_tensors, dcoord_, ntypes, + datype_, dbox, cell_size, fparam, + aparam, atommap, "", aparam_nall); + run_model(dener, dforce_, dvirial, datom_energy_, datom_virial_, + session, input_tensors, atommap, nframes); + } +} + +template void DeepPotTF::compute( + ENERGYTYPE& dener, + std::vector& dforce_, + std::vector& dvirial, + std::vector& datom_energy_, + std::vector& datom_virial_, + const std::vector& dcoord_, + const std::vector& datype_, + const std::vector& dbox, + const std::vector& fparam, + const std::vector& aparam); + +template void DeepPotTF::compute( + ENERGYTYPE& dener, + std::vector& dforce_, + std::vector& dvirial, + std::vector& datom_energy_, + std::vector& datom_virial_, + const std::vector& dcoord_, + const std::vector& datype_, + const std::vector& dbox, + const std::vector& fparam, + const std::vector& aparam); + +template void DeepPotTF::compute>( + std::vector& dener, + std::vector& dforce_, + std::vector& dvirial, + std::vector& datom_energy_, + std::vector& datom_virial_, + const std::vector& dcoord_, + const std::vector& datype_, + const std::vector& dbox, + const std::vector& fparam, + const std::vector& aparam); + +template void DeepPotTF::compute>( + std::vector& dener, + std::vector& dforce_, + std::vector& dvirial, + std::vector& datom_energy_, + std::vector& datom_virial_, + const std::vector& dcoord_, + const std::vector& datype_, + const std::vector& dbox, + const std::vector& fparam, + const std::vector& aparam); + +template +void DeepPotTF::compute(ENERGYVTYPE& dener, + std::vector& dforce_, + std::vector& dvirial, + std::vector& datom_energy_, + std::vector& datom_virial_, + const std::vector& dcoord_, + const std::vector& datype_, + const std::vector& dbox, + const int nghost, + const InputNlist& lmp_list, + const int& ago, + const std::vector& fparam_, + const std::vector& aparam__) { + int nall = datype_.size(); + // if nall==0, unclear nframes, but 1 is ok + int nframes = nall > 0 ? (dcoord_.size() / nall / 3) : 1; + int nloc = nall - nghost; + std::vector fparam; + std::vector aparam_; + validate_fparam_aparam(nframes, (aparam_nall ? nall : nloc), fparam_, + aparam__); + tile_fparam_aparam(fparam, nframes, dfparam, fparam_); + tile_fparam_aparam(aparam_, nframes, (aparam_nall ? nall : nloc) * daparam, + aparam__); + std::vector> input_tensors; + // select real atoms + std::vector dcoord, dforce, aparam, datom_energy, datom_virial; + std::vector datype, fwd_map, bkw_map; + int nghost_real, nall_real, nloc_real; + select_real_atoms_coord(dcoord, datype, aparam, nghost_real, fwd_map, bkw_map, + nall_real, nloc_real, dcoord_, datype_, aparam_, + nghost, ntypes, nframes, daparam, nall, aparam_nall); + + if (ago == 0) { + atommap = deepmd::AtomMap(datype.begin(), datype.begin() + nloc_real); + assert(nloc_real == atommap.get_type().size()); + + nlist_data.copy_from_nlist(lmp_list); + nlist_data.shuffle_exclude_empty(fwd_map); + nlist_data.shuffle(atommap); + nlist_data.make_inlist(nlist); + } + + if (dtype == tensorflow::DT_DOUBLE) { + int ret = session_input_tensors( + input_tensors, dcoord, ntypes, datype, dbox, nlist, fparam, aparam, + atommap, nghost_real, ago, "", aparam_nall); + assert(nloc_real == ret); + run_model(dener, dforce, dvirial, datom_energy, datom_virial, + session, input_tensors, atommap, nframes, nghost_real); + } else { + int ret = session_input_tensors( + input_tensors, dcoord, ntypes, datype, dbox, nlist, fparam, aparam, + atommap, nghost_real, ago, "", aparam_nall); + assert(nloc_real == ret); + run_model(dener, dforce, dvirial, datom_energy, datom_virial, + session, input_tensors, atommap, nframes, nghost_real); + } + + // bkw map + dforce_.resize(static_cast(nframes) * fwd_map.size() * 3); + datom_energy_.resize(static_cast(nframes) * fwd_map.size()); + datom_virial_.resize(static_cast(nframes) * fwd_map.size() * 9); + select_map(dforce_, dforce, bkw_map, 3, nframes, fwd_map.size(), + nall_real); + select_map(datom_energy_, datom_energy, bkw_map, 1, nframes, + fwd_map.size(), nall_real); + select_map(datom_virial_, datom_virial, bkw_map, 9, nframes, + fwd_map.size(), nall_real); +} + +template void DeepPotTF::compute( + ENERGYTYPE& dener, + std::vector& dforce_, + std::vector& dvirial, + std::vector& datom_energy_, + std::vector& datom_virial_, + const std::vector& dcoord_, + const std::vector& datype_, + const std::vector& dbox, + const int nghost, + const InputNlist& lmp_list, + const int& ago, + const std::vector& fparam, + const std::vector& aparam_); + +template void DeepPotTF::compute( + ENERGYTYPE& dener, + std::vector& dforce_, + std::vector& dvirial, + std::vector& datom_energy_, + std::vector& datom_virial_, + const std::vector& dcoord_, + const std::vector& datype_, + const std::vector& dbox, + const int nghost, + const InputNlist& lmp_list, + const int& ago, + const std::vector& fparam, + const std::vector& aparam_); + +template void DeepPotTF::compute>( + std::vector& dener, + std::vector& dforce_, + std::vector& dvirial, + std::vector& datom_energy_, + std::vector& datom_virial_, + const std::vector& dcoord_, + const std::vector& datype_, + const std::vector& dbox, + const int nghost, + const InputNlist& lmp_list, + const int& ago, + const std::vector& fparam, + const std::vector& aparam_); + +template void DeepPotTF::compute>( + std::vector& dener, + std::vector& dforce_, + std::vector& dvirial, + std::vector& datom_energy_, + std::vector& datom_virial_, + const std::vector& dcoord_, + const std::vector& datype_, + const std::vector& dbox, + const int nghost, + const InputNlist& lmp_list, + const int& ago, + const std::vector& fparam, + const std::vector& aparam_); + +// mixed type +template +void DeepPotTF::compute_mixed_type(ENERGYVTYPE& dener, + std::vector& dforce_, + std::vector& dvirial, + const int& nframes, + const std::vector& dcoord_, + const std::vector& datype_, + const std::vector& dbox, + const std::vector& fparam_, + const std::vector& aparam_) { + int nloc = datype_.size() / nframes; + // here atommap only used to get nloc + atommap = deepmd::AtomMap(datype_.begin(), datype_.begin() + nloc); + std::vector fparam; + std::vector aparam; + validate_fparam_aparam(nframes, nloc, fparam_, aparam_); + tile_fparam_aparam(fparam, nframes, dfparam, fparam_); + tile_fparam_aparam(aparam, nframes, nloc * daparam, aparam_); + + std::vector> input_tensors; + + if (dtype == tensorflow::DT_DOUBLE) { + int ret = session_input_tensors_mixed_type( + input_tensors, nframes, dcoord_, ntypes, datype_, dbox, cell_size, + fparam, aparam, atommap, "", aparam_nall); + assert(ret == nloc); + run_model(dener, dforce_, dvirial, session, input_tensors, atommap, + nframes); + } else { + int ret = session_input_tensors_mixed_type( + input_tensors, nframes, dcoord_, ntypes, datype_, dbox, cell_size, + fparam, aparam, atommap, "", aparam_nall); + assert(ret == nloc); + run_model(dener, dforce_, dvirial, session, input_tensors, atommap, + nframes); + } +} + +template void DeepPotTF::compute_mixed_type( + ENERGYTYPE& dener, + std::vector& dforce_, + std::vector& dvirial, + const int& nframes, + const std::vector& dcoord_, + const std::vector& datype_, + const std::vector& dbox, + const std::vector& fparam, + const std::vector& aparam); + +template void DeepPotTF::compute_mixed_type( + ENERGYTYPE& dener, + std::vector& dforce_, + std::vector& dvirial, + const int& nframes, + const std::vector& dcoord_, + const std::vector& datype_, + const std::vector& dbox, + const std::vector& fparam, + const std::vector& aparam); + +template void DeepPotTF::compute_mixed_type>( + std::vector& dener, + std::vector& dforce_, + std::vector& dvirial, + const int& nframes, + const std::vector& dcoord_, + const std::vector& datype_, + const std::vector& dbox, + const std::vector& fparam, + const std::vector& aparam); + +template void DeepPotTF::compute_mixed_type>( + std::vector& dener, + std::vector& dforce_, + std::vector& dvirial, + const int& nframes, + const std::vector& dcoord_, + const std::vector& datype_, + const std::vector& dbox, + const std::vector& fparam, + const std::vector& aparam); + +template +void DeepPotTF::compute_mixed_type(ENERGYVTYPE& dener, + std::vector& dforce_, + std::vector& dvirial, + std::vector& datom_energy_, + std::vector& datom_virial_, + const int& nframes, + const std::vector& dcoord_, + const std::vector& datype_, + const std::vector& dbox, + const std::vector& fparam_, + const std::vector& aparam_) { + int nloc = datype_.size() / nframes; + // here atommap only used to get nloc + atommap = deepmd::AtomMap(datype_.begin(), datype_.begin() + nloc); + std::vector fparam; + std::vector aparam; + validate_fparam_aparam(nframes, nloc, fparam_, aparam_); + tile_fparam_aparam(fparam, nframes, dfparam, fparam_); + tile_fparam_aparam(aparam, nframes, nloc * daparam, aparam_); + + std::vector> input_tensors; + + if (dtype == tensorflow::DT_DOUBLE) { + int nloc = session_input_tensors_mixed_type( + input_tensors, nframes, dcoord_, ntypes, datype_, dbox, cell_size, + fparam, aparam, atommap, "", aparam_nall); + run_model(dener, dforce_, dvirial, datom_energy_, datom_virial_, + session, input_tensors, atommap, nframes); + } else { + int nloc = session_input_tensors_mixed_type( + input_tensors, nframes, dcoord_, ntypes, datype_, dbox, cell_size, + fparam, aparam, atommap, "", aparam_nall); + run_model(dener, dforce_, dvirial, datom_energy_, datom_virial_, + session, input_tensors, atommap, nframes); + } +} + +template void DeepPotTF::compute_mixed_type( + ENERGYTYPE& dener, + std::vector& dforce_, + std::vector& dvirial, + std::vector& datom_energy_, + std::vector& datom_virial_, + const int& nframes, + const std::vector& dcoord_, + const std::vector& datype_, + const std::vector& dbox, + const std::vector& fparam, + const std::vector& aparam); + +template void DeepPotTF::compute_mixed_type( + ENERGYTYPE& dener, + std::vector& dforce_, + std::vector& dvirial, + std::vector& datom_energy_, + std::vector& datom_virial_, + const int& nframes, + const std::vector& dcoord_, + const std::vector& datype_, + const std::vector& dbox, + const std::vector& fparam, + const std::vector& aparam); + +template void DeepPotTF::compute_mixed_type>( + std::vector& dener, + std::vector& dforce_, + std::vector& dvirial, + std::vector& datom_energy_, + std::vector& datom_virial_, + const int& nframes, + const std::vector& dcoord_, + const std::vector& datype_, + const std::vector& dbox, + const std::vector& fparam, + const std::vector& aparam); + +template void DeepPotTF::compute_mixed_type>( + std::vector& dener, + std::vector& dforce_, + std::vector& dvirial, + std::vector& datom_energy_, + std::vector& datom_virial_, + const int& nframes, + const std::vector& dcoord_, + const std::vector& datype_, + const std::vector& dbox, + const std::vector& fparam, + const std::vector& aparam); + +void DeepPotTF::get_type_map(std::string& type_map) { + type_map = get_scalar("model_attr/tmap"); +} + +// forward to template method +void DeepPotTF::computew(double& ener, + std::vector& force, + std::vector& virial, + const std::vector& coord, + const std::vector& atype, + const std::vector& box, + const std::vector& fparam, + const std::vector& aparam) { + compute(ener, force, virial, coord, atype, box, fparam, aparam); +} + +void DeepPotTF::computew(double& ener, + std::vector& force, + std::vector& virial, + const std::vector& coord, + const std::vector& atype, + const std::vector& box, + const std::vector& fparam, + const std::vector& aparam) { + compute(ener, force, virial, coord, atype, box, fparam, aparam); +} +void DeepPotTF::computew(std::vector& ener, + std::vector& force, + std::vector& virial, + const std::vector& coord, + const std::vector& atype, + const std::vector& box, + const std::vector& fparam, + const std::vector& aparam) { + compute(ener, force, virial, coord, atype, box, fparam, aparam); +} +void DeepPotTF::computew(std::vector& ener, + std::vector& force, + std::vector& virial, + const std::vector& coord, + const std::vector& atype, + const std::vector& box, + const std::vector& fparam, + const std::vector& aparam) { + compute(ener, force, virial, coord, atype, box, fparam, aparam); +} +void DeepPotTF::computew(double& ener, + std::vector& force, + std::vector& virial, + const std::vector& coord, + const std::vector& atype, + const std::vector& box, + const int nghost, + const InputNlist& inlist, + const int& ago, + const std::vector& fparam, + const std::vector& aparam) { + compute(ener, force, virial, coord, atype, box, nghost, inlist, ago, fparam, + aparam); +} + +void DeepPotTF::computew(double& ener, + std::vector& force, + std::vector& virial, + const std::vector& coord, + const std::vector& atype, + const std::vector& box, + const int nghost, + const InputNlist& inlist, + const int& ago, + const std::vector& fparam, + const std::vector& aparam) { + compute(ener, force, virial, coord, atype, box, nghost, inlist, ago, fparam, + aparam); +} +void DeepPotTF::computew(std::vector& ener, + std::vector& force, + std::vector& virial, + const std::vector& coord, + const std::vector& atype, + const std::vector& box, + const int nghost, + const InputNlist& inlist, + const int& ago, + const std::vector& fparam, + const std::vector& aparam) { + compute(ener, force, virial, coord, atype, box, nghost, inlist, ago, fparam, + aparam); +} +void DeepPotTF::computew(std::vector& ener, + std::vector& force, + std::vector& virial, + const std::vector& coord, + const std::vector& atype, + const std::vector& box, + const int nghost, + const InputNlist& inlist, + const int& ago, + const std::vector& fparam, + const std::vector& aparam) { + compute(ener, force, virial, coord, atype, box, nghost, inlist, ago, fparam, + aparam); +} +void DeepPotTF::computew(double& ener, + std::vector& force, + std::vector& virial, + std::vector& atom_energy, + std::vector& atom_virial, + const std::vector& coord, + const std::vector& atype, + const std::vector& box, + const std::vector& fparam, + const std::vector& aparam) { + compute(ener, force, virial, atom_energy, atom_virial, coord, atype, box, + fparam, aparam); +} + +void DeepPotTF::computew(double& ener, + std::vector& force, + std::vector& virial, + std::vector& atom_energy, + std::vector& atom_virial, + const std::vector& coord, + const std::vector& atype, + const std::vector& box, + const std::vector& fparam, + const std::vector& aparam) { + compute(ener, force, virial, atom_energy, atom_virial, coord, atype, box, + fparam, aparam); +} +void DeepPotTF::computew(std::vector& ener, + std::vector& force, + std::vector& virial, + std::vector& atom_energy, + std::vector& atom_virial, + const std::vector& coord, + const std::vector& atype, + const std::vector& box, + const std::vector& fparam, + const std::vector& aparam) { + compute(ener, force, virial, atom_energy, atom_virial, coord, atype, box, + fparam, aparam); +} +void DeepPotTF::computew(std::vector& ener, + std::vector& force, + std::vector& virial, + std::vector& atom_energy, + std::vector& atom_virial, + const std::vector& coord, + const std::vector& atype, + const std::vector& box, + const std::vector& fparam, + const std::vector& aparam) { + compute(ener, force, virial, atom_energy, atom_virial, coord, atype, box, + fparam, aparam); +} +void DeepPotTF::computew(double& ener, + std::vector& force, + std::vector& virial, + std::vector& atom_energy, + std::vector& atom_virial, + const std::vector& coord, + const std::vector& atype, + const std::vector& box, + const int nghost, + const InputNlist& inlist, + const int& ago, + const std::vector& fparam, + const std::vector& aparam) { + compute(ener, force, virial, atom_energy, atom_virial, coord, atype, box, + nghost, inlist, ago, fparam, aparam); +} + +void DeepPotTF::computew(double& ener, + std::vector& force, + std::vector& virial, + std::vector& atom_energy, + std::vector& atom_virial, + const std::vector& coord, + const std::vector& atype, + const std::vector& box, + const int nghost, + const InputNlist& inlist, + const int& ago, + const std::vector& fparam, + const std::vector& aparam) { + compute(ener, force, virial, atom_energy, atom_virial, coord, atype, box, + nghost, inlist, ago, fparam, aparam); +} +void DeepPotTF::computew(std::vector& ener, + std::vector& force, + std::vector& virial, + std::vector& atom_energy, + std::vector& atom_virial, + const std::vector& coord, + const std::vector& atype, + const std::vector& box, + const int nghost, + const InputNlist& inlist, + const int& ago, + const std::vector& fparam, + const std::vector& aparam) { + compute(ener, force, virial, atom_energy, atom_virial, coord, atype, box, + nghost, inlist, ago, fparam, aparam); +} +void DeepPotTF::computew(std::vector& ener, + std::vector& force, + std::vector& virial, + std::vector& atom_energy, + std::vector& atom_virial, + const std::vector& coord, + const std::vector& atype, + const std::vector& box, + const int nghost, + const InputNlist& inlist, + const int& ago, + const std::vector& fparam, + const std::vector& aparam) { + compute(ener, force, virial, atom_energy, atom_virial, coord, atype, box, + nghost, inlist, ago, fparam, aparam); +} +void DeepPotTF::computew_mixed_type(double& ener, + std::vector& force, + std::vector& virial, + const int& nframes, + const std::vector& coord, + const std::vector& atype, + const std::vector& box, + const std::vector& fparam, + const std::vector& aparam) { + compute_mixed_type(ener, force, virial, nframes, coord, atype, box, fparam, + aparam); +} +void DeepPotTF::computew_mixed_type(double& ener, + std::vector& force, + std::vector& virial, + const int& nframes, + const std::vector& coord, + const std::vector& atype, + const std::vector& box, + const std::vector& fparam, + const std::vector& aparam) { + compute_mixed_type(ener, force, virial, nframes, coord, atype, box, fparam, + aparam); +} +void DeepPotTF::computew_mixed_type(std::vector& ener, + std::vector& force, + std::vector& virial, + const int& nframes, + const std::vector& coord, + const std::vector& atype, + const std::vector& box, + const std::vector& fparam, + const std::vector& aparam) { + compute_mixed_type(ener, force, virial, nframes, coord, atype, box, fparam, + aparam); +} +void DeepPotTF::computew_mixed_type(std::vector& ener, + std::vector& force, + std::vector& virial, + const int& nframes, + const std::vector& coord, + const std::vector& atype, + const std::vector& box, + const std::vector& fparam, + const std::vector& aparam) { + compute_mixed_type(ener, force, virial, nframes, coord, atype, box, fparam, + aparam); +} +void DeepPotTF::computew_mixed_type(double& ener, + std::vector& force, + std::vector& virial, + std::vector& atom_energy, + std::vector& atom_virial, + const int& nframes, + const std::vector& coord, + const std::vector& atype, + const std::vector& box, + const std::vector& fparam, + const std::vector& aparam) { + compute_mixed_type(ener, force, virial, atom_energy, atom_virial, nframes, + coord, atype, box, fparam, aparam); +} +void DeepPotTF::computew_mixed_type(double& ener, + std::vector& force, + std::vector& virial, + std::vector& atom_energy, + std::vector& atom_virial, + const int& nframes, + const std::vector& coord, + const std::vector& atype, + const std::vector& box, + const std::vector& fparam, + const std::vector& aparam) { + compute_mixed_type(ener, force, virial, atom_energy, atom_virial, nframes, + coord, atype, box, fparam, aparam); +} +void DeepPotTF::computew_mixed_type(std::vector& ener, + std::vector& force, + std::vector& virial, + std::vector& atom_energy, + std::vector& atom_virial, + const int& nframes, + const std::vector& coord, + const std::vector& atype, + const std::vector& box, + const std::vector& fparam, + const std::vector& aparam) { + compute_mixed_type(ener, force, virial, atom_energy, atom_virial, nframes, + coord, atype, box, fparam, aparam); +} +void DeepPotTF::computew_mixed_type(std::vector& ener, + std::vector& force, + std::vector& virial, + std::vector& atom_energy, + std::vector& atom_virial, + const int& nframes, + const std::vector& coord, + const std::vector& atype, + const std::vector& box, + const std::vector& fparam, + const std::vector& aparam) { + compute_mixed_type(ener, force, virial, atom_energy, atom_virial, nframes, + coord, atype, box, fparam, aparam); +}