Merge branch 'devel' into pt-load-jit

deepmd/dpmodel/model/dp_atomic_model.py

@@ -62,6 +62,10 @@ def get_sel(self) -> List[int]:
         """Get the neighbor selection."""
         return self.descriptor.get_sel()
 
+    def get_type_map(self) -> Optional[List[str]]:
+        """Get the type map."""
+        return self.type_map
+
     def distinguish_types(self) -> bool:
         """Returns if model requires a neighbor list that distinguish different
         atomic types or not.

deepmd/dpmodel/model/linear_atomic_model.py

@@ -62,6 +62,10 @@ def get_rcut(self) -> float:
         """Get the cut-off radius."""
         return max(self.get_model_rcuts())
 
+    def get_type_map(self) -> Optional[List[str]]:
+        """Get the type map."""
+        raise NotImplementedError("TODO: get_type_map should be implemented")
+
     def get_model_rcuts(self) -> List[float]:
         """Get the cut-off radius for each individual models."""
         return [model.get_rcut() for model in self.models]

deepmd/dpmodel/model/make_base_atomic_model.py

@@ -44,6 +44,10 @@ def get_rcut(self) -> float:
             """Get the cut-off radius."""
             pass
 
+        @abstractmethod
+        def get_type_map(self) -> Optional[List[str]]:
+            """Get the type map."""
+
         @abstractmethod
         def get_sel(self) -> List[int]:
             """Returns the number of selected atoms for each type."""

deepmd/dpmodel/model/pairtab_atomic_model.py

@@ -82,6 +82,9 @@ def fitting_output_def(self) -> FittingOutputDef:
     def get_rcut(self) -> float:
         return self.rcut
 
+    def get_type_map(self) -> Optional[List[str]]:
+        raise NotImplementedError("TODO: get_type_map should be implemented")
+
     def get_sel(self) -> List[int]:
         return [self.sel]
 

deepmd/main.py

@@ -275,7 +275,7 @@ def main_parser() -> argparse.ArgumentParser:
         "--checkpoint",
         type=str,
         default=".",
-        help="Path to checkpoint. TensorFlow backend: a folder; PyTorch backend: either a folder containing checkpoint, or a pt file",
+        help="Path to checkpoint, either a folder containing checkpoint or the checkpoint prefix",
     )
     parser_frz.add_argument(
         "-o",

deepmd/pt/model/model/dp_atomic_model.py

@@ -67,10 +67,16 @@ def fitting_output_def(self) -> FittingOutputDef:
             else self.coord_denoise_net.output_def()
         )
 
+    @torch.jit.export
     def get_rcut(self) -> float:
         """Get the cut-off radius."""
         return self.rcut
 
+    @torch.jit.export
+    def get_type_map(self) -> List[str]:
+        """Get the type map."""
+        return self.type_map
+
     def get_sel(self) -> List[int]:
         """Get the neighbor selection."""
         return self.sel

deepmd/pt/model/model/linear_atomic_model.py

@@ -61,10 +61,16 @@ def distinguish_types(self) -> bool:
         """If distinguish different types by sorting."""
         return False
 
+    @torch.jit.export
     def get_rcut(self) -> float:
         """Get the cut-off radius."""
         return max(self.get_model_rcuts())
 
+    @torch.jit.export
+    def get_type_map(self) -> List[str]:
+        """Get the type map."""
+        raise NotImplementedError("TODO: implement this method")
+
     def get_model_rcuts(self) -> List[float]:
         """Get the cut-off radius for each individual models."""
         return [model.get_rcut() for model in self.models]

deepmd/pt/model/model/pairtab_atomic_model.py

@@ -95,9 +95,14 @@ def fitting_output_def(self) -> FittingOutputDef:
             ]
         )
 
+    @torch.jit.export
     def get_rcut(self) -> float:
         return self.rcut
 
+    @torch.jit.export
+    def get_type_map(self) -> Optional[List[str]]:
+        raise NotImplementedError("TODO: implement this method")
+
     def get_sel(self) -> List[int]:
         return [self.sel]
 

deepmd/tf/entrypoints/freeze.py

@@ -12,6 +12,9 @@
 from os.path import (
     abspath,
 )
+from pathlib import (
+    Path,
+)
 from typing import (
     List,
     Optional,
@@ -479,7 +482,7 @@ def freeze(
     Parameters
     ----------
     checkpoint_folder : str
-        location of the folder with model
+        location of either the folder with checkpoint or the checkpoint prefix
     output : str
         output file name
     node_names : Optional[str], optional
@@ -492,8 +495,11 @@ def freeze(
         other arguments
     """
     # We retrieve our checkpoint fullpath
-    checkpoint = tf.train.get_checkpoint_state(checkpoint_folder)
-    input_checkpoint = checkpoint.model_checkpoint_path
+    if Path(checkpoint_folder).is_dir():
+        checkpoint = tf.train.get_checkpoint_state(checkpoint_folder)
+        input_checkpoint = checkpoint.model_checkpoint_path
+    else:
+        input_checkpoint = checkpoint_folder
 
     # expand the output file to full path
     output_graph = abspath(output)

doc/install/install-from-c-library.md

@@ -14,7 +14,7 @@ tar xzf libdeepmd_c.tar.gz
 
 The library is built in Linux (GLIBC 2.17) with CUDA 12.2 (`libdeepmd_c.tar.gz`) or 11.8 (`libdeepmd_c_cu11.tar.gz`). It's noted that this package does not contain CUDA Toolkit and cuDNN, so one needs to download them from the NVIDIA website.
 
-## Use Pre-compiled C Library to build the LAMMPS plugin and GROMACS patch
+## Use Pre-compiled C Library to build the LAMMPS plugin, i-PI driver, and GROMACS patch
 
 When one [installs DeePMD-kit's C++ interface](./install-from-source.md#install-deepmd-kits-c-interface), one can use the CMake argument `DEEPMD_C_ROOT` to the path `libdeepmd_c`.
 
@@ -27,4 +27,5 @@ make -j8
 make install
 ```
 
-Then one can follow the manual [Install LAMMPS](./install-lammps.md) and/or [Install GROMACS](./install-gromacs.md).
+Then the i-PI driver `dp_ipi` will be built and installed.
+One can also follow the manual [Install LAMMPS](./install-lammps.md) and/or [Install GROMACS](./install-gromacs.md).

source/CMakeLists.txt

@@ -297,10 +297,7 @@ if(BUILD_CPP_IF)
   endif()
   if(CMAKE_CXX_COMPILER_VERSION VERSION_GREATER 4.8)
     # add_subdirectory (md/)
-    if(ENABLE_IPI
-       OR NOT BUILD_PY_IF
-       AND NOT DEEPMD_C_ROOT)
-      # ipi has a dependency on libdeepmd
+    if(ENABLE_IPI OR NOT BUILD_PY_IF)
       add_subdirectory(ipi/)
     endif()
     if(NOT BUILD_PY_IF)

source/ipi/CMakeLists.txt

@@ -14,7 +14,7 @@ add_executable(${ipiname} ${DRIVER_SOURCE_FILES})
 # link: libdeepmd_cc
 if(DP_USING_C_API)
   # SimulationRegion.h
-  target_link_libraries(${ipiname} PRIVATE ${LIB_DEEPMD_C} ${LIB_DEEPMD})
+  target_link_libraries(${ipiname} PRIVATE ${LIB_DEEPMD_C})
   target_precompile_headers(${ipiname} PRIVATE [["deepmd.hpp"]])
   remove_definitions(-D_GLIBCXX_USE_CXX11_ABI=${OP_CXX_ABI})
 else()

source/ipi/driver.cc

@@ -12,7 +12,6 @@ namespace deepmd_compat = deepmd;
 #include "deepmd.hpp"
 namespace deepmd_compat = deepmd::hpp;
 #endif
-#include "SimulationRegion.h"
 #include "XyzFileManager.h"
 #include "json.hpp"
 #include "sockets.h"
@@ -49,25 +48,6 @@ char *trimwhitespace(char *str) {
   return str;
 }
 
-void normalize_coord(std::vector<double> &coord,
-                     const SimulationRegion<double> &region) {
-  int natoms = coord.size() / 3;
-
-  for (int ii = 0; ii < natoms; ++ii) {
-    double inter[3];
-    region.phys2Inter(inter, &coord[3 * ii]);
-    for (int dd = 0; dd < 3; ++dd) {
-      inter[dd] -= int(floor(inter[dd]));
-      if (inter[dd] < 0) {
-        inter[dd] += 1.;
-      } else if (inter[dd] >= 1) {
-        inter[dd] -= 1.;
-      }
-    }
-    region.inter2Phys(&coord[3 * ii], inter);
-  }
-}
-
 int main(int argc, char *argv[]) {
   if (argc == 1) {
     std::cerr << "usage " << std::endl;
@@ -122,7 +102,6 @@ int main(int argc, char *argv[]) {
   std::vector<double> dcoord_tmp;
   std::vector<int> dtype = cvt.get_type();
   std::vector<double> dbox(9, 0);
-  SimulationRegion<double> region;
   double *msg_buff = NULL;
   double ener;
   double virial[9];
@@ -179,7 +158,6 @@ int main(int argc, char *argv[]) {
       for (int dd = 0; dd < 9; ++dd) {
         dbox[dd] = cell_h[(dd % 3) * 3 + (dd / 3)] * cvt_len;
       }
-      region.reinitBox(&dbox[0]);
 
       // get number of atoms
       readbuffer_(&socket, (char *)(&cbuf), sizeof(int32_t));
@@ -203,7 +181,6 @@ int main(int argc, char *argv[]) {
         dcoord_tmp[ii] = msg_buff[ii] * cvt_len;
       }
       cvt.forward(dcoord, dcoord_tmp, 3);
-      normalize_coord(dcoord, region);
 
       // nnp over writes ener, force and virial
       nnp_inter.compute(dener, dforce_tmp, dvirial, dcoord, dtype, dbox);

source/ipi/tests/test_driver.py

@@ -53,7 +53,7 @@ def write_input(self, atoms, **kwargs):
         atoms.write(self.xyz_file, format="xyz")
 
 
-class TestDeepPotALargeBoxNoPBC(unittest.TestCase):
+class TestDPIPI(unittest.TestCase):
     # copy from test_deeppot_a.py
     @classmethod
     def setUpClass(cls):
@@ -193,8 +193,8 @@ def test_ase_unix(self):
                 cell=self.box.reshape((3, 3)),
                 calculator=calc,
             )
-        ee = water.get_potential_energy()
-        ff = water.get_forces()
+            ee = water.get_potential_energy()
+            ff = water.get_forces()
         nframes = 1
         np.testing.assert_almost_equal(
             ff.ravel(), self.expected_f.ravel(), default_places
@@ -213,8 +213,38 @@ def test_ase_nounix(self):
                 cell=self.box.reshape((3, 3)),
                 calculator=calc,
             )
-        ee = water.get_potential_energy()
-        ff = water.get_forces()
+            ee = water.get_potential_energy()
+            ff = water.get_forces()
+        nframes = 1
+        np.testing.assert_almost_equal(
+            ff.ravel(), self.expected_f.ravel(), default_places
+        )
+        expected_se = np.sum(self.expected_e.reshape([nframes, -1]), axis=1)
+        np.testing.assert_almost_equal(ee.ravel(), expected_se.ravel(), default_places)
+
+    def test_normalize_coords(self):
+        # coordinate nomarlization should happen inside the interface
+        cell = self.box.reshape((3, 3))
+        coord = self.coords.reshape((-1, 3))
+        # random unwrap coords
+        coord[0] += np.array([3, 0, 0]) @ cell
+        coord[1] += np.array([0, -3, 0]) @ cell
+        coord[2] += np.array([0, 0, 3]) @ cell
+        coord[3] += np.array([-3, 0, 0]) @ cell
+        coord[4] += np.array([0, 3, 0]) @ cell
+        coord[5] += np.array([0, 0, -3]) @ cell
+        with SocketIOCalculator(
+            DPiPICalculator(self.model_file, use_unix=False),
+            log=sys.stdout,
+        ) as calc:
+            water = Atoms(
+                "OHHOHH",
+                positions=coord,
+                cell=cell,
+                calculator=calc,
+            )
+            ee = water.get_potential_energy()
+            ff = water.get_forces()
         nframes = 1
         np.testing.assert_almost_equal(
             ff.ravel(), self.expected_f.ravel(), default_places

source/tests/pt/model/test_dp_atomic_model.py

@@ -107,4 +107,6 @@ def test_jit(self):
         ).to(env.DEVICE)
         type_map = ["foo", "bar"]
         md0 = DPAtomicModel(ds, ft, type_map=type_map).to(env.DEVICE)
-        torch.jit.script(md0)
+        md0 = torch.jit.script(md0)
+        self.assertEqual(md0.get_rcut(), self.rcut)
+        self.assertEqual(md0.get_type_map(), type_map)

source/tests/pt/model/test_dp_model.py

@@ -284,7 +284,9 @@ def test_jit(self):
         ).to(env.DEVICE)
         type_map = ["foo", "bar"]
         md0 = DPModel(ds, ft, type_map=type_map).to(env.DEVICE)
-        torch.jit.script(md0)
+        md0 = torch.jit.script(md0)
+        md0.get_rcut()
+        md0.get_type_map()
 
 
 class TestDPModelFormatNlist(unittest.TestCase):
@@ -521,4 +523,6 @@ def test_jit(self):
         ).to(env.DEVICE)
         type_map = ["foo", "bar"]
         md0 = EnergyModel(ds, ft, type_map=type_map).to(env.DEVICE)
-        torch.jit.script(md0)
+        md0 = torch.jit.script(md0)
+        self.assertEqual(md0.get_rcut(), self.rcut)
+        self.assertEqual(md0.get_type_map(), type_map)

source/tests/pt/model/test_linear_atomic_model.py

@@ -171,8 +171,14 @@ def test_self_consistency(self):
         )
 
     def test_jit(self):
-        torch.jit.script(self.md1)
-        torch.jit.script(self.md3)
+        md1 = torch.jit.script(self.md1)
+        self.assertEqual(md1.get_rcut(), self.rcut)
+        with self.assertRaises(torch.jit.Error):
+            self.assertEqual(md1.get_type_map(), ["foo", "bar"])
+        md3 = torch.jit.script(self.md3)
+        self.assertEqual(md3.get_rcut(), self.rcut)
+        with self.assertRaises(torch.jit.Error):
+            self.assertEqual(md3.get_type_map(), ["foo", "bar"])
 
 
 if __name__ == "__main__":

source/tests/pt/model/test_pairtab_atomic_model.py

@@ -82,6 +82,9 @@ def test_with_mask(self):
 
     def test_jit(self):
         model = torch.jit.script(self.model)
+        self.assertEqual(model.get_rcut(), 0.02)
+        with self.assertRaises(torch.jit.Error):
+            self.assertEqual(model.get_type_map(), None)
 
     def test_deserialize(self):
         model1 = PairTabModel.deserialize(self.model.serialize())
@@ -101,6 +104,9 @@ def test_deserialize(self):
         )
 
         model1 = torch.jit.script(model1)
+        self.assertEqual(model1.get_rcut(), 0.02)
+        with self.assertRaises(torch.jit.Error):
+            self.assertEqual(model1.get_type_map(), None)
 
     def test_cross_deserialize(self):
         model_dict = self.model.serialize()  # pytorch model to dict