Merge pull request #18 from deepmodeling/api

api branche update

deepmd/entrypoints/freeze.py

@@ -43,6 +43,7 @@ def _make_node_names(model_type: str, modifier_type: Optional[str] = None) -> Li
         "descrpt_attr/ntypes",
         "model_attr/tmap",
         "model_attr/model_type",
+        "model_attr/model_version",
     ]
 
     if model_type == "ener":
@@ -59,6 +60,7 @@ def _make_node_names(model_type: str, modifier_type: Optional[str] = None) -> Li
         nodes += [
             "o_wfc",
             "model_attr/sel_type",
+            "model_attr/output_dim",
         ]
     elif model_type == "dipole":
         nodes += [

deepmd/infer/deep_dipole.py

@@ -1,6 +1,6 @@
-from typing import TYPE_CHECKING
+from deepmd.infer.deep_tensor import DeepTensor
 
-from deepmd.infer.deep_eval import DeepTensor
+from typing import TYPE_CHECKING
 
 if TYPE_CHECKING:
     from pathlib import Path
@@ -33,7 +33,6 @@ def __init__(
         # instance namespace
         self.tensors = dict(
             {
-                "t_sel_type": "model_attr/sel_type:0",
                 # output tensor
                 "t_tensor": "o_dipole:0",
             },
@@ -43,7 +42,6 @@ def __init__(
         DeepTensor.__init__(
             self,
             model_file,
-            3,
             load_prefix=load_prefix,
             default_tf_graph=default_tf_graph,
         )

deepmd/infer/deep_eval.py

@@ -4,6 +4,7 @@
 import numpy as np
 from deepmd.common import make_default_mesh
 from deepmd.env import default_tf_session_config, tf
+from deepmd.run_options import MODEL_VERSION
 
 if TYPE_CHECKING:
     from pathlib import Path
@@ -13,6 +14,7 @@ class DeepEval:
     """Common methods for DeepPot, DeepWFC, DeepPolar, ..."""
 
     _model_type: Optional[str] = None
+    _model_version: Optional[str] = None
     load_prefix: str  # set by subclass
 
     def __init__(
@@ -26,6 +28,13 @@ def __init__(
         )
         self.load_prefix = load_prefix
 
+        # graph_compatable should be called after graph and prefix are set
+        if not self._graph_compatable():
+            raise RuntimeError(
+                f"model in graph (version {self.model_version}) is incompatible"
+                f"with the model (version {MODEL_VERSION}) supported by the current code."
+            )
+
     @property
     def model_type(self) -> str:
         """Get type of model.
@@ -37,9 +46,44 @@ def model_type(self) -> str:
             sess = tf.Session(graph=self.graph, config=default_tf_session_config)
             [mt] = sess.run([t_mt], feed_dict={})
             self._model_type = mt.decode("utf-8")
-
         return self._model_type
 
+    @property
+    def model_version(self) -> str:
+        """Get type of model.
+
+        :type:str
+        """
+        if not self._model_version:
+            try:
+                t_mt = self._get_tensor("model_attr/model_version:0")
+                sess = tf.Session(graph=self.graph, config=default_tf_session_config)
+                [mt] = sess.run([t_mt], feed_dict={})
+                self._model_version = mt.decode("utf-8")
+            except KeyError:
+                # For deepmd-kit version 0.x - 1.x, set model version to 0.0
+                self._model_version = "0.0"
+        return self._model_version    
+
+    def _graph_compatable(
+        self
+    ) -> bool :
+        """ Check the model compatability
+        
+        Return
+        bool
+            If the model stored in the graph file is compatable with the current code
+        """
+        model_version_major = int(self.model_version.split('.')[0])
+        model_version_minor = int(self.model_version.split('.')[1])
+        MODEL_VERSION_MAJOR = int(MODEL_VERSION.split('.')[0])
+        MODEL_VERSION_MINOR = int(MODEL_VERSION.split('.')[1])
+        if (model_version_major != MODEL_VERSION_MAJOR) or \
+           (model_version_minor >  MODEL_VERSION_MINOR) :
+            return False
+        else:
+            return True
+
     def _get_tensor(
         self, tensor_name: str, attr_name: Optional[str] = None
     ) -> tf.Tensor:
@@ -70,8 +114,6 @@ def _get_tensor(
     def _load_graph(
         frozen_graph_filename: "Path", prefix: str = "load", default_tf_graph: bool = False
     ):
-
-
         # We load the protobuf file from the disk and parse it to retrieve the
         # unserialized graph_def
         with tf.gfile.GFile(str(frozen_graph_filename), "rb") as f:
@@ -101,168 +143,6 @@ def _load_graph(
 
             return graph
 
-class DeepTensor(DeepEval):
-    """Evaluates a tensor model.
-
-    Constructor
-
-    Parameters
-    ----------
-    model_file: str
-        The name of the frozen model file.
-    variable_dof: int
-        The DOF of the variable to evaluate.
-    load_prefix: str
-        The prefix in the load computational graph
-    default_tf_graph : bool
-        If uses the default tf graph, otherwise build a new tf graph for evaluation
-    """
-
-    tensors = {
-        "t_ntypes": "descrpt_attr/ntypes:0",
-        "t_rcut": "descrpt_attr/rcut:0",
-        "t_tmap": "model_attr/tmap:0",
-        # inputs
-        "t_coord": "t_coord:0",
-        "t_type": "t_type:0",
-        "t_natoms": "t_natoms:0",
-        "t_box": "t_box:0",
-        "t_mesh": "t_mesh:0",
-    }
-
-    def __init__(
-        self,
-        model_file: "Path",
-        variable_dof: Optional[int],
-        load_prefix: str = 'load',
-        default_tf_graph: bool = False
-    ) -> None:
-        DeepEval.__init__(
-            self,
-            model_file,
-            load_prefix=load_prefix,
-            default_tf_graph=default_tf_graph
-        )
-        self.variable_dof = variable_dof
-
-        # now load tensors to object attributes
-        for attr_name, tensor_name in self.tensors.items():
-            self._get_tensor(tensor_name, attr_name)
-
-        # start a tf session associated to the graph
-        self.sess = tf.Session(graph=self.graph, config=default_tf_session_config)
-        self._run_default_sess()
-        self.tmap = self.tmap.decode('UTF-8').split()
-
-    def _run_default_sess(self):
-        [self.ntypes, self.rcut, self.tmap, self.tselt] = self.sess.run(
-            [self.t_ntypes, self.t_rcut, self.t_tmap, self.t_sel_type]
-        )
-
-    def get_ntypes(self) -> int:
-        """Get the number of atom types of this model."""
-        return self.ntypes
-
-    def get_rcut(self) -> float:
-        """Get the cut-off radius of this model."""
-        return self.rcut
-
-    def get_type_map(self) -> List[int]:
-        """Get the type map (element name of the atom types) of this model."""
-        return self.tmap
-
-    def get_sel_type(self) -> List[int]:
-        """Get the selected atom types of this model."""        
-        return self.tselt
-
-    def get_dim_fparam(self) -> int:
-        """Get the number (dimension) of frame parameters of this DP."""
-        return self.dfparam
-
-    def get_dim_aparam(self) -> int:
-        """Get the number (dimension) of atomic parameters of this DP."""
-        return self.daparam
-
-    def eval(
-        self,
-        coords: np.array,
-        cells: np.array,
-        atom_types: List[int],
-        atomic: bool = True,
-        fparam: Optional[np.array] = None,
-        aparam: Optional[np.array] = None,
-        efield: Optional[np.array] = None
-    ) -> np.array:
-        """Evaluate the model.
-
-        Parameters
-        ----------
-        coords
-            The coordinates of atoms. 
-            The array should be of size nframes x natoms x 3
-        cells
-            The cell of the region. 
-            If None then non-PBC is assumed, otherwise using PBC. 
-            The array should be of size nframes x 9
-        atom_types
-            The atom types
-            The list should contain natoms ints
-        atomic
-            Calculate the atomic energy and virial
-        fparam
-            Not used in this model
-        aparam
-            Not used in this model
-        efield
-            Not used in this model
-
-        Returns
-        -------
-        tensor
-                The returned tensor
-                If atomic == False then of size nframes x variable_dof
-                else of size nframes x natoms x variable_dof
-        """
-        # standarize the shape of inputs
-        coords = np.array(coords)
-        cells = np.array(cells)
-        atom_types = np.array(atom_types, dtype = int)
-
-        # reshape the inputs 
-        cells = np.reshape(cells, [-1, 9])
-        nframes = cells.shape[0]
-        coords = np.reshape(coords, [nframes, -1])
-        natoms = coords.shape[1] // 3
-
-        # sort inputs
-        coords, atom_types, imap, sel_at, sel_imap = self.sort_input(coords, atom_types, sel_atoms = self.get_sel_type())
-
-        # make natoms_vec and default_mesh
-        natoms_vec = self.make_natoms_vec(atom_types)
-        assert(natoms_vec[0] == natoms)
-
-        # evaluate
-        tensor = []
-        feed_dict_test = {}
-        feed_dict_test[self.t_natoms] = natoms_vec
-        feed_dict_test[self.t_type  ] = np.tile(atom_types, [nframes,1]).reshape([-1])
-        t_out = [self.t_tensor]
-        feed_dict_test[self.t_coord] = np.reshape(coords, [-1])
-        feed_dict_test[self.t_box  ] = np.reshape(cells , [-1])
-        feed_dict_test[self.t_mesh ] = make_default_mesh(cells)
-        v_out = self.sess.run (t_out, feed_dict = feed_dict_test)
-        tensor = v_out[0]
-
-        # reverse map of the outputs
-        if atomic:
-            tensor = np.array(tensor)
-            tensor = self.reverse_map(np.reshape(tensor, [nframes,-1,self.variable_dof]), sel_imap)
-            tensor = np.reshape(tensor, [nframes, len(sel_at), self.variable_dof])
-        else:
-            tensor = np.reshape(tensor, [nframes, self.variable_dof])
-
-        return tensor
-
     @staticmethod
     def sort_input(
         coord : np.array, atom_type : np.array, sel_atoms : List[int] = None
@@ -339,6 +219,7 @@ def reverse_map(vec : np.ndarray, imap : List[int]) -> np.ndarray:
             ret[:,ii,:] = vec[:,idx,:]
         return ret
 
+
     def make_natoms_vec(self, atom_types : np.ndarray) -> np.ndarray :
         """Make the natom vector used by deepmd-kit.
 
@@ -363,4 +244,3 @@ def make_natoms_vec(self, atom_types : np.ndarray) -> np.ndarray :
         for ii in range (self.ntypes) :
             natoms_vec[ii+2] = np.count_nonzero(atom_types == ii)
         return natoms_vec
-

deepmd/infer/deep_polar.py

@@ -1,7 +1,7 @@
-from typing import TYPE_CHECKING, List, Optional
-
+from deepmd.infer.deep_tensor import DeepTensor
 import numpy as np
-from deepmd.infer.deep_eval import DeepTensor
+
+from typing import TYPE_CHECKING, List, Optional
 
 if TYPE_CHECKING:
     from pathlib import Path
@@ -34,7 +34,6 @@ def __init__(
         # instance namespace
         self.tensors = dict(
             {
-                "t_sel_type": "model_attr/sel_type:0",
                 # output tensor
                 "t_tensor": "o_polar:0",
             },
@@ -44,7 +43,6 @@ def __init__(
         DeepTensor.__init__(
             self,
             model_file,
-            9,
             load_prefix=load_prefix,
             default_tf_graph=default_tf_graph,
         )