feat(jax): passing mapping from LAMMPS to DPA-2

.pre-commit-config.yaml

@@ -15,7 +15,8 @@ repos:
         exclude: |
           (?x)^(
               source/tests/infer/dipolecharge_e.pbtxt|
-              source/tests/infer/deeppolar_new.pbtxt
+              source/tests/infer/deeppolar_new.pbtxt|
+              source/tests/infer/deeppot_dpa.savedmodel/saved_model.pb
           )$
       - id: check-merge-conflict
       - id: check-symlinks

doc/backend.md

@@ -31,7 +31,9 @@ While `.pth` and `.pt` are the same in the PyTorch package, they have different
 [JAX](https://jax.readthedocs.io/) 0.4.33 (which requires Python 3.10 or above) or above is required.
 Both `.xlo` and `.jax` are customized format extensions defined in DeePMD-kit, since JAX has no convention for file extensions.
 `.savedmodel` is the TensorFlow [SavedModel format](https://www.tensorflow.org/guide/saved_model) generated by [JAX2TF](https://www.tensorflow.org/guide/jax2tf), which needs the installation of TensorFlow.
-Currently, this backend is developed actively, and has no support for training and the C++ interface.
+Only the `.savedmodel` format supports C++ inference, which needs the TensorFlow C++ interface.
+The model is device-specific, so that the model generated on the GPU device cannot be run on the CPUs.
+Currently, this backend is developed actively, and has no support for training.
 
 ### DP {{ dpmodel_icon }}
 

doc/install/install-from-source.md

@@ -297,7 +297,9 @@ If one does not need to use DeePMD-kit with LAMMPS or i-PI, then the python inte
 
 ::::{tab-set}
 
-:::{tab-item} TensorFlow {{ tensorflow_icon }}
+:::{tab-item} TensorFlow {{ tensorflow_icon }} / JAX {{ jax_icon }}
+
+The C++ interfaces of both TensorFlow and JAX backends are based on the TensorFlow C++ library.
 
 Since TensorFlow 2.12, TensorFlow C++ library (`libtensorflow_cc`) is packaged inside the Python library. Thus, you can skip building TensorFlow C++ library manually. If that does not work for you, you can still build it manually.
 
@@ -338,7 +340,7 @@ We recommend using [conda packages](https://docs.deepmodeling.org/faq/conda.html
 
 ::::{tab-set}
 
-:::{tab-item} TensorFlow {{ tensorflow_icon }}
+:::{tab-item} TensorFlow {{ tensorflow_icon }} / JAX {{ jax_icon }}
 
 I assume you have activated the TensorFlow Python environment and want to install DeePMD-kit into path `$deepmd_root`, then execute CMake
 
@@ -375,7 +377,7 @@ One may add the following CMake variables to `cmake` using the [`-D <var>=<value
 
 **Type**: `BOOL` (`ON`/`OFF`), Default: `OFF`
 
-{{ tensorflow_icon }} Whether building the TensorFlow backend.
+{{ tensorflow_icon }} {{ jax_icon }} Whether building the TensorFlow backend and the JAX backend.
 
 :::
 
@@ -391,7 +393,7 @@ One may add the following CMake variables to `cmake` using the [`-D <var>=<value
 
 **Type**: `PATH`
 
-{{ tensorflow_icon }} The Path to TensorFlow's C++ interface.
+{{ tensorflow_icon }} {{ jax_icon }} The Path to TensorFlow's C++ interface.
 
 :::
 

doc/model/dpa2.md

@@ -18,6 +18,14 @@ If one runs LAMMPS with MPI, the customized OP library for the C++ interface sho
 If one runs LAMMPS with MPI and CUDA devices, it is recommended to compile the customized OP library for the C++ interface with a [CUDA-Aware MPI](https://developer.nvidia.com/mpi-solutions-gpus) library and CUDA,
 otherwise the communication between GPU cards falls back to the slower CPU implementation.
 
+## Limiations of the JAX backend with LAMMPS {{ jax_icon }}
+
+When using the JAX backend, 2 or more MPI ranks are not supported. One must set `map` to `yes` using the [`atom_modify`](https://docs.lammps.org/atom_modify.html) command.
+
+```lammps
+atom_modify map yes
+```
+
 ## Data format
 
 DPA-2 supports both the [standard data format](../data/system.md) and the [mixed type data format](../data/system.md#mixed-type).

source/api_c/include/c_api.h

@@ -12,7 +12,7 @@ extern "C" {
 /** C API version. Bumped whenever the API is changed.
  * @since API version 22
  */
-#define DP_C_API_VERSION 23
+#define DP_C_API_VERSION 24
 
 /**
  * @brief Neighbor list.
@@ -31,7 +31,7 @@ extern DP_Nlist* DP_NewNlist(int inum_,
                              int* ilist_,
                              int* numneigh_,
                              int** firstneigh_);
-/*
+/**
  * @brief Create a new neighbor list with communication capabilities.
  * @details This function extends DP_NewNlist by adding support for parallel
  * communication, allowing the neighbor list to be used in distributed
@@ -68,7 +68,7 @@ extern DP_Nlist* DP_NewNlist_comm(int inum_,
                                   int* recvproc,
                                   void* world);
 
-/*
+/**
  * @brief Set mask for a neighbor list.
  *
  * @param nl Neighbor list.
@@ -78,6 +78,16 @@ extern DP_Nlist* DP_NewNlist_comm(int inum_,
  **/
 extern void DP_NlistSetMask(DP_Nlist* nl, int mask);
 
+/**
+ * @brief Set mapping for a neighbor list.
+ *
+ * @param nl Neighbor list.
+ * @param mapping mapping from all atoms to real atoms, in size nall.
+ * @since API version 24
+ *
+ **/
+extern void DP_NlistSetMapping(DP_Nlist* nl, int* mapping);
+
 /**
  * @brief Delete a neighbor list.
  *

source/api_c/include/deepmd.hpp

@@ -615,6 +615,11 @@ struct InputNlist {
    * @brief Set mask for this neighbor list.
    */
   void set_mask(int mask) { DP_NlistSetMask(nl, mask); };
+  /**
+   * @brief Set mapping for this neighbor list.
+   * @param mapping mapping from all atoms to real atoms, in size nall.
+   */
+  void set_mapping(int *mapping) { DP_NlistSetMapping(nl, mapping); };
 };
 
 /**

source/api_c/src/c_api.cc

@@ -43,6 +43,9 @@ DP_Nlist* DP_NewNlist_comm(int inum_,
   return new_nl;
 }
 void DP_NlistSetMask(DP_Nlist* nl, int mask) { nl->nl.set_mask(mask); }
+void DP_NlistSetMapping(DP_Nlist* nl, int* mapping) {
+  nl->nl.set_mapping(mapping);
+}
 void DP_DeleteNlist(DP_Nlist* nl) { delete nl; }
 
 DP_DeepPot::DP_DeepPot() {}

source/api_cc/include/DeepPotJAX.h

@@ -0,0 +1,249 @@
+// SPDX-License-Identifier: LGPL-3.0-or-later
+#pragma once
+
+#include <tensorflow/c/c_api.h>
+#include <tensorflow/c/eager/c_api.h>
+
+#include "DeepPot.h"
+#include "common.h"
+#include "neighbor_list.h"
+
+namespace deepmd {
+/**
+ * @brief TensorFlow implementation for Deep Potential.
+ **/
+class DeepPotJAX : public DeepPotBase {
+ public:
+  /**
+   * @brief DP constructor without initialization.
+   **/
+  DeepPotJAX();
+  virtual ~DeepPotJAX();
+  /**
+   * @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.
+   **/
+  DeepPotJAX(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 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 0;
+  };
+  /**
+   * @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 false;
+  };
+
+  // forward to template class
+  void computew(std::vector<double>& ener,
+                std::vector<double>& force,
+                std::vector<double>& virial,
+                std::vector<double>& atom_energy,
+                std::vector<double>& atom_virial,
+                const std::vector<double>& coord,
+                const std::vector<int>& atype,
+                const std::vector<double>& box,
+                const std::vector<double>& fparam,
+                const std::vector<double>& aparam,
+                const bool atomic);
+  void computew(std::vector<double>& ener,
+                std::vector<float>& force,
+                std::vector<float>& virial,
+                std::vector<float>& atom_energy,
+                std::vector<float>& atom_virial,
+                const std::vector<float>& coord,
+                const std::vector<int>& atype,
+                const std::vector<float>& box,
+                const std::vector<float>& fparam,
+                const std::vector<float>& aparam,
+                const bool atomic);
+  void computew(std::vector<double>& ener,
+                std::vector<double>& force,
+                std::vector<double>& virial,
+                std::vector<double>& atom_energy,
+                std::vector<double>& atom_virial,
+                const std::vector<double>& coord,
+                const std::vector<int>& atype,
+                const std::vector<double>& box,
+                const int nghost,
+                const InputNlist& inlist,
+                const int& ago,
+                const std::vector<double>& fparam,
+                const std::vector<double>& aparam,
+                const bool atomic);
+  void computew(std::vector<double>& ener,
+                std::vector<float>& force,
+                std::vector<float>& virial,
+                std::vector<float>& atom_energy,
+                std::vector<float>& atom_virial,
+                const std::vector<float>& coord,
+                const std::vector<int>& atype,
+                const std::vector<float>& box,
+                const int nghost,
+                const InputNlist& inlist,
+                const int& ago,
+                const std::vector<float>& fparam,
+                const std::vector<float>& aparam,
+                const bool atomic);
+  void computew_mixed_type(std::vector<double>& ener,
+                           std::vector<double>& force,
+                           std::vector<double>& virial,
+                           std::vector<double>& atom_energy,
+                           std::vector<double>& atom_virial,
+                           const int& nframes,
+                           const std::vector<double>& coord,
+                           const std::vector<int>& atype,
+                           const std::vector<double>& box,
+                           const std::vector<double>& fparam,
+                           const std::vector<double>& aparam,
+                           const bool atomic);
+  void computew_mixed_type(std::vector<double>& ener,
+                           std::vector<float>& force,
+                           std::vector<float>& virial,
+                           std::vector<float>& atom_energy,
+                           std::vector<float>& atom_virial,
+                           const int& nframes,
+                           const std::vector<float>& coord,
+                           const std::vector<int>& atype,
+                           const std::vector<float>& box,
+                           const std::vector<float>& fparam,
+                           const std::vector<float>& aparam,
+                           const bool atomic);
+
+ private:
+  bool inited;
+  // device
+  std::string device;
+  // the cutoff radius
+  double rcut;
+  // the number of types
+  int ntypes;
+  // the dimension of the frame parameter
+  int dfparam;
+  // the dimension of the atomic parameter
+  int daparam;
+  // type map
+  std::string type_map;
+  // sel
+  std::vector<int64_t> sel;
+  // number of neighbors
+  int nnei;
+  /**  TF C API objects.
+   * @{
+   */
+  TF_Graph* graph;
+  TF_Status* status;
+  TF_Session* session;
+  TF_SessionOptions* sessionopts;
+  TFE_ContextOptions* ctx_opts;
+  TFE_Context* ctx;
+  std::vector<TF_Function*> func_vector;
+  /**
+   * @}
+   */
+  // neighbor list data
+  NeighborListData nlist_data;
+  /**
+   * @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.
+   * @param[in] atomic Whether to compute atomic energy and virial.
+   **/
+  template <typename VALUETYPE>
+  void compute(std::vector<ENERGYTYPE>& ener,
+               std::vector<VALUETYPE>& force,
+               std::vector<VALUETYPE>& virial,
+               std::vector<VALUETYPE>& atom_energy,
+               std::vector<VALUETYPE>& atom_virial,
+               const std::vector<VALUETYPE>& coord,
+               const std::vector<int>& atype,
+               const std::vector<VALUETYPE>& box,
+               const int nghost,
+               const InputNlist& lmp_list,
+               const int& ago,
+               const std::vector<VALUETYPE>& fparam,
+               const std::vector<VALUETYPE>& aparam,
+               const bool atomic);
+};
+}  // namespace deepmd

source/api_cc/include/common.h

@@ -13,7 +13,7 @@
 namespace deepmd {
 
 typedef double ENERGYTYPE;
-enum DPBackend { TensorFlow, PyTorch, Paddle, Unknown };
+enum DPBackend { TensorFlow, PyTorch, Paddle, JAX, Unknown };
 
 struct NeighborListData {
   /// Array stores the core region atom's index