first working version

source/CMakeLists.txt

@@ -201,6 +201,7 @@ include_directories(${CMAKE_BINARY_DIR}/lib/)
 add_subdirectory (op/)
 if (BUILD_PY_IF)
   add_subdirectory (train/)
+  add_subdirectory (api/)
   add_subdirectory (scripts/)
   add_subdirectory (tests/)
 endif (BUILD_PY_IF)

source/api/CMakeLists.txt

@@ -0,0 +1,8 @@
+# train
+
+file(GLOB LIB_PY  descrpt_*.py network.py)
+
+install(
+  FILES		${LIB_PY} 
+  DESTINATION	deepmd
+)

source/train/DescrptSeA.py → source/api/descrpt_se_a.py

@@ -1,48 +1,81 @@
 import numpy as np
+from typing import Tuple, List
+
 from deepmd.env import tf
-from deepmd.common import ClassArg, get_activation_func, get_precision
+from deepmd.common import ClassArg, get_activation_func, get_precision, activation_fn_dict, precision_dict, docstring_parameter
+from deepmd.argcheck import list_to_doc
 from deepmd.RunOptions import global_tf_float_precision
 from deepmd.RunOptions import global_np_float_precision
 from deepmd.env import op_module
 from deepmd.env import default_tf_session_config
-from deepmd.Network import embedding_net
+from deepmd.network import embedding_net
 
 
 class DescrptSeA ():
-    def __init__ (self, jdata):
-        args = ClassArg()\
-               .add('sel',      list,   must = True) \
-               .add('rcut',     float,  default = 6.0) \
-               .add('rcut_smth',float,  default = 0.5) \
-               .add('neuron',   list,   default = [10, 20, 40]) \
-               .add('axis_neuron', int, default = 4, alias = 'n_axis_neuron') \
-               .add('resnet_dt',bool,   default = False) \
-               .add('trainable',bool,   default = True) \
-               .add('seed',     int) \
-               .add('type_one_side', bool, default = False) \
-               .add('exclude_types', list, default = []) \
-               .add('set_davg_zero', bool, default = False) \
-               .add('activation_function', str,    default = 'tanh') \
-               .add('precision', str, default = "default")
-        class_data = args.parse(jdata)
-        self.sel_a = class_data['sel']
-        self.rcut_r = class_data['rcut']
-        self.rcut_r_smth = class_data['rcut_smth']
-        self.filter_neuron = class_data['neuron']
-        self.n_axis_neuron = class_data['axis_neuron']
-        self.filter_resnet_dt = class_data['resnet_dt']
-        self.seed = class_data['seed']
-        self.trainable = class_data['trainable']
-        self.filter_activation_fn = get_activation_func(class_data['activation_function'])
-        self.filter_precision = get_precision(class_data['precision'])
-        exclude_types = class_data['exclude_types']
+    @docstring_parameter(list_to_doc(activation_fn_dict.keys()), list_to_doc(precision_dict.keys()))
+    def __init__ (self, 
+                  rcut: float,
+                  rcut_smth: float,
+                  sel: List[str],
+                  neuron: List[int],
+                  axis_neuron: int,
+                  resnet_dt: bool = False,
+                  trainable: bool = True,
+                  seed: int = 1,
+                  type_one_side: bool = True,
+                  exclude_types: List[int] = [],
+                  set_davg_zero: bool = False,
+                  activation_function: str = 'tanh',
+                  precision: str = 'default'
+    ) -> None:
+        """
+        Constructor
+
+        Parameters
+        ----------
+        rcut
+                The cut-off radius
+        rcut_smth
+                From where the environment matrix should be smoothed
+        sel : list[str]
+                sel[i] specifies the maxmum number of type i atoms in the cut-off radius
+        neuron : list[int]
+                Number of neurons in each hidden layers of the embedding net
+        axis_neuron
+                Number of the axis neuron (number of columns of the sub-matrix of the embedding matrix)
+        resnet_dt
+                Time-step `dt` in the resnet construction:
+                y = x + dt * \phi (Wx + b)
+        trainable
+                If the weights of embedding net are trainable.
+        seed
+                Random seed for initializing the network parameters.
+        type_one_side
+                Try to build N_types embedding nets. Otherwise, building N_types^2 embedding nets
+        exclude_types : list[int]
+                The Excluded types
+        activation_function
+                The activation function in the embedding net. Supported options are {0}
+        precision
+                The precision of the embedding net parameters. Supported options are {1}
+        """
+        self.sel_a = sel
+        self.rcut_r = rcut
+        self.rcut_r_smth = rcut_smth
+        self.filter_neuron = neuron
+        self.n_axis_neuron = axis_neuron
+        self.filter_resnet_dt = resnet_dt
+        self.seed = seed
+        self.trainable = trainable
+        self.filter_activation_fn = get_activation_func(activation_function)
+        self.filter_precision = get_precision(precision)
         self.exclude_types = set()
         for tt in exclude_types:
             assert(len(tt) == 2)
             self.exclude_types.add((tt[0], tt[1]))
             self.exclude_types.add((tt[1], tt[0]))
-        self.set_davg_zero = class_data['set_davg_zero']
-        self.type_one_side = class_data['type_one_side']
+        self.set_davg_zero = set_davg_zero
+        self.type_one_side = type_one_side
         if self.type_one_side and len(exclude_types) != 0:
             raise RuntimeError('"type_one_side" is not compatible with "exclude_types"')
 
@@ -89,28 +122,70 @@ def __init__ (self, jdata):
         self.sub_sess = tf.Session(graph = sub_graph, config=default_tf_session_config)
 
 
-    def get_rcut (self) :
+    def get_rcut (self) -> float:
+        """
+        Returns the cut-off radisu
+        """
         return self.rcut_r
 
-    def get_ntypes (self) :
+    def get_ntypes (self) -> int:
+        """
+        Returns the number of atom types
+        """
         return self.ntypes
 
-    def get_dim_out (self) :
+    def get_dim_out (self) -> int:
+        """
+        Returns the output dimension of this descriptor
+        """
         return self.filter_neuron[-1] * self.n_axis_neuron
 
-    def get_dim_rot_mat_1 (self) :
+    def get_dim_rot_mat_1 (self) -> int:
+        """
+        Returns the first dimension of the rotation matrix. The rotation is of shape dim_1 x 3
+        """
         return self.filter_neuron[-1]
 
-    def get_nlist (self) :
+    def get_nlist (self) -> Tuple[tf.Tensor, tf.Tensor, List[int], List[int]]:
+        """
+        Returns
+        -------
+        nlist
+                Neighbor list
+        rij
+                The relative distance between the neighbor and the center atom.
+        sel_a
+                The number of neighbors with full information
+        sel_r
+                The number of neighbors with only radial information
+        """
         return self.nlist, self.rij, self.sel_a, self.sel_r
 
     def compute_input_stats (self,
-                             data_coord, 
-                             data_box, 
-                             data_atype, 
-                             natoms_vec,
-                             mesh, 
-                             input_dict) :
+                             data_coord : list, 
+                             data_box : list, 
+                             data_atype : list, 
+                             natoms_vec : list,
+                             mesh : list, 
+                             input_dict : dict) -> None :
+        """
+        Compute the statisitcs (avg and std) of the training data. The input will be normalized by the statistics.
+        
+        Parameters
+        ----------
+        data_coord
+                The coordinates. Can be generated by deepmd.model.make_stat_input
+        data_box
+                The box. Can be generated by deepmd.model.make_stat_input
+        data_atype
+                The atom types. Can be generated by deepmd.model.make_stat_input
+        natoms_vec
+                The vector for the number of atoms of the system and different types of atoms. Can be generated by deepmd.model.make_stat_input
+        mesh
+                The mesh for neighbor searching. Can be generated by deepmd.model.make_stat_input
+        input_dict
+                Dictionary for additional input
+        """
         all_davg = []
         all_dstd = []
         if True:
@@ -150,14 +225,44 @@ def compute_input_stats (self,
 
 
     def build (self, 
-               coord_, 
-               atype_,
-               natoms,
-               box_, 
-               mesh,
-               input_dict,
-               suffix = '', 
-               reuse = None):
+               coord_ : tf.Tensor, 
+               atype_ : tf.Tensor,
+               natoms : tf.Tensor,
+               box_ : tf.Tensor, 
+               mesh : tf.Tensor,
+               input_dict : dict,
+               suffix : str = '', 
+               reuse : bool = None) -> tf.Tensor:
+        """
+        Build the computational graph for the descriptor
+
+        Parameters
+        ----------
+        coord_
+                The coordinate of atoms
+        atype_
+                The type of atoms
+        natoms
+                The number of atoms. This tensor has the length of Ntypes + 2
+                natoms[0]: number of local atoms
+                natoms[1]: total number of atoms held by this processor
+                natoms[i]: 2 <= i < Ntypes+2, number of type i atoms
+        mesh
+                For historical reasons, only the length of the Tensor matters.
+                if size of mesh == 6, pbc is assumed. 
+                if size of mesh == 0, no-pbc is assumed. 
+        input_dict
+                Dictionary for additional inputs
+        suffix
+                Name suffix to identify this descriptor
+        reuse
+                The weights in the networks should be reused when get the variable.
+
+        Returns
+        -------
+        descriptor
+                The output descriptor
+        """
         davg = self.davg
         dstd = self.dstd
         with tf.variable_scope('descrpt_attr' + suffix, reuse = reuse) :
@@ -229,11 +334,38 @@ def build (self,
         return self.dout
 
 
-    def get_rot_mat(self) :
+    def get_rot_mat(self) -> tf.Tensor:
+        """
+        Get rotational matrix
+        """
         return self.qmat
 
 
-    def prod_force_virial(self, atom_ener, natoms) :
+    def prod_force_virial(self, 
+                          atom_ener : tf.Tensor, 
+                          natoms : tf.Tensor
+    ) -> Tuple[tf.Tensor, tf.Tensor, tf.Tensor]:
+        """
+        Compute force and virial
+
+        Parameters
+        ----------
+        atom_ener
+                The atomic energy
+        natoms
+                The number of atoms. This tensor has the length of Ntypes + 2
+                natoms[0]: number of local atoms
+                natoms[1]: total number of atoms held by this processor
+                natoms[i]: 2 <= i < Ntypes+2, number of type i atoms
+        Return
+        ------
+        force
+                The force on atoms
+        virial
+                The total virial
+        atom_virial
+                The atomic virial
+        """
         [net_deriv] = tf.gradients (atom_ener, self.descrpt_reshape)
         tf.summary.histogram('net_derivative', net_deriv)
         net_deriv_reshape = tf.reshape (net_deriv, [-1, natoms[0] * self.ndescrpt])        

source/train/DescrptSeAEbd.py → source/api/descrpt_se_a_ebd.py

@@ -1,13 +1,13 @@
 import numpy as np
 from deepmd.env import tf
 from deepmd.common import ClassArg, get_activation_func, get_precision, add_data_requirement
-from deepmd.Network import one_layer
+from deepmd.network import one_layer
 from deepmd.RunOptions import global_tf_float_precision
 from deepmd.RunOptions import global_np_float_precision
 from deepmd.env import op_module
 from deepmd.env import default_tf_session_config
-from deepmd.DescrptSeA import DescrptSeA
-from deepmd.Network import embedding_net
+from deepmd.descrpt_se_a import DescrptSeA
+from deepmd.network import embedding_net
 
 class DescrptSeAEbd (DescrptSeA):
     def __init__ (self, jdata):

source/train/DescrptSeAEf.py → source/api/descrpt_se_a_ef.py

@@ -5,7 +5,7 @@
 from deepmd.RunOptions import global_np_float_precision
 from deepmd.env import op_module
 from deepmd.env import default_tf_session_config
-from deepmd.DescrptSeA import DescrptSeA
+from deepmd.descrpt_se_a import DescrptSeA
 
 class DescrptSeAEf ():
     def __init__(self, jdata):

source/train/DescrptSeAT.py → source/api/descrpt_se_a_t.py

@@ -5,7 +5,7 @@
 from deepmd.RunOptions import global_np_float_precision
 from deepmd.env import op_module
 from deepmd.env import default_tf_session_config
-from deepmd.Network import embedding_net
+from deepmd.network import embedding_net
 
 class DescrptSeAT ():
     def __init__ (self, jdata):

source/train/DescrptSeAR.py → source/api/descrpt_se_ar.py

@@ -2,8 +2,8 @@
 from deepmd.env import tf
 from deepmd.common import ClassArg
 
-from deepmd.DescrptSeA import DescrptSeA
-from deepmd.DescrptSeR import DescrptSeR
+from deepmd.descrpt_se_a import DescrptSeA
+from deepmd.descrpt_se_r import DescrptSeR
 from deepmd.env import op_module
 
 class DescrptSeAR ():

source/train/DescrptSeR.py → source/api/descrpt_se_r.py

@@ -5,7 +5,7 @@
 from deepmd.RunOptions import global_np_float_precision
 from deepmd.env import op_module
 from deepmd.env import default_tf_session_config
-from deepmd.Network import embedding_net
+from deepmd.network import embedding_net
 
 class DescrptSeR ():
     def __init__ (self, jdata):

source/train/DescrptHybrid.py

@@ -5,11 +5,11 @@
 from deepmd.RunOptions import global_tf_float_precision
 from deepmd.RunOptions import global_np_float_precision
 from deepmd.DescrptLocFrame import DescrptLocFrame
-from deepmd.DescrptSeA import DescrptSeA
-from deepmd.DescrptSeAT import DescrptSeAT
-from deepmd.DescrptSeAEbd import DescrptSeAEbd
-from deepmd.DescrptSeAEf import DescrptSeAEf
-from deepmd.DescrptSeR import DescrptSeR
+from deepmd.descrpt_se_a import DescrptSeA
+from deepmd.descrpt_se_a_t import DescrptSeAT
+from deepmd.descrpt_se_a_ebd import DescrptSeAEbd
+from deepmd.descrpt_se_a_ef import DescrptSeAEf
+from deepmd.descrpt_se_r import DescrptSeR
 
 class DescrptHybrid ():
     def __init__ (self, jdata):
@@ -25,7 +25,7 @@ def __init__ (self, jdata):
             if this_type == 'loc_frame':
                 this_descrpt = DescrptLocFrame(ii)
             elif this_type == 'se_a' :
-                this_descrpt = DescrptSeA(ii)
+                this_descrpt = DescrptSeA(**ii)
             elif this_type == 'se_at' :
                 this_descrpt = DescrptSeAT(ii)
             elif this_type == 'se_a_ebd' :

source/train/Fitting.py

@@ -3,9 +3,9 @@
 
 from deepmd.env import tf
 from deepmd.common import ClassArg, add_data_requirement, get_activation_func, get_precision
-from deepmd.Network import one_layer
+from deepmd.network import one_layer
 from deepmd.DescrptLocFrame import DescrptLocFrame
-from deepmd.DescrptSeA import DescrptSeA
+from deepmd.descrpt_se_a import DescrptSeA
 
 from deepmd.RunOptions import global_cvt_2_tf_float
 from deepmd.RunOptions import global_tf_float_precision